From d70f24a80141384ddfcff443687faec89ee3e2b0 Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Tue, 14 Apr 2026 06:09:17 -0700 Subject: [PATCH 1/2] Add penalization term to calibration loss for alpha/beta uncertainty (#419) When alpha_uncertainty / beta_uncertainty are provided in the experiment config, the calibration training loss now includes penalty terms that constrain inferred alpha and beta to stay near their guess values, weighted by the specified uncertainties. Made-with: Cursor --- ml/Neural_Net_Classes.py | 93 +++++++++++++++++++++++++++++++++++++++- ml/train_model.py | 78 +++++++++++++++++++++++++++++---- 2 files changed, 161 insertions(+), 10 deletions(-) diff --git a/ml/Neural_Net_Classes.py b/ml/Neural_Net_Classes.py index 1b2f3b6..2de393f 100644 --- a/ml/Neural_Net_Classes.py +++ b/ml/Neural_Net_Classes.py @@ -156,10 +156,62 @@ def train_model( break +def _calibration_penalty( + c_normcal, + o_normcal, + c_guess, + o_guess, + c_norm, + o_norm, + alpha_uncertainty, + beta_uncertainty, +): + """Compute penalty that keeps inferred alpha/beta near their guess values. + + The inferred calibration is obtained by composing the guess calibration, + normalization, and learned normalized calibration (see build_inferred_calibration). + From those compositions: + alpha_inferred = 1 / (c_guess * c_normcal) + beta_inferred = o_guess + c_guess*o_norm + c_guess*c_norm*o_normcal + - c_guess*c_normcal*o_norm + + The penalty is sum((alpha_I - alpha_G)^2 / alpha_U^2) + + sum((beta_I - beta_G )^2 / beta_U^2) + where alpha_G = 1/c_guess and beta_G = o_guess. + Dimensions with infinite uncertainty contribute zero penalty. + """ + c_inferred = c_guess * c_normcal + alpha_inferred = 1.0 / c_inferred + alpha_guess = 1.0 / c_guess + + beta_inferred = ( + o_guess + c_guess * o_norm + c_guess * c_norm * o_normcal - c_inferred * o_norm + ) + beta_guess = o_guess + + penalty_alpha = torch.sum( + (alpha_inferred - alpha_guess) ** 2 / alpha_uncertainty**2 + ) + penalty_beta = torch.sum((beta_inferred - beta_guess) ** 2 / beta_uncertainty**2) + return penalty_alpha + penalty_beta + + def train_calibration( model, exp_inputs, exp_targets, + c_guess_input, + o_guess_input, + c_norm_input, + o_norm_input, + alpha_uncertainty_input, + beta_uncertainty_input, + c_guess_output, + o_guess_output, + c_norm_output, + o_norm_output, + alpha_uncertainty_output, + beta_uncertainty_output, num_epochs=5000, lr=0.001, ): @@ -174,10 +226,26 @@ def train_calibration( calibrated_input = (1 / c_normcal_input) * (x - o_normcal_input) calibrated_output = c_normcal_output * model(calibrated_input) + o_normcal_output + A penalization term is added to the loss to keep the inferred alpha/beta + close to their guess values (see _calibration_penalty). Dimensions with + infinite uncertainty contribute zero penalty. + Args: model: frozen callable that maps exp_inputs -> predictions exp_inputs: experimental input tensor exp_targets: experimental target values (may contain NaN) + c_guess_input: guess calibration coefficients for inputs + o_guess_input: guess calibration offsets for inputs + c_norm_input: normalization coefficients for inputs + o_norm_input: normalization offsets for inputs + alpha_uncertainty_input: uncertainty on alpha for inputs (inf = no penalty) + beta_uncertainty_input: uncertainty on beta for inputs (inf = no penalty) + c_guess_output: guess calibration coefficients for outputs + o_guess_output: guess calibration offsets for outputs + c_norm_output: normalization coefficients for outputs + o_norm_output: normalization offsets for outputs + alpha_uncertainty_output: uncertainty on alpha for outputs (inf = no penalty) + beta_uncertainty_output: uncertainty on beta for outputs (inf = no penalty) num_epochs: number of training epochs lr: learning rate @@ -217,7 +285,30 @@ def train_calibration( base_predictions = model(calibrated_inputs) calibrated_outputs = c_normcal_output * base_predictions + o_normcal_output - loss = nan_mse_loss(exp_targets, calibrated_outputs) + loss = ( + nan_mse_loss(exp_targets, calibrated_outputs) + + _calibration_penalty( + c_normcal_input, + o_normcal_input, + c_guess_input, + o_guess_input, + c_norm_input, + o_norm_input, + alpha_uncertainty_input, + beta_uncertainty_input, + ) + + _calibration_penalty( + c_normcal_output, + o_normcal_output, + c_guess_output, + o_guess_output, + c_norm_output, + o_norm_output, + alpha_uncertainty_output, + beta_uncertainty_output, + ) + ) + loss.backward() optimizer.step() diff --git a/ml/train_model.py b/ml/train_model.py index f6a5847..c0f510b 100644 --- a/ml/train_model.py +++ b/ml/train_model.py @@ -200,31 +200,47 @@ def build_guess_calibration(config_dict, input_variables, output_variables): def _get_calibration(exp_name): if exp_name in depends_on_lookup: - # Experimental variables is part of the "simulation_calibration" section entry = depends_on_lookup[exp_name] - return entry["name"], entry["alpha_guess"], entry["beta_guess"] + return ( + entry["name"], + entry["alpha_guess"], + entry["beta_guess"], + entry.get("alpha_uncertainty", float("inf")), + entry.get("beta_uncertainty", float("inf")), + ) else: - # Experimental variable is not part of the "simulation_calibration" section - # In this case, no calibration is needed ; the simulation variable is identical - return exp_name, 1.0, 0.0 + # No calibration needed; the simulation variable is identical + return exp_name, 1.0, 0.0, float("inf"), float("inf") # Build the list of simulation variables sim_input_names = [] alpha_input_list = [] beta_input_list = [] + alpha_uncertainty_input_list = [] + beta_uncertainty_input_list = [] for key in input_variables: - sim_name, alpha, beta = _get_calibration(input_variables[key]["name"]) + sim_name, alpha, beta, alpha_u, beta_u = _get_calibration( + input_variables[key]["name"] + ) sim_input_names.append(sim_name) alpha_input_list.append(alpha) beta_input_list.append(beta) + alpha_uncertainty_input_list.append(alpha_u) + beta_uncertainty_input_list.append(beta_u) sim_output_names = [] alpha_output_list = [] beta_output_list = [] + alpha_uncertainty_output_list = [] + beta_uncertainty_output_list = [] for key in output_variables: - sim_name, alpha, beta = _get_calibration(output_variables[key]["name"]) + sim_name, alpha, beta, alpha_u, beta_u = _get_calibration( + output_variables[key]["name"] + ) sim_output_names.append(sim_name) alpha_output_list.append(alpha) beta_output_list.append(beta) + alpha_uncertainty_output_list.append(alpha_u) + beta_uncertainty_output_list.append(beta_u) # Build the AffineInputTransforms for the guess calibration alpha_inputs = torch.tensor(alpha_input_list, dtype=torch.float) @@ -240,11 +256,22 @@ def _get_calibration(exp_name): n_outputs, coefficient=1.0 / alpha_outputs, offset=beta_outputs ) + uncertainty_inputs = { + "alpha": torch.tensor(alpha_uncertainty_input_list, dtype=torch.float), + "beta": torch.tensor(beta_uncertainty_input_list, dtype=torch.float), + } + uncertainty_outputs = { + "alpha": torch.tensor(alpha_uncertainty_output_list, dtype=torch.float), + "beta": torch.tensor(beta_uncertainty_output_list, dtype=torch.float), + } + return ( input_guess_calibration, output_guess_calibration, sim_input_names, sim_output_names, + uncertainty_inputs, + uncertainty_outputs, ) @@ -307,11 +334,18 @@ def train_calibration_phase( input_names, output_names, device, + input_guess_calibration, + output_guess_calibration, + input_normalization, + output_normalization, + uncertainty_inputs, + uncertainty_outputs, ): """Phase 2: Train calibration layers on experimental data. Passes the frozen model to train_calibration(), which re-evaluates it at - each iteration. + each iteration. A penalization term constrains inferred alpha/beta toward + their guess values, weighted by the provided uncertainties. Returns an AffineInputTransform representing the learned calibration. """ @@ -336,7 +370,25 @@ def predict_fn(x): # Train calibration c_normcal_input, o_normcal_input, c_normcal_output, o_normcal_output = ( - train_calibration(predict_fn, exp_X, exp_y, num_epochs=5000, lr=0.001) + train_calibration( + predict_fn, + exp_X, + exp_y, + c_guess_input=input_guess_calibration.coefficient.to(device), + o_guess_input=input_guess_calibration.offset.to(device), + c_norm_input=input_normalization.coefficient.to(device), + o_norm_input=input_normalization.offset.to(device), + alpha_uncertainty_input=uncertainty_inputs["alpha"].to(device), + beta_uncertainty_input=uncertainty_inputs["beta"].to(device), + c_guess_output=output_guess_calibration.coefficient.to(device), + o_guess_output=output_guess_calibration.offset.to(device), + c_norm_output=output_normalization.coefficient.to(device), + o_norm_output=output_normalization.offset.to(device), + alpha_uncertainty_output=uncertainty_outputs["alpha"].to(device), + beta_uncertainty_output=uncertainty_outputs["beta"].to(device), + num_epochs=5000, + lr=0.001, + ) ) # Build calibration transforms @@ -564,6 +616,8 @@ def register_model_to_mlflow(model, model_type, experiment, config_dict): output_guess_calibration, sim_input_names, sim_output_names, + uncertainty_inputs, + uncertainty_outputs, ) = build_guess_calibration(config_dict, input_variables, output_variables) # Convert experimental data to simulation variable space @@ -649,6 +703,12 @@ def register_model_to_mlflow(model, model_type, experiment, config_dict): sim_input_names, sim_output_names, device, + input_guess_calibration=input_guess_calibration, + output_guess_calibration=output_guess_calibration, + input_normalization=input_normalization, + output_normalization=output_normalization, + uncertainty_inputs=uncertainty_inputs, + uncertainty_outputs=uncertainty_outputs, ) ) From 034a48804ac42ecb0bb120d6843bff6f51ce4ddd Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Tue, 14 Apr 2026 07:07:41 -0700 Subject: [PATCH 2/2] Rename alpha/beta lists to alpha_guess/beta_guess for clarity Made-with: Cursor --- ml/train_model.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/ml/train_model.py b/ml/train_model.py index c0f510b..529e51f 100644 --- a/ml/train_model.py +++ b/ml/train_model.py @@ -214,8 +214,8 @@ def _get_calibration(exp_name): # Build the list of simulation variables sim_input_names = [] - alpha_input_list = [] - beta_input_list = [] + alpha_guess_input_list = [] + beta_guess_input_list = [] alpha_uncertainty_input_list = [] beta_uncertainty_input_list = [] for key in input_variables: @@ -223,13 +223,13 @@ def _get_calibration(exp_name): input_variables[key]["name"] ) sim_input_names.append(sim_name) - alpha_input_list.append(alpha) - beta_input_list.append(beta) + alpha_guess_input_list.append(alpha) + beta_guess_input_list.append(beta) alpha_uncertainty_input_list.append(alpha_u) beta_uncertainty_input_list.append(beta_u) sim_output_names = [] - alpha_output_list = [] - beta_output_list = [] + alpha_guess_output_list = [] + beta_guess_output_list = [] alpha_uncertainty_output_list = [] beta_uncertainty_output_list = [] for key in output_variables: @@ -237,23 +237,23 @@ def _get_calibration(exp_name): output_variables[key]["name"] ) sim_output_names.append(sim_name) - alpha_output_list.append(alpha) - beta_output_list.append(beta) + alpha_guess_output_list.append(alpha) + beta_guess_output_list.append(beta) alpha_uncertainty_output_list.append(alpha_u) beta_uncertainty_output_list.append(beta_u) # Build the AffineInputTransforms for the guess calibration - alpha_inputs = torch.tensor(alpha_input_list, dtype=torch.float) - beta_inputs = torch.tensor(beta_input_list, dtype=torch.float) - alpha_outputs = torch.tensor(alpha_output_list, dtype=torch.float) - beta_outputs = torch.tensor(beta_output_list, dtype=torch.float) + alpha_guess_inputs = torch.tensor(alpha_guess_input_list, dtype=torch.float) + beta_guess_inputs = torch.tensor(beta_guess_input_list, dtype=torch.float) + alpha_guess_outputs = torch.tensor(alpha_guess_output_list, dtype=torch.float) + beta_guess_outputs = torch.tensor(beta_guess_output_list, dtype=torch.float) n_inputs = len(input_variables) n_outputs = len(output_variables) input_guess_calibration = AffineInputTransform( - n_inputs, coefficient=1.0 / alpha_inputs, offset=beta_inputs + n_inputs, coefficient=1.0 / alpha_guess_inputs, offset=beta_guess_inputs ) output_guess_calibration = AffineInputTransform( - n_outputs, coefficient=1.0 / alpha_outputs, offset=beta_outputs + n_outputs, coefficient=1.0 / alpha_guess_outputs, offset=beta_guess_outputs ) uncertainty_inputs = {