Implement eigh() fallback

CHANGELOG.rst

@@ -9,6 +9,8 @@ New Features
 - Integrated MESMER-M into the code base, see `Integration of MESMER-M`_.
 - Added number of observations to the output of the AR process (`#395 <https://github.com/MESMER-group/mesmer/pull/395>`_).
   By `Victoria Bauer`_.
+- Implemented option to allow for singular covariance matrices in the crossvalidation of localization radii (`#493 <https://github.com/MESMER-group/mesmer/pull/493>`_).
+  By `Victoria Bauer`_.
 
 Breaking changes
 ^^^^^^^^^^^^^^^^
@@ -72,6 +74,9 @@ Internal Changes
   By `Victoria Bauer`_.
 - Use ruff instead of isort and flake8 to lint the code base (`#490 <https://github.com/MESMER-group/mesmer/pull/490>`_).
   By `Mathias Hauser`_.
+- Allow singular covariance matrices for localization radius selection. For this purpose, eigenvalue decomposition is implemented
+  as fallback for singular matrices also in the crossvalidation of localization radii (`#493 <https://github.com/MESMER-group/mesmer/pull/493>`_).
+  By `Victoria Bauer`_.
 
 Integration of MESMER-X
 ^^^^^^^^^^^^^^^^^^^^^^^

mesmer/core/utils.py

@@ -53,29 +53,36 @@ def _minimize_local_discrete(func, sequence, **kwargs):
     current_min = float("inf")
     # ensure it's a list because we cannot get an item from an iterable
     sequence = list(sequence)
+    last_valid_element = None
 
     for i, element in enumerate(sequence):
-
         res = func(element, **kwargs)
 
         if np.isneginf(res):
-            raise ValueError("`fun` returned `-inf`")
-        # skip element if inf is returned - not sure about this?
+            raise ValueError(f"`func` returned `-inf` for element {element}")
         elif np.isinf(res):
-            warnings.warn("`fun` returned `inf`", OptimizeWarning)
+            warnings.warn(
+                f"`func` returned `inf` for element {element}, skipping.",
+                OptimizeWarning,
+            )
+            continue  # Skip this element and move to the next
 
         if res < current_min:
             current_min = res
+            last_valid_element = element
         else:
-            # need to return element from the previous iteration
-            sel = i - 1
-            if sel == 0:
-                warnings.warn("First element is local minimum.", OptimizeWarning)
-            return sequence[sel]
-
-    warnings.warn("No local minimum found, returning the last element", OptimizeWarning)
-
-    return element
+            if last_valid_element is not None:
+                if last_valid_element == sequence[0]:
+                    warnings.warn("First element is local minimum.", OptimizeWarning)
+                return last_valid_element
+
+    if last_valid_element is None:
+        raise ValueError("No valid values were found in the sequence.")
+
+    warnings.warn(
+        "No local minimum found, returning the last valid element.", OptimizeWarning
+    )
+    return last_valid_element
 
 
 def _to_set(arg):

mesmer/stats/_auto_regression.py

@@ -500,16 +500,17 @@ def _draw_innovations_correlated_np(
 ):
     # NOTE: 'innovations' is the error or noise term.
     # innovations has shape (n_samples, n_ts + buffer, n_coeffs)
+
     try:
         cov = scipy.stats.Covariance.from_cholesky(np.linalg.cholesky(covariance))
-    except np.linalg.LinAlgError as e:
-        if "Matrix is not positive definite" in str(e):
-            w, v = np.linalg.eigh(covariance)
-            cov = scipy.stats.Covariance.from_eigendecomposition((w, v))
-            warnings.warn(
-                "Covariance matrix is not positive definite, using eigh instead of cholesky.",
-                LinAlgWarning,
-            )
+
+    except np.linalg.LinAlgError:
+        w, v = np.linalg.eigh(covariance)
+        cov = scipy.stats.Covariance.from_eigendecomposition((w, v))
+        warnings.warn(
+            "Covariance matrix is not positive definite, using eigh instead of cholesky.",
+            LinAlgWarning,
+        )
 
     innovations = scipy.stats.multivariate_normal.rvs(
         mean=np.zeros(n_gridcells),

mesmer/stats/_localized_covariance.py

@@ -77,7 +77,13 @@ def _adjust_ecov_ar1_np(covariance, ar_coefs):
 
 
 def find_localized_empirical_covariance(
-    data, weights, localizer, dim, k_folds, equal_dim_suffixes=("_i", "_j")
+    data,
+    weights,
+    localizer,
+    dim,
+    k_folds,
+    equal_dim_suffixes=("_i", "_j"),
+    allow_singluar=False,
 ):
     """determine localized empirical covariance by cross validation
 
@@ -98,6 +104,9 @@ def find_localized_empirical_covariance(
     equal_dim_suffixes : tuple of str, default: ("_i", "_j")
         Suffixes to add to the the name of ``dim`` for the covariance array
         (xr.DataArray cannot have two dimensions with the same name).
+    allow_singluar : bool, default: False
+        If True, allow singular matrices to be used in the cross validation. In this case,
+        the method of decomposition is switched from cholesky to eigh and a warning is emitted.
 
     Returns
     -------
@@ -129,7 +138,11 @@ def find_localized_empirical_covariance(
         _find_localized_empirical_covariance_np,
         data,
         weights,
-        kwargs={"localizer": localizer, "k_folds": k_folds},
+        kwargs={
+            "localizer": localizer,
+            "k_folds": k_folds,
+            "allow_singluar": allow_singluar,
+        },
         input_core_dims=[all_dims, [dim]],
         output_core_dims=([], out_dims, out_dims),
     )
@@ -145,7 +158,13 @@ def find_localized_empirical_covariance(
 
 
 def find_localized_empirical_covariance_monthly(
-    data, weights, localizer, dim, k_folds, equal_dim_suffixes=("_i", "_j")
+    data,
+    weights,
+    localizer,
+    dim,
+    k_folds,
+    equal_dim_suffixes=("_i", "_j"),
+    allow_singluar=False,
 ):
     """determine localized empirical covariance by cross validation for each month. `data`
     should be the residuals of the cyclo-stationary AR(1) process, see
@@ -170,6 +189,9 @@ def find_localized_empirical_covariance_monthly(
     equal_dim_suffixes : tuple of str, default: ("_i", "_j")
         Suffixes to add to the the name of ``dim`` for the covariance array
         (xr.DataArray cannot have two dimensions with the same name).
+    allow_singluar : bool, default: False
+        If True, allow singular matrices to be used in the cross validation. In this case,
+        the method of decomposition is switched from cholesky to eigh and a warning is emitted.
 
     Returns
     -------
@@ -199,14 +221,17 @@ def find_localized_empirical_covariance_monthly(
             dim=dim,
             k_folds=k_folds,
             equal_dim_suffixes=equal_dim_suffixes,
+            allow_singluar=allow_singluar,
         )
         localized_ecov.append(res)
 
     month = xr.Variable("month", range(1, 13))
     return xr.concat(localized_ecov, dim=month)
 
 
-def _find_localized_empirical_covariance_np(data, weights, localizer, k_folds):
+def _find_localized_empirical_covariance_np(
+    data, weights, localizer, k_folds, allow_singluar
+):
     """determine localized empirical covariance by cross validation
 
     Parameters
@@ -221,6 +246,9 @@ def _find_localized_empirical_covariance_np(data, weights, localizer, k_folds):
         Currently only the Gaspari-Cohn localizer is implemented in MESMER.
     k_folds : int
         Number of folds to use for cross validation.
+    allow_singluar : bool
+        If True, allow singular matrices to be used in the cross validation. In this case,
+        the method of decomposition is switched from cholesky to eigh and a warning is emitted.
 
     Returns
     -------
@@ -250,13 +278,16 @@ def _find_localized_empirical_covariance_np(data, weights, localizer, k_folds):
     # start again. Better to stop once min is reached (to limit computational effort
     # and singular matrices).
 
+    # start with cholesky decomposition
+    # function returns method and can switch to eigh() if cov is singular
     localization_radius = _minimize_local_discrete(
-        _ecov_crossvalidation,
+        _EcovCrossvalidation(method="cholesky").crossvalidate,
         localization_radii,
         data=data,
         weights=weights,
         localizer=localizer,
         k_folds=k_folds,
+        allow_singluar=allow_singluar,
     )
 
     covariance = np.cov(data, rowvar=False, aweights=weights)
@@ -265,45 +296,71 @@ def _find_localized_empirical_covariance_np(data, weights, localizer, k_folds):
     return localization_radius, covariance, localized_covariance
 
 
-def _ecov_crossvalidation(localization_radius, *, data, weights, localizer, k_folds):
-    """k-fold crossvalidation for a single localization radius"""
+class _EcovCrossvalidation:
+    # we organize this as a class so we can store `method` as state
+    # ensures we use `eigh` after `cholesky` fails for one localization
+    # radius
+    def __init__(self, method=None):
+
+        self.method = method or "cholesky"
+
+    def crossvalidate(
+        self, localization_radius, *, data, weights, localizer, k_folds, allow_singluar
+    ):
+        """k-fold crossvalidation for a single localization radius"""
+
+        n_samples, __ = data.shape
+        n_iterations = min(n_samples, k_folds)
 
-    n_samples, __ = data.shape
-    n_iterations = min(n_samples, k_folds)
+        nll = 0  # negative log likelihood
 
-    nll = 0  # negative log likelihood
+        for it in range(n_iterations):
 
-    for it in range(n_iterations):
+            # every `k_folds` element for validation such that each is used exactly once
+            sel = np.ones(n_samples, dtype=bool)
+            sel[it::k_folds] = False
 
-        # every `k_folds` element for validation such that each is used exactly once
-        sel = np.ones(n_samples, dtype=bool)
-        sel[it::k_folds] = False
+            # extract training set
+            data_train, weights_train = data[sel, :], weights[sel]
 
-        # extract training set
-        data_train, weights_train = data[sel, :], weights[sel]
+            # extract validation set
+            data_cv, weights_cv = data[~sel, :], weights[~sel]
 
-        # extract validation set
-        data_cv, weights_cv = data[~sel, :], weights[~sel]
+            # compute (localized) empirical covariance
+            cov = np.cov(data_train, rowvar=False, aweights=weights_train)
+            localized_cov = localizer[localization_radius] * cov
 
-        # compute (localized) empirical covariance
-        cov = np.cov(data_train, rowvar=False, aweights=weights_train)
-        localized_cov = localizer[localization_radius] * cov
+            try:
+                # sum log likelihood of all crossvalidation folds
+                nll += _get_neg_loglikelihood(
+                    data_cv, localized_cov, weights_cv, self.method
+                )
 
-        try:
-            # sum log likelihood of all crossvalidation folds
-            nll += _get_neg_loglikelihood(data_cv, localized_cov, weights_cv)
-        except np.linalg.LinAlgError:
-            warnings.warn(
-                f"Singular matrix for localization_radius of {localization_radius}."
-                " Skipping this radius.",
-                LinAlgWarning,
-            )
-            return float("inf")
+            except np.linalg.LinAlgError as e:
+                if not allow_singluar:
+                    raise np.linalg.LinAlgError(
+                        f"Singular matrix for localization_radius of {localization_radius}."
+                    )
 
-    return nll
+                if self.method == "eigh":
+                    raise e
 
+                # switch to eigh from now on
+                self.method = "eigh"
+                warnings.warn(
+                    f"Singular matrix for localization_radius of {localization_radius}."
+                    "\n Switching to eigh().",
+                    LinAlgWarning,
+                )
 
-def _get_neg_loglikelihood(data, covariance, weights):
+                nll += _get_neg_loglikelihood(
+                    data_cv, localized_cov, weights_cv, self.method
+                )
+
+        return nll
+
+
+def _get_neg_loglikelihood(data, covariance, weights, method):
     """calculate weighted log likelihood for multivariate normal distribution
 
     Parameters
@@ -330,11 +387,16 @@ def _get_neg_loglikelihood(data, covariance, weights):
     The mean is assumed to be zero for all points.
     """
 
-    # NOTE: 90 % of time is spent in multivariate_normal.logpdf - not much point
-    # optimizing the rest
+    if method == "cholesky":
+        L = np.linalg.cholesky(covariance)
+        cov = scipy.stats.Covariance.from_cholesky(L)
+    else:
+        w, v = np.linalg.eigh(covariance)
+        cov = scipy.stats.Covariance.from_eigendecomposition((w, v))
 
-    cov = scipy.stats.Covariance.from_cholesky(np.linalg.cholesky(covariance))
-    log_likelihood = scipy.stats.multivariate_normal.logpdf(data, cov=cov)
+    log_likelihood = scipy.stats.multivariate_normal.logpdf(
+        data, cov=cov, allow_singular=True
+    )
 
     # logpdf can return a scalar, which np.average does not like
     log_likelihood = np.atleast_1d(log_likelihood)

tests/unit/test_auto_regression.py

@@ -433,7 +433,7 @@ def test_draw_auto_regression_deterministic_coefs_buffer():
         np.testing.assert_allclose(result, expected[:, i : i + 4])
 
 
-def test_draw_auto_regression_random():
+def test_draw_auto_regression_numeric():
 
     result = mesmer.stats._auto_regression._draw_auto_regression_correlated_np(
         intercept=1,
@@ -674,6 +674,7 @@ def test_fit_auto_regression_monthly():
         mesmer.stats.fit_auto_regression_monthly(data.values)
 
 
+@pytest.mark.filterwarnings("ignore:Covariance matrix is not positive definite")
 @pytest.mark.parametrize("buffer", [1, 10, 20])
 def test_draw_auto_regression_monthly_np_buffer(buffer):
     n_realisations = 1

tests/unit/test_harmonic_model.py

@@ -65,6 +65,7 @@ def test_predict_harmonic_model():
     )
 
 
+@pytest.mark.filterwarnings("ignore:divide by zero encountered in log")
 @pytest.mark.parametrize(
     "coefficients",
     [