Adding local stepping capability

mirgecom/simutil.py

@@ -96,20 +96,28 @@ def check_step(step, interval):
     return False
 
 
-def get_sim_timestep(discr, state, t, dt, cfl, t_final, constant_cfl=False):
+from mirgecom.viscous import get_viscous_timestep
+from grudge.op import nodal_min
+def get_sim_timestep(discr, state, t=0.0, dt=None, cfl=None, t_final=0.0,
+                            constant_cfl=False, local_dt=False):
     """Return the maximum stable timestep for a typical fluid simulation.
 
-    This routine returns *dt*, the users defined constant timestep, or *max_dt*, the
-    maximum domain-wide stability-limited timestep for a fluid simulation.
+    If local_dt == False, this routine returns *dt*, the users defined constant
+    timestep, or *max_dt*, the maximum domain-wide stability-limited timestep
+    for a fluid simulation.
+    If local_dt == True, this routine returns *dt* as a DOFArray, and time is
+    advanced locally such that the cells are not at the same time instant. This is
+    useful for steady state convergence.
 
     .. important::
         This routine calls the collective: :func:`~grudge.op.nodal_min` on the inside
         which makes it domain-wide regardless of parallel domain decomposition. Thus
         this routine must be called *collectively* (i.e. by all ranks).
 
-    Two modes are supported:
+    Three modes are supported:
         - Constant DT mode: returns the minimum of (t_final-t, dt)
         - Constant CFL mode: returns (cfl * max_dt)
+        - Local DT mode: returns local dt
 
     Parameters
     ----------
@@ -127,22 +135,26 @@ def get_sim_timestep(discr, state, t, dt, cfl, t_final, constant_cfl=False):
         The current CFL number
     constant_cfl: bool
         True if running constant CFL mode
+    local_dt: bool
+        True if running local DT mode. False by default.
 
     Returns
     -------
     float
         The maximum stable DT based on a viscous fluid.
     """
-    t_remaining = max(0, t_final - t)
-    mydt = dt
-    if constant_cfl:
-        from mirgecom.viscous import get_viscous_timestep
-        from grudge.op import nodal_min
-        mydt = state.array_context.to_numpy(
-            cfl * nodal_min(
-                discr, "vol",
-                get_viscous_timestep(discr=discr, state=state)))[()]
-    return min(t_remaining, mydt)
+    if local_dt:
+        my_local_dt = cfl * get_viscous_timestep(discr=discr, state=state)
+        return my_local_dt
+    else:
+        mydt = dt
+        t_remaining = max(0, t_final - t)
+        if constant_cfl:
+            mydt = state.array_context.to_numpy(
+                cfl * nodal_min(
+                    discr, "vol",
+                    get_viscous_timestep(discr=discr, state=state)))[()]
+        return min(t_remaining, mydt)
 
 
 def write_visfile(discr, io_fields, visualizer, vizname,

mirgecom/steppers.py

@@ -138,6 +138,82 @@ def _advance_state_stepper_func(rhs, timestepper, state, t_final, dt=0,
     return istep, t, state
 
 
+def _advance_locally_state_stepper_func(rhs, timestepper, state, t, dt, nsteps,
+                                istep=0, pre_step_callback=None,
+                                post_step_callback=None, force_eval=True):
+    """Advance state for specific number of iterations using local time stepping.
+       Not to be used in time accurate simulations, only for convergence towards
+       steady state regime.
+
+    Parameters
+    ----------
+    rhs
+        Function that should return the time derivative of the state.
+        This function should take time and state as arguments, with
+        a call with signature ``rhs(t, state)``.
+    timestepper
+        Function that advances the state from t=time to t=(time+dt), and
+        returns the advanced state. Has a call with signature
+        ``timestepper(state, t, dt, rhs)``.
+    state: numpy.ndarray
+        Agglomerated object array containing at least the state variables that
+        will be advanced by this stepper
+    t: numpy.ndarray
+        Time at which to start. For compatibility of arguments, it is the same
+        as the step number.
+    dt: numpy.ndarray
+        Initial timestep size to use. Each cell has its own timestep and the
+        solution is advanced locally.
+    nsteps:
+        Number of iterations to be performed.
+    istep: int
+        Step number from which to start
+    pre_step_callback
+        An optional user-defined function, with signature:
+        ``state, dt = pre_step_callback(step, t, dt, state)``,
+        to be called before the timestepper is called for that particular step.
+    post_step_callback
+        An optional user-defined function, with signature:
+        ``state, dt = post_step_callback(step, t, dt, state)``,
+        to be called after the timestepper is called for that particular step.
+    force_eval
+        An optional boolean indicating whether to force lazy evaluation between
+        timesteps. Defaults to True.
+
+    Returns
+    -------
+    istep: int
+        the current step number
+    t: float
+        the current time
+    state: numpy.ndarray
+    """
+    if nsteps <= istep:
+        return istep, t, state
+
+    actx = get_container_context_recursively(state)
+
+    compiled_rhs = _compile_rhs(actx, rhs)
+
+    while istep < nsteps:
+
+        if pre_step_callback is not None:
+            state, dt = pre_step_callback(state=state, step=istep, t=t, dt=dt)
+
+        if force_eval:
+            state = _force_evaluation(actx, state)
+
+        state = timestepper(state=state, t=t, dt=dt, rhs=compiled_rhs)
+
+        t += 1.0
+        istep += 1
+
+        if post_step_callback is not None:
+            state, dt = post_step_callback(state=state, step=istep, t=t, dt=dt)
+
+    return istep, t, state
+
+
 def _advance_state_leap(rhs, timestepper, state, t_final, dt=0,
                         component_id="state", t=0.0, istep=0,
                         pre_step_callback=None, post_step_callback=None,
@@ -260,7 +336,7 @@ def generate_singlerate_leap_advancer(timestepper, component_id, rhs, t, dt,
 
 def advance_state(rhs, timestepper, state, t_final, t=0, istep=0, dt=0,
                   component_id="state", pre_step_callback=None,
-                  post_step_callback=None, force_eval=True):
+                  post_step_callback=None, force_eval=True, local_dt=False):
     """Determine what stepper we're using and advance the state from (t) to (t_final).
 
     Parameters
@@ -334,13 +410,23 @@ def advance_state(rhs, timestepper, state, t_final, t=0, istep=0, dt=0,
                 force_eval=force_eval
             )
     else:
-        (current_step, current_t, current_state) = \
-            _advance_state_stepper_func(
-                rhs=rhs, timestepper=timestepper,
-                state=state, t=t, t_final=t_final, dt=dt,
-                pre_step_callback=pre_step_callback,
-                post_step_callback=post_step_callback,
-                istep=istep, force_eval=force_eval
-            )
+        if local_dt:
+            (current_step, current_t, current_state) = \
+                _advance_locally_state_stepper_func(
+                    rhs=rhs, timestepper=timestepper,
+                    state=state, t=t, dt=dt, nsteps=nsteps,
+                    pre_step_callback=pre_step_callback,
+                    post_step_callback=post_step_callback,
+                    istep=istep, force_eval=force_eval
+                )
+        else:
+            (current_step, current_t, current_state) = \
+                _advance_state_stepper_func(
+                    rhs=rhs, timestepper=timestepper,
+                    state=state, t=t, t_final=t_final, dt=dt,
+                    pre_step_callback=pre_step_callback,
+                    post_step_callback=post_step_callback,
+                    istep=istep, force_eval=force_eval
+                )
 
     return current_step, current_t, current_state