pardax 1.0.0

JAX-native solver for initial value problems

pardax

A JAX-native solver for initial value problems.

pardax provides a variety of composable time-stepping schemes that work seamlessly with JAX transformations, along with a familiar interface inspired by scipy.integrate.solve_ivp

Check out the documentation for more details.

Installation

pip install pardax

or with uv:

uv add pardax

Quick start

import pardax as pdx

# 1. Define your ODE
# NOTE: must be functionally pure (JAX-compatible)
def my_pde_fun(t: float, y: jax.Array, params: PyTree):
    """dy/dt = f(t, y, ...)"""
    ...

# 2. Set initial condition
y0 = ...

# 3. Choose time-stepping method
method = pdx.RK4()

# 4. Integrate
t, y = pdx.solve_ivp(
    my_pde_fun,
    t_span=(0.0, 1.0),
    y0=y0,
    stepper=method,
    step_size=0.001,
    params=params
)

Contributing

See CONTRIBUTING.md for instructions on setting up a development environment, running tests, linting, and building the documentation locally.