Coverage for pySDC / tutorial / step_7 / F_pySDC_with_Gusto.py: 0%
112 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-13 09:00 +0000
1"""
2Example for running pySDC together with Gusto. This test runs a shallow water equation and may take a considerable
3amount of time. After you have run it, move on to step F_2, which includes a plotting script.
5This is Test Case 5 (flow over a mountain) of Williamson et al, 1992:
6``A standard test set for numerical approximations to the shallow water
7equations in spherical geometry'', JCP.
9This script is adapted from the Gusto example: https://github.com/firedrakeproject/gusto/blob/main/examples/shallow_water/williamson_5.py
11The pySDC coupling works by setting up pySDC as a time integrator within Gusto.
12To this end, you need to construct a pySDC description and controller parameters as usual and pass them when
13constructing the pySDC time discretization.
15After passing this to a Gusto timestepper, you have two choices:
16 - Access the `.scheme.controller` variable of the timestepper, which is the pySDC controller and use pySDC for
17 running
18 - Use the Gusto timestepper for running
19You may wonder why it is necessary to construct a Gusto timestepper if you don't want to use it. The reason is the
20setup of spatial methods, such as upwinding. These are passed to the Gusto timestepper and modify the residual of the
21equations during its instantiation. Once the residual is modified, we can choose whether to continue in Gusto or pySDC.
23This script supports space-time parallelism, as well as running the Gusto SDC implementation or the pySDC-Gusto coupling.
24Please run with `--help` to learn how to configure this script.
25"""
27import firedrake as fd
28from pySDC.helpers.pySDC_as_gusto_time_discretization import pySDC_integrator
29from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator
30from gusto import SDC, BackwardEuler
31from gusto.core.labels import implicit, time_derivative
32from gusto.core.logging import logger, INFO
34logger.setLevel(INFO)
37from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
38from firedrake import SpatialCoordinate, as_vector, pi, sqrt, min_value, Function
39from gusto import (
40 Domain,
41 IO,
42 OutputParameters,
43 DGUpwind,
44 ShallowWaterParameters,
45 ShallowWaterEquations,
46 Sum,
47 lonlatr_from_xyz,
48 GeneralIcosahedralSphereMesh,
49 ZonalComponent,
50 MeridionalComponent,
51 RelativeVorticity,
52 Timestepper,
53)
55williamson_5_defaults = {
56 'ncells_per_edge': 12, # number of cells per icosahedron edge
57 'dt': 900.0,
58 'tmax': 50.0 * 24.0 * 60.0 * 60.0, # 50 days
59 'dumpfreq': 10, # output every <dumpfreq> steps
60 'dirname': 'williamson_5', # results will go into ./results/<dirname>
61 'time_parallelism': False, # use parallel diagonal SDC or not
62 'QI': 'MIN-SR-S', # implicit preconditioner
63 'M': '3', # number of collocation nodes
64 'kmax': '5', # use fixed number of iteration up to this value
65 'use_pySDC': True, # whether to use pySDC for time integration
66 'use_adaptivity': True, # whether to use adaptive step size selection
67 'Nlevels': 1, # number of levels in SDC
68 'logger_level': 15, # pySDC logger level
69}
72def williamson_5(
73 ncells_per_edge=williamson_5_defaults['ncells_per_edge'],
74 dt=williamson_5_defaults['dt'],
75 tmax=williamson_5_defaults['tmax'],
76 dumpfreq=williamson_5_defaults['dumpfreq'],
77 dirname=williamson_5_defaults['dirname'],
78 time_parallelism=williamson_5_defaults['time_parallelism'],
79 QI=williamson_5_defaults['QI'],
80 M=williamson_5_defaults['M'],
81 kmax=williamson_5_defaults['kmax'],
82 use_pySDC=williamson_5_defaults['use_pySDC'],
83 use_adaptivity=williamson_5_defaults['use_adaptivity'],
84 Nlevels=williamson_5_defaults['Nlevels'],
85 logger_level=williamson_5_defaults['logger_level'],
86 mesh=None,
87 _ML_is_setup=True,
88):
89 """
90 Run the Williamson 5 test case.
92 Args:
93 ncells_per_edge (int): number of cells per icosahedron edge
94 dt (float): Initial step size
95 tmax (float): Time to integrate to
96 dumpfreq (int): Output every <dumpfreq> time steps
97 dirname (str): Output will go into ./results/<dirname>
98 time_parallelism (bool): True for parallel SDC, False for serial
99 M (int): Number of collocation nodes
100 kmax (int): Max number of SDC iterations
101 use_pySDC (bool): Use pySDC as Gusto time integrator or Gusto SDC implementation
102 Nlevels (int): Number of SDC levels
103 logger_level (int): Logger level
104 """
105 if not use_pySDC and use_adaptivity:
106 raise NotImplementedError('Adaptive step size selection not yet implemented in Gusto')
107 if not use_pySDC and Nlevels > 1:
108 raise NotImplementedError('Multi-level SDC not yet implemented in Gusto')
109 if time_parallelism and Nlevels > 1:
110 raise NotImplementedError('Multi-level SDC does not work with MPI parallel sweeper yet')
112 # ------------------------------------------------------------------------ #
113 # Parameters for test case
114 # ------------------------------------------------------------------------ #
116 radius = 6371220.0 # planetary radius (m)
117 mean_depth = 5960 # reference depth (m)
118 g = 9.80616 # acceleration due to gravity (m/s^2)
119 u_max = 20.0 # max amplitude of the zonal wind (m/s)
120 mountain_height = 2000.0 # height of mountain (m)
121 R0 = pi / 9.0 # radius of mountain (rad)
122 lamda_c = -pi / 2.0 # longitudinal centre of mountain (rad)
123 phi_c = pi / 6.0 # latitudinal centre of mountain (rad)
125 # ------------------------------------------------------------------------ #
126 # Our settings for this set up
127 # ------------------------------------------------------------------------ #
129 element_order = 1
131 # ------------------------------------------------------------------------ #
132 # Set up model objects
133 # ------------------------------------------------------------------------ #
135 # parallelism
136 if time_parallelism:
137 ensemble_comm = FiredrakeEnsembleCommunicator(fd.COMM_WORLD, fd.COMM_WORLD.size // M)
138 space_comm = ensemble_comm.space_comm
139 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI as sweeper_class
141 if ensemble_comm.time_comm.rank > 0:
142 dirname = f'{dirname}-{ensemble_comm.time_comm.rank}'
143 else:
144 ensemble_comm = None
145 space_comm = fd.COMM_WORLD
146 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
148 # Domain
149 mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2, comm=space_comm) if mesh is None else mesh
150 if Nlevels > 1:
151 hierarchy = fd.MeshHierarchy(mesh, Nlevels - 1)
152 mesh = hierarchy[-1]
153 domain = Domain(mesh, dt, 'BDM', element_order)
154 x, y, z = SpatialCoordinate(mesh)
155 lamda, phi, _ = lonlatr_from_xyz(x, y, z)
157 # Equation: topography
158 rsq = min_value(R0**2, (lamda - lamda_c) ** 2 + (phi - phi_c) ** 2)
159 r = sqrt(rsq)
160 tpexpr = mountain_height * (1 - r / R0)
161 parameters = ShallowWaterParameters(mesh, H=mean_depth, g=g, topog_expr=tpexpr)
162 eqns = ShallowWaterEquations(domain, parameters)
164 eqns.label_terms(lambda t: not t.has_label(time_derivative), implicit)
166 # I/O
167 output = OutputParameters(
168 dirname=dirname,
169 dumplist_latlon=['D'],
170 dumpfreq=dumpfreq,
171 dump_vtus=True,
172 dump_nc=True,
173 dumplist=['D', 'topography'],
174 )
175 diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), MeridionalComponent('u'), ZonalComponent('u')]
176 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
178 # Transport schemes
179 transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
181 # ------------------------------------------------------------------------ #
182 # pySDC parameters: description and controller parameters
183 # ------------------------------------------------------------------------ #
185 level_params = dict()
186 level_params['restol'] = -1
187 level_params['dt'] = dt
188 level_params['residual_type'] = 'full_rel'
190 step_params = dict()
191 step_params['maxiter'] = kmax
193 sweeper_params = dict()
194 sweeper_params['quad_type'] = 'RADAU-RIGHT'
195 sweeper_params['node_type'] = 'LEGENDRE'
196 sweeper_params['num_nodes'] = M
197 sweeper_params['QI'] = QI
198 sweeper_params['QE'] = 'PIC'
199 sweeper_params['comm'] = ensemble_comm
200 sweeper_params['initial_guess'] = 'copy'
202 problem_params = dict()
204 convergence_controllers = {}
205 if use_adaptivity:
206 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
207 from pySDC.implementations.convergence_controller_classes.spread_step_sizes import (
208 SpreadStepSizesBlockwiseNonMPI,
209 )
211 convergence_controllers[Adaptivity] = {'e_tol': 1e-6, 'rel_error': True, 'dt_max': 1e4, 'dt_rel_min_slope': 0.5}
212 # this is needed because the coupling runs on the controller level and this will almost always overwrite
213 convergence_controllers[SpreadStepSizesBlockwiseNonMPI] = {'overwrite_to_reach_Tend': False}
215 controller_params = dict()
216 controller_params['logger_level'] = logger_level if fd.COMM_WORLD.rank == 0 else 30
217 controller_params['mssdc_jac'] = False
219 description = dict()
220 description['problem_params'] = problem_params
221 description['sweeper_class'] = sweeper_class
222 description['sweeper_params'] = sweeper_params
223 description['level_params'] = level_params
224 description['step_params'] = step_params
225 description['convergence_controllers'] = convergence_controllers
227 from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrakeHierarchy
229 description['space_transfer_class'] = MeshToMeshFiredrakeHierarchy
231 # ------------------------------------------------------------------------ #
232 # petsc solver parameters
233 # ------------------------------------------------------------------------ #
235 solver_parameters = {
236 'snes_type': 'newtonls',
237 'ksp_type': 'gmres',
238 'pc_type': 'bjacobi',
239 'sub_pc_type': 'ilu',
240 'ksp_rtol': 1e-12,
241 'snes_rtol': 1e-12,
242 'ksp_atol': 1e-30,
243 'snes_atol': 1e-30,
244 'ksp_divtol': 1e30,
245 'snes_divtol': 1e30,
246 'snes_max_it': 999,
247 'ksp_max_it': 999,
248 }
250 # ------------------------------------------------------------------------ #
251 # Set Gusto SDC parameters to match the pySDC ones
252 # ------------------------------------------------------------------------ #
254 SDC_params = {
255 'base_scheme': BackwardEuler(domain, solver_parameters=solver_parameters),
256 'M': sweeper_params['num_nodes'],
257 'maxk': step_params['maxiter'],
258 'quad_type': sweeper_params['quad_type'],
259 'node_type': sweeper_params['node_type'],
260 'qdelta_imp': sweeper_params['QI'],
261 'qdelta_exp': sweeper_params['QE'],
262 'formulation': 'Z2N',
263 'initial_guess': 'copy',
264 'nonlinear_solver_parameters': solver_parameters,
265 'linear_solver_parameters': solver_parameters,
266 'final_update': False,
267 }
269 # ------------------------------------------------------------------------ #
270 # Setup time stepper
271 # ------------------------------------------------------------------------ #
273 if use_pySDC:
274 method = pySDC_integrator(description, controller_params, domain=domain, solver_parameters=solver_parameters)
275 else:
276 method = SDC(**SDC_params, domain=domain)
278 stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
280 # ------------------------------------------------------------------------ #
281 # Setup multi-level SDC
282 # ------------------------------------------------------------------------ #
284 if not _ML_is_setup:
285 return stepper
287 if Nlevels > 1:
288 steppers = [
289 None,
290 ] * (Nlevels)
291 steppers[0] = stepper
293 # get different steppers on the different levels
294 # recall that the setup of the problems is only finished when the stepper is setup
295 for i in range(1, Nlevels):
296 steppers[i] = williamson_5(
297 ncells_per_edge=ncells_per_edge,
298 dt=dt,
299 tmax=tmax,
300 dumpfreq=dumpfreq,
301 dirname=f'{dirname}_unused_{i}',
302 time_parallelism=time_parallelism,
303 QI=QI,
304 M=M,
305 kmax=kmax,
306 use_pySDC=use_pySDC,
307 use_adaptivity=use_adaptivity,
308 Nlevels=1,
309 mesh=hierarchy[-i - 1], # mind that the finest level in pySDC is 0, but -1 in hierarchy
310 logger_level=50,
311 _ML_is_setup=False,
312 )
314 # update description and setup pySDC again with the discretizations from different steppers
315 description['problem_params']['residual'] = [me.scheme.residual for me in steppers]
316 description['problem_params']['equation'] = [me.scheme.equation for me in steppers]
317 method = pySDC_integrator(description, controller_params, domain=domain, solver_parameters=solver_parameters)
318 stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
320 # ------------------------------------------------------------------------ #
321 # Initial conditions
322 # ------------------------------------------------------------------------ #
324 u0 = stepper.fields('u')
325 D0 = stepper.fields('D')
326 uexpr = as_vector([-u_max * y / radius, u_max * x / radius, 0.0])
327 Dexpr = mean_depth - tpexpr - (radius * parameters.Omega * u_max + 0.5 * u_max**2) * (z / radius) ** 2 / g
329 u0.project(uexpr)
330 D0.interpolate(Dexpr)
332 Dbar = Function(D0.function_space()).assign(mean_depth)
333 stepper.set_reference_profiles([('D', Dbar)])
335 # ------------------------------------------------------------------------ #
336 # Run
337 # ------------------------------------------------------------------------ #
339 if use_pySDC and use_adaptivity:
340 # we have to do this for adaptive time stepping, because it is a bit of a mess
341 method.timestepper = stepper
343 stepper.run(t=0, tmax=tmax)
344 return stepper, mesh
347# ---------------------------------------------------------------------------- #
348# MAIN
349# ---------------------------------------------------------------------------- #
352if __name__ == "__main__":
354 parser = ArgumentParser(description=__doc__, formatter_class=ArgumentDefaultsHelpFormatter)
355 parser.add_argument(
356 '--ncells_per_edge',
357 help="The number of cells per edge of icosahedron",
358 type=int,
359 default=williamson_5_defaults['ncells_per_edge'],
360 )
361 parser.add_argument('--dt', help="The time step in seconds.", type=float, default=williamson_5_defaults['dt'])
362 parser.add_argument(
363 "--tmax", help="The end time for the simulation in seconds.", type=float, default=williamson_5_defaults['tmax']
364 )
365 parser.add_argument(
366 '--dumpfreq',
367 help="The frequency at which to dump field output.",
368 type=int,
369 default=williamson_5_defaults['dumpfreq'],
370 )
371 parser.add_argument(
372 '--dirname', help="The name of the directory to write to.", type=str, default=williamson_5_defaults['dirname']
373 )
374 parser.add_argument(
375 '--time_parallelism',
376 help="Whether to use parallel diagonal SDC or not.",
377 type=str,
378 default=williamson_5_defaults['time_parallelism'],
379 )
380 parser.add_argument('--kmax', help='SDC iteration count', type=int, default=williamson_5_defaults['kmax'])
381 parser.add_argument('-M', help='SDC node count', type=int, default=williamson_5_defaults['M'])
382 parser.add_argument(
383 '--use_pySDC',
384 help='whether to use pySDC or Gusto SDC implementation',
385 type=str,
386 default=williamson_5_defaults['use_pySDC'],
387 )
388 parser.add_argument(
389 '--use_adaptivity',
390 help='whether to use adaptive step size selection',
391 type=str,
392 default=williamson_5_defaults['use_adaptivity'],
393 )
394 parser.add_argument('--QI', help='Implicit preconditioner', type=str, default=williamson_5_defaults['QI'])
395 parser.add_argument(
396 '--Nlevels',
397 help="Number of SDC levels.",
398 type=int,
399 default=williamson_5_defaults['Nlevels'],
400 )
401 args, unknown = parser.parse_known_args()
403 options = vars(args)
404 for key in ['use_pySDC', 'use_adaptivity', 'time_parallelism']:
405 options[key] = options[key] not in ['False', 0, False, 'false']
407 williamson_5(**options)