Changelog

Best viewed here.

All notable changes to PhiJAX are documented in this file.

The format follows Keep a Changelog, and PhiJAX follows Semantic Versioning.

Unreleased

Added no_slip_boundary with configurable velocity components and residual/output streams for stationary and moving walls. (#39)
Added configurable weight, projection, and period batch keys to fidelity equations, and a configurable normals key to free_slip_boundary. (#38)

phase_wrapped_fidelity now uses batch["projection"] when present. Set projection_key=None to preserve direct supervision when that field contains unrelated data. (#38)

Updated the optional W&B dependency to >=0.30,<0.31. (#37)
Added an optimizer guide covering Optax, the optional SOAP_JAX compatibility fork, and checkpoint restoration. (#40)
Clarified installation workflows with uv and pip. (#36)

Log all term losses before loss-balancer weights for consistent metric ordering. (#35)

Added public type aliases for supported accelerator, device, precision, monitoring, activation, and initialization options. (#25)

Removed the inert Trainer.deterministic option. Deterministic GPU experiments must configure XLA before JAX initializes. (#30)

Reorganized the documentation into task-oriented guides and focused API reference pages. (#26, #27, #28)

Added a target.name OmegaConf resolver for deriving display labels from Hydra _target_ paths. (#22)

Made LearningRateMonitor require an explicit optimizer name and changed its default metric group from train/ to optimizer/, using Lightning-style keys such as optimizer/lr-Adam. (#22)

Made the Trainer report its precision and accelerator environment automatically when constructed on global rank zero. (#21)
Expanded the balancer documentation with the exact metrics, meanings, update behavior, and default destinations for static, gradient-norm, and exact-NTK balancing. (#23)

Added to_hyperparameters() for converting composed Hydra configs before passing them to Trainer.fit(). (#19)

Removed callback and logger enabled configuration options. Omit a service entry to disable it. (#16)
Removed Trainer.set_logger(). Construct loggers first and pass them through Trainer(logger=...) or instantiate_trainer(..., logger=...); instantiate_loggers() no longer accepts a Trainer. (#19)
Renamed the experiment logger method log_hyperparameters() to the Lightning-style log_hyperparams(). (#19)

Made CSV, TensorBoard, and W&B logger resource setup lazy, idempotent, and restricted to global rank zero during Trainer tasks. Direct built-in logger use still starts resources automatically on first use. (#19)

Concise Trainer.fit() and Trainer.predict() APIs that initialize models, state, balancing, batch routing, and compiled plans from a module blueprint, DataModule, optimizer, and seed. (#15)
ModelFactory, balancer-owned adaptive update plans, objective equation composition, finite-pool batch-source construction, and safe HostPool defaults. (#15)
Independent persistent PRNG streams for model execution, sampling, and adaptive-balancer diagnostics. (#15)
Structured training diagnostics with host-side PhiModule.log() control for loggers and progress displays. (#15)
Automatic TQDM progress, plain model summaries, default TensorBoard or CSV logging, and versioned run directories. (#15)
Monitored top-k checkpoints with persistent callback state and experimental TPU installation support. (#15)

Renamed the explicit execution APIs to Trainer.fit_state() and Trainer.predict_state(). (#15)
Changed PhiModule into an immutable model-factory blueprint; the bound module is returned as FitResult.module. (#15)
Changed ResidualTerm to accept its residual function first and infer its default NTK stream from equation metadata. (#15)
Raised the checkpoint schema to version 2 for the expanded TrainState layout. (#15)
Moved adaptive cadence and diagnostic sampling settings into adaptive-balancer constructors. (#15)
Replaced the adaptive skip_first_step flag with cadence-anchored update_start_step scheduling. (#15)
Removed separate sampling and balancer key arguments from Trainer methods; both now live in TrainState. (#15)
Removed resume_latest() in favor of fit(..., ckpt_path="last"). (#15)
Removed the low-level with_balancer_updates() helper; advanced schedules now belong to TrainingPlan. (#15)
Removed BalancerUpdateConfig and BalancerUpdateSchedule; TrainingPlan now accepts one BalancerUpdatePlan. (#15)
Removed the Hydra balancer.factory and balancer.update layers; balancer groups are directly instantiable. (#15)
Removed MetricRoute; modules now customize metric destinations with self.log() from on_train_batch_end(). (#15)
Removed the phijax.module deep import path; use top-level imports or phijax.core for custom module contracts. (#15)

Made DataModule input normalization opt-in through input_statistics(). (#15)
Made LearningRateMonitor require a configured experiment logger before fitting begins. (#15)
Moved adaptive-balancer scheduling into the Trainer host loop and removed unconditional per-step transfers of TrainState.step. (#15)
Split module contracts into phijax.core and moved fit, prediction, logging, and signal orchestration into private loops and connectors while keeping the top-level API unchanged. (#15)

Initial beta API for building, fitting, and evaluating physics-informed neural networks with JAX. (#1)
Lightning-inspired Trainer, PhiModule, DataModule, callback, logger, and checkpoint lifecycles. (#1)
Explicit model, optimizer, loss-balancer, PRNG, and training state for compiled updates. (#1)
Selective coordinate derivatives and reusable PDE, boundary, data-fidelity, and objective components. (#1)
Static, normalized-static, gradient-norm, and exact-NTK loss balancing. (#1)
MLP, Modified MLP, and PirateNet architectures with a common model initialization contract. (#1)
CPU, CUDA 12, and CUDA 13 installation options, configurable precision, and centralized device placement. (#1)
Versioned checkpoints and portable prediction artifacts for downstream evaluation. (#1)
Optional Hydra integration and Weights & Biases and TensorBoard logging. (#1)