Jaxion: differentiable astrophysics simulations

In today’s blog post, we’re going behind the scenes of a new software release: Jaxion, a differentiable astrophysics code I’ve been developing. We will look at the physics, algorithms, JAX implementation, and details of putting together the Python library. Below is a splash page showcasing a range of simulations carried out by Jaxion.

• repo: https://github.com/JaxionProject/jaxion
• docs: https://jaxion.readthedocs.io/

This package grew out of the need to rewrite and release as open-source a bunch of old code I wrote during my time as a graduate student at Harvard (2012 — 2017) and postdoc at Princeton (2017 — 2021), which I’ve been sharing with graduate students, postdocs, and professors interested in computational research on fuzzy dark matter, and has led to over a dozen publications (examples: [1], [2], [3]).

I’ve also been loving the simplicity of JAX to write numpy-like differentiable code that is deployable on GPUs.

So I decided to rewrite and release my code in JAX. The code remains clean and simple, and as a bonus, it can process billions of cell updates per second on an Nvidia H100, so it is blazing fast! Furthermore, Jaxion is a differentiable physics code. This means the simulation can be plugged in as a blackbox function to inverse problems, optimization, and inference, making use of forward- and backward-mode gradient calculations. This is an incredibly powerful concept. Differentiable physics has recently been gaining traction in the computational sciences thanks to GPUs having more memory and rapid throughput, although its roots go back to the late 1960s.

For examples of simple tutorials on writing your own differentiable simulations in Python/JAX, see my previous blog posts on the subject:

Create Your Own Differentiable N-body Simulation (With Python/JAX)

Create Your Own Automatically Differentiable Simulation (With Python/JAX)

Physics

Jaxion evolves fuzzy dark matter, isothermal gas, and collisionless stars coupled to each other by gravity.

Fuzzy dark matter is represented by a wave function, ψ, normalized such that its density is ρ_dm=|ψ|². The field has a boson mass of m and is evolved according to the Schrödinger–Poisson equations:

where the gravitational potential V is sourced by the total density of matter (dark matter + gas + stars).

For the gas, we consider the compressible isothermal Euler equations, with density ρ, velocity v, and (fixed) sound-speed c_s:

Finally, star particles (position r_s, velocity v_s, mass m) evolve according to the collisionless Boltzmann equation:

Algorithms

The Schrödinger–Poisson equations that describe fuzzy dark matter are solved using a second-order spectral solver in a periodic box, the very same algorithm I’ve presented in my “Create Your Own …” series:

Create Your Own Quantum Mechanics Simulation (With Python)

To recap at a high level, we solve the equations using a “kick-drift-kick” scheme, where a “kick” refers to a unitary operator accounting for evolution due to the potential (half-timestep):

and a “drift” refers to solving for the action of the kinetic operator in Fourier space:

For gas, the compressible isothermal Euler equations are solved using a second-order Finite Volume scheme, which we’ve also covered in a previous blog post:

Create Your Own Finite Volume Fluid Simulation (With Python)

The main idea of the solver is to update the fluid variables (Q) conservatively by calculating second-order accurate estimates of fluxes (F) through cell faces.

Finally, stars are evolved (collisionlessly) according to gravitational forces using a Particle-Mesh approach. Collisionless means we are not solving the N-body problem of pair-wise interactions, but rather use a particle-in-cell approach to calculate the gravitational potential field and use it to update particle accelerations. This is, in fact, the same method as when we created the Particle-in-Cell (PIC) solver for the electric potential (except with opposite sign):

Create Your Own Plasma PIC Simulation (With Python)

At a high level, we deposit the mass of star particles into cells (using cloud-in-cell binning) and solve the Poisson equation on the grid to update the acceleration of the particles.

Implementation

The core engine of Jaxion is in contained in the file jaxion/simulation.py . The main simulation loop likes something like:

@jax.jit
def _update(_, carry):
    # Update the simulation state by one timestep
    # according to a 2nd-order `kick-drift-kick` scheme
    state, kx, ky, kz, k_sq = carry
    _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
    _drift(state, k_sq, dt)
    _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
    state["t"] += dt
    return state, kx, ky, kz, k_sq

which gets jax jit compiled, and carries info on the simulation state (dark matter [wave function], gas [density, velocity], star particles [position, velocity]) and Fourier variables

Get Philip Mocz’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The kick step accelerates the wave function, fluid velocity, and star particles according to self-gravity:

def _kick(state, kx, ky, kz, k_sq, dt):
    # Kick (half-step)
    V = self._calc_grav_potential(state, k_sq)
    if self.params["physics"]["quantum"]:
        state["psi"] = quantum_kick(state["psi"], V, m_per_hbar, dt)
    if self.params["physics"]["hydro"]:
        state["vx"], state["vy"], state["vz"] = hydro_accelerate(
            state["vx"], state["vy"], state["vz"], V, kx, ky, kz, dt
        )
    if self.params["physics"]["particles"]:
        state["vel"] = particles_accelerate(
            state["vel"], state["pos"], V, kx, ky, kz, dx, dt
        )

And the drift step:

def _drift(state, k_sq, dt):
    # Drift (full-step)
    if self.params["physics"]["quantum"]:
        state["psi"] = quantum_drift(state["psi"], k_sq, m_per_hbar, dt)
    if self.params["physics"]["hydro"]:
        state["rho"], state["vx"], state["vy"], state["vz"] = hydro_fluxes(
            state["rho"], state["vx"], state["vy"], state["vz"], dt, dx, c_sound
        )
    if self.params["physics"]["particles"]:
        state["pos"] = particles_drift(state["pos"], state["vel"], dt, box_size)

applies the kinetic operator to the wave function, the hydro fluxes to the hydro state, and simple drifts to the star particle positions.

Implementation of these operators can be found in physics files: quantum.py, hydro.py, and particles.py.

Sharding

To make code run on distributed GPU clusters, JAX uses the concept of sharding arrays/matrices/tensors across GPUs. We won’t go over details here, but present a simple example of how to obtain, for example, a sharded meshgrid function in JAX to avoid running out of memory:

def xmeshgrid(x_lin):
    xx, yy, zz = jnp.meshgrid(x_lin, x_lin, x_lin, indexing="ij")
    return xx, yy, zz

xmeshgrid_jit = jax.jit(
    xmeshgrid, in_shardings=None, out_shardings=sharding
)

Making a Python Package

Creating the Python package was a lot of fun. The package is hosted PyPI, so anyone can obtain the Python library with:

pip install jaxion

Code quality and style are enforced with Ruff, the fastest Python linter out there.

The project unit and integration tests are tested with Pytest. A coverage report is generated and uploaded to codecov.io. GitHub Actions run tests on every commit to the code.

Benchmarking (keeping track of memory usage and speed) is done with Airspeed Velocity (asv).

Documentation is hosted at readthedocs.

Example: Soliton Binary Merger

Jaxion is easy to use. It involves [1] setting up simple parameters in a Python dictionary (only params that overwrite defaults need to be specified), and [2] writing initial conditions.

Below is a fun example of two fuzzy dark matter solitons, initialized with opposite phases, colliding. Having opposite phases, the quantum “pressure” initially repels the two, causing them to bounce.

The Python script to set up and run the simulation and obtain these figures looks like:

import jax
import jax.numpy as jnp
import os
import jaxion


"""
Collide two solitons with opposite phases
Philip Mocz (2025)
"""


def set_up_simulation():
    # Parameters changed from default values
    params = {}

    # Initialize the simulation
    sim = jaxion.Simulation(params)

    # Set initial conditions (two solitons)
    # Jaxion uses units of [L] = kpc, [V] = km/s, [M] = M_sun
    m_22 = sim.params["quantum"]["m_22"]
    box_size = sim.params["domain"]["box_size"]
    r_soliton = 0.02 * box_size
    xx, yy, zz = sim.grid

    def rho_soliton(r, r_soliton, m_22):
        return (
            1.9e7 * m_22**-2 * r_soliton**-4 / (1.0 + 0.091 * (r / r_soliton) ** 2) ** 8
        )

    rho = 0.0
    r_soliton = 0.02 * box_size
    for i in range(2):
        x_soliton = (0.5 + 0.3 * (i - 0.5)) * box_size
        y_soliton = 0.5 * box_size
        z_soliton = 0.5 * box_size

        r = jnp.sqrt(
            (xx - x_soliton) ** 2 + (yy - y_soliton) ** 2 + (zz - z_soliton) ** 2
        )
        rho += rho_soliton(r, r_soliton, m_22)

    half = (xx - 0.5 * box_size) > 0
    sim.state["psi"] = jnp.array(jnp.sqrt(rho) + 0j) * (
        1.0 * half + -1.0 * (1.0 - half)
    )

    return sim


def main():
    sim = set_up_simulation()
    sim.run()


if __name__ == "__main__":
    main()

Example: Inverse Problem

Next, we show off the power of differentiable simulations by solving an example inverse problem. The problem setup includes uniformly distributed dark matter and stars (coupled via gravity). The question posed is what initial conditions for the star velocities evolve the system to reach a target density (“J” shape, for Jaxion) in the dark matter at time t=1?

We can in fact find a solution in a few seconds:

The simulation can be plugged into Google Deepmind’s optax library’s L-BFGS solver to solve the optimization problem efficiently using gradients. The code to do so looks something like:

import jax
import jax.numpy as jnp
import jaxion
import chex
from typing import NamedTuple
import optax
import matplotlib.image as img


"""
Inverse problem: Find initial conditions for velocity (dm+stars) to achieve target density at t=1
Philip Mocz (2025)
"""


def set_up_simulation(save=False):
    # Parameters added/changed from default values
    params = {
        "physics": {
            "particles": True,
        },
        "domain": {
            "resolution_base": 32,
            "box_size": 20.0,  # kpc
        },
        "time": {
            "end": 1.0,
        },
        "output": {
            "save": save,
            "plot_dynamic_range": 10.0,
        },
        "particles": {
            "num_particles": 64,
            "particle_mass": 1.0e7,  # Msun
        },
    }

    # Initialize the simulation
    sim = jaxion.Simulation(params)

    # dark matter
    rho_dm = 1.0e5  # (Msun / kpc^3)

    box_size = sim.box_size
    nx = sim.resolution

    # dark matter
    psi = jnp.sqrt(rho_dm) * jnp.ones((nx, nx, nx)) + 0j

    # stars are initially uniformly distributed
    num_stars = params["particles"]["num_particles"]
    # Arrange positions in a uniform grid in x-y plane, z=box_size/2
    side = int(num_stars ** (1 / 2))
    xlin = jnp.linspace(0, box_size, side, endpoint=False) + box_size / (2 * side)
    ylin = jnp.linspace(0, box_size, side, endpoint=False) + box_size / (2 * side)
    zlin = jnp.array([box_size / 2.0])
    xx, yy, zz = jnp.meshgrid(xlin, ylin, zlin, indexing="ij")
    pos = jnp.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T
    vel = jnp.zeros((num_stars, 3))

    sim.state["psi"] = psi
    sim.state["pos"] = pos
    sim.state["vel"] = vel

    return sim


class InfoState(NamedTuple):
    iter_num: chex.Numeric


def print_info():
    def init_fn(params):
        del params
        return InfoState(iter_num=0)

    def update_fn(updates, state, params, *, value, grad, **extra_args):
        del params, extra_args

        jax.debug.print(
            "Iteration: {i}, Loss: {v:.2e}, |grad|: {e:.2e}",
            i=state.iter_num,
            v=value,
            e=optax.tree_utils.tree_norm(grad),
        )
        return updates, InfoState(iter_num=state.iter_num + 1)

    return optax.GradientTransformationExtraArgs(init_fn, update_fn)


def run_opt(init_params, fun, opt, max_iter, tol):
    value_and_grad_fun = optax.value_and_grad_from_state(fun)

    def step(carry):
        params, state = carry
        value, grad = value_and_grad_fun(params, state=state)
        updates, state = opt.update(
            grad, state, params, value=value, grad=grad, value_fn=fun
        )
        params = optax.apply_updates(params, updates)
        return params, state

    def continuing_criterion(carry):
        _, state = carry
        iter_num = optax.tree_utils.tree_get(state, "count")
        grad = optax.tree_utils.tree_get(state, "grad")
        err = optax.tree_utils.tree_norm(grad)
        return (iter_num == 0) | ((iter_num < max_iter) & (err >= tol))

    init_carry = (init_params, opt.init(init_params))
    final_params, final_state = jax.lax.while_loop(
        continuing_criterion, step, init_carry
    )
    return final_params, final_state


def solve_inverse_problem():
    """Optimize the initial star positions to recreate the logo"""
    # Load the target density field
    target_data = img.imread("target.png")[:, :, 0]
    target_data = target_data[::2, ::2]  # downsample
    target = jnp.flipud(jnp.array(target_data, dtype=float)).T
    target = 1.0 - 1.6 * (target - 0.5)
    target /= jnp.mean(target)

    @jax.jit
    def loss_function(theta):
        sim = set_up_simulation()
        sim.state["vel"] = theta

        sim.run()

        projected_density = jnp.mean(jnp.abs(sim.state["psi"]) ** 2, axis=2)
        norm = jnp.mean(projected_density)
        projected_density /= norm

        error_norm = jnp.mean((projected_density - target) ** 2)

        return error_norm

    # opt = optax.lbfgs()
    opt = optax.chain(print_info(), optax.lbfgs())

    sim = set_up_simulation()
    init_params = sim.state["vel"]

    print(
        f"Initial value: {loss_function(init_params):.2e} "
        f"Initial gradient norm: {optax.tree_utils.tree_norm(jax.grad(loss_function)(init_params)):.2e}"
    )
    final_params, _ = run_opt(init_params, loss_function, opt, max_iter=100, tol=1e-5)
    print(
        f"Final value: {loss_function(final_params):.2e}, "
        f"Final gradient norm: {optax.tree_utils.tree_norm(jax.grad(loss_function)(final_params)):.2e}"
    )

    return final_params


def main():
    optimized_ics = solve_inverse_problem()

    sim = set_up_simulation(save=True)
    sim.state["vel"] = optimized_ics
    sim.run()


if __name__ == "__main__":
    main()

Try out Jaxion

Try out Jaxion today ( just pip install jaxion), and explore how interesting configurations of dark matter, gas, and stars evolve!

If you’d like to use Jaxion for your research, you may cite it as:

@software{Mocz_Jaxion_2025,
   author = {Mocz, Philip},
      doi = {10.5281/zenodo.17438467},
    month = oct,
    title = {{Jaxion}},
      url = {https://github.com/JaxionProject/jaxion},
  version = {0.0.4},
     year = {2025}
}

Enjoy!

About the author

Philip Mocz is a scientific software engineer at the Flatiron Institute. He received his Ph.D. (’17) and A.B. (’12) from Harvard University, where he developed high-performance computing algorithms for simulations of astrophysical systems. If you’ve enjoyed this post, give the story a clap 👏 and 👉 follow me for more scientific computing tutorials with Python!