# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Data structures for GAM smooth terms specification."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Literal
[docs]
@dataclass
class SmoothTerm:
"""Specification of a smooth term s(x, ...)
This class defines how a predictor variable should be smoothed
in a GAM model.
Parameters
----------
variable : str or int
Variable name (if using DataFrame) or column index (if using array).
basis_type : {'bspline', 'cubic', 'tp', 'cr', 'ps'}, default='bspline'
Type of spline basis:
- 'bspline': B-spline basis (default)
- 'cubic': Natural cubic splines
- 'tp': Thin plate splines (future)
- 'cr': Cubic regression splines (future)
- 'ps': P-splines (future)
n_basis : int, default=10
Number of basis functions (k in mgcv terminology).
penalty_order : int, default=2
Order of difference penalty (usually 2 for second derivative).
by : str or int, optional
Factor variable for by-factor smooths (not yet implemented).
lambda_ : float, optional
Fixed smoothing parameter. If None, selected automatically.
knot_method : {'quantile', 'uniform'}, default='quantile'
Method for placing knots.
Examples
--------
>>> # Smooth term with default settings
>>> s1 = SmoothTerm(variable='x1')
>>> # Smooth term with custom settings
>>> s2 = SmoothTerm(
... variable='x2',
... basis_type='cubic',
... n_basis=15,
... lambda_=0.1
... )
>>> # Using column index instead of name
>>> s3 = SmoothTerm(variable=0, n_basis=12)
Notes
-----
This is inspired by mgcv's s() function in R but simplified for
initial implementation.
"""
variable: str | int
basis_type: Literal["bspline", "cubic", "tp", "cr", "ps"] = "bspline"
n_basis: int = 10
penalty_order: int = 2
by: str | int | None = None
lambda_: float | None = None
knot_method: Literal["quantile", "uniform"] = "quantile"
[docs]
def __post_init__(self):
"""Validate parameters after initialization."""
if self.n_basis < 3:
raise ValueError(f"n_basis must be at least 3, got {self.n_basis}")
if self.penalty_order < 1:
raise ValueError(f"penalty_order must be positive, got {self.penalty_order}")
if self.basis_type not in ["bspline", "cubic", "tp", "cr", "ps"]:
raise ValueError(f"Unknown basis_type: {self.basis_type}")
if self.basis_type not in ["bspline", "cubic"]:
raise NotImplementedError(
f"basis_type='{self.basis_type}' not yet implemented. Use 'bspline' or 'cubic'."
)
if self.lambda_ is not None and self.lambda_ < 0:
raise ValueError(f"lambda_ must be non-negative, got {self.lambda_}")
if self.by is not None:
raise NotImplementedError("by-factor smooths not yet implemented")
[docs]
@dataclass
class ParametricTerm:
"""Specification of a parametric (linear) term.
Parameters
----------
variable : str or int
Variable name or column index.
Examples
--------
>>> p1 = ParametricTerm(variable='age')
>>> p2 = ParametricTerm(variable=2) # column index
"""
variable: str | int
[docs]
@dataclass
class TensorTerm:
"""Specification of a tensor product smooth te(x1, x2, ...).
Tensor product smooths model interactions between variables using
the tensor product of marginal basis functions.
Parameters
----------
variables : tuple of (str or int)
Tuple of variable names or column indices (length 2 or more).
basis_types : tuple of str, optional
Basis type for each variable. If None, uses 'bspline' for all.
n_basis : tuple of int, optional
Number of basis functions for each variable. If None, uses 10 for all.
lambdas : tuple of float, optional
Fixed smoothing parameters for each direction. If None, selected automatically.
Examples
--------
>>> # 2D tensor product with defaults
>>> te1 = TensorTerm(variables=(0, 1))
>>> # Custom basis sizes
>>> te2 = TensorTerm(variables=('x1', 'x2'), n_basis=(12, 12))
>>> # Different basis types
>>> te3 = TensorTerm(
... variables=(0, 1),
... basis_types=('bspline', 'cubic'),
... n_basis=(10, 8)
... )
Notes
-----
Tensor products are useful for modeling interactions. They are more
flexible than additive models but require more data and computation.
The tensor product te(x1, x2) creates a surface f(x1, x2) rather than
the sum f1(x1) + f2(x2).
"""
variables: tuple[str | int, ...]
basis_types: tuple[str, ...] | None = None
n_basis: tuple[int, ...] | None = None
lambdas: tuple[float, ...] | None = None
[docs]
def __post_init__(self):
"""Validate parameters after initialization."""
if len(self.variables) < 2:
raise ValueError("TensorTerm requires at least 2 variables")
# Set defaults
n_vars = len(self.variables)
if self.basis_types is None:
self.basis_types = ("bspline",) * n_vars
elif len(self.basis_types) != n_vars:
raise ValueError(
f"basis_types must have same length as variables, "
f"got {len(self.basis_types)} vs {n_vars}"
)
if self.n_basis is None:
self.n_basis = (10,) * n_vars
elif len(self.n_basis) != n_vars:
raise ValueError(
f"n_basis must have same length as variables, got {len(self.n_basis)} vs {n_vars}"
)
if self.lambdas is not None and len(self.lambdas) != n_vars:
raise ValueError(
f"lambdas must have same length as variables, got {len(self.lambdas)} vs {n_vars}"
)
# Validate basis types
for bt in self.basis_types:
if bt not in ["bspline", "cubic"]:
raise NotImplementedError(
f"basis_type='{bt}' not supported in tensor products yet. "
f"Use 'bspline' or 'cubic'."
)
# Validate n_basis
for nb in self.n_basis:
if nb < 3:
raise ValueError(f"Each n_basis must be at least 3, got {nb}")
# Validate lambdas
if self.lambdas is not None:
for lam in self.lambdas:
if lam < 0:
raise ValueError(f"All lambdas must be non-negative, got {lam}")
__all__ = ["SmoothTerm", "ParametricTerm", "TensorTerm"]