feat(gam): add derivative-penalized B-splines

README.md

@@ -125,7 +125,7 @@ result, and prediction interfaces.
 
 | Formula surface    | Supported terms                                                                                      |
 | ------------------ | ---------------------------------------------------------------------------------------------------- |
-| Univariate smooths | `s(..., bs='cr')`, `cs`, `cc`, `cp`, `ps`, `tp`, `ts`                                                |
+| Univariate smooths | `s(..., bs='bs')`, `cr`, `cs`, `cc`, `cp`, `ps`, `tp`, `ts`                                          |
 | Structured smooths | random effects `re`, factor smooths `fs`, sum-to-zero factor smooths `sz`                            |
 | Tensor products    | `te(...)` and `ti(...)` over supported numeric marginals                                             |
 | Parametric terms   | numeric and factor terms, supported interactions, intercept policies, and formula offsets            |

nampy/gam/compiler/factory.py

@@ -15,13 +15,15 @@
 from ..smooths.registry import make_smooth_term
 from ..smooths.shape.bivariate import BivariateShapePSplineTerm
 from ..smooths.shape.scop import ShapeConstrainedPSplineTerm
+from ..smooths.univariate.bs import DerivativeBSplineTerm1D
 from ..smooths.univariate.cr import CubicSplineTerm
 from ..smooths.univariate.ps import PSplineTerm1D
 from ..specs import LinearPredictorSpec, PenaltyGroupSpec, TermSpec
 from ..specs.smooth import (
     CubicRegressionSmoothSpec,
     CubicShrinkageSmoothSpec,
     CyclicCubicRegressionSmoothSpec,
+    DerivativeBSplineSmoothSpec,
     FactorSmoothInteractionSpec,
     PSplineSmoothSpec,
     RandomEffectSmoothSpec,
@@ -162,6 +164,28 @@ def instantiate_term(term_like: TermSpec | Any):
             metadata=metadata,
         )
 
+    if isinstance(smooth_spec, DerivativeBSplineSmoothSpec):
+        if len(features) != 1:
+            raise NotImplementedError(
+                "Current runtime only materializes 1D s(..., bs='bs') terms."
+            )
+        return DerivativeBSplineTerm1D(
+            feature=features[0],
+            k=smooth_spec.k,
+            m=smooth_spec.m,
+            label=label,
+            term_id=term_like.term_id,
+            smoothing_id=smoothing_id,
+            by=by,
+            sp=smooth_spec.sp,
+            select=smooth_spec.select,
+            fixed=smooth_spec.fx,
+            constraint_mode=smooth_spec.constraint_mode,
+            pc=smooth_spec.pc,
+            knots=smooth_spec.knots,
+            metadata=metadata,
+        )
+
     if isinstance(smooth_spec, ShapeConstrainedSmoothSpec):
         if len(features) == 2:
             return BivariateShapePSplineTerm(
@@ -242,6 +266,7 @@ def instantiate_term(term_like: TermSpec | Any):
             by=by,
             sp=smooth_spec.sp,
             select=smooth_spec.select,
+            m=smooth_spec.m,
             xt=smooth_spec.xt,
             fixed=smooth_spec.fx,
             knots=smooth_spec.knots,
@@ -258,6 +283,7 @@ def instantiate_term(term_like: TermSpec | Any):
             by=by,
             sp=smooth_spec.sp,
             select=smooth_spec.select,
+            m=smooth_spec.m,
             xt=smooth_spec.xt,
             fixed=smooth_spec.fx,
             knots=smooth_spec.knots,
@@ -304,9 +330,15 @@ def _expected_penalty_group_size(runtime_term):
     if bool(getattr(runtime_term, "fixed", False)):
         return 0
 
+    if hasattr(runtime_term, "expected_linked_penalty_count"):
+        value = runtime_term.expected_linked_penalty_count
+        return None if value is None else int(value)
+
     fixed_flags = getattr(runtime_term, "fixed_flags", None)
     if fixed_flags is not None:
         n_penalties = int(np.sum(~np.asarray(fixed_flags, dtype=bool)))
+    elif getattr(runtime_term, "n_main_penalties", None) is not None:
+        n_penalties = int(runtime_term.n_main_penalties)
     else:
         term_type = str(getattr(runtime_term, "term_type", "smooth"))
         if term_type in {"tensor_smooth", "tensor_interaction"}:

nampy/gam/smooths/__init__.py

@@ -19,6 +19,7 @@
 )
 from .tensor.te import TensorProductSplineTerm
 from .tensor.ti import InteractionTensorProductSplineTerm
+from .univariate.bs import DerivativeBSplineTerm1D
 from .univariate.cr import CubicSplineTerm
 from .univariate.ps import PSplineTerm1D
 from .univariate.tp import ThinPlateSplineTerm
@@ -28,6 +29,7 @@
 te = TensorProductSplineTerm
 ti = InteractionTensorProductSplineTerm
 
+bs = DerivativeBSplineTerm1D
 cr = cs = cc = CubicSplineTerm
 cp = ps = PSplineTerm1D
 tp = ts = ThinPlateSplineTerm
@@ -55,6 +57,7 @@
     "sync_by_state_attributes",
     "build_penalty_definition",
     "CubicSplineTerm",
+    "DerivativeBSplineTerm1D",
     "PSplineTerm1D",
     "ThinPlateSplineTerm",
     "TensorProductSplineTerm",
@@ -65,6 +68,7 @@
     "ShapeConstrainedPSplineTerm",
     "te",
     "ti",
+    "bs",
     "cr",
     "cs",
     "cc",

nampy/gam/smooths/categorical/fs.py

@@ -15,6 +15,7 @@
 from ..algebra import rowwise_kronecker
 from ..registry import make_smooth_term
 from ..smooth_base import BaseSmoothTerm, by_values_from_new_data, column_as_object
+from ..univariate.bs import DerivativeBSplineTerm1D
 from ..univariate.cr import CubicSplineTerm
 from ..univariate.ps import PSplineTerm1D
 from .categorical_utils import (
@@ -97,6 +98,7 @@ def _build_base_smooth_term(
     by,
     knots,
     xt_rest,
+    outer_m,
     mode,  # "fs" or "sz"
     select,
     constraint_mode,
@@ -106,7 +108,7 @@ def _build_base_smooth_term(
     Build the per-level base smooth used inside fs/sz.
 
     Supported base smooth classes in the current codebase:
-    cr, cs, cc, cp, ps, tp, ts
+    bs, cr, cs, cc, cp, ps, tp, ts
     """
     base_bs = str(base_bs).lower()
     metric_features = list(metric_features)
@@ -123,9 +125,9 @@ def _build_base_smooth_term(
             f"for bs in {{'tp','ts'}}, got base bs={base_bs!r}."
         )
 
-    if xt_rest is not None and base_bs not in {"tp", "ts", "ps", "cp"}:
+    if xt_rest is not None and base_bs not in {"bs", "tp", "ts", "ps", "cp"}:
         raise NotImplementedError(
-            "Extra xt options are currently only supported for tp/ts/ps/cp base "
+            "Extra xt options are currently only supported for bs/tp/ts/ps/cp base "
             "smooths, "
             f"got xt={xt_rest!r} with base bs={base_bs!r}."
         )
@@ -149,7 +151,9 @@ def _build_base_smooth_term(
         )
 
     if base_bs in {"ps", "cp"}:
-        ps_m = None if xt_rest is None else xt_rest.get("m", None)
+        ps_m = outer_m
+        if ps_m is None and xt_rest is not None:
+            ps_m = xt_rest.get("m", None)
         # For fs/sz, mgcv keeps the outer basis dimension and uses xt mainly to
         # choose the base smoother family / order parameters.
         ps_k = k
@@ -170,6 +174,26 @@ def _build_base_smooth_term(
             metadata=metadata,
         )
 
+    if base_bs == "bs":
+        bs_m = outer_m
+        if bs_m is None and xt_rest is not None:
+            bs_m = xt_rest.get("m", None)
+        return DerivativeBSplineTerm1D(
+            feature=metric_features[0],
+            k=k,
+            m=bs_m,
+            label=label,
+            smoothing_id=None,
+            by=by,
+            sp=None,
+            select=bool(select),
+            fixed=bool(fixed),
+            constraint_mode=str(constraint_mode),
+            pc=None,
+            knots=knots,
+            metadata=metadata,
+        )
+
     if base_bs in {"tp", "ts"}:
         return make_smooth_term(
             base_bs,
@@ -192,11 +216,13 @@ def _build_base_smooth_term(
 
     raise NotImplementedError(
         f"Current {mode} implementation supports base bs in "
-        f"{{'cr','cs','cc','cp','ps','tp','ts'}}, got {base_bs!r}."
+        f"{{'bs','cr','cs','cc','cp','ps','tp','ts'}}, got {base_bs!r}."
     )
 
 
 def _penalty_rank_from_base_term(base_term, basis_matrix, penalty_matrix) -> int:
+    if isinstance(base_term, DerivativeBSplineTerm1D):
+        return int(base_term._setup.ranks[0])
     if isinstance(base_term, PSplineTerm1D) and len(base_term.penalties) > 0:
         if str(base_term.basis_name).lower() == "cp":
             return int(base_term._setup.rank)
@@ -286,6 +312,7 @@ def __init__(
         by=None,
         sp=None,
         select=False,
+        m=None,
         xt=None,
         fixed=False,
         knots=None,
@@ -307,6 +334,7 @@ def __init__(
         self.term_type = term_type
         self.k = int(k)
         self.select = bool(select)
+        self.m = m
         self.xt = xt
         self.fixed = bool(fixed)
         self.knots = knots
@@ -438,6 +466,7 @@ def _build_delegate_base_or_re(self, X, feature_names, *, default_bs, mode):
                 by=self.by,
                 knots=self.knots,
                 xt_rest=base_spec.xt_rest,
+                outer_m=self.m,
                 mode=mode,
                 select=self.select,
                 constraint_mode=("auto" if mode == "fs" else "never"),
@@ -504,6 +533,7 @@ def __init__(
         by=None,
         sp=None,
         select=False,
+        m=None,
         xt=None,
         fixed=False,
         knots=None,
@@ -520,6 +550,7 @@ def __init__(
             by=by,
             sp=sp,
             select=select,
+            m=m,
             xt=xt,
             fixed=fixed,
             knots=knots,
@@ -558,6 +589,7 @@ def fit(self, X, feature_names):
             by=None,
             knots=self.knots,
             xt_rest=base_spec.xt_rest,
+            outer_m=self.m,
             mode="fs",
             select=False,
             constraint_mode=base_constraint,
@@ -567,7 +599,7 @@ def fit(self, X, feature_names):
 
         if len(base_term.penalties) > 1:
             raise NotImplementedError(
-                'bs="fs" currently requires a singly penalized base smooth.'
+                '"fs" smooth cannot use a multiply penalized basis (wrong basis in xt)'
             )
 
         self._base_term = base_term
@@ -789,6 +821,7 @@ def __init__(
         by=None,
         sp=None,
         select=False,
+        m=None,
         xt=None,
         fixed=False,
         knots=None,
@@ -805,6 +838,7 @@ def __init__(
             by=by,
             sp=sp,
             select=select,
+            m=m,
             xt=xt,
             fixed=fixed,
             knots=knots,
@@ -834,6 +868,7 @@ def fit(self, X, feature_names):
             by=None,
             knots=self.knots,
             xt_rest=base_spec.xt_rest,
+            outer_m=self.m,
             mode="sz",
             select=False,
             constraint_mode="never",
@@ -843,7 +878,7 @@ def fit(self, X, feature_names):
 
         if len(base_term.penalties) > 1:
             raise NotImplementedError(
-                'bs="sz" currently requires a singly penalized base smooth.'
+                '"sz" smooth cannot use a multiply penalized basis (wrong basis in xt)'
             )
 
         self._base_term = base_term

nampy/gam/smooths/tensor/marginals.py

@@ -7,11 +7,12 @@
 from ...penalties.tensor import normalize_tensor_marginal_penalty
 from ..algebra import rowwise_kronecker
 from ..smooth_base import column_as_float
+from ..univariate.bs import DerivativeBSplineTerm1D
 from ..univariate.cr import CubicSplineTerm
 from ..univariate.ps import PSplineTerm1D
 from ..univariate.tp import ThinPlateSplineTerm
 
-TENSOR_MARGINAL_BASES = frozenset({"cr", "cs", "cc", "cp", "ps", "tp", "ts"})
+TENSOR_MARGINAL_BASES = frozenset({"bs", "cr", "cs", "cc", "cp", "ps", "tp", "ts"})
 
 
 def _as_marginal_features(feature):
@@ -93,6 +94,23 @@ def make_tensor_marginal_term(
             metadata=metadata,
         )
 
+    if basis == "bs":
+        if len(marginal_features) != 1:
+            raise ValueError("Tensor marginal basis 'bs' only handles one feature.")
+        return DerivativeBSplineTerm1D(
+            feature=marginal_features[0],
+            k=k,
+            m=m,
+            label=str(feature),
+            smoothing_id=None,
+            by=None,
+            select=False,
+            fixed=False,
+            constraint_mode=constraint_mode,
+            knots=knots,
+            metadata=metadata,
+        )
+
     if basis in {"tp", "ts"}:
         return ThinPlateSplineTerm(
             feature=marginal_features,

nampy/gam/smooths/univariate/__init__.py

@@ -1,10 +1,17 @@
+from .bs import DerivativeBSplineTerm1D
 from .cr import CubicSplineTerm
 from .ps import PSplineTerm1D
 from .tp import ThinPlateSplineTerm
 
+bs = DerivativeBSplineTerm1D
 cr = cs = cc = CubicSplineTerm
 cp = ps = PSplineTerm1D
 tp = ts = ThinPlateSplineTerm
 
-__all__ = ["CubicSplineTerm", "PSplineTerm1D", "ThinPlateSplineTerm"]
-__all__ += ["cr", "cs", "cc", "cp", "ps", "tp", "ts"]
+__all__ = [
+    "DerivativeBSplineTerm1D",
+    "CubicSplineTerm",
+    "PSplineTerm1D",
+    "ThinPlateSplineTerm",
+]
+__all__ += ["bs", "cr", "cs", "cc", "cp", "ps", "tp", "ts"]