diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 110405818..4bb39fb7e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -11,6 +11,13 @@ on: # Scheduled build at 0330 UTC on Monday mornings to detect bitrot. - cron: '30 3 * * 1' +concurrency: + # Cancels jobs running if new commits are pushed + group: > + ${{ github.workflow }}- + ${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + jobs: build: name: "Build Gusto" @@ -39,6 +46,10 @@ jobs: firedrake-clean export GUSTO_PARALLEL_LOG=FILE export PYOP2_CFLAGS=-O0 + python -m pip uninstall -y netCDF4 + export HDF5_DIR=$PETSC_DIR/packages + export NETCDF4_DIR=$PETSC_DIR/packages + python -m pip install --no-binary netCDF4 --no-build-isolation netCDF4 python -m pytest \ -n 12 --dist worksteal \ --durations=100 \ @@ -53,6 +64,17 @@ jobs: mkdir logs cd /tmp/pytest-of-firedrake/pytest-0/ find . -name "*.log" -exec cp --parents {} /__w/gusto/gusto/logs/ \; + - name: Test serial netCDF + run: | + . /home/firedrake/firedrake/bin/activate + python -m pip uninstall -y netCDF4 + python -m pip cache remove netCDF4 + python -m pip install --only-binary netCDF4 netCDF4 + firedrake-clean + export GUSTO_PARALLEL_LOG=FILE + export PYOP2_CFLAGS=-O0 + python -m pytest -n 3 -v integration-tests/model/test_nc_outputting.py + timeout-minutes: 10 - name: Upload artifact if: always() uses: actions/upload-pages-artifact@v3 diff --git a/docs/about_gusto.md b/docs/about_gusto.md index bf1719fd2..3b301211c 100644 --- a/docs/about_gusto.md +++ b/docs/about_gusto.md @@ -20,6 +20,30 @@ cd firedrake/src/gusto make test ``` +#### Parallel output with netCDF + +The [`netCDF4`](https://pypi.org/project/netCDF4/) package installed by Gusto does not support parallel usage. +This means that, when Gusto is run in parallel, distributed data structures must first be gathered onto rank 0 before they can be output. +This is *extremely inefficient* at high levels of parallelism. + +To avoid this it is possible to build a parallel-aware version of `netCDF4`. +The steps to do this are as follows: + +1. Activate the Firedrake virtual environment. +2. Uninstall the existing `netCDF4` package: + ``` + $ pip uninstall netCDF4 + ``` +3. Set necessary environment variables (note that this assumes that PETSc was built by Firedrake): + ``` + $ export HDF5_DIR=$VIRTUAL_ENV/src/petsc/default + $ export NETCDF4_DIR=$HDF5_DIR + ``` +4. Install the parallel version of `netCDF4`: + ``` + $ pip install --no-build-isolation --no-binary netCDF4 netCDF4 + ``` + ### Getting Started Once you have a working installation, the best way to get started with Gusto is to play with some of the examples in the `gusto/examples` directory. diff --git a/examples/compressible_euler/mountain_hydrostatic.py b/examples/compressible_euler/schaer_mountain.py similarity index 62% rename from examples/compressible_euler/mountain_hydrostatic.py rename to examples/compressible_euler/schaer_mountain.py index 1fc113a71..d60679d1d 100644 --- a/examples/compressible_euler/mountain_hydrostatic.py +++ b/examples/compressible_euler/schaer_mountain.py @@ -1,68 +1,71 @@ """ -The hydrostatic 1 metre high mountain test case from Melvin et al, 2010: -``An inherently mass-conserving iterative semi-implicit semi-Lagrangian -discretization of the non-hydrostatic vertical-slice equations.'', QJRMS. +The Schär mountain test case of Schär et al, 2002: +``A new terrain-following vertical coordinate formulation for atmospheric +prediction models.'', MWR. -This test describes a wave over a mountain in a hydrostatic atmosphere. +This test describes a wave over a set of idealised mountains, testing how the +discretisation handles orography. The setup used here uses the order 1 finite elements. """ - from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter + from firedrake import ( as_vector, VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh, - SpatialCoordinate, exp, pi, cos, Function, conditional, Mesh, Constant + SpatialCoordinate, exp, pi, cos, Function, Mesh, Constant ) from gusto import ( - Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - TrapeziumRule, SUPGOptions, ZComponent, Perturbation, - CompressibleParameters, HydrostaticCompressibleEulerEquations, - CompressibleSolver, compressible_hydrostatic_balance, HydrostaticImbalance, - SpongeLayerParameters, MinKernel, MaxKernel, logger + Domain, CompressibleParameters, CompressibleSolver, logger, + OutputParameters, IO, SSPRK3, DGUpwind, SemiImplicitQuasiNewton, + compressible_hydrostatic_balance, SpongeLayerParameters, Exner, ZComponent, + Perturbation, SUPGOptions, TrapeziumRule, MaxKernel, MinKernel, + CompressibleEulerEquations, SubcyclingOptions, RungeKuttaFormulation ) -mountain_hydrostatic_defaults = { - 'ncolumns': 200, - 'nlayers': 120, - 'dt': 5.0, - 'tmax': 15000., - 'dumpfreq': 1500, - 'dirname': 'mountain_hydrostatic' +schaer_mountain_defaults = { + 'ncolumns': 100, + 'nlayers': 50, + 'dt': 8.0, + 'tmax': 5*60*60., # 5 hours + 'dumpfreq': 2250, # dump at end with default settings + 'dirname': 'schaer_mountain' } -def mountain_hydrostatic( - ncolumns=mountain_hydrostatic_defaults['ncolumns'], - nlayers=mountain_hydrostatic_defaults['nlayers'], - dt=mountain_hydrostatic_defaults['dt'], - tmax=mountain_hydrostatic_defaults['tmax'], - dumpfreq=mountain_hydrostatic_defaults['dumpfreq'], - dirname=mountain_hydrostatic_defaults['dirname'] +def schaer_mountain( + ncolumns=schaer_mountain_defaults['ncolumns'], + nlayers=schaer_mountain_defaults['nlayers'], + dt=schaer_mountain_defaults['dt'], + tmax=schaer_mountain_defaults['tmax'], + dumpfreq=schaer_mountain_defaults['dumpfreq'], + dirname=schaer_mountain_defaults['dirname'] ): # ------------------------------------------------------------------------ # # Parameters for test case # ------------------------------------------------------------------------ # - domain_width = 240000. # width of domain in x direction, in m - domain_height = 50000. # height of model top, in m - a = 10000. # scale width of mountain, in m - hm = 1. # height of mountain, in m - zh = 5000. # height at which mesh is no longer distorted, in m - Tsurf = 250. # temperature of surface, in K - initial_wind = 20.0 # initial horizontal wind, in m/s - sponge_depth = 20000.0 # depth of sponge layer, in m - g = 9.80665 # acceleration due to gravity, in m/s^2 - cp = 1004. # specific heat capacity at constant pressure - sponge_mu = 0.15 # parameter for strength of sponge layer, in J/kg/K + domain_width = 100000. # width of domain in x direction, in m + domain_height = 30000. # height of model top, in m + a = 5000. # scale width of mountain profile, in m + lamda = 4000. # scale width of individual mountains, in m + hm = 250. # height of mountain, in m + Tsurf = 288. # temperature of surface, in K + initial_wind = 10.0 # initial horizontal wind, in m/s + sponge_depth = 10000.0 # depth of sponge layer, in m + g = 9.810616 # acceleration due to gravity, in m/s^2 + cp = 1004.5 # specific heat capacity at constant pressure + mu_dt = 1.2 # strength of sponge layer, no units exner_surf = 1.0 # maximum value of Exner pressure at surface - max_iterations = 10 # maximum number of hydrostatic balance iterations - tolerance = 1e-7 # tolerance for hydrostatic balance iteration + max_iterations = 20 # maximum number of hydrostatic balance iterations + tolerance = 1e-8 # tolerance for hydrostatic balance iteration # ------------------------------------------------------------------------ # # Our settings for this set up # ------------------------------------------------------------------------ # + spinup_steps = 5 # Not necessary but helps balance initial conditions + alpha = 0.51 # Necessary to absorb grid scale waves element_order = 1 u_eqn_type = 'vector_invariant_form' @@ -81,9 +84,9 @@ def mountain_hydrostatic( # Describe the mountain xc = domain_width/2. x, z = SpatialCoordinate(ext_mesh) - zs = hm * a**2 / ((x - xc)**2 + a**2) + zs = hm * exp(-((x - xc)/a)**2) * (cos(pi*(x - xc)/lamda))**2 xexpr = as_vector( - [x, conditional(z < zh, z + cos(0.5 * pi * z / zh)**6 * zs, z)] + [x, z + ((domain_height - z) / domain_height) * zs] ) # Make new mesh @@ -95,64 +98,50 @@ def mountain_hydrostatic( # Equation parameters = CompressibleParameters(g=g, cp=cp) sponge = SpongeLayerParameters( - H=domain_height, z_level=domain_height-sponge_depth, mubar=sponge_mu/dt + H=domain_height, z_level=domain_height-sponge_depth, mubar=mu_dt/dt ) - eqns = HydrostaticCompressibleEulerEquations( + eqns = CompressibleEulerEquations( domain, parameters, sponge_options=sponge, u_transport_option=u_eqn_type ) # I/O output = OutputParameters( - dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=False + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True ) diagnostic_fields = [ - ZComponent('u'), HydrostaticImbalance(eqns), - Perturbation('theta'), Perturbation('rho') + Exner(parameters), ZComponent('u'), Perturbation('theta'), + Perturbation('rho') ] io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes + subcycling_opts = SubcyclingOptions(subcycle_by_courant=0.25) theta_opts = SUPGOptions() transported_fields = [ - TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho"), - SSPRK3(domain, "theta", options=theta_opts) + TrapeziumRule(domain, "u", subcycling_options=subcycling_opts), + SSPRK3( + domain, "rho", rk_formulation=RungeKuttaFormulation.predictor, + subcycling_options=subcycling_opts + ), + SSPRK3( + domain, "theta", options=theta_opts, + subcycling_options=subcycling_opts + ) ] transport_methods = [ DGUpwind(eqns, "u"), - DGUpwind(eqns, "rho"), + DGUpwind(eqns, "rho", advective_then_flux=True), DGUpwind(eqns, "theta", ibp=theta_opts.ibp) ] # Linear solver - params = {'mat_type': 'matfree', - 'ksp_type': 'preonly', - 'pc_type': 'python', - 'pc_python_type': 'firedrake.SCPC', - # Velocity mass operator is singular in the hydrostatic case. - # So for reconstruction, we eliminate rho into u - 'pc_sc_eliminate_fields': '1, 0', - 'condensed_field': {'ksp_type': 'fgmres', - 'ksp_rtol': 1.0e-8, - 'ksp_atol': 1.0e-8, - 'ksp_max_it': 100, - 'pc_type': 'gamg', - 'pc_gamg_sym_graph': True, - 'mg_levels': {'ksp_type': 'gmres', - 'ksp_max_it': 5, - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'}}} - - alpha = 0.51 # off-centering parameter - linear_solver = CompressibleSolver( - eqns, alpha, solver_parameters=params, - overwrite_solver_parameters=True - ) + tau_values = {'rho': 1.0, 'theta': 1.0} + linear_solver = CompressibleSolver(eqns, alpha, tau_values=tau_values) # Time stepper stepper = SemiImplicitQuasiNewton( eqns, io, transported_fields, transport_methods, - linear_solver=linear_solver, alpha=alpha + linear_solver=linear_solver, alpha=alpha, spinup_steps=spinup_steps ) # ------------------------------------------------------------------------ # @@ -189,7 +178,8 @@ def mountain_hydrostatic( bottom_boundary = Constant(exner_surf, domain=mesh) logger.info(f'Solving hydrostatic with bottom Exner of {exner_surf}') compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary + eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary, + solve_for_rho=True ) # Solve hydrostatic balance again, but now use minimum value from first @@ -243,7 +233,7 @@ def mountain_hydrostatic( theta0.assign(theta_b) rho0.assign(rho_b) - u0.project(as_vector([initial_wind, 0.0]), bcs=eqns.bcs['u']) + u0.project(as_vector([initial_wind, 0.0])) stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) @@ -268,38 +258,38 @@ def mountain_hydrostatic( '--ncolumns', help="The number of columns in the vertical slice mesh.", type=int, - default=mountain_hydrostatic_defaults['ncolumns'] + default=schaer_mountain_defaults['ncolumns'] ) parser.add_argument( '--nlayers', help="The number of layers for the mesh.", type=int, - default=mountain_hydrostatic_defaults['nlayers'] + default=schaer_mountain_defaults['nlayers'] ) parser.add_argument( '--dt', help="The time step in seconds.", type=float, - default=mountain_hydrostatic_defaults['dt'] + default=schaer_mountain_defaults['dt'] ) parser.add_argument( "--tmax", help="The end time for the simulation in seconds.", type=float, - default=mountain_hydrostatic_defaults['tmax'] + default=schaer_mountain_defaults['tmax'] ) parser.add_argument( '--dumpfreq', help="The frequency at which to dump field output.", type=int, - default=mountain_hydrostatic_defaults['dumpfreq'] + default=schaer_mountain_defaults['dumpfreq'] ) parser.add_argument( '--dirname', help="The name of the directory to write to.", type=str, - default=mountain_hydrostatic_defaults['dirname'] + default=schaer_mountain_defaults['dirname'] ) args, unknown = parser.parse_known_args() - mountain_hydrostatic(**vars(args)) + schaer_mountain(**vars(args)) diff --git a/examples/compressible_euler/skamarock_klemp_hydrostatic.py b/examples/compressible_euler/skamarock_klemp_hydrostatic.py deleted file mode 100644 index f75ad9b10..000000000 --- a/examples/compressible_euler/skamarock_klemp_hydrostatic.py +++ /dev/null @@ -1,204 +0,0 @@ -""" -This example uses the hydrostatic compressible Euler equations to solve the -vertical slice gravity wave test case of Skamarock and Klemp, 1994: -``Efficiency and Accuracy of the Klemp-Wilhelmson Time-Splitting Technique'', -MWR. - -Potential temperature is transported using SUPG, and the degree 1 elements are -used. This also uses a mesh which is one cell thick in the y-direction. -""" - -from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter -from firedrake import ( - as_vector, SpatialCoordinate, PeriodicRectangleMesh, ExtrudedMesh, exp, sin, - Function, pi -) -from gusto import ( - Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - TrapeziumRule, SUPGOptions, CourantNumber, Perturbation, - CompressibleParameters, HydrostaticCompressibleEulerEquations, - CompressibleSolver, compressible_hydrostatic_balance -) - -skamarock_klemp_hydrostatic_defaults = { - 'ncolumns': 150, - 'nlayers': 10, - 'dt': 25.0, - 'tmax': 60000., - 'dumpfreq': 1200, - 'dirname': 'skamarock_klemp_hydrostatic' -} - - -def skamarock_klemp_hydrostatic( - ncolumns=skamarock_klemp_hydrostatic_defaults['ncolumns'], - nlayers=skamarock_klemp_hydrostatic_defaults['nlayers'], - dt=skamarock_klemp_hydrostatic_defaults['dt'], - tmax=skamarock_klemp_hydrostatic_defaults['tmax'], - dumpfreq=skamarock_klemp_hydrostatic_defaults['dumpfreq'], - dirname=skamarock_klemp_hydrostatic_defaults['dirname'] -): - - # ------------------------------------------------------------------------ # - # Test case parameters - # ------------------------------------------------------------------------ # - - domain_width = 6.0e6 # Width of domain in x direction (m) - domain_length = 1.0e4 # Length of domain in y direction (m) - domain_height = 1.0e4 # Height of domain (m) - Tsurf = 300. # Temperature at surface (K) - wind_initial = 20. # Initial wind in x direction (m/s) - pert_width = 5.0e3 # Width parameter of perturbation (m) - deltaTheta = 1.0e-2 # Magnitude of theta perturbation (K) - N = 0.01 # Brunt-Vaisala frequency (1/s) - Omega = 0.5e-4 # Planetary rotation rate (1/s) - pressure_gradient_y = -1.0e-4*20 # Prescribed force in y direction (m/s^2) - - # ------------------------------------------------------------------------ # - # Our settings for this set up - # ------------------------------------------------------------------------ # - - element_order = 1 - - # ------------------------------------------------------------------------ # - # Set up model objects - # ------------------------------------------------------------------------ # - - # Domain -- 3D volume mesh - base_mesh = PeriodicRectangleMesh( - ncolumns, 1, domain_width, domain_length, quadrilateral=True - ) - mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) - domain = Domain(mesh, dt, "RTCF", element_order) - - # Equation - parameters = CompressibleParameters(Omega=Omega) - balanced_pg = as_vector((0., pressure_gradient_y, 0.)) - eqns = HydrostaticCompressibleEulerEquations( - domain, parameters, extra_terms=[("u", balanced_pg)] - ) - - # I/O - output = OutputParameters( - dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=False, - dumplist=['u'], - ) - diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] - io = IO(domain, output, diagnostic_fields=diagnostic_fields) - - # Transport schemes - theta_opts = SUPGOptions() - transported_fields = [ - TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho"), - SSPRK3(domain, "theta", options=theta_opts) - ] - transport_methods = [ - DGUpwind(eqns, "u"), - DGUpwind(eqns, "rho"), - DGUpwind(eqns, "theta", ibp=theta_opts.ibp) - ] - - # Linear solver - linear_solver = CompressibleSolver(eqns) - - # Time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, transport_methods, - linear_solver=linear_solver - ) - - # ------------------------------------------------------------------------ # - # Initial conditions - # ------------------------------------------------------------------------ # - - u0 = stepper.fields("u") - rho0 = stepper.fields("rho") - theta0 = stepper.fields("theta") - - # spaces - Vt = domain.spaces("theta") - Vr = domain.spaces("DG") - - # Thermodynamic constants required for setting initial conditions - # and reference profiles - g = parameters.g - - x, _, z = SpatialCoordinate(mesh) - - # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) - thetab = Tsurf*exp(N**2*z/g) - - theta_b = Function(Vt).interpolate(thetab) - rho_b = Function(Vr) - - theta_pert = ( - deltaTheta * sin(pi*z/domain_height) - / (1 + (x - domain_width/2)**2 / pert_width**2) - ) - theta0.interpolate(theta_b + theta_pert) - - compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) - - rho0.assign(rho_b) - u0.project(as_vector([wind_initial, 0.0, 0.0])) - - stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - - # ------------------------------------------------------------------------ # - # Run - # ------------------------------------------------------------------------ # - - stepper.run(t=0, tmax=tmax) - -# ---------------------------------------------------------------------------- # -# MAIN -# ---------------------------------------------------------------------------- # - - -if __name__ == "__main__": - - parser = ArgumentParser( - description=__doc__, - formatter_class=ArgumentDefaultsHelpFormatter - ) - parser.add_argument( - '--ncolumns', - help="The number of columns in the vertical slice mesh.", - type=int, - default=skamarock_klemp_hydrostatic_defaults['ncolumns'] - ) - parser.add_argument( - '--nlayers', - help="The number of layers for the mesh.", - type=int, - default=skamarock_klemp_hydrostatic_defaults['nlayers'] - ) - parser.add_argument( - '--dt', - help="The time step in seconds.", - type=float, - default=skamarock_klemp_hydrostatic_defaults['dt'] - ) - parser.add_argument( - "--tmax", - help="The end time for the simulation in seconds.", - type=float, - default=skamarock_klemp_hydrostatic_defaults['tmax'] - ) - parser.add_argument( - '--dumpfreq', - help="The frequency at which to dump field output.", - type=int, - default=skamarock_klemp_hydrostatic_defaults['dumpfreq'] - ) - parser.add_argument( - '--dirname', - help="The name of the directory to write to.", - type=str, - default=skamarock_klemp_hydrostatic_defaults['dirname'] - ) - args, unknown = parser.parse_known_args() - - skamarock_klemp_hydrostatic(**vars(args)) diff --git a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py index 057d34ade..5ec8ab12b 100644 --- a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py +++ b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py @@ -4,6 +4,10 @@ ``Efficiency and Accuracy of the Klemp-Wilhelmson Time-Splitting Technique'', MWR. +The domain is smaller than the "hydrostatic" gravity wave test, so that there +is difference between the hydrostatic and non-hydrostatic solutions. The test +can be run with and without a hydrostatic switch. + Potential temperature is transported using SUPG, and the degree 1 elements are used. """ @@ -19,19 +23,20 @@ import numpy as np from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - SUPGOptions, CourantNumber, Perturbation, Gradient, - CompressibleParameters, CompressibleEulerEquations, CompressibleSolver, - compressible_hydrostatic_balance, logger, RichardsonNumber, - RungeKuttaFormulation, SubcyclingOptions + logger, SUPGOptions, Perturbation, CompressibleParameters, + CompressibleEulerEquations, HydrostaticCompressibleEulerEquations, + compressible_hydrostatic_balance, RungeKuttaFormulation, CompressibleSolver, + SubcyclingOptions, hydrostatic_parameters ) skamarock_klemp_nonhydrostatic_defaults = { 'ncolumns': 150, 'nlayers': 10, 'dt': 6.0, - 'tmax': 3600., - 'dumpfreq': 300, - 'dirname': 'skamarock_klemp_nonhydrostatic' + 'tmax': 3000., + 'dumpfreq': 250, + 'dirname': 'skamarock_klemp_nonhydrostatic', + 'hydrostatic': False } @@ -41,7 +46,8 @@ def skamarock_klemp_nonhydrostatic( dt=skamarock_klemp_nonhydrostatic_defaults['dt'], tmax=skamarock_klemp_nonhydrostatic_defaults['tmax'], dumpfreq=skamarock_klemp_nonhydrostatic_defaults['dumpfreq'], - dirname=skamarock_klemp_nonhydrostatic_defaults['dirname'] + dirname=skamarock_klemp_nonhydrostatic_defaults['dirname'], + hydrostatic=skamarock_klemp_nonhydrostatic_defaults['hydrostatic'] ): # ------------------------------------------------------------------------ # @@ -73,18 +79,25 @@ def skamarock_klemp_nonhydrostatic( # Equation parameters = CompressibleParameters() - eqns = CompressibleEulerEquations(domain, parameters) + if hydrostatic: + eqns = HydrostaticCompressibleEulerEquations(domain, parameters) + else: + eqns = CompressibleEulerEquations(domain, parameters) # I/O points_x = np.linspace(0., domain_width, 100) points_z = [domain_height/2.] points = np.array([p for p in itertools.product(points_x, points_z)]) + # Adjust default directory name + if hydrostatic and dirname == skamarock_klemp_nonhydrostatic_defaults['dirname']: + dirname = f'hyd_switch_{dirname}' + # Dumping point data using legacy PointDataOutput is not supported in parallel if COMM_WORLD.size == 1: output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq, - dump_vtus=True, dump_nc=False, + dump_vtus=False, dump_nc=True, point_data=[('theta_perturbation', points)], ) else: @@ -94,14 +107,10 @@ def skamarock_klemp_nonhydrostatic( ) output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq, - dump_vtus=True, dump_nc=True, + dump_vtus=False, dump_nc=True, ) - diagnostic_fields = [ - CourantNumber(), Gradient('u'), Perturbation('theta'), - Gradient('theta_perturbation'), Perturbation('rho'), - RichardsonNumber('theta', parameters.g/Tsurf), Gradient('theta') - ] + diagnostic_fields = [Perturbation('theta')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes @@ -125,7 +134,13 @@ def skamarock_klemp_nonhydrostatic( ] # Linear solver - linear_solver = CompressibleSolver(eqns) + if hydrostatic: + linear_solver = CompressibleSolver( + eqns, solver_parameters=hydrostatic_parameters, + overwrite_solver_parameters=True + ) + else: + linear_solver = CompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton( @@ -223,6 +238,16 @@ def skamarock_klemp_nonhydrostatic( type=str, default=skamarock_klemp_nonhydrostatic_defaults['dirname'] ) + parser.add_argument( + '--hydrostatic', + help=( + "Whether to use the hydrostatic switch to emulate the " + + "hydrostatic equations. Otherwise use the full non-hydrostatic" + + "equations." + ), + action="store_true", + default=skamarock_klemp_nonhydrostatic_defaults['hydrostatic'] + ) args, unknown = parser.parse_known_args() skamarock_klemp_nonhydrostatic(**vars(args)) diff --git a/examples/compressible_euler/test_compressible_euler_examples.py b/examples/compressible_euler/test_compressible_euler_examples.py index 384824e12..d74d5dade 100644 --- a/examples/compressible_euler/test_compressible_euler_examples.py +++ b/examples/compressible_euler/test_compressible_euler_examples.py @@ -46,66 +46,46 @@ def test_dry_bryan_fritsch_parallel(): test_dry_bryan_fritsch() -# Hydrostatic equations not currently working -@pytest.mark.xfail -def test_mountain_hydrostatic(): - from mountain_hydrostatic import mountain_hydrostatic - test_name = 'mountain_hydrostatic' - mountain_hydrostatic( - ncolumns=20, - nlayers=10, - dt=5.0, - tmax=50.0, - dumpfreq=10, - dirname=make_dirname(test_name) - ) - - -# Hydrostatic equations not currently working -@pytest.mark.xfail -@pytest.mark.parallel(nprocs=4) -def test_mountain_hydrostatic_parallel(): - test_mountain_hydrostatic() - - -# Hydrostatic equations not currently working -@pytest.mark.xfail -def test_skamarock_klemp_hydrostatic(): - from skamarock_klemp_hydrostatic import skamarock_klemp_hydrostatic - test_name = 'skamarock_klemp_hydrostatic' - skamarock_klemp_hydrostatic( +def test_skamarock_klemp_nonhydrostatic(): + from skamarock_klemp_nonhydrostatic import skamarock_klemp_nonhydrostatic + test_name = 'skamarock_klemp_nonhydrostatic' + skamarock_klemp_nonhydrostatic( ncolumns=30, nlayers=5, dt=6.0, tmax=60.0, dumpfreq=10, - dirname=make_dirname(test_name) + dirname=make_dirname(test_name), + hydrostatic=False ) -# Hydrostatic equations not currently working -@pytest.mark.xfail @pytest.mark.parallel(nprocs=2) -def test_skamarock_klemp_hydrostatic_parallel(): - test_skamarock_klemp_hydrostatic() +def test_skamarock_klemp_nonhydrostatic_parallel(): + test_skamarock_klemp_nonhydrostatic() -def test_skamarock_klemp_nonhydrostatic(): +# Hydrostatic equations not currently working +@pytest.mark.xfail +def test_hyd_switch_skamarock_klemp_nonhydrostatic(): from skamarock_klemp_nonhydrostatic import skamarock_klemp_nonhydrostatic - test_name = 'skamarock_klemp_nonhydrostatic' + test_name = 'hyd_switch_skamarock_klemp_nonhydrostatic' skamarock_klemp_nonhydrostatic( ncolumns=30, nlayers=5, dt=6.0, tmax=60.0, dumpfreq=10, - dirname=make_dirname(test_name) + dirname=make_dirname(test_name), + hydrostatic=True ) +# Hydrostatic equations not currently working +@pytest.mark.xfail @pytest.mark.parallel(nprocs=2) -def test_skamarock_klemp_nonhydrostatic_parallel(): - test_skamarock_klemp_nonhydrostatic() +def test_hyd_switch_skamarock_klemp_nonhydrostatic_parallel(): + test_hyd_switch_skamarock_klemp_nonhydrostatic() def test_straka_bubble(): diff --git a/examples/compressible_euler/unsaturated_bubble.py b/examples/compressible_euler/unsaturated_bubble.py index 394e15f97..59a911cb9 100644 --- a/examples/compressible_euler/unsaturated_bubble.py +++ b/examples/compressible_euler/unsaturated_bubble.py @@ -216,16 +216,11 @@ def unsaturated_bubble( R_v = eqns.parameters.R_v epsilon = R_d / R_v - # make expressions for determining water_v0 - exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) - p = thermodynamics.p(eqns.parameters, exner) - T = thermodynamics.T(eqns.parameters, theta0, exner, water_v0) - r_v_expr = thermodynamics.r_v(eqns.parameters, rel_hum, T, p) - # make expressions to evaluate residual exner_expr = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) p_expr = thermodynamics.p(eqns.parameters, exner_expr) T_expr = thermodynamics.T(eqns.parameters, theta0, exner_expr, water_v0) + water_v_expr = thermodynamics.r_v(eqns.parameters, rel_hum, T_expr, p_expr) rel_hum_expr = thermodynamics.RH(eqns.parameters, water_v0, T_expr, p_expr) rel_hum_eval = Function(Vt) @@ -247,7 +242,7 @@ def unsaturated_bubble( # first solve for r_v for _ in range(max_inner_solve_count): - water_v_eval.interpolate(r_v_expr) + water_v_eval.interpolate(water_v_expr) water_v0.assign(water_v0 * (1 - delta) + delta * water_v_eval) # compute theta_vd diff --git a/examples/shallow_water/linear_thermal_galewsky_jet.py b/examples/shallow_water/linear_thermal_galewsky_jet.py new file mode 100644 index 000000000..a4fd68631 --- /dev/null +++ b/examples/shallow_water/linear_thermal_galewsky_jet.py @@ -0,0 +1,220 @@ +""" +A linearised form of the steady thermal Galewsky jet. The initial conditions are +taken from Hartney et al, 2024: ``A compatible finite element discretisation +for moist shallow water equations'' (without the perturbation). + +This uses an icosahedral mesh of the sphere. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + SpatialCoordinate, pi, assemble, dx, Constant, exp, conditional, Function, + cos +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, DefaultTransport, + DGUpwind, ForwardEuler, ShallowWaterParameters, NumericalIntegral, + LinearThermalShallowWaterEquations, GeneralIcosahedralSphereMesh, + ZonalComponent, ThermalSWSolver, lonlatr_from_xyz, xyz_vector_from_lonlatr, + RelativeVorticity, MeridionalComponent +) + +import numpy as np + +linear_thermal_galewsky_jet_defaults = { + 'ncells_per_edge': 12, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 6.*24.*60.*60., # 6 days + 'dumpfreq': 96, # once per day with default options + 'dirname': 'linear_thermal_galewsky' +} + + +def linear_thermal_galewsky_jet( + ncells_per_edge=linear_thermal_galewsky_jet_defaults['ncells_per_edge'], + dt=linear_thermal_galewsky_jet_defaults['dt'], + tmax=linear_thermal_galewsky_jet_defaults['tmax'], + dumpfreq=linear_thermal_galewsky_jet_defaults['dumpfreq'], + dirname=linear_thermal_galewsky_jet_defaults['dirname'] +): + # ----------------------------------------------------------------- # + # Parameters for test case + # ----------------------------------------------------------------- # + + R = 6371220. # planetary radius (m) + H = 10000. # reference depth (m) + u_max = 80.0 # Max amplitude of the zonal wind (m/s) + phi0 = pi/7. # latitude of the southern boundary of the jet (rad) + phi1 = pi/2. - phi0 # latitude of the northern boundary of the jet (rad) + db = 1.0 # diff in buoyancy between equator and poles (m/s^2) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ----------------------------------------------------------------- # + # Set up model objects + # ----------------------------------------------------------------- # + + # Domain + mesh = GeneralIcosahedralSphereMesh(R, ncells_per_edge, degree=2) + xyz = SpatialCoordinate(mesh) + domain = Domain(mesh, dt, 'BDM', element_order) + + # Equation + parameters = ShallowWaterParameters(H=H) + Omega = parameters.Omega + fexpr = 2*Omega*xyz[2]/R + eqns = LinearThermalShallowWaterEquations(domain, parameters, fexpr=fexpr) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dump_vtus=False, + dumplist=['D', 'b'] + ) + diagnostic_fields = [ + ZonalComponent('u'), MeridionalComponent('u'), RelativeVorticity() + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transport_schemes = [ForwardEuler(domain, "D")] + transport_methods = [DefaultTransport(eqns, "D"), DGUpwind(eqns, "b")] + + # Linear solver + linear_solver = ThermalSWSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transport_schemes, transport_methods, + linear_solver=linear_solver + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0_field = stepper.fields("u") + D0_field = stepper.fields("D") + b0_field = stepper.fields("b") + + # Parameters + g = parameters.g + Omega = parameters.Omega + e_n = np.exp(-4./((phi1-phi0)**2)) + + _, lat, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) + lat_VD = Function(D0_field.function_space()).interpolate(lat) + + # ------------------------------------------------------------------------ # + # Obtain u and D (by integration of analytic expression) + # ------------------------------------------------------------------------ # + + # Buoyancy + bexpr = g - db*cos(lat) + + # Wind -- UFL expression + u_zonal = conditional( + lat <= phi0, 0.0, + conditional( + lat >= phi1, 0.0, + u_max / e_n * exp(1.0 / ((lat - phi0) * (lat - phi1))) + ) + ) + uexpr = xyz_vector_from_lonlatr(u_zonal, Constant(0.0), Constant(0.0), xyz) + + # Numpy function + def u_func(y): + u_array = np.where( + y <= phi0, 0.0, + np.where( + y >= phi1, 0.0, + u_max / e_n * np.exp(1.0 / ((y - phi0) * (y - phi1))) + ) + ) + return u_array + + # Function for depth field in terms of u function + def h_func(y): + h_array = u_func(y)*R/g*( + 2*Omega*np.sin(y) + ) + + return h_array + + # Find h from numerical integral + D0_integral = Function(D0_field.function_space()) + h_integral = NumericalIntegral(-pi/2, pi/2) + h_integral.tabulate(h_func) + D0_integral.dat.data[:] = h_integral.evaluate_at(lat_VD.dat.data[:]) + Dexpr = H - D0_integral + + # Obtain fields + u0_field.project(uexpr) + D0_field.interpolate(Dexpr) + + # Adjust mean value of initial D + C = Function(D0_field.function_space()).assign(Constant(1.0)) + area = assemble(C*dx) + Dmean = assemble(D0_field*dx)/area + D0_field -= Dmean + D0_field += Constant(H) + b0_field.interpolate(bexpr) + + # Set reference profiles + Dbar = Function(D0_field.function_space()).assign(H) + bbar = Function(b0_field.function_space()).interpolate(g) + stepper.set_reference_profiles([('D', Dbar), ('b', bbar)]) + + # ----------------------------------------------------------------- # + # Run + # ----------------------------------------------------------------- # + + stepper.run(t=0, tmax=tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncells_per_edge', + help="The number of cells per edge of icosahedron", + type=int, + default=linear_thermal_galewsky_jet_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=linear_thermal_galewsky_jet_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=linear_thermal_galewsky_jet_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=linear_thermal_galewsky_jet_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=linear_thermal_galewsky_jet_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + linear_thermal_galewsky_jet(**vars(args)) diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py index 07fc7e96c..376a76c64 100644 --- a/examples/shallow_water/moist_convective_williamson_2.py +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -103,7 +103,7 @@ def moist_convect_williamson_2( # define saturation function def sat_func(x_in): - h = x_in.split()[1] + h = x_in.subfunctions[1] numerator = ( theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma)) diff --git a/examples/shallow_water/moist_thermal_equivb_gw.py b/examples/shallow_water/moist_thermal_equivb_gw.py new file mode 100644 index 000000000..912e9deba --- /dev/null +++ b/examples/shallow_water/moist_thermal_equivb_gw.py @@ -0,0 +1,230 @@ +""" +A gravity wave on the sphere, solved with the moist thermal shallow water +equations. The initial conditions are saturated and cloudy everywhere. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + SpatialCoordinate, pi, sqrt, min_value, cos, Constant, Function, exp, sin +) +from gusto import ( + Domain, IO, OutputParameters, DGUpwind, ShallowWaterParameters, + ThermalShallowWaterEquations, lonlatr_from_xyz, MeridionalComponent, + GeneralIcosahedralSphereMesh, SubcyclingOptions, ZonalComponent, + PartitionedCloud, RungeKuttaFormulation, SSPRK3, ThermalSWSolver, + SemiImplicitQuasiNewton, xyz_vector_from_lonlatr +) + +moist_thermal_gw_defaults = { + 'ncells_per_edge': 12, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 5.*24.*60.*60., # 5 days + 'dumpfreq': 48, # dump twice per day + 'dirname': 'moist_thermal_equivb_gw' +} + + +def moist_thermal_gw( + ncells_per_edge=moist_thermal_gw_defaults['ncells_per_edge'], + dt=moist_thermal_gw_defaults['dt'], + tmax=moist_thermal_gw_defaults['tmax'], + dumpfreq=moist_thermal_gw_defaults['dumpfreq'], + dirname=moist_thermal_gw_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + radius = 6371220. # planetary radius (m) + mean_depth = 5960. # reference depth (m) + q0 = 0.0115 # saturation curve coefficient (kg/kg) + beta2 = 9.80616*10 # thermal feedback coefficient (m/s^2) + nu = 1.5 # dimensionless parameter in saturation curve + R0 = pi/9. # radius of perturbation (rad) + lamda_c = -pi/2. # longitudinal centre of perturbation (rad) + phi_c = pi/6. # latitudinal centre of perturbation (rad) + phi_0 = 3.0e4 # scale factor for poleward buoyancy gradient + epsilon = 1/300 # linear air expansion coeff (1/K) + u_max = 20. # max amplitude of the zonal wind (m/s) + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + degree = 1 + domain = Domain(mesh, dt, "BDM", degree) + xyz = SpatialCoordinate(mesh) + + # Equation parameters + parameters = ShallowWaterParameters(H=mean_depth, q0=q0, nu=nu, beta2=beta2) + Omega = parameters.Omega + fexpr = 2*Omega*xyz[2]/radius + + # Equation + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, + equivalent_buoyancy=True + ) + + # IO + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dump_vtus=False, + dumplist=['b_e', 'D'] + ) + diagnostic_fields = [ + ZonalComponent('u'), MeridionalComponent('u'), PartitionedCloud(eqns) + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + transport_methods = [ + DGUpwind(eqns, field_name) for field_name in eqns.field_names + ] + + linear_solver = ThermalSWSolver(eqns) + + # ------------------------------------------------------------------------ # + # Timestepper + # ------------------------------------------------------------------------ # + + subcycling_opts = SubcyclingOptions(subcycle_by_courant=0.25) + transported_fields = [ + SSPRK3(domain, "u", subcycling_options=subcycling_opts), + SSPRK3( + domain, "D", subcycling_options=subcycling_opts, + rk_formulation=RungeKuttaFormulation.linear + ), + SSPRK3(domain, "b_e", subcycling_options=subcycling_opts), + SSPRK3(domain, "q_t", subcycling_options=subcycling_opts), + ] + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + D0 = stepper.fields("D") + b0 = stepper.fields("b_e") + qt0 = stepper.fields("q_t") + + lamda, phi, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) + + # Velocity -- a solid body rotation + uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, xyz) + + # Buoyancy -- dependent on latitude + g = parameters.g + w = Omega*radius*u_max + (u_max**2)/2 + sigma = w/10 + theta_0 = epsilon*phi_0**2 + numerator = ( + theta_0 + sigma*((cos(phi))**2) * ( + (w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma) + ) + ) + denominator = ( + phi_0**2 + (w + sigma)**2 + * (sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2 + ) + theta = numerator / denominator + b_guess = parameters.g * (1 - theta) + + # Depth -- in balance with the contribution of a perturbation + Dexpr = mean_depth - (1/g)*(w + sigma)*((sin(phi))**2) + + # Perturbation + lsq = (lamda - lamda_c)**2 + thsq = (phi - phi_c)**2 + rsq = min_value(R0**2, lsq+thsq) + r = sqrt(rsq) + pert = 2000 * (1 - r/R0) + Dexpr += pert + + # Actual initial buoyancy is specified through equivalent buoyancy + q_t = 0.03 + b_init = Function(b0.function_space()).interpolate(b_guess) + b_e_init = Function(b0.function_space()).interpolate(b_init - beta2*q_t) + q_v_init = Function(qt0.function_space()).interpolate(q_t) + + # Iterate to obtain equivalent buoyancy and saturation water vapour + n_iterations = 10 + + for _ in range(n_iterations): + q_sat_expr = q0*mean_depth/Dexpr * exp(nu*(1-b_e_init/g)) + dq_sat_dq_v_expr = nu*beta2/g*q_sat_expr + q_v_init.interpolate(q_v_init - (q_sat_expr - q_v_init)/(dq_sat_dq_v_expr - 1.0)) + b_e_init.interpolate(b_init - beta2*q_v_init) + + # Water vapour set to saturation amount + vexpr = q0*mean_depth/Dexpr * exp(nu*(1-b_e_init/g)) + + # Back out the initial buoyancy using b_e and q_v + bexpr = b_e_init + beta2*vexpr + + u0.project(uexpr) + D0.interpolate(Dexpr) + b0.interpolate(bexpr) + qt0.interpolate(Constant(0.03)) + + # Set reference profiles + Dbar = Function(D0.function_space()).assign(mean_depth) + bbar = Function(b0.function_space()).interpolate(bexpr) + stepper.set_reference_profiles([('D', Dbar), ('b_e', bbar)]) + + # ----------------------------------------------------------------- # + # Run + # ----------------------------------------------------------------- # + + stepper.run(t=0, tmax=tmax) + + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncells_per_edge', + help="The number of cells per edge of icosahedron", + type=int, + default=moist_thermal_gw_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=moist_thermal_gw_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=moist_thermal_gw_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=moist_thermal_gw_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=moist_thermal_gw_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + moist_thermal_gw(**vars(args)) diff --git a/examples/shallow_water/moist_thermal_williamson_5.py b/examples/shallow_water/moist_thermal_williamson_5.py index 9ce7da77e..376cc94ad 100644 --- a/examples/shallow_water/moist_thermal_williamson_5.py +++ b/examples/shallow_water/moist_thermal_williamson_5.py @@ -21,7 +21,7 @@ ) from gusto import ( Domain, IO, OutputParameters, Timestepper, RK4, DGUpwind, - ShallowWaterParameters, ShallowWaterEquations, Sum, + ShallowWaterParameters, ThermalShallowWaterEquations, Sum, lonlatr_from_xyz, InstantRain, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, GeneralIcosahedralSphereMesh, RelativeVorticity, ZonalComponent, MeridionalComponent @@ -99,8 +99,8 @@ def moist_thermal_williamson_5( tracers = [ WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG') ] - eqns = ShallowWaterEquations( - domain, parameters, fexpr=fexpr, bexpr=tpexpr, thermal=True, + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, topog_expr=tpexpr, active_tracers=tracers, u_transport_option=u_eqn_type ) @@ -121,14 +121,14 @@ def q_sat(b, D): # Function to pass to physics (takes mixed function as argument) def phys_sat_func(x_in): - D = x_in.split()[1] - b = x_in.split()[2] + D = x_in.subfunctions[1] + b = x_in.subfunctions[2] return q_sat(b, D) # Feedback proportionality is dependent on D and b def gamma_v(x_in): - D = x_in.split()[1] - b = x_in.split()[2] + D = x_in.subfunctions[1] + b = x_in.subfunctions[2] return 1.0 / (1.0 + nu*beta2/g*q_sat(b, D)) SWSaturationAdjustment( diff --git a/examples/shallow_water/test_shallow_water_examples.py b/examples/shallow_water/test_shallow_water_examples.py index eb64e94e0..7a96c6316 100644 --- a/examples/shallow_water/test_shallow_water_examples.py +++ b/examples/shallow_water/test_shallow_water_examples.py @@ -127,3 +127,37 @@ def test_williamson_5(): @pytest.mark.parallel(nprocs=4) def test_williamson_5_parallel(): test_williamson_5() + + +def test_linear_thermal_galewsky_jet(): + from linear_thermal_galewsky_jet import linear_thermal_galewsky_jet + test_name = 'linear_thermal_galewsky_jet' + linear_thermal_galewsky_jet( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def test_linear_thermal_galewsky_jet_parallel(): + test_linear_thermal_galewsky_jet() + + +def test_moist_thermal_eqiuvb_gw(): + from moist_thermal_equivb_gw import moist_thermal_gw + test_name = 'moist_thermal_gw' + moist_thermal_gw( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def moist_thermal_eqiuvb_gw_parallel(): + test_moist_thermal_eqiuvb_gw() diff --git a/examples/shallow_water/thermal_williamson_2.py b/examples/shallow_water/thermal_williamson_2.py index d2c4e4e29..7adc305dc 100644 --- a/examples/shallow_water/thermal_williamson_2.py +++ b/examples/shallow_water/thermal_williamson_2.py @@ -15,7 +15,7 @@ from firedrake import Function, SpatialCoordinate, sin, cos from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, + TrapeziumRule, ShallowWaterParameters, ThermalShallowWaterEquations, RelativeVorticity, PotentialVorticity, SteadyStateError, ZonalComponent, MeridionalComponent, ThermalSWSolver, xyz_vector_from_lonlatr, lonlatr_from_xyz, GeneralIcosahedralSphereMesh, @@ -71,8 +71,8 @@ def thermal_williamson_2( params = ShallowWaterParameters(H=mean_depth, g=g) Omega = params.Omega fexpr = 2*Omega*z/radius - eqns = ShallowWaterEquations( - domain, params, fexpr=fexpr, u_transport_option=u_eqn_type, thermal=True + eqns = ThermalShallowWaterEquations( + domain, params, fexpr=fexpr, u_transport_option=u_eqn_type ) # IO diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index b67e5b079..efe41b183 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -73,7 +73,8 @@ def williamson_5( rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2) r = sqrt(rsq) tpexpr = mountain_height * (1 - r/R0) - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + topog_expr=tpexpr) # I/O output = OutputParameters( diff --git a/figures/compressible_euler/schaer_mountain_final.png b/figures/compressible_euler/schaer_mountain_final.png new file mode 100644 index 000000000..a5289cca3 Binary files /dev/null and b/figures/compressible_euler/schaer_mountain_final.png differ diff --git a/figures/compressible_euler/schaer_mountain_initial.png b/figures/compressible_euler/schaer_mountain_initial.png new file mode 100644 index 000000000..5e78c469a Binary files /dev/null and b/figures/compressible_euler/schaer_mountain_initial.png differ diff --git a/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png b/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png index c1187a1cd..e76a82582 100644 Binary files a/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png and b/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png differ diff --git a/figures/shallow_water/linear_thermal_galewsky_final.png b/figures/shallow_water/linear_thermal_galewsky_final.png new file mode 100644 index 000000000..8400dc6fd Binary files /dev/null and b/figures/shallow_water/linear_thermal_galewsky_final.png differ diff --git a/figures/shallow_water/linear_thermal_galewsky_initial.png b/figures/shallow_water/linear_thermal_galewsky_initial.png new file mode 100644 index 000000000..5fdc323a0 Binary files /dev/null and b/figures/shallow_water/linear_thermal_galewsky_initial.png differ diff --git a/figures/shallow_water/moist_thermal_equivb_gw_final.png b/figures/shallow_water/moist_thermal_equivb_gw_final.png new file mode 100644 index 000000000..e5aaa2330 Binary files /dev/null and b/figures/shallow_water/moist_thermal_equivb_gw_final.png differ diff --git a/figures/shallow_water/moist_thermal_equivb_gw_initial.png b/figures/shallow_water/moist_thermal_equivb_gw_initial.png new file mode 100644 index 000000000..94468de03 Binary files /dev/null and b/figures/shallow_water/moist_thermal_equivb_gw_initial.png differ diff --git a/gusto/core/configuration.py b/gusto/core/configuration.py index 133c2c3f6..40852033a 100644 --- a/gusto/core/configuration.py +++ b/gusto/core/configuration.py @@ -79,7 +79,11 @@ def __setattr__(self, name, value): # Almost all parameters should be Constants -- but there are some # specific exceptions which should be kept as integers - if type(value) in [float, int] and name not in ['dumpfreq', 'pddumpfreq', 'chkptfreq']: + non_constants = [ + 'dumpfreq', 'pddumpfreq', 'chkptfreq', + 'fixed_subcycles', 'max_subcycles', 'subcycle_by_courant' + ] + if type(value) in [float, int] and name not in non_constants: object.__setattr__(self, name, Constant(value)) else: object.__setattr__(self, name, value) @@ -149,6 +153,15 @@ class ShallowWaterParameters(Configuration): g = 9.80616 Omega = 7.292e-5 # rotation rate H = None # mean depth + # Factor that multiplies the vapour in the equivalent buoyancy + # formulation of the thermal shallow water equations + beta2 = None + # Scaling factor for the saturation function in the equivalent buoyancy + # formulation of the thermal shallow water equations + nu = None + # Scaling factor for the saturation function in the equivalent buoyancy + # formulation of the thermal shallow water equations + q0 = None class WrapperOptions(Configuration, metaclass=ABCMeta): @@ -256,6 +269,19 @@ class BoundaryLayerParameters(Configuration): mu = 100. # Interior penalty coefficient for vertical diffusion +class HeldSuarezParameters(Configuration): + """ + Parameters used in the default configuration for the Held Suarez test case. + """ + T0stra = 200 # Stratosphere temp + T0surf = 315 # Surface temperature at equator + T0horiz = 60 # Equator to pole temperature difference + T0vert = 10 # Stability parameter + sigmab = 0.7 # Height of the boundary layer + tau_d = 40 * 24 * 60 * 60 # 40 day time scale + tau_u = 4 * 24 * 60 * 60 # 4 day timescale + + class SubcyclingOptions(Configuration): """ Describes the process of subcycling a time discretisation, by dividing the diff --git a/gusto/core/coordinates.py b/gusto/core/coordinates.py index 28862e176..a2eedd488 100644 --- a/gusto/core/coordinates.py +++ b/gusto/core/coordinates.py @@ -67,7 +67,6 @@ def __init__(self, mesh, on_sphere=False, rotated_pole=None, radius=None): # -------------------------------------------------------------------- # self.chi_coords = {} # Dict of natural coords by space - self.global_chi_coords = {} # Dict of whole coords stored on first proc self.parallel_array_lims = {} # Dict of array lengths for each proc def register_space(self, domain, space_name): @@ -87,8 +86,6 @@ def register_space(self, domain, space_name): """ comm = self.mesh.comm - comm_size = comm.Get_size() - my_rank = comm.Get_rank() topological_dimension = self.mesh.topological_dimension() if space_name in self.chi_coords.keys(): @@ -112,49 +109,11 @@ def register_space(self, domain, space_name): for i in range(topological_dimension): self.chi_coords[space_name].append(Function(space).interpolate(self.coords[i])) - # -------------------------------------------------------------------- # - # Code for settings up coordinates for parallel-serial IO - # -------------------------------------------------------------------- # - - len_coords = space.dim() - my_num_dofs = len(self.chi_coords[space_name][0].dat.data_ro[:]) - - if my_rank != 0: - # Do not make full coordinate array - self.global_chi_coords[space_name] = None - self.parallel_array_lims[space_name] = None - # Find number of DoFs on this processor - comm.send(my_num_dofs, dest=0) - else: - # First processor has one large array of the global chi data - self.global_chi_coords[space_name] = np.zeros((topological_dimension, len_coords)) - # Store the limits inside this array telling us how data is partitioned - self.parallel_array_lims[space_name] = np.zeros((comm_size, 2), dtype=int) - # First processor has the first bit of data - self.parallel_array_lims[space_name][my_rank][0] = 0 - self.parallel_array_lims[space_name][my_rank][1] = my_num_dofs - 1 - # Receive number of DoFs on other processors - for procid in range(1, comm_size): - other_num_dofs = comm.recv(source=procid) - self.parallel_array_lims[space_name][procid][0] = self.parallel_array_lims[space_name][procid-1][1] + 1 - self.parallel_array_lims[space_name][procid][1] = self.parallel_array_lims[space_name][procid][0] + other_num_dofs - 1 - - # Now move coordinates to first processor - for i in range(topological_dimension): - if my_rank != 0: - # Send information to rank 0 - my_tag = comm_size*i + my_rank - comm.send(self.chi_coords[space_name][i].dat.data_ro[:], dest=0, tag=my_tag) - else: - # Rank 0 -- the receiver - (low_lim, up_lim) = self.parallel_array_lims[space_name][my_rank][:] - self.global_chi_coords[space_name][i][low_lim:up_lim+1] = self.chi_coords[space_name][i].dat.data_ro[:] - # Receive coords from each processor and put them into array - for procid in range(1, comm_size): - my_tag = comm_size*i + procid - new_coords = comm.recv(source=procid, tag=my_tag) - (low_lim, up_lim) = self.parallel_array_lims[space_name][procid][:] - self.global_chi_coords[space_name][i, low_lim:up_lim+1] = new_coords + # Determine the offsets of the local piece of data into the global array + nlocal_dofs = len(self.chi_coords[space_name][0].dat.data_ro) + start = comm.exscan(nlocal_dofs) or 0 + stop = start + nlocal_dofs + self.parallel_array_lims[space_name] = (start, stop) def get_column_data(self, field, domain): """ @@ -177,9 +136,9 @@ def get_column_data(self, field, domain): coords = self.chi_coords[space_name] data_is_3d = (len(coords) == 3) - coords_X = coords[0].dat.data - coords_Y = coords[1].dat.data if data_is_3d else None - coords_Z = coords[-1].dat.data + coords_X = coords[0].dat.data_ro + coords_Y = coords[1].dat.data_ro if data_is_3d else None + coords_Z = coords[-1].dat.data_ro # ------------------------------------------------------------------------ # # Round data to ensure sorting in dataframe is OK diff --git a/gusto/core/domain.py b/gusto/core/domain.py index 0cf23fb7f..27addd9fa 100644 --- a/gusto/core/domain.py +++ b/gusto/core/domain.py @@ -6,10 +6,12 @@ from gusto.core.coordinates import Coordinates from gusto.core.function_spaces import Spaces, check_degree_args -from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, - inner, grad, VectorFunctionSpace, Function, FunctionSpace, - perp) +from firedrake import ( + Constant, SpatialCoordinate, sqrt, CellNormal, cross, inner, grad, + VectorFunctionSpace, Function, FunctionSpace, perp, curl +) import numpy as np +from mpi4py import MPI class Domain(object): @@ -26,7 +28,7 @@ class Domain(object): """ def __init__(self, mesh, dt, family, degree=None, horizontal_degree=None, vertical_degree=None, - rotated_pole=None): + rotated_pole=None, max_quad_degree=None): """ Args: mesh (:class:`Mesh`): the model's mesh. @@ -47,6 +49,11 @@ def __init__(self, mesh, dt, family, degree=None, system. These are expressed in the original coordinate system. The longitude and latitude must be expressed in radians. Defaults to None. This is unused for non-spherical domains. + max_quad_degree (int, optional): the maximum quadrature degree to + use in certain non-linear terms (e.g. when using an expression + for the Exner pressure). Defaults to None, in which case this + will be set to the 2*p+3, where p is the maximum polynomial + degree for the DG space. Raises: ValueError: if incompatible degrees are specified (e.g. specifying @@ -78,6 +85,12 @@ def __init__(self, mesh, dt, family, degree=None, self.horizontal_degree = degree if horizontal_degree is None else horizontal_degree self.vertical_degree = degree if vertical_degree is None else vertical_degree + if max_quad_degree is None: + max_degree = max(self.horizontal_degree, self.vertical_degree) + self.max_quad_degree = 2*max_degree + 3 + else: + self.max_quad_degree = max_quad_degree + self.mesh = mesh self.family = family self.spaces = Spaces(mesh) @@ -113,12 +126,14 @@ def __init__(self, mesh, dt, family, degree=None, V = VectorFunctionSpace(mesh, "DG", sphere_degree) self.outward_normals = Function(V).interpolate(CellNormal(mesh)) self.perp = lambda u: cross(self.outward_normals, u) + self.divperp = lambda u: inner(self.outward_normals, curl(u)) else: kvec = [0.0]*dim kvec[dim-1] = 1.0 self.k = Constant(kvec) if dim == 2: self.perp = perp + self.divperp = lambda u: -u[0].dx(1) + u[1].dx(0) # -------------------------------------------------------------------- # # Construct information relating to height/radius @@ -218,25 +233,26 @@ def construct_domain_metadata(mesh, coords, on_sphere): else: raise ValueError('Unable to determine domain type') + # Determine domain properties + chi = coords.chi_coords['DG1_equispaced'] comm = mesh.comm - my_rank = comm.Get_rank() - - # Properties of domain will be determined from full coords, so need - # doing on the first processor then broadcasting to others - - if my_rank == 0: - chi = coords.global_chi_coords['DG1_equispaced'] - if not on_sphere: - metadata['domain_extent_x'] = np.max(chi[0, :]) - np.min(chi[0, :]) - if metadata['domain_type'] in ['plane', 'extruded_plane']: - metadata['domain_extent_y'] = np.max(chi[1, :]) - np.min(chi[1, :]) - if mesh.extruded: - metadata['domain_extent_z'] = np.max(chi[-1, :]) - np.min(chi[-1, :]) - - else: - metadata = {} - - # Send information to other processors - metadata = comm.bcast(metadata, root=0) + if not on_sphere: + _min_x = np.min(chi[0].dat.data_ro) + _max_x = np.max(chi[0].dat.data_ro) + min_x = comm.allreduce(_min_x, MPI.MIN) + max_x = comm.allreduce(_max_x, MPI.MAX) + metadata['domain_extent_x'] = max_x - min_x + if metadata['domain_type'] in ['plane', 'extruded_plane']: + _min_y = np.min(chi[1].dat.data_ro) + _max_y = np.max(chi[1].dat.data_ro) + min_y = comm.allreduce(_min_y, MPI.MIN) + max_y = comm.allreduce(_max_y, MPI.MAX) + metadata['domain_extent_y'] = max_y - min_y + if mesh.extruded: + _min_z = np.min(chi[-1].dat.data_ro) + _max_z = np.max(chi[-1].dat.data_ro) + min_z = comm.allreduce(_min_z, MPI.MIN) + max_z = comm.allreduce(_max_z, MPI.MAX) + metadata['domain_extent_z'] = max_z - min_z return metadata diff --git a/gusto/core/fields.py b/gusto/core/fields.py index 4d204642c..5fdf011d4 100644 --- a/gusto/core/fields.py +++ b/gusto/core/fields.py @@ -28,6 +28,12 @@ def add_field(self, name, space, subfield_names=None): variables. Defaults to None. """ + if name == 'X': + raise ValueError( + 'As "X" is used for the generic field state, it is not allowed ' + + 'as the name of a field -- please choose a different name!' + ) + value = Function(space, name=name) setattr(self, name, value) self.fields.append(value) @@ -35,6 +41,14 @@ def add_field(self, name, space, subfield_names=None): if len(space) > 1: assert len(space) == len(subfield_names) for field_name, field in zip(subfield_names, value.subfunctions): + + if field_name == 'X': + raise ValueError( + 'As "X" is used for the generic field state, it is not ' + + 'allowed as the name of a field -- please choose a ' + + 'different name!' + ) + setattr(self, field_name, field) field.rename(field_name) self.fields.append(field) @@ -114,6 +128,7 @@ def __init__(self, prognostic_fields, prescribed_fields, *fields_to_dump): self.to_pick_up = [] self._field_types = [] self._field_names = [] + self.X = prognostic_fields.np1 # Add pointers to prognostic fields for field in prognostic_fields.np1.fields: diff --git a/gusto/core/function_spaces.py b/gusto/core/function_spaces.py index 97820d454..c83a15cd6 100644 --- a/gusto/core/function_spaces.py +++ b/gusto/core/function_spaces.py @@ -4,8 +4,8 @@ """ from gusto.core.logging import logger -from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement, - TensorProductElement, interval) +from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement, VectorElement, + TensorProductElement, BrokenElement, interval) __all__ = ["Spaces", "check_degree_args"] @@ -71,6 +71,7 @@ def __init__(self, mesh): self.mesh = mesh self.extruded_mesh = hasattr(mesh, "_base_mesh") self.de_rham_complex = {} + self.continuity = {} def __call__(self, name): """ @@ -89,7 +90,7 @@ def __call__(self, name): else: raise ValueError(f'The space container has no space {name}') - def add_space(self, name, space, overwrite_space=False): + def add_space(self, name, space, overwrite_space=False, continuity=None): """ Adds a function space to the container. @@ -100,6 +101,10 @@ def add_space(self, name, space, overwrite_space=False): overwrite_space (bool, optional): Logical to allow space existing in container to be overwritten by an incoming space. Defaults to False. + continuity: Whether the space is continuous or not. For spaces on + extruded meshes, must be a dictionary with entries for the + 'horizontal' and 'vertical' continuity. + If None, defaults to value of is_cg(space). """ if hasattr(self, name) and not overwrite_space: @@ -109,9 +114,28 @@ def add_space(self, name, space, overwrite_space=False): setattr(self, name, space) - def create_space(self, name, family, degree, overwrite_space=False): + # set the continuity of the space - for extruded meshes specify both directions + if continuity is None: + continuity = is_cg(space) + if self.extruded_mesh: + self.continuity[name] = { + 'horizontal': continuity, + 'vertical': continuity + } + else: + self.continuity[name] = continuity + else: + if self.extruded_mesh: + self.continuity[name] = { + 'horizontal': continuity['horizontal'], + 'vertical': continuity['vertical'] + } + else: + self.continuity[name] = continuity + + def create_space(self, name, family, degree, overwrite_space=False, continuity=None): """ - Creates a space and adds it to the container.. + Creates a space and adds it to the container. Args: name (str): the name to give to the space. @@ -120,18 +144,18 @@ def create_space(self, name, family, degree, overwrite_space=False): overwrite_space (bool, optional): Logical to allow space existing in container to be overwritten by an incoming space. Defaults to False. + continuity: Whether the space is continuous or not. For spaces on + extruded meshes, must be a tuple for the (horizontal, vertical) + continuity. If None, defaults to value of is_cg(space). Returns: :class:`FunctionSpace`: the desired function space. """ - - if hasattr(self, name) and not overwrite_space: - raise RuntimeError(f'Space {name} already exists. If you really ' - + 'to create it then set `overwrite_space` as ' - + 'to be True') - space = FunctionSpace(self.mesh, family, degree, name=name) - setattr(self, name, space) + + self.add_space(name, space, + overwrite_space=overwrite_space, + continuity=continuity) return space def build_compatible_spaces(self, family, horizontal_degree, @@ -181,9 +205,15 @@ def build_compatible_spaces(self, family, horizontal_degree, setattr(self, "L2"+complex_name, de_rham_complex.L2) # Register L2 space as DG also setattr(self, "DG"+complex_name, de_rham_complex.L2) - if hasattr(de_rham_complex, "theta"+complex_name): + if hasattr(de_rham_complex, "theta"): setattr(self, "theta"+complex_name, de_rham_complex.theta) + # Grab the continuity information from the complex + for space_type in ("H1, HCurl", "HDiv", "L2", "DG", "theta"): + space_name = space_type + complex_name + if hasattr(de_rham_complex, space_type): + self.continuity[space_name] = de_rham_complex.continuity[space_type] + def build_dg1_equispaced(self): """ Builds the equispaced variant of the DG1 function space, which is used in @@ -198,12 +228,15 @@ def build_dg1_equispaced(self): hori_elt = FiniteElement('DG', cell, 1, variant='equispaced') vert_elt = FiniteElement('DG', interval, 1, variant='equispaced') V_elt = TensorProductElement(hori_elt, vert_elt) + continuity = {'horizontal': False, 'vertical': False} else: cell = self.mesh.ufl_cell().cellname() V_elt = FiniteElement('DG', cell, 1, variant='equispaced') + continuity = False space = FunctionSpace(self.mesh, V_elt, name='DG1_equispaced') - setattr(self, 'DG1_equispaced', space) + + self.add_space('DG1_equispaced', space, continuity=continuity) return space @@ -234,6 +267,7 @@ def __init__(self, mesh, family, horizontal_degree, vertical_degree=None, self.extruded_mesh = hasattr(mesh, '_base_mesh') self.family = family self.complex_name = complex_name + self.continuity = {} self.build_base_spaces(family, horizontal_degree, vertical_degree) self.build_compatible_spaces() @@ -303,15 +337,24 @@ def build_compatible_spaces(self): if self.extruded_mesh: # Horizontal and vertical degrees # need specifying separately. Vtheta needs returning. - Vcg = self.build_h1_space() + Vcg, continuity = self.build_h1_space() + self.continuity["H1"] = continuity setattr(self, "H1", Vcg) - Vcurl = self.build_hcurl_space() + + Vcurl, continuity = self.build_hcurl_space() + self.continuity["HCurl"] = continuity setattr(self, "HCurl", Vcurl) - Vu = self.build_hdiv_space() + + Vu, continuity = self.build_hdiv_space() + self.continuity["HDiv"] = continuity setattr(self, "HDiv", Vu) - Vdg = self.build_l2_space() + + Vdg, continuity = self.build_l2_space() + self.continuity["L2"] = continuity setattr(self, "L2", Vdg) - Vth = self.build_theta_space() + + Vth, continuity = self.build_theta_space() + self.continuity["theta"] = continuity setattr(self, "theta", Vth) return Vcg, Vcurl, Vu, Vdg, Vth @@ -320,13 +363,20 @@ def build_compatible_spaces(self): # 2D: two de Rham complexes (hcurl or hdiv) with 3 spaces # 3D: one de Rham complexes with 4 spaces # either way, build all spaces - Vcg = self.build_h1_space() + Vcg, continuity = self.build_h1_space() + self.continuity["H1"] = continuity setattr(self, "H1", Vcg) - Vcurl = self.build_hcurl_space() + + Vcurl, continuity = self.build_hcurl_space() + self.continuity["HCurl"] = continuity setattr(self, "HCurl", Vcurl) - Vu = self.build_hdiv_space() + + Vu, continuity = self.build_hdiv_space() + self.continuity["HDiv"] = continuity setattr(self, "HDiv", Vu) - Vdg = self.build_l2_space() + + Vdg, continuity = self.build_l2_space() + self.continuity["L2"] = continuity setattr(self, "L2", Vdg) return Vcg, Vcurl, Vu, Vdg @@ -334,11 +384,18 @@ def build_compatible_spaces(self): else: # 1D domain, de Rham complex has 2 spaces # CG, hdiv and hcurl spaces should be the same - Vcg = self.build_h1_space() + Vcg, continuity = self.build_h1_space() + + self.continuity["H1"] = continuity setattr(self, "H1", Vcg) + setattr(self, "HCurl", None) + + self.continuity["HDiv"] = continuity setattr(self, "HDiv", Vcg) - Vdg = self.build_l2_space() + + Vdg, continuity = self.build_l2_space() + self.continuity["L2"] = continuity setattr(self, "L2", Vdg) return Vcg, Vdg @@ -360,7 +417,7 @@ def build_hcurl_space(self): """ if hdiv_hcurl_dict[self.family] is None: logger.warning('There is no HCurl space for this family. Not creating one') - return None + return None, None if self.extruded_mesh: Vh_elt = HCurl(TensorProductElement(self.base_elt_hori_hcurl, @@ -368,10 +425,13 @@ def build_hcurl_space(self): Vv_elt = HCurl(TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_dg)) V_elt = Vh_elt + Vv_elt + continuity = {'horizontal': True, 'vertical': True} else: V_elt = self.base_elt_hori_hcurl + continuity = True - return FunctionSpace(self.mesh, V_elt, name='HCurl'+self.complex_name) + space_name = 'HCurl'+self.complex_name + return FunctionSpace(self.mesh, V_elt, name=space_name), continuity def build_hdiv_space(self): """ @@ -387,9 +447,13 @@ def build_hdiv_space(self): self.base_elt_vert_cg) Vv_elt = HDiv(Vt_elt) V_elt = Vh_elt + Vv_elt + continuity = {'horizontal': True, 'vertical': True} else: V_elt = self.base_elt_hori_hdiv - return FunctionSpace(self.mesh, V_elt, name='HDiv'+self.complex_name) + continuity = True + + space_name = 'HDiv'+self.complex_name + return FunctionSpace(self.mesh, V_elt, name=space_name), continuity def build_l2_space(self): """ @@ -401,10 +465,13 @@ def build_l2_space(self): if self.extruded_mesh: V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_dg) + continuity = {'horizontal': False, 'vertical': False} else: V_elt = self.base_elt_hori_dg + continuity = False - return FunctionSpace(self.mesh, V_elt, name='L2'+self.complex_name) + space_name = 'L2'+self.complex_name + return FunctionSpace(self.mesh, V_elt, name=space_name), continuity def build_theta_space(self): """ @@ -423,8 +490,10 @@ def build_theta_space(self): assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded' V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg) + continuity = {'horizontal': False, 'vertical': True} - return FunctionSpace(self.mesh, V_elt, name='theta'+self.complex_name) + space_name = 'theta'+self.complex_name + return FunctionSpace(self.mesh, V_elt, name=space_name), continuity def build_h1_space(self): """ @@ -440,11 +509,13 @@ def build_h1_space(self): if self.extruded_mesh: V_elt = TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_cg) - + continuity = {'horizontal': True, 'vertical': True} else: V_elt = self.base_elt_hori_cg + continuity = True - return FunctionSpace(self.mesh, V_elt, name='H1'+self.complex_name) + space_name = 'H1'+self.complex_name + return FunctionSpace(self.mesh, V_elt, name=space_name), continuity def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): @@ -476,3 +547,24 @@ def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): raise ValueError(f'Cannot pass both "degree" and "vertical_degree" to {name}') if not extruded_mesh and vertical_degree is not None: raise ValueError(f'Cannot pass "vertical_degree" to {name} if mesh is not extruded') + + +def is_cg(V): + """ + Checks if a :class:`FunctionSpace` is continuous. + + Function to check if a given space, V, is CG. Broken elements are always + discontinuous; for vector elements we check the names of the Sobolev spaces + of the subelements and for all other elements we just check the Sobolev + space name. + + Args: + V (:class:`FunctionSpace`): the space to check. + """ + ele = V.ufl_element() + if isinstance(ele, BrokenElement): + return False + elif type(ele) == VectorElement: + return all([e.sobolev_space.name == "H1" for e in ele._sub_elements]) + else: + return V.ufl_element().sobolev_space.name == "H1" diff --git a/gusto/core/io.py b/gusto/core/io.py index 710495828..bd80f86d5 100644 --- a/gusto/core/io.py +++ b/gusto/core/io.py @@ -2,7 +2,7 @@ from os import path, makedirs import itertools -from netCDF4 import Dataset +from netCDF4 import Dataset, stringtochar import sys import time from gusto.diagnostics import Diagnostics, CourantNumber @@ -395,7 +395,7 @@ def setup_dump(self, state_fields, t, pick_up=False): running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname # Raising exceptions needs to be done in parallel - if self.mesh.comm.Get_rank() == 0: + if self.mesh.comm.rank == 0: # Create results directory if it doesn't already exist if not path.exists(self.dumpdir): try: @@ -446,7 +446,7 @@ def setup_dump(self, state_fields, t, pick_up=False): if pick_up: # Pick up t idx - if self.mesh.comm.Get_rank() == 0: + if self.mesh.comm.rank == 0: nc_field_file = Dataset(self.nc_filename, 'r') self.field_t_idx = len(nc_field_file['time'][:]) nc_field_file.close() @@ -763,75 +763,93 @@ def dump(self, state_fields, time_data): self.dumpfile_ll.write(*self.to_dump_latlon) def create_nc_dump(self, filename, space_names): - my_rank = self.mesh.comm.Get_rank() self.field_t_idx = 0 - if my_rank == 0: - nc_field_file = Dataset(filename, 'w') + comm = self.mesh.comm + nc_field_file, nc_supports_parallel = make_nc_dataset(filename, 'w', comm) + + if nc_field_file: nc_field_file.createDimension('time', None) nc_field_file.createVariable('time', float, ('time',)) # Add mesh metadata + nc_field_file.createDimension("dim_one", 1) + nc_field_file.createDimension("dim_string", 256) for metadata_key, metadata_value in self.domain.metadata.items(): # If the metadata is None or a Boolean, try converting to string # This is because netCDF can't take these types as options - if type(metadata_value) in [type(None), type(True)]: + if metadata_value is None or isinstance(metadata_value, bool): output_value = str(metadata_value) else: output_value = metadata_value - # Get the type from the metadata itself - nc_field_file.createVariable(metadata_key, type(output_value), []) - nc_field_file.variables[metadata_key][0] = output_value - - # Add coordinates if they are not already in the file - for space_name in space_names: - if space_name not in self.domain.coords.chi_coords.keys(): - # Space not registered - # TODO: we should fail here, but currently there are some spaces - # that we can't output for so instead just skip outputting - pass + # At present parallel netCDF crashes when saving strings. To get around this + # we instead save string metadata as char arrays of fixed length. + if isinstance(output_value, str): + nc_field_file.createVariable(metadata_key, 'S1', ('dim_one', 'dim_string')) + output_char_array = np.array([output_value], dtype='S256') + nc_field_file[metadata_key][:] = stringtochar(output_char_array) else: - coord_fields = self.domain.coords.global_chi_coords[space_name] - num_points = len(self.domain.coords.global_chi_coords[space_name][0]) - - nc_field_file.createDimension('coords_'+space_name, num_points) - - for (coord_name, coord_field) in zip(self.domain.coords.coords_name, coord_fields): - nc_field_file.createVariable(coord_name+'_'+space_name, float, 'coords_'+space_name) - nc_field_file.variables[coord_name+'_'+space_name][:] = coord_field[:] - - # Create variable for storing the field values - for field in self.to_dump: - field_name = field.name() - space_name = field.function_space().name - if space_name not in self.domain.coords.chi_coords.keys(): - # Space not registered - # TODO: we should fail here, but currently there are some spaces - # that we can't output for so instead just skip outputting - logger.warning(f'netCDF outputting for space {space_name} ' - + 'not yet implemented, so unable to output ' - + f'{field_name} field') - else: - nc_field_file.createGroup(field_name) - nc_field_file[field_name].createVariable('field_values', float, ('coords_'+space_name, 'time')) + nc_field_file.createVariable(metadata_key, type(output_value), ('dim_one',)) + nc_field_file[metadata_key][0] = output_value + + # Add coordinates if they are not already in the file + for space_name in space_names: + if space_name not in self.domain.coords.chi_coords.keys(): + # Space not registered + # TODO: we should fail here, but currently there are some spaces + # that we can't output for so instead just skip outputting + pass + else: + coord_fields = self.domain.coords.chi_coords[space_name] + ndofs = coord_fields[0].function_space().dim() + if nc_field_file: + nc_field_file.createDimension(f'coords_{space_name}', ndofs) + + for coord_name, coord_field in zip(self.domain.coords.coords_name, coord_fields): + if nc_field_file: + nc_field_file.createVariable(f'{coord_name}_{space_name}', float, f'coords_{space_name}') + + if nc_supports_parallel: + start, stop = self.domain.coords.parallel_array_lims[space_name] + nc_field_file.variables[f'{coord_name}_{space_name}'][start:stop] = coord_field.dat.data_ro + else: + global_coord_field = gather_field_data(coord_field, self.domain) + if comm.rank == 0: + nc_field_file.variables[f'{coord_name}_{space_name}'][...] = global_coord_field + + # Create variable for storing the field values + for field in self.to_dump: + field_name = field.name() + space_name = field.function_space().name + if space_name not in self.domain.coords.chi_coords.keys(): + # Space not registered + # TODO: we should fail here, but currently there are some spaces + # that we can't output for so instead just skip outputting + logger.warning(f'netCDF outputting for space {space_name} ' + + 'not yet implemented, so unable to output ' + + f'{field_name} field') + else: + if nc_field_file: + nc_field_file.createGroup(field_name) + nc_field_file[field_name].createVariable('field_values', float, (f'coords_{space_name}', 'time')) + if nc_field_file: nc_field_file.close() def write_nc_dump(self, t): - comm = self.mesh.comm - my_rank = comm.Get_rank() - comm_size = comm.Get_size() + nc_field_file, nc_supports_parallel = make_nc_dataset(self.nc_filename, 'a', comm) - # Open file to add time - if my_rank == 0: - nc_field_file = Dataset(self.nc_filename, 'a') + if nc_field_file and 'time' in nc_field_file.variables.keys(): + # Unbounded variables need to be accessed collectively + # (https://unidata.github.io/netcdf4-python/#parallel-io) + if nc_supports_parallel: + nc_field_file['time'].set_collective(True) nc_field_file['time'][self.field_t_idx] = t # Loop through output field data here - num_fields = len(self.to_dump) - for i, field in enumerate(self.to_dump): + for field in self.to_dump: field_name = field.name() space_name = field.function_space().name @@ -845,29 +863,16 @@ def write_nc_dump(self, t): # Scalar elements # -------------------------------------------------------- # else: - j = 0 - # For most processors send data to first processor - if my_rank != 0: - # Make a tag to uniquely identify this call - my_tag = comm_size*(num_fields*j + i) + my_rank - comm.send(field.dat.data_ro[:], dest=0, tag=my_tag) + if nc_supports_parallel: + nc_field_file[field_name]['field_values'].set_collective(True) + start, stop = self.domain.coords.parallel_array_lims[space_name] + nc_field_file[field_name]['field_values'][start:stop, self.field_t_idx] = field.dat.data_ro else: - # Set up array to store full data in - total_data_size = self.domain.coords.parallel_array_lims[space_name][comm_size-1][1]+1 - single_proc_data = np.zeros(total_data_size) - # Get data for this processor first - (low_lim, up_lim) = self.domain.coords.parallel_array_lims[space_name][my_rank][:] - single_proc_data[low_lim:up_lim+1] = field.dat.data_ro[:] - # Receive data from other processors - for procid in range(1, comm_size): - my_tag = comm_size*(num_fields*j + i) + procid - incoming_data = comm.recv(source=procid, tag=my_tag) - (low_lim, up_lim) = self.domain.coords.parallel_array_lims[space_name][procid][:] - single_proc_data[low_lim:up_lim+1] = incoming_data[:] - # Store whole field data - nc_field_file[field_name].variables['field_values'][:, self.field_t_idx] = single_proc_data[:] - - if my_rank == 0: + global_field_data = gather_field_data(field, self.domain) + if comm.rank == 0: + nc_field_file[field_name]['field_values'][:, self.field_t_idx] = global_field_data + + if nc_field_file: nc_field_file.close() self.field_t_idx += 1 @@ -917,3 +922,71 @@ def topo_sort(field_deps): if f not in schedule) raise RuntimeError("Field dependencies have a cycle:\n\n%s" % cycle) return list(map(name2field.__getitem__, schedule)) + + +def make_nc_dataset(filename, access, comm): + """Create a netCDF data set, possibly in parallel. + + Args: + filename (str): The filename. + access (str): The access descriptor - ``r``, ``w`` or ``a``. + comm: The communicator. + + Returns: + tuple: 2-tuple of :class:`netCDF4_netCDF4.Dataset` (or `None`) and `bool` + indicating whether netCDF supports parallel usage. If parallel is not + supported then the dataset will be `None` on all but rank 0. + + A warning will be thrown if this function is called in parallel and a + non-parallel netCDF4 is installed. + + """ + try: + nc_field_file = Dataset(filename, access, parallel=True) + nc_supports_parallel = True + except ValueError: + # parallel netCDF not available, use the serial version instead + if comm.size > 1: + import warnings + warnings.warn( + "Serial netCDF in use even though you are running in parallel. This " + "is especially inefficient at high core counts. Please refer to the " + "documentation for information about installing a parallel version " + "of netCDF.", + UserWarning, + ) + + if comm.rank == 0: + nc_field_file = Dataset(filename, access, parallel=False) + else: + nc_field_file = None + nc_supports_parallel = False + return nc_field_file, nc_supports_parallel + + +def gather_field_data(field, domain): + """Gather global field data into a single array on rank 0. + + Args: + field (:class:`firedrake.Function`): The field to gather. + domain (:class:`Domain`): The domain. + + Returns: + :class:`numpy.ndarray` that is the concatenation of all DoFs on + all ranks. + + Note that this operation is *extremely inefficient* when run with large + amounts of parallelism. Avoid calling if at all possible. + + """ + comm = domain.mesh.comm + + if comm.size == 1: + return field.dat.data_ro.copy() + + gathered_data = comm.gather(field.dat.data_ro) + if comm.rank == 0: + return np.concatenate(gathered_data) + else: + assert gathered_data is None + return None diff --git a/gusto/core/labels.py b/gusto/core/labels.py index 962b60c50..c325e0aa0 100644 --- a/gusto/core/labels.py +++ b/gusto/core/labels.py @@ -94,6 +94,7 @@ def __call__(self, target, value=None): explicit = Label("explicit") horizontal = Label("horizontal") vertical = Label("vertical") +source_label = Label("source_label") transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor, ufl.indexed.Indexed]) prognostic = Label("prognostic", validator=lambda value: type(value) == str) linearisation = Label("linearisation", validator=lambda value: type(value) in [LabelledForm, Term]) diff --git a/gusto/diagnostics/compressible_euler_diagnostics.py b/gusto/diagnostics/compressible_euler_diagnostics.py index b625e6846..bbacc549e 100644 --- a/gusto/diagnostics/compressible_euler_diagnostics.py +++ b/gusto/diagnostics/compressible_euler_diagnostics.py @@ -718,20 +718,24 @@ def setup(self, domain, state_fields): cp = Constant(self.parameters.cp) n = FacetNormal(domain.mesh) + dx_qp = dx(degree=domain.max_quad_degree) + dS_v_qp = dS_v(degree=domain.max_quad_degree) + # TODO: not sure about this expression! # Gravity does not appear, and why are there reference profiles? F = TrialFunction(Vu) w = TestFunction(Vu) imbalance = Function(Vu) a = inner(w, F)*dx - L = (- cp*div((theta-thetabar)*w)*exnerbar*dx - + cp*jump((theta-thetabar)*w, n)*avg(exnerbar)*dS_v - - cp*div(thetabar*w)*(exner-exnerbar)*dx - + cp*jump(thetabar*w, n)*avg(exner-exnerbar)*dS_v) + L = (- cp*div((theta-thetabar)*w)*exnerbar*dx_qp + + cp*jump((theta-thetabar)*w, n)*avg(exnerbar)*dS_v_qp + - cp*div(thetabar*w)*(exner-exnerbar)*dx_qp + + cp*jump(thetabar*w, n)*avg(exner-exnerbar)*dS_v_qp) bcs = self.equations.bcs['u'] - imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs) + imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs, + constant_jacobian=True) self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) self.expr = dot(imbalance, domain.k) super().setup(domain, state_fields) @@ -786,12 +790,14 @@ def setup(self, domain, state_fields): eqn_rhs = domain.dt * self.phi * (rain * dot(- v, domain.k) * rho / area) * ds_b # Compute area normalisation - area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area) + area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area, + constant_jacobian=True) area_solver = LinearVariationalSolver(area_prob) area_solver.solve() # setup solver - rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux) + rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux, + constant_jacobian=True) self.solver = LinearVariationalSolver(rain_prob) self.field = state_fields(self.name, space=DG0, dump=True, pick_up=True) # Initialise field to zero, if picking up this will be overridden diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py index b287d6d36..982c2a5e9 100644 --- a/gusto/diagnostics/diagnostics.py +++ b/gusto/diagnostics/diagnostics.py @@ -1,14 +1,15 @@ """Common diagnostic fields.""" -from firedrake import (assemble, dot, dx, Function, sqrt, TestFunction, +from firedrake import (dot, dx, Function, sqrt, TestFunction, TrialFunction, Constant, grad, inner, FacetNormal, LinearVariationalProblem, LinearVariationalSolver, ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, pi, TensorFunctionSpace, SpatialCoordinate, as_vector, - Projector, Interpolator, FunctionSpace, FiniteElement, - TensorProductElement) + Projector, assemble, FunctionSpace, FiniteElement, + TensorProductElement, CellVolume, Cofunction) from firedrake.assign import Assigner +from firedrake.__future__ import interpolate from ufl.domain import extract_unique_domain from abc import ABCMeta, abstractmethod, abstractproperty @@ -21,7 +22,8 @@ "XComponent", "YComponent", "ZComponent", "MeridionalComponent", "ZonalComponent", "RadialComponent", "Energy", "KineticEnergy", "Sum", "Difference", "SteadyStateError", "Perturbation", - "Divergence", "TracerDensity", "IterativeDiagnosticField"] + "Divergence", "TracerDensity", "IterativeDiagnosticField", + "CumulativeSum"] class Diagnostics(object): @@ -192,7 +194,7 @@ def setup(self, domain, state_fields, space=None): # Solve method must be declared in diagnostic's own setup routine if self.method == 'interpolate': - self.evaluator = Interpolator(self.expr, self.field) + self.evaluator = interpolate(self.expr, self.space) elif self.method == 'project': self.evaluator = Projector(self.expr, self.field) elif self.method == 'assign': @@ -206,7 +208,7 @@ def compute(self): logger.debug(f'Computing diagnostic {self.name} with {self.method} method') if self.method == 'interpolate': - self.evaluator.interpolate() + self.field.assign(assemble(self.evaluator)) elif self.method == 'assign': self.evaluator.assign() elif self.method == 'project': @@ -293,7 +295,7 @@ def setup(self, domain, state_fields, space=None): # Solve method must be declared in diagnostic's own setup routine if self.method == 'interpolate': - self.evaluator = Interpolator(self.expr, self.field) + self.evaluator = interpolate(self.expr, self.space) elif self.method == 'project': self.evaluator = Projector(self.expr, self.field) elif self.method == 'assign': @@ -408,12 +410,8 @@ def setup(self, domain, state_fields): V = FunctionSpace(domain.mesh, "DG", 0) test = TestFunction(V) - cell_volume = Function(V) - self.cell_flux = Function(V) - - # Calculate cell volumes - One = Function(V).assign(1) - assemble(One*test*dx, tensor=cell_volume) + cell_volume = Function(V).interpolate(CellVolume(domain.mesh)) + self.cell_flux = Cofunction(V.dual()) # Get the velocity that is being used if type(self.velocity) is str: @@ -448,7 +446,10 @@ def setup(self, domain, state_fields): self.cell_flux_form = 2*avg(un*test)*dS_calc + un*test*ds_calc # Final Courant number expression - self.expr = self.cell_flux * domain.dt / cell_volume + cell_flux = self.cell_flux.riesz_representation( + 'l2', solver_options={'function_space': V} + ) + self.expr = cell_flux * domain.dt / cell_volume super().setup(domain, state_fields) @@ -513,7 +514,8 @@ def setup(self, domain, state_fields): L = -inner(div(test), f)*dx if space.extruded: L += dot(dot(test, n), f)*(ds_t + ds_b) - prob = LinearVariationalProblem(a, L, self.field) + prob = LinearVariationalProblem(a, L, self.field, + constant_jacobian=True) self.evaluator = LinearVariationalSolver(prob) @@ -1040,3 +1042,43 @@ def setup(self, domain, state_fields): else: super().setup(domain, state_fields) + + +class CumulativeSum(DiagnosticField): + """Base diagnostic for cumulatively summing a field in time.""" + def __init__(self, name): + """ + Args: + name (str): name of the field to take the cumulative sum of. + """ + self.field_name = name + self.integral_name = name+"_cumulative" + super().__init__(self, method='assign', required_fields=(self.field_name,)) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + + # Gather the field to be summed + self.integrand = state_fields(self.field_name) + self.space = self.integrand.function_space() + + # Create a new field to store the cumulative sum + self.field = state_fields(self.integral_name, space=self.space, dump=True, pick_up=True) + # Initialise the new field to zero, if picking up from a checkpoint + # file the original cumulative field will load and not be overwritten. + self.field.assign(0.0) + + def compute(self): + """Increment the cumulative sum diagnostic.""" + self.field.assign(self.field + self.integrand) + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.integral_name diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index dce6fac69..024e58ecf 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -1,13 +1,17 @@ """Common diagnostic fields for the Shallow Water equations.""" -from firedrake import (dx, TestFunction, TrialFunction, grad, inner, curl, - LinearVariationalProblem, LinearVariationalSolver) +from firedrake import ( + dx, TestFunction, TrialFunction, grad, inner, curl, Function, assemble, + LinearVariationalProblem, LinearVariationalSolver, conditional +) +from firedrake.__future__ import interpolate from gusto.diagnostics.diagnostics import DiagnosticField, Energy __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", "ShallowWaterPotentialEnstrophy", "PotentialVorticity", - "RelativeVorticity", "AbsoluteVorticity"] + "RelativeVorticity", "AbsoluteVorticity", "PartitionedVapour", + "PartitionedCloud"] class ShallowWaterKineticEnergy(Energy): @@ -179,15 +183,18 @@ def setup(self, domain, state_fields, vorticity_type=None): if vorticity_type == "potential": a = q*gamma*D*dx + constant_jacobian = False else: a = q*gamma*dx + constant_jacobian = True L = (- inner(domain.perp(grad(gamma)), u))*dx if vorticity_type != "relative": f = state_fields("coriolis") L += gamma*f*dx - problem = LinearVariationalProblem(a, L, self.field) + problem = LinearVariationalProblem(a, L, self.field, + constant_jacobian=constant_jacobian) self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) @@ -274,3 +281,104 @@ def setup(self, domain, state_fields): state_fields (:class:`StateFields`): the model's field container. """ super().setup(domain, state_fields, vorticity_type="relative") + + +class PartitionedVapour(DiagnosticField): + """ + Diagnostic for computing the vapour in the equivalent buoyancy formulation + of the moist thermal shallow water equations. + """ + name = "PartitionedVapour" + + def __init__(self, equation, name='q_t', space=None, + method='interpolate'): + """ + Args: + equation (:class:`PrognosticEquation`): the model's equation. + name (str, optional): name of the total moisture field to use to + compute the vapour from. Defaults to total moisture, q_t. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.fname = name + self.equation = equation + super().__init__(space=space, method=method, required_fields=(self.fname,)) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + q_t = state_fields(self.fname) + space = domain.spaces("DG") + self.qsat_func = Function(space) + + qsat_expr = self.equation.compute_saturation(state_fields.X( + self.equation.field_name)) + self.qsat_interpolate = interpolate(qsat_expr, space) + self.expr = conditional(q_t < self.qsat_func, q_t, self.qsat_func) + + super().setup(domain, state_fields, space=space) + + def compute(self): + """Performs the computation of the diagnostic field.""" + self.qsat_func.assign(assemble(self.qsat_interpolate)) + super().compute() + + +class PartitionedCloud(DiagnosticField): + """ + Diagnostic for computing the cloud in the equivalent buoyancy formulation + of the moist thermal shallow water equations. + """ + name = "PartitionedCloud" + + def __init__(self, equation, name='q_t', space=None, + method='interpolate'): + """ + Args: + equation (:class:`PrognosticEquation`): the model's equation. + name (str, optional): name of the total moisture field to use to + compute the vapour from. Defaults to total moisture, q_t. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.fname = name + self.equation = equation + super().__init__(space=space, method=method, required_fields=(self.fname,)) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + q_t = state_fields(self.fname) + space = domain.spaces("DG") + self.qsat_func = Function(space) + + qsat_expr = self.equation.compute_saturation(state_fields.X( + self.equation.field_name)) + self.qsat_interpolate = interpolate(qsat_expr, space) + vapour = conditional(q_t < self.qsat_func, q_t, self.qsat_func) + self.expr = q_t - vapour + + super().setup(domain, state_fields, space=space) + + def compute(self): + """Performs the computation of the diagnostic field.""" + self.qsat_func.assign(assemble(self.qsat_interpolate)) + super().compute() diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index 84b187e47..73bc76d45 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -112,6 +112,10 @@ def __init__(self, domain, parameters, sponge_options=None, exner = exner_pressure(parameters, rho, theta) n = FacetNormal(domain.mesh) + # Specify quadrature degree to use for pressure gradient term + dx_qp = dx(degree=(domain.max_quad_degree)) + dS_v_qp = dS_v(degree=(domain.max_quad_degree)) + # -------------------------------------------------------------------- # # Time Derivative Terms # -------------------------------------------------------------------- # @@ -165,8 +169,8 @@ def __init__(self, domain, parameters, sponge_options=None, theta_v = theta / (Constant(1.0) + tracer_mr_total) pressure_gradient_form = pressure_gradient(subject(prognostic( - cp*(-div(theta_v*w)*exner*dx - + jump(theta_v*w, n)*avg(exner)*dS_v), 'u'), self.X)) + cp*(-div(theta_v*w)*exner*dx_qp + + jump(theta_v*w, n)*avg(exner)*dS_v_qp), 'u'), self.X)) # -------------------------------------------------------------------- # # Gravitational Term @@ -209,7 +213,7 @@ def __init__(self, domain, parameters, sponge_options=None, residual += subject(prognostic( gamma * theta * div(u) - * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml))*dx, 'theta'), self.X) + * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml))*dx_qp, 'theta'), self.X) # -------------------------------------------------------------------- # # Extra Terms (Coriolis, Sponge, Diffusion and others) @@ -236,7 +240,7 @@ def __init__(self, domain, parameters, sponge_options=None, self.mu = self.prescribed_fields("sponge", W_DG).interpolate(muexpr) residual += sponge(subject(prognostic( - self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, 'u'), self.X)) + self.mu*inner(w, domain.k)*inner(u, domain.k)*dx_qp, 'u'), self.X)) if diffusion_options is not None: for field, diffusion in diffusion_options: diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index b2df68e2c..a562795d6 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -94,7 +94,7 @@ def __init__(self, field_names, domain, space_names, None. active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included - in the equations.. Defaults to None. + in the equations. Defaults to None. """ self.field_names = field_names diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 854266952..fd61cedf7 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -1,19 +1,22 @@ """Classes for defining variants of the shallow-water equations.""" from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, - dS, split) -from firedrake.fml import subject + dS, split, conditional, exp) +from firedrake.fml import subject, drop from gusto.core.labels import (time_derivative, transport, prognostic, linearisation, pressure_gradient, coriolis) from gusto.equations.common_forms import ( advection_form, advection_form_1d, continuity_form, continuity_form_1d, vector_invariant_form, kinetic_energy_form, advection_equation_circulation_form, diffusion_form_1d, - linear_continuity_form + linear_continuity_form, linear_advection_form ) from gusto.equations.prognostic_equations import PrognosticEquationSet + __all__ = ["ShallowWaterEquations", "LinearShallowWaterEquations", + "ThermalShallowWaterEquations", + "LinearThermalShallowWaterEquations", "ShallowWaterEquations_1d", "LinearShallowWaterEquations_1d"] @@ -21,16 +24,15 @@ class ShallowWaterEquations(PrognosticEquationSet): u""" Class for the (rotating) shallow-water equations, which evolve the velocity 'u' and the depth field 'D', via some variant of: \n - ∂u/∂t + (u.∇)u + f×u + g*∇(D+b) = 0, \n + ∂u/∂t + (u.∇)u + f×u + g*∇(D+B) = 0, \n ∂D/∂t + ∇.(D*u) = 0, \n - for Coriolis parameter 'f' and bottom surface 'b'. + for Coriolis parameter 'f' and bottom surface 'B'. """ - def __init__(self, domain, parameters, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, topog_expr=None, space_names=None, linearisation_map='default', u_transport_option='vector_invariant_form', - no_normal_flow_bc_ids=None, active_tracers=None, - thermal=False): + no_normal_flow_bc_ids=None, active_tracers=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -39,8 +41,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, the model's physical parameters. fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis parameter. Defaults to None. - bexpr (:class:`ufl.Expr`, optional): an expression for the bottom - surface of the fluid. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. space_names (dict, optional): a dictionary of strings for names of the function spaces to use for the spatial discretisation. The keys are the names of the prognostic variables. Defaults to None @@ -62,47 +64,66 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations. Defaults to None. - thermal (flag, optional): specifies whether the equations have a - thermal or buoyancy variable that feeds back on the momentum. - Defaults to False. - - Raises: - NotImplementedError: active tracers are not yet implemented. """ - self.thermal = thermal - field_names = ['u', 'D'] - - if space_names is None: - space_names = {'u': 'HDiv', 'D': 'L2'} - if active_tracers is None: active_tracers = [] - if self.thermal: - field_names.append('b') - if 'b' not in space_names.keys(): - space_names['b'] = 'L2' - if linearisation_map == 'default': # Default linearisation is time derivatives, pressure gradient and # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ - t.get(prognostic) in ['u', 'D', 'b'] \ - and (any(t.has_label(time_derivative, pressure_gradient)) - or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) + t.get(prognostic) in ['u', 'D'] \ + and (any(t.has_label(time_derivative, coriolis, pressure_gradient)) + or (t.get(prognostic) in ['D'] and t.has_label(transport))) + + field_names = ['u', 'D'] + space_names = {'u': 'HDiv', 'D': 'L2'} + super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) self.parameters = parameters - g = parameters.g - H = parameters.H + self.domain = domain + self.active_tracers = active_tracers + + self._setup_residual(fexpr, topog_expr, u_transport_option) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms( + self.residual, self.linearisation_map) + + def _setup_residual(self, fexpr, topog_expr, u_transport_option): + """ + Sets up the residual for the shallow water equations. This + is separate from the __init__ method because the thermal + shallow water equation class expands on this equation set by + adding additional fields that depend on the formulation. This + increases the size of the mixed function space and the + residual must be setup after this has happened. + + Args: + fexpr (:class:`ufl.Expr`): an expression for the Coroilis + parameter. + topog_expr (:class:`ufl.Expr`): an expression for the + bottom surface of the fluid. + u_transport_option (str): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + """ + + g = self.parameters.g w, phi = self.tests[0:2] u, D = split(self.X)[0:2] u_trial = split(self.trials)[0] + Dbar = split(self.X_ref)[1] # -------------------------------------------------------------------- # # Time Derivative Terms @@ -114,116 +135,84 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(domain, w, u, u), 'u') + u_adv = prognostic(vector_invariant_form(self.domain, w, u, u), 'u') elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), 'u') elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') - u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form + u_adv = prognostic(advection_equation_circulation_form(self.domain, w, u, u), 'u') + ke_form else: - raise ValueError("Invalid u_transport_option: %s" % u_transport_option) + raise ValueError("Invalid u_transport_option: %s" % self.u_transport_option) # Depth transport term D_adv = prognostic(continuity_form(phi, D, u), 'D') # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): - linear_D_adv = linear_continuity_form(phi, H, u_trial) + linear_D_adv = linear_continuity_form(phi, Dbar, u_trial) # Add linearisation to D_adv D_adv = linearisation(D_adv, linear_D_adv) adv_form = subject(u_adv + D_adv, self.X) # Add transport of tracers - if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(active_tracers) - # Add transport of buoyancy, if thermal shallow water equations - if self.thermal: - gamma = self.tests[2] - b = split(self.X)[2] - b_adv = prognostic(advection_form(gamma, b, u), 'b') - adv_form += subject(b_adv, self.X) + if len(self.active_tracers) > 0: + adv_form += self.generate_tracer_transport_terms( + self.active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term # -------------------------------------------------------------------- # - # Add pressure gradient only if not doing thermal - if self.thermal: - residual = (mass_form + adv_form) - else: - pressure_gradient_form = pressure_gradient( - subject(prognostic(-g*div(w)*D*dx, 'u'), self.X)) + pressure_gradient_form = pressure_gradient( + subject(prognostic(-g*div(w)*D*dx, 'u'), self.X)) - residual = (mass_form + adv_form + pressure_gradient_form) + residual = (mass_form + adv_form + pressure_gradient_form) # -------------------------------------------------------------------- # - # Extra Terms (Coriolis, Topography and Thermal) + # Extra Terms (Coriolis, Topography) # -------------------------------------------------------------------- # # TODO: Is there a better way to store the Coriolis / topography fields? # The current approach is that these are prescribed fields, stored in # the equation, and initialised when the equation is if fexpr is not None: - V = FunctionSpace(domain.mesh, 'CG', 1) + V = FunctionSpace(self.domain.mesh, 'CG', 1) f = self.prescribed_fields('coriolis', V).interpolate(fexpr) coriolis_form = coriolis(subject( - prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X)) + prognostic(f*inner(self.domain.perp(u), w)*dx, "u"), self.X)) # Add linearisation if self.linearisation_map(coriolis_form.terms[0]): linear_coriolis = coriolis( - subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, 'u'), self.X)) + subject(prognostic(f*inner(self.domain.perp(u_trial), w)*dx, 'u'), self.X)) coriolis_form = linearisation(coriolis_form, linear_coriolis) residual += coriolis_form - if bexpr is not None: - topography = self.prescribed_fields('topography', domain.spaces('DG')).interpolate(bexpr) - if self.thermal: - n = FacetNormal(domain.mesh) - topography_form = subject(prognostic - (-topography*div(b*w)*dx - + jump(b*w, n)*avg(topography)*dS, - 'u'), self.X) - else: - topography_form = subject(prognostic - (-g*div(w)*topography*dx, 'u'), - self.X) + if topog_expr is not None: + topography = self.prescribed_fields('topography', self.domain.spaces('DG')).interpolate(topog_expr) + topography_form = subject(prognostic + (-g*div(w)*topography*dx, 'u'), + self.X) residual += topography_form - # thermal source terms not involving topography. - # label these as the equivalent pressure gradient term - if self.thermal: - n = FacetNormal(domain.mesh) - source_form = pressure_gradient(subject(prognostic(-D*div(b*w)*dx - - 0.5*b*div(D*w)*dx - + jump(b*w, n)*avg(D)*dS - + 0.5*jump(D*w, n)*avg(b)*dS, - 'u'), self.X)) - residual += source_form - - # -------------------------------------------------------------------- # - # Linearise equations - # -------------------------------------------------------------------- # - # Add linearisations to equations - self.residual = self.generate_linear_terms(residual, self.linearisation_map) + self.residual = residual class LinearShallowWaterEquations(ShallowWaterEquations): u""" Class for the linear (rotating) shallow-water equations, which describe the velocity 'u' and the depth field 'D', solving some variant of: \n - ∂u/∂t + f×u + g*∇(D+b) = 0, \n + ∂u/∂t + f×u + g*∇(D+B) = 0, \n ∂D/∂t + H*∇.(u) = 0, \n - for mean depth 'H', Coriolis parameter 'f' and bottom surface 'b'. + for mean depth 'H', Coriolis parameter 'f' and bottom surface 'B'. - This is set up the from the underlying :class:`ShallowWaterEquations`, + This is set up from the underlying :class:`ShallowWaterEquations`, which is then linearised. """ - def __init__(self, domain, parameters, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, topog_expr=None, space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", - no_normal_flow_bc_ids=None, active_tracers=None, - thermal=False): + no_normal_flow_bc_ids=None, active_tracers=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -232,8 +221,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, the model's physical parameters. fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis parameter. Defaults to None. - bexpr (:class:`ufl.Expr`, optional): an expression for the bottom - surface of the fluid. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. space_names (dict, optional): a dictionary of strings for names of the function spaces to use for the spatial discretisation. The keys are the names of the prognostic variables. Defaults to None @@ -255,24 +244,345 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations. Defaults to None. - thermal (flag, optional): specifies whether the equations have a - thermal or buoyancy variable that feeds back on the momentum. - Defaults to False. """ + super().__init__(domain, parameters, + fexpr=fexpr, topog_expr=topog_expr, + space_names=space_names, + linearisation_map=linearisation_map, + u_transport_option=u_transport_option, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + # Use the underlying routine to do a first linearisation of the equations + self.linearise_equation_set() + + +class ThermalShallowWaterEquations(ShallowWaterEquations): + u""" + Class for the (rotating) shallow-water equations, which evolve the velocity + 'u' and the depth field 'D' via some variant of either: \n + ∂u/∂t + (u.∇)u + f×u + b*∇(D+B) + 0.5*D*∇b= 0, \n + ∂D/∂t + ∇.(D*u) = 0, \n + ∂b/∂t + u.∇(b) = 0, \n + for Coriolis parameter 'f', bottom surface 'B' and buoyancy field b, + or, if equivalent_buoyancy=True: + ∂u/∂t + (u.∇)u + f×u + b_e*∇(D+B) + 0.5*D*∇(b_e + beta_2 q_v)= 0, \n + ∂D/∂t + ∇.(D*u) = 0, \n + ∂b_e/∂t + u.∇(b_e) = 0, \n + ∂q_t/∂t + u.∇(q_t) = 0, \n + for Coriolis parameter 'f', bottom surface 'B', equivalent buoyancy field \n + `b_e`=b-beta_2 q_v, and total moisture `q_t`=q_v+q_c, i.e. the sum of \n + water vapour and cloud water. + """ + + def __init__(self, domain, parameters, equivalent_buoyancy=False, + fexpr=None, topog_expr=None, + space_names=None, linearisation_map='default', + u_transport_option='vector_invariant_form', + no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + equivalent_buoyancy (bool, optional): switch to specify formulation + (see comments above). Defaults to False to give standard + thermal shallow water. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. If None is specified + then no terms are linearised. Defaults to the string 'default', + in which case the linearisation includes both time derivatives, + the 'D' transport term and the pressure gradient term. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + """ + + self.equivalent_buoyancy = equivalent_buoyancy + field_names = ['u', 'D'] + space_names = {'u': 'HDiv', 'D': 'L2'} + self.b_name = 'b_e' if equivalent_buoyancy else 'b' + + if equivalent_buoyancy: + for new_field in [self.b_name, 'q_t']: + field_names.append(new_field) + space_names[new_field] = 'L2' + else: + field_names.append(self.b_name) + space_names[self.b_name] = 'L2' + + if active_tracers is None: + active_tracers = [] + if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient, - # Coriolis and transport term from depth equation + # Default linearisation is time derivatives, pressure + # gradient and transport terms from depth and buoyancy + # equations. Include q_t if equivalent buoyancy. Don't include + # active tracers. + linear_transported = ['D', self.b_name] + if equivalent_buoyancy: + linear_transported.append('q_t') linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis)) - or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) + t.get(prognostic) in field_names \ + and (any(t.has_label(time_derivative, pressure_gradient, + coriolis)) + or (t.get(prognostic) in linear_transported + and t.has_label(transport))) + + # Bypass ShallowWaterEquations.__init__ to avoid having to + # define the field_names separately + PrognosticEquationSet.__init__( + self, field_names, domain, space_names, + linearisation_map=linearisation_map, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + self.parameters = parameters + self.domain = domain + self.active_tracers = active_tracers + + self._setup_residual(fexpr, topog_expr, u_transport_option) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms( + self.residual, self.linearisation_map) + + def _setup_residual(self, fexpr, topog_expr, u_transport_option): + """ + Sets up the residual for the thermal shallow water + equations, first calling the shallow water equation + _setup_residual method to get the standard shallow water forms. + + Args: + fexpr (:class:`ufl.Expr`): an expression for the Coroilis + parameter. + topog_expr (:class:`ufl.Expr`): an expression for the + bottom surface of the fluid. + u_transport_option (str): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + + """ + # don't pass topography to super class as we deal with those + # terms here later + super()._setup_residual(fexpr=fexpr, topog_expr=None, + u_transport_option=u_transport_option) + + w = self.tests[0] + gamma = self.tests[2] + u, D, b = split(self.X)[0:3] + Dbar, bbar = split(self.X_ref)[1:3] + u_trial, D_trial, b_trial = split(self.trials)[0:3] + n = FacetNormal(self.domain.mesh) + topog = self.prescribed_fields('topography', self.domain.spaces('DG')).interpolate(topog_expr) if topog_expr else None + self.topog = topog + if self.equivalent_buoyancy: + gamma_qt = self.tests[3] + qt = split(self.X)[3] + qtbar = split(self.X_ref)[3] + qt_trial = split(self.trials)[3] + + # -------------------------------------------------------------------- # + # Add pressure gradient-like terms to residual + # -------------------------------------------------------------------- # + # drop usual pressure gradient term + residual = self.residual.label_map( + lambda t: t.has_label(pressure_gradient), + drop) + + # add (moist) thermal source terms not involving topography - + # label these as the equivalent pressure gradient term and + # provide linearisation + if self.equivalent_buoyancy: + beta2 = self.parameters.beta2 + qsat_expr = self.compute_saturation(self.X) + qv = conditional(qt < qsat_expr, qt, qsat_expr) + qvbar = conditional(qtbar < qsat_expr, qtbar, qsat_expr) + source_form = pressure_gradient(subject(prognostic( + -D * div(b*w) * dx - 0.5 * b * div(D*w) * dx + + jump(b*w, n) * avg(D) * dS + 0.5 * jump(D*w, n) * avg(b) * dS + - beta2 * D * div(qv*w)*dx - 0.5 * beta2 * qv * div(D*w) * dx + + beta2 * jump(qv*w, n) * avg(D) * dS + + 0.5 * beta2 * jump(D*w, n) * avg(qv) * dS, + 'u'), self.X)) + linear_source_form = pressure_gradient(subject(prognostic( + -D_trial * div(bbar*w) * dx + - 0.5 * b_trial * div(Dbar*w) * dx + + jump(bbar*w, n) * avg(D_trial) * dS + + 0.5 * jump(Dbar*w, n) * avg(b_trial) * dS + - beta2 * D_trial * div(qvbar*w)*dx + - 0.5 * beta2 * qvbar * div(Dbar*w) * dx + + beta2 * jump(qvbar*w, n) * avg(D_trial) * dS + + 0.5 * beta2 * jump(Dbar*w, n) * avg(qvbar) * dS + - 0.5 * bbar * div(Dbar*w) * dx + + 0.5 * jump(Dbar*w, n) * avg(bbar) * dS + - 0.5 * bbar * div(D_trial*w) * dx + + 0.5 * jump(D_trial*w, n) * avg(bbar) * dS + - beta2 * 0.5 * qvbar * div(D_trial*w) * dx + + beta2 * 0.5 * jump(D_trial*w, n) * avg(qvbar) * dS + - beta2 * 0.5 * qt_trial * div(Dbar*w) * dx + + beta2 * 0.5 * jump(Dbar*w, n) * avg(qt_trial) * dS, + 'u'), self.X)) + else: + source_form = pressure_gradient( + subject(prognostic(-D * div(b*w) * dx + + jump(b*w, n) * avg(D) * dS + - 0.5 * b * div(D*w) * dx + + 0.5 * jump(D*w, n) * avg(b) * dS, + 'u'), self.X)) + linear_source_form = pressure_gradient( + subject(prognostic(-D_trial * div(bbar*w) * dx + + jump(bbar*w, n) * avg(D_trial) * dS + - 0.5 * b_trial * div(Dbar*w) * dx + + 0.5 * jump(Dbar*w, n) * avg(b_trial) * dS + - 0.5 * bbar * div(Dbar*w) * dx + + 0.5 * jump(Dbar*w, n) * avg(bbar) * dS + - 0.5 * bbar * div(D_trial*w) * dx + + 0.5 * jump(D_trial*w, n) * avg(bbar) * dS, + 'u'), self.X)) + source_form = linearisation(source_form, linear_source_form) + residual += source_form + + # -------------------------------------------------------------------- # + # add transport terms and their linearisations: + # -------------------------------------------------------------------- # + b_adv = prognostic(advection_form(gamma, b, u), self.b_name) + if self.linearisation_map(b_adv.terms[0]): + linear_b_adv = linear_advection_form(gamma, bbar, u_trial) + b_adv = linearisation(b_adv, linear_b_adv) + residual += subject(b_adv, self.X) + + if self.equivalent_buoyancy: + qt_adv = prognostic(advection_form(gamma_qt, qt, u), "q_t") + if self.linearisation_map(qt_adv.terms[0]): + linear_qt_adv = linear_advection_form(gamma_qt, qtbar, u_trial) + qt_adv = linearisation(qt_adv, linear_qt_adv) + residual += subject(qt_adv, self.X) + + # -------------------------------------------------------------------- # + # add topography terms: + # -------------------------------------------------------------------- # + if topog_expr is not None: + if self.equivalent_buoyancy: + topography_form = subject(prognostic( + - topog * div(b*w) * dx + + jump(b*w, n) * avg(topog) * dS + - beta2 * topog * div(qv*w) * dx + + beta2 * jump(qv*w, n) * avg(topog) * dS, + 'u'), self.X) + else: + topography_form = subject(prognostic( + - topog * div(b*w) * dx + + jump(b*w, n) * avg(topog) * dS, + 'u'), self.X) + residual += topography_form + + self.residual = residual + + def compute_saturation(self, X): + # Returns the saturation expression as a function of the + # parameters specified in self.parameters and the input + # functions X. The latter are left as inputs to the + # function so that it can also be used for initialisation + q0 = self.parameters.q0 + nu = self.parameters.nu + g = self.parameters.g + H = self.parameters.H + D, b = split(X)[1:3] + topog = self.topog + if topog is None: + sat_expr = q0*H/(D) * exp(nu*(1-b/g)) + else: + sat_expr = q0*H/(D+topog) * exp(nu*(1-b/g)) + return sat_expr + + +class LinearThermalShallowWaterEquations(ThermalShallowWaterEquations): + u""" + Class for the linear (rotating) thermal shallow-water equations, which + describe the velocity 'u' and depth field 'D', solving some variant of: \n + ∂u/∂t + f×u + bbar*∇D + 0.5*H*∇b = 0, \n + ∂D/∂t + H*∇.(u) = 0, \n + ∂b/∂t + u.∇bbar = 0, \n + ∂q_t/∂t + u.∇(q_tbar) = 0, \n + for mean depth 'H', mean buoyancy `bbar`, Coriolis parameter 'f' + + This is set up from the underlying :class:`ThermalShallowWaterEquations`, + which is then linearised. + """ + + def __init__(self, domain, parameters, equivalent_buoyancy=False, + fexpr=None, topog_expr=None, + space_names=None, linearisation_map='default', + u_transport_option="vector_invariant_form", + no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + equivalent_buoyancy (bool, optional): switch to specify formulation + (see comments above). Defaults to False to give standard + thermal shallow water. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. If None is specified + then no terms are linearised. Defaults to the string 'default', + in which case the linearisation includes both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + """ super().__init__(domain, parameters, - fexpr=fexpr, bexpr=bexpr, space_names=space_names, + equivalent_buoyancy=equivalent_buoyancy, + fexpr=fexpr, topog_expr=topog_expr, + space_names=space_names, linearisation_map=linearisation_map, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers, thermal=thermal) + active_tracers=active_tracers) # Use the underlying routine to do a first linearisation of the equations self.linearise_equation_set() @@ -441,13 +751,6 @@ def __init__(self, domain, parameters, fexpr=None, in the equations. Defaults to None. """ - if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient, - # Coriolis and transport term from depth equation - linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis)) - or (t.get(prognostic) == 'D' and t.has_label(transport))) - super().__init__(domain, parameters, fexpr=fexpr, space_names=space_names, linearisation_map=linearisation_map, diff --git a/gusto/initialisation/hydrostatic_initialisation.py b/gusto/initialisation/hydrostatic_initialisation.py index 53bf537e8..4cddd00de 100644 --- a/gusto/initialisation/hydrostatic_initialisation.py +++ b/gusto/initialisation/hydrostatic_initialisation.py @@ -8,7 +8,6 @@ NonlinearVariationalProblem, NonlinearVariationalSolver, split, solve, \ FunctionSpace, errornorm, zero from gusto import thermodynamics -from gusto.core import logger from gusto.recovery import Recoverer, BoundaryMethod @@ -144,9 +143,17 @@ def compressible_hydrostatic_balance(equation, theta0, rho0, exner0=None, dv, dexner = TestFunctions(W) n = FacetNormal(equation.domain.mesh) - cp = parameters.cp + # measures + dx_qp = dx(degree=domain.max_quad_degree) + if top: + bmeasure = ds_t(degree=domain.max_quad_degree) + bstring = "bottom" + else: + bmeasure = ds_b(degree=domain.max_quad_degree) + bstring = "top" + # add effect of density of water upon theta theta = theta0 @@ -154,17 +161,10 @@ def compressible_hydrostatic_balance(equation, theta0, rho0, exner0=None, theta = theta0 / (1 + mr_t) alhs = ( - (cp*inner(v, dv) - cp*div(dv*theta)*exner)*dx - + dexner*div(theta*v)*dx + (cp*inner(v, dv) - cp*div(dv*theta)*exner)*dx_qp + + dexner*div(theta*v)*dx_qp ) - if top: - bmeasure = ds_t - bstring = "bottom" - else: - bmeasure = ds_b - bstring = "top" - arhs = -cp*inner(dv, n)*theta*exner_boundary*bmeasure # Possibly make g vary with spatial coordinates? @@ -205,8 +205,8 @@ def compressible_hydrostatic_balance(equation, theta0, rho0, exner0=None, dv, dexner = TestFunctions(W) exner = thermodynamics.exner_pressure(parameters, rho, theta0) F = ( - (cp*inner(v, dv) - cp*div(dv*theta)*exner)*dx - + dexner*div(theta0*v)*dx + (cp*inner(v, dv) - cp*div(dv*theta)*exner)*dx_qp + + dexner*div(theta0*v)*dx_qp + cp*inner(dv, n)*theta*exner_boundary*bmeasure ) F += g*inner(dv, equation.domain.k)*dx @@ -274,15 +274,10 @@ def saturated_hydrostatic_balance(equation, state_fields, theta_e, mr_t, mr_v0 = state_fields('water_vapour') # Calculate hydrostatic exner pressure - domain = equation.domain parameters = equation.parameters Vt = theta0.function_space() Vr = rho0.function_space() - VDG = domain.spaces("DG") - if any(deg > 2 for deg in VDG.ufl_element().degree()): - logger.warning("default quadrature degree most likely not sufficient for this degree element") - theta0.interpolate(theta_e) mr_v0.interpolate(mr_t) @@ -402,7 +397,6 @@ def unsaturated_hydrostatic_balance(equation, state_fields, theta_d, H, mr_v0 = state_fields('water_vapour') # Calculate hydrostatic exner pressure - domain = equation.domain parameters = equation.parameters Vt = theta0.function_space() Vr = rho0.function_space() @@ -410,10 +404,6 @@ def unsaturated_hydrostatic_balance(equation, state_fields, theta_d, H, R_v = parameters.R_v epsilon = R_d / R_v - VDG = domain.spaces("DG") - if any(deg > 2 for deg in VDG.ufl_element().degree()): - logger.warning("default quadrature degree most likely not sufficient for this degree element") - # apply first guesses theta0.assign(theta_d * 1.01) mr_v0.assign(0.01) diff --git a/gusto/physics/__init__.py b/gusto/physics/__init__.py index 91f09c2ec..b1407a299 100644 --- a/gusto/physics/__init__.py +++ b/gusto/physics/__init__.py @@ -2,4 +2,5 @@ from gusto.physics.chemistry import * # noqa from gusto.physics.boundary_and_turbulence import * # noqa from gusto.physics.microphysics import * # noqa -from gusto.physics.shallow_water_microphysics import * # noqa \ No newline at end of file +from gusto.physics.shallow_water_microphysics import * # noqa +from gusto.physics.held_suarez_forcing import * # noqa diff --git a/gusto/physics/boundary_and_turbulence.py b/gusto/physics/boundary_and_turbulence.py index 317cc4ac1..dc19c7d75 100644 --- a/gusto/physics/boundary_and_turbulence.py +++ b/gusto/physics/boundary_and_turbulence.py @@ -4,15 +4,14 @@ from firedrake import ( Interpolator, conditional, Function, dx, sqrt, dot, Constant, grad, - TestFunctions, split, inner, TestFunction, exp, avg, outer, FacetNormal, - SpatialCoordinate, dS_v, NonlinearVariationalProblem, - NonlinearVariationalSolver + TestFunctions, split, inner, Projector, exp, avg, outer, FacetNormal, + SpatialCoordinate, dS_v ) from firedrake.fml import subject from gusto.core.configuration import BoundaryLayerParameters from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations -from gusto.core.labels import prognostic +from gusto.core.labels import prognostic, source_label from gusto.core.logging import logger from gusto.equations import thermodynamics from gusto.physics.physics_parametrisation import PhysicsParametrisation @@ -70,15 +69,17 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) self.dt = Constant(0.0) + self.source = Function(equation.X.function_space()) # -------------------------------------------------------------------- # # Extract prognostic variables # -------------------------------------------------------------------- # u_idx = equation.field_names.index('u') T_idx = equation.field_names.index('theta') + self.T_idx = T_idx rho_idx = equation.field_names.index('rho') if vapour_name is not None: - m_v_idx = equation.field_names.index(vapour_name) + self.m_v_idx = equation.field_names.index(vapour_name) X = self.X tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests @@ -89,8 +90,8 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, test_theta = tests[T_idx] if vapour_name is not None: - m_v = split(X)[m_v_idx] - test_m_v = tests[m_v_idx] + m_v = split(X)[self.m_v_idx] + test_m_v = tests[self.m_v_idx] else: m_v = None @@ -127,22 +128,23 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, self.source_interpolators = [] # First specify T_np1 expression - Vtheta = equation.spaces[T_idx] T_np1_expr = ((T + C_H*u_hori_mag*T_surface_expr*self.dt/z_a) / (1 + C_H*u_hori_mag*self.dt/z_a)) # If moist formulation, determine next vapour value if vapour_name is not None: - source_mv = Function(Vtheta) + self.source_mv_int = self.source.subfunctions[self.m_v_idx] + self.source_mv = split(self.source)[self.m_v_idx] mv_sat = thermodynamics.r_sat(equation.parameters, T, p) mv_np1_expr = ((m_v + C_E*u_hori_mag*mv_sat*self.dt/z_a) / (1 + C_E*u_hori_mag*self.dt/z_a)) dmv_expr = surface_expr * (mv_np1_expr - m_v) / self.dt - source_mv_expr = test_m_v * source_mv * dx + source_mv_expr = test_m_v * self.source_mv * dx - self.source_interpolators.append(Interpolator(dmv_expr, source_mv)) - equation.residual -= self.label(subject(prognostic(source_mv_expr, vapour_name), - X), self.evaluate) + self.source_interpolators.append(Interpolator(dmv_expr, self.source_mv_int)) + equation.residual -= source_label( + self.label(subject(prognostic(source_mv_expr, vapour_name), self.source), self.evaluate) + ) # Moisture needs including in theta_vd expression # NB: still using old pressure here, which implies constant p? @@ -153,12 +155,14 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, else: theta_np1_expr = thermodynamics.theta(equation.parameters, T_np1_expr, p) - source_theta_vd = Function(Vtheta) + self.source_theta_vd = split(self.source)[self.T_idx] + self.source_theta_vd_int = self.source.subfunctions[self.T_idx] dtheta_vd_expr = surface_expr * (theta_np1_expr - theta_vd) / self.dt - source_theta_expr = test_theta * source_theta_vd * dx - self.source_interpolators.append(Interpolator(dtheta_vd_expr, source_theta_vd)) - equation.residual -= self.label(subject(prognostic(source_theta_expr, 'theta'), - X), self.evaluate) + source_theta_expr = test_theta * self.source_theta_vd * dx + self.source_interpolators.append(Interpolator(dtheta_vd_expr, self.source_theta_vd_int)) + equation.residual -= source_label( + self.label(subject(prognostic(source_theta_expr, 'theta'), self.source), self.evaluate) + ) # General formulation ------------------------------------------------ # else: @@ -187,7 +191,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, equation.residual -= self.label( subject(prognostic(source_theta_expr, 'theta'), X), self.evaluate) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source term generated by the physics. This does nothing if the implicit formulation is not used. @@ -196,6 +200,8 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') @@ -206,6 +212,9 @@ def evaluate(self, x_in, dt): self.rho_recoverer.project() for source_interpolator in self.source_interpolators: source_interpolator.interpolate() + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) class WindDrag(PhysicsParametrisation): @@ -280,22 +289,22 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): if implicit_formulation: # First specify T_np1 expression - Vu = equation.spaces[u_idx] - source_u = Function(Vu) + self.source = Function(equation.X.function_space()) + source_u = split(self.source)[u_idx] + source_u_proj = self.source.subfunctions[u_idx] u_np1_expr = u_hori / (1 + C_D*u_hori_mag*self.dt/z_a) du_expr = surface_expr * (u_np1_expr - u_hori) / self.dt - # TODO: introduce reduced projector - test_Vu = TestFunction(Vu) - dx_reduced = dx(degree=4) - proj_eqn = inner(test_Vu, source_u - du_expr)*dx_reduced - proj_prob = NonlinearVariationalProblem(proj_eqn, source_u) - self.source_projector = NonlinearVariationalSolver(proj_prob) + self.source_projector = Projector( + du_expr, source_u_proj, + quadrature_degree=equation.domain.max_quad_degree + ) source_expr = inner(test, source_u - k*dot(source_u, k)) * dx - equation.residual -= self.label(subject(prognostic(source_expr, 'u'), - X), self.evaluate) + equation.residual -= source_label( + self.label(subject(prognostic(source_expr, 'u'), self.source), self.evaluate) + ) # General formulation ------------------------------------------------ # else: @@ -307,7 +316,7 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): equation.residual -= self.label(subject(prognostic(source_expr, 'u'), X), self.evaluate) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source term generated by the physics. This does nothing if the implicit formulation is not used. @@ -316,6 +325,8 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') @@ -323,7 +334,10 @@ def evaluate(self, x_in, dt): if self.implicit_formulation: self.X.assign(x_in) self.dt.assign(dt) - self.source_projector.solve() + self.source_projector.project() + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) class StaticAdjustment(PhysicsParametrisation): @@ -402,7 +416,7 @@ def __init__(self, equation, theta_variable='theta_vd'): equation.residual -= self.label(subject(prognostic(source_expr, 'theta'), equation.X), self.evaluate) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source term generated by the physics. This does nothing if the implicit formulation is not used. @@ -411,6 +425,8 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') @@ -425,6 +441,11 @@ def evaluate(self, x_in, dt): self.set_column_data(self.theta_to_sort, theta_column_data, index_data) self.set_theta_variable.interpolate() + if x_out is not None: + raise NotImplementedError("Static adjustment does not output a source term, " + "or a non-interpolated/projected expression and hence " + "cannot be used in a nonsplit physics formulation.") + class SuppressVerticalWind(PhysicsParametrisation): """ @@ -482,7 +503,7 @@ def __init__(self, equation, spin_up_period): equation.residual -= self.label(subject(prognostic(source_expr, 'u'), equation.X), self.evaluate) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source term generated by the physics. This does nothing if the implicit formulation is not used. @@ -491,6 +512,9 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. + This is unused. """ if float(self.t) < float(self.spin_up_period): @@ -628,7 +652,7 @@ def __init__(self, equation, field_name, parameters=None): equation.residual += self.label( subject(prognostic(source_expr, field_name), X), self.evaluate) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source term generated by the physics. This only recovers the density field. @@ -637,9 +661,16 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') self.X.assign(x_in) self.rho_recoverer.project() + + if x_out is not None: + raise NotImplementedError("Boundary layer mixing does not output a source term, " + "or a non-interpolated/projected expression and hence " + "cannot be used in a nonsplit physics formulation.") diff --git a/gusto/physics/chemistry.py b/gusto/physics/chemistry.py index 2304e64a3..79dde5ef0 100644 --- a/gusto/physics/chemistry.py +++ b/gusto/physics/chemistry.py @@ -53,10 +53,9 @@ def __init__(self, equation, k1=1, k2=1, "The function spaces for the two species need to be the same" self.Xq = Function(equation.X.function_space()) - Xq = self.Xq - species1 = split(Xq)[self.species1_idx] - species2 = split(Xq)[self.species2_idx] + species1 = split(self.Xq)[self.species1_idx] + species2 = split(self.Xq)[self.species2_idx] test_1 = equation.tests[self.species1_idx] test_2 = equation.tests[self.species2_idx] @@ -66,16 +65,18 @@ def __init__(self, equation, k1=1, k2=1, source1_expr = test_1 * 2*Kx * dx source2_expr = test_2 * -Kx * dx - equation.residual -= self.label(subject(prognostic(source1_expr, 'X'), Xq), self.evaluate) - equation.residual -= self.label(subject(prognostic(source2_expr, 'X2'), Xq), self.evaluate) + equation.residual -= self.label(subject(prognostic(source1_expr, 'X'), self.Xq), self.evaluate) + equation.residual -= self.label(subject(prognostic(source2_expr, 'X2'), self.Xq), self.evaluate) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source/sink for the coalescence process. Args: x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') diff --git a/gusto/physics/held_suarez_forcing.py b/gusto/physics/held_suarez_forcing.py new file mode 100644 index 000000000..d7573d009 --- /dev/null +++ b/gusto/physics/held_suarez_forcing.py @@ -0,0 +1,188 @@ +import numpy as np +from firedrake import (Interpolator, Function, dx, pi, SpatialCoordinate, + split, conditional, ge, sin, dot, ln, cos, inner, Projector) +from firedrake.fml import subject +from gusto.core.coord_transforms import lonlatr_from_xyz +from gusto.recovery import Recoverer, BoundaryMethod +from gusto.physics.physics_parametrisation import PhysicsParametrisation +from gusto.core.labels import prognostic +from gusto.equations import thermodynamics +from gusto.core.configuration import HeldSuarezParameters +from gusto.core import logger + + +class Relaxation(PhysicsParametrisation): + """ + Relaxation term for Held Suarez + """ + + def __init__(self, equation, variable_name, parameters, hs_parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + variable_name (str): the name of the variable + hs_parameters (:class'Configuration'): contains the parameters for the Held-suariez test case + + """ + label_name = f'relaxation_{variable_name}' + if hs_parameters is None: + hs_parameters = HeldSuarezParameters() + logger.warning('Using default Held-Suarez parameters') + super().__init__(equation, label_name, hs_parameters) + + if equation.domain.on_sphere: + x, y, z = SpatialCoordinate(equation.domain.mesh) + _, lat, _ = lonlatr_from_xyz(x, y, z) + else: + # TODO: this could be determined some other way + # Take a mid-latitude + lat = pi / 4 + + self.X = Function(equation.X.function_space()) + X = self.X + self.domain = equation.domain + theta_idx = equation.field_names.index('theta') + self.theta = X.subfunctions[theta_idx] + Vt = equation.domain.spaces('theta') + rho_idx = equation.field_names.index('rho') + rho = split(X)[rho_idx] + + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + self.rho_averaged = Function(Vt) + self.rho_recoverer = Recoverer(rho, self.rho_averaged, boundary_method=boundary_method) + self.exner = Function(Vt) + self.exner_interpolator = Interpolator( + thermodynamics.exner_pressure(equation.parameters, + self.rho_averaged, self.theta), self.exner) + self.sigma = Function(Vt) + kappa = equation.parameters.kappa + + T0surf = hs_parameters.T0surf + T0horiz = hs_parameters.T0horiz + T0vert = hs_parameters.T0vert + T0stra = hs_parameters.T0stra + + sigma_b = hs_parameters.sigmab + tau_d = hs_parameters.tau_d + tau_u = hs_parameters.tau_u + + theta_condition = (T0surf - T0horiz * sin(lat)**2 - (T0vert * ln(self.exner) * cos(lat)**2)/kappa) + Theta_eq = conditional(T0stra/self.exner >= theta_condition, T0stra/self.exner, theta_condition) + + # timescale of temperature forcing + tau_cond = (self.sigma**(1/kappa) - sigma_b) / (1 - sigma_b) + newton_freq = 1 / tau_d + (1/tau_u - 1/tau_d) * conditional(0 >= tau_cond, 0, tau_cond) * cos(lat)**4 + forcing_expr = newton_freq * (self.theta - Theta_eq) + + # Create source for forcing + self.source_relaxation = Function(Vt) + self.source_interpolator = Interpolator(forcing_expr, self.source_relaxation) + + # Add relaxation term to residual + test = equation.tests[theta_idx] + dx_reduced = dx(degree=equation.domain.max_quad_degree) + forcing_form = test * self.source_relaxation * dx_reduced + equation.residual += self.label(subject(prognostic(forcing_form, 'theta'), X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evalutes the source term generated by the physics. + + Args: + x_in: (:class:`Function`): the (mixed) field to be evolved. + dt: (:class:`Constant`): the timestep, which can be the time + interval for the scheme. + """ + self.X.assign(x_in) + self.rho_recoverer.project() + self.exner_interpolator.interpolate() + + # Determine sigma:= exner / exner_surf + exner_columnwise, index_data = self.domain.coords.get_column_data(self.exner, self.domain) + sigma_columnwise = np.zeros_like(exner_columnwise) + for col in range(len(exner_columnwise[:, 0])): + sigma_columnwise[col, :] = exner_columnwise[col, :] / exner_columnwise[col, 0] + self.domain.coords.set_field_from_column_data(self.sigma, sigma_columnwise, index_data) + + self.source_interpolator.interpolate() + + +class RayleighFriction(PhysicsParametrisation): + """ + Forcing term on the velocity of the form + F_u = -u / a, + where a is some friction factor + """ + def __init__(self, equation, hs_parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + hs_parameters (:class'Configuration'): contains the parameters for the Held-suariez test case + """ + label_name = 'rayleigh_friction' + if hs_parameters is None: + hs_parameters = HeldSuarezParameters() + logger.warning('Using default Held-Suarez parameters') + super().__init__(equation, label_name, hs_parameters) + + self.domain = equation.domain + self.X = Function(equation.X.function_space()) + X = self.X + k = equation.domain.k + u_idx = equation.field_names.index('u') + u = split(X)[u_idx] + theta_idx = equation.field_names.index('theta') + self.theta = X.subfunctions[theta_idx] + rho_idx = equation.field_names.index('rho') + rho = split(X)[rho_idx] + Vt = equation.domain.spaces('theta') + Vu = equation.domain.spaces('HDiv') + u_hori = u - k*dot(u, k) + + boundary_method = BoundaryMethod.extruded if self.domain == 0 else None + self.rho_averaged = Function(Vt) + self.exner = Function(Vt) + self.rho_recoverer = Recoverer(rho, self.rho_averaged, boundary_method=boundary_method) + self.exner_interpolator = Interpolator( + thermodynamics.exner_pressure(equation.parameters, + self.rho_averaged, self.theta), self.exner) + + self.sigma = Function(Vt) + sigmab = hs_parameters.sigmab + kappa = equation.parameters.kappa + tau_fric = 24 * 60 * 60 + + tau_cond = (self.sigma**(1/kappa) - sigmab) / (1 - sigmab) + wind_timescale = conditional(ge(0, tau_cond), 0, tau_cond) / tau_fric + forcing_expr = u_hori * wind_timescale + + self.source_friction = Function(Vu) + self.source_projector = Projector(forcing_expr, self.source_friction) + + tests = equation.tests + test = tests[u_idx] + dx_reduced = dx(degree=equation.domain.max_quad_degree) + source_form = inner(test, self.source_friction) * dx_reduced + equation.residual += self.label(subject(prognostic(source_form, 'u'), X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This does nothing if + the implicit formulation is not used. + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + """ + self.X.assign(x_in) + self.rho_recoverer.project() + self.exner_interpolator.interpolate() + # Determine sigma:= exner / exner_surf + exner_columnwise, index_data = self.domain.coords.get_column_data(self.exner, self.domain) + sigma_columnwise = np.zeros_like(exner_columnwise) + for col in range(len(exner_columnwise[:, 0])): + sigma_columnwise[col, :] = exner_columnwise[col, :] / exner_columnwise[col, 0] + self.domain.coords.set_field_from_column_data(self.sigma, sigma_columnwise, index_data) + + self.source_projector.project() diff --git a/gusto/physics/microphysics.py b/gusto/physics/microphysics.py index c1d3a440d..cd620b885 100644 --- a/gusto/physics/microphysics.py +++ b/gusto/physics/microphysics.py @@ -4,14 +4,15 @@ """ from firedrake import ( - Interpolator, conditional, Function, dx, min_value, max_value, Constant, pi, - inner, TestFunction, NonlinearVariationalProblem, NonlinearVariationalSolver + conditional, Function, dx, min_value, max_value, Constant, pi, + Projector, assemble, split ) +from firedrake.__future__ import interpolate from firedrake.fml import identity, Term, subject from gusto.equations import Phases, TracerVariableType from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations -from gusto.core.labels import transporting_velocity, transport, prognostic +from gusto.core.labels import transporting_velocity, transport, prognostic, source_label from gusto.core.logging import logger from gusto.equations import thermodynamics from gusto.physics.physics_parametrisation import PhysicsParametrisation @@ -79,6 +80,7 @@ def __init__(self, equation, vapour_name='water_vapour', parameters = self.parameters self.X = Function(equation.X.function_space()) self.latent_heat = latent_heat + W = equation.function_space # Vapour and cloud variables are needed for every form of this scheme cloud_idx = equation.field_names.index(cloud_name) @@ -169,24 +171,28 @@ def __init__(self, equation, vapour_name='water_vapour', # -------------------------------------------------------------------- # # Add terms to equations and make interpolators # -------------------------------------------------------------------- # - self.source = [Function(V) for factor in factors] - self.source_interpolators = [Interpolator(sat_adj_expr*factor, source) - for factor, source in zip(factors, self.source)] + self.source = Function(W) + self.source_expr = [split(self.source)[V_idx] for V_idx in V_idxs] + self.source_int = [self.source.subfunctions[V_idx] for V_idx in V_idxs] + self.source_interpolate = [interpolate(sat_adj_expr*factor, source) + for source, factor in zip(self.source_int, factors)] tests = [equation.tests[idx] for idx in V_idxs] # Add source terms to residual - for test, source in zip(tests, self.source): - equation.residual += self.label(subject(test * source * dx, - equation.X), self.evaluate) + for test, source in zip(tests, self.source_expr): + equation.residual += source_label(self.label(subject(test * source * dx, + equation.X), self.evaluate)) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source/sink for the saturation adjustment process. Args: x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') # Update the values of internal variables @@ -195,8 +201,12 @@ def evaluate(self, x_in, dt): if isinstance(self.equation, CompressibleEulerEquations): self.rho_recoverer.project() # Evaluate the source - for interpolator in self.source_interpolators: - interpolator.interpolate() + for interpolator, src in zip(self.source_interpolate, self.source_int): + src.assign(assemble(interpolator)) + + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) class AdvectedMoments(Enum): @@ -337,25 +347,30 @@ def __init__(self, equation, rain_name, domain, transport_method, + 'AdvectedMoments.M0 and AdvectedMoments.M3') if moments != AdvectedMoments.M0: - # TODO: introduce reduced projector - test = TestFunction(Vu) - dx_reduced = dx(degree=4) - proj_eqn = inner(test, v + v_expression*domain.k)*dx_reduced - proj_prob = NonlinearVariationalProblem(proj_eqn, v) - self.determine_v = NonlinearVariationalSolver(proj_prob) - - def evaluate(self, x_in, dt): + self.determine_v = Projector( + -v_expression*domain.k, v, + quadrature_degree=domain.max_quad_degree + ) + + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source/sink corresponding to the fallout process. Args: x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') self.X.assign(x_in) if self.moments != AdvectedMoments.M0: - self.determine_v.solve() + self.determine_v.project() + + if x_out is not None: + raise NotImplementedError("Fallout does not output a source term, " + "or a non-interpolated/projected expression and hence " + "cannot be used in a nonsplit physics formulation.") class Coalescence(PhysicsParametrisation): @@ -405,12 +420,14 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', self.rain_idx = equation.field_names.index(rain_name) Vcl = equation.function_space.sub(self.cloud_idx) Vr = equation.function_space.sub(self.rain_idx) + W = equation.function_space self.cloud_water = Function(Vcl) self.rain = Function(Vr) # declare function space and source field - Vt = self.cloud_water.function_space() - self.source = Function(Vt) + self.source = Function(W) + self.source_expr = split(self.source)[self.cloud_idx] + self.source_int = self.source.subfunctions[self.cloud_idx] # define some parameters as attributes self.dt = Constant(0.0) @@ -442,31 +459,41 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', min_value(accu_rate, self.cloud_water / self.dt), min_value(accr_rate + accu_rate, self.cloud_water / self.dt)))) - self.source_interpolator = Interpolator(rain_expr, self.source) + self.source_interpolate = interpolate(rain_expr, self.source_int) # Add term to equation's residual test_cl = equation.tests[self.cloud_idx] test_r = equation.tests[self.rain_idx] - equation.residual += self.label(subject(test_cl * self.source * dx - - test_r * self.source * dx, - equation.X), - self.evaluate) - - def evaluate(self, x_in, dt): + equation.residual += source_label( + self.label( + subject( + test_cl * self.source_expr * dx - test_r * self.source_expr * dx, + self.source + ), + self.evaluate + ) + ) + + def evaluate(self, x_in, dt, x_out=None): """ Evaluates the source/sink for the coalescence process. Args: x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') # Update the values of internal variables self.dt.assign(dt) self.rain.assign(x_in.subfunctions[self.rain_idx]) self.cloud_water.assign(x_in.subfunctions[self.cloud_idx]) - # Evaluate the source - self.source.assign(self.source_interpolator.interpolate()) + + self.source_int.assign(assemble(self.source_interpolate)) + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) class EvaporationOfRain(PhysicsParametrisation): @@ -528,6 +555,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', # Indices of variables in mixed function space V_idxs = [rain_idx, vap_idx] V = equation.function_space.sub(rain_idx) # space in which to do the calculation + W = equation.function_space # Get variables used to calculate saturation curve if isinstance(equation, CompressibleEulerEquations): @@ -610,24 +638,30 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', # -------------------------------------------------------------------- # # Add terms to equations and make interpolators # -------------------------------------------------------------------- # - self.source = [Function(V) for factor in factors] - self.source_interpolators = [Interpolator(evap_rate*factor, source) - for factor, source in zip(factors, self.source)] + self.source = Function(W) + self.source_expr = [split(self.source)[V_idx] for V_idx in V_idxs] + self.source_int = [self.source.subfunctions[V_idx] for V_idx in V_idxs] + self.source_interpolate = [ + interpolate(evap_rate * factor, source) + for source, factor in zip(self.source_int, factors) + ] tests = [equation.tests[idx] for idx in V_idxs] # Add source terms to residual - for test, source in zip(tests, self.source): - equation.residual += self.label(subject(test * source * dx, - equation.X), self.evaluate) + for test, source in zip(tests, self.source_expr): + equation.residual += source_label(self.label(subject(test * source * dx, + self.source), self.evaluate)) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Applies the process to evaporate rain droplets. Args: x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') # Update the values of internal variables @@ -635,6 +669,10 @@ def evaluate(self, x_in, dt): self.X.assign(x_in) if isinstance(self.equation, CompressibleEulerEquations): self.rho_recoverer.project() - # Evaluate the source - for interpolator in self.source_interpolators: - interpolator.interpolate() + + for interpolator, src in zip(self.source_interpolate, self.source_int): + src.assign(assemble(interpolator)) + + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) diff --git a/gusto/physics/physics_parametrisation.py b/gusto/physics/physics_parametrisation.py index 2d43dae48..e429f1499 100644 --- a/gusto/physics/physics_parametrisation.py +++ b/gusto/physics/physics_parametrisation.py @@ -8,9 +8,10 @@ """ from abc import ABCMeta, abstractmethod -from firedrake import Interpolator, Function, dx, Projector +from firedrake import Function, dx, Projector, assemble, split +from firedrake.__future__ import interpolate from firedrake.fml import subject -from gusto.core.labels import PhysicsLabel +from gusto.core.labels import PhysicsLabel, source_label from gusto.core.logging import logger @@ -96,18 +97,25 @@ def __init__(self, equation, variable_name, rate_expression, # Work out the appropriate function space if hasattr(equation, "field_names"): - V_idx = equation.field_names.index(variable_name) + self.V_idx = equation.field_names.index(variable_name) W = equation.function_space - V = W.sub(V_idx) - test = equation.tests[V_idx] + V = W.sub(self.V_idx) + test = equation.tests[self.V_idx] + self.V_idx = self.V_idx + self.source = Function(W) + self.source_expr = split(self.source)[self.V_idx] + self.source_int = self.source.subfunctions[self.V_idx] else: V = equation.function_space test = equation.test + self.source = Function(V) + self.source_expr = self.source + self.source_int = self.source # Make source/sink term - self.source = Function(V) - equation.residual += self.label(subject(test * self.source * dx, equation.X), - self.evaluate) + equation.residual += source_label( + self.label(subject(test * self.source_expr * dx, self.source), self.evaluate) + ) # Handle whether the expression is time-varying or not if self.time_varying: @@ -117,18 +125,18 @@ def __init__(self, equation, variable_name, rate_expression, # Handle method of evaluating source/sink if self.method == 'interpolate': - self.source_interpolator = Interpolator(expression, V) + self.source_interpolate = interpolate(expression, self.source_int) else: - self.source_projector = Projector(expression, V) + self.source_projector = Projector(expression, self.source_int) # If not time-varying, evaluate for the first time here if not self.time_varying: if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) + self.source_int.assign(assemble(self.source_interpolate)) else: - self.source.assign(self.source_projector.project()) + self.source_projector.project() - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evalutes the source term generated by the physics. @@ -136,12 +144,17 @@ def evaluate(self, x_in, dt): x_in: (:class:`Function`): the (mixed) field to be evolved. Unused. dt: (:class:`Constant`): the timestep, which can be the time interval for the scheme. Unused. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ if self.time_varying: logger.info(f'Evaluating physics parametrisation {self.label.label}') if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) + self.source_int.assign(assemble(self.source_interpolate)) else: - self.source.assign(self.source_projector.project()) + self.source_projector.project() else: pass + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) diff --git a/gusto/physics/shallow_water_microphysics.py b/gusto/physics/shallow_water_microphysics.py index 90cb5328d..29b3637bd 100644 --- a/gusto/physics/shallow_water_microphysics.py +++ b/gusto/physics/shallow_water_microphysics.py @@ -3,13 +3,16 @@ """ from firedrake import ( - Interpolator, conditional, Function, dx, min_value, max_value, Constant + conditional, Function, dx, min_value, max_value, Constant, assemble, split ) +from firedrake.__future__ import interpolate from firedrake.fml import subject from gusto.core.logging import logger from gusto.physics.physics_parametrisation import PhysicsParametrisation +from gusto.core.labels import source_label from types import FunctionType + __all__ = ["InstantRain", "SWSaturationAdjustment"] @@ -90,7 +93,9 @@ def __init__(self, equation, saturation_curve, # the source function is the difference between the water vapour and # the saturation function self.water_v = Function(Vv) - self.source = Function(Vv) + self.source = Function(W) + self.source_expr = split(self.source)[self.Vv_idx] + self.source_int = self.source.subfunctions[self.Vv_idx] # tau is the timescale for conversion (may or may not be the timestep) if tau is not None: @@ -113,32 +118,33 @@ def __init__(self, equation, saturation_curve, self.saturation_curve = saturation_curve # lose proportion of vapour above the saturation curve - equation.residual += self.label(subject(test_v * self.source * dx, - equation.X), - self.evaluate) + equation.residual += source_label( + self.label(subject(test_v * self.source_expr * dx, self.source), self.evaluate) + ) # if rain is not none then the excess vapour is being tracked and is # added to rain if rain_name is not None: - Vr_idx = equation.field_names.index(rain_name) - test_r = equation.tests[Vr_idx] - equation.residual -= self.label(subject(test_r * self.source * dx, - equation.X), - self.evaluate) + self.Vr_idx = equation.field_names.index(rain_name) + test_r = equation.tests[self.Vr_idx] + equation.residual -= source_label( + self.label(subject(test_r * self.source_expr * dx, self.source), self.evaluate) + ) # if feeding back on the height adjust the height equation if convective_feedback: - equation.residual += self.label(subject(test_D * beta1 * self.source * dx, - equation.X), - self.evaluate) + self.VD_idx = equation.field_names.index("D") + equation.residual += source_label( + self.label(subject(test_D * beta1 * self.source * dx, equation.X), self.evaluate) + ) # interpolator does the conversion of vapour to rain - self.source_interpolator = Interpolator(conditional( + self.source_interpolate = interpolate(conditional( self.water_v > self.saturation_curve, (1/self.tau)*gamma_r*(self.water_v - self.saturation_curve), - 0), Vv) + 0), self.source_int) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ Evalutes the source term generated by the physics. @@ -149,6 +155,8 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') if self.convective_feedback: @@ -158,7 +166,11 @@ def evaluate(self, x_in, dt): if self.set_tau_to_dt: self.tau.assign(dt) self.water_v.assign(x_in.subfunctions[self.Vv_idx]) - self.source.assign(self.source_interpolator.interpolate()) + + self.source_int.assign(assemble(self.source_interpolate)) + + if x_out is not None: + x_out.assign(self.source) class SWSaturationAdjustment(PhysicsParametrisation): @@ -319,9 +331,11 @@ def __init__(self, equation, saturation_curve, # Add terms to equations and make interpolators # sources have the same order as V_idxs and factors - self.source = [Function(Vc) for factor in factors] - self.source_interpolators = [Interpolator(sat_adj_expr*factor, source) - for factor, source in zip(factors, self.source)] + self.source = Function(W) + self.source_expr = [split(self.source)[V_idx] for V_idx in V_idxs] + self.source_int = [self.source.subfunctions[V_idx] for V_idx in V_idxs] + self.source_interpolate = [interpolate(sat_adj_expr*factor, source) + for source, factor in zip(self.source_int, factors)] # test functions have the same order as factors and sources (vapour, # cloud, depth, buoyancy) so that the correct test function multiplies @@ -329,13 +343,14 @@ def __init__(self, equation, saturation_curve, tests = [equation.tests[idx] for idx in V_idxs] # Add source terms to residual - for test, source in zip(tests, self.source): - equation.residual += self.label(subject(test * source * dx, - equation.X), self.evaluate) + for test, source_val in zip(tests, self.source_expr): + equation.residual += source_label( + self.label(subject(test * source_val * dx, self.source), self.evaluate) + ) - def evaluate(self, x_in, dt): + def evaluate(self, x_in, dt, x_out=None): """ - Evaluates the source term generated by the physics. + Evaluates the source_label term generated by the physics. Computes the phyiscs contributions to water vapour and cloud water at each timestep. @@ -344,6 +359,8 @@ def evaluate(self, x_in, dt): x_in: (:class: 'Function'): the (mixed) field to be evolved. dt: (:class: 'Constant'): the timestep, which can be the time interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. """ logger.info(f'Evaluating physics parametrisation {self.label.label}') if self.convective_feedback: @@ -358,5 +375,9 @@ def evaluate(self, x_in, dt): self.cloud.assign(x_in.subfunctions[self.Vc_idx]) if self.time_varying_gamma_v: self.gamma_v.interpolate(self.gamma_v_computation(x_in)) - for interpolator in self.source_interpolators: - interpolator.interpolate() + + for interpolator, src in zip(self.source_interpolate, self.source_int): + src.assign(assemble(interpolator)) + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) diff --git a/gusto/recovery/recovery.py b/gusto/recovery/recovery.py index b88c6ad6a..9d186dd48 100644 --- a/gusto/recovery/recovery.py +++ b/gusto/recovery/recovery.py @@ -13,7 +13,8 @@ Function, FunctionSpace, Interpolator, Projector, SpatialCoordinate, TensorProductElement, VectorFunctionSpace, as_vector, function, interval, - VectorElement) + VectorElement, assemble) +from firedrake.__future__ import interpolate from gusto.recovery import Averager from .recovery_kernels import (BoundaryRecoveryExtruded, BoundaryRecoveryHCurl, BoundaryGaussianElimination) @@ -144,14 +145,14 @@ def __init__(self, x_inout, method=BoundaryMethod.extruded, eff_coords=None): V_broken = FunctionSpace(mesh, BrokenElement(V_inout.ufl_element())) self.x_DG1_wrong = Function(V_broken) self.x_DG1_correct = Function(V_broken) - self.interpolator = Interpolator(self.x_inout, self.x_DG1_wrong) + self.interpolate = interpolate(self.x_inout, V_broken) self.averager = Averager(self.x_DG1_correct, self.x_inout) self.kernel = BoundaryGaussianElimination(V_broken) def apply(self): """Applies the boundary recovery process.""" if self.method == BoundaryMethod.taylor: - self.interpolator.interpolate() + self.x_DG1_wrong.assign(assemble(self.interpolate)) self.kernel.apply(self.x_DG1_wrong, self.x_DG1_correct, self.act_coords, self.eff_coords, self.num_ext) self.averager.project() @@ -275,7 +276,7 @@ def __init__(self, x_in, x_out, method='interpolate', boundary_method=None): self.boundary_recoverers.append(BoundaryRecoverer(x_out_scalars[i], method=BoundaryMethod.taylor, eff_coords=eff_coords[i])) - self.interpolate_to_vector = Interpolator(as_vector(x_out_scalars), self.x_out) + self.interpolate_to_vector = interpolate(as_vector(x_out_scalars), V_out) def project(self): """Perform the whole recovery step.""" @@ -294,7 +295,7 @@ def project(self): # Correct at boundaries boundary_recoverer.apply() # Combine the components to obtain the vector field - self.interpolate_to_vector.interpolate() + self.x_out.assign(assemble(self.interpolate_to_vector)) else: # Extrapolate at boundaries self.boundary_recoverer.apply() diff --git a/gusto/recovery/reversible_recovery.py b/gusto/recovery/reversible_recovery.py index fc8fbd332..d87449817 100644 --- a/gusto/recovery/reversible_recovery.py +++ b/gusto/recovery/reversible_recovery.py @@ -4,7 +4,8 @@ """ from gusto.core.conservative_projection import ConservativeProjector -from firedrake import (Projector, Function, Interpolator) +from firedrake import (Projector, Function, assemble) +from firedrake.__future__ import interpolate from .recovery import Recoverer __all__ = ["ReversibleRecoverer", "ConservativeRecoverer"] @@ -52,7 +53,7 @@ def __init__(self, source_field, target_field, reconstruct_opts): elif self.opts.project_high_method == 'project': self.projector_high = Projector(self.q_recovered, self.q_rec_high) elif self.opts.project_high_method == 'interpolate': - self.projector_high = Interpolator(self.q_recovered, self.q_rec_high) + self.projector_high = interpolate(self.q_recovered, target_field.function_space()) self.interp_high = True else: raise ValueError(f'Method {self.opts.project_high_method} ' @@ -68,7 +69,7 @@ def __init__(self, source_field, target_field, reconstruct_opts): elif self.opts.project_low_method == 'project': self.projector_low = Projector(self.q_rec_high, self.q_corr_low) elif self.opts.project_low_method == 'interpolate': - self.projector_low = Interpolator(self.q_rec_high, self.q_corr_low) + self.projector_low = interpolate(self.q_rec_high, source_field.function_space()) self.interp_low = True else: raise ValueError(f'Method {self.opts.project_low_method} ' @@ -84,17 +85,17 @@ def __init__(self, source_field, target_field, reconstruct_opts): elif self.opts.injection_method == 'project': self.injector = Projector(self.q_corr_low, self.q_corr_high) elif self.opts.injection_method == 'interpolate': - self.injector = Interpolator(self.q_corr_low, self.q_corr_high) + self.injector = interpolate(self.q_corr_low, target_field.function_space()) self.interp_inj = True else: raise ValueError(f'Method {self.opts.injection_method} for injection not valid') def project(self): self.recoverer.project() - self.projector_high.interpolate() if self.interp_high else self.projector_high.project() - self.projector_low.interpolate() if self.interp_low else self.projector_low.project() + self.q_rec_high.assign(assemble(self.projector_high)) if self.interp_high else self.projector_high.project() + self.q_corr_low.assign(assemble(self.projector_low)) if self.interp_low else self.projector_low.project() self.q_corr_low.assign(self.q_low - self.q_corr_low) - self.injector.interpolate() if self.interp_inj else self.injector.project() + self.q_corr_high.assign(assemble(self.injector)) if self.interp_inj else self.injector.project() self.q_high.assign(self.q_corr_high + self.q_rec_high) diff --git a/gusto/solvers/__init__.py b/gusto/solvers/__init__.py index efb4a5af4..fc4c8a7d2 100644 --- a/gusto/solvers/__init__.py +++ b/gusto/solvers/__init__.py @@ -1,2 +1,3 @@ +from gusto.solvers.parameters import * # noqa from gusto.solvers.linear_solvers import * # noqa -from gusto.solvers.preconditioners import * # noqa \ No newline at end of file +from gusto.solvers.preconditioners import * # noqa diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 91361cc24..213d09df2 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -10,10 +10,12 @@ TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs, rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b, ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, cross, - BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector + BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector, + assemble, conditional ) from firedrake.fml import Term, drop from firedrake.petsc import flatten_parameters +from firedrake.__future__ import interpolate from pyop2.profiling import timed_function, timed_region from gusto.equations.active_tracers import TracerVariableType @@ -53,7 +55,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None, """ self.equations = equations self.dt = equations.domain.dt - self.alpha = alpha + self.alpha = Constant(alpha) self.tau_values = tau_values if tau_values is not None else {} if solver_parameters is not None: @@ -83,6 +85,17 @@ def solver_parameters(self): def _setup_solver(self): pass + @abstractmethod + def update_reference_profiles(self): + """ + Update the solver when the reference profiles have changed. + + This typically includes forcing any Jacobians that depend on + the reference profiles to be reassembled, and recalculating + any values derived from the reference values. + """ + pass + @abstractmethod def solve(self): pass @@ -140,8 +153,7 @@ class CompressibleSolver(TimesteppingSolver): 'sub_pc_type': 'ilu'}}} def __init__(self, equations, alpha=0.5, tau_values=None, - quadrature_degree=None, solver_parameters=None, - overwrite_solver_parameters=False): + solver_parameters=None, overwrite_solver_parameters=False): """ Args: equations (:class:`PrognosticEquation`): the model's equation. @@ -149,9 +161,6 @@ def __init__(self, equations, alpha=0.5, tau_values=None, Defaults to 0.5. A value of 1 is fully-implicit. tau_values (dict, optional): contains the semi-implicit relaxation parameters. Defaults to None, in which case the value of alpha is used. - quadrature_degree (tuple, optional): a tuple (q_h, q_v) where q_h is - the required quadrature degree in the horizontal direction and - q_v is that in the vertical direction. Defaults to None. solver_parameters (dict, optional): contains the options to be passed to the underlying :class:`LinearVariationalSolver`. Defaults to None. @@ -161,14 +170,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None, passed in. Defaults to False. """ self.equations = equations - - if quadrature_degree is not None: - self.quadrature_degree = quadrature_degree - else: - dgspace = equations.domain.spaces("DG") - if any(deg > 2 for deg in dgspace.ufl_element().degree()): - logger.warning("default quadrature degree most likely not sufficient for this degree element") - self.quadrature_degree = (5, 5) + self.quadrature_degree = equations.domain.max_quad_degree super().__init__(equations, alpha, tau_values, solver_parameters, overwrite_solver_parameters) @@ -230,12 +232,12 @@ def V(u): h_project = lambda u: u - k*inner(u, k) # Specify degree for some terms as estimated degree is too large - dxp = dx(degree=(self.quadrature_degree)) - dS_vp = dS_v(degree=(self.quadrature_degree)) - dS_hp = dS_h(degree=(self.quadrature_degree)) - ds_vp = ds_v(degree=(self.quadrature_degree)) - ds_tbp = (ds_t(degree=(self.quadrature_degree)) - + ds_b(degree=(self.quadrature_degree))) + dx_qp = dx(degree=(equations.domain.max_quad_degree)) + dS_v_qp = dS_v(degree=(equations.domain.max_quad_degree)) + dS_h_qp = dS_h(degree=(equations.domain.max_quad_degree)) + ds_v_qp = ds_v(degree=(equations.domain.max_quad_degree)) + ds_tb_qp = (ds_t(degree=(equations.domain.max_quad_degree)) + + ds_b(degree=(equations.domain.max_quad_degree))) # Add effect of density of water upon theta, using moisture reference profiles # TODO: Explore if this is the right thing to do for the linear problem @@ -258,10 +260,10 @@ def V(u): _l0 = TrialFunction(Vtrace) _dl = TestFunction(Vtrace) - a_tr = _dl('+')*_l0('+')*(dS_vp + dS_hp) + _dl*_l0*ds_vp + _dl*_l0*ds_tbp + a_tr = _dl('+')*_l0('+')*(dS_v_qp + dS_h_qp) + _dl*_l0*ds_v_qp + _dl*_l0*ds_tb_qp def L_tr(f): - return _dl('+')*avg(f)*(dS_vp + dS_hp) + _dl*f*ds_vp + _dl*f*ds_tbp + return _dl('+')*avg(f)*(dS_v_qp + dS_h_qp) + _dl*f*ds_v_qp + _dl*f*ds_tb_qp cg_ilu_parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', @@ -271,8 +273,10 @@ def L_tr(f): rhobar_avg = Function(Vtrace) exnerbar_avg = Function(Vtrace) - rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) - exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg) + rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg, + constant_jacobian=True) + exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg, + constant_jacobian=True) self.rho_avg_solver = LinearVariationalSolver(rho_avg_prb, solver_parameters=cg_ilu_parameters, @@ -292,16 +296,16 @@ def L_tr(f): eqn = ( # momentum equation u_mass - - beta_u*cp*div(theta_w*V(w))*exnerbar*dxp + - beta_u*cp*div(theta_w*V(w))*exnerbar*dx_qp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). - # + beta*cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp - + beta_u*cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_hp - + beta_u*cp*dot(theta_w*V(w), n)*exnerbar_avg*ds_tbp - - beta_u*cp*div(thetabar_w*w)*exner*dxp + # + beta*cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_v_qp + + beta_u*cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_h_qp + + beta_u*cp*dot(theta_w*V(w), n)*exnerbar_avg*ds_tb_qp + - beta_u*cp*div(thetabar_w*w)*exner*dx_qp # trace terms appearing after integrating momentum equation - + beta_u*cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_vp + dS_hp) - + beta_u*cp*dot(thetabar_w*w, n)*l0*(ds_tbp + ds_vp) + + beta_u*cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_v_qp + dS_h_qp) + + beta_u*cp*dot(thetabar_w*w, n)*l0*(ds_tb_qp + ds_v_qp) # mass continuity equation + (phi*(rho - rho_in) - beta_r*inner(grad(phi), u)*rhobar)*dx + beta_r*jump(phi*u, n=n)*rhobar_avg('+')*(dS_v + dS_h) @@ -310,13 +314,13 @@ def L_tr(f): # constraint equation to enforce continuity of the velocity # through the interior facets and weakly impose the no-slip # condition - + dl('+')*jump(u, n=n)*(dS_vp + dS_hp) - + dl*dot(u, n)*(ds_tbp + ds_vp) + + dl('+')*jump(u, n=n)*(dS_v + dS_h) + + dl*dot(u, n)*(ds_t + ds_b + ds_v) ) # TODO: can we get this term using FML? # contribution of the sponge term if hasattr(self.equations, "mu"): - eqn += dt*self.equations.mu*inner(w, k)*inner(u, k)*dx + eqn += dt*self.equations.mu*inner(w, k)*inner(u, k)*dx_qp if equations.parameters.Omega is not None: Omega = as_vector([0, 0, equations.parameters.Omega]) @@ -328,7 +332,8 @@ def L_tr(f): # Function for the hybridized solutions self.urhol0 = Function(M) - hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0) + hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0, + constant_jacobian=True) hybridized_solver = LinearVariationalSolver(hybridized_prb, solver_parameters=self.solver_parameters, options_prefix='ImplicitSolver') @@ -354,7 +359,8 @@ def L_tr(f): theta_eqn = gamma*(theta - theta_in + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta_t)*dx - theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) + theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta, + constant_jacobian=True) self.theta_solver = LinearVariationalSolver(theta_problem, solver_parameters=cg_ilu_parameters, options_prefix='thetabacksubstitution') @@ -371,10 +377,6 @@ def L_tr(f): @timed_function("Gusto:UpdateReferenceProfiles") def update_reference_profiles(self): - """ - Updates the reference profiles. - """ - with timed_region("Gusto:HybridProjectRhobar"): logger.info('Compressible linear solver: rho average solve') self.rho_avg_solver.solve() @@ -383,6 +385,13 @@ def update_reference_profiles(self): logger.info('Compressible linear solver: Exner average solve') self.exner_avg_solver.solve() + # Because the left hand side of the hybridised problem depends + # on the reference profile, the Jacobian matrix should change + # when the reference profiles are updated. This call will tell + # the hybridized_solver to reassemble the Jacobian next time + # `solve` is called. + self.hybridized_solver.invalidate_jacobian() + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ @@ -521,7 +530,8 @@ def V(u): bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] # Solver for u, p - up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs) + up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs, + constant_jacobian=True) # Provide callback for the nullspace of the trace system def trace_nullsp(T): @@ -544,12 +554,17 @@ def trace_nullsp(T): b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), - self.b) + self.b, + constant_jacobian=True) self.b_solver = LinearVariationalSolver(b_problem) # Log residuals on hybridized solver self.log_ksp_residuals(self.up_solver.snes.ksp) + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + self.up_solver.invalidate_jacobian() + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ @@ -580,17 +595,18 @@ def solve(self, xrhs, dy): class ThermalSWSolver(TimesteppingSolver): - """ - Linear solver object for the thermal shallow water equations. + """Linear solver object for the thermal shallow water equations. - This solves a linear problem for the thermal shallow water equations with - prognostic variables u (velocity), D (depth) and b (buoyancy). It follows - the following strategy: + This solves a linear problem for the thermal shallow water + equations with prognostic variables u (velocity), D (depth) and + either b (buoyancy) or b_e (equivalent buoyancy). It follows the + following strategy: (1) Eliminate b (2) Solve the resulting system for (u, D) using a hybrid-mixed method (3) Reconstruct b - """ + + """ solver_parameters = { 'ksp_type': 'preonly', @@ -609,6 +625,8 @@ class ThermalSWSolver(TimesteppingSolver): @timed_function("Gusto:SolverSetup") def _setup_solver(self): equation = self.equations # just cutting down line length a bit + equivalent_buoyancy = equation.equivalent_buoyancy + dt = self.dt beta_u = dt*self.tau_values.get("u", self.alpha) beta_d = dt*self.tau_values.get("D", self.alpha) @@ -618,14 +636,12 @@ def _setup_solver(self): Vb = equation.domain.spaces("DG") # Check that the third field is buoyancy - if not equation.field_names[2] == 'b': - raise NotImplementedError("Field 'b' must exist to use the thermal linear solver in the SIQN scheme") + if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): + raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") # Split up the rhs vector - self.xrhs = Function(self.equations.function_space) - u_in = split(self.xrhs)[0] - D_in = split(self.xrhs)[1] - b_in = split(self.xrhs)[2] + self.xrhs = Function(equation.function_space) + u_in, D_in, b_in = split(self.xrhs)[0:3] # Build the reduced function space for u, D M = MixedFunctionSpace((Vu, VD)) @@ -633,12 +649,27 @@ def _setup_solver(self): u, D = TrialFunctions(M) # Get background buoyancy and depth - Dbar = split(equation.X_ref)[1] - bbar = split(equation.X_ref)[2] + Dbar, bbar = split(equation.X_ref)[1:3] # Approximate elimination of b b = -dot(u, grad(bbar))*beta_b + b_in + if equivalent_buoyancy: + # compute q_v using q_sat to partition q_t into q_v and q_c + self.q_sat_func = Function(VD) + self.qvbar = Function(VD) + qtbar = split(equation.X_ref)[3] + + # set up interpolators that use the X_ref values for D and b_e + self.q_sat_expr_interpolate = interpolate( + equation.compute_saturation(equation.X_ref), VD) + self.q_v_interpolate = interpolate( + conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), + VD) + + # bbar was be_bar and here we correct to become bbar + bbar += equation.parameters.beta2 * self.qvbar + n = FacetNormal(equation.domain.mesh) eqn = ( @@ -667,7 +698,8 @@ def _setup_solver(self): bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] # Solver for u, D - uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs) + uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs, + constant_jacobian=True) # Provide callback for the nullspace of the trace system def trace_nullsp(T): @@ -688,12 +720,19 @@ def trace_nullsp(T): b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), - self.b) + self.b, + constant_jacobian=True) self.b_solver = LinearVariationalSolver(b_problem) # Log residuals on hybridized solver self.log_ksp_residuals(self.uD_solver.snes.ksp) + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + if self.equations.equivalent_buoyancy: + self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) + self.qvbar.assign(assemble(self.q_v_interpolate)) + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ @@ -756,21 +795,26 @@ class LinearTimesteppingSolver(object): 'sub_pc_type': 'ilu'}} } - def __init__(self, equation, alpha): + def __init__(self, equation, alpha, reference_dependent=True): """ Args: equation (:class:`PrognosticEquation`): the model's equation object. alpha (float): the semi-implicit off-centring factor. A value of 1 is fully-implicit. + reference_dependent: this indicates that the solver Jacobian should + be rebuilt if the reference profiles have been updated. """ + self.reference_dependent = reference_dependent + residual = equation.residual.label_map( lambda t: t.has_label(linearisation), lambda t: Term(t.get(linearisation).form, t.labels), drop) + self.alpha = Constant(alpha) dt = equation.domain.dt W = equation.function_space - beta = dt*alpha + beta = dt*self.alpha # Split up the rhs vector (symbolically) self.xrhs = Function(W) @@ -789,12 +833,18 @@ def __init__(self, equation, alpha): bcs = [DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u']] problem = LinearVariationalProblem(aeqn.form, action(Leqn.form, self.xrhs), - self.dy, bcs=bcs) + self.dy, bcs=bcs, + constant_jacobian=True) self.solver = LinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix='linear_solver') + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + if self.reference_dependent: + self.solver.invalidate_jacobian() + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ @@ -879,7 +929,8 @@ def _setup_solver(self): bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] # Solver for u, D - uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs) + uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs, + constant_jacobian=True) # Provide callback for the nullspace of the trace system def trace_nullsp(T): @@ -893,6 +944,10 @@ def trace_nullsp(T): # Log residuals on hybridized solver self.log_ksp_residuals(self.uD_solver.snes.ksp) + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + self.uD_solver.invalidate_jacobian() + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ diff --git a/gusto/solvers/parameters.py b/gusto/solvers/parameters.py new file mode 100644 index 000000000..f7f0d93dd --- /dev/null +++ b/gusto/solvers/parameters.py @@ -0,0 +1,124 @@ +""" +This module provides some parameters sets that are good defaults +for particular kinds of system. +""" +from gusto.core.function_spaces import is_cg + +__all__ = ['mass_parameters', 'hydrostatic_parameters'] + + +def mass_parameters(V, spaces=None, ignore_vertical=True): + """ + PETSc solver parameters for mass matrices. + + Currently this sets to a monolithic CG+ILU. + + TODO: implement field-by-field parameters that choose + preonly for discontinuous fields and CG for continuous + fields - see docstring below. + + ================= FUTURE DOCSTRING ================= + Any fields which are discontinuous will have block diagonal + mass matrices, so are solved directly using: + 'ksp_type': 'preonly' + 'pc_type': 'ilu' + + All continuous fields are solved with CG, with the preconditioner + being ILU independently on each field. By solving all continuous fields + "monolithically", the total number of inner products is minimised, which + is beneficial for scaling to large core counts because it minimises the + total number of MPI_Allreduce calls. + 'ksp_type': 'cg' + 'pc_type': 'fieldsplit' + 'pc_fieldsplit_type': 'additive' + 'fieldsplit_ksp_type': 'preonly' + 'fieldsplit_pc_type': 'ilu' + + Args: + spaces: Optional `Spaces` object. If present, any subspace + of V that came from the `Spaces` object will use the + continuity information from `spaces`. + If not present, continuity is checked with `is_cg`. + + ignore_vertical: whether to include the vertical direction when checking + field continuity on extruded meshes. If True, only the horizontal + continuity will be considered, e.g. the standard theta space will + be treated as discontinuous. + """ + return { + 'ksp_type': 'cg', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu' + } + + extruded = hasattr(V.mesh, "_base_mesh") + + continuous_fields = set() + for i, Vsub in enumerate(V.subfunctions): + # field = Vsub.name or str(i) + field = str(i) + + if spaces is not None: + continuous = spaces.continuity.get(field, is_cg(Vsub)) + else: + continuous = is_cg(Vsub) + + # For extruded meshes the continuity is recorded + # separately for the horizontal and vertical directions. + if extruded and spaces is not None: + if ignore_vertical: + continuous = continuous['horizontal'] + else: + continuous = (continuous['horizontal'] + or continuous['vertical']) + + if continuous: + continuous_fields.add(field) + + if len(V.subfunctions) == 1: + parameters = { + 'ksp_type': 'cg' if all(continuous_fields) else 'preonly', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu', + } + else: + + parameters = { + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'additive', + 'pc_fieldsplit_0_fields': ','.join(continuous_fields), + 'fieldsplit': { + 'ksp_type': 'preonly', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu' + }, + 'fieldsplit_0_ksp_type': 'cg', + } + + return parameters + + +hydrostatic_parameters = { + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.SCPC', + # Velocity mass operator is singular in the hydrostatic case. + # So for reconstruction, we eliminate rho into u + 'pc_sc_eliminate_fields': '1, 0', + 'condensed_field': { + 'ksp_type': 'fgmres', + 'ksp_rtol': 1.0e-8, + 'ksp_atol': 1.0e-8, + 'ksp_max_it': 100, + 'pc_type': 'gamg', + 'pc_gamg_sym_graph': True, + 'mg_levels': { + 'ksp_type': 'gmres', + 'ksp_max_it': 5, + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu' + } + } +} diff --git a/gusto/spatial_methods/__init__.py b/gusto/spatial_methods/__init__.py index baa1e0788..349094cdf 100644 --- a/gusto/spatial_methods/__init__.py +++ b/gusto/spatial_methods/__init__.py @@ -2,3 +2,4 @@ from gusto.spatial_methods.diffusion_methods import * # noqa from gusto.spatial_methods.transport_methods import * # noqa from gusto.spatial_methods.limiters import * # noqa +from gusto.spatial_methods.augmentation import * # noqa \ No newline at end of file diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py new file mode 100644 index 000000000..0a893394d --- /dev/null +++ b/gusto/spatial_methods/augmentation.py @@ -0,0 +1,240 @@ +""" +A module defining objects for temporarily augmenting an equation with another. +""" + + +from abc import ABCMeta, abstractmethod +from firedrake import ( + MixedFunctionSpace, Function, TestFunctions, split, inner, dx, grad, + LinearVariationalProblem, LinearVariationalSolver, lhs, rhs, dot, + ds_b, ds_v, ds_t, ds, FacetNormal, TestFunction, TrialFunction, + transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div, + Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC +) +from firedrake.fml import subject +from gusto import ( + time_derivative, transport, transporting_velocity, TransportEquationType, + logger +) + + +class Augmentation(object, metaclass=ABCMeta): + """ + Augments an equation with another equation to be solved simultaneously. + """ + + @abstractmethod + def pre_apply(self, x_in): + """ + Steps to take at the beginning of an apply method, for instance to + assign the input field to the internal mixed function. + """ + + pass + + @abstractmethod + def post_apply(self, x_out): + """ + Steps to take at the end of an apply method, for instance to assign the + internal mixed function to the output field. + """ + + pass + + @abstractmethod + def update(self, x_in_mixed): + """ + Any intermediate update steps, depending on the current mixed function. + """ + + pass + + +class VorticityTransport(Augmentation): + """ + Solves the transport of a vector field, simultaneously with the vorticity + as a mixed proble, as described in Bendall and Wimmer (2022). + + Note that this is most effective with implicit time discretisations. The + residual-based SUPG option provides a dissipation method. + + Args: + domain (:class:`Domain`): The domain object. + eqns (:class:`PrognosticEquationSet`): The overarching equation set. + transpose_commutator (bool, optional): Whether to include the commutator + of the transpose gradient terms. This is necessary for solving the + general vector transport equation, but is not necessary when the + transporting and transported fields are the same. Defaults to True. + supg (bool, optional): Whether to include dissipation through a + residual-based SUPG scheme. Defaults to False. + """ + + def __init__( + self, domain, eqns, transpose_commutator=True, supg=False + ): + + V_vel = domain.spaces('HDiv') + V_vort = domain.spaces('H1') + + self.fs = MixedFunctionSpace((V_vel, V_vort)) + self.X = Function(self.fs) + self.tests = TestFunctions(self.fs) + + u = Function(V_vel) + F, Z = split(self.X) + test_F, test_Z = self.tests + + quad = domain.max_quad_degree + + if hasattr(domain.mesh, "_base_mesh"): + self.ds = ds_b(degree=quad) + ds_t(degree=quad) + ds_v(degree=quad) + self.dS = dS_v(degree=quad) + dS_h(degree=quad) + else: + self.ds = ds(degree=quad) + self.dS = dS(degree=quad) + + # Add boundary conditions + self.bcs = [] + if 'u' in eqns.bcs.keys(): + for bc in eqns.bcs['u']: + self.bcs.append( + DirichletBC(self.fs.sub(0), bc.function_arg, bc.sub_domain) + ) + + # Set up test function and the vorticity term + n = FacetNormal(domain.mesh) + sign_u = 0.5*(sign(dot(u, n)) + 1) + upw = lambda f: (sign_u('+')*f('+') + sign_u('-')*f('-')) + + if domain.mesh.topological_dimension() == 2: + mix_test = test_F - domain.perp(grad(test_Z)) + F_cross_u = Z*domain.perp(u) + elif domain.mesh.topological_dimension == 3: + mix_test = test_F - curl(test_Z) + F_cross_u = cross(Z, u) + + time_deriv_form = inner(F, test_F)*dx + inner(Z, test_Z)*dx + + # Standard vector invariant transport form ----------------------------- + transport_form = ( + # vorticity term + inner(mix_test, F_cross_u)*dx + + inner(n, test_Z*Z*u)*self.ds + # 0.5*grad(v . F) + - 0.5 * div(mix_test) * inner(u, F)*dx + + 0.5 * inner(mix_test, n) * inner(u, F)*self.ds + ) + + # Communtator of tranpose gradient terms ------------------------------- + # This is needed for general vector transport + if transpose_commutator: + u_dot_nabla_F = dot(u, transpose(nabla_grad(F))) + transport_form += ( + - inner(n, test_Z*domain.perp(u_dot_nabla_F))*self.ds + # + 0.5*grad(F).v + - 0.5 * dot(F, div(outer(u, mix_test)))*dx + + 0.5 * inner(mix_test('+'), n('+'))*dot(jump(u), upw(F))*self.dS + # - 0.5*grad(v).F + + 0.5 * dot(u, div(outer(F, mix_test)))*dx + - 0.5 * inner(mix_test('+'), n('+'))*dot(jump(F), upw(u))*self.dS + ) + + # SUPG terms ----------------------------------------------------------- + # Add the vorticity residual to the transported vorticity, + # which damps enstrophy + if supg: + + # Determine SUPG coefficient --------------------------------------- + tau = 0.5*domain.dt + + # Find mean grid spacing to determine a Courant number + DG0 = FunctionSpace(domain.mesh, 'DG', 0) + ones = Function(DG0).interpolate(Constant(1.0)) + area = assemble(ones*dx) + mean_dx = (area/DG0.dof_count)**(1/domain.mesh.geometric_dimension()) + + # Divide by approximately (1 + c) + tau /= (1.0 + sqrt(dot(u, u))*domain.dt/Constant(mean_dx)) + + dxqp = dx(degree=3) + + if domain.mesh.topological_dimension() == 2: + time_deriv_form -= inner(mix_test, tau*Z*domain.perp(u)/domain.dt)*dxqp + transport_form -= inner( + mix_test, tau*domain.perp(u)*domain.divperp(Z*domain.perp(u)) + )*dxqp + if transpose_commutator: + transport_form -= inner( + mix_test, + tau*domain.perp(u)*domain.divperp(u_dot_nabla_F) + )*dxqp + elif domain.mesh.topological_dimension() == 3: + time_deriv_form -= inner(mix_test, tau*cross(Z, u)/domain.dt)*dxqp + transport_form -= inner( + mix_test, tau*cross(curl(Z*u), u) + )*dxqp + if transpose_commutator: + transport_form -= inner( + mix_test, + tau*cross(curl(u_dot_nabla_F), u) + )*dxqp + + # Assemble the residual ------------------------------------------------ + residual = ( + time_derivative(time_deriv_form) + + transport( + transport_form, TransportEquationType.vector_invariant + ) + ) + residual = transporting_velocity(residual, u) + + self.residual = subject(residual, self.X) + + self.x_in = Function(self.fs) + self.Z_in = Function(V_vort) + self.x_out = Function(self.fs) + + vort_test = TestFunction(V_vort) + vort_trial = TrialFunction(V_vort) + + F_in, _ = split(self.x_in) + + eqn = ( + inner(vort_trial, vort_test)*dx + + inner(domain.perp(grad(vort_test)), F_in)*dx + + vort_test*inner(n, domain.perp(F_in))*self.ds + ) + problem = LinearVariationalProblem( + lhs(eqn), rhs(eqn), self.Z_in, constant_jacobian=True + ) + self.solver = LinearVariationalSolver(problem) + + def pre_apply(self, x_in): + """ + Sets the velocity field for the local mixed function. + + Args: + x_in (:class:`Function`): The input velocity field + """ + self.x_in.subfunctions[0].assign(x_in) + + def post_apply(self, x_out): + """ + Sets the output velocity field from the local mixed function. + + Args: + x_out (:class:`Function`): the output velocity field. + """ + x_out.assign(self.x_out.subfunctions[0]) + + def update(self, x_in_mixed): + """ + Performs the solve to determine the vorticity function. + + Args: + x_in_mixed (:class:`Function`): The mixed function to update. + """ + self.x_in.subfunctions[0].assign(x_in_mixed.subfunctions[0]) + logger.debug('Vorticity solve') + self.solver.solve() + self.x_in.subfunctions[1].assign(self.Z_in) diff --git a/gusto/spatial_methods/diffusion_methods.py b/gusto/spatial_methods/diffusion_methods.py index c36fe1be3..768b7532e 100644 --- a/gusto/spatial_methods/diffusion_methods.py +++ b/gusto/spatial_methods/diffusion_methods.py @@ -44,7 +44,8 @@ def interior_penalty_diffusion_form(domain, test, q, parameters): :class:`ufl.Form`: the diffusion form. """ - dS_ = (dS_v + dS_h) if domain.mesh.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if domain.mesh.extruded else dS(degree=quad) kappa = parameters.kappa mu = parameters.mu diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index fc07fcbd6..a800a8812 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -472,7 +472,8 @@ def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=Fa if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS(degree=quad) ubar = Function(Vu) if ibp == IntegrateByParts.ONCE: @@ -534,10 +535,12 @@ def upwind_advection_form_1d(domain, test, q, ibp=IntegrateByParts.ONCE, ubar = Function(Vu) n = FacetNormal(domain.mesh) un = 0.5*(ubar * n[0] + abs(ubar * n[0])) + quad = domain.max_quad_degree + dS_ = dS(degree=quad) if ibp == IntegrateByParts.ONCE: L = -(test * ubar).dx(0) * q * dx + \ - jump(test) * (un('+')*q('+') - un('-')*q('-'))*dS + jump(test) * (un('+')*q('+') - un('-')*q('-'))*dS_ else: raise NotImplementedError("1d advection form only implemented with option ibp=IntegrateByParts.ONCE") @@ -576,7 +579,8 @@ def upwind_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=F if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS(degree=quad) ubar = Function(Vu) if ibp == IntegrateByParts.ONCE: @@ -638,10 +642,12 @@ def upwind_continuity_form_1d(domain, test, q, ibp=IntegrateByParts.ONCE, ubar = Function(Vu) n = FacetNormal(domain.mesh) un = 0.5*(ubar * n[0] + abs(ubar * n[0])) + quad = domain.max_quad_degree + dS_ = dS(degree=quad) if ibp == IntegrateByParts.ONCE: L = -test.dx(0) * q * ubar * dx \ - + jump(test) * (un('+')*q('+') - un('-')*q('-')) * dS + + jump(test) * (un('+')*q('+') - un('-')*q('-')) * dS_ else: raise NotImplementedError("1d continuity form only implemented with option ibp=IntegrateByParts.ONCE") @@ -682,7 +688,8 @@ def upwind_tracer_conservative_form(domain, test, q, rho, if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS(degree=quad) ubar = Function(Vu) if ibp == IntegrateByParts.ONCE: @@ -738,7 +745,8 @@ def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, o # TODO: there should maybe be a restriction on IBP here Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS(degree=quad) ubar = Function(Vu) n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) @@ -775,7 +783,8 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, L = upwind_continuity_form(domain, test, q, ibp, outflow) Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS(degree=quad) ubar = Function(Vu) n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) @@ -815,7 +824,8 @@ def upwind_circulation_form(domain, test, q, ibp=IntegrateByParts.ONCE): """ Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS(degree=quad) ubar = Function(Vu) n = FacetNormal(domain.mesh) Upwind = 0.5*(sign(dot(ubar, n))+1) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 0055ddb89..187339df8 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -5,18 +5,18 @@ from enum import Enum from firedrake import (Function, Constant, NonlinearVariationalProblem, NonlinearVariationalSolver) -from firedrake.fml import replace_subject, all_terms, drop, keep, Term +from firedrake.fml import replace_subject, drop, keep, Term from firedrake.utils import cached_property from firedrake.formmanipulation import split_form -from gusto.core.labels import time_derivative, all_but_last +from gusto.core.labels import time_derivative, all_but_last, source_label from gusto.core.logging import logger from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation __all__ = [ - "ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun", - "RungeKuttaFormulation" + "ForwardEuler", "ExplicitRungeKutta", "SSPRK2", "SSPRK3", "SSPRK4", + "RK4", "Heun", "RungeKuttaFormulation" ] @@ -89,7 +89,8 @@ class ExplicitRungeKutta(ExplicitTimeDiscretisation): def __init__(self, domain, butcher_matrix, field_name=None, subcycling_options=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None): + solver_parameters=None, limiter=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -112,19 +113,19 @@ def __init__(self, domain, butcher_matrix, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, field_name=field_name, subcycling_options=subcycling_options, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) self.butcher_matrix = butcher_matrix - self.nbutcher = int(np.shape(self.butcher_matrix)[0]) + self.nStages = int(np.shape(self.butcher_matrix)[0]) self.rk_formulation = rk_formulation - @property - def nStages(self): - return self.nbutcher - def setup(self, equation, apply_bcs=True, *active_labels): """ Set up the time discretisation based on the equation. @@ -138,6 +139,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.rk_formulation == RungeKuttaFormulation.predictor: self.field_i = [Function(self.fs) for _ in range(self.nStages+1)] + self.source_i = [Function(self.fs) for _ in range(self.nStages+1)] elif self.rk_formulation == RungeKuttaFormulation.increment: self.k = [Function(self.fs) for _ in range(self.nStages)] elif self.rk_formulation == RungeKuttaFormulation.linear: @@ -158,7 +160,7 @@ def solver(self): for stage in range(self.nStages): # setup linear solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem( - self.lhs[stage].form - self.rhs[stage].form, + self.res[stage].form, self.field_i[stage+1], bcs=self.bcs ) solver_name = self.field_name+self.__class__.__name__+str(stage) @@ -171,7 +173,7 @@ def solver(self): elif self.rk_formulation == RungeKuttaFormulation.linear: problem = NonlinearVariationalProblem( - self.lhs - self.rhs[0], self.x1, bcs=self.bcs + self.res[0], self.x1, bcs=self.bcs ) solver_name = self.field_name+self.__class__.__name__ solver = NonlinearVariationalSolver( @@ -181,7 +183,7 @@ def solver(self): # Set up problem for final step problem_last = NonlinearVariationalProblem( - self.lhs - self.rhs[1], self.x1, bcs=self.bcs + self.res[1], self.x1, bcs=self.bcs ) solver_name = self.field_name+self.__class__.__name__+'_last' solver_last = NonlinearVariationalSolver( @@ -197,54 +199,21 @@ def solver(self): ) @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" + def res(self): + """Set up the discretisation's residual.""" if self.rk_formulation == RungeKuttaFormulation.increment: - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, self.idx), - map_if_false=drop) - - return l.form - - elif self.rk_formulation == RungeKuttaFormulation.predictor: - lhs_list = [] - for stage in range(self.nStages): - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.field_i[stage+1], self.idx), - map_if_false=drop) - lhs_list.append(l) - - return lhs_list - - if self.rk_formulation == RungeKuttaFormulation.linear: - l = self.residual.label_map( + residual = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x1, self.idx), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), map_if_false=drop) - - return l.form - - else: - raise NotImplementedError( - 'Runge-Kutta formulation is not implemented' - ) - - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side.""" - - if self.rk_formulation == RungeKuttaFormulation.increment: r = self.residual.label_map( - all_terms, + lambda t: not t.has_label(source_label), map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map( + residual += r.label_map( lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=lambda t: -1*t) + map_if_true=drop) # If there are no active labels, we may have no terms at this point # So that we can still do xnp1 = xn, put in a zero term here @@ -256,38 +225,51 @@ def rhs(self): # Drop label from this map_if_true=lambda t: time_derivative.remove(t), map_if_false=drop) - r += null_term + residual += null_term - return r.form + return residual.form elif self.rk_formulation == RungeKuttaFormulation.predictor: - rhs_list = [] - + residual_list = [] for stage in range(self.nStages): + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.field_i[stage+1], self.idx), + map_if_false=drop) r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) + lambda t: not t.has_label(source_label), + map_if_true=replace_subject(self.field_i[0], old_idx=self.idx), + map_if_false=drop) - r = r.label_map( + residual -= r.label_map( lambda t: t.has_label(time_derivative), map_if_true=keep, map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t) for i in range(1, stage+1): r_i = self.residual.label_map( - lambda t: t.has_label(time_derivative), + lambda t: any(t.has_label(time_derivative, source_label)), map_if_true=drop, map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) ) - r -= self.butcher_matrix[stage, i]*self.dt*r_i + residual += self.butcher_matrix[stage, i]*self.dt*r_i + # Add on any source terms + for i in range(0, stage+1): + r_source = self.residual.label_map( + lambda t: t.has_label(source_label), + map_if_true=replace_subject(self.source_i[i], old_idx=self.idx), + map_if_false=drop) + residual += self.butcher_matrix[stage, i]*self.dt*r_source + residual_list.append(residual) - rhs_list.append(r) - - return rhs_list - - elif self.rk_formulation == RungeKuttaFormulation.linear: + return residual_list + if self.rk_formulation == RungeKuttaFormulation.linear: + time_term = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x1, self.idx), + map_if_false=drop) r = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x0, old_idx=self.idx), @@ -326,8 +308,9 @@ def rhs(self): map_if_true=keep, map_if_false=lambda t: -self.dt*t ) - - return r_all_but_last.form, r.form + res = time_term - r + res_all_but_last = time_term - r_all_but_last + return res_all_but_last.form, res.form else: raise NotImplementedError( @@ -369,17 +352,12 @@ def solve_stage(self, x0, stage): # Use previous stage value as a first guess (otherwise may not converge) self.field_i[stage+1].assign(self.field_i[stage]) - - # Update field_i for physics / limiters - for evaluate in self.evaluate_source: - # TODO: not implemented! Here we need to evaluate the m-th term - # in the i-th RHS with field_m - raise NotImplementedError( - 'Physics not implemented with RK schemes that use the ' - + 'predictor form') if self.limiter is not None: self.limiter.apply(self.field_i[stage]) + for evaluate in self.evaluate_source: + evaluate(self.field_i[stage], self.dt, x_out=self.source_i[stage]) + # Obtain field_ip1 = field_n - dt* sum_m{a_im*F[field_m]} self.solver[stage].solve() @@ -461,6 +439,9 @@ def apply_cycle(self, x_out, x_in): x_out (:class:`Function`): the output field to be computed. """ + if self.augmentation is not None: + self.augmentation.update(x_in) + # TODO: is this limiter application necessary? if self.limiter is not None: self.limiter.apply(x_in) @@ -484,7 +465,8 @@ class ForwardEuler(ExplicitRungeKutta): def __init__( self, domain, field_name=None, subcycling_options=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None + solver_parameters=None, limiter=None, options=None, + augmentation=None ): """ Args: @@ -506,6 +488,9 @@ def __init__( options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ butcher_matrix = np.array([1.]).reshape(1, 1) @@ -514,23 +499,136 @@ def __init__( subcycling_options=subcycling_options, rk_formulation=rk_formulation, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) + + +class SSPRK2(ExplicitRungeKutta): + u""" + Implements 2nd order Strong-Stability-Preserving Runge-Kutta methods + for solving ∂y/∂t = F(y). \n + + Spiteri & Ruuth, 2002, SIAM J. Numer. Anal. \n + "A new class of optimal high-order strong-stability-preserving time \n + discretisation methods". \n + + The 2-stage method (Heun's method) can be written as: \n + + k0 = F[y^n] \n + k1 = F{y^n + dt*k0] \n + y^(n+1) = y^n + (1/2)*dt*(k0+k1) \n + + The 3-stage method can be written as: \n + + k0 = F[y^n] \n + k1 = F[y^n + (1/2*dt*k0] \n + k2 = F[y^n + (1/2)*dt*(k0+k1)] \n + y^(n+1) = y^n + (1/3)*dt*(k0 + k1 + k2) \n + + The 4-stage method can be written as: \n + + k0 = F[y^n] \n + k1 = F[y^n + (1/3)*dt*k1] \n + k2 = F[y^n + (1/3)*dt*(k0+k1)] \n + k3 = F[y^n + (1/3)*dt*(k0+k1+k2)] \n + y^(n+1) = y^n + (1/4)*dt*(k0 + k1 + k2 + k3) \n + """ + def __init__( + self, domain, field_name=None, subcycling_options=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None, + augmentation=None, stages=2 + ): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + subcycling_options(:class:`SubcyclingOptions`, optional): an object + containing options for subcycling the time discretisation. + Defaults to None. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. + stages (int, optional): number of stages: (2, 3, 4). Defaults to 2. + """ + + if stages == 2: + butcher_matrix = np.array([ + [1., 0.], + [0.5, 0.5] + ]) + self.cfl_limit = 1 + + elif stages == 3: + butcher_matrix = np.array([ + [1./2., 0., 0.], + [1./2., 1./2., 0.], + [1./3., 1./3., 1./3.] + ]) + self.cfl_limit = 2 + elif stages == 4: + butcher_matrix = np.array([ + [1./3., 0., 0., 0.], + [1./3., 1./3., 0., 0.], + [1./3., 1./3., 1./3., 0.], + [1./4., 1./4., 1./4., 1./4.] + ]) + self.cfl_limit = 3 + else: + raise ValueError(f"{stages} stage 2rd order SSPRK not implemented") + + super().__init__(domain, butcher_matrix, field_name=field_name, + subcycling_options=subcycling_options, + rk_formulation=rk_formulation, + solver_parameters=solver_parameters, + limiter=limiter, options=options, + augmentation=augmentation) class SSPRK3(ExplicitRungeKutta): u""" - Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method - for solving ∂y/∂t = F(y). It can be written as: \n + Implements 3rd order Strong-Stability-Preserving Runge-Kutta methods + for solving ∂y/∂t = F(y). \n + + Spiteri & Ruuth, 2002, SIAM J. Numer. Anal. \n + "A new class of optimal high-order strong-stability-preserving time \n + discretisation methods". \n + + The 3-stage method can be written as: \n k0 = F[y^n] \n k1 = F[y^n + dt*k1] \n k2 = F[y^n + (1/4)*dt*(k0+k1)] \n y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) \n + + The 4-stage method can be written as: \n + + k0 = F[y^n] \n + k1 = F[y^n + (1/2)*dt*k1] \n + k2 = F[y^n + (1/2)*dt*(k0+k1)] \n + k3 = F[y^n + (1/6)*dt*(k0+k1+k2)] \n + y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + k2 + 3*k3) \n + + The 5-stage method has numerically optimised coefficients. \n """ def __init__( self, domain, field_name=None, subcycling_options=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None + solver_parameters=None, limiter=None, options=None, + augmentation=None, stages=3 ): """ Args: @@ -552,39 +650,63 @@ def __init__( options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. + stages (int, optional): number of stages: (3, 4, 5). Defaults to 3. """ - butcher_matrix = np.array([ - [1., 0., 0.], - [1./4., 1./4., 0.], - [1./6., 1./6., 2./3.] - ]) + if stages == 3: + butcher_matrix = np.array([ + [1., 0., 0.], + [1./4., 1./4., 0.], + [1./6., 1./6., 2./3.] + ]) + self.cfl_limit = 1 + elif stages == 4: + butcher_matrix = np.array([ + [1./2., 0., 0., 0.], + [1./2., 1./2., 0., 0.], + [1./6., 1./6., 1./6., 0.], + [1./6., 1./6., 1./6., 1./2.] + ]) + self.cfl_limit = 2 + elif stages == 5: + self.cfl_limit = 2.65062919294483 + butcher_matrix = np.array([ + [0.37726891511710, 0., 0., 0., 0.], + [0.37726891511710, 0.37726891511710, 0., 0., 0.], + [0.16352294089771, 0.16352294089771, 0.16352294089771, 0., 0.], + [0.14904059394856, 0.14831273384724, 0.14831273384724, 0.34217696850008, 0.], + [0.19707596384481, 0.11780316509765, 0.11709725193772, 0.27015874934251, 0.29786487010104] + ]) + else: + raise ValueError(f"{stages} stage 3rd order SSPRK not implemented") + super().__init__(domain, butcher_matrix, field_name=field_name, subcycling_options=subcycling_options, rk_formulation=rk_formulation, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) -class RK4(ExplicitRungeKutta): +class SSPRK4(ExplicitRungeKutta): u""" - Implements the classic 4-stage Runge-Kutta method. + Implements 4th order Strong-Stability-Preserving Runge-Kutta methods + for solving ∂y/∂t = F(y). \n - The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be - written as: \n + Spiteri & Ruuth, 2002, SIAM J. Numer. Anal. \n + "A new class of optimal high-order strong-stability-preserving time \n + discretisation methods". \n - k0 = F[y^n] \n - k1 = F[y^n + 1/2*dt*k1] \n - k2 = F[y^n + 1/2*dt*k2] \n - k3 = F[y^n + dt*k3] \n - y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3) \n - - where superscripts indicate the time-level. \n + The 5-stage method has numerically optimised coefficients. \n """ def __init__( self, domain, field_name=None, subcycling_options=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None + solver_parameters=None, limiter=None, options=None, + augmentation=None, stages=5 ): """ Args: @@ -606,37 +728,52 @@ def __init__( options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. + stages (int, optional): number of stages: (5,). Defaults to 5. """ - butcher_matrix = np.array([ - [0.5, 0., 0., 0.], - [0., 0.5, 0., 0.], - [0., 0., 1., 0.], - [1./6., 1./3., 1./3., 1./6.] - ]) + + if stages == 5: + self.cfl_limit = 1.50818004975927 + butcher_matrix = np.array([ + [0.39175222700392, 0., 0., 0., 0.], + [0.21766909633821, 0.36841059262959, 0., 0., 0.], + [0.08269208670950, 0.13995850206999, 0.25189177424738, 0., 0.], + [0.06796628370320, 0.11503469844438, 0.20703489864929, 0.54497475021237, 0.], + [0.14681187618661, 0.24848290924556, 0.10425883036650, 0.27443890091960, 0.22600748319395] + ]) + else: + raise ValueError(f"{stages} stage 4rd order SSPRK not implemented") + super().__init__(domain, butcher_matrix, field_name=field_name, subcycling_options=subcycling_options, rk_formulation=rk_formulation, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) -class Heun(ExplicitRungeKutta): +class RK4(ExplicitRungeKutta): u""" - Implements Heun's method. + Implements the classic 4-stage Runge-Kutta method. - The 2-stage Runge-Kutta scheme known as Heun's method,for solving - ∂y/∂t = F(y). It can be written as: \n + The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be + written as: \n - y_1 = F[y^n] \n - y^(n+1) = (1/2)y^n + (1/2)F[y_1] \n + k0 = F[y^n] \n + k1 = F[y^n + 1/2*dt*k1] \n + k2 = F[y^n + 1/2*dt*k2] \n + k3 = F[y^n + dt*k3] \n + y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3) \n - where superscripts indicate the time-level and subscripts indicate the stage - number. + where superscripts indicate the time-level. \n """ def __init__( self, domain, field_name=None, subcycling_options=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None + solver_parameters=None, limiter=None, options=None, + augmentation=None ): """ Args: @@ -658,14 +795,40 @@ def __init__( options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ - butcher_matrix = np.array([ - [1., 0.], - [0.5, 0.5] + [0.5, 0., 0., 0.], + [0., 0.5, 0., 0.], + [0., 0., 1., 0.], + [1./6., 1./3., 1./3., 1./6.] ]) super().__init__(domain, butcher_matrix, field_name=field_name, subcycling_options=subcycling_options, rk_formulation=rk_formulation, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) + + +def Heun(domain, field_name=None, subcycling_options=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None, + augmentation=None): + u""" + Implements Heun's method. + + The 2-stage Runge-Kutta scheme known as Heun's method,for solving + ∂y/∂t = F(y). It can be written as: \n + + y_1 = F[y^n] \n + y^(n+1) = (1/2)y^n + (1/2)F[y_1] \n + + where superscripts indicate the time-level and subscripts indicate the stage + number. + """ + return SSPRK2(domain, field_name=field_name, subcycling_options=subcycling_options, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, + limiter=limiter, options=options, augmentation=augmentation, stages=2) diff --git a/gusto/time_discretisation/imex_runge_kutta.py b/gusto/time_discretisation/imex_runge_kutta.py index fcaefe1ec..ba3911fd4 100644 --- a/gusto/time_discretisation/imex_runge_kutta.py +++ b/gusto/time_discretisation/imex_runge_kutta.py @@ -4,7 +4,7 @@ NonlinearVariationalSolver) from firedrake.fml import replace_subject, all_terms, drop from firedrake.utils import cached_property -from gusto.core.labels import time_derivative, implicit, explicit +from gusto.core.labels import time_derivative, implicit, explicit, source_label from gusto.time_discretisation.time_discretisation import ( TimeDiscretisation, wrapper_apply ) @@ -61,7 +61,7 @@ class IMEXRungeKutta(TimeDiscretisation): def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -82,10 +82,13 @@ def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, field_name=field_name, solver_parameters=nonlinear_solver_parameters, - options=options) + options=options, augmentation=augmentation) self.butcher_imp = butcher_imp self.butcher_exp = butcher_exp self.nStages = int(np.shape(self.butcher_imp)[1]) @@ -122,20 +125,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Check all terms are labeled implicit, exlicit for t in self.residual: if ((not t.has_label(implicit)) and (not t.has_label(explicit)) - and (not t.has_label(time_derivative))): - raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit") + and (not t.has_label(time_derivative)) and (not t.has_label(source_label))): + raise NotImplementedError("Non time-derivative or source terms must be labeled as implicit or explicit") self.xs = [Function(self.fs) for i in range(self.nStages)] - - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(IMEXRungeKutta, self).lhs - - @cached_property - def rhs(self): - """Set up the discretisation's right hand side (the time derivative).""" - return super(IMEXRungeKutta, self).rhs + self.source = [Function(self.fs) for i in range(self.nStages)] def res(self, stage): """Set up the discretisation's residual for a given stage.""" @@ -168,6 +162,18 @@ def res(self, stage): map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) residual += r_imp residual += r_exp + + # Calculate source terms + r_source = self.residual.label_map( + lambda t: t.has_label(source_label), + map_if_true=replace_subject(self.source[i], old_idx=self.idx), + map_if_false=drop) + r_source = r_source.label_map( + all_terms, + map_if_true=lambda t: Constant(self.butcher_exp[stage, i]) * self.dt * t + ) + residual += r_source + # Calculate and add on dt*a_ss*F(y_s) r_imp = self.residual.label_map( lambda t: t.has_label(implicit), @@ -210,6 +216,15 @@ def final_res(self): map_if_false=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) residual += r_imp residual += r_exp + # Calculate source terms + r_source = self.residual.label_map( + lambda t: t.has_label(source_label), + map_if_true=replace_subject(self.source[i], old_idx=self.idx), + map_if_false=drop) + r_source = r_source.label_map( + all_terms, + map_if_true=lambda t: Constant(self.butcher_exp[self.nStages, i])*self.dt*t) + residual += r_source return residual.form @cached_property @@ -226,7 +241,7 @@ def solvers(self): @cached_property def final_solver(self): """Set up a solver for the final solve to evaluate time level n+1.""" - # setup solver using lhs and rhs defined in derived class + # setup solver using residual (res) defined in derived class problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.linear_solver_parameters, options_prefix=solver_name) @@ -242,12 +257,19 @@ def apply(self, x_out, x_in): # Set initial solver guess if (stage > 0): self.x_out.assign(self.xs[stage-1]) + # Evaluate source terms + for evaluate in self.evaluate_source: + evaluate(self.xs[stage-1], self.dt, x_out=self.source[stage-1]) self.solver.solve() # Apply limiter if self.limiter is not None: self.limiter.apply(self.x_out) self.xs[stage].assign(self.x_out) + + # Solve final stage + for evaluate in self.evaluate_source: + evaluate(self.xs[-1], self.dt, x_out=self.source[-1]) self.final_solver.solve() # Apply limiter @@ -269,7 +291,7 @@ class IMEX_Euler(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -286,13 +308,16 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_ARS3(IMEXRungeKutta): @@ -313,7 +338,7 @@ class IMEX_ARS3(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -330,6 +355,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ g = (3. + np.sqrt(3.))/6. butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) @@ -338,7 +366,7 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_ARK2(IMEXRungeKutta): @@ -359,7 +387,7 @@ class IMEX_ARK2(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -376,6 +404,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ g = 1. - 1./np.sqrt(2.) d = 1./(2.*np.sqrt(2.)) @@ -385,7 +416,7 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_SSP3(IMEXRungeKutta): @@ -404,7 +435,7 @@ class IMEX_SSP3(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -421,6 +452,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ g = 1. - (1./np.sqrt(2.)) butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]]) @@ -428,7 +462,7 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_Trap2(IMEXRungeKutta): @@ -447,7 +481,7 @@ class IMEX_Trap2(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -464,6 +498,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ e = 0. butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) @@ -471,4 +508,4 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 8eb1c75ab..b31b6e6a6 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -3,14 +3,15 @@ import numpy as np from firedrake import (Function, split, NonlinearVariationalProblem, - NonlinearVariationalSolver) + NonlinearVariationalSolver, Constant) from firedrake.fml import replace_subject, all_terms, drop from firedrake.utils import cached_property -from gusto.core.labels import time_derivative +from gusto.core.labels import time_derivative, source_label from gusto.time_discretisation.time_discretisation import ( TimeDiscretisation, wrapper_apply ) +from gusto.time_discretisation.explicit_runge_kutta import RungeKuttaFormulation __all__ = ["ImplicitRungeKutta", "ImplicitMidpoint", "QinZhang"] @@ -30,7 +31,9 @@ class ImplicitRungeKutta(TimeDiscretisation): For each i = 1, s in an s stage method we have the intermediate solutions: \n y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i) \n - We compute the gradient at the intermediate location, k_i = F(y_i) \n + For the increment form we compute the gradient at the \n + intermediate location, k_i = F(y_i), whilst for the \n + predictor form we solve for each intermediate solution y_i. \n At the last stage, compute the new solution by: \n y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) @@ -56,7 +59,8 @@ class ImplicitRungeKutta(TimeDiscretisation): # --------------------------------------------------------------------------- def __init__(self, domain, butcher_matrix, field_name=None, - solver_parameters=None, options=None,): + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -66,18 +70,25 @@ def __init__(self, domain, butcher_matrix, field_name=None, discretisation. field_name (str, optional): name of the field to be evolved. Defaults to None. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. options (:class:`AdvectionOptions`, optional): an object containing options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, field_name=field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) self.butcher_matrix = butcher_matrix self.nStages = int(np.shape(self.butcher_matrix)[1]) + self.rk_formulation = rk_formulation def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -93,29 +104,110 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.k = [Function(self.fs) for i in range(self.nStages)] - def lhs(self): - return super().lhs - - def rhs(self): - return super().rhs + # Check that we do not have source terms + for t in self.residual: + if (t.has_label(source_label)): + raise NotImplementedError("Source terms have not been implemented with implicit RK schemes") + + if self.rk_formulation == RungeKuttaFormulation.predictor: + self.xs = [Function(self.fs) for _ in range(self.nStages)] + elif self.rk_formulation == RungeKuttaFormulation.increment: + self.k = [Function(self.fs) for _ in range(self.nStages)] + elif self.rk_formulation == RungeKuttaFormulation.linear: + raise NotImplementedError( + 'Linear Implicit Runge-Kutta formulation is not implemented' + ) + else: + raise NotImplementedError( + 'Runge-Kutta formulation is not implemented' + ) - def solver(self, stage): - residual = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.xnph, self.idx), - ) + def res(self, stage): + """Set up the residual for the predictor formulation for a given stage.""" + # Add time derivative terms y_s - y^n for stage s mass_form = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_false=drop) - residual += mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) + residual = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + residual -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + # Loop through stages up to s-1 and calculate sum + # dt*(a_s1*F(y_1) + a_s2*F(y_2)+ ... + a_{s,s-1}*F(y_{s-1})) + for i in range(stage): + r_imp = self.residual.label_map( + lambda t: not t.has_label(time_derivative), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + all_terms, + map_if_true=lambda t: Constant(self.butcher_matrix[stage, i])*self.dt*t) + residual += r_imp + # Calculate and add on dt*a_ss*F(y_s) + r_imp = self.residual.label_map( + lambda t: not t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + all_terms, + map_if_true=lambda t: Constant(self.butcher_matrix[stage, stage])*self.dt*t) + residual += r_imp + return residual.form + + @property + def final_res(self): + """Set up the final residual for the predictor formulation.""" + # Add time derivative terms y^{n+1} - y^n + mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + residual -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + # Loop through stages up to s-1 and calculate/sum + # dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) + for i in range(self.nStages): + r_imp = self.residual.label_map( + lambda t: not t.has_label(time_derivative), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + all_terms, + map_if_true=lambda t: Constant(self.butcher_matrix[self.nStages, i])*self.dt*t) + residual += r_imp + return residual.form - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + def solver(self, stage): + if self.rk_formulation == RungeKuttaFormulation.increment: + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.xnph, self.idx), + ) + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual += mass_form.label_map(all_terms, + replace_subject(self.x_out, self.idx)) + + problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + + elif self.rk_formulation == RungeKuttaFormulation.predictor: + problem = NonlinearVariationalProblem(self.res(stage), self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) + + @cached_property + def final_solver(self): + """ + Set up a solver for the final solve for the predictor + formulation to evaluate time level n+1. + """ + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @cached_property def solvers(self): @@ -126,32 +218,48 @@ def solvers(self): def solve_stage(self, x0, stage): self.x1.assign(x0) - for i in range(stage): - self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) + if self.rk_formulation == RungeKuttaFormulation.increment: + for i in range(stage): + self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) - if self.idx is None and len(self.fs) > 1: - self.xnph = tuple([self.dt*self.butcher_matrix[stage, stage]*a + b - for a, b in zip(split(self.x_out), split(self.x1))]) - else: - self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out - solver = self.solvers[stage] - # Set initial guess for solver - if (stage > 0): - self.x_out.assign(self.k[stage-1]) + if self.idx is None and len(self.fs) > 1: + self.xnph = tuple( + self.dt * self.butcher_matrix[stage, stage] * a + b + for a, b in zip(split(self.x_out), split(self.x1)) + ) + else: + self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out + + solver = self.solvers[stage] - solver.solve() + # Set initial guess for solver + if (stage > 0): + self.x_out.assign(self.k[stage-1]) - self.k[stage].assign(self.x_out) + solver.solve() + self.k[stage].assign(self.x_out) + + elif self.rk_formulation == RungeKuttaFormulation.predictor: + if (stage > 0): + self.x_out.assign(self.xs[stage-1]) + solver = self.solvers[stage] + solver.solve() + + self.xs[stage].assign(self.x_out) @wrapper_apply def apply(self, x_out, x_in): - + self.x_out.assign(x_in) for i in range(self.nStages): self.solve_stage(x_in, i) - x_out.assign(x_in) - for i in range(self.nStages): - x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) + if self.rk_formulation == RungeKuttaFormulation.increment: + x_out.assign(x_in) + for i in range(self.nStages): + x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) + elif self.rk_formulation == RungeKuttaFormulation.predictor: + self.final_solver.solve() + x_out.assign(self.x_out) class ImplicitMidpoint(ImplicitRungeKutta): @@ -164,25 +272,33 @@ class ImplicitMidpoint(ImplicitRungeKutta): k0 = F[y^n + 0.5*dt*k0] \n y^(n+1) = y^n + dt*k0 \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, - options=None): + def __init__(self, domain, field_name=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. options (:class:`AdvectionOptions`, optional): an object containing options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ butcher_matrix = np.array([[0.5], [1.]]) super().__init__(domain, butcher_matrix, field_name, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) class QinZhang(ImplicitRungeKutta): @@ -196,22 +312,30 @@ class QinZhang(ImplicitRungeKutta): k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] \n y^(n+1) = y^n + 0.5*dt*(k0 + k1) \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, - options=None): + def __init__(self, domain, field_name=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. options (:class:`AdvectionOptions`, optional): an object containing options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]]) super().__init__(domain, butcher_matrix, field_name, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) diff --git a/gusto/time_discretisation/multi_level_schemes.py b/gusto/time_discretisation/multi_level_schemes.py index a11671cc4..ff8b212fd 100644 --- a/gusto/time_discretisation/multi_level_schemes.py +++ b/gusto/time_discretisation/multi_level_schemes.py @@ -6,7 +6,7 @@ NonlinearVariationalSolver) from firedrake.fml import replace_subject, all_terms, drop from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions -from gusto.core.labels import time_derivative +from gusto.core.labels import time_derivative, source_label from gusto.time_discretisation.time_discretisation import TimeDiscretisation @@ -17,7 +17,7 @@ class MultilevelTimeDiscretisation(TimeDiscretisation): """Base class for multi-level timesteppers""" def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -32,12 +32,16 @@ def __init__(self, domain, field_name=None, solver_parameters=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") super().__init__(domain=domain, field_name=field_name, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) self.initial_timesteps = 0 @abstractproperty @@ -49,6 +53,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): for n in range(self.nlevels, 1, -1): setattr(self, "xnm%i" % (n-1), Function(self.fs)) + # Check that we do not have source terms + for t in self.residual: + if (t.has_label(source_label)): + raise NotImplementedError("Source terms have not been implemented with multilevel schemes") + class BDF2(MultilevelTimeDiscretisation): """ @@ -63,58 +72,57 @@ def nlevels(self): return 2 @property - def lhs0(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res0(self): + """Set up the discretisation's residual for initial BDF step.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt * t + ) - @property - def rhs0(self): - """Set up the time discretisation's right hand side for inital BDF step.""" r = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) + residual -= r - return r.form + return residual.form @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: (2/3)*self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: (2 / 3) * self.dt * t + ) - @property - def rhs(self): - """Set up the time discretisation's right hand side for BDF2 steps.""" xn = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) xnm1 = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.xnm1, old_idx=self.idx), - map_if_false=drop) - - r = (4/3.) * xn - (1/3.) * xnm1 + map_if_false=drop + ) - return r.form + residual -= (4 / 3) * xn - (1 / 3) * xnm1 + return residual.form @property def solver0(self): """Set up the problem and the solver for initial BDF step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -122,8 +130,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for BDF2 steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -162,38 +170,53 @@ def nlevels(self): return 2 @property - def rhs0(self): - """Set up the discretisation's right hand side for initial forward euler step.""" + def res0(self): + """Set up the discretisation's residual for initial forward euler step.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) - - return r.form + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * t + ) - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(Leapfrog, self).lhs + residual -= r + return residual.form @property def rhs(self): - """Set up the discretisation's right hand side for leapfrog steps.""" + """Set up the discretisation's residual for leapfrog steps.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + r = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_false=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.xnm1, old_idx=self.idx), - map_if_false=lambda t: -2.0*self.dt*t) + map_if_false=replace_subject(self.x1, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.xnm1, old_idx=self.idx), + map_if_false=lambda t: -2.0 * self.dt * t + ) - return r.form + residual -= r + return residual.form @property def solver0(self): """Set up the problem and the solver for initial forward euler step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -201,8 +224,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for leapfrog steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -223,6 +246,7 @@ def apply(self, x_out, *x_in): self.xnm1.assign(x_in[0]) self.x1.assign(x_in[1]) + # Set initial solver guess self.x_out.assign(x_in[1]) solver.solve() @@ -238,7 +262,7 @@ class AdamsBashforth(MultilevelTimeDiscretisation): y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) """ def __init__(self, domain, order, field_name=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -252,6 +276,9 @@ def __init__(self, domain, order, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. Raises: ValueError: if order is not provided, or is in incorrect range. @@ -264,7 +291,7 @@ def __init__(self, domain, order, field_name=None, super().__init__(domain, field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) self.order = order @@ -291,42 +318,59 @@ def nlevels(self): return self.order @property - def rhs0(self): - """Set up the discretisation's right hand side for initial forward euler step.""" + def res0(self): + """Set up the discretisation's residual for initial forward euler step.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) - - return r.form + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * t + ) - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(AdamsBashforth, self).lhs + residual -= r + return residual.form @property - def rhs(self): - """Set up the discretisation's right hand side for Adams Bashforth steps.""" - r = self.residual.label_map(all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.b[-1]*self.dt*t) - for n in range(self.nlevels-1): - rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.x[n], old_idx=self.idx)) - rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*self.b[n]*t) - r += rtemp - return r.form + def res(self): + """Set up the discretisation's residual for Adams Bashforth steps.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + residual -= r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.b[-1] * self.dt * t + ) + for n in range(self.nlevels - 1): + rtemp = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.x[n], old_idx=self.idx) + ) + rtemp = rtemp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * self.b[n] * t + ) + residual -= rtemp + return residual.form @property def solver0(self): """Set up the problem and the solverfor initial forward euler step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -334,8 +378,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for Adams Bashforth steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -371,7 +415,8 @@ class AdamsMoulton(MultilevelTimeDiscretisation): y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) """ def __init__(self, domain, order, field_name=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -385,6 +430,9 @@ def __init__(self, domain, order, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. Raises: ValueError: if order is not provided, or is in incorrect range. @@ -400,7 +448,7 @@ def __init__(self, domain, order, field_name=None, super().__init__(domain, field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) self.order = order @@ -427,63 +475,67 @@ def nlevels(self): return self.order @property - def rhs0(self): + def res0(self): """ - Set up the discretisation's right hand side for initial trapezoidal - step. + Set up the discretisation's residual for initial trapezoidal step. """ + residual = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: 0.5 * self.dt * t + ) r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -0.5*self.dt*t) + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -0.5 * self.dt * t + ) - return r.form + residual -= r + return residual.form @property - def lhs0(self): - """ - Set up the time discretisation's right hand side for initial - trapezoidal step. - """ - l = self.residual.label_map( + def res(self): + """Set up the time discretisation's residual for Adams Moulton steps.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: 0.5*self.dt*t) - return l.form - - @property - def lhs(self): - """Set up the time discretisation's right hand side for Adams Moulton steps.""" - l = self.residual.label_map( + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.bl * self.dt * t + ) + r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.bl*self.dt*t) - return l.form - - @property - def rhs(self): - """Set up the discretisation's right hand side for Adams Moulton steps.""" - r = self.residual.label_map(all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.br[-1]*self.dt*t) - for n in range(self.nlevels-1): - rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.x[n], old_idx=self.idx)) - rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*self.br[n]*t) - r += rtemp - return r.form + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + residual -= r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.br[-1] * self.dt * t + ) + for n in range(self.nlevels - 1): + rtemp = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.x[n], old_idx=self.idx) + ) + rtemp = rtemp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * self.br[n] * t + ) + residual -= rtemp + return residual.form @property def solver0(self): """Set up the problem and the solver for initial trapezoidal step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -491,8 +543,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for Adams Moulton steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) diff --git a/gusto/time_discretisation/sdc.py b/gusto/time_discretisation/sdc.py index 0fe0c9f29..3e634f17a 100644 --- a/gusto/time_discretisation/sdc.py +++ b/gusto/time_discretisation/sdc.py @@ -46,16 +46,14 @@ from abc import ABCMeta import numpy as np from firedrake import ( - Function, NonlinearVariationalProblem, TestFunction, TestFunctions, - NonlinearVariationalSolver, Constant + Function, NonlinearVariationalProblem, NonlinearVariationalSolver, Constant ) from firedrake.fml import ( - replace_subject, replace_test_function, all_terms, drop + replace_subject, all_terms, drop ) from firedrake.utils import cached_property -from gusto.time_discretisation.wrappers import * from gusto.time_discretisation.time_discretisation import wrapper_apply -from gusto.core.labels import (time_derivative, implicit, explicit) +from gusto.core.labels import (time_derivative, implicit, explicit, source_label) from qmat import genQCoeffs, genQDeltaCoeffs @@ -118,36 +116,8 @@ def __init__(self, base_scheme, domain, M, maxk, quad_type, node_type, qdelta_im self.final_update = final_update self.formulation = formulation self.limiter = limiter - - # Initialise wrappers - if options is not None: - self.wrapper_name = options.name - if self.wrapper_name == "mixed_options": - self.wrapper = MixedFSWrapper() - - for field, suboption in options.suboptions.items(): - if suboption.name == 'embedded_dg': - self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)}) - elif suboption.name == "recovered": - self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)}) - elif suboption.name == "supg": - raise RuntimeError( - 'SDC: suboption SUPG is currently not implemented within MixedOptions') - else: - raise RuntimeError( - f'SDC: suboption wrapper {wrapper_name} not implemented') - elif self.wrapper_name == "embedded_dg": - self.wrapper = EmbeddedDGWrapper(self, options) - elif self.wrapper_name == "recovered": - self.wrapper = RecoveryWrapper(self, options) - elif self.wrapper_name == "supg": - self.wrapper = SUPGWrapper(self, options) - else: - raise RuntimeError( - f'SDC: wrapper {self.wrapper_name} not implemented') - else: - self.wrapper = None - self.wrapper_name = None + self.augmentation = self.base.augmentation + self.wrapper = self.base.wrapper # Get quadrature nodes and weights self.nodes, self.weights, self.Q = genQCoeffs("Collocation", nNodes=M, @@ -160,6 +130,10 @@ def __init__(self, base_scheme, domain, M, maxk, quad_type, node_type, qdelta_im self.dtau = np.diff(np.append(0, self.nodes)) self.Q = float(self.dt_coarse)*self.Q self.Qfin = float(self.dt_coarse)*self.weights + self.qdelta_imp_type = qdelta_imp + self.formulation = formulation + self.node_type = node_type + self.quad_type = quad_type # Get Q_delta matrices self.Qdelta_imp = genQDeltaCoeffs(qdelta_imp, form=formulation, @@ -206,68 +180,13 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.base.setup(equation, apply_bcs, *active_labels) self.equation = self.base.equation self.residual = self.base.residual + self.evaluate_source = self.base.evaluate_source for t in self.residual: # Check all terms are labeled implicit or explicit if ((not t.has_label(implicit)) and (not t.has_label(explicit)) - and (not t.has_label(time_derivative))): - raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit") - - # Check we are not using wrappers for implicit schemes - if (t.has_label(implicit) and self.wrapper is not None): - raise NotImplementedError("Implicit terms not supported with wrappers") - - # Check we are not using limiters for implicit schemes - if (t.has_label(implicit) and self.limiter is not None): - raise NotImplementedError("Implicit terms not supported with limiters") - - # Set up Wrappers - if self.wrapper is not None: - if self.wrapper_name == "mixed_options": - - self.wrapper.wrapper_spaces = equation.spaces - self.wrapper.field_names = equation.field_names - - for field, subwrapper in self.wrapper.subwrappers.items(): - - if field not in equation.field_names: - raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") - - field_idx = equation.field_names.index(field) - subwrapper.setup(equation.spaces[Wrappersfield_idx]) - - # Update the function space to that needed by the wrapper - self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space - - self.wrapper.setup() - self.fs = self.wrapper.function_space - new_test_mixed = TestFunctions(self.fs) - - # Replace the original test function with one from the new - # function space defined by the subwrappers - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test_mixed)) - - else: - if self.wrapper_name == "supg": - self.wrapper.setup() - else: - self.wrapper.setup(self.fs) - self.fs = self.wrapper.function_space - if self.solver_parameters is None: - self.solver_parameters = self.wrapper.solver_parameters - new_test = TestFunction(self.wrapper.test_space) - # SUPG has a special wrapper - if self.wrapper_name == "supg": - new_test = self.wrapper.test - - # Replace the original test function with the one from the wrapper - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test)) - - self.residual = self.wrapper.label_terms(self.residual) + and (not t.has_label(time_derivative)) and (not t.has_label(source_label))): + raise NotImplementedError("Non time-derivative or source terms must be labeled as implicit or explicit") # Set up bcs self.bcs = self.base.bcs @@ -285,6 +204,8 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.Unodes1 = [Function(W) for _ in range(self.M+1)] self.fUnodes = [Function(W) for _ in range(self.M)] self.quad = [Function(W) for _ in range(self.M)] + self.source_Uk = [Function(W) for _ in range(self.M+1)] + self.source_Ukp1 = [Function(W) for _ in range(self.M+1)] self.U_SDC = Function(W) self.U_start = Function(W) self.Un = Function(W) @@ -293,6 +214,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.U_fin = Function(W) self.Urhs = Function(W) self.Uin = Function(W) + self.source_in = Function(W) @property def nlevels(self): @@ -322,10 +244,13 @@ def res_rhs(self): replace_subject(self.Urhs, old_idx=self.idx), drop) # F(y) - L = self.residual.label_map(lambda t: t.has_label(time_derivative), + L = self.residual.label_map(lambda t: any(t.has_label(time_derivative, source_label)), drop, replace_subject(self.Uin, old_idx=self.idx)) - residual_rhs = a - L + L_source = self.residual.label_map(lambda t: t.has_label(source_label), + replace_subject(self.source_in, old_idx=self.idx), + drop) + residual_rhs = a - (L + L_source) return residual_rhs.form @property @@ -401,6 +326,25 @@ def res(self, m): lambda t: Constant(self.Qdelta_exp[m, i])*t) residual -= r_exp_k + # Calculate source terms + r_source_kp1 = self.residual.label_map( + lambda t: t.has_label(source_label), + map_if_true=replace_subject(self.source_Ukp1[i+1], old_idx=self.idx), + map_if_false=drop) + r_source_kp1 = r_source_kp1.label_map( + all_terms, + lambda t: Constant(self.Qdelta_exp[m, i])*t) + residual += r_source_kp1 + + r_source_k = self.residual.label_map( + lambda t: t.has_label(source_label), + map_if_true=replace_subject(self.source_Ukp1[i+1], old_idx=self.idx), + map_if_false=drop) + r_source_k = r_source_k.label_map( + all_terms, + map_if_true=lambda t: Constant(self.Qdelta_exp[m, i])*t) + residual -= r_source_k + # Add on final implicit terms # Qdelta_imp[m,m]*(F(y_(m)^(k+1)) - F(y_(m)^k)) r_imp_kp1 = self.residual.label_map( @@ -474,22 +418,43 @@ def apply(self, x_out, x_in): else: for m in range(self.M): self.Unodes[m+1].assign(self.Un) + for m in range(self.M+1): + for evaluate in self.evaluate_source: + evaluate(self.Unodes[m], self.base.dt, x_out=self.source_Uk[m]) # Iterate through correction sweeps k = 0 while k < self.maxk: k += 1 + if self.qdelta_imp_type == "MIN-SR-FLEX": + # Recompute Implicit Q_delta matrix for each iteration k + self.Qdelta_imp = genQDeltaCoeffs( + self.qdelta_imp_type, + form=self.formulation, + nodes=self.nodes, + Q=self.Q, + nNodes=self.M, + nodeType=self.node_type, + quadType=self.quad_type, + k=k + ) + # Compute for N2N: sum(j=1,M) (s_mj*F(y_m^k) + s_mj*S(y_m^k)) # for Z2N: sum(j=1,M) (q_mj*F(y_m^k) + q_mj*S(y_m^k)) for m in range(1, self.M+1): self.Uin.assign(self.Unodes[m]) + # Include source terms + for evaluate in self.evaluate_source: + evaluate(self.Uin, self.base.dt, x_out=self.source_in) self.solver_rhs.solve() self.fUnodes[m-1].assign(self.Urhs) self.compute_quad() # Loop through quadrature nodes and solve self.Unodes1[0].assign(self.Unodes[0]) + for evaluate in self.evaluate_source: + evaluate(self.Unodes[0], self.base.dt, x_out=self.source_Uk[0]) for m in range(1, self.M+1): # Set Q or S matrix self.Q_.assign(self.quad[m-1]) @@ -511,21 +476,28 @@ def apply(self, x_out, x_in): self.solver.solve() self.Unodes1[m].assign(self.U_SDC) + # Evaluate source terms + for evaluate in self.evaluate_source: + evaluate(self.Unodes1[m], self.base.dt, x_out=self.source_Ukp1[m]) + # Apply limiter if required if self.limiter is not None: self.limiter.apply(self.Unodes1[m]) for m in range(1, self.M+1): self.Unodes[m].assign(self.Unodes1[m]) + self.source_Uk[m].assign(self.source_Ukp1[m]) if self.maxk > 0: # Compute value at dt rather than final quadrature node tau_M if self.final_update: for m in range(1, self.M+1): self.Uin.assign(self.Unodes1[m]) + self.source_in.assign(self.source_Ukp1[m]) self.solver_rhs.solve() self.fUnodes[m-1].assign(self.Urhs) self.compute_quad_final() # Compute y_(n+1) = y_n + sum(j=1,M) q_j*F(y_j) + self.U_fin.assign(self.Unodes[-1]) self.solver_fin.solve() # Apply limiter if required diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index d7b4d4aab..db71412ca 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -5,7 +5,7 @@ operator F. """ -from abc import ABCMeta, abstractmethod, abstractproperty +from abc import ABCMeta, abstractmethod import math from firedrake import (Function, TestFunction, TestFunctions, DirichletBC, @@ -21,6 +21,7 @@ mass_weighted, nonlinear_time_derivative) from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.time_discretisation.wrappers import * +from gusto.solvers import mass_parameters __all__ = ["TimeDiscretisation", "ExplicitTimeDiscretisation", "BackwardEuler", "ThetaMethod", "TrapeziumRule", "TR_BDF2"] @@ -30,7 +31,17 @@ def wrapper_apply(original_apply): """Decorator to add steps for using a wrapper around the apply method.""" def get_apply(self, x_out, x_in): - if self.wrapper is not None: + if self.augmentation is not None: + + def new_apply(self, x_out, x_in): + + self.augmentation.pre_apply(x_in) + original_apply(self, self.augmentation.x_out, self.augmentation.x_in) + self.augmentation.post_apply(x_out) + + return new_apply(self, x_out, x_in) + + elif self.wrapper is not None: def new_apply(self, x_out, x_in): @@ -51,13 +62,17 @@ class TimeDiscretisation(object, metaclass=ABCMeta): """Base class for time discretisation schemes.""" def __init__(self, domain, field_name=None, subcycling_options=None, - solver_parameters=None, limiter=None, options=None): + solver_parameters=None, limiter=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. + subcycling_options(:class:`SubcyclingOptions`, optional): an object + containing options for subcycling the time discretisation. + Defaults to None. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -66,6 +81,9 @@ def __init__(self, domain, field_name=None, subcycling_options=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ self.domain = domain self.field_name = field_name @@ -78,7 +96,8 @@ def __init__(self, domain, field_name=None, subcycling_options=None, self.options = options self.limiter = limiter self.courant_max = None - self.subcycling_options = None + self.augmentation = augmentation + self.subcycling_options = subcycling_options if self.subcycling_options is not None: self.subcycling_options.check_options() @@ -159,6 +178,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.fs = equation.function_space self.idx = None + if self.augmentation is not None: + self.fs = self.augmentation.fs + self.residual = self.augmentation.residual + self.idx = None + if len(active_labels) > 0: if isinstance(self.field_name, list): # Multiple fields are being solved for simultaneously. @@ -179,6 +203,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): map_if_false=drop) self.residual = residual + else: self.residual = self.residual.label_map( lambda t: any(t.has_label(time_derivative, *active_labels)), @@ -188,7 +213,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): if isinstance(self.field_name, list): self.field_name = equation.field_name - bcs = equation.bcs[self.field_name] + if self.augmentation is not None: + # Transfer BCs from appropriate function space + bcs = self.augmentation.bcs if hasattr(self.augmentation, "bcs") else None + else: + bcs = equation.bcs[self.field_name] self.evaluate_source = [] self.physics_names = [] @@ -204,8 +233,16 @@ def setup(self, equation, apply_bcs=True, *active_labels): for field in equation.field_names: # Check if the mass term for this prognostic is mass-weighted - if len(self.residual.label_map(lambda t: t.get(prognostic) == field and t.has_label(time_derivative) and t.has_label(mass_weighted), map_if_false=drop)) == 1: - field_terms = self.residual.label_map(lambda t: t.get(prognostic) == field and not t.has_label(time_derivative), map_if_false=drop) + if len(self.residual.label_map(( + lambda t: t.get(prognostic) == field + and t.has_label(time_derivative) + and t.has_label(mass_weighted) + ), map_if_false=drop)) == 1: + + field_terms = self.residual.label_map( + lambda t: t.get(prognostic) == field and not t.has_label(time_derivative), + map_if_false=drop + ) # Check that the equation for this prognostic does not involve # both mass-weighted and non-mass-weighted terms; if so, a split @@ -328,34 +365,33 @@ def setup(self, equation, apply_bcs=True, *active_labels): def nlevels(self): return 1 - @abstractproperty - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + @property + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) - - return l.form - - @abstractproperty - def rhs(self): - """Set up the time discretisation's right hand side.""" + map_if_false=drop + ) r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) r = r.label_map( lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) + map_if_false=lambda t: -self.dt * t + ) + + residual -= r - return r.form + return residual.form @cached_property def solver(self): """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ solver = NonlinearVariationalSolver( problem, @@ -382,11 +418,12 @@ def update_subcycling(self): # Cap number of subcycles if max_subcycles is not None: - self.ncycles = min(self.ncycles, max_subcycles) - logger.warning( - 'Adaptive subcycling: capping number of subcycles at ' - f'{max_subcycles}' - ) + if self.ncycles > max_subcycles: + logger.warning( + 'Adaptive subcycling: capping number of subcycles at ' + f'{max_subcycles}' + ) + self.ncycles = max_subcycles logger.debug(f'Performing {self.ncycles} subcycles') self.dt.assign(self.original_dt/self.ncycles) @@ -407,7 +444,8 @@ class ExplicitTimeDiscretisation(TimeDiscretisation): """Base class for explicit time discretisations.""" def __init__(self, domain, field_name=None, subcycling_options=None, - solver_parameters=None, limiter=None, options=None): + solver_parameters=None, limiter=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -425,20 +463,15 @@ def __init__(self, domain, field_name=None, subcycling_options=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, field_name, subcycling_options=subcycling_options, solver_parameters=solver_parameters, - limiter=limiter, options=options) - - # get default solver options if none passed in - if solver_parameters is None: - self.solver_parameters = {'snes_type': 'ksponly', - 'ksp_type': 'cg', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - else: - self.solver_parameters = solver_parameters + limiter=limiter, options=options, + augmentation=augmentation) def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -453,20 +486,25 @@ def setup(self, equation, apply_bcs=True, *active_labels): """ super().setup(equation, apply_bcs, *active_labels) + # get default solver options if none passed in + self.solver_parameters.update(mass_parameters( + self.fs, equation.domain.spaces)) + self.solver_parameters['snes_type'] = 'ksponly' + # if user has specified a number of fixed subcycles, then save this # and rescale dt accordingly; else perform just one cycle using dt if (self.subcycling_options is not None and self.subcycling_options.fixed_subcycles is not None): - fixed_subcycles = self.subcycling_options.fixed_subcycles - self.dt.assign(self.dt/fixed_subcycles) - self.ncycles = self.fixed_subcycles + self.ncycles = self.subcycling_options.fixed_subcycles + self.dt.assign(self.dt/self.ncycles) else: - self.dt = self.dt self.ncycles = 1 + self.dt = self.dt self.x0 = Function(self.fs) self.x1 = Function(self.fs) - # If the time_derivative term is nonlinear, we must use a nonlinear solver + # If the time_derivative term is nonlinear, we must use a nonlinear solver, + # but if the time_derivative term is linear, we can reuse the factorisations. if ( len(self.residual.label_map( lambda t: t.has_label(nonlinear_time_derivative), @@ -478,22 +516,40 @@ def setup(self, equation, apply_bcs=True, *active_labels): + ' as the time derivative term is nonlinear') logger.warning(message) self.solver_parameters['snes_type'] = 'newtonls' + else: + self.solver_parameters.setdefault('snes_lag_jacobian', -2) + self.solver_parameters.setdefault('snes_lag_jacobian_persists', None) + self.solver_parameters.setdefault('snes_lag_preconditioner', -2) + self.solver_parameters.setdefault('snes_lag_preconditioner_persists', None) @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual""" + residual = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, self.idx), - map_if_false=drop) + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) - return l.form + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) + + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * t + ) + + residual -= r + + return residual.form @cached_property def solver(self): """Set up the problem and the solver.""" - # setup linear solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs - self.rhs, self.x_out, bcs=self.bcs) + # setup linear solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -536,7 +592,8 @@ class BackwardEuler(TimeDiscretisation): y^(n+1) = y^n + dt*F[y^(n+1)]. \n """ def __init__(self, domain, field_name=None, subcycling_options=None, - solver_parameters=None, limiter=None, options=None): + solver_parameters=None, limiter=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -553,7 +610,10 @@ def __init__(self, domain, field_name=None, subcycling_options=None, options (:class:`AdvectionOptions`, optional): an object containing options to either be passed to the spatial discretisation, or to control the "wrapper" methods. Defaults to None. - """ + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. + """ if not solver_parameters: # default solver parameters solver_parameters = {'ksp_type': 'gmres', @@ -562,7 +622,8 @@ def __init__(self, domain, field_name=None, subcycling_options=None, super().__init__(domain=domain, field_name=field_name, subcycling_options=subcycling_options, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -589,25 +650,27 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.x1 = Function(self.fs) @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - return l.form - - @property - def rhs(self): - """Set up the time discretisation's right hand side.""" + map_if_true=replace_subject( + self.x_out, old_idx=self.idx + ) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt*t + ) r = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) - return r.form + residual -= r + + return residual.form def apply_cycle(self, x_out, x_in): """ @@ -658,7 +721,7 @@ class ThetaMethod(TimeDiscretisation): """ def __init__(self, domain, theta, field_name=None, subcycling_options=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -676,6 +739,9 @@ def __init__(self, domain, theta, field_name=None, subcycling_options=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. Raises: ValueError: if theta is not provided. @@ -694,7 +760,8 @@ def __init__(self, domain, theta, field_name=None, subcycling_options=None, super().__init__(domain, field_name, subcycling_options=subcycling_options, solver_parameters=solver_parameters, - options=options) + options=options, + augmentation=augmentation) self.theta = theta @@ -723,26 +790,27 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.x1 = Function(self.fs) @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.theta*self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.theta * self.dt * t + ) - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side.""" r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -(1-self.theta)*self.dt*t) - - return r.form + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -(1 - self.theta) * self.dt * t + ) + residual -= r + return residual.form def apply_cycle(self, x_out, x_in): """ @@ -771,6 +839,8 @@ def apply(self, x_out, x_in): x_in (:class:`Function`): the input field. """ self.update_subcycling() + if self.augmentation is not None: + self.augmentation.update(x_in) self.x0.assign(x_in) for i in range(self.ncycles): @@ -791,7 +861,7 @@ class TrapeziumRule(ThetaMethod): """ def __init__(self, domain, field_name=None, subcycling_options=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -804,11 +874,14 @@ def __init__(self, domain, field_name=None, subcycling_options=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, 0.5, field_name, subcycling_options=subcycling_options, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) class TR_BDF2(TimeDiscretisation): @@ -861,60 +934,63 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.xn = Function(self.fs) @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative) for the TR stage.""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual for the TR stage.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.xnpg, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: 0.5*self.gamma*self.dt*t) - - return l.form + map_if_true=replace_subject(self.xnpg, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: 0.5 * self.gamma * self.dt * t + ) - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side for the TR stage.""" r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.xn, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -0.5*self.gamma*self.dt*t) + map_if_true=replace_subject(self.xn, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -0.5 * self.gamma * self.dt * t + ) + residual -= r - return r.form + return residual.form @cached_property - def lhs_bdf2(self): - """Set up the discretisation's left hand side (the time derivative) for - the BDF2 stage.""" - l = self.residual.label_map( + def res_bdf2(self): + """Set up the discretisation's residual for the BDF2 stage.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: ((1.0-self.gamma)/(2.0-self.gamma))*self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: (1.0 - self.gamma) / (2.0 - self.gamma) * self.dt * t + ) - @cached_property - def rhs_bdf2(self): - """Set up the time discretisation's right hand side for the BDF2 stage.""" xn = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.xn, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) xnpg = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.xnpg, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) - r = (1.0/(self.gamma*(2.0-self.gamma)))*xnpg - ((1.0-self.gamma)**2/(self.gamma*(2.0-self.gamma)))*xn + r = (1.0 / (self.gamma * (2.0 - self.gamma))) * xnpg - \ + ((1.0 - self.gamma) ** 2 / (self.gamma * (2.0 - self.gamma))) * xn - return r.form + residual -= r + return residual.form @cached_property def solver_tr(self): """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.xnpg, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.xnpg, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"_tr" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -922,8 +998,8 @@ def solver_tr(self): @cached_property def solver_bdf2(self): """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs_bdf2-self.rhs_bdf2, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res_bdf2, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"_bdf2" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) diff --git a/gusto/time_discretisation/wrappers.py b/gusto/time_discretisation/wrappers.py index 2d388e998..18381bc1f 100644 --- a/gusto/time_discretisation/wrappers.py +++ b/gusto/time_discretisation/wrappers.py @@ -6,14 +6,16 @@ from abc import ABCMeta, abstractmethod from firedrake import ( - FunctionSpace, Function, BrokenElement, Projector, Interpolator, - VectorElement, Constant, as_ufl, dot, grad, TestFunction, MixedFunctionSpace + FunctionSpace, Function, BrokenElement, Projector, Constant, as_ufl, dot, grad, + TestFunction, MixedFunctionSpace, assemble ) +from firedrake.__future__ import interpolate from firedrake.fml import Term from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions from gusto.recovery import Recoverer, ReversibleRecoverer, ConservativeRecoverer from gusto.core.labels import transporting_velocity from gusto.core.conservative_projection import ConservativeProjector +from gusto.core.function_spaces import is_cg import ufl __all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedFSWrapper"] @@ -263,7 +265,7 @@ def setup(self, original_space, post_apply_bcs): # Operators for projecting back self.interp_back = (self.options.project_low_method == 'interpolate') if self.options.project_low_method == 'interpolate': - self.x_out_projector = Interpolator(self.x_out, self.x_projected) + self.x_out_projector = interpolate(self.x_out, self.original_space) elif self.options.project_low_method == 'project': self.x_out_projector = Projector(self.x_out, self.x_projected, bcs=post_apply_bcs) @@ -301,33 +303,12 @@ def post_apply(self, x_out): """ if self.interp_back: - self.x_out_projector.interpolate() + self.x_projected.assign(assemble(self.x_out_projector)) else: self.x_out_projector.project() x_out.assign(self.x_projected) -def is_cg(V): - """ - Checks if a :class:`FunctionSpace` is continuous. - - Function to check if a given space, V, is CG. Broken elements are always - discontinuous; for vector elements we check the names of the Sobolev spaces - of the subelements and for all other elements we just check the Sobolev - space name. - - Args: - V (:class:`FunctionSpace`): the space to check. - """ - ele = V.ufl_element() - if isinstance(ele, BrokenElement): - return False - elif type(ele) == VectorElement: - return all([e.sobolev_space.name == "H1" for e in ele._sub_elements]) - else: - return V.ufl_element().sobolev_space.name == "H1" - - class SUPGWrapper(Wrapper): """ Wrapper for computing a time discretisation with SUPG, which adjusts the diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py index 7c4de236b..51da84f2c 100644 --- a/gusto/timestepping/semi_implicit_quasi_newton.py +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -4,16 +4,17 @@ """ from firedrake import ( - Function, Constant, TrialFunctions, DirichletBC, div, Interpolator, + Function, Constant, TrialFunctions, DirichletBC, div, assemble, LinearVariationalProblem, LinearVariationalSolver ) from firedrake.fml import drop, replace_subject +from firedrake.__future__ import interpolate from pyop2.profiling import timed_stage from gusto.core import TimeLevelFields, StateFields from gusto.core.labels import (transport, diffusion, time_derivative, - linearisation, prognostic, hydrostatic, - physics_label, sponge, incompressible) -from gusto.solvers import LinearTimesteppingSolver + hydrostatic, physics_label, sponge, + incompressible) +from gusto.solvers import LinearTimesteppingSolver, mass_parameters from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation from gusto.timestepping.timestepper import BaseTimestepper @@ -38,7 +39,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, slow_physics_schemes=None, fast_physics_schemes=None, alpha=Constant(0.5), off_centred_u=False, num_outer=2, num_inner=2, accelerator=False, - predictor=None, reference_update_freq=None): + predictor=None, reference_update_freq=None, + spinup_steps=0): """ Args: equation_set (:class:`PrognosticEquationSet`): the prognostic @@ -105,15 +107,24 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, time step. Setting it to None turns off the update, and reference profiles will remain at their initial values. Defaults to None. + spinup_steps (int, optional): the number of steps to run the model + in "spin-up" mode, where the alpha parameter is set to 1.0. + Defaults to 0, which corresponds to no spin-up. """ self.num_outer = num_outer self.num_inner = num_inner - self.alpha = alpha + self.alpha = Constant(alpha) self.predictor = predictor self.accelerator = accelerator + + # Options relating to reference profiles and spin-up + self._alpha_original = Constant(alpha) self.reference_update_freq = reference_update_freq self.to_update_ref_profile = False + self.spinup_steps = spinup_steps + self.spinup_begun = False + self.spinup_done = False # Flag for if we have simultaneous transport self.simult = False @@ -209,14 +220,9 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.setup_transporting_velocity(aux_scheme) self.tracers_to_copy = [] - for name in equation_set.field_names: - # Extract time derivative for that prognostic - mass_form = equation_set.residual.label_map( - lambda t: (t.has_label(time_derivative) and t.get(prognostic) == name), - map_if_false=drop) - # Copy over field if the time derivative term has no linearisation - if not mass_form.terms[0].has_label(linearisation): - self.tracers_to_copy.append(name) + if equation_set.active_tracers is not None: + for active_tracer in equation_set.active_tracers: + self.tracers_to_copy.append(active_tracer.name) self.field_name = equation_set.field_name W = equation_set.function_space @@ -234,9 +240,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, V_DG = equation_set.domain.spaces('DG') self.predictor_field_in = Function(V_DG) div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.n('u')) - self.predictor_interpolator = Interpolator( - self.x.star(predictor)*div_factor, self.predictor_field_in - ) + self.predictor_interpolate = interpolate( + self.x.star(predictor)*div_factor, V_DG) def _apply_bcs(self): """ @@ -334,7 +339,7 @@ def transport_fields(self, outer, xstar, xp): # Pre-multiply this variable by (1 - dt*beta*div(u)) V = xstar(name).function_space() field_out = Function(V) - self.predictor_interpolator.interpolate() + self.predictor_field_in.assign(assemble(self.predictor_interpolate)) scheme.apply(field_out, self.predictor_field_in) # xp is xstar plus the increment from the transported predictor @@ -353,13 +358,43 @@ def update_reference_profiles(self): if float(self.t) + self.reference_update_freq > self.last_ref_update_time: self.equation.X_ref.assign(self.x.n(self.field_name)) self.last_ref_update_time = float(self.t) - if hasattr(self.linear_solver, 'update_reference_profiles'): - self.linear_solver.update_reference_profiles() + self.linear_solver.update_reference_profiles() elif self.to_update_ref_profile: - if hasattr(self.linear_solver, 'update_reference_profiles'): - self.linear_solver.update_reference_profiles() - self.to_update_ref_profile = False + self.linear_solver.update_reference_profiles() + self.to_update_ref_profile = False + + def start_spinup(self): + """ + Initialises the spin-up period, so that the scheme is implicit by + setting the off-centering parameter alpha to be 1. + """ + logger.debug('Starting spin-up period') + # Update alpha + self.alpha.assign(1.0) + self.linear_solver.alpha.assign(1.0) + # We need to tell solvers that they may need rebuilding + self.linear_solver.update_reference_profiles() + self.forcing.solvers['explicit'].invalidate_jacobian() + self.forcing.solvers['implicit'].invalidate_jacobian() + # This only needs doing once, so update the flag + self.spinup_begun = True + + def finish_spinup(self): + """ + Finishes the spin-up period, returning the off-centering parameter + to its original value. + """ + logger.debug('Finishing spin-up period') + # Update alpha + self.alpha.assign(self._alpha_original) + self.linear_solver.alpha.assign(self._alpha_original) + # We need to tell solvers that they may need rebuilding + self.linear_solver.update_reference_profiles() + self.forcing.solvers['explicit'].invalidate_jacobian() + self.forcing.solvers['implicit'].invalidate_jacobian() + # This only needs doing once, so update the flag + self.spinup_done = True def timestep(self): """Defines the timestep""" @@ -376,6 +411,13 @@ def timestep(self): # Update reference profiles -------------------------------------------- self.update_reference_profiles() + # Are we in spin-up period? -------------------------------------------- + # Note: steps numbered from 1 onwards + if self.step < self.spinup_steps + 1 and not self.spinup_begun: + self.start_spinup() + elif self.step >= self.spinup_steps + 1 and not self.spinup_done: + self.finish_spinup() + # Slow physics --------------------------------------------------------- x_after_slow(self.field_name).assign(xn(self.field_name)) if len(self.slow_physics_schemes) > 0: @@ -526,7 +568,7 @@ def __init__(self, equation, alpha): map_if_false=drop) # the explicit forms are multiplied by (1-alpha) and moved to the rhs - L_explicit = -(1-alpha)*dt*residual.label_map( + L_explicit = -(Constant(1)-alpha)*dt*residual.label_map( lambda t: any(t.has_label(time_derivative, hydrostatic, *implicit_terms, return_tuple=True)), @@ -571,13 +613,17 @@ def __init__(self, equation, alpha): constant_jacobian=True ) + self.solver_parameters = mass_parameters(W, equation.domain.spaces) + self.solvers = {} self.solvers["explicit"] = LinearVariationalSolver( explicit_forcing_problem, + solver_parameters=self.solver_parameters, options_prefix="ExplicitForcingSolver" ) self.solvers["implicit"] = LinearVariationalSolver( implicit_forcing_problem, + solver_parameters=self.solver_parameters, options_prefix="ImplicitForcingSolver" ) diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py index 3c619753b..24a2058f7 100644 --- a/gusto/timestepping/timestepper.py +++ b/gusto/timestepping/timestepper.py @@ -8,10 +8,11 @@ from gusto.core import TimeLevelFields, StateFields from gusto.core.io import TimeData from gusto.core.labels import transport, diffusion, prognostic, transporting_velocity -from gusto.core.logging import logger +from gusto.core.logging import logger, DEBUG from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation from gusto.spatial_methods.transport_methods import TransportMethod import ufl +import numpy as np __all__ = ["BaseTimestepper", "Timestepper", "PrescribedTransport"] @@ -162,6 +163,23 @@ def log_timestep(self): logger.info('='*40) logger.info(f'at start of timestep {self.step}, t={float(self.t)}, dt={float(self.dt)}') + def log_field_stats(self): + """ + Logs some field stats, which can be useful for debugging. + """ + current_log_level = logger.getEffectiveLevel() + if current_log_level > DEBUG: + return + for field_name in self.fields._field_names: + field_data = self.fields(field_name).dat.data_ro + # Mixed functions don't have min or max routines, and are less + # useful, so try to eliminate these by only logging fields with + # a 1-dimension array of data + if type(field_data) is np.ndarray and len(np.shape(field_data)) == 1: + min_val = field_data.min() + max_val = field_data.max() + logger.debug(f'{field_name}, min: {min_val}, max: {max_val}') + def run(self, t, tmax, pick_up=False): """ Runs the model for the specified time, from t to tmax @@ -200,6 +218,8 @@ def run(self, t, tmax, pick_up=False): logger.debug('Dumping output to disk') self.io.setup_dump(self.fields, t, pick_up) + self.log_field_stats() + self.t.assign(t) # Time loop @@ -226,6 +246,8 @@ def run(self, t, tmax, pick_up=False): ) self.io.dump(self.fields, time_data) + self.log_field_stats() + if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint': self.io.chkpt.close() @@ -406,7 +428,8 @@ def setup_prescribed_expr(self, expr_func): raise RuntimeError('Prescribed velocity already set up!') self.velocity_projection = Projector( - expr_func(self.t), self.fields('u') + expr_func(self.t), self.fields('u'), + quadrature_degree=self.equation.domain.max_quad_degree ) self.is_velocity_setup = True diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index b0508aae3..02b1addfc 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -25,7 +25,7 @@ def tracer_sphere(tmpdir, degree, small_dt): # Parameters chosen so that dt != 1 # Gaussian is translated from (lon=pi/2, lat=0) to (lon=0, lat=0) # to demonstrate that transport is working correctly - if (small_dt): + if small_dt: dt = pi/3. * 0.005 else: dt = pi/3. * 0.02 @@ -47,7 +47,7 @@ def tracer_sphere(tmpdir, degree, small_dt): uexpr, umax, radius, tol) -def tracer_slice(tmpdir, degree): +def tracer_slice(tmpdir, degree, small_dt): n = 30 if degree == 0 else 15 m = PeriodicIntervalMesh(n, 1.) mesh = ExtrudedMesh(m, layers=n, layer_height=1./n) @@ -56,7 +56,10 @@ def tracer_slice(tmpdir, degree): # Gaussian is translated by 1.5 times width of domain to demonstrate # that transport is working correctly - dt = 0.01 + if small_dt: + dt = 0.002 + else: + dt = 0.01 tmax = 0.75 output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) domain = Domain(mesh, dt, family="CG", degree=degree) @@ -80,8 +83,11 @@ def tracer_slice(tmpdir, degree): return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, tol=tol) -def tracer_blob_slice(tmpdir, degree): - dt = 0.01 +def tracer_blob_slice(tmpdir, degree, small_dt): + if small_dt: + dt = 0.002 + else: + dt = 0.01 L = 10. m = PeriodicIntervalMesh(10, L) mesh = ExtrudedMesh(m, layers=10, layer_height=1.) @@ -105,10 +111,9 @@ def _tracer_setup(tmpdir, geometry, blob=False, degree=1, small_dt=False): assert not blob return tracer_sphere(tmpdir, degree, small_dt) elif geometry == "slice": - assert not small_dt if blob: - return tracer_blob_slice(tmpdir, degree) + return tracer_blob_slice(tmpdir, degree, small_dt) else: - return tracer_slice(tmpdir, degree) + return tracer_slice(tmpdir, degree, small_dt) return _tracer_setup diff --git a/integration-tests/data/boussinesq_compressible_chkpt.h5 b/integration-tests/data/boussinesq_compressible_chkpt.h5 index 78331ca96..13819691c 100644 Binary files a/integration-tests/data/boussinesq_compressible_chkpt.h5 and b/integration-tests/data/boussinesq_compressible_chkpt.h5 differ diff --git a/integration-tests/data/boussinesq_incompressible_chkpt.h5 b/integration-tests/data/boussinesq_incompressible_chkpt.h5 index 22f269d22..a4612c676 100644 Binary files a/integration-tests/data/boussinesq_incompressible_chkpt.h5 and b/integration-tests/data/boussinesq_incompressible_chkpt.h5 differ diff --git a/integration-tests/data/dry_compressible_chkpt.h5 b/integration-tests/data/dry_compressible_chkpt.h5 index f55213d26..c10668ea1 100644 Binary files a/integration-tests/data/dry_compressible_chkpt.h5 and b/integration-tests/data/dry_compressible_chkpt.h5 differ diff --git a/integration-tests/data/linear_sw_wave_chkpt.h5 b/integration-tests/data/linear_sw_wave_chkpt.h5 index c7cdb5c21..70c7c0244 100644 Binary files a/integration-tests/data/linear_sw_wave_chkpt.h5 and b/integration-tests/data/linear_sw_wave_chkpt.h5 differ diff --git a/integration-tests/data/moist_compressible_chkpt.h5 b/integration-tests/data/moist_compressible_chkpt.h5 index 3e45478f3..60381d2b6 100644 Binary files a/integration-tests/data/moist_compressible_chkpt.h5 and b/integration-tests/data/moist_compressible_chkpt.h5 differ diff --git a/integration-tests/data/simult_SIQN_order0_chkpt.h5 b/integration-tests/data/simult_SIQN_order0_chkpt.h5 index dd8b74f75..32e48d376 100644 Binary files a/integration-tests/data/simult_SIQN_order0_chkpt.h5 and b/integration-tests/data/simult_SIQN_order0_chkpt.h5 differ diff --git a/integration-tests/data/simult_SIQN_order1_chkpt.h5 b/integration-tests/data/simult_SIQN_order1_chkpt.h5 index 66b2fbadc..f13913a48 100644 Binary files a/integration-tests/data/simult_SIQN_order1_chkpt.h5 and b/integration-tests/data/simult_SIQN_order1_chkpt.h5 differ diff --git a/integration-tests/data/sw_fplane_chkpt.h5 b/integration-tests/data/sw_fplane_chkpt.h5 index d09afda93..d340df20d 100644 Binary files a/integration-tests/data/sw_fplane_chkpt.h5 and b/integration-tests/data/sw_fplane_chkpt.h5 differ diff --git a/integration-tests/equations/test_linear_sw_wave.py b/integration-tests/equations/test_linear_sw_wave.py index 81b6745e8..95eecdca4 100644 --- a/integration-tests/equations/test_linear_sw_wave.py +++ b/integration-tests/equations/test_linear_sw_wave.py @@ -11,7 +11,7 @@ def run_linear_sw_wave(tmpdir): - # Paramerers + # Parameters dt = 0.001 tmax = 30*dt H = 1 @@ -57,9 +57,11 @@ def run_linear_sw_wave(tmpdir): u0 = stepper.fields("u") D0 = stepper.fields("D") + Dbar = eqns.X_ref.subfunctions[1] u0.project(uexpr) D0.interpolate(Dexpr) + Dbar.interpolate(H) # --------------------------------------------------------------------- # # Run diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py index 8e0779c2b..3d0acff5f 100644 --- a/integration-tests/equations/test_sw_fplane.py +++ b/integration-tests/equations/test_sw_fplane.py @@ -38,10 +38,12 @@ def run_sw_fplane(tmpdir): io = IO(domain, output, diagnostic_fields=[CourantNumber()]) # Transport schemes - transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "D")] - transport_methods = [DGUpwind(eqns, "u"), - DGUpwind(eqns, "D")] + vorticity_transport = VorticityTransport(domain, eqns, supg=True) + transported_fields = [ + TrapeziumRule(domain, "u", augmentation=vorticity_transport), + SSPRK3(domain, "D") + ] + transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, diff --git a/integration-tests/equations/test_thermal_sw.py b/integration-tests/equations/test_thermal_sw.py index 8328c308d..f201e78c6 100644 --- a/integration-tests/equations/test_thermal_sw.py +++ b/integration-tests/equations/test_thermal_sw.py @@ -40,9 +40,9 @@ def setup_sw(dirname, dt, u_transport_option): parameters = ShallowWaterParameters(H=H, g=g) Omega = parameters.Omega fexpr = 2*Omega*x[2]/R - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option=u_transport_option, - thermal=True) + eqns = ThermalShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option=u_transport_option) + # I/O diagnostic_fields = [SteadyStateError('D'), SteadyStateError('u'), diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py index 0f8d7f112..55d63264d 100644 --- a/integration-tests/model/test_nc_outputting.py +++ b/integration-tests/model/test_nc_outputting.py @@ -11,8 +11,17 @@ ForwardEuler, OutputParameters, XComponent, YComponent, ZComponent, MeridionalComponent, ZonalComponent, RadialComponent, DGUpwind) -from netCDF4 import Dataset +from mpi4py import MPI +from netCDF4 import Dataset, chartostring import pytest +from pytest_mpi import parallel_assert + + +def make_dirname(test_name, suffix=""): + if MPI.COMM_WORLD.size > 1: + return f"pytest_{test_name}_parallel" + suffix + else: + return f"pytest_{test_name}" + suffix @pytest.fixture @@ -53,17 +62,19 @@ def domain_and_mesh_details(geometry): return (domain, mesh_details) -# TODO: make parallel configurations of this test +@pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("geometry", ["interval", "vertical_slice", "plane", "extruded_plane", "spherical_shell", "extruded_spherical_shell"]) -def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): +def test_nc_outputting(geometry, domain_and_mesh_details): # ------------------------------------------------------------------------ # # Make model objects # ------------------------------------------------------------------------ # - dirname = str(tmpdir) + # Make sure all ranks use the same file + dirname = make_dirname("nc_outputting", suffix=f"_{geometry}_{MPI.COMM_WORLD.size}") + domain, mesh_details = domain_and_mesh_details V = domain.spaces('DG') if geometry == "interval": @@ -136,16 +147,28 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): # ------------------------------------------------------------------------ # # Check that metadata is correct - output_data = Dataset(f'{dirname}/field_output.nc', 'r') + try: + output_data = Dataset(f'results/{dirname}/field_output.nc', 'r', parallel=True) + except ValueError: + # serial netCDF4, do everything on rank 0 + if MPI.COMM_WORLD.rank == 0: + output_data = Dataset(f'results/{dirname}/field_output.nc', 'r', parallel=False) + else: + output_data = None + for metadata_key, metadata_value in mesh_details.items(): # Convert None or booleans to strings - if type(metadata_value) in [type(None), type(True)]: + if metadata_value is None or isinstance(metadata_value, bool): output_value = str(metadata_value) else: output_value = metadata_value error_message = f'Metadata {metadata_key} for geometry {geometry} is incorrect' if type(output_value) == float: - assert output_data[metadata_key][0] - output_value < 1e-14, error_message + def assertion(): + return output_data[metadata_key][0] - output_value < 1e-14 else: - assert output_data[metadata_key][0] == output_value, error_message + def assertion(): + return str(chartostring(output_data[metadata_key][0])) == output_value + + parallel_assert(assertion, participating=output_data is not None, msg=error_message) diff --git a/integration-tests/model/test_nonsplit_physics.py b/integration-tests/model/test_nonsplit_physics.py new file mode 100644 index 000000000..cf2c7bd3d --- /dev/null +++ b/integration-tests/model/test_nonsplit_physics.py @@ -0,0 +1,229 @@ +""" +This script tests the non-split timestepper against the split timestepper +using a coupled wave equation, where one of the spatial derivatives +is a physics parametrisation. +One split method is tested, whilst different nonsplit IMEX and explicit time +discretisations are used for the dynamics and physics. +""" + +from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, split, + norm, Constant, cos, sin, Function, Projector, inner, dx) +from gusto import * +import pytest +from math import pi + + +class WaveEquation(PrognosticEquationSet): + u"""Discretises the wave equation as ∂u/∂t + ∇v = 0 and ∂v/∂t + ∇u = 0""" + + def __init__(self, domain, space_names=None, field_names=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + space_names (:class:`FunctionSpace`): the function spaces that the + equation's prognostic fields are defined on. + field_names (str): names of the prognostic fields. + """ + + field_names = ['u', 'v'] + + if space_names is None: + space_names = {'u': 'HDiv', 'v': 'DG'} + + super().__init__(field_names, domain, space_names) + + w, phi = self.tests + u, v = split(self.X) + + # Non time derivative term of u equation is a physics parametrisation + mass_form_u = time_derivative(inner(u, w)*dx) + u_eqn = prognostic(subject(mass_form_u, self.X), 'u') + + # v equation includes a time derivative term and gradient term + mass_form_v = time_derivative(inner(v, phi)*dx) + rhs_v = u.dx(0)*phi*dx + v_eqn = prognostic(subject(mass_form_v + rhs_v, self.X), 'v') + + self.residual = u_eqn + v_eqn + + +class WaveEquationForcing(PhysicsParametrisation): + """ + A physics parametrisation for the coupled wave equation. The source term is ∇v in + the u equation. + """ + + def __init__(self, equation): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + """ + + label_name = 'wave_forcing' + super().__init__(equation, label_name, parameters=None) + + self.dt = Constant(0.0) + self.V = equation.X.function_space() + self.fields = Function(self.V) + + # -------------------------------------------------------------------- # + # Extract prognostic variables and test functions + # -------------------------------------------------------------------- # + u, v = self.fields.subfunctions + Vu = u.function_space() + Vv = v.function_space() + test_u, test_v = equation.tests + + # Set up funtions for projectors + self.v_cg = Function(Vu) + self.grad_v = Function(Vv) + + # Set up source term in mixed function space + self.source = Function(self.V) + + # First project from DG0 into CG1 + self.projector_cg = Projector(v, self.v_cg) + + # Then project the derivative of v back into DG0 + self.projector_grad = Projector(self.v_cg.dx(0), self.grad_v) + + # Finally project the gradient into the source term in CG1. This must be projected + # into the u subfunction of the mixed source function + self.projector_source = Projector(self.grad_v, self.source.subfunctions[0]) + + # Set up the expression of source term in the u equation. The u component of the + # source term must be split using the ufl split function + source_u_expr = inner(split(self.source)[0], test_u)*dx + + # Add the source term to the equation + equation.residual += source_label(self.label(subject(source_u_expr, + self.source), self.evaluate)) + + def evaluate(self, x_in, dt, x_out=None): + """ + Evaluates the source term generated by the wave equation physics forcing + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + x_out: (:class:`Function`, optional): the (mixed) source + field to be outputed. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + self.fields.assign(x_in) + + # Do the projections + self.projector_cg.project() + self.projector_grad.project() + self.projector_source.project() + + # If a source output is provided, assign the source term to it + if x_out is not None: + x_out.assign(self.source) + + +def run_nonsplit_physics(tmpdir, timestepper): + """ + Runs the coupled wave equation with a ∇v physics parametrisation using different timesteppers. + """ + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + dt = 0.001 + tmax = 0.3 + L = 1 + mesh = PeriodicIntervalMesh(100, L) + domain = Domain(mesh, dt, "CG", 1) + + # Equation + x = SpatialCoordinate(mesh)[0] + eqn = WaveEquation(domain) + + # I/O + output = OutputParameters(dirname=str(tmpdir), dumpfreq=30) + io = IO(domain, output) + + # Time stepper + if timestepper == 'split': + physics_schemes = [(WaveEquationForcing(eqn), + RK4(domain))] + stepper = SplitPhysicsTimestepper(eqn, RK4(domain), + io, + physics_schemes=physics_schemes) + elif timestepper == 'nonsplit_exp_rk_predictor': + physics_parametrisation = [WaveEquationForcing(eqn)] + scheme = RK4(domain, rk_formulation=RungeKuttaFormulation.predictor) + stepper = Timestepper(eqn, scheme, + io, physics_parametrisations=physics_parametrisation) + elif timestepper == 'nonsplit_exp_rk_increment': + physics_parametrisation = [WaveEquationForcing(eqn)] + scheme = RK4(domain, rk_formulation=RungeKuttaFormulation.increment) + stepper = Timestepper(eqn, scheme, + io, physics_parametrisations=physics_parametrisation) + elif timestepper == 'nonsplit_imex_rk': + physics_parametrisation = [WaveEquationForcing(eqn)] + eqn.label_terms(lambda t: not any(t.has_label(time_derivative, source_label)), implicit) + scheme = IMEX_SSP3(domain) + stepper = Timestepper(eqn, scheme, + io, physics_parametrisations=physics_parametrisation) + elif timestepper == 'nonsplit_imex_sdc': + physics_parametrisation = [WaveEquationForcing(eqn)] + quad_type = "GAUSS" + M = 2 + k = 2 + qdelta_imp = "LU" + node_type = "LEGENDRE" + qdelta_exp = "FE" + eqn.label_terms(lambda t: not any(t.has_label(time_derivative, source_label)), implicit) + base_scheme = IMEX_Euler(domain) + scheme = SDC(base_scheme, domain, M, k, quad_type, node_type, qdelta_imp, + qdelta_exp, final_update=True, initial_guess="base") + stepper = Timestepper(eqn, scheme, + io, physics_parametrisations=physics_parametrisation) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + c = 1 + uexp = sin(2*pi*x/L) + u = stepper.fields("u") + v = stepper.fields("v") + u.interpolate(uexp) + v.interpolate(Constant(0)) + u_final_exp = sin(2*pi*x/L)*cos(2*pi*c*tmax/L) + u_exact = Function(u.function_space()).project(u_final_exp) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(0, tmax=tmax) + + error = norm(stepper.fields('u') - u_exact) / norm(u_exact) + return error + + +@pytest.mark.parametrize("timestepper", ["split", "nonsplit_exp_rk_predictor", + "nonsplit_exp_rk_increment", + "nonsplit_imex_rk", "nonsplit_imex_sdc"]) +def test_nonsplit_physics(tmpdir, timestepper): + """ + Test the nonsplit timestepper in the wave equation with source physics. + """ + if timestepper == 'split': + # Split has a looser tolerance due to the physics + # parametrisation not being tightly coupled to the dynamics + tol = 1e-2 + else: + tol = 1e-4 + error = run_nonsplit_physics(tmpdir, timestepper) + assert error < tol, 'The nonsplit timestepper in the coupled wave ' + \ + 'equation with source physics has an error greater than ' + \ + 'the permitted tolerance' diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 6d484bfd1..b33a34300 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -10,9 +10,19 @@ def run(timestepper, tmax, f_end): @pytest.mark.parametrize( "scheme", [ - "ssprk3_increment", "TrapeziumRule", "ImplicitMidpoint", "QinZhang", + "ssprk2_increment_2", "ssprk2_predictor_2", "ssprk2_linear_2", + "ssprk2_increment_3", "ssprk2_predictor_3", "ssprk2_linear_3", + "ssprk2_increment_4", "ssprk2_predictor_4", "ssprk2_linear_4", + + "ssprk3_increment_3", "ssprk3_predictor_3", "ssprk3_linear_3", + "ssprk3_increment_4", "ssprk3_predictor_4", "ssprk3_linear_4", + "ssprk3_increment_5", "ssprk3_predictor_5", "ssprk3_linear_5", + + "ssprk4_increment_5", "ssprk4_predictor_5", "ssprk4_linear_5", + + "TrapeziumRule", "ImplicitMidpoint", "QinZhang_increment", "QinZhang_predictor", "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", - "AdamsMoulton", "AdamsMoulton", "ssprk3_predictor", "ssprk3_linear" + "AdamsMoulton", "AdamsMoulton" ] ) def test_time_discretisation(tmpdir, scheme, tracer_setup): @@ -30,18 +40,60 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): V = domain.spaces("DG") eqn = AdvectionEquation(domain, V, "f") - if scheme == "ssprk3_increment": + if scheme == "ssprk2_increment_2": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment) + elif scheme == "ssprk2_predictor_2": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor) + elif scheme == "ssprk2_linear_2": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear) + elif scheme == "ssprk2_increment_3": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment, stages=3) + elif scheme == "ssprk2_predictor_3": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=3) + elif scheme == "ssprk2_linear_3": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear, stages=3) + elif scheme == "ssprk2_increment_4": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment, stages=4) + elif scheme == "ssprk2_predictor_4": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=4) + elif scheme == "ssprk2_linear_4": + transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear, stages=4) + + elif scheme == "ssprk3_increment_3": transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment) - elif scheme == "ssprk3_predictor": + elif scheme == "ssprk3_predictor_3": transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor) - elif scheme == "ssprk3_linear": + elif scheme == "ssprk3_linear_3": transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear) + elif scheme == "ssprk3_increment_4": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment, stages=4) + elif scheme == "ssprk3_predictor_4": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=4) + elif scheme == "ssprk3_linear_4": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear, stages=4) + + elif scheme == "ssprk3_increment_5": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment, stages=5) + elif scheme == "ssprk3_predictor_5": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=5) + elif scheme == "ssprk3_linear_5": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear, stages=5) + + elif scheme == "ssprk4_increment_5": + transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.increment) + elif scheme == "ssprk4_predictor_5": + transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.predictor) + elif scheme == "ssprk4_linear_5": + transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.linear) + elif scheme == "TrapeziumRule": transport_scheme = TrapeziumRule(domain) elif scheme == "ImplicitMidpoint": transport_scheme = ImplicitMidpoint(domain) - elif scheme == "QinZhang": - transport_scheme = QinZhang(domain) + elif scheme == "QinZhang_increment": + transport_scheme = QinZhang(domain, rk_formulation=RungeKuttaFormulation.increment) + elif scheme == "QinZhang_predictor": + transport_scheme = QinZhang(domain, rk_formulation=RungeKuttaFormulation.predictor) elif scheme == "RK4": transport_scheme = RK4(domain) elif scheme == "Heun": diff --git a/integration-tests/physics/test_held_suarez_friction.py b/integration-tests/physics/test_held_suarez_friction.py new file mode 100644 index 000000000..dfe7e34f9 --- /dev/null +++ b/integration-tests/physics/test_held_suarez_friction.py @@ -0,0 +1,114 @@ +""" +This tests the Rayleigh friction term used in the Held Suarez test case. +""" + +from gusto import * +import gusto.equations.thermodynamics as td +from gusto.core.labels import physics_label +from firedrake import (Constant, PeriodicIntervalMesh, as_vector, norm, + ExtrudedMesh, Function, dot) +from firedrake.fml import identity, drop + + +def run_apply_rayleigh_friction(dirname): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 3600.0 + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = int(H / 5.) + ncolumns = int(L / 5.) + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 0) + + # Set up equation + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters) + + # I/O + output = OutputParameters(dirname=dirname+"/held_suarez_friction", + dumpfreq=1, + dumplist=['u']) + io = IO(domain, output) + + # Physics scheme + physics_parametrisation = RayleighFriction(eqn) + + time_discretisation = BackwardEuler(domain) + + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vu = domain.spaces("HDiv") + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Declare prognostic fields + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set a background state with constant pressure and temperature + pressure = Function(Vr).interpolate(Constant(100000.)) + temperature = Function(Vt).interpolate(Constant(295.)) + theta_d = td.theta(parameters, temperature, pressure) + + theta0.project(theta_d) + rho0.interpolate(pressure / (temperature*parameters.R_d)) + + # Constant horizontal wind + u0.project(as_vector([864, 0.0])) + + # Answer: slower winds than initially + u_true = Function(Vu) + u_true.project(as_vector([828, 0.0])) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return mesh, stepper, u_true + + +def test_rayleigh_friction(tmpdir): + + dirname = str(tmpdir) + mesh, stepper, u_true = run_apply_rayleigh_friction(dirname) + + u_final = stepper.fields('u') + + # Project into CG1 to get sensible values + e_x = as_vector([1.0, 0.0]) + e_z = as_vector([0.0, 1.0]) + + DG0 = FunctionSpace(mesh, "DG", 0) + u_x_final = Function(DG0).project(dot(u_final, e_x)) + u_x_true = Function(DG0).project(dot(u_true, e_x)) + u_z_final = Function(DG0).project(dot(u_final, e_z)) + u_z_true = Function(DG0).project(dot(u_true, e_z)) + + denom = norm(u_x_true) + assert norm(u_x_final - u_x_true) / denom < 0.0001, 'Final horizontal wind is incorrect' + assert norm(u_z_final - u_z_true) < 1e-12, 'Final vertical wind is incorrect' diff --git a/integration-tests/physics/test_held_suarez_relaxation.py b/integration-tests/physics/test_held_suarez_relaxation.py new file mode 100644 index 000000000..615ba8d29 --- /dev/null +++ b/integration-tests/physics/test_held_suarez_relaxation.py @@ -0,0 +1,107 @@ +""" +This tests the Held-Suarez physics routine to apply Rayleigh friction. +""" +from gusto import * +import gusto.equations.thermodynamics as td +from gusto.core.labels import physics_label +from firedrake import (Constant, PeriodicIntervalMesh, as_vector, + ExtrudedMesh, Function) +from firedrake.fml import identity, drop +import pytest + + +def run_held_suarez_relaxation(dirname, temp): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 3600.0 + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = int(H / 5.) + ncolumns = int(L / 5.) + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 0) + + # Set up equation + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters) + + # I/O + output = OutputParameters(dirname=dirname+"/held_suarez_friction", + dumpfreq=1, + dumplist=['u']) + io = IO(domain, output) + + # Physics scheme + physics_parametrisation = Relaxation(eqn, 'theta', parameters) + + time_discretisation = ForwardEuler(domain) + + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Declare prognostic fields + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set a background state with constant pressure and temperature + p0 = 100000. + pressure = Function(Vr).interpolate(Constant(p0)) + temperature = Function(Vt).interpolate(Constant(temp)) + theta_d = td.theta(parameters, temperature, pressure) + + theta0.project(theta_d) + rho0.interpolate(p0 / (theta_d*parameters.R_d)) # This ensures that exner = 1 + + # Constant horizontal wind + u0.project(as_vector([1.0, 1.0])) + theta_initial = Function(Vt) + theta_initial.interpolate(theta_d) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return stepper, theta_initial + + +@pytest.mark.parametrize('temp', [280, 290]) +def test_held_suarez_relaxation(tmpdir, temp): + # By configuring the fields we have set the equilibrium temperature to 285K + # We test a temperature value eith side to check it moves in the right direction + dirname = str(tmpdir) + stepper, theta_initial = run_held_suarez_relaxation(dirname, temp) + + theta_final = stepper.fields('theta') + final_data = theta_final.dat.data + initial_data = theta_initial.dat.data + if temp == 280: + assert np.mean(final_data) > np.mean(initial_data) + if temp == 290: + assert np.mean(final_data) < np.mean(initial_data) diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index 075af40f2..3cb42d839 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -10,9 +10,10 @@ from gusto import * from firedrake import (Constant, PeriodicSquareMesh, SpatialCoordinate, sqrt, conditional, cos, pi, FunctionSpace) +import pytest -def run_instant_rain(dirname): +def run_instant_rain(dirname, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -53,12 +54,18 @@ def run_instant_rain(dirname): # Physics schemes # define saturation function saturation = Constant(0.5) - physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain"), - RK4(domain))] - - # Time stepper - stepper = SplitPhysicsTimestepper(eqns, RK4(domain), io, transport_method, - physics_schemes=physics_schemes) + if physics_coupling == "split": + physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain"), + RK4(domain))] + # Time stepper + stepper = SplitPhysicsTimestepper(eqns, RK4(domain), io, transport_method, + physics_schemes=physics_schemes) + else: + physics_parametrisation = [InstantRain(eqns, saturation, rain_name="rain")] + scheme = RK4(domain, rk_formulation=RungeKuttaFormulation.predictor) + # Time stepper + stepper = Timestepper(eqns, scheme, io, transport_method, + physics_parametrisations=physics_parametrisation) # ------------------------------------------------------------------------ # # Initial conditions @@ -91,9 +98,11 @@ def run_instant_rain(dirname): return stepper, saturation, initial_vapour, vapour_true, rain_true -def test_instant_rain_setup(tmpdir): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_instant_rain_setup(tmpdir, physics_coupling): dirname = str(tmpdir) - stepper, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) + stepper, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname, + physics_coupling) v = stepper.fields("water_vapour") r = stepper.fields("rain") diff --git a/integration-tests/physics/test_saturation_adjustment.py b/integration-tests/physics/test_saturation_adjustment.py index 74e6e129d..b39ee0345 100644 --- a/integration-tests/physics/test_saturation_adjustment.py +++ b/integration-tests/physics/test_saturation_adjustment.py @@ -15,7 +15,7 @@ import pytest -def run_cond_evap(dirname, process): +def run_cond_evap(dirname, process, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -47,13 +47,21 @@ def run_cond_evap(dirname, process): dumplist=['u']) io = IO(domain, output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) - # Physics scheme - physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain))] + if physics_coupling == "split": + # Physics scheme + physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain))] - # Time stepper - scheme = ForwardEuler(domain) - stepper = SplitPhysicsTimestepper(eqn, scheme, io, - physics_schemes=physics_schemes) + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + else: + # Physics scheme + physics_parametrisation = [SaturationAdjustment(eqn, parameters=parameters)] + + # Time stepper + scheme = ForwardEuler(domain, rk_formulation=RungeKuttaFormulation.predictor) + stepper = Timestepper(eqn, scheme, io, physics_parametrisations=physics_parametrisation) # ------------------------------------------------------------------------ # # Initial conditions @@ -116,10 +124,13 @@ def run_cond_evap(dirname, process): @pytest.mark.parametrize("process", ["evaporation", "condensation"]) -def test_cond_evap(tmpdir, process): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_cond_evap(tmpdir, process, physics_coupling): dirname = str(tmpdir) - eqn, stepper, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) + eqn, stepper, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, + process, + physics_coupling) water_v = stepper.fields('water_vapour') water_c = stepper.fields('cloud_water') diff --git a/integration-tests/physics/test_source_sink.py b/integration-tests/physics/test_source_sink.py index 97aa1ce8b..b64e9262e 100644 --- a/integration-tests/physics/test_source_sink.py +++ b/integration-tests/physics/test_source_sink.py @@ -8,7 +8,7 @@ import pytest -def run_source_sink(dirname, process, time_varying): +def run_source_sink(dirname, process, time_varying, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -60,11 +60,15 @@ def time_varying_expression(t): # Time stepper time_varying_velocity = False + if physics_coupling == "split": + scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment) + else: + scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor) + stepper = PrescribedTransport( - eqn, SSPRK3(domain), io, time_varying_velocity, + eqn, scheme, io, time_varying_velocity, transport_method, physics_parametrisations=physics_parametrisations ) - # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # @@ -94,9 +98,10 @@ def time_varying_expression(t): @pytest.mark.parametrize("process", ["source", "sink"]) @pytest.mark.parametrize("time_varying", [False, True]) -def test_source_sink(tmpdir, process, time_varying): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_source_sink(tmpdir, process, time_varying, physics_coupling): dirname = str(tmpdir) - stepper, initial_ash = run_source_sink(dirname, process, time_varying) + stepper, initial_ash = run_source_sink(dirname, process, time_varying, physics_coupling) final_ash = stepper.fields("ash") initial_total_ash = assemble(initial_ash*dx) diff --git a/integration-tests/physics/test_suppress_vertical_wind.py b/integration-tests/physics/test_suppress_vertical_wind.py index 8765b3985..fee46b84a 100644 --- a/integration-tests/physics/test_suppress_vertical_wind.py +++ b/integration-tests/physics/test_suppress_vertical_wind.py @@ -8,9 +8,10 @@ from firedrake import (Constant, PeriodicIntervalMesh, as_vector, sin, norm, SpatialCoordinate, ExtrudedMesh, Function, dot, pi) from firedrake.fml import identity, drop +import pytest -def run_suppress_vertical_wind(dirname): +def run_suppress_vertical_wind(dirname, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -46,19 +47,20 @@ def run_suppress_vertical_wind(dirname): # Physics scheme physics_parametrisation = SuppressVerticalWind(eqn, spin_up_period) - time_discretisation = ForwardEuler(domain) - - # time_discretisation = ForwardEuler(domain) - physics_schemes = [(physics_parametrisation, time_discretisation)] - # Only want time derivatives and physics terms in equation, so drop the rest eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), map_if_true=identity, map_if_false=drop) - - # Time stepper - scheme = ForwardEuler(domain) - stepper = SplitPhysicsTimestepper(eqn, scheme, io, - physics_schemes=physics_schemes) + if physics_coupling == "split": + time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + else: + # Time stepper + scheme = ForwardEuler(domain, rk_formulation=RungeKuttaFormulation.predictor) + stepper = Timestepper(eqn, scheme, io, physics_parametrisations=[physics_parametrisation]) # ------------------------------------------------------------------------ # # Initial conditions @@ -83,10 +85,11 @@ def run_suppress_vertical_wind(dirname): return domain, stepper -def test_suppress_vertical_wind(tmpdir): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_suppress_vertical_wind(tmpdir, physics_coupling): dirname = str(tmpdir) - domain, stepper = run_suppress_vertical_wind(dirname) + domain, stepper = run_suppress_vertical_wind(dirname, physics_coupling) u = stepper.fields('u') vertical_wind = Function(domain.spaces('theta')) diff --git a/integration-tests/physics/test_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py index 35d8e9b89..14990941c 100644 --- a/integration-tests/physics/test_surface_fluxes.py +++ b/integration-tests/physics/test_surface_fluxes.py @@ -13,7 +13,7 @@ import pytest -def run_surface_fluxes(dirname, moist, implicit_formulation): +def run_surface_fluxes(dirname, moist, implicit_formulation, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -49,22 +49,34 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): # Physics scheme surf_params = BoundaryLayerParameters() T_surf = Constant(300.0) - physics_parametrisation = SurfaceFluxes(eqn, T_surf, vapour_name, - implicit_formulation, surf_params) - time_discretisation = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) + if physics_coupling == "split": + physics_parametrisation = SurfaceFluxes(eqn, T_surf, vapour_name, + implicit_formulation, surf_params) - # time_discretisation = ForwardEuler(domain) - physics_schemes = [(physics_parametrisation, time_discretisation)] + time_discretisation = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) - # Only want time derivatives and physics terms in equation, so drop the rest - eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), - map_if_true=identity, map_if_false=drop) + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] - # Time stepper - scheme = ForwardEuler(domain) - stepper = SplitPhysicsTimestepper(eqn, scheme, io, - physics_schemes=physics_schemes) + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + else: + physics_parametrisation = [SurfaceFluxes(eqn, T_surf, vapour_name, + implicit_formulation, surf_params)] + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) + stepper = Timestepper(eqn, scheme, io, physics_parametrisations=physics_parametrisation) # ------------------------------------------------------------------------ # # Initial conditions @@ -120,10 +132,11 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): @pytest.mark.parametrize("moist", [False, True]) @pytest.mark.parametrize("implicit_formulation", [False, True]) -def test_surface_fluxes(tmpdir, moist, implicit_formulation): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_surface_fluxes(tmpdir, moist, implicit_formulation, physics_coupling): dirname = str(tmpdir) - eqn, stepper, T_true, mv_true = run_surface_fluxes(dirname, moist, implicit_formulation) + eqn, stepper, T_true, mv_true = run_surface_fluxes(dirname, moist, implicit_formulation, physics_coupling) # Back out temperature from prognostic fields theta_vd = stepper.fields('theta') diff --git a/integration-tests/physics/test_sw_saturation_adjustment.py b/integration-tests/physics/test_sw_saturation_adjustment.py index 66c971519..5b430a45c 100644 --- a/integration-tests/physics/test_sw_saturation_adjustment.py +++ b/integration-tests/physics/test_sw_saturation_adjustment.py @@ -17,7 +17,7 @@ import pytest -def run_sw_cond_evap(dirname, process): +def run_sw_cond_evap(dirname, process, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -49,26 +49,35 @@ def run_sw_cond_evap(dirname, process): tracers = [WaterVapour(space='DG'), CloudWater(space='DG')] - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option='vector_advection_form', - thermal=True, active_tracers=tracers) + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, + u_transport_option='vector_advection_form', + active_tracers=tracers) # I/O output = OutputParameters(dirname=dirname+"/sw_cond_evap", dumpfreq=1) io = IO(domain, output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) - - # Physics schemes - physics_schemes = [(SWSaturationAdjustment(eqns, sat, - parameters=parameters, - thermal_feedback=True, - beta2=beta2), - ForwardEuler(domain))] - - # Timestepper - stepper = SplitPhysicsTimestepper(eqns, RK4(domain), io, - physics_schemes=physics_schemes) + if physics_coupling == "split": + # Physics schemes + physics_schemes = [(SWSaturationAdjustment(eqns, sat, + parameters=parameters, + thermal_feedback=True, + beta2=beta2), + ForwardEuler(domain))] + + # Timestepper + stepper = SplitPhysicsTimestepper(eqns, RK4(domain), io, + physics_schemes=physics_schemes) + else: + SWSaturationAdjustment(eqns, sat, + parameters=parameters, + thermal_feedback=True, + beta2=beta2) + stepper = Timestepper(eqns, + ForwardEuler(domain, rk_formulation=RungeKuttaFormulation.predictor), + io) # Initial conditions b0 = stepper.fields("b") @@ -118,10 +127,11 @@ def run_sw_cond_evap(dirname, process): @pytest.mark.parametrize("process", ["evaporation", "condensation"]) -def test_cond_evap(tmpdir, process): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_cond_evap(tmpdir, process, physics_coupling): dirname = str(tmpdir) - eqns, stepper, v_true, c_true, b_true, c_init = run_sw_cond_evap(dirname, process) + eqns, stepper, v_true, c_true, b_true, c_init = run_sw_cond_evap(dirname, process, physics_coupling) vapour = stepper.fields("water_vapour") cloud = stepper.fields("cloud_water") diff --git a/integration-tests/physics/test_terminator_toy.py b/integration-tests/physics/test_terminator_toy.py index 1ae7b69df..2ff6e5445 100644 --- a/integration-tests/physics/test_terminator_toy.py +++ b/integration-tests/physics/test_terminator_toy.py @@ -8,9 +8,10 @@ sin, SpatialCoordinate, Function, max_value, as_vector, \ errornorm, norm import numpy as np +import pytest -def run_terminator_toy(dirname): +def run_terminator_toy(dirname, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -32,15 +33,15 @@ def run_terminator_toy(dirname): domain = Domain(mesh, dt, 'BDM', 1) # Define the interacting species - X = ActiveTracer(name='X', space='DG', + Y = ActiveTracer(name='Y', space='DG', variable_type=TracerVariableType.mixing_ratio, transport_eqn=TransportEquationType.advective) - X2 = ActiveTracer(name='X2', space='DG', + Y2 = ActiveTracer(name='Y2', space='DG', variable_type=TracerVariableType.mixing_ratio, transport_eqn=TransportEquationType.advective) - tracers = [X, X2] + tracers = [Y, Y2] # Equation V = domain.spaces("HDiv") @@ -56,18 +57,27 @@ def run_terminator_toy(dirname): k1 = max_value(0, sin(theta)*sin(theta_c) + cos(theta)*cos(theta_c)*cos(lamda-lamda_c)) k2 = 1 - - physics_schemes = [(TerminatorToy(eqn, k1=k1, k2=k2, species1_name='X', - species2_name='X2'), BackwardEuler(domain))] - - transport_scheme = SSPRK3(domain) - transport_method = [DGUpwind(eqn, 'X'), DGUpwind(eqn, 'X2')] - - time_varying_velocity = True - stepper = SplitPrescribedTransport( - eqn, transport_scheme, io, time_varying_velocity, - spatial_methods=transport_method, physics_schemes=physics_schemes - ) + transport_method = [DGUpwind(eqn, 'Y'), DGUpwind(eqn, 'Y2')] + if physics_coupling == "split": + physics_schemes = [(TerminatorToy(eqn, k1=k1, k2=k2, species1_name='Y', + species2_name='Y2'), BackwardEuler(domain))] + + transport_scheme = SSPRK3(domain) + time_varying_velocity = True + stepper = SplitPrescribedTransport( + eqn, transport_scheme, io, time_varying_velocity, + spatial_methods=transport_method, physics_schemes=physics_schemes + ) + else: + physics_parametrisation = [TerminatorToy(eqn, k1=k1, k2=k2, species1_name='Y', + species2_name='Y2')] + eqn.label_terms(lambda t: not t.has_label(time_derivative), implicit) + transport_scheme = IMEX_SSP3(domain) + time_varying_velocity = True + stepper = PrescribedTransport( + eqn, transport_scheme, io, time_varying_velocity, transport_method, + physics_parametrisations=physics_parametrisation + ) # Set up a non-divergent, time-varying, velocity field def u_t(t): @@ -75,39 +85,40 @@ def u_t(t): stepper.setup_prescribed_expr(u_t) - X_T_0 = 4e-6 - X_0 = X_T_0 + 0*lamda - X2_0 = 0*lamda + Y_T_0 = 4e-6 + Y_0 = Y_T_0 + 0*lamda + Y2_0 = 0*lamda - stepper.fields("X").interpolate(X_0) - stepper.fields("X2").interpolate(X2_0) + stepper.fields("Y").interpolate(Y_0) + stepper.fields("Y2").interpolate(Y2_0) stepper.run(t=0, tmax=10*dt) # Compute the steady state solution to compare to steady_space = domain.spaces('DG') - X_steady = Function(steady_space) - X2_steady = Function(steady_space) + Y_steady = Function(steady_space) + Y2_steady = Function(steady_space) r = k1/(4*k2) - D_val = sqrt(r**2 + 2*X_T_0*r) + D_val = sqrt(r**2 + 2*Y_T_0*r) - X_steady.interpolate(D_val - r) - X2_steady.interpolate(0.5*(X_T_0 - D_val + r)) + Y_steady.interpolate(D_val - r) + Y2_steady.interpolate(0.5*(Y_T_0 - D_val + r)) - return stepper, X_steady, X2_steady + return stepper, Y_steady, Y2_steady -def test_terminator_toy_setup(tmpdir): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_terminator_toy_setup(tmpdir, physics_coupling): dirname = str(tmpdir) - stepper, X_steady, X2_steady = run_terminator_toy(dirname) - X_field = stepper.fields("X") - X2_field = stepper.fields("X2") + stepper, Y_steady, Y2_steady = run_terminator_toy(dirname, physics_coupling) + Y_field = stepper.fields("Y") + Y2_field = stepper.fields("Y2") - print(errornorm(X_field, X_steady)/norm(X_steady)) - print(errornorm(X2_field, X2_steady)/norm(X2_steady)) + print(errornorm(Y_field, Y_steady)/norm(Y_steady)) + print(errornorm(Y2_field, Y2_steady)/norm(Y2_steady)) # Assert that the physics scheme has sufficiently moved # the species fields near their steady state solutions - assert errornorm(X_field, X_steady)/norm(X_steady) < 0.4, "The X field is not sufficiently close to the steady state profile" - assert errornorm(X2_field, X2_steady)/norm(X2_steady) < 0.4, "The X2 field is not sufficiently close to the steady state profile" + assert errornorm(Y_field, Y_steady)/norm(Y_steady) < 0.4, "The Y field is not sufficiently close to the steady state profile" + assert errornorm(Y2_field, Y2_steady)/norm(Y2_steady) < 0.4, "The Y2 field is not sufficiently close to the steady state profile" diff --git a/integration-tests/physics/test_wind_drag.py b/integration-tests/physics/test_wind_drag.py index cc70ef356..4214ac1e5 100644 --- a/integration-tests/physics/test_wind_drag.py +++ b/integration-tests/physics/test_wind_drag.py @@ -11,7 +11,7 @@ import pytest -def run_wind_drag(dirname, implicit_formulation): +def run_wind_drag(dirname, implicit_formulation, physics_coupling): # ------------------------------------------------------------------------ # # Set up model objects @@ -46,19 +46,29 @@ def run_wind_drag(dirname, implicit_formulation): surf_params = BoundaryLayerParameters() physics_parametrisation = WindDrag(eqn, implicit_formulation, surf_params) - time_discretisation = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) + if physics_coupling == "split": + time_discretisation = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) - # time_discretisation = ForwardEuler(domain) - physics_schemes = [(physics_parametrisation, time_discretisation)] + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] - # Only want time derivatives and physics terms in equation, so drop the rest - eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), - map_if_true=identity, map_if_false=drop) + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) - # Time stepper - scheme = ForwardEuler(domain) - stepper = SplitPhysicsTimestepper(eqn, scheme, io, - physics_schemes=physics_schemes) + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + else: + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) + stepper = Timestepper(eqn, scheme, io, + physics_parametrisations=[physics_parametrisation]) # ------------------------------------------------------------------------ # # Initial conditions @@ -101,10 +111,11 @@ def run_wind_drag(dirname, implicit_formulation): @pytest.mark.parametrize("implicit_formulation", [False, True]) -def test_wind_drag(tmpdir, implicit_formulation): +@pytest.mark.parametrize("physics_coupling", ["split", "nonsplit"]) +def test_wind_drag(tmpdir, implicit_formulation, physics_coupling): dirname = str(tmpdir) - mesh, stepper, u_true = run_wind_drag(dirname, implicit_formulation) + mesh, stepper, u_true = run_wind_drag(dirname, implicit_formulation, physics_coupling) u_final = stepper.fields('u') diff --git a/integration-tests/rexi/test_linear_sw.py b/integration-tests/rexi/test_linear_sw.py index c7a29aab9..77d5eafe2 100644 --- a/integration-tests/rexi/test_linear_sw.py +++ b/integration-tests/rexi/test_linear_sw.py @@ -57,9 +57,11 @@ def run_rexi_sw(tmpdir, ensemble=None): U_in = Function(eqns.function_space, name="U_in") Uexpl = Function(eqns.function_space, name="Uexpl") u, D = U_in.subfunctions + Dbar = eqns.X_ref.subfunctions[1] u.project(uexpr) D.interpolate(Dexpr) + Dbar.interpolate(H) if write_output: rexi_output.write(u, D) @@ -70,6 +72,7 @@ def run_rexi_sw(tmpdir, ensemble=None): uexpl, Dexpl = Uexpl.subfunctions u.assign(uexpl) D.assign(Dexpl) + if write_output: rexi_output.write(u, D) diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index 54c951cd7..9f6dd773f 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -15,6 +15,9 @@ def run(timestepper, tmax, f_end): return norm(timestepper.fields("f") - f_end) / norm(f_end) +# NB: The default vector transport test is not valid on the sphere as it is +# designed for a 3-component vector function space, and not the space of tangent +# vectors on the sphere @pytest.mark.parametrize("geometry", ["slice"]) def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): diff --git a/integration-tests/transport/test_vorticity_transport.py b/integration-tests/transport/test_vorticity_transport.py new file mode 100644 index 000000000..1cd9bdcbb --- /dev/null +++ b/integration-tests/transport/test_vorticity_transport.py @@ -0,0 +1,111 @@ +""" +This tests the transport of a vector-valued field using vorticity augmentation. +The computed solution is compared with a true one to check that the transport +is working correctly. +""" + +from gusto import * +from firedrake import ( + as_vector, norm, exp, PeriodicRectangleMesh, SpatialCoordinate, min_value +) +import pytest + + +def run(timestepper, tmax, f_end): + timestepper.run(0, tmax) + + return norm(timestepper.fields("f") - f_end) / norm(f_end) + + +@pytest.mark.parametrize("supg", [False, True]) +def test_vorticity_transport_setup(tmpdir, supg): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + Lx = 2000. # length of domain in x direction, in m + Ly = 2000. # width of domain in y direction, in m + ncells_1d = 20 # number of points in x and y directions + tmax = 500. + degree = 1 + + if supg: + # Smaller time steps for RK scheme + dt = 5.0 + dumpfreq = 50 + else: + dt = 25.0 + dumpfreq = 10 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + mesh = PeriodicRectangleMesh(ncells_1d, ncells_1d, Lx, Ly, quadrilateral=True) + domain = Domain(mesh, dt, "RTCF", degree) + x, y = SpatialCoordinate(mesh) + + Vu = domain.spaces("HDiv") + + # Equation + eqn = AdvectionEquation(domain, Vu, "f") + + # I/O + dirname = f'{tmpdir}/vorticity_plane' + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=False, dump_vtus=True + ) + io = IO(domain, output) + + augmentation = VorticityTransport(domain, eqn, supg=supg) + + # Make equation + if supg: + transport_scheme = SSPRK3( + domain, augmentation=augmentation, + rk_formulation=RungeKuttaFormulation.predictor + ) + else: + transport_scheme = TrapeziumRule(domain, augmentation=augmentation) + transport_method = DGUpwind(eqn, "f") + + time_varying_velocity = False + timestepper = PrescribedTransport( + eqn, transport_scheme, io, time_varying_velocity, transport_method + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + # Specify locations of the two Gaussians + xc = Lx/2. + xend = 3.*Lx/4. + yc = Ly/2. + + def l2_dist(xc, yc): + return min_value(abs(x - xc), Lx - abs(x - xc))**2 + (y - yc)**2 + + f0 = 1. + lc = 4.*Lx/25. + + init_scalar_expr = f0*exp(-l2_dist(xc, yc)/lc**2) + uexpr = as_vector([1.0, 0.0]) + + # Set fields + f = timestepper.fields("f") + f.project(as_vector([init_scalar_expr, init_scalar_expr])) + + u0 = timestepper.fields("u") + u0.project(uexpr) + + final_scalar_expr = f0*exp(-l2_dist(xend, yc)/lc**2) + final_vector_expr = as_vector([final_scalar_expr, final_scalar_expr]) + + # Run and check error + error = run(timestepper, tmax, final_vector_expr) + + tol = 1e-1 + assert error < tol, \ + 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_zero_limiter.py b/integration-tests/transport/test_zero_limiter.py index b4f2086a3..0604dd3d3 100644 --- a/integration-tests/transport/test_zero_limiter.py +++ b/integration-tests/transport/test_zero_limiter.py @@ -52,9 +52,10 @@ def setup_zero_limiter(dirname, limiter=False, rain=False): else: tracers = [WaterVapour(space='DG'), CloudWater(space='DG')] - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option='vector_advection_form', - thermal=True, active_tracers=tracers) + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, + u_transport_option='vector_advection_form', + active_tracers=tracers) output = OutputParameters(dirname=dirname, dumpfreq=1) @@ -66,14 +67,14 @@ def setup_zero_limiter(dirname, limiter=False, rain=False): # Saturation function def sat_func(x_in): - D = x_in.split()[1] - b = x_in.split()[2] + D = x_in.subfunctions[1] + b = x_in.subfunctions[2] return q0/(g*D) * exp(20*(1 - b/g)) # Feedback proportionality is dependent on h and b def gamma_v(x_in): - D = x_in.split()[1] - b = x_in.split()[2] + D = x_in.subfunctions[1] + b = x_in.subfunctions[2] return (1 + 10*(20*q0/g*D * exp(20*(1 - b/g))))**(-1) transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] diff --git a/plotting/compressible_euler/plot_schaer_mountain.py b/plotting/compressible_euler/plot_schaer_mountain.py new file mode 100644 index 000000000..f3d3f7e3d --- /dev/null +++ b/plotting/compressible_euler/plot_schaer_mountain.py @@ -0,0 +1,181 @@ +""" +Plots the Schär mountain test case. + +This plots: +(a) w @ t = 5 hr, (b) theta perturbation @ t = 5 hr +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +from netCDF4 import Dataset +import numpy as np +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, tomplot_field_title, tomplot_contours, + extract_gusto_coords, extract_gusto_field, reshape_gusto_data +) + +test = 'schaer_mountain' + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file +results_file_name = f'{abspath(dirname(__file__))}/../../results/{test}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/compressible_euler/{test}' + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['u_z', 'theta_perturbation', 'u_z', 'theta_perturbation'] +final_colour_schemes = ['PiYG', 'RdBu_r', 'PiYG', 'RdBu_r'] +final_field_labels = [ + r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)', + r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)' +] +final_contours = [ + np.linspace(-1.1, 1.1, 23), np.linspace(-1.4, 1.4, 15), + np.linspace(-1.1, 1.1, 23), np.linspace(-1.4, 1.4, 31) +] + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +initial_field_names = ['Exner', 'theta'] +initial_colour_schemes = ['PuBu', 'Reds'] +initial_field_labels = [r'$\Pi$', r'$\theta$ (K)'] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'contour' # Need to use this method to show mountains! +xlims = [0., 100.] +ylims = [0., 30.] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # + +fig, axarray = plt.subplots(1, 2, figsize=(18, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, colour_scheme, field_label) in \ + enumerate(zip( + axarray.flatten(), initial_field_names, initial_colour_schemes, + initial_field_labels + )): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + field_data, coords_X, coords_Y = \ + reshape_gusto_data(field_data, coords_X, coords_Y) + time = data_file['time'][time_idx] + + contours = tomplot_contours(field_data) + cmap, lines = tomplot_cmap(contours, colour_scheme) + + # Plot data ---------------------------------------------------------------- + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom') + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.suptitle(f't = {time:.1f} s') +fig.subplots_adjust(wspace=0.25) +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +xlims_zoom = [30., 70.] +ylims_zoom = [0., 12.] + +fig, axarray = plt.subplots(2, 2, figsize=(18, 12), sharex='row', sharey='row') +time_idx = 1 + +for i, (ax, field_name, colour_scheme, field_label, contours) in \ + enumerate(zip( + axarray.flatten(), final_field_names, final_colour_schemes, + final_field_labels, final_contours + )): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] + + field_data, coords_X, coords_Y = \ + reshape_gusto_data(field_data, coords_X, coords_Y) + + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=0.0) + + # Plot data ---------------------------------------------------------------- + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom') + + if i in [0, 1]: + # Only print min/max for top plots + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, minmax_format='.3f' + ) + + # Labels ------------------------------------------------------------------- + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + + if i in [0, 1]: + ax.set_xlim(xlims) + ax.set_ylim(ylims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + else: + ax.set_xlim(xlims_zoom) + ax.set_ylim(ylims_zoom) + ax.set_xticks(xlims_zoom) + ax.set_xticklabels(xlims_zoom) + + if i == 0: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + elif i == 2: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_yticks(ylims_zoom) + ax.set_yticklabels(ylims_zoom) + + +# Save figure ------------------------------------------------------------------ +fig.suptitle(f't = {time:.1f} s') +fig.subplots_adjust(wspace=0.25) +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/compressible_euler/plot_skamarock_klemp_nonhydrostatic.py b/plotting/compressible_euler/plot_skamarock_klemp_nonhydrostatic.py index 668d75dc8..566948b41 100644 --- a/plotting/compressible_euler/plot_skamarock_klemp_nonhydrostatic.py +++ b/plotting/compressible_euler/plot_skamarock_klemp_nonhydrostatic.py @@ -3,7 +3,7 @@ This plots the initial conditions @ t = 0 s, with (a) theta perturbation, (b) theta -and the final state @ t = 3600 s, with +and the final state @ t = 3000 s, with (a) theta perturbation, (b) a 1D slice through the wave """ @@ -11,6 +11,7 @@ import matplotlib.pyplot as plt import numpy as np from netCDF4 import Dataset +import pandas as pd from tomplot import ( set_tomplot_style, tomplot_cmap, plot_contoured_field, add_colorbar_ax, tomplot_field_title, extract_gusto_coords, @@ -49,7 +50,6 @@ # General options # ---------------------------------------------------------------------------- # contour_method = 'tricontour' -xlims = [0, 300.0] ylims = [0, 10.0] # Things that are likely the same for all plots -------------------------------- @@ -59,6 +59,8 @@ # ---------------------------------------------------------------------------- # # INITIAL PLOTTING # ---------------------------------------------------------------------------- # +xlims = [0, 300.0] + fig, axarray = plt.subplots(1, 2, figsize=(12, 6), sharex='all', sharey='all') time_idx = 0 @@ -107,6 +109,9 @@ # ---------------------------------------------------------------------------- # # FINAL PLOTTING # ---------------------------------------------------------------------------- # +x_offset = -3000.0*20/1000.0 +xlims = [-x_offset, 300.0-x_offset] + fig, axarray = plt.subplots(2, 1, figsize=(8, 8), sharex='all') time_idx = -1 @@ -115,6 +120,21 @@ coords_X, coords_Y = extract_gusto_coords(data_file, final_field_name) time = data_file['time'][time_idx] +# Wave has wrapped around periodic boundary, so shift the coordinates +coords_X = np.where(coords_X < xlims[0], coords_X + 300.0, coords_X) + +# Sort data given the change in coordinates +data_dict = { + 'X': coords_X, + 'Y': coords_Y, + 'field': field_data +} +data_frame = pd.DataFrame(data_dict) +data_frame.sort_values(by=['X', 'Y'], inplace=True) +coords_X = data_frame['X'].values[:] +coords_Y = data_frame['Y'].values[:] +field_data = data_frame['field'].values[:] + # Plot 2D data ----------------------------------------------------------------- ax = axarray[0] diff --git a/plotting/shallow_water/plot_linear_thermal_galewsky_jet.py b/plotting/shallow_water/plot_linear_thermal_galewsky_jet.py new file mode 100644 index 000000000..de4d446a9 --- /dev/null +++ b/plotting/shallow_water/plot_linear_thermal_galewsky_jet.py @@ -0,0 +1,173 @@ +""" +Plots the linear thermal Galewsky jet test case. + """ +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file + +results_file_name = f'{abspath(dirname(__file__))}/../../results/linear_thermal_galewsky/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/linear_thermal_galewsky' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'RelativeVorticity', 'b'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'RdBu_r', 'PuRd_r'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', + r'$\zeta$ (s$^{-1})$', r'$b$ (m s$^{-2}$)'] +init_contours_to_remove = [None, None, 0.0, None] +init_contours = [np.linspace(0.0, 80.0, 9), + np.linspace(8900.0, 10200.0+1e-3, 12), + np.linspace(-2e-4, 2e-4, 17), + np.linspace(8.8, 9.8, 11)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'b'] +final_colour_schemes = ['RdBu_r', 'PuRd_r'] +final_field_labels = [r'$\zeta$ (s$^{-1}$)', r'$b$ (m s$^{-2}$)'] +final_contours_to_remove = [0.0, None] +final_contours = [np.linspace(-2e-4, 2e-4, 17), + np.linspace(8.8, 9.8, 11)] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [-180, 180] +ylims = [10, 80] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(2, 2, figsize=(16, 12), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray.flatten(), init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + lon_1d = np.linspace(-180.0, 180.0, 91) + lat_1d = np.linspace(-90.0, 90.0, 81) + lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=(12, 1), magnitude_filter=1.0, + ) + + # Labels ------------------------------------------------------------------- + if i in [0, 2]: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + if i in [2, 3]: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(16, 8), sharex='all', sharey='all') +time_idx = -1 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray, final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove, final_contours)): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_moist_thermal_equivb_gravity_wave.py b/plotting/shallow_water/plot_moist_thermal_equivb_gravity_wave.py new file mode 100644 index 000000000..b771cc89b --- /dev/null +++ b/plotting/shallow_water/plot_moist_thermal_equivb_gravity_wave.py @@ -0,0 +1,187 @@ +""" +Plots the moist thermal gravity wave test case. +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file + +results_file_name = f'{abspath(dirname(__file__))}/../../results/moist_thermal_equivb_gw/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/moist_thermal_equivb_gw' + +beta2 = 9.80616*10 + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'PartitionedCloud', 'b_e'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'cividis', 'PuRd_r'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', + r'$q_{cl}$ (kg kg$^{-1})$', r'$b_e$ (m s$^{-2}$)'] +init_contours_to_remove = [None, None, None, None] +init_contours = [np.linspace(0.0, 20.0, 9), + np.linspace(4800.0, 8000.0, 19), + np.linspace(0.01, 0.02, 11), + np.linspace(9, 10, 15)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['PartitionedCloud', 'b_e'] +final_colour_schemes = ['cividis', 'PuRd_r'] +final_field_labels = [r'$q_{cl}$ (kg kg$^{-1}$)', r'$b_e$ (m s$^{-2}$)'] +final_contours_to_remove = [None, None] +final_contours = [np.linspace(0.01, 0.02, 11), + np.linspace(9, 10, 15)] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [-180, 180] +ylims = [-90, 90] +minmax_format = { + 'PartitionedCloud': '.1e', + 'b_e': '.2f', + 'u': '.1f', + 'D': '.0f' +} + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(2, 2, figsize=(16, 12), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray.flatten(), init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, + minmax_format=minmax_format[field_name] + ) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + lon_1d = np.linspace(-180.0, 180.0, 91) + lat_1d = np.linspace(-90.0, 90.0, 81) + lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + ) + + # Labels ------------------------------------------------------------------- + if i in [0, 2]: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + if i in [2, 3]: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(16, 8), sharex='all', sharey='all') +time_idx = -1 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray, final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove, final_contours)): + + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, + minmax_format=minmax_format[field_name] + ) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/pyproject.toml b/pyproject.toml index f226c0660..031d28e95 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,17 +24,5 @@ classifiers = [ Homepage = "http://www.firedrakeproject.org/gusto/" Repository = "https://github.com/firedrakeproject/gusto.git" -[tool.setuptools] -packages = [ - "gusto", - "gusto.core", - "gusto.diagnostics", - "gusto.equations", - "gusto.initialisation", - "gusto.physics", - "gusto.recovery", - "gusto.solvers", - "gusto.spatial_methods", - "gusto.time_discretisation", - "gusto.timestepping", -] +[tool.setuptools.packages.find] +include = ["gusto"] diff --git a/unit-tests/diagnostic_tests/test_cumulative_sum.py b/unit-tests/diagnostic_tests/test_cumulative_sum.py new file mode 100644 index 000000000..295e522ca --- /dev/null +++ b/unit-tests/diagnostic_tests/test_cumulative_sum.py @@ -0,0 +1,48 @@ + +from gusto.diagnostics import CumulativeSum +from gusto.core.fields import StateFields, PrescribedFields, TimeLevelFields +from gusto import (Domain, CompressibleParameters, CompressibleEulerEquations, + WaterVapour) +from firedrake import PeriodicIntervalMesh, ExtrudedMesh +import numpy as np + + +def test_dewpoint(): + + L = 10 + H = 10 + ncol = 3 + nlayers = 3 + + m = PeriodicIntervalMesh(ncol, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + + domain = Domain(mesh, 0.1, 'CG', 1) + params = CompressibleParameters() + active_tracers = [WaterVapour()] + eqn = CompressibleEulerEquations(domain, params, active_tracers=active_tracers) + prog_fields = TimeLevelFields(eqn) + + Vtheta = domain.spaces('theta') + + # Setting up prognostic fields for the diagnostic to use + prescribed_fields = PrescribedFields() + state_fields = StateFields(prog_fields, prescribed_fields) + + theta = state_fields('theta', Vtheta) + + # Initial conditions + theta.interpolate(300.0) + + diagnostic = CumulativeSum(name='theta') + diagnostic.setup(domain, state_fields) + diagnostic.compute() + + assert np.allclose(diagnostic.field.dat.data, 300.0, atol=0.0), \ + 'The cumulative sum diagnostic does not seem to be correct after 1 iteration' + + theta.interpolate(250.0) + diagnostic.compute() + + assert np.allclose(diagnostic.field.dat.data, 550.0, atol=0.0), \ + 'The cumulative sum diagnostic does not seem to be correct after 2 iterations' diff --git a/unit-tests/test_function_spaces.py b/unit-tests/test_function_spaces.py index b39aa8cc5..bcf105e35 100644 --- a/unit-tests/test_function_spaces.py +++ b/unit-tests/test_function_spaces.py @@ -100,6 +100,27 @@ def test_de_rham_spaces(domain, family): assert elt.degree() == degree, '"L2" space does not seem to be degree ' \ + f'{degree}. Found degree {elt.degree()}' + # Check that continuities have been recorded correctly + if hasattr(mesh, "_base_mesh"): + expected_continuity = { + "H1": {'horizontal': True, 'vertical': True}, + "L2": {'horizontal': False, 'vertical': False}, + "HDiv": {'horizontal': True, 'vertical': True}, + "HCurl": {'horizontal': True, 'vertical': True}, + "theta": {'horizontal': False, 'vertical': True} + } + else: + expected_continuity = { + "H1": True, + "L2": False, + "HDiv": True, + "HCurl": True + } + + for space, continuity in expected_continuity.items(): + if space in spaces.continuity: + assert spaces.continuity[space] == continuity + # ---------------------------------------------------------------------------- # # Test creation of DG1 equispaced @@ -118,6 +139,13 @@ def test_dg_equispaced(domain, family): assert elt.variant() == "equispaced", '"DG1 equispaced" does not seem to ' \ + f'be equispaced variant. Found variant {elt.variant()}' + if hasattr(mesh, "_base_mesh"): + expected_continuity = {'horizontal': False, 'vertical': False} + else: + expected_continuity = False + + assert spaces.continuity['DG1_equispaced'] == expected_continuity, "DG is discontinuous" + # ---------------------------------------------------------------------------- # # Test creation of DG0 space @@ -133,6 +161,13 @@ def test_dg0(domain, family): assert elt.degree() in [0, (0, 0)], '"DG0" space does not seem to be' \ + f'degree 0. Found degree {elt.degree()}' + if hasattr(mesh, "_base_mesh"): + expected_continuity = {'horizontal': False, 'vertical': False} + else: + expected_continuity = False + + assert spaces.continuity['DG0'] == expected_continuity, "DG is discontinuous" + # ---------------------------------------------------------------------------- # # Test creation of a general CG space @@ -150,3 +185,10 @@ def test_cg(domain, family): assert elt.degree() == degree or elt.degree() == (degree, degree), \ (f'"CG" space does not seem to be degree {degree}. ' + f'Found degree {elt.degree()}') + + if hasattr(mesh, "_base_mesh"): + expected_continuity = {'horizontal': True, 'vertical': True} + else: + expected_continuity = True + + assert spaces.continuity['CG'] == expected_continuity, "CG is continuous"