update

Author: fernanvr

devito/types/multistage.py

@@ -339,7 +339,6 @@ def source_inclusion(self, u, k, src_index, src_deriv, e_p, t, dt, mu, n_eq):
 
         return src_lhs, e_p
 
-
     def _evaluate(self, **kwargs):
         """
         Generate the stage-wise equations for a Runge-Kutta time integration method.
@@ -356,7 +355,7 @@ def _evaluate(self, **kwargs):
             - 1 final update equation of the form `u.forward = u + dt * sum(b_i * k_i)`
         """
 
-        n_eq=len(self.eq)
+        n_eq = len(self.eq)
         u = [i.function for i in self.lhs]
         t = u[0].grid.time_dim
         dt = t.spacing
@@ -386,19 +385,26 @@ def _evaluate(self, **kwargs):
 
         # Build each stage
         for i in range(1, self.deg-1):
-            [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
-            src_lhs, e_p = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
-            [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+            # [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
+            # src_lhs, e_p = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+            # [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+            # [stage_eqs.append(Eq(u[j].forward, u[j].forward+k[j]*alpha[i])) for j in range(n_eq)]
+            src_lhs, e_p = self.source_inclusion(u, k, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+            [stage_eqs.append(Eq(k[j], k[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
             [stage_eqs.append(Eq(u[j].forward, u[j].forward+k[j]*alpha[i])) for j in range(n_eq)]
 
         # Final Runge-Kutta updates
-        [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
-        src_lhs, e_p = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
-        [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
-
-        [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
-        src_lhs, _ = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
-        [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+        # [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
+        # src_lhs, e_p = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+        # [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+        src_lhs, e_p = self.source_inclusion(u, k, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+        [stage_eqs.append(Eq(k[j], k[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+
+        # [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
+        # src_lhs, _ = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+        # [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+        src_lhs, _ = self.source_inclusion(u, k, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+        [stage_eqs.append(Eq(k[j], k[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
 
         # Compute final approximation
         [stage_eqs.append(Eq(u[j].forward, u[j].forward+k[j]*alpha[self.deg-1])) for j in range(n_eq)]

tests/test_multistage.py

@@ -91,7 +91,7 @@ def test_multistage_coupled_op_computing_exp(time_int='HORK_EXP'):
     src_spatial.data[100, 100] = 1
     import sympy as sym
     src_temporal = sym.exp(- 100 * (t - 0.01)**2)
-    # import matplotlib.pyplot as plt
+    import matplotlib.pyplot as plt
     # import numpy as np
     # t=np.linspace(0,2000,1000)
     # plt.plot(t,np.exp(1 - 2 * (t - 1)**2))
@@ -107,13 +107,13 @@ def test_multistage_coupled_op_computing_exp(time_int='HORK_EXP'):
            [src_spatial, src_temporal, U_multi_stage[1]]]
 
     # Time integration scheme
-    pdes = resolve_method(time_int)(U_multi_stage, system_eqs_rhs, source=src, degree=4)
+    pdes = resolve_method(time_int)(U_multi_stage, system_eqs_rhs, source=src, degree=3)
     op = Operator(pdes, subs=grid.spacing_map)
     op(dt=0.001, time=2000)
 
-    # plt.imshow(U_multi_stage[0].data[0,:])
-    # plt.colorbar()
-    # plt.show()
+    plt.imshow(U_multi_stage[0].data[0,:])
+    plt.colorbar()
+    plt.show()
 
 
 def test_multistage_object(time_int='RK44'):