misc: PR comments part 1

Author: JDBetteridge

devito/timestepping/superstep.py

@@ -1,17 +1,18 @@
 from devito.types import Eq, Function, TimeFunction
 
 
-def superstep_generator_iterative(field, stencil, k, tn=0):
-    ''' Generate superstep iteratively:
+def superstep_generator_iterative(field, stencil, k, nt=0):
+    """
+    Generate superstep iteratively:
     Aʲ⁺¹ = A·Aʲ
-    '''
+    """
     # New fields, for vector formulation both current and previous timestep are needed
     name = field.name
     grid = field.grid
     u = TimeFunction(name=f'{name}_ss', grid=grid, time_order=2, space_order=2*k)
     u_prev = TimeFunction(name=f'{name}_ss_p', grid=grid, time_order=2, space_order=2*k)
 
-    superstep_solution_transfer(field, u, u_prev, tn)
+    superstep_solution_transfer(field, u, u_prev, nt)
 
     # Substitute new fields into stencil
     ss_stencil = stencil.subs({field: u, field.backward: u_prev}, postprocess=False)
@@ -40,25 +41,27 @@ def superstep_generator_iterative(field, stencil, k, tn=0):
     return u, u_prev, Eq(u.forward, stencil_next), Eq(u_prev.forward, current)
 
 
-def superstep_generator(field, stencil, k, tn=0):
-    ''' Generate superstep using a binary decomposition:
+def superstep_generator(field, stencil, k, nt=0):
+    """
+    Generate superstep using a binary decomposition:
     A^k = aⱼ A^2ʲ × ... × a₂ A^2² × a₁ A² × a₀ A
     where k = aⱼ·2ʲ + ... + a₂·2² + a₁·2¹ + a₀·2⁰
-    '''
+    """
     # New fields, for vector formulation both current and previous timestep are needed
     name = field.name
-    grid = field.grid
     # time_order of `field` needs to be 2
-    u = TimeFunction(name=f'{name}_ss', grid=grid, time_order=1, space_order=2*k)
-    u_prev = TimeFunction(name=f'{name}_ss_p', grid=grid, time_order=1, space_order=2*k)
+    u = field._rebuild(name=f'{name}_ss', time_order=1, space_order=2*k)
+    u_prev = field._rebuild(name=f'{name}_ss', time_order=1, space_order=2*k)
 
-    superstep_solution_transfer(field, u, u_prev, tn)
+    superstep_solution_transfer(field, u, u_prev, nt)
 
     # Substitute new fields into stencil
     ss_stencil = stencil.subs({field: u, field.backward: u_prev}, postprocess=False)
     ss_stencil = ss_stencil.expand().expand(add=True, nest=True)
 
-    # Binary decomposition algorithm
+    # Binary decomposition algorithm (see docstring):
+    # Calculate the binary decomposition of the exponent (k) and accumulate the
+    # resultant operator
     current = (ss_stencil, u)
     q, r = divmod(k, 2)
     accumulate = current if r else (1, 1)
@@ -71,23 +74,26 @@ def superstep_generator(field, stencil, k, tn=0):
     return u, u_prev, Eq(u.forward, accumulate[0]), Eq(u_prev.forward, accumulate[1])
 
 
-def superstep_solution_transfer(old, new, new_p, tn):
-    ''' Transfer state from a previous TimeFunction to a 2 field superstep
+def superstep_solution_transfer(old, new, new_p, nt):
+    """
+    Transfer state from a previous TimeFunction to a 2 field superstep
     Used after injecting source using standard timestepping.
-    '''
-    # 3 should be replaced with `old.time_order + 1` although this needs some thought
-    idx = tn % 3 if old.save is None else -1
-    new.data[0, :] = old.data[idx - 1]
-    new.data[1, :] = old.data[idx]
-    new_p.data[0, :] = old.data[idx - 2]
-    new_p.data[1, :] = old.data[idx - 1]
+    """
+    # This method is completely generic for future development, but currently
+    # only time_order == 2 is implemented!
+    idx = nt % (old.time_order + 1) if old.save is None else -1
+    for ii in range(old.time_order + 1):
+        new.data[ii, :] = old.data[idx - ii - 1]
+        new_p.data[ii, :] = old.data[idx - ii - 2]
 
 
 def _combine_superstep(stencil_a, stencil_b, u, u_prev, k):
-    ''' Combine two arbitrary order supersteps
-    '''
+    """
+    Combine two arbitrary order supersteps
+    """
     # Placeholder fields for forming the superstep
     grid = u.grid
+    # Can I use a TempFunction here?
     a_tmp = Function(name="a_tmp", grid=grid, space_order=2*k)
     b_tmp = Function(name="b_tmp", grid=grid, space_order=2*k)
 

examples/timestepping/acoustic_superstep_2d.py

@@ -10,6 +10,8 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from scipy.interpolate import interpn
+
 from devito import (
     ConditionalDimension,
     Eq,
@@ -21,7 +23,6 @@
     solve,
 )
 from devito.timestepping.superstep import superstep_generator
-from scipy.interpolate import interpn
 
 
 @dataclass
@@ -99,24 +100,24 @@ def acoustic_model(model, step=1, snapshots=1):
     pde = (1/velocity**2)*u.dt2 - u.laplace
     stencil = Eq(u.forward, solve(pde, u.forward))
 
-    tn1 = int(np.ceil((t1 - t0)/model.critical_dt))
-    dt = (t1 - t0)/tn1
+    nt1 = int(np.ceil((t1 - t0)/model.critical_dt))
+    dt = (t1 - t0)/nt1
 
     # Source
-    t = np.linspace(t0, t1, tn1)
+    t = np.linspace(t0, t1, nt1)
     rick = ricker(t)
     source = SparseTimeFunction(
-        name="ricker", npoint=1, coordinates=[model.source], nt=tn1, grid=grid,
+        name="ricker", npoint=1, coordinates=[model.source], nt=nt1, grid=grid,
         time_order=2, space_order=4
     )
     source.data[:, 0] = rick
     src_term = source.inject(field=u.forward, expr=source*velocity*velocity*dt*dt)
 
     op1 = Operator([stencil] + src_term)
-    op1(time=tn1 - 1, dt=dt)
+    op1(time=nt1 - 1, dt=dt)
 
     # Stencil and operator
-    idx = tn1 % 3
+    idx = nt1 % 3
     if step == 1:
         # Non-superstep case
         new_u = TimeFunction(name="new_u", grid=grid, time_order=2, space_order=2)
@@ -127,13 +128,13 @@ def acoustic_model(model, step=1, snapshots=1):
         new_u.data[1, :] = u.data[idx - 1]
         new_u.data[2, :] = u.data[idx]
     else:
-        new_u, new_u_p, *stencil = superstep_generator(u, stencil.rhs, step, tn=tn1)
+        new_u, new_u_p, *stencil = superstep_generator(u, stencil.rhs, step, nt=nt1)
 
-    tn2 = int(np.ceil((t2 - t1)/model.critical_dt))
-    dt = (t2 - t1)/tn2
+    nt2 = int(np.ceil((t2 - t1)/model.critical_dt))
+    dt = (t2 - t1)/nt2
 
     # Snapshot the solution
-    factor = int(np.ceil(tn2/(snapshots + 1)))
+    factor = int(np.ceil(nt2/(snapshots + 1)))
     t_sub = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
     u_save = TimeFunction(
         name='usave', grid=grid,

examples/timestepping/superstep_1d.py

@@ -3,6 +3,7 @@
 '''
 import matplotlib.pyplot as plt
 import numpy as np
+
 from devito import Eq, Function, Grid, Operator, TimeFunction, solve
 from devito.timestepping.superstep import superstep_generator
 
@@ -64,12 +65,12 @@ def wave_on_string(step=1):
             stencil2,
         ], opt='noop')
 
-    tn = int(np.ceil(t1/critical_dt))
-    dt = t1/tn
+    nt = int(np.ceil(t1/critical_dt))
+    dt = t1/nt
 
-    op(time=tn, dt=dt)
+    op(time=nt, dt=dt)
 
-    idx = tn % 3
+    idx = nt % 3
     return newu.data[idx]
 
 

examples/timestepping/superstep_2d.py

@@ -3,6 +3,7 @@
 '''
 import matplotlib.pyplot as plt
 import numpy as np
+
 from devito import ConditionalDimension, Eq, Function, Grid, Operator, TimeFunction, solve
 from devito.timestepping.superstep import superstep_generator
 
@@ -60,11 +61,11 @@ def ripple_on_pond(step=1, snapshots=1):
         new_u_p.data[0, :] = ic
         new_u_p.data[1, :] = ic
 
-    tn = int(np.ceil((t1 - t0)/critical_dt))
-    dt = t1/tn
+    nt = int(np.ceil((t1 - t0)/critical_dt))
+    dt = t1/nt
 
     # Snapshot the solution
-    factor = int(np.ceil(tn/(snapshots + 1)))
+    factor = int(np.ceil(nt/(snapshots + 1)))
     t_sub = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
     u_save = TimeFunction(
         name='usave', grid=grid,
@@ -103,7 +104,7 @@ def ripple_on_pond(step=1, snapshots=1):
                 )
                 idx += 1
                 if step == 1:
-                    ax.set_title(f't = {(ii*t1)/(m - 1) :0.3f}')
+                    ax.set_title(f't = {(ii*t1)/(m - 1):0.3f}')
                     if ii != 0:
                         ax.set_yticklabels([])
                     if ii % 2 == 1: