jaxion 0.0.3__tar.gz → 0.0.5__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.3
+Version: 0.0.5
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -32,8 +32,12 @@ Dynamic: license-file
 [![Repo Status][status-badge]][status-link]
 [![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
@@ -41,10 +45,19 @@ Dynamic: license-file
 [pypi-badge]:          https://img.shields.io/pypi/v/jaxion?label=PyPI&logo=pypi
 [workflow-test-link]:  https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml
 [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!
 
@@ -63,34 +76,86 @@ Install with:
 pip install jaxion
 ```
 
-or, for GPU support use:
+or, for GPU support, use:
 
 ```console
 pip install jaxion[cuda12]
 ```
 
-Check out the `examples/` directory for demonstrations of using Jaxion.
+See the docs for more info on how to [build from source](https://jaxion.readthedocs.io/en/latest/pages/installation.html).
 
 
 ## Examples
 
+Check out the `examples/` directory for demonstrations of using Jaxion.
+
 <p align="center">
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">
+    <img src="examples/cosmological_box/movie.gif" alt="cosmological_box" width="128"/>
+  </a>
   <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>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_stars">
+    <img src="examples/heating_stars/movie.gif" alt="heating_stars" width="128"/>
+  </a>
+  <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_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">
+    <img src="examples/soliton_binary_merger/movie.gif" alt="soliton_binary_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_merger">
+    <img src="examples/soliton_merger/movie.gif" alt="soliton_merger" 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>
 
+
+## High-Performance
+
+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>
+
+
+## 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 running jaxion.
+* [Documentation](https://jaxion.readthedocs.io) for up-to-date information about installing and using Jaxion.
 
 
 ## Cite this repository
 
+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.5/README.md

@@ -0,0 +1,138 @@
+<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]
+[![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
+[pypi-link]:           https://pypi.org/project/jaxion
+[pypi-badge]:          https://img.shields.io/pypi/v/jaxion?label=PyPI&logo=pypi
+[workflow-test-link]:  https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml
+[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!
+
+Author: [Philip Mocz (@pmocz)](https://github.com/pmocz/)
+
+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 the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
+
+
+## Getting started
+
+Install with:
+
+```console
+pip install jaxion
+```
+
+or, for GPU support, use:
+
+```console
+pip install jaxion[cuda12]
+```
+
+See the docs for more info on how to [build from source](https://jaxion.readthedocs.io/en/latest/pages/installation.html).
+
+
+## Examples
+
+Check out the `examples/` directory for demonstrations of using Jaxion.
+
+<p align="center">
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">
+    <img src="examples/cosmological_box/movie.gif" alt="cosmological_box" width="128"/>
+  </a>
+  <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>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_stars">
+    <img src="examples/heating_stars/movie.gif" alt="heating_stars" width="128"/>
+  </a>
+  <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_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">
+    <img src="examples/soliton_binary_merger/movie.gif" alt="soliton_binary_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_merger">
+    <img src="examples/soliton_merger/movie.gif" alt="soliton_merger" 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>
+
+
+## High-Performance
+
+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>
+
+
+## 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
+
+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.3 → jaxion-0.0.5}/jaxion/analysis.py

@@ -4,10 +4,25 @@ 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
+    Computes the radially averaged power spectral density of a 3D datacube.
+
+    Parameters
+    ----------
+        data_cube : jnp.ndarray
+            3D data cube, must be cubic
+        kx, ky, kz: jnp.ndarray
+            wavenumber grids in each dimension
+        box_size: float
+            physical size of box
+
+    Returns
+    -------
+        Pf: jnp.ndarray
+            radial power spectrum
+        k: jnp.ndarray
+            wavenumbers
+        total_power: float
+            total power
     """
     dim = data_cube.ndim
     nx = data_cube.shape[0]

{jaxion-0.0.3 → jaxion-0.0.5}/jaxion/constants.py

@@ -1,13 +1,13 @@
 # Copyright (c) 2025 The Jaxion Team.
 
-# Physical Constants
-#
-# Jaxion uses a unit system of:
-# [L] = kpc
-# [V] = km/s
-# [M] = Msun
-#
-# other units are derived from these base units, e.g., [T] = [L]/[V] = kpc / (km/s) ~= 0.978 Gyr
+"""
+Physical constants in units of:
+[L] = kpc,
+[V] = km/s,
+[M] = Msun
+
+note: other units are derived from these base units, e.g., [T] = [L]/[V] = kpc / (km/s) ~= 0.978 Gyr
+"""
 
 constants = {
     "gravitational_constant": 4.30241002e-6,  # G / (kpc * (km/s)^2 / Msun)

{jaxion-0.0.3 → jaxion-0.0.5}/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,
@@ -28,8 +29,11 @@ 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.
+        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.
+        sharding : jax.sharding.NamedSharding, optional
+            jax sharding used for distributed (multi-GPU) simulations
 
     """
 
@@ -69,10 +73,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(
@@ -126,6 +132,9 @@ class Simulation:
 
     @property
     def resolution(self):
+        """
+        Return the (linear) resolution of the simulation
+        """
         return (
             self.params["domain"]["resolution_base"]
             * self.params["domain"]["resolution_multiplier"]
@@ -133,18 +142,30 @@ class Simulation:
 
     @property
     def num_particles(self):
+        """
+        Return the number of particles in the simulation
+        """
         return self.params["particles"]["num_particles"]
 
     @property
     def box_size(self):
+        """
+        Return the box size of the simulation (kpc)
+        """
         return self.params["domain"]["box_size"]
 
     @property
     def dx(self):
+        """
+        Return the cell size size of the simulation (kpc)
+        """
         return self.box_size / self.resolution
 
     @property
     def axion_mass(self):
+        """
+        Return the axion particle mass in the simulation (M_sun)
+        """
         return (
             self.params["quantum"]["m_22"]
             * 1.0e-22
@@ -152,16 +173,32 @@ class Simulation:
             / constants["speed_of_light"] ** 2
         )
 
+    @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):
+        """
+        Return the isothermal gas sound speed in the simulation (km/s)
+        """
         return self.params["hydro"]["sound_speed"]
 
     @property
     def params(self):
+        """
+        Return the parameters of the simulation
+        """
         return self._params
 
     @property
     def grid(self):
+        """
+        Return the simulation grid
+        """
         hx = 0.5 * self.dx
         x_lin = jnp.linspace(hx, self.box_size - hx, self.resolution)
         xx, yy, zz = self.xmeshgrid_jit(x_lin)
@@ -169,6 +206,9 @@ class Simulation:
 
     @property
     def kgrid(self):
+        """
+        Return the simulation spectral grid
+        """
         nx = self.resolution
         k_lin = (2.0 * jnp.pi / self.box_size) * jnp.arange(-nx / 2, nx / 2)
         kx, ky, kz = self.xmeshgrid_transpose_jit(k_lin)
@@ -179,8 +219,10 @@ class Simulation:
 
     @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)
+        """
+        Return the dark matter velocity field from the wavefunction
+        """
+        return quantum_velocity(self.state["psi"], self.box_size, self.m_per_hbar)
 
     def _calc_rho_bar(self, state):
         rho_bar = 0.0
@@ -214,6 +256,9 @@ class Simulation:
 
     @property
     def potential(self):
+        """
+        Return the gravitational potential
+        """
         kx, ky, kz = self.kgrid
         k_sq = kx**2 + ky**2 + kz**2
         return self._calc_grav_potential(self.state, k_sq)
@@ -235,7 +280,7 @@ class Simulation:
 
         # Simulation parameters
         dx = self.dx
-        m_per_hbar = self.axion_mass / constants["reduced_planck_constant"]
+        m_per_hbar = self.m_per_hbar
         box_size = self.box_size
 
         dt_fac = 1.0
@@ -372,14 +417,17 @@ 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}%: 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)
                 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, kx, ky, kz, k_sq = carry
         jax.block_until_ready(state)
         if jax.process_index() == 0:
             print("Simulation Run Time (s): ", time.time() - t_start_timer)
@@ -387,5 +435,8 @@ class Simulation:
         return state
 
     def run(self):
+        """
+        Run the simulation
+        """
         self.state = self._evolve(self.state)
         jax.block_until_ready(self.state)

{jaxion-0.0.3 → jaxion-0.0.5}/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,6 +13,22 @@ 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"

{jaxion-0.0.3 → jaxion-0.0.5}/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(jnp.abs(state["rho"]) ** 2, 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:
@@ -48,7 +52,7 @@ def plot_sim(state, checkpoint_dir, i, params):
                 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
+                    sx, sy, color="cyan", marker=".", linestyle="None", markersize=5
                 )
             ax.set_aspect("equal")
             ax.get_xaxis().set_visible(False)

{jaxion-0.0.3 → jaxion-0.0.5}/jaxion.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.3
+Version: 0.0.5
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -32,8 +32,12 @@ Dynamic: license-file
 [![Repo Status][status-badge]][status-link]
 [![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
@@ -41,10 +45,19 @@ Dynamic: license-file
 [pypi-badge]:          https://img.shields.io/pypi/v/jaxion?label=PyPI&logo=pypi
 [workflow-test-link]:  https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml
 [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!
 
@@ -63,34 +76,86 @@ Install with:
 pip install jaxion
 ```
 
-or, for GPU support use:
+or, for GPU support, use:
 
 ```console
 pip install jaxion[cuda12]
 ```
 
-Check out the `examples/` directory for demonstrations of using Jaxion.
+See the docs for more info on how to [build from source](https://jaxion.readthedocs.io/en/latest/pages/installation.html).
 
 
 ## Examples
 
+Check out the `examples/` directory for demonstrations of using Jaxion.
+
 <p align="center">
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/cosmological_box">
+    <img src="examples/cosmological_box/movie.gif" alt="cosmological_box" width="128"/>
+  </a>
   <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>
+  <br>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/heating_stars">
+    <img src="examples/heating_stars/movie.gif" alt="heating_stars" width="128"/>
+  </a>
+  <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_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">
+    <img src="examples/soliton_binary_merger/movie.gif" alt="soliton_binary_merger" width="128"/>
+  </a>
+  <a href="https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_merger">
+    <img src="examples/soliton_merger/movie.gif" alt="soliton_merger" 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>
 
+
+## High-Performance
+
+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>
+
+
+## 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 running jaxion.
+* [Documentation](https://jaxion.readthedocs.io) for up-to-date information about installing and using Jaxion.
 
 
 ## Cite this repository
 
+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.3 → jaxion-0.0.5}/jaxion.egg-info/SOURCES.txt

@@ -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.5/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.5/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.3 → jaxion-0.0.5}/tests/test_examples.py

@@ -7,7 +7,18 @@ 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)
+
+
+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

jaxion-0.0.3/README.md

@@ -1,73 +0,0 @@
-<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]
-[![PyPI Version Status][pypi-badge]][pypi-link]
-[![Test Status][workflow-test-badge]][workflow-test-link]
-[![Readthedocs Status][docs-badge]][docs-link]
-[![License][license-badge]][license-link]
-
-[status-link]:         https://www.repostatus.org/#active
-[status-badge]:        https://www.repostatus.org/badges/latest/active.svg
-[pypi-link]:           https://pypi.org/project/jaxion
-[pypi-badge]:          https://img.shields.io/pypi/v/jaxion?label=PyPI&logo=pypi
-[workflow-test-link]:  https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml
-[workflow-test-badge]: https://github.com/JaxionProject/jaxion/actions/workflows/test-package.yml/badge.svg?event=push
-[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
-
-A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
-
-Author: [Philip Mocz (@pmocz)](https://github.com/pmocz/)
-
-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 the simpler companion project to differentiable astrophysics code [Adirondax](https://github.com/AdirondaxProject/adirondax)
-
-
-## Getting started
-
-Install with:
-
-```console
-pip install jaxion
-```
-
-or, for GPU support use:
-
-```console
-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).
-* [Documentation](https://jaxion.readthedocs.io) for up-to-date information about installing and running jaxion.
-
-
-## Cite this repository
-