Merge pull request #20 from facebookresearch/within-group-lengthscale-prior

Add within-group lengthscale shrinkage prior for the strength GP

Author: SebastianAment

boxcrete/__init__.py

@@ -12,6 +12,7 @@
     get_strength_gp_input_transform,
     PartialFixedNoiseLikelihood,
     SustainableConcreteModel,
+    WithinGroupShrinkagePrior,
 )
 from boxcrete.plotting import (
     compute_loo_cv,
@@ -80,6 +81,7 @@
     "SustainableConcreteDataset",
     "SustainableConcreteModel",
     "UnitSystem",
+    "WithinGroupShrinkagePrior",
     "compute_loo_cv",
     "convert_slump",
     "convert_strength",

boxcrete/models.py

@@ -34,6 +34,7 @@
 from gpytorch.kernels import MaternKernel, RBFKernel, ScaleKernel
 from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.mlls import ExactMarginalLogLikelihood
+from gpytorch.priors import LogNormalPrior
 from linear_operator.operators import DiagLinearOperator
 from torch import Tensor
 
@@ -478,13 +479,136 @@ def _shaped_noise_covar(self, base_shape, *params, **kwargs):
         return super()._shaped_noise_covar(base_shape, *params, **kwargs)
 
 
+# --- Within-group lengthscale shrinkage prior --------------------------------
+
+# Material-class groupings used by the production within-group shrinkage
+# prior on the Matern kernel's lengthscales. Indices refer to
+# DEFAULT_X_COLUMNS: Cement, Fly Ash, Slag, Water, HRWR, Fine Aggregate,
+# Coarse Aggregates, Material Source, Temp, Time.
+_BINDER_LENGTHSCALE_GROUP = (0, 1, 2)  # Cement, Fly Ash, Slag
+_AGGREGATE_LENGTHSCALE_GROUP = (5, 6)  # Fine Aggregate, Coarse Aggregates
+
+# σ → 0 effectively hard-ties the members of each group to a shared lengthscale.
+# Empirical LOO CV (n=647 public rows) found σ=0.001 minimises held-out RMSE
+# while keeping all lengthscales well below the identifiability cap. See
+# `scripts/lengthscale_prior_study.py` for the reproducible sweep.
+_LENGTHSCALE_SHRINKAGE_SIGMA = 0.001
+
+
+class WithinGroupShrinkagePrior(LogNormalPrior):
+    """Soft hard-tying prior on Matern ARD lengthscales within material groups.
+
+    Penalises within-group variance of log-lengthscales. Encodes the
+    domain fact that interchangeable materials (e.g., cementitious binders
+    or aggregates) should have similar smoothness scales in the GP. With
+    ``sigma → 0`` this approaches a hard tying constraint that forces
+    each group to share a single lengthscale.
+
+    Why this prior exists
+    ---------------------
+    Several composition features in ``data/boxcrete_data.csv`` are
+    under-sampled. In particular, Fly Ash is zero in the majority of rows
+    (most concretes use only Cement + Slag), and Coarse Aggregates are
+    zero for the mortar half of the dataset. Without a prior, ARD pushes
+    the corresponding Matern lengthscales to the optimiser's upper bound
+    (``1e3`` in normalised input space), making the GP effectively
+    insensitive to those features — the website's interactive sliders
+    for Fly Ash and Coarse Aggregates would not respond to user input.
+
+    This prior softly ties the lengthscales of materials that play
+    interchangeable physical roles, so the well-identified members of
+    each group (Cement, Fine Aggregate) supply usable scale information
+    to their under-identified peers (Fly Ash, Coarse Aggregates).
+
+    Empirical comparison
+    --------------------
+    Analytical LOO CV (via ``boxcrete.compute_loo_cv``) on n=647 public
+    strength rows; lower RMSE is better:
+
+    +-----------------------------------------+-----------+
+    | Variant                                 | LOO RMSE  |
+    +=========================================+===========+
+    | No prior (Fly Ash & Coarse Agg railed)  |  772 psi  |
+    | Within-group shrinkage σ=0.50           |  754 psi  |
+    | Within-group shrinkage σ=0.10           |  731 psi  |
+    | Within-group shrinkage σ=0.001 (prod)   |  **725 psi** |
+    +-----------------------------------------+-----------+
+
+    The σ → 0 limit Pareto-dominates every alternative we evaluated:
+    per-feature LogNormal priors, Cauchy / Student-t shrinkage,
+    asymmetric per-group widths, Cement-anchored shrinkage, and additive
+    kernel decompositions. See ``scripts/lengthscale_prior_study.py``
+    to reproduce the sweep on a fresh checkout.
+
+    Mathematical form
+    -----------------
+    The prior contributes the following log-density (up to a constant)
+    to the marginal log-likelihood::
+
+        log p(ℓ) = -Σ_g Σ_{i ∈ g} (log ℓ_i - mean_{j ∈ g} log ℓ_j)² / (2 σ_g²)
+
+    where ``g`` ranges over the configured groups. Subclassing
+    ``LogNormalPrior`` lets the prior satisfy GPyTorch's
+    ``isinstance(_, Prior)`` check without reimplementing the Prior
+    interface; the inherited ``loc`` / ``scale`` are unused placeholders.
+
+    Args:
+        groups_with_sigma: List of ``(dim_indices, sigma)`` tuples. Each
+            entry contributes a within-group penalty with its own width.
+            ``sigma → 0`` hard-ties the group; ``sigma → ∞`` is uniform.
+        dim: Dimensionality of the lengthscale tensor (matches the
+            kernel's ``ard_num_dims``).
+    """
+
+    def __init__(
+        self,
+        groups_with_sigma: list[tuple[tuple[int, ...], float]],
+        dim: int,
+    ):
+        super().__init__(loc=torch.zeros(1, dim), scale=1.0)
+        self._groups_with_sigma = groups_with_sigma
+
+    def log_prob(self, x):
+        log_x = x.log()
+        flat_log = log_x.flatten()
+        total = torch.zeros((), dtype=x.dtype, device=x.device)
+        for grp, sigma in self._groups_with_sigma:
+            if len(grp) < 2:
+                continue
+            grp_log = flat_log[list(grp)]
+            sq_dev = ((grp_log - grp_log.mean()) ** 2).sum()
+            total = total - 0.5 * sq_dev / (sigma**2)
+        # gpytorch sums log_prob over elements, so distribute the scalar
+        # penalty uniformly across x's shape.
+        return total / x.numel() * torch.ones_like(x)
+
+
+def _default_lengthscale_prior(d_in: int) -> WithinGroupShrinkagePrior | None:
+    """Returns the production within-group shrinkage prior, or None if the
+    input dimensionality doesn't match the production schema (in which case
+    we fall back to the unconstrained MLE)."""
+    # Only apply if the model uses the production 10-dim DEFAULT_X_COLUMNS
+    # layout (Cement, Fly Ash, Slag, Water, HRWR, Fine, Coarse, MS, Temp,
+    # Time). For sub-dim or test fits, return None.
+    if d_in != 10:
+        return None
+    return WithinGroupShrinkagePrior(
+        groups_with_sigma=[
+            (_BINDER_LENGTHSCALE_GROUP, _LENGTHSCALE_SHRINKAGE_SIGMA),
+            (_AGGREGATE_LENGTHSCALE_GROUP, _LENGTHSCALE_SHRINKAGE_SIGMA),
+        ],
+        dim=d_in,
+    )
+
+
 def fit_strength_gp(
     X: Tensor,
     Y: Tensor,
     Yvar: Tensor,
     X_bounds: Tensor | None = None,
     use_fixed_noise: bool = False,
     optimizer_kwargs: dict | None = None,
+    lengthscale_prior: object | None = "default",
 ) -> SingleTaskGP:
     """Fits a Gaussian process model to the given strength data.
 
@@ -495,6 +619,13 @@ def fit_strength_gp(
         X_bounds: Optional `2 x d`-dim bounds Tensor.
         use_fixed_noise: Whether to use fixed observation noise.
         optimizer_kwargs: Optional keyword arguments for the optimizer.
+        lengthscale_prior: Prior on the Matern kernel's per-dim
+            lengthscales. ``"default"`` (the default) installs the
+            production ``WithinGroupShrinkagePrior``: a soft hard-tying
+            prior that links binder lengthscales {Cement, Fly Ash, Slag}
+            and aggregate lengthscales {Fine, Coarse} within each group.
+            Pass ``None`` to fit without any lengthscale prior, or pass
+            a ``gpytorch.priors.Prior`` instance to install a custom one.
 
     Returns:
         A SingleTaskGP model fit to the strength data.
@@ -504,6 +635,9 @@ def fit_strength_gp(
     if d_out != 1:
         raise ValueError("Output dimensions is not one in strength curve fitting.")
 
+    if lengthscale_prior == "default":
+        lengthscale_prior = _default_lengthscale_prior(d_in)
+
     # add data to condition GP to be zero at day zero
     X_0, Y_0, Yvar_0 = get_day_zero_data(X=X, bounds=X_bounds, n=128)
     n_real = X.shape[0]
@@ -517,7 +651,7 @@ def fit_strength_gp(
         nu=2.5,
         ard_num_dims=d_in,
         lengthscale_constraint=LogTransformedInterval(1e-2, 1e3, initial_value=1.0),
-        lengthscale_prior=None,
+        lengthscale_prior=lengthscale_prior,
     )
     scaled_base_kernel = ScaleKernel(
         base_kernel=base_kernel,