jaxion 0.0.4__tar.gz → 0.0.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.4
+Version: 0.0.6
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -33,8 +33,11 @@ Dynamic: license-file
 [![PyPI Version Status][pypi-badge]][pypi-link]
 [![Test Status][workflow-test-badge]][workflow-test-link]
 [![Coverage][coverage-badge]][coverage-link]
+[![Ruff][ruff-badge]][ruff-link]
+[![asv][asv-badge]][asv-link]
 [![Readthedocs Status][docs-badge]][docs-link]
 [![License][license-badge]][license-link]
+[![Software DOI][software-doi-badge]][software-doi-link]
 
 [status-link]:         https://www.repostatus.org/#active
 [status-badge]:        https://www.repostatus.org/badges/latest/active.svg
@@ -44,10 +47,17 @@ Dynamic: license-file
 [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
+[ruff-link]:           https://github.com/astral-sh/ruff
+[ruff-badge]:          https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
+[asv-link]:            https://jaxionproject.github.io/jaxion-benchmarks/
+[asv-badge]:           https://img.shields.io/badge/benchmarked%20by-asv-blue.svg?style=flat
 [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
+[software-doi-link]:   https://doi.org/10.5281/zenodo.17438467
+[software-doi-badge]:  https://zenodo.org/badge/1072645376.svg
+
 
 A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
 
@@ -58,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:
 
@@ -77,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">
@@ -96,8 +106,8 @@ Check out the `examples/` directory for demonstrations of using Jaxion.
   <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 href="https://github.com/JaxionProject/jaxion/tree/main/examples/logo_inverse_problem">
+    <img src="examples/logo_inverse_problem/movie.gif" alt="logo_inverse_problem" width="128"/>
   </a>
   <br>
   <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_binary_merger">
@@ -112,22 +122,40 @@ Check out the `examples/` directory for demonstrations of using Jaxion.
 </p>
 
 
-## Links
+## High-Performance
 
-* [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.
+Jaxion is scalable on multiple GPUs!
 
+<p align="center">
+  <a href="https://jaxion.readthedocs.io">
+    <img src="examples/soliton_binary_merger/timing.png" alt="timing" width="400"/>
+  </a>
+</p>
 
-## 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)
 
-## 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)
+## 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.
 
 
 ## Cite this repository
 
-TODO XXX
+If you use this software, please cite it as below.
+
+```bibtex
+@software{Mocz_Jaxion_2025,
+   author = {Mocz, Philip},
+      doi = {10.5281/zenodo.17438467},
+    month = oct,
+    title = {{Jaxion}},
+      url = {https://github.com/JaxionProject/jaxion},
+  version = {0.0.4},
+     year = {2025}
+}
+```

{jaxion-0.0.4 → jaxion-0.0.6}/README.md

@@ -10,8 +10,11 @@
 [![PyPI Version Status][pypi-badge]][pypi-link]
 [![Test Status][workflow-test-badge]][workflow-test-link]
 [![Coverage][coverage-badge]][coverage-link]
+[![Ruff][ruff-badge]][ruff-link]
+[![asv][asv-badge]][asv-link]
 [![Readthedocs Status][docs-badge]][docs-link]
 [![License][license-badge]][license-link]
+[![Software DOI][software-doi-badge]][software-doi-link]
 
 [status-link]:         https://www.repostatus.org/#active
 [status-badge]:        https://www.repostatus.org/badges/latest/active.svg
@@ -21,10 +24,17 @@
 [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
+[ruff-link]:           https://github.com/astral-sh/ruff
+[ruff-badge]:          https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
+[asv-link]:            https://jaxionproject.github.io/jaxion-benchmarks/
+[asv-badge]:           https://img.shields.io/badge/benchmarked%20by-asv-blue.svg?style=flat
 [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
+[software-doi-link]:   https://doi.org/10.5281/zenodo.17438467
+[software-doi-badge]:  https://zenodo.org/badge/1072645376.svg
+
 
 A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
 
@@ -35,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:
 
@@ -54,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">
@@ -73,8 +83,8 @@ Check out the `examples/` directory for demonstrations of using Jaxion.
   <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 href="https://github.com/JaxionProject/jaxion/tree/main/examples/logo_inverse_problem">
+    <img src="examples/logo_inverse_problem/movie.gif" alt="logo_inverse_problem" width="128"/>
   </a>
   <br>
   <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_binary_merger">
@@ -89,22 +99,40 @@ Check out the `examples/` directory for demonstrations of using Jaxion.
 </p>
 
 
-## Links
+## High-Performance
 
-* [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.
+Jaxion is scalable on multiple GPUs!
 
+<p align="center">
+  <a href="https://jaxion.readthedocs.io">
+    <img src="examples/soliton_binary_merger/timing.png" alt="timing" width="400"/>
+  </a>
+</p>
 
-## 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)
 
-## 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)
+## 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.
 
 
 ## Cite this repository
 
-TODO XXX
+If you use this software, please cite it as below.
+
+```bibtex
+@software{Mocz_Jaxion_2025,
+   author = {Mocz, Philip},
+      doi = {10.5281/zenodo.17438467},
+    month = oct,
+    title = {{Jaxion}},
+      url = {https://github.com/JaxionProject/jaxion},
+  version = {0.0.4},
+     year = {2025}
+}
+```

jaxion-0.0.4/jaxion/params_default.json → jaxion-0.0.6/jaxion/defaults.json

@@ -75,6 +75,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": {

{jaxion-0.0.4 → jaxion-0.0.6}/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.4 → jaxion-0.0.6}/jaxion/simulation.py

@@ -14,6 +14,7 @@ from .cosmology import get_supercomoving_time_interval, get_next_scale_factor
 from .utils import (
     set_up_parameters,
     print_parameters,
+    print_distributed_info,
     xmeshgrid,
     xmeshgrid_transpose,
     xzeros,
@@ -61,9 +62,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"]:
@@ -72,10 +74,12 @@ class Simulation:
                     "hydro/particles sharding is not yet implemented."
                 )
 
-        if self.params["output"]["save"]:
-            if jax.process_index() == 0:
-                print("Simulation parameters:")
-                print_parameters(self.params)
+        # print info
+        if jax.process_index() == 0:
+            print("Simulation parameters:")
+            print_parameters(self.params)
+            if sharding is not None:
+                print_distributed_info()
 
         # jitted functions
         self.xmeshgrid_jit = jax.jit(
@@ -91,6 +95,12 @@ class Simulation:
             xones, static_argnums=0, in_shardings=None, out_shardings=sharding
         )
 
+        # customfunctions
+        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
@@ -170,6 +180,30 @@ 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):
+        """
+        Return the mass per hbar in the simulation (M_sun / hbar)
+        """
+        return self.axion_mass / constants["reduced_planck_constant"]
+
     @property
     def sound_speed(self):
         """
@@ -212,8 +246,14 @@ class Simulation:
         """
         Return the dark matter velocity field from the wavefunction
         """
-        m_per_hbar = self.axion_mass / constants["reduced_planck_constant"]
-        return quantum_velocity(self.state["psi"], self.box_size, m_per_hbar)
+        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
@@ -226,6 +266,8 @@ class Simulation:
             n_particles = self.num_particles
             box_size = self.box_size
             rho_bar += m_particle * n_particles / box_size
+        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):
@@ -239,6 +281,8 @@ 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.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
@@ -271,8 +315,9 @@ class Simulation:
 
         # Simulation parameters
         dx = self.dx
-        m_per_hbar = self.axion_mass / constants["reduced_planck_constant"]
         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)
@@ -281,8 +326,20 @@ class Simulation:
         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"]
+        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"]
@@ -294,6 +351,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"]
 
@@ -308,16 +376,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)
@@ -327,80 +393,86 @@ 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)
+
+            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))
@@ -408,14 +480,18 @@ class Simulation:
                 elapsed = time.time() - t_start_timer
                 est_total = elapsed / i * num_checkpoints
                 est_remaining = est_total - elapsed
+                mcups = (num_cells * (i * nt_sub)) / (elapsed * 1.0e6)
                 if jax.process_index() == 0:
                     print(
-                        f"{percent:.1f}%: estimated time remaining (s): {est_remaining:.1f}"
+                        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:
-            state = jax.lax.fori_loop(0, nt, _update, init_val=state)
+            carry = jax.lax.fori_loop(0, nt, _update, init_val=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.4 → jaxion-0.0.6}/jaxion/utils.py

@@ -5,6 +5,7 @@ from pathlib import Path
 import importlib.resources
 import json
 from importlib.metadata import version
+import jax
 import jax.numpy as jnp
 
 
@@ -12,9 +13,25 @@ def print_parameters(params):
     print(json.dumps(params, indent=2))
 
 
+def print_distributed_info():
+    for env_var in [
+        "SLURM_JOB_ID",
+        "SLURM_NTASKS",
+        "SLURM_NODELIST",
+        "SLURM_STEP_NODELIST",
+        "SLURM_STEP_GPUS",
+        "SLURM_GPUS",
+    ]:
+        print(f"{env_var}: {os.getenv(env_var, '')}")
+    print("Total number of processes: ", jax.process_count())
+    print("Total number of devices: ", jax.device_count())
+    print("List of devices: ", jax.devices())
+    print("Number of devices on this process: ", jax.local_device_count())
+
+
 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.4 → jaxion-0.0.6}/jaxion/visualization.py

@@ -13,15 +13,19 @@ def plot_sim(state, checkpoint_dir, i, params):
     if params["physics"]["quantum"]:
         nx = state["psi"].shape[0]
         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)
+        rho_proj_dm = jnp.log10(
+            jax.experimental.multihost_utils.process_allgather(
+                jnp.mean(jnp.abs(state["psi"]) ** 2, axis=2), tiled=True
+            )
+        ).T
     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(state["rho"], axis=2).T)
-        ).reshape(nx, nx)
+        rho_proj_gas = jnp.log10(
+            jax.experimental.multihost_utils.process_allgather(
+                jnp.mean(state["rho"], axis=2), tiled=True
+            )
+        ).T
 
     # create plot on process 0
     if jax.process_index() == 0:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.4
+Version: 0.0.6
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -33,8 +33,11 @@ Dynamic: license-file
 [![PyPI Version Status][pypi-badge]][pypi-link]
 [![Test Status][workflow-test-badge]][workflow-test-link]
 [![Coverage][coverage-badge]][coverage-link]
+[![Ruff][ruff-badge]][ruff-link]
+[![asv][asv-badge]][asv-link]
 [![Readthedocs Status][docs-badge]][docs-link]
 [![License][license-badge]][license-link]
+[![Software DOI][software-doi-badge]][software-doi-link]
 
 [status-link]:         https://www.repostatus.org/#active
 [status-badge]:        https://www.repostatus.org/badges/latest/active.svg
@@ -44,10 +47,17 @@ Dynamic: license-file
 [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
+[ruff-link]:           https://github.com/astral-sh/ruff
+[ruff-badge]:          https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
+[asv-link]:            https://jaxionproject.github.io/jaxion-benchmarks/
+[asv-badge]:           https://img.shields.io/badge/benchmarked%20by-asv-blue.svg?style=flat
 [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
+[software-doi-link]:   https://doi.org/10.5281/zenodo.17438467
+[software-doi-badge]:  https://zenodo.org/badge/1072645376.svg
+
 
 A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
 
@@ -58,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:
 
@@ -77,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">
@@ -96,8 +106,8 @@ Check out the `examples/` directory for demonstrations of using Jaxion.
   <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 href="https://github.com/JaxionProject/jaxion/tree/main/examples/logo_inverse_problem">
+    <img src="examples/logo_inverse_problem/movie.gif" alt="logo_inverse_problem" width="128"/>
   </a>
   <br>
   <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_binary_merger">
@@ -112,22 +122,40 @@ Check out the `examples/` directory for demonstrations of using Jaxion.
 </p>
 
 
-## Links
+## High-Performance
 
-* [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.
+Jaxion is scalable on multiple GPUs!
 
+<p align="center">
+  <a href="https://jaxion.readthedocs.io">
+    <img src="examples/soliton_binary_merger/timing.png" alt="timing" width="400"/>
+  </a>
+</p>
 
-## 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)
 
-## 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)
+## 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.
 
 
 ## Cite this repository
 
-TODO XXX
+If you use this software, please cite it as below.
+
+```bibtex
+@software{Mocz_Jaxion_2025,
+   author = {Mocz, Philip},
+      doi = {10.5281/zenodo.17438467},
+    month = oct,
+    title = {{Jaxion}},
+      url = {https://github.com/JaxionProject/jaxion},
+  version = {0.0.4},
+     year = {2025}
+}
+```

{jaxion-0.0.4 → jaxion-0.0.6}/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
@@ -19,4 +19,6 @@ jaxion.egg-info/SOURCES.txt
 jaxion.egg-info/dependency_links.txt
 jaxion.egg-info/requires.txt
 jaxion.egg-info/top_level.txt
+tests/test_analysis.py
+tests/test_cosmology.py
 tests/test_examples.py

{jaxion-0.0.4 → jaxion-0.0.6}/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.6/tests/test_analysis.py

@@ -0,0 +1,24 @@
+import jax.numpy as jnp
+import jaxion
+from jaxion.quantum import quantum_velocity
+from jaxion.analysis import radial_power_spectrum
+import pytest
+
+
+def test_quantum_velocity_and_radial_power_spectrum():
+    sim = jaxion.Simulation({})
+    box_size = sim.box_size
+    m_per_hbar = sim.m_per_hbar
+    xx, yy, _ = sim.grid
+    kx, ky, kz = sim.kgrid
+
+    psi = (
+        jnp.cos(2.0 * jnp.pi * xx / box_size) ** 2
+        + jnp.cos(2.0 * jnp.pi * yy / box_size) * 1j
+    )
+
+    vx, _, _ = quantum_velocity(psi, box_size, m_per_hbar)
+
+    Pf, _, _ = radial_power_spectrum(vx, kx, ky, kz, box_size)
+
+    assert jnp.max(Pf) == pytest.approx(8244.606, rel=1e-4)

jaxion-0.0.6/tests/test_cosmology.py

@@ -0,0 +1,29 @@
+from jaxion.cosmology import (
+    get_physical_time_interval,
+    get_supercomoving_time_interval,
+    get_scale_factor,
+    get_next_scale_factor,
+)
+import pytest
+
+
+def test_cosmology_functions():
+    z_start = 127.0
+    z_end = 0.0
+    omega_matter = 0.3
+    omega_lambda = 0.7
+    little_h = 0.7
+    dt_hat = 10.0
+
+    assert get_physical_time_interval(
+        z_start, z_end, omega_matter, omega_lambda, little_h
+    ) == pytest.approx(13.76084)
+    assert get_supercomoving_time_interval(
+        z_start, z_end, omega_matter, omega_lambda, little_h
+    ) == pytest.approx(530.44415)
+    assert get_scale_factor(
+        z_start, dt_hat, omega_matter, omega_lambda, little_h
+    ) == pytest.approx(0.008084139320999384)
+    assert get_next_scale_factor(
+        z_start, dt_hat, omega_matter, omega_lambda, little_h
+    ) == pytest.approx(0.00808401)

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

@@ -7,11 +7,22 @@ rel_tol = 1e-4
 
 def test_tidal_stripping():
     sim = run_example_main(
-        "examples/tidal_stripping/tidal_stripping.py", argv=["--res", "1"]
+        "examples/tidal_stripping/tidal_stripping.py",
+        argv=["--res", "1", "--save", "False"],
     )
     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():
+    sim = run_example_main(
+        "examples/tidal_stripping/tidal_stripping.py",
+        argv=["--res", "1", "--distributed", "--emulate"],
+    )
+    assert sim.resolution == 32
+    assert sim.state["t"] > 0.0
+    assert jnp.mean(jnp.abs(sim.state["psi"])) == pytest.approx(162.028, rel=rel_tol)
 
 
 def test_heating_gas():