Adaptive stepper

dedalus/core/solvers.py

@@ -709,6 +709,8 @@ def step(self, dt):
         # Update iteration
         self.iteration += 1
         self.dt = dt
+
+        return dt
 
     def evolve(self, timestep_function, log_cadence=100):
         """Advance system until stopping criterion is reached."""

dedalus/core/timesteppers.py

@@ -85,7 +85,7 @@ def step(self, dt, wall_time):
 
         # Solver references
         solver = self.solver
-        subproblems = [sp for sp in solver.subproblems if sp.size] # Skip empty subproblems
+        subproblems = solver.subproblems
         evaluator = solver.evaluator
         state_fields = solver.state
         F_fields = solver.F
@@ -542,7 +542,7 @@ def step(self, dt, wall_time):
 
         # Solver references
         solver = self.solver
-        subproblems = [sp for sp in solver.subproblems if sp.size] # Skip empty subproblems
+        subproblems = solver.subproblems
         evaluator = solver.evaluator
         state_fields = solver.state
         F_fields = solver.F
@@ -629,6 +629,7 @@ def step(self, dt, wall_time):
                 spRHS = RHS.get_subdata(sp)
                 spX = sp.LHS_solvers[i].solve(spRHS)  # CREATES TEMPORARY
                 sp.scatter_inputs(spX, state_fields)
+
             solver.sim_time = sim_time_0 + k*c[i]
 
 
@@ -673,6 +674,9 @@ class RK443(RungeKuttaIMEX):
 
     stages = 4
 
+    # A = explicit
+    # H = Implicit
+
     c = np.array([0, 1/2, 2/3, 1/2, 1])
 
     A = np.array([[  0  ,   0  ,  0 ,   0 , 0],
@@ -727,3 +731,282 @@ class RKGFY(RungeKuttaIMEX):
                   [0.5 , 0.5,   0],
                   [0.5 , 0  , 0.5]])
 
+class RungeKuttaIMEX_Adapt:
+    """
+    Base class for implicit-explicit multistep methods.
+
+    Parameters
+    ----------
+    nfields : int
+        Number of fields in problem
+    domain : domain object
+        Problem domain
+
+    Notes
+    -----
+    These timesteppers discretize the system
+        M.dt(X) + L.X = F
+    by constructing s stages
+        M.X(n,i) - M.X(n,0) + k Hij L.X(n,j) = k Aij F(n,j)
+    where j runs from {0, 0} to {i, i-1}, and F(n,i) is evaluated at time
+        t(n,i) = t(n,0) + k ci
+
+    The s stages are solved as
+        (M + k Hii L).X(n,i) = M.X(n,0) + k Aij F(n,j) - k Hij L.X(n,j)
+    where j runs from {0, 0} to {i-1, i-1}.
+
+    The final stage is used as the advanced solution*:
+        X(n+1,0) = X(n,s)
+        t(n+1,0) = t(n,s) = t(n,0) + k
+
+    * Equivalently the Butcher tableaus must follow
+        b_im = H[s, :]
+        b_ex = A[s, :]
+        c[s] = 1
+
+    References
+    ----------
+    U. M. Ascher, S. J. Ruuth, and R. J. Spiteri, Applied Numerical Mathematics (1997).
+
+    """
+
+    steps = 1
+
+    def __init__(self, solver):
+
+        self.solver = solver
+        self.RHS = CoeffSystem(solver.subproblems, dtype=solver.dtype)
+
+        # Create coefficient systems for multistep history
+        self.MX0 = CoeffSystem(solver.subproblems, dtype=solver.dtype)
+        self.LX = [CoeffSystem(solver.subproblems, dtype=solver.dtype) for i in range(self.stages)]
+        self.F = [CoeffSystem(solver.subproblems, dtype=solver.dtype) for i in range(self.stages)]
+
+        self._LHS_params = None
+        self.axpy = blas.get_blas_funcs('axpy', dtype=solver.dtype)
+
+    def step(self, dt, wall_time):
+        """Advance solver by one timestep."""
+
+        # Solver references
+        solver = self.solver
+        subproblems = solver.subproblems
+        evaluator = solver.evaluator
+        state_fields = solver.state
+        F_fields = solver.F
+        sim_time_0 = solver.sim_time
+        iteration = solver.iteration
+        STORE_EXPANDED_MATRICES = solver.store_expanded_matrices
+
+        error_tolerance = 1e-04
+        growth_factor = 0.9
+        max_dt = 1e-03
+        min_dt = 1e-14
+
+        # Other references
+        RHS = self.RHS
+        MX0 = self.MX0
+        LX = self.LX
+        LX0 = LX[0]
+        F = self.F
+        A = self.A
+        H = self.H
+        b_hat_EX = self.b_hat_EX
+        b_hat_IM = self.b_hat_IM
+        c = self.c
+        k = dt
+        axpy = self.axpy
+
+        # Check on updating LHS
+        update_LHS = (k != self._LHS_params)
+        self._LHS_params = k
+        if update_LHS:
+            # Remove old solver references
+            for sp in subproblems:
+                sp.LHS_solvers = [None] * (self.stages+1)
+
+        # Compute M.X(n,0) and L.X(n,0)
+        # Ensure coeff space before subsystem gathers
+        evaluator.require_coeff_space(state_fields)
+        for sp in subproblems:
+            spX = sp.gather_inputs(state_fields)
+            apply_sparse(sp.M_min, spX, axis=0, out=MX0.get_subdata(sp))
+            apply_sparse(sp.L_min, spX, axis=0, out=LX0.get_subdata(sp))
+
+        # Compute stages
+        # (M + k Hii L).X(n,i) = M.X(n,0) + k Aij F(n,j) - k Hij L.X(n,j)
+        for i in range(1, self.stages + 1):
+            # Compute L.X(n,i-1), already done for i=1
+            if i > 1:
+                LXi = LX[i-1]
+                # Ensure coeff space before subsystem gathers
+                evaluator.require_coeff_space(state_fields)
+                for sp in subproblems:
+                    spX = sp.gather_inputs(state_fields)
+                    apply_sparse(sp.L_min, spX, axis=0, out=LXi.get_subdata(sp))
+
+            # Compute F(n,i-1), only doing output on first evaluation
+            if i == 1:
+                evaluator.evaluate_scheduled(iteration=iteration, wall_time=wall_time, sim_time=solver.sim_time, timestep=dt)
+            else:
+                evaluator.evaluate_group('F')
+            Fi = F[i-1]
+            for sp in subproblems:
+                # F fields should be in coeff space from evaluator
+                sp.gather_outputs(F_fields, out=Fi.get_subdata(sp))
+
+            # Construct RHS(n,i)
+            if RHS.data.size:
+                np.copyto(RHS.data, MX0.data)
+                for j in range(0, i):
+                    # RHS.data += (k * A[i,j]) * F[j].data
+                    axpy(a=(k * A[i, j]), x=F[j].data, y=RHS.data)
+                    # RHS.data -= (k * H[i,j]) * LX[j].data
+                    axpy(a=-(k * H[i, j]), x=LX[j].data, y=RHS.data)
+
+            # Solve for stage (Compute X)
+            k_Hii = k * H[i, i]
+            # Ensure coeff space before subsystem scatters
+            for field in state_fields:
+                field.preset_layout('c')
+            for sp in subproblems:
+                # Construct LHS(n,i) using old H
+                if update_LHS:
+                    if STORE_EXPANDED_MATRICES:
+                        # sp.LHS.data[:] = sp.M_exp.data + k_Hii*sp.L_exp.data
+                        np.copyto(sp.LHS.data, sp.M_exp.data)
+                        axpy(a=k_Hii, x=sp.L_exp.data, y=sp.LHS.data)
+                    else:
+                        sp.LHS = (sp.M_min + k_Hii * sp.L_min)  # CREATES TEMPORARY
+                    sp.LHS_solvers[i] = solver.matsolver(sp.LHS, solver)
+                # Slice out valid subdata, skipping invalid components
+                spRHS = RHS.get_subdata(sp)
+                spX = sp.LHS_solvers[i].solve(spRHS)  # CREATES TEMPORARY
+                sp.scatter_inputs(spX, state_fields)
+
+            # If this is the last iteration, swap H's last row with b_hat and compute X_hat
+            if i == self.stages-1:
+                # Swap the last row of H with b_hat
+                H[-1, :] = b_hat_IM
+                # A[-1, :] = b_hat_EX <- not sure if this is needed
+
+                # Reconstruct RHS(n,i) with updated H
+                if RHS.data.size:
+                    np.copyto(RHS.data, MX0.data)
+                    for j in range(0, i):
+                        # RHS.data += (k * A[i,j]) * F[j].data
+                        axpy(a=(k * A[i, j]), x=F[j].data, y=RHS.data)
+                        # RHS.data -= (k * H[i,j]) * LX[j].data with updated H
+                        axpy(a=-(k * H[i, j]), x=LX[j].data, y=RHS.data)
+
+                # Solve again with updated H to compute X_hat
+                for sp in subproblems:
+                    k_Hii = k * H[i, i]
+                    if update_LHS:
+                        if STORE_EXPANDED_MATRICES:
+                            np.copyto(sp.LHS.data, sp.M_exp.data)
+                            axpy(a=k_Hii, x=sp.L_exp.data, y=sp.LHS.data)
+                        else:
+                            sp.LHS = (sp.M_min + k_Hii * sp.L_min)
+                        sp.LHS_solvers[i] = solver.matsolver(sp.LHS, solver)
+
+                    # Solve for X_hat using the updated H
+                    spX_hat = sp.LHS_solvers[i].solve(spRHS)
+                    #sp.scatter_inputs(spX_hat, state_fields)
+
+            solver.sim_time = sim_time_0 + k * c[i]
+
+        # Maximum error could be something else
+        spX_diff = np.max(np.abs(spX - spX_hat))  
+        #print(spX_diff)
+
+        # Evaluating the error and adjusting the step size
+        adapt_fac = 0.9
+        if spX_diff > error_tolerance:
+            dt = dt * adapt_fac
+        else:
+            dt = min(dt / adapt_fac, max_dt)
+        dt = max(min(dt,max_dt), min_dt)
+        return dt
+
+
+@add_scheme
+class ARK437L2SA(RungeKuttaIMEX_Adapt):
+    """4th-order 6-stage scheme from Higher-order additive Runge–Kutta schemes for ordinary differential equations, Christopher A. Kennedy, Mark H. Carpenter"""
+
+    stages = 6
+    γ = 1235/10000
+    c = np.array([0, 247/2000, 4276536705230/10142255878289, 67/200, 3/40, 7/10, 1.0])
+    #Explicit
+    A = np.array([[  0,  0, 0, 0, 0, 0, 0],
+              [247/1000, 0, 0, 0, 0, 0, 0],
+              [247/4000, 2694949928731/7487940209513, 0, 0, 0, 0, 0], 
+              [464650059369/8764239774964, 878889893998/2444806327765, -952945855348/12294611323341, 0, 0, 0, 0],
+              [476636172619/8159180917465, -1271469283451/7793814740893, -859560642026/4356155882851, 1723805262919/4571918432560, 0, 0, 0],
+              [6338158500785/11769362343261, -4970555480458/10924838743837, 3326578051521/2647936831840, -880713585975/1841400956686, -1428733748635/8843423958496, 0, 0],
+              [760814592956/3276306540349, 760814592956/3276306540349, -47223648122716/6934462133451, 71187472546993/9669769126921, -13330509492149/9695768672337, 11565764226357/8513123442827, 0]])
+
+    #Implicit
+    H = np.array([[0,  0, 0, 0, 0, 0, 0],
+                  [γ , γ, 0, 0, 0, 0, 0],
+                  [624185399699/4186980696204 , 624185399699/4186980696204, γ, 0, 0, 0, 0],
+                  [1258591069120/10082082980243, 1258591069120/10082082980243, -322722984531/8455138723562, γ, 0, 0, 0],
+                  [-436103496990/5971407786587, -436103496990/5971407786587, -2689175662187/11046760208243, 4431412449334/12995360898505, γ, 0, 0],
+                  [-2207373168298/14430576638973, -2207373168298/14430576638973, 242511121179/3358618340039, 3145666661981/7780404714551, 5882073923981/14490790706663, γ, 0],
+                  [0, 0, 9164257142617/17756377923965, -10812980402763/74029279521829, 1335994250573/5691609445217, 2273837961795/8368240463276, γ]])
+
+    b_hat_IM = np.array([[0, 0, 4469248916618/8635866897933, -621260224600/4094290005349, 696572312987/2942599194819, 1532940081127/5565293938103, 2441/20000]
+    ])
+    b_hat_EX = np.array([[0, 0, 4469248916618/8635866897933, -621260224600/4094290005349, 696572312987/2942599194819, 1532940081127/5565293938103, 2441/20000]
+    ])
+
+class ARK548L2SA(RungeKuttaIMEX_Adapt):
+    """5th-order 7-stage scheme from Higher-order additive Runge–Kutta schemes for ordinary differential equations, Christopher A. Kennedy, Mark H. Carpenter"""
+
+    stages = 7
+    γ = 2/9
+    c = np.array([0, 4/9, 6456083330201/8509243623797, 1632083962415/14158861528103, 6365430648612/17842476412687, 18/25, 191/200, 1.0])
+    #Explicit a^[E] <- for comparing with the paper referenced
+    A = np.array([[  0,  0, 0, 0, 0, 0, 0, 0],
+              [1/9, 0, 0, 0, 0, 0, 0, 0],
+              [1/9, 1183333538310/1827251437969, 0, 0, 0, 0, 0, 0], 
+              [895379019517/9750411845327, 477606656805/13473228687314, -112564739183/9373365219272, 0, 0, 0, 0, 0],
+              [-4458043123994/13015289567637, -2500665203865/9342069639922, 983347055801/8893519644487, 2185051477207/2551468980502, 0, 0, 0, 0],
+              [-167316361917/17121522574472, 1605541814917/7619724128744, 991021770328/13052792161721, 2342280609577/11279663441611, 3012424348531/12792462456678, 0, 0, 0],
+              [6680998715867/14310383562358, 5029118570809/3897454228471, 2415062538259/6382199904604, -3924368632305/6964820224454, -4331110370267/15021686902756, -3944303808049/11994238218192, 0, 0],
+              [2193717860234/3570523412979, 2193717860234/3570523412979, 5952760925747/18750164281544, -4412967128996/6196664114337, 4151782504231/36106512998704, 572599549169/6265429158920, -457874356192/11306498036315, 0]])
+
+    #Implicit a^[I] <- for comparing with the paper referenced
+    H = np.array([[0,  0, 0, 0, 0, 0, 0, 0],
+                  [γ , γ, 0, 0, 0, 0, 0, 0],
+                  [2366667076620/8822750406821 , 2366667076620/8822750406821, γ, 0, 0, 0, 0, 0],
+                  [-257962897183/4451812247028, -257962897183/4451812247028, 128530224461/14379561246022, γ, 0, 0, 0, 0],
+                  [-486229321650/11227943450093, -486229321650/11227943450093, -225633144460/6633558740617, 1741320951451/6824444397158, γ, 0, 0, 0],
+                  [621307788657/4714163060173, 621307788657/4714163060173, -125196015625/3866852212004, 940440206406/7593089888465, 961109811699/6734810228204, γ, 0, 0],
+                  [2036305566805/6583108094622, 2036305566805/6583108094622, -3039402635899/4450598839912, -1829510709469/31102090912115, -286320471013/6931253422520, 8651533662697/9642993110008, γ, 0],
+                  [0 , 0, 3517720773327/20256071687669, 4569610470461/17934693873752, 2819471173109/11655438449929, 3296210113763/10722700128969, -1142099968913/5710983926999, γ]])
+
+    b_hat_IM = np.array([[0, 0, 520639020421/8300446712847, 4550235134915/17827758688493, 1482366381361/6201654941325, 5551607622171/13911031047899, -5266607656330/36788968843917, 1074053359553/5740751784926]
+    ])
+    b_hat_EX = np.array([[0, 0, 520639020421/8300446712847, 4550235134915/17827758688493, 1482366381361/6201654941325, 5551607622171/13911031047899, -5266607656330/36788968843917, 1074053359553/5740751784926]
+    ])
+
+# Third order scheme
+@add_scheme
+class IMEXRKCB3f(RungeKuttaIMEX_Adapt):
+    """3rd order scheme from Low-storage implicit/explicit Runge–Kutta schemes for the simulation of stiff high-dimensional ODE systems, by Daniele Cavaglieri, Thomas Bewley"""
+    stages = 4
+    c = np.array([0, 49/50, 1/25, 1])
+
+    A = np.array([[0, 0, 0, 0],
+                  [49/50, 0, 0, 0],
+                  [13244205847/647648310246, 13419997131/686433909488, 0, 0],
+                  [-2179897048956/603118880443, 231677526244/1085522130027, 3007879347537/683461566472, 0]])
+
+    H = np.array([[0, 0, 0, 0],
+                  [49/100, 49/100, 0, 0],
+                  [-785157464198/1093480182337, -30736234873/978681420651, 983779726483/1246172347126, 0],
+                  [-2179897048956/603118880443, 99189146040/891495457793, 6064140186914/1415701440113, 146791865627/668377518349]])
+
+    b_hat_IM = np.array([[0, 337712514207/759004992869, 311412265155/608745789881, 52826596233/1214539205236]])
+    b_hat_EX = np.array([[0, 0, 25/48, 23/48]])

dedalus/tests/test_ivp.py

@@ -42,7 +42,7 @@ def test_heat_periodic(basis, N,  dtype, dealias, timestepper):
     dt = 1e-5
     iter = 20
     for i in range(iter):
-        solver.step(dt)
+        dt = solver.step(dt)
     # Check solution
     amp = 1 - np.exp(-solver.sim_time)
     u_true = amp * np.sin(x)