TL: continued refactoring

Author: tlunet

pySDC/playgrounds/dedalus/demo_timestepper_burger.py

@@ -4,7 +4,7 @@
 Quick demo to solve the 1D advection-diffusion with Dedalus
 using the pySDC interface
 """
-# Base user imports
+# Base python imports
 import numpy as np
 import matplotlib.pyplot as plt
 

pySDC/playgrounds/dedalus/demos/demo_interface_advDiff.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Demo script for the KdV-Burgers equation
+Demo script for the KdV-Burgers equation, using the pySDC interface
 """
 import numpy as np
 import matplotlib.pyplot as plt
@@ -12,20 +12,28 @@
 from pySDC.playgrounds.dedalus.interface import DedalusProblem, DedalusSweeperIMEX
 from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
 
-# Space parameters
-xEnd = 10
-nX = 512
-nu = 1e-4
-b = 2e-4
+# -----------------------------------------------------------------------------
+# User parameters
+# -----------------------------------------------------------------------------
+xEnd = 10   # space domain size
+nX = 512    # number of point in space
+nu = 1e-4   # diffusion coefficient
+b = 2e-4    # hyper-diffusion coefficient
 
-# Time-integration parameters
+
+# -- time integration
 nSweeps = 4
 nNodes = 4
 tEnd = 10
 nSteps = 5000
+
+# -----------------------------------------------------------------------------
+# Solver setup
+# -----------------------------------------------------------------------------
 timeStep = tEnd / nSteps
 
 pData = buildKdVBurgerProblem(nX, xEnd, nu, b)
+problem, u, x = [pData[key] for key in ["problem", "u", "x"]]
 
 description = {
     # Sweeper and its parameters
@@ -36,7 +44,7 @@
         "node_type": "LEGENDRE",
         "initial_guess": "copy",
         "do_coll_update": False,
-        "QI": "MIN-SR-S",
+        "QI": "MIN-SR-FLEX",
         "QE": "PIC",
         'skip_residual_computation':
             ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
@@ -53,13 +61,14 @@
     },
     "problem_class": DedalusProblem,
     "problem_params": {
-        'problem': pData["problem"],
+        'problem': problem,
         'nNodes': nNodes,
     }
 }
 
-# Main loop
-u, x = [pData[key] for key in ["u", "x"]]
+# -----------------------------------------------------------------------------
+# Simulation run
+# -----------------------------------------------------------------------------
 u.change_scales(1)
 u_list = [np.copy(u['g'])]
 t_list = [0]
@@ -84,8 +93,9 @@
         u_list.append(np.copy(u['g']))
         t_list.append(tVals[i])
 
-
-# Plot
+# -----------------------------------------------------------------------------
+# Plotting solution in real space
+# -----------------------------------------------------------------------------
 plt.figure(figsize=(6, 4))
 plt.pcolormesh(
     x.ravel(), np.array(t_list), np.array(u_list), cmap='RdBu_r',

pySDC/playgrounds/dedalus/demo_interface_burger.py renamed to pySDC/playgrounds/dedalus/demos/demo_interface_burger.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Demo script for the KdV-Burgers equation, using the pySDC timestepper
+"""
+# Base python imports
+import numpy as np
+import matplotlib.pyplot as plt
+
+import logging
+logger = logging.getLogger(__name__)
+
+# pySDC imports
+from pySDC.playgrounds.dedalus.problems import buildKdVBurgerProblem
+from pySDC.playgrounds.dedalus.timestepper import SpectralDeferredCorrectionIMEX
+
+# Dedalus import (for alternative time-stepper)
+import dedalus.public as d3
+
+# -----------------------------------------------------------------------------
+# User parameters
+# -----------------------------------------------------------------------------
+xEnd = 10   # space domain size
+nX = 512    # number of point in space
+nu = 1e-4   # diffusion coefficient
+b = 2e-4    # hyper-diffusion coefficient
+
+# -- time integration
+tEnd = 10
+nSteps = 5000
+SpectralDeferredCorrectionIMEX.setParameters(
+    nSweeps=4,
+    nNodes=4,
+    implSweep="MIN-SR-FLEX",
+    explSweep="PIC")
+useSDC = True
+
+# -----------------------------------------------------------------------------
+# Solver setup
+# -----------------------------------------------------------------------------
+timestepper = SpectralDeferredCorrectionIMEX if useSDC else d3.RK443
+timestep = tEnd/nSteps
+
+pData = buildKdVBurgerProblem(nX, xEnd, nu, b)
+problem, u, x = [pData[key] for key in ["problem", "u", "x"]]
+
+solver = problem.build_solver(timestepper)
+solver.stop_sim_time = tEnd
+
+# -----------------------------------------------------------------------------
+# Simulation run
+# -----------------------------------------------------------------------------
+u.change_scales(1)
+u_list = [np.copy(u['g'])]
+t_list = [solver.sim_time]
+i = 0
+while solver.proceed:
+    solver.step(timestep)
+    if solver.iteration % 100 == 0:
+        print(f"step {solver.iteration}/{nSteps}")
+        logger.info('Iteration=%i, Time=%e, dt=%e' %(solver.iteration, solver.sim_time, timestep))
+    if solver.iteration % 25 == 0:
+        u.change_scales(1)
+        u_list.append(np.copy(u['g']))
+        t_list.append(solver.sim_time)
+    i += 1
+solver.log_stats()
+
+# -----------------------------------------------------------------------------
+# Plotting solution in real space
+# -----------------------------------------------------------------------------
+plt.figure(figsize=(6, 4))
+plt.pcolormesh(x.ravel(), np.array(t_list), np.array(u_list), cmap='RdBu_r', shading='gouraud', rasterized=True, clim=(-0.8, 0.8))
+plt.xlim(0, xEnd)
+plt.ylim(0, tEnd)
+plt.xlabel('x')
+plt.ylabel('t')
+plt.title(r'KdV-Burgers, $(\nu,b)='f'({nu},{b})$')
+plt.tight_layout()
+plt.savefig(f"demo_timestepper_burger{'_SDC' if useSDC else ''}.png")

pySDC/playgrounds/dedalus/demos/demo_timestepper_burger.py

@@ -96,8 +96,7 @@ def update_nodes(self):
 
             # Add quadrature terms
             for j in range(M):
-                axpy(a=dt*q[m, j], x=Fk[j].data, y=RHS.data)
-                axpy(a=-dt*q[m, j], x=LXk[j].data, y=RHS.data)
+                axpy(a=dt*q[m, j], x=(Fk[j].data - LXk[j].data), y=RHS.data)
 
             # Add F and LX terms from iteration k+1
             for j in range(m):
@@ -111,6 +110,8 @@ def update_nodes(self):
             axpy(a=dt*qI[m, m], x=LXk[m].data, y=RHS.data)
 
             # Solve system and store node solution in solver state
+            if self.genQI.isKDependent():
+                P.updateLHS(dt, qI)
             P.solveAndStoreState(m)
 
             # Evaluate and store LX with current state

pySDC/playgrounds/dedalus/interface/sweeper.py

@@ -221,6 +221,10 @@ def leftIsNode(cls):
     def doProlongation(cls):
         return not cls.rightIsNode or cls.forceProl
 
+    @classmethod
+    def isKDependent(cls):
+        return not np.all(cls.QDeltaI == cls.QDeltaI[0])
+
 
 # -----------------------------------------------------------------------------
 # Dedalus based IMEX timeintegrator class
@@ -312,15 +316,26 @@ def _updateLHS(self, dt, init=False):
             if init:
                 # Potentially instantiate list of solver (ony first time step)
                 sp.LHS_solvers = [[None for _ in range(self.M)] for _ in range(self.nSweeps)]
-            for k in range(self.nSweeps):
+            if self.isKDependent():
+                for k in range(self.nSweeps):
+                    for m in range(self.M):
+                        if solver.store_expanded_matrices:
+                            raise NotImplementedError("code correction required")
+                        else:
+                            sp.LHS = (sp.M_min + dt*qI[k, m, m]*sp.L_min)
+                        sp.LHS_solvers[k][m] = solver.matsolver(sp.LHS, solver)
+            else:
                 for m in range(self.M):
-                    if solver.store_expanded_matrices:
-                        raise NotImplementedError("code correction required")
-                        np.copyto(sp.LHS.data, sp.M_exp.data)
-                        self.axpy(a=dt*qI[k, m, m], x=sp.L_exp.data, y=sp.LHS.data)
-                    else:
-                        sp.LHS = (sp.M_min + dt*qI[k, m, m]*sp.L_min)
-                    sp.LHS_solvers[k][m] = solver.matsolver(sp.LHS, solver)
+                    for k in range(self.nSweeps):
+                        if k == 0:
+                            if solver.store_expanded_matrices:
+                                raise NotImplementedError("code correction required")
+                            else:
+                                sp.LHS = (sp.M_min + dt*qI[k, m, m]*sp.L_min)
+                            sp.LHS_solvers[k][m] = solver.matsolver(sp.LHS, solver)
+                        else:
+                            sp.LHS_solvers[k][m] = sp.LHS_solvers[0][m]
+
             if self.initSweep == "QDELTA":
                 raise NotImplementedError()