#368 modify iterate method to allow timestep control

demos/demo_2d_tracer.py

@@ -66,9 +66,11 @@
 # computation functionality is switched off and the simulation time is
 # split into 600 steps, giving a timestep close to the CFL limit. ::
 
+t_end = 2*pi
+timestep = pi/300.0
 options.tracer_timestepper_type = 'SSPRK33'
-options.timestep = pi/300.0
-options.simulation_end_time = 2*pi
+options.timestep = timestep
+options.simulation_end_time = t_end
 options.simulation_export_time = pi/15.0
 options.tracer_timestepper_options.use_automatic_timestep = False
 
@@ -87,6 +89,7 @@
 
 # The velocity field is set up using a simple analytic expression. ::
 
+x, y = SpatialCoordinate(mesh2d)
 uv_init = as_vector([0.5 - y, x - 0.5])
 
 # Now, we set up the cosine-bell--cone--slotted-cylinder initial condition. The
@@ -116,9 +119,15 @@
 q_init = Function(P1_2d).interpolate(1.0 + bell + cone + slot_cyl)
 solver_obj.assign_initial_conditions(uv=uv_init, tracer_2d=q_init)
 
-# Now we are in a position to run the time loop. ::
+# Now we are in a position to run the time loop. We can either automatically run
+# the time loop to the end time as in :doc:`2D channel example <demo_2d_channel.py>`
+# or call a time iterator to explicitly run the time loop, by issuing ::
 
-solver_obj.iterate()
+thetis_timestepper = solver_obj.create_iterator()
+t_Thetis = 0
+
+while t_Thetis < t_end - timestep:
+    t_Thetis = next(thetis_timestepper)
 
 # Finally, we display the normalised :math:`L^2` error, by comparing to the initial condition. ::
 

thetis/solver2d.py

@@ -942,20 +942,56 @@ def print_state(self, cputime, print_header=False):
         sys.stdout.flush()
 
     @PETSc.Log.EventDecorator("thetis.FlowSolver2d.iterate")
-    def iterate(self, update_forcings=None,
-                export_func=None):
+    def iterate(self, update_forcings=None, export_func=None):
         """
         Runs the simulation
 
+        Wrapper for function for `create_iterator` generator and automatically
+        iterates over the time loop until time ``options.simulation_end_time`` is reached.
+        Exports fields to disk on ``options.simulation_export_time`` intervals.
+
+        :kwarg update_forcings: User-defined function that takes simulation
+            time as an argument and updates time-dependent boundary conditions
+            (if any).
+        :kwarg export_func: User-defined function (with no arguments) that will
+            be called on every export.
+        """
+        for _ in self.create_iterator(update_forcings=update_forcings,
+                                      export_func=export_func):
+            pass
+
+    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_iterator")
+    def create_iterator(self, update_forcings=None, export_func=None):
+        """
+        Creates a generator to iterate through the simulation and return access
+        to time advancing function when time control is handled externally.
+
         Iterates over the time loop until time ``options.simulation_end_time`` is reached.
         Exports fields to disk on ``options.simulation_export_time`` intervals.
 
+        For example:
+
+        .. code-block:: python
+
+            for t in solver_obj.create_iterator():
+                # user code
+
+        or, to get per time-step control:
+
+        .. code-block:: python
+
+            thetis_timestepper = solver_obj.create_iterator()
+
+            while t_Thetis<t_end - timestep and .... :
+                t_Thetis = next(thetis_timestepper)
+
         :kwarg update_forcings: User-defined function that takes simulation
             time as an argument and updates time-dependent boundary conditions
             (if any).
         :kwarg export_func: User-defined function (with no arguments) that will
             be called on every export.
         """
+
         if not self._initialized:
             self.initialize()
 
@@ -1043,6 +1079,9 @@ def iterate(self, update_forcings=None,
         while self.simulation_time <= self.options.simulation_end_time - t_epsilon:
             self.timestepper.advance(self.simulation_time, update_forcings)
 
+            # returns internal simulation time
+            yield self.simulation_time
+
             # Move to next time step
             self.iteration += 1
             internal_iteration += 1