Metadata-Version: 2.4
Name: phijax
Version: 0.2.0b4
Summary: A JAX framework for reproducible physics-informed neural network training.
Keywords: jax,physics-informed-neural-networks,pinn,scientific-machine-learning
Author: Hang Jung Ling
Author-email: Hang Jung Ling <hangjung.ling@gmail.com>
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Dist: colorlog>=6.10,<7
Requires-Dist: flax==0.12.9
Requires-Dist: hydra-core==1.3.6
Requires-Dist: h5py>=3.16,<4
Requires-Dist: jax==0.11.1
Requires-Dist: mat73>=0.65,<1
Requires-Dist: numpy>=2.4,<3
Requires-Dist: omegaconf==2.3.1
Requires-Dist: optax==0.2.8
Requires-Dist: orbax-checkpoint==0.12.4
Requires-Dist: pyyaml>=6.0.3,<7
Requires-Dist: rich>=15,<16
Requires-Dist: rootutils>=1.0.7,<2
Requires-Dist: scipy>=1.17,<2
Requires-Dist: tqdm>=4.67,<5
Requires-Dist: jax[cuda12]==0.11.1 ; extra == 'cuda12'
Requires-Dist: jax[cuda13]==0.11.1 ; extra == 'cuda13'
Requires-Dist: tensorboard>=2.20,<3 ; extra == 'tensorboard'
Requires-Dist: jax[tpu]==0.11.1 ; extra == 'tpu'
Requires-Dist: wandb>=0.29,<0.30 ; extra == 'wandb'
Requires-Python: >=3.12
Project-URL: Changelog, https://github.com/HangJung97/PhiJAX/blob/main/CHANGELOG.md
Project-URL: Documentation, https://hangjung97.github.io/PhiJAX
Project-URL: Issues, https://github.com/HangJung97/PhiJAX/issues
Project-URL: Repository, https://github.com/HangJung97/PhiJAX
Provides-Extra: cuda12
Provides-Extra: cuda13
Provides-Extra: tensorboard
Provides-Extra: tpu
Provides-Extra: wandb
Description-Content-Type: text/markdown

<div align="center">

# PhiJAX

[![Code Quality](https://github.com/HangJung97/PhiJAX/actions/workflows/code-quality-main.yaml/badge.svg)](https://github.com/HangJung97/PhiJAX/actions/workflows/code-quality-main.yaml)
[![Tests](https://github.com/HangJung97/PhiJAX/actions/workflows/tests.yaml/badge.svg)](https://github.com/HangJung97/PhiJAX/actions/workflows/tests.yaml)
[![Documentation](https://github.com/HangJung97/PhiJAX/actions/workflows/docs.yaml/badge.svg)](https://github.com/HangJung97/PhiJAX/actions/workflows/docs.yaml)
[![Codecov](https://codecov.io/gh/HangJung97/PhiJAX/graph/badge.svg)](https://codecov.io/gh/HangJung97/PhiJAX)
<br>
[![Python](https://img.shields.io/pypi/pyversions/phijax?color=blue&logo=python&logoColor=white)](https://pypi.org/project/phijax/)
[![PyPI](https://img.shields.io/pypi/v/phijax?include_prereleases)](https://pypi.org/project/phijax/)
<br>
[![License](https://img.shields.io/github/license/HangJung97/PhiJAX?color=blue)](LICENSE)

</div>

PhiJAX is a typed [JAX](https://docs.jax.dev/) framework for physics-informed neural networks (PINNs). It provides the
core tools for training PINNs: functional state, explicit-key sampling, reusable equations and objectives, adaptive
loss balancing, callbacks, logging, checkpointing, and prediction artifacts. Its `Trainer` and `PhiModule` lifecycles
are inspired by Lightning.

For a configuration-first project layout, start from
[`phijax-hydra-template`](https://github.com/HangJung97/phijax-hydra-template) and customize its entrypoints, Hydra
configs, and application DataModules.

> **Beta API:** PhiJAX is under active development. Breaking changes may occur between major versions and during the
> beta period before the stable release. The supported public API is documented in package `__all__` declarations.
> Checkpoints can be restored by compatible PhiJAX releases from the same major and minor version.

[Quickstart](https://hangjung97.github.io/PhiJAX/getting-started/quickstart/) |
[Guides](https://hangjung97.github.io/PhiJAX/guides/datasets/) |
[API reference](https://hangjung97.github.io/PhiJAX/api/) |
[Changelog](CHANGELOG.md) |
[Contributing](CONTRIBUTING.md)

## Why PhiJAX?

JAX provides automatic differentiation, vectorization, compilation, and accelerator support. PhiJAX adds the
structure needed to build reproducible, maintainable PINN experiments:

- selective coordinate derivatives and reusable PDE, boundary, and data-fidelity equations;
- named objective terms with static, gradient-norm, and exact-NTK loss balancing;
- standard MLP, Modified MLP, and adaptive-residual PirateNet architectures with a custom NNX adapter;
- explicit model, optimizer, balancer, and PRNG state for reproducible compiled updates;
- familiar, Lightning-inspired lifecycles for trainers, modules, DataModules, callbacks, loggers, and checkpoints; and
- scientific data preparation on the CPU, automatic batch placement by the Trainer, and prediction files with a
  stable, versioned format.

PhiJAX reduces boilerplate for training, prediction, logging, and checkpointing without hiding JAX transformations or
functional state. Hydra support is optional, and the core Trainer API does not depend on a configuration framework.

## Installation

PhiJAX requires Python 3.12 or newer. Python 3.12, 3.13, and 3.14 are covered by CI.

```bash
pip install phijax
```

The default installation includes CPU-capable JAX. Select one mutually exclusive GPU extra when the matching NVIDIA
runtime is available:

```bash
pip install "phijax[cuda12]"
# or
pip install "phijax[cuda13]"
```

Optional experiment loggers are installed separately:

```bash
pip install "phijax[tensorboard]"
pip install "phijax[wandb]"
```

Extras can be combined when an environment needs GPU support and both logging integrations:

```bash
pip install "phijax[cuda13,wandb,tensorboard]"
```

Replace `cuda13` with `cuda12` for a CUDA 12 environment; do not install both CUDA extras together.

Experimental TPU environments can use `pip install "phijax[tpu]"`. The TPU, CUDA 12, and CUDA 13 extras are mutually
exclusive.

See the [installation guide](https://hangjung97.github.io/PhiJAX/getting-started/installation/) for environment
verification and development setup.

## Quickstart

The first runnable example solves a one-dimensional heat equation with initial, boundary, and PDE losses. It uses no
external data and runs on CPU after JAX completes its initial compilation:

```bash
git clone https://github.com/HangJung97/PhiJAX.git
cd PhiJAX
uv sync
JAX_PLATFORMS=cpu uv run --no-sync python examples/quickstart.py
```

`JAX_PLATFORMS=cpu` keeps the whole process strictly CPU-only. This is stronger than selecting
`Trainer(accelerator="cpu")`, which controls PhiJAX placement without changing JAX's process-wide default backend.

To run the same example on an available NVIDIA GPU:

```bash
uv sync --extra cuda13
uv run --no-sync python examples/quickstart.py --accelerator gpu
```

Replace `cuda13` with `cuda12` when using the CUDA 12 extra. If JAX reports a GPU allocation error, see the
[troubleshooting guide](https://hangjung97.github.io/PhiJAX/guides/troubleshooting/#cuda-runs-out-of-memory-during-startup).

Read the [annotated quickstart](https://hangjung97.github.io/PhiJAX/getting-started/quickstart/) for the equation,
complete source, and an explanation of each runtime object.

## Core workflow

PhiJAX keeps each part of a PINN workflow explicit while assembling the common training path for you:

```text
Model factory + objective ---> PhiModule blueprint ----+
DataModule -------------------> named batches ---------+---> Trainer.fit() ---> FitResult ---> Trainer.predict()
Optimizer + seed --------------------------------------+
LossBalancer + update policy (optional) ---------------+
```

The Trainer initializes the model, splits independent PRNG streams, prepares data sources, compiles the update, moves
batches to devices, runs hooks, restores checkpoints, and cleans up resources. `FitResult` returns the bound module and
explicit `TrainState` for prediction or advanced workflows.

Application DataModules implement `setup()`, `train_batch_source()`, and optionally `predict_batch_source()`.
Prediction is skipped cleanly when a DataModule has no prediction source.

`Trainer` manages setup and cleanup automatically across fitting and prediction.

## Core APIs

- `Trainer.fit()` assembles the common application workflow; `fit_state()` exposes explicit state and plan control.
- `Trainer.predict()` consumes a `FitResult`; `predict_state()` supports checkpoint templates and custom state.
- `TrainingPlan` and `TrainState` remain public advanced contracts for custom compiled execution.
- `BasePhiModule` and `PhiModule` provide overridable fit and prediction hooks without owning optimizers or balancers.
- `PhiDataModule` owns host data and explicit-key batch sources. The Trainer owns device placement.
- `ModelFactory` and `InitializedModel` let any JAX architecture expose a pure apply callable, state, and summary.
- `LossBalancer` supports arbitrary JAX-compatible state and exposes diagnostics without a prescribed state layout.
- Derivative, equation, objective, callback, logger, evaluation, and artifact APIs can be reused across projects.
- Hydra instantiation helpers and OmegaConf resolvers are available under `phijax.integrations` without coupling the
  Trainer to Hydra.

Only names documented in package `__all__` declarations are supported public imports.

## Checkpoints and prediction artifacts

PhiJAX saves complete training state in versioned checkpoints and produces portable prediction artifacts for
downstream evaluation. Restoration raises an error when a checkpoint version is incompatible.

## Development

Install all development dependencies, then enable the repository hooks:

```bash
uv sync --group all
uv run --no-sync pre-commit install
```

Run the local validation suite before submitting a change:

```bash
JAX_PLATFORMS=cpu uv run --no-sync pytest
uv run --no-sync ruff check .
uv run --no-sync ruff format --check .
uv run --no-sync pyright
uv run --no-sync mkdocs build --strict
uv run --no-sync pre-commit run --all-files
```

See the [documentation](https://hangjung97.github.io/PhiJAX/) for the complete API and extension guides.

## Contributing

Contributions are welcome. Read the
[contribution guide](https://github.com/HangJung97/PhiJAX/blob/main/CONTRIBUTING.md) for the development workflow,
testing expectations, documentation requirements, and pull-request checklist.

## License

PhiJAX is distributed under the [Apache License 2.0](LICENSE).
