jaxion 0.0.7__tar.gz → 0.0.9__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.7
+Version: 0.0.9
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -9,16 +9,17 @@ Project-URL: Homepage, https://github.com/JaxionProject/jaxion
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: jax==0.5.3
-Requires-Dist: jaxdecomp==0.2.7
+Requires-Dist: jax==0.6.0
+Requires-Dist: jaxdecomp==0.2.9
 Requires-Dist: tensorflow
-Requires-Dist: orbax-checkpoint==0.11.18
-Requires-Dist: optax==0.2.5
+Requires-Dist: orbax-checkpoint==0.11.25
+Requires-Dist: optax==0.2.6
 Requires-Dist: numpy
 Requires-Dist: matplotlib
 Requires-Dist: setuptools>=70.1.1
+Requires-Dist: chex==0.1.90
 Provides-Extra: cuda12
-Requires-Dist: jax[cuda12]==0.5.3; extra == "cuda12"
+Requires-Dist: jax[cuda12]==0.6.0; extra == "cuda12"
 Dynamic: license-file
 
 <p align="center">

{jaxion-0.0.7 → jaxion-0.0.9}/jaxion/simulation.py

@@ -42,7 +42,7 @@ class Simulation:
 
     """
 
-    def __init__(self, params, sharding=None):
+    def __init__(self, params, sharding=None, checkpoint_number=None):
         # allow loading directly from a checkpoint path
         load_from_checkpoint = False
         checkpoint_dir = ""
@@ -105,6 +105,7 @@ class Simulation:
         self.custom_drift = None
         self.custom_density = None
         self.custom_plot = None
+        self.callback = None
 
         # simulation state
         self.state = {}
@@ -142,7 +143,10 @@ class Simulation:
             async_checkpoint_manager = ocp.CheckpointManager(
                 checkpoint_dir, options=options
             )
-            step = async_checkpoint_manager.latest_step()
+            if checkpoint_number is not None:
+                step = checkpoint_number
+            else:
+                step = async_checkpoint_manager.latest_step()
             self.state = async_checkpoint_manager.restore(
                 step, args=ocp.args.StandardRestore(self.state)
             )
@@ -178,6 +182,23 @@ class Simulation:
         """
         return self.box_size / self.resolution
 
+    @property
+    def nt(self):
+        """
+        Return the number of timesteps
+        """
+        dx = self.dx
+        m_per_hbar = self.m_per_hbar
+        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
+        num_checkpoints = self.params["output"]["num_checkpoints"]
+        nt = int(round(round(t_span / dt_kin) / num_checkpoints) * num_checkpoints)
+
+        return nt
+
     @property
     def axion_mass(self):
         """
@@ -341,8 +362,6 @@ class Simulation:
         num_cells = self.resolution**3
         m_per_hbar = self.m_per_hbar
 
-        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
@@ -360,6 +379,9 @@ class Simulation:
         if use_custom:
             custom_kick = self.custom_kick
             custom_drift = self.custom_drift
+        use_callback = self.callback is not None
+        if use_callback:
+            custom_callback = self.callback
 
         # cosmology
         if use_cosmology:
@@ -390,7 +412,7 @@ class Simulation:
 
         # round up to the nearest multiple of num_checkpoints
         num_checkpoints = self.params["output"]["num_checkpoints"]
-        nt = int(round(round(t_span / dt_kin) / num_checkpoints) * num_checkpoints)
+        nt = self.nt
         nt_sub = int(round(nt / num_checkpoints))
         dt = t_span / nt
 
@@ -466,7 +488,7 @@ class Simulation:
 
             return state
 
-        def _update(_, carry):
+        def _update(i, carry):
             # Update the simulation state by one timestep
             # according to a 2nd-order `kick-drift-kick` scheme
             state, kx, ky, kz, k_sq = carry
@@ -481,6 +503,9 @@ class Simulation:
                 state["redshift"] = 1.0 / scale_factor - 1.0
             state = _kick(state, kx, ky, kz, k_sq, 0.5 * dt)
 
+            if use_callback:
+                custom_callback(i, state)
+
             return state, kx, ky, kz, k_sq
 
         # save initial state
@@ -490,16 +515,19 @@ class Simulation:
             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))
+            async_checkpoint_manager.wait_until_finished()
             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 save:
             for i in range(1, num_checkpoints + 1):
-                carry = jax.lax.fori_loop(0, nt_sub, _update, init_val=carry)
+                i_timestep = (i - 1) * nt_sub
+                carry = jax.lax.fori_loop(
+                    i_timestep, i_timestep + nt_sub, _update, init_val=carry
+                )
                 state, _, _, _, _ = carry
                 jax.block_until_ready(state)
                 # save state
@@ -513,10 +541,10 @@ class Simulation:
                     print(
                         f"{percent:.1f}%: mcups={mcups:.1f}, estimated time left (s): {est_remaining:.1f}"
                     )
+                async_checkpoint_manager.wait_until_finished()
                 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, _, _, _, _ = carry

{jaxion-0.0.7 → jaxion-0.0.9}/jaxion.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxion
-Version: 0.0.7
+Version: 0.0.9
 Summary: A differentiable simulation library for fuzzy dark matter in JAX
 Author-email: Philip Mocz <philip.mocz@gmail.com>
 License-Expression: Apache-2.0
@@ -9,16 +9,17 @@ Project-URL: Homepage, https://github.com/JaxionProject/jaxion
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: jax==0.5.3
-Requires-Dist: jaxdecomp==0.2.7
+Requires-Dist: jax==0.6.0
+Requires-Dist: jaxdecomp==0.2.9
 Requires-Dist: tensorflow
-Requires-Dist: orbax-checkpoint==0.11.18
-Requires-Dist: optax==0.2.5
+Requires-Dist: orbax-checkpoint==0.11.25
+Requires-Dist: optax==0.2.6
 Requires-Dist: numpy
 Requires-Dist: matplotlib
 Requires-Dist: setuptools>=70.1.1
+Requires-Dist: chex==0.1.90
 Provides-Extra: cuda12
-Requires-Dist: jax[cuda12]==0.5.3; extra == "cuda12"
+Requires-Dist: jax[cuda12]==0.6.0; extra == "cuda12"
 Dynamic: license-file
 
 <p align="center">

jaxion-0.0.9/jaxion.egg-info/requires.txt

@@ -0,0 +1,12 @@
+jax==0.6.0
+jaxdecomp==0.2.9
+tensorflow
+orbax-checkpoint==0.11.25
+optax==0.2.6
+numpy
+matplotlib
+setuptools>=70.1.1
+chex==0.1.90
+
+[cuda12]
+jax[cuda12]==0.6.0

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

@@ -25,7 +25,7 @@ dependencies = {file = ["requirements.txt"]}
 
 [project.optional-dependencies]
 cuda12 = [
-    "jax[cuda12]==0.5.3",
+    "jax[cuda12]==0.6.0",
 ]
 
 [tool.setuptools-git-versioning]

jaxion-0.0.9/requirements.txt

@@ -0,0 +1,9 @@
+jax==0.6.0
+jaxdecomp==0.2.9
+tensorflow
+orbax-checkpoint==0.11.25
+optax==0.2.6
+numpy
+matplotlib
+setuptools>=70.1.1
+chex==0.1.90

jaxion-0.0.7/jaxion.egg-info/requires.txt

@@ -1,11 +0,0 @@
-jax==0.5.3
-jaxdecomp==0.2.7
-tensorflow
-orbax-checkpoint==0.11.18
-optax==0.2.5
-numpy
-matplotlib
-setuptools>=70.1.1
-
-[cuda12]
-jax[cuda12]==0.5.3

jaxion-0.0.7/requirements.txt

@@ -1,8 +0,0 @@
-jax==0.5.3
-jaxdecomp==0.2.7
-tensorflow
-orbax-checkpoint==0.11.18
-optax==0.2.5
-numpy
-matplotlib
-setuptools>=70.1.1