Coverage for pySDC/tutorial/step_7/F_pySDC_with_Gusto.py: 0%
114 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 10:09 +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: coriolis
158 parameters = ShallowWaterParameters(H=mean_depth, g=g)
159 Omega = parameters.Omega
160 fexpr = 2 * Omega * z / radius
162 # Equation: topography
163 rsq = min_value(R0**2, (lamda - lamda_c) ** 2 + (phi - phi_c) ** 2)
164 r = sqrt(rsq)
165 tpexpr = mountain_height * (1 - r / R0)
166 eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, topog_expr=tpexpr)
168 eqns.label_terms(lambda t: not t.has_label(time_derivative), implicit)
170 # I/O
171 output = OutputParameters(
172 dirname=dirname,
173 dumplist_latlon=['D'],
174 dumpfreq=dumpfreq,
175 dump_vtus=True,
176 dump_nc=True,
177 dumplist=['D', 'topography'],
178 )
179 diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), MeridionalComponent('u'), ZonalComponent('u')]
180 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
182 # Transport schemes
183 transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
185 # ------------------------------------------------------------------------ #
186 # pySDC parameters: description and controller parameters
187 # ------------------------------------------------------------------------ #
189 level_params = dict()
190 level_params['restol'] = -1
191 level_params['dt'] = dt
192 level_params['residual_type'] = 'full_rel'
194 step_params = dict()
195 step_params['maxiter'] = kmax
197 sweeper_params = dict()
198 sweeper_params['quad_type'] = 'RADAU-RIGHT'
199 sweeper_params['node_type'] = 'LEGENDRE'
200 sweeper_params['num_nodes'] = M
201 sweeper_params['QI'] = QI
202 sweeper_params['QE'] = 'PIC'
203 sweeper_params['comm'] = ensemble_comm
204 sweeper_params['initial_guess'] = 'copy'
206 problem_params = dict()
208 convergence_controllers = {}
209 if use_adaptivity:
210 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
211 from pySDC.implementations.convergence_controller_classes.spread_step_sizes import (
212 SpreadStepSizesBlockwiseNonMPI,
213 )
215 convergence_controllers[Adaptivity] = {'e_tol': 1e-6, 'rel_error': True, 'dt_max': 1e4, 'dt_rel_min_slope': 0.5}
216 # this is needed because the coupling runs on the controller level and this will almost always overwrite
217 convergence_controllers[SpreadStepSizesBlockwiseNonMPI] = {'overwrite_to_reach_Tend': False}
219 controller_params = dict()
220 controller_params['logger_level'] = logger_level if fd.COMM_WORLD.rank == 0 else 30
221 controller_params['mssdc_jac'] = False
223 description = dict()
224 description['problem_params'] = problem_params
225 description['sweeper_class'] = sweeper_class
226 description['sweeper_params'] = sweeper_params
227 description['level_params'] = level_params
228 description['step_params'] = step_params
229 description['convergence_controllers'] = convergence_controllers
231 from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrakeHierarchy
233 description['space_transfer_class'] = MeshToMeshFiredrakeHierarchy
235 # ------------------------------------------------------------------------ #
236 # petsc solver parameters
237 # ------------------------------------------------------------------------ #
239 solver_parameters = {
240 'snes_type': 'newtonls',
241 'ksp_type': 'gmres',
242 'pc_type': 'bjacobi',
243 'sub_pc_type': 'ilu',
244 'ksp_rtol': 1e-12,
245 'snes_rtol': 1e-12,
246 'ksp_atol': 1e-30,
247 'snes_atol': 1e-30,
248 'ksp_divtol': 1e30,
249 'snes_divtol': 1e30,
250 'snes_max_it': 999,
251 'ksp_max_it': 999,
252 }
254 # ------------------------------------------------------------------------ #
255 # Set Gusto SDC parameters to match the pySDC ones
256 # ------------------------------------------------------------------------ #
258 SDC_params = {
259 'base_scheme': BackwardEuler(domain, solver_parameters=solver_parameters),
260 'M': sweeper_params['num_nodes'],
261 'maxk': step_params['maxiter'],
262 'quad_type': sweeper_params['quad_type'],
263 'node_type': sweeper_params['node_type'],
264 'qdelta_imp': sweeper_params['QI'],
265 'qdelta_exp': sweeper_params['QE'],
266 'formulation': 'Z2N',
267 'initial_guess': 'copy',
268 'nonlinear_solver_parameters': solver_parameters,
269 'linear_solver_parameters': solver_parameters,
270 'final_update': False,
271 }
273 # ------------------------------------------------------------------------ #
274 # Setup time stepper
275 # ------------------------------------------------------------------------ #
277 if use_pySDC:
278 method = pySDC_integrator(description, controller_params, domain=domain, solver_parameters=solver_parameters)
279 else:
280 method = SDC(**SDC_params, domain=domain)
282 stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
284 # ------------------------------------------------------------------------ #
285 # Setup multi-level SDC
286 # ------------------------------------------------------------------------ #
288 if not _ML_is_setup:
289 return stepper
291 if Nlevels > 1:
292 steppers = [
293 None,
294 ] * (Nlevels)
295 steppers[0] = stepper
297 # get different steppers on the different levels
298 # recall that the setup of the problems is only finished when the stepper is setup
299 for i in range(1, Nlevels):
300 steppers[i] = williamson_5(
301 ncells_per_edge=ncells_per_edge,
302 dt=dt,
303 tmax=tmax,
304 dumpfreq=dumpfreq,
305 dirname=f'{dirname}_unused_{i}',
306 time_parallelism=time_parallelism,
307 QI=QI,
308 M=M,
309 kmax=kmax,
310 use_pySDC=use_pySDC,
311 use_adaptivity=use_adaptivity,
312 Nlevels=1,
313 mesh=hierarchy[-i - 1], # mind that the finest level in pySDC is 0, but -1 in hierarchy
314 logger_level=50,
315 _ML_is_setup=False,
316 )
318 # update description and setup pySDC again with the discretizations from different steppers
319 description['problem_params']['residual'] = [me.scheme.residual for me in steppers]
320 description['problem_params']['equation'] = [me.scheme.equation for me in steppers]
321 method = pySDC_integrator(description, controller_params, domain=domain, solver_parameters=solver_parameters)
322 stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
324 # ------------------------------------------------------------------------ #
325 # Initial conditions
326 # ------------------------------------------------------------------------ #
328 u0 = stepper.fields('u')
329 D0 = stepper.fields('D')
330 uexpr = as_vector([-u_max * y / radius, u_max * x / radius, 0.0])
331 Dexpr = mean_depth - tpexpr - (radius * Omega * u_max + 0.5 * u_max**2) * (z / radius) ** 2 / g
333 u0.project(uexpr)
334 D0.interpolate(Dexpr)
336 Dbar = Function(D0.function_space()).assign(mean_depth)
337 stepper.set_reference_profiles([('D', Dbar)])
339 # ------------------------------------------------------------------------ #
340 # Run
341 # ------------------------------------------------------------------------ #
343 if use_pySDC and use_adaptivity:
344 # we have to do this for adaptive time stepping, because it is a bit of a mess
345 method.timestepper = stepper
347 stepper.run(t=0, tmax=tmax)
348 return stepper, mesh
351# ---------------------------------------------------------------------------- #
352# MAIN
353# ---------------------------------------------------------------------------- #
356if __name__ == "__main__":
358 parser = ArgumentParser(description=__doc__, formatter_class=ArgumentDefaultsHelpFormatter)
359 parser.add_argument(
360 '--ncells_per_edge',
361 help="The number of cells per edge of icosahedron",
362 type=int,
363 default=williamson_5_defaults['ncells_per_edge'],
364 )
365 parser.add_argument('--dt', help="The time step in seconds.", type=float, default=williamson_5_defaults['dt'])
366 parser.add_argument(
367 "--tmax", help="The end time for the simulation in seconds.", type=float, default=williamson_5_defaults['tmax']
368 )
369 parser.add_argument(
370 '--dumpfreq',
371 help="The frequency at which to dump field output.",
372 type=int,
373 default=williamson_5_defaults['dumpfreq'],
374 )
375 parser.add_argument(
376 '--dirname', help="The name of the directory to write to.", type=str, default=williamson_5_defaults['dirname']
377 )
378 parser.add_argument(
379 '--time_parallelism',
380 help="Whether to use parallel diagonal SDC or not.",
381 type=str,
382 default=williamson_5_defaults['time_parallelism'],
383 )
384 parser.add_argument('--kmax', help='SDC iteration count', type=int, default=williamson_5_defaults['kmax'])
385 parser.add_argument('-M', help='SDC node count', type=int, default=williamson_5_defaults['M'])
386 parser.add_argument(
387 '--use_pySDC',
388 help='whether to use pySDC or Gusto SDC implementation',
389 type=str,
390 default=williamson_5_defaults['use_pySDC'],
391 )
392 parser.add_argument(
393 '--use_adaptivity',
394 help='whether to use adaptive step size selection',
395 type=str,
396 default=williamson_5_defaults['use_adaptivity'],
397 )
398 parser.add_argument('--QI', help='Implicit preconditioner', type=str, default=williamson_5_defaults['QI'])
399 parser.add_argument(
400 '--Nlevels',
401 help="Number of SDC levels.",
402 type=int,
403 default=williamson_5_defaults['Nlevels'],
404 )
405 args, unknown = parser.parse_known_args()
407 options = vars(args)
408 for key in ['use_pySDC', 'use_adaptivity', 'time_parallelism']:
409 options[key] = options[key] not in ['False', 0, False, 'false']
411 williamson_5(**options)