Skip to content

Commit

Permalink
try to get Eady code working with all of the new objects
Browse files Browse the repository at this point in the history
  • Loading branch information
@tommbendall
tommbendall committed Oct 8, 2023
1 parent 3d6f38c commit 2f32f67
Show file tree
Hide file tree
Showing 8 changed files with 632 additions and 432 deletions.
145 changes: 90 additions & 55 deletions examples/compressible/compressible_eady.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
118 changes: 78 additions & 40 deletions examples/incompressible/incompressible_eady.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -16,7 +20,7 @@
tmax = dt
tdump = dt
columns = 10
nlayers = 10
nlayers = 5
else:
tmax = 30*day
tdump = 2*hour
Expand All @@ -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)
Expand Down Expand Up @@ -109,32 +130,49 @@ 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
u = -dbdy/f*(z-H/2)

# 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)
1 change: 0 additions & 1 deletion gusto/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit 2f32f67

Please sign in to comment.