From c1a1854b6b947879869d3cf66cf9766bc190082e Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Thu, 3 Oct 2024 16:01:49 +0100
Subject: [PATCH 01/11] Adding the embedded RK scheme from Kennedy and
 Carpenter

---
 dedalus/core/timesteppers.py | 234 ++++++++++++++++++++++++++++++++++-
 1 file changed, 232 insertions(+), 2 deletions(-)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index 4937c810..e3568593 100644
--- a/dedalus/core/timesteppers.py
+++ b/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,229 @@ 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 = self.b_hat
+        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
+
+                # 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:
+                    # Recompute LHS with updated H
+                    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]
+
+        spX_diff = np.max(np.abs(spX - spX_hat))  # Element-wise difference
+        #print(spX_diff)
+
+        adapt_fac = 0.9
+        if spX_diff > error_tolerance:
+            # If the error exceeds the tolerance, decrease the time step proportionally
+            dt = dt * adapt_fac
+        else:
+            # If the error is within tolerance, increase the time step but cap it at max_dt
+            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 ..."""
+
+    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],  # Added comma here
+              [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]])  # Fixed the list
+
+    #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 
+    
+    b_hat = np.array([[0, 0, 4469248916618/8635866897933, -621260224600/4094290005349, 696572312987/2942599194819, 1532940081127/5565293938103, 2441/20000]
+    ])
+

From b12cb59397a69c94dec0a157f9d862d10a142b5f Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Thu, 3 Oct 2024 16:06:10 +0100
Subject: [PATCH 02/11] Update to the reference

---
 dedalus/core/timesteppers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index e3568593..8a6547d6 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -931,7 +931,7 @@ def step(self, dt, wall_time):
 
 @add_scheme
 class ARK437L2SA(RungeKuttaIMEX_Adapt):
-    """4th-order 6-stage scheme from ..."""
+    """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

From 463e41522ae76a20e68e755ecb446f2561c8bd57 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Thu, 3 Oct 2024 16:37:24 +0100
Subject: [PATCH 03/11] Removing redundant comments

---
 dedalus/core/timesteppers.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index 8a6547d6..693c5768 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -899,7 +899,6 @@ def step(self, dt, wall_time):
 
                 # Solve again with updated H to compute X_hat
                 for sp in subproblems:
-                    # Recompute LHS with updated H
                     k_Hii = k * H[i, i]
                     if update_LHS:
                         if STORE_EXPANDED_MATRICES:
@@ -920,10 +919,8 @@ def step(self, dt, wall_time):
 
         adapt_fac = 0.9
         if spX_diff > error_tolerance:
-            # If the error exceeds the tolerance, decrease the time step proportionally
             dt = dt * adapt_fac
         else:
-            # If the error is within tolerance, increase the time step but cap it at max_dt
             dt = min(dt / adapt_fac, max_dt)
         dt = max(min(dt,max_dt), min_dt)
         return dt

From f760cd260a70f7dd17b4130ed78596972f1b1780 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Thu, 3 Oct 2024 20:22:02 +0100
Subject: [PATCH 04/11] Removing further redundant comments

---
 dedalus/core/timesteppers.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index 693c5768..b893d186 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -936,11 +936,11 @@ class ARK437L2SA(RungeKuttaIMEX_Adapt):
     #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],  # Added comma here
+              [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]])  # Fixed the list
+              [760814592956/3276306540349, 760814592956/3276306540349, -47223648122716/6934462133451, 71187472546993/9669769126921, -13330509492149/9695768672337, 11565764226357/8513123442827, 0]])
 
     #Implicit
     H = np.array([[0,  0, 0, 0, 0, 0, 0],
@@ -949,7 +949,7 @@ class ARK437L2SA(RungeKuttaIMEX_Adapt):
                   [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 
+                  [0, 0, 9164257142617/17756377923965, -10812980402763/74029279521829, 1335994250573/5691609445217, 2273837961795/8368240463276, γ]])
     
     b_hat = np.array([[0, 0, 4469248916618/8635866897933, -621260224600/4094290005349, 696572312987/2942599194819, 1532940081127/5565293938103, 2441/20000]
     ])

From 00ae284059bb3f22dee2e29abab3a9c72af0edf5 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 08:50:57 +0100
Subject: [PATCH 05/11] Addition to the functions so dt can be varried

---
 dedalus/core/solvers.py | 121 +++++++++++-----------------------------
 1 file changed, 32 insertions(+), 89 deletions(-)

diff --git a/dedalus/core/solvers.py b/dedalus/core/solvers.py
index ddd27155..06981f02 100644
--- a/dedalus/core/solvers.py
+++ b/dedalus/core/solvers.py
@@ -6,11 +6,7 @@
 import h5py
 import pathlib
 import scipy.linalg
-import cProfile
-import pstats
 from math import prod
-from collections import defaultdict
-import pickle
 
 from . import subsystems
 from . import timesteppers
@@ -18,11 +14,6 @@
 from ..libraries.matsolvers import matsolvers
 from ..tools.config import config
 from ..tools.array import scipy_sparse_eigs
-from ..tools.parallel import ProfileWrapper, parallel_mkdir
-
-PROFILE_DEFAULT = config['profiling'].getboolean('PROFILE_DEFAULT')
-PARALLEL_PROFILE_DEFAULT = config['profiling'].getboolean('PARALLEL_PROFILE_DEFAULT')
-PROFILE_DIRECTORY = pathlib.Path(config['profiling'].get('PROFILE_DIRECTORY'))
 
 import logging
 logger = logging.getLogger(__name__.split('.')[-1])
@@ -172,8 +163,8 @@ def print_subproblem_ranks(self, subproblems=None, target=0):
         for i, sp in enumerate(subproblems):
             if not hasattr(sp, 'L_min'):
                 continue
-            L = sp.L_min.toarray()
-            M = sp.M_min.toarray()
+            L = sp.L_min.A
+            M = sp.M_min.A
             A = L + target * M
             print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(A)}/{A.shape[0]}, cond: {np.linalg.cond(A):.1e}")
 
@@ -205,15 +196,15 @@ def solve_dense(self, subproblem, rebuild_matrices=False, left=False, normalize_
         if rebuild_matrices or not hasattr(sp, 'L_min'):
             self.build_matrices([sp], ['M', 'L'])
         # Solve as dense general eigenvalue problem
-        A = sp.L_min.toarray()
-        B = - sp.M_min.toarray()
+        A = sp.L_min.A
+        B = - sp.M_min.A
         eig_output = scipy.linalg.eig(A, b=B, left=left, **kw)
         # Unpack output
         if left:
             self.eigenvalues, pre_left_evecs, pre_right_evecs = eig_output
             self.right_eigenvectors = self.eigenvectors = sp.pre_right @ pre_right_evecs
-            self.left_eigenvectors = sp.pre_left.conj().T @ pre_left_evecs
-            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).conj().T @ pre_left_evecs
+            self.left_eigenvectors = sp.pre_left.H @ pre_left_evecs
+            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).H @ pre_left_evecs
             if normalize_left:
                 norms = np.diag(pre_left_evecs.T.conj() @ sp.M_min @ pre_right_evecs)
                 self.left_eigenvectors /= np.conj(norms)
@@ -222,7 +213,7 @@ def solve_dense(self, subproblem, rebuild_matrices=False, left=False, normalize_
             self.eigenvalues, pre_eigenvectors = eig_output
             self.eigenvectors = sp.pre_right @ pre_eigenvectors
 
-    def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False, normalize_left=True, raise_on_mismatch=True, v0=None, **kw):
+    def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False, normalize_left=True, raise_on_mismatch=True, **kw):
         """
         Perform targeted sparse eigenvector search for selected subproblem.
         This routine finds a subset of eigenvectors near the specified target.
@@ -249,8 +240,6 @@ def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False
             eigenvectors (default: True).
         raise_on_mismatch : bool, optional
             Raise a RuntimeError if the left and right eigenvalues do not match (default: True).
-        v0 : ndarray, optional
-            Initial guess for eigenvector, e.g. from subsystem.gather (default: None).
         **kw :
             Other keyword options passed to scipy.sparse.linalg.eig.
         """
@@ -261,17 +250,15 @@ def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False
         # Solve as sparse general eigenvalue problem
         A = sp.L_min
         B = - sp.M_min
-        # Precondition starting guess if provided
-        if v0 is not None:
-            v0 = sp.pre_right_pinv @ v0
         # Solve for the right (and optionally left) eigenvectors
-        eig_output = scipy_sparse_eigs(A=A, B=B, left=left, N=N, target=target, matsolver=self.matsolver, v0=v0, **kw)
+        eig_output = scipy_sparse_eigs(A=A, B=B, left=left, N=N, target=target, matsolver=self.matsolver, **kw)
+
         if left:
             # Note: this definition of "left eigenvectors" is consistent with the documentation for scipy.linalg.eig
             self.eigenvalues, pre_right_evecs, self.left_eigenvalues, pre_left_evecs = eig_output
             self.right_eigenvectors = self.eigenvectors = sp.pre_right @ pre_right_evecs
-            self.left_eigenvectors = sp.pre_left.conj().T @ pre_left_evecs
-            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).conj().T @ pre_left_evecs
+            self.left_eigenvectors = sp.pre_left.H @ pre_left_evecs
+            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).H @ pre_left_evecs
             # Check that eigenvalues match
             if not np.allclose(self.eigenvalues, np.conjugate(self.left_eigenvalues)):
                 if raise_on_mismatch:
@@ -293,7 +280,7 @@ def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False
             self.eigenvalues, pre_right_evecs = eig_output
             self.right_eigenvectors = self.eigenvectors = sp.pre_right @ pre_right_evecs
 
-    def set_state(self, index, subsystem=0):
+    def set_state(self, index, subsystem):
         """
         Set state vector to the specified eigenmode.
 
@@ -301,10 +288,10 @@ def set_state(self, index, subsystem=0):
         ----------
         index : int
             Index of desired eigenmode.
-        subsystem : Subsystem object or int, optional
+        subsystem : Subsystem object or int
             Subsystem that will be set to the corresponding eigenmode.
             If an integer, the corresponding subsystem of the last specified
-            eigenvalue_subproblem will be used. Default: 0.
+            eigenvalue_subproblem will be used.
         """
         # TODO: allow setting left modified eigenvectors?
         subproblem = self.eigenvalue_subproblem
@@ -317,8 +304,6 @@ def set_state(self, index, subsystem=0):
         for var in self.state:
             var['c'] = 0
         subsystem.scatter(self.eigenvectors[:, index], self.state)
-        # Set eigenvalue
-        self.problem.eigenvalue['g'] = self.eigenvalues[index]
 
 
 class LinearBoundaryValueSolver(SolverBase):
@@ -363,7 +348,7 @@ def print_subproblem_ranks(self, subproblems=None):
         for i, sp in enumerate(subproblems):
             if not hasattr(sp, 'L_min'):
                 continue
-            L = sp.L_min.toarray()
+            L = sp.L_min.A
             print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(L)}/{L.shape[0]}, cond: {np.linalg.cond(L):.1e}")
 
     def solve(self, subproblems=None, rebuild_matrices=False):
@@ -444,9 +429,6 @@ def __init__(self, problem, **kw):
         logger.debug('Beginning NLBVP instantiation')
         super().__init__(problem, **kw)
         self.perturbations = problem.perturbations
-        # Copy valid modes from variables to perturbations (may have been changed after problem instantiation)
-        for pert, var in zip(problem.perturbations, problem.variables):
-            pert.valid_modes[:] = var.valid_modes
         self.iteration = 0
         # Create RHS handler
         F_handler = self.evaluator.add_system_handler(iter=1, group='F')
@@ -456,17 +438,6 @@ def __init__(self, problem, **kw):
         self.F = F_handler.fields
         logger.debug('Finished NLBVP instantiation')
 
-    def print_subproblem_ranks(self, subproblems=None):
-        """Print rank of each subproblem LHS."""
-        if subproblems is None:
-            subproblems = self.subproblems
-        # Check matrix rank
-        for i, sp in enumerate(subproblems):
-            if not hasattr(sp, 'dF_min'):
-                continue
-            dF = sp.dF_min.toarray()
-            print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(dF)}/{dF.shape[0]}, cond: {np.linalg.cond(dF):.1e}")
-
     def newton_iteration(self, damping=1):
         """Update solution with a Newton iteration."""
         # Compute RHS
@@ -514,10 +485,6 @@ class InitialValueSolver(SolverBase):
         Iteration cadence for enforcing Hermitian symmetry on real variables (default: 100).
     warmup_iterations : int, optional
         Number of warmup iterations to disregard when computing runtime statistics (default: 10).
-    profile : bool, optional
-        Save accumulated profiles with cProfile (default: False).
-    parallel_profile : bool, optional
-        Save per-process and accumulated profiles with cProfile (default: False).
     **kw :
         Other options passed to ProblemBase.
 
@@ -543,24 +510,15 @@ class InitialValueSolver(SolverBase):
     matsolver_default = 'MATRIX_FACTORIZER'
     matrices = ['M', 'L']
 
-    def __init__(self, problem, timestepper, enforce_real_cadence=100, warmup_iterations=10, profile=PROFILE_DEFAULT, parallel_profile=PARALLEL_PROFILE_DEFAULT, **kw):
+    def __init__(self, problem, timestepper, enforce_real_cadence=100, warmup_iterations=10, **kw):
         logger.debug('Beginning IVP instantiation')
-        # Setup timing and profiling
-        self.dist = problem.dist
+        super().__init__(problem, **kw)
+        if np.isrealobj(self.dtype.type()):
+            self.enforce_real_cadence = enforce_real_cadence
+        else:
+            self.enforce_real_cadence = None
         self._bcast_array = np.zeros(1, dtype=float)
         self.init_time = self.world_time
-        if profile or parallel_profile:
-            parallel_mkdir(PROFILE_DIRECTORY, comm=self.dist.comm)
-            self.profile = True
-            self.parallel_profile = parallel_profile
-            self.setup_profiler = cProfile.Profile()
-            self.warmup_profiler = cProfile.Profile()
-            self.run_profiler = cProfile.Profile()
-            self.setup_profiler.enable()
-        else:
-            self.profile = False
-        # Build subsystems and subproblems
-        super().__init__(problem, **kw)
         # Build LHS matrices
         self.build_matrices(self.subproblems, ['M', 'L'])
         # Compute total modes
@@ -580,10 +538,6 @@ def __init__(self, problem, timestepper, enforce_real_cadence=100, warmup_iterat
         self.sim_time = self.initial_sim_time = problem.time.allreduce_data_max(layout='g')
         self.iteration = self.initial_iteration = 0
         self.warmup_iterations = warmup_iterations
-        if np.isrealobj(self.dtype.type()):
-            self.enforce_real_cadence = enforce_real_cadence
-        else:
-            self.enforce_real_cadence = None
         # Default integration parameters
         self.stop_sim_time = np.inf
         self.stop_wall_time = np.inf
@@ -693,22 +647,17 @@ def step(self, dt):
         # Record times
         wall_time = self.wall_time
         if self.iteration == self.initial_iteration:
-            self.start_time_end = wall_time
-            if self.profile:
-                self.dump_profiles(self.setup_profiler, "setup")
-                self.warmup_profiler.enable()
-            self.warmup_time_start = self.wall_time
+            self.start_time = wall_time
         if self.iteration == self.initial_iteration + self.warmup_iterations:
-            self.warmup_time_end = self.wall_time
-            if self.profile:
-                self.dump_profiles(self.warmup_profiler, "warmup")
-                self.run_profiler.enable()
-            self.run_time_start = self.wall_time
+            self.warmup_time = wall_time
         # Advance using timestepper
-        self.timestepper.step(dt, wall_time)
+        dt = self.timestepper.step(dt, wall_time)
         # 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."""
@@ -739,7 +688,7 @@ def print_subproblem_ranks(self, subproblems=None, dt=1):
         for i, sp in enumerate(subproblems):
             M = sp.M_min
             L = sp.L_min
-            A = (M + dt*L).toarray()
+            A = (M + dt*L).A
             print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(A)}/{A.shape[0]}, cond: {np.linalg.cond(A):.1e}")
 
     def evaluate_handlers_now(self, dt, handlers=None):
@@ -754,23 +703,17 @@ def evaluate_handlers(self, handlers=None, dt=0):
 
     def log_stats(self, format=".4g"):
         """Log timing statistics with specified string formatting (optional)."""
-        self.run_time_end = self.wall_time
-        start_time = self.start_time_end
+        log_time = self.wall_time
         logger.info(f"Final iteration: {self.iteration}")
         logger.info(f"Final sim time: {self.sim_time}")
-        logger.info(f"Setup time (init - iter 0): {start_time:{format}} sec")
-        if self.profile:
-            self.dump_profiles(self.run_profiler, "runtime")
+        logger.info(f"Setup time (init - iter 0): {self.start_time:{format}} sec")
         if self.iteration >= self.initial_iteration + self.warmup_iterations:
-            warmup_time = self.warmup_time_end - self.warmup_time_start
-            run_time = self.run_time_end - self.run_time_start
-            profile_output_time = (self.warmup_time_start-self.start_time_end) + (self.run_time_start-self.warmup_time_end)
+            warmup_time = self.warmup_time - self.start_time
+            run_time = log_time - self.warmup_time
             cpus = self.dist.comm.size
             modes = self.total_modes
             stages = (self.iteration - self.warmup_iterations - self.initial_iteration) * self.timestepper.stages
             logger.info(f"Warmup time (iter 0-{self.warmup_iterations}): {warmup_time:{format}} sec")
-            if self.profile:
-                logger.info(f"Profile output time: {profile_output_time:{format}} sec")
             logger.info(f"Run time (iter {self.warmup_iterations}-end): {run_time:{format}} sec")
             logger.info(f"CPU time (iter {self.warmup_iterations}-end): {run_time*cpus/3600:{format}} cpu-hr")
             logger.info(f"Speed: {(modes*stages/cpus/run_time):{format}} mode-stages/cpu-sec")

From 3309437711331d5fe8768d0642c8df034b22422b Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 08:57:45 +0100
Subject: [PATCH 06/11] Changes to get the calculated dt from step(...)

---
 dedalus/core/solvers.py | 119 ++++++++++++++++++++++++++++++----------
 1 file changed, 89 insertions(+), 30 deletions(-)

diff --git a/dedalus/core/solvers.py b/dedalus/core/solvers.py
index 06981f02..dff6363f 100644
--- a/dedalus/core/solvers.py
+++ b/dedalus/core/solvers.py
@@ -6,7 +6,11 @@
 import h5py
 import pathlib
 import scipy.linalg
+import cProfile
+import pstats
 from math import prod
+from collections import defaultdict
+import pickle
 
 from . import subsystems
 from . import timesteppers
@@ -14,6 +18,11 @@
 from ..libraries.matsolvers import matsolvers
 from ..tools.config import config
 from ..tools.array import scipy_sparse_eigs
+from ..tools.parallel import ProfileWrapper, parallel_mkdir
+
+PROFILE_DEFAULT = config['profiling'].getboolean('PROFILE_DEFAULT')
+PARALLEL_PROFILE_DEFAULT = config['profiling'].getboolean('PARALLEL_PROFILE_DEFAULT')
+PROFILE_DIRECTORY = pathlib.Path(config['profiling'].get('PROFILE_DIRECTORY'))
 
 import logging
 logger = logging.getLogger(__name__.split('.')[-1])
@@ -163,8 +172,8 @@ def print_subproblem_ranks(self, subproblems=None, target=0):
         for i, sp in enumerate(subproblems):
             if not hasattr(sp, 'L_min'):
                 continue
-            L = sp.L_min.A
-            M = sp.M_min.A
+            L = sp.L_min.toarray()
+            M = sp.M_min.toarray()
             A = L + target * M
             print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(A)}/{A.shape[0]}, cond: {np.linalg.cond(A):.1e}")
 
@@ -196,15 +205,15 @@ def solve_dense(self, subproblem, rebuild_matrices=False, left=False, normalize_
         if rebuild_matrices or not hasattr(sp, 'L_min'):
             self.build_matrices([sp], ['M', 'L'])
         # Solve as dense general eigenvalue problem
-        A = sp.L_min.A
-        B = - sp.M_min.A
+        A = sp.L_min.toarray()
+        B = - sp.M_min.toarray()
         eig_output = scipy.linalg.eig(A, b=B, left=left, **kw)
         # Unpack output
         if left:
             self.eigenvalues, pre_left_evecs, pre_right_evecs = eig_output
             self.right_eigenvectors = self.eigenvectors = sp.pre_right @ pre_right_evecs
-            self.left_eigenvectors = sp.pre_left.H @ pre_left_evecs
-            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).H @ pre_left_evecs
+            self.left_eigenvectors = sp.pre_left.conj().T @ pre_left_evecs
+            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).conj().T @ pre_left_evecs
             if normalize_left:
                 norms = np.diag(pre_left_evecs.T.conj() @ sp.M_min @ pre_right_evecs)
                 self.left_eigenvectors /= np.conj(norms)
@@ -213,7 +222,7 @@ def solve_dense(self, subproblem, rebuild_matrices=False, left=False, normalize_
             self.eigenvalues, pre_eigenvectors = eig_output
             self.eigenvectors = sp.pre_right @ pre_eigenvectors
 
-    def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False, normalize_left=True, raise_on_mismatch=True, **kw):
+    def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False, normalize_left=True, raise_on_mismatch=True, v0=None, **kw):
         """
         Perform targeted sparse eigenvector search for selected subproblem.
         This routine finds a subset of eigenvectors near the specified target.
@@ -240,6 +249,8 @@ def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False
             eigenvectors (default: True).
         raise_on_mismatch : bool, optional
             Raise a RuntimeError if the left and right eigenvalues do not match (default: True).
+        v0 : ndarray, optional
+            Initial guess for eigenvector, e.g. from subsystem.gather (default: None).
         **kw :
             Other keyword options passed to scipy.sparse.linalg.eig.
         """
@@ -250,15 +261,17 @@ def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False
         # Solve as sparse general eigenvalue problem
         A = sp.L_min
         B = - sp.M_min
+        # Precondition starting guess if provided
+        if v0 is not None:
+            v0 = sp.pre_right_pinv @ v0
         # Solve for the right (and optionally left) eigenvectors
-        eig_output = scipy_sparse_eigs(A=A, B=B, left=left, N=N, target=target, matsolver=self.matsolver, **kw)
-
+        eig_output = scipy_sparse_eigs(A=A, B=B, left=left, N=N, target=target, matsolver=self.matsolver, v0=v0, **kw)
         if left:
             # Note: this definition of "left eigenvectors" is consistent with the documentation for scipy.linalg.eig
             self.eigenvalues, pre_right_evecs, self.left_eigenvalues, pre_left_evecs = eig_output
             self.right_eigenvectors = self.eigenvectors = sp.pre_right @ pre_right_evecs
-            self.left_eigenvectors = sp.pre_left.H @ pre_left_evecs
-            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).H @ pre_left_evecs
+            self.left_eigenvectors = sp.pre_left.conj().T @ pre_left_evecs
+            self.modified_left_eigenvectors = (sp.M_min @ sp.pre_right_pinv).conj().T @ pre_left_evecs
             # Check that eigenvalues match
             if not np.allclose(self.eigenvalues, np.conjugate(self.left_eigenvalues)):
                 if raise_on_mismatch:
@@ -280,7 +293,7 @@ def solve_sparse(self, subproblem, N, target, rebuild_matrices=False, left=False
             self.eigenvalues, pre_right_evecs = eig_output
             self.right_eigenvectors = self.eigenvectors = sp.pre_right @ pre_right_evecs
 
-    def set_state(self, index, subsystem):
+    def set_state(self, index, subsystem=0):
         """
         Set state vector to the specified eigenmode.
 
@@ -288,10 +301,10 @@ def set_state(self, index, subsystem):
         ----------
         index : int
             Index of desired eigenmode.
-        subsystem : Subsystem object or int
+        subsystem : Subsystem object or int, optional
             Subsystem that will be set to the corresponding eigenmode.
             If an integer, the corresponding subsystem of the last specified
-            eigenvalue_subproblem will be used.
+            eigenvalue_subproblem will be used. Default: 0.
         """
         # TODO: allow setting left modified eigenvectors?
         subproblem = self.eigenvalue_subproblem
@@ -304,6 +317,8 @@ def set_state(self, index, subsystem):
         for var in self.state:
             var['c'] = 0
         subsystem.scatter(self.eigenvectors[:, index], self.state)
+        # Set eigenvalue
+        self.problem.eigenvalue['g'] = self.eigenvalues[index]
 
 
 class LinearBoundaryValueSolver(SolverBase):
@@ -348,7 +363,7 @@ def print_subproblem_ranks(self, subproblems=None):
         for i, sp in enumerate(subproblems):
             if not hasattr(sp, 'L_min'):
                 continue
-            L = sp.L_min.A
+            L = sp.L_min.toarray()
             print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(L)}/{L.shape[0]}, cond: {np.linalg.cond(L):.1e}")
 
     def solve(self, subproblems=None, rebuild_matrices=False):
@@ -429,6 +444,9 @@ def __init__(self, problem, **kw):
         logger.debug('Beginning NLBVP instantiation')
         super().__init__(problem, **kw)
         self.perturbations = problem.perturbations
+        # Copy valid modes from variables to perturbations (may have been changed after problem instantiation)
+        for pert, var in zip(problem.perturbations, problem.variables):
+            pert.valid_modes[:] = var.valid_modes
         self.iteration = 0
         # Create RHS handler
         F_handler = self.evaluator.add_system_handler(iter=1, group='F')
@@ -438,6 +456,17 @@ def __init__(self, problem, **kw):
         self.F = F_handler.fields
         logger.debug('Finished NLBVP instantiation')
 
+    def print_subproblem_ranks(self, subproblems=None):
+        """Print rank of each subproblem LHS."""
+        if subproblems is None:
+            subproblems = self.subproblems
+        # Check matrix rank
+        for i, sp in enumerate(subproblems):
+            if not hasattr(sp, 'dF_min'):
+                continue
+            dF = sp.dF_min.toarray()
+            print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(dF)}/{dF.shape[0]}, cond: {np.linalg.cond(dF):.1e}")
+
     def newton_iteration(self, damping=1):
         """Update solution with a Newton iteration."""
         # Compute RHS
@@ -485,6 +514,10 @@ class InitialValueSolver(SolverBase):
         Iteration cadence for enforcing Hermitian symmetry on real variables (default: 100).
     warmup_iterations : int, optional
         Number of warmup iterations to disregard when computing runtime statistics (default: 10).
+    profile : bool, optional
+        Save accumulated profiles with cProfile (default: False).
+    parallel_profile : bool, optional
+        Save per-process and accumulated profiles with cProfile (default: False).
     **kw :
         Other options passed to ProblemBase.
 
@@ -510,15 +543,24 @@ class InitialValueSolver(SolverBase):
     matsolver_default = 'MATRIX_FACTORIZER'
     matrices = ['M', 'L']
 
-    def __init__(self, problem, timestepper, enforce_real_cadence=100, warmup_iterations=10, **kw):
+    def __init__(self, problem, timestepper, enforce_real_cadence=100, warmup_iterations=10, profile=PROFILE_DEFAULT, parallel_profile=PARALLEL_PROFILE_DEFAULT, **kw):
         logger.debug('Beginning IVP instantiation')
-        super().__init__(problem, **kw)
-        if np.isrealobj(self.dtype.type()):
-            self.enforce_real_cadence = enforce_real_cadence
-        else:
-            self.enforce_real_cadence = None
+        # Setup timing and profiling
+        self.dist = problem.dist
         self._bcast_array = np.zeros(1, dtype=float)
         self.init_time = self.world_time
+        if profile or parallel_profile:
+            parallel_mkdir(PROFILE_DIRECTORY, comm=self.dist.comm)
+            self.profile = True
+            self.parallel_profile = parallel_profile
+            self.setup_profiler = cProfile.Profile()
+            self.warmup_profiler = cProfile.Profile()
+            self.run_profiler = cProfile.Profile()
+            self.setup_profiler.enable()
+        else:
+            self.profile = False
+        # Build subsystems and subproblems
+        super().__init__(problem, **kw)
         # Build LHS matrices
         self.build_matrices(self.subproblems, ['M', 'L'])
         # Compute total modes
@@ -538,6 +580,10 @@ def __init__(self, problem, timestepper, enforce_real_cadence=100, warmup_iterat
         self.sim_time = self.initial_sim_time = problem.time.allreduce_data_max(layout='g')
         self.iteration = self.initial_iteration = 0
         self.warmup_iterations = warmup_iterations
+        if np.isrealobj(self.dtype.type()):
+            self.enforce_real_cadence = enforce_real_cadence
+        else:
+            self.enforce_real_cadence = None
         # Default integration parameters
         self.stop_sim_time = np.inf
         self.stop_wall_time = np.inf
@@ -647,17 +693,24 @@ def step(self, dt):
         # Record times
         wall_time = self.wall_time
         if self.iteration == self.initial_iteration:
-            self.start_time = wall_time
+            self.start_time_end = wall_time
+            if self.profile:
+                self.dump_profiles(self.setup_profiler, "setup")
+                self.warmup_profiler.enable()
+            self.warmup_time_start = self.wall_time
         if self.iteration == self.initial_iteration + self.warmup_iterations:
-            self.warmup_time = wall_time
+            self.warmup_time_end = self.wall_time
+            if self.profile:
+                self.dump_profiles(self.warmup_profiler, "warmup")
+                self.run_profiler.enable()
+            self.run_time_start = self.wall_time
         # Advance using timestepper
-        dt = self.timestepper.step(dt, wall_time)
+        self.timestepper.step(dt, wall_time)
         # 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."""
@@ -688,7 +741,7 @@ def print_subproblem_ranks(self, subproblems=None, dt=1):
         for i, sp in enumerate(subproblems):
             M = sp.M_min
             L = sp.L_min
-            A = (M + dt*L).A
+            A = (M + dt*L).toarray()
             print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(A)}/{A.shape[0]}, cond: {np.linalg.cond(A):.1e}")
 
     def evaluate_handlers_now(self, dt, handlers=None):
@@ -703,17 +756,23 @@ def evaluate_handlers(self, handlers=None, dt=0):
 
     def log_stats(self, format=".4g"):
         """Log timing statistics with specified string formatting (optional)."""
-        log_time = self.wall_time
+        self.run_time_end = self.wall_time
+        start_time = self.start_time_end
         logger.info(f"Final iteration: {self.iteration}")
         logger.info(f"Final sim time: {self.sim_time}")
-        logger.info(f"Setup time (init - iter 0): {self.start_time:{format}} sec")
+        logger.info(f"Setup time (init - iter 0): {start_time:{format}} sec")
+        if self.profile:
+            self.dump_profiles(self.run_profiler, "runtime")
         if self.iteration >= self.initial_iteration + self.warmup_iterations:
-            warmup_time = self.warmup_time - self.start_time
-            run_time = log_time - self.warmup_time
+            warmup_time = self.warmup_time_end - self.warmup_time_start
+            run_time = self.run_time_end - self.run_time_start
+            profile_output_time = (self.warmup_time_start-self.start_time_end) + (self.run_time_start-self.warmup_time_end)
             cpus = self.dist.comm.size
             modes = self.total_modes
             stages = (self.iteration - self.warmup_iterations - self.initial_iteration) * self.timestepper.stages
             logger.info(f"Warmup time (iter 0-{self.warmup_iterations}): {warmup_time:{format}} sec")
+            if self.profile:
+                logger.info(f"Profile output time: {profile_output_time:{format}} sec")
             logger.info(f"Run time (iter {self.warmup_iterations}-end): {run_time:{format}} sec")
             logger.info(f"CPU time (iter {self.warmup_iterations}-end): {run_time*cpus/3600:{format}} cpu-hr")
             logger.info(f"Speed: {(modes*stages/cpus/run_time):{format}} mode-stages/cpu-sec")

From ad15039266f6b37ed5ac47f5314267c538b995d6 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 09:01:12 +0100
Subject: [PATCH 07/11] Adjust for test to allow for adaptivity

---
 dedalus/tests/test_ivp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dedalus/tests/test_ivp.py b/dedalus/tests/test_ivp.py
index 1c225629..e53e7f5d 100644
--- a/dedalus/tests/test_ivp.py
+++ b/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)

From 15e3118d9bd3d5aa7d725623c0f0a9417fa98bf1 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 09:06:15 +0100
Subject: [PATCH 08/11] Additional comments

---
 dedalus/core/timesteppers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index b893d186..80fb6865 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -914,9 +914,11 @@ def step(self, dt, wall_time):
 
             solver.sim_time = sim_time_0 + k * c[i]
 
-        spX_diff = np.max(np.abs(spX - spX_hat))  # Element-wise difference
+        # 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

From 0e397c8af459ac3b00dc758e905cd0c24c3243e4 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 10:52:10 +0100
Subject: [PATCH 09/11] Additional scheme from Kennedy and Carpenter paper

---
 dedalus/core/timesteppers.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index 80fb6865..ce73a8e8 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -955,4 +955,33 @@ class ARK437L2SA(RungeKuttaIMEX_Adapt):
     
     b_hat = 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 = np.array([[0, 0, 520639020421/8300446712847, 4550235134915/17827758688493, 1482366381361/6201654941325, 5551607622171/13911031047899, -5266607656330/36788968843917, 1074053359553/5740751784926]
+    ])
 

From 8c291c8c8cca720c5ad11c248be656e1bfeba806 Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 14:06:54 +0100
Subject: [PATCH 10/11] Addition the final row in A swapped out

---
 dedalus/core/timesteppers.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index ce73a8e8..32e6271a 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -887,6 +887,7 @@ def step(self, dt, wall_time):
             if i == self.stages-1:
                 # Swap the last row of H with b_hat
                 H[-1, :] = b_hat
+                A[-1, :] = b_hat
 
                 # Reconstruct RHS(n,i) with updated H
                 if RHS.data.size:

From eb6e4c01809eaf56f8a1820dc1e3293a7c50997e Mon Sep 17 00:00:00 2001
From: danvau98 <vaughandan180@gmail.com>
Date: Fri, 4 Oct 2024 17:33:36 +0100
Subject: [PATCH 11/11] Addition of a 3rd order scheme

---
 dedalus/core/timesteppers.py | 34 +++++++++++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/dedalus/core/timesteppers.py b/dedalus/core/timesteppers.py
index 32e6271a..54b6a27a 100644
--- a/dedalus/core/timesteppers.py
+++ b/dedalus/core/timesteppers.py
@@ -811,7 +811,8 @@ def step(self, dt, wall_time):
         F = self.F
         A = self.A
         H = self.H
-        b_hat = self.b_hat
+        b_hat_EX = self.b_hat_EX
+        b_hat_IM = self.b_hat_IM
         c = self.c
         k = dt
         axpy = self.axpy
@@ -886,8 +887,8 @@ def step(self, dt, wall_time):
             # 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
-                A[-1, :] = 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:
@@ -954,7 +955,9 @@ class ARK437L2SA(RungeKuttaIMEX_Adapt):
                   [-2207373168298/14430576638973, -2207373168298/14430576638973, 242511121179/3358618340039, 3145666661981/7780404714551, 5882073923981/14490790706663, γ, 0],
                   [0, 0, 9164257142617/17756377923965, -10812980402763/74029279521829, 1335994250573/5691609445217, 2273837961795/8368240463276, γ]])
     
-    b_hat = np.array([[0, 0, 4469248916618/8635866897933, -621260224600/4094290005349, 696572312987/2942599194819, 1532940081127/5565293938103, 2441/20000]
+    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):
@@ -983,6 +986,27 @@ class ARK548L2SA(RungeKuttaIMEX_Adapt):
                   [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 = np.array([[0, 0, 520639020421/8300446712847, 4550235134915/17827758688493, 1482366381361/6201654941325, 5551607622171/13911031047899, -5266607656330/36788968843917, 1074053359553/5740751784926]
+    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]])
\ No newline at end of file