Implement eigh() fallback

mesmer/core/utils.py

@@ -23,7 +23,7 @@ def create_equal_dim_names(dim, suffixes):
     return tuple(f"{dim}{suffix}" for suffix in suffixes)
 
 
-def _minimize_local_discrete(func, sequence, **kwargs):
+def _minimize_local_discrete(func, sequence, changable, **kwargs):
     """find the local minimum for a function that consumes discrete input
 
     Parameters
@@ -33,6 +33,9 @@ def _minimize_local_discrete(func, sequence, **kwargs):
         as input and return a float that is to be minimized.
     sequence : iterable
         An iterable with discrete values to evaluate func for.
+    changable : str
+        changable parameter for the function, enables to feed back function
+        output at input
     **kwargs : Mapping
         Keyword arguments passed to `func`.
 
@@ -57,7 +60,7 @@ def _minimize_local_discrete(func, sequence, **kwargs):
 
     for i, element in enumerate(sequence):
 
-        res = func(element, **kwargs)
+        res, changable = func(element, changable, **kwargs)
 
         if np.isneginf(res):
             raise ValueError("`fun` returned `-inf`")
@@ -72,6 +75,9 @@ def _minimize_local_discrete(func, sequence, **kwargs):
             sel = i - 1
             if sel == 0:
                 warnings.warn("First element is local minimum.", OptimizeWarning)
+            elif sel == -1:
+                raise ValueError("First radius already inf.")
+
             return sequence[sel]
 
     warnings.warn("No local minimum found, returning the last element", OptimizeWarning)

mesmer/stats/_auto_regression.py

@@ -475,11 +475,8 @@ def _draw_auto_regression_correlated_np(
     # arbitrary lags? no, see: https://github.com/MESMER-group/mesmer/issues/164
     ar_lags = np.arange(1, ar_order + 1, dtype=int)
 
-    # ensure reproducibility (TODO: https://github.com/MESMER-group/mesmer/issues/35)
-    np.random.seed(seed)
-
     innovations = _draw_innovations_correlated_np(
-        covariance, n_coeffs, n_samples, n_ts, buffer
+        seed, covariance, n_coeffs, n_samples, n_ts, buffer
     )
 
     out = np.zeros([n_samples, n_ts + buffer, n_coeffs])
@@ -492,25 +489,32 @@ def _draw_auto_regression_correlated_np(
     return out[:, buffer:, :]
 
 
-def _draw_innovations_correlated_np(covariance, n_gridcells, n_samples, n_ts, buffer):
+def _draw_innovations_correlated_np(
+    seed, covariance, n_gridcells, n_samples, n_ts, buffer
+):
     # NOTE: 'innovations' is the error or noise term.
     # innovations has shape (n_samples, n_ts + buffer, n_coeffs)
+    # ensure reproducibility (TODO: https://github.com/MESMER-group/mesmer/issues/35)
+    rng = np.random.default_rng(seed=seed)
     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,
-            )
-
-    innovations = scipy.stats.multivariate_normal.rvs(
-        mean=np.zeros(n_gridcells),
-        cov=cov,
-        size=[n_samples, n_ts + buffer],
-    ).reshape(n_samples, n_ts + buffer, n_gridcells)
+        innovations = rng.multivariate_normal(
+            mean=np.zeros(n_gridcells),
+            cov=covariance,
+            size=[n_samples, n_ts + buffer],
+            method="cholesky",
+        ).reshape(n_samples, n_ts + buffer, n_gridcells)
+
+    except np.linalg.LinAlgError:
+        warnings.warn(
+            "Covariance matrix is not positive definite, using eigh instead of cholesky.",
+            LinAlgWarning,
+        )
+        innovations = rng.multivariate_normal(
+            mean=np.zeros(n_gridcells),
+            cov=covariance,
+            size=[n_samples, n_ts + buffer],
+            method="eigh",
+        ).reshape(n_samples, n_ts + buffer, n_gridcells)
 
     return innovations
 
@@ -879,16 +883,13 @@ def _draw_auto_regression_monthly_np(
 
     _, n_gridcells = intercept.shape
 
-    # ensure reproducibility (TODO: https://github.com/MESMER-group/mesmer/issues/35)
-    np.random.seed(seed)
-
     # draw innovations for each month
     innovations = np.zeros([n_samples, n_ts // 12 + buffer, 12, n_gridcells])
     if covariance is not None:
         for month in range(12):
             cov_month = covariance[month, :, :]
             innovations[:, :, month, :] = _draw_innovations_correlated_np(
-                cov_month, n_gridcells, n_samples, n_ts // 12, buffer
+                seed, cov_month, n_gridcells, n_samples, n_ts // 12, buffer
             )
     # reshape innovations into continuous time series
     innovations = innovations.reshape(n_samples, n_ts + buffer * 12, n_gridcells)

mesmer/stats/_localized_covariance.py

@@ -250,9 +250,13 @@ 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
+    method = "cholesky"
     localization_radius = _minimize_local_discrete(
         _ecov_crossvalidation,
         localization_radii,
+        method,
         data=data,
         weights=weights,
         localizer=localizer,
@@ -265,7 +269,9 @@ 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):
+def _ecov_crossvalidation(
+    localization_radius, method, *, data, weights, localizer, k_folds
+):
     """k-fold crossvalidation for a single localization radius"""
 
     n_samples, __ = data.shape
@@ -291,19 +297,22 @@ def _ecov_crossvalidation(localization_radius, *, data, weights, localizer, k_fo
 
         try:
             # sum log likelihood of all crossvalidation folds
-            nll += _get_neg_loglikelihood(data_cv, localized_cov, weights_cv)
+            nll += _get_neg_loglikelihood(data_cv, localized_cov, weights_cv, method)
         except np.linalg.LinAlgError:
+            # NOTE: this error is thrown by np.linalg.cholesky not by the logpdf anymore
             warnings.warn(
                 f"Singular matrix for localization_radius of {localization_radius}."
-                " Skipping this radius.",
+                "\n Switching to eigh().",
                 LinAlgWarning,
             )
-            return float("inf")
+            # switch to eigh from now on
+            method = "eigh"
+            nll += _get_neg_loglikelihood(data_cv, localized_cov, weights_cv, method)
 
-    return nll
+    return nll, method
 
 
-def _get_neg_loglikelihood(data, covariance, weights):
+def _get_neg_loglikelihood(data, covariance, weights, method):
     """calculate weighted log likelihood for multivariate normal distribution
 
     Parameters
@@ -330,11 +339,15 @@ 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":
+        cov = scipy.stats.Covariance.from_cholesky(np.linalg.cholesky(covariance))
+    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

@@ -432,7 +432,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,
@@ -444,7 +444,7 @@ def test_draw_auto_regression_random():
         buffer=3,
     )
 
-    expected = np.array([2.58455078, 3.28976946, 1.86569258, 2.78266986])
+    expected = np.array([1.07417558, 1.0240404, 1.77397341, 2.71531235])
     expected = expected.reshape(1, 4, 1)
 
     np.testing.assert_allclose(result, expected)