[WIP] Add penalization term to calibration loss for alpha/beta uncertainty

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,51 +200,78 @@ 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_guess_input_list = []
+    beta_guess_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_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:
-        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_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 = {
+        "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,
             )
         )