ENH: Make regularized least squares as defualt reconstruction method, compute prior and noise if not given

Author: yb6599

pysensors/optimizers/_tpgr.py

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 from sklearn.base import BaseEstimator
 from sklearn.utils.validation import check_is_fitted
@@ -15,19 +17,30 @@ class TPGR(BaseEstimator):
 
     """
 
-    def __init__(self, prior, n_sensors=None, noise=1):
-        self.prior = prior
+    def __init__(self, n_sensors=None, noise=None, prior="decreasing"):
         self.n_sensors = n_sensors
         self.noise = noise
         self.sensors_ = None
-        self.G = None
+        self.prior = prior
 
-    def fit(self, basis_matrix):
-        if self.n_sensors is None:
-            self.n_sensors = basis_matrix.shape[
-                1
-            ]  # Set number of sensors to number of basis modes if unspecified
-        G = basis_matrix @ np.diag(self.prior)
+    def fit(self, basis_matrix, singular_values=None):
+        if isinstance(self.prior, str) and self.prior == "decreasing":
+            computed_prior = singular_values
+        elif isinstance(self.prior, np.ndarray):
+            if self.prior.ndim != 1:
+                raise ValueError("prior must be a 1D array")
+            if self.prior.shape[0] != basis_matrix.shape[1]:
+                raise ValueError(
+                    f"prior must be of shape {(basis_matrix.shape[1],)},"
+                    f" but got {self.prior.shape[0]}"
+                )
+            computed_prior = self.prior
+        if self.noise is None:
+            warnings.warn(
+                "noise is None. noise will be set to the " "average of the prior"
+            )
+            self.noise = computed_prior.mean()
+        G = basis_matrix @ np.diag(computed_prior)
         self.G = G
         n = G.shape[0]
         mask = np.ones(n, dtype=bool)

pysensors/reconstruction/_sspor.py

@@ -6,7 +6,7 @@
 from sklearn.utils.validation import check_is_fitted
 
 from ..basis import Identity
-from ..optimizers import CCQR, QR
+from ..optimizers import CCQR, QR, TPGR
 from ..utils import validate_input
 
 INT_DTYPES = (int, np.int64, np.int32, np.int16, np.int8)
@@ -150,11 +150,19 @@ def fit(self, x, quiet=False, prefit_basis=False, seed=None, **optimizer_kws):
         # Check that n_sensors doesn't exceed dimension of basis vectors and
         # that it doesn't exceed the number of samples when using the CCQR optimizer.
         self._validate_n_sensors()
-
+        # Calculate the singular values
+        X_proj = x @ self.basis_matrix_
+        # Normalized singular values
+        self.singular_values = np.linalg.norm(X_proj, axis=0) / np.sqrt(x.shape[0])
         # Find sparse sensor locations
-        self.ranked_sensors_ = self.optimizer.fit(
-            self.basis_matrix_, **optimizer_kws
-        ).get_sensors()
+        if isinstance(self.optimizer, TPGR):
+            self.ranked_sensors_ = self.optimizer.fit(
+                self.basis_matrix_, self.singular_values, **optimizer_kws
+            ).get_sensors()
+        else:
+            self.ranked_sensors_ = self.optimizer.fit(
+                self.basis_matrix_, **optimizer_kws
+            ).get_sensors()
 
         # Randomly shuffle sensors after self.basis.n_basis_modes
         rng = np.random.default_rng(seed)
@@ -165,7 +173,7 @@ def fit(self, x, quiet=False, prefit_basis=False, seed=None, **optimizer_kws):
 
         return self
 
-    def predict(self, x, method=None, noise=None, prior=None, **solve_kws):
+    def predict(self, x, method=None, noise=None, prior="decreasing", **solve_kws):
         """
         Predict values at all positions given measurements at sensor locations.
 
@@ -200,6 +208,25 @@ def predict(self, x, method=None, noise=None, prior=None, **solve_kws):
             )
 
         if method is None:
+            if isinstance(prior, str) and prior == "decreasing":
+                computed_prior = self.singular_values
+            elif isinstance(prior, np.ndarray):
+                if prior.ndim != 1:
+                    raise ValueError("prior must be a 1D array")
+                if prior.shape[0] != self.basis_matrix_.shape[1]:
+                    raise ValueError(
+                        f"prior must be of shape {(self.basis_matrix_.shape[1],)},"
+                        f" but got {prior.shape}"
+                    )
+                computed_prior = prior
+            if noise is None:
+                warnings.warn(
+                    "noise is None. noise will be set to the "
+                    "average of the normalized prior"
+                )
+                noise = computed_prior.mean()
+            return self._regularized_reconstruction(x, computed_prior, noise)
+        elif method == "unregularized":
             # Square matrix
             if self.n_sensors == self.basis_matrix_.shape[1]:
                 return self._square_predict(
@@ -210,12 +237,12 @@ def predict(self, x, method=None, noise=None, prior=None, **solve_kws):
                 return self._rectangular_predict(
                     x, self.ranked_sensors_[: self.n_sensors], **solve_kws
                 )
-        if method == "mle":
-            return self._max_likelihood_reconstruction(x, prior, noise)
+        else:
+            raise NotImplementedError("Method not implemented")
 
-    def _max_likelihood_reconstruction(self, x, prior, noise):
+    def _regularized_reconstruction(self, x, prior, noise):
         """
-        Reconstruct the state using maximum likelihood reconstruction
+        Reconstruct the state using regularized reconstruction
 
         See the following reference for more information
 
@@ -234,15 +261,11 @@ def _max_likelihood_reconstruction(self, x, prior, noise):
         """
         if noise is None:
             noise = 1
-        prior_cov = np.diag(prior**2)
-        sensor_selection_matrix = np.zeros(
-            (len(self.selected_sensors), self.basis_matrix_.shape[0])
-        )
-        sensor_selection_matrix[
-            np.arange(len(self.selected_sensors)), self.selected_sensors
-        ] = 1
-        low_rank_selection_matrix = sensor_selection_matrix @ self.basis_matrix_
-        composite_matrix = np.linalg.inv(prior_cov) + (
+        if not isinstance(prior, np.ndarray):
+            raise ValueError("prior must be a numpy array")
+        prior_cov = 1 / (prior**2)
+        low_rank_selection_matrix = self.basis_matrix_[self.selected_sensors, :]
+        composite_matrix = np.diag(prior_cov) + (
             low_rank_selection_matrix.T @ low_rank_selection_matrix
         ) / (noise**2)
         rhs = low_rank_selection_matrix.T @ x
@@ -583,15 +606,20 @@ def std(self, prior, noise=None):
             noise = 1
         if noise <= 0:
             raise ValueError("Noise must be positive")
-        prior_sq = prior**2
-        sensor_selection_matrix = np.zeros(
-            (len(self.selected_sensors), self.basis_matrix_.shape[0])
-        )
-        sensor_selection_matrix[
-            np.arange(len(self.selected_sensors)), self.selected_sensors
-        ] = 1
-        low_rank_selection_matrix = sensor_selection_matrix @ self.basis_matrix_
-        composite_matrix = np.diag(prior_sq) + (
+        if isinstance(prior, str) and prior == "decreasing":
+            computed_prior = self.singular_values
+        elif isinstance(prior, np.ndarray):
+            if prior.ndim != 1:
+                raise ValueError("prior must be a 1D array")
+            if prior.shape[0] != self.basis_matrix_.shape[1]:
+                raise ValueError(
+                    f"prior must be of shape {(self.basis_matrix_.shape[1],)},"
+                    f" but got {prior.shape}"
+                )
+            computed_prior = prior
+        sq_inv_prior = 1.0 / (computed_prior**2)
+        low_rank_selection_matrix = self.basis_matrix_[self.selected_sensors, :]
+        composite_matrix = np.diag(sq_inv_prior) + (
             low_rank_selection_matrix.T @ low_rank_selection_matrix
         ) / (noise**2)
         diag_cov_matrix = (
@@ -603,14 +631,48 @@ def std(self, prior, noise=None):
         sigma = noise * np.sqrt(np.sum(diag_cov_matrix**2, axis=1))
         return sigma
 
-    def one_pt_energy_landscape(self, prior, noise):
+    def one_pt_energy_landscape(self, prior="decreasing", noise=None):
         check_is_fitted(self, "optimizer")
-        G = self.basis_matrix_ @ np.diag(prior)
+        if isinstance(prior, str) and prior == "decreasing":
+            computed_prior = self.singular_values
+        elif isinstance(prior, np.ndarray):
+            if prior.ndim != 1:
+                raise ValueError("prior must be a 1D array")
+            if prior.shape[0] != self.basis_matrix_.shape[1]:
+                raise ValueError(
+                    f"prior must be of shape {(self.basis_matrix_.shape[1],)},"
+                    f" but got {prior.shape}"
+                )
+            computed_prior = prior
+        if noise is None:
+            warnings.warn(
+                "noise is None. noise will be set to the "
+                "average of the normalized prior"
+            )
+            noise = computed_prior.mean()
+        G = self.basis_matrix_ @ np.diag(computed_prior)
         return -np.log(1 + np.einsum("ij,ij->i", G, G) / noise**2)
 
-    def two_pt_energy_landscape(self, prior, noise, selected_sensors):
+    def two_pt_energy_landscape(self, selected_sensors, prior="decreasing", noise=None):
         check_is_fitted(self, "optimizer")
-        G = self.basis_matrix_ @ np.diag(prior)
+        if isinstance(prior, str) and prior == "decreasing":
+            computed_prior = self.singular_values
+        elif isinstance(prior, np.ndarray):
+            if prior.ndim != 1:
+                raise ValueError("prior must be a 1D array")
+            if prior.shape[0] != self.basis_matrix_.shape[1]:
+                raise ValueError(
+                    f"prior must be of shape {(self.basis_matrix_.shape[1],)},"
+                    f" but got {prior.shape}"
+                )
+            computed_prior = prior
+        if noise is None:
+            warnings.warn(
+                "noise is None. noise will be set to the "
+                "average of the normalized prior"
+            )
+            noise = computed_prior.mean()
+        G = self.basis_matrix_ @ np.diag(computed_prior)
         mask = np.ones(G.shape[0], dtype=bool)
         mask[selected_sensors] = False
         G_selected = G[selected_sensors, :]

tests/reconstruction/test_sspor.py

@@ -147,11 +147,24 @@ def test_predict_accuracy(data_vandermonde_testing):
     model.fit(data, seed=1)
     model.set_number_of_sensors(8)
     sensors = model.get_selected_sensors()
-    assert sqrt(mean((x_test - model.predict(x_test[sensors])) ** 2)) <= 1.0e-3
+    assert (
+        sqrt(
+            mean((x_test - model.predict(x_test[sensors], method="unregularized")) ** 2)
+        )
+        <= 1.0e-3
+    )
 
     # Should also work for row vectors
     x_test = x_test.reshape(1, -1)
-    assert sqrt(mean((x_test - model.predict(x_test[:, sensors])) ** 2)) <= 1.0e-3
+    assert (
+        sqrt(
+            mean(
+                (x_test - model.predict(x_test[:, sensors], method="unregularized"))
+                ** 2
+            )
+        )
+        <= 1.0e-3
+    )
 
 
 def test_reconstruction_error(data_vandermonde_testing):
@@ -224,7 +237,7 @@ def test_update_n_basis_modes(data_random):
     assert model.basis_matrix_.shape[1] == data.shape[0]
 
     n_basis_modes = 5
-    model.update_n_basis_modes(n_basis_modes)
+    model.update_n_basis_modes(n_basis_modes, x=data)
     assert model.basis.n_basis_modes == data.shape[0]
     assert model.basis_matrix_.shape[1] == n_basis_modes
 
@@ -350,7 +363,7 @@ def test_predict_warns_when_n_sensors_exceeds_basis_dimension():
             UserWarning, match="n_sensors exceeds dimension of basis modes"
         ):
             predictor.predict.__globals__["validate_input"] = lambda x, y: X.T
-            predictor.predict(X)
+            predictor.predict(X, method="unregularized")
     finally:
         if "validate_input" in predictor.predict.__globals__:
             predictor.predict.__globals__["validate_input"] = original_validate_input
@@ -376,7 +389,7 @@ def test_predict_no_warning_when_n_sensors_not_exceeds_basis_dimension():
         X = np.random.rand(2, 3)
         with warnings.catch_warnings(record=True) as recorded_warnings:
             predictor.predict.__globals__["validate_input"] = lambda x, y: X.T
-            predictor.predict(X)
+            predictor.predict(X, method="unregularized")
             relevant_warnings = [
                 w
                 for w in recorded_warnings
@@ -733,7 +746,7 @@ def test_maximal_likelihood_reconstruction():
     selected = model.get_selected_sensors()
     x_sensors = X[:, selected]
 
-    y_pred = model.predict(x_sensors, method="mle", prior=prior, noise=0.1)
+    y_pred = model.predict(x_sensors, method=None, prior=prior, noise=0.1)
 
     assert isinstance(y_pred, np.ndarray)
     assert y_pred.shape == (n_samples, n_features)