API Reference
The public surface is split across the pages below, grouped by concern. The modelling workflow is re-exported at the top level, so from jaxhybridmodels import MLPPredictor works like from jaxhybridmodels.predictors import MLPPredictor. Lower-level pieces (kernels, penalties, losses, transforms, registries, UI) are imported from their submodule; each entry shows its import line.
Models and Results
HybridModel, FitResult, FitHistory, PredictionResult, ExperimentPrediction
Data: Experiments, Channels, Datasets
ChannelObs, Experiment, make_experiment, experiment_from_array, experiments_from_dataframe, BucketPayload, Dataset, make_dataset, make_bootstrap_dataset, split_dataset, describe_buckets
Predictors: Trainable Components
Predictor, BoundScaler, BoundedPredictor, Parameter, Constant, MLPPredictor, KANPredictor, NeuralNPolynomial, reinitialize_with_key, reinitialize_pytree_with_key
Penalties: Gradient-Safe Bound Handling
PenaltyPointSource, bounded_predictor_paths, resolve_penalty_points, 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
Objective, OptaxTrainingConfig, TrainingPhase, train_with_optax, 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_mask, count_trainable_params
Prediction: Forward Simulation
predict_ensemble, 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, load, SavedRun, save_predictors, load_predictors
UI: Training Dashboards
TrainingUI, EvosaxUI, SilentUI, RichTrainingUI, RichEvosaxUI