jaxion 0.0.6__tar.gz → 0.0.8__tar.gz

{jaxion-0.0.6 → jaxion-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.6
+Version: 0.0.8
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0

{jaxion-0.0.6 → jaxion-0.0.8}/jaxion/defaults.json

@@ -48,6 +48,10 @@
       "default": 1.0,
       "description": "simulation end time [kpc/(km/s)] or [redshift] (cosmology=true)."
     },
+    "safety_factor": { 
+      "default": 1.0,
+      "description": "safety factor for time stepping."
+    },
     "adaptive": { 
       "default": false,
       "description": "switch on for adaptive time stepping."
@@ -95,6 +99,10 @@
     "particle_mass": {
       "default": 1.0,
       "description": "particle mass [M_sun]."
+    },
+    "accrete_gas": {
+      "default": false,
+      "description": "switch on to accrete gas."
     }
   },
   "cosmology": {

{jaxion-0.0.6 → jaxion-0.0.8}/jaxion/particles.py

@@ -18,7 +18,7 @@ def get_cic_indices_and_weights(pos, dx, resolution):
     return i, ip1, weight_i, weight_ip1
 
 
-def bin_particles(pos, dx, resolution, m_particle):
+def bin_particles(pos, m_particles, dx, resolution, multiple_masses):
     """Bin the particles into the grid using cloud-in-cell weights."""
     nx = resolution
     n_particle = pos.shape[0]
@@ -27,6 +27,10 @@ def bin_particles(pos, dx, resolution, m_particle):
 
     def deposit_particle(s, rho):
         """Deposit the particle mass into the grid."""
+        if multiple_masses:
+            m_particle = m_particles[s]
+        else:
+            m_particle = m_particles
         fac = m_particle / (dx * dx * dx)
         rho = rho.at[i[s, 0], i[s, 1], i[s, 2]].add(
             w_i[s, 0] * w_i[s, 1] * w_i[s, 2] * fac
@@ -112,3 +116,102 @@ def particles_drift(pos, vel, dt, box_size):
     pos = jnp.mod(pos, jnp.array([box_size, box_size, box_size]))
 
     return pos
+
+
+def particles_accrete_gas(mass, rho, pos, G, sound_speed, dx, dt):
+    """Accrete gas onto particles (Bondi)."""
+    n_particle = pos.shape[0]
+    resolution = rho.shape[0]
+    i, ip1, w_i, w_ip1 = get_cic_indices_and_weights(pos, dx, resolution)
+    d_mass = jnp.zeros_like(mass)
+    d_rho = jnp.zeros_like(rho)
+    lam = jnp.exp(1.5) / 4.0  # ≈ 1.12
+    vol = dx**3
+
+    def accrete(s, deltas):
+        """Deposit the particle mass into the grid."""
+        d_mass, d_rho = deltas
+        dM_fac = dt * 4.0 * jnp.pi * lam * (G * mass[s]) ** 2 / sound_speed**3
+        # dM = dM_fac * rho
+
+        dm = w_i[s, 0] * w_i[s, 1] * w_i[s, 2] * dM_fac * rho[i[s, 0], i[s, 1], i[s, 2]]
+        d_rho = d_rho.at[i[s, 0], i[s, 1], i[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_ip1[s, 0]
+            * w_i[s, 1]
+            * w_i[s, 2]
+            * dM_fac
+            * rho[ip1[s, 0], i[s, 1], i[s, 2]]
+        )
+        d_rho = d_rho.at[ip1[s, 0], i[s, 1], i[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_i[s, 0]
+            * w_ip1[s, 1]
+            * w_i[s, 2]
+            * dM_fac
+            * rho[i[s, 0], ip1[s, 1], i[s, 2]]
+        )
+        d_rho = d_rho.at[i[s, 0], ip1[s, 1], i[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_i[s, 0]
+            * w_i[s, 1]
+            * w_ip1[s, 2]
+            * dM_fac
+            * rho[i[s, 0], i[s, 1], ip1[s, 2]]
+        )
+        d_rho = d_rho.at[i[s, 0], i[s, 1], ip1[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_ip1[s, 0]
+            * w_ip1[s, 1]
+            * w_i[s, 2]
+            * dM_fac
+            * rho[ip1[s, 0], ip1[s, 1], i[s, 2]]
+        )
+        d_rho = d_rho.at[ip1[s, 0], ip1[s, 1], i[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_ip1[s, 0]
+            * w_i[s, 1]
+            * w_ip1[s, 2]
+            * dM_fac
+            * rho[ip1[s, 0], i[s, 1], ip1[s, 2]]
+        )
+        d_rho = d_rho.at[ip1[s, 0], i[s, 1], ip1[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_i[s, 0]
+            * w_ip1[s, 1]
+            * w_ip1[s, 2]
+            * dM_fac
+            * rho[i[s, 0], ip1[s, 1], ip1[s, 2]]
+        )
+        d_rho = d_rho.at[i[s, 0], ip1[s, 1], ip1[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        dm = (
+            w_ip1[s, 0]
+            * w_ip1[s, 1]
+            * w_ip1[s, 2]
+            * dM_fac
+            * rho[ip1[s, 0], ip1[s, 1], ip1[s, 2]]
+        )
+        d_rho = d_rho.at[ip1[s, 0], ip1[s, 1], ip1[s, 2]].add(-dm / vol)
+        d_mass = d_mass.at[s].add(dm)
+
+        return d_mass, d_rho
+
+    d_mass, d_rho = jax.lax.fori_loop(0, n_particle, accrete, (d_mass, d_rho))
+    mass = mass + d_mass
+    rho = rho + d_rho
+
+    return mass, rho

{jaxion-0.0.6 → jaxion-0.0.8}/jaxion/simulation.py

@@ -9,7 +9,12 @@ from .constants import constants
 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 .particles import (
+    particles_accelerate,
+    particles_drift,
+    particles_accrete_gas,
+    bin_particles,
+)
 from .cosmology import get_supercomoving_time_interval, get_next_scale_factor
 from .utils import (
     set_up_parameters,
@@ -37,7 +42,7 @@ class Simulation:
 
     """
 
-    def __init__(self, params, sharding=None):
+    def __init__(self, params, sharding=None, checkpoint_number=None):
         # allow loading directly from a checkpoint path
         load_from_checkpoint = False
         checkpoint_dir = ""
@@ -95,11 +100,12 @@ class Simulation:
             xones, static_argnums=0, in_shardings=None, out_shardings=sharding
         )
 
-        # customfunctions
+        # custom functions
         self.custom_kick = None
         self.custom_drift = None
         self.custom_density = None
         self.custom_plot = None
+        self.callback = None
 
         # simulation state
         self.state = {}
@@ -126,13 +132,21 @@ class Simulation:
         if self.params["physics"]["particles"]:
             self.state["pos"] = jnp.zeros((self.num_particles, 3))
             self.state["vel"] = jnp.zeros((self.num_particles, 3))
+            if self.params["particles"]["accrete_gas"]:
+                self.state["mass"] = (
+                    jnp.zeros(self.num_particles)
+                    + self.params["particles"]["particle_mass"]
+                )
 
         if load_from_checkpoint:
             options = ocp.CheckpointManagerOptions()
             async_checkpoint_manager = ocp.CheckpointManager(
                 checkpoint_dir, options=options
             )
-            step = async_checkpoint_manager.latest_step()
+            if checkpoint_number is not None:
+                step = checkpoint_number
+            else:
+                step = async_checkpoint_manager.latest_step()
             self.state = async_checkpoint_manager.restore(
                 step, args=ocp.args.StandardRestore(self.state)
             )
@@ -168,6 +182,23 @@ class Simulation:
         """
         return self.box_size / self.resolution
 
+    @property
+    def nt(self):
+        """
+        Return the number of timesteps
+        """
+        dx = self.dx
+        m_per_hbar = self.m_per_hbar
+        safety = self.params["time"]["safety_factor"]
+        dt_kin = safety * (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
+        num_checkpoints = self.params["output"]["num_checkpoints"]
+        nt = int(round(round(t_span / dt_kin) / num_checkpoints) * num_checkpoints)
+
+        return nt
+
     @property
     def axion_mass(self):
         """
@@ -262,17 +293,19 @@ class Simulation:
         if self.params["physics"]["hydro"]:
             rho_bar += jnp.mean(state["rho"])
         if self.params["physics"]["particles"]:
-            m_particle = self.params["particles"]["particle_mass"]
-            n_particles = self.num_particles
             box_size = self.box_size
-            rho_bar += m_particle * n_particles / box_size
+            if self.params["particles"]["accrete_gas"]:
+                rho_bar += jnp.sum(state["mass"]) / box_size**3
+            else:
+                m_particle = self.params["particles"]["particle_mass"]
+                n_particles = self.num_particles
+                rho_bar += m_particle * n_particles / box_size**3
         if self.custom_density is not None:
             rho_bar += jnp.mean(self.custom_density(state))
         return rho_bar
 
     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(state)
         rho_tot = 0.0
         if self.params["physics"]["quantum"]:
@@ -280,7 +313,17 @@ class Simulation:
         if self.params["physics"]["hydro"]:
             rho_tot += state["rho"]
         if self.params["physics"]["particles"]:
-            rho_tot += bin_particles(state["pos"], self.dx, self.resolution, m_particle)
+            multiple_masses = self.params["particles"]["accrete_gas"]
+            if multiple_masses:
+                m_particles = state["mass"]
+                rho_tot += bin_particles(
+                    state["pos"], m_particles, self.dx, self.resolution, multiple_masses
+                )
+            else:
+                m_particles = self.params["particles"]["particle_mass"]
+                rho_tot += bin_particles(
+                    state["pos"], m_particles, self.dx, self.resolution, multiple_masses
+                )
         if self.custom_density is not None:
             rho_tot += self.custom_density(state)
         if self.params["physics"]["cosmology"]:
@@ -319,8 +362,6 @@ class Simulation:
         num_cells = self.resolution**3
         m_per_hbar = self.m_per_hbar
 
-        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
@@ -332,11 +373,15 @@ class Simulation:
         use_particles = self.params["physics"]["particles"]
         use_cosmology = self.params["physics"]["cosmology"]
         use_external_potential = self.params["physics"]["external_potential"]
+        accrete_gas = self.params["particles"]["accrete_gas"]
         save = self.params["output"]["save"]
         use_custom = self.custom_kick is not None or self.custom_drift is not None
         if use_custom:
             custom_kick = self.custom_kick
             custom_drift = self.custom_drift
+        use_callback = self.callback is not None
+        if use_callback:
+            custom_callback = self.callback
 
         # cosmology
         if use_cosmology:
@@ -367,7 +412,7 @@ class Simulation:
 
         # round up to the nearest multiple of num_checkpoints
         num_checkpoints = self.params["output"]["num_checkpoints"]
-        nt = int(round(round(t_span / dt_kin) / num_checkpoints) * num_checkpoints)
+        nt = self.nt
         nt_sub = int(round(nt / num_checkpoints))
         dt = t_span / nt
 
@@ -435,10 +480,15 @@ class Simulation:
                 state["pos"] = particles_drift(state["pos"], state["vel"], dt, box_size)
             if use_custom:
                 state = custom_drift(state, k_sq, dt)
+            if use_hydro and accrete_gas:
+                G = constants["gravitational_constant"]
+                state["mass"], state["rho"] = particles_accrete_gas(
+                    state["mass"], state["rho"], state["pos"], G, c_sound, dx, dt
+                )
 
             return state
 
-        def _update(_, carry):
+        def _update(i, carry):
             # Update the simulation state by one timestep
             # according to a 2nd-order `kick-drift-kick` scheme
             state, kx, ky, kz, k_sq = carry
@@ -453,6 +503,9 @@ class Simulation:
                 state["redshift"] = 1.0 / scale_factor - 1.0
             state = _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
 
+            if use_callback:
+                custom_callback(i, state)
+
             return state, kx, ky, kz, k_sq
 
         # save initial state
@@ -462,16 +515,19 @@ class Simulation:
             with open(os.path.join(checkpoint_dir, "params.json"), "w") as f:
                 json.dump(self.params, f, indent=2)
             async_checkpoint_manager.save(0, args=ocp.args.StandardSave(state))
+            async_checkpoint_manager.wait_until_finished()
             plot_sim(state, checkpoint_dir, 0, self.params)
             if self.custom_plot is not None:
                 self.custom_plot(state, checkpoint_dir, 0, self.params)
-            async_checkpoint_manager.wait_until_finished()
 
         # Simulation Main Loop
         t_start_timer = time.time()
         if save:
             for i in range(1, num_checkpoints + 1):
-                carry = jax.lax.fori_loop(0, nt_sub, _update, init_val=carry)
+                i_timestep = (i - 1) * nt_sub
+                carry = jax.lax.fori_loop(
+                    i_timestep, i_timestep + nt_sub, _update, init_val=carry
+                )
                 state, _, _, _, _ = carry
                 jax.block_until_ready(state)
                 # save state
@@ -485,10 +541,10 @@ class Simulation:
                     print(
                         f"{percent:.1f}%: mcups={mcups:.1f}, estimated time left (s): {est_remaining:.1f}"
                     )
+                async_checkpoint_manager.wait_until_finished()
                 plot_sim(state, checkpoint_dir, i, self.params)
                 if self.custom_plot is not None:
                     self.custom_plot(state, checkpoint_dir, i, self.params)
-                async_checkpoint_manager.wait_until_finished()
         else:
             carry = jax.lax.fori_loop(0, nt, _update, init_val=carry)
             state, _, _, _, _ = carry

{jaxion-0.0.6 → jaxion-0.0.8}/jaxion/visualization.py

@@ -82,6 +82,14 @@ def plot_sim(state, checkpoint_dir, i, params):
                 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="red", marker=".", linestyle="None", markersize=5
+                )
             ax.set_aspect("equal")
             ax.get_xaxis().set_visible(False)
             ax.get_yaxis().set_visible(False)

{jaxion-0.0.6 → jaxion-0.0.8}/jaxion.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.6
+Version: 0.0.8
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0

{jaxion-0.0.6 → jaxion-0.0.8}/tests/test_examples.py

@@ -41,13 +41,13 @@ def test_heating_stars():
     )
     assert sim.resolution == 32
     assert sim.state["t"] > 0.0
-    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(2574.4248, rel=rel_tol)
+    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(2574.395, rel=rel_tol)
     assert jnp.mean(jnp.abs(sim.state["vel"][:, 0])) == pytest.approx(
-        16.625286, rel=rel_tol
+        16.625353, rel=rel_tol
     )
     assert jnp.mean(jnp.abs(sim.state["vel"][:, 1])) == pytest.approx(
-        17.345531, rel=rel_tol
+        17.345486, rel=rel_tol
     )
     assert jnp.mean(jnp.abs(sim.state["vel"][:, 2])) == pytest.approx(
-        18.218365, rel=rel_tol
+        18.218296, rel=rel_tol
     )