From 95b86187c9552c4c86f7e877e529ec7ab1a1ddfb Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 6 Oct 2023 08:28:15 +0100 Subject: [PATCH 1/7] add Eady problems back in (old code copied across so not working!) --- examples/compressible/compressible_eady.py | 148 +++++++++++ .../incompressible/incompressible_eady.py | 140 ++++++++++ gusto/__init__.py | 1 + gusto/configuration.py | 33 ++- gusto/eady_diagnostics.py | 241 ++++++++++++++++++ gusto/equations.py | 74 +++++- gusto/initialisation_tools.py | 64 ++++- 7 files changed, 695 insertions(+), 6 deletions(-) create mode 100644 examples/compressible/compressible_eady.py create mode 100644 examples/incompressible/incompressible_eady.py create mode 100644 gusto/eady_diagnostics.py diff --git a/examples/compressible/compressible_eady.py b/examples/compressible/compressible_eady.py new file mode 100644 index 000000000..c1760b44b --- /dev/null +++ b/examples/compressible/compressible_eady.py @@ -0,0 +1,148 @@ +""" +This solves the Eady problem using the compressible Euler equations. +""" + +from gusto import * +from gusto import thermodynamics +from firedrake import (as_vector, SpatialCoordinate, + PeriodicRectangleMesh, ExtrudedMesh, + exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt) +import sys + +day = 24.*60.*60. +hour = 60.*60. +dt = 30. +if '--running-tests' in sys.argv: + tmax = dt + tdump = dt + columns = 10 # number of columns + nlayers = 5 # horizontal layers +else: + tmax = 30*day + tdump = 5*day + columns = 30 # number of columns + nlayers = 30 # horizontal layers + +# set up mesh +L = 1000000. +m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True) + +# build 2D mesh by extruding the base mesh +H = 10000. # Height position of the model top +mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + +# Coriolis expression +f = 1.e-04 +Omega = as_vector([0., 0., f*0.5]) + +dirname = 'compressible_eady' +output = OutputParameters(dirname=dirname, + dumpfreq=int(tdump/dt), + dumplist=['u', 'rho', 'theta'], + perturbation_fields=['rho', 'theta', 'ExnerPi'], + log_level='INFO') + +parameters = CompressibleEadyParameters(H=H, f=f) + +diagnostic_fields = [CourantNumber(), VelocityY(), + ExnerPi(), ExnerPi(reference=True), + CompressibleKineticEnergy(), + CompressibleKineticEnergyY(), + CompressibleEadyPotentialEnergy(), + Sum("CompressibleKineticEnergy", + "CompressibleEadyPotentialEnergy"), + Difference("CompressibleKineticEnergy", + "CompressibleKineticEnergyY")] + +state = State(mesh, + dt=dt, + output=output, + parameters=parameters, + diagnostic_fields=diagnostic_fields) + +eqns = CompressibleEadyEquations(state, "RTCF", 1) + +# Initial conditions +u0 = state.fields("u") +rho0 = state.fields("rho") +theta0 = state.fields("theta") + +# spaces +Vu = state.spaces("HDiv") +Vt = state.spaces("theta") +Vr = state.spaces("DG") + +# first setup the background buoyancy profile +# z.grad(bref) = N**2 +# the following is symbolic algebra, using the default buoyancy frequency +# from the parameters class. +x, y, z = SpatialCoordinate(mesh) +g = parameters.g +Nsq = parameters.Nsq +theta_surf = parameters.theta_surf + +# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) +theta_ref = theta_surf*exp(Nsq*(z-H/2)/g) +theta_b = Function(Vt).interpolate(theta_ref) + + +# set theta_pert +def coth(x): + return cosh(x)/sinh(x) + + +def Z(z): + return Bu*((z/H)-0.5) + + +def n(): + return Bu**(-1)*sqrt((Bu*0.5-tanh(Bu*0.5))*(coth(Bu*0.5)-Bu*0.5)) + + +a = -4.5 +Bu = 0.5 +theta_exp = a*theta_surf/g*sqrt(Nsq)*(-(1.-Bu*0.5*coth(Bu*0.5))*sinh(Z(z))*cos(pi*(x-L)/L) + - n()*Bu*cosh(Z(z))*sin(pi*(x-L)/L)) +theta_pert = Function(Vt).interpolate(theta_exp) + +# set theta0 +theta0.interpolate(theta_b + theta_pert) + +# calculate hydrostatic Pi +rho_b = Function(Vr) +compressible_hydrostatic_balance(state, theta_b, rho_b) +compressible_hydrostatic_balance(state, theta0, rho0) + +# set Pi0 +Pi0 = calculate_Pi0(state, theta0, rho0) +state.parameters.Pi0 = Pi0 + +# set x component of velocity +cp = state.parameters.cp +dthetady = state.parameters.dthetady +Pi = thermodynamics.pi(state.parameters, rho0, theta0) +u = cp*dthetady/f*(Pi-Pi0) + +# set y component of velocity +v = Function(Vr).assign(0.) +compressible_eady_initial_v(state, theta0, rho0, 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 = CrankNicolson(state, eqns, transported_fields, + linear_solver=linear_solver) + +stepper.run(t=0, tmax=tmax) diff --git a/examples/incompressible/incompressible_eady.py b/examples/incompressible/incompressible_eady.py new file mode 100644 index 000000000..1ba9f82ed --- /dev/null +++ b/examples/incompressible/incompressible_eady.py @@ -0,0 +1,140 @@ +""" +The Eady problem solved using the incompressible Boussinesq equations. +""" + +from gusto import * +from firedrake import (as_vector, SpatialCoordinate, + PeriodicRectangleMesh, ExtrudedMesh, + cos, sin, cosh, sinh, tanh, pi, Function, sqrt) +import sys + +day = 24.*60.*60. +hour = 60.*60. +dt = 100. + +if '--running-tests' in sys.argv: + tmax = dt + tdump = dt + columns = 10 + nlayers = 10 +else: + tmax = 30*day + tdump = 2*hour + columns = 30 + nlayers = 30 + +H = 10000. +L = 1000000. +f = 1.e-04 + +# rescaling +beta = 1.0 +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) + + +output = OutputParameters(dirname='incompressible_eady', + dumpfreq=int(tdump/dt), + perturbation_fields=['p', 'b'], + log_level='INFO') + +parameters = EadyParameters(H=H, L=L, f=f, + deltax=2.*L/float(columns), + deltaz=H/float(nlayers), + fourthorder=True) + +diagnostic_fields = [CourantNumber(), VelocityY(), + KineticEnergy(), KineticEnergyY(), + EadyPotentialEnergy(), + Sum("KineticEnergy", "EadyPotentialEnergy"), + Difference("KineticEnergy", "KineticEnergyY"), + GeostrophicImbalance(), TrueResidualV()] + +state = State(mesh, + dt=dt, + output=output, + parameters=parameters, + diagnostic_fields=diagnostic_fields) + +# Coriolis expression +Omega = as_vector([0., 0., f*0.5]) +eqns = IncompressibleEadyEquations(state, "RTCF", 1, Omega=Omega) + +# Initial conditions +u0 = state.fields("u") +b0 = state.fields("b") +p0 = state.fields("p") + +# spaces +Vu = state.spaces("HDiv") +Vb = state.spaces("theta") +Vp = state.spaces("DG") + +# parameters +x, y, z = SpatialCoordinate(mesh) +Nsq = parameters.Nsq + +# background buoyancy +bref = (z-H/2)*Nsq +b_b = Function(Vb).project(bref) + + +# buoyancy perturbation +def coth(x): + return cosh(x)/sinh(x) + + +def Z(z): + return Bu*((z/H)-0.5) + + +def n(): + return Bu**(-1)*sqrt((Bu*0.5-tanh(Bu*0.5))*(coth(Bu*0.5)-Bu*0.5)) + + +a = -4.5 +Bu = 0.5 +b_exp = a*sqrt(Nsq)*(-(1.-Bu*0.5*coth(Bu*0.5))*sinh(Z(z))*cos(pi*(x-L)/L) + - n()*Bu*cosh(Z(z))*sin(pi*(x-L)/L)) +b_pert = Function(Vb).interpolate(b_exp) + +# set total buoyancy +b0.project(b_b + b_pert) + +# calculate hydrostatic pressure +p_b = Function(Vp) +incompressible_hydrostatic_balance(state, b_b, p_b) +incompressible_hydrostatic_balance(state, b0, p0) + +# set x component of velocity +dbdy = parameters.dbdy +u = -dbdy/f*(z-H/2) + +# set y component of velocity +v = Function(Vp).assign(0.) +eady_initial_v(state, p0, 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)] + +linear_solver = IncompressibleSolver(state, eqns) + +stepper = CrankNicolson(state, eqns, transported_fields, + linear_solver=linear_solver) + +stepper.run(t=0, tmax=tmax) diff --git a/gusto/__init__.py b/gusto/__init__.py index d7dc21faa..b25efa62b 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -9,6 +9,7 @@ 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 039b32740..b8307d959 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -6,9 +6,9 @@ __all__ = [ "IntegrateByParts", "TransportEquationType", "OutputParameters", - "CompressibleParameters", "ShallowWaterParameters", - "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", - "SpongeLayerParameters", "DiffusionParameters" + "CompressibleParameters", "ShallowWaterParameters", "EadyParameters", + "CompressibleEadyParameters", "EmbeddedDGOptions", "RecoveryOptions", + "SUPGOptions", "SpongeLayerParameters", "DiffusionParameters" ] @@ -134,6 +134,33 @@ class ShallowWaterParameters(Configuration): H = None # mean depth +class EadyParameters(Configuration): + + """ + Physical parameters for Incompressible Eady + """ + Nsq = 2.5e-05 # squared Brunt-Vaisala frequency (1/s) + dbdy = -1.0e-07 + H = None + L = None + f = None + deltax = None + deltaz = None + fourthorder = False + + +class CompressibleEadyParameters(CompressibleParameters, EadyParameters): + + """ + Physical parameters for Compressible Eady + """ + g = 10. + N = sqrt(EadyParameters.Nsq) + theta_surf = 300. + dthetady = theta_surf/g*EadyParameters.dbdy + Pi0 = 0.0 + + class WrapperOptions(Configuration, metaclass=ABCMeta): """Base class for specifying options for a transport scheme.""" diff --git a/gusto/eady_diagnostics.py b/gusto/eady_diagnostics.py new file mode 100644 index 000000000..778fb57b0 --- /dev/null +++ b/gusto/eady_diagnostics.py @@ -0,0 +1,241 @@ +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 1c50bd4e6..9dc7a24eb 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -5,7 +5,7 @@ FunctionSpace, MixedFunctionSpace, TestFunctions, TrialFunction, FacetNormal, jump, avg, dS_v, dS, DirichletBC, conditional, SpatialCoordinate, - split, Constant, action) + split, Constant, action, as_vector) from gusto.fields import PrescribedFields from gusto.fml import (Term, all_terms, keep, drop, Label, subject, name, replace_subject, replace_trial_function) @@ -1179,6 +1179,44 @@ def hydrostatic_projection(self, t): return replace_subject(new_subj)(t) +class CompressibleEadyEquations(CompressibleEulerEquations): + + def __init__(self, state, family, degree, Omega=None, sponge=None, + terms_to_linearise={'u': [time_derivative], + 'rho': [time_derivative, transport], + 'theta': [time_derivative, transport]}, + u_transport_option="vector_invariant_form", + diffusion_options=None, + no_normal_flow_bc_ids=None, + active_tracers=None): + + super().__init__(state, family, degree, + terms_to_linearise=terms_to_linearise, + Omega=Omega, sponge=sponge, + 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 + y_vec = as_vector([0., 1., 0.]) + + W = self.function_space + w, _, gamma = TestFunctions(W) + X = self.X + u, rho, theta = split(X) + + exner_pi = exner_pressure(state.parameters, rho, theta) + + self.residual -= subject(prognostic( + cp*dthetady*(exner_pi-Pi0)*inner(w, y_vec)*dx, "u"), X) + + self.residual += subject(prognostic( + gamma*(dthetady*inner(u, y_vec))*dx, "theta"), X) + + class IncompressibleBoussinesqEquations(PrognosticEquationSet): """ Class for the incompressible Boussinesq equations, which evolve the velocity @@ -1328,3 +1366,37 @@ def __init__(self, domain, parameters, Omega=None, # -------------------------------------------------------------------- # # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) + + +class IncompressibleEadyEquations(IncompressibleBoussinesqEquations): + def __init__(self, state, family, degree, Omega=None, + terms_to_linearise={'u': [time_derivative], + 'p': [time_derivative], + 'b': [time_derivative, transport]}, + u_transport_option="vector_invariant_form", + no_normal_flow_bc_ids=None, + active_tracers=None): + + super().__init__(state, family, degree, + Omega=Omega, + terms_to_linearise=terms_to_linearise, + 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.) + y_vec = as_vector([0., 1., 0.]) + + W = self.function_space + w, _, gamma = TestFunctions(W) + X = self.X + u, _, b = split(X) + + self.residual += subject(prognostic( + dbdy*eady_exp*inner(w, y_vec)*dx, "u"), X) + + self.residual += subject(prognostic( + gamma*dbdy*inner(u, y_vec)*dx, "b"), X) diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index 1a5c5db93..fdd285e40 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, \ + Function, Constant, SpatialCoordinate, \ LinearVariationalProblem, LinearVariationalSolver, \ NonlinearVariationalProblem, NonlinearVariationalSolver, split, solve, \ FunctionSpace, errornorm, zero @@ -14,7 +14,8 @@ __all__ = ["incompressible_hydrostatic_balance", "compressible_hydrostatic_balance", "remove_initial_w", - "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance"] + "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance", + "eady_initial_v", "compressible_eady_initial_v"] def incompressible_hydrostatic_balance(equation, b0, p0, top=False, params=None): @@ -498,3 +499,62 @@ 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 From 2f32f672e56162328e331c9d6bb7c624875ffd76 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 8 Oct 2023 18:46:06 +0100 Subject: [PATCH 2/7] 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 From a92947ec9b77fd202d3751fe8aeb7c991b458628 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 8 Oct 2023 20:01:57 +0100 Subject: [PATCH 3/7] fix incompressible eady example --- .../incompressible/incompressible_eady.py | 8 +++--- gusto/diagnostics.py | 26 ++++++++++++------- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/examples/incompressible/incompressible_eady.py b/examples/incompressible/incompressible_eady.py index a7ad6abcf..f9d9fad36 100644 --- a/examples/incompressible/incompressible_eady.py +++ b/examples/incompressible/incompressible_eady.py @@ -64,8 +64,8 @@ IncompressibleEadyPotentialEnergy(parameters), Sum("KineticEnergy", "EadyPotentialEnergy"), Difference("KineticEnergy", "KineticEnergyY"), - IncompressibleGeostrophicImbalance(parameters), - TrueResidualV(parameters), + IncompressibleGeostrophicImbalance(eqns), + TrueResidualV(parameters), SawyerEliassenU(eqns), Perturbation('p'), Perturbation('b')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) @@ -73,7 +73,7 @@ # 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)] +transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "b", ibp=b_opts.ibp)] # Linear solve linear_solver = IncompressibleSolver(eqns) @@ -168,7 +168,7 @@ def n(): ('b', b_b)]) # The residual diagnostic needs to have u_n added to stepper.fields -u_n = stepper.xn('u') +u_n = stepper.x.n('u') stepper.fields('u_n', field=u_n) # ---------------------------------------------------------------------------- # diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index 769a8f139..dcf98b27b 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1724,7 +1724,7 @@ def setup(self, domain, state_fields): """ _, _, z = SpatialCoordinate(domain.mesh) b = state_fields("b") - bbar = state_fields("bbar") + bbar = state_fields("b_bar") H = self.parameters.H self.expr = -(z-H/2)*(b-bbar) super().setup(domain, state_fields) @@ -1774,11 +1774,11 @@ class IncompressibleGeostrophicImbalance(DiagnosticField): """Diagnostic for the amount of geostrophic imbalance.""" name = "GeostrophicImbalance" - def __init__(self, parameters, space=None, method='interpolate'): + def __init__(self, equations, space=None, method='interpolate'): """ Args: - parameters (:class:`EadyParameters`): the configuration - object containing the physical parameters for this equation. + equations (:class:`IncompressibleEadyEquations`): the equation set + being solved by the model. 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. @@ -1786,7 +1786,8 @@ def __init__(self, parameters, space=None, method='interpolate'): for this diagnostic. Valid options are 'interpolate', 'project', 'assign' and 'solve'. Defaults to 'interpolate'. """ - self.parameters = parameters + self.equations = equations + self.parameters = equations.parameters super().__init__(space=space, method=method, required_fields=("u", "b", "p")) def setup(self, domain, state_fields): @@ -1808,14 +1809,14 @@ def setup(self, domain, state_fields): 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")] + bcs = self.equations.bcs['u'] 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'}) + super().setup(domain, state_fields) def compute(self): """Compute the diagnostic field from the current state.""" @@ -1863,17 +1864,19 @@ def setup(self, domain, state_fields): dt = domain.dt _, _, z = SpatialCoordinate(domain.mesh) - wv = TestFunction(self.space) - v = TrialFunction(self.space) + V = self.space if self.space is not None else FunctionSpace(domain.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 - vtres = Function(self.space) + vtres = Function(V) self.expr = vtres vtresproblem = LinearVariationalProblem(vlhs, vrhs, vtres) self.v_residual_solver = LinearVariationalSolver( vtresproblem, solver_parameters={'ksp_type': 'cg'}) + super().setup(domain, state_fields) def compute(self): """Compute the diagnostic field from the current state.""" @@ -1897,6 +1900,7 @@ def __init__(self, equations): """ space = equations.domain.spaces('HDiv') self.parameters = equations.parameters + self.solve_implemented = True super().__init__(space=space, method='solve', required_fields=("u", "b", "p")) def setup(self, domain, state_fields): @@ -1907,6 +1911,8 @@ def setup(self, domain, state_fields): state_fields (:class:`StateFields`): the model's field container. """ + super().setup(domain, state_fields) + u = state_fields("u") b = state_fields("b") v = inner(u, as_vector([0., 1., 0.])) From 0ed07d4132dfe8da170e04397f32fb915d862b97 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 9 Oct 2023 17:09:30 +0100 Subject: [PATCH 4/7] use correct Pi0 variable --- examples/compressible/compressible_eady.py | 13 +++++++------ gusto/diagnostics.py | 2 +- gusto/equations.py | 6 ++---- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/examples/compressible/compressible_eady.py b/examples/compressible/compressible_eady.py index 029700009..64503a70a 100644 --- a/examples/compressible/compressible_eady.py +++ b/examples/compressible/compressible_eady.py @@ -5,7 +5,7 @@ from gusto import * from gusto import thermodynamics from firedrake import (as_vector, SpatialCoordinate, solve, ds_b, ds_t, - PeriodicRectangleMesh, ExtrudedMesh, + PeriodicRectangleMesh, ExtrudedMesh, assemble, exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt) import sys @@ -44,7 +44,7 @@ # Equation Omega = as_vector([0., 0., f*0.5]) -parameters = CompressibleEadyParameters(H=H, f=f) +parameters = CompressibleEadyParameters(H=H, f=f, Pi0=Constant(1.0)) eqns = CompressibleEadyEquations(domain, parameters, Omega=Omega) # I/O @@ -136,14 +136,15 @@ def n(): compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) compressible_hydrostatic_balance(eqns, theta0, rho0, solve_for_rho=True) -# set Pi0 -- have to get this from the equations -Pi0 = eqns.prescribed_fields('Pi0') -Pi0.interpolate(exner_pressure(parameters, rho0, theta0)) +# set Pi0 -- set this in configuration so that everything sees the same value +Pi = thermodynamics.exner_pressure(parameters, rho0, theta0) +Pi0 = parameters.Pi0 +Pi0_value = assemble(Pi*dx) / assemble(Constant(1.0)*dx(domain=mesh)) +Pi0.assign(Pi0_value) # set x component of velocity cp = parameters.cp dthetady = parameters.dthetady -Pi = thermodynamics.exner_pressure(parameters, rho0, theta0) u = cp*dthetady/f*(Pi-Pi0) # set y component of velocity by solving a problem diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index dcf98b27b..af45e5e8f 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1760,7 +1760,7 @@ def setup(self, domain, state_fields): g = self.parameters.g cp = self.parameters.cp cv = self.parameters.cv - Pi0 = state_fields('Pi0') + Pi0 = self.parameters.Pi0 rho = state_fields("rho") theta = state_fields("theta") diff --git a/gusto/equations.py b/gusto/equations.py index d0e39b9f8..8e72891df 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -1250,12 +1250,10 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, u, rho, theta = split(X) 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') + Pi0 = self.parameters.Pi0 self.residual -= subject(prognostic( - cp*dthetady*(exner_pi-Pi0_field)*inner(w, y_vec)*dx, "u"), X) + cp*dthetady*(exner_pi-Pi0)*inner(w, y_vec)*dx, "u"), X) self.residual += subject(prognostic( gamma*(dthetady*inner(u, y_vec))*dx, "theta"), X) From f8fe35b1c2e9711f6f00629ba75170c5f4e9f825 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 24 Apr 2024 08:58:46 +0100 Subject: [PATCH 5/7] change parameters in an attempt to avoid conflicts between N and Nsq --- examples/boussinesq/incompressible_eady.py | 9 +++++---- examples/compressible/compressible_eady.py | 2 +- gusto/configuration.py | 17 +++++++++++------ 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/examples/boussinesq/incompressible_eady.py b/examples/boussinesq/incompressible_eady.py index 365547963..711096f14 100644 --- a/examples/boussinesq/incompressible_eady.py +++ b/examples/boussinesq/incompressible_eady.py @@ -49,10 +49,11 @@ # 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) +parameters = IncompressibleEadyParameters(H=H, L=L, f=f, + deltax=2.*L/float(columns), + deltaz=H/float(nlayers), + fourthorder=True, + N=sqrt(2.5e-5)) eqns = IncompressibleEadyEquations(domain, parameters, Omega=Omega) # I/O diff --git a/examples/compressible/compressible_eady.py b/examples/compressible/compressible_eady.py index 20c3b6d9a..1ccccd601 100644 --- a/examples/compressible/compressible_eady.py +++ b/examples/compressible/compressible_eady.py @@ -44,7 +44,7 @@ # Equation Omega = as_vector([0., 0., f*0.5]) -parameters = CompressibleEadyParameters(H=H, f=f, Pi0=Constant(1.0)) +parameters = CompressibleEadyParameters(N=sqrt(2.5e-5), H=H, f=f, Pi0=Constant(1.0)) eqns = CompressibleEadyEquations(domain, parameters, Omega=Omega) # I/O diff --git a/gusto/configuration.py b/gusto/configuration.py index ba47ef92e..27e47e7c1 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -6,8 +6,8 @@ __all__ = [ "IntegrateByParts", "TransportEquationType", "OutputParameters", - "BoussinesqParameters", "CompressibleParameters", - "ShallowWaterParameters", "EadyParameters", "CompressibleEadyParameters", + "BoussinesqParameters", "CompressibleParameters", "ShallowWaterParameters", + "IncompressibleEadyParameters", "CompressibleEadyParameters", "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", "MixedFSOptions", "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters" ] @@ -147,11 +147,9 @@ class ShallowWaterParameters(Configuration): class EadyParameters(Configuration): - """ - Physical parameters for Incompressible Eady + Base class of physical parameters for Eady problems """ - Nsq = 2.5e-05 # squared Brunt-Vaisala frequency (1/s) dbdy = -1.0e-07 H = None L = None @@ -161,13 +159,20 @@ class EadyParameters(Configuration): fourthorder = False +class IncompressibleEadyParameters(BoussinesqParameters, EadyParameters): + """ + Base class of physical parameters for incompressible Eady problems + """ + Nsq = BoussinesqParameters.N**2 + + class CompressibleEadyParameters(CompressibleParameters, EadyParameters): """ Physical parameters for Compressible Eady """ g = 10. - N = sqrt(EadyParameters.Nsq) + Nsq = CompressibleParameters.N**2 theta_surf = 300. dthetady = theta_surf/g*EadyParameters.dbdy Pi0 = 0.0 From df1659e13b7094a217df283799e7d3d70d3f06ee Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 1 May 2024 14:45:05 +0100 Subject: [PATCH 6/7] switch to vector advection form and make stepper 4x1 for compressible Eady --- examples/compressible/compressible_eady.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/compressible/compressible_eady.py b/examples/compressible/compressible_eady.py index 1ccccd601..3a6752f2f 100644 --- a/examples/compressible/compressible_eady.py +++ b/examples/compressible/compressible_eady.py @@ -45,7 +45,8 @@ # Equation Omega = as_vector([0., 0., f*0.5]) parameters = CompressibleEadyParameters(N=sqrt(2.5e-5), H=H, f=f, Pi0=Constant(1.0)) -eqns = CompressibleEadyEquations(domain, parameters, Omega=Omega) +eqns = CompressibleEadyEquations(domain, parameters, Omega=Omega, + u_transport_option='vector_advection_form') # I/O output = OutputParameters(dirname=dirname, @@ -84,7 +85,8 @@ # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes, transport_methods, - linear_solver=linear_solver) + linear_solver=linear_solver, + num_outer=4, num_inner=1) # ---------------------------------------------------------------------------- # # Initial conditions From 58f259f1a7f16fa16855dda7e7039322c0f72f51 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 24 Feb 2025 08:20:45 +0000 Subject: [PATCH 7/7] fix lint and add boussinesq diagnostics file --- .../compressible_euler/compressible_eady.py | 4 +- gusto/diagnostics/boussinesq_diagnostics.py | 197 ++++++++++++++++++ 2 files changed, 199 insertions(+), 2 deletions(-) create mode 100644 gusto/diagnostics/boussinesq_diagnostics.py diff --git a/examples/compressible_euler/compressible_eady.py b/examples/compressible_euler/compressible_eady.py index 675f39509..4b859a96e 100644 --- a/examples/compressible_euler/compressible_eady.py +++ b/examples/compressible_euler/compressible_eady.py @@ -16,8 +16,8 @@ from gusto import thermodynamics as tde from firedrake import ( as_vector, SpatialCoordinate, solve, ds_b, ds_t, PeriodicRectangleMesh, - ExtrudedMesh, assemble, exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt, - TrialFunction, TestFunction, inner, dx, div, FacetNormal, FunctionSpace + ExtrudedMesh, exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt, + TrialFunction, TestFunction, inner, dx, div, FacetNormal ) compressible_eady_defaults = { diff --git a/gusto/diagnostics/boussinesq_diagnostics.py b/gusto/diagnostics/boussinesq_diagnostics.py new file mode 100644 index 000000000..03e017e5c --- /dev/null +++ b/gusto/diagnostics/boussinesq_diagnostics.py @@ -0,0 +1,197 @@ +"""Some diagnostic fields for the Boussinesq equations.""" + +from firedrake import ( + as_vector, inner, dx, div, as_matrix, TrialFunction, TestFunction, + LinearVariationalProblem, LinearVariationalSolver, DirichletBC, Function, + grad, dot, FacetNormal, SpatialCoordinate, avg, jump, lhs, rhs, sqrt, + FunctionSpace, dS_v, dS_h +) +from gusto.diagnostics.diagnostics import DiagnosticField + +__all__ = [ + 'IncompressibleGeostrophicImbalance', 'SawyerEliassenU'] + + +class IncompressibleGeostrophicImbalance(DiagnosticField): + """Diagnostic for the amount of geostrophic imbalance.""" + name = "GeostrophicImbalance" + + def __init__(self, equations, space=None, method='interpolate'): + """ + Args: + equations (:class:`IncompressibleEadyEquations`): the equation set + being solved by the model. + 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.equations = equations + self.parameters = equations.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 = self.equations.bcs['u'] + + 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'}) + super().setup(domain, state_fields) + + def compute(self): + """Compute the diagnostic field from the current state.""" + self.imbalance_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 + self.solve_implemented = True + 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. + """ + + super().setup(domain, state_fields) + + 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: + R = FunctionSpace(domain.mesh, "R", 0) + eps = Function(R, val=0.0001) + brennersigma = Function(R, val=10.0) + n = FacetNormal(domain.mesh) + deltax = self.parameters.deltax + deltaz = self.parameters.deltaz + + nn = as_matrix([[sqrt(brennersigma/deltax), 0., 0.], + [0., 0., 0.], + [0., 0., sqrt(brennersigma/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()