jaxion 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

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/analysis.py

@@ -0,0 +1,43 @@
+import jax.numpy as jnp
+import jaxdecomp as jd
+
+
+def radial_power_spectrum(data_cube, kx, ky, kz, box_size):
+    """
+    Computes radially averaged power spectral density of data_cube (3D).
+    data_cube: jnp.ndarray (3D, must be cubic)
+    boxsize: float (physical size of box)
+    Returns: Pf (radial power spectrum), k (wavenumbers), 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/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/gravity.py

@@ -1,9 +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 = -jnp.fft.fftn(4.0 * jnp.pi * G * (rho - rho_bar)) / (k_sq + (k_sq == 0))
-    V = jnp.real(jnp.fft.ifftn(Vhat))
+    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/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/params_default.json

@@ -28,7 +28,7 @@
   "domain": {
     "box_size": { 
       "default": 10.0,
-      "description": "periodic domain box size (kpc)."
+      "description": "periodic domain box size [kpc]."
     },
     "resolution_base": { 
       "default": 32,
@@ -40,38 +40,75 @@
     }
   },
   "time": {
-    "start": 0.0,
-    "end": 1.0,
-    "adaptive": false
+    "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": "./checkpoints",
-    "num_checkpoints": 100,
-    "save": true,
-    "plot_dynamic_range": 100.0
+    "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 in units of 10^{-22} eV."
+      "description": "axion mass [10^{-22} eV]."
     }
   },
   "hydro": {
     "sound_speed": { 
       "default": 1.0,
-      "description": "Isothermal sound speed (km/s)."
+      "description": "isothermal sound speed [km/s]."
     }
   },
   "particles": {
-    "num_particles": 0,
-    "particle_mass": 1.0
+    "num_particles": {
+      "default": 0,
+      "description": "number of particles."
+    },
+    "particle_mass": {
+      "default": 1.0,
+      "description": "particle mass [M_sun]."
+    }
   },
   "cosmology": {
-    "Omega_m": 0.3,
-    "Omega_Lambda": 0.7
+    "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": "version of the jaxion library used to create these parameters (automatically detected)."
+    "description": "jaxion version used (auto detected)."
   }
 }

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/quantum.py

@@ -1,4 +1,5 @@
 import jax.numpy as jnp
+import jaxdecomp as jd
 
 # Pure functions for quantum simulation
 
@@ -9,7 +10,86 @@ def quantum_kick(psi, V, m_per_hbar, dt):
 
 
 def quantum_drift(psi, k_sq, m_per_hbar, dt):
-    psi_hat = jnp.fft.fftn(psi)
+    psi_hat = jd.fft.pfft3d(psi)
     psi_hat = jnp.exp(dt * (-1.0j * k_sq / m_per_hbar / 2.0)) * psi_hat
-    psi = jnp.fft.ifftn(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

jaxion/simulation.py

@@ -6,11 +6,19 @@ import json
 import time
 
 from .constants import constants
-from .quantum import quantum_kick, quantum_drift
+from .quantum import quantum_kick, quantum_drift, quantum_velocity
 from .gravity import calculate_gravitational_potential
 from .hydro import hydro_fluxes, hydro_accelerate
 from .particles import particles_accelerate, particles_drift, bin_particles
-from .utils import set_up_parameters, print_parameters
+from .cosmology import get_supercomoving_time_interval, get_next_scale_factor
+from .utils import (
+    set_up_parameters,
+    print_parameters,
+    xmeshgrid,
+    xmeshgrid_transpose,
+    xzeros,
+    xones,
+)
 from .visualization import plot_sim
 
 
@@ -21,10 +29,20 @@ class Simulation:
     Parameters
     ----------
       params (dict): The Python dictionary that contains the simulation parameters.
+                     Params can also be a string path to a checkpoint directory to load a saved simulation.
 
     """
 
-    def __init__(self, params):
+    def __init__(self, params, sharding=None):
+        # allow loading directly from a checkpoint path
+        load_from_checkpoint = False
+        checkpoint_dir = ""
+        if isinstance(params, str):
+            load_from_checkpoint = True
+            checkpoint_dir = os.path.join(os.getcwd(), params)
+            with open(os.path.join(checkpoint_dir, "params.json"), "r") as f:
+                params = json.load(f)
+
         # start from default simulation parameters and update with user params
         self._params = set_up_parameters(params)
 
@@ -32,38 +50,80 @@ class Simulation:
         if self.resolution % 2 != 0:
             raise ValueError("Resolution must be divisible by 2.")
 
+        if self.params["time"]["adaptive"]:
+            raise NotImplementedError("Adaptive time stepping is not yet implemented.")
+
+        if self.params["physics"]["cosmology"]:
+            if (
+                self.params["physics"]["hydro"]
+                or self.params["physics"]["particles"]
+                or self.params["physics"]["external_potential"]
+            ):
+                raise NotImplementedError(
+                    "Cosmological hydro/particles/external_potential physics is not yet implemented."
+                )
+
+        if self.params["physics"]["hydro"] or self.params["physics"]["particles"]:
+            if sharding is not None:
+                raise NotImplementedError(
+                    "hydro/particles sharding is not yet implemented."
+                )
+
         if self.params["output"]["save"]:
-            print("Simulation initialized with parameters:")
-            print_parameters(self.params)
+            if jax.process_index() == 0:
+                print("Simulation parameters:")
+                print_parameters(self.params)
+
+        # jitted functions
+        self.xmeshgrid_jit = jax.jit(
+            xmeshgrid, in_shardings=None, out_shardings=sharding
+        )
+        self.xmeshgrid_transpose_jit = jax.jit(
+            xmeshgrid_transpose, in_shardings=None, out_shardings=sharding
+        )
+        self.xzeros_jit = jax.jit(
+            xzeros, static_argnums=0, in_shardings=None, out_shardings=sharding
+        )
+        self.xones_jit = jax.jit(
+            xones, static_argnums=0, in_shardings=None, out_shardings=sharding
+        )
 
         # simulation state
         self.state = {}
         self.state["t"] = 0.0
+        if self.params["physics"]["cosmology"]:
+            self.state["redshift"] = 0.0
         if self.params["physics"]["quantum"]:
-            self.state["psi"] = (
-                jnp.zeros((self.resolution, self.resolution, self.resolution)) * 1j
-            )
+            self.state["psi"] = self.xzeros_jit(self.resolution) * 1j
         if self.params["physics"]["external_potential"]:
-            self.state["V_ext"] = jnp.zeros(
-                (self.resolution, self.resolution, self.resolution)
-            )
+            self.state["V_ext"] = self.xzeros_jit(self.resolution)
         if self.params["physics"]["hydro"]:
             self.state["rho"] = jnp.zeros(
-                (self.resolution, self.resolution, self.resolution)
+                (self.resolution, self.resolution, self.resolution),
             )
             self.state["vx"] = jnp.zeros(
-                (self.resolution, self.resolution, self.resolution)
+                (self.resolution, self.resolution, self.resolution),
             )
             self.state["vy"] = jnp.zeros(
-                (self.resolution, self.resolution, self.resolution)
+                (self.resolution, self.resolution, self.resolution),
             )
             self.state["vz"] = jnp.zeros(
-                (self.resolution, self.resolution, self.resolution)
+                (self.resolution, self.resolution, self.resolution),
             )
         if self.params["physics"]["particles"]:
             self.state["pos"] = jnp.zeros((self.num_particles, 3))
             self.state["vel"] = jnp.zeros((self.num_particles, 3))
 
+        if load_from_checkpoint:
+            options = ocp.CheckpointManagerOptions()
+            async_checkpoint_manager = ocp.CheckpointManager(
+                checkpoint_dir, options=options
+            )
+            step = async_checkpoint_manager.latest_step()
+            self.state = async_checkpoint_manager.restore(
+                step, args=ocp.args.StandardRestore(self.state)
+            )
+
     @property
     def resolution(self):
         return (
@@ -104,19 +164,24 @@ class Simulation:
     def grid(self):
         hx = 0.5 * self.dx
         x_lin = jnp.linspace(hx, self.box_size - hx, self.resolution)
-        X, Y, Z = jnp.meshgrid(x_lin, x_lin, x_lin, indexing="ij")
-        return X, Y, Z
+        xx, yy, zz = self.xmeshgrid_jit(x_lin)
+        return xx, yy, zz
 
     @property
     def kgrid(self):
         nx = self.resolution
         k_lin = (2.0 * jnp.pi / self.box_size) * jnp.arange(-nx / 2, nx / 2)
-        kx, ky, kz = jnp.meshgrid(k_lin, k_lin, k_lin, indexing="ij")
+        kx, ky, kz = self.xmeshgrid_transpose_jit(k_lin)
         kx = jnp.fft.ifftshift(kx)
         ky = jnp.fft.ifftshift(ky)
         kz = jnp.fft.ifftshift(kz)
         return kx, ky, kz
 
+    @property
+    def quantum_velocity(self):
+        m_per_hbar = self.axion_mass / constants["reduced_planck_constant"]
+        return quantum_velocity(self.state["psi"], self.box_size, m_per_hbar)
+
     def _calc_rho_bar(self, state):
         rho_bar = 0.0
         if self.params["physics"]["quantum"]:
@@ -133,7 +198,7 @@ class Simulation:
     def _calc_grav_potential(self, state, k_sq):
         G = constants["gravitational_constant"]
         m_particle = self.params["particles"]["particle_mass"]
-        rho_bar = self._calc_rho_bar(self.state)
+        rho_bar = self._calc_rho_bar(state)
         rho_tot = 0.0
         if self.params["physics"]["quantum"]:
             rho_tot += jnp.abs(state["psi"]) ** 2
@@ -141,6 +206,10 @@ class Simulation:
             rho_tot += state["rho"]
         if self.params["physics"]["particles"]:
             rho_tot += bin_particles(state["pos"], self.dx, self.resolution, m_particle)
+        if self.params["physics"]["cosmology"]:
+            scale_factor = 1.0 / (1.0 + state["redshift"])
+            rho_bar *= scale_factor
+            rho_tot *= scale_factor
         return calculate_gravitational_potential(rho_tot, k_sq, G, rho_bar)
 
     @property
@@ -167,19 +236,40 @@ class Simulation:
         # Simulation parameters
         dx = self.dx
         m_per_hbar = self.axion_mass / constants["reduced_planck_constant"]
+        box_size = self.box_size
 
         dt_fac = 1.0
         dt_kin = dt_fac * (m_per_hbar / 6.0) * (dx * dx)
+        t_start = self.params["time"]["start"]
         t_end = self.params["time"]["end"]
+        t_span = t_end - t_start
+        state["t"] = t_start
+
+        # cosmology
+        if self.params["physics"]["cosmology"]:
+            z_start = self.params["time"]["start"]
+            z_end = self.params["time"]["end"]
+            omega_matter = self.params["cosmology"]["omega_matter"]
+            omega_lambda = self.params["cosmology"]["omega_lambda"]
+            little_h = self.params["cosmology"]["little_h"]
+            t_span = get_supercomoving_time_interval(
+                z_start, z_end, omega_matter, omega_lambda, little_h
+            )
+            state["t"] = 0.0
+            state["redshift"] = z_start
 
+        # hydro
         c_sound = self.params["hydro"]["sound_speed"]
-        box_size = self.box_size
 
         # round up to the nearest multiple of num_checkpoints
         num_checkpoints = self.params["output"]["num_checkpoints"]
-        nt = int(round(round(t_end / dt_kin) / num_checkpoints) * num_checkpoints)
+        nt = int(round(round(t_span / dt_kin) / num_checkpoints) * num_checkpoints)
         nt_sub = int(round(nt / num_checkpoints))
-        dt = t_end / nt
+        dt = t_span / nt
+
+        # distributed arrays (fixed) needed for calculations
+        kx, ky, kz = None, None, None
+        k_sq = None
 
         # Fourier space variables
         if self.params["physics"]["gravity"] or self.params["physics"]["quantum"]:
@@ -192,10 +282,14 @@ class Simulation:
             checkpoint_dir = checkpoint_dir = os.path.join(
                 os.getcwd(), self.params["output"]["path"]
             )
-            path = ocp.test_utils.erase_and_create_empty(checkpoint_dir)
+            path = os.path.join(os.getcwd(), checkpoint_dir)
+            if jax.process_index() == 0:
+                path = ocp.test_utils.erase_and_create_empty(checkpoint_dir)
             async_checkpoint_manager = ocp.CheckpointManager(path, options=options)
 
-        def _kick(state, dt):
+        carry = (state, kx, ky, kz, k_sq)
+
+        def _kick(state, kx, ky, kz, k_sq, dt):
             # Kick (half-step)
             if (
                 self.params["physics"]["gravity"]
@@ -222,7 +316,7 @@ class Simulation:
                         state["vel"], state["pos"], V, kx, ky, kz, dx, dt
                     )
 
-        def _drift(state, dt):
+        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)
@@ -234,20 +328,30 @@ class Simulation:
                 state["pos"] = particles_drift(state["pos"], state["vel"], dt, box_size)
 
         @jax.jit
-        def _update(_, state):
+        def _update(_, carry):
             # Update the simulation state by one timestep
             # according to a 2nd-order `kick-drift-kick` scheme
-            _kick(state, 0.5 * dt)
-            _drift(state, dt)
-            _kick(state, 0.5 * dt)
-
-            # update time
+            state, kx, ky, kz, k_sq = carry
+            _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
+            _drift(state, k_sq, dt)
+            # update time & redshift
             state["t"] += dt
+            if self.params["physics"]["cosmology"]:
+                scale_factor = get_next_scale_factor(
+                    state["redshift"],
+                    dt,
+                    self.params["cosmology"]["omega_matter"],
+                    self.params["cosmology"]["omega_lambda"],
+                    self.params["cosmology"]["little_h"],
+                )
+                state["redshift"] = 1.0 / scale_factor - 1.0
+            _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
 
-            return state
+            return state, kx, ky, kz, k_sq
 
         # save initial state
-        print("Starting simulation ...")
+        if jax.process_index() == 0:
+            print(f"Starting simulation (res={self.resolution}, nt={nt}) ...")
         if self.params["output"]["save"]:
             with open(os.path.join(checkpoint_dir, "params.json"), "w") as f:
                 json.dump(self.params, f, indent=2)
@@ -259,7 +363,8 @@ class Simulation:
         t_start_timer = time.time()
         if self.params["output"]["save"]:
             for i in range(1, num_checkpoints + 1):
-                state = jax.lax.fori_loop(0, nt_sub, _update, init_val=state)
+                carry = jax.lax.fori_loop(0, nt_sub, _update, init_val=carry)
+                state, kx, ky, kz, k_sq = carry
                 jax.block_until_ready(state)
                 # save state
                 async_checkpoint_manager.save(i, args=ocp.args.StandardSave(state))
@@ -267,15 +372,17 @@ class Simulation:
                 elapsed = time.time() - t_start_timer
                 est_total = elapsed / i * num_checkpoints
                 est_remaining = est_total - elapsed
-                print(
-                    f"{percent:.1f}%: estimated time remaining (s): {est_remaining:.1f}"
-                )
+                if jax.process_index() == 0:
+                    print(
+                        f"{percent:.1f}%: estimated time remaining (s): {est_remaining:.1f}"
+                    )
                 plot_sim(state, checkpoint_dir, i, self.params)
                 async_checkpoint_manager.wait_until_finished()
         else:
             state = jax.lax.fori_loop(0, nt, _update, init_val=state)
         jax.block_until_ready(state)
-        print("Simulation Run Time (s): ", time.time() - t_start_timer)
+        if jax.process_index() == 0:
+            print("Simulation Run Time (s): ", time.time() - t_start_timer)
 
         return state
 

jaxion/utils.py

@@ -5,6 +5,7 @@ from pathlib import Path
 import importlib.resources
 import json
 from importlib.metadata import version
+import jax.numpy as jnp
 
 
 def print_parameters(params):
@@ -90,3 +91,29 @@ def run_example_main(example_path, argv=None):
         os.chdir(old_cwd)
         sys.argv = old_argv
     return result
+
+
+# Make a distributed meshgrid function
+def xmeshgrid(x_lin):
+    xx, yy, zz = jnp.meshgrid(x_lin, x_lin, x_lin, indexing="ij")
+    return xx, yy, zz
+
+
+# NOTE: jaxdecomp (jd) has pfft3d transpose the axis (X, Y, Z) --> (Y, Z, X), and pifft3d undo it
+# so the fourier space variables (e.g. kx, ky, kz) all need to be transposed
+
+
+def xmeshgrid_transpose(x_lin):
+    xx, yy, zz = jnp.meshgrid(x_lin, x_lin, x_lin, indexing="ij")
+    xx = jnp.transpose(xx, (1, 2, 0))
+    yy = jnp.transpose(yy, (1, 2, 0))
+    zz = jnp.transpose(zz, (1, 2, 0))
+    return xx, yy, zz
+
+
+def xzeros(nx):
+    return jnp.zeros((nx, nx, nx))
+
+
+def xones(nx):
+    return jnp.ones((nx, nx, nx))

jaxion/visualization.py

@@ -1,3 +1,4 @@
+import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
 import os
@@ -8,67 +9,82 @@ def plot_sim(state, checkpoint_dir, i, params):
 
     dynamic_range = params["output"]["plot_dynamic_range"]
 
+    # process distributed data
     if params["physics"]["quantum"]:
-        plt.clf()
-
-        # DM projection
         nx = state["psi"].shape[0]
-        rho_bar = jnp.mean(jnp.abs(state["psi"]) ** 2)
-        vmin = jnp.log10(rho_bar / dynamic_range)
-        vmax = jnp.log10(rho_bar * dynamic_range)
+        rho_bar_dm = jnp.mean(jnp.abs(state["psi"]) ** 2)
+        rho_proj_dm = jax.experimental.multihost_utils.process_allgather(
+            jnp.log10(jnp.mean(jnp.abs(state["psi"]) ** 2, axis=2).T)
+        ).reshape(nx, nx)
+    if params["physics"]["hydro"]:
+        nx = state["rho"].shape[0]
+        rho_bar_gas = jnp.mean(state["rho"])
+        rho_proj_gas = jax.experimental.multihost_utils.process_allgather(
+            jnp.log10(jnp.mean(jnp.abs(state["rho"]) ** 2, axis=2).T)
+        ).reshape(nx, nx)
 
-        rho_proj_dm = jnp.log10(jnp.mean(jnp.abs(state["psi"]) ** 2, axis=2)).T
-        ax = plt.gca()
-        ax.imshow(
-            rho_proj_dm,
-            cmap="inferno",
-            origin="lower",
-            vmin=vmin,
-            vmax=vmax,
-            extent=[0, nx, 0, nx],
-        )
-        if params["physics"]["particles"]:
-            # draw particles
-            box_size = params["domain"]["box_size"]
-            sx = (state["pos"][:, 0] / box_size) * nx
-            sy = (state["pos"][:, 1] / box_size) * nx
-            plt.plot(sx, sy, color="cyan", marker=".", linestyle="None", markersize=2)
-        ax.set_aspect("equal")
-        ax.get_xaxis().set_visible(False)
-        ax.get_yaxis().set_visible(False)
+    # create plot on process 0
+    if jax.process_index() == 0:
+        if params["physics"]["quantum"]:
+            plt.clf()
 
-        plt.savefig(
-            os.path.join(checkpoint_dir, f"dm{i:03d}.png"),
-            bbox_inches="tight",
-            pad_inches=0,
-        )
-        plt.close()
+            # DM projection
+            nx = state["psi"].shape[0]
+            vmin = jnp.log10(rho_bar_dm / dynamic_range)
+            vmax = jnp.log10(rho_bar_dm * dynamic_range)
 
-    if params["physics"]["hydro"]:
-        plt.clf()
+            ax = plt.gca()
+            ax.imshow(
+                rho_proj_dm,
+                cmap="inferno",
+                origin="lower",
+                vmin=vmin,
+                vmax=vmax,
+                extent=[0, nx, 0, nx],
+            )
+            if params["physics"]["particles"]:
+                # draw particles
+                box_size = params["domain"]["box_size"]
+                sx = (state["pos"][:, 0] / box_size) * nx
+                sy = (state["pos"][:, 1] / box_size) * nx
+                plt.plot(
+                    sx, sy, color="cyan", marker=".", linestyle="None", markersize=2
+                )
+            ax.set_aspect("equal")
+            ax.get_xaxis().set_visible(False)
+            ax.get_yaxis().set_visible(False)
 
-        # gas projection
-        nx = state["rho"].shape[0]
-        rho_bar = jnp.mean(state["rho"])
-        vmin = jnp.log10(rho_bar / dynamic_range)
-        vmax = jnp.log10(rho_bar * dynamic_range)
-        rho_proj_gas = jnp.log10(jnp.mean(state["rho"], axis=2)).T
-        ax = plt.gca()
-        ax.imshow(
-            rho_proj_gas,
-            cmap="viridis",
-            origin="lower",
-            vmin=vmin,
-            vmax=vmax,
-            extent=[0, nx, 0, nx],
-        )
-        ax.set_aspect("equal")
-        ax.get_xaxis().set_visible(False)
-        ax.get_yaxis().set_visible(False)
+            plt.savefig(
+                os.path.join(checkpoint_dir, f"dm{i:03d}.png"),
+                bbox_inches="tight",
+                pad_inches=0,
+            )
+            plt.close()
+
+        if params["physics"]["hydro"]:
+            plt.clf()
+
+            # gas projection
+            nx = state["rho"].shape[0]
+            vmin = jnp.log10(rho_bar_gas / dynamic_range)
+            vmax = jnp.log10(rho_bar_gas * dynamic_range)
+
+            ax = plt.gca()
+            ax.imshow(
+                rho_proj_gas,
+                cmap="viridis",
+                origin="lower",
+                vmin=vmin,
+                vmax=vmax,
+                extent=[0, nx, 0, nx],
+            )
+            ax.set_aspect("equal")
+            ax.get_xaxis().set_visible(False)
+            ax.get_yaxis().set_visible(False)
 
-        plt.savefig(
-            os.path.join(checkpoint_dir, f"gas{i:03d}.png"),
-            bbox_inches="tight",
-            pad_inches=0,
-        )
-        plt.close()
+            plt.savefig(
+                os.path.join(checkpoint_dir, f"gas{i:03d}.png"),
+                bbox_inches="tight",
+                pad_inches=0,
+            )
+            plt.close()

{jaxion-0.0.2.dist-info → jaxion-0.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.2
+Version: 0.0.3
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -21,6 +21,12 @@ 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]
@@ -44,9 +50,9 @@ A simple JAX-powered simulation library for numerical experiments of fuzzy dark
 
 Author: [Philip Mocz (@pmocz)](https://github.com/pmocz/)
 
-⚠️ Jaxion is currently being developed and is not yet ready for use. Check back later ⚠️
+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 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 seemlessly integrate with piplines 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
@@ -66,6 +72,20 @@ pip install jaxion[cuda12]
 Check out the `examples/` directory for demonstrations of using Jaxion.
 
 
+## Examples
+
+<p align="center">
+  <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>
+  <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).

jaxion-0.0.3.dist-info/RECORD

@@ -0,0 +1,17 @@
+jaxion/__init__.py,sha256=Hdji1UQ47lG24Pqcy6UUq9L0-qy6m9Ax41L0vIYzBto,164
+jaxion/analysis.py,sha256=UZpCvDr2ISDAgRoZFtuhvDk0QtN3mn5NxIoYYY8wzs8,1304
+jaxion/constants.py,sha256=b0vgqtezIEGp4O5JRJo3YK3eQC8AN6hmV84-4pYWgmc,528
+jaxion/cosmology.py,sha256=UC1McXNTXGoPRYXn0nI2-csVkJWL-ZBNoCa44oU1b4w,2681
+jaxion/gravity.py,sha256=3brRZelKm-soXqk_Lt3SqhbZ00woJCraqwdMuR-KooA,291
+jaxion/hydro.py,sha256=KoJ02tRpAc4V3Ofzw4zbHLRaE2GdIatbOBE04_LsSRw,6980
+jaxion/params_default.json,sha256=9CJrhEPsv5zGEs7_WqFyuccCDipPCDhXgKzVdqOsOWE,2775
+jaxion/particles.py,sha256=pMopGvoZ0J_3EviD0WnTMmiebU9h2_8IO-p6I-E5DEU,3980
+jaxion/quantum.py,sha256=GWOpN6ipfEw-6Ah2zQpxS3oqeSt_iHMDSgnVYSjXY5E,3321
+jaxion/simulation.py,sha256=3UrrHPJr0_39BiP86nQSl4ui8EsPHf9RrZ5hDobit5k,14884
+jaxion/utils.py,sha256=rT7NM0FNEgFwN7oTgTb-jkR66Iw0xYTHHxcoikYd1ag,3572
+jaxion/visualization.py,sha256=MJ3iqrvAem4RbozM79Ri2cNTIwKUV8o5X8M32FiDfbo,2928
+jaxion-0.0.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+jaxion-0.0.3.dist-info/METADATA,sha256=1C3fwJHJEBZYMEhxmTkKd7z_ZUrF15E3P0oS5XxrdjA,3579
+jaxion-0.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jaxion-0.0.3.dist-info/top_level.txt,sha256=S1OV2VdlDG_9UwpKOIji4itQGOS-VWUOWUi3GeXWzt0,7
+jaxion-0.0.3.dist-info/RECORD,,

jaxion-0.0.2.dist-info/RECORD

@@ -1,15 +0,0 @@
-jaxion/__init__.py,sha256=gHrU_SvIBIbuYt-MeEeyiMyuJKwLjTTiz513OZ7CDHQ,95
-jaxion/constants.py,sha256=b0vgqtezIEGp4O5JRJo3YK3eQC8AN6hmV84-4pYWgmc,528
-jaxion/gravity.py,sha256=DdVOhAaoIPLXtW7ALrMK2umLPkStzxJ778l1SD2yMkY,266
-jaxion/hydro.py,sha256=3ZFNDhIaPBw4uTBqFbjJotCLkuKG8AvRcY1T78QQ96U,6953
-jaxion/params_default.json,sha256=4-XNPN52phGxYsP6FOeIJH0ALtICjrQxApotz-NhJEc,1782
-jaxion/particles.py,sha256=7oJBVyddfOXUDYxQuGIs8_ByL2Hu6jkuWE1MWtshFjk,3953
-jaxion/quantum.py,sha256=91tEotQhX1CB3HCaRCDARRPX_dEx58GIUCluY5RyPvA,377
-jaxion/simulation.py,sha256=KcIi91LHJO5K4M0gxU4qc7e4JEeOyW1HUZOYwt8FOG0,10385
-jaxion/utils.py,sha256=LCVx5_N9-6HsbENENOEz8tI6Czl0Z8Bit2icf1k4024,2880
-jaxion/visualization.py,sha256=WKID2_LFbbqXwmW_TLTnVznKYHfjiDKNbjD9xS4GfLc,2209
-jaxion-0.0.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-jaxion-0.0.2.dist-info/METADATA,sha256=uUe69xVa678F2uV28-CczduNdfoOtBpwhfuM_mJKXpM,2811
-jaxion-0.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jaxion-0.0.2.dist-info/top_level.txt,sha256=S1OV2VdlDG_9UwpKOIji4itQGOS-VWUOWUi3GeXWzt0,7
-jaxion-0.0.2.dist-info/RECORD,,