abs - surrogate-aided Bayesian calibration updates

modules/performUQ/common/adaptive_doe.py

@@ -1,6 +1,14 @@
 """Implements an adaptive design of experiments strategy using Gaussian Process (GP) and Principal Component Analysis (PCA)."""
 
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
 import numpy as np
+import uq_utilities
+from gp_model import remove_duplicate_inputs
+from scipy.stats import qmc
 
 
 class AdaptiveDesignOfExperiments:
@@ -28,6 +36,9 @@ def __init__(self, gp_model):
         self._kernel_for_doe()
         self._gp_for_doe()
 
+    def _scale(self, x):
+        return self.gp_model.apply_input_scaling(x, fit=False)
+
     def _hyperparameters_for_doe(self):
         """
         Compute the weighted average of kernel hyperparameters for DoE.
@@ -85,56 +96,6 @@ def _gp_for_doe(self):
         self.gp_model_for_doe.kern = self.kernel
         return self.gp_model_for_doe
 
-    def _imse_w_approximation(self, x_train, mci_samples, weights=None):
-        """
-        Compute the IMSE approximation for candidate training points.
-
-        Parameters
-        ----------
-            X_train (array-like): The current training data.
-            mci_samples (array-like): Monte Carlo integration samples.
-
-        Returns
-        -------
-            array: The IMSE values for the candidate training points.
-        """
-        candidate_training_points = mci_samples
-        self.gp_model_for_doe.set_XY(
-            x_train,
-            np.zeros((x_train.shape[0], 1)),
-        )
-        _, pred_var = self.gp_model_for_doe.predict(mci_samples)
-        n_theta = x_train.shape[1]
-        beta = 2.0 * n_theta
-        imse = np.zeros((candidate_training_points.shape[0], 1))
-        for i, candidate in enumerate(candidate_training_points):
-            correlation_vector = self.gp_model_for_doe.kern.K(
-                mci_samples, np.atleast_2d(candidate)
-            )
-            imse[i] = (1 / mci_samples.shape[0]) * np.sum(
-                (correlation_vector**beta) * pred_var
-            )
-        return imse
-
-    def _mse_approximation(self, x_train, mci_samples):
-        self.gp_model_for_doe.set_XY(
-            x_train,
-            np.zeros((x_train.shape[0], 1)),
-        )
-        _, pred_var = self.gp_model_for_doe.predict(mci_samples)
-        return np.reshape(pred_var, (-1, 1))
-
-    def _mse_w_approximation(self, x_train, mci_samples, weights):
-        self.gp_model_for_doe.set_XY(
-            x_train,
-            np.zeros((x_train.shape[0], 1)),
-        )
-        _, pred_var = self.gp_model_for_doe.predict(mci_samples)
-        if weights is None:
-            weights = np.ones_like(pred_var)
-        mse_w = pred_var.flatten() * weights.flatten()
-        return np.reshape(mse_w, (-1, 1))
-
     def select_training_points(
         self,
         x_train,
@@ -143,28 +104,208 @@ def select_training_points(
         *,
         use_mse_w=True,
         weights=None,
+        n_samples=4000,
+        seed=None,
     ):
         """
-        Select new training points based on the IMSE criterion.
+        Efficient sequential DoE using MSEw or IMSEw with vectorized acquisition.
 
         Parameters
         ----------
-            X_train (array-like): The current training data.
-            n_points (int): The number of new training points to select.
-            mci_samples (array-like): Monte Carlo integration samples.
+        x_train : array-like
+            Current training data (original space).
+        n_points : int
+            Number of new training points to select.
+        mci_samples : array-like
+            Monte Carlo integration samples (original space).
+        use_mse_w : bool
+            Whether to use MSEw (True) or IMSEw (False).
+        weights : array-like or None
+            Optional importance weights.
+        n_samples : int
+            Number of candidate points to generate (IMSEw only).
+        seed : int or None
+            Random seed for LHS sampling.
 
         Returns
         -------
-            array: The selected new training points.
+        selected_points : np.ndarray of shape (n_points, d)
+            Selected new training points.
         """
+        x_train = np.atleast_2d(x_train)
+        x_train, _ = remove_duplicate_inputs(
+            x_train, np.zeros((x_train.shape[0], 1))
+        )
+
+        selected_points = []
+
+        # 1. Fixed candidate pool
         if use_mse_w:
-            acquisition_function = self._mse_w_approximation
+            candidate_pool = mci_samples.copy()
         else:
-            acquisition_function = self._imse_w_approximation
+            bounds = compute_lhs_bounds(x_train, mci_samples, padding=0)
+            candidate_pool = generate_lhs_candidates(n_samples, bounds, seed=seed)
+
         for _ in range(n_points):
-            acquisition_function_values = acquisition_function(
-                x_train, mci_samples, weights
+            # 2. Scale everything
+            scaled_x_train = self._scale(x_train)
+            scaled_candidates = self._scale(candidate_pool)
+            scaled_mci_samples = self._scale(mci_samples)
+
+            # 3. Fit GP with dummy outputs
+            self.gp_model_for_doe.set_XY(
+                scaled_x_train, np.zeros((scaled_x_train.shape[0], 1))
             )
-            next_training_point = mci_samples[np.argmax(acquisition_function_values)]
-            x_train = np.vstack((x_train, next_training_point))
-        return x_train[-n_points:, :]
+
+            # 4. Predict variance at integration points
+            _, pred_var = self.gp_model_for_doe.predict(scaled_mci_samples)
+            if weights is not None:
+                pred_var *= weights.reshape(-1, 1)  # shape (n_mci, 1)
+
+            # 5. Vectorized acquisition
+            if use_mse_w:
+                # Just return the average variance (same for all candidates)
+                acquisition_values = np.mean(pred_var) * np.ones(
+                    (scaled_candidates.shape[0], 1)
+                )
+            else:
+                # Covariance matrix: shape (n_mci, n_candidates)
+                k_matrix = self.gp_model_for_doe.kern.K(
+                    scaled_mci_samples, scaled_candidates
+                )
+
+                # Normalize to get correlation matrix
+                kernel_var = self.gp_model_for_doe.kern.variance[0]  # scalar
+                corr_matrix = k_matrix / kernel_var  # element-wise division
+
+                beta = 2.0 * scaled_x_train.shape[1]
+                # beta = 1.0 * scaled_x_train.shape[1]
+
+                # Vectorized IMSEw: (corr^beta) * pred_var → mean over MCI samples
+                weighted_var = (
+                    corr_matrix**beta
+                ).T @ pred_var  # shape (n_candidates, 1)
+                acquisition_values = weighted_var / scaled_mci_samples.shape[0]
+
+            # 6. Select best candidate
+            idx = np.argmax(acquisition_values)
+            next_point = candidate_pool[idx]
+            selected_points.append(next_point)
+
+            # 7. Update training set and candidate pool
+            x_train = np.vstack([x_train, next_point])
+            candidate_pool = np.delete(candidate_pool, idx, axis=0)
+
+        return np.array(selected_points)
+
+    def write_gp_for_doe_to_json(self, filepath: str | Path):
+        """
+        Write DoE GP kernel hyperparameters and contributing model param_arrays to JSON.
+
+        Parameters
+        ----------
+        filepath : str or Path
+            Output file path.
+        """
+        if not hasattr(self, 'gp_model_for_doe'):
+            msg = 'gp_model_for_doe has not been initialized.'
+            raise RuntimeError(msg)
+
+        kernel = self.gp_model_for_doe.kern
+
+        # Detailed hyperparameters for the DoE model
+        doe_hyperparams = {
+            param.name: {
+                'value': param.values.tolist()  # noqa: PD011
+                if param.size > 1
+                else float(param.values),
+                'shape': param.shape,
+            }
+            for param in kernel.parameters
+        }
+
+        # Aggregation weights
+        if self.gp_model.use_pca:
+            weights = np.asarray(self.gp_model.pca_info['explained_variance_ratio'])[
+                : self.gp_model.pca_info['n_components']
+            ]
+            weights = (weights / np.sum(weights)).tolist()
+        else:
+            weights = (
+                np.ones(len(self.gp_model.model)) / len(self.gp_model.model)
+            ).tolist()
+
+        # Contributing models' param_arrays only
+        contributing_param_arrays = [
+            gp.kern.param_array.tolist() for gp in self.gp_model.model
+        ]
+
+        # Output structure
+        output = {
+            'doe_kernel_type': kernel.name,
+            'doe_ARD': kernel.ARD,
+            'doe_hyperparameters': doe_hyperparams,
+            'weighted_param_array': self.weighted_hyperparameters.tolist(),
+            'aggregation_weights': weights,
+            'contributing_param_arrays': contributing_param_arrays,
+        }
+
+        # Write JSON file
+        filepath = Path(filepath)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        with filepath.open('w') as f:
+            json.dump(uq_utilities.make_json_serializable(output), f, indent=4)
+
+
+def generate_lhs_candidates(n_samples, input_bounds, seed=None):
+    """
+    Generate LHS candidate points using scipy's QMC module with an optional random seed.
+
+    Parameters
+    ----------
+    n_samples : int
+        Number of candidate points to generate.
+    input_bounds : array-like of shape (d, 2)
+        Lower and upper bounds for each input dimension.
+    seed : int or None
+        Random seed for reproducibility (default: None).
+
+    Returns
+    -------
+    candidates : np.ndarray of shape (n_samples, d)
+        Generated candidate points in the original input space.
+    """
+    input_bounds = np.asarray(input_bounds)
+    d = input_bounds.shape[0]
+
+    sampler = qmc.LatinHypercube(d, seed=seed)
+    lhs_unit = sampler.random(n=n_samples)
+    candidates = qmc.scale(lhs_unit, input_bounds[:, 0], input_bounds[:, 1])
+    return candidates  # noqa: RET504
+
+
+def compute_lhs_bounds(x_train, mci_samples, padding=0):
+    """
+    Compute input bounds for LHS based on x_train and mci_samples.
+
+    Parameters
+    ----------
+    x_train : array-like, shape (n_train, d)
+    mci_samples : array-like, shape (n_mci, d)
+    padding : float
+        Relative padding (e.g., 0.05 adds ±5% range to each side).
+
+    Returns
+    -------
+    bounds : np.ndarray of shape (d, 2)
+        Array of (min, max) bounds for each dimension.
+    """
+    x_all = np.vstack([x_train, mci_samples])
+    min_vals = np.min(x_all, axis=0)
+    max_vals = np.max(x_all, axis=0)
+
+    ranges = max_vals - min_vals
+    min_vals_padded = min_vals - padding * ranges
+    max_vals_padded = max_vals + padding * ranges
+
+    return np.vstack([min_vals_padded, max_vals_padded]).T