Pyomo. DoE: Sensitivity initialization

pyomo/contrib/doe/doe.py

@@ -46,11 +46,14 @@
 
 import pyomo.environ as pyo
 from pyomo.contrib.doe.utils import (
-    check_FIM,
+    check_matrix,
     compute_FIM_metrics,
     _SMALL_TOLERANCE_DEFINITENESS,
+    snake_traversal_grid_sampling,
+    update_model_from_suffix,
 )
 
+
 from pyomo.opt import SolverStatus
 
 
@@ -67,6 +70,11 @@ class FiniteDifferenceStep(Enum):
     backward = "backward"
 
 
+class SensitivityInitialization(Enum):
+    snake_traversal = "snake_traversal"
+    nested_for_loop = "nested_for_loop"
+
+
 class DesignOfExperiments:
     def __init__(
         self,
@@ -621,6 +629,9 @@ def _sequential_FIM(self, model=None):
             if self.scale_nominal_param_value:
                 scale_factor *= v
 
+            # TODO: scale by response values (i.e., measurement values)
+            # if self.scale_response_values:
+            #     scale_factor /= measurement_vals_np[:, col_1].mean()
             # Calculate the column of the sensitivity matrix
             self.seq_jac[:, i] = (
                 measurement_vals_np[:, col_1] - measurement_vals_np[:, col_2]
@@ -1376,7 +1387,7 @@ def check_model_FIM(self, model=None, FIM=None):
             )
 
         # Check FIM is positive definite and symmetric
-        check_FIM(FIM)
+        check_matrix(FIM)
 
         self.logger.info(
             "FIM provided matches expected dimensions from model and is approximately positive (semi) definite."
@@ -1611,6 +1622,350 @@ def compute_FIM_full_factorial(
 
         return self.fim_factorial_results
 
+    def compute_FIM_factorial(
+        self,
+        model=None,
+        design_vals: dict = None,
+        abs_change: list = None,
+        rel_change: list = None,
+        n_designs: int = 5,
+        method="sequential",
+        df_settings=(True, None, None, 500),
+        initialization_scheme="snake_traversal",
+        file_name: str = None,
+    ):
+        """Will run a simulation-based factorial exploration of the experimental input
+        space (i.e., a ``grid search`` or ``parameter sweep``) to understand how the
+        FIM metrics change as a function of the experimental design space. This method
+        can be used for both full factorial and fractional factorial designs.
+
+        Parameters
+        ----------
+        model : DoE model, optional
+            The model to perform the full factorial exploration on. Default: None
+        design_vals : dict,
+            dict of lists or other array-like objects, of the form
+            {"var_name": <var_values>}. Default: None.
+            The `design_values` should have the key(s) passed as strings that is a
+            subset of the `experiment_inputs`. If one or more design variables are not
+            to be changed, then they should not be passed in the `design_values`
+            dictionary, but if they are passed in the dictionary, then they must be a
+            array-like object of floats. For example, if our experiment has 4 design variables
+            (i.e., `experiment_inputs`): model.x1, model.x2, model.x3, and model.x4,
+            their values may be passed as, design_values= {"x1": [1, 2, 3], "x3": [7],
+            "x4": [-10, 20]}. In this case, x2 will not be changed and will be fixed at
+            the value in the model.
+            If design_values is provided, then `abs_change` and `rel_change` are ignored.
+        abs_change : list, optional
+            Absolute change in the design variable values. Default: None.
+            If provided, will use this value to generate the design values.
+            If `abs_change` is provided, but `rel_change` is not provided, `rel_change`
+            will be set to zero.
+            Formula to calculate the design values:
+                change_in_value = lower_bound * rel_change + abs_change`
+                design_value += design_value + change_in_value
+        rel_change : list, optional
+            Relative change in the design variable values. Default: None.
+            If provided, will use this value to generate the design values.
+            If `rel_change` is provided, but `abs_change` is not provided, `abs_change`
+            will be set to zero.
+        n_designs : int, optional
+            Number of designs to generate for each design variable. Default: 5.
+            If `abs_change` and/or `rel_change` are provided, this value will be ignored.
+        method : str, optional
+            string to specify which method should be used. options are ``kaug`` and
+            ``sequential`. Default: "sequential"
+        df_settings : tuple, optional
+            A tuple containing the settings for set_option() method of the pandas
+            DataFrame. Default: (True, None, None, 500)
+            - first element: whether to return a pandas DataFrame (True/False)
+            - second element: number of max_columns for the DataFrame. Default: None,
+                i.e., no limit on the number of columns.
+            - third element: number of max_rows for the DataFrame. Default: None,
+                i.e., no limit on the number of rows.
+            - fourth element: display width for the DataFrame. Default: 500.
+        initialization_scheme : str, optional
+            Which scheme to use for initializing the design variables.
+            Options are ``"snake_traversal"`` and ``"nested_for_loop"``.
+            Default: "snake_traversal"
+        file_name : str, optional
+            if provided, will save the results to a json file. Default: None
+
+        Returns
+        -------
+        factorial_results: dict
+            A dictionary containing the results of the factorial design with the
+            following keys:
+            - keys of model's experiment_inputs
+            - "log10(D-opt)": list of D-optimality values
+            - "log10(A-opt)": list of A-optimality values
+            - "log10(E-opt)": list of E-optimality values
+            - "log10(ME-opt)": list of ME-optimality values
+            - "eigval_min": list of minimum eigenvalues
+            - "eigval_max": list of maximum eigenvalues
+            - "det_FIM": list of determinants of the FIM
+            - "trace_FIM": list of traces of the FIM
+            - "solve_time": list of solve times
+            - "total_points": total number of points in the factorial design
+            - "success_count": number of successful runs
+            - "failure_count": number of failed runs
+            - "FIM_all": list of all FIMs computed for each point in the factorial
+        """
+
+        # Start timer
+        sp_timer = TicTocTimer()
+        sp_timer.tic(msg=None)
+        self.logger.info("Beginning Factorial Design.")
+
+        # Make new model for factorial design
+        self.factorial_model = self.experiment.get_labeled_model(
+            **self.get_labeled_model_args
+        ).clone()
+        model = self.factorial_model
+
+        if design_vals is not None:
+            # Check whether the design_ranges keys are in the experiment_inputs
+            design_keys = set(design_vals.keys())
+            map_keys = set([k.name for k in model.experiment_inputs.keys()])
+            if not design_keys.issubset(map_keys):
+                incorrect_given_keys = design_keys - map_keys
+                suggested_keys = map_keys - design_keys
+                raise ValueError(
+                    f"design_values keys: {incorrect_given_keys} are incorrect."
+                    f"The keys should be from the following keys: {suggested_keys}."
+                )
+
+            # Get the design map keys that match the design_values keys
+            design_map_keys = [
+                k
+                for k in model.experiment_inputs.keys()
+                if k.name in design_vals.keys()
+            ]
+            # This ensures that the order of the design_values keys matches the order of the
+            # design_map_keys so that design_point can be constructed correctly in the loop.
+            design_values = [design_vals[k.name] for k in design_map_keys]
+
+            # Create a temporary suffix to pass in `update_model_from_suffix`
+            design_suff = pyo.Suffix(direction=pyo.Suffix.LOCAL)
+            design_suff.update((k, None) for k in design_map_keys)
+
+        else:
+            design_keys = [k for k in model.experiment_inputs.keys()]
+            if abs_change:
+                if len(abs_change) != len(design_keys):
+                    raise ValueError(
+                        "`abs_change` must have the same length of "
+                        f"`{len(design_keys)}` as `design_keys`."
+                    )
+
+            if rel_change:
+                if len(rel_change) != len(design_keys):
+                    raise ValueError(
+                        "`rel_change` must have the same length of "
+                        f"`{len(design_keys)}` as `design_keys`."
+                    )
+
+            # if either abs_change or rel_change is not provided, set it to list of
+            # zeros
+            if abs_change or rel_change:
+                if not abs_change:
+                    abs_change = [0] * len(design_keys)
+                elif not rel_change:
+                    rel_change = [0] * len(design_keys)
+
+            design_values = []
+            # loop over design keys and generate design values
+            for i, comp in enumerate(design_keys):
+                lb = comp.lb
+                ub = comp.ub
+                # Check if the component has finite lower and upper bounds
+                if lb is None or ub is None:
+                    raise ValueError(
+                        f"{comp.name} does not have a lower or upper bound."
+                    )
+
+                if abs_change is None and rel_change is None:
+                    des_val = np.linspace(lb, ub, n_designs)
+
+                # if abs_change and/or rel_change is provided, generate design values
+                # using the formula:
+                # change_in_value = lower_bound * rel_change + abs_change
+                elif abs_change or rel_change:
+                    des_val = []
+                    del_val = comp.lb * rel_change[i] + abs_change[i]
+                    if del_val == 0:
+                        raise ValueError(
+                            f"Design variable {comp.name} has no change in value - "
+                            "check abs_change and rel_change values."
+                        )
+                    val = lb
+                    while val <= ub:
+                        des_val.append(val)
+                        val += del_val
+
+                else:
+                    raise ValueError(
+                        "Unexpected error in generating design values. Please check the"
+                        " input parameters."
+                    )
+
+                design_values.append(des_val)
+
+        # generate the factorial points based on the initialization scheme
+        try:
+            scheme_enum = SensitivityInitialization(initialization_scheme)
+        except ValueError:
+            self.logger.warning(
+                f"Initialization scheme '{initialization_scheme}' is not recognized. "
+                "Using `snake_traversal` as the default initialization scheme."
+            )
+            scheme_enum = SensitivityInitialization.snake_traversal
+
+        if scheme_enum == SensitivityInitialization.snake_traversal:
+            factorial_points = snake_traversal_grid_sampling(*design_values)
+        elif scheme_enum == SensitivityInitialization.nested_for_loop:
+            factorial_points = product(*design_values)
+
+        # TODO: Add more initialization schemes
+
+        factorial_points_list = list(factorial_points)
+
+        factorial_results = {k.name: [] for k in model.experiment_inputs.keys()}
+        factorial_results.update(
+            {
+                "log10(D-opt)": [],
+                "log10(A-opt)": [],
+                "log10(E-opt)": [],
+                "log10(ME-opt)": [],
+                "eigval_min": [],
+                "eigval_max": [],
+                "det_FIM": [],
+                "trace_FIM": [],
+                "solve_time": [],
+            }
+        )
+
+        success_count = 0
+        failure_count = 0
+        total_points = len(factorial_points_list)
+
+        # save the FIM for each point in the factorial design
+        self.n_parameters = len(model.unknown_parameters)
+        FIM_all = np.zeros((total_points, self.n_parameters, self.n_parameters))
+
+        time_set = []
+        curr_point = 1  # Initial current point
+        for design_point in factorial_points_list:
+            if design_vals:
+                update_model_from_suffix(design_suff, design_point)
+            else:
+                update_model_from_suffix(model.experiment_inputs, design_point)
+
+            # Timing and logging objects
+            self.logger.info(f"=======Iteration Number: {curr_point} =======")
+            iter_timer = TicTocTimer()
+            iter_timer.tic(msg=None)
+
+            try:
+                curr_point = success_count + failure_count + 1
+                self.logger.info(f"This is run {curr_point} out of {total_points}.")
+                self.compute_FIM(model=model, method=method)
+                success_count += 1
+                # iteration time
+                iter_t = iter_timer.toc(msg=None)
+                time_set.append(iter_t)
+
+                # More logging
+                self.logger.info(
+                    f"The code has run for {round(sum(time_set), 2)} seconds."
+                )
+                self.logger.info(
+                    "Estimated remaining time:  %s seconds",
+                    round(
+                        sum(time_set) / (curr_point) * (total_points - curr_point + 1),
+                        2,
+                    ),
+                )
+            except:
+                self.logger.warning(
+                    ":::::::::::Warning: Cannot converge this run.::::::::::::"
+                )
+                failure_count += 1
+                self.logger.warning("failed count:", failure_count)
+
+                self._computed_FIM = np.zeros(self.prior_FIM.shape)
+
+                iter_t = iter_timer.toc(msg=None)
+                time_set.append(iter_t)
+
+            FIM = self._computed_FIM
+
+            # Save FIM for the current design point
+            FIM_all[curr_point - 1, :, :] = FIM
+
+            # Compute and record metrics on FIM
+            det_FIM, trace_FIM, E_vals, E_vecs, D_opt, A_opt, E_opt, ME_opt = (
+                compute_FIM_metrics(FIM)
+            )
+
+            for k in model.experiment_inputs.keys():
+                factorial_results[k.name].append(pyo.value(k))
+
+            factorial_results["log10(D-opt)"].append(D_opt)
+            factorial_results["log10(A-opt)"].append(A_opt)
+            factorial_results["log10(E-opt)"].append(E_opt)
+            factorial_results["log10(ME-opt)"].append(ME_opt)
+            factorial_results["eigval_min"].append(np.min(E_vals))
+            factorial_results["eigval_max"].append(np.max(E_vals))
+            factorial_results["det_FIM"].append(det_FIM)
+            factorial_results["trace_FIM"].append(trace_FIM)
+            factorial_results["solve_time"].append(time_set[-1])
+
+        factorial_results.update(
+            {
+                "total_points": total_points,
+                "success_count": success_count,
+                "failure_count": failure_count,
+                "FIM_all": FIM_all.tolist(),  # Save all FIMs
+            }
+        )
+        self.factorial_results = factorial_results
+
+        # Set pandas DataFrame options
+        if df_settings[0]:
+            with pd.option_context(
+                "display.max_columns",
+                df_settings[1],
+                "display.max_rows",
+                df_settings[2],
+                "display.width",
+                df_settings[3],
+            ):
+                exclude_keys = {
+                    "total_points",
+                    "success_count",
+                    "failure_count",
+                    "FIM_all",
+                }
+                dict_for_df = {
+                    k: v for k, v in factorial_results.items() if k not in exclude_keys
+                }
+                res_df = pd.DataFrame(dict_for_df)
+                print("\n\n=========Factorial results===========")
+                print("Total points:", total_points)
+                print("Success count:", success_count)
+                print("Failure count:", failure_count)
+                print("\n")
+                print(res_df)
+
+        # Save the results to a json file based on the file_name provided
+        if file_name is not None:
+            with open(file_name + ".json", "w") as f:
+                json.dump(self.factorial_results, f, indent=4)
+                self.logger.info(f"Results saved to {file_name}.json.")
+
+        return self.factorial_results
+
     # TODO: Overhaul plotting functions to not use strings
     # TODO: Make the plotting functionalities work for >2 design features
     def draw_factorial_figure(