.

Author: comane

colibri/analytic_fit.py

@@ -14,6 +14,8 @@
 import jax.lax.linalg as jlinalg
 import numpy as np
 import scipy.special as special
+from sklearn.linear_model import Ridge, RidgeCV, ElasticNet, ElasticNetCV
+from sklearn.preprocessing import StandardScaler
 
 from colibri.core import AnalyticFit
 from colibri.export_results import write_replicas, export_bayes_results
@@ -151,28 +153,173 @@ def analytic_fit(
 
     t0 = time.time()
 
+    ridge_alpha = analytic_settings.get("ridge_alpha", 0.0)
+    ridge_cv_alphas = analytic_settings.get("ridge_cv_alphas", None)
+    ridge_cv_folds = analytic_settings.get("ridge_cv_folds", 5)
+
+    elasticnet_alpha = analytic_settings.get("elasticnet_alpha", 0.0)
+    elasticnet_l1_ratio = analytic_settings.get("elasticnet_l1_ratio", 0.5)
+    elasticnet_cv_alphas = analytic_settings.get("elasticnet_cv_alphas", None)
+    elasticnet_cv_l1_ratios = analytic_settings.get("elasticnet_cv_l1_ratios", None)
+    elasticnet_cv_folds = analytic_settings.get("elasticnet_cv_folds", 5)
+
     # Cholesky factorization: S = L L^T
     # upper False means that we want the lower triangular matrix L
     L = jla.cholesky(covmat, upper=False)
 
     # Whiten the problem: Y' = L^-1 Y, X' = L^-1 X
+    # (X^T Sigma^-1 X) = (X'^T X'), (X^T Sigma^-1 Y) = (X'^T Y')
     Y_tilde = jlinalg.triangular_solve(L, Y, left_side=True, lower=True)
     X_tilde = jlinalg.triangular_solve(L, X, left_side=True, lower=True)
 
-    if jnp.any(jla.eigh(X_tilde.T @ X_tilde)[0] <= 0.0):
-        raise ValueError(
-            "The obtained covariance matrix for the analytic solution is not positive definite."
+    if ridge_cv_alphas is not None or ridge_alpha > 0.0:
+        # Ridge regression with sklearn StandardScaler for scale-invariant regularisation.
+        #
+        # StandardScaler (with_mean=False) divides each column of X_tilde by its
+        # standard deviation s_j, so the Ridge penalty ||v||^2 on the scaled
+        # coefficients v treats every parameter equally regardless of feature scale.
+        # After fitting, we back-transform to the original parameter space:
+        #   sol_mean   = v_ridge / s
+        #   sol_covmat = D^{-1} (X_scaled^T X_scaled + alpha I)^{-1} D^{-1},  D = diag(s)
+        # Equivalent to a Bayesian MAP with Gaussian prior N(0, (1/alpha) I) on v.
+        log.warning(
+            "With Ridge regularisation the Bayesian evidence metrics assume a "
+            "Gaussian prior N(0, (1/alpha)*I) on the scaled parameters rather than a uniform prior."
         )
-
-    # Compute QR decomposition of X_tilde for numerical stability in the inversion
-    Q, R = jla.qr(X_tilde)
-
-    # NOTE: R is upper triangular in QR decomposition, so we need to set lower=False
-    sol_mean = jlinalg.triangular_solve(R, Q.T @ Y_tilde, left_side=True, lower=False)
-
-    I_R = jnp.eye(R.shape[0])
-    R_inv = jlinalg.triangular_solve(R, I_R, left_side=True, lower=False)
-    sol_covmat = R_inv @ R_inv.T
+        n_params = len(parameters)
+        X_tilde_np = np.array(X_tilde)
+        Y_tilde_np = np.array(Y_tilde)
+
+        scaler = StandardScaler(with_mean=False)
+        X_scaled = scaler.fit_transform(X_tilde_np)
+        col_scales = scaler.scale_  # per-feature standard deviations, shape (n_params,)
+
+        if ridge_cv_alphas is not None:
+            # Select alpha via k-fold cross-validation on the whitened, scaled problem.
+            log.info(
+                f"Selecting Ridge alpha via {ridge_cv_folds}-fold CV "
+                f"over candidates {ridge_cv_alphas}."
+            )
+            ridge_cv = RidgeCV(
+                alphas=ridge_cv_alphas,
+                cv=ridge_cv_folds,
+                fit_intercept=False,
+            )
+            ridge_cv.fit(X_scaled, Y_tilde_np)
+            best_alpha = float(ridge_cv.alpha_)
+            v_ridge = ridge_cv.coef_
+            log.info(f"RidgeCV selected alpha={best_alpha}.")
+        else:
+            best_alpha = ridge_alpha
+            log.info(f"Using Ridge regression with alpha={best_alpha}.")
+            ridge = Ridge(alpha=best_alpha, fit_intercept=False)
+            ridge.fit(X_scaled, Y_tilde_np)
+            v_ridge = ridge.coef_
+
+        # Posterior covariance in the scaled space: (X_scaled^T X_scaled + alpha I)^{-1}
+        A_scaled = X_scaled.T @ X_scaled + best_alpha * np.eye(n_params)
+        v_covmat = np.linalg.inv(A_scaled)
+
+        # Back-transform to original parameter space
+        sol_mean = jnp.array(v_ridge / col_scales)
+        sol_covmat = jnp.array(v_covmat / col_scales[:, None] / col_scales[None, :])
+    elif elasticnet_alpha > 0.0 or elasticnet_cv_alphas is not None:
+        # Elastic Net regularisation with sklearn StandardScaler for scale-invariant
+        # regularisation. As with Ridge we work in the whitened+scaled space so that
+        # the L1 and L2 penalties treat every feature on equal footing, and we
+        # back-transform afterwards.
+        #
+        # Elastic Net minimises (in the scaled space)
+        #   (1/(2n)) ||Y_tilde - X_scaled v||^2
+        #     + alpha * l1_ratio * ||v||_1
+        #     + 0.5 * alpha * (1 - l1_ratio) * ||v||_2^2
+        # which has no closed-form posterior covariance, so we use the Laplace
+        # approximation around the MAP solution v* (see below).
+        log.warning(
+            "With Elastic Net regularisation the Bayesian evidence metrics assume a "
+            "Laplace approximation around the MAP solution; this can be inaccurate "
+            "for very sparse solutions (many near-zero coefficients)."
+        )
+        n_params = len(parameters)
+        X_tilde_np = np.array(X_tilde)
+        Y_tilde_np = np.array(Y_tilde)
+
+        scaler = StandardScaler(with_mean=False)
+        X_scaled = scaler.fit_transform(X_tilde_np)
+        col_scales = scaler.scale_  # per-feature standard deviations, shape (n_params,)
+
+        if elasticnet_cv_alphas is not None:
+            # Select alpha (and possibly l1_ratio) via k-fold CV on the whitened,
+            # scaled problem.
+            log.info(
+                f"Selecting Elastic Net alpha via {elasticnet_cv_folds}-fold CV "
+                f"over alpha candidates {elasticnet_cv_alphas} "
+                f"and l1_ratio candidates {elasticnet_cv_l1_ratios}."
+            )
+            # ElasticNetCV defaults l1_ratio to 0.5 if None is passed; sklearn requires
+            # a list/scalar so fall back to the fixed-alpha default.
+            cv_l1_ratios = (
+                elasticnet_cv_l1_ratios
+                if elasticnet_cv_l1_ratios is not None
+                else elasticnet_l1_ratio
+            )
+            enet_cv = ElasticNetCV(
+                alphas=elasticnet_cv_alphas,
+                l1_ratio=cv_l1_ratios,
+                cv=elasticnet_cv_folds,
+                fit_intercept=False,
+            )
+            enet_cv.fit(X_scaled, Y_tilde_np)
+            best_alpha = float(enet_cv.alpha_)
+            best_l1_ratio = float(enet_cv.l1_ratio_)
+            v_enet = enet_cv.coef_
+            log.info(
+                f"ElasticNetCV selected alpha={best_alpha}, l1_ratio={best_l1_ratio}."
+            )
+        else:
+            best_alpha = elasticnet_alpha
+            best_l1_ratio = elasticnet_l1_ratio
+            log.info(
+                f"Using Elastic Net regression with alpha={best_alpha}, "
+                f"l1_ratio={best_l1_ratio}."
+            )
+            enet = ElasticNet(
+                alpha=best_alpha,
+                l1_ratio=best_l1_ratio,
+                fit_intercept=False,
+            )
+            enet.fit(X_scaled, Y_tilde_np)
+            v_enet = enet.coef_
+
+        # Posterior covariance via the Laplace approximation around v*.
+        # The Hessian of the Elastic Net objective in the scaled space at v* is
+        #   H = X_scaled^T X_scaled + diag(lambda2 + lambda1 / |v*_i|)
+        # with lambda1 = alpha * l1_ratio and lambda2 = alpha * (1 - l1_ratio).
+        # Near-zero coefficients are clamped (|v*_i| < 1e-10 -> 1e-10) to avoid
+        # division by zero; the penalty there becomes very large, which correctly
+        # reflects that the L1 term sharply pins these coefficients to zero.
+        lambda1 = best_alpha * best_l1_ratio
+        lambda2 = best_alpha * (1.0 - best_l1_ratio)
+        abs_v = np.maximum(np.abs(v_enet), 1e-10)
+        diag_penalty = lambda2 + lambda1 / abs_v
+        H = X_scaled.T @ X_scaled + np.diag(diag_penalty)
+        v_covmat = np.linalg.inv(H)
+
+        # Back-transform to original parameter space (same formula as Ridge)
+        sol_mean = jnp.array(v_enet / col_scales)
+        sol_covmat = jnp.array(v_covmat / col_scales[:, None] / col_scales[None, :])
+    else:
+        if jnp.any(jla.eigh(X_tilde.T @ X_tilde)[0] <= 0.0):
+            raise ValueError(
+                "The obtained covariance matrix for the analytic solution is not positive definite."
+            )
+        # OLS: use QR decomposition of X_tilde for numerical stability
+        Q, R = jla.qr(X_tilde)
+        # NOTE: R is upper triangular in QR decomposition, so lower=False
+        sol_mean = jlinalg.triangular_solve(R, Q.T @ Y_tilde, left_side=True, lower=False)
+        I_R = jnp.eye(R.shape[0])
+        R_inv = jlinalg.triangular_solve(R, I_R, left_side=True, lower=False)
+        sol_covmat = R_inv @ R_inv.T
 
     key = jax.random.PRNGKey(analytic_settings["sampling_seed"])
 

colibri/config.py

@@ -479,6 +479,14 @@ def parse_analytic_settings(
             "n_posterior_samples",
             "sampling_seed",
             "full_sample_size",
+            "ridge_alpha",
+            "ridge_cv_alphas",
+            "ridge_cv_folds",
+            "elasticnet_alpha",
+            "elasticnet_l1_ratio",
+            "elasticnet_cv_alphas",
+            "elasticnet_cv_l1_ratios",
+            "elasticnet_cv_folds",
         }
 
         kdiff = settings.keys() - known_keys
@@ -502,6 +510,132 @@ def parse_analytic_settings(
         # Set the full sample size
         analytic_settings["full_sample_size"] = settings.get("full_sample_size", 1000)
 
+        # Ridge regularisation: fixed alpha or CV-selected alpha are mutually exclusive.
+        ridge_alpha = settings.get("ridge_alpha", 0.0)
+        ridge_cv_alphas = settings.get("ridge_cv_alphas", None)
+
+        if ridge_alpha != 0.0 and ridge_cv_alphas is not None:
+            raise ConfigError(
+                "Specify either 'ridge_alpha' (fixed) or 'ridge_cv_alphas' (CV), not both.",
+                "ridge_alpha",
+                known_keys,
+            )
+
+        if not isinstance(ridge_alpha, (int, float)) or ridge_alpha < 0:
+            raise ConfigError(
+                "ridge_alpha must be a non-negative number.", "ridge_alpha", known_keys
+            )
+        analytic_settings["ridge_alpha"] = float(ridge_alpha)
+
+        if ridge_cv_alphas is not None:
+            if (
+                not isinstance(ridge_cv_alphas, list)
+                or len(ridge_cv_alphas) < 2
+                or not all(isinstance(a, (int, float)) and a > 0 for a in ridge_cv_alphas)
+            ):
+                raise ConfigError(
+                    "ridge_cv_alphas must be a list of at least two positive numbers.",
+                    "ridge_cv_alphas",
+                    known_keys,
+                )
+        analytic_settings["ridge_cv_alphas"] = ridge_cv_alphas
+
+        ridge_cv_folds = settings.get("ridge_cv_folds", 5)
+        if not isinstance(ridge_cv_folds, int) or ridge_cv_folds < 2:
+            raise ConfigError(
+                "ridge_cv_folds must be an integer >= 2.", "ridge_cv_folds", known_keys
+            )
+        analytic_settings["ridge_cv_folds"] = ridge_cv_folds
+
+        # Elastic Net regularisation: fixed alpha or CV-selected alpha are mutually
+        # exclusive, and cannot be combined with Ridge regularisation.
+        elasticnet_alpha = settings.get("elasticnet_alpha", 0.0)
+        elasticnet_l1_ratio = settings.get("elasticnet_l1_ratio", 0.5)
+        elasticnet_cv_alphas = settings.get("elasticnet_cv_alphas", None)
+        elasticnet_cv_l1_ratios = settings.get("elasticnet_cv_l1_ratios", None)
+        elasticnet_cv_folds = settings.get("elasticnet_cv_folds", 5)
+
+        # Fixed-alpha and CV are mutually exclusive within Elastic Net
+        if elasticnet_alpha != 0.0 and (
+            elasticnet_cv_alphas is not None or elasticnet_cv_l1_ratios is not None
+        ):
+            raise ConfigError(
+                "Specify either 'elasticnet_alpha' (fixed) or "
+                "'elasticnet_cv_alphas'/'elasticnet_cv_l1_ratios' (CV), not both.",
+                "elasticnet_alpha",
+                known_keys,
+            )
+
+        # Ridge and Elastic Net cannot be mixed
+        elasticnet_requested = (
+            elasticnet_alpha != 0.0
+            or elasticnet_cv_alphas is not None
+            or elasticnet_cv_l1_ratios is not None
+        )
+        ridge_requested = ridge_alpha != 0.0 or ridge_cv_alphas is not None
+        if elasticnet_requested and ridge_requested:
+            raise ConfigError(
+                "Cannot combine Ridge ('ridge_alpha'/'ridge_cv_alphas') with "
+                "Elastic Net ('elasticnet_alpha'/'elasticnet_cv_alphas'/'elasticnet_cv_l1_ratios').",
+                "elasticnet_alpha",
+                known_keys,
+            )
+
+        if not isinstance(elasticnet_alpha, (int, float)) or elasticnet_alpha < 0:
+            raise ConfigError(
+                "elasticnet_alpha must be a non-negative number.",
+                "elasticnet_alpha",
+                known_keys,
+            )
+        analytic_settings["elasticnet_alpha"] = float(elasticnet_alpha)
+
+        if (
+            not isinstance(elasticnet_l1_ratio, (int, float))
+            or elasticnet_l1_ratio <= 0
+            or elasticnet_l1_ratio > 1
+        ):
+            raise ConfigError(
+                "elasticnet_l1_ratio must be a number in (0, 1].",
+                "elasticnet_l1_ratio",
+                known_keys,
+            )
+        analytic_settings["elasticnet_l1_ratio"] = float(elasticnet_l1_ratio)
+
+        if elasticnet_cv_alphas is not None:
+            if (
+                not isinstance(elasticnet_cv_alphas, list)
+                or len(elasticnet_cv_alphas) < 2
+                or not all(
+                    isinstance(a, (int, float)) and a > 0 for a in elasticnet_cv_alphas
+                )
+            ):
+                raise ConfigError(
+                    "elasticnet_cv_alphas must be a list of at least two positive numbers.",
+                    "elasticnet_cv_alphas",
+                    known_keys,
+                )
+        analytic_settings["elasticnet_cv_alphas"] = elasticnet_cv_alphas
+
+        if elasticnet_cv_l1_ratios is not None:
+            if not isinstance(elasticnet_cv_l1_ratios, list) or not all(
+                isinstance(r, (int, float)) and 0 < r <= 1
+                for r in elasticnet_cv_l1_ratios
+            ):
+                raise ConfigError(
+                    "elasticnet_cv_l1_ratios must be a list of numbers in (0, 1].",
+                    "elasticnet_cv_l1_ratios",
+                    known_keys,
+                )
+        analytic_settings["elasticnet_cv_l1_ratios"] = elasticnet_cv_l1_ratios
+
+        if not isinstance(elasticnet_cv_folds, int) or elasticnet_cv_folds < 2:
+            raise ConfigError(
+                "elasticnet_cv_folds must be an integer >= 2.",
+                "elasticnet_cv_folds",
+                known_keys,
+            )
+        analytic_settings["elasticnet_cv_folds"] = elasticnet_cv_folds
+
         return analytic_settings
 
     def parse_hessian_settings(self, settings):

colibri/tests/test_analytic_fit.py

@@ -165,6 +165,46 @@ def test_analytic_fit_different_priors(caplog):
     )
 
 
+def test_analytic_fit_elasticnet():
+    """
+    Tests that the analytic fit runs with Elastic Net regularisation
+    (fixed alpha, l1_ratio).
+    """
+    MOCK_PDF_MODEL.pred_and_pdf_func = lambda xgrid, forward_map: (
+        lambda params, fkarrs: (params, TEST_PDF_GRID)
+    )
+
+    _pred_data = TEST_FORWARD_MAP_DIS
+
+    enet_settings = {
+        "sampling_seed": 123,
+        "full_sample_size": 100,
+        "n_posterior_samples": 10,
+        "elasticnet_alpha": 0.1,
+        "elasticnet_l1_ratio": 0.5,
+        "elasticnet_cv_alphas": None,
+        "elasticnet_cv_l1_ratios": None,
+        "elasticnet_cv_folds": 5,
+    }
+
+    result = analytic_fit(
+        MOCK_CENTRAL_COVMAT_INDEX,
+        _pred_data,
+        MOCK_PDF_MODEL,
+        enet_settings,
+        TEST_PRIOR_SETTINGS_UNIFORM,
+        TEST_XGRID,
+        TEST_FK_ARRAYS,
+    )
+
+    assert isinstance(result, AnalyticFit)
+    assert result.analytic_specs == enet_settings
+    assert (
+        result.resampled_posterior.shape[0] == enet_settings["n_posterior_samples"]
+    )
+    assert len(result.param_names) == len(MOCK_PDF_MODEL.param_names)
+
+
 @patch("colibri.export_results.write_exportgrid")
 def test_run_analytic_fit(mock_write_exportgrid, tmp_path):