Skip to content

Commit 453d024

Browse files
authored
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
1 parent b08df67 commit 453d024

File tree

12 files changed

+1737
-0
lines changed

12 files changed

+1737
-0
lines changed
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
*.pdf
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
# Playground to use Dedalus with pySDC
2+
3+
:scroll: Interface for [Dedalus code](https://dedalus-project.readthedocs.io/en/latest/) so it can be used within the pySDC framework.
4+
5+
> :warning: This is only compatible with the latest version of Dedalus
6+
7+
## Usage Example
8+
9+
See [demo.py](./scratch.py) for a first demo script using pySDC to apply SDC on the Advection equation.
10+
11+
1. Define the problem, as it would have been done with Dedalus
12+
13+
```python
14+
import numpy as np
15+
import dedalus.public as d3
16+
17+
# Coordonates, distributor and basis
18+
coords = d3.CartesianCoordinates('x')
19+
dist = d3.Distributor(coords, dtype=np.float64)
20+
xbasis = d3.RealFourier(coords['x'], size=64, bounds=(0, 2*np.pi))
21+
u = dist.Field(name='u', bases=xbasis)
22+
23+
# Initial solution
24+
x = xbasis.local_grid(dist, scale=1)
25+
listK = [0, 1, 2]
26+
u0 = np.sum([np.cos(k*x) for k in listK], axis=0)
27+
np.copyto(u['g'], u0)
28+
29+
# Problem
30+
dx = lambda f: d3.Differentiate(f, coords['x'])
31+
problem = d3.IVP([u], namespace=locals())
32+
problem.add_equation("dt(u) + dx(u) = 0")
33+
```
34+
35+
2. Define the pySDC parameters
36+
37+
```python
38+
from problem import DedalusProblem
39+
from sweeper import DedalusSweeperIMEX
40+
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
41+
42+
nSweeps = 3
43+
nNodes = 4
44+
tEnd = 1
45+
nSteps = 10
46+
dt = tEnd / nSteps
47+
48+
# pySDC controller settings
49+
description = {
50+
# Sweeper and its parameters
51+
"sweeper_class": DedalusSweeperIMEX,
52+
"sweeper_params": {
53+
"quad_type": "RADAU-RIGHT",
54+
"num_nodes": nNodes,
55+
"node_type": "LEGENDRE",
56+
"initial_guess": "copy",
57+
"do_coll_update": False,
58+
"QI": "IE",
59+
"QE": "EE",
60+
'skip_residual_computation': ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
61+
},
62+
# Step parameters
63+
"step_params": {
64+
"maxiter": 1,
65+
},
66+
# Level parameters
67+
"level_params": {
68+
"dt": dt,
69+
"restol": -1,
70+
"nsweeps": nSweeps,
71+
},
72+
"problem_class": DedalusProblem,
73+
"problem_params": {
74+
'problem': problem,
75+
'nNodes': nNodes,
76+
}
77+
}
78+
```
79+
80+
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.
81+
82+
3. Instantiate the controller and run the simulation
83+
84+
```python
85+
controller = controller_nonMPI(
86+
num_procs=1, controller_params={'logger_level': 30},
87+
description=description)
88+
89+
prob = controller.MS[0].levels[0].prob
90+
uSol = prob.solver.state
91+
tVals = np.linspace(0, tEnd, nSteps + 1)
92+
93+
for i in range(nSteps):
94+
uSol, _ = controller.run(u0=uSol, t0=tVals[i], Tend=tVals[i + 1])
95+
```
96+
97+
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 :
98+
99+
<p align="center">
100+
<img src="./demo_advection.png" width="500"/>
101+
</p>
102+
103+
See an other example with the [Burger equation](./burger.py)
104+
105+
106+
## Use a pySDC based time-integrator within Dedalus
107+
108+
This playground also provide a standalone SDC solver that can be used directly,
109+
see the [demo file for the Burger equation](./burger_ref.py).
110+
111+
To use this standalone time-integrator, simply do :
112+
113+
```python
114+
# Base import
115+
from pySDC.playgrounds.dedalus.sdc import SpectralDeferredCorrectionIMEX
116+
117+
# Setup SDC parameters (non set parameters use default values ...)
118+
SpectralDeferredCorrectionIMEX.setParameters(
119+
nSweeps=4,
120+
nNodes=4,
121+
implSweep="MIN-SR-S",
122+
explSweep="PIC")
123+
```
124+
125+
then you can use this class when instantiating and using a Dedalus solver simply like this :
126+
127+
```python
128+
solver = problem.build_solver(SpectralDeferredCorrectionIMEX)
129+
timestep = 2e-3 # dummy value for example ...
130+
while solver.proceed:
131+
solver.step(timestep)
132+
```
133+
134+
A full example script for the 2D Rayleigh-Benard Convection problem can be found [here](./rayleighBenardSDC.py).
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
#!/usr/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
"""
4+
Basic script to run the Advection problem with Dedalus
5+
and the SpectralDeferredCorrectionIMEX time integrator
6+
"""
7+
import numpy as np
8+
import dedalus.public as d3
9+
from dedalus_dev import ERK4
10+
from dedalus_dev import SpectralDeferredCorrectionIMEX
11+
from utils import plt # import matplotlib with improved graph settings
12+
13+
# Bases and field
14+
coords = d3.CartesianCoordinates('x')
15+
dist = d3.Distributor(coords, dtype=np.float64)
16+
xbasis = d3.RealFourier(coords['x'], size=16, bounds=(0, 2*np.pi))
17+
u = dist.Field(name='u', bases=xbasis)
18+
19+
# Initial solution
20+
x = xbasis.local_grid()
21+
listK = [0, 1, 2]
22+
u0 = np.sum([np.cos(k*x) for k in listK], axis=0)
23+
np.copyto(u['g'], u0)
24+
25+
plt.figure('Initial solution')
26+
plt.plot(u['g'], label='Real space')
27+
plt.plot(u['c'], 'o', label='Coefficient space')
28+
plt.legend()
29+
plt.grid()
30+
31+
# Problem
32+
dx = lambda f: d3.Differentiate(f, coords['x'])
33+
problem = d3.IVP([u], namespace=locals())
34+
problem.add_equation("dt(u) + dx(u) = 0")
35+
36+
# Prepare plots
37+
orderPlot = {'RK111': 1,
38+
'RK222': 2,
39+
'RK443': 3,
40+
'ERK4': 4}
41+
plt.figure('Error')
42+
43+
SpectralDeferredCorrectionIMEX.setParameters(
44+
M=3, quadType='RADAU-RIGHT', nodeDistr='LEGENDRE',
45+
implSweep='OPT-QmQd-0', explSweep='PIC', initSweep='COPY',
46+
forceProl=True)
47+
48+
for timeStepper in [d3.RK111, ERK4, 1, 2]:
49+
50+
# For integer number, use SDC with given number of sweeps
51+
useSDC = False
52+
nSweep = None
53+
if isinstance(timeStepper, int):
54+
# Using SDC with a given number of sweeps
55+
nSweep = timeStepper
56+
timeStepper = SpectralDeferredCorrectionIMEX
57+
timeStepper.nSweep = nSweep
58+
useSDC = True
59+
60+
# Build solver
61+
solver = problem.build_solver(timeStepper)
62+
solver.stop_sim_time = 2*np.pi
63+
name = timeStepper.__name__
64+
65+
# Function to run the simulation with one given time step
66+
def getErr(nStep):
67+
np.copyto(u['g'], u0)
68+
solver.sim_time = 0
69+
dt = 2*np.pi/nStep
70+
for _ in range(nStep):
71+
solver.step(dt)
72+
err = np.linalg.norm(u0-u['g'], ord=np.inf)
73+
return dt, err
74+
75+
# Run all simulations
76+
listNumStep = [2**(i+2) for i in range(11)]
77+
dt, err = np.array([getErr(n) for n in listNumStep]).T
78+
79+
# Plot error VS time step
80+
lbl = f'SDC, nSweep={nSweep}' if useSDC else name
81+
sym = '^-' if useSDC else 'o-'
82+
plt.loglog(dt, err, sym, label=lbl)
83+
84+
# Eventually plot order curve
85+
if name in orderPlot:
86+
order = orderPlot[name]
87+
c = err[-1]/dt[-1]**order * 2
88+
plt.plot(dt, c*dt**order, '--', color='gray')
89+
90+
plt.xlabel(r'$\Delta{t}$')
91+
plt.ylabel(r'error ($L_{inf}$)')
92+
plt.ylim(1e-9, 1e2)
93+
plt.legend()
94+
plt.grid(True)
95+
plt.tight_layout()
Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
#!/usr/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
"""
4+
Demo script for the KdV-Burgers equation
5+
"""
6+
7+
import numpy as np
8+
import matplotlib.pyplot as plt
9+
import dedalus.public as d3
10+
import logging
11+
logger = logging.getLogger(__name__)
12+
13+
from problem import DedalusProblem
14+
from sweeper import DedalusSweeperIMEX
15+
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
16+
17+
18+
# Parameters
19+
Lx = 10
20+
Nx = 1024
21+
a = 1e-4
22+
b = 2e-4
23+
dealias = 3/2
24+
dtype = np.float64
25+
26+
tEnd = 10
27+
nSteps = 10000
28+
29+
30+
# Bases
31+
xcoord = d3.Coordinate('x')
32+
dist = d3.Distributor(xcoord, dtype=dtype)
33+
xbasis = d3.RealFourier(xcoord, size=Nx, bounds=(0, Lx), dealias=dealias)
34+
35+
# Fields
36+
u = dist.Field(name='u', bases=xbasis)
37+
38+
# Substitutions
39+
dx = lambda A: d3.Differentiate(A, xcoord)
40+
41+
# Problem
42+
problem = d3.IVP([u], namespace=locals())
43+
problem.add_equation("dt(u) - a*dx(dx(u)) - b*dx(dx(dx(u))) = - u*dx(u)")
44+
45+
# Initial conditions
46+
x = dist.local_grid(xbasis)
47+
n = 20
48+
u['g'] = np.log(1 + np.cosh(n)**2/np.cosh(n*(x-0.2*Lx))**2) / (2*n)
49+
50+
# pySDC parameters
51+
dt = tEnd / nSteps
52+
nSweeps = 1
53+
nNodes = 4
54+
55+
description = {
56+
# Sweeper and its parameters
57+
"sweeper_class": DedalusSweeperIMEX,
58+
"sweeper_params": {
59+
"quad_type": "RADAU-RIGHT",
60+
"num_nodes": nNodes,
61+
"node_type": "LEGENDRE",
62+
"initial_guess": "copy",
63+
"do_coll_update": False,
64+
"QI": "IE",
65+
"QE": "EE",
66+
'skip_residual_computation':
67+
('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE'),
68+
},
69+
# Step parameters
70+
"step_params": {
71+
"maxiter": 1,
72+
},
73+
# Level parameters
74+
"level_params": {
75+
"dt": dt,
76+
"restol": -1,
77+
"nsweeps": nSweeps,
78+
},
79+
"problem_class": DedalusProblem,
80+
"problem_params": {
81+
'problem': problem,
82+
'nNodes': nNodes,
83+
}
84+
}
85+
86+
# Main loop
87+
u.change_scales(1)
88+
u_list = [np.copy(u['g'])]
89+
t_list = [0]
90+
91+
size = u_list[0].size
92+
93+
controller = controller_nonMPI(
94+
num_procs=1, controller_params={'logger_level': 30},
95+
description=description)
96+
97+
prob = controller.MS[0].levels[0].prob
98+
uSol = prob.solver.state
99+
100+
tVals = np.linspace(0, tEnd, nSteps + 1)
101+
102+
for i in range(nSteps):
103+
uSol, _ = controller.run(u0=uSol, t0=tVals[i], Tend=tVals[i + 1])
104+
if (i+1) % 100 == 0:
105+
print(f"step {i+1}/{nSteps}")
106+
if (i+1) % 25 == 0:
107+
108+
u.change_scales(1)
109+
u_list.append(np.copy(u['g']))
110+
t_list.append(tVals[i])
111+
112+
113+
# Plot
114+
plt.figure(figsize=(6, 4))
115+
plt.pcolormesh(
116+
x.ravel(), np.array(t_list), np.array(u_list), cmap='RdBu_r',
117+
shading='gouraud', rasterized=True, clim=(-0.8, 0.8))
118+
plt.xlim(0, Lx)
119+
plt.ylim(0, tEnd)
120+
plt.xlabel('x')
121+
plt.ylabel('t')
122+
plt.title(f'KdV-Burgers, (a,b)=({a},{b})')
123+
plt.tight_layout()
124+
plt.savefig("KdV_Burgers_pySDC.pdf")

0 commit comments

Comments
 (0)