jaxion 0.0.6__tar.gz → 0.0.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.6
+Version: 0.0.7
 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.7}/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.7}/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.7}/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,
@@ -95,7 +100,7 @@ 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
@@ -126,6 +131,11 @@ 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()
@@ -262,17 +272,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 +292,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 +341,8 @@ 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)
+        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
@@ -332,6 +354,7 @@ 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:
@@ -435,6 +458,11 @@ 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
 

{jaxion-0.0.6 → jaxion-0.0.7}/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.7}/jaxion.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.6
+Version: 0.0.7
 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.7}/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
     )