Pyomo.DoE: fix trace/Cholesky initialization consistency

pyomo/contrib/doe/doe.py

@@ -52,9 +52,10 @@
 
 import pyomo.environ as pyo
 from pyomo.contrib.doe.utils import (
+    _SMALL_TOLERANCE_DEFINITENESS,
     check_FIM,
     compute_FIM_metrics,
-    _SMALL_TOLERANCE_DEFINITENESS,
+    regularize_fim_for_cholesky,
 )
 
 from pyomo.opt import SolverStatus
@@ -281,7 +282,6 @@ def run_doe(self, model=None, results_file=None):
         results_file: string name of the file path to save the results
                       to in the form of a .json file
                       default: None --> don't save
-
         """
         # Check results file name
         if results_file is not None:
@@ -376,6 +376,11 @@ def run_doe(self, model=None, results_file=None):
             if self.objective_option == ObjectiveLib.trace:
                 trace_val = np.trace(np.linalg.pinv(self.get_FIM()))
                 model.obj_cons.egb_fim_block.outputs["A-opt"].set_value(trace_val)
+            elif self.objective_option == ObjectiveLib.pseudo_trace:
+                pseudo_trace_val = np.trace(np.array(self.get_FIM()))
+                model.obj_cons.egb_fim_block.outputs["pseudo-A-opt"].set_value(
+                    pseudo_trace_val
+                )
             elif self.objective_option == ObjectiveLib.determinant:
                 det_val = np.linalg.det(np.array(self.get_FIM()))
                 model.obj_cons.egb_fim_block.outputs["log-D-opt"].set_value(
@@ -389,63 +394,8 @@ def run_doe(self, model=None, results_file=None):
                 cond_number = np.log(np.abs(np.max(eig) / np.min(eig)))
                 model.obj_cons.egb_fim_block.outputs["ME-opt"].set_value(cond_number)
 
-        # If the model has L, initialize it with the solved FIM
-        if hasattr(model, "L"):
-            # Get the FIM values
-            fim_vals = [
-                pyo.value(model.fim[i, j])
-                for i in model.parameter_names
-                for j in model.parameter_names
-            ]
-            fim_np = np.array(fim_vals).reshape(
-                (len(model.parameter_names), len(model.parameter_names))
-            )
-
-            # Need to compute the full FIM before
-            # initializing the Cholesky factorization
-            if self.only_compute_fim_lower:
-                fim_np = fim_np + fim_np.T - np.diag(np.diag(fim_np))
-
-            # Check if the FIM is positive definite
-            # If not, add jitter to the diagonal
-            # to ensure positive definiteness
-            min_eig = np.min(np.linalg.eigvals(fim_np))
-
-            if min_eig < _SMALL_TOLERANCE_DEFINITENESS:
-                # Raise the minimum eigenvalue to at
-                # least _SMALL_TOLERANCE_DEFINITENESS
-                jitter = np.min(
-                    [
-                        -min_eig + _SMALL_TOLERANCE_DEFINITENESS,
-                        _SMALL_TOLERANCE_DEFINITENESS,
-                    ]
-                )
-            else:
-                # No jitter needed
-                jitter = 0
-
-            # Add jitter to the diagonal to ensure positive definiteness
-            L_vals_sq = np.linalg.cholesky(
-                fim_np + jitter * np.eye(len(model.parameter_names))
-            )
-            for i, c in enumerate(model.parameter_names):
-                for j, d in enumerate(model.parameter_names):
-                    model.L[c, d].value = L_vals_sq[i, j]
-
-            # Initialize the inverse of L if it exists
-            if hasattr(model, "L_inv"):
-                L_inv_vals = np.linalg.inv(L_vals_sq)
-
-                for i, c in enumerate(model.parameter_names):
-                    for j, d in enumerate(model.parameter_names):
-                        if i >= j:
-                            model.L_inv[c, d].value = L_inv_vals[i, j]
-                        else:
-                            model.L_inv[c, d].value = 0.0
-                # Initialize the cov_trace if it exists
-                if hasattr(model, "cov_trace"):
-                    initial_cov_trace = np.sum(L_inv_vals**2)
-                    model.cov_trace.value = initial_cov_trace
+        # Keep Cholesky-related variables synchronized with current FIM values
+        self._initialize_cholesky_from_fim(model=model)
 
         if hasattr(model, "determinant"):
             model.determinant.value = np.linalg.det(np.array(self.get_FIM()))
@@ -537,6 +487,92 @@ def run_doe(self, model=None, results_file=None):
             with open(results_file, "w") as file:
                 json.dump(self.results, file)
 
+    def _get_fim_numpy(self, model):
+        """
+        Assemble the current FIM variable values into a NumPy array.
+
+        Parameters
+        ----------
+        model: ConcreteModel
+            DoE model containing variable ``fim``.
+
+        Returns
+        -------
+        ndarray
+            Dense FIM array. If ``only_compute_fim_lower`` is True, the
+            returned array is symmetrized from the lower triangle.
+        """
+        fim_vals = [
+            pyo.value(model.fim[i, j])
+            for i in model.parameter_names
+            for j in model.parameter_names
+        ]
+        fim_np = np.array(fim_vals, dtype=float).reshape(
+            (len(model.parameter_names), len(model.parameter_names))
+        )
+        if self.only_compute_fim_lower:
+            fim_np = fim_np + fim_np.T - np.diag(np.diag(fim_np))
+        return fim_np
+
+    def _initialize_cholesky_from_fim(self, model=None):
+        """
+        Synchronize Cholesky-related variables using the current FIM.
+
+        Parameters
+        ----------
+        model: ConcreteModel, optional
+            DoE model to update. Defaults to ``self.model``.
+
+        Returns
+        -------
+        None
+            Updates model values in place for available variables:
+            ``L``, ``L_inv``, ``fim_inv``, and ``cov_trace``.
+        """
+        if model is None:
+            model = self.model
+        if not hasattr(model, "L"):
+            # The model doesn't have the Cholesky variables, so we can't initialize them.
+            # This happens if the function is called with a model using GreyBox.
+            return
+
+        fim_np = self._get_fim_numpy(model)
+        fim_pd, _ = regularize_fim_for_cholesky(fim_np)
+
+        L_vals = np.linalg.cholesky(fim_pd)
+        for i, c in enumerate(model.parameter_names):
+            for j, d in enumerate(model.parameter_names):
+                if i >= j:
+                    model.L[c, d].value = L_vals[i, j]
+                else:
+                    model.L[c, d].value = 0.0
+
+        if hasattr(model, "L_inv"):
+            L_inv_vals = np.linalg.inv(L_vals)
+            for i, c in enumerate(model.parameter_names):
+                for j, d in enumerate(model.parameter_names):
+                    if i >= j:
+                        model.L_inv[c, d].value = L_inv_vals[i, j]
+                    else:
+                        model.L_inv[c, d].value = 0.0
+
+        if hasattr(model, "fim_inv"):
+            # Use the pseudo-inverse here rather than the strict inverse.
+            # The jittered matrix should be positive definite, but ``pinv``
+            # is safer for borderline ill-conditioned cases and matches the
+            # defensive approach already used when initializing ``fim_inv``
+            # from user-provided starting values.
+            fim_inv_vals = np.linalg.pinv(fim_pd)
+            for i, c in enumerate(model.parameter_names):
+                for j, d in enumerate(model.parameter_names):
+                    if self.only_compute_fim_lower and i < j:
+                        model.fim_inv[c, d].value = 0.0
+                    else:
+                        model.fim_inv[c, d].value = fim_inv_vals[i, j]
+
+        if hasattr(model, "cov_trace"):
+            model.cov_trace.value = np.trace(fim_inv_vals)
+
     # Perform multi-experiment doe (sequential, or ``greedy`` approach)
     def run_multi_doe_sequential(self, N_exp=1):
         raise NotImplementedError("Multiple experiment optimization not yet supported.")
@@ -898,6 +934,7 @@ def create_doe_model(self, model=None):
             self.only_compute_fim_lower
             and self.objective_option == ObjectiveLib.determinant
             and not self.Cholesky_option
+            and not self.use_grey_box
         ):
             raise ValueError(
                 "Cannot compute determinant with explicit formula "
@@ -1372,29 +1409,14 @@ def create_objective_function(self, model=None):
         model.obj_cons = pyo.Block()
 
         # Assemble the FIM matrix. This is helpful for initialization!
-        fim_vals = [
-            model.fim[bu, un].value
-            for i, bu in enumerate(model.parameter_names)
-            for j, un in enumerate(model.parameter_names)
-        ]
-        fim = np.array(fim_vals).reshape(
-            len(model.parameter_names), len(model.parameter_names)
-        )
+        fim = self._get_fim_numpy(model)
 
         ### Initialize the Cholesky decomposition matrix
         if self.Cholesky_option and self.objective_option in (
             ObjectiveLib.determinant,
             ObjectiveLib.trace,
         ):
-            # Calculate the eigenvalues of the FIM matrix
-            eig = np.linalg.eigvals(fim)
-
-            # If the smallest eigenvalue is (practically) negative,
-            # add a diagonal matrix to make it positive definite
-            if min(eig) < small_number:
-                fim = fim + np.eye(len(model.parameter_names)) * (
-                    small_number - min(eig)
-                )
+            fim, _ = regularize_fim_for_cholesky(fim)
 
             # Compute the Cholesky decomposition of the FIM matrix
             L = np.linalg.cholesky(fim)
@@ -1663,6 +1685,11 @@ def FIM_egb_cons(m, p1, p2):
             model.objective = pyo.Objective(
                 expr=model.obj_cons.egb_fim_block.outputs["A-opt"], sense=pyo.minimize
             )
+        elif self.objective_option == ObjectiveLib.pseudo_trace:
+            model.objective = pyo.Objective(
+                expr=model.obj_cons.egb_fim_block.outputs["pseudo-A-opt"],
+                sense=pyo.maximize,
+            )
         elif self.objective_option == ObjectiveLib.determinant:
             model.objective = pyo.Objective(
                 expr=model.obj_cons.egb_fim_block.outputs["log-D-opt"],

pyomo/contrib/doe/grey_box_utilities.py

@@ -108,9 +108,11 @@ def __init__(self, doe_object, objective_option="determinant", logger_level=None
     def _get_FIM(self):
         # Grabs the current FIM subject
         # to the input values.
-        # This function currently assumes
-        # that we use a lower triangular
-        # FIM.
+        # Inputs store one triangular half
+        # of a symmetric FIM. Reconstruct
+        # the full symmetric matrix here,
+        # consistent with manuscript equation S5.
+        # https://arxiv.org/abs/2604.03354v1
         upt_FIM = self._input_values
 
         # Create FIM in the correct way
@@ -164,6 +166,8 @@ def output_names(self):
 
         if self.objective_option == ObjectiveLib.trace:
             obj_name = "A-opt"
+        elif self.objective_option == ObjectiveLib.pseudo_trace:
+            obj_name = "pseudo-A-opt"
         elif self.objective_option == ObjectiveLib.determinant:
             obj_name = "log-D-opt"
         elif self.objective_option == ObjectiveLib.minimum_eigenvalue:
@@ -196,6 +200,8 @@ def evaluate_outputs(self):
 
         if self.objective_option == ObjectiveLib.trace:
             obj_value = np.trace(np.linalg.pinv(M))
+        elif self.objective_option == ObjectiveLib.pseudo_trace:
+            obj_value = np.trace(M)
         elif self.objective_option == ObjectiveLib.determinant:
             sign, logdet = np.linalg.slogdet(M)
             obj_value = logdet
@@ -231,6 +237,8 @@ def finalize_block_construction(self, pyomo_block):
         # objective function.
         if self.objective_option == ObjectiveLib.trace:
             pyomo_block.outputs["A-opt"] = output_value
+        elif self.objective_option == ObjectiveLib.pseudo_trace:
+            pyomo_block.outputs["pseudo-A-opt"] = output_value
         elif self.objective_option == ObjectiveLib.determinant:
             pyomo_block.outputs["log-D-opt"] = output_value
         elif self.objective_option == ObjectiveLib.minimum_eigenvalue:
@@ -268,6 +276,8 @@ def evaluate_jacobian_outputs(self):
             # is -inv(FIM) @ inv(FIM). Add reference to
             # pyomo.DoE 2.0 manuscript S.I.
             jac_M = -Minv @ Minv
+        elif self.objective_option == ObjectiveLib.pseudo_trace:
+            jac_M = np.eye(self._n_params, dtype=np.float64)
         elif self.objective_option == ObjectiveLib.determinant:
             Minv = np.linalg.pinv(M)
             # Derivative formula derived using tensor

pyomo/contrib/doe/tests/test_doe_build.py

@@ -25,9 +25,6 @@
 
 if scipy_available:
     from pyomo.contrib.doe import DesignOfExperiments
-    from pyomo.contrib.doe.examples.reactor_example import (
-        ReactorExperiment as FullReactorExperiment,
-    )
     from pyomo.contrib.parmest.examples.rooney_biegler.rooney_biegler import (
         RooneyBieglerExperiment,
     )
@@ -98,10 +95,6 @@ def get_FIM_FIMPrior_Q_L(doe_obj=None):
         for j in model.parameter_names
     ]
     sigma_inv = [1 / v for k, v in model.scenario_blocks[0].measurement_error.items()]
-    param_vals = np.array(
-        [[v for k, v in model.scenario_blocks[0].unknown_parameters.items()]]
-    )
-
     FIM_vals_np = np.array(FIM_vals).reshape((n_param, n_param))
     FIM_prior_vals_np = np.array(FIM_prior_vals).reshape((n_param, n_param))
 

pyomo/contrib/doe/tests/test_doe_errors.py

@@ -27,9 +27,6 @@
         BadExperiment,
         RooneyBieglerExperimentFlag,
     )
-    from pyomo.contrib.parmest.examples.rooney_biegler.rooney_biegler import (
-        RooneyBieglerExperiment,
-    )
 
 from pyomo.contrib.doe.examples.rooney_biegler_doe_example import run_rooney_biegler_doe
 from pyomo.opt import SolverFactory
@@ -95,7 +92,7 @@ def test_experiment_none_error(self):
             # Experiment provided as None
             DoE_args = get_standard_args(None, fd_method, obj_used, flag_val)
 
-            doe_obj = DesignOfExperiments(**DoE_args)
+            DesignOfExperiments(**DoE_args)
 
     def test_reactor_check_no_get_labeled_model(self):
         fd_method = "central"
@@ -110,7 +107,7 @@ def test_reactor_check_no_get_labeled_model(self):
         ):
             DoE_args = get_standard_args(experiment, fd_method, obj_used, flag_val)
 
-            doe_obj = DesignOfExperiments(**DoE_args)
+            DesignOfExperiments(**DoE_args)
 
     def test_reactor_check_no_experiment_outputs(self):
         fd_method = "central"
@@ -781,6 +778,25 @@ def test_invalid_trace_without_cholesky(self):
         ):
             doe_obj.create_objective_function()
 
+    def test_invalid_determinant_without_cholesky(self):
+        fd_method = "central"
+        obj_used = "determinant"
+
+        experiment = get_rooney_biegler_experiment_flag()
+
+        DoE_args = get_standard_args(experiment, fd_method, obj_used, flag=None)
+        DoE_args["_Cholesky_option"] = False
+
+        doe_obj = DesignOfExperiments(**DoE_args)
+
+        # The explicit determinant formulation needs the full FIM, so we keep
+        # this regression test focused on the early validation guard.
+        with self.assertRaisesRegex(
+            ValueError,
+            "Cannot compute determinant with explicit formula if only_compute_fim_lower is True.",
+        ):
+            doe_obj.create_doe_model()
+
 
 if __name__ == "__main__":
     unittest.main()