API Reference
The public surface is split across the pages below, grouped by concern. Every symbol below is also re-exported at the top level — from jaxhybridmodels import MLPPredictor works exactly like from jaxhybridmodels.predictors import MLPPredictor.
Data: Experiments, Channels, Datasets
ChannelObs, Experiment, make_experiment, BucketPayload, Dataset, make_dataset, make_bootstrap_dataset, split_dataset, describe_buckets
Predictors: Trainable Components
Predictor, BoundScaler, BoundedPredictor, MLPPredictor, KANPredictor, NeuralNPolynomial, reinitialize_with_key, reinitialize_pytree_with_key
Penalties: Gradient-Safe Bound Handling
PenaltyPointSource, attach_penalty_state, bound_penalty, box_grid, box_violation, clip_ste, data_penalty_points, length_mask_keep, penalty_integral, penalty_vector_field, select_penalty_points, soft_inverse, soft_logit, softclip, strip_penalty_state, trajectory_saturation_penalty, validate_penalty_points
Transforms: Squash Shapes and Axis Warps
BoundTransform, BOUND_TRANSFORMS, register_bound_transform, Warp, WARPS, register_warp
Solver: ODE Integration
SolverConfig, SOLVER_REGISTRY, register_solver, ADJOINT_REGISTRY, register_adjoint
Training: Optax & Evosax Loops
OptaxTrainingConfig, train_with_optax, train_seed_ensemble, train_bootstrap_ensemble, EvosaxTrainingConfig, train_with_evosax, register_algorithm
Training Kernels: Build Your Own Loop
apply_length_mask, predict_bucket_obs, build_bucket_step, build_score_bucket, build_penalty_step, build_apply_update
Losses: Masked & Balanced Objectives
masked_mse, masked_mle, bal_mse, bal_mle, LOSS_REGISTRY, resolve_loss_fn
Trainable Masks: Freezing Leaves
default_trainable, trainable_mask, freeze_paths, freeze_modules_of_type, freeze_where, frozen_default_mask, count_trainable_params
Prediction: Forward Simulation
predict_bucket, predict_dataset, predict_dense, ensemble_predictions, evaluate_predictor
Metrics: Per-Channel Evaluation
ChannelMetrics, compute_metrics, print_metrics
Profiles: Time-Varying Inputs
constant_profile, step_profile, ramp_profile, piecewise_linear_profile
Schedules: Epoch-Scaled Annealing
Serialise: Save & Load Runs
save_predictors, load_predictors, save_run, load_run
UI: Training Dashboards
TrainingUI, EvosaxUI, SilentUI, RichTrainingUI, RichEvosaxUI