From f6372737c07a19528b3e3dd4aedb587f5d4f94be Mon Sep 17 00:00:00 2001 From: Philipp Kopper Date: Sun, 14 Jun 2026 18:44:29 +0200 Subject: [PATCH] Terms implementation --- src/glum/_glm.py | 4 + src/glum/_term_layout.py | 118 ++++++++++++++ tests/glm/test_term_layout.py | 280 ++++++++++++++++++++++++++++++++++ 3 files changed, 402 insertions(+) create mode 100644 src/glum/_term_layout.py create mode 100644 tests/glm/test_term_layout.py diff --git a/src/glum/_glm.py b/src/glum/_glm.py index a88e90da..6c8beb75 100644 --- a/src/glum/_glm.py +++ b/src/glum/_glm.py @@ -48,6 +48,7 @@ _tikhonov_solver, _trust_constr_solver, ) +from ._term_layout import TermLayout from ._typing import ( ArrayLike, ShapedArrayLike, @@ -2162,6 +2163,9 @@ def _set_up_and_check_fit_args( self.feature_names_ = X.get_names(type="column", missing_prefix="_col_") # type: ignore self.term_names_ = X.get_names(type="term", missing_prefix="_col_") + self.term_layout_ = TermLayout.from_term_names( + self.term_names_, self.fit_intercept + ) return X, y, sample_weight, offset, weights_sum, P1, P2 # type: ignore diff --git a/src/glum/_term_layout.py b/src/glum/_term_layout.py new file mode 100644 index 00000000..1dbb1aa6 --- /dev/null +++ b/src/glum/_term_layout.py @@ -0,0 +1,118 @@ +"""Term-to-coefficient layout for fitted GLM models. + +A :class:`TermLayout` is a lightweight, immutable view that records which +contiguous slice of the coefficient vector each model term occupies. It is +populated alongside ``term_names_`` after fitting and is intended to be the +single source of truth for code that needs to reason about *terms* rather than +individual coefficients (penalty assembly, term-level diagnostics, smooth +construction, etc.). + +This module is intentionally small. It contains no formula logic and no +reference to ``formulaic`` or ``tabmat``; it builds the layout from the +per-coefficient term-name list that ``GeneralizedLinearRegressorBase`` already +constructs. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class TermSlice: + """One term and the coefficient indices it owns. + + ``start`` / ``stop`` are half-open and index the full coefficient vector, + i.e. with the intercept at position 0 when ``fit_intercept`` is true. + """ + + name: str + start: int + stop: int + kind: str + + @property + def n_coefs(self) -> int: + return self.stop - self.start + + +@dataclass(frozen=True, slots=True) +class TermLayout: + """Ordered view of terms and the coefficient slice each one owns. + + The layout is derived from the per-coefficient term-name list produced + during fit and the intercept flag. Construction is O(n_features); after + that all lookups are O(n_terms). + """ + + terms: tuple[TermSlice, ...] + has_intercept: bool + _by_name: dict[str, TermSlice] + + @classmethod + def from_term_names( + cls, + term_names: Sequence[str], + fit_intercept: bool, + ) -> TermLayout: + """Build a layout from a per-coefficient term-name list. + + ``term_names`` must have one entry per non-intercept coefficient, in + coefficient order. Adjacent equal entries are collapsed into a single + term slice. If ``fit_intercept`` is true, an ``"intercept"`` slice is + prepended at index 0. + """ + slices: list[TermSlice] = [] + offset = 1 if fit_intercept else 0 + if fit_intercept: + slices.append(TermSlice("intercept", 0, 1, "intercept")) + + i = 0 + n = len(term_names) + while i < n: + j = i + 1 + current = term_names[i] + while j < n and term_names[j] == current: + j += 1 + slices.append( + TermSlice( + name=str(current) if current is not None else f"_term_{i}", + start=offset + i, + stop=offset + j, + kind="term", + ) + ) + i = j + + return cls( + terms=tuple(slices), + has_intercept=fit_intercept, + _by_name={s.name: s for s in slices}, + ) + + def __iter__(self) -> Iterator[TermSlice]: + return iter(self.terms) + + def __len__(self) -> int: + return len(self.terms) + + def __getitem__(self, key): + if isinstance(key, str): + return self._by_name[key] + return self.terms[key] + + def __contains__(self, key: object) -> bool: + if not isinstance(key, str): + return False + return key in self._by_name + + @property + def n_coefs(self) -> int: + """Total number of coefficients spanned by the layout.""" + return self.terms[-1].stop if self.terms else 0 + + @property + def names(self) -> tuple[str, ...]: + """Term names in coefficient order.""" + return tuple(term.name for term in self.terms) diff --git a/tests/glm/test_term_layout.py b/tests/glm/test_term_layout.py new file mode 100644 index 00000000..0798ceff --- /dev/null +++ b/tests/glm/test_term_layout.py @@ -0,0 +1,280 @@ +import time + +import numpy as np +import pandas as pd +import pytest + +from glum._glm import GeneralizedLinearRegressor +from glum._term_layout import TermLayout + + +def _slices_partition(layout: TermLayout) -> bool: + """True iff layout slices form a contiguous, ordered partition of [0, n_coefs).""" + expected_start = 0 + for term in layout: + if term.start != expected_start or term.stop <= term.start: + return False + expected_start = term.stop + return expected_start == layout.n_coefs + + +def test_consecutive_duplicate_term_names_collapse_into_one_slice(): + # Arrange + term_names = ["x1", "c1", "c1", "c1", "x2"] + + # Act + layout = TermLayout.from_term_names(term_names, fit_intercept=False) + + # Assert + assert [t.name for t in layout] == ["x1", "c1", "x2"] + assert layout["c1"].n_coefs == 3 + + +def test_intercept_owns_first_coefficient_when_fit_intercept_true(): + # Arrange + term_names = ["x1", "x2"] + + # Act + layout = TermLayout.from_term_names(term_names, fit_intercept=True) + + # Assert + assert layout.has_intercept is True + assert "intercept" in layout + assert layout["intercept"].start == 0 + assert layout["intercept"].stop == 1 + assert layout["x1"].start == 1 + + +def test_layout_without_intercept_starts_at_index_zero(): + # Arrange + term_names = ["x1", "x2"] + + # Act + layout = TermLayout.from_term_names(term_names, fit_intercept=False) + + # Assert + assert layout.has_intercept is False + assert "intercept" not in layout + assert layout["x1"].start == 0 + + +def test_empty_term_names_with_intercept_yields_intercept_only_layout(): + # Arrange / Act + layout = TermLayout.from_term_names([], fit_intercept=True) + + # Assert + assert layout.n_coefs == 1 + assert [t.name for t in layout] == ["intercept"] + + +def test_empty_term_names_without_intercept_yields_empty_layout(): + # Arrange / Act + layout = TermLayout.from_term_names([], fit_intercept=False) + + # Assert + assert layout.n_coefs == 0 + assert list(layout) == [] + + +@pytest.mark.parametrize( + "term_names,fit_intercept,expected_n_coefs", + [ + (["x1", "x2", "x3"], False, 3), + (["x1", "x2", "x3"], True, 4), + (["c1", "c1", "c1"], True, 4), + (["a", "b", "b", "c", "c", "c"], False, 6), + ], +) +def test_slices_partition_the_full_coefficient_vector( + term_names, fit_intercept, expected_n_coefs +): + # Arrange / Act + layout = TermLayout.from_term_names(term_names, fit_intercept=fit_intercept) + + # Assert + assert layout.n_coefs == expected_n_coefs + assert _slices_partition(layout) + + +def test_lookup_by_name_returns_slice_covering_that_terms_coefficients(): + # Arrange + term_names = ["x1", "c1", "c1", "c1", "x2"] + layout = TermLayout.from_term_names(term_names, fit_intercept=True) + + # Act + c1_slice = layout["c1"] + + # Assert: slice covers exactly the c1 entries in term_names (intercept offset 1) + coef_indices = range(c1_slice.start, c1_slice.stop) + assert [term_names[i - 1] for i in coef_indices] == ["c1", "c1", "c1"] + + +def test_lookup_by_unknown_name_raises_key_error(): + # Arrange + layout = TermLayout.from_term_names(["x1"], fit_intercept=False) + + # Act / Assert + with pytest.raises(KeyError): + layout["does_not_exist"] + + +def test_membership_check_distinguishes_known_from_unknown_terms(): + # Arrange + layout = TermLayout.from_term_names(["x1", "x2"], fit_intercept=True) + + # Assert (no separate Act; membership is the operation under test) + assert "x1" in layout + assert "intercept" in layout + assert "x3" not in layout + assert 42 not in layout # non-string membership returns False + + +def test_iteration_yields_terms_in_coefficient_order(): + # Arrange + term_names = ["a", "b", "b", "c"] + layout = TermLayout.from_term_names(term_names, fit_intercept=True) + + # Act + iterated = [(t.name, t.start) for t in layout] + + # Assert: each successive term starts at or after the previous one + starts = [start for _, start in iterated] + assert starts == sorted(starts) + assert [name for name, _ in iterated] == ["intercept", "a", "b", "c"] + + +def test_term_slice_n_coefs_equals_stop_minus_start(): + # Arrange + layout = TermLayout.from_term_names(["c1", "c1", "c1"], fit_intercept=False) + + # Act + term = layout["c1"] + + # Assert + assert term.n_coefs == term.stop - term.start + + +# --- Integration: end-to-end after .fit() --------------------------------- + + +def test_fit_layout_n_coefs_matches_total_coefficients(): + # Arrange + rng = np.random.default_rng(0) + X = rng.standard_normal((50, 3)) + y = rng.standard_normal(50) + + # Act + model = GeneralizedLinearRegressor(family="normal", alpha=0.1).fit(X, y) + + # Assert + expected_total = (1 if model.fit_intercept else 0) + len(model.coef_) + assert model.term_layout_.n_coefs == expected_total + + +def test_fit_without_intercept_layout_excludes_intercept_slice(): + # Arrange + rng = np.random.default_rng(0) + X = rng.standard_normal((40, 2)) + y = rng.standard_normal(40) + + # Act + model = GeneralizedLinearRegressor( + family="normal", alpha=0.1, fit_intercept=False + ).fit(X, y) + + # Assert + assert "intercept" not in model.term_layout_ + assert model.term_layout_.has_intercept is False + + +def test_fit_layout_slices_partition_the_coefficient_vector(): + # Arrange + rng = np.random.default_rng(1) + df = pd.DataFrame( + { + "y": rng.standard_normal(60), + "x1": rng.standard_normal(60), + "c1": pd.Categorical(rng.choice(["a", "b", "c"], 60)), + } + ) + + # Act + model = GeneralizedLinearRegressor( + family="normal", formula="y ~ x1 + c1", alpha=0.1, drop_first=False + ).fit(df) + + # Assert + assert _slices_partition(model.term_layout_) + + +def test_fit_layout_term_lookup_yields_correct_design_matrix_columns(): + """The strongest behavioral test: layout slices must point at the right + columns of the design matrix that produced ``coef_``. + """ + # Arrange + rng = np.random.default_rng(2) + df = pd.DataFrame( + { + "y": rng.standard_normal(80), + "x1": rng.standard_normal(80), + "c1": pd.Categorical(rng.choice(["a", "b", "c"], 80)), + } + ) + model = GeneralizedLinearRegressor( + family="normal", formula="y ~ x1 + c1", alpha=0.1, drop_first=False + ).fit(df) + intercept_offset = 1 if model.fit_intercept else 0 + + # Act: pick the categorical term and read back the column names it spans + c1_slice = model.term_layout_["c1"] + coef_indices = range(c1_slice.start, c1_slice.stop) + feature_indices = [i - intercept_offset for i in coef_indices] + spanned_features = [model.feature_names_[i] for i in feature_indices] + + # Assert: every spanned feature name carries the c1 prefix + assert len(spanned_features) >= 2 # at least two levels remain + assert all(name.startswith("c1") for name in spanned_features) + + +def test_fit_layout_is_deterministic_for_identical_fits(): + # Arrange + rng = np.random.default_rng(3) + X = rng.standard_normal((30, 4)) + y = rng.standard_normal(30) + + # Act + layout_a = ( + GeneralizedLinearRegressor(family="normal", alpha=0.1).fit(X, y).term_layout_ + ) + layout_b = ( + GeneralizedLinearRegressor(family="normal", alpha=0.1).fit(X, y).term_layout_ + ) + + # Assert + assert [t.name for t in layout_a] == [t.name for t in layout_b] + assert [(t.start, t.stop) for t in layout_a] == [ + (t.start, t.stop) for t in layout_b + ] + + +# --- Performance regression guard ----------------------------------------- + + +def test_layout_construction_overhead_is_bounded(): + """Construction is O(n_terms) and runs once per .fit(). Lock in a generous + upper bound so an accidental quadratic regression would fail loudly. + """ + # Arrange + n = 1000 + term_names = [f"x{i}" for i in range(n)] + + # Act + start = time.perf_counter() + for _ in range(100): + TermLayout.from_term_names(term_names, fit_intercept=True) + elapsed_per_build_ms = (time.perf_counter() - start) / 100 * 1000 + + # Assert: 1000 single-coef terms should build in well under 5 ms + assert elapsed_per_build_ms < 5.0, ( + f"TermLayout build regressed: {elapsed_per_build_ms:.2f} ms/build" + )