jaxion 0.0.2__tar.gz → 0.0.4__tar.gz

jaxion-0.0.4/PKG-INFO

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: jaxion
+Version: 0.0.4
+Summary: A differentiable simulation library for fuzzy dark matter in JAX
+Author-email: Philip Mocz <philip.mocz@gmail.com>
+License-Expression: Apache-2.0
+Project-URL: Documentation, https://jaxion.readthedocs.io
+Project-URL: Homepage, https://github.com/JaxionProject/jaxion
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax==0.5.3
+Requires-Dist: jaxdecomp==0.2.7
+Requires-Dist: tensorflow
+Requires-Dist: orbax-checkpoint==0.11.18
+Requires-Dist: optax==0.2.5
+Requires-Dist: numpy
+Requires-Dist: matplotlib
+Requires-Dist: setuptools>=70.1.1
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]==0.5.3; extra == "cuda12"
+Dynamic: license-file
+
+<p align="center">
+  <a href="https://jaxion.readthedocs.io">
+    <img src="docs/_static/jaxion-logo.svg" alt="jaxion logo" width="128"/>
+  </a>
+</p>
+
+# jaxion
+
+[![Repo Status][status-badge]][status-link]
+[![PyPI Version Status][pypi-badge]][pypi-link]
+[![Test Status][workflow-test-badge]][workflow-test-link]
+[![Coverage][coverage-badge]][coverage-link]
+[![Readthedocs Status][docs-badge]][docs-link]
+[![License][license-badge]][license-link]
+
+[status-link]:         https://www.repostatus.org/#active
+[status-badge]:        https://www.repostatus.org/badges/latest/active.svg
+[pypi-link]:           https://pypi.org/project/jaxion
+[pypi-badge]:          https://img.shields.io/pypi/v/jaxion?label=PyPI&logo=pypi
+[workflow-test-link]:  https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml
+[workflow-test-badge]: https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml/badge.svg?event=push
+[coverage-link]:       https://app.codecov.io/gh/JaxionProject/jaxion
+[coverage-badge]:      https://codecov.io/github/jaxionproject/jaxion/graph/jaxion-server/badge.svg
+[docs-link]:           https://jaxion.readthedocs.io
+[docs-badge]:          https://readthedocs.org/projects/jaxion/badge
+[license-link]:        https://opensource.org/licenses/Apache-2.0
+[license-badge]:       https://img.shields.io/badge/License-Apache_2.0-blue.svg
+
+A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
+
+Author: [Philip Mocz (@pmocz)](https://github.com/pmocz/)
+
+Jaxion is built for multi-GPU scalability and is fully differentiable. It is a high-performance JAX-based simulation library for modeling fuzzy dark matter alongside stars, gas, and cosmological dynamics. Being differentiable, Jaxion can seamlessly integrate with pipelines for inverse-problems, inference, optimization, and coupling to ML models.
+
+Jaxion is the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
+
+
+## Getting started
+
+Install with:
+
+```console
+pip install jaxion
+```
+
+or, for GPU support, use:
+
+```console
+pip install jaxion[cuda12]
+```
+
+See the docs for more info on how to [build from source](https://jaxion.readthedocs.io/en/latest/pages/installation.html).
+
+
+## Examples
+
+Check out the `examples/` directory for demonstrations of using Jaxion.
+
+<p align="center">
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">
+    <img src="examples/cosmological_box/movie.gif" alt="cosmological_box" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/dynamical_friction">
+    <img src="examples/dynamical_friction/movie.gif" alt="dynamical_friction" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_gas">
+    <img src="examples/heating_gas/movie.gif" alt="heating_gas" width="128"/>
+  </a>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_stars">
+    <img src="examples/heating_stars/movie.gif" alt="heating_stars" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/kinetic_condensation">
+    <img src="examples/kinetic_condensation/movie.gif" alt="kinetic_condensation" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/logo">
+    <img src="examples/logo/movie.gif" alt="logo" width="128"/>
+  </a>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_binary_merger">
+    <img src="examples/soliton_binary_merger/movie.gif" alt="soliton_binary_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_merger">
+    <img src="examples/soliton_merger/movie.gif" alt="soliton_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/tidal_stripping">
+    <img src="examples/tidal_stripping/movie.gif" alt="tidal_stripping" width="128"/>
+  </a>
+</p>
+
+
+## Links
+
+* [Code repository](https://github.com/JaxionProject/jaxion) on GitHub (this page).
+* [Documentation](https://jaxion.readthedocs.io) for up-to-date information about installing and using jaxion.
+
+
+## Testing
+
+Jaxion is tested with `pytest`. Tests are included in the `tests/` folder.
+
+
+## Contributing
+
+Jaxion welcomes community contributions of all kinds. Open an issue or fork the code and submit a Pull request. Please check out the [Contributing Guidelines](CONTRIBUTING.md)
+
+
+## Cite this repository
+
+TODO XXX

jaxion-0.0.4/README.md

@@ -0,0 +1,110 @@
+<p align="center">
+  <a href="https://jaxion.readthedocs.io">
+    <img src="docs/_static/jaxion-logo.svg" alt="jaxion logo" width="128"/>
+  </a>
+</p>
+
+# jaxion
+
+[![Repo Status][status-badge]][status-link]
+[![PyPI Version Status][pypi-badge]][pypi-link]
+[![Test Status][workflow-test-badge]][workflow-test-link]
+[![Coverage][coverage-badge]][coverage-link]
+[![Readthedocs Status][docs-badge]][docs-link]
+[![License][license-badge]][license-link]
+
+[status-link]:         https://www.repostatus.org/#active
+[status-badge]:        https://www.repostatus.org/badges/latest/active.svg
+[pypi-link]:           https://pypi.org/project/jaxion
+[pypi-badge]:          https://img.shields.io/pypi/v/jaxion?label=PyPI&logo=pypi
+[workflow-test-link]:  https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml
+[workflow-test-badge]: https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml/badge.svg?event=push
+[coverage-link]:       https://app.codecov.io/gh/JaxionProject/jaxion
+[coverage-badge]:      https://codecov.io/github/jaxionproject/jaxion/graph/jaxion-server/badge.svg
+[docs-link]:           https://jaxion.readthedocs.io
+[docs-badge]:          https://readthedocs.org/projects/jaxion/badge
+[license-link]:        https://opensource.org/licenses/Apache-2.0
+[license-badge]:       https://img.shields.io/badge/License-Apache_2.0-blue.svg
+
+A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
+
+Author: [Philip Mocz (@pmocz)](https://github.com/pmocz/)
+
+Jaxion is built for multi-GPU scalability and is fully differentiable. It is a high-performance JAX-based simulation library for modeling fuzzy dark matter alongside stars, gas, and cosmological dynamics. Being differentiable, Jaxion can seamlessly integrate with pipelines for inverse-problems, inference, optimization, and coupling to ML models.
+
+Jaxion is the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
+
+
+## Getting started
+
+Install with:
+
+```console
+pip install jaxion
+```
+
+or, for GPU support, use:
+
+```console
+pip install jaxion[cuda12]
+```
+
+See the docs for more info on how to [build from source](https://jaxion.readthedocs.io/en/latest/pages/installation.html).
+
+
+## Examples
+
+Check out the `examples/` directory for demonstrations of using Jaxion.
+
+<p align="center">
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">
+    <img src="examples/cosmological_box/movie.gif" alt="cosmological_box" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/dynamical_friction">
+    <img src="examples/dynamical_friction/movie.gif" alt="dynamical_friction" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_gas">
+    <img src="examples/heating_gas/movie.gif" alt="heating_gas" width="128"/>
+  </a>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_stars">
+    <img src="examples/heating_stars/movie.gif" alt="heating_stars" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/kinetic_condensation">
+    <img src="examples/kinetic_condensation/movie.gif" alt="kinetic_condensation" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/logo">
+    <img src="examples/logo/movie.gif" alt="logo" width="128"/>
+  </a>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_binary_merger">
+    <img src="examples/soliton_binary_merger/movie.gif" alt="soliton_binary_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_merger">
+    <img src="examples/soliton_merger/movie.gif" alt="soliton_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/tidal_stripping">
+    <img src="examples/tidal_stripping/movie.gif" alt="tidal_stripping" width="128"/>
+  </a>
+</p>
+
+
+## Links
+
+* [Code repository](https://github.com/JaxionProject/jaxion) on GitHub (this page).
+* [Documentation](https://jaxion.readthedocs.io) for up-to-date information about installing and using jaxion.
+
+
+## Testing
+
+Jaxion is tested with `pytest`. Tests are included in the `tests/` folder.
+
+
+## Contributing
+
+Jaxion welcomes community contributions of all kinds. Open an issue or fork the code and submit a Pull request. Please check out the [Contributing Guidelines](CONTRIBUTING.md)
+
+
+## Cite this repository
+
+TODO XXX

{jaxion-0.0.2 → jaxion-0.0.4}/jaxion/__init__.py

@@ -1,2 +1,3 @@
 from .simulation import Simulation as Simulation
 from .constants import constants as constants
+from .analysis import radial_power_spectrum as radial_power_spectrum

jaxion-0.0.4/jaxion/analysis.py

@@ -0,0 +1,58 @@
+import jax.numpy as jnp
+import jaxdecomp as jd
+
+
+def radial_power_spectrum(data_cube, kx, ky, kz, box_size):
+    """
+    Computes the radially averaged power spectral density of a 3D datacube.
+
+    Parameters
+    ----------
+        data_cube : jnp.ndarray
+            3D data cube, must be cubic
+        kx, ky, kz: jnp.ndarray
+            wavenumber grids in each dimension
+        box_size: float
+            physical size of box
+
+    Returns
+    -------
+        Pf: jnp.ndarray
+            radial power spectrum
+        k: jnp.ndarray
+            wavenumbers
+        total_power: float
+            total power
+    """
+    dim = data_cube.ndim
+    nx = data_cube.shape[0]
+    dx = box_size / nx
+
+    # Compute power spectrum
+    data_cube_hat = jd.fft.pfft3d(data_cube)
+    total_power = 0.5 * jnp.sum(jnp.abs(data_cube_hat) ** 2) / nx**dim * dx**dim
+    phi_k = 0.5 * jnp.abs(data_cube_hat) ** 2 / nx**dim * dx**dim
+    half_size = nx // 2 + 1
+
+    # Compute radially-averaged power spectrum
+    # if dim == 2:
+    #    k_r = jnp.sqrt(kx**2 + ky**2)
+    k_r = jnp.sqrt(kx**2 + ky**2 + kz**2)
+
+    Pf, _ = jnp.histogram(
+        k_r, range=(-0.5, half_size - 0.5), bins=half_size, weights=phi_k
+    )
+    norm, _ = jnp.histogram(k_r, range=(-0.5, half_size - 0.5), bins=half_size)
+    Pf /= norm + (norm == 0)
+
+    k = 2.0 * jnp.pi * jnp.arange(half_size) / box_size
+    dk = 2.0 * jnp.pi / box_size
+
+    Pf /= dk**dim
+
+    # Add geometrical factor
+    # if dim == 2:
+    #     Pf = Pf * 2.0 * jnp.pi * k
+    Pf *= 4.0 * jnp.pi * k**2
+
+    return Pf, k, total_power

{jaxion-0.0.2 → jaxion-0.0.4}/jaxion/constants.py

@@ -1,13 +1,13 @@
 # Copyright (c) 2025 The Jaxion Team.
 
-# Physical Constants
-#
-# Jaxion uses a unit system of:
-# [L] = kpc
-# [V] = km/s
-# [M] = Msun
-#
-# other units are derived from these base units, e.g., [T] = [L]/[V] = kpc / (km/s) ~= 0.978 Gyr
+"""
+Physical constants in units of:
+[L] = kpc,
+[V] = km/s,
+[M] = Msun
+
+note: other units are derived from these base units, e.g., [T] = [L]/[V] = kpc / (km/s) ~= 0.978 Gyr
+"""
 
 constants = {
     "gravitational_constant": 4.30241002e-6,  # G / (kpc * (km/s)^2 / Msun)

jaxion-0.0.4/jaxion/cosmology.py

@@ -0,0 +1,74 @@
+import jax.numpy as jnp
+
+# Pure functions for cosmology simulation
+
+
+def get_physical_time_interval(z_start, z_end, omega_matter, omega_lambda, little_h):
+    """Compute the total physical time between two redshifts."""
+    #  da/dt = H0 * sqrt(omega_matter / a + omega_lambda * a^2)
+    a_start = 1.0 / (1.0 + z_start)
+    a_end = 1.0 / (1.0 + z_end)
+    H0 = 0.1 * little_h  # Hubble constant in (km/s/kpc)
+    n_quad = 10000
+    a_lin = jnp.linspace(a_start, a_end, n_quad)
+    a_dot = H0 * jnp.sqrt((omega_matter / a_lin) + (omega_lambda * a_lin**2))
+    dt_hat = jnp.trapezoid(1.0 / a_dot, a_lin)
+    return dt_hat
+
+
+def get_supercomoving_time_interval(
+    z_start, z_end, omega_matter, omega_lambda, little_h
+):
+    """Compute the total supercomoving time (dt_hat = a^-2 dt) between two redshifts."""
+    #  da/dt = H0 * sqrt(omega_matter / a + omega_lambda * a^2)
+    #  da/dt_hat = a^2 * da/dt
+    #  dt_hat/da = a^-2 / (da/dt)
+    a_start = 1.0 / (1.0 + z_start)
+    a_end = 1.0 / (1.0 + z_end)
+    H0 = 0.1 * little_h  # Hubble constant in (km/s/kpc)
+    n_quad = 10000
+    a_lin = jnp.linspace(a_start, a_end, n_quad)
+    a_dot = H0 * jnp.sqrt((omega_matter / a_lin) + (omega_lambda * a_lin**2))
+    dt_hat_da = a_lin**-2 / a_dot
+    dt_hat = jnp.trapezoid(dt_hat_da, a_lin)
+    return dt_hat
+
+
+def get_scale_factor(z_start, dt_hat, omega_matter, omega_lambda, little_h):
+    """Compute the scale factor corresponding to a given supercomoving time,
+    by root finding."""
+    a_start = 1.0 / (1.0 + z_start)
+    a = a_start
+    tolerance = 1e-6
+    max_iterations = 100
+    lower_bound = 0.0
+    upper_bound = 2.0  # Set a reasonable upper bound for the scale factor
+    for _ in range(max_iterations):
+        dt_hat_guess = get_supercomoving_time_interval(
+            z_start, 1.0 / a - 1.0, omega_matter, omega_lambda, little_h
+        )
+        error = dt_hat_guess - dt_hat
+        if jnp.abs(error) < tolerance:
+            break
+        # Use bisection method
+        if error > 0:
+            upper_bound = a  # Move upper bound down
+        else:
+            lower_bound = a  # Move lower bound up
+        a = (lower_bound + upper_bound) / 2  # Update scale factor to midpoint
+    return a
+
+
+def get_next_scale_factor(z_start, dt_hat, omega_matter, omega_lambda, little_h):
+    """Advance scale factor by dt_hat using RK2 (midpoint) on da/dt_hat = a^2 * da/dt."""
+    a_start = 1.0 / (1.0 + z_start)
+    H0 = 0.1 * little_h
+
+    def g(a):
+        a_dot = H0 * jnp.sqrt((omega_matter / a) + (omega_lambda * a**2))
+        return a**2 * a_dot  # da/dt_hat
+
+    k1 = g(a_start)
+    k2 = g(a_start + 0.5 * dt_hat * k1)
+    a = a_start + dt_hat * k2
+    return a

jaxion-0.0.4/jaxion/gravity.py

@@ -0,0 +1,10 @@
+import jax.numpy as jnp
+import jaxdecomp as jd
+
+# Pure functions for gravity calculations
+
+
+def calculate_gravitational_potential(rho, k_sq, G, rho_bar):
+    Vhat = -jd.fft.pfft3d(4.0 * jnp.pi * G * (rho - rho_bar)) / (k_sq + (k_sq == 0))
+    V = jnp.real(jd.fft.pifft3d(Vhat))
+    return V

{jaxion-0.0.2 → jaxion-0.0.4}/jaxion/hydro.py

@@ -1,4 +1,5 @@
 import jax.numpy as jnp
+import jaxdecomp as jd
 
 # Pure functions for hydro simulation
 
@@ -89,11 +90,11 @@ def get_flux(rho_L, vx_L, vy_L, vz_L, rho_R, vx_R, vy_R, vz_R, cs):
 
 
 def hydro_accelerate(vx, vy, vz, V, kx, ky, kz, dt):
-    V_hat = jnp.fft.fftn(V)
+    V_hat = jd.fft.pfft3d(V)
 
-    ax = -jnp.real(jnp.fft.ifftn(1.0j * kx * V_hat))
-    ay = -jnp.real(jnp.fft.ifftn(1.0j * ky * V_hat))
-    az = -jnp.real(jnp.fft.ifftn(1.0j * kz * V_hat))
+    ax = -jnp.real(jd.fft.pifft3d(1.0j * kx * V_hat))
+    ay = -jnp.real(jd.fft.pifft3d(1.0j * ky * V_hat))
+    az = -jnp.real(jd.fft.pifft3d(1.0j * kz * V_hat))
 
     vx += ax * dt
     vy += ay * dt

jaxion-0.0.4/jaxion/params_default.json

@@ -0,0 +1,114 @@
+{
+  "physics": {
+    "quantum": { 
+      "default": true,
+      "description": "switch on for fuzzy dark matter."
+    },
+    "gravity": { 
+      "default": true,
+      "description": "switch on for self-gravity."
+    },
+    "hydro": { 
+      "default": false,
+      "description": "switch on for hydrodynamics (isothermal)."
+    },
+    "particles": { 
+      "default": false,
+      "description": "switch on for particle-mesh particles."
+    },
+    "cosmology": { 
+      "default": false,
+      "description": "switch on for cosmological factors/comoving units."
+    },
+    "external_potential": { 
+      "default": false,
+      "description": "switch on for external gravitational potential."
+    }
+  },
+  "domain": {
+    "box_size": { 
+      "default": 10.0,
+      "description": "periodic domain box size [kpc]."
+    },
+    "resolution_base": { 
+      "default": 32,
+      "description": "base resolution per linear dimension."
+    },
+    "resolution_multiplier": { 
+      "default": 1,
+      "description": "resolution multiplier."
+    }
+  },
+  "time": {
+    "start": { 
+      "default": 0.0,
+      "description": "simulation start time [kpc/(km/s)] or [redshift] (cosmology=true)."
+    },
+    "end": { 
+      "default": 1.0,
+      "description": "simulation end time [kpc/(km/s)] or [redshift] (cosmology=true)."
+    },
+    "adaptive": { 
+      "default": false,
+      "description": "switch on for adaptive time stepping."
+    }
+  },
+  "output": {
+    "path": {
+      "default": "./checkpoints",
+      "description": "path to output directory."
+    },
+    "num_checkpoints": {
+      "default": 100,
+      "description": "number of checkpoints to save."
+    },
+    "save": {
+      "default": true,
+      "description": "switch on to save checkpoints."
+    },
+    "plot_dynamic_range": {
+      "default": 100.0,
+      "description": "dynamic range for plotting."
+    }
+  },
+  "quantum": {
+    "m_22": { 
+      "default": 1.0,
+      "description": "axion mass [10^{-22} eV]."
+    }
+  },
+  "hydro": {
+    "sound_speed": { 
+      "default": 1.0,
+      "description": "isothermal sound speed [km/s]."
+    }
+  },
+  "particles": {
+    "num_particles": {
+      "default": 0,
+      "description": "number of particles."
+    },
+    "particle_mass": {
+      "default": 1.0,
+      "description": "particle mass [M_sun]."
+    }
+  },
+  "cosmology": {
+    "omega_matter": {
+      "default": 0.3,
+      "description": "matter density parameter."
+    },
+    "omega_lambda": {
+      "default": 0.7,
+      "description": "dark energy density parameter."
+    },
+    "little_h": {
+      "default": 0.7,
+      "description": "Hubble parameter little h (H_0 = 100 h km/s/Mpc)."
+    }
+  },
+  "version": {
+    "default": "unknown",
+    "description": "jaxion version used (auto detected)."
+  }
+}

{jaxion-0.0.2 → jaxion-0.0.4}/jaxion/particles.py

@@ -1,5 +1,6 @@
 import jax
 import jax.numpy as jnp
+import jaxdecomp as jd
 
 # Pure functions for particle-mesh calculations (stars, BHs, ...)
 
@@ -64,10 +65,10 @@ def get_acceleration(pos, V, kx, ky, kz, dx):
     i, ip1, w_i, w_ip1 = get_cic_indices_and_weights(pos, dx, resolution)
 
     # find accelerations on the grid
-    V_hat = jnp.fft.fftn(V)
-    ax = -jnp.real(jnp.fft.ifftn(1.0j * kx * V_hat))
-    ay = -jnp.real(jnp.fft.ifftn(1.0j * ky * V_hat))
-    az = -jnp.real(jnp.fft.ifftn(1.0j * kz * V_hat))
+    V_hat = jd.fft.pfft3d(V)
+    ax = -jnp.real(jd.fft.pifft3d(1.0j * kx * V_hat))
+    ay = -jnp.real(jd.fft.pifft3d(1.0j * ky * V_hat))
+    az = -jnp.real(jd.fft.pifft3d(1.0j * kz * V_hat))
     a_grid = jnp.stack((ax, ay, az), axis=-1)
 
     # interpolate the accelerations to the particle positions

jaxion-0.0.4/jaxion/quantum.py

@@ -0,0 +1,95 @@
+import jax.numpy as jnp
+import jaxdecomp as jd
+
+# Pure functions for quantum simulation
+
+
+def quantum_kick(psi, V, m_per_hbar, dt):
+    psi = jnp.exp(-1.0j * m_per_hbar * dt * V) * psi
+    return psi
+
+
+def quantum_drift(psi, k_sq, m_per_hbar, dt):
+    psi_hat = jd.fft.pfft3d(psi)
+    psi_hat = jnp.exp(dt * (-1.0j * k_sq / m_per_hbar / 2.0)) * psi_hat
+    psi = jd.fft.pifft3d(psi_hat)
+    return psi
+
+
+def get_gradient(psi, kx, ky, kz):
+    """
+    Computes gradient of wavefunction psi
+    psi: jnp.ndarray (3D)
+    Returns: grad_psi_x, grad_psi_y, grad_psi_z
+    """
+    psi_hat = jd.fft.pfft3d(psi)
+    grad_psi_x_hat = 1.0j * kx * psi_hat
+    grad_psi_y_hat = 1.0j * ky * psi_hat
+    grad_psi_z_hat = 1.0j * kz * psi_hat
+
+    grad_psi_x = jd.fft.pifft3d(grad_psi_x_hat)
+    grad_psi_y = jd.fft.pifft3d(grad_psi_y_hat)
+    grad_psi_z = jd.fft.pifft3d(grad_psi_z_hat)
+
+    return grad_psi_x, grad_psi_y, grad_psi_z
+
+
+def quantum_velocity(psi, box_size, m_per_hbar):
+    """
+    Compute the velocity from the wave-function
+    v = nabla S / m
+    psi = sqrt(rho) exp(i S / hbar)
+    """
+    N = psi.shape[0]
+    dx = box_size / N
+
+    S_per_hbar = jnp.angle(psi)
+
+    # Central differences with phase unwrapping
+    # vx = jnp.roll(S_per_hbar, -1, axis=0) - jnp.roll(S_per_hbar, 1, axis=0)
+    # vy = jnp.roll(S_per_hbar, -1, axis=1) - jnp.roll(S_per_hbar, 1, axis=1)
+    # vz = jnp.roll(S_per_hbar, -1, axis=2) - jnp.roll(S_per_hbar, 1, axis=2)
+    # vx = jnp.where(vx > jnp.pi, vx - 2 * jnp.pi, vx)
+    # vx = jnp.where(vx <= -jnp.pi, vx + 2 * jnp.pi, vx)
+    # vy = jnp.where(vy > jnp.pi, vy - 2 * jnp.pi, vy)
+    # vy = jnp.where(vy <= -jnp.pi, vy + 2 * jnp.pi, vy)
+    # vz = jnp.where(vz > jnp.pi, vz - 2 * jnp.pi, vz)
+    # vz = jnp.where(vz <= -jnp.pi, vz + 2 * jnp.pi, vz)
+    # vx = vx / (2.0 * dx) / m_per_hbar
+    # vy = vy / (2.0 * dx) / m_per_hbar
+    # vz = vz / (2.0 * dx) / m_per_hbar
+
+    # Forward differences
+    vx1 = jnp.roll(S_per_hbar, -1, axis=0) - S_per_hbar
+    vy1 = jnp.roll(S_per_hbar, -1, axis=1) - S_per_hbar
+    vz1 = jnp.roll(S_per_hbar, -1, axis=2) - S_per_hbar
+    vx1 = jnp.where(vx1 > jnp.pi, vx1 - 2 * jnp.pi, vx1)
+    vx1 = jnp.where(vx1 <= -jnp.pi, vx1 + 2 * jnp.pi, vx1)
+    vy1 = jnp.where(vy1 > jnp.pi, vy1 - 2 * jnp.pi, vy1)
+    vy1 = jnp.where(vy1 <= -jnp.pi, vy1 + 2 * jnp.pi, vy1)
+    vz1 = jnp.where(vz1 > jnp.pi, vz1 - 2 * jnp.pi, vz1)
+    vz1 = jnp.where(vz1 <= -jnp.pi, vz1 + 2 * jnp.pi, vz1)
+    vx1 = vx1 / dx / m_per_hbar
+    vy1 = vy1 / dx / m_per_hbar
+    vz1 = vz1 / dx / m_per_hbar
+
+    # Backward differences
+    vx2 = S_per_hbar - jnp.roll(S_per_hbar, 1, axis=0)
+    vy2 = S_per_hbar - jnp.roll(S_per_hbar, 1, axis=1)
+    vz2 = S_per_hbar - jnp.roll(S_per_hbar, 1, axis=2)
+    vx2 = jnp.where(vx2 > jnp.pi, vx2 - 2 * jnp.pi, vx2)
+    vx2 = jnp.where(vx2 <= -jnp.pi, vx2 + 2 * jnp.pi, vx2)
+    vy2 = jnp.where(vy2 > jnp.pi, vy2 - 2 * jnp.pi, vy2)
+    vy2 = jnp.where(vy2 <= -jnp.pi, vy2 + 2 * jnp.pi, vy2)
+    vz2 = jnp.where(vz2 > jnp.pi, vz2 - 2 * jnp.pi, vz2)
+    vz2 = jnp.where(vz2 <= -jnp.pi, vz2 + 2 * jnp.pi, vz2)
+    vx2 = vx2 / dx / m_per_hbar
+    vy2 = vy2 / dx / m_per_hbar
+    vz2 = vz2 / dx / m_per_hbar
+
+    # Average forward and backward
+    vx = 0.5 * (vx1 + vx2)
+    vy = 0.5 * (vy1 + vy2)
+    vz = 0.5 * (vz1 + vz2)
+
+    return vx, vy, vz