jaxion 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

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/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/simulation.py

@@ -28,8 +28,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
 
     """
 
@@ -126,6 +129,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 +139,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
@@ -154,14 +172,23 @@ class Simulation:
 
     @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 +196,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,6 +209,9 @@ class Simulation:
 
     @property
     def quantum_velocity(self):
+        """
+        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)
 
@@ -214,6 +247,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)
@@ -387,5 +423,8 @@ class Simulation:
         return state
 
     def run(self):
+        """
+        Run the simulation
+        """
         self.state = self._evolve(self.state)
         jax.block_until_ready(self.state)

jaxion/visualization.py

@@ -20,7 +20,7 @@ def plot_sim(state, checkpoint_dir, i, params):
         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)
+            jnp.log10(jnp.mean(state["rho"], axis=2).T)
         ).reshape(nx, nx)
 
     # create plot on process 0
@@ -48,7 +48,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.dist-info → jaxion-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.3
+Version: 0.0.4
 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,6 +32,7 @@ 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]
 [![Readthedocs Status][docs-badge]][docs-link]
 [![License][license-badge]][license-link]
 
@@ -41,6 +42,8 @@ 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
 [docs-link]:           https://jaxion.readthedocs.io
 [docs-badge]:          https://readthedocs.org/projects/jaxion/badge
 [license-link]:        https://opensource.org/licenses/Apache-2.0
@@ -63,34 +66,68 @@ 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">
+    <img src="examples/logo/movie.gif" alt="logo" 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>
 
+
 ## 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.
+
+
+## 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)
 
 
 ## Cite this repository
 
+TODO XXX

{jaxion-0.0.3.dist-info → jaxion-0.0.4.dist-info}/RECORD

@@ -1,17 +1,17 @@
 jaxion/__init__.py,sha256=Hdji1UQ47lG24Pqcy6UUq9L0-qy6m9Ax41L0vIYzBto,164
-jaxion/analysis.py,sha256=UZpCvDr2ISDAgRoZFtuhvDk0QtN3mn5NxIoYYY8wzs8,1304
-jaxion/constants.py,sha256=b0vgqtezIEGp4O5JRJo3YK3eQC8AN6hmV84-4pYWgmc,528
+jaxion/analysis.py,sha256=4YT9Z2dkFoXwft3fQM1HyynVPlIdtRd80VtI2vWTyq4,1568
+jaxion/constants.py,sha256=HyY2ktKQakv78jD1yQvFdM3sklUJcPgDMYlTsSPQTxI,512
 jaxion/cosmology.py,sha256=UC1McXNTXGoPRYXn0nI2-csVkJWL-ZBNoCa44oU1b4w,2681
 jaxion/gravity.py,sha256=3brRZelKm-soXqk_Lt3SqhbZ00woJCraqwdMuR-KooA,291
 jaxion/hydro.py,sha256=KoJ02tRpAc4V3Ofzw4zbHLRaE2GdIatbOBE04_LsSRw,6980
 jaxion/params_default.json,sha256=9CJrhEPsv5zGEs7_WqFyuccCDipPCDhXgKzVdqOsOWE,2775
 jaxion/particles.py,sha256=pMopGvoZ0J_3EviD0WnTMmiebU9h2_8IO-p6I-E5DEU,3980
 jaxion/quantum.py,sha256=GWOpN6ipfEw-6Ah2zQpxS3oqeSt_iHMDSgnVYSjXY5E,3321
-jaxion/simulation.py,sha256=3UrrHPJr0_39BiP86nQSl4ui8EsPHf9RrZ5hDobit5k,14884
+jaxion/simulation.py,sha256=s6gCAt-gAoN5d46vcdxoqtn4TwsrfNGb4Cq-2p_JxsI,15927
 jaxion/utils.py,sha256=rT7NM0FNEgFwN7oTgTb-jkR66Iw0xYTHHxcoikYd1ag,3572
-jaxion/visualization.py,sha256=MJ3iqrvAem4RbozM79Ri2cNTIwKUV8o5X8M32FiDfbo,2928
-jaxion-0.0.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-jaxion-0.0.3.dist-info/METADATA,sha256=1C3fwJHJEBZYMEhxmTkKd7z_ZUrF15E3P0oS5XxrdjA,3579
-jaxion-0.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jaxion-0.0.3.dist-info/top_level.txt,sha256=S1OV2VdlDG_9UwpKOIji4itQGOS-VWUOWUi3GeXWzt0,7
-jaxion-0.0.3.dist-info/RECORD,,
+jaxion/visualization.py,sha256=K5EQOHPfj7LF29fW_naWH8a7TEyEa3wIaQw7rpebx0w,2914
+jaxion-0.0.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+jaxion-0.0.4.dist-info/METADATA,sha256=66lU0x1ZofP-uCwD4hF0C45Ao-pNhPvMzr6Krs__Hws,5305
+jaxion-0.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jaxion-0.0.4.dist-info/top_level.txt,sha256=S1OV2VdlDG_9UwpKOIji4itQGOS-VWUOWUi3GeXWzt0,7
+jaxion-0.0.4.dist-info/RECORD,,