Starting playground for dedalus+pySDC (#472)

* TL: started dedalus playground

* TL: main coupling done, need to adapt to new dedalus version

* TL: first working version

* TL: still need to solve problem with non-linear term

* TL: working version for one-level sweeps

* TL: minor changes

* TL: starting cleaning

* TL: adapted implementation to qmat switch

* TL: updated SDC version to be more efficient and allow variable sweeps

* TL: merged all sdc files into one

* TL: added mpi space-time communicator

* TL: minor debug

* TL: preparing for release

* TL: added RBC script

* TL: correcting a minor mistake

pySDC/playgrounds/dedalus/.gitignore

@@ -0,0 +1 @@
+*.pdf

pySDC/playgrounds/dedalus/README.md

@@ -0,0 +1,134 @@
+# Playground to use Dedalus with pySDC
+
+:scroll: Interface for [Dedalus code](https://dedalus-project.readthedocs.io/en/latest/) so it can be used within the pySDC framework.
+
+> :warning: This is only compatible with the latest version of Dedalus
+
+## Usage Example
+
+See [demo.py](./scratch.py) for a first demo script using pySDC to apply SDC on the Advection equation.
+
+1. Define the problem, as it would have been done with Dedalus
+
+```python
+import numpy as np
+import dedalus.public as d3
+
+# Coordonates, distributor and basis
+coords = d3.CartesianCoordinates('x')
+dist = d3.Distributor(coords, dtype=np.float64)
+xbasis = d3.RealFourier(coords['x'], size=64, bounds=(0, 2*np.pi))
+u = dist.Field(name='u', bases=xbasis)
+
+# Initial solution
+x = xbasis.local_grid(dist, scale=1)
+listK = [0, 1, 2]
+u0 = np.sum([np.cos(k*x) for k in listK], axis=0)
+np.copyto(u['g'], u0)
+
+# Problem
+dx = lambda f: d3.Differentiate(f, coords['x'])
+problem = d3.IVP([u], namespace=locals())
+problem.add_equation("dt(u) + dx(u) = 0")
+```
+
+2. Define the pySDC parameters
+
+```python
+from problem import DedalusProblem
+from sweeper import DedalusSweeperIMEX
+from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+
+nSweeps = 3
+nNodes = 4
+tEnd = 1
+nSteps = 10
+dt = tEnd / nSteps
+
+# pySDC controller settings
+description = {
+    # Sweeper and its parameters
+    "sweeper_class": DedalusSweeperIMEX,
+    "sweeper_params": {
+        "quad_type": "RADAU-RIGHT",
+        "num_nodes": nNodes,
+        "node_type": "LEGENDRE",
+        "initial_guess": "copy",
+        "do_coll_update": False,
+        "QI": "IE",
+        "QE": "EE",
+        'skip_residual_computation': ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
+    },
+    # Step parameters
+    "step_params": {
+        "maxiter": 1,
+    },
+    # Level parameters
+    "level_params": {
+        "dt": dt,
+        "restol": -1,
+        "nsweeps": nSweeps,
+    },
+    "problem_class": DedalusProblem,
+    "problem_params": {
+        'problem': problem,
+        'nNodes': nNodes,
+    }
+}
+```
+
+Here the `DedalusProblem` (defined in [`problem.py`](problem.py)) and the `DedalusSweeperIMEX` (defined in [`sweeper.py`](./sweeper.py)) are the main interfaces between Dedalus and pySDC.
+
+3. Instantiate the controller and run the simulation
+
+```python
+controller = controller_nonMPI(
+    num_procs=1, controller_params={'logger_level': 30},
+    description=description)
+
+prob = controller.MS[0].levels[0].prob
+uSol = prob.solver.state
+tVals = np.linspace(0, tEnd, nSteps + 1)
+
+for i in range(nSteps):
+    uSol, _ = controller.run(u0=uSol, t0=tVals[i], Tend=tVals[i + 1])
+```
+
+Then `uSol` contains a list of `Fields` that represent the final solution of the simulation. Running the [`demo.py`](./demo.py) script produce the following output :
+
+<p align="center">
+  <img src="./demo_advection.png" width="500"/>
+</p>
+
+See an other example with the [Burger equation](./burger.py)
+
+
+## Use a pySDC based time-integrator within Dedalus
+
+This playground also provide a standalone SDC solver that can be used directly,
+see the [demo file for the Burger equation](./burger_ref.py).
+
+To use this standalone time-integrator, simply do :
+
+```python
+# Base import
+from pySDC.playgrounds.dedalus.sdc import SpectralDeferredCorrectionIMEX
+
+# Setup SDC parameters (non set parameters use default values ...)
+SpectralDeferredCorrectionIMEX.setParameters(
+    nSweeps=4,
+    nNodes=4,
+    implSweep="MIN-SR-S",
+    explSweep="PIC")
+```
+
+then you can use this class when instantiating and using a Dedalus solver simply like this :
+
+```python
+solver = problem.build_solver(SpectralDeferredCorrectionIMEX)
+timestep = 2e-3  # dummy value for example ...
+while solver.proceed:
+    solver.step(timestep)
+```
+
+A full example script for the 2D Rayleigh-Benard Convection problem can be found [here](./rayleighBenardSDC.py).

pySDC/playgrounds/dedalus/advection.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Basic script to run the Advection problem with Dedalus
+and the SpectralDeferredCorrectionIMEX time integrator
+"""
+import numpy as np
+import dedalus.public as d3
+from dedalus_dev import ERK4
+from dedalus_dev import SpectralDeferredCorrectionIMEX
+from utils import plt  # import matplotlib with improved graph settings
+
+# Bases and field
+coords = d3.CartesianCoordinates('x')
+dist = d3.Distributor(coords, dtype=np.float64)
+xbasis = d3.RealFourier(coords['x'], size=16, bounds=(0, 2*np.pi))
+u = dist.Field(name='u', bases=xbasis)
+
+# Initial solution
+x = xbasis.local_grid()
+listK = [0, 1, 2]
+u0 = np.sum([np.cos(k*x) for k in listK], axis=0)
+np.copyto(u['g'], u0)
+
+plt.figure('Initial solution')
+plt.plot(u['g'], label='Real space')
+plt.plot(u['c'], 'o', label='Coefficient space')
+plt.legend()
+plt.grid()
+
+# Problem
+dx = lambda f: d3.Differentiate(f, coords['x'])
+problem = d3.IVP([u], namespace=locals())
+problem.add_equation("dt(u) + dx(u) = 0")
+
+# Prepare plots
+orderPlot = {'RK111': 1,
+             'RK222': 2,
+             'RK443': 3,
+             'ERK4': 4}
+plt.figure('Error')
+
+SpectralDeferredCorrectionIMEX.setParameters(
+    M=3, quadType='RADAU-RIGHT', nodeDistr='LEGENDRE',
+    implSweep='OPT-QmQd-0', explSweep='PIC', initSweep='COPY',
+    forceProl=True)
+
+for timeStepper in [d3.RK111, ERK4, 1, 2]:
+
+    # For integer number, use SDC with given number of sweeps
+    useSDC = False
+    nSweep = None
+    if isinstance(timeStepper, int):
+        # Using SDC with a given number of sweeps
+        nSweep = timeStepper
+        timeStepper = SpectralDeferredCorrectionIMEX
+        timeStepper.nSweep = nSweep
+        useSDC = True
+
+    # Build solver
+    solver = problem.build_solver(timeStepper)
+    solver.stop_sim_time = 2*np.pi
+    name = timeStepper.__name__
+
+    # Function to run the simulation with one given time step
+    def getErr(nStep):
+        np.copyto(u['g'], u0)
+        solver.sim_time = 0
+        dt = 2*np.pi/nStep
+        for _ in range(nStep):
+            solver.step(dt)
+        err = np.linalg.norm(u0-u['g'], ord=np.inf)
+        return dt, err
+
+    # Run all simulations
+    listNumStep = [2**(i+2) for i in range(11)]
+    dt, err = np.array([getErr(n) for n in listNumStep]).T
+
+    # Plot error VS time step
+    lbl = f'SDC, nSweep={nSweep}' if useSDC else name
+    sym = '^-' if useSDC else 'o-'
+    plt.loglog(dt, err, sym, label=lbl)
+
+    # Eventually plot order curve
+    if name in orderPlot:
+        order = orderPlot[name]
+        c = err[-1]/dt[-1]**order * 2
+        plt.plot(dt, c*dt**order, '--', color='gray')
+
+plt.xlabel(r'$\Delta{t}$')
+plt.ylabel(r'error ($L_{inf}$)')
+plt.ylim(1e-9, 1e2)
+plt.legend()
+plt.grid(True)
+plt.tight_layout()

pySDC/playgrounds/dedalus/burger.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Demo script for the KdV-Burgers equation
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import dedalus.public as d3
+import logging
+logger = logging.getLogger(__name__)
+
+from problem import DedalusProblem
+from sweeper import DedalusSweeperIMEX
+from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+
+
+# Parameters
+Lx = 10
+Nx = 1024
+a = 1e-4
+b = 2e-4
+dealias = 3/2
+dtype = np.float64
+
+tEnd = 10
+nSteps = 10000
+
+
+# Bases
+xcoord = d3.Coordinate('x')
+dist = d3.Distributor(xcoord, dtype=dtype)
+xbasis = d3.RealFourier(xcoord, size=Nx, bounds=(0, Lx), dealias=dealias)
+
+# Fields
+u = dist.Field(name='u', bases=xbasis)
+
+# Substitutions
+dx = lambda A: d3.Differentiate(A, xcoord)
+
+# Problem
+problem = d3.IVP([u], namespace=locals())
+problem.add_equation("dt(u) - a*dx(dx(u)) - b*dx(dx(dx(u))) = - u*dx(u)")
+
+# Initial conditions
+x = dist.local_grid(xbasis)
+n = 20
+u['g'] = np.log(1 + np.cosh(n)**2/np.cosh(n*(x-0.2*Lx))**2) / (2*n)
+
+# pySDC parameters
+dt = tEnd / nSteps
+nSweeps = 1
+nNodes = 4
+
+description = {
+    # Sweeper and its parameters
+    "sweeper_class": DedalusSweeperIMEX,
+    "sweeper_params": {
+        "quad_type": "RADAU-RIGHT",
+        "num_nodes": nNodes,
+        "node_type": "LEGENDRE",
+        "initial_guess": "copy",
+        "do_coll_update": False,
+        "QI": "IE",
+        "QE": "EE",
+        'skip_residual_computation':
+            ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
+    },
+    # Step parameters
+    "step_params": {
+        "maxiter": 1,
+    },
+    # Level parameters
+    "level_params": {
+        "dt": dt,
+        "restol": -1,
+        "nsweeps": nSweeps,
+    },
+    "problem_class": DedalusProblem,
+    "problem_params": {
+        'problem': problem,
+        'nNodes': nNodes,
+    }
+}
+
+# Main loop
+u.change_scales(1)
+u_list = [np.copy(u['g'])]
+t_list = [0]
+
+size = u_list[0].size
+
+controller = controller_nonMPI(
+    num_procs=1, controller_params={'logger_level': 30},
+    description=description)
+
+prob = controller.MS[0].levels[0].prob
+uSol = prob.solver.state
+
+tVals = np.linspace(0, tEnd, nSteps + 1)
+
+for i in range(nSteps):
+    uSol, _ = controller.run(u0=uSol, t0=tVals[i], Tend=tVals[i + 1])
+    if (i+1) % 100 == 0:
+        print(f"step {i+1}/{nSteps}")
+    if (i+1) % 25 == 0:
+
+        u.change_scales(1)
+        u_list.append(np.copy(u['g']))
+        t_list.append(tVals[i])
+
+
+# Plot
+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, Lx)
+plt.ylim(0, tEnd)
+plt.xlabel('x')
+plt.ylabel('t')
+plt.title(f'KdV-Burgers, (a,b)=({a},{b})')
+plt.tight_layout()
+plt.savefig("KdV_Burgers_pySDC.pdf")