optimize the compressible sdc solver

pyro/compressible_sdc/simulation.py

@@ -3,15 +3,26 @@
 
 """
 
-
 from pyro import compressible_fv4
-from pyro.mesh import patch
+from pyro.mesh import fv, patch
 from pyro.util import msg
 
 
 class Simulation(compressible_fv4.Simulation):
     """Drive the 4th-order compressible solver with SDC time integration"""
 
+    def __init__(self, solver_name, problem_name, problem_func, rp, *,
+                 problem_finalize_func=None, problem_source_func=None,
+                 timers=None, data_class=fv.FV2d):
+
+        super().__init__(solver_name, problem_name, problem_func, rp,
+                         problem_finalize_func=problem_finalize_func,
+                         problem_source_func=problem_source_func,
+                         timers=timers, data_class=data_class)
+
+        self.n_nodes = 3  # Gauss-Lobotto temporal nodes
+        self.n_iter = 4   # 4 SDC iterations for 4th order
+
     def sdc_integral(self, m_start, m_end, As):
         """Compute the integral over the sources from m to m+1 with a
         Simpson's rule"""
@@ -58,38 +69,50 @@ def evolve(self):
         U_knew.append(patch.cell_center_data_clone(self.cc_data))
         U_knew.append(patch.cell_center_data_clone(self.cc_data))
 
-        # loop over iterations
-        for _ in range(1, 5):
+        # we need the advective term at all time nodes at the old
+        # iteration -- we'll compute this for the initial state
+        # now
+        A_kold = []
+        A_kold.append(self.substep(U_kold[0]))
+        A_kold.append(A_kold[-1].copy())
+        A_kold.append(A_kold[-1].copy())
+
+        A_knew = []
+        for adv in A_kold:
+            A_knew.append(adv.copy())
 
-            # we need the advective term at all time nodes at the old
-            # iteration -- we'll compute this now
-            A_kold = []
-            for m in range(3):
-                _tmp = self.substep(U_kold[m])
-                A_kold.append(_tmp)
+        # loop over iterations
+        for _ in range(self.n_iter):
 
             # loop over the time nodes and update
-            for m in range(2):
+            for m in range(self.n_nodes):
 
                 # update m to m+1 for knew
 
                 # compute A(U_m^{k+1})
-                A_knew = self.substep(U_knew[m])
+                # note: for m = 0, the advective term doesn't change
+                if m > 0:
+                    A_knew[m] = self.substep(U_knew[m])
+
+                # only do an update if we are not at the last node
+                if m < self.n_nodes-1:
 
-                # compute the integral over A at the old iteration
-                integral = self.sdc_integral(m, m+1, A_kold)
+                    # compute the integral over A at the old iteration
+                    integral = self.sdc_integral(m, m+1, A_kold)
 
-                # and the final update
-                for n in range(self.ivars.nvar):
-                    U_knew[m+1].data.v(n=n)[:, :] = U_knew[m].data.v(n=n) + \
-                        0.5*self.dt * (A_knew.v(n=n) - A_kold[m].v(n=n)) + integral.v(n=n)
+                    # and the final update
+                    for n in range(self.ivars.nvar):
+                        U_knew[m+1].data.v(n=n)[:, :] = U_knew[m].data.v(n=n) + \
+                            0.5*self.dt * (A_knew[m].v(n=n) - A_kold[m].v(n=n)) + integral.v(n=n)
 
-                # fill ghost cells
-                U_knew[m+1].fill_BC_all()
+                    # fill ghost cells
+                    U_knew[m+1].fill_BC_all()
 
             # store the current iteration as the old iteration
-            for m in range(1, 3):
+            # node m = 0 data doesn't change
+            for m in range(1, self.n_nodes):
                 U_kold[m].data[:, :, :] = U_knew[m].data[:, :, :]
+                A_kold[m][:, :, :] = A_knew[m][:, :, :]
 
         # store the new solution
         self.cc_data.data[:, :, :] = U_knew[-1].data[:, :, :]