diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 117257e3c..c86631274 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -38,6 +38,7 @@ jobs: . /home/firedrake/firedrake/bin/activate firedrake-clean export GUSTO_PARALLEL_LOG=FILE + export PYOP2_CFLAGS=-O0 python -m pytest \ -n 12 --dist worksteal \ --durations=100 \ diff --git a/Makefile b/Makefile index 4ffd2902f..2b50a2dcb 100644 --- a/Makefile +++ b/Makefile @@ -32,4 +32,8 @@ integration_test: example: @echo " Running all examples" - @python3 -m pytest examples $(PYTEST_ARGS) + @python3 -m pytest examples -v -m "not parallel" $(PYTEST_ARGS) + +parallel_example: + @echo " Running all parallel examples" + @python3 -m pytest examples -v -m "parallel" $(PYTEST_ARGS) diff --git a/examples/boussinesq/skamarock_klemp_compressible.py b/examples/boussinesq/skamarock_klemp_compressible.py index b6e6d3bf7..c15daff84 100644 --- a/examples/boussinesq/skamarock_klemp_compressible.py +++ b/examples/boussinesq/skamarock_klemp_compressible.py @@ -1,114 +1,194 @@ """ -The gravity wave test case of Skamarock and Klemp (1994), solved using the -incompressible Boussinesq equations. +This example uses the compressible Boussinesq 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. -Buoyancy is transported using SUPG. +Buoyancy is transported using SUPG, and the degree 1 elements are used. """ -from gusto import * -from firedrake import (as_vector, PeriodicIntervalMesh, ExtrudedMesh, - sin, SpatialCoordinate, Function, pi) -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 6. -L = 3.0e5 # Domain length -H = 1.0e4 # Height position of the model top - -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 - columns = 30 # number of columns - nlayers = 5 # horizontal layers - -else: - tmax = 3600. - dumpfreq = int(tmax / (2*dt)) - columns = 300 # number of columns - nlayers = 10 # horizontal layers - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -m = PeriodicIntervalMesh(columns, L) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -domain = Domain(mesh, dt, 'CG', 1) - -# Equation -parameters = BoussinesqParameters(cs=300) -eqns = BoussinesqEquations(domain, parameters) - -# I/O -output = OutputParameters( - dirname='skamarock_klemp_compressible', - dumpfreq=dumpfreq, - dumplist=['u'], +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + as_vector, PeriodicIntervalMesh, ExtrudedMesh, sin, SpatialCoordinate, + Function, pi ) -# list of diagnostic fields, each defined in a class in diagnostics.py -diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - -# Transport schemes -b_opts = SUPGOptions() -transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "p"), - SSPRK3(domain, "b", options=b_opts)] -transport_methods = [DGUpwind(eqns, "u"), - DGUpwind(eqns, "p"), - DGUpwind(eqns, "b", ibp=b_opts.ibp)] - -# Linear solver -linear_solver = BoussinesqSolver(eqns) - -# Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, - transport_methods, - linear_solver=linear_solver) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, Divergence, Perturbation, CourantNumber, + BoussinesqParameters, BoussinesqEquations, BoussinesqSolver, + boussinesq_hydrostatic_balance +) + +skamarock_klemp_compressible_bouss_defaults = { + 'ncolumns': 300, + 'nlayers': 10, + 'dt': 6.0, + 'tmax': 3600., + 'dumpfreq': 300, + 'dirname': 'skamarock_klemp_compressible_bouss' +} + + +def skamarock_klemp_compressible_bouss( + ncolumns=skamarock_klemp_compressible_bouss_defaults['ncolumns'], + nlayers=skamarock_klemp_compressible_bouss_defaults['nlayers'], + dt=skamarock_klemp_compressible_bouss_defaults['dt'], + tmax=skamarock_klemp_compressible_bouss_defaults['tmax'], + dumpfreq=skamarock_klemp_compressible_bouss_defaults['dumpfreq'], + dirname=skamarock_klemp_compressible_bouss_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 3.0e5 # Width of domain (m) + domain_height = 1.0e4 # Height of domain (m) + wind_initial = 20. # Initial wind in x direction (m/s) + pert_width = 5.0e3 # Width parameter of perturbation (m) + deltab = 1.0e-2 # Magnitude of buoyancy perturbation (m/s^2) + N = 0.01 # Brunt-Vaisala frequency (1/s) + cs = 300. # Speed of sound (m/s) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, 'CG', element_order) + + # Equation + parameters = BoussinesqParameters(cs=cs) + eqns = BoussinesqEquations(domain, parameters) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=False, + ) + # list of diagnostic fields, each defined in a class in diagnostics.py + diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + b_opts = SUPGOptions() + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "p"), + SSPRK3(domain, "b", options=b_opts) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "p"), + DGUpwind(eqns, "b", ibp=b_opts.ibp) + ] + + # Linear solver + linear_solver = BoussinesqSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + b0 = stepper.fields("b") + p0 = stepper.fields("p") + + # spaces + Vb = b0.function_space() + Vp = p0.function_space() + + x, z = SpatialCoordinate(mesh) + + # first setup the background buoyancy profile + # z.grad(bref) = N**2 + bref = z*(N**2) + # interpolate the expression to the function + b_b = Function(Vb).interpolate(bref) + + # setup constants + b_pert = ( + deltab * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + # interpolate the expression to the function + b0.interpolate(b_b + b_pert) + + p_b = Function(Vp) + boussinesq_hydrostatic_balance(eqns, b_b, p_b) + p0.assign(p_b) + + uinit = (as_vector([wind_initial, 0.0])) + u0.project(uinit) + + # set the background buoyancy + stepper.set_reference_profiles([('p', p_b), ('b', b_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # -# Initial conditions +# MAIN # ---------------------------------------------------------------------------- # -u0 = stepper.fields("u") -b0 = stepper.fields("b") -p0 = stepper.fields("p") - -# spaces -Vb = b0.function_space() -Vp = p0.function_space() - -x, z = SpatialCoordinate(mesh) - -# first setup the background buoyancy profile -# z.grad(bref) = N**2 -N = parameters.N -bref = z*(N**2) -# interpolate the expression to the function -b_b = Function(Vb).interpolate(bref) - -# setup constants -a = 5.0e3 -deltab = 1.0e-2 -b_pert = deltab*sin(pi*z/H)/(1 + (x - L/2)**2/a**2) -# interpolate the expression to the function -b0.interpolate(b_b + b_pert) -p_b = Function(Vp) -boussinesq_hydrostatic_balance(eqns, b_b, p_b) -p0.assign(p_b) - -uinit = (as_vector([20.0, 0.0])) -u0.project(uinit) - -# set the background buoyancy -stepper.set_reference_profiles([('p', p_b), ('b', b_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # -stepper.run(t=0, tmax=tmax) +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_compressible_bouss_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=skamarock_klemp_compressible_bouss_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=skamarock_klemp_compressible_bouss_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=skamarock_klemp_compressible_bouss_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=skamarock_klemp_compressible_bouss_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=skamarock_klemp_compressible_bouss_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + skamarock_klemp_compressible_bouss(**vars(args)) diff --git a/examples/boussinesq/skamarock_klemp_incompressible.py b/examples/boussinesq/skamarock_klemp_incompressible.py index b71e89830..b7bb5fd5e 100644 --- a/examples/boussinesq/skamarock_klemp_incompressible.py +++ b/examples/boussinesq/skamarock_klemp_incompressible.py @@ -1,110 +1,190 @@ """ -The gravity wave test case of Skamarock and Klemp (1994), solved using the -incompressible Boussinesq equations. +This example uses the incompressible Boussinesq 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. -Buoyancy is transported using SUPG. +Buoyancy is transported using SUPG, and the degree 1 elements are used. """ -from gusto import * -from firedrake import (as_vector, PeriodicIntervalMesh, ExtrudedMesh, - sin, SpatialCoordinate, Function, pi) -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 6. -L = 3.0e5 # Domain length -H = 1.0e4 # Height position of the model top - -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 - columns = 30 # number of columns - nlayers = 5 # horizontal layers - -else: - tmax = 3600. - dumpfreq = int(tmax / (2*dt)) - columns = 300 # number of columns - nlayers = 10 # horizontal layers - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -m = PeriodicIntervalMesh(columns, L) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -domain = Domain(mesh, dt, 'CG', 1) - -# Equation -parameters = BoussinesqParameters() -eqns = BoussinesqEquations(domain, parameters, compressible=False) - -# I/O -output = OutputParameters( - dirname='skamarock_klemp_incompressible', - dumpfreq=dumpfreq, - dumplist=['u'], +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + as_vector, PeriodicIntervalMesh, ExtrudedMesh, sin, SpatialCoordinate, + Function, pi +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, Divergence, Perturbation, CourantNumber, + BoussinesqParameters, BoussinesqEquations, BoussinesqSolver, + boussinesq_hydrostatic_balance ) -# list of diagnostic fields, each defined in a class in diagnostics.py -diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - -# Transport schemes -b_opts = SUPGOptions() -transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "b", options=b_opts)] -transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "b", ibp=b_opts.ibp)] - -# Linear solver -linear_solver = BoussinesqSolver(eqns) -# Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, - transport_methods, - linear_solver=linear_solver) +skamarock_klemp_incompressible_bouss_defaults = { + 'ncolumns': 300, + 'nlayers': 10, + 'dt': 6.0, + 'tmax': 3600., + 'dumpfreq': 300, + 'dirname': 'skamarock_klemp_incompressible_bouss' +} + + +def skamarock_klemp_incompressible_bouss( + ncolumns=skamarock_klemp_incompressible_bouss_defaults['ncolumns'], + nlayers=skamarock_klemp_incompressible_bouss_defaults['nlayers'], + dt=skamarock_klemp_incompressible_bouss_defaults['dt'], + tmax=skamarock_klemp_incompressible_bouss_defaults['tmax'], + dumpfreq=skamarock_klemp_incompressible_bouss_defaults['dumpfreq'], + dirname=skamarock_klemp_incompressible_bouss_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 3.0e5 # Width of domain (m) + domain_height = 1.0e4 # Height of domain (m) + wind_initial = 20. # Initial wind in x direction (m/s) + pert_width = 5.0e3 # Width parameter of perturbation (m) + deltab = 1.0e-2 # Magnitude of buoyancy perturbation (m/s^2) + N = 0.01 # Brunt-Vaisala frequency (1/s) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, 'CG', element_order) + + # Equation + parameters = BoussinesqParameters() + eqns = BoussinesqEquations(domain, parameters, compressible=False) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=True, + ) + # list of diagnostic fields, each defined in a class in diagnostics.py + diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + b_opts = SUPGOptions() + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "b", options=b_opts) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "b", ibp=b_opts.ibp) + ] + + # Linear solver + linear_solver = BoussinesqSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + b0 = stepper.fields("b") + p0 = stepper.fields("p") + + # spaces + Vb = b0.function_space() + + x, z = SpatialCoordinate(mesh) + + # first setup the background buoyancy profile + # z.grad(bref) = N**2 + bref = z*(N**2) + # interpolate the expression to the function + b_b = Function(Vb).interpolate(bref) + + # setup constants + b_pert = ( + deltab * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + # interpolate the expression to the function + b0.interpolate(b_b + b_pert) + + boussinesq_hydrostatic_balance(eqns, b_b, p0) + + uinit = (as_vector([wind_initial, 0.0])) + u0.project(uinit) + + # set the background buoyancy + stepper.set_reference_profiles([('b', b_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + # Run! + stepper.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # -# Initial conditions +# MAIN # ---------------------------------------------------------------------------- # -u0 = stepper.fields("u") -b0 = stepper.fields("b") -p0 = stepper.fields("p") - -# spaces -Vb = b0.function_space() - -x, z = SpatialCoordinate(mesh) - -# first setup the background buoyancy profile -# z.grad(bref) = N**2 -N = parameters.N -bref = z*(N**2) -# interpolate the expression to the function -b_b = Function(Vb).interpolate(bref) - -# setup constants -a = 5.0e3 -deltab = 1.0e-2 -b_pert = deltab*sin(pi*z/H)/(1 + (x - L/2)**2/a**2) -# interpolate the expression to the function -b0.interpolate(b_b + b_pert) - -boussinesq_hydrostatic_balance(eqns, b_b, p0) - -uinit = (as_vector([20.0, 0.0])) -u0.project(uinit) - -# set the background buoyancy -stepper.set_reference_profiles([('b', b_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # -# Run! -stepper.run(t=0, tmax=tmax) +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_incompressible_bouss_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=skamarock_klemp_incompressible_bouss_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=skamarock_klemp_incompressible_bouss_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=skamarock_klemp_incompressible_bouss_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=skamarock_klemp_incompressible_bouss_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=skamarock_klemp_incompressible_bouss_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + skamarock_klemp_incompressible_bouss(**vars(args)) diff --git a/examples/boussinesq/skamarock_klemp_linear.py b/examples/boussinesq/skamarock_klemp_linear.py index bb64c9b3a..e8145548d 100644 --- a/examples/boussinesq/skamarock_klemp_linear.py +++ b/examples/boussinesq/skamarock_klemp_linear.py @@ -1,98 +1,177 @@ """ -The gravity wave test case of Skamarock and Klemp (1994), solved using the -incompressible Boussinesq equations. +This example uses the linear Boussinesq 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. -Buoyancy is transported using SUPG. +The degree 1 elements are used, with an explicit RK4 time stepper. """ -from gusto import * -from firedrake import (PeriodicIntervalMesh, ExtrudedMesh, - sin, SpatialCoordinate, Function, pi) -import sys +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + PeriodicIntervalMesh, ExtrudedMesh, sin, SpatialCoordinate, Function, pi +) +from gusto import ( + Domain, IO, OutputParameters, RK4, DGUpwind, SUPGOptions, Divergence, + Timestepper, Perturbation, CourantNumber, BoussinesqParameters, + LinearBoussinesqEquations, boussinesq_hydrostatic_balance +) + +skamarock_klemp_linear_bouss_defaults = { + 'ncolumns': 300, + 'nlayers': 10, + 'dt': 0.5, + 'tmax': 3600., + 'dumpfreq': 3600, + 'dirname': 'skamarock_klemp_linear_bouss' +} + + +def skamarock_klemp_linear_bouss( + ncolumns=skamarock_klemp_linear_bouss_defaults['ncolumns'], + nlayers=skamarock_klemp_linear_bouss_defaults['nlayers'], + dt=skamarock_klemp_linear_bouss_defaults['dt'], + tmax=skamarock_klemp_linear_bouss_defaults['tmax'], + dumpfreq=skamarock_klemp_linear_bouss_defaults['dumpfreq'], + dirname=skamarock_klemp_linear_bouss_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 3.0e5 # Width of domain (m) + domain_height = 1.0e4 # Height of domain (m) + pert_width = 5.0e3 # Width parameter of perturbation (m) + deltab = 1.0e-2 # Magnitude of buoyancy perturbation (m/s^2) + N = 0.01 # Brunt-Vaisala frequency (1/s) + cs = 300. # Speed of sound (m/s) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, 'CG', element_order) + + # Equation + parameters = BoussinesqParameters(cs=cs) + eqns = LinearBoussinesqEquations(domain, parameters) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=True, + ) + # list of diagnostic fields, each defined in a class in diagnostics.py + diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + b_opts = SUPGOptions() + transport_methods = [ + DGUpwind(eqns, "p"), + DGUpwind(eqns, "b", ibp=b_opts.ibp) + ] + + # Time stepper + stepper = Timestepper( + eqns, RK4(domain), io, spatial_methods=transport_methods + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + b0 = stepper.fields("b") + p0 = stepper.fields("p") + + # spaces + Vb = b0.function_space() + Vp = p0.function_space() + + x, z = SpatialCoordinate(mesh) + + # first setup the background buoyancy profile + # z.grad(bref) = N**2 + bref = z*(N**2) + # interpolate the expression to the function + b_b = Function(Vb).interpolate(bref) + + # setup constants + b_pert = ( + deltab * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + # interpolate the expression to the function + b0.interpolate(b_b + b_pert) + + p_b = Function(Vp) + boussinesq_hydrostatic_balance(eqns, b_b, p_b) + p0.assign(p_b) + + # set the background buoyancy + stepper.set_reference_profiles([('p', p_b), ('b', b_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # -# Test case parameters +# MAIN # ---------------------------------------------------------------------------- # -dt = 0.5 -L = 3.0e5 # Domain length -H = 1.0e4 # Height position of the model top -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 - columns = 30 # number of columns - nlayers = 5 # horizontal layers - -else: - tmax = 3600. - dumpfreq = int(tmax / (2*dt)) - columns = 300 # number of columns - nlayers = 10 # horizontal layers - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -m = PeriodicIntervalMesh(columns, L) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -domain = Domain(mesh, dt, 'CG', 1) - -# Equation -parameters = BoussinesqParameters(cs=300) -eqns = LinearBoussinesqEquations(domain, parameters) - -# I/O -output = OutputParameters(dirname='skamarock_klemp_linear') -# list of diagnostic fields, each defined in a class in diagnostics.py -diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - -# Transport schemes -b_opts = SUPGOptions() -transport_methods = [DGUpwind(eqns, "p"), - DGUpwind(eqns, "b", ibp=b_opts.ibp)] - - -# Time stepper -stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) - -# ---------------------------------------------------------------------------- # -# Initial conditions -# ---------------------------------------------------------------------------- # - -b0 = stepper.fields("b") -p0 = stepper.fields("p") - -# spaces -Vb = b0.function_space() -Vp = p0.function_space() - -x, z = SpatialCoordinate(mesh) - -# first setup the background buoyancy profile -# z.grad(bref) = N**2 -N = parameters.N -bref = z*(N**2) -# interpolate the expression to the function -b_b = Function(Vb).interpolate(bref) - -# setup constants -a = 5.0e3 -deltab = 1.0e-2 -b_pert = deltab*sin(pi*z/H)/(1 + (x - L/2)**2/a**2) -# interpolate the expression to the function -b0.interpolate(b_b + b_pert) - -p_b = Function(Vp) -boussinesq_hydrostatic_balance(eqns, b_b, p_b) -p0.assign(p_b) - -# set the background buoyancy -stepper.set_reference_profiles([('p', p_b), ('b', b_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # -stepper.run(t=0, tmax=tmax) +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_linear_bouss_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=skamarock_klemp_linear_bouss_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=skamarock_klemp_linear_bouss_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=skamarock_klemp_linear_bouss_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=skamarock_klemp_linear_bouss_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=skamarock_klemp_linear_bouss_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + skamarock_klemp_linear_bouss(**vars(args)) diff --git a/examples/boussinesq/test_boussinesq_examples.py b/examples/boussinesq/test_boussinesq_examples.py new file mode 100644 index 000000000..53d3d479d --- /dev/null +++ b/examples/boussinesq/test_boussinesq_examples.py @@ -0,0 +1,64 @@ +import pytest + + +def make_dirname(test_name): + from mpi4py import MPI + comm = MPI.COMM_WORLD + if comm.size > 1: + return f'pytest_{test_name}_parallel' + else: + return f'pytest_{test_name}' + + +def test_skamarock_klemp_compressible_bouss(): + from skamarock_klemp_compressible import skamarock_klemp_compressible_bouss + test_name = 'skamarock_klemp_compressible_bouss' + skamarock_klemp_compressible_bouss( + ncolumns=30, + nlayers=5, + dt=6.0, + tmax=60.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_skamarock_klemp_compressible_bouss_parallel(): + test_skamarock_klemp_compressible_bouss() + + +def test_skamarock_klemp_incompressible_bouss(): + from skamarock_klemp_incompressible import skamarock_klemp_incompressible_bouss + test_name = 'skamarock_klemp_incompressible_bouss' + skamarock_klemp_incompressible_bouss( + ncolumns=30, + nlayers=5, + dt=60.0, + tmax=12.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_skamarock_klemp_incompressible_bouss_parallel(): + test_skamarock_klemp_incompressible_bouss() + + +def test_skamarock_klemp_linear_bouss(): + from skamarock_klemp_linear import skamarock_klemp_linear_bouss + test_name = 'skamarock_klemp_linear_bouss' + skamarock_klemp_linear_bouss( + ncolumns=30, + nlayers=5, + dt=60.0, + tmax=12.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_skamarock_klemp_linear_bouss_parallel(): + test_skamarock_klemp_linear_bouss() diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py deleted file mode 100644 index 06701a0ba..000000000 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ /dev/null @@ -1,166 +0,0 @@ -""" -The non-orographic gravity wave test case (3-1) from the DCMIP test case -document of Ullrich et al (2012). - -This uses a cubed-sphere mesh. -""" - -from gusto import * -from firedrake import (CubedSphereMesh, ExtrudedMesh, FunctionSpace, - Function, SpatialCoordinate, as_vector) -from firedrake import exp, acos, cos, sin, pi, sqrt, asin, atan2 -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 100.0 # Time-step size (s) - -if '--running-tests' in sys.argv: - nlayers = 4 # Number of vertical layers - refinements = 2 # Number of horiz. refinements - tmax = dt - dumpfreq = 1 -else: - nlayers = 10 # Number of vertical layers - refinements = 3 # Number of horiz. refinements - tmax = 3600.0 - dumpfreq = int(tmax / (4*dt)) - -parameters = CompressibleParameters() -a_ref = 6.37122e6 # Radius of the Earth (m) -X = 125.0 # Reduced-size Earth reduction factor -a = a_ref/X # Scaled radius of planet (m) -g = parameters.g # Acceleration due to gravity (m/s^2) -N = parameters.N # Brunt-Vaisala frequency (1/s) -p_0 = parameters.p_0 # Reference pressure (Pa, not hPa) -c_p = parameters.cp # SHC of dry air at constant pressure (J/kg/K) -R_d = parameters.R_d # Gas constant for dry air (J/kg/K) -kappa = parameters.kappa # R_d/c_p -T_eq = 300.0 # Isothermal atmospheric temperature (K) -p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator -u_max = 20.0 # Maximum amplitude of the zonal wind (m/s) -d = 5000.0 # Width parameter for Theta' -lamda_c = 2.0*pi/3.0 # Longitudinal centerpoint of Theta' -phi_c = 0.0 # Latitudinal centerpoint of Theta' (equator) -deltaTheta = 1.0 # Maximum amplitude of Theta' (K) -L_z = 20000.0 # Vertical wave length of the Theta' perturb. -z_top = 1.0e4 # Height position of the model top - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -# Cubed-sphere horizontal mesh -m = CubedSphereMesh(radius=a, - refinement_level=refinements, - degree=2) -# Build volume mesh -mesh = ExtrudedMesh(m, layers=nlayers, - layer_height=z_top/nlayers, - extrusion_type="radial") -domain = Domain(mesh, dt, "RTCF", 1) -x = SpatialCoordinate(mesh) - -# Create polar coordinates: -# Since we use a CG1 field, this is constant on layers -W_Q1 = FunctionSpace(mesh, "CG", 1) -z_expr = sqrt(x[0]*x[0] + x[1]*x[1] + x[2]*x[2]) - a -z = Function(W_Q1).interpolate(z_expr) -lat_expr = asin(x[2]/sqrt(x[0]*x[0] + x[1]*x[1] + x[2]*x[2])) -lat = Function(W_Q1).interpolate(lat_expr) -lon = Function(W_Q1).interpolate(atan2(x[1], x[0])) - -# Equation -eqns = CompressibleEulerEquations(domain, parameters) - -# I/O -dirname = 'dcmip_3_1_meanflow' -output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, -) -diagnostic_fields = [Perturbation('theta'), Perturbation('rho'), - CompressibleKineticEnergy(), PotentialEnergy(eqns)] - -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - -# Transport schemes -transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho", fixed_subcycles=2), - SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2)] -transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta"]] - -# 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') -theta0 = stepper.fields('theta') -rho0 = stepper.fields('rho') - -# spaces -Vu = domain.spaces("HDiv") -Vt = domain.spaces("theta") -Vr = domain.spaces("DG") - -# Initial conditions with u0 -uexpr = as_vector([-u_max*x[1]/a, u_max*x[0]/a, 0.0]) -u0.project(uexpr) - -# Surface temperature -G = g**2/(N**2*c_p) -Ts_expr = G + (T_eq-G)*exp(-(u_max*N**2/(4*g*g))*u_max*(cos(2.0*lat)-1.0)) -Ts = Function(W_Q1).interpolate(Ts_expr) - -# Surface pressure -ps_expr = p_eq*exp((u_max/(4.0*G*R_d))*u_max*(cos(2.0*lat)-1.0))*(Ts/T_eq)**(1.0/kappa) -ps = Function(W_Q1).interpolate(ps_expr) - -# Background pressure -p_expr = ps*(1 + G/Ts*(exp(-N**2*z/g)-1))**(1.0/kappa) -p = Function(W_Q1).interpolate(p_expr) - -# Background temperature -Tb_expr = G*(1 - exp(N**2*z/g)) + Ts*exp(N**2*z/g) -Tb = Function(W_Q1).interpolate(Tb_expr) - -# Background potential temperature -thetab_expr = Tb*(p_0/p)**kappa -thetab = Function(W_Q1).interpolate(thetab_expr) -theta_b = Function(theta0.function_space()).interpolate(thetab) -rho_b = Function(rho0.function_space()) -sin_tmp = sin(lat) * sin(phi_c) -cos_tmp = cos(lat) * cos(phi_c) -r = a*acos(sin_tmp + cos_tmp*cos(lon-lamda_c)) -s = (d**2)/(d**2 + r**2) -theta_pert = deltaTheta*s*sin(2*pi*z/L_z) -theta0.interpolate(theta_b) - -# Compute the balanced density -compressible_hydrostatic_balance(eqns, - theta_b, - rho_b, - top=False, - exner_boundary=(p/p_0)**kappa) -theta0.interpolate(theta_pert) -theta0 += theta_b -rho0.assign(rho_b) - -stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # - -# Run! -stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/dry_bryan_fritsch.py b/examples/compressible/dry_bryan_fritsch.py deleted file mode 100644 index 653fc70bb..000000000 --- a/examples/compressible/dry_bryan_fritsch.py +++ /dev/null @@ -1,138 +0,0 @@ -""" -The dry rising bubble test from Bryan & Fritsch (2002). - -This uses the lowest-order function spaces, with the recovered methods for -transporting the fields. The test also uses a non-periodic base mesh. -""" - -from gusto import * -from firedrake import (IntervalMesh, ExtrudedMesh, - SpatialCoordinate, conditional, cos, pi, sqrt, - TestFunction, dx, TrialFunction, Constant, Function, - LinearVariationalProblem, LinearVariationalSolver) -import sys -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 1.0 -L = 10000. -H = 10000. - -if '--running-tests' in sys.argv: - deltax = 1000. - tmax = 5. -else: - deltax = 100. - tmax = 1000. - -degree = 0 -dirname = 'dry_bryan_fritsch' - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -nlayers = int(H/deltax) -ncolumns = int(L/deltax) -m = IntervalMesh(ncolumns, L) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -domain = Domain(mesh, dt, "CG", degree) - -# Equation -params = CompressibleParameters() -u_transport_option = "vector_advection_form" -eqns = CompressibleEulerEquations(domain, params, - u_transport_option=u_transport_option, - no_normal_flow_bc_ids=[1, 2]) - -# I/O -output = OutputParameters( - dirname=dirname, - dumpfreq=int(tmax / (5*dt)), - dumplist=['rho'], - dump_vtus=False, - dump_nc=True, -) -diagnostic_fields = [Perturbation('theta')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - -# Transport schemes -- set up options for using recovery wrapper -boundary_methods = {'DG': BoundaryMethod.taylor, - 'HDiv': BoundaryMethod.taylor} - -recovery_spaces = RecoverySpaces(domain, boundary_methods, use_vector_spaces=True) - -u_opts = recovery_spaces.HDiv_options -rho_opts = recovery_spaces.DG_options -theta_opts = recovery_spaces.theta_options - -transported_fields = [SSPRK3(domain, "rho", options=rho_opts), - SSPRK3(domain, "theta", options=theta_opts), - SSPRK3(domain, "u", options=u_opts)] - -transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta"]] - -# 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 -Vu = domain.spaces("HDiv") -Vt = domain.spaces("theta") -Vr = domain.spaces("DG") -x, z = SpatialCoordinate(mesh) - -# Define constant theta_e and water_t -Tsurf = 300.0 -theta_b = Function(Vt).interpolate(Constant(Tsurf)) - -# Calculate hydrostatic fields -compressible_hydrostatic_balance(eqns, theta_b, rho0, solve_for_rho=True) - -# make mean fields -rho_b = Function(Vr).assign(rho0) - -# define perturbation -xc = L / 2 -zc = 2000. -rc = 2000. -Tdash = 2.0 -r = sqrt((x - xc) ** 2 + (z - zc) ** 2) -theta_pert = Function(Vt).interpolate(conditional(r > rc, - 0.0, - Tdash * (cos(pi * r / (2.0 * rc))) ** 2)) - -# define initial theta -theta0.interpolate(theta_b * (theta_pert / 300.0 + 1.0)) - -# find perturbed rho -gamma = TestFunction(Vr) -rho_trial = TrialFunction(Vr) -lhs = gamma * rho_trial * dx -rhs = gamma * (rho_b * theta_b / theta0) * dx -rho_problem = LinearVariationalProblem(lhs, rhs, rho0) -rho_solver = LinearVariationalSolver(rho_problem) -rho_solver.solve() - -stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/mountain_hydrostatic.py b/examples/compressible/mountain_hydrostatic.py deleted file mode 100644 index f1b7ac015..000000000 --- a/examples/compressible/mountain_hydrostatic.py +++ /dev/null @@ -1,222 +0,0 @@ -""" -The 1 metre high mountain test case. This is solved with the hydrostatic -compressible Euler equations. -""" - -from gusto import * -from firedrake import (as_vector, VectorFunctionSpace, - PeriodicIntervalMesh, ExtrudedMesh, SpatialCoordinate, - exp, pi, cos, Function, conditional, Mesh, sqrt) -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 5.0 -L = 240000. # Domain length -H = 50000. # Height position of the model top - -if '--running-tests' in sys.argv: - tmax = dt - res = 1 - dumpfreq = 1 -else: - tmax = 15000. - res = 10 - dumpfreq = int(tmax / (5*dt)) - - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -# Make an normal extruded mesh which will be distorted to describe the mountain -nlayers = res*20 # horizontal layers -columns = res*12 # number of columns -m = PeriodicIntervalMesh(columns, L) -ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -Vc = VectorFunctionSpace(ext_mesh, "DG", 2) - -# Describe the mountain -a = 10000. -xc = L/2. -x, z = SpatialCoordinate(ext_mesh) -hm = 1. -zs = hm*a**2/((x-xc)**2 + a**2) -zh = 5000. -xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)]) - -# Make new mesh -new_coords = Function(Vc).interpolate(xexpr) -mesh = Mesh(new_coords) -mesh._base_mesh = m # Force new mesh to inherit original base mesh -domain = Domain(mesh, dt, "CG", 1) - -# Equation -parameters = CompressibleParameters(g=9.80665, cp=1004.) -sponge = SpongeLayerParameters(H=H, z_level=H-20000, mubar=0.3) -eqns = CompressibleEulerEquations(domain, parameters, sponge_options=sponge) - -# I/O -dirname = 'hydrostatic_mountain' -output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - dumplist=['u'], -) -diagnostic_fields = [CourantNumber(), ZComponent('u'), HydrostaticImbalance(eqns), - 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 -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) - -# Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, - transport_methods, - linear_solver=linear_solver, - alpha=alpha) - -# ---------------------------------------------------------------------------- # -# Initial conditions -# ---------------------------------------------------------------------------- # - -u0 = stepper.fields("u") -rho0 = stepper.fields("rho") -theta0 = stepper.fields("theta") - -# spaces -Vu = domain.spaces("HDiv") -Vt = domain.spaces("theta") -Vr = domain.spaces("DG") - -# Thermodynamic constants required for setting initial conditions -# and reference profiles -g = parameters.g -p_0 = parameters.p_0 -c_p = parameters.cp -R_d = parameters.R_d -kappa = parameters.kappa - -# Hydrostatic case: Isothermal with T = 250 -x, z = SpatialCoordinate(mesh) -Tsurf = 250. -N = g/sqrt(c_p*Tsurf) - -# 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) - -# Calculate hydrostatic exner -exner = Function(Vr) -rho_b = Function(Vr) - -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 - -# Set up kernels to evaluate global minima and maxima of fields -min_kernel = MinKernel() -max_kernel = MaxKernel() - -# First solve hydrostatic balance that gives Exner = 1 at bottom boundary -# This gives us a guess for the top boundary condition -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 -) - -# Solve hydrostatic balance again, but now use minimum value from first -# solve as the *top* boundary condition for Exner -top_value = min_kernel.apply(exner) -top_boundary = Constant(top_value, domain=mesh) -logger.info(f'Solving hydrostatic with top Exner of {top_value}') -compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary -) - -max_bottom_value = max_kernel.apply(exner) - -# Now we iterate, adjusting the top boundary condition, until this gives -# a maximum value of 1.0 at the surface -lower_top_guess = 0.9*top_value -upper_top_guess = 1.2*top_value -for i in range(max_iterations): - # If max bottom Exner value is equal to desired value, stop iteration - if abs(max_bottom_value - exner_surf) < tolerance: - break - - # Make new guess by average of previous guesses - top_guess = 0.5*(lower_top_guess + upper_top_guess) - top_boundary.assign(top_guess) - - logger.info( - f'Solving hydrostatic balance iteration {i}, with top Exner value ' - + f'of {top_guess}' - ) - - compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary - ) - - max_bottom_value = max_kernel.apply(exner) - - # Adjust guesses based on new value - if max_bottom_value < exner_surf: - lower_top_guess = top_guess - else: - upper_top_guess = top_guess - -logger.info(f'Final max bottom Exner value of {max_bottom_value}') - -# Perform a final solve to obtain hydrostatically balanced rho -compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary, - solve_for_rho=True -) - -theta0.assign(theta_b) -rho0.assign(rho_b) -u0.project(as_vector([20.0, 0.0])) -remove_initial_w(u0) - -stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/skamarock_klemp_hydrostatic.py b/examples/compressible/skamarock_klemp_hydrostatic.py deleted file mode 100644 index 9988a6d72..000000000 --- a/examples/compressible/skamarock_klemp_hydrostatic.py +++ /dev/null @@ -1,124 +0,0 @@ -""" -The non-linear gravity wave test case of Skamarock and Klemp (1994), but solved -with the hydrostatic Compressible Euler equations. - -Potential temperature is transported using SUPG. -""" - -from gusto import * -from firedrake import (as_vector, SpatialCoordinate, PeriodicRectangleMesh, - ExtrudedMesh, exp, sin, Function, pi) -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 25. -if '--running-tests' in sys.argv: - nlayers = 5 # horizontal layers - columns = 10 # number of columns - tmax = dt - dumpfreq = 1 -else: - nlayers = 10 # horizontal layers - columns = 150 # number of columns - tmax = 60000.0 - dumpfreq = int(tmax / (2*dt)) - -L = 6.0e6 # Length of domain -H = 1.0e4 # Height position of the model top - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -- 3D volume mesh -m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -domain = Domain(mesh, dt, "RTCF", 1) - -# Equation -parameters = CompressibleParameters(Omega=0.5e-4) -balanced_pg = as_vector((0., -1.0e-4*20, 0.)) -eqns = CompressibleEulerEquations(domain, parameters, - extra_terms=[("u", balanced_pg)]) - -# I/O -dirname = 'skamarock_klemp_hydrostatic' -output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - 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 -Vu = domain.spaces("HDiv") -Vt = domain.spaces("theta") -Vr = domain.spaces("DG") - -# Thermodynamic constants required for setting initial conditions -# and reference profiles -g = parameters.g -N = parameters.N -p_0 = parameters.p_0 -c_p = parameters.cp -R_d = parameters.R_d -kappa = parameters.kappa - -x, y, z = SpatialCoordinate(mesh) - -# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) -Tsurf = 300. -thetab = Tsurf*exp(N**2*z/g) - -theta_b = Function(Vt).interpolate(thetab) -rho_b = Function(Vr) - -a = 1.0e5 -deltaTheta = 1.0e-2 -theta_pert = deltaTheta*sin(pi*z/H)/(1 + (x - L/2)**2/a**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([20.0, 0.0, 0.0])) - -stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/skamarock_klemp_nonhydrostatic.py b/examples/compressible/skamarock_klemp_nonhydrostatic.py deleted file mode 100644 index e347472aa..000000000 --- a/examples/compressible/skamarock_klemp_nonhydrostatic.py +++ /dev/null @@ -1,141 +0,0 @@ -""" -The non-linear gravity wave test case of Skamarock and Klemp (1994). - -Potential temperature is transported using SUPG. -""" - -from petsc4py import PETSc -PETSc.Sys.popErrorHandler() -from gusto import * -import itertools -from firedrake import (as_vector, SpatialCoordinate, PeriodicIntervalMesh, - ExtrudedMesh, exp, sin, Function, pi, COMM_WORLD) -import numpy as np -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 6. -L = 3.0e5 # Domain length -H = 1.0e4 # Height position of the model top - -if '--running-tests' in sys.argv: - nlayers = 5 - columns = 30 - tmax = dt - dumpfreq = 1 -else: - nlayers = 10 - columns = 150 - tmax = 3600 - dumpfreq = int(tmax / (2*dt)) - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -- 3D volume mesh -m = PeriodicIntervalMesh(columns, L) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) -domain = Domain(mesh, dt, "CG", 1) - -# Equation -Tsurf = 300. -parameters = CompressibleParameters() -eqns = CompressibleEulerEquations(domain, parameters) - -# I/O -points_x = np.linspace(0., L, 100) -points_z = [H/2.] -points = np.array([p for p in itertools.product(points_x, points_z)]) -dirname = 'skamarock_klemp_nonlinear' - -# Dumping point data using legacy PointDataOutput is not supported in parallel -if COMM_WORLD.size == 1: - output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - pddumpfreq=dumpfreq, - dumplist=['u'], - point_data=[('theta_perturbation', points)], - ) -else: - logger.warning( - 'Dumping point data using legacy PointDataOutput is not' - ' supported in parallel\nDisabling PointDataOutput' - ) - output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - pddumpfreq=dumpfreq, - dumplist=['u'], - ) - -diagnostic_fields = [CourantNumber(), Gradient('u'), Perturbation('theta'), - Gradient('theta_perturbation'), Perturbation('rho'), - RichardsonNumber('theta', parameters.g/Tsurf), Gradient('theta')] -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 -Vu = domain.spaces("HDiv") -Vt = domain.spaces("theta") -Vr = domain.spaces("DG") - -# Thermodynamic constants required for setting initial conditions -# and reference profiles -g = parameters.g -N = parameters.N - -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) - -# Calculate hydrostatic exner -compressible_hydrostatic_balance(eqns, theta_b, rho_b) - -a = 5.0e3 -deltaTheta = 1.0e-2 -theta_pert = deltaTheta*sin(pi*z/H)/(1 + (x - L/2)**2/a**2) -theta0.interpolate(theta_b + theta_pert) -rho0.assign(rho_b) -u0.project(as_vector([20.0, 0.0])) - -stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/straka_bubble.py b/examples/compressible/straka_bubble.py deleted file mode 100644 index 2d3388ccc..000000000 --- a/examples/compressible/straka_bubble.py +++ /dev/null @@ -1,130 +0,0 @@ -""" -The falling cold density current test of Straka et al (1993). - -This example runs at a series of resolutions with different time steps. -""" - -from gusto import * -from firedrake import (PeriodicIntervalMesh, ExtrudedMesh, SpatialCoordinate, - Constant, pi, cos, Function, sqrt, - conditional) -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -if '--running-tests' in sys.argv: - res_dt = {800.: 4.} - tmax = 4. - ndumps = 1 -else: - res_dt = {800.: 4., 400.: 2., 200.: 1., 100.: 0.5, 50.: 0.25} - tmax = 15.*60. - ndumps = 4 - - -L = 51200. - -# build volume mesh -H = 6400. # Height position of the model top - -for delta, dt in res_dt.items(): - - # ------------------------------------------------------------------------ # - # Set up model objects - # ------------------------------------------------------------------------ # - - # Domain - nlayers = int(H/delta) # horizontal layers - columns = int(L/delta) # number of columns - m = PeriodicIntervalMesh(columns, L) - mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - domain = Domain(mesh, dt, "CG", 1) - - # Equation - parameters = CompressibleParameters() - u_diffusion_opts = DiffusionParameters(kappa=75., mu=10./delta) - theta_diffusion_opts = DiffusionParameters(kappa=75., mu=10./delta) - diffusion_options = [("u", u_diffusion_opts), ("theta", theta_diffusion_opts)] - eqns = CompressibleEulerEquations(domain, parameters, - diffusion_options=diffusion_options) - - # I/O - dirname = "straka_dx%s_dt%s" % (delta, dt) - dumpfreq = int(tmax / (ndumps*dt)) - output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - 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) - - # Diffusion schemes - diffusion_schemes = [BackwardEuler(domain, "u"), - BackwardEuler(domain, "theta")] - diffusion_methods = [InteriorPenaltyDiffusion(eqns, "u", u_diffusion_opts), - InteriorPenaltyDiffusion(eqns, "theta", theta_diffusion_opts)] - - # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, - spatial_methods=transport_methods+diffusion_methods, - linear_solver=linear_solver, - diffusion_schemes=diffusion_schemes) - - # ------------------------------------------------------------------------ # - # Initial conditions - # ------------------------------------------------------------------------ # - - u0 = stepper.fields("u") - rho0 = stepper.fields("rho") - theta0 = stepper.fields("theta") - - # spaces - Vu = domain.spaces("HDiv") - Vt = domain.spaces("theta") - Vr = domain.spaces("DG") - - # Isentropic background state - Tsurf = Constant(300.) - - theta_b = Function(Vt).interpolate(Tsurf) - rho_b = Function(Vr) - exner = Function(Vr) - - # Calculate hydrostatic exner - compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner0=exner, - solve_for_rho=True) - - x = SpatialCoordinate(mesh) - a = 5.0e3 - xc = 0.5*L - xr = 4000. - zc = 3000. - zr = 2000. - r = sqrt(((x[0]-xc)/xr)**2 + ((x[1]-zc)/zr)**2) - T_pert = conditional(r > 1., 0., -7.5*(1.+cos(pi*r))) - theta0.interpolate(theta_b + T_pert*exner) - rho0.assign(rho_b) - - stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - - # ------------------------------------------------------------------------ # - # Run - # ------------------------------------------------------------------------ # - - stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py deleted file mode 100644 index 23e38624f..000000000 --- a/examples/compressible/unsaturated_bubble.py +++ /dev/null @@ -1,236 +0,0 @@ -""" -A moist thermal in an unsaturated atmosphere. This test is based on that of -Grabowski and Clark (1991), and is described in Bendall et al (2020). - -As the thermal rises, water vapour condenses into cloud and forms rain. -Limiters are applied to the transport of the water species. -""" -from gusto import * -from gusto.equations import thermodynamics -from firedrake import (PeriodicIntervalMesh, ExtrudedMesh, - SpatialCoordinate, conditional, cos, pi, sqrt, exp, - TestFunction, dx, TrialFunction, Constant, Function, - LinearVariationalProblem, LinearVariationalSolver, - errornorm) -from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter -import sys - -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -dt = 1.0 -if '--running-tests' in sys.argv: - deltax = 240. - tmax = 10. - tdump = tmax -else: - deltax = 20. - tmax = 600. - tdump = 100. - -L = 3600. -h = 2400. -nlayers = int(h/deltax) -ncolumns = int(L/deltax) -degree = 0 - -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # - -# Domain -m = PeriodicIntervalMesh(ncolumns, L) -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=h/nlayers) -domain = Domain(mesh, dt, "CG", degree) - -# Equation -params = CompressibleParameters() -tracers = [WaterVapour(), CloudWater(), Rain()] -eqns = CompressibleEulerEquations(domain, params, - active_tracers=tracers) - -# I/O -dirname = 'unsaturated_bubble' -output = OutputParameters( - dirname=dirname, - dumpfreq=tdump, - dump_nc=True, - dumplist=['cloud_water', 'rain'], - checkpoint=False -) -diagnostic_fields = [RelativeHumidity(eqns), Perturbation('theta'), - Perturbation('water_vapour'), Perturbation('rho'), Perturbation('RelativeHumidity')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) -# Transport schemes -- specify options for using recovery wrapper -boundary_methods = {'DG': BoundaryMethod.taylor, - 'HDiv': BoundaryMethod.taylor} - -recovery_spaces = RecoverySpaces(domain, boundary_method=boundary_methods, use_vector_spaces=True) - -u_opts = recovery_spaces.HDiv_options -rho_opts = recovery_spaces.DG_options -theta_opts = recovery_spaces.theta_options - -VDG1 = domain.spaces("DG1_equispaced") -limiter = VertexBasedLimiter(VDG1) - -transported_fields = [SSPRK3(domain, "u", options=u_opts), - SSPRK3(domain, "rho", options=rho_opts), - SSPRK3(domain, "theta", options=theta_opts), - SSPRK3(domain, "water_vapour", options=theta_opts, limiter=limiter), - SSPRK3(domain, "cloud_water", options=theta_opts, limiter=limiter), - SSPRK3(domain, "rain", options=theta_opts, limiter=limiter)] - -transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta", "water_vapour", "cloud_water", "rain"]] - -# Linear solver -linear_solver = CompressibleSolver(eqns) - -# Physics schemes -# NB: can't yet use wrapper or limiter options with physics -Vt = domain.spaces('theta') -rainfall_method = DGUpwind(eqns, 'rain', outflow=True) -zero_limiter = MixedFSLimiter(eqns, {'water_vapour': ZeroLimiter(Vt), - 'cloud_water': ZeroLimiter(Vt)}) -physics_schemes = [(Fallout(eqns, 'rain', domain, rainfall_method), SSPRK3(domain)), - (Coalescence(eqns), ForwardEuler(domain)), - (EvaporationOfRain(eqns), ForwardEuler(domain)), - (SaturationAdjustment(eqns), ForwardEuler(domain, limiter=zero_limiter))] - -# Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, - transport_methods, - linear_solver=linear_solver, - physics_schemes=physics_schemes) - -# ---------------------------------------------------------------------------- # -# Initial conditions -# ---------------------------------------------------------------------------- # - -u0 = stepper.fields("u") -rho0 = stepper.fields("rho") -theta0 = stepper.fields("theta") -water_v0 = stepper.fields("water_vapour") -water_c0 = stepper.fields("cloud_water") -rain0 = stepper.fields("rain") - -# spaces -Vu = domain.spaces("HDiv") -Vr = domain.spaces("DG") -x, z = SpatialCoordinate(mesh) -quadrature_degree = (4, 4) -dxp = dx(degree=(quadrature_degree)) - -physics_boundary_method = BoundaryMethod.extruded - -# Define constant theta_e and water_t -Tsurf = 283.0 -psurf = 85000. -exner_surf = (psurf / eqns.parameters.p_0) ** eqns.parameters.kappa -humidity = 0.2 -S = 1.3e-5 -theta_surf = thermodynamics.theta(eqns.parameters, Tsurf, psurf) -theta_d = Function(Vt).interpolate(theta_surf * exp(S*z)) -H = Function(Vt).assign(humidity) - -# Calculate hydrostatic fields -unsaturated_hydrostatic_balance(eqns, stepper.fields, theta_d, H, - exner_boundary=Constant(exner_surf)) - -# make mean fields -theta_b = Function(Vt).assign(theta0) -rho_b = Function(Vr).assign(rho0) -water_vb = Function(Vt).assign(water_v0) - -# define perturbation to RH -xc = L / 2 -zc = 800. -r1 = 300. -r2 = 200. -r = sqrt((x - xc) ** 2 + (z - zc) ** 2) - -H_expr = conditional( - r > r1, 0.0, - conditional(r > r2, - (1 - humidity) * cos(pi * (r - r2) - / (2 * (r1 - r2))) ** 2, - 1 - humidity)) -H_pert = Function(Vt).interpolate(H_expr) -H.assign(H + H_pert) - -# now need to find perturbed rho, theta_vd and r_v -# follow approach used in unsaturated hydrostatic setup -rho_averaged = Function(Vt) -rho_recoverer = Recoverer(rho0, rho_averaged, boundary_method=physics_boundary_method) -rho_h = Function(Vr) -w_h = Function(Vt) -delta = 1.0 - -R_d = eqns.parameters.R_d -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, H, T, p) - -# make expressions to evaluate residual -exner_ev = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) -p_ev = thermodynamics.p(eqns.parameters, exner_ev) -T_ev = thermodynamics.T(eqns.parameters, theta0, exner_ev, water_v0) -RH_ev = thermodynamics.RH(eqns.parameters, water_v0, T_ev, p_ev) -RH = Function(Vt) - -# set-up rho problem to keep exner constant -gamma = TestFunction(Vr) -rho_trial = TrialFunction(Vr) -a = gamma * rho_trial * dxp -L = gamma * (rho_b * theta_b / theta0) * dxp -rho_problem = LinearVariationalProblem(a, L, rho_h) -rho_solver = LinearVariationalSolver(rho_problem) - -max_outer_solve_count = 20 -max_inner_solve_count = 10 - -for i in range(max_outer_solve_count): - # calculate averaged rho - rho_recoverer.project() - - RH.interpolate(RH_ev) - if errornorm(RH, H) < 1e-10: - break - - # first solve for r_v - for j in range(max_inner_solve_count): - w_h.interpolate(r_v_expr) - water_v0.assign(water_v0 * (1 - delta) + delta * w_h) - - # compute theta_vd - theta0.interpolate(theta_d * (1 + water_v0 / epsilon)) - - # test quality of solution by re-evaluating expression - RH.interpolate(RH_ev) - if errornorm(RH, H) < 1e-10: - break - - # now solve for rho with theta_vd and w_v guesses - rho_solver.solve() - - # damp solution - rho0.assign(rho0 * (1 - delta) + delta * rho_h) - - if i == max_outer_solve_count: - raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') - -# initialise fields -stepper.set_reference_profiles([('rho', rho_b), - ('theta', theta_b), - ('water_vapour', water_vb)]) -# ---------------------------------------------------------------------------- # -# Run -# ---------------------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible_euler/dcmip_3_1_gravity_wave.py b/examples/compressible_euler/dcmip_3_1_gravity_wave.py new file mode 100644 index 000000000..2fccfbe3f --- /dev/null +++ b/examples/compressible_euler/dcmip_3_1_gravity_wave.py @@ -0,0 +1,231 @@ +""" +The non-orographic gravity wave test case (3-1) from the DCMIP test case +document of Ullrich et al, 2012: +``Dynamical core model intercomparison project (DCMIP) test case document''. + +This uses a cubed-sphere mesh, the degree 1 finite element spaces and tests +substepping the transport schemes. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + ExtrudedMesh, Function, SpatialCoordinate, as_vector, + exp, acos, cos, sin, pi +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, lonlatr_from_xyz, CompressibleParameters, + CompressibleEulerEquations, CompressibleSolver, ZonalComponent, + compressible_hydrostatic_balance, Perturbation, GeneralCubedSphereMesh, +) + +dcmip_3_1_gravity_wave_defaults = { + 'ncells_per_edge': 8, + 'nlayers': 10, + 'dt': 50.0, + 'tmax': 3600., + 'dumpfreq': 9, + 'dirname': 'dcmip_3_1_gravity_wave' +} + + +def dcmip_3_1_gravity_wave( + ncells_per_edge=dcmip_3_1_gravity_wave_defaults['ncells_per_edge'], + nlayers=dcmip_3_1_gravity_wave_defaults['nlayers'], + dt=dcmip_3_1_gravity_wave_defaults['dt'], + tmax=dcmip_3_1_gravity_wave_defaults['tmax'], + dumpfreq=dcmip_3_1_gravity_wave_defaults['dumpfreq'], + dirname=dcmip_3_1_gravity_wave_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + parameters = CompressibleParameters() + a_ref = 6.37122e6 # Radius of the Earth (m) + X = 125.0 # Reduced-size Earth reduction factor + a = a_ref/X # Scaled radius of planet (m) + g = parameters.g # Acceleration due to gravity (m/s^2) + N = 0.01 # Brunt-Vaisala frequency (1/s) + p_0 = parameters.p_0 # Reference pressure (Pa, not hPa) + c_p = parameters.cp # SHC of dry air at constant pressure (J/kg/K) + R_d = parameters.R_d # Gas constant for dry air (J/kg/K) + kappa = parameters.kappa # R_d/c_p + T_eq = 300.0 # Isothermal atmospheric temperature (K) + p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator + u_max = 20.0 # Maximum amplitude of the zonal wind (m/s) + d = 5000.0 # Width parameter for Theta' + lamda_c = 2.0*pi/3.0 # Longitudinal centerpoint of Theta' + phi_c = 0.0 # Latitudinal centerpoint of Theta' (equator) + deltaTheta = 1.0 # Maximum amplitude of Theta' (K) + L_z = 20000.0 # Vertical wave length of the Theta' perturb. + z_top = 1.0e4 # Height position of the model top + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + u_eqn_type = 'vector_invariant_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = GeneralCubedSphereMesh(a, ncells_per_edge, degree=2) + mesh = ExtrudedMesh( + base_mesh, nlayers, layer_height=z_top/nlayers, + extrusion_type="radial" + ) + domain = Domain(mesh, dt, "RTCF", element_order) + + # Equation + eqns = CompressibleEulerEquations( + domain, parameters, u_transport_option=u_eqn_type + ) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True + ) + diagnostic_fields = [ + Perturbation('theta'), Perturbation('rho'), ZonalComponent('u'), + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "rho", fixed_subcycles=2), + SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2) + ] + transport_methods = [ + DGUpwind(eqns, field) for field in ["u", "rho", "theta"] + ] + + # 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') + theta0 = stepper.fields('theta') + rho0 = stepper.fields('rho') + + # spaces + Vr = domain.spaces("DG") + + x, y, z = SpatialCoordinate(mesh) + lon, lat, r = lonlatr_from_xyz(x, y, z) + h = r - a + + # Initial conditions with u0 + uexpr = as_vector([-u_max*y/a, u_max*x/a, 0.0]) + + # Surface temperature + G = g**2/(N**2*c_p) + Ts_expr = ( + G + (T_eq - G) * exp(-(u_max*N**2/(4*g*g)) * u_max*(cos(2.0*lat)-1.0)) + ) + + # Surface pressure + ps_expr = ( + p_eq * exp((u_max/(4.0*G*R_d)) * u_max*(cos(2.0*lat)-1.0)) + * (Ts_expr / T_eq)**(1.0/kappa) + ) + + # Background pressure + p_expr = ps_expr*(1 + G/Ts_expr*(exp(-N**2*h/g)-1))**(1.0/kappa) + p = Function(Vr).interpolate(p_expr) + + # Background temperature + Tb_expr = G*(1 - exp(N**2*h/g)) + Ts_expr*exp(N**2*h/g) + + # Background potential temperature + thetab_expr = Tb_expr*(p_0/p)**kappa + theta_b = Function(theta0.function_space()).interpolate(thetab_expr) + rho_b = Function(rho0.function_space()) + sin_tmp = sin(lat) * sin(phi_c) + cos_tmp = cos(lat) * cos(phi_c) + l = a*acos(sin_tmp + cos_tmp*cos(lon-lamda_c)) + s = (d**2)/(d**2 + l**2) + theta_pert = deltaTheta*s*sin(2*pi*h/L_z) + + # Compute the balanced density + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, top=False, exner_boundary=(p/p_0)**kappa + ) + + u0.project(uexpr) + theta0.interpolate(theta_b + theta_pert) + rho0.assign(rho_b) + + stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + # 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 panel edge of the cubed-sphere.", + type=int, + default=dcmip_3_1_gravity_wave_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=dcmip_3_1_gravity_wave_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=dcmip_3_1_gravity_wave_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=dcmip_3_1_gravity_wave_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=dcmip_3_1_gravity_wave_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=dcmip_3_1_gravity_wave_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + dcmip_3_1_gravity_wave(**vars(args)) diff --git a/examples/compressible_euler/dry_bryan_fritsch.py b/examples/compressible_euler/dry_bryan_fritsch.py new file mode 100644 index 000000000..dc4c92e9e --- /dev/null +++ b/examples/compressible_euler/dry_bryan_fritsch.py @@ -0,0 +1,220 @@ +""" +The dry rising bubble test from Bryan & Fritsch, 2002: +``A Benchmark Simulation for Moist Nonhydrostatic Numerical Models'', GMD. + +This uses the lowest-order function spaces, with the recovered methods for +transporting the fields. The test also uses a non-periodic base mesh. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + IntervalMesh, ExtrudedMesh, SpatialCoordinate, conditional, cos, pi, sqrt, + TestFunction, dx, TrialFunction, Constant, Function, as_vector, + LinearVariationalProblem, LinearVariationalSolver +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + RecoverySpaces, BoundaryMethod, Perturbation, CompressibleParameters, + CompressibleEulerEquations, CompressibleSolver, + compressible_hydrostatic_balance +) + +dry_bryan_fritsch_defaults = { + 'ncolumns': 100, + 'nlayers': 100, + 'dt': 2.0, + 'tmax': 1000., + 'dumpfreq': 500, + 'dirname': 'dry_bryan_fritsch' +} + + +def dry_bryan_fritsch( + ncolumns=dry_bryan_fritsch_defaults['ncolumns'], + nlayers=dry_bryan_fritsch_defaults['nlayers'], + dt=dry_bryan_fritsch_defaults['dt'], + tmax=dry_bryan_fritsch_defaults['tmax'], + dumpfreq=dry_bryan_fritsch_defaults['dumpfreq'], + dirname=dry_bryan_fritsch_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + domain_width = 10000. # domain width (m) + domain_height = 10000. # domain height (m) + zc = 2000. # vertical centre of bubble (m) + rc = 2000. # radius of bubble (m) + Tdash = 2.0 # strength of temperature perturbation (K) + Tsurf = 300.0 # background theta value (K) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 0 + u_eqn_type = 'vector_advection_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = IntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + params = CompressibleParameters() + eqns = CompressibleEulerEquations( + domain, params, u_transport_option=u_eqn_type, + no_normal_flow_bc_ids=[1, 2] + ) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True, + dumplist=['rho'] + ) + diagnostic_fields = [Perturbation('theta')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes -- set up options for using recovery wrapper + boundary_methods = {'DG': BoundaryMethod.taylor, + 'HDiv': BoundaryMethod.taylor} + + recovery_spaces = RecoverySpaces( + domain, boundary_methods, use_vector_spaces=True + ) + + u_opts = recovery_spaces.HDiv_options + rho_opts = recovery_spaces.DG_options + theta_opts = recovery_spaces.theta_options + + transported_fields = [ + SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "u", options=u_opts) + ] + + transport_methods = [ + DGUpwind(eqns, field) for field in ["u", "rho", "theta"] + ] + + # 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") + x, z = SpatialCoordinate(mesh) + + # Define constant theta_e and water_t + theta_b = Function(Vt).interpolate(Constant(Tsurf)) + + # Set initial wind to be zero + zero = Constant(0.0, domain=mesh) + u0.project(as_vector([zero, zero])) + + # Calculate hydrostatic fields + compressible_hydrostatic_balance(eqns, theta_b, rho0, solve_for_rho=True) + + # make mean fields + rho_b = Function(Vr).assign(rho0) + + # define perturbation + xc = domain_width / 2 + r = sqrt((x - xc) ** 2 + (z - zc) ** 2) + theta_pert = Function(Vt).interpolate( + conditional( + r > rc, + 0.0, + Tdash * (cos(pi * r / (2.0 * rc))) ** 2 + ) + ) + + # define initial theta + theta0.interpolate(theta_b * (theta_pert / 300.0 + 1.0)) + + # find perturbed rho + gamma = TestFunction(Vr) + rho_trial = TrialFunction(Vr) + lhs = gamma * rho_trial * dx + rhs = gamma * (rho_b * theta_b / theta0) * dx + rho_problem = LinearVariationalProblem(lhs, rhs, rho0) + rho_solver = LinearVariationalSolver(rho_problem) + rho_solver.solve() + + 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=dry_bryan_fritsch_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=dry_bryan_fritsch_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=dry_bryan_fritsch_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=dry_bryan_fritsch_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=dry_bryan_fritsch_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=dry_bryan_fritsch_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + dry_bryan_fritsch(**vars(args)) diff --git a/examples/compressible_euler/mountain_hydrostatic.py b/examples/compressible_euler/mountain_hydrostatic.py new file mode 100644 index 000000000..8d9b0f8a6 --- /dev/null +++ b/examples/compressible_euler/mountain_hydrostatic.py @@ -0,0 +1,306 @@ +""" +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. + +This test describes a wave over a mountain in a hydrostatic atmosphere. + +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 +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, ZComponent, Perturbation, + CompressibleParameters, HydrostaticCompressibleEulerEquations, + CompressibleSolver, compressible_hydrostatic_balance, HydrostaticImbalance, + SpongeLayerParameters, MinKernel, MaxKernel, remove_initial_w, logger +) + +mountain_hydrostatic_defaults = { + 'ncolumns': 200, + 'nlayers': 120, + 'dt': 5.0, + 'tmax': 15000., + 'dumpfreq': 1500, + 'dirname': 'mountain_hydrostatic' +} + + +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'] +): + + # ------------------------------------------------------------------------ # + # 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 + 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 + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + u_eqn_type = 'vector_invariant_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + # Make normal extruded mesh which will be distorted to describe the mountain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + ext_mesh = ExtrudedMesh( + base_mesh, layers=nlayers, layer_height=domain_height/nlayers + ) + Vc = VectorFunctionSpace(ext_mesh, "DG", 2) + + # Describe the mountain + xc = domain_width/2. + x, z = SpatialCoordinate(ext_mesh) + zs = hm * a**2 / ((x - xc)**2 + a**2) + xexpr = as_vector( + [x, conditional(z < zh, z + cos(0.5 * pi * z / zh)**6 * zs, z)] + ) + + # Make new mesh + new_coords = Function(Vc).interpolate(xexpr) + mesh = Mesh(new_coords) + mesh._base_mesh = base_mesh # Force new mesh to inherit original base mesh + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + parameters = CompressibleParameters(g=g, cp=cp) + sponge = SpongeLayerParameters( + H=domain_height, z_level=domain_height-sponge_depth, mubar=sponge_mu/dt + ) + eqns = HydrostaticCompressibleEulerEquations( + 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 + ) + diagnostic_fields = [ + ZComponent('u'), HydrostaticImbalance(eqns), + 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 + 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 + ) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, alpha=alpha + ) + + # ------------------------------------------------------------------------ # + # 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 + N = parameters.N + + # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) + x, z = SpatialCoordinate(mesh) + thetab = Tsurf*exp(N**2*z/g) + theta_b = Function(Vt).interpolate(thetab) + + # Calculate hydrostatic exner + exner = Function(Vr) + rho_b = Function(Vr) + + # Set up kernels to evaluate global minima and maxima of fields + min_kernel = MinKernel() + max_kernel = MaxKernel() + + # First solve hydrostatic balance that gives Exner = 1 at bottom boundary + # This gives us a guess for the top boundary condition + 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 + ) + + # Solve hydrostatic balance again, but now use minimum value from first + # solve as the *top* boundary condition for Exner + top_value = min_kernel.apply(exner) + top_boundary = Constant(top_value, domain=mesh) + logger.info(f'Solving hydrostatic with top Exner of {top_value}') + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary + ) + + max_bottom_value = max_kernel.apply(exner) + + # Now we iterate, adjusting the top boundary condition, until this gives + # a maximum value of 1.0 at the surface + lower_top_guess = 0.9*top_value + upper_top_guess = 1.2*top_value + for i in range(max_iterations): + # If max bottom Exner value is equal to desired value, stop iteration + if abs(max_bottom_value - exner_surf) < tolerance: + break + + # Make new guess by average of previous guesses + top_guess = 0.5*(lower_top_guess + upper_top_guess) + top_boundary.assign(top_guess) + + logger.info( + f'Solving hydrostatic balance iteration {i}, with top Exner value ' + + f'of {top_guess}' + ) + + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary + ) + + max_bottom_value = max_kernel.apply(exner) + + # Adjust guesses based on new value + if max_bottom_value < exner_surf: + lower_top_guess = top_guess + else: + upper_top_guess = top_guess + + logger.info(f'Final max bottom Exner value of {max_bottom_value}') + + # Perform a final solve to obtain hydrostatically balanced rho + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary, + solve_for_rho=True + ) + + theta0.assign(theta_b) + rho0.assign(rho_b) + u0.project(as_vector([initial_wind, 0.0])) + remove_initial_w(u0) + + 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=mountain_hydrostatic_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=mountain_hydrostatic_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=mountain_hydrostatic_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=mountain_hydrostatic_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=mountain_hydrostatic_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=mountain_hydrostatic_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + mountain_hydrostatic(**vars(args)) diff --git a/examples/compressible_euler/skamarock_klemp_hydrostatic.py b/examples/compressible_euler/skamarock_klemp_hydrostatic.py new file mode 100644 index 000000000..f75ad9b10 --- /dev/null +++ b/examples/compressible_euler/skamarock_klemp_hydrostatic.py @@ -0,0 +1,204 @@ +""" +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 new file mode 100644 index 000000000..777134d61 --- /dev/null +++ b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py @@ -0,0 +1,220 @@ +""" +This example uses the non-linear 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. +""" +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter + +from petsc4py import PETSc +PETSc.Sys.popErrorHandler() +import itertools +from firedrake import ( + as_vector, SpatialCoordinate, PeriodicIntervalMesh, ExtrudedMesh, exp, sin, + Function, pi, COMM_WORLD +) +import numpy as np +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, CourantNumber, Perturbation, Gradient, + CompressibleParameters, CompressibleEulerEquations, CompressibleSolver, + compressible_hydrostatic_balance, logger, RichardsonNumber +) + +skamarock_klemp_nonhydrostatic_defaults = { + 'ncolumns': 150, + 'nlayers': 10, + 'dt': 6.0, + 'tmax': 3600., + 'dumpfreq': 300, + 'dirname': 'skamarock_klemp_nonhydrostatic' +} + + +def skamarock_klemp_nonhydrostatic( + ncolumns=skamarock_klemp_nonhydrostatic_defaults['ncolumns'], + nlayers=skamarock_klemp_nonhydrostatic_defaults['nlayers'], + dt=skamarock_klemp_nonhydrostatic_defaults['dt'], + tmax=skamarock_klemp_nonhydrostatic_defaults['tmax'], + dumpfreq=skamarock_klemp_nonhydrostatic_defaults['dumpfreq'], + dirname=skamarock_klemp_nonhydrostatic_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 3.0e5 # Width of domain (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) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain -- 3D volume mesh + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + parameters = CompressibleParameters() + 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)]) + + # 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, + point_data=[('theta_perturbation', points)], + ) + else: + logger.warning( + 'Dumping point data using legacy PointDataOutput is not' + ' supported in parallel\nDisabling PointDataOutput' + ) + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq, + dump_vtus=True, dump_nc=True, + ) + + diagnostic_fields = [ + CourantNumber(), Gradient('u'), Perturbation('theta'), + Gradient('theta_perturbation'), Perturbation('rho'), + RichardsonNumber('theta', parameters.g/Tsurf), Gradient('theta') + ] + 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) + + # Calculate hydrostatic exner + compressible_hydrostatic_balance(eqns, theta_b, rho_b) + + theta_pert = ( + deltaTheta * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + theta0.interpolate(theta_b + theta_pert) + rho0.assign(rho_b) + u0.project(as_vector([wind_initial, 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_nonhydrostatic_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=skamarock_klemp_nonhydrostatic_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=skamarock_klemp_nonhydrostatic_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=skamarock_klemp_nonhydrostatic_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=skamarock_klemp_nonhydrostatic_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=skamarock_klemp_nonhydrostatic_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + skamarock_klemp_nonhydrostatic(**vars(args)) diff --git a/examples/compressible_euler/straka_bubble.py b/examples/compressible_euler/straka_bubble.py new file mode 100644 index 000000000..fed55a1ec --- /dev/null +++ b/examples/compressible_euler/straka_bubble.py @@ -0,0 +1,213 @@ +""" +The falling cold density current test of Straka et al, 1993: +``Numerical solutions of a non‐linear density current: A benchmark solution and +comparisons'', MiF. + +Diffusion is included in the velocity and potential temperature equations. The +degree 1 finite elements are used in this configuration. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + PeriodicIntervalMesh, ExtrudedMesh, SpatialCoordinate, Constant, pi, cos, + Function, sqrt, conditional, as_vector +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, CourantNumber, Perturbation, + DiffusionParameters, InteriorPenaltyDiffusion, BackwardEuler, + CompressibleParameters, CompressibleEulerEquations, CompressibleSolver, + compressible_hydrostatic_balance +) + +straka_bubble_defaults = { + 'nlayers': 32, + 'dt': 1.0, + 'tmax': 900., + 'dumpfreq': 225, + 'dirname': 'straka_bubble' +} + + +def straka_bubble( + nlayers=straka_bubble_defaults['nlayers'], + dt=straka_bubble_defaults['dt'], + tmax=straka_bubble_defaults['tmax'], + dumpfreq=straka_bubble_defaults['dumpfreq'], + dirname=straka_bubble_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + domain_width = 51200. # domain width (m) + domain_height = 6400. # domain height (m) + zc = 3000. # vertical centre of perturbation (m) + xr = 4000. # horizontal radius of perturbation (m) + zr = 2000. # vertical radius of perturbation (m) + T_pert = -7.5 # strength of temperature perturbation (K) + Tsurf = 300.0 # background theta value (K) + kappa = 75. # diffusivity parameter (m^2/s) + mu0 = 10. # interior penalty parameter (1/m) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + delta = domain_height/nlayers + ncolumns = 8 * nlayers + element_order = 1 + u_eqn_type = 'vector_advection_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=delta) + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + parameters = CompressibleParameters() + diffusion_params = DiffusionParameters(kappa=kappa, mu=mu0/delta) + diffusion_options = [("u", diffusion_params), ("theta", diffusion_params)] + eqns = CompressibleEulerEquations( + domain, parameters, u_transport_option=u_eqn_type, + diffusion_options=diffusion_options + ) + + # 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) + + # Diffusion schemes + diffusion_schemes = [ + BackwardEuler(domain, "u"), + BackwardEuler(domain, "theta") + ] + diffusion_methods = [ + InteriorPenaltyDiffusion(eqns, "u", diffusion_params), + InteriorPenaltyDiffusion(eqns, "theta", diffusion_params) + ] + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, + spatial_methods=transport_methods+diffusion_methods, + linear_solver=linear_solver, diffusion_schemes=diffusion_schemes + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # spaces + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Isentropic background state + theta_b = Function(Vt).interpolate(Tsurf) + rho_b = Function(Vr) + exner = Function(Vr) + + # Calculate hydrostatic exner + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner0=exner, solve_for_rho=True + ) + + x, z = SpatialCoordinate(mesh) + xc = 0.5*domain_width + r = sqrt(((x - xc)/xr)**2 + ((z - zc)/zr)**2) + T_pert_expr = conditional( + r > 1., + 0., + 0.5*T_pert*(1. + cos(pi*r)) + ) + + # Set initial fields + zero = Constant(0.0, domain=mesh) + u0.project(as_vector([zero, zero])) + theta0.interpolate(theta_b + T_pert_expr*exner) + rho0.assign(rho_b) + + # Reference profiles + 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( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=straka_bubble_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=straka_bubble_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=straka_bubble_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=straka_bubble_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=straka_bubble_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + straka_bubble(**vars(args)) diff --git a/examples/compressible_euler/test_compressible_euler_examples.py b/examples/compressible_euler/test_compressible_euler_examples.py new file mode 100644 index 000000000..384824e12 --- /dev/null +++ b/examples/compressible_euler/test_compressible_euler_examples.py @@ -0,0 +1,143 @@ +import pytest + + +def make_dirname(test_name): + from mpi4py import MPI + comm = MPI.COMM_WORLD + if comm.size > 1: + return f'pytest_{test_name}_parallel' + else: + return f'pytest_{test_name}' + + +def test_dcmip_3_1_gravity_wave(): + from dcmip_3_1_gravity_wave import dcmip_3_1_gravity_wave + test_name = 'dcmip_3_1_gravity_wave' + dcmip_3_1_gravity_wave( + ncells_per_edge=4, + nlayers=4, + dt=100, + tmax=200, + dumpfreq=2, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_dcmip_3_1_gravity_wave_parallel(): + test_dcmip_3_1_gravity_wave() + + +def test_dry_bryan_fritsch(): + from dry_bryan_fritsch import dry_bryan_fritsch + test_name = 'dry_bryan_fritsch' + dry_bryan_fritsch( + ncolumns=20, + nlayers=20, + dt=2.0, + tmax=20.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +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( + ncolumns=30, + nlayers=5, + dt=6.0, + tmax=60.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +# 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(): + 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) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_skamarock_klemp_nonhydrostatic_parallel(): + test_skamarock_klemp_nonhydrostatic() + + +def test_straka_bubble(): + from straka_bubble import straka_bubble + test_name = 'straka_bubble' + straka_bubble( + nlayers=6, + dt=4.0, + tmax=40.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=3) +def test_straka_bubble_parallel(): + test_straka_bubble() + + +def test_unsaturated_bubble(): + from unsaturated_bubble import unsaturated_bubble + test_name = 'unsaturated_bubble' + unsaturated_bubble( + ncolumns=20, + nlayers=20, + dt=1.0, + tmax=10.0, + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_unsaturated_bubble_parallel(): + test_unsaturated_bubble() diff --git a/examples/compressible_euler/unsaturated_bubble.py b/examples/compressible_euler/unsaturated_bubble.py new file mode 100644 index 000000000..394e15f97 --- /dev/null +++ b/examples/compressible_euler/unsaturated_bubble.py @@ -0,0 +1,338 @@ +""" +A moist thermal in an unsaturated atmosphere, including a rain species. This +test is based on that of Grabowski and Clark, 1991: +``Cloud–environment interface instability: Rising thermal calculations in two +spatial dimensions'', JAS. + +and is described in Bendall et al, 2020: +``A compatible finite‐element discretisation for the moist compressible Euler +equations'', QJRMS. + +As the thermal rises, water vapour condenses into cloud and forms rain. +Limiters are applied to the transport of the water species. + +This configuration uses the lowest-order finite elements, and the recovery +wrapper to provide higher-order accuracy. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + PeriodicIntervalMesh, ExtrudedMesh, SpatialCoordinate, conditional, cos, pi, + sqrt, exp, TestFunction, dx, TrialFunction, Constant, Function, errornorm, + LinearVariationalProblem, LinearVariationalSolver, as_vector +) +from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + Perturbation, RecoverySpaces, BoundaryMethod, Recoverer, Fallout, + Coalescence, SaturationAdjustment, EvaporationOfRain, thermodynamics, + CompressibleParameters, CompressibleEulerEquations, CompressibleSolver, + unsaturated_hydrostatic_balance, WaterVapour, CloudWater, Rain, + RelativeHumidity, ForwardEuler, MixedFSLimiter, ZeroLimiter +) + +unsaturated_bubble_defaults = { + 'ncolumns': 180, + 'nlayers': 120, + 'dt': 1.0, + 'tmax': 600., + 'dumpfreq': 300, + 'dirname': 'unsaturated_bubble' +} + + +def unsaturated_bubble( + ncolumns=unsaturated_bubble_defaults['ncolumns'], + nlayers=unsaturated_bubble_defaults['nlayers'], + dt=unsaturated_bubble_defaults['dt'], + tmax=unsaturated_bubble_defaults['tmax'], + dumpfreq=unsaturated_bubble_defaults['dumpfreq'], + dirname=unsaturated_bubble_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + domain_width = 3600. # domain width (m) + domain_height = 2400. # domain height (m) + zc = 800. # height of centre of perturbation (m) + r1 = 300. # outer radius of perturbation (m) + r2 = 200. # inner radius of perturbation (m) + Tsurf = 283.0 # surface temperature (K) + psurf = 85000. # surface pressure (Pa) + rel_hum_background = 0.2 # background relative humidity (dimensionless) + S = 1.3e-5 # height factor for theta profile (1/m) + max_outer_solve_count = 20 # max num outer iterations for initialisation + max_inner_solve_count = 10 # max num inner iterations for initialisation + tol_initialisation = 1e-10 # tolerance for initialisation + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 0 + u_eqn_type = 'vector_advection_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + params = CompressibleParameters() + tracers = [WaterVapour(), CloudWater(), Rain()] + eqns = CompressibleEulerEquations( + domain, params, active_tracers=tracers, u_transport_option=u_eqn_type + ) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True, + dumplist=['cloud_water', 'rain'] + ) + diagnostic_fields = [ + RelativeHumidity(eqns), Perturbation('theta'), Perturbation('rho'), + Perturbation('water_vapour'), Perturbation('RelativeHumidity') + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes -- specify options for using recovery wrapper + boundary_methods = {'DG': BoundaryMethod.taylor, + 'HDiv': BoundaryMethod.taylor} + + recovery_spaces = RecoverySpaces(domain, boundary_method=boundary_methods, use_vector_spaces=True) + + u_opts = recovery_spaces.HDiv_options + rho_opts = recovery_spaces.DG_options + theta_opts = recovery_spaces.theta_options + + VDG1 = domain.spaces("DG1_equispaced") + limiter = VertexBasedLimiter(VDG1) + + transported_fields = [ + SSPRK3(domain, "u", options=u_opts), + SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "water_vapour", options=theta_opts, limiter=limiter), + SSPRK3(domain, "cloud_water", options=theta_opts, limiter=limiter), + SSPRK3(domain, "rain", options=theta_opts, limiter=limiter) + ] + + transport_methods = [ + DGUpwind(eqns, field) for field in + ["u", "rho", "theta", "water_vapour", "cloud_water", "rain"] + ] + + # Linear solver + linear_solver = CompressibleSolver(eqns) + + # Physics schemes + Vt = domain.spaces('theta') + rainfall_method = DGUpwind(eqns, 'rain', outflow=True) + zero_limiter = MixedFSLimiter( + eqns, + {'water_vapour': ZeroLimiter(Vt), 'cloud_water': ZeroLimiter(Vt)} + ) + physics_schemes = [ + (Fallout(eqns, 'rain', domain, rainfall_method), SSPRK3(domain)), + (Coalescence(eqns), ForwardEuler(domain)), + (EvaporationOfRain(eqns), ForwardEuler(domain)), + (SaturationAdjustment(eqns), ForwardEuler(domain, limiter=zero_limiter)) + ] + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, physics_schemes=physics_schemes + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") + water_r0 = stepper.fields("rain") + + # spaces + Vr = domain.spaces("DG") + x, z = SpatialCoordinate(mesh) + quadrature_degree = (4, 4) + dxp = dx(degree=(quadrature_degree)) + + physics_boundary_method = BoundaryMethod.extruded + + # Define constant theta_e and water_t + exner_surf = (psurf / eqns.parameters.p_0) ** eqns.parameters.kappa + theta_surf = thermodynamics.theta(eqns.parameters, Tsurf, psurf) + theta_d = Function(Vt).interpolate(theta_surf * exp(S*z)) + rel_hum = Function(Vt).assign(rel_hum_background) + + # Calculate hydrostatic fields + unsaturated_hydrostatic_balance( + eqns, stepper.fields, theta_d, rel_hum, + exner_boundary=Constant(exner_surf) + ) + + # make mean fields + theta_b = Function(Vt).assign(theta0) + rho_b = Function(Vr).assign(rho0) + water_vb = Function(Vt).assign(water_v0) + + # define perturbation to RH + xc = domain_width / 2 + r = sqrt((x - xc) ** 2 + (z - zc) ** 2) + + rel_hum_pert_expr = conditional( + r > r1, + 0.0, + conditional( + r > r2, + (1 - rel_hum_background) * cos(pi*(r - r2) / (2*(r1 - r2)))**2, + 1 - rel_hum_background + ) + ) + rel_hum.interpolate(rel_hum_background + rel_hum_pert_expr) + + # now need to find perturbed rho, theta_vd and r_v + # follow approach used in unsaturated hydrostatic setup + rho_averaged = Function(Vt) + rho_recoverer = Recoverer( + rho0, rho_averaged, boundary_method=physics_boundary_method + ) + rho_eval = Function(Vr) + water_v_eval = Function(Vt) + delta = 1.0 + + R_d = eqns.parameters.R_d + 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) + rel_hum_expr = thermodynamics.RH(eqns.parameters, water_v0, T_expr, p_expr) + rel_hum_eval = Function(Vt) + + # set-up rho problem to keep exner constant + gamma = TestFunction(Vr) + rho_trial = TrialFunction(Vr) + lhs = gamma * rho_trial * dxp + rhs = gamma * (rho_b * theta_b / theta0) * dxp + rho_problem = LinearVariationalProblem(lhs, rhs, rho_eval) + rho_solver = LinearVariationalSolver(rho_problem) + + for i in range(max_outer_solve_count): + # calculate averaged rho + rho_recoverer.project() + + rel_hum_eval.interpolate(rel_hum_expr) + if errornorm(rel_hum_eval, rel_hum) < tol_initialisation: + break + + # first solve for r_v + for _ in range(max_inner_solve_count): + water_v_eval.interpolate(r_v_expr) + water_v0.assign(water_v0 * (1 - delta) + delta * water_v_eval) + + # compute theta_vd + theta0.interpolate(theta_d * (1 + water_v0 / epsilon)) + + # test quality of solution by re-evaluating expression + rel_hum_eval.interpolate(rel_hum_expr) + if errornorm(rel_hum_eval, rel_hum) < tol_initialisation: + break + + # now solve for rho with theta_vd and w_v guesses + rho_solver.solve() + + # damp solution + rho0.assign(rho0 * (1 - delta) + delta * rho_eval) + + if i == max_outer_solve_count: + raise RuntimeError( + f'Hydrostatic balance solve has not converged within {i} iterations' + ) + + # Set wind, cloud and rain to be zero + zero = Constant(0.0, domain=mesh) + u0.project(as_vector([zero, zero])) + water_c0.interpolate(zero) + water_r0.interpolate(zero) + + # initialise reference profiles + stepper.set_reference_profiles( + [('rho', rho_b), ('theta', theta_b), ('water_vapour', water_vb)] + ) + + # ------------------------------------------------------------------------ # + # 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=unsaturated_bubble_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=unsaturated_bubble_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=unsaturated_bubble_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=unsaturated_bubble_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=unsaturated_bubble_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=unsaturated_bubble_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + unsaturated_bubble(**vars(args)) diff --git a/examples/shallow_water/linear_williamson_2.py b/examples/shallow_water/linear_williamson_2.py index e32a92e40..9d022bf11 100644 --- a/examples/shallow_water/linear_williamson_2.py +++ b/examples/shallow_water/linear_williamson_2.py @@ -1,82 +1,150 @@ """ -The Williamson 2 shallow-water test case (solid-body rotation), solved with a -discretisation of the linear shallow-water equations. +A linearised form of Test Case 2 (solid-body rotation) of Williamson et al 1992: +``A standard test set for numerical approximations to the shallow water +equations in spherical geometry'', JCP. -This uses an icosahedral mesh of the sphere. +This uses an icosahedral mesh of the sphere, and the linear shallow water +equations. """ -from gusto import * -from firedrake import IcosahedralSphereMesh, SpatialCoordinate, as_vector, pi -import sys +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import Function, SpatialCoordinate, as_vector, pi +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, DefaultTransport, + ForwardEuler, SteadyStateError, ShallowWaterParameters, + LinearShallowWaterEquations, GeneralIcosahedralSphereMesh, + ZonalComponent, MeridionalComponent, RelativeVorticity +) -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # +linear_williamson_2_defaults = { + 'ncells_per_edge': 16, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 5.*24.*60.*60., # 5 days + 'dumpfreq': 96, # once per day with default options + 'dirname': 'linear_williamson_2' +} -dt = 3600. -day = 24.*60.*60. -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 -else: - tmax = 5*day - dumpfreq = int(tmax / (5*dt)) -refinements = 3 # number of horizontal cells = 20*(4^refinements) +def linear_williamson_2( + ncells_per_edge=linear_williamson_2_defaults['ncells_per_edge'], + dt=linear_williamson_2_defaults['dt'], + tmax=linear_williamson_2_defaults['tmax'], + dumpfreq=linear_williamson_2_defaults['dumpfreq'], + dirname=linear_williamson_2_defaults['dirname'] +): -R = 6371220. -H = 2000. + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # -# ---------------------------------------------------------------------------- # -# Set up model objects -# ---------------------------------------------------------------------------- # + radius = 6371220. # planetary radius (m) + mean_depth = 2000. # reference depth (m) + u_max = 2*pi*radius/(12*24*60*60) # Max amplitude of the zonal wind (m/s) -# Domain -mesh = IcosahedralSphereMesh(radius=R, - refinement_level=refinements, degree=3) -x = SpatialCoordinate(mesh) -domain = Domain(mesh, dt, 'BDM', 1) - -# Equation -parameters = ShallowWaterParameters(H=H) -Omega = parameters.Omega -x = SpatialCoordinate(mesh) -fexpr = 2*Omega*x[2]/R -eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) - -# I/O -output = OutputParameters( - dirname='linear_williamson_2', - dumpfreq=dumpfreq, -) -diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # -# Transport schemes -transport_schemes = [ForwardEuler(domain, "D")] -transport_methods = [DefaultTransport(eqns, "D")] + element_order = 1 -# Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes, transport_methods) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # -# ---------------------------------------------------------------------------- # -# Initial conditions -# ---------------------------------------------------------------------------- # + # Domain + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + x, y, z = SpatialCoordinate(mesh) + domain = Domain(mesh, dt, 'BDM', element_order) + + # Equation + parameters = ShallowWaterParameters(H=mean_depth) + Omega = parameters.Omega + fexpr = 2*Omega*z/radius + eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=False, dump_vtus=True + ) + diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D'), + ZonalComponent('u'), MeridionalComponent('u'), + RelativeVorticity()] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) -u0 = stepper.fields("u") -D0 = stepper.fields("D") -u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) -uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) -g = parameters.g -Dexpr = - ((R * Omega * u_max)*(x[2]*x[2]/(R*R)))/g -u0.project(uexpr) -D0.interpolate(Dexpr) + # Transport schemes + transport_schemes = [ForwardEuler(domain, "D")] + transport_methods = [DefaultTransport(eqns, "D")] -Dbar = Function(D0.function_space()).assign(H) -stepper.set_reference_profiles([('D', Dbar)]) + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transport_schemes, transport_methods + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + g = parameters.g + + u0 = stepper.fields("u") + D0 = stepper.fields("D") + + uexpr = as_vector([-u_max*y/radius, u_max*x/radius, 0.0]) + Dexpr = - ((radius*Omega*u_max) * (z/radius)**2) / g + + u0.project(uexpr) + D0.interpolate(Dexpr) + + Dbar = Function(D0.function_space()).assign(mean_depth) + stepper.set_reference_profiles([('D', Dbar)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # -# Run +# MAIN # ---------------------------------------------------------------------------- # -stepper.run(t=0, tmax=tmax) + +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_williamson_2_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=linear_williamson_2_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=linear_williamson_2_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=linear_williamson_2_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=linear_williamson_2_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + linear_williamson_2(**vars(args)) diff --git a/examples/shallow_water/moist_convective_williamson2.py b/examples/shallow_water/moist_convective_williamson2.py deleted file mode 100644 index 22ca6b78e..000000000 --- a/examples/shallow_water/moist_convective_williamson2.py +++ /dev/null @@ -1,163 +0,0 @@ -""" -A moist convective version of the Williamson 2 shallow water test (steady state -geostrophically-balanced flow). The saturation function depends on height, -with a constant background buoyancy/temperature field. -Vapour is initialised very close to saturation and small overshoots will -generate clouds. -""" -from gusto import * -from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, sin, cos, exp) -import sys - -# ----------------------------------------------------------------- # -# Test case parameters -# ----------------------------------------------------------------- # - -dt = 120 - -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 -else: - day = 24*60*60 - tmax = 5*day - ndumps = 5 - dumpfreq = int(tmax / (ndumps*dt)) - -R = 6371220. -u_max = 20 -phi_0 = 3e4 -epsilon = 1/300 -theta_0 = epsilon*phi_0**2 -g = 9.80616 -H = phi_0/g -xi = 0 -q0 = 200 -beta1 = 110 -alpha = 16 -gamma_v = 0.98 -qprecip = 1e-4 -gamma_r = 1e-3 - -# ----------------------------------------------------------------- # -# Set up model objects -# ----------------------------------------------------------------- # - -# Domain -mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2) -degree = 1 -domain = Domain(mesh, dt, 'BDM', degree) -x = SpatialCoordinate(mesh) - -# Equations -parameters = ShallowWaterParameters(H=H, g=g) -Omega = parameters.Omega -fexpr = 2*Omega*x[2]/R - -tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')] - -eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option='vector_advection_form', - active_tracers=tracers) - -# IO -dirname = "moist_convective_williamson2" -output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - dumplist_latlon=['D', 'D_error'], - dump_nc=True, - dump_vtus=True) - -diagnostic_fields = [CourantNumber(), RelativeVorticity(), - PotentialVorticity(), - ShallowWaterKineticEnergy(), - ShallowWaterPotentialEnergy(parameters), - ShallowWaterPotentialEnstrophy(), - SteadyStateError('u'), SteadyStateError('D'), - SteadyStateError('water_vapour'), - SteadyStateError('cloud_water')] - -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - - -# define saturation function -def sat_func(x_in): - h = x_in.split()[1] - lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[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 - return q0/(g*h) * exp(20*(theta)) - - -transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] - -limiter = DG1Limiter(domain.spaces('DG')) - -transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "D"), - SSPRK3(domain, "water_vapour", limiter=limiter), - SSPRK3(domain, "cloud_water", limiter=limiter), - SSPRK3(domain, "rain", limiter=limiter) - ] - -linear_solver = MoistConvectiveSWSolver(eqns) - -sat_adj = SWSaturationAdjustment(eqns, sat_func, - time_varying_saturation=True, - convective_feedback=True, beta1=beta1, - gamma_v=gamma_v, time_varying_gamma_v=False, - parameters=parameters) - -inst_rain = InstantRain(eqns, qprecip, vapour_name="cloud_water", - rain_name="rain", gamma_r=gamma_r) - -physics_schemes = [(sat_adj, ForwardEuler(domain)), - (inst_rain, ForwardEuler(domain))] - -stepper = SemiImplicitQuasiNewton(eqns, io, - transport_schemes=transported_fields, - spatial_methods=transport_methods, - linear_solver=linear_solver, - physics_schemes=physics_schemes) - -# ----------------------------------------------------------------- # -# Initial conditions -# ----------------------------------------------------------------- # - -u0 = stepper.fields("u") -D0 = stepper.fields("D") -v0 = stepper.fields("water_vapour") - -lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) - -uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x) -g = parameters.g -w = Omega*R*u_max + (u_max**2)/2 -sigma = 0 - -Dexpr = H - (1/g)*(w)*((sin(phi))**2) -D_for_v = H - (1/g)*(w + sigma)*((sin(phi))**2) - -# though this set-up has no buoyancy, we use the expression for theta to set up -# the initial vapour -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 - -initial_msat = q0/(g*Dexpr) * exp(20*theta) -vexpr = (1 - xi) * initial_msat - -u0.project(uexpr) -D0.interpolate(Dexpr) -v0.interpolate(vexpr) - -# Set reference profiles -Dbar = Function(D0.function_space()).assign(H) -stepper.set_reference_profiles([('D', Dbar)]) - -# ----------------------------------------------------------------- # -# Run -# ----------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py new file mode 100644 index 000000000..07fc7e96c --- /dev/null +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -0,0 +1,240 @@ +""" +A moist convective form of Test Case 2 (solid-body rotation with flow in +geostrophic balance) of Williamson 2 et al, 1992: +``A standard test set for numerical approximations to the shallow water +equations in spherical geometry'', JCP. + +Three moist variables (vapour, cloud liquid and rain) are used. The saturation +function depends on height, with a temporally-constant background buoyancy/ +temperature field. Vapour is initialised very close to saturation and +small overshoots in will generate clouds. + +This example uses the icosahedral sphere mesh and degree 1 spaces. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import SpatialCoordinate, sin, cos, exp, Function +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, + ZonalComponent, MeridionalComponent, SteadyStateError, lonlatr_from_xyz, + DG1Limiter, InstantRain, MoistConvectiveSWSolver, ForwardEuler, + RelativeVorticity, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, + GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr +) + +moist_convect_williamson_2_defaults = { + 'ncells_per_edge': 16, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 5.*24.*60.*60., # 5 days + 'dumpfreq': 96, # once per day with default options + 'dirname': 'moist_convective_williamson_2' +} + + +def moist_convect_williamson_2( + ncells_per_edge=moist_convect_williamson_2_defaults['ncells_per_edge'], + dt=moist_convect_williamson_2_defaults['dt'], + tmax=moist_convect_williamson_2_defaults['tmax'], + dumpfreq=moist_convect_williamson_2_defaults['dumpfreq'], + dirname=moist_convect_williamson_2_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + radius = 6371220. # planetary radius (m) + u_max = 20. # max amplitude of the zonal wind (m/s) + phi_0 = 3.0e4 # reference geopotential height (m^2/s^2) + epsilon = 1/300 # linear air expansion coeff (1/K) + theta_0 = epsilon*phi_0**2 # ref depth-integrated temperature (no units) + g = 9.80616 # acceleration due to gravity (m/s^2) + mean_depth = phi_0/g # reference depth (m) + xi = 0 # fraction of excess vapour/cloud not converted + q0 = 200 # saturation mixing ratio scaling (kg/kg) + beta1 = 1600 # depth-vaporisation factor (m) + gamma_v = 0.98 # vaporisation implicit factor + qprecip = 1e-4 # cloud to rain conversion threshold (kg/kg) + gamma_r = 1e-3 # rain-coalescence implicit factor + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + u_eqn_type = 'vector_advection_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + domain = Domain(mesh, dt, 'BDM', element_order) + x, y, z = SpatialCoordinate(mesh) + _, phi, _ = lonlatr_from_xyz(x, y, z) + + # Equations + parameters = ShallowWaterParameters(H=mean_depth, g=g) + Omega = parameters.Omega + fexpr = 2*Omega*z/radius + + tracers = [ + WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG') + ] + + eqns = ShallowWaterEquations( + domain, parameters, fexpr=fexpr, u_transport_option=u_eqn_type, + active_tracers=tracers + ) + + # IO + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=False, dump_vtus=True, + dumplist_latlon=['D', 'D_error'] + ) + diagnostic_fields = [ + SteadyStateError('u'), SteadyStateError('D'), + SteadyStateError('water_vapour'), ZonalComponent('u'), + MeridionalComponent('u'), RelativeVorticity() + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # define saturation function + def sat_func(x_in): + h = x_in.split()[1] + 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 + return q0/(g*h) * exp(20*(theta)) + + transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] + + limiter = DG1Limiter(domain.spaces('DG')) + + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "D"), + SSPRK3(domain, "water_vapour", limiter=limiter), + SSPRK3(domain, "cloud_water", limiter=limiter), + SSPRK3(domain, "rain", limiter=limiter) + ] + + linear_solver = MoistConvectiveSWSolver(eqns) + + # Physics schemes + sat_adj = SWSaturationAdjustment( + eqns, sat_func, time_varying_saturation=True, + convective_feedback=True, beta1=beta1, gamma_v=gamma_v, + time_varying_gamma_v=False, parameters=parameters + ) + inst_rain = InstantRain( + eqns, qprecip, vapour_name="cloud_water", rain_name="rain", + gamma_r=gamma_r + ) + + physics_schemes = [ + (sat_adj, ForwardEuler(domain)), (inst_rain, ForwardEuler(domain)) + ] + + stepper = SemiImplicitQuasiNewton( + eqns, io, transport_schemes=transported_fields, + spatial_methods=transport_methods, linear_solver=linear_solver, + physics_schemes=physics_schemes + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + D0 = stepper.fields("D") + v0 = stepper.fields("water_vapour") + + uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, (x, y, z)) + g = parameters.g + w = Omega*radius*u_max + (u_max**2)/2 + sigma = 0 + + Dexpr = mean_depth - (1/g)*(w)*((sin(phi))**2) + + # though this set-up has no buoyancy, we use the expression for theta to + # set up the initial vapour + 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 + + initial_msat = q0/(g*Dexpr) * exp(20*theta) + vexpr = (1 - xi) * initial_msat + + u0.project(uexpr) + D0.interpolate(Dexpr) + v0.interpolate(vexpr) + + # Set reference profiles + Dbar = Function(D0.function_space()).assign(mean_depth) + stepper.set_reference_profiles([('D', Dbar)]) + + # ------------------------------------------------------------------------ # + # 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_convect_williamson_2_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=moist_convect_williamson_2_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=moist_convect_williamson_2_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=moist_convect_williamson_2_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=moist_convect_williamson_2_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + moist_convect_williamson_2(**vars(args)) diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py deleted file mode 100644 index f022e1058..000000000 --- a/examples/shallow_water/moist_thermal_williamson5.py +++ /dev/null @@ -1,148 +0,0 @@ -""" -Moist flow over a mountain test case from Zerroukat and Allen (2015). Similar -to the Williamson et al test 5 but with additional thermodynamic equations. -""" -from gusto import * -from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, - as_vector, pi, sqrt, min_value, exp, cos, sin) -import sys - -# ----------------------------------------------------------------- # -# Test case parameters -# ----------------------------------------------------------------- # - -dt = 300 - -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 -else: - day = 24*60*60 - tmax = 50*day - ndumps = 50 - dumpfreq = int(tmax / (ndumps*dt)) - -R = 6371220. -H = 5960. -u_max = 20. -# moist shallow water parameters -epsilon = 1/300 -SP = -40*epsilon -EQ = 30*epsilon -NP = -20*epsilon -mu1 = 0.05 -mu2 = 0.98 -q0 = 135 # chosen to give an initial max vapour of approx 0.02 -beta2 = 10 -qprecip = 1e-4 -gamma_r = 1e-3 -# topography parameters -R0 = pi/9. -R0sq = R0**2 -lamda_c = -pi/2. -phi_c = pi/6. - -# ----------------------------------------------------------------- # -# Set up model objects -# ----------------------------------------------------------------- # - -# Domain -mesh = IcosahedralSphereMesh(radius=R, - refinement_level=4, degree=1) -degree = 1 -domain = Domain(mesh, dt, "BDM", degree) -x = SpatialCoordinate(mesh) - -# Equation -parameters = ShallowWaterParameters(H=H) -Omega = parameters.Omega -fexpr = 2*Omega*x[2]/R - -# Topography -lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) -lsq = (lamda - lamda_c)**2 -thsq = (phi - phi_c)**2 -rsq = min_value(R0sq, lsq+thsq) -r = sqrt(rsq) -tpexpr = 2000 * (1 - r/R0) - -tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')] -eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr, - thermal=True, - active_tracers=tracers) - -# I/O -dirname = "moist_thermal_williamson5" -output = OutputParameters( - dirname=dirname, - dumplist_latlon=['D'], - dumpfreq=dumpfreq, -) -diagnostic_fields = [Sum('D', 'topography'), CourantNumber()] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - - -# Saturation function -def sat_func(x_in): - h = x_in.split()[1] - b = x_in.split()[2] - return (q0/(g*h + g*tpexpr)) * exp(20*(1 - b/g)) - - -# Feedback proportionality is dependent on h and b -def gamma_v(x_in): - h = x_in.split()[1] - b = x_in.split()[2] - return (1 + beta2*(20*q0/(g*h + g*tpexpr) * exp(20*(1 - b/g))))**(-1) - - -SWSaturationAdjustment(eqns, sat_func, time_varying_saturation=True, - parameters=parameters, thermal_feedback=True, - beta2=beta2, gamma_v=gamma_v, - time_varying_gamma_v=True) - -InstantRain(eqns, qprecip, vapour_name="cloud_water", rain_name="rain", - gamma_r=gamma_r) - -transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] - -# Timestepper -stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) - -# ----------------------------------------------------------------- # -# Initial conditions -# ----------------------------------------------------------------- # - -u0 = stepper.fields("u") -D0 = stepper.fields("D") -b0 = stepper.fields("b") -v0 = stepper.fields("water_vapour") -c0 = stepper.fields("cloud_water") -r0 = stepper.fields("rain") - -uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) - -g = parameters.g -Rsq = R**2 -Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - tpexpr - -# Expression for initial buoyancy - note the bracket around 1-mu -F = (2/(pi**2))*(phi*(phi-pi/2)*SP - 2*(phi+pi/2)*(phi-pi/2)*(1-mu1)*EQ + phi*(phi+pi/2)*NP) -theta_expr = F + mu1*EQ*cos(phi)*sin(lamda) -bexpr = g * (1 - theta_expr) - -# Expression for initial vapour depends on initial saturation -initial_msat = q0/(g*D0 + g*tpexpr) * exp(20*theta_expr) -vexpr = mu2 * initial_msat - -# Initialise (cloud and rain initially zero) -u0.project(uexpr) -D0.interpolate(Dexpr) -b0.interpolate(bexpr) -v0.interpolate(vexpr) - -# ----------------------------------------------------------------- # -# Run -# ----------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/moist_thermal_williamson_5.py b/examples/shallow_water/moist_thermal_williamson_5.py new file mode 100644 index 000000000..9ce7da77e --- /dev/null +++ b/examples/shallow_water/moist_thermal_williamson_5.py @@ -0,0 +1,243 @@ +""" +The moist thermal form of Test Case 5 (flow over a mountain) of Williamson et +al, 1992: +``A standard test set for numerical approximations to the shallow water +equations in spherical geometry'', JCP. + +The initial conditions are taken from Zerroukat & Allen, 2015: +``A moist Boussinesq shallow water equations set for testing atmospheric +models'', JCP. + +Three moist variables (vapour, cloud liquid and rain) are used. This set of +equations involves an active buoyancy field. + +The example here uses the icosahedral sphere mesh and degree 1 spaces. An +explicit RK4 timestepper is used. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + SpatialCoordinate, as_vector, pi, sqrt, min_value, exp, cos, sin +) +from gusto import ( + Domain, IO, OutputParameters, Timestepper, RK4, DGUpwind, + ShallowWaterParameters, ShallowWaterEquations, Sum, + lonlatr_from_xyz, InstantRain, SWSaturationAdjustment, WaterVapour, + CloudWater, Rain, GeneralIcosahedralSphereMesh, RelativeVorticity, + ZonalComponent, MeridionalComponent +) + +moist_thermal_williamson_5_defaults = { + 'ncells_per_edge': 16, # number of cells per icosahedron edge + 'dt': 300.0, # 5 minutes + 'tmax': 50.*24.*60.*60., # 50 days + 'dumpfreq': 2880, # once per 10 days with default options + 'dirname': 'moist_thermal_williamson_5' +} + + +def moist_thermal_williamson_5( + ncells_per_edge=moist_thermal_williamson_5_defaults['ncells_per_edge'], + dt=moist_thermal_williamson_5_defaults['dt'], + tmax=moist_thermal_williamson_5_defaults['tmax'], + dumpfreq=moist_thermal_williamson_5_defaults['dumpfreq'], + dirname=moist_thermal_williamson_5_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + radius = 6371220. # planetary radius (m) + mean_depth = 5960 # reference depth (m) + g = 9.80616 # acceleration due to gravity (m/s^2) + u_max = 20. # max amplitude of the zonal wind (m/s) + epsilon = 1/300 # linear air expansion coeff (1/K) + theta_SP = -40*epsilon # value of theta at south pole (no units) + theta_EQ = 30*epsilon # value of theta at equator (no units) + theta_NP = -20*epsilon # value of theta at north pole (no units) + mu1 = 0.05 # scaling for theta with longitude (no units) + mu2 = 0.98 # proportion of qsat to make init qv (no units) + q0 = 135 # qsat scaling, gives init q_v of ~0.02, (kg/kg) + beta2 = 10*g # buoyancy-vaporisation factor (m/s^2) + nu = 20. # qsat factor in exponent (no units) + qprecip = 1e-4 # cloud to rain conversion threshold (kg/kg) + gamma_r = 1e-3 # rain-coalescence implicit factor + mountain_height = 2000. # height of mountain (m) + R0 = pi/9. # radius of mountain (rad) + lamda_c = -pi/2. # longitudinal centre of mountain (rad) + phi_c = pi/6. # latitudinal centre of mountain (rad) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + u_eqn_type = 'vector_invariant_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + domain = Domain(mesh, dt, "BDM", element_order) + x, y, z = SpatialCoordinate(mesh) + lamda, phi, _ = lonlatr_from_xyz(x, y, z) + + # Equation: coriolis + parameters = ShallowWaterParameters(H=mean_depth, g=g) + Omega = parameters.Omega + fexpr = 2*Omega*z/radius + + # Equation: topography + rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2) + r = sqrt(rsq) + tpexpr = mountain_height * (1 - r/R0) + + # Equation: moisture + tracers = [ + WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG') + ] + eqns = ShallowWaterEquations( + domain, parameters, fexpr=fexpr, bexpr=tpexpr, thermal=True, + active_tracers=tracers, u_transport_option=u_eqn_type + ) + + # I/O + output = OutputParameters( + dirname=dirname, dumplist_latlon=['D'], dumpfreq=dumpfreq, + dump_vtus=True, dump_nc=False, + dumplist=['D', 'b', 'water_vapour', 'cloud_water'] + ) + diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), + ZonalComponent('u'), MeridionalComponent('u')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Physics ------------------------------------------------------------------ + # Saturation function -- first define simple expression + def q_sat(b, D): + return (q0/(g*D + g*tpexpr)) * exp(nu*(1 - b/g)) + + # 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] + 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] + return 1.0 / (1.0 + nu*beta2/g*q_sat(b, D)) + + SWSaturationAdjustment( + eqns, phys_sat_func, time_varying_saturation=True, + parameters=parameters, thermal_feedback=True, + beta2=beta2, gamma_v=gamma_v, time_varying_gamma_v=True + ) + + InstantRain( + eqns, qprecip, vapour_name="cloud_water", rain_name="rain", + gamma_r=gamma_r + ) + + transport_methods = [ + DGUpwind(eqns, field_name) for field_name in eqns.field_names + ] + + # Timestepper + stepper = Timestepper( + eqns, RK4(domain), io, spatial_methods=transport_methods + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + D0 = stepper.fields("D") + b0 = stepper.fields("b") + v0 = stepper.fields("water_vapour") + c0 = stepper.fields("cloud_water") + r0 = stepper.fields("rain") + + uexpr = as_vector([-u_max*y/radius, u_max*x/radius, 0.0]) + + Dexpr = ( + mean_depth - tpexpr + - (radius * Omega * u_max + 0.5*u_max**2)*(z/radius)**2/g + ) + + # Expression for initial buoyancy - note the bracket around 1-mu + theta_expr = ( + 2/(pi**2) * ( + phi*(phi - pi/2)*theta_SP + - 2*(phi + pi/2) * (phi - pi/2)*(1 - mu1)*theta_EQ + + phi*(phi + pi/2)*theta_NP + ) + + mu1*theta_EQ*cos(phi)*sin(lamda) + ) + bexpr = g * (1 - theta_expr) + + # Expression for initial vapour depends on initial saturation + vexpr = mu2 * q_sat(bexpr, Dexpr) + + # Initialise (cloud and rain initially zero) + u0.project(uexpr) + D0.interpolate(Dexpr) + b0.interpolate(bexpr) + v0.interpolate(vexpr) + c0.assign(0.0) + r0.assign(0.0) + + # ----------------------------------------------------------------- # + # 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_williamson_5_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=moist_thermal_williamson_5_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=moist_thermal_williamson_5_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=moist_thermal_williamson_5_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=moist_thermal_williamson_5_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + moist_thermal_williamson_5(**vars(args)) diff --git a/examples/shallow_water/shallow_water_1d.py b/examples/shallow_water/shallow_water_1d.py deleted file mode 100644 index ab50bd524..000000000 --- a/examples/shallow_water/shallow_water_1d.py +++ /dev/null @@ -1,69 +0,0 @@ -import numpy as np -import sys - -from firedrake import * -from gusto import * -from pyop2.mpi import MPI - -L = 2*pi -n = 128 -delta = L/n -mesh = PeriodicIntervalMesh(128, L) -dt = 0.0001 -if '--running-tests' in sys.argv: - T = 0.0005 -else: - T = 1 - -domain = Domain(mesh, dt, 'CG', 1) - -epsilon = 0.1 -parameters = ShallowWaterParameters(H=1/epsilon, g=1/epsilon) - -u_diffusion_opts = DiffusionParameters(kappa=1e-2) -v_diffusion_opts = DiffusionParameters(kappa=1e-2, mu=10/delta) -D_diffusion_opts = DiffusionParameters(kappa=1e-2, mu=10/delta) -diffusion_options = [("u", u_diffusion_opts), - ("v", v_diffusion_opts), - ("D", D_diffusion_opts)] - -eqns = ShallowWaterEquations_1d(domain, parameters, - fexpr=Constant(1/epsilon), - diffusion_options=diffusion_options) - -output = OutputParameters(dirname="1dsw_%s" % str(epsilon), - dumpfreq=50) -io = IO(domain, output) - -transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "v"), - DGUpwind(eqns, "D")] - -diffusion_methods = [CGDiffusion(eqns, "u", u_diffusion_opts), - InteriorPenaltyDiffusion(eqns, "v", v_diffusion_opts), - InteriorPenaltyDiffusion(eqns, "D", D_diffusion_opts)] - -stepper = Timestepper(eqns, RK4(domain), io, - spatial_methods=transport_methods+diffusion_methods) - -D = stepper.fields("D") -x = SpatialCoordinate(mesh)[0] -hexpr = ( - sin(x - pi/2) * exp(-4*(x - pi/2)**2) - + sin(8*(x - pi)) * exp(-2*(x - pi)**2) -) -h = Function(D.function_space()).interpolate(hexpr) - -A = assemble(h*dx) - -# B must be the maximum value of h (across all ranks) -B = np.zeros(1) -COMM_WORLD.Allreduce(h.dat.data_ro.max(), B, MPI.MAX) - -C0 = 1/(1 - 2*pi*B[0]/A) -C1 = (1 - C0)/B[0] -H = parameters.H -D.interpolate(C1*hexpr + C0) - -D += parameters.H - -stepper.run(0, T) diff --git a/examples/shallow_water/shallow_water_1d_wave.py b/examples/shallow_water/shallow_water_1d_wave.py new file mode 100644 index 000000000..3e0b1582b --- /dev/null +++ b/examples/shallow_water/shallow_water_1d_wave.py @@ -0,0 +1,177 @@ +""" +A shallow water wave on a 1D periodic domain. The test is taken from +Haut & Wingate, 2014: +``An asymptotic parallel-in-time method for highly oscillatory PDEs'', SIAM JSC. + +The velocity includes a component normal to the domain, and diffusion terms are +included in the equations. + +This example uses an explicit RK4 timestepper to solve the equations. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +import numpy as np +from pyop2.mpi import MPI +from firedrake import ( + PeriodicIntervalMesh, Function, assemble, SpatialCoordinate, COMM_WORLD, + Constant, pi, sin, exp, dx +) +from gusto import ( + Domain, IO, OutputParameters, Timestepper, RK4, DGUpwind, + ShallowWaterParameters, ShallowWaterEquations_1d, CGDiffusion, + InteriorPenaltyDiffusion, DiffusionParameters +) + +shallow_water_1d_wave_defaults = { + 'ncells': 128, + 'dt': 0.0001, + 'tmax': 1.0, + 'dumpfreq': 1000, # 10 outputs with default options + 'dirname': 'shallow_water_1d_wave' +} + + +def shallow_water_1d_wave( + ncells=shallow_water_1d_wave_defaults['ncells'], + dt=shallow_water_1d_wave_defaults['dt'], + tmax=shallow_water_1d_wave_defaults['tmax'], + dumpfreq=shallow_water_1d_wave_defaults['dumpfreq'], + dirname=shallow_water_1d_wave_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + domain_length = 2*pi # length of domain (m) + kappa = 1.e-2 # diffusivity (m^2/s^2) + epsilon = 0.1 # scaling factor for depth, gravity and rotation + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + mesh = PeriodicIntervalMesh(ncells, domain_length) + domain = Domain(mesh, dt, 'CG', element_order) + + # Diffusion + delta = domain_length / ncells + u_diffusion_opts = DiffusionParameters(kappa=kappa) + v_diffusion_opts = DiffusionParameters(kappa=kappa, mu=10/delta) + D_diffusion_opts = DiffusionParameters(kappa=kappa, mu=10/delta) + diffusion_options = [ + ("u", u_diffusion_opts), + ("v", v_diffusion_opts), + ("D", D_diffusion_opts) + ] + + # Equation + parameters = ShallowWaterParameters(H=1/epsilon, g=1/epsilon) + eqns = ShallowWaterEquations_1d( + domain, parameters, fexpr=Constant(1/epsilon), + diffusion_options=diffusion_options + ) + + output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq) + io = IO(domain, output) + + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "v"), + DGUpwind(eqns, "D") + ] + + diffusion_methods = [ + CGDiffusion(eqns, "u", u_diffusion_opts), + InteriorPenaltyDiffusion(eqns, "v", v_diffusion_opts), + InteriorPenaltyDiffusion(eqns, "D", D_diffusion_opts) + ] + + stepper = Timestepper( + eqns, RK4(domain), io, + spatial_methods=transport_methods+diffusion_methods + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + x = SpatialCoordinate(mesh)[0] + D = stepper.fields("D") + + # Spatially-varying part of initial condition + hexpr = ( + sin(x - pi/2) * exp(-4*(x - pi/2)**2) + + sin(8*(x - pi)) * exp(-2*(x - pi)**2) + ) + + # Make a function to include spatially-varying part + h = Function(D.function_space()).interpolate(hexpr) + + A = assemble(h*dx) + + # B must be the maximum value of h (across all ranks) + B = np.zeros(1) + COMM_WORLD.Allreduce(h.dat.data_ro.max(), B, MPI.MAX) + + # D is normalised form of h + C0 = 1/(1 - 2*pi*B[0]/A) + C1 = (1 - C0)/B[0] + D.interpolate(parameters.H + C1*hexpr + C0) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(0, tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncells', + help="The number of cells in the 1D domain", + type=int, + default=shallow_water_1d_wave_defaults['ncells'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=shallow_water_1d_wave_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=shallow_water_1d_wave_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=shallow_water_1d_wave_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=shallow_water_1d_wave_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + shallow_water_1d_wave(**vars(args)) diff --git a/examples/shallow_water/test_shallow_water_examples.py b/examples/shallow_water/test_shallow_water_examples.py new file mode 100644 index 000000000..eb64e94e0 --- /dev/null +++ b/examples/shallow_water/test_shallow_water_examples.py @@ -0,0 +1,129 @@ +import pytest + + +def make_dirname(test_name): + from mpi4py import MPI + comm = MPI.COMM_WORLD + if comm.size > 1: + return f'pytest_{test_name}_parallel' + else: + return f'pytest_{test_name}' + + +def test_linear_williamson_2(): + from linear_williamson_2 import linear_williamson_2 + test_name = 'linear_williamson_2' + linear_williamson_2( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_linear_williamson_2_parallel(): + test_linear_williamson_2() + + +def test_moist_convective_williamson_2(): + from moist_convective_williamson_2 import moist_convect_williamson_2 + test_name = 'moist_convective_williamson_2' + moist_convect_williamson_2( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_moist_convective_williamson_2_parallel(): + test_moist_convective_williamson_2() + + +def test_moist_thermal_williamson_5(): + from moist_thermal_williamson_5 import moist_thermal_williamson_5 + test_name = 'moist_thermal_williamson_5' + moist_thermal_williamson_5( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=2) +def test_moist_thermal_williamson_5_parallel(): + test_moist_thermal_williamson_5() + + +def test_shallow_water_1d_wave(): + from shallow_water_1d_wave import shallow_water_1d_wave + test_name = 'shallow_water_1d_wave' + shallow_water_1d_wave( + ncells=20, + dt=1.0e-4, + tmax=1.0e-3, + dumpfreq=2, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def test_shallow_water_1d_wave_parallel(): + test_shallow_water_1d_wave() + + +def test_thermal_williamson_2(): + from thermal_williamson_2 import thermal_williamson_2 + test_name = 'thermal_williamson_2' + thermal_williamson_2( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def test_thermal_williamson_2_parallel(): + test_thermal_williamson_2() + + +def test_williamson_2(): + from williamson_2 import williamson_2 + test_name = 'williamson_2' + williamson_2( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def test_williamson_2_parallel(): + test_williamson_2() + + +def test_williamson_5(): + from williamson_5 import williamson_5 + test_name = 'williamson_5' + williamson_5( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def test_williamson_5_parallel(): + test_williamson_5() diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py deleted file mode 100644 index 342ee8b07..000000000 --- a/examples/shallow_water/thermal_williamson2.py +++ /dev/null @@ -1,115 +0,0 @@ -from gusto import * -from firedrake import IcosahedralSphereMesh, SpatialCoordinate, sin, cos -import sys - -# ----------------------------------------------------------------- # -# Test case parameters -# ----------------------------------------------------------------- # - -dt = 4000 - -if '--running-tests' in sys.argv: - tmax = dt - dumpfreq = 1 -else: - day = 24*60*60 - tmax = 5*day - ndumps = 5 - dumpfreq = int(tmax / (ndumps*dt)) - -R = 6371220. -u_max = 20 -phi_0 = 3e4 -epsilon = 1/300 -theta_0 = epsilon*phi_0**2 -g = 9.80616 -H = phi_0/g - -# ----------------------------------------------------------------- # -# Set up model objects -# ----------------------------------------------------------------- # - -# Domain -mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2) -degree = 1 -domain = Domain(mesh, dt, 'BDM', degree) -x = SpatialCoordinate(mesh) - -# Equations -params = ShallowWaterParameters(H=H, g=g) -Omega = params.Omega -fexpr = 2*Omega*x[2]/R -eqns = ShallowWaterEquations(domain, params, fexpr=fexpr, u_transport_option='vector_advection_form', thermal=True) - -# IO -dirname = "thermal_williamson2" -output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - dumplist_latlon=['D', 'D_error'], -) - -diagnostic_fields = [RelativeVorticity(), PotentialVorticity(), - ShallowWaterKineticEnergy(), - ShallowWaterPotentialEnergy(params), - ShallowWaterPotentialEnstrophy(), - SteadyStateError('u'), SteadyStateError('D'), - SteadyStateError('b'), MeridionalComponent('u'), - ZonalComponent('u')] -io = IO(domain, output, diagnostic_fields=diagnostic_fields) - -# Transport schemes -transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "D", fixed_subcycles=2), - SSPRK3(domain, "b", fixed_subcycles=2)] -transport_methods = [DGUpwind(eqns, "u"), - DGUpwind(eqns, "D"), - DGUpwind(eqns, "b")] - -# Linear solver -linear_solver = ThermalSWSolver(eqns) - -# Time stepper -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") - -lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) - -uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x) -g = params.g -w = Omega*R*u_max + (u_max**2)/2 -sigma = w/10 - -Dexpr = H - (1/g)*(w + sigma)*((sin(phi))**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 - -bexpr = params.g * (1 - theta) - -u0.project(uexpr) -D0.interpolate(Dexpr) -b0.interpolate(bexpr) - -# Set reference profiles -Dbar = Function(D0.function_space()).assign(H) -bbar = Function(b0.function_space()).interpolate(bexpr) -stepper.set_reference_profiles([('D', Dbar), ('b', bbar)]) - -# ----------------------------------------------------------------- # -# Run -# ----------------------------------------------------------------- # - -stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/thermal_williamson_2.py b/examples/shallow_water/thermal_williamson_2.py new file mode 100644 index 000000000..fe489ee32 --- /dev/null +++ b/examples/shallow_water/thermal_williamson_2.py @@ -0,0 +1,197 @@ +""" +The thermal form of Test Case 2 (solid-body rotation with geostrophically +balanced flow) of Williamson et al, 1992: +``A standard test set for numerical approximations to the shallow water +equations in spherical geometry'', JCP. + +The initial conditions are taken from Zerroukat & Allen, 2015: +``A moist Boussinesq shallow water equations set for testing atmospheric +models'', JCP. + +The example here uses the icosahedral sphere mesh and degree 1 spaces. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import Function, SpatialCoordinate, sin, cos +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, + RelativeVorticity, PotentialVorticity, SteadyStateError, + ZonalComponent, MeridionalComponent, ThermalSWSolver, + xyz_vector_from_lonlatr, lonlatr_from_xyz, GeneralIcosahedralSphereMesh +) + +thermal_williamson_2_defaults = { + 'ncells_per_edge': 16, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 5.*24.*60.*60., # 5 days + 'dumpfreq': 96, # once per day with default options + 'dirname': 'thermal_williamson_2' +} + + +def thermal_williamson_2( + ncells_per_edge=thermal_williamson_2_defaults['ncells_per_edge'], + dt=thermal_williamson_2_defaults['dt'], + tmax=thermal_williamson_2_defaults['tmax'], + dumpfreq=thermal_williamson_2_defaults['dumpfreq'], + dirname=thermal_williamson_2_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + radius = 6371220. # planetary radius (m) + u_max = 20. # max amplitude of the zonal wind (m/s) + phi_0 = 3.0e4 # reference geopotential height (m^2/s^2) + epsilon = 1/300 # linear air expansion coeff (1/K) + theta_0 = epsilon*phi_0**2 # ref depth-integrated temperature (no units) + g = 9.80616 # acceleration due to gravity (m/s^2) + mean_depth = phi_0/g # reference depth (m) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + u_eqn_type = 'vector_advection_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + domain = Domain(mesh, dt, 'BDM', element_order) + x, y, z = SpatialCoordinate(mesh) + + # Equations + 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 + ) + + # IO + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dumplist_latlon=['D', 'D_error'], + dump_vtus=False, dump_nc=True + ) + + diagnostic_fields = [ + RelativeVorticity(), PotentialVorticity(), + SteadyStateError('u'), SteadyStateError('D'), SteadyStateError('b'), + MeridionalComponent('u'), ZonalComponent('u') + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "D", fixed_subcycles=2), + SSPRK3(domain, "b", fixed_subcycles=2) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "D"), + DGUpwind(eqns, "b") + ] + + # Linear solver + linear_solver = ThermalSWSolver(eqns) + + # Time stepper + 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") + + _, phi, _ = lonlatr_from_xyz(x, y, z) + + uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, (x, y, z)) + w = Omega*radius*u_max + (u_max**2)/2 + sigma = w/10 + + Dexpr = mean_depth - (1/g)*(w + sigma)*((sin(phi))**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 + bexpr = params.g * (1 - theta) + + u0.project(uexpr) + D0.interpolate(Dexpr) + b0.interpolate(bexpr) + + # 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', 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=thermal_williamson_2_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=thermal_williamson_2_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=thermal_williamson_2_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=thermal_williamson_2_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=thermal_williamson_2_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + thermal_williamson_2(**vars(args)) diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py index 16057b5b1..0bbd82851 100644 --- a/examples/shallow_water/williamson_2.py +++ b/examples/shallow_water/williamson_2.py @@ -1,101 +1,121 @@ """ -The Williamson 2 shallow-water test case (solid-body rotation), solved with a -discretisation of the non-linear shallow-water equations. +Test Case 2 (solid-body rotation with geostrophically-balanced flow) of +Williamson et al, 1992: +``A standard test set for numerical approximations to the shallow water +equations in spherical geometry'', JCP. -This uses an icosahedral mesh of the sphere, and runs a series of resolutions -to act as a convergence test. +The example here uses the icosahedral sphere mesh and degree 1 spaces. """ -from gusto import * -from firedrake import IcosahedralSphereMesh, SpatialCoordinate, sin, cos, pi -import sys +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import SpatialCoordinate, sin, cos, pi, Function +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, + RelativeVorticity, PotentialVorticity, SteadyStateError, + ShallowWaterKineticEnergy, ShallowWaterPotentialEnergy, + ShallowWaterPotentialEnstrophy, rotated_lonlatr_coords, + ZonalComponent, MeridionalComponent, rotated_lonlatr_vectors, + GeneralIcosahedralSphereMesh +) + +williamson_2_defaults = { + 'ncells_per_edge': 16, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 5.*24.*60.*60., # 5 days + 'dumpfreq': 96, # once per day with default options + 'dirname': 'williamson_2' +} + + +def williamson_2( + ncells_per_edge=williamson_2_defaults['ncells_per_edge'], + dt=williamson_2_defaults['dt'], + tmax=williamson_2_defaults['tmax'], + dumpfreq=williamson_2_defaults['dumpfreq'], + dirname=williamson_2_defaults['dirname'] +): -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -day = 24.*60.*60. -if '--running-tests' in sys.argv: - ref_dt = {3: 3000.} - tmax = 3000. - ndumps = 1 -else: - # setup resolution and timestepping parameters for convergence test - ref_dt = {3: 4000., 4: 2000., 5: 1000., 6: 500.} - tmax = 5*day - ndumps = 5 + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # -# setup shallow water parameters -R = 6371220. -H = 5960. -rotated_pole = (0.0, pi/3) + radius = 6371220. # planetary radius (m) + mean_depth = 5960. # reference depth (m) + rotate_pole_to = (0.0, pi/3) # location of North pole of mesh + u_max = 2*pi*radius/(12*24*60*60) # Max amplitude of the zonal wind (m/s) -# setup input that doesn't change with ref level or dt -parameters = ShallowWaterParameters(H=H) + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # -for ref_level, dt in ref_dt.items(): + element_order = 1 + u_eqn_type = 'vector_invariant_form' # ------------------------------------------------------------------------ # # Set up model objects # ------------------------------------------------------------------------ # # Domain - mesh = IcosahedralSphereMesh(radius=R, - refinement_level=ref_level, degree=2) - x = SpatialCoordinate(mesh) - domain = Domain(mesh, dt, 'BDM', 1, rotated_pole=rotated_pole) + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + domain = Domain(mesh, dt, 'BDM', element_order, rotated_pole=rotate_pole_to) + xyz = SpatialCoordinate(mesh) # Equation + parameters = ShallowWaterParameters(H=mean_depth) Omega = parameters.Omega - _, lat, _ = rotated_lonlatr_coords(x, rotated_pole) - e_lon, _, _ = rotated_lonlatr_vectors(x, rotated_pole) + _, lat, _ = rotated_lonlatr_coords(xyz, rotate_pole_to) + e_lon, _, _ = rotated_lonlatr_vectors(xyz, rotate_pole_to) fexpr = 2*Omega*sin(lat) - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) + eqns = ShallowWaterEquations( + domain, parameters, fexpr=fexpr, u_transport_option=u_eqn_type) # I/O - dirname = "williamson_2_ref%s_dt%s" % (ref_level, dt) - dumpfreq = int(tmax / (ndumps*dt)) output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, + dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dumplist_latlon=['D', 'D_error'], - dump_nc=True, ) - - diagnostic_fields = [RelativeVorticity(), SteadyStateError('RelativeVorticity'), - PotentialVorticity(), - ShallowWaterKineticEnergy(), - ShallowWaterPotentialEnergy(parameters), - ShallowWaterPotentialEnstrophy(), - SteadyStateError('u'), SteadyStateError('D'), - MeridionalComponent('u', rotated_pole), - ZonalComponent('u', rotated_pole)] + diagnostic_fields = [ + RelativeVorticity(), SteadyStateError('RelativeVorticity'), + PotentialVorticity(), ShallowWaterKineticEnergy(), + ShallowWaterPotentialEnergy(parameters), + ShallowWaterPotentialEnstrophy(), + SteadyStateError('u'), SteadyStateError('D'), + MeridionalComponent('u', rotate_pole_to), + ZonalComponent('u', rotate_pole_to) + ] io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes - transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "D", fixed_subcycles=2)] - transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "D", fixed_subcycles=2)] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "D") + ] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods) + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods + ) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # + g = parameters.g + u0 = stepper.fields("u") D0 = stepper.fields("D") - x = SpatialCoordinate(mesh) - u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) + uexpr = u_max*cos(lat)*e_lon - g = parameters.g - Dexpr = H - (R * Omega * u_max + u_max*u_max/2.0)*(sin(lat))**2/g + Dexpr = mean_depth - (radius * Omega * u_max + 0.5*u_max**2)*(sin(lat))**2/g u0.project(uexpr) D0.interpolate(Dexpr) - Dbar = Function(D0.function_space()).assign(H) + Dbar = Function(D0.function_space()).assign(mean_depth) stepper.set_reference_profiles([('D', Dbar)]) # ------------------------------------------------------------------------ # @@ -103,3 +123,48 @@ # ------------------------------------------------------------------------ # 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=williamson_2_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=williamson_2_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=williamson_2_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=williamson_2_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=williamson_2_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + williamson_2(**vars(args)) diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 366a5a903..cdb2b58a9 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -1,82 +1,95 @@ """ -The Williamson 5 shallow-water test case (flow over topography), solved with a -discretisation of the non-linear shallow-water equations. +Test Case 5 (flow over a mountain) of Williamson et al, 1992: +``A standard test set for numerical approximations to the shallow water +equations in spherical geometry'', JCP. -This uses an icosahedral mesh of the sphere, and runs a series of resolutions. +The example here uses the icosahedral sphere mesh and degree 1 spaces. """ -from gusto import * -from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, - as_vector, pi, sqrt, min_value) -import sys +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + SpatialCoordinate, as_vector, pi, sqrt, min_value, Function +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, Sum, + lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent, + MeridionalComponent, RelativeVorticity +) + +williamson_5_defaults = { + 'ncells_per_edge': 16, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 50.*24.*60.*60., # 50 days + 'dumpfreq': 960, # once per 10 days with default options + 'dirname': 'williamson_5' +} + + +def williamson_5( + ncells_per_edge=williamson_5_defaults['ncells_per_edge'], + dt=williamson_5_defaults['dt'], + tmax=williamson_5_defaults['tmax'], + dumpfreq=williamson_5_defaults['dumpfreq'], + dirname=williamson_5_defaults['dirname'] +): -# ---------------------------------------------------------------------------- # -# Test case parameters -# ---------------------------------------------------------------------------- # - -day = 24.*60.*60. -if '--running-tests' in sys.argv: - ref_dt = {3: 3000.} - tmax = 3000. - ndumps = 1 -else: - # setup resolution and timestepping parameters for convergence test - ref_dt = {3: 900., 4: 450., 5: 225., 6: 112.5} - tmax = 50*day - ndumps = 5 - -# setup shallow water parameters -R = 6371220. -H = 5960. + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # -# setup input that doesn't change with ref level or dt -parameters = ShallowWaterParameters(H=H) + radius = 6371220. # planetary radius (m) + mean_depth = 5960 # reference depth (m) + g = 9.80616 # acceleration due to gravity (m/s^2) + u_max = 20. # max amplitude of the zonal wind (m/s) + mountain_height = 2000. # height of mountain (m) + R0 = pi/9. # radius of mountain (rad) + lamda_c = -pi/2. # longitudinal centre of mountain (rad) + phi_c = pi/6. # latitudinal centre of mountain (rad) -for ref_level, dt in ref_dt.items(): + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + element_order = 1 # ------------------------------------------------------------------------ # # Set up model objects # ------------------------------------------------------------------------ # # Domain - mesh = IcosahedralSphereMesh(radius=R, - refinement_level=ref_level, degree=2) - x = SpatialCoordinate(mesh) - domain = Domain(mesh, dt, 'BDM', 1) + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + domain = Domain(mesh, dt, 'BDM', element_order) + x, y, z = SpatialCoordinate(mesh) + lamda, phi, _ = lonlatr_from_xyz(x, y, z) - # Equation + # Equation: coriolis + parameters = ShallowWaterParameters(H=mean_depth, g=g) Omega = parameters.Omega - fexpr = 2*Omega*x[2]/R - lamda, theta, _ = lonlatr_from_xyz(x[0], x[1], x[2]) - R0 = pi/9. - R0sq = R0**2 - lamda_c = -pi/2. - lsq = (lamda - lamda_c)**2 - theta_c = pi/6. - thsq = (theta - theta_c)**2 - rsq = min_value(R0sq, lsq+thsq) + fexpr = 2*Omega*z/radius + + # Equation: topography + rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2) r = sqrt(rsq) - bexpr = 2000 * (1 - r/R0) - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=bexpr) + tpexpr = mountain_height * (1 - r/R0) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr) # I/O - dirname = "williamson_5_ref%s_dt%s" % (ref_level, dt) - dumpfreq = int(tmax / (ndumps*dt)) output = OutputParameters( - dirname=dirname, - dumplist_latlon=['D'], - dumpfreq=dumpfreq, + dirname=dirname, dumplist_latlon=['D'], dumpfreq=dumpfreq, + dump_vtus=True, dump_nc=False, dumplist=['D', 'topography'] ) - diagnostic_fields = [Sum('D', 'topography')] + diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), + MeridionalComponent('u'), ZonalComponent('u')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes - transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "D")] + transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")] transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods) + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods + ) # ------------------------------------------------------------------------ # # Initial conditions @@ -84,16 +97,16 @@ u0 = stepper.fields('u') D0 = stepper.fields('D') - u_max = 20. # Maximum amplitude of the zonal wind (m/s) - uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) - g = parameters.g - Rsq = R**2 - Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - bexpr + uexpr = as_vector([-u_max*y/radius, u_max*x/radius, 0.0]) + Dexpr = ( + mean_depth - tpexpr + - (radius*Omega*u_max + 0.5*u_max**2)*(z/radius)**2/g + ) u0.project(uexpr) D0.interpolate(Dexpr) - Dbar = Function(D0.function_space()).assign(H) + Dbar = Function(D0.function_space()).assign(mean_depth) stepper.set_reference_profiles([('D', Dbar)]) # ------------------------------------------------------------------------ # @@ -101,3 +114,48 @@ # ------------------------------------------------------------------------ # 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=williamson_5_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=williamson_5_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=williamson_5_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=williamson_5_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=williamson_5_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + williamson_5(**vars(args)) diff --git a/examples/test_examples_run.py b/examples/test_examples_run.py deleted file mode 100644 index f69dec270..000000000 --- a/examples/test_examples_run.py +++ /dev/null @@ -1,38 +0,0 @@ -import pytest -from os.path import abspath, basename, dirname -import subprocess -import glob -import sys -import os - - -examples_dir = abspath(dirname(__file__)) -example_files = glob.glob("%s/*/*.py" % examples_dir) - - -@pytest.fixture(params=glob.glob("%s/*/*.py" % examples_dir), - ids=lambda x: basename(x)) -def example_file(request): - return abspath(request.param) - - -def test_example_runs(example_file, tmpdir, monkeypatch): - # This ensures that the test writes output in a temporary - # directory, rather than where pytest was run from. - monkeypatch.chdir(tmpdir) - subprocess.run( - [sys.executable, example_file, "--running-tests"], - check=True, - env=os.environ | {"PYOP2_CFLAGS": "-O0"} - ) - - -def test_example_runs_parallel(example_file, tmpdir, monkeypatch): - # This ensures that the test writes output in a temporary - # directory, rather than where pytest was run from. - monkeypatch.chdir(tmpdir) - subprocess.run( - ["mpiexec", "-n", "4", sys.executable, example_file, "--running-tests"], - check=True, - env=os.environ | {"PYOP2_CFLAGS": "-O0"} - ) diff --git a/figures/boussinesq/skamarock_klemp_compressible_bouss_final.png b/figures/boussinesq/skamarock_klemp_compressible_bouss_final.png new file mode 100644 index 000000000..4107d44ec Binary files /dev/null and b/figures/boussinesq/skamarock_klemp_compressible_bouss_final.png differ diff --git a/figures/boussinesq/skamarock_klemp_compressible_bouss_initial.png b/figures/boussinesq/skamarock_klemp_compressible_bouss_initial.png new file mode 100644 index 000000000..ef961cc01 Binary files /dev/null and b/figures/boussinesq/skamarock_klemp_compressible_bouss_initial.png differ diff --git a/figures/boussinesq/skamarock_klemp_incompressible_bouss_final.png b/figures/boussinesq/skamarock_klemp_incompressible_bouss_final.png new file mode 100644 index 000000000..622266ca1 Binary files /dev/null and b/figures/boussinesq/skamarock_klemp_incompressible_bouss_final.png differ diff --git a/figures/boussinesq/skamarock_klemp_incompressible_bouss_initial.png b/figures/boussinesq/skamarock_klemp_incompressible_bouss_initial.png new file mode 100644 index 000000000..ef961cc01 Binary files /dev/null and b/figures/boussinesq/skamarock_klemp_incompressible_bouss_initial.png differ diff --git a/figures/boussinesq/skamarock_klemp_linear_bouss_final.png b/figures/boussinesq/skamarock_klemp_linear_bouss_final.png new file mode 100644 index 000000000..c893cf013 Binary files /dev/null and b/figures/boussinesq/skamarock_klemp_linear_bouss_final.png differ diff --git a/figures/boussinesq/skamarock_klemp_linear_bouss_initial.png b/figures/boussinesq/skamarock_klemp_linear_bouss_initial.png new file mode 100644 index 000000000..ef961cc01 Binary files /dev/null and b/figures/boussinesq/skamarock_klemp_linear_bouss_initial.png differ diff --git a/figures/compressible_euler/dcmip_3_1_gravity_wave_final.png b/figures/compressible_euler/dcmip_3_1_gravity_wave_final.png new file mode 100644 index 000000000..d9864b820 Binary files /dev/null and b/figures/compressible_euler/dcmip_3_1_gravity_wave_final.png differ diff --git a/figures/compressible_euler/dcmip_3_1_gravity_wave_initial.png b/figures/compressible_euler/dcmip_3_1_gravity_wave_initial.png new file mode 100644 index 000000000..f6febb2e8 Binary files /dev/null and b/figures/compressible_euler/dcmip_3_1_gravity_wave_initial.png differ diff --git a/figures/compressible_euler/dry_bryan_fritsch.png b/figures/compressible_euler/dry_bryan_fritsch.png new file mode 100644 index 000000000..b763f4275 Binary files /dev/null and b/figures/compressible_euler/dry_bryan_fritsch.png differ diff --git a/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png b/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png new file mode 100644 index 000000000..c1187a1cd Binary files /dev/null and b/figures/compressible_euler/skamarock_klemp_nonhydrostatic_final.png differ diff --git a/figures/compressible_euler/skamarock_klemp_nonhydrostatic_initial.png b/figures/compressible_euler/skamarock_klemp_nonhydrostatic_initial.png new file mode 100644 index 000000000..94c1eb09f Binary files /dev/null and b/figures/compressible_euler/skamarock_klemp_nonhydrostatic_initial.png differ diff --git a/figures/compressible_euler/straka_bubble.png b/figures/compressible_euler/straka_bubble.png new file mode 100644 index 000000000..395bb9105 Binary files /dev/null and b/figures/compressible_euler/straka_bubble.png differ diff --git a/figures/compressible_euler/unsaturated_bubble_final.png b/figures/compressible_euler/unsaturated_bubble_final.png new file mode 100644 index 000000000..4f32dc099 Binary files /dev/null and b/figures/compressible_euler/unsaturated_bubble_final.png differ diff --git a/figures/compressible_euler/unsaturated_bubble_initial.png b/figures/compressible_euler/unsaturated_bubble_initial.png new file mode 100644 index 000000000..993a75dfb Binary files /dev/null and b/figures/compressible_euler/unsaturated_bubble_initial.png differ diff --git a/figures/shallow_water/linear_williamson_2_final.png b/figures/shallow_water/linear_williamson_2_final.png new file mode 100644 index 000000000..bd326d3e2 Binary files /dev/null and b/figures/shallow_water/linear_williamson_2_final.png differ diff --git a/figures/shallow_water/linear_williamson_2_initial.png b/figures/shallow_water/linear_williamson_2_initial.png new file mode 100644 index 000000000..992d3b5b3 Binary files /dev/null and b/figures/shallow_water/linear_williamson_2_initial.png differ diff --git a/figures/shallow_water/moist_convective_williamson_2_final.png b/figures/shallow_water/moist_convective_williamson_2_final.png new file mode 100644 index 000000000..c85b82415 Binary files /dev/null and b/figures/shallow_water/moist_convective_williamson_2_final.png differ diff --git a/figures/shallow_water/moist_convective_williamson_2_initial.png b/figures/shallow_water/moist_convective_williamson_2_initial.png new file mode 100644 index 000000000..07e171d25 Binary files /dev/null and b/figures/shallow_water/moist_convective_williamson_2_initial.png differ diff --git a/figures/shallow_water/moist_thermal_williamson_5_final.png b/figures/shallow_water/moist_thermal_williamson_5_final.png new file mode 100644 index 000000000..cc63024af Binary files /dev/null and b/figures/shallow_water/moist_thermal_williamson_5_final.png differ diff --git a/figures/shallow_water/moist_thermal_williamson_5_initial.png b/figures/shallow_water/moist_thermal_williamson_5_initial.png new file mode 100644 index 000000000..3d9c8486f Binary files /dev/null and b/figures/shallow_water/moist_thermal_williamson_5_initial.png differ diff --git a/figures/shallow_water/shallow_water_1d_wave.png b/figures/shallow_water/shallow_water_1d_wave.png new file mode 100644 index 000000000..8496e625e Binary files /dev/null and b/figures/shallow_water/shallow_water_1d_wave.png differ diff --git a/figures/shallow_water/thermal_williamson_2_final.png b/figures/shallow_water/thermal_williamson_2_final.png new file mode 100644 index 000000000..705a1f436 Binary files /dev/null and b/figures/shallow_water/thermal_williamson_2_final.png differ diff --git a/figures/shallow_water/thermal_williamson_2_initial.png b/figures/shallow_water/thermal_williamson_2_initial.png new file mode 100644 index 000000000..0f53c1da5 Binary files /dev/null and b/figures/shallow_water/thermal_williamson_2_initial.png differ diff --git a/figures/shallow_water/williamson_2_final.png b/figures/shallow_water/williamson_2_final.png new file mode 100644 index 000000000..77fbdc8b3 Binary files /dev/null and b/figures/shallow_water/williamson_2_final.png differ diff --git a/figures/shallow_water/williamson_2_initial.png b/figures/shallow_water/williamson_2_initial.png new file mode 100644 index 000000000..dee716023 Binary files /dev/null and b/figures/shallow_water/williamson_2_initial.png differ diff --git a/figures/shallow_water/williamson_5_final.png b/figures/shallow_water/williamson_5_final.png new file mode 100644 index 000000000..dc23560b7 Binary files /dev/null and b/figures/shallow_water/williamson_5_final.png differ diff --git a/figures/shallow_water/williamson_5_initial.png b/figures/shallow_water/williamson_5_initial.png new file mode 100644 index 000000000..ae18f58f6 Binary files /dev/null and b/figures/shallow_water/williamson_5_initial.png differ diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index 73b80b7d6..84b187e47 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -279,7 +279,7 @@ class HydrostaticCompressibleEulerEquations(CompressibleEulerEquations): equations. """ - def __init__(self, domain, parameters, sponge=None, + def __init__(self, domain, parameters, sponge_options=None, extra_terms=None, space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, @@ -291,8 +291,9 @@ def __init__(self, domain, parameters, sponge=None, mesh and the compatible function spaces. parameters (:class:`Configuration`, optional): an object containing the model's physical parameters. - sponge (:class:`ufl.Expr`, optional): an expression for a sponge - layer. Defaults to None. + sponge_options (:class:`SpongeLayerParameters`, optional): any + parameters for applying a sponge layer to the upper boundary. + Defaults to None. extra_terms (:class:`ufl.Expr`, optional): any extra terms to be included in the equation set. Defaults to None. space_names (dict, optional): a dictionary of strings for names of @@ -323,7 +324,7 @@ def __init__(self, domain, parameters, sponge=None, NotImplementedError: only mixing ratio tracers are implemented. """ - super().__init__(domain, parameters, sponge=sponge, + super().__init__(domain, parameters, sponge_options=sponge_options, extra_terms=extra_terms, space_names=space_names, linearisation_map=linearisation_map, u_transport_option=u_transport_option, @@ -331,42 +332,49 @@ def __init__(self, domain, parameters, sponge=None, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) + # Replace self.residual = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=lambda t: hydrostatic(t, self.hydrostatic_projection(t)) + map_if_true=lambda t: self.hydrostatic_projection(t, 'u') ) + # Add an extra hydrostatic term + u_idx = self.field_names.index('u') + u = split(self.X)[u_idx] k = self.domain.k - u = split(self.X)[0] self.residual += hydrostatic( subject( prognostic( - -inner(k, self.tests[0]) * inner(k, u) * dx, "u"), - self.X)) + -inner(k, self.tests[u_idx]) * inner(k, u) * dx, "u"), + self.X + ) + ) - def hydrostatic_projection(self, t): + def hydrostatic_projection(self, term, field_name): """ Performs the 'hydrostatic' projection. Takes a term involving a vector prognostic variable and replaces the - prognostic with only its horizontal components. + prognostic with only its horizontal components. It also adds the + 'hydrostatic' label to that term. Args: - t (:class:`Term`): the term to perform the projection upon. + term (:class:`Term`): the term to perform the projection upon. + field_name (str): the prognostic field to make hydrostatic. Returns: :class:`LabelledForm`: the labelled form containing the new term. - - Raises: - AssertionError: spherical geometry is not yet implemented. """ - # TODO: make this more general, i.e. should work on the sphere - if self.domain.on_sphere: - raise NotImplementedError("The hydrostatic projection is not yet " - + "implemented for spherical geometry") - k = Constant((*self.domain.k, 0, 0)) - X = t.get(subject) - - new_subj = X - k * inner(X, k) - return replace_subject(new_subj)(t) + f_idx = self.field_names.index(field_name) + k = self.domain.k + X = term.get(subject) + field = split(X)[f_idx] + + new_subj = field - inner(field, k) * k + # In one step: + # - set up the replace_subject routine (which returns a function) + # - call that function on the supplied `term` argument, + # to replace the subject with the new hydrostatic subject + # - add the hydrostatic label + return replace_subject(new_subj, old_idx=f_idx)(term) diff --git a/gusto/physics/shallow_water_microphysics.py b/gusto/physics/shallow_water_microphysics.py index 8ffdcf7da..27f5c09b2 100644 --- a/gusto/physics/shallow_water_microphysics.py +++ b/gusto/physics/shallow_water_microphysics.py @@ -313,7 +313,7 @@ def __init__(self, equation, saturation_curve, if convective_feedback: factors.append(self.gamma_v*beta1) if thermal_feedback: - factors.append(parameters.g*self.gamma_v*beta2) + factors.append(self.gamma_v*beta2) # Add terms to equations and make interpolators self.source = [Function(Vc) for factor in factors] diff --git a/integration-tests/physics/test_sw_saturation_adjustment.py b/integration-tests/physics/test_sw_saturation_adjustment.py index 3bec4bb0f..66c971519 100644 --- a/integration-tests/physics/test_sw_saturation_adjustment.py +++ b/integration-tests/physics/test_sw_saturation_adjustment.py @@ -88,13 +88,12 @@ def run_sw_cond_evap(dirname, process): v_true = Function(v0.function_space()).interpolate(sat*(0.96+0.005*pert)) c_true = Function(c0.function_space()).interpolate(Constant(0.0)) # gain buoyancy - factor = parameters.g*beta2 sat_adj_expr = (v0 - sat) / dt sat_adj_expr = conditional(sat_adj_expr < 0, max_value(sat_adj_expr, -c0 / dt), min_value(sat_adj_expr, v0 / dt)) # include factor of -1 in true solution to compare term to LHS in Gusto - b_true = Function(b0.function_space()).interpolate(-dt*sat_adj_expr*factor) + b_true = Function(b0.function_space()).interpolate(-dt*sat_adj_expr*beta2) elif process == "condensation": # vapour is above saturation @@ -103,13 +102,12 @@ def run_sw_cond_evap(dirname, process): v_true = Function(v0.function_space()).interpolate(Constant(sat)) c_true = Function(c0.function_space()).interpolate(v0 - sat) # lose buoyancy - factor = parameters.g*beta2 sat_adj_expr = (v0 - sat) / dt sat_adj_expr = conditional(sat_adj_expr < 0, max_value(sat_adj_expr, -c0 / dt), min_value(sat_adj_expr, v0 / dt)) # include factor of -1 in true solution to compare term to LHS in Gusto - b_true = Function(b0.function_space()).interpolate(-dt*sat_adj_expr*factor) + b_true = Function(b0.function_space()).interpolate(-dt*sat_adj_expr*beta2) c_init = Function(c0.function_space()).interpolate(c0) diff --git a/plotting/boussinesq/plot_skamarock_klemp_boussinesq.py b/plotting/boussinesq/plot_skamarock_klemp_boussinesq.py new file mode 100644 index 000000000..6ac9a11cd --- /dev/null +++ b/plotting/boussinesq/plot_skamarock_klemp_boussinesq.py @@ -0,0 +1,175 @@ +""" +Plots the Boussinesq Skamarock-Klemp gravity wave in a vertical slice. This can +plot the compressible, incompressible and linear cases. + +This plots the initial conditions @ t = 0 s, with +(a) buoyancy perturbation, (b) buoyancy +and the final state @ t = 3600 s, with +(a) buoyancy perturbation, +(b) a 1D slice through the wave +""" +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, tomplot_field_title, extract_gusto_coords, + extract_gusto_field, reshape_gusto_data, add_colorbar_fig +) + +# Can be incompressible/compressible/linear +test = 'skamarock_klemp_linear_bouss' + +# ---------------------------------------------------------------------------- # +# 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/boussinesq/{test}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['b_perturbation', 'b'] +init_colour_schemes = ['YlOrRd', 'Purples'] +init_field_labels = [r'$\Delta b$ (m s$^{-2}$)', r'$b$ (m s$^{-2}$)'] +init_contours = [np.linspace(0.0, 0.01, 11), np.linspace(0, 1, 11)] +init_contours_to_remove = [None, None] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_name = 'b_perturbation' +final_colour_scheme = 'RdBu_r' +final_field_label = r'$\Delta b$ (m s$^{-2}$)' +final_contours = np.linspace(-3.0e-3, 3.0e-3, 13) +final_contour_to_remove = 0.0 + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [0, 300.0] +ylims = [0, 10.0] + +# 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=(12, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, field_label, colour_scheme, contours, to_remove) in \ + enumerate(zip(axarray.flatten(), init_field_names, init_field_labels, + init_colour_schemes, init_contours, init_contours_to_remove)): + + # 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] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10 + ) + tomplot_field_title( + ax, f't = {time:.1f} s', 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.subplots_adjust(wspace=0.2) +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(2, 1, figsize=(8, 8), sharex='all') +time_idx = -1 + +# Data extraction ---------------------------------------------------------- +field_data = extract_gusto_field(data_file, final_field_name, time_idx=time_idx) +coords_X, coords_Y = extract_gusto_coords(data_file, final_field_name) +time = data_file['time'][time_idx] + +# Plot 2D data ----------------------------------------------------------------- +ax = axarray[0] + +cmap, lines = tomplot_cmap( + final_contours, final_colour_scheme, remove_contour=final_contour_to_remove +) +cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, final_contours, + cmap=cmap, line_contours=lines +) + +add_colorbar_fig( + fig, cf, final_field_label, ax_idxs=[0], location='right', cbar_labelpad=-40 +) +tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data +) + +ax.set_ylabel(r'$z$ (km)', labelpad=-20) +ax.set_ylim(ylims) +ax.set_yticks(ylims) +ax.set_yticklabels(ylims) + +# Plot 1D data ----------------------------------------------------------------- +ax = axarray[1] + +field_data, coords_X, coords_Y = reshape_gusto_data(field_data, coords_X, coords_Y) + +# Determine midpoint index +mid_idx = np.floor_divide(np.shape(field_data)[1], 2) +slice_height = coords_Y[0, mid_idx] + +ax.plot(coords_X[:, mid_idx], field_data[:, mid_idx], color='black') + +tomplot_field_title( + ax, r'$z$' + f' = {slice_height} km' +) + +b_lims = [np.min(final_contours), np.max(final_contours)] + +ax.set_ylabel(final_field_label, labelpad=-20) +ax.set_ylim(b_lims) +ax.set_yticks(b_lims) +ax.set_yticklabels(b_lims) + +ax.set_xlabel(r'$x$ (km)', labelpad=-10) +ax.set_xlim(xlims) +ax.set_xticks(xlims) +ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(hspace=0.2) +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_dry_bryan_fritsch.py b/plotting/compressible_euler/plot_dry_bryan_fritsch.py new file mode 100644 index 000000000..f696ac23f --- /dev/null +++ b/plotting/compressible_euler/plot_dry_bryan_fritsch.py @@ -0,0 +1,93 @@ +""" +Plots the dry Bryan and Fritsch rising bubble test case. + +This plots: +(a) theta perturbation @ t = 0 s, (b) theta perturbation @ t = 1000 s +""" +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_fig, tomplot_field_title, extract_gusto_coords, + extract_gusto_field +) + +test = 'dry_bryan_fritsch' + +# ---------------------------------------------------------------------------- # +# 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}' + +# ---------------------------------------------------------------------------- # +# Plot details +# ---------------------------------------------------------------------------- # +field_names = ['theta_perturbation', 'theta_perturbation'] +time_idxs = [0, -1] +cbars = [False, True] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contours = np.linspace(-0.5, 2.5, 13) +colour_scheme = 'OrRd' +field_label = r'$\Delta \theta$ (K)' +contour_method = 'tricontour' +xlims = [0, 10] +ylims = [0, 10] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(16, 6), sharex='all', sharey='all') + +for i, (ax, time_idx, field_name, cbar) in \ + enumerate(zip(axarray.flatten(), time_idxs, field_names, cbars)): + + # 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] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=0.0) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + if cbar: + add_colorbar_fig( + fig, cf, field_label, ax_idxs=[i], location='right' + ) + tomplot_field_title( + ax, f't = {time:.1f} s', 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.subplots_adjust(wspace=0.15) +plot_name = f'{plot_stem}.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 new file mode 100644 index 000000000..668d75dc8 --- /dev/null +++ b/plotting/compressible_euler/plot_skamarock_klemp_nonhydrostatic.py @@ -0,0 +1,173 @@ +""" +Plots the Skamarock-Klemp gravity wave in a vertical slice. + +This plots the initial conditions @ t = 0 s, with +(a) theta perturbation, (b) theta +and the final state @ t = 3600 s, with +(a) theta perturbation, +(b) a 1D slice through the wave +""" +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, tomplot_field_title, extract_gusto_coords, + extract_gusto_field, reshape_gusto_data, add_colorbar_fig +) + +test = 'skamarock_klemp_nonhydrostatic' + +# ---------------------------------------------------------------------------- # +# 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}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['theta_perturbation', 'theta'] +init_colour_schemes = ['YlOrRd', 'Purples'] +init_field_labels = [r'$\Delta\theta$ (K)', r'$\theta$ (K)'] +init_contours = [np.linspace(0.0, 0.01, 11), np.linspace(300, 335, 8)] +init_contours_to_remove = [None, None] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_name = 'theta_perturbation' +final_colour_scheme = 'RdBu_r' +final_field_label = r'$\Delta\theta$ (K)' +final_contours = np.linspace(-3.0e-3, 3.0e-3, 13) +final_contour_to_remove = 0.0 + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [0, 300.0] +ylims = [0, 10.0] + +# 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=(12, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, field_label, colour_scheme, contours, to_remove) in \ + enumerate(zip(axarray.flatten(), init_field_names, init_field_labels, + init_colour_schemes, init_contours, init_contours_to_remove)): + + # 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] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10 + ) + tomplot_field_title( + ax, f't = {time:.1f} s', 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.subplots_adjust(wspace=0.2) +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(2, 1, figsize=(8, 8), sharex='all') +time_idx = -1 + +# Data extraction ---------------------------------------------------------- +field_data = extract_gusto_field(data_file, final_field_name, time_idx=time_idx) +coords_X, coords_Y = extract_gusto_coords(data_file, final_field_name) +time = data_file['time'][time_idx] + +# Plot 2D data ----------------------------------------------------------------- +ax = axarray[0] + +cmap, lines = tomplot_cmap( + final_contours, final_colour_scheme, remove_contour=final_contour_to_remove +) +cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, final_contours, + cmap=cmap, line_contours=lines +) + +add_colorbar_fig( + fig, cf, final_field_label, ax_idxs=[0], location='right', cbar_labelpad=-40 +) +tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data +) + +ax.set_ylabel(r'$z$ (km)', labelpad=-20) +ax.set_ylim(ylims) +ax.set_yticks(ylims) +ax.set_yticklabels(ylims) + +# Plot 1D data ----------------------------------------------------------------- +ax = axarray[1] + +field_data, coords_X, coords_Y = reshape_gusto_data(field_data, coords_X, coords_Y) + +# Determine midpoint index +mid_idx = np.floor_divide(np.shape(field_data)[1], 2) +slice_height = coords_Y[0, mid_idx] + +ax.plot(coords_X[:, mid_idx], field_data[:, mid_idx], color='black') + +tomplot_field_title( + ax, r'$z$' + f' = {slice_height} km' +) + +theta_lims = [np.min(final_contours), np.max(final_contours)] + +ax.set_ylabel(final_field_label, labelpad=-20) +ax.set_ylim(theta_lims) +ax.set_yticks(theta_lims) +ax.set_yticklabels(theta_lims) + +ax.set_xlabel(r'$x$ (km)', labelpad=-10) +ax.set_xlim(xlims) +ax.set_xticks(xlims) +ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(hspace=0.2) +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_straka_bubble.py b/plotting/compressible_euler/plot_straka_bubble.py new file mode 100644 index 000000000..557de8bb6 --- /dev/null +++ b/plotting/compressible_euler/plot_straka_bubble.py @@ -0,0 +1,93 @@ +""" +Plots the Straka bubble test case. + +This plots: +(a) theta perturbation @ t = 0 s, (b) theta perturbation @ t = 900 s +""" +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_fig, tomplot_field_title, extract_gusto_coords, + extract_gusto_field +) + +test = 'straka_bubble' + +# ---------------------------------------------------------------------------- # +# 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}' + +# ---------------------------------------------------------------------------- # +# Plot details +# ---------------------------------------------------------------------------- # +field_names = ['theta_perturbation', 'theta_perturbation'] +time_idxs = [0, -1] +cbars = [False, True] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contours = np.linspace(-7.5, 0.5, 17) +colour_scheme = 'Blues_r' +field_label = r'$\Delta \theta$ (K)' +contour_method = 'tricontour' +xlims = [25.6, 38.4] +ylims = [0., 5.0] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(18, 6), sharex='all', sharey='all') + +for i, (ax, time_idx, field_name, cbar) in \ + enumerate(zip(axarray.flatten(), time_idxs, field_names, cbars)): + + # 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] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=0.0) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, negative_linestyles='solid' + ) + + if cbar: + add_colorbar_fig( + fig, cf, field_label, ax_idxs=[i], location='right' + ) + tomplot_field_title( + ax, f't = {time:.1f} s', 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.subplots_adjust(wspace=0.15) +plot_name = f'{plot_stem}.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/compressible_euler/plot_unsaturated_bubble.py b/plotting/compressible_euler/plot_unsaturated_bubble.py new file mode 100644 index 000000000..3d1240c02 --- /dev/null +++ b/plotting/compressible_euler/plot_unsaturated_bubble.py @@ -0,0 +1,155 @@ +""" +Plots the unsaturated moist rising bubble test case, which features rain. + +This plots the initial conditions @ t = 0 s, with +(a) theta perturbation, (b) relative humidity, +and the final state @ t = 600 s, with +(a) theta perturbation, (b) relative humidity and (c) rain mixing ratio +""" +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, tomplot_field_title, extract_gusto_coords, + extract_gusto_field +) + +test = 'unsaturated_bubble' + +# ---------------------------------------------------------------------------- # +# 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}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['theta_perturbation', 'RelativeHumidity'] +init_colour_schemes = ['Reds', 'Blues'] +init_field_labels = [r'$\Delta\theta$ (K)', 'Relative Humidity'] +init_contours = [np.linspace(-0.25, 3.0, 14), np.linspace(0.0, 1.1, 12)] +init_contours_to_remove = [0.0, 0.2] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['theta_perturbation', 'RelativeHumidity', 'rain'] +final_colour_schemes = ['RdBu_r', 'Blues', 'Purples'] +final_field_labels = [r'$\Delta\theta$ (K)', 'Relative Humidity', r'$m_r$ (kg/kg)'] +final_contours = [np.linspace(-3.5, 3.5, 15), + np.linspace(0.0, 1.1, 12), + np.linspace(-2.5e-6, 5.0e-5, 12)] +final_contours_to_remove = [0.0, None, None] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [0, 3.6] +ylims = [0, 2.4] + +# 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=(12, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, field_label, colour_scheme, contours, to_remove) in \ + enumerate(zip(axarray.flatten(), init_field_names, init_field_labels, + init_colour_schemes, init_contours, init_contours_to_remove)): + + # 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] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10 + ) + tomplot_field_title( + ax, f't = {time:.1f} s', 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.subplots_adjust(wspace=0.15) +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, 3, figsize=(18, 6), sharex='all', sharey='all') +time_idx = -1 + +for i, (ax, field_name, field_label, colour_scheme, contours, to_remove) in \ + enumerate(zip(axarray.flatten(), final_field_names, final_field_labels, + final_colour_schemes, final_contours, + final_contours_to_remove)): + + # 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] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10 + ) + tomplot_field_title( + ax, f't = {time:.1f} s', 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.subplots_adjust(wspace=0.18) +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close()