add adaptive timesteps

Author: pmocz

adirondax/hydro/euler2d.py

@@ -61,15 +61,21 @@ def get_flux(rho_L, rho_R, vx_L, vx_R, vy_L, vy_R, P_L, P_R, gamma):
     return flux_Mass, flux_Momx, flux_Momy, flux_Energy
 
 
+def hydro_euler2d_timestep(rho, vx, vy, P, gamma, dx):
+    """Calculate the simulation timestep based on CFL condition"""
+
+    # get time step (CFL) = dx / max signal speed
+    dt = jnp.min(dx / (jnp.sqrt(gamma * P / rho) + jnp.sqrt(vx**2 + vy**2)))
+
+    return dt
+
+
 def hydro_euler2d_fluxes(rho, vx, vy, P, gamma, dx, dt):
     """Take a simulation timestep"""
 
     # get Conserved variables
     Mass, Momx, Momy, Energy = get_conserved(rho, vx, vy, P, gamma, dx**2)
 
-    # get time step (CFL) = dx / max signal speed
-    # dt = courant_fac * jnp.min(dx / (jnp.sqrt(gamma * P / rho) + jnp.sqrt(vx**2 + vy**2)))
-
     # calculate gradients
     rho_dx, rho_dy = get_gradient(rho, dx)
     vx_dx, vx_dy = get_gradient(vx, dx)

adirondax/hydro/mhd2d.py

@@ -421,16 +421,26 @@ def get_flux_hlld(
     return flux_Mass, flux_Momx, flux_Momy, flux_Energy, flux_By
 
 
+def hydro_mhd2d_timestep(rho, vx, vy, P, bx, by, gamma, dx):
+    """Calculate the simulation timestep based on CFL condition"""
+
+    # get time step (CFL) = dx / max signal speed
+    Bx, By = get_avg(bx, by)
+    dt = jnp.min(
+        dx
+        / (jnp.sqrt(gamma * P / rho) + jnp.sqrt(vx**2 + vy**2 + (Bx**2 + By**2) / rho))
+    )
+
+    return dt
+
+
 def hydro_mhd2d_fluxes(rho, vx, vy, P, bx, by, gamma, dx, dt):
     """Take a simulation timestep"""
 
     # get Conserved variables
     Bx, By = get_avg(bx, by)
     Mass, Momx, Momy, Energy = get_conserved(rho, vx, vy, P, Bx, By, gamma, dx**2)
 
-    # get time step (CFL) = dx / max signal speed
-    # dt = courant_fac * jnp.min(dx / (jnp.sqrt(gamma * P / rho) + jnp.sqrt(vx**2 + vy**2)))
-
     # calculate gradients
     rho_dx, rho_dy = get_gradient(rho, dx)
     vx_dx, vx_dy = get_gradient(vx, dx)

adirondax/params_default.json

@@ -85,6 +85,14 @@
         "default": "./eos_table.h5",
         "description": "path to tabular EOS data file for 'tabular' gas."
       }
+    },
+    "riemann_solver": {
+      "default": "llf",
+      "description": "options: 'llf', 'hlld'."
+    },
+    "cfl": {
+      "default": 0.5,
+      "description": "CFL number for hydrodynamics timestep."
     }
   },
   "quantum": {

adirondax/quantum.py

@@ -13,3 +13,7 @@ def quantum_drift(psi, k_sq, m_per_hbar, dt):
     psi_hat = jnp.exp(dt * (-1.0j * k_sq / m_per_hbar / 2.0)) * psi_hat
     psi = jnp.fft.ifftn(psi_hat)
     return psi
+
+
+def quantum_timestep(m_per_hbar, dx):
+    return (m_per_hbar / 6.0) * (dx * dx)

adirondax/simulation.py

@@ -2,9 +2,9 @@
 import jax.numpy as jnp
 
 from .constants import constants
-from .hydro.euler2d import hydro_euler2d_fluxes
-from .hydro.mhd2d import hydro_mhd2d_fluxes
-from .quantum import quantum_kick, quantum_drift
+from .hydro.euler2d import hydro_euler2d_fluxes, hydro_euler2d_timestep
+from .hydro.mhd2d import hydro_mhd2d_fluxes, hydro_mhd2d_timestep
+from .quantum import quantum_kick, quantum_drift, quantum_timestep
 from .gravity import calculate_gravitational_potential
 from .utils import set_up_parameters, print_parameters
 
@@ -37,7 +37,7 @@ def __init__(self, params):
 
         # simulation state
         self.state = {}
-        self.state["t"] = jnp.array(0) + jnp.nan
+        self.state["t"] = jnp.array(0.0) + jnp.nan
         if self.params["physics"]["hydro"]:
             self.state["rho"] = jnp.zeros(self.resolution) + jnp.nan
             self.state["vx"] = jnp.zeros(self.resolution) + jnp.nan
@@ -51,6 +51,9 @@ def __init__(self, params):
                 jnp.zeros(self.resolution, dtype=jnp.complex64) + jnp.nan
             )
 
+        # extra info to keep track of
+        self.state["steps_taken"] = jnp.array(0) + jnp.nan
+
     @property
     def resolution(self):
         """
@@ -72,6 +75,13 @@ def dim(self):
         """
         return len(self.resolution)
 
+    @property
+    def steps_taken(self):
+        """
+        Return the number of steps taken in the simulation
+        """
+        return self.state["steps_taken"]
+
     @property
     def params(self):
         """
@@ -155,11 +165,8 @@ def _evolve(self, state):
         nt = self.params["time"]["num_timesteps"]
         t_span = self.params["time"]["span"]
 
-        fixed_timestepping = True if nt > 0 else False
-        if fixed_timestepping:
-            dt = t_span / nt
-
-        assert fixed_timestepping  # XXX for now
+        use_adaptive_timesteps = True if nt < 1 else False
+        dt = 0.0 if use_adaptive_timesteps else t_span / nt
 
         # Physics flags
         use_hydro = self.params["physics"]["hydro"]
@@ -168,6 +175,9 @@ def _evolve(self, state):
         use_gravity = self.params["physics"]["gravity"]
 
         gamma = self.params["hydro"]["eos"]["gamma"]
+        cfl = self.params["hydro"]["cfl"]
+
+        m_per_hbar = 1.0  # XXX
 
         # Precompute Fourier space variables
         k_sq = None
@@ -180,27 +190,56 @@ def _evolve(self, state):
         if use_gravity:
             V = self._calc_grav_potential(state, k_sq, use_quantum, use_hydro)
 
-        # Build the carry: (state, V, k_sq)
-        carry = (state, V, k_sq)
+        # Build the carry:
+        carry = (state, dt, V, k_sq)
 
-        def step_fn(carry, _):
+        def step_fn(carry):
             """
             Pure step function: advances state by one timestep.
-            Returns new carry and None (no stacked outputs).
             """
-            state, V, k_sq = carry
+            state, dt, V, k_sq = carry
 
             # Create new state dict to avoid mutation
             new_state = {}
 
+            # Get the timestep
+            if use_adaptive_timesteps:
+                dt = jnp.inf
+                if use_hydro:
+                    if use_magnetic:
+                        dt_hydro = hydro_mhd2d_timestep(
+                            state["rho"],
+                            state["vx"],
+                            state["vy"],
+                            state["P"],
+                            state["bx"],
+                            state["by"],
+                            gamma,
+                            dx,
+                        )
+                    else:
+                        dt_hydro = hydro_euler2d_timestep(
+                            state["rho"],
+                            state["vx"],
+                            state["vy"],
+                            state["P"],
+                            gamma,
+                            dx,
+                        )
+                    dt = jnp.minimum(dt, cfl * dt_hydro)
+                if use_quantum:
+                    dt_quantum = quantum_timestep(m_per_hbar, dx)
+                    dt = jnp.minimum(dt, dt_quantum)
+                dt = jnp.minimum(dt, t_span - state["t"])
+
             # Kick (half-step) - quantum + gravity
             psi = state.get("psi")
             if use_quantum and use_gravity and psi is not None:
-                psi = quantum_kick(psi, V, 1.0, dt / 2.0)
+                psi = quantum_kick(psi, V, m_per_hbar, dt / 2.0)
 
             # Drift (full-step) - quantum
             if use_quantum and psi is not None:
-                psi = quantum_drift(psi, k_sq, 1.0, dt)
+                psi = quantum_drift(psi, k_sq, m_per_hbar, dt)
 
             if use_quantum:
                 new_state["psi"] = psi
@@ -254,21 +293,47 @@ def step_fn(carry, _):
             # Update time
             new_state["t"] = state["t"] + dt
 
-            return (new_state, new_V, k_sq), None
+            # Update diagnostics
+            new_state["steps_taken"] = state["steps_taken"] + 1
+
+            return (new_state, dt, new_V, k_sq)
 
         # Run the entire loop as a single JIT-compiled function
         def run_loop(carry):
-            final_carry, _ = jax.lax.scan(step_fn, carry, xs=None, length=nt)
+            if use_adaptive_timesteps:
+                # def cond_fn(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
+                while state["t"] < t_span * (1.0 - 1e-10):
+                    carry = step_fn(carry)
+                    state, _, _, _ = carry
+                final_carry = carry
+            else:
+
+                def step_fn_stacked(carry, _):
+                    # Returns new carry and None (no stacked outputs) for jax.lax.scan()
+                    return step_fn(carry), None
+
+                final_carry, _ = jax.lax.scan(
+                    step_fn_stacked, carry, xs=None, length=nt
+                )
             return final_carry
 
         # Execute the compiled loop
-        state, _, _ = run_loop(carry)
+        state, _, _, _ = run_loop(carry)
 
         return state
 
     def run(self):
         """
         Run the simulation
         """
+        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"

examples/kelvin_helmholtz_instability/output.png

@@ -23,7 +23,6 @@
 def set_up_simulation():
     # Define the parameters for the simulation
     n = 512
-    nt = 100 * int(n / 32)
     t_stop = 0.5
     gamma = 5.0 / 3.0
     box_size = 1.0
@@ -41,10 +40,10 @@ def set_up_simulation():
         },
         "time": {
             "span": t_stop,
-            "num_timesteps": nt,
         },
         "hydro": {
             "eos": {"type": "ideal", "gamma": gamma},
+            "cfl": 0.6,
         },
     }
 
@@ -90,7 +89,8 @@ def main():
     # Evolve the system
     t0 = time.time()
     sim.run()
-    print("Run time (s): ", time.time() - t0)
+    print("Steps taken:", sim.steps_taken)
+    print("Run time (s):", time.time() - t0)
 
     make_plot(sim)