jaxsr.basis

jaxsr.basis#

Basis Function Library for JAXSR.

Provides a flexible, extensible system for defining candidate basis functions for symbolic regression.

class jaxsr.basis.BasisFunction(name: str, func: ~collections.abc.Callable[[~jax.jaxlib._jax.Array], ~jax.jaxlib._jax.Array], complexity: int = 1, feature_indices: tuple[int, ...] = (), func_type: str = 'custom', func_config: dict[str, ~typing.Any] = <factory>, block: str | None = None)#

Bases: object

A single basis function with metadata.

Parameters:

name (str) – Human-readable name for the basis function.
func (Callable) – Function that takes X of shape (n_samples, n_features) and returns array of shape (n_samples,).
complexity (int) – Complexity score for Pareto optimization (higher = more complex).
feature_indices (tuple) – Indices of features used by this basis function.
func_type (str) – Type of function (for serialization): “constant”, “linear”, “polynomial”, “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 BasisLibrary.add_block(). None for unlabelled functions.

name: str#

func: Callable[[Array], Array]#

complexity: int = 1#

feature_indices: tuple[int, ...] = ()#

func_type: str = 'custom'#

func_config: dict[str, Any]#

block: str | None = None#

evaluate(X: Array) → Array#: Evaluate the basis function on input data.

to_dict() → dict[str, Any]#: Serialize to dictionary (excluding func).

class jaxsr.basis.ParametricBasisInfo(basis_index: int, name: str, func: Callable, param_bounds: dict[str, tuple[float, float]], initial_params: dict[str, float], log_scale: bool = False, resolved_params: dict[str, float] | None = None)#

Bases: object

Metadata for a parametric basis function with free nonlinear parameters.

Parameters:

basis_index (int) – Index of the corresponding BasisFunction in the library.
name (str) – Name template with parameter symbols, e.g. "exp(-a*x)".
func (Callable) – Function (X, **params) -> array of shape (n_samples,).
param_bounds (dict) – {param_name: (lower, upper)} bounds for optimisation.
initial_params (dict) – {param_name: initial_value} midpoint (or geometric mean) guesses.
log_scale (bool) – If True, optimise in log-space (useful when the parameter spans orders of magnitude).
resolved_params (dict or None) – Optimised parameter values, set after fitting.

basis_index: int#

name: str#

func: Callable#

param_bounds: dict[str, tuple[float, float]]#

initial_params: dict[str, float]#

log_scale: bool = False#

resolved_params: dict[str, float] | None = None#

class jaxsr.basis.BasisLibrary(n_features: int, feature_names: list[str] | None = None, feature_bounds: list[tuple[float, float]] | None = None, feature_types: list[str] | None = None, categories: dict[int, list] | None = None)#

Bases: object

Library of candidate basis functions for symbolic regression.

Generates basis functions up to specified complexity from input features. Supports method chaining for convenient library construction.

Parameters:

n_features (int) – Number of input features.
feature_names (list of str, optional) – Names for each feature. Defaults to [“x0”, “x1”, …].
feature_bounds (list of tuple, optional) – Bounds (lower, upper) for each feature. Used for constraint-aware basis generation and adaptive sampling.

Examples

>>> library = (BasisLibrary(n_features=2, feature_names=["x", "y"])
...     .add_constant()
...     .add_linear()
...     .add_polynomials(max_degree=3)
...     .add_interactions(max_order=2)
...     .add_transcendental(["log", "exp"])
... )
>>> Phi = library.evaluate(X)

property continuous_indices: list[int]#: Indices of continuous features.

property categorical_indices: list[int]#: Indices of categorical features.

add_constant() → BasisLibrary#: Add constant (intercept) term.

add_linear() → BasisLibrary#: Add linear terms: x_i for each continuous feature.

add_polynomials(max_degree: int = 2) → BasisLibrary#

Add polynomial terms: x_i^d for d in 2..max_degree.

Parameters:: max_degree (int) – Maximum polynomial degree (default 2).

add_interactions(max_order: int = 2) → BasisLibrary#

Add interaction terms: products of distinct features.

Parameters:: max_order (int) – Maximum interaction order (default 2 for pairwise).

add_transcendental(funcs: list[str] | None = None, safe: bool = True) → BasisLibrary#

Add transcendental terms: log, exp, sqrt, inv, sin, cos.

Parameters:

funcs (list of str, optional) – Which functions to include. Defaults to [“log”, “exp”, “sqrt”, “inv”].
safe (bool) – If True, use safe versions that return NaN for invalid inputs.

add_ratios(safe: bool = True) → BasisLibrary#

Add ratio terms: x_i / x_j for all distinct pairs.

Parameters:: safe (bool) – If True, use safe division that returns NaN for zero denominator.

add_custom(name: str, func: Callable[[Array], Array], complexity: int = 3, feature_indices: tuple[int, ...] | None = None) → BasisLibrary#

Add a custom basis function.

Parameters:

name (str) – Human-readable name.
func (callable) – Function that takes X of shape (n_samples, n_features) and returns array of shape (n_samples,).
complexity (int) – Complexity score for Pareto optimization.
feature_indices (tuple of int, optional) – Indices of features used by this function.

Notes

Custom functions cannot be serialized automatically. The library can still be saved, but the custom function will need to be re-added after loading.

add_block(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#

Add every function of another library, optionally times a data column.

This builds design-matrix blocks of the form \(\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 blocks and can be selected with filter_by_block() or dropped with 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 – For method chaining.

Return type:

BasisLibrary

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 "<basis>*<column>", with the constant term collapsing to just "<column>". 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']

property blocks: dict[str, list[int]]#: Mapping from block label to the indices of that block’s functions.

filter_by_block(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 – Indices of basis functions meeting the criteria.

Return type:

list of int

Raises:

ValueError – If a named block is not present in the library.

without_blocks(*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 – New library holding copies of the remaining basis functions, with parametric bookkeeping re-indexed. The original is unchanged.
Return type:: BasisLibrary
Raises:: ValueError – If a named block is not present in the library.

Examples

>>> reduced = library.without_blocks("vertical")  

add_categorical_indicators(features: list[int] | None = None) → BasisLibrary#

Add indicator (dummy variable) basis functions for categorical features.

For each categorical feature with K categories, adds K-1 indicator functions using reference encoding (first category dropped) to avoid multicollinearity with the intercept.

Parameters:: features (list of int, optional) – Indices of categorical features to encode. Defaults to all categorical features.
Returns:: self – For method chaining.
Return type:: BasisLibrary
Raises:: ValueError – If a specified feature is not categorical.

add_categorical_interactions(cat_features: list[int] | None = None, cont_features: list[int] | None = None) → BasisLibrary#

Add interactions between categorical indicators and continuous features.

For each (categorical feature, continuous feature) pair, creates indicator * continuous terms. This allows the model to learn different slopes per category.

Parameters:

cat_features (list of int, optional) – Categorical feature indices. Defaults to all categorical.
cont_features (list of int, optional) – Continuous feature indices. Defaults to all continuous.

Returns:

self – For method chaining.

Return type:

BasisLibrary

property has_parametric: bool#: Whether the library contains parametric basis functions.

add_parametric(name: str, func: Callable, param_bounds: dict[str, tuple[float, float]], complexity: int = 3, feature_indices: tuple[int, ...] | None = None, log_scale: bool = False) → BasisLibrary#

Add a parametric basis function with free nonlinear parameters.

The nonlinear parameter(s) are optimised via profile likelihood during model selection: for each candidate value the linear coefficients are solved exactly via OLS.

Parameters:

name (str) – Name with parameter symbols, e.g. "exp(-a*x)".
func (callable) – func(X, **params) -> array of shape (n_samples,).
param_bounds (dict) – {param_name: (lower, upper)} search bounds for each parameter.
complexity (int) – Complexity score for Pareto optimisation.
feature_indices (tuple of int, optional) – Indices of input features used by this function.
log_scale (bool) – If True, search in log-space (useful for parameters that span orders of magnitude).

Returns:

self – For method chaining.

Return type:

BasisLibrary

Examples

>>> library.add_parametric(
...     name="exp(-a*x)",
...     func=lambda X, a: jnp.exp(-a * X[:, 0]),
...     param_bounds={"a": (0.01, 10.0)},
...     feature_indices=(0,),
... )

add_polynomial_interactions(max_total_degree: int = 3, max_individual_degree: int = 2) → BasisLibrary#

Add polynomial terms with mixed powers: x_i^a * x_j^b * …

Parameters:

max_total_degree (int) – Maximum sum of all exponents (default 3).
max_individual_degree (int) – Maximum exponent for any single variable (default 2).

build_default(max_poly_degree: int = 3, include_transcendental: bool = True, include_ratios: bool = False) → BasisLibrary#

Build a default library similar to ALAMO’s standard set.

Parameters:

max_poly_degree (int) – Maximum polynomial degree.
include_transcendental (bool) – Whether to include transcendental functions.
include_ratios (bool) – Whether to include ratio terms.

add_compositions(outer_funcs: list[str] | None = None, inner_forms: list[str] | None = None) → BasisLibrary#

Add compositions: outer_func(inner_form).

Creates functions like log(x*y), exp(x/y), sqrt(x+y), etc.

Parameters:

outer_funcs (list of str, optional) – Outer functions to apply. Defaults to [“log”, “exp”, “sqrt”].
inner_forms (list of str, optional) – Inner forms: “product”, “ratio”, “sum”, “diff”. Defaults to [“product”, “ratio”].

Examples

>>> library.add_compositions(["log", "exp"], ["product", "ratio"])
# Adds: log(x*y), log(x/y), exp(x*y), exp(x/y), etc.

add_rational_forms(numerator_degree: int = 1, denominator_degree: int = 1) → BasisLibrary#

Add rational function templates: (a + b*x) / (1 + c*y).

Common in chemical kinetics (Langmuir-Hinshelwood, Michaelis-Menten).

Parameters:

numerator_degree (int) – Max polynomial degree in numerator.
denominator_degree (int) – Max polynomial degree in denominator.

Examples

>>> library.add_rational_forms()
# Adds: x/(1+y), x*y/(1+x), x/(1+x+y), etc.

add_power_laws(exponents: list[float] | None = None) → BasisLibrary#

Add power law terms with fractional exponents.

Common in empirical correlations (Nu = Re^0.8 * Pr^0.33).

Parameters:: exponents (list of float, optional) – Exponents to use. Defaults to [0.25, 0.33, 0.5, 0.67, 0.75, 1.5, 2.0].

expand_sisso_style(operations: list[str] | None = None, max_depth: int = 2) → BasisLibrary#

SISSO-style recursive feature expansion.

Iteratively applies operations to existing features to create new composite features.

Parameters:

operations (list of str, optional) – Operations to apply: “add”, “sub”, “mul”, “div”, “exp”, “log”, “sqrt”, “sq”, “inv”. Defaults to [“mul”, “div”, “sq”, “sqrt”].
max_depth (int) – Maximum recursion depth.

Notes

This can create a very large library. Use with caution.

evaluate(X: Array) → Array#

Evaluate all basis functions on input data.

Parameters:: X (jnp.ndarray) – Input array of shape (n_samples, n_features).
Returns:: Phi – Design matrix of shape (n_samples, n_basis).
Return type:: jnp.ndarray

evaluate_subset(X: Array, indices: list[int] | Array) → Array#

Evaluate a subset of basis functions.

Parameters:

X (jnp.ndarray) – Input array of shape (n_samples, n_features).
indices (array-like) – Indices of basis functions to evaluate.

Returns:

Phi – Design matrix of shape (n_samples, len(indices)).

Return type:

jnp.ndarray

property names: list[str]#: List of basis function names.

property canonical_names: list[str]#

Basis function names with parametric parameters left as symbols.

Identical to names except for parametric basis functions, which keep the template they were registered with ("exp(-a*x)") instead of the fitted rendering ("exp(-0.4913*x)").

Returns:: names – One name per basis function, in library order.
Return type:: list of str

canonical_name(index: int) → str#

Name identifying a basis function independently of any fitted values.

Fitting a parametric basis rewrites its name to embed the optimised parameter value, so names is not a stable identity across refits: "exp(-a*x)" becomes "exp(-0.4913*x)", and a different value on the next fit. Anything that aggregates across refits (bootstrap stability, selection frequencies) must key on this name instead.

Parameters:: index (int) – Index of the basis function in the library.
Returns:: name – The registered template name for a parametric basis function, otherwise the ordinary basis function name.
Return type:: str
Raises:: IndexError – If index is out of range for the library.

copy() → BasisLibrary#

Return an independent copy of this library.

Basis function callables are shared (they are stateless), but every piece of mutable metadata is duplicated. Fitting a model on a parametric library rewrites basis names and rebinds evaluation closures in place, so a repeated-refit procedure (bootstrap, cross-validation) must fit on a copy or it will leave the caller’s library — and therefore the caller’s fitted model — pinned to the last refit’s parameter values.

Returns:: library – A new library holding the same basis functions in the same order.
Return type:: BasisLibrary

property complexities: Array#: Array of complexity scores.

summary() → str#: Return a summary of the library contents.

to_dict() → dict[str, Any]#

Serialize library configuration to dictionary.

Notes

Custom functions are not fully serializable. The library config is saved, but custom functions will need to be re-added manually after loading.

classmethod from_dict(config: dict[str, Any]) → BasisLibrary#

Reconstruct library from configuration dictionary.

Parameters:: config (dict) – Configuration from to_dict().
Returns:: library – Reconstructed library. Custom functions will raise an error and need to be re-added manually.
Return type:: BasisLibrary

save(filepath: str) → None#

Save library configuration to JSON file.

Parameters:: filepath (str) – Path to save the configuration.

classmethod load(filepath: str) → BasisLibrary#

Load library from JSON file.

Parameters:: filepath (str) – Path to the configuration file.

filter_by_complexity(max_complexity: int | None = None, min_complexity: int | None = None) → list[int]#

Get indices of basis functions within complexity bounds.

Parameters:

max_complexity (int, optional) – Maximum allowed complexity.
min_complexity (int, optional) – Minimum required complexity.

Returns:

indices – Indices of basis functions meeting the criteria.

Return type:

list of int

filter_by_features(required_features: list[int] | None = None, excluded_features: list[int] | None = None) → list[int]#

Get indices of basis functions using specified features.

Parameters:

required_features (list of int, optional) – Only include basis functions using these features.
excluded_features (list of int, optional) – Exclude basis functions using these features.

Returns:

indices – Indices of basis functions meeting the criteria.

Return type:

list of int

jaxsr.basis

Contents

jaxsr.basis#