do not carry dt

Author: pmocz

adirondax/hydro/common2d.py

@@ -49,47 +49,40 @@ def slope_limit(f, dx, f_dx, f_dy):
     """
     Apply slope limiter to slopes
     """
-
-    f_dx = (
-        jnp.maximum(
+    f_new_dx = (
+        f_dx
+        * jnp.maximum(
             0.0,
             jnp.minimum(
                 1.0, ((f - jnp.roll(f, 1, axis=0)) / dx) / (f_dx + 1.0e-8 * (f_dx == 0))
             ),
         )
-        * f_dx
-    )
-    f_dx = (
-        jnp.maximum(
+        * jnp.maximum(
             0.0,
             jnp.minimum(
                 1.0,
                 (-(f - jnp.roll(f, -1, axis=0)) / dx) / (f_dx + 1.0e-8 * (f_dx == 0)),
             ),
         )
-        * f_dx
     )
-    f_dy = (
-        jnp.maximum(
+    f_new_dy = (
+        f_dy
+        * jnp.maximum(
             0.0,
             jnp.minimum(
                 1.0, ((f - jnp.roll(f, 1, axis=1)) / dx) / (f_dy + 1.0e-8 * (f_dy == 0))
             ),
         )
-        * f_dy
-    )
-    f_dy = (
-        jnp.maximum(
+        * jnp.maximum(
             0.0,
             jnp.minimum(
                 1.0,
                 (-(f - jnp.roll(f, -1, axis=1)) / dx) / (f_dy + 1.0e-8 * (f_dy == 0)),
             ),
         )
-        * f_dy
     )
 
-    return f_dx, f_dy
+    return f_new_dx, f_new_dy
 
 
 def extrapolate_to_face(f, f_dx, f_dy, dx):
@@ -110,9 +103,11 @@ def apply_fluxes(F, flux_F_X, flux_F_Y, dx, dt):
     """
     Apply fluxes to conserved variables
     """
-    F += -dt * dx * flux_F_X
-    F += dt * dx * jnp.roll(flux_F_X, 1, axis=0)
-    F += -dt * dx * flux_F_Y
-    F += dt * dx * jnp.roll(flux_F_Y, 1, axis=1)
+    F_new = F + (dt * dx) * (
+        -flux_F_X
+        + jnp.roll(flux_F_X, 1, axis=0)
+        - flux_F_Y
+        + jnp.roll(flux_F_Y, 1, axis=1)
+    )
 
-    return F
+    return F_new

adirondax/simulation.py

@@ -175,7 +175,7 @@ def _evolve(self, state):
         t_span = self.params["time"]["span"]
 
         use_adaptive_timesteps = True if nt < 1 else False
-        dt = 0.0 if use_adaptive_timesteps else t_span / nt
+        dt_ref = jnp.nan if use_adaptive_timesteps else t_span / nt
 
         # Physics flags
         use_hydro = self.params["physics"]["hydro"]
@@ -201,18 +201,19 @@ def _evolve(self, state):
             V = self._calc_grav_potential(state, k_sq, use_quantum, use_hydro)
 
         # Build the carry:
-        carry = (state, dt, V, k_sq)
+        carry = (state, V, k_sq)
 
         def step_fn(carry):
             """
             Pure step function: advances state by one timestep.
             """
-            state, dt, V, k_sq = carry
+            state, V, k_sq = carry
 
             # Create new state dict to avoid mutation
             new_state = {}
 
             # Get the timestep
+            dt = dt_ref
             if use_adaptive_timesteps:
                 dt = jnp.inf
                 if use_hydro:
@@ -307,22 +308,22 @@ def step_fn(carry):
             # Update diagnostics
             new_state["steps_taken"] = state["steps_taken"] + 1
 
-            return (new_state, dt, new_V, k_sq)
+            return (new_state, new_V, k_sq)
 
         # Run the entire loop as a single JIT-compiled function
         def run_loop(carry):
             if use_adaptive_timesteps:
                 # def cond_fn(carry):
-                #    state, _, _, _ = carry
+                #    state, _, _ = carry
                 #    return state["t"] < t_span * (1.0 - 1e-10)
 
                 # final_carry = jax.lax.while_loop(cond_fn, step_fn, carry)
 
                 # do a simple while loop
-                state, _, _, _ = carry
+                state, _, _ = carry
                 while state["t"] < t_span * (1.0 - 1e-10):
                     carry = step_fn(carry)
-                    state, _, _, _ = carry
+                    state, _, _ = carry
                 final_carry = carry
             else:
 
@@ -336,7 +337,7 @@ def step_fn_stacked(carry, _):
             return final_carry
 
         # Execute the compiled loop
-        state, _, _, _ = run_loop(carry)
+        state, _, _ = run_loop(carry)
 
         return state
 
@@ -347,4 +348,3 @@ def run(self):
         self.state["steps_taken"] = 0
         self.state = self._evolve(self.state)
         jax.block_until_ready(self.state)
-        # assert jnp.isfinite(self.state["t"]), "state['t'] is NaN/infinity"