From 2f32f672e56162328e331c9d6bb7c624875ffd76 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 8 Oct 2023 18:46:06 +0100 Subject: [PATCH] try to get Eady code working with all of the new objects --- examples/compressible/compressible_eady.py | 145 ++++--- .../incompressible/incompressible_eady.py | 118 ++++-- gusto/__init__.py | 1 - gusto/configuration.py | 5 +- gusto/diagnostics.py | 361 +++++++++++++++++- gusto/eady_diagnostics.py | 241 ------------ gusto/equations.py | 129 +++++-- gusto/initialisation_tools.py | 64 +--- 8 files changed, 632 insertions(+), 432 deletions(-) delete mode 100644 gusto/eady_diagnostics.py diff --git a/examples/compressible/compressible_eady.py b/examples/compressible/compressible_eady.py index c1760b44b..029700009 100644 --- a/examples/compressible/compressible_eady.py +++ b/examples/compressible/compressible_eady.py @@ -4,14 +4,22 @@ from gusto import * from gusto import thermodynamics -from firedrake import (as_vector, SpatialCoordinate, +from firedrake import (as_vector, SpatialCoordinate, solve, ds_b, ds_t, PeriodicRectangleMesh, ExtrudedMesh, exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + day = 24.*60.*60. hour = 60.*60. dt = 30. +L = 1000000. +H = 10000. # Height position of the model top +f = 1.e-04 + if '--running-tests' in sys.argv: tmax = dt tdump = dt @@ -23,54 +31,69 @@ columns = 30 # number of columns nlayers = 30 # horizontal layers -# set up mesh -L = 1000000. -m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True) +dirname = 'compressible_eady' -# build 2D mesh by extruding the base mesh -H = 10000. # Height position of the model top +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain -- 2D periodic base mesh which is one cell thick +m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "RTCF", 1) -# Coriolis expression -f = 1.e-04 +# Equation Omega = as_vector([0., 0., f*0.5]) +parameters = CompressibleEadyParameters(H=H, f=f) +eqns = CompressibleEadyEquations(domain, parameters, Omega=Omega) -dirname = 'compressible_eady' +# I/O output = OutputParameters(dirname=dirname, - dumpfreq=int(tdump/dt), - dumplist=['u', 'rho', 'theta'], - perturbation_fields=['rho', 'theta', 'ExnerPi'], - log_level='INFO') + dumpfreq=int(tdump/dt)) -parameters = CompressibleEadyParameters(H=H, f=f) - -diagnostic_fields = [CourantNumber(), VelocityY(), - ExnerPi(), ExnerPi(reference=True), +diagnostic_fields = [CourantNumber(), YComponent('u'), + Exner(parameters), Exner(parameters, reference=True), CompressibleKineticEnergy(), CompressibleKineticEnergyY(), - CompressibleEadyPotentialEnergy(), + CompressibleEadyPotentialEnergy(parameters), Sum("CompressibleKineticEnergy", "CompressibleEadyPotentialEnergy"), Difference("CompressibleKineticEnergy", - "CompressibleKineticEnergyY")] - -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) - -eqns = CompressibleEadyEquations(state, "RTCF", 1) - + "CompressibleKineticEnergyY"), + Perturbation('rho'), Perturbation('theta'), + Perturbation('Exner')] + +io = IO(domain, output, diagnostic_fields=diagnostic_fields) + +# Transport schemes and methods +theta_opts = SUPGOptions() +transport_schemes = [SSPRK3(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, transport_schemes, + transport_methods, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # first setup the background buoyancy profile # z.grad(bref) = N**2 @@ -110,39 +133,51 @@ def n(): # calculate hydrostatic Pi rho_b = Function(Vr) -compressible_hydrostatic_balance(state, theta_b, rho_b) -compressible_hydrostatic_balance(state, theta0, rho0) +compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) +compressible_hydrostatic_balance(eqns, theta0, rho0, solve_for_rho=True) -# set Pi0 -Pi0 = calculate_Pi0(state, theta0, rho0) -state.parameters.Pi0 = Pi0 +# set Pi0 -- have to get this from the equations +Pi0 = eqns.prescribed_fields('Pi0') +Pi0.interpolate(exner_pressure(parameters, rho0, theta0)) # set x component of velocity -cp = state.parameters.cp -dthetady = state.parameters.dthetady -Pi = thermodynamics.pi(state.parameters, rho0, theta0) +cp = parameters.cp +dthetady = parameters.dthetady +Pi = thermodynamics.exner_pressure(parameters, rho0, theta0) u = cp*dthetady/f*(Pi-Pi0) -# set y component of velocity +# set y component of velocity by solving a problem v = Function(Vr).assign(0.) -compressible_eady_initial_v(state, theta0, rho0, v) + +# get Pi gradient +g = TrialFunction(Vu) +wg = TestFunction(Vu) + +n = FacetNormal(mesh) + +a = inner(wg, g)*dx +L = -div(wg)*Pi*dx + inner(wg, n)*Pi*(ds_t + ds_b) +pgrad = Function(Vu) +solve(a == L, pgrad) + +# get initial v +m = TrialFunction(Vr) +phi = TestFunction(Vr) + +a = phi*f*m*dx +L = phi*cp*theta0*pgrad[0]*dx +solve(a == L, v) # set initial u u_exp = as_vector([u, v, 0.]) u0.project(u_exp) # set the background profiles -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -transported_fields = [SSPRK3(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta")] - -linear_solver = CompressibleSolver(state, eqns) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -stepper = CrankNicolson(state, eqns, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/incompressible/incompressible_eady.py b/examples/incompressible/incompressible_eady.py index 1ba9f82ed..a7ad6abcf 100644 --- a/examples/incompressible/incompressible_eady.py +++ b/examples/incompressible/incompressible_eady.py @@ -3,11 +3,15 @@ """ from gusto import * -from firedrake import (as_vector, SpatialCoordinate, +from firedrake import (as_vector, SpatialCoordinate, solve, ds_t, ds_b, PeriodicRectangleMesh, ExtrudedMesh, cos, sin, cosh, sinh, tanh, pi, Function, sqrt) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + day = 24.*60.*60. hour = 60.*60. dt = 100. @@ -16,7 +20,7 @@ tmax = dt tdump = dt columns = 10 - nlayers = 10 + nlayers = 5 else: tmax = 30*day tdump = 2*hour @@ -32,49 +36,66 @@ f = f/beta L = beta*L -# Construct 2D periodic base mesh -m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True) - -# build 3D mesh by extruding the base mesh -mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +dirname = 'incompressible_eady' +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -output = OutputParameters(dirname='incompressible_eady', - dumpfreq=int(tdump/dt), - perturbation_fields=['p', 'b'], - log_level='INFO') +# Domain -- 2D periodic base mesh which is one cell thick +m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True) +mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "RTCF", 1) +# Equation +Omega = as_vector([0., 0., f*0.5]) parameters = EadyParameters(H=H, L=L, f=f, deltax=2.*L/float(columns), deltaz=H/float(nlayers), fourthorder=True) +eqns = IncompressibleEadyEquations(domain, parameters, Omega=Omega) + +# I/O +output = OutputParameters(dirname=dirname, + dumpfreq=int(tdump/dt)) -diagnostic_fields = [CourantNumber(), VelocityY(), +diagnostic_fields = [CourantNumber(), YComponent('u'), KineticEnergy(), KineticEnergyY(), - EadyPotentialEnergy(), + IncompressibleEadyPotentialEnergy(parameters), Sum("KineticEnergy", "EadyPotentialEnergy"), Difference("KineticEnergy", "KineticEnergyY"), - GeostrophicImbalance(), TrueResidualV()] + IncompressibleGeostrophicImbalance(parameters), + TrueResidualV(parameters), + Perturbation('p'), Perturbation('b')] -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -# Coriolis expression -Omega = as_vector([0., 0., f*0.5]) -eqns = IncompressibleEadyEquations(state, "RTCF", 1, Omega=Omega) +# Transport schemes and methods +b_opts = SUPGOptions() +transport_schemes = [SSPRK3(domain, "u"), SSPRK3(domain, "b", options=b_opts)] +transport_methods = [DGUpwind(eqns, "u"), DGUpwind("b", ibp=b_opts.ibp)] + +# Linear solve +linear_solver = IncompressibleSolver(eqns) + +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes, + transport_methods, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # +# Initial conditions +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -b0 = state.fields("b") -p0 = state.fields("p") +u0 = stepper.fields("u") +b0 = stepper.fields("b") +p0 = stepper.fields("p") # spaces -Vu = state.spaces("HDiv") -Vb = state.spaces("theta") -Vp = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vb = domain.spaces("theta") +Vp = domain.spaces("DG") # parameters x, y, z = SpatialCoordinate(mesh) @@ -109,8 +130,8 @@ def n(): # calculate hydrostatic pressure p_b = Function(Vp) -incompressible_hydrostatic_balance(state, b_b, p_b) -incompressible_hydrostatic_balance(state, b0, p0) +incompressible_hydrostatic_balance(eqns, b_b, p_b) +incompressible_hydrostatic_balance(eqns, b0, p0) # set x component of velocity dbdy = parameters.dbdy @@ -118,23 +139,40 @@ def n(): # set y component of velocity v = Function(Vp).assign(0.) -eady_initial_v(state, p0, v) + +g = TrialFunction(Vu) +wg = TestFunction(Vu) + +n = FacetNormal(mesh) + +a = inner(wg, g)*dx +L = -div(wg)*p0*dx + inner(wg, n)*p0*(ds_t + ds_b) +pgrad = Function(Vu) +solve(a == L, pgrad) + +# get initial v +Vp = p0.function_space() +phi = TestFunction(Vp) +m = TrialFunction(Vp) + +a = f*phi*m*dx +L = phi*pgrad[0]*dx +solve(a == L, v) # set initial u u_exp = as_vector([u, v, 0.]) u0.project(u_exp) # set the background profiles -state.set_reference_profiles([('p', p_b), - ('b', b_b)]) - -# Set up transport schemes -b_opts = SUPGOptions() -transported_fields = [SSPRK3(state, "u"), SSPRK3(state, "b", options=b_opts)] +stepper.set_reference_profiles([('p', p_b), + ('b', b_b)]) -linear_solver = IncompressibleSolver(state, eqns) +# The residual diagnostic needs to have u_n added to stepper.fields +u_n = stepper.xn('u') +stepper.fields('u_n', field=u_n) -stepper = CrankNicolson(state, eqns, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/gusto/__init__.py b/gusto/__init__.py index b25efa62b..d7dc21faa 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -9,7 +9,6 @@ from gusto.coord_transforms import * # noqa from gusto.domain import * # noqa from gusto.diagnostics import * # noqa -from gusto.eady_diagnostics import * # noqa from gusto.diffusion_methods import * # noqa from gusto.equations import * # noqa from gusto.fml import * # noqa diff --git a/gusto/configuration.py b/gusto/configuration.py index b8307d959..6dea8f3ed 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -6,8 +6,9 @@ __all__ = [ "IntegrateByParts", "TransportEquationType", "OutputParameters", - "CompressibleParameters", "ShallowWaterParameters", "EadyParameters", - "CompressibleEadyParameters", "EmbeddedDGOptions", "RecoveryOptions", + "CompressibleParameters", "ShallowWaterParameters", + "EadyParameters", "CompressibleEadyParameters", + "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", "SpongeLayerParameters", "DiffusionParameters" ] diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index d6e1be4cc..769a8f139 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1,11 +1,12 @@ """Common diagnostic fields.""" -from firedrake import op2, assemble, dot, dx, Function, sqrt, \ - TestFunction, TrialFunction, Constant, grad, inner, curl, \ - LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \ - ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, \ - TensorFunctionSpace, SpatialCoordinate, as_vector, \ - Projector, Interpolator, FunctionSpace +from firedrake import (op2, assemble, dot, dx, Function, sqrt, TestFunction, + TrialFunction, Constant, grad, inner, curl, FacetNormal, + LinearVariationalProblem, LinearVariationalSolver, + ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, + TensorFunctionSpace, SpatialCoordinate, as_vector, + as_matrix, Projector, Interpolator, FunctionSpace, + DirichletBC, lhs, rhs) from firedrake.assign import Assigner from abc import ABCMeta, abstractmethod, abstractproperty @@ -25,7 +26,9 @@ "ThermodynamicKineticEnergy", "Dewpoint", "Temperature", "Theta_d", "RelativeHumidity", "Pressure", "Exner_Vt", "HydrostaticImbalance", "Precipitation", "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence", - "TracerDensity"] + "TracerDensity", "KineticEnergyY", "CompressibleKineticEnergyY", + "IncompressibleEadyPotentialEnergy", "CompressibleEadyPotentialEnergy", + "IncompressibleGeostrophicImbalance", "TrueResidualV", "SawyerEliassenU"] class Diagnostics(object): @@ -1658,3 +1661,347 @@ def setup(self, domain, state_fields): rho_d = state_fields(self.density_name) self.expr = m_X*rho_d super().setup(domain, state_fields) + + +class KineticEnergyY(KineticEnergy): + """Kinetic energy associated with the Y-component of the velocity.""" + name = "KineticEnergyY" + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + u = state_fields("u") + self.expr = self.kinetic(u[1]) + super().setup(domain, state_fields) + + +class CompressibleKineticEnergyY(CompressibleKineticEnergy): + """Kinetic energy associated with the Y-component of the velocity.""" + name = "CompressibleKineticEnergyY" + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + u = state_fields("u") + rho = state_fields("rho") + self.expr = self.kinetic(u[1], rho) + super().setup(domain, state_fields) + + +class IncompressibleEadyPotentialEnergy(DiagnosticField): + """Potential energy in the incompressible Eady equations.""" + name = "EadyPotentialEnergy" + + def __init__(self, parameters, space=None, method='interpolate'): + """ + Args: + parameters (:class:`EadyParameters`): the configuration + object containing the physical parameters for this equation. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.parameters = parameters + super().__init__(space=space, method=method, required_fields=("b")) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + _, _, z = SpatialCoordinate(domain.mesh) + b = state_fields("b") + bbar = state_fields("bbar") + H = self.parameters.H + self.expr = -(z-H/2)*(b-bbar) + super().setup(domain, state_fields) + + +class CompressibleEadyPotentialEnergy(DiagnosticField): + """Potential energy in the compressible Eady equations.""" + name = "CompressibleEadyPotentialEnergy" + + def __init__(self, parameters, space=None, method='interpolate'): + """ + Args: + parameters (:class:`CompressibleParameters`): the configuration + object containing the physical parameters for this equation. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.parameters = parameters + super().__init__(space=space, method=method, required_fields=("rho", "theta")) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + _, _, z = SpatialCoordinate(domain.mesh) + g = self.parameters.g + cp = self.parameters.cp + cv = self.parameters.cv + Pi0 = state_fields('Pi0') + + rho = state_fields("rho") + theta = state_fields("theta") + Pi = tde.exner_pressure(self.parameters, rho, theta) + + self.expr = rho*(g*z + cv*Pi*theta - cp*Pi0*theta) + super().setup(domain, state_fields) + + +class IncompressibleGeostrophicImbalance(DiagnosticField): + """Diagnostic for the amount of geostrophic imbalance.""" + name = "GeostrophicImbalance" + + def __init__(self, parameters, space=None, method='interpolate'): + """ + Args: + parameters (:class:`EadyParameters`): the configuration + object containing the physical parameters for this equation. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.parameters = parameters + super().__init__(space=space, method=method, required_fields=("u", "b", "p")) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + + u = state_fields("u") + b = state_fields("b") + p = state_fields("p") + f = self.parameters.f + Vu = u.function_space() + + v = TrialFunction(Vu) + w = TestFunction(Vu) + a = inner(w, v)*dx + L = (div(w)*p+inner(w, as_vector([f*u[1], 0.0, b])))*dx + + bcs = [DirichletBC(Vu, 0.0, "bottom"), + DirichletBC(Vu, 0.0, "top")] + + imbalance = Function(Vu) + self.expr = imbalance[0]/f + imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs) + self.imbalance_solver = LinearVariationalSolver( + imbalanceproblem, solver_parameters={'ksp_type': 'cg'}) + + def compute(self): + """Compute the diagnostic field from the current state.""" + self.imbalance_solver.solve() + super().compute() + + +class TrueResidualV(DiagnosticField): + """ + Residual in the meridional velocity equation for the Eady equations. + + NB: the user is required to manually add "u_n" to the state fields. + """ + name = "TrueResidualV" + + def __init__(self, parameters, space=None, method='interpolate'): + """ + Args: + parameters (:class:`EadyParameters`): the configuration + object containing the physical parameters for this equation. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.parameters = parameters + super().__init__(space=space, method=method, required_fields=("u")) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + + unew = state_fields("u") + uold = state_fields("u_n") + ubar = 0.5*(unew+uold) + H = self.parameters.H + f = self.parameters.f + dbdy = self.parameters.dbdy + dt = domain.dt + _, _, z = SpatialCoordinate(domain.mesh) + + wv = TestFunction(self.space) + v = TrialFunction(self.space) + vlhs = wv*v*dx + vrhs = wv*((unew[1]-uold[1])/dt + ubar[0]*ubar[1].dx(0) + + ubar[2]*ubar[1].dx(2) + + f*ubar[0] + dbdy*(z-H/2))*dx + vtres = Function(self.space) + self.expr = vtres + vtresproblem = LinearVariationalProblem(vlhs, vrhs, vtres) + self.v_residual_solver = LinearVariationalSolver( + vtresproblem, solver_parameters={'ksp_type': 'cg'}) + + def compute(self): + """Compute the diagnostic field from the current state.""" + self.v_residual_solver.solve() + super().compute() + + +class SawyerEliassenU(DiagnosticField): + """ + Velocity associated with the Sawyer-Eliassen balance equation: the + secondary circulation associated with a stream function that ensures thermal + wind balance. + """ + name = "SawyerEliassenU" + + def __init__(self, equations): + """ + Args: + equations (:class:`IncompressibleEadyEquations`): the equation set + being solved by the model. + """ + space = equations.domain.spaces('HDiv') + self.parameters = equations.parameters + super().__init__(space=space, method='solve', required_fields=("u", "b", "p")) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + + u = state_fields("u") + b = state_fields("b") + v = inner(u, as_vector([0., 1., 0.])) + + # spaces + V0 = domain.spaces('H1') + Vu = domain.spaces('HDiv') + + # project b to V0 + b_v0 = Function(V0) + btri = TrialFunction(V0) + btes = TestFunction(V0) + a = inner(btes, btri) * dx + L = inner(btes, b) * dx + projectbproblem = LinearVariationalProblem(a, L, b_v0) + self.project_b_solver = LinearVariationalSolver( + projectbproblem, solver_parameters={'ksp_type': 'cg'}) + + # project v to V0 + v_v0 = Function(V0) + vtri = TrialFunction(V0) + vtes = TestFunction(V0) + a = inner(vtes, vtri) * dx + L = inner(vtes, v) * dx + projectvproblem = LinearVariationalProblem(a, L, v_v0) + self.project_v_solver = LinearVariationalSolver( + projectvproblem, solver_parameters={'ksp_type': 'cg'}) + + # stm/psi is a stream function + stm = Function(V0) + psi = TrialFunction(V0) + xsi = TestFunction(V0) + + f = self.parameters.f + H = self.parameters.H + L = self.parameters.L + dbdy = self.parameters.dbdy + _, _, z = SpatialCoordinate(domain.mesh) + + bcs = [DirichletBC(V0, 0., "bottom"), + DirichletBC(V0, 0., "top")] + + Mat = as_matrix([[b.dx(2), 0., -f*v_v0.dx(2)], + [0., 0., 0.], + [-b_v0.dx(0), 0., f**2+f*v_v0.dx(0)]]) + + Equ = ( + inner(grad(xsi), dot(Mat, grad(psi))) + - dbdy*inner(grad(xsi), as_vector([-v, 0., f*(z-H/2)])) + )*dx + + # fourth-order terms + if self.parameters.fourthorder: + eps = Constant(0.0001) + brennersigma = Constant(10.0) + n = FacetNormal(domain.mesh) + deltax = Constant(self.parameters.deltax) + deltaz = Constant(self.parameters.deltaz) + + nn = as_matrix([[sqrt(brennersigma/Constant(deltax)), 0., 0.], + [0., 0., 0.], + [0., 0., sqrt(brennersigma/Constant(deltaz))]]) + + mu = as_matrix([[1., 0., 0.], + [0., 0., 0.], + [0., 0., H/L]]) + + # anisotropic form + Equ += eps*( + div(dot(mu, grad(psi)))*div(dot(mu, grad(xsi)))*dx + - ( + avg(dot(dot(grad(grad(psi)), n), n))*jump(grad(xsi), n=n) + + avg(dot(dot(grad(grad(xsi)), n), n))*jump(grad(psi), n=n) + - jump(nn*grad(psi), n=n)*jump(nn*grad(xsi), n=n) + )*(dS_h + dS_v) + ) + + Au = lhs(Equ) + Lu = rhs(Equ) + stmproblem = LinearVariationalProblem(Au, Lu, stm, bcs=bcs) + self.stream_function_solver = LinearVariationalSolver( + stmproblem, solver_parameters={'ksp_type': 'cg'}) + + # solve for sawyer_eliassen u + utrial = TrialFunction(Vu) + w = TestFunction(Vu) + a = inner(w, utrial)*dx + L = (w[0]*(-stm.dx(2))+w[2]*(stm.dx(0)))*dx + ugproblem = LinearVariationalProblem(a, L, self.field) + self.sawyer_eliassen_u_solver = LinearVariationalSolver( + ugproblem, solver_parameters={'ksp_type': 'cg'}) + + def compute(self): + """Compute the diagnostic field from the current state.""" + self.project_b_solver.solve() + self.project_v_solver.solve() + self.stream_function_solver.solve() + self.sawyer_eliassen_u_solver.solve() diff --git a/gusto/eady_diagnostics.py b/gusto/eady_diagnostics.py deleted file mode 100644 index 778fb57b0..000000000 --- a/gusto/eady_diagnostics.py +++ /dev/null @@ -1,241 +0,0 @@ -from firedrake import (SpatialCoordinate, TrialFunction, - TestFunction, Function, DirichletBC, - LinearVariationalProblem, LinearVariationalSolver, - FunctionSpace, lhs, rhs, inner, div, dx, grad, dot, - as_vector, as_matrix, dS_h, dS_v, Constant, avg, - sqrt, jump, FacetNormal) -from gusto import thermodynamics -from gusto.diagnostics import DiagnosticField, Energy - - -__all__ = ["KineticEnergyY", "CompressibleKineticEnergyY", "EadyPotentialEnergy", - "CompressibleEadyPotentialEnergy", "GeostrophicImbalance", - "TrueResidualV", "SawyerEliassenU"] - - -class KineticEnergyY(Energy): - name = "KineticEnergyY" - - def compute(self, state): - """ - Computes the kinetic energy of the y component - """ - u = state.fields("u") - energy = self.kinetic(u[1]) - return self.field.interpolate(energy) - - -class CompressibleKineticEnergyY(Energy): - name = "CompressibleKineticEnergyY" - - def compute(self, state): - """ - Computes the kinetic energy of the y component - """ - u = state.fields("u") - rho = state.fields("rho") - energy = self.kinetic(u[1], rho) - return self.field.interpolate(energy) - - -class EadyPotentialEnergy(Energy): - name = "EadyPotentialEnergy" - - def compute(self, state): - x, y, z = SpatialCoordinate(state.mesh) - b = state.fields("b") - bbar = state.fields("bbar") - H = state.parameters.H - potential = -(z-H/2)*(b-bbar) - return self.field.interpolate(potential) - - -class CompressibleEadyPotentialEnergy(Energy): - name = "CompressibleEadyPotentialEnergy" - - def compute(self, state): - x, y, z = SpatialCoordinate(state.mesh) - g = state.parameters.g - cp = state.parameters.cp - cv = state.parameters.cv - Pi0 = state.parameters.Pi0 - - rho = state.fields("rho") - theta = state.fields("theta") - Pi = thermodynamics.pi(state.parameters, rho, theta) - - potential = rho*(g*z + cv*Pi*theta - cp*Pi0*theta) - return self.field.interpolate(potential) - - -class GeostrophicImbalance(DiagnosticField): - name = "GeostrophicImbalance" - - def setup(self, state): - super(GeostrophicImbalance, self).setup(state) - u = state.fields("u") - b = state.fields("b") - p = state.fields("p") - f = state.parameters.f - Vu = state.spaces("HDiv") - - v = TrialFunction(Vu) - w = TestFunction(Vu) - a = inner(w, v)*dx - L = (div(w)*p+inner(w, as_vector([f*u[1], 0.0, b])))*dx - - bcs = [DirichletBC(Vu, 0.0, "bottom"), - DirichletBC(Vu, 0.0, "top")] - - self.imbalance = Function(Vu) - imbalanceproblem = LinearVariationalProblem(a, L, self.imbalance, bcs=bcs) - self.imbalance_solver = LinearVariationalSolver( - imbalanceproblem, solver_parameters={'ksp_type': 'cg'}) - - def compute(self, state): - f = state.parameters.f - self.imbalance_solver.solve() - geostrophic_imbalance = self.imbalance[0]/f - return self.field.interpolate(geostrophic_imbalance) - - -class TrueResidualV(DiagnosticField): - name = "TrueResidualV" - - def setup(self, state): - super(TrueResidualV, self).setup(state) - unew = state.fields("u") - uold = state.xb("u") - ubar = 0.5*(unew+uold) - H = state.parameters.H - f = state.parameters.f - dbdy = state.parameters.dbdy - dt = state.dt - x, y, z = SpatialCoordinate(state.mesh) - V = FunctionSpace(state.mesh, "DG", 0) - - wv = TestFunction(V) - v = TrialFunction(V) - vlhs = wv*v*dx - vrhs = wv*((unew[1]-uold[1])/dt + ubar[0]*ubar[1].dx(0) - + ubar[2]*ubar[1].dx(2) - + f*ubar[0] + dbdy*(z-H/2))*dx - self.vtres = Function(V) - vtresproblem = LinearVariationalProblem(vlhs, vrhs, self.vtres) - self.v_residual_solver = LinearVariationalSolver( - vtresproblem, solver_parameters={'ksp_type': 'cg'}) - - def compute(self, state): - self.v_residual_solver.solve() - v_residual = self.vtres - return self.field.interpolate(v_residual) - - -class SawyerEliassenU(DiagnosticField): - name = "SawyerEliassenU" - - def setup(self, state): - - space = state.spaces("HDiv") - super(SawyerEliassenU, self).setup(state, space=space) - - u = state.fields("u") - b = state.fields("b") - v = inner(u, as_vector([0., 1., 0.])) - - # spaces - V0 = FunctionSpace(state.mesh, "CG", 2) - Vu = u.function_space() - - # project b to V0 - self.b_v0 = Function(V0) - btri = TrialFunction(V0) - btes = TestFunction(V0) - a = inner(btes, btri) * dx - L = inner(btes, b) * dx - projectbproblem = LinearVariationalProblem(a, L, self.b_v0) - self.project_b_solver = LinearVariationalSolver( - projectbproblem, solver_parameters={'ksp_type': 'cg'}) - - # project v to V0 - self.v_v0 = Function(V0) - vtri = TrialFunction(V0) - vtes = TestFunction(V0) - a = inner(vtes, vtri) * dx - L = inner(vtes, v) * dx - projectvproblem = LinearVariationalProblem(a, L, self.v_v0) - self.project_v_solver = LinearVariationalSolver( - projectvproblem, solver_parameters={'ksp_type': 'cg'}) - - # stm/psi is a stream function - self.stm = Function(V0) - psi = TrialFunction(V0) - xsi = TestFunction(V0) - - f = state.parameters.f - H = state.parameters.H - L = state.parameters.L - dbdy = state.parameters.dbdy - x, y, z = SpatialCoordinate(state.mesh) - - bcs = [DirichletBC(V0, 0., "bottom"), - DirichletBC(V0, 0., "top")] - - Mat = as_matrix([[b.dx(2), 0., -f*self.v_v0.dx(2)], - [0., 0., 0.], - [-self.b_v0.dx(0), 0., f**2+f*self.v_v0.dx(0)]]) - - Equ = ( - inner(grad(xsi), dot(Mat, grad(psi))) - - dbdy*inner(grad(xsi), as_vector([-v, 0., f*(z-H/2)])) - )*dx - - # fourth-order terms - if state.parameters.fourthorder: - eps = Constant(0.0001) - brennersigma = Constant(10.0) - n = FacetNormal(state.mesh) - deltax = Constant(state.parameters.deltax) - deltaz = Constant(state.parameters.deltaz) - - nn = as_matrix([[sqrt(brennersigma/Constant(deltax)), 0., 0.], - [0., 0., 0.], - [0., 0., sqrt(brennersigma/Constant(deltaz))]]) - - mu = as_matrix([[1., 0., 0.], - [0., 0., 0.], - [0., 0., H/L]]) - - # anisotropic form - Equ += eps*( - div(dot(mu, grad(psi)))*div(dot(mu, grad(xsi)))*dx - - ( - avg(dot(dot(grad(grad(psi)), n), n))*jump(grad(xsi), n=n) - + avg(dot(dot(grad(grad(xsi)), n), n))*jump(grad(psi), n=n) - - jump(nn*grad(psi), n=n)*jump(nn*grad(xsi), n=n) - )*(dS_h + dS_v) - ) - - Au = lhs(Equ) - Lu = rhs(Equ) - stmproblem = LinearVariationalProblem(Au, Lu, self.stm, bcs=bcs) - self.stream_function_solver = LinearVariationalSolver( - stmproblem, solver_parameters={'ksp_type': 'cg'}) - - # solve for sawyer_eliassen u - self.u = Function(Vu) - utrial = TrialFunction(Vu) - w = TestFunction(Vu) - a = inner(w, utrial)*dx - L = (w[0]*(-self.stm.dx(2))+w[2]*(self.stm.dx(0)))*dx - ugproblem = LinearVariationalProblem(a, L, self.u) - self.sawyer_eliassen_u_solver = LinearVariationalSolver( - ugproblem, solver_parameters={'ksp_type': 'cg'}) - - def compute(self, state): - self.project_b_solver.solve() - self.project_v_solver.solve() - self.stream_function_solver.solve() - self.sawyer_eliassen_u_solver.solve() - sawyer_eliassen_u = self.u - return self.field.project(sawyer_eliassen_u) diff --git a/gusto/equations.py b/gusto/equations.py index 9dc7a24eb..d0e39b9f8 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -1180,27 +1180,68 @@ def hydrostatic_projection(self, t): class CompressibleEadyEquations(CompressibleEulerEquations): + """ + The compressible equations in a vertical slice, augmented with a third + velocity component normal to the slice, as used in the Eady problem. + """ - def __init__(self, state, family, degree, Omega=None, sponge=None, - terms_to_linearise={'u': [time_derivative], - 'rho': [time_derivative, transport], - 'theta': [time_derivative, transport]}, + def __init__(self, domain, parameters, Omega=None, sponge=None, + extra_terms=None, space_names=None, + linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + Omega (:class:`ufl.Expr`, optional): an expression for the planet's + rotation vector. Defaults to None. + sponge (:class:`ufl.Expr`, optional): an expression for a sponge + layer. 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 + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. If None is specified + then no terms are linearised. Defaults to the string 'default', + in which case the linearisation includes time derivatives and + scalar transport terms. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + diffusion_options (:class:`DiffusionOptions`, optional): any options + to specify for applying diffusion terms to variables. Defaults + to None. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. + + Raises: + NotImplementedError: only mixing ratio tracers are implemented. + """ - super().__init__(state, family, degree, - terms_to_linearise=terms_to_linearise, - Omega=Omega, sponge=sponge, + super().__init__(domain, parameters, Omega=Omega, sponge=sponge, + extra_terms=extra_terms, space_names=space_names, + linearisation_map=linearisation_map, u_transport_option=u_transport_option, diffusion_options=diffusion_options, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) - dthetady = state.parameters.dthetady - Pi0 = state.parameters.Pi0 - cp = state.parameters.cp + dthetady = parameters.dthetady + cp = parameters.cp y_vec = as_vector([0., 1., 0.]) W = self.function_space @@ -1208,14 +1249,19 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, X = self.X u, rho, theta = split(X) - exner_pi = exner_pressure(state.parameters, rho, theta) + exner_pi = exner_pressure(parameters, rho, theta) + DG0 = FunctionSpace(domain.mesh, "DG", 0) + self.prescribed_fields.add_field('Pi0', DG0) + Pi0_field = self.prescribed_fields('Pi0') self.residual -= subject(prognostic( - cp*dthetady*(exner_pi-Pi0)*inner(w, y_vec)*dx, "u"), X) + cp*dthetady*(exner_pi-Pi0_field)*inner(w, y_vec)*dx, "u"), X) self.residual += subject(prognostic( gamma*(dthetady*inner(u, y_vec))*dx, "theta"), X) + # TODO: should there be moist contributions here? + class IncompressibleBoussinesqEquations(PrognosticEquationSet): """ @@ -1369,31 +1415,66 @@ def __init__(self, domain, parameters, Omega=None, class IncompressibleEadyEquations(IncompressibleBoussinesqEquations): - def __init__(self, state, family, degree, Omega=None, - terms_to_linearise={'u': [time_derivative], - 'p': [time_derivative], - 'b': [time_derivative, transport]}, + """ + The incompressible equations in a vertical slice, augmented with a third + velocity component normal to the slice, as used in the Eady problem. + """ + + def __init__(self, domain, parameters, Omega=None, + space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + Omega (:class:`ufl.Expr`, optional): an expression for the planet's + rotation vector. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. If None is specified + then no terms are linearised. Defaults to the string 'default', + in which case the linearisation includes time derivatives and + scalar transport terms. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. - super().__init__(state, family, degree, - Omega=Omega, - terms_to_linearise=terms_to_linearise, + Raises: + NotImplementedError: active tracers are not implemented. + """ + + super().__init__(domain, parameters, Omega=Omega, + space_names=space_names, + linearisation_map=linearisation_map, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) - dbdy = state.parameters.dbdy - H = state.parameters.H - _, _, z = SpatialCoordinate(state.mesh) - eady_exp = Function(state.spaces("DG")).interpolate(z-H/2.) + dbdy = parameters.dbdy + H = parameters.H + _, _, z = SpatialCoordinate(domain.mesh) + eady_exp = Function(domain.spaces("DG")).interpolate(z-H/2.) y_vec = as_vector([0., 1., 0.]) W = self.function_space w, _, gamma = TestFunctions(W) X = self.X - u, _, b = split(X) + u, _, _ = split(X) self.residual += subject(prognostic( dbdy*eady_exp*inner(w, y_vec)*dx, "u"), X) diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index fdd285e40..1a5c5db93 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation_tools.py @@ -3,7 +3,7 @@ from firedrake import MixedFunctionSpace, TrialFunctions, TestFunctions, \ TestFunction, TrialFunction, \ FacetNormal, inner, div, dx, ds_b, ds_t, DirichletBC, \ - Function, Constant, SpatialCoordinate, \ + Function, Constant, \ LinearVariationalProblem, LinearVariationalSolver, \ NonlinearVariationalProblem, NonlinearVariationalSolver, split, solve, \ FunctionSpace, errornorm, zero @@ -14,8 +14,7 @@ __all__ = ["incompressible_hydrostatic_balance", "compressible_hydrostatic_balance", "remove_initial_w", - "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance", - "eady_initial_v", "compressible_eady_initial_v"] + "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance"] def incompressible_hydrostatic_balance(equation, b0, p0, top=False, params=None): @@ -499,62 +498,3 @@ def unsaturated_hydrostatic_balance(equation, state_fields, theta_d, H, compressible_hydrostatic_balance(equation, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True) - - -def eady_initial_v(state, p0, v): - f = state.parameters.f - x, y, z = SpatialCoordinate(state.mesh) - - # get pressure gradient - Vu = state.spaces("HDiv") - g = TrialFunction(Vu) - wg = TestFunction(Vu) - - n = FacetNormal(state.mesh) - - a = inner(wg, g)*dx - L = -div(wg)*p0*dx + inner(wg, n)*p0*(ds_t + ds_b) - pgrad = Function(Vu) - solve(a == L, pgrad) - - # get initial v - Vp = p0.function_space() - phi = TestFunction(Vp) - m = TrialFunction(Vp) - - a = f*phi*m*dx - L = phi*pgrad[0]*dx - solve(a == L, v) - - return v - - -def compressible_eady_initial_v(state, theta0, rho0, v): - f = state.parameters.f - cp = state.parameters.cp - - # exner function - Vr = rho0.function_space() - Pi = Function(Vr).interpolate(thermodynamics.exner_pressure(state.parameters, rho0, theta0)) - - # get Pi gradient - Vu = state.spaces("HDiv") - g = TrialFunction(Vu) - wg = TestFunction(Vu) - - n = FacetNormal(state.mesh) - - a = inner(wg, g)*dx - L = -div(wg)*Pi*dx + inner(wg, n)*Pi*(ds_t + ds_b) - pgrad = Function(Vu) - solve(a == L, pgrad) - - # get initial v - m = TrialFunction(Vr) - phi = TestFunction(Vr) - - a = phi*f*m*dx - L = phi*cp*theta0*pgrad[0]*dx - solve(a == L, v) - - return v