Coverage for pySDC/tutorial/step_7/F_pySDC_with_Gusto.py: 91%
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 F: pySDC and Gusto
14#
15# [Gusto](https://github.com/firedrakeproject/gusto) is a toolkit for geophysical simulations that uses
16# [Firedrake](https://github.com/firedrakeproject/firedrake) for the spatial discretization. This example is an
17# adaptation of the Williamson 5 test case as implemented in
18# [Gusto](https://github.com/firedrakeproject/gusto/blob/main/examples/shallow_water/williamson_5.py): the shallow
19# water equations, with flow over a mountain (Test Case 5 of Williamson et al., 1992, *A standard test set for
20# numerical approximations to the shallow water equations in spherical geometry*, JCP). It may take a considerable
21# amount of time.
22#
23# This coupling works slightly differently from the other examples of this step, as time stepping is part of Gusto.
24# The aim of the coupling is not a spatial discretization, but to use the equations that are implemented in Gusto. A
25# Gusto equation includes the basic form of the equation set, but a crucial part is to modify terms in the discretized
26# equations with spatial methods, such as upwinding schemes.
27#
28# We get the finished equation set into pySDC by setting up pySDC as a time discretization for Gusto: a pySDC
29# description and controller parameters as usual, passed to the pySDC time discretization, which then goes into a
30# Gusto time stepper. During this instantiation, the equation, and the residual used for solving systems, is modified
31# with all the spatial methods. Afterwards, there are two choices:
32#
33# - use the Gusto time stepper for running, or
34# - take its `.scheme.controller`, which is the pySDC controller, and run with pySDC.
35#
36# That is why a Gusto time stepper is needed even to run with pySDC: the spatial methods modify the residual of the
37# equations when it is created. Once they have, we can continue in Gusto or in pySDC.
39# %%
40import firedrake as fd
41from pySDC.helpers.pySDC_as_gusto_time_discretization import pySDC_integrator
42from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator
43from gusto import SDC, BackwardEuler
44from gusto.core.labels import implicit, time_derivative
45from gusto.core.logging import logger, INFO
47logger.setLevel(INFO)
50from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
51from firedrake import SpatialCoordinate, as_vector, pi, sqrt, min_value, Function
52from gusto import (
53 Domain,
54 IO,
55 OutputParameters,
56 DGUpwind,
57 ShallowWaterParameters,
58 ShallowWaterEquations,
59 Sum,
60 lonlatr_from_xyz,
61 GeneralIcosahedralSphereMesh,
62 ZonalComponent,
63 MeridionalComponent,
64 RelativeVorticity,
65 Timestepper,
66)
68williamson_5_defaults = {
69 'ncells_per_edge': 12, # number of cells per icosahedron edge
70 'dt': 900.0,
71 'tmax': 50.0 * 24.0 * 60.0 * 60.0, # 50 days
72 'dumpfreq': 10, # output every <dumpfreq> steps
73 'dirname': 'williamson_5', # results will go into ./results/<dirname>
74 'time_parallelism': False, # use parallel diagonal SDC or not
75 'QI': 'MIN-SR-S', # implicit preconditioner
76 'M': '3', # number of collocation nodes
77 'kmax': '5', # use fixed number of iteration up to this value
78 'use_pySDC': True, # whether to use pySDC for time integration
79 'use_adaptivity': True, # whether to use adaptive step size selection
80 'Nlevels': 1, # number of levels in SDC
81 'logger_level': 15, # pySDC logger level
82}
84# %% [markdown]
85# ## Williamson 5
86#
87# The whole test case, with the Gusto SDC implementation or the pySDC-Gusto coupling, optionally with adaptivity,
88# several levels in space, or space-time parallelism.
91# %%
92def williamson_5(
93 ncells_per_edge=williamson_5_defaults['ncells_per_edge'],
94 dt=williamson_5_defaults['dt'],
95 tmax=williamson_5_defaults['tmax'],
96 dumpfreq=williamson_5_defaults['dumpfreq'],
97 dirname=williamson_5_defaults['dirname'],
98 time_parallelism=williamson_5_defaults['time_parallelism'],
99 QI=williamson_5_defaults['QI'],
100 M=williamson_5_defaults['M'],
101 kmax=williamson_5_defaults['kmax'],
102 use_pySDC=williamson_5_defaults['use_pySDC'],
103 use_adaptivity=williamson_5_defaults['use_adaptivity'],
104 Nlevels=williamson_5_defaults['Nlevels'],
105 logger_level=williamson_5_defaults['logger_level'],
106 mesh=None,
107 _ML_is_setup=True,
108):
109 """
110 Run the Williamson 5 test case.
112 Args:
113 ncells_per_edge (int): number of cells per icosahedron edge
114 dt (float): Initial step size
115 tmax (float): Time to integrate to
116 dumpfreq (int): Output every <dumpfreq> time steps
117 dirname (str): Output will go into ./results/<dirname>
118 time_parallelism (bool): True for parallel SDC, False for serial
119 M (int): Number of collocation nodes
120 kmax (int): Max number of SDC iterations
121 use_pySDC (bool): Use pySDC as Gusto time integrator or Gusto SDC implementation
122 Nlevels (int): Number of SDC levels
123 logger_level (int): Logger level
124 """
125 if not use_pySDC and use_adaptivity:
126 raise NotImplementedError('Adaptive step size selection not yet implemented in Gusto')
127 if not use_pySDC and Nlevels > 1:
128 raise NotImplementedError('Multi-level SDC not yet implemented in Gusto')
129 if time_parallelism and Nlevels > 1:
130 raise NotImplementedError('Multi-level SDC does not work with MPI parallel sweeper yet')
132 # ------------------------------------------------------------------------ #
133 # Parameters for test case
134 # ------------------------------------------------------------------------ #
136 radius = 6371220.0 # planetary radius (m)
137 mean_depth = 5960 # reference depth (m)
138 g = 9.80616 # acceleration due to gravity (m/s^2)
139 u_max = 20.0 # max amplitude of the zonal wind (m/s)
140 mountain_height = 2000.0 # height of mountain (m)
141 R0 = pi / 9.0 # radius of mountain (rad)
142 lamda_c = -pi / 2.0 # longitudinal centre of mountain (rad)
143 phi_c = pi / 6.0 # latitudinal centre of mountain (rad)
145 # ------------------------------------------------------------------------ #
146 # Our settings for this set up
147 # ------------------------------------------------------------------------ #
149 element_order = 1
151 # ------------------------------------------------------------------------ #
152 # Set up model objects
153 # ------------------------------------------------------------------------ #
155 # parallelism
156 if time_parallelism:
157 ensemble_comm = FiredrakeEnsembleCommunicator(fd.COMM_WORLD, fd.COMM_WORLD.size // M)
158 space_comm = ensemble_comm.space_comm
159 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI as sweeper_class
161 if ensemble_comm.time_comm.rank > 0:
162 dirname = f'{dirname}-{ensemble_comm.time_comm.rank}'
163 else:
164 ensemble_comm = None
165 space_comm = fd.COMM_WORLD
166 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
168 # Domain
169 mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2, comm=space_comm) if mesh is None else mesh
170 if Nlevels > 1:
171 hierarchy = fd.MeshHierarchy(mesh, Nlevels - 1)
172 mesh = hierarchy[-1]
173 domain = Domain(mesh, dt, 'BDM', element_order)
174 x, y, z = SpatialCoordinate(mesh)
175 lamda, phi, _ = lonlatr_from_xyz(x, y, z)
177 # Equation: topography
178 rsq = min_value(R0**2, (lamda - lamda_c) ** 2 + (phi - phi_c) ** 2)
179 r = sqrt(rsq)
180 tpexpr = mountain_height * (1 - r / R0)
181 parameters = ShallowWaterParameters(mesh, H=mean_depth, g=g, topog_expr=tpexpr)
182 eqns = ShallowWaterEquations(domain, parameters)
184 eqns.label_terms(lambda t: not t.has_label(time_derivative), implicit)
186 # I/O
187 output = OutputParameters(
188 dirname=dirname,
189 dumplist_latlon=['D'],
190 dumpfreq=dumpfreq,
191 dump_vtus=True,
192 dump_nc=True,
193 dumplist=['D', 'topography'],
194 )
195 diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), MeridionalComponent('u'), ZonalComponent('u')]
196 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
198 # Transport schemes
199 transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
201 # ------------------------------------------------------------------------ #
202 # pySDC parameters: description and controller parameters
203 # ------------------------------------------------------------------------ #
205 level_params = dict()
206 level_params['restol'] = -1
207 level_params['dt'] = dt
208 level_params['residual_type'] = 'full_rel'
210 step_params = dict()
211 step_params['maxiter'] = kmax
213 sweeper_params = dict()
214 sweeper_params['quad_type'] = 'RADAU-RIGHT'
215 sweeper_params['node_type'] = 'LEGENDRE'
216 sweeper_params['num_nodes'] = M
217 sweeper_params['QI'] = QI
218 sweeper_params['QE'] = 'PIC'
219 sweeper_params['comm'] = ensemble_comm
220 sweeper_params['initial_guess'] = 'copy'
222 problem_params = dict()
224 convergence_controllers = {}
225 if use_adaptivity:
226 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
227 from pySDC.implementations.convergence_controller_classes.spread_step_sizes import (
228 SpreadStepSizesBlockwiseNonMPI,
229 )
231 convergence_controllers[Adaptivity] = {'e_tol': 1e-6, 'rel_error': True, 'dt_max': 1e4, 'dt_rel_min_slope': 0.5}
232 # this is needed because the coupling runs on the controller level and this will almost always overwrite
233 convergence_controllers[SpreadStepSizesBlockwiseNonMPI] = {'overwrite_to_reach_Tend': False}
235 controller_params = dict()
236 controller_params['logger_level'] = logger_level if fd.COMM_WORLD.rank == 0 else 30
237 controller_params['mssdc_jac'] = False
239 description = dict()
240 description['problem_params'] = problem_params
241 description['sweeper_class'] = sweeper_class
242 description['sweeper_params'] = sweeper_params
243 description['level_params'] = level_params
244 description['step_params'] = step_params
245 description['convergence_controllers'] = convergence_controllers
247 from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrakeHierarchy
249 description['space_transfer_class'] = MeshToMeshFiredrakeHierarchy
251 # ------------------------------------------------------------------------ #
252 # petsc solver parameters
253 # ------------------------------------------------------------------------ #
255 solver_parameters = {
256 'snes_type': 'newtonls',
257 'ksp_type': 'gmres',
258 'pc_type': 'bjacobi',
259 'sub_pc_type': 'ilu',
260 'ksp_rtol': 1e-12,
261 'snes_rtol': 1e-12,
262 'ksp_atol': 1e-30,
263 'snes_atol': 1e-30,
264 'ksp_divtol': 1e30,
265 'snes_divtol': 1e30,
266 'snes_max_it': 999,
267 'ksp_max_it': 999,
268 }
270 # ------------------------------------------------------------------------ #
271 # Set Gusto SDC parameters to match the pySDC ones
272 # ------------------------------------------------------------------------ #
274 SDC_params = {
275 'base_scheme': BackwardEuler(domain, solver_parameters=solver_parameters),
276 'M': sweeper_params['num_nodes'],
277 'maxk': step_params['maxiter'],
278 'quad_type': sweeper_params['quad_type'],
279 'node_type': sweeper_params['node_type'],
280 'qdelta_imp': sweeper_params['QI'],
281 'qdelta_exp': sweeper_params['QE'],
282 'formulation': 'Z2N',
283 'initial_guess': 'copy',
284 'nonlinear_solver_parameters': solver_parameters,
285 'linear_solver_parameters': solver_parameters,
286 'final_update': False,
287 }
289 # ------------------------------------------------------------------------ #
290 # Setup time stepper
291 # ------------------------------------------------------------------------ #
293 if use_pySDC:
294 method = pySDC_integrator(description, controller_params, domain=domain, solver_parameters=solver_parameters)
295 else:
296 method = SDC(**SDC_params, domain=domain)
298 stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
300 # ------------------------------------------------------------------------ #
301 # Setup multi-level SDC
302 # ------------------------------------------------------------------------ #
304 if not _ML_is_setup:
305 return stepper
307 if Nlevels > 1:
308 steppers = [
309 None,
310 ] * (Nlevels)
311 steppers[0] = stepper
313 # get different steppers on the different levels
314 # recall that the setup of the problems is only finished when the stepper is setup
315 for i in range(1, Nlevels):
316 steppers[i] = williamson_5(
317 ncells_per_edge=ncells_per_edge,
318 dt=dt,
319 tmax=tmax,
320 dumpfreq=dumpfreq,
321 dirname=f'{dirname}_unused_{i}',
322 time_parallelism=time_parallelism,
323 QI=QI,
324 M=M,
325 kmax=kmax,
326 use_pySDC=use_pySDC,
327 use_adaptivity=use_adaptivity,
328 Nlevels=1,
329 mesh=hierarchy[-i - 1], # mind that the finest level in pySDC is 0, but -1 in hierarchy
330 logger_level=50,
331 _ML_is_setup=False,
332 )
334 # update description and setup pySDC again with the discretizations from different steppers
335 description['problem_params']['residual'] = [me.scheme.residual for me in steppers]
336 description['problem_params']['equation'] = [me.scheme.equation for me in steppers]
337 method = pySDC_integrator(description, controller_params, domain=domain, solver_parameters=solver_parameters)
338 stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
340 # ------------------------------------------------------------------------ #
341 # Initial conditions
342 # ------------------------------------------------------------------------ #
344 u0 = stepper.fields('u')
345 D0 = stepper.fields('D')
346 uexpr = as_vector([-u_max * y / radius, u_max * x / radius, 0.0])
347 Dexpr = mean_depth - tpexpr - (radius * parameters.Omega * u_max + 0.5 * u_max**2) * (z / radius) ** 2 / g
349 u0.project(uexpr)
350 D0.interpolate(Dexpr)
352 Dbar = Function(D0.function_space()).assign(mean_depth)
353 stepper.set_reference_profiles([('D', Dbar)])
355 # ------------------------------------------------------------------------ #
356 # Run
357 # ------------------------------------------------------------------------ #
359 if use_pySDC and use_adaptivity:
360 # we have to do this for adaptive time stepping, because it is a bit of a mess
361 method.timestepper = stepper
363 stepper.run(t=0, tmax=tmax)
364 return stepper, mesh
367# ---------------------------------------------------------------------------- #
368# MAIN
369# ---------------------------------------------------------------------------- #
372# %% [markdown]
373# ## Running it
374#
375# The script supports space-time parallelism, as well as running the Gusto SDC implementation or the pySDC-Gusto
376# coupling. Run it with `--help` to learn how to configure it. Afterwards,
377# [`F_2_plot_pySDC_with_Gusto_result.py`](https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_7/F_2_plot_pySDC_with_Gusto_result.py),
378# next to this tutorial, plots the results.
380# %%
381if __name__ == "__main__":
383 parser = ArgumentParser(
384 description='Williamson 5 (flow over a mountain) with pySDC as the time discretization of Gusto, or with '
385 "Gusto's own SDC implementation for comparison.",
386 formatter_class=ArgumentDefaultsHelpFormatter,
387 )
388 parser.add_argument(
389 '--ncells_per_edge',
390 help="The number of cells per edge of icosahedron",
391 type=int,
392 default=williamson_5_defaults['ncells_per_edge'],
393 )
394 parser.add_argument('--dt', help="The time step in seconds.", type=float, default=williamson_5_defaults['dt'])
395 parser.add_argument(
396 "--tmax", help="The end time for the simulation in seconds.", type=float, default=williamson_5_defaults['tmax']
397 )
398 parser.add_argument(
399 '--dumpfreq',
400 help="The frequency at which to dump field output.",
401 type=int,
402 default=williamson_5_defaults['dumpfreq'],
403 )
404 parser.add_argument(
405 '--dirname', help="The name of the directory to write to.", type=str, default=williamson_5_defaults['dirname']
406 )
407 parser.add_argument(
408 '--time_parallelism',
409 help="Whether to use parallel diagonal SDC or not.",
410 type=str,
411 default=williamson_5_defaults['time_parallelism'],
412 )
413 parser.add_argument('--kmax', help='SDC iteration count', type=int, default=williamson_5_defaults['kmax'])
414 parser.add_argument('-M', help='SDC node count', type=int, default=williamson_5_defaults['M'])
415 parser.add_argument(
416 '--use_pySDC',
417 help='whether to use pySDC or Gusto SDC implementation',
418 type=str,
419 default=williamson_5_defaults['use_pySDC'],
420 )
421 parser.add_argument(
422 '--use_adaptivity',
423 help='whether to use adaptive step size selection',
424 type=str,
425 default=williamson_5_defaults['use_adaptivity'],
426 )
427 parser.add_argument('--QI', help='Implicit preconditioner', type=str, default=williamson_5_defaults['QI'])
428 parser.add_argument(
429 '--Nlevels',
430 help="Number of SDC levels.",
431 type=int,
432 default=williamson_5_defaults['Nlevels'],
433 )
434 args, unknown = parser.parse_known_args()
436 options = vars(args)
437 for key in ['use_pySDC', 'use_adaptivity', 'time_parallelism']:
438 options[key] = options[key] not in ['False', 0, False, 'false']
440 williamson_5(**options)
442# %% [markdown]
443# Firedrake and Gusto do not run in the browser, nor in the environment this website is built in. Our CI runs this part
444# in a Firedrake container with Gusto: it checks that the pySDC coupling gives the same result as Gusto's own SDC
445# after a few time steps, and that the multilevel version converges faster than the single-level one. These are the
446# results of the run that built this page, for a few time steps only, to keep the CI fast:
447#
448# :::{literalinclude} /../../data_firedrake/step_7_F_out.txt
449# :language: text
450# :::