Custom predictor

The library ships an MLP, a KAN and a neural polynomial, but a predictor is only a trainable module mapping an array to an array. This example writes a fourth family from scratch, random Fourier features, and trains it inside a hybrid ODE. Read it before writing your own; the procedure is written up in Custom predictors.

Run it

uv run python examples/custom_predictor/train_custom_predictor.py

About a minute on CPU. Prints the learned rate against the truth before and after training, the save/load drift and metrics; writes parity.png and trajectories.png to examples/custom_predictor/figures/. --steps 5 --no-plot is a quick smoke run.

The model

First-order decay whose rate depends on temperature:

$$\frac{dy}{dt}=-k(T)y,k(T)=A\exp (-\frac{E_a}{RT})$$

Each experiment is one isothermal run: a temperature between 290 and 350 K, 15 noisy points, and $y_0=1$. All runs share a time grid, so the dataset is one fully observed bucket. The Arrhenius form is the answer, not an input: the predictor sees a temperature and must produce a rate. The true rates span 1.5 decades, so the output scaler uses a log10 warp.

Random Fourier features project the input through a fixed random bank of cosines and fit linear weights on top:

$$\varphi (x)=\sqrt{\frac{2}{m}}\cos (x\mathrm{\Omega }+b),f(x)=\varphi (x)W+c$$

The bank $(\mathrm{\Omega },b)$ is drawn once and never updated; only the readout $(W,c)$ learns. That split between fixed and trained arrays is why this architecture was chosen. It forces every question the predictor contract asks:

• Which fields are static. in_size, out_size, n_features and bandwidth are eqx.field(static=True). The optimiser never sees them and hm.save does not write them, so the load site has to be able to rebuild them. bandwidth is the standard deviation of the frequency draw; because BoundScaler hands the inner predictor a latent input of roughly unit scale, a value near 1 is the sane start.
• How it re-initialises. The restart tournament calls initialized_with_key. Without it, every float array is resampled from a standard normal: the phases stop being uniform on $[0,2\pi )$, the readout loses its $1/\sqrt{m}$ scale and bandwidth is ignored.
• How the bank is held fixed. A module cannot declare its own field untrainable. Trainability is a mask over the pytree, set by the caller: here fit(freeze=...) with the bank's paths, which raises on a path that matches nothing.
• How it round-trips. hm.load(path, like=model) needs a model with the same class, container and static fields. The untrained model, built by the same build_predictors used for training, is one.

Bounds stay outside: RandomFourierPredictor has no bound, covariate name or unit. BoundedPredictor supplies all three. The rate is constant along a trajectory, so simulate_fn calls the predictor once, above the solve, and the network never touches the solver tape. See embedded and parallel predictors for the case where it cannot.

Code

The predictor subclasses hm.Predictor. Its class body declares the trainable arrays, the fixed feature bank and the static shape metadata, then implements __call__ and the restart hook:

weights: Float[Array, "n_features out_size"]
bias: Float[Array, " out_size"]
frequencies: Float[Array, "in_size n_features"]
phases: Float[Array, " n_features"]
in_size: int = eqx.field(static=True)
out_size: int = eqx.field(static=True)
n_features: int = eqx.field(static=True)
bandwidth: float = eqx.field(static=True)

def __init__(
    self,
    *,
    in_size: int,
    out_size: int,
    n_features: int = 64,
    bandwidth: float = 1.0,
    key: Array,
) -> None:
    k_freq, k_phase, k_weight = jr.split(key, 3)
    self.in_size = in_size
    self.out_size = out_size
    self.n_features = n_features
    self.bandwidth = bandwidth
    self.frequencies = bandwidth * jr.normal(k_freq, (in_size, n_features))
    self.phases = 2.0 * jnp.pi * jr.uniform(k_phase, (n_features,))
    # Scaling the readout by 1/sqrt(n_features) keeps the output
    # variance independent of the feature count, so changing
    # ``n_features`` does not silently change the effective learning
    # rate.
    self.weights = jr.normal(k_weight, (n_features, out_size)) / jnp.sqrt(n_features)
    self.bias = jnp.zeros((out_size,))

def __call__(self, x: Float[Array, " in_size"]) -> Float[Array, " out_size"]:
    features = jnp.sqrt(2.0 / self.n_features) * jnp.cos(x @ self.frequencies + self.phases)
    return features @ self.weights + self.bias

def initialized_with_key(self, key: Array) -> RandomFourierPredictor:
    """Re-init for the restart tournament: a fresh bank and readout.

    The default in ``reinitialize_with_key`` replaces every float
    leaf with a standard normal, which would give the bank the wrong
    scale and drop ``bandwidth`` on the floor. Rebuilding through
    ``__init__`` keeps the construction scheme in one place.

    Shapes are unchanged, so the trainability mask built from
    ``freeze=`` still lines up after a restart.
    """
    return RandomFourierPredictor(
        in_size=self.in_size,
        out_size=self.out_size,
        n_features=self.n_features,
        bandwidth=self.bandwidth,
        key=key,
    )

Wrapping it with bounds:

rate = hm.BoundedPredictor(
    input_keys=("temperature",),
    input_bounds=TEMPERATURE_BOUNDS,  # default sigmoid input scaler
    inner=RandomFourierPredictor(
        in_size=1, out_size=1, n_features=n_features, bandwidth=bandwidth, key=key
    ),
    out_scaler=hm.BoundScaler(bounds=K_BOUNDS, transform="sigmoid", warp="log10"),
)
return {"rate": rate}

The bank is frozen by path. The paths read from the root of the predictors dict, and fit takes them through freeze=:

# Dot-joined paths of the two fixed arrays inside the predictors dict: the
# "rate" entry, the BoundedPredictor's ``inner``, then the field.
BANK_PATHS: tuple[str, ...] = (
    "rate.inner.frequencies",
    "rate.inner.phases",
)

# By default every array of the custom predictor trains, the bank
# included. ``fit(freeze=...)`` turns the bank off by name, and raises
# on a path that matches nothing, so a renamed field fails loudly
# instead of training an array you believed was fixed.
freeze = None if args.train_bank else BANK_PATHS

The fit runs a three-restart tournament, which is what exercises initialized_with_key:

# The tournament restarts the predictor three times through
# ``initialized_with_key`` and keeps the best start.
config = hm.OptaxTrainingConfig(
    steps=args.steps,
    lr=args.lr,
    tournament_attempts=3,
    tournament_steps=40,
    tournament_lr=args.lr,
    verbose=False,
)
result = model.fit(dataset, config, key=k_train, freeze=freeze)

Results

Default settings, seed 0, 400 steps after the tournament. "Worst rate error" is the largest relative error in $k(T)$ across the seven sampled temperatures. The save and reload changes the learned rate by 0.000e+00.

run	final data loss	$R^2$	worst rate error
default: 64 features, bandwidth 1, bank frozen	8.55e-05	0.9992	3.92%
--n-features 16	8.63e-05	0.9992	4.25%
--train-bank	8.46e-05	0.9993	3.34%
--bandwidth 8	8.44e-05	0.9993	3.35%

Every row fits the trajectories to the noise floor (noise variance ${10}^{-4}$), and the rate law comes back within about 4%. The differences between rows are within what a different seed would move. Below 330 K every run is within about 1%; the error sits at 340 and 350 K, where the decay is over in a few samples and the data says little about the rate.

The bottom panels are the hot runs: the concentration reaches zero within three or four samples, which is why the rate law is least constrained there.

Things to try

• --train-bank trains the bank too, turning the model into an ordinary cosine-activation network.
• --n-features 16 is enough for this smooth target; 64 leaves room for a harder rate law.
• --bandwidth 8 resolves shorter length scales and overfits sparse data sooner.