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) → 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.

Returns:

self

Return type:

Constraints

add_convex(feature: str, weight: float = 1.0, hard: bool = False) → Constraints#

Add convexity constraint (positive second derivative).

Parameters:

feature (str) – Feature name.
weight (float) – Penalty weight.
hard (bool) – If True, enforce strictly.

Returns:

self

Return type:

Constraints

add_concave(feature: str, weight: float = 1.0, hard: bool = False) → Constraints#

Add concavity constraint (negative second derivative).

Parameters:

feature (str) – Feature name.
weight (float) – Penalty weight.
hard (bool) – If True, enforce strictly.

Returns:

self

Return type:

Constraints

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.
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.
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.
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

jaxsr.constraints

Contents

jaxsr.constraints#