Coverage for pySDC/tutorial/step_7/E_pySDC_with_Firedrake.py: 100%
112 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 13:41 +0000
1# ---
2# jupyter:
3# jupytext:
4# formats: py:percent
5# kernelspec:
6# display_name: Python 3
7# name: python3
8# language_info:
9# name: python
10# ---
12# %% [markdown]
13# # Part E: pySDC and Firedrake
14#
15# [Firedrake](https://github.com/firedrakeproject/firedrake) is a finite element library with similar features as
16# FEniCS. This example runs the same forced heat equation as the [FEniCS example](A_pySDC_with_FEniCS), but
17# implemented in Firedrake. The setup proceeds very much as in earlier tutorials; the interesting part is rather the
18# problem class. See
19# [pySDC/implementations/problem_classes/HeatFiredrake.py](https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/implementations/problem_classes/HeatFiredrake.py)
20# as a blueprint for how to implement problems with Firedrake in a way that pySDC understands.
21#
22# ## Single level, serial or parallel across the nodes
23#
24# SDC with the diagonal preconditioner `MIN-SR-S`, whose nodes can be solved in parallel. The space-time parallelism
25# comes from a Firedrake ensemble: with `--useMPIsweeper` and a multiple of 3 ranks, the ranks are split across the
26# 3 collocation nodes, and the ranks of each node share the spatial problem. See the
27# [Firedrake documentation on parallelism](https://www.firedrakeproject.org/firedrake/parallelism.html).
29# %%
30from pathlib import Path
32import numpy as np
33from mpi4py import MPI
36def setup(useMPIsweeper):
37 """
38 Helper routine to set up parameters
40 Returns:
41 description and controller_params parameter dictionaries
42 """
43 from pySDC.implementations.problem_classes.HeatFiredrake import Heat1DForcedFiredrake
44 from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
45 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI
46 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun
47 from pySDC.implementations.hooks.log_work import LogWork
48 from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator
50 # setup space-time parallelism via ensemble for Firedrake, see https://www.firedrakeproject.org/firedrake/parallelism.html
51 num_nodes = 3
52 ensemble = FiredrakeEnsembleCommunicator(MPI.COMM_WORLD, max([MPI.COMM_WORLD.size // num_nodes, 1]))
54 level_params = dict()
55 level_params['restol'] = 5e-10
56 level_params['dt'] = 0.2
58 step_params = dict()
59 step_params['maxiter'] = 20
61 sweeper_params = dict()
62 sweeper_params['quad_type'] = 'RADAU-RIGHT'
63 sweeper_params['num_nodes'] = num_nodes
64 sweeper_params['QI'] = 'MIN-SR-S'
65 sweeper_params['QE'] = 'PIC'
66 sweeper_params['comm'] = ensemble
68 problem_params = dict()
69 problem_params['nu'] = 0.1
70 problem_params['n'] = 128
71 problem_params['c'] = 1.0
72 problem_params['comm'] = ensemble.space_comm
74 controller_params = dict()
75 controller_params['logger_level'] = 15 if MPI.COMM_WORLD.rank == 0 else 30
76 controller_params['hook_class'] = [LogGlobalErrorPostRun, LogWork]
78 description = dict()
79 description['problem_class'] = Heat1DForcedFiredrake
80 description['problem_params'] = problem_params
81 description['sweeper_class'] = imex_1st_order_MPI if useMPIsweeper else imex_1st_order
82 description['sweeper_params'] = sweeper_params
83 description['level_params'] = level_params
84 description['step_params'] = step_params
86 return description, controller_params
89# %% [markdown]
90# ## Three levels
91#
92# Multilevel SDC, serial, coarsened in space from 128 to 32 to 4 elements, with interpolation of the right-hand side.
95# %%
96def setup_ML():
97 """
98 Helper routine to set up parameters
100 Returns:
101 description and controller_params parameter dictionaries
102 """
103 from pySDC.implementations.problem_classes.HeatFiredrake import Heat1DForcedFiredrake
104 from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
105 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI
106 from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrake
107 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun
108 from pySDC.implementations.hooks.log_work import LogWork
109 from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator
111 level_params = dict()
112 level_params['restol'] = 5e-10
113 level_params['dt'] = 0.2
115 step_params = dict()
116 step_params['maxiter'] = 20
118 sweeper_params = dict()
119 sweeper_params['quad_type'] = 'RADAU-RIGHT'
120 sweeper_params['num_nodes'] = 3
121 sweeper_params['QI'] = 'MIN-SR-S'
122 sweeper_params['QE'] = 'PIC'
124 problem_params = dict()
125 problem_params['nu'] = 0.1
126 problem_params['n'] = [128, 32, 4]
127 problem_params['c'] = 1.0
129 base_transfer_params = dict()
130 base_transfer_params['finter'] = True
132 controller_params = dict()
133 controller_params['logger_level'] = 15 if MPI.COMM_WORLD.rank == 0 else 30
134 controller_params['hook_class'] = [LogGlobalErrorPostRun, LogWork]
136 description = dict()
137 description['problem_class'] = Heat1DForcedFiredrake
138 description['problem_params'] = problem_params
139 description['sweeper_class'] = imex_1st_order
140 description['sweeper_params'] = sweeper_params
141 description['level_params'] = level_params
142 description['step_params'] = step_params
143 description['space_transfer_class'] = MeshToMeshFiredrake
144 description['base_transfer_params'] = base_transfer_params
146 return description, controller_params
149# %% [markdown]
150# ## Running it
151#
152# The run prints the error and the work: SDC iterations, solver setups, solves and right-hand side evaluations on the
153# finest level, and checks them against what we got last time.
156# %%
157def runHeatFiredrake(useMPIsweeper=False, ML=False):
158 """
159 Run the example defined by the above parameters
160 """
161 from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
162 from pySDC.helpers.stats_helper import get_sorted
164 Tend = 1.0
165 t0 = 0.0
167 if ML:
168 assert not useMPIsweeper, 'MPI parallel diagonal SDC and ML SDC are not compatible at the moment'
169 description, controller_params = setup_ML()
170 else:
171 description, controller_params = setup(useMPIsweeper)
173 controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description)
175 # get initial values
176 P = controller.MS[0].levels[0].prob
177 uinit = P.u_exact(0.0)
179 # call main function to get things done...
180 uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
182 # see what we get
183 error = get_sorted(stats, type='e_global_post_run')
184 work_solver_setup = get_sorted(stats, type='work_solver_setup')
185 work_solves = get_sorted(stats, type='work_solves')
186 work_rhs = get_sorted(stats, type='work_rhs')
187 niter = get_sorted(stats, type='niter')
189 tot_iter = np.sum([me[1] for me in niter])
190 tot_solver_setup = np.sum([me[1] for me in work_solver_setup])
191 tot_solves = np.sum([me[1] for me in work_solves])
192 tot_rhs = np.sum([me[1] for me in work_rhs])
194 time_rank = description["sweeper_params"]["comm"].rank if useMPIsweeper else 0
195 print(
196 f'Finished with error {error[0][1]:.2e}. Used {tot_iter} SDC iterations, with {tot_solver_setup} solver setups, {tot_solves} solves and {tot_rhs} right hand side evaluations on the finest level of time task {time_rank}.'
197 )
199 # the results the website shows, from the first rank
200 if MPI.COMM_WORLD.rank == 0:
201 timing = get_sorted(stats, type='timing_run')[0][1]
202 variant = 'three-level SDC' if ML else ('SDC parallel across the nodes' if useMPIsweeper else 'serial SDC')
203 Path('data').mkdir(parents=True, exist_ok=True)
204 with open('data/step_7_E_out.txt', 'a') as file:
205 file.write(
206 f'{variant:30s} error {error[0][1]:.2e}, {tot_iter:2d} iterations, {tot_solves:3d} solves on the '
207 f'finest level, time to solution {timing:.2f} s\n'
208 )
210 # do tests that we got the same as last time
211 n_nodes = 1 if useMPIsweeper else description['sweeper_params']['num_nodes']
212 assert error[0][1] < 2e-7
213 assert tot_iter == (10 if ML else 29)
214 assert tot_solver_setup == n_nodes
215 assert tot_solves == n_nodes * tot_iter
216 assert tot_rhs == n_nodes * tot_iter + (n_nodes + 1) * len(niter)
219# %% [markdown]
220# The script runs in three different ways:
221#
222# - `python E_pySDC_with_Firedrake.py` for single-level serial SDC,
223# - `mpiexec -np 3 python E_pySDC_with_Firedrake.py --useMPIsweeper` for single-level SDC, parallel across the nodes,
224# - `python E_pySDC_with_Firedrake.py --ML` for three-level serial SDC.
225#
226# Multilevel SDC reduces the number of SDC iterations quite a bit, but not necessarily the time to solution, partly
227# because more solvers are constructed on the coarse levels, and the parallel variant does the work of one node per
228# rank. The results below show both; mind that they come from a CI machine, not from a dedicated one. We do not claim
229# to have found the best parameters: this is just an example of how to use it.
231# %%
232if __name__ == "__main__":
233 from argparse import ArgumentParser
235 parser = ArgumentParser()
236 parser.add_argument(
237 '--ML',
238 help='Whether you want to run multi-level',
239 default=False,
240 required=False,
241 action='store_const',
242 const=True,
243 )
244 parser.add_argument(
245 '--useMPIsweeper',
246 help='Whether you want to use MPI parallel diagonal SDC',
247 default=False,
248 required=False,
249 action='store_const',
250 const=True,
251 )
253 args = parser.parse_args()
255 runHeatFiredrake(**vars(args))
257# %% [markdown]
258# Firedrake does not run in the browser, nor in the environment this website is built in. Our CI runs all three
259# variants in a Firedrake container, and checks their errors and work counts. These are its results, in the run that
260# built this page:
261#
262# :::{literalinclude} /../../data_firedrake/step_7_E_out.txt
263# :language: text
264# :::