jaxion 0.0.5__tar.gz → 0.0.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.5
+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
@@ -68,7 +68,7 @@ Jaxion is built for multi-GPU scalability and is fully differentiable. It is a h
 Jaxion is the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
 
 
-## Getting started
+## Install Jaxion
 
 Install with:
 
@@ -87,7 +87,7 @@ See the docs for more info on how to [build from source](https://jaxion.readthed
 
 ## Examples
 
-Check out the `examples/` directory for demonstrations of using Jaxion.
+Check out the [`examples/`](https://github.com/JaxionProject/jaxion/tree/main/examples/) directory for demonstrations of using Jaxion.
 
 <p align="center">
   <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">

{jaxion-0.0.5 → jaxion-0.0.7}/README.md

@@ -45,7 +45,7 @@ Jaxion is built for multi-GPU scalability and is fully differentiable. It is a h
 Jaxion is the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
 
 
-## Getting started
+## Install Jaxion
 
 Install with:
 
@@ -64,7 +64,7 @@ See the docs for more info on how to [build from source](https://jaxion.readthed
 
 ## Examples
 
-Check out the `examples/` directory for demonstrations of using Jaxion.
+Check out the [`examples/`](https://github.com/JaxionProject/jaxion/tree/main/examples/) directory for demonstrations of using Jaxion.
 
 <p align="center">
   <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">

jaxion-0.0.5/jaxion/params_default.json → 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."
@@ -75,6 +79,10 @@
     "m_22": { 
       "default": 1.0,
       "description": "axion mass [10^{-22} eV]."
+    },
+    "f_15": { 
+      "default": 0.0,
+      "description": "self-interaction strength [10^{15} GeV]."
     }
   },
   "hydro": {
@@ -91,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.5 → jaxion-0.0.7}/jaxion/gravity.py

@@ -5,6 +5,6 @@ import jaxdecomp as jd
 
 
 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))
+    V_hat = -jd.fft.pfft3d(4.0 * jnp.pi * G * (rho - rho_bar)) / (k_sq + (k_sq == 0))
+    V = jnp.real(jd.fft.pifft3d(V_hat))
     return V

{jaxion-0.0.5 → 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.5 → 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,
@@ -62,9 +67,10 @@ class Simulation:
                 self.params["physics"]["hydro"]
                 or self.params["physics"]["particles"]
                 or self.params["physics"]["external_potential"]
+                or self.params["quantum"]["f_15"] != 0.0
             ):
                 raise NotImplementedError(
-                    "Cosmological hydro/particles/external_potential physics is not yet implemented."
+                    "Cosmological hydro/particles/external_potential/SI is not yet implemented."
                 )
 
         if self.params["physics"]["hydro"] or self.params["physics"]["particles"]:
@@ -94,6 +100,12 @@ class Simulation:
             xones, static_argnums=0, in_shardings=None, out_shardings=sharding
         )
 
+        # custom functions
+        self.custom_kick = None
+        self.custom_drift = None
+        self.custom_density = None
+        self.custom_plot = None
+
         # simulation state
         self.state = {}
         self.state["t"] = 0.0
@@ -119,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()
@@ -173,6 +190,23 @@ class Simulation:
             / constants["speed_of_light"] ** 2
         )
 
+    @property
+    def scattering_length(self):
+        """
+        Return the axion self-interaction scattering length in the simulation (kpc)
+        """
+        f_15 = self.params["quantum"]["f_15"]
+        if f_15 == 0.0:
+            return 0.0
+        else:
+            f = f_15 * 1.0e24 * constants["electron_volt"]
+            sign = 1.0 if f > 0 else -1.0
+            hbar = constants["reduced_planck_constant"]
+            c = constants["speed_of_light"]
+            m = self.axion_mass
+            a_s = sign * (hbar * c**3 * m) / (32.0 * jnp.pi * (f**2))
+            return a_s
+
     @property
     def m_per_hbar(self):
         """
@@ -224,6 +258,13 @@ class Simulation:
         """
         return quantum_velocity(self.state["psi"], self.box_size, self.m_per_hbar)
 
+    @property
+    def rho_bar(self):
+        """
+        Return the mean density of the simulation
+        """
+        return self._calc_rho_bar(self.state)
+
     def _calc_rho_bar(self, state):
         rho_bar = 0.0
         if self.params["physics"]["quantum"]:
@@ -231,15 +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"]:
@@ -247,7 +292,19 @@ 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"]:
             scale_factor = 1.0 / (1.0 + state["redshift"])
             rho_bar *= scale_factor
@@ -280,18 +337,32 @@ class Simulation:
 
         # Simulation parameters
         dx = self.dx
-        m_per_hbar = self.m_per_hbar
         box_size = self.box_size
+        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
         state["t"] = t_start
 
+        use_quantum = self.params["physics"]["quantum"]
+        use_gravity = self.params["physics"]["gravity"]
+        use_hydro = self.params["physics"]["hydro"]
+        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
+
         # cosmology
-        if self.params["physics"]["cosmology"]:
+        if use_cosmology:
             z_start = self.params["time"]["start"]
             z_end = self.params["time"]["end"]
             omega_matter = self.params["cosmology"]["omega_matter"]
@@ -303,6 +374,17 @@ class Simulation:
             state["t"] = 0.0
             state["redshift"] = z_start
 
+        # self-interaction
+        a_s = self.scattering_length
+        c = constants["speed_of_light"]
+        m = self.axion_mass
+        si_coeff = None
+        si_coeff2 = None
+        do_self_interaction = a_s != 0.0
+        if do_self_interaction:
+            si_coeff = (4.0 * jnp.pi) * (a_s / m) / m_per_hbar**2
+            si_coeff2 = (32.0 * jnp.pi**2 / 3.0) * (a_s / (m * c)) ** 2 / m_per_hbar**5
+
         # hydro
         c_sound = self.params["hydro"]["sound_speed"]
 
@@ -317,16 +399,14 @@ class Simulation:
         k_sq = None
 
         # Fourier space variables
-        if self.params["physics"]["gravity"] or self.params["physics"]["quantum"]:
+        if use_gravity or use_quantum:
             kx, ky, kz = self.kgrid
             k_sq = kx**2 + ky**2 + kz**2
 
         # Checkpointer
-        if self.params["output"]["save"]:
+        if save:
             options = ocp.CheckpointManagerOptions()
-            checkpoint_dir = checkpoint_dir = os.path.join(
-                os.getcwd(), self.params["output"]["path"]
-            )
+            checkpoint_dir = os.path.join(os.getcwd(), self.params["output"]["path"])
             path = os.path.join(os.getcwd(), checkpoint_dir)
             if jax.process_index() == 0:
                 path = ocp.test_utils.erase_and_create_empty(checkpoint_dir)
@@ -336,80 +416,91 @@ class Simulation:
 
         def _kick(state, kx, ky, kz, k_sq, dt):
             # Kick (half-step)
-            if (
-                self.params["physics"]["gravity"]
-                and self.params["physics"]["external_potential"]
-            ):
+            if use_gravity and use_external_potential:
                 V = self._calc_grav_potential(state, k_sq) + state["V_ext"]
-            elif self.params["physics"]["gravity"]:
+            elif use_gravity:
                 V = self._calc_grav_potential(state, k_sq)
-            elif self.params["physics"]["external_potential"]:
+            elif use_external_potential:
                 V = state["V_ext"]
 
-            if (
-                self.params["physics"]["gravity"]
-                or self.params["physics"]["external_potential"]
-            ):
-                if self.params["physics"]["quantum"]:
-                    state["psi"] = quantum_kick(state["psi"], V, m_per_hbar, dt)
-                if self.params["physics"]["hydro"]:
+            if use_gravity or use_external_potential:
+                if use_quantum:
+                    if do_self_interaction:
+                        rho = jnp.abs(state["psi"]) ** 2
+                        V_prime = V + si_coeff * rho + si_coeff2 * rho**2
+                        state["psi"] = quantum_kick(
+                            state["psi"], V_prime, m_per_hbar, dt
+                        )
+                    else:
+                        state["psi"] = quantum_kick(state["psi"], V, m_per_hbar, dt)
+                if use_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"]:
+                if use_particles:
                     state["vel"] = particles_accelerate(
                         state["vel"], state["pos"], V, kx, ky, kz, dx, dt
                     )
+                if use_custom:
+                    state = custom_kick(state, V, dt)
+
+            return state
 
         def _drift(state, k_sq, dt):
             # Drift (full-step)
-            if self.params["physics"]["quantum"]:
+            if use_quantum:
                 state["psi"] = quantum_drift(state["psi"], k_sq, m_per_hbar, dt)
-            if self.params["physics"]["hydro"]:
+            if use_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"]:
+            if use_particles:
                 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
 
-        @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)
+            state = _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
+            state = _drift(state, k_sq, dt)
             # update time & redshift
             state["t"] += dt
-            if self.params["physics"]["cosmology"]:
+            if use_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"], dt, omega_matter, omega_lambda, little_h
                 )
                 state["redshift"] = 1.0 / scale_factor - 1.0
-            _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
+            state = _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
 
             return state, kx, ky, kz, k_sq
 
         # save initial state
         if jax.process_index() == 0:
             print(f"Starting simulation (res={self.resolution}, nt={nt}) ...")
-        if self.params["output"]["save"]:
+        if save:
             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))
             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 self.params["output"]["save"]:
+        if save:
             for i in range(1, num_checkpoints + 1):
                 carry = jax.lax.fori_loop(0, nt_sub, _update, init_val=carry)
-                state, kx, ky, kz, k_sq = carry
+                state, _, _, _, _ = carry
                 jax.block_until_ready(state)
                 # save state
                 async_checkpoint_manager.save(i, args=ocp.args.StandardSave(state))
@@ -417,17 +508,18 @@ class Simulation:
                 elapsed = time.time() - t_start_timer
                 est_total = elapsed / i * num_checkpoints
                 est_remaining = est_total - elapsed
-                num_cells = self.resolution**3
                 mcups = (num_cells * (i * nt_sub)) / (elapsed * 1.0e6)
                 if jax.process_index() == 0:
                     print(
                         f"{percent:.1f}%: mcups={mcups:.1f}, estimated time left (s): {est_remaining:.1f}"
                     )
                 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, kx, ky, kz, k_sq = carry
+            state, _, _, _, _ = carry
         jax.block_until_ready(state)
         if jax.process_index() == 0:
             print("Simulation Run Time (s): ", time.time() - t_start_timer)

{jaxion-0.0.5 → jaxion-0.0.7}/jaxion/utils.py

@@ -31,7 +31,7 @@ def print_distributed_info():
 
 def set_up_parameters(user_overwrites):
     # first load the default params
-    params_path = importlib.resources.files("jaxion") / "params_default.json"
+    params_path = importlib.resources.files("jaxion") / "defaults.json"
     with params_path.open("r", encoding="utf-8") as f:
         params = json.load(f)
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.5
+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
@@ -68,7 +68,7 @@ Jaxion is built for multi-GPU scalability and is fully differentiable. It is a h
 Jaxion is the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
 
 
-## Getting started
+## Install Jaxion
 
 Install with:
 
@@ -87,7 +87,7 @@ See the docs for more info on how to [build from source](https://jaxion.readthed
 
 ## Examples
 
-Check out the `examples/` directory for demonstrations of using Jaxion.
+Check out the [`examples/`](https://github.com/JaxionProject/jaxion/tree/main/examples/) directory for demonstrations of using Jaxion.
 
 <p align="center">
   <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">

{jaxion-0.0.5 → jaxion-0.0.7}/jaxion.egg-info/SOURCES.txt

@@ -6,9 +6,9 @@ jaxion/__init__.py
 jaxion/analysis.py
 jaxion/constants.py
 jaxion/cosmology.py
+jaxion/defaults.json
 jaxion/gravity.py
 jaxion/hydro.py
-jaxion/params_default.json
 jaxion/particles.py
 jaxion/quantum.py
 jaxion/simulation.py

{jaxion-0.0.5 → jaxion-0.0.7}/pyproject.toml

@@ -10,7 +10,7 @@ license = "Apache-2.0"
 dynamic = ["version", "dependencies"]
 
 [tool.setuptools.packages.find]
-include = ["jaxion"]
+include = ["jaxio*"]
 exclude = ["paper"]
 
 [tool.setuptools.package-data]

{jaxion-0.0.5 → jaxion-0.0.7}/tests/test_examples.py

@@ -12,7 +12,7 @@ def test_tidal_stripping():
     )
     assert sim.resolution == 32
     assert sim.state["t"] > 0.0
-    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(639.0479, rel=rel_tol)
+    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(162.028, rel=rel_tol)
 
 
 def test_tidal_stripping_distributed_emulate():
@@ -22,7 +22,7 @@ def test_tidal_stripping_distributed_emulate():
     )
     assert sim.resolution == 32
     assert sim.state["t"] > 0.0
-    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(639.0479, rel=rel_tol)
+    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(162.028, rel=rel_tol)
 
 
 def test_heating_gas():
@@ -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
     )