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

Author: RemiLehe

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()
 

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,
             )
         )