diff --git a/pySDC/playgrounds/dedalus/.gitignore b/pySDC/playgrounds/dedalus/.gitignore new file mode 100644 index 0000000000..f08278d85a --- /dev/null +++ b/pySDC/playgrounds/dedalus/.gitignore @@ -0,0 +1 @@ +*.pdf \ No newline at end of file diff --git a/pySDC/playgrounds/dedalus/README.md b/pySDC/playgrounds/dedalus/README.md new file mode 100644 index 0000000000..1163e1f5bd --- /dev/null +++ b/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 : + +

+ +

+ +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). \ No newline at end of file diff --git a/pySDC/playgrounds/dedalus/advection.py b/pySDC/playgrounds/dedalus/advection.py new file mode 100644 index 0000000000..94169cfbcc --- /dev/null +++ b/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() diff --git a/pySDC/playgrounds/dedalus/burger.py b/pySDC/playgrounds/dedalus/burger.py new file mode 100644 index 0000000000..c1c32abdf8 --- /dev/null +++ b/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") diff --git a/pySDC/playgrounds/dedalus/burger_ref.py b/pySDC/playgrounds/dedalus/burger_ref.py new file mode 100644 index 0000000000..d04db05046 --- /dev/null +++ b/pySDC/playgrounds/dedalus/burger_ref.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Dedalus script simulating the 1D Korteweg-de Vries / Burgers equation. +This script demonstrates solving a 1D initial value problem and produces +a space-time plot of the solution. It should take just a few seconds to +run (serial only). + +We use a Fourier basis to solve the IVP: + dt(u) + u*dx(u) = a*dx(dx(u)) + b*dx(dx(dx(u))) + +To run and plot: + $ python3 kdv_burgers.py +""" + +import numpy as np +import matplotlib.pyplot as plt +import dedalus.public as d3 +import logging +logger = logging.getLogger(__name__) + +from pySDC.playgrounds.dedalus.sdc import SpectralDeferredCorrectionIMEX + +SpectralDeferredCorrectionIMEX.setParameters( + nSweeps=4, + nNodes=4, + implSweep="MIN-SR-S", + explSweep="PIC") + +useSDC = True + +# Parameters +Lx = 10 +Nx = 1024 +a = 1e-4 +b = 2e-4 +dealias = 3/2 +stop_sim_time = 10 +timestepper = SpectralDeferredCorrectionIMEX if useSDC else d3.SBDF2 +timestep = 2e-3 +dtype = np.float64 + +# 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) + +# Solver +solver = problem.build_solver(timestepper) +solver.stop_sim_time = stop_sim_time + +# Main loop +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}/...") + 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 + +# 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, stop_sim_time) +plt.xlabel('x') +plt.ylabel('t') +plt.title(f'KdV-Burgers, (a,b)=({a},{b})') +plt.tight_layout() +plt.savefig(f"KdV_Burgers_ref_useSDC{useSDC}.pdf") diff --git a/pySDC/playgrounds/dedalus/demo.py b/pySDC/playgrounds/dedalus/demo.py new file mode 100644 index 0000000000..d3c555e219 --- /dev/null +++ b/pySDC/playgrounds/dedalus/demo.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Experiments with dedalus and pySDC +""" +# Base user imports +import numpy as np +import matplotlib.pyplot as plt +import dedalus.public as d3 + +from problem import DedalusProblem +from sweeper import DedalusSweeperIMEX +from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI + + +coords = d3.CartesianCoordinates('x') +dist = d3.Distributor(coords, dtype=np.float64) +xbasis = d3.RealFourier(coords['x'], size=32, 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) + +plt.figure('Initial solution') +plt.plot(u['g'], label='Real space (u0)') +plt.plot(u['c'], 'o', label='Coefficient space (u0)') +plt.legend() +plt.grid(True) + +# Problem +dx = lambda f: d3.Differentiate(f, coords['x']) +problem = d3.IVP([u], namespace=locals()) +problem.add_equation("dt(u) + dx(u) - dx(dx(u)) = 0") + +nSweeps = 4 +nNodes = 4 +tEnd = 2*np.pi +nSteps = 500 +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, + } +} + +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]) + +plt.plot(uSol[0]['g'], label='pySDC solution (u(t=1))') +plt.legend() diff --git a/pySDC/playgrounds/dedalus/demo_advection.png b/pySDC/playgrounds/dedalus/demo_advection.png new file mode 100644 index 0000000000..9d181b11b6 Binary files /dev/null and b/pySDC/playgrounds/dedalus/demo_advection.png differ diff --git a/pySDC/playgrounds/dedalus/mpi.py b/pySDC/playgrounds/dedalus/mpi.py new file mode 100644 index 0000000000..db01503a34 --- /dev/null +++ b/pySDC/playgrounds/dedalus/mpi.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +MPI space-time decomposition +""" +from mpi4py import MPI +import sys + +def initSpaceTimeCommunicators(nProcSpace=None, nProcTime=None, groupSpace=True): + + gComm = MPI.COMM_WORLD + gRank = gComm.Get_rank() + gSize = gComm.Get_size() + + def log(msg): + if gRank == 0: + print(msg) + + if (nProcTime is None) and (nProcSpace is None): + # No parallelization specified, space takes everything + nProcTime = 1 + nProcSpace = gSize + elif nProcSpace is None: + # Only time parallelization specified, space takes the rest + nProcSpace = gSize//nProcTime + elif nProcTime is None: + # Only space parallelization specified, time takes the rest + nProcTime = gSize//nProcSpace + + log("-- Starting MPI initialization ...") + log(f" - nProcTotal = {gSize}") + log(f" - nProcSpace = {nProcSpace}") + log(f" - nProcTime = {nProcTime}") + + # Check for inadequate decomposition + if (gSize != nProcSpace*nProcTime) and (gSize != 1): + if gRank == 0: + raise ValueError( + f' product of nProcSpace ({nProcSpace}) with nProcTime ({nProcTime})' + f' is not equal to the total number of MPI processes ({gSize})') + else: + sys.exit(0) + + + # Information message + if gSize == 1: + log(" - no parallelization at all") + nProcSpace = 1 + nProcSpace = 1 + else: + if nProcSpace != 1: + log(f" - space parallelization activated : {nProcSpace} mpi processes") + else: + log(" - no space parallelization") + if nProcSpace != 1: + log(f" - time parallelization activated : {nProcSpace} mpi processes") + else: + log(" - no time parallelization") + log('-- finished MPI initialization --') + + # Construction of MPI communicators + if groupSpace: + tColor = gRank % nProcSpace + tComm = gComm.Split(tColor, gRank) + gComm.Barrier() + sColor = (gRank-gRank % nProcSpace)/nProcSpace + sComm = gComm.Split(sColor, gRank) + gComm.Barrier() + else: + sColor = gRank % nProcTime + sComm = gComm.Split(sColor, gRank) + gComm.Barrier() + tColor = (gRank-gRank % nProcTime)/nProcSpace + tComm = gComm.Split(tColor, gRank) + gComm.Barrier() + + return gComm, sComm, tComm + + +if __name__ == "__main__": + + gComm, sComm, tComm = initSpaceTimeCommunicators(nProcTime=4) + + gRank, gSize = gComm.Get_rank(), gComm.Get_size() + sRank, sSize = sComm.Get_rank(), sComm.Get_size() + tRank, tSize = tComm.Get_rank(), tComm.Get_size() + + from time import sleep + sleep(0.1*gRank) + + import psutil + coreNum = psutil.Process().cpu_num() + + print(f"Global rank {gRank} ({gSize}), space rank {sRank} ({sSize})," + f" time rank {tRank} ({tRank}) running on CPU core {coreNum}") + \ No newline at end of file diff --git a/pySDC/playgrounds/dedalus/problem.py b/pySDC/playgrounds/dedalus/problem.py new file mode 100644 index 0000000000..fce924a454 --- /dev/null +++ b/pySDC/playgrounds/dedalus/problem.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Problem class for dedalus +""" +from pySDC.core.problem import Problem + +import numpy as np +from scipy.linalg import blas +from collections import deque + +import dedalus.public as d3 + +from dedalus.core.system import CoeffSystem +from dedalus.core.evaluator import Evaluator +from dedalus.tools.array import apply_sparse +from dedalus.core.field import Field + + +def State(cls, fields): + return [f.copy() for f in fields] + + +class DedalusProblem(Problem): + + dtype_u = State + dtype_f = State + + # Dummy class to trick Dedalus + class DedalusTimeStepper: + steps = 1 + stages = 1 + def __init__(self, solver): + self.solver = solver + + def __init__(self, problem:d3.IVP, nNodes, collUpdate=False): + + self.DedalusTimeStepper.stages = nNodes + solver = problem.build_solver(self.DedalusTimeStepper) + self.solver = solver + + # From new version + self.subproblems = [sp for sp in solver.subproblems if sp.size] + self.evaluator:Evaluator = solver.evaluator + self.F_fields = solver.F + + self.M = nNodes + + c = lambda: CoeffSystem(solver.subproblems, dtype=solver.dtype) + self.MX0, self.RHS = c(), c() + self.LX = deque([[c() for _ in range(self.M)] for _ in range(2)]) + self.F = deque([[c() for _ in range(self.M)] for _ in range(2)]) + + # Attributes + self.axpy = blas.get_blas_funcs('axpy', dtype=solver.dtype) + self.dt = None + self.firstEval = True + self.init = True + + # Instantiate M solver, needed only for collocation update + if collUpdate: + for sp in solver.subproblems: + if solver.store_expanded_matrices: + np.copyto(sp.LHS.data, sp.M_exp.data) + else: + sp.LHS = sp.M_min @ sp.pre_right + sp.M_solver = solver.matsolver(sp.LHS, solver) + self.collUpdate = collUpdate + + @property + def state(self): + return self.solver.state + + def stateCopy(self): + return [f.copy() for f in self.solver.state] + + + def computeMX0(self): + """ + Compute MX0 term used in RHS of both initStep and sweep methods + + Update the MX0 attribute of the timestepper object. + """ + MX0 = self.MX0 + state:list[Field] = self.solver.state + + # Require coefficient space + for f in state: + f.require_coeff_space() + + # Compute and store MX0 + for sp in self.subproblems: + spX = sp.gather_inputs(state) + apply_sparse(sp.M_min, spX, axis=0, out=MX0.get_subdata(sp)) + + + def updateLHS(self, dt, qI, init=False): + """Update LHS and LHS solvers for each subproblem + + Parameters + ---------- + dt : float + Time-step for the updated LHS. + qI : 2darray + QDeltaI coefficients. + init : bool, optional + Wether or not initialize the LHS_solvers attribute for each + subproblem. The default is False. + """ + # Update only if different dt + if self.dt == dt: + return + + # Attribute references + solver = self.solver + + # Update LHS and LHS solvers for each subproblems + for sp in solver.subproblems: + if self.init: + # Eventually instanciate list of solvers (ony first time step) + sp.LHS_solvers = [None] * self.M + self.init = False + for i in range(self.M): + if solver.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) + self.axpy(a=dt*qI[i, i], x=sp.L_exp.data, y=sp.LHS.data) + else: + sp.LHS = (sp.M_min + dt*qI[i, i]*sp.L_min) # CREATES TEMPORARY + sp.LHS_solvers[i] = solver.matsolver(sp.LHS, solver) + + + def evalLX(self, LX): + """ + Evaluate LX using the current state, and store it + + Parameters + ---------- + LX : dedalus.core.system.CoeffSystem + Where to store the evaluated fields. + + Returns + ------- + None. + + """ + state:list[Field] = self.solver.state + # Require coefficient space + for f in state: + f.require_coeff_space() + + # Evaluate matrix vector product and store + for sp in self.solver.subproblems: + spX = sp.gather_inputs(self.solver.state) + apply_sparse(sp.L_min, spX, axis=0, out=LX.get_subdata(sp)) + + + def evalF(self, time, F): + """ + Evaluate the F operator from the current solver state + + Note + ---- + After evaluation, state fields are left in grid space + + Parameters + ---------- + time : float + Time of evaluation. + F : dedalus.core.system.CoeffSystem + Where to store the evaluated fields. + """ + solver = self.solver + # Evaluate non linear term on current state + t0 = solver.sim_time + solver.sim_time = time + if self.firstEval: + solver.evaluator.evaluate_scheduled( + sim_time=time, timestep=solver.dt, + iteration=0, wall_time=0) + self.firstEval = False + else: + solver.evaluator.evaluate_group('F') + # Store F evaluation + for sp in solver.subproblems: + sp.gather_outputs(solver.F, out=F.get_subdata(sp)) + # Put back initial solver simulation time + solver.sim_time = t0 + + def solveAndStoreState(self, m): + """ + Solve LHS * X = RHS using the LHS associated to a given node, + and store X into the solver state. + It uses the current RHS attribute of the object. + + Parameters + ---------- + m : int + Index of the nodes. + """ + # Attribute references + solver = self.solver + RHS = self.RHS + + for field in solver.state: + field.preset_layout('c') + + # Solve and store for each subproblem + for sp in solver.subproblems: + # Slice out valid subdata, skipping invalid components + spRHS = RHS.get_subdata(sp) + spX = sp.LHS_solvers[m].solve(spRHS) # CREATES TEMPORARY + sp.scatter_inputs(spX, solver.state) diff --git a/pySDC/playgrounds/dedalus/rayleighBenardSDC.py b/pySDC/playgrounds/dedalus/rayleighBenardSDC.py new file mode 100644 index 0000000000..ca09c5e81e --- /dev/null +++ b/pySDC/playgrounds/dedalus/rayleighBenardSDC.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Dedalus script simulating 2D horizontally-periodic Rayleigh-Benard convection +using Spectral Deferred Corrections. +This script demonstrates solving a 2D Cartesian initial value problem. It can +be ran serially or in parallel, and uses the built-in analysis framework to save +data snapshots to HDF5 files. The `plot_snapshots.py` script can be used to +produce plots from the saved data. It should take about 5 cpu-minutes to run. + +The problem is non-dimensionalized using the box height and freefall time, so +the resulting thermal diffusivity and viscosity are related to the Prandtl +and Rayleigh numbers as: + + kappa = (Rayleigh * Prandtl)**(-1/2) + nu = (Rayleigh / Prandtl)**(-1/2) + +For incompressible hydro with two boundaries, we need two tau terms for each the +velocity and buoyancy. Here we choose to use a first-order formulation, putting +one tau term each on auxiliary first-order gradient variables and the others in +the PDE, and lifting them all to the first derivative basis. This formulation puts +a tau term in the divergence constraint, as required for this geometry. + +To run using e.g. 4 processes: + $ mpiexec -n 4 python rayleighBenardSDC.py + +Then use https://github.com/DedalusProject/dedalus/blob/master/examples/ivp_2d_rayleigh_benard/plot_snapshots.py +for plotting. +""" +import numpy as np +import dedalus.public as d3 +from pySDC.playgrounds.dedalus.sdc import SpectralDeferredCorrectionIMEX +import logging +logger = logging.getLogger(__name__) + +useSDC = True +SpectralDeferredCorrectionIMEX.setParameters( + nSweeps=4, + nNodes=4, + implSweep="MIN-SR-FLEX", + explSweep="PIC") + +# Parameters +Lx, Lz = 4, 1 +Nx, Nz = 512, 128 +Rayleigh = 2e6 +Prandtl = 1 +dealias = 3/2 +stop_sim_time = 50 +timestepper = SpectralDeferredCorrectionIMEX if useSDC else d3.RK222 +timestep = 1e-2/4 +dtype = np.float64 + +# Bases +coords = d3.CartesianCoordinates('x', 'z') +dist = d3.Distributor(coords, dtype=dtype) +xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(0, Lx), dealias=dealias) +zbasis = d3.ChebyshevT(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias) + +# Fields +p = dist.Field(name='p', bases=(xbasis,zbasis)) +b = dist.Field(name='b', bases=(xbasis,zbasis)) +u = dist.VectorField(coords, name='u', bases=(xbasis,zbasis)) +tau_p = dist.Field(name='tau_p') +tau_b1 = dist.Field(name='tau_b1', bases=xbasis) +tau_b2 = dist.Field(name='tau_b2', bases=xbasis) +tau_u1 = dist.VectorField(coords, name='tau_u1', bases=xbasis) +tau_u2 = dist.VectorField(coords, name='tau_u2', bases=xbasis) + +# Substitutions +kappa = (Rayleigh * Prandtl)**(-1/2) +nu = (Rayleigh / Prandtl)**(-1/2) +x, z = dist.local_grids(xbasis, zbasis) +ex, ez = coords.unit_vector_fields(dist) +lift_basis = zbasis.derivative_basis(1) +lift = lambda A: d3.Lift(A, lift_basis, -1) +grad_u = d3.grad(u) + ez*lift(tau_u1) # First-order reduction +grad_b = d3.grad(b) + ez*lift(tau_b1) # First-order reduction + +# Problem +# First-order form: "div(f)" becomes "trace(grad_f)" +# First-order form: "lap(f)" becomes "div(grad_f)" +problem = d3.IVP([p, b, u, tau_p, tau_b1, tau_b2, tau_u1, tau_u2], namespace=locals()) +problem.add_equation("trace(grad_u) + tau_p = 0") +problem.add_equation("dt(b) - kappa*div(grad_b) + lift(tau_b2) = - u@grad(b)") +problem.add_equation("dt(u) - nu*div(grad_u) + grad(p) - b*ez + lift(tau_u2) = - u@grad(u)") +problem.add_equation("b(z=0) = Lz") +problem.add_equation("u(z=0) = 0") +problem.add_equation("b(z=Lz) = 0") +problem.add_equation("u(z=Lz) = 0") +problem.add_equation("integ(p) = 0") # Pressure gauge + +# Solver +solver = problem.build_solver(timestepper) +solver.stop_sim_time = stop_sim_time + +# Initial conditions +b.fill_random('g', seed=42, distribution='normal', scale=1e-3) # Random noise +b['g'] *= z * (Lz - z) # Damp noise at walls +b['g'] += Lz - z # Add linear background + +# Analysis +snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.25, max_writes=50) +snapshots.add_task(b, name='buoyancy') +snapshots.add_task(-d3.div(d3.skew(u)), name='vorticity') + +# Main loop +try: + logger.info('Starting main loop') + while solver.proceed: + solver.step(timestep) + if (solver.iteration-1) % 10 == 0: + logger.info('Iteration=%i, Time=%e, dt=%e' %(solver.iteration, solver.sim_time, timestep)) +except: + logger.error('Exception raised, triggering end of main loop.') + raise +finally: + solver.log_stats() diff --git a/pySDC/playgrounds/dedalus/sdc.py b/pySDC/playgrounds/dedalus/sdc.py new file mode 100644 index 0000000000..15895c5670 --- /dev/null +++ b/pySDC/playgrounds/dedalus/sdc.py @@ -0,0 +1,616 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +SDC time-integrator for Dedalus +""" +# Python imports +import numpy as np +from scipy.linalg import blas +from collections import deque + +# Dedalus import +from dedalus.core.system import CoeffSystem +from dedalus.tools.array import apply_sparse, csr_matvecs + +# QMat imports +from qmat.qcoeff.collocation import Collocation +from qmat import genQDeltaCoeffs + +# ----------------------------------------------------------------------------- +# Main SDC parameters +# ----------------------------------------------------------------------------- +DEFAULT = { + 'nNodes': 4, + 'nSweeps': 3, + 'nodeDistr': 'LEGENDRE', + 'quadType': 'RADAU-RIGHT', + 'implSweep': 'BE', + 'explSweep': 'FE', + 'initSweep': 'COPY', + 'initSweepQDelta': 'BE', + 'forceProl': False, + } + +PARAMS = { + ('nNodes', '-sM', '--sdcNNodes'): + dict(help='number of nodes for SDC', type=int, + default=DEFAULT['nNodes']), + ('quadType', '-sqt', '--sdcQuadType'): + dict(help='quadrature type for SDC', + default=DEFAULT['quadType']), + ('nodeDistr', '-snd', '--sdcNodeDistr'): + dict(help='node distribution for SDC', + default=DEFAULT['nodeDistr']), + ('nSweeps', '-sK', '--sdcNSweeps'): + dict(help='number of sweeps for SDC', type=int, + default=DEFAULT['nSweeps']), + ('initSweep', '-si', '--sdcInitSweep'): + dict(help='initial sweep to get initialized nodes values', + default=DEFAULT['initSweep']), + ('initSweepQDelta', '-siq', '--sdcInitSweepQDelta'): + dict(help='QDelta matrix used with initSweep=QDELTA', + default=DEFAULT['initSweepQDelta']), + ('implSweep', '-sis', '--sdcImplSweep'): + dict(help='type of QDelta matrix for implicit sweep', + default=DEFAULT['implSweep']), + ('explSweep', '-ses', '--sdcExplSweep'): + dict(help='type of QDelta matrix for explicit sweep', + default=DEFAULT['explSweep']), + ('forceProl', '-sfp', '--sdcForceProl'): + dict(help='if specified force the prolongation stage ' + '(ignored for quadType=GAUSS or RADAU-LEFT)', + action='store_true') + } + +# Printing function +def sdcInfos(nNodes, quadType, nodeDistr, nSweeps, + implSweep, explSweep, initSweep, forceProl, + **kwargs): + return f""" +-- nNodes : {nNodes} +-- quadType : {quadType} +-- nodeDistr : {nodeDistr} +-- nSweeps : {nSweeps} +-- implSweep : {implSweep} +-- explSweep : {explSweep} +-- initSweep : {initSweep} +-- forceProl : {forceProl} +""".strip() + +# ----------------------------------------------------------------------------- +# Base class implementation +# ----------------------------------------------------------------------------- +class IMEXSDCCore(object): + + # Initialize parameters with default values + nSweeps:int = DEFAULT['nSweeps'] + nodeType:str = DEFAULT['nodeDistr'] + quadType:str = DEFAULT['quadType'] + implSweep = DEFAULT['implSweep'] + explSweep = DEFAULT['explSweep'] + initSweep = DEFAULT['initSweep'] + initSweepQDelta = DEFAULT['initSweepQDelta'] + forceProl = DEFAULT['forceProl'] + + # Collocation method attributes + coll = Collocation( + nNodes=DEFAULT['nNodes'], nodeType=nodeType, quadType=quadType) + nodes, weights, Q = coll.genCoeffs() + # IMEX SDC attributes, QDelta matrices are 3D with shape (K, M, M) + QDeltaI, dtauI = genQDeltaCoeffs( + implSweep, dTau=True, nSweeps=nSweeps, + Q=Q, nodes=nodes, nNodes=DEFAULT['nNodes'], + nodeType=nodeType, quadType=quadType) + QDeltaE, dtauE = genQDeltaCoeffs( + explSweep, dTau=True, nSweeps=nSweeps, + Q=Q, nodes=nodes, nNodes=DEFAULT['nNodes'], + nodeType=nodeType, quadType=quadType) + QDelta0 = genQDeltaCoeffs( + initSweepQDelta, + Q=Q, nodes=nodes, nNodes=DEFAULT['nNodes'], + nodeType=nodeType, quadType=quadType) + + diagonal = np.all([np.diag(np.diag(qD)) == qD for qD in QDeltaI]) + diagonal *= np.all([np.diag(np.diag(qD)) == 0 for qD in QDeltaE]) + if initSweep == "QDelta": + diagonal *= np.all(np.diag(np.diag(QDelta0)) == QDelta0) + + # Default attributes, used later as instance attributes + # => should be defined in inherited class + dt = None + axpy = None + + @classmethod + def setParameters(cls, nNodes=None, nodeType=None, quadType=None, + implSweep=None, explSweep=None, initSweep=None, + initSweepQDelta=None, nSweeps=None, forceProl=None): + + # Get non-changing parameters + keepNNodes = nNodes is None + keepNodeDistr = nodeType is None + keepQuadType = quadType is None + keepImplSweep = implSweep is None + keepExplSweep = explSweep is None + keepNSweeps = nSweeps is None + keepInitSweepQDelta = initSweepQDelta is None + + # Set parameter values + nNodes = len(cls.nodes) if keepNNodes else nNodes + nodeType = cls.nodeType if keepNodeDistr else nodeType + quadType = cls.quadType if keepQuadType else quadType + implSweep = cls.implSweep if keepImplSweep else implSweep + explSweep = cls.explSweep if keepExplSweep else explSweep + nSweeps = cls.nSweeps if keepNSweeps else nSweeps + initSweepQDelta = cls.initSweepQDelta if keepInitSweepQDelta else initSweepQDelta + + # Determine updated parts + updateNode = (not keepNNodes) or (not keepNodeDistr) or (not keepQuadType) + updateQDeltaI = (not keepImplSweep) or updateNode or (not keepNSweeps) + updateQDeltaE = (not keepExplSweep) or updateNode or (not keepNSweeps) + updateQDelta0 = (not keepInitSweepQDelta) or updateNode + + # Update parameters + if updateNode: + cls.coll = Collocation( + nNodes=nNodes, nodeType=nodeType, quadType=quadType) + cls.nodes, cls.weights, cls.Q = cls.coll.genCoeffs() + cls.nodeType, cls.quadType = nodeType, quadType + if updateQDeltaI: + cls.QDeltaI, cls.dtauI = genQDeltaCoeffs( + implSweep, dTau=True, nSweeps=nSweeps, + Q=cls.Q, nodes=cls.nodes, nNodes=nNodes, + nodeType=nodeType, quadType=quadType) + cls.implSweep = implSweep + if updateQDeltaE: + cls.QDeltaE, cls.dtauE = genQDeltaCoeffs( + explSweep, dTau=True, nSweeps=nSweeps, + Q=cls.Q, nodes=cls.nodes, nNodes=nNodes, + nodeType=nodeType, quadType=quadType) + cls.explSweep = explSweep + if updateQDelta0: + cls.QDelta0 = genQDeltaCoeffs( + initSweepQDelta, + Q=cls.Q, nodes=cls.nodes, nNodes=nNodes, + nodeType=nodeType, quadType=quadType) + + # Eventually update additional parameters + if forceProl is not None: cls.forceProl = forceProl + if initSweep is not None: cls.initSweep = initSweep + if not keepNSweeps: + cls.nSweeps = nSweeps + + diagonal = np.all([np.diag(np.diag(qD)) == qD for qD in cls.QDeltaI]) + diagonal *= np.all([np.diag(np.diag(qD)) == 0 for qD in cls.QDeltaE]) + if cls.initSweep == "QDELTA": + diagonal *= np.all(np.diag(np.diag(cls.QDelta0)) == cls.QDelta0) + cls.diagonal = diagonal + + # ------------------------------------------------------------------------- + # Class properties + # ------------------------------------------------------------------------- + @classmethod + def getMaxOrder(cls): + # TODO : adapt to non-LEGENDRE node distributions + M = len(cls.nodes) + return 2*M if cls.quadType == 'GAUSS' else \ + 2*M-1 if cls.quadType.startswith('RADAU') else \ + 2*M-2 # LOBATTO + + @classmethod + def getInfos(cls): + return sdcInfos( + len(cls.nodes), cls.quadType, cls.nodeType, cls.nSweeps, + cls.implSweep, cls.explSweep, cls.initSweep, cls.forceProl) + + @classmethod + def getM(cls): + return len(cls.nodes) + + @classmethod + def rightIsNode(cls): + return np.isclose(cls.nodes[-1], 1.0) + + @classmethod + def leftIsNode(cls): + return np.isclose(cls.nodes[0], 0.0) + + @classmethod + def doProlongation(cls): + return not cls.rightIsNode or cls.forceProl + + +# ----------------------------------------------------------------------------- +# Dedalus based IMEX timeintegrator class +# ----------------------------------------------------------------------------- +class SpectralDeferredCorrectionIMEX(IMEXSDCCore): + + steps = 1 + + # ------------------------------------------------------------------------- + # Instance methods + # ------------------------------------------------------------------------- + def __init__(self, solver): + + # Store class attributes as instance attributes + self.infos = self.getInfos() + + # Store solver as attribute + self.solver = solver + self.subproblems = [sp for sp in solver.subproblems if sp.size] + self.stages = self.M # need this for solver.log_stats() + + # Create coefficient systems for steps history + c = lambda: CoeffSystem(solver.subproblems, dtype=solver.dtype) + self.MX0, self.RHS = c(), c() + self.LX = deque([[c() for _ in range(self.M)] for _ in range(2)]) + self.F = deque([[c() for _ in range(self.M)] for _ in range(2)]) + + if not self.leftIsNode and self.initSweep == "QDelta": + self.F0, self.LX0 = c(), c() + + # Attributes + self.axpy = blas.get_blas_funcs('axpy', dtype=solver.dtype) + self.dt = None + self.firstEval = True + + @property + def M(self): + return len(self.nodes) + + @property + def rightIsNode(self): + return np.allclose(self.nodes[-1], 1.0) + + @property + def leftIsNode(self): + return np.allclose(self.nodes[0], 0.0) + + @property + def doProlongation(self): + return not self.rightIsNode or self.forceProl + + def _computeMX0(self, MX0): + """ + Compute MX0 term used in RHS of both initStep and sweep methods + + Update the MX0 attribute of the timestepper object. + """ + state = self.solver.state + # Assert coefficient space + self._requireStateCoeffSpace(state) + # Compute and store MX0 + MX0.data.fill(0) + for sp in self.subproblems: + spX = sp.gather_inputs(state) + apply_sparse(sp.M_min, spX, axis=0, out=MX0.get_subdata(sp)) + + def _updateLHS(self, dt, init=False): + """Update LHS and LHS solvers for each subproblem + + Parameters + ---------- + dt : float + Time-step for the updated LHS. + init : bool, optional + Wether or not initialize the LHS_solvers attribute for each + subproblem. The default is False. + """ + # Attribute references + qI = self.QDeltaI + solver = self.solver + + # Update LHS and LHS solvers for each subproblems + for sp in solver.subproblems: + if init: + # Eventually 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): + 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) + if self.initSweep == "QDELTA": + raise NotImplementedError() + + def _evalLX(self, LX): + """ + Evaluate LX using the current state, and store it + + Parameters + ---------- + LX : dedalus.core.system.CoeffSystem + Where to store the evaluated fields. + + Returns + ------- + None. + + """ + # Attribute references + solver = self.solver + # Assert coefficient spacec + self._requireStateCoeffSpace(solver.state) + # Evaluate matrix vector product and store + for sp in solver.subproblems: + spX = sp.gather_inputs(solver.state) + apply_sparse(sp.L_min, spX, axis=0, out=LX.get_subdata(sp)) + + def _evalF(self, F, time, dt, wall_time): + """ + Evaluate the F operator from the current solver state + + Note + ---- + After evaluation, state fields are left in grid space + + Parameters + ---------- + time : float + Time of evaluation. + F : dedalus.core.system.CoeffSystem + Where to store the evaluated fields. + dt : float + Current time step. + wall_time : float + Current wall time. + """ + + solver = self.solver + # Evaluate non linear term on current state + t0 = solver.sim_time + solver.sim_time = time + if self.firstEval: + solver.evaluator.evaluate_scheduled( + wall_time=wall_time, timestep=dt, sim_time=time, + iteration=solver.iteration) + self.firstEval = False + else: + solver.evaluator.evaluate_group('F') + # Store F evaluation + for sp in solver.subproblems: + sp.gather_outputs(solver.F, out=F.get_subdata(sp)) + # Put back initial solver simulation time + solver.sim_time = t0 + + def _solveAndStoreState(self, k, m): + """ + Solve LHS * X = RHS using the LHS associated to a given node, + and store X into the solver state. + It uses the current RHS attribute of the object. + + Parameters + ---------- + k : int + Sweep index (0 for the first sweep). + m : int + Index of the nodes. + """ + # Attribute references + solver = self.solver + RHS = self.RHS + + self._presetStateCoeffSpace(solver.state) + + # Solve and store for each subproblem + for sp in solver.subproblems: + # Slice out valid subdata, skipping invalid components + spRHS = RHS.get_subdata(sp) + spX = sp.LHS_solvers[k][m].solve(spRHS) # CREATES TEMPORARY + sp.scatter_inputs(spX, solver.state) + + def _requireStateCoeffSpace(self, state): + """Transform current state fields in coefficient space. + If already in coefficient space, doesn't do anything.""" + for field in state: + field.require_coeff_space() + + def _presetStateCoeffSpace(self, state): + """Allow to write fields in coefficient space into current state + fields, without transforming current state in coefficient space.""" + for field in state: + field.preset_layout('c') + + def _initSweep(self): + """ + Initialize node terms for one given time-step + + Parameters + ---------- + iType : str, optional + Type of initialization, can be : + - iType="QDELTA" : use QDelta[I,E] for coarse time integration. + - iType="COPY" : just copy the values from the initial solution. + - iType="FNO" : use an FNO-ML model to predict values (incoming ...) + """ + # Attribute references + tau, qI, qE = self.nodes, self.QDeltaI, self.QDeltaE + solver = self.solver + t0, dt, wall_time = solver.sim_time, self.dt, self.wall_time + RHS, MX0, Fk, LXk = self.RHS, self.MX0, self.F[0], self.LX[0] + if not self.leftIsNode and self.initSweep == "QDELTA": + F0, LX0 = self.F0, self.LX0 + axpy = self.axpy + + if self.initSweep == 'QDELTA': + + # Eventual initial field evaluation + if not self.leftIsNode: + if np.any(self.dtauE): + self._evalF(F0, t0, dt, wall_time) + if np.any(self.dtauI): + self._evalLX(LX0) + + # Loop on all quadrature nodes + for m in range(self.M): + + # Build RHS + if RHS.data.size: + + # Initialize with MX0 term + np.copyto(RHS.data, MX0.data) + + # Add eventual initial field evaluation + if not self.leftIsNode: + if self.dtauE[m]: + axpy(a=dt*self.dtauE[m], x=F0.data, y=RHS.data) + if self.dtauI[m]: + axpy(a=-dt*self.dtauI[m], x=LX0.data, y=RHS.data) + + # Add F and LX terms (already computed) + for i in range(m): + axpy(a=dt*qE[m, i], x=Fk[i].data, y=RHS.data) + axpy(a=-dt*qI[m, i], x=LXk[i].data, y=RHS.data) + + # Solve system and store node solution in solver state + self._solveAndStoreState(m) + + # Evaluate and store LX with current state + self._evalLX(LXk[m]) + + # Evaluate and store F(X, t) with current state + self._evalF(Fk[m], t0+dt*tau[m], dt, wall_time) + + elif self.initSweep == 'COPY': + self._evalLX(LXk[0]) + self._evalF(Fk[0], t0, dt, wall_time) + for m in range(1, self.M): + np.copyto(LXk[m].data, LXk[0].data) + np.copyto(Fk[m].data, Fk[0].data) + + else: + raise NotImplementedError(f'initSweep={self.initSweep}') + + def _sweep(self, k): + """Perform a sweep for the current time-step""" + # Attribute references + tau, qI, qE, q = self.nodes, self.QDeltaI, self.QDeltaE, self.Q + solver = self.solver + t0, dt, wall_time = solver.sim_time, self.dt, self.wall_time + RHS, MX0 = self.RHS, self.MX0 + Fk, LXk, Fk1, LXk1 = self.F[0], self.LX[0], self.F[1], self.LX[1] + axpy = self.axpy + + # Loop on all quadrature nodes + for m in range(self.M): + + # Build RHS + if RHS.data.size: + + # Initialize with MX0 term + np.copyto(RHS.data, MX0.data) + + # Add quadrature terms + for i in range(self.M): + axpy(a=dt*q[m, i], x=Fk[i].data, y=RHS.data) + axpy(a=-dt*q[m, i], x=LXk[i].data, y=RHS.data) + + if not self.diagonal: + # Add F and LX terms from iteration k+1 and previous nodes + for i in range(m): + axpy(a=dt*qE[k, m, i], x=Fk1[i].data, y=RHS.data) + axpy(a=-dt*qI[k, m, i], x=LXk1[i].data, y=RHS.data) + # Add F and LX terms from iteration k and previous nodes + for i in range(m): + axpy(a=-dt*qE[k, m, i], x=Fk[i].data, y=RHS.data) + axpy(a=dt*qI[k, m, i], x=LXk[i].data, y=RHS.data) + + # Add LX terms from iteration k+1 and current nodes + axpy(a=dt*qI[k, m, m], x=LXk[m].data, y=RHS.data) + + # Solve system and store node solution in solver state + self._solveAndStoreState(k, m) + + # Avoid non necessary RHS evaluations work + if not self.forceProl and k == self.nSweeps-1: + if self.diagonal: + continue + elif m == self.M-1: + continue + + # Evaluate and store LX with current state + self._evalLX(LXk1[m]) + # Evaluate and store F(X, t) with current state + self._evalF(Fk1[m], t0+dt*tau[m], dt, wall_time) + + # Inverse position for iterate k and k+1 in storage + # ie making the new evaluation the old for next iteration + self.F.rotate() + self.LX.rotate() + + def _prolongation(self): + """Compute prolongation (needed if last node != 1)""" + # Attribute references + solver = self.solver + w, dt = self.weights, self.dt + RHS, MX0, Fk, LXk = self.RHS, self.MX0, self.F[0], self.LX[0] + axpy = self.axpy + + # Build RHS + if RHS.data.size: + # Initialize with MX0 term + np.copyto(RHS.data, MX0.data) + # Add quadrature terms + for i in range(self.M): + axpy(a=dt*w[i], x=Fk[i].data, y=RHS.data) + axpy(a=-dt*w[i], x=LXk[i].data, y=RHS.data) + + self._presetStateCoeffSpace(solver.state) + + # Solve and store for each subproblem + for sp in solver.subproblems: + # Slice out valid subdata, skipping invalid components + spRHS = RHS.get_subdata(sp)[:sp.LHS.shape[0]] + # Solve using LHS of the node + spX = sp.M_solver.solve(spRHS) + # Make output buffer including invalid components for scatter + spX2 = np.zeros( + (sp.pre_right.shape[0], len(sp.subsystems)), + dtype=spX.dtype) + # Store X to state_fields + csr_matvecs(sp.pre_right, spX, spX2) + sp.scatter(spX2, solver.state) + + def step(self, dt, wall_time): + """ + Compute the next time-step solution using the time-stepper method, + and modify to state field of the solver + + Note + ---- + State fields should be left in grid space after at the end of the step. + + Parameters + ---------- + dt : float + Length of the current time-step. + wall_time : float + Current wall time for the simulation. + """ + self.wall_time = wall_time + + # Initialize and/or update LHS terms, depending on dt + if dt != self.dt: + self._updateLHS(dt, init=self.dt is None) + self.dt = dt + + # Compute MX0 for the whole time step + self._computeMX0(self.MX0) + + # Initialize node values + self._initSweep() + + # Performs sweeps + for k in range(self.nSweeps): + self._sweep(k) + + # Compute prolongation if needed + if self.doProlongation: + self._prolongation() + + # Update simulation time and reset evaluation tag + self.solver.sim_time += dt + self.firstEval = True diff --git a/pySDC/playgrounds/dedalus/sweeper.py b/pySDC/playgrounds/dedalus/sweeper.py new file mode 100644 index 0000000000..13c41a337e --- /dev/null +++ b/pySDC/playgrounds/dedalus/sweeper.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Sweeper class for dedalus +""" +import numpy as np + +from problem import DedalusProblem +from pySDC.core.sweeper import Sweeper + + +class DedalusSweeperIMEX(Sweeper): + + def __init__(self, params): + if 'QI' not in params: params['QI'] = 'IE' + if 'QE' not in params: params['QE'] = 'EE' + # call parent's initialization routine + super().__init__(params) + # IMEX integration matrices + self.QI = self.get_Qdelta_implicit(qd_type=self.params.QI) + self.QE = self.get_Qdelta_explicit(qd_type=self.params.QE) + + def predict(self): + """Copy for now ...""" + + L = self.level + t0, dt = L.time, L.dt + qI = self.QI[1:, 1:] + + P:DedalusProblem = L.prob + assert type(P) == DedalusProblem + assert self.coll.num_nodes == P.M + + P.firstEval = True + P.solver.sim_time = t0 + P.solver.dt = dt + + # for f0, f in zip(L.u[0], P.solver.state): + # np.copyto(f.data, f0.data) + + # for f in P.solver.state: + # f.change_scales(f.domain.dealias) + # P.solver.evaluator.require_grid_space(P.solver.state) + # P.solver.evaluator.require_coeff_space(P.solver.state) + + P.updateLHS(dt, qI) + P.computeMX0() + + Fk, LXk = P.F[0], P.LX[0] + P.evalLX(LXk[0]) + P.evalF(t0, Fk[0]) + for m in range(1, P.M): + np.copyto(LXk[m].data, LXk[0].data) + np.copyto(Fk[m].data, Fk[0].data) + + # additional stuff for pySDC, which stores solution at each nodes + for m in range(1, self.coll.num_nodes + 1): + L.u[m] = P.stateCopy() + + # indicate that this level is now ready for sweeps + L.status.unlocked = True + L.status.updated = True + + + def update_nodes(self): + """ + TODO + """ + # get current level and problem description + L = self.level + # only if the level has been touched before + assert L.status.unlocked + + t0, dt = L.time, L.dt + tau, qI, qE, q = self.coll.nodes, self.QI[1:, 1:], self.QE[1:, 1:], self.coll.Qmat[1:, 1:] + + P:DedalusProblem = L.prob + assert type(P) == DedalusProblem + + # get number of collocation nodes for easier access + M = self.coll.num_nodes + assert M == P.M + + # Attribute references + RHS, MX0 = P.RHS, P.MX0 + Fk, LXk, Fk1, LXk1 = P.F[0], P.LX[0], P.F[1], P.LX[1] + axpy = P.axpy + + # Loop on all quadrature nodes + for m in range(M): + + # Initialize with MX0 term + np.copyto(RHS.data, MX0.data) + + # 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) + + # Add F and LX terms from iteration k+1 + for j in range(m): + axpy(a=dt*qE[m, j], x=Fk1[j].data, y=RHS.data) + axpy(a=-dt*qI[m, j], x=LXk1[j].data, y=RHS.data) + + # Add F and LX terms from iteration k + for j in range(m): + axpy(a=-dt*qE[m, j], x=Fk[j].data, y=RHS.data) + axpy(a=dt*qI[m, j], x=LXk[j].data, y=RHS.data) + axpy(a=dt*qI[m, m], x=LXk[m].data, y=RHS.data) + + # Solve system and store node solution in solver state + P.solveAndStoreState(m) + + # Evaluate and store LX with current state + P.evalLX(LXk1[m]) + # Evaluate and store F(X, t) with current state + P.evalF(t0+dt*tau[m], Fk1[m]) + + # Update u for pySDC + L.u[m+1] = P.stateCopy() + + # Inverse position for iterate k and k+1 in storage + # ie making the new evaluation the old for next iteration + P.F.rotate() + P.LX.rotate() + + # indicate presence of new values at this level + L.status.updated = True + + + def compute_end_point(self): + """ + Compute u at the right point of the interval + + The value uend computed here is a full evaluation of the Picard formulation unless do_full_update==False + + Returns: + None + """ + + # get current level and problem description + L = self.level + t0, dt = L.time, L.dt + P:DedalusProblem = L.prob + + # check if Mth node is equal to right point and do_coll_update is false, perform a simple copy + if self.coll.right_is_node and not self.params.do_coll_update: + # a copy is sufficient + L.uend = L.u[-1] + else: + raise NotImplementedError() + # start with u0 and add integral over the full interval (using coll.weights) + L.uend = P.dtype_u(L.u[0]) + for m in range(self.coll.num_nodes): + L.uend += L.dt * self.coll.weights[m] * (L.f[m + 1].impl + L.f[m + 1].expl) + # add up tau correction of the full interval (last entry) + if L.tau[-1] is not None: + L.uend += L.tau[-1] + + P.solver.sim_time = t0 + dt + P.firstEval = True