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

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,
@@ -72,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(
@@ -170,6 +173,13 @@ 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):
         """
@@ -212,8 +222,7 @@ 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)
 
     def _calc_rho_bar(self, state):
         rho_bar = 0.0
@@ -271,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
@@ -408,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)

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/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.dist-info → jaxion-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.4
+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
@@ -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!
 
@@ -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.dist-info → jaxion-0.0.5.dist-info}/RECORD

@@ -7,11 +7,11 @@ 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=s6gCAt-gAoN5d46vcdxoqtn4TwsrfNGb4Cq-2p_JxsI,15927
-jaxion/utils.py,sha256=rT7NM0FNEgFwN7oTgTb-jkR66Iw0xYTHHxcoikYd1ag,3572
-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,,
+jaxion/simulation.py,sha256=2YkHh3tUVvg5tUNnOxf4s7wGeuMntYUJcJWV0M-3Pl8,16267
+jaxion/utils.py,sha256=f8SvJjqzcW2K91qbPNqrsjfVjyPShuf50yoSHc0YqYE,4093
+jaxion/visualization.py,sha256=Vx_xuEE7BB1AkEGaY9KHNSIIDxJzUVT104o-3uglW8o,2966
+jaxion-0.0.5.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+jaxion-0.0.5.dist-info/METADATA,sha256=Ot_IWpsCk6l7KX94fCs4j1LLcJQLlQFmp-jYCyj_yWY,6378
+jaxion-0.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jaxion-0.0.5.dist-info/top_level.txt,sha256=S1OV2VdlDG_9UwpKOIji4itQGOS-VWUOWUi3GeXWzt0,7
+jaxion-0.0.5.dist-info/RECORD,,