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..529e51f 100644 --- a/ml/train_model.py +++ b/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, ) )