def tracer_sphere(tmpdir, degree):
radius = 1
mesh = IcosahedralSphereMesh(radius=radius, refinement_level=3, degree=1)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
# Parameters chosen so that dt != 1
# Gaussian is translated from (lon=pi/2, lat=0) to (lon=0, lat=0)
# to demonstrate that transport is working correctly
dt = pi / 3. * 0.02
output = OutputParameters(dirname=str(tmpdir), dumpfreq=15)
state = State(mesh, dt=dt, output=output)
umax = 1.0
uexpr = as_vector([-umax * x[1] / radius, umax * x[0] / radius, 0.0])
tmax = pi / 2
f_init = exp(-x[2]**2 - x[0]**2)
f_end = exp(-x[2]**2 - x[1]**2)
tol = 0.05
return TracerSetup(state, tmax, f_init, f_end, "BDM", degree, uexpr, umax,
radius, tol)
def setup_DGadvection(element, vector=False):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 1.0
dt = pi / 3 * 0.01
mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ["u", "D"]
timestepping = TimesteppingParameters(dt=dt)
state = ShallowWaterState(
mesh, vertical_degree=None, horizontal_degree=2, family="BDM", timestepping=timestepping, fieldlist=fieldlist
)
# interpolate initial conditions
u0 = Function(state.V[0], name="velocity")
x = SpatialCoordinate(mesh)
uexpr = as_vector([-x[1], x[0], 0.0])
u0.project(uexpr)
if vector:
if element is "BDM":
BrokenSpace = FunctionSpace(mesh, element, 2)
else:
BrokenSpace = VectorFunctionSpace(mesh, "DG", 1)
Space = VectorFunctionSpace(mesh, "CG", 1)
f = Function(Space, name="f")
fexpr = Expression(("exp(-pow(x[2],2) - pow(x[1],2))", "0.0", "0.0"))
f_end = Function(Space)
f_end_expr = Expression(("exp(-pow(x[2],2) - pow(x[0],2))", "0", "0"))
else:
if element is "BDM":
BrokenSpace = FunctionSpace(mesh, element, 2)
else:
BrokenSpace = FunctionSpace(mesh, "DG", 1)
Space = FunctionSpace(mesh, "CG", 1)
f = Function(Space, name="f")
fexpr = Expression("exp(-pow(x[2],2) - pow(x[1],2))")
f_end = Function(Space)
f_end_expr = Expression("exp(-pow(x[2],2) - pow(x[0],2))")
f.interpolate(fexpr)
f_end.interpolate(f_end_expr)
return state, BrokenSpace, u0, f, f_end
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24. * 60. * 60.
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements,
degree=3)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
dt = 3600.
output = OutputParameters(dirname=dirname + "/sw_linear_w2",
steady_state_error_fields=['u', 'D'],
dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
state = State(mesh, dt=dt, output=output, parameters=parameters)
# Coriolis
Omega = parameters.Omega
fexpr = 2 * Omega * x[2] / R
eqns = LinearShallowWaterEquations(state, "BDM", 1, fexpr=fexpr)
# interpolate initial conditions
# Initial/current conditions
u0 = state.fields("u")
D0 = state.fields("D")
u_max = 2 * pi * R / (12 * day
) # Maximum amplitude of the zonal wind (m/s)
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
g = parameters.g
Dexpr = H - ((R * Omega * u_max) * (x[2] * x[2] / (R * R))) / g
u0.project(uexpr)
D0.interpolate(Dexpr)
transport_schemes = [ForwardEuler(state, "D")]
# build time stepper
stepper = CrankNicolson(state, eqns, transport_schemes)
return stepper, 2 * day
def state(tmpdir, geometry):
"""
returns an instance of the State class, having set up either spherical
geometry or 2D vertical slice geometry
"""
output = OutputParameters(dirname=str(tmpdir), dumplist=["f"], dumpfreq=15)
if geometry == "sphere":
mesh = IcosahedralSphereMesh(radius=1, refinement_level=3, degree=1)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
family = "BDM"
vertical_degree = None
fieldlist = ["u", "D"]
dt = pi / 3. * 0.01
uexpr = as_vector([-x[1], x[0], 0.0])
if geometry == "slice":
m = PeriodicIntervalMesh(15, 1.)
mesh = ExtrudedMesh(m, layers=15, layer_height=1. / 15.)
family = "CG"
vertical_degree = 1
fieldlist = ["u", "rho", "theta"]
dt = 0.01
x = SpatialCoordinate(mesh)
uexpr = as_vector([1.0, 0.0])
timestepping = TimesteppingParameters(dt=dt)
state = State(mesh,
vertical_degree=vertical_degree,
horizontal_degree=1,
family=family,
timestepping=timestepping,
output=output,
fieldlist=fieldlist)
u0 = state.fields("u")
u0.project(uexpr)
return state
# setup shallow water parameters
R = 6371220.
H = 5960.
# setup input that doesn't change with ref level or dt
fieldlist = ['u', 'D']
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
for ref_level, dt in ref_dt.items():
dirname = "sw_W5_ref%s_dt%s" % (ref_level, dt)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=ref_level, degree=3)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname, dumplist_latlon=['D'], dumpfreq=100)
diagnostic_fields = [Sum('D', 'topography')]
state = State(mesh, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = state.fields('u')
def setup_sw(dirname, euler_poincare):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=1500.)
output = OutputParameters(dirname=dirname + "/sw",
dumplist_latlon=['D', 'D_error'],
steady_state_error_fields=['D', 'u'])
parameters = ShallowWaterParameters(H=H)
diagnostic_fields = [
RelativeVorticity(),
AbsoluteVorticity(),
PotentialVorticity(),
ShallowWaterPotentialEnstrophy('RelativeVorticity'),
ShallowWaterPotentialEnstrophy('AbsoluteVorticity'),
ShallowWaterPotentialEnstrophy('PotentialVorticity'),
Difference('RelativeVorticity', 'AnalyticalRelativeVorticity'),
Difference('AbsoluteVorticity', 'AnalyticalAbsoluteVorticity'),
Difference('PotentialVorticity', 'AnalyticalPotentialVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_RelativeVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_AbsoluteVorticity'),
MeridionalComponent('u'),
ZonalComponent('u'),
RadialComponent('u')
]
state = State(mesh,
vertical_degree=None,
horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
Omega = parameters.Omega
g = parameters.g
Dexpr = H - ((R * Omega * u_max + u_max * u_max / 2.0) * (x[2] * x[2] /
(R * R))) / g
# Coriolis
fexpr = 2 * Omega * x[2] / R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", Function(V))
f.interpolate(fexpr) # Coriolis frequency (1/s)
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0), ('D', D0)])
if euler_poincare:
ueqn = EulerPoincare(state, u0.function_space())
sw_forcing = ShallowWaterForcing(state, euler_poincare=True)
else:
ueqn = VectorInvariant(state, u0.function_space())
sw_forcing = ShallowWaterForcing(state, euler_poincare=False)
Deqn = AdvectionEquation(state,
D0.function_space(),
equation_form="continuity")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("D", SSPRK3(state, D0, Deqn)))
linear_solver = ShallowWaterSolver(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, sw_forcing)
vspace = FunctionSpace(state.mesh, "CG", 3)
vexpr = (2 * u_max / R) * x[2] / R
vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace)
vrel_analytical.interpolate(vexpr)
vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace)
vabs_analytical.interpolate(vexpr + f)
pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace)
pv_analytical.interpolate((vexpr + f) / D0)
return stepper, 0.25 * day
def setup_sw(dirname, dt, u_transport_option):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u'])
parameters = ShallowWaterParameters(H=H)
diagnostic_fields = [RelativeVorticity(), AbsoluteVorticity(),
PotentialVorticity(),
ShallowWaterPotentialEnstrophy('RelativeVorticity'),
ShallowWaterPotentialEnstrophy('AbsoluteVorticity'),
ShallowWaterPotentialEnstrophy('PotentialVorticity'),
Difference('RelativeVorticity',
'AnalyticalRelativeVorticity'),
Difference('AbsoluteVorticity',
'AnalyticalAbsoluteVorticity'),
Difference('PotentialVorticity',
'AnalyticalPotentialVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_RelativeVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_AbsoluteVorticity'),
MeridionalComponent('u'),
ZonalComponent('u'),
RadialComponent('u')]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
Omega = parameters.Omega
fexpr = 2*Omega*x[2]/R
eqns = ShallowWaterEquations(state, family="BDM", degree=1,
fexpr=fexpr,
u_transport_option=u_transport_option)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = parameters.g
Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
vspace = FunctionSpace(state.mesh, "CG", 3)
vexpr = (2*u_max/R)*x[2]/R
f = state.fields("coriolis")
vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace)
vrel_analytical.interpolate(vexpr)
vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace)
vabs_analytical.interpolate(vexpr + f)
pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace)
pv_analytical.interpolate((vexpr+f)/D0)
return state, eqns
from gusto import *
from firedrake import IcosahedralSphereMesh, Expression, SpatialCoordinate, \
Constant, as_vector
from math import pi
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24.*60.*60.
u_0 = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements, degree=3)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname='sw_linear_w2', steady_state_dump_err={'u':True, 'D':True})
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = ShallowWaterState(mesh, vertical_degree=None, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
write_latlon = True
nc_diag = {'Energy': 0., 'Enstrophy': 0., 'Mass': 0.}
h5_safe, nc_safe = True, True
if COMM_WORLD.Get_rank() == 0:
print("Simulation for Galewsky test case, model u-ad_flux", "\n",
"| Discr details: Poisson implicit, no EP, TM=0.5", "\n", "| dt", dt,
"| tmax", tmax, " | refinement level", ref_level, "\n", "| maxk",
maxk, "| dfr field, nc_h5:", field_dumpfreq, nc_h5_dumpfreq, "\n",
"| pickup", pickup, "| Diagnostics:", tuple(nc_diag.keys()))
print("Starting Initial condition, and function setup at", ctime())
R, Omega = 6371220., 7.292e-5
mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2)
mesh.init_cell_orientations(SpatialCoordinate(mesh))
x = SpatialCoordinate(mesh)
f = Function(FunctionSpace(mesh, "CG", 1))
f.interpolate(2 * Omega * x[2] / R)
g, H = 9.810616, 5960.
def latlon_coords(mesh):
"""Compute latitude-longitude coordinates given Cartesian ones"""
x0, y0, z0 = SpatialCoordinate(mesh)
unsafe = z0 / sqrt(x0 * x0 + y0 * y0 + z0 * z0)
safe = Min(Max(unsafe, -1.0), 1.0) # avoid silly roundoff errors
theta = asin(safe) # latitude
lamda = atan_2(y0, x0) # longitude
return theta, lamda
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24.*60.*60.
u_0 = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements, degree=3)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_dump_err={'u':True,'D':True}, dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = ShallowWaterState(mesh, vertical_degree=None, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
g = parameters.g
Omega = parameters.Omega
# Coriolis expression
R = Constant(R)
Omega = Constant(parameters.Omega)
x = SpatialCoordinate(mesh)
fexpr = 2*Omega*x[2]/R
V = FunctionSpace(mesh, "CG", 1)
state.f = Function(V).interpolate(fexpr) # Coriolis frequency (1/s)
u_max = Constant(u_0)
# interpolate initial conditions
# Initial/current conditions
u0, D0 = Function(state.V[0]), Function(state.V[1])
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = Constant(parameters.g)
Dexpr = - ((R * Omega * u_max)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([u0, D0])
advection_dict = {}
advection_dict["u"] = NoAdvection(state)
advection_dict["D"] = NoAdvection(state)
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state, linear=True)
# build time stepper
stepper = Timestepper(state, advection_dict, linear_solver,
sw_forcing)
return stepper, 2*day
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24. * 60. * 60.
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements,
degree=3)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname=dirname + "/sw_linear_w2",
steady_state_error_fields=['u', 'D'],
dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = State(mesh,
horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# Coriolis
Omega = parameters.Omega
fexpr = 2 * Omega * x[2] / R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", Function(V))
f.interpolate(fexpr) # Coriolis frequency (1/s)
# interpolate initial conditions
# Initial/current conditions
u0 = state.fields("u")
D0 = state.fields("D")
u_max = 2 * pi * R / (12 * day
) # Maximum amplitude of the zonal wind (m/s)
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
g = parameters.g
Dexpr = -((R * Omega * u_max) * (x[2] * x[2] / (R * R))) / g
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0), ('D', D0)])
Deqn = LinearAdvection(state,
D0.function_space(),
state.parameters.H,
ibp=IntegrateByParts.ONCE,
equation_form="continuity")
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("D", ForwardEuler(state, D0, Deqn)))
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state, linear=True)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, sw_forcing)
return stepper, 2 * day