jaxsr.constraints

jaxsr.constraints#

Physical Constraints for JAXSR.

Provides functionality to incorporate domain knowledge through: - Output bounds - Monotonicity constraints - Convexity constraints - Sign constraints on coefficients - Linear constraints on coefficients

class jaxsr.constraints.ConstraintType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)#

Bases: Enum

Types of constraints supported.

BOUND = 'bound'#
MONOTONIC = 'monotonic'#
CONVEX = 'convex'#
CONCAVE = 'concave'#
SIGN = 'sign'#
LINEAR = 'linear'#
FIXED = 'fixed'#
CUSTOM = 'custom'#
class jaxsr.constraints.Constraint(constraint_type: ~jaxsr.constraints.ConstraintType, target: str, params: dict[str, ~typing.Any] = <factory>, weight: float = 1.0, hard: bool = False)#

Bases: object

Base constraint specification.

Parameters:

  • constraint_type (ConstraintType) – Type of constraint.

  • target (str) – Target of constraint (“y” for output, feature name for input).

  • params (dict) – Constraint-specific parameters.

  • weight (float) – Weight for soft constraint penalty.

  • hard (bool) – If True, enforce as hard constraint.

constraint_type: ConstraintType#
target: str#
params: dict[str, Any]#
weight: float = 1.0#
hard: bool = False#
to_dict() → dict[str, Any]#

Serialize to dictionary.

classmethod from_dict(data: dict[str, Any]) → Constraint#

Deserialize from dictionary.

class jaxsr.constraints.Constraints#

Bases: object

Builder for constraint specifications.

Supports method chaining for convenient constraint construction.

Examples

>>> constraints = (Constraints()
...     .add_bounds("y", lower=0)
...     .add_monotonic("T", direction="increasing")
...     .add_sign_constraint("T", sign="positive")
... )
add_bounds(target: str = 'y', lower: float | None = None, upper: float | None = None, weight: float = 1.0, hard: bool = False) → Constraints#

Add bounds constraint.

Parameters:

  • target (str) – “y” for output bounds, or feature name.

  • lower (float, optional) – Lower bound.

  • upper (float, optional) – Upper bound.

  • weight (float) – Penalty weight for soft constraint.

  • hard (bool) – If True, project to satisfy constraint.

Returns:

self – For method chaining.

Return type:

Constraints

add_monotonic(feature: str, direction: str = 'increasing', weight: float = 1.0, hard: bool = False, grid: int | None = None, grid_range: tuple[float, float] | None = None) → Constraints#

Add monotonicity constraint.

Parameters:

  • feature (str) – Feature name for which output should be monotonic.

  • direction (str) – “increasing” or “decreasing”.

  • weight (float) – Penalty weight.

  • hard (bool) – If True, enforce strictly.

  • grid (int, optional) – Also impose the constraint on grid evenly spaced values of feature (in addition to the training rows). Without it the constraint only holds at the training rows and the fitted curve can violate it between them. With several features, the grid is laid along feature for every distinct combination of the other features in the training data; if that would add more than 1000 rows, those combinations are evenly subsampled (with a warning). enforcement="exact" handles many rows cheaply; "constrained" slows down noticeably as they grow.

  • grid_range (tuple of (float, float), optional) – (low, high) span of the grid. Defaults to the range of feature in the training data; widen it to constrain the model where you intend to extrapolate.

Returns:

self

Return type:

Constraints

Raises:

ValueError – If direction is invalid, grid is not an integer >= 2 or grid_range is not an increasing pair of finite numbers.

add_convex(feature: str, weight: float = 1.0, hard: bool = False, grid: int | None = None, grid_range: tuple[float, float] | None = None) → Constraints#

Add convexity constraint (positive second derivative).

Parameters:

  • feature (str) – Feature name.

  • weight (float) – Penalty weight.

  • hard (bool) – If True, enforce strictly.

  • grid (int, optional) – Also impose the constraint on grid evenly spaced values of feature (in addition to the training rows). Without it the constraint only holds at the training rows and the fitted curve can violate it between them. With several features, the grid is laid along feature for every distinct combination of the other features in the training data; if that would add more than 1000 rows, those combinations are evenly subsampled (with a warning). enforcement="exact" handles many rows cheaply; "constrained" slows down noticeably as they grow.

  • grid_range (tuple of (float, float), optional) – (low, high) span of the grid. Defaults to the range of feature in the training data; widen it to constrain the model where you intend to extrapolate.

Returns:

self

Return type:

Constraints

Raises:

ValueError – If grid is not an integer >= 2 or grid_range is not an increasing pair of finite numbers.

add_concave(feature: str, weight: float = 1.0, hard: bool = False, grid: int | None = None, grid_range: tuple[float, float] | None = None) → Constraints#

Add concavity constraint (negative second derivative).

Parameters:

  • feature (str) – Feature name.

  • weight (float) – Penalty weight.

  • hard (bool) – If True, enforce strictly.

  • grid (int, optional) – Also impose the constraint on grid evenly spaced values of feature (in addition to the training rows). Without it the constraint only holds at the training rows and the fitted curve can violate it between them. With several features, the grid is laid along feature for every distinct combination of the other features in the training data; if that would add more than 1000 rows, those combinations are evenly subsampled (with a warning). enforcement="exact" handles many rows cheaply; "constrained" slows down noticeably as they grow.

  • grid_range (tuple of (float, float), optional) – (low, high) span of the grid. Defaults to the range of feature in the training data; widen it to constrain the model where you intend to extrapolate.

Returns:

self

Return type:

Constraints

Raises:

ValueError – If grid is not an integer >= 2 or grid_range is not an increasing pair of finite numbers.

add_sign_constraint(basis_name: str, sign: str = 'positive', weight: float = 1.0, hard: bool = True) → Constraints#

Add sign constraint on coefficient.

Parameters:

  • basis_name (str) – Name of basis function whose coefficient is constrained.

  • sign (str) – “positive” or “negative”.

  • weight (float) – Penalty weight.

  • hard (bool) – If True, project coefficient to satisfy constraint.

Returns:

self

Return type:

Constraints

add_linear_constraint(A: Array, b: Array, weight: float = 1.0, hard: bool = False) → Constraints#

Add linear constraint: A @ coefficients <= b.

Parameters:

  • A (jnp.ndarray) – Constraint matrix of shape (n_constraints, n_basis).

  • b (jnp.ndarray) – Constraint bounds of shape (n_constraints,).

  • weight (float) – Penalty weight.

  • hard (bool) – If True, project to satisfy constraint.

Returns:

self

Return type:

Constraints

add_known_coefficient(basis_name: str, value: float, fixed: bool = True) → Constraints#

Fix a coefficient to a known value.

Parameters:

  • basis_name (str) – Name of basis function.

  • value (float) – Value to fix coefficient to.

  • fixed (bool) – If True, coefficient is fixed during fitting.

Returns:

self

Return type:

Constraints

add_custom(name: str, constraint_fn: Callable[[Array, Array, Array], float], weight: float = 1.0) → Constraints#

Add a custom nonlinear constraint.

The constraint function should return a penalty value (0 if satisfied, positive otherwise). It receives (coefficients, X, y_pred) as arguments.

Parameters:

  • name (str) – Name for this constraint (for reporting).

  • constraint_fn (callable) – Function (coefficients, X, y_pred) -> float penalty. Should return 0 when satisfied, positive when violated.

  • weight (float) – Penalty weight.

Returns:

self

Return type:

Constraints

Examples

>>> # Constraint: sum of coefficients must equal 1
>>> def sum_to_one(coeffs, X, y_pred):
...     return (jnp.sum(coeffs) - 1.0) ** 2
>>> constraints = Constraints().add_custom("sum_to_one", sum_to_one)

>>> # Constraint: prediction at x=0 should be near 0
>>> def zero_at_origin(coeffs, X, y_pred):
...     origin = jnp.zeros((1, X.shape[1]))
...     # Evaluate at origin requires basis evaluation
...     return 0.0  # Placeholder
>>> constraints = Constraints().add_custom("zero_origin", zero_at_origin)

>>> # Constraint: ratio of coefficients
>>> def coeff_ratio(coeffs, X, y_pred):
...     # coeffs[0] / coeffs[1] should be approximately 2
...     if len(coeffs) < 2 or abs(coeffs[1]) < 1e-8:
...         return 0.0
...     ratio = coeffs[0] / coeffs[1]
...     return (ratio - 2.0) ** 2
add_physics_constraint(name: str, constraint_type: str, params: dict[str, Any], weight: float = 1.0) → Constraints#

Add a physics-based constraint using predefined templates.

Parameters:

  • name (str) – Name for the constraint.

  • constraint_type (str) – Type of physics constraint: - “asymptotic”: y -> value as x -> inf - “periodic”: y(x) = y(x + period) - “symmetric”: y(x) = y(-x) or y(x) = -y(-x) - “scaling”: y(ax) = a^n * y(x) for some n - “passthrough”: y(x0) = y0 (passes through a specific point)

  • params (dict) – Parameters for the constraint type.

  • weight (float) – Penalty weight.

Returns:

self

Return type:

Constraints

Examples

>>> # y -> 0 as x -> infinity
>>> constraints = Constraints().add_physics_constraint(
...     "asymptotic_zero", "asymptotic",
...     {"feature": "x", "value": 0.0, "at": "infinity"}
... )

>>> # y passes through (0, 1)
>>> constraints = Constraints().add_physics_constraint(
...     "initial_condition", "passthrough",
...     {"point": [0.0], "value": 1.0}
... )
to_dict() → dict[str, Any]#

Serialize to dictionary.

classmethod from_dict(data: dict[str, Any]) → Constraints#

Deserialize from dictionary.

class jaxsr.constraints.ConstraintEvaluator(constraints: Constraints, basis_names: list[str], feature_names: list[str])#

Bases: object

Evaluates constraints for a given model.

Parameters:

  • constraints (Constraints) – Constraint specifications.

  • basis_names (list of str) – Names of basis functions.

  • feature_names (list of str) – Names of input features.

compute_penalty(coefficients: Array, predict_fn: Callable[[Array], Array], X: Array, y: Array | None = None) → float#

Compute total penalty for constraint violations.

Parameters:

  • coefficients (jnp.ndarray) – Current coefficients.

  • predict_fn (callable) – Function that predicts y given X. Monotonic/convex/concave constraints differentiate it with jax.jvp, so it must be written with jax.numpy operations.

  • X (jnp.ndarray) – Input points to check constraints.

  • y (jnp.ndarray, optional) – Target values (for output bounds).

Returns:

penalty – Total weighted penalty.

Return type:

float

apply_hard_constraints(coefficients: Array) → Array#

Project coefficients to satisfy hard constraints.

Parameters:

coefficients (jnp.ndarray) – Current coefficients.

Returns:

coefficients – Projected coefficients.

Return type:

jnp.ndarray

get_fixed_indices() → list[tuple[int, float]]#

Get indices and values of fixed coefficients.

Returns:

fixed – Indices and fixed values.

Return type:

list of (int, float)

compute_hard_penalty(coefficients: Array, predict_fn: Callable[[Array], Array], X: Array) → float#

Compute penalty for hard shape constraints (MONOTONIC, CONVEX, CONCAVE, BOUND, LINEAR).

These are constraints marked as hard=True that cannot be enforced via simple coefficient projection (unlike SIGN/FIXED). They are penalized heavily during optimization to approximate hard enforcement.

Parameters:

  • coefficients (jnp.ndarray) – Current coefficients.

  • predict_fn (callable) – Function that predicts y given X. Monotonic/convex/concave constraints differentiate it with jax.jvp, so it must be written with jax.numpy operations.

  • X (jnp.ndarray) – Input points to check constraints.

Returns:

penalty – Total weighted penalty for hard shape constraints.

Return type:

float

check_satisfaction(coefficients: Array, predict_fn: Callable[[Array], Array], X: Array, tolerance: float = 1e-06) → dict[str, bool]#

Check which constraints are satisfied.

Parameters:

  • coefficients (jnp.ndarray) – Current coefficients.

  • predict_fn (callable) – Prediction function. Monotonic/convex/concave constraints differentiate it with jax.jvp, so it must be written with jax.numpy operations.

  • X (jnp.ndarray) – Test points.

  • tolerance (float) – Tolerance for constraint satisfaction.

Returns:

satisfied – Constraint names mapped to satisfaction status.

Return type:

dict

jaxsr.constraints.fit_constrained_ols(Phi: Array, y: Array, constraints: Constraints, basis_names: list[str], feature_names: list[str], X: Array, max_iter: int = 100, tol: float = 1e-06, penalty_weight: float = 1.0, basis_library: Any | None = None, selected_indices: Any | None = None, enforcement: str = 'penalty', sample_weight: Array | None = None) → tuple[Array, float]#

Fit least squares with constraints.

For simple SIGN/FIXED constraints, uses OLS + projection (fast path). For shape constraints (monotonic, convex, bounds, etc.), the solver is chosen by the enforcement parameter.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • constraints (Constraints) – Constraint specifications.

  • basis_names (list of str) – Names of basis functions.

  • feature_names (list of str) – Names of input features.

  • X (jnp.ndarray) – Input data (for evaluating constraints).

  • max_iter (int) – Maximum iterations for optimizer.

  • tol (float) – Convergence tolerance for optimizer.

  • penalty_weight (float) – Weight for soft constraint penalties.

  • basis_library (BasisLibrary, optional) – Basis library for evaluating predictions at arbitrary X points.

  • selected_indices (array-like, optional) – Indices of selected basis functions in the full library.

  • enforcement (str) –

    Constraint enforcement level:

    • "penalty" (default) – L-BFGS-B with penalty terms. Approximate.

    • "constrained" – scipy trust-constr with LinearConstraint. Solver-tolerance guarantee (~1e-8).

    • "exact" – cvxpy QP. Mathematical guarantee. Requires cvxpy.

  • sample_weight (jnp.ndarray, optional) – Per-sample weights of shape (n_samples,). Only the least-squares part of the objective is weighted; the constraints are properties of the fitted function over the design space and are evaluated on the raw X, so a down-weighted row still has to obey them. The returned MSE is the weighted MSE sum_i w_i r_i^2 / n.

Returns:

  • coefficients (jnp.ndarray) – Fitted coefficients.

  • mse (float) – (Weighted) mean squared error.

Raises:

  • ValueError – If enforcement is not one of the valid levels, if enforcement='exact' is used with hard=True CUSTOM constraints, or if sample_weight is invalid.

  • ImportError – If enforcement='exact' and cvxpy is not installed.

  • RuntimeError – If enforcement='exact' and the QP is infeasible.

jaxsr.constraints.build_constraint_scorer(constraints: Constraints, X: Array, basis_library: Any, feature_names: list[str], penalty_weight: float = 1.0, hard_penalty_weight: float = 1000.0) → Callable#

Build a callable that scores constraint violations for a candidate model.

The returned scorer can be passed to selection strategies via the constraint_scorer keyword argument. During model selection each candidate’s information criterion is augmented by the penalty returned by the scorer, biasing selection towards models that can better satisfy the constraints.

Parameters:

  • constraints (Constraints) – Physical constraints to evaluate.

  • X (jnp.ndarray) – Training data used for constraint evaluation (shape (n, d)).

  • basis_library (BasisLibrary) – Full basis library (needed for evaluate_subset).

  • feature_names (list of str) – Names of input features (e.g. ["x1", "x2"]).

  • penalty_weight (float) – Weight applied to soft constraint penalties.

  • hard_penalty_weight (float) – Weight applied to hard constraint penalties.

Returns:

scorer – scorer(result, indices) -> float where result is a SelectionResult and indices is the list of selected basis function indices. Returns 0.0 when there are no violations.

Return type:

callable