Add within-group lengthscale shrinkage prior for the strength GP

In the strength GP fit on the public dataset, the Matern ARD
lengthscales for Fly Ash and Coarse Aggregates rail to the upper
cap (~1000): both features are under-sampled relative to Cement and
Fine Aggregate, so the marginal likelihood is essentially flat in
their lengthscale direction and the optimiser drifts to the constraint
boundary. This makes the strength surface look insensitive to those
two ingredients in the explorer UI, even though concrete chemistry
clearly says it isn't.

This PR introduces `WithinGroupShrinkagePrior` — a soft hard-tying
prior that ties Matern ARD lengthscales within material groups. The
production setting (`sigma=0.001`) ties members of {Cement, Fly Ash,
Slag} to a single binder lengthscale, and {Fine, Coarse} aggregates
to a single aggregate lengthscale, encoding the domain prior that
ingredients of the same material class should share kernel smoothness.
The structural prior:

  * Beats the no-prior baseline on held-out LOO CV RMSE — 725 vs
    772 psi on the 647 public strength rows, via the GP's closed-form
    LOO identity.
  * Encodes only group structure (a public, ASTM-defined fact about
    cementitious vs aggregate materials), not any numeric values.
  * Makes Fly Ash and Coarse Aggregate sliders responsive in the
    explorer UI: both lengthscales now sit at 4.79 and 6.67
    respectively, well below the cap and within the same range as
    the previously-identifiable features.

Adds 11 regression tests in `test/test_lengthscale_identifiability.py`
covering identifiability against the cap, within-group tying on both
the committed and freshly-fit GP, prior installation by
`fit_strength_gp`, an LOO-CV ceiling that catches the no-prior baseline,
and unit tests for the prior class. A reproducible empirical sweep
(`scripts/lengthscale_prior_study.py`) lets users re-run the variant
comparison on their own machine. `docs/model/{strength,compositions,
test_vectors}.json` are regenerated with the new fit; the JS GP
inference and JS<->Python parity tests stay green. All LOO computation
re-uses `boxcrete.compute_loo_cv` (now also re-exported from the
package root).

See `WithinGroupShrinkagePrior` docstring in `boxcrete/models.py` for
the full empirical comparison and methodology.

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,