diff --git a/.claude/skills/jaxsr/SKILL.md b/.claude/skills/jaxsr/SKILL.md index 5b1c908..1956e0f 100644 --- a/.claude/skills/jaxsr/SKILL.md +++ b/.claude/skills/jaxsr/SKILL.md @@ -43,12 +43,15 @@ Based on the answers, recommend a basis library configuration: | Large feature space (screening) | `add_constant + add_linear + add_interactions(2)` then use `lasso_path` strategy | | Response surface (DOE) | `add_constant + add_linear + add_polynomials(2) + add_interactions(2)` — or use `ResponseSurface` directly | | Categorical factors present | Add `add_categorical_indicators() + add_categorical_interactions()` to any of the above | +| Unknown coefficient *function* multiplying a data column (superposition, implicit dynamics) | `add_block(theta, multiply_by="", block_name=...)` — see `guides/basis-library.md` | **Key guidance:** - Start simple. You can always add complexity. - `add_transcendental(safe=True)` guards against log(0), 1/0, sqrt(<0). Always use `safe=True`. - `add_ratios(safe=True)` adds x_i/x_j terms. Doubles the library size — only use when ratios are physically meaningful. - `add_parametric()` enables nonlinear parameters (e.g., `exp(-a*x)`). Powerful but slower to fit. +- `add_block()` multiplies a whole basis by a data column, so a selected coefficient is a term of an + unknown coefficient function. Drop a block with `without_blocks()` to test whether it earned its place. - If n_features > 5, avoid `add_polynomials(degree>2)` — the library becomes enormous. ### Step 3: Recommend Selection Strategy diff --git a/.claude/skills/jaxsr/guides/basis-library.md b/.claude/skills/jaxsr/guides/basis-library.md index 031b625..cba7f24 100644 --- a/.claude/skills/jaxsr/guides/basis-library.md +++ b/.claude/skills/jaxsr/guides/basis-library.md @@ -172,6 +172,69 @@ library.add_parametric( - Parametric fits are slower (requires nonlinear optimization per candidate). - Only add parametric terms when you have physical motivation. +### `add_block(library, multiply_by=None, block_name=None, complexity_offset=0, feature_map=None)` + +Adds every function of another library, optionally multiplied by a column of the +data. This builds design-matrix blocks of the form `Θ(a) ⊙ b`, where `Θ` is a +basis over one variable and `b` is another *column* — typically a measured or +estimated derivative. A coefficient selected in such a block is literally a term +of the unknown coefficient *function* multiplying `b`. + +```python +from jaxsr import BasisLibrary + +# Basis over the coefficient function's argument +theta = (BasisLibrary(n_features=1, feature_names=["c"]) + .add_constant() + .add_linear() + .add_polynomials(max_degree=2)) + +# y_c = s'(c)*y_x + v'(c): one block per unknown function +library = (BasisLibrary(n_features=2, feature_names=["c", "y_x"]) + .add_block(theta, multiply_by="y_x", block_name="horizontal") + .add_block(theta, block_name="vertical")) + +library.names +# ['y_x', 'c*y_x', 'c^2*y_x', '1', 'c', 'c^2'] +``` + +**What it does for you:** + +- **Names** are generated as `*`, consistently; the constant term + collapses to just the column name (`1*y_x` → `y_x`), and a source name that is + a bare sum is parenthesized (`1+c` → `(1+c)*y_x`). +- **Complexity** is inherited from the source, plus 1 for the multiplication, + plus `complexity_offset`. +- **Feature indices** are remapped: the source is written against its own + columns, and `add_block` re-expresses it on this library's feature space, + matching features by name (use `feature_map={"src": "target"}` when the names + differ). +- **Parametric terms pass through unchanged** — bounds, `log_scale` and the name + template are preserved, so profile-likelihood optimization still applies inside + the block. +- The source library is **copied, not shared**, so one `theta` can seed several + blocks. + +**Working with blocks:** + +```python +library.blocks +# {'horizontal': [0, 1, 2], 'vertical': [3, 4, 5]} + +library.filter_by_block(include="horizontal") # -> [0, 1, 2] +library.filter_by_block(exclude=["vertical"]) # -> [0, 1, 2] + +# First diagnostic for a structured library: did the block earn its place? +reduced = library.without_blocks("vertical") # new library, original untouched +``` + +Comparing the fit of `library` against `library.without_blocks("vertical")` is +the fastest way to check a horizontal/vertical identifiability trade-off, which +is a real hazard whenever two blocks can explain the same variation. + +Block functions are not deserializable — like `add_custom`, the library config +saves but the block must be re-added after `load()`. + ### `add_categorical_indicators(features=None)` For categorical features, adds binary indicator (dummy) variables. diff --git a/src/jaxsr/basis.py b/src/jaxsr/basis.py index c74cea3..97605f0 100644 --- a/src/jaxsr/basis.py +++ b/src/jaxsr/basis.py @@ -10,7 +10,7 @@ import itertools import json from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import partial from typing import Any @@ -36,9 +36,12 @@ class BasisFunction: Indices of features used by this basis function. func_type : str Type of function (for serialization): "constant", "linear", "polynomial", - "interaction", "transcendental", "ratio", "custom". + "interaction", "transcendental", "ratio", "block", "custom". func_config : dict Configuration for reconstructing the function (for serialization). + block : str or None + Label of the structured block this function belongs to, set by + :meth:`BasisLibrary.add_block`. ``None`` for unlabelled functions. """ name: str @@ -47,6 +50,7 @@ class BasisFunction: feature_indices: tuple[int, ...] = () func_type: str = "custom" func_config: dict[str, Any] = field(default_factory=dict) + block: str | None = None def evaluate(self, X: jnp.ndarray) -> jnp.ndarray: """Evaluate the basis function on input data.""" @@ -54,13 +58,16 @@ def evaluate(self, X: jnp.ndarray) -> jnp.ndarray: def to_dict(self) -> dict[str, Any]: """Serialize to dictionary (excluding func).""" - return { + d = { "name": self.name, "complexity": self.complexity, "feature_indices": self.feature_indices, "func_type": self.func_type, "func_config": self.func_config, } + if self.block is not None: + d["block"] = self.block + return d @dataclass @@ -136,6 +143,126 @@ def _safe_div(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: return jnp.where(jnp.abs(y) > 1e-10, x / y, jnp.nan) +def _make_block_func( + func: Callable, + gather: jnp.ndarray | None, + multiply_index: int | None, +) -> Callable[[jnp.ndarray], jnp.ndarray]: + """ + Re-express a basis function of a source library on the current feature space. + + Parameters + ---------- + func : callable + ``func(X_source) -> array`` written against the source library's columns. + gather : jnp.ndarray or None + Column indices of the current feature space that reproduce the source + feature space, or None when the two coincide. + multiply_index : int or None + Column of the current feature space to multiply the result by, or None. + + Returns + ------- + wrapped : callable + Function ``(X) -> array`` of shape ``(n_samples,)``. + """ + if gather is None and multiply_index is None: + return func + + def wrapped(X: jnp.ndarray) -> jnp.ndarray: + values = func(X if gather is None else X[:, gather]) + if multiply_index is None: + return values + return values * X[:, multiply_index] + + return wrapped + + +def _make_block_parametric_func( + func: Callable, + gather: jnp.ndarray | None, + multiply_index: int | None, +) -> Callable: + """ + Parametric counterpart of :func:`_make_block_func`. + + Parameters + ---------- + func : callable + ``func(X_source, **params) -> array`` written against the source + library's columns. + gather : jnp.ndarray or None + Column indices of the current feature space that reproduce the source + feature space, or None when the two coincide. + multiply_index : int or None + Column of the current feature space to multiply the result by, or None. + + Returns + ------- + wrapped : callable + Function ``(X, **params) -> array`` of shape ``(n_samples,)``. + """ + + def wrapped(X: jnp.ndarray, **params: float) -> jnp.ndarray: + values = func(X if gather is None else X[:, gather], **params) + if multiply_index is None: + return values + return values * X[:, multiply_index] + + return wrapped + + +def _parenthesize(name: str) -> str: + """ + Wrap a basis-function name in parentheses if it reads as a bare sum. + + ``"q^2"`` is safe to concatenate into ``"q^2*y_x"``, but ``"1+q"`` is not -- + ``"1+q*y_x"`` would name a different function. + + Parameters + ---------- + name : str + Basis function name. + + Returns + ------- + name : str + The name, parenthesised if it contains a top-level ``+`` or ``-``. + """ + depth = 0 + for i, ch in enumerate(name): + if ch in "([{": + depth += 1 + elif ch in ")]}": + depth -= 1 + elif ch in "+-" and depth == 0 and i > 0 and name[i - 1] not in "^*/+-eE(": + return f"({name})" + return name + + +def _block_term_name(basis_name: str, column_name: str | None) -> str: + """ + Name a block term: the source basis name times the multiplying column. + + Parameters + ---------- + basis_name : str + Name of the basis function in the source library. + column_name : str or None + Name of the multiplying feature, or None for an unmultiplied block. + + Returns + ------- + name : str + Generated name, e.g. ``"q^2"`` and ``"y_x"`` give ``"q^2*y_x"``. + """ + if column_name is None: + return basis_name + if basis_name == "1": + return column_name + return f"{_parenthesize(basis_name)}*{column_name}" + + class BasisLibrary: """ Library of candidate basis functions for symbolic regression. @@ -451,6 +578,296 @@ def add_custom( self._compiled_evaluate = None return self + # ------------------------------------------------------------------ + # Structured blocks + # ------------------------------------------------------------------ + + def _resolve_block_columns( + self, + library: BasisLibrary, + feature_map: dict[str, str] | None, + ) -> list[int]: + """Map each source feature onto a column of this library.""" + columns = [] + for name in library.feature_names: + target = (feature_map or {}).get(name, name) + if target not in self.feature_names: + raise ValueError( + f"Source feature '{name}' has no counterpart in this library " + f"(features: {self.feature_names}). Rename it, or pass " + f"feature_map={{'{name}': ''}}." + ) + columns.append(self.feature_names.index(target)) + return columns + + def _resolve_feature_index(self, feature: str | int, argument: str) -> int: + """Resolve a feature given by name or index to a column index.""" + if isinstance(feature, str): + if feature not in self.feature_names: + raise ValueError( + f"{argument}='{feature}' is not a feature of this library " + f"(features: {self.feature_names})" + ) + return self.feature_names.index(feature) + if isinstance(feature, (int, np.integer)) and not isinstance(feature, bool): + index = int(feature) + if not 0 <= index < self.n_features: + raise ValueError( + f"{argument}={index} is out of range for {self.n_features} features" + ) + return index + raise TypeError(f"{argument} must be a feature name or index, got {type(feature).__name__}") + + def add_block( + self, + library: BasisLibrary, + multiply_by: str | int | None = None, + block_name: str | None = None, + complexity_offset: int = 0, + feature_map: dict[str, str] | None = None, + ) -> BasisLibrary: + r""" + Add every function of another library, optionally times a data column. + + This builds design-matrix blocks of the form :math:`\Theta(a) \odot b`, + where ``\Theta`` is a basis over one (or a few) variables and ``b`` is + another column of the data -- typically a measured or estimated + derivative. A coefficient selected in such a block is then literally a + term of the unknown coefficient *function* multiplying ``b``. + + Parameters + ---------- + library : BasisLibrary + Source library. Its features are matched to this library's features + by name (see ``feature_map``); its basis functions are copied, not + shared, so the source stays reusable across blocks. + multiply_by : str or int, optional + Feature of *this* library (name or index) to multiply every function + of the block by. If None, the block is added unmultiplied. + block_name : str, optional + Label recorded on every function of the block. Blocks are reported + by :attr:`blocks` and can be selected with :meth:`filter_by_block` + or dropped with :meth:`without_blocks`. + complexity_offset : int + Added to every inherited complexity score. Multiplying by a column + already costs 1 on its own. + feature_map : dict, optional + ``{source_feature_name: target_feature_name}`` for source features + whose names differ from this library's. Unlisted features are + matched by name. + + Returns + ------- + self : BasisLibrary + For method chaining. + + Raises + ------ + ValueError + If the source library is empty, a source feature has no counterpart + in this library, or ``multiply_by`` names an unknown feature. + TypeError + If ``library`` is not a BasisLibrary, or ``multiply_by`` is neither + a feature name nor an index. + + Notes + ----- + Names are generated as ``"*"``, with the constant term + collapsing to just ``""``. Parametric basis functions are + carried over as parametric: their bounds, log-scale flag and name + template pass through unchanged, so profile-likelihood optimisation + still applies inside the block. + + Like custom functions, block functions cannot be deserialized -- the + library config saves, but the block must be re-added after loading. + + Examples + -------- + >>> theta = (BasisLibrary(n_features=1, feature_names=["q"]) + ... .add_constant().add_linear().add_polynomials(max_degree=2)) + >>> library = (BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + ... .add_block(theta, multiply_by="y_x", block_name="horizontal") + ... .add_block(theta, block_name="vertical")) + >>> library.names[:4] + ['y_x', 'q*y_x', 'q^2*y_x', '1'] + >>> sorted(library.blocks) + ['horizontal', 'vertical'] + """ + if not isinstance(library, BasisLibrary): + raise TypeError(f"library must be a BasisLibrary, got {type(library).__name__}") + if len(library.basis_functions) == 0: + raise ValueError("Source library has no basis functions to add") + + columns = self._resolve_block_columns(library, feature_map) + multiply_index = ( + None if multiply_by is None else self._resolve_feature_index(multiply_by, "multiply_by") + ) + multiply_name = None if multiply_index is None else self.feature_names[multiply_index] + + # Skip the gather when the source feature space is this one, so that the + # copied functions stay as cheap (and as jit-friendly) as the originals. + gather = None if columns == list(range(self.n_features)) else jnp.asarray(columns) + extra_complexity = complexity_offset + (0 if multiply_index is None else 1) + + parametric_by_index = {p.basis_index: p for p in library._parametric_info} + + for source_index, bf in enumerate(library.basis_functions): + feature_indices = {columns[i] for i in bf.feature_indices} + if multiply_index is not None: + feature_indices.add(multiply_index) + feature_indices = tuple(sorted(feature_indices)) + complexity = bf.complexity + extra_complexity + + p_info = parametric_by_index.get(source_index) + if p_info is not None: + self.add_parametric( + name=_block_term_name(p_info.name, multiply_name), + func=_make_block_parametric_func(p_info.func, gather, multiply_index), + param_bounds=dict(p_info.param_bounds), + complexity=complexity, + feature_indices=feature_indices, + log_scale=p_info.log_scale, + ) + self.basis_functions[-1].block = block_name + continue + + name = _block_term_name(bf.name, multiply_name) + self.basis_functions.append( + BasisFunction( + name=name, + func=_make_block_func(bf.func, gather, multiply_index), + complexity=complexity, + feature_indices=feature_indices, + func_type="block", + func_config={ + "name": name, + "source_name": bf.name, + "source_type": bf.func_type, + "multiply_by": multiply_name, + "block": block_name, + }, + block=block_name, + ) + ) + + self._compiled_evaluate = None + return self + + @property + def blocks(self) -> dict[str, list[int]]: + """Mapping from block label to the indices of that block's functions.""" + blocks: dict[str, list[int]] = {} + for i, bf in enumerate(self.basis_functions): + if bf.block is not None: + blocks.setdefault(bf.block, []).append(i) + return blocks + + def filter_by_block( + self, + include: str | list[str] | None = None, + exclude: str | list[str] | None = None, + ) -> list[int]: + """ + Get indices of basis functions by block membership. + + Parameters + ---------- + include : str or list of str, optional + Only keep functions in these blocks. Unlabelled functions are + dropped when this is given. + exclude : str or list of str, optional + Drop functions in these blocks. + + Returns + ------- + indices : list of int + Indices of basis functions meeting the criteria. + + Raises + ------ + ValueError + If a named block is not present in the library. + """ + blocks = set(self.blocks) + + def _as_set(value: str | list[str] | None) -> set[str] | None: + if value is None: + return None + names = {value} if isinstance(value, str) else set(value) + unknown = names - blocks + if unknown: + raise ValueError(f"Unknown block(s) {sorted(unknown)}. Available: {sorted(blocks)}") + return names + + included = _as_set(include) + excluded = _as_set(exclude) or set() + + indices = [] + for i, bf in enumerate(self.basis_functions): + if included is not None and bf.block not in included: + continue + if bf.block in excluded: + continue + indices.append(i) + return indices + + def without_blocks(self, *block_names: str) -> BasisLibrary: + """ + Return a copy of this library with whole blocks removed. + + Dropping a block and refitting is the first diagnostic for a structured + library: it answers whether the block earned its place at all. + + Parameters + ---------- + *block_names : str + Labels of the blocks to drop. + + Returns + ------- + library : BasisLibrary + New library holding copies of the remaining basis functions, with + parametric bookkeeping re-indexed. The original is unchanged. + + Raises + ------ + ValueError + If a named block is not present in the library. + + Examples + -------- + >>> reduced = library.without_blocks("vertical") # doctest: +SKIP + """ + keep = set(self.filter_by_block(exclude=list(block_names))) + + reduced = BasisLibrary( + n_features=self.n_features, + feature_names=list(self.feature_names), + feature_bounds=self.feature_bounds, + feature_types=list(self.feature_types), + categories={k: list(v) for k, v in self.categories.items()} or None, + ) + + old_to_new = {} + for i, bf in enumerate(self.basis_functions): + if i not in keep: + continue + old_to_new[i] = len(reduced.basis_functions) + reduced.basis_functions.append(replace(bf, func_config=dict(bf.func_config))) + + for p_info in self._parametric_info: + if p_info.basis_index in old_to_new: + reduced._parametric_info.append( + replace( + p_info, + basis_index=old_to_new[p_info.basis_index], + param_bounds=dict(p_info.param_bounds), + initial_params=dict(p_info.initial_params), + ) + ) + + return reduced + # ------------------------------------------------------------------ # Categorical basis functions # ------------------------------------------------------------------ @@ -1601,6 +2018,11 @@ def from_dict(cls, config: dict[str, Any]) -> BasisLibrary: f"Cannot deserialize parametric function '{bf_config['name']}'. " "Re-add it manually using add_parametric()." ) + elif func_type == "block": + raise ValueError( + f"Cannot deserialize block function '{bf_config['name']}'. " + "Re-add the block manually using add_block()." + ) elif func_type == "custom": raise ValueError( f"Cannot deserialize custom function '{bf_config['name']}'. " diff --git a/src/jaxsr/skill/SKILL.md b/src/jaxsr/skill/SKILL.md index 5b1c908..1956e0f 100644 --- a/src/jaxsr/skill/SKILL.md +++ b/src/jaxsr/skill/SKILL.md @@ -43,12 +43,15 @@ Based on the answers, recommend a basis library configuration: | Large feature space (screening) | `add_constant + add_linear + add_interactions(2)` then use `lasso_path` strategy | | Response surface (DOE) | `add_constant + add_linear + add_polynomials(2) + add_interactions(2)` — or use `ResponseSurface` directly | | Categorical factors present | Add `add_categorical_indicators() + add_categorical_interactions()` to any of the above | +| Unknown coefficient *function* multiplying a data column (superposition, implicit dynamics) | `add_block(theta, multiply_by="", block_name=...)` — see `guides/basis-library.md` | **Key guidance:** - Start simple. You can always add complexity. - `add_transcendental(safe=True)` guards against log(0), 1/0, sqrt(<0). Always use `safe=True`. - `add_ratios(safe=True)` adds x_i/x_j terms. Doubles the library size — only use when ratios are physically meaningful. - `add_parametric()` enables nonlinear parameters (e.g., `exp(-a*x)`). Powerful but slower to fit. +- `add_block()` multiplies a whole basis by a data column, so a selected coefficient is a term of an + unknown coefficient function. Drop a block with `without_blocks()` to test whether it earned its place. - If n_features > 5, avoid `add_polynomials(degree>2)` — the library becomes enormous. ### Step 3: Recommend Selection Strategy diff --git a/src/jaxsr/skill/guides/basis-library.md b/src/jaxsr/skill/guides/basis-library.md index 031b625..cba7f24 100644 --- a/src/jaxsr/skill/guides/basis-library.md +++ b/src/jaxsr/skill/guides/basis-library.md @@ -172,6 +172,69 @@ library.add_parametric( - Parametric fits are slower (requires nonlinear optimization per candidate). - Only add parametric terms when you have physical motivation. +### `add_block(library, multiply_by=None, block_name=None, complexity_offset=0, feature_map=None)` + +Adds every function of another library, optionally multiplied by a column of the +data. This builds design-matrix blocks of the form `Θ(a) ⊙ b`, where `Θ` is a +basis over one variable and `b` is another *column* — typically a measured or +estimated derivative. A coefficient selected in such a block is literally a term +of the unknown coefficient *function* multiplying `b`. + +```python +from jaxsr import BasisLibrary + +# Basis over the coefficient function's argument +theta = (BasisLibrary(n_features=1, feature_names=["c"]) + .add_constant() + .add_linear() + .add_polynomials(max_degree=2)) + +# y_c = s'(c)*y_x + v'(c): one block per unknown function +library = (BasisLibrary(n_features=2, feature_names=["c", "y_x"]) + .add_block(theta, multiply_by="y_x", block_name="horizontal") + .add_block(theta, block_name="vertical")) + +library.names +# ['y_x', 'c*y_x', 'c^2*y_x', '1', 'c', 'c^2'] +``` + +**What it does for you:** + +- **Names** are generated as `*`, consistently; the constant term + collapses to just the column name (`1*y_x` → `y_x`), and a source name that is + a bare sum is parenthesized (`1+c` → `(1+c)*y_x`). +- **Complexity** is inherited from the source, plus 1 for the multiplication, + plus `complexity_offset`. +- **Feature indices** are remapped: the source is written against its own + columns, and `add_block` re-expresses it on this library's feature space, + matching features by name (use `feature_map={"src": "target"}` when the names + differ). +- **Parametric terms pass through unchanged** — bounds, `log_scale` and the name + template are preserved, so profile-likelihood optimization still applies inside + the block. +- The source library is **copied, not shared**, so one `theta` can seed several + blocks. + +**Working with blocks:** + +```python +library.blocks +# {'horizontal': [0, 1, 2], 'vertical': [3, 4, 5]} + +library.filter_by_block(include="horizontal") # -> [0, 1, 2] +library.filter_by_block(exclude=["vertical"]) # -> [0, 1, 2] + +# First diagnostic for a structured library: did the block earn its place? +reduced = library.without_blocks("vertical") # new library, original untouched +``` + +Comparing the fit of `library` against `library.without_blocks("vertical")` is +the fastest way to check a horizontal/vertical identifiability trade-off, which +is a real hazard whenever two blocks can explain the same variation. + +Block functions are not deserializable — like `add_custom`, the library config +saves but the block must be re-added after `load()`. + ### `add_categorical_indicators(features=None)` For categorical features, adds binary indicator (dummy) variables. diff --git a/tests/test_basis.py b/tests/test_basis.py index fce2e55..a6e0415 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -397,3 +397,309 @@ def test_pole_bearing_composition_is_excluded(self): assert "exp(x0*x1)" in model.selected_features_ assert np.all(np.isfinite(np.asarray(model.predict(X)))) + + +class TestAddBlock: + """Tests for structured basis blocks (Theta(a) times a data column).""" + + @pytest.fixture + def theta(self): + """A small basis over a single variable q.""" + return ( + BasisLibrary(n_features=1, feature_names=["q"]) + .add_constant() + .add_linear() + .add_polynomials(max_degree=2) + ) + + @pytest.fixture + def X(self): + """Data over (q, y_x).""" + rng = np.random.default_rng(0) + return jnp.array(rng.uniform(0.5, 2.0, size=(20, 2))) + + def test_names_are_generated(self, theta): + """Block names are the source names times the multiplying column.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + assert library.names == ["y_x", "q*y_x", "q^2*y_x"] + + def test_unmultiplied_block_keeps_source_names(self, theta): + """Without multiply_by the block is a plain copy of the source.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block(theta) + assert library.names == theta.names + + def test_evaluates_as_product(self, theta, X): + """Each column is the source column times the data column.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + Phi = np.asarray(library.evaluate(X)) + q, y_x = np.asarray(X[:, 0]), np.asarray(X[:, 1]) + np.testing.assert_allclose(Phi[:, 0], y_x, rtol=1e-6) + np.testing.assert_allclose(Phi[:, 1], q * y_x, rtol=1e-6) + np.testing.assert_allclose(Phi[:, 2], q**2 * y_x, rtol=1e-6) + + def test_multiply_by_index(self, theta, X): + """multiply_by accepts a column index as well as a name.""" + by_name = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + by_index = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by=1 + ) + assert by_index.names == by_name.names + np.testing.assert_allclose( + np.asarray(by_index.evaluate(X)), np.asarray(by_name.evaluate(X)), rtol=1e-6 + ) + + def test_features_are_remapped(self, X): + """A source written against its own columns is re-expressed on ours.""" + theta = BasisLibrary(n_features=1, feature_names=["y_x"]).add_linear() + # 'y_x' is column 0 of the source but column 1 here + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block(theta) + Phi = np.asarray(library.evaluate(X)) + np.testing.assert_allclose(Phi[:, 0], np.asarray(X[:, 1]), rtol=1e-6) + + def test_complexity_inherited_plus_one_for_the_product(self, theta): + """Multiplying by a column costs one, on top of the inherited score.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + assert list(np.asarray(library.complexities)) == [1, 2, 3] + + def test_complexity_offset(self, theta): + """complexity_offset shifts the whole block.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, complexity_offset=2 + ) + assert list(np.asarray(library.complexities)) == [2, 3, 4] + + def test_feature_indices_include_the_multiplier(self, theta): + """Feature indices are mapped to this library and include the column.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + assert library.basis_functions[0].feature_indices == (1,) + assert library.basis_functions[1].feature_indices == (0, 1) + + def test_source_library_is_reusable(self, theta, X): + """Adding a block copies functions; the source is untouched.""" + library = ( + BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + .add_block(theta, multiply_by="y_x", block_name="horizontal") + .add_block(theta, block_name="vertical") + ) + assert len(theta) == 3 + assert theta.names == ["1", "q", "q^2"] + assert len(library) == 6 + # The unmultiplied half still evaluates on q alone + Phi = np.asarray(library.evaluate(X)) + np.testing.assert_allclose(Phi[:, 4], np.asarray(X[:, 0]), rtol=1e-6) + + def test_feature_map(self, X): + """feature_map matches a source feature whose name differs.""" + theta = BasisLibrary(n_features=1, feature_names=["x"]).add_linear() + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x", feature_map={"x": "q"} + ) + assert library.names == ["x*y_x"] + Phi = np.asarray(library.evaluate(X)) + np.testing.assert_allclose(Phi[:, 0], np.asarray(X[:, 0] * X[:, 1]), rtol=1e-6) + + def test_sum_names_are_parenthesized(self, X): + """A bare sum is parenthesized so the generated name stays correct.""" + theta = BasisLibrary(n_features=1, feature_names=["q"]).add_custom( + "1+q", lambda Xs: 1.0 + Xs[:, 0], complexity=2, feature_indices=(0,) + ) + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + assert library.names == ["(1+q)*y_x"] + Phi = np.asarray(library.evaluate(X)) + np.testing.assert_allclose(Phi[:, 0], np.asarray((1.0 + X[:, 0]) * X[:, 1]), rtol=1e-6) + + def test_custom_source_functions_are_carried_over(self, X): + """add_custom terms in the source work in the block.""" + theta = BasisLibrary(n_features=1, feature_names=["q"]).add_custom( + "1/(1+q)^2", lambda Xs: 1.0 / (1.0 + Xs[:, 0]) ** 2, complexity=3, feature_indices=(0,) + ) + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x" + ) + assert library.names == ["1/(1+q)^2*y_x"] + assert library.basis_functions[0].complexity == 4 + Phi = np.asarray(library.evaluate(X)) + np.testing.assert_allclose(Phi[:, 0], np.asarray(X[:, 1] / (1.0 + X[:, 0]) ** 2), rtol=1e-6) + + def test_parametric_passes_through(self, X): + """Parametric source terms stay parametric inside the block.""" + theta = BasisLibrary(n_features=1, feature_names=["q"]).add_parametric( + name="1/(c2+q)^2", + func=lambda Xs, c2: 1.0 / (c2 + Xs[:, 0]) ** 2, + param_bounds={"c2": (0.2, 0.8)}, + complexity=3, + feature_indices=(0,), + ) + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x", block_name="horizontal" + ) + assert library.has_parametric + info = library._parametric_info[0] + assert info.name == "1/(c2+q)^2*y_x" + assert info.param_bounds == {"c2": (0.2, 0.8)} + assert info.basis_index == 0 + assert library.basis_functions[0].block == "horizontal" + assert library.basis_functions[0].complexity == 4 + # The registered function multiplies by the data column + got = np.asarray(info.func(X, c2=0.5)) + np.testing.assert_allclose(got, np.asarray(X[:, 1] / (0.5 + X[:, 0]) ** 2), rtol=1e-6) + + def test_blocks_property(self, theta): + """Blocks report the indices they own.""" + library = ( + BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + .add_block(theta, multiply_by="y_x", block_name="horizontal") + .add_block(theta, block_name="vertical") + ) + assert library.blocks == {"horizontal": [0, 1, 2], "vertical": [3, 4, 5]} + + def test_unlabelled_functions_have_no_block(self, theta): + """Functions added outside a block are not part of one.""" + library = ( + BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + .add_constant() + .add_block(theta, multiply_by="y_x", block_name="horizontal") + ) + assert library.blocks == {"horizontal": [1, 2, 3]} + assert library.basis_functions[0].block is None + + def test_filter_by_block(self, theta): + """filter_by_block selects and drops whole blocks.""" + library = ( + BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + .add_constant() + .add_block(theta, multiply_by="y_x", block_name="horizontal") + .add_block(theta, block_name="vertical") + ) + assert library.filter_by_block(include="horizontal") == [1, 2, 3] + assert library.filter_by_block(exclude=["vertical"]) == [0, 1, 2, 3] + assert library.filter_by_block() == list(range(7)) + + def test_filter_by_block_unknown_name(self, theta): + """A typo in a block name is an error, not an empty result.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, block_name="vertical" + ) + with pytest.raises(ValueError, match="Unknown block"): + library.filter_by_block(include="verticle") + + def test_without_blocks(self, theta, X): + """Dropping a block leaves a working library and the original intact.""" + library = ( + BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + .add_block(theta, multiply_by="y_x", block_name="horizontal") + .add_block(theta, block_name="vertical") + ) + reduced = library.without_blocks("vertical") + + assert len(library) == 6 + assert reduced.names == ["y_x", "q*y_x", "q^2*y_x"] + assert reduced.blocks == {"horizontal": [0, 1, 2]} + np.testing.assert_allclose( + np.asarray(reduced.evaluate(X)), np.asarray(library.evaluate(X))[:, :3], rtol=1e-6 + ) + + def test_without_blocks_reindexes_parametric(self, X): + """Parametric bookkeeping follows the surviving functions.""" + theta = BasisLibrary(n_features=1, feature_names=["q"]).add_parametric( + name="1/(c2+q)^2", + func=lambda Xs, c2: 1.0 / (c2 + Xs[:, 0]) ** 2, + param_bounds={"c2": (0.2, 0.8)}, + feature_indices=(0,), + ) + library = ( + BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + .add_block(theta, block_name="vertical") + .add_block(theta, multiply_by="y_x", block_name="horizontal") + ) + assert [p.basis_index for p in library._parametric_info] == [0, 1] + + reduced = library.without_blocks("vertical") + assert len(reduced) == 1 + assert [p.basis_index for p in reduced._parametric_info] == [0] + assert reduced._parametric_info[0].name == "1/(c2+q)^2*y_x" + + def test_empty_source_raises(self): + """An empty source library is a mistake worth reporting.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + with pytest.raises(ValueError, match="no basis functions"): + library.add_block(BasisLibrary(n_features=1, feature_names=["q"])) + + def test_unmatched_source_feature_raises(self, theta): + """A source feature with no counterpart names the fix.""" + library = BasisLibrary(n_features=2, feature_names=["c", "y_x"]) + with pytest.raises(ValueError, match="feature_map"): + library.add_block(theta, multiply_by="y_x") + + def test_unknown_multiply_by_raises(self, theta): + """multiply_by must name a feature of this library.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + with pytest.raises(ValueError, match="multiply_by"): + library.add_block(theta, multiply_by="y_c") + + def test_out_of_range_multiply_by_raises(self, theta): + """An index past the end of the feature space is an error.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + with pytest.raises(ValueError, match="out of range"): + library.add_block(theta, multiply_by=5) + + def test_bad_multiply_by_type_raises(self, theta): + """multiply_by is a name or an index, nothing else.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + with pytest.raises(TypeError, match="multiply_by"): + library.add_block(theta, multiply_by=1.5) + + def test_bad_source_type_raises(self): + """The source must be a BasisLibrary.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]) + with pytest.raises(TypeError, match="BasisLibrary"): + library.add_block(["1", "q"]) + + def test_block_survives_to_dict(self, theta): + """The block label is serialized; the function itself is not.""" + library = BasisLibrary(n_features=2, feature_names=["q", "y_x"]).add_block( + theta, multiply_by="y_x", block_name="horizontal" + ) + d = library.to_dict() + assert d["basis_functions"][0]["block"] == "horizontal" + with pytest.raises(ValueError, match="add_block"): + BasisLibrary.from_dict(d) + + def test_recovers_a_coefficient_function(self): + """The motivating case: y_c = s'(c)*y_x with s'(c) = 1 + 2c.""" + from jaxsr import SymbolicRegressor + + rng = np.random.default_rng(0) + c = rng.uniform(0.1, 1.0, size=200) + y_x = rng.uniform(0.5, 2.0, size=200) + y_c = (1.0 + 2.0 * c) * y_x + + theta = ( + BasisLibrary(n_features=1, feature_names=["c"]) + .add_constant() + .add_linear() + .add_polynomials(max_degree=3) + ) + library = BasisLibrary(n_features=2, feature_names=["c", "y_x"]).add_block( + theta, multiply_by="y_x", block_name="horizontal" + ) + + X = jnp.array(np.column_stack([c, y_x])) + model = SymbolicRegressor(basis_library=library, max_terms=4).fit(X, jnp.array(y_c)) + + assert set(model.selected_features_) == {"y_x", "c*y_x"} + coefs = dict(zip(model.selected_features_, np.asarray(model.coefficients_), strict=False)) + np.testing.assert_allclose(coefs["y_x"], 1.0, atol=1e-4) + np.testing.assert_allclose(coefs["c*y_x"], 2.0, atol=1e-4)