jaxhybridmodels

Appearance

Penalties: Gradient-Safe Bound Handling ​

PenaltyPointSource ​

from jaxhybridmodels.penalties import PenaltyPointSource  ·  also re-exported as jaxhybridmodels.PenaltyPointSource

python
PenaltyPointSource(
    points: ForwardRef('Array'),
    cell_ts: ForwardRef('Array'),
    cell_T: ForwardRef('Array'),
)

One BoundedPredictor leaf's gathered measured points.

Produced by data_penalty_points for leaves whose input_keys all resolve to dataset covariates. One entry per observed cell: the input vector the loss actually sees at that cell, in input_keys column order, in physical units.

Attributes

FieldTypeDescription
pointsFloat[Array, "G n_inputs"]The measured input vectors.
cell_tsInt[Array, " G"]Timestamp index of each cell inside its own experiment.
cell_TInt[Array, " G"]Length of that experiment's time grid. cell_ts and cell_T let length_mask_keep apply the loss's prefix mask to the penalty, so the two never disagree about which points are live.

Source

attach_penalty_state() ​

from jaxhybridmodels.penalties import attach_penalty_state  ·  also re-exported as jaxhybridmodels.attach_penalty_state

python
attach_penalty_state(y0: 'Array', n: 'int' = 1) -> 'Array'

Append n zero-valued penalty accumulators to y0.

The first step of the trajectory-penalty recipe: the ODE state widens from [S] to [S + n], where the trailing components are integrated penalty rates supplied by penalty_vector_field. y0_fn should return attach_penalty_state(physics_y0, n).

Source

bound_penalty() ​

from jaxhybridmodels.penalties import bound_penalty  ·  also re-exported as jaxhybridmodels.bound_penalty

python
bound_penalty(predictors: 'Any', points: 'tuple[Array, ...]') -> 'Array'

Mean output-squash saturation over every BoundedPredictor in a pytree.

For each leaf, evaluates inner(in_scaler.to_latent(x)) across the leaf's point set and charges BoundScaler.saturation on the resulting latents. Leaves are summed.

points is positional, one [G, n_inputs] array per leaf in traversal order: measured points from data_penalty_points, user-supplied penalty-only points, or a box_grid sweep. An empty per-leaf array contributes exactly zero, so a leaf the penalty cannot reach does not NaN a run.

The default way to penalise bound behaviour here, and deliberately evaluated at fixed input points rather than along a solve: saturation is a property of the predictor as a function, whatever any particular solve does. So it needs no cooperation from simulate_fn, predict_bucket or the loss protocol, behaves the same inside a vector field or above one, and handles arbitrary nesting by walking leaves.

That cuts both ways. It reports saturation wherever the point set reaches, including regions no training trajectory visited when the points say so (a box_grid sweep, user extras); for "did this solve push an input out of range" the trajectory-aware penalty is the right instrument.

Parameters

ParameterTypeDescription
predictorsPyTree[eqx.Module]Any pytree shape. Only BoundedPredictor leaves contribute.
pointstuple[Array, ...]One [G, n_inputs] array per leaf, in traversal order: the output of data_penalty_points (with any user extras concatenated), user penalty-only points, or a box_grid sweep. An empty per-leaf array contributes zero.

Returns

ItemTypeDescription
ArrayNon-negative scalar. Exactly zero when no leaf saturates.

Source

box_grid() ​

from jaxhybridmodels.penalties import box_grid  ·  also re-exported as jaxhybridmodels.box_grid

python
box_grid(in_scaler: 'Any', n_per_dim: 'int' = 5) -> 'Array'

Build a deterministic sweep of an input box, uniform in warped coordinates.

The collocation-as-extension recipe: a tensor product of n_per_dim evenly spaced points along every dimension of in_scaler's box, in the box's warped coordinates, mapped back to physical units. A linear warp therefore reproduces a plain physical linspace sweep bit-for-bit, while log / log10 warps cover each decade evenly instead of starving the low end of the box.

n_per_dim must be at least 2 so both edges of every input box are represented. Deterministic: it depends only on the scaler's static bounds and warp, so restore_best compares raw loss values across steps with no resampling noise.

Point count grows exponentially in input dimension. Predictors here take two or three inputs, so n_per_dim=5 is 25 to 125 forward passes and negligible beside an ODE solve. Lower it if that stops holding.

Parameters

ParameterTypeDescription
in_scalerBoundScalerThe input scaler of the predictor the sweep is for (leaf.in_scaler).

Returns

ItemTypeDescription
Array[n_per_dim ** n_inputs, n_inputs] physical points, positional in the scaler's input dimension order.

Source

box_violation() ​

from jaxhybridmodels.penalties import box_violation  ·  also re-exported as jaxhybridmodels.box_violation

python
box_violation(x: 'Array', lows: 'Array', highs: 'Array') -> 'Array'

Width-normalised squared hinge measuring how far x falls outside its box.

Zero in value and gradient strictly inside the box, so it never perturbs the feasible interior. Outside it grows quadratically, giving a restoring gradient linear in the overshoot, which does not vanish the way a reparameterised bound's does.

Each component is normalised by its own width high - low so one penalty weight works across channels. Bounds here run from fractions of a unit to hundreds of kelvin, and an unnormalised hinge would let the widest channel dominate on units alone.

Parameters

ParameterTypeDescription
xArrayValues in physical units; broadcast against lows / highs.
lows, highsArrayPer-component box edges, as produced by BoundScaler._lows_highs.

Returns

ItemTypeDescription
ArrayScalar sum of squared fractional violations.

Source

clip_ste() ​

from jaxhybridmodels.penalties import clip_ste  ·  also re-exported as jaxhybridmodels.clip_ste

python
clip_ste(x: 'Array', lo: 'float | Array', hi: 'float | Array') -> 'Array'

Hard-clip on the forward pass, identity on the backward pass.

The straight-through estimator. Use it when downstream code genuinely requires a feasible number, say a concentration that must not go negative before a log, while the task loss keeps flowing as though the clip were not there.

The identity gradient is a deliberate fiction: it propagates whatever the data loss asks for, including "go further out of bounds", forever. Pair it with box_violation on the pre-clip value for the restoring force.

Source

data_penalty_points() ​

from jaxhybridmodels.penalties import data_penalty_points  ·  also re-exported as jaxhybridmodels.data_penalty_points

python
data_penalty_points(
    predictors: 'Any',
    dataset: 'Dataset',
) -> tuple[PenaltyPointSource | None, ...]

Gather one PenaltyPointSource per BoundedPredictor leaf.

A leaf is resolvable when every input_keys name is a dataset covariate (in every bucket). Its measured points are then the input vectors at the observed cells — the cells the loss actually charges — with a None entry for an unresolvable leaf. A resolvable leaf whose dataset holds no observed cells (all probes) yields a source with zero points, which contributes nothing until extras are added.

Call this once on the host before the training loop: the points depend only on the dataset and the leaf's static input_keys.

Covariates must be scalar per experiment (constant in time, the v1 contract); a per-experiment covariate with extra dimensions has no unambiguous column to feed an input key and raises.

Source

length_mask_keep() ​

from jaxhybridmodels.penalties import length_mask_keep  ·  also re-exported as jaxhybridmodels.length_mask_keep

python
length_mask_keep(source: 'PenaltyPointSource', length_mask_fraction: 'float') -> 'Array'

Boolean keep-vector over a source's cells, matching the loss's prefix mask.

Applies the same prefix semantics apply_length_mask gives the loss: a cell is kept when its timestamp index is below ceil(T * fraction), clamped at 1 so a phase never scores nothing. The penalty therefore follows the curriculum exactly, disagreeing with the loss about which points are live only by mistake.

length_mask_fraction must be a Python float, not a traced array: the keep-vector is used to index the gathered points, and JAX rejects boolean indexing with tracers. The penalty kernel retraces only when the fraction changes (a phase boundary), never per step.

Source

penalty_integral() ​

from jaxhybridmodels.penalties import penalty_integral  ·  also re-exported as jaxhybridmodels.penalty_integral

python
penalty_integral(state: 'Array', n: 'int' = 1) -> 'Array'

The accumulated (time-integrated) penalty values at the trajectory's end.

state is the full state trajectory [..., T, S + n] as produced by a penalty_vector_field solve. Returns the trailing n components at the final time, [..., n]. The training hook charges these; divide by the time span to get the time-mean instead of the integral.

Source

penalty_vector_field() ​

from jaxhybridmodels.penalties import penalty_vector_field  ·  also re-exported as jaxhybridmodels.penalty_vector_field

python
penalty_vector_field(
    base_rhs: 'Callable[[Array, Array, Any], Array]',
    penalty_rhs: 'Callable[[Array, Array, Any], Array]',
) -> Callable[[Array, Array, Any], Array]

Wrap a physics vector field with per-call penalty rates.

Returns (t, y, args) -> [physics_dot, penalty_rates]. base_rhs is the user's vector field on the physical components; penalty_rhs returns the n penalty rates (e.g. [saturation(z), input_violation(x)]) at the current state. The solver integrates both, so the trailing accumulators carry the time-integral of the penalty along the trajectory.

The rates must read the physical components (typically y[:-n]) and are closed over the user's predictors — only the user knows the latent z or input x of an embedded predictor at call time.

Source

select_penalty_points() ​

from jaxhybridmodels.penalties import select_penalty_points  ·  also re-exported as jaxhybridmodels.select_penalty_points

python
select_penalty_points(
    sources: 'tuple[PenaltyPointSource | None, ...]',
    extras: 'tuple[Array, ...]',
    length_mask_fraction: 'float',
) -> tuple[Array, ...]

Per-leaf point arrays for one phase: length-mask-kept cells ∪ extras.

sources is the per-leaf output of data_penalty_points (None for unresolvable leaves); extras the user-supplied penalty-only points, positional per leaf. Returns one [G, n_inputs] array per leaf in traversal order: the measured cells kept by length_mask_keep concatenated with that leaf's extras.

This is host-side selection: the keep-vector must be concrete to index the gathered points, so call it with the phase's Python float outside any jit (the stock trainers do, once per phase). The returned arrays are what a penalty_step receives.

Source

soft_inverse() ​

from jaxhybridmodels.penalties import soft_inverse  ·  also re-exported as jaxhybridmodels.soft_inverse

python
soft_inverse(
    s: 'Array',
    inverse: 'Callable[[Array], Array]',
    inverse_slope: 'Callable[[Array], Array]',
    eps: 'float' = 0.001,
) -> Array

inverse(s), extended linearly outside [eps, 1 - eps].

Generalises soft_logit to the inverse of any squashing function. Every such inverse has a pole at each end of the unit interval, where a hard clip would zero the derivative and silently drop state-derived sensitivities from the ODE adjoint (R-P2).

Exact in value and derivative inside the band, and C^1 across the junction, since the continuation uses the inverse's own slope there. The stop_gradient on the clamp is load-bearing: without it the correction term picks up a contribution through the clip and the interior derivative comes out wrong.

Source

soft_logit() ​

from jaxhybridmodels.penalties import soft_logit  ·  also re-exported as jaxhybridmodels.soft_logit

python
soft_logit(s: 'Array', eps: 'float' = 0.001) -> 'Array'

logit(s), extended linearly outside [eps, 1 - eps].

s is a physical value already normalised into [0, 1] across its declared box, so this is the step that turns a bounded quantity into an unbounded latent. Outside the band the result grows linearly instead of blowing up at the pole, and the derivative is a finite constant instead of the exact zero logit(jnp.clip(s, eps, 1 - eps)) would give.

softclip cannot do this job: its interior error is O(1 / beta) and s spans only [0, 1], so any beta gentle enough to keep gradient far outside the box also distorts the middle.

eps sets the continuation slope, roughly 1 / eps, and is the one tuning knob. The default 1e-3 maps a 1% overshoot to |z| ~ 10, a number a network can still consume; 1e-6 would map it to |z| ~ 1e4.

Source

softclip() ​

from jaxhybridmodels.penalties import softclip  ·  also re-exported as jaxhybridmodels.softclip

python
softclip(
    x: 'Array',
    lo: 'float | Array',
    hi: 'float | Array',
    beta: 'float' = 20.0,
) -> Array

Smoothly clamp x into [lo, hi] with a derivative that never hits zero.

Built from two softplus shoulders, so the derivative is sigmoid(beta * (x - lo)) - sigmoid(beta * (x - hi)), analytically in (0, 1). The interior is reproduced to O(1 / beta). Larger beta tracks a hard clip more closely but underflows sooner outside the box; the default 20 holds interior error below about 0.05 box widths and keeps usable gradient roughly one width out.

Repairs the near field only. Several widths out the derivative underflows as a hard clip's does, so pair it with box_violation for unbounded push-back.

Source

strip_penalty_state() ​

from jaxhybridmodels.penalties import strip_penalty_state  ·  also re-exported as jaxhybridmodels.strip_penalty_state

python
strip_penalty_state(state: 'Array', n: 'int' = 1) -> 'Array'

Drop the trailing n penalty accumulators from a full-state trajectory.

Use in state_to_output: the loss and the observed channels should see only the physics, not the integrated penalty components. state is [..., S + n]; the result is [..., S].

Source

trajectory_saturation_penalty() ​

from jaxhybridmodels.penalties import trajectory_saturation_penalty  ·  also re-exported as jaxhybridmodels.trajectory_saturation_penalty

python
trajectory_saturation_penalty(state: 'Array', out_scaler: 'Any') -> 'Array'

Sum over time of output saturation for a predictor whose output is the state.

For a parallel hybrid model — the predictor is outside the solver and its output is a predicted channel — the full state already holds the physical outputs. Invert them back to latents with the predictor's out_scaler and charge BoundScaler.saturation at every time step, summed over time. This is the trajectory-aware counterpart of bound_penalty for the hoisted case: it fires only where the model actually predicted, not across a synthetic grid.

state is [..., T, D] (or [..., T, S] projected to the predictor's channel); out_scaler is the BoundScaler whose from_latent produced those outputs.

Source

validate_penalty_points() ​

from jaxhybridmodels.penalties import validate_penalty_points  ·  also re-exported as jaxhybridmodels.validate_penalty_points

python
validate_penalty_points(
    predictors: 'Any',
    sources: 'tuple[PenaltyPointSource | None, ...]',
    extras: 'tuple[Array, ...]',
    enabled: 'bool',
) -> None

Raise when an enabled penalty has a leaf no point set reaches.

sources is the per-leaf output of data_penalty_points (None for unresolvable leaves); extras the user-supplied penalty-only points, positional per leaf. enabled is whether any phase charges the penalty.

With the penalty on, every BoundedPredictor leaf must be covered by measured points, extras, or both; an uncovered leaf would otherwise be a silent no-penalty, the failure mode this penalty exists to prevent. For embedded predictors (state-derived inputs) the trajectory penalty is the instrument, and the error says so.

Point-shape mismatches are checked unconditionally, so a user cannot carry a silently-wrong extras array into a run that later enables it.

Source

Released under the BSD-3-Clause License.