def test_physics_recovery_kernels(boundary):
m = IntervalMesh(3, 3)
mesh = ExtrudedMesh(m, layers=3, layer_height=1.0)
cell = m.ufl_cell().cellname()
hori_elt = FiniteElement("DG", cell, 0)
vert_elt = FiniteElement("CG", interval, 1)
theta_elt = TensorProductElement(hori_elt, vert_elt)
Vt = FunctionSpace(mesh, theta_elt)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_elt))
initial_field = Function(Vt)
true_field = Function(Vt_brok)
new_field = Function(Vt_brok)
initial_field, true_field = setup_values(boundary, initial_field,
true_field)
kernel = kernels.PhysicsRecoveryTop(
) if boundary == "top" else kernels.PhysicsRecoveryBottom()
kernel.apply(new_field, initial_field)
tolerance = 1e-12
index = 11 if boundary == "top" else 6
assert abs(true_field.dat.data[index] - new_field.dat.data[index]) < tolerance, \
"Value at %s from physics recovery is not correct" % boundary
def mesh(geometry, element):
Lx = 100.
Ly = 100.
H = 100.
deltax = Lx / 5.
deltay = Ly / 5.
deltaz = H / 5.
ncolumnsy = int(Ly/deltay)
ncolumnsx = int(Lx/deltax)
nlayers = int(H/deltaz)
quadrilateral = True if element == "quadrilateral" else False
if geometry == "periodic-in-both":
m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly,
direction='both', quadrilateral=quadrilateral)
elif geometry == "periodic-in-x":
m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly,
direction='x', quadrilateral=quadrilateral)
elif geometry == "periodic-in-y":
m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly,
direction='y', quadrilateral=quadrilateral)
elif geometry == "non-periodic":
m = RectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly, quadrilateral=quadrilateral)
extruded_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=deltaz)
return extruded_mesh
def run(problem, tensor, size_factor, degree):
formname = problem.__name__
cellname = 'cube' if tensor else 'simplex'
PETSc.Sys.Print("%s: %s, degree=%d" % (formname, cellname, degree))
num_cells = COMM_WORLD.size * max(1, 1e7 * size_factor / (degree + 1)**{'spectral': 4, 'coffee': 6}[args.mode])
h = int(floor(cbrt(num_cells / COMM_WORLD.size)))
w = int(floor(sqrt(num_cells / h)))
d = int(round(num_cells / (w * h)))
num_cells = w * d * h
if tensor:
mesh = ExtrudedMesh(UnitSquareMesh(w, d, quadrilateral=True), h)
else:
mesh = UnitCubeMesh(w, d, h)
comm = mesh.comm
J = problem(mesh, int(degree))
# Warmup and allocate
A = assemble(J, mat_type="matfree", form_compiler_parameters={'mode': args.mode})
A.force_evaluation()
Ap = A.petscmat
x, y = Ap.createVecs()
assert x.size == y.size
num_dofs = x.size
Ap.mult(x, y)
stage = PETSc.Log.Stage("%s(%d) %s" % (formname, degree, cellname))
with stage:
assemble(J, mat_type="matfree", form_compiler_parameters={'mode': args.mode}, tensor=A)
A.force_evaluation()
Ap = A.petscmat
for _ in range(num_matvecs):
Ap.mult(x, y)
parloop = parloop_event.getPerfInfo()
matmult = matmult_event.getPerfInfo()
assert parloop["count"] == 2 * num_matvecs
assert matmult["count"] == num_matvecs
parloop_time = comm.allreduce(parloop["time"], op=MPI.SUM) / (comm.size * num_matvecs)
matmult_time = comm.allreduce(matmult["time"], op=MPI.SUM) / (comm.size * num_matvecs)
matfree_overhead = (1 - parloop_time / matmult_time)
PETSc.Sys.Print("Matrix-free action overhead: %.1f%%" % (matfree_overhead * 100,))
if COMM_WORLD.rank == 0:
header = not os.path.exists(filepath)
data = {"num_cells": num_cells,
"num_dofs": num_dofs,
"num_procs": comm.size,
"tsfc_mode": args.mode,
"problem": formname,
"cell_type": cellname,
"degree": degree,
"matmult_time": matmult_time,
"parloop_time": parloop_time}
df = pandas.DataFrame(data, index=[0])
df.to_csv(filepath, index=False, mode='a', header=header)
def setup_balance(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 5.
deltax = 400
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
output = OutputParameters(dirname=dirname + '/dry_balance',
dumpfreq=10,
dumplist=['u'])
parameters = CompressibleParameters()
state = State(mesh, dt=dt, output=output, parameters=parameters)
eqns = CompressibleEulerEquations(state, "CG", 1)
# Initial conditions
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# Isentropic background state
Tsurf = Constant(300.)
theta0.interpolate(Tsurf)
# Calculate hydrostatic exner
compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True)
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up transport schemes
transported_fields = [
ImplicitMidpoint(state, "u"),
SSPRK3(state, "rho"),
SSPRK3(state, "theta", options=EmbeddedDGOptions())
]
# Set up linear solver
linear_solver = CompressibleSolver(state, eqns)
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver)
return stepper, tmax
def setup_IPdiffusion(dirname, vector, DG):
dt = 0.01
L = 10.
m = PeriodicIntervalMesh(50, L)
mesh = ExtrudedMesh(m, layers=50, layer_height=0.2)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=dt)
parameters = CompressibleParameters()
output = OutputParameters(dirname=dirname)
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
parameters=parameters,
output=output,
fieldlist=fieldlist)
x = SpatialCoordinate(mesh)
if vector:
kappa = Constant([[0.05, 0.], [0., 0.05]])
if DG:
Space = VectorFunctionSpace(mesh, "DG", 1)
else:
Space = state.spaces("HDiv")
fexpr = as_vector([exp(-(L / 2. - x[0])**2 - (L / 2. - x[1])**2), 0.])
else:
kappa = 0.05
if DG:
Space = state.spaces("DG")
else:
Space = state.spaces("HDiv_v")
fexpr = exp(-(L / 2. - x[0])**2 - (L / 2. - x[1])**2)
f = state.fields("f", Space)
try:
f.interpolate(fexpr)
except NotImplementedError:
# finat elements raise NotImplementedError if they don't
# advertise a dual for interpolation.
f.project(fexpr)
mu = 5.
f_diffusion = InteriorPenalty(state,
f.function_space(),
kappa=kappa,
mu=mu)
diffused_fields = [("f", f_diffusion)]
stepper = AdvectionDiffusion(state, diffused_fields=diffused_fields)
return stepper
def tracer_blob_slice(tmpdir, degree):
dt = 0.01
L = 10.
m = PeriodicIntervalMesh(10, L)
mesh = ExtrudedMesh(m, layers=10, layer_height=1.)
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
state = State(mesh, dt=dt, output=output)
tmax = 1.
x = SpatialCoordinate(mesh)
f_init = exp(-((x[0] - 0.5 * L)**2 + (x[1] - 0.5 * L)**2))
return TracerSetup(state, tmax, f_init, family="CG", degree=degree)
def setup_gw(dirname):
nlayers = 10 # horizontal layers
columns = 30 # number of columns
L = 1.e5
m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True)
dt = 6.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
fieldlist = ['u', 'p', 'b']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname + "/gw_incompressible",
dumplist=['u'],
dumpfreq=5)
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="RTCF",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# Initial conditions
u0 = state.fields("u")
p0 = state.fields("p")
b0 = state.fields("b")
# z.grad(bref) = N**2
x, y, z = SpatialCoordinate(mesh)
N = parameters.N
bref = z * (N**2)
b_b = Function(b0.function_space()).interpolate(bref)
b0.interpolate(b_b)
incompressible_hydrostatic_balance(state, b0, p0)
state.initialise([('u', u0), ('p', p0), ('b', b0)])
# Set up forcing
forcing = IncompressibleForcing(state)
return state, forcing
def setup_fallout(dirname):
# declare grid shape, with length L and height H
L = 10.
H = 10.
nlayers = 10
ncolumns = 10
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
dt = 0.1
output = OutputParameters(dirname=dirname + "/fallout",
dumpfreq=10,
dumplist=['rain_mixing_ratio'])
parameters = CompressibleParameters()
diagnostic_fields = [Precipitation()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
Vrho = state.spaces("DG1_equispaced")
problem = [(AdvectionEquation(state, Vrho, "rho", ufamily="CG",
udegree=1), ForwardEuler(state))]
state.fields("rho").assign(1.)
physics_list = [Fallout(state)]
rain0 = state.fields("rain_mixing_ratio")
# set up rain
xc = L / 2
zc = H / 2
rc = H / 4
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2)
rain0.interpolate(rain_expr)
# build time stepper
stepper = PrescribedTransport(state, problem, physics_list=physics_list)
return stepper, 10.0
def mesh(geometry):
L = 100.
H = 100.
deltax = L / 5.
deltaz = H / 5.
nlayers = int(H / deltaz)
ncolumns = int(L / deltax)
if geometry == "periodic":
m = PeriodicIntervalMesh(ncolumns, L)
elif geometry == "non-periodic":
m = IntervalMesh(ncolumns, L)
extruded_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=deltaz)
return extruded_mesh
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
def setup_checkpointing(dirname):
nlayers = 5 # horizontal layers
columns = 15 # number of columns
L = 3.e5
m = PeriodicIntervalMesh(columns, L)
dt = 2.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
output = OutputParameters(dirname=dirname,
dumpfreq=1,
chkptfreq=2,
log_level='INFO')
parameters = CompressibleParameters()
state = State(mesh, dt=dt, output=output, parameters=parameters)
eqns = CompressibleEulerEquations(state, "CG", 1)
# Set up transport schemes
transported_fields = []
transported_fields.append(SSPRK3(state, "u"))
transported_fields.append(SSPRK3(state, "rho"))
transported_fields.append(SSPRK3(state, "theta"))
# Set up linear solver
linear_solver = CompressibleSolver(state, eqns)
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver)
return state, stepper, dt
def tracer_slice(tmpdir, degree):
n = 30 if degree == 0 else 15
m = PeriodicIntervalMesh(n, 1.)
mesh = ExtrudedMesh(m, layers=n, layer_height=1. / n)
# Parameters chosen so that dt != 1 and u != 1
# Gaussian is translated by 1.5 times width of domain to demonstrate
# that transport is working correctly
dt = 0.01
tmax = 0.75
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
state = State(mesh, dt=dt, output=output)
uexpr = as_vector([2.0, 0.0])
x = SpatialCoordinate(mesh)
width = 1. / 10.
f0 = 0.5
fmax = 2.0
xc_init = 0.25
xc_end = 0.75
r_init = sqrt((x[0] - xc_init)**2 + (x[1] - 0.5)**2)
r_end = sqrt((x[0] - xc_end)**2 + (x[1] - 0.5)**2)
f_init = f0 + (fmax - f0) * exp(-(r_init / width)**2)
f_end = f0 + (fmax - f0) * exp(-(r_end / width)**2)
tol = 0.12
return TracerSetup(state,
tmax,
f_init,
f_end,
"CG",
degree,
uexpr,
tol=tol)
def setup_fallout(dirname):
# declare grid shape, with length L and height H
L = 10.
H = 10.
nlayers = 10
ncolumns = 10
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta', 'rain']
timestepping = TimesteppingParameters(dt=0.1, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/fallout",
dumpfreq=10,
dumplist=['rain'])
parameters = CompressibleParameters()
diagnostic_fields = [Precipitation()]
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vt = theta0.function_space()
# declare tracer field and a background field
rain0 = state.fields("rain", Vt)
# set up rain
xc = L / 2
zc = H / 2
rc = H / 4
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2)
rain0.interpolate(rain_expr)
rho0.assign(1.0)
state.initialise([('u', u0), ('rho', rho0), ('rain', rain0)])
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", NoAdvection(state, rho0, None)))
advected_fields.append(("rain", NoAdvection(state, rain0, None)))
physics_list = [Fallout(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 10.0
def test_limit_midpoints(profile):
# ------------------------------------------------------------------------ #
# Set up meshes and spaces
# ------------------------------------------------------------------------ #
m = IntervalMesh(3, 3)
mesh = ExtrudedMesh(m, layers=1, layer_height=3.0)
cell = m.ufl_cell().cellname()
DG_hori_elt = FiniteElement("DG", cell, 1, variant='equispaced')
DG_vert_elt = FiniteElement("DG", interval, 1, variant='equispaced')
Vt_vert_elt = FiniteElement("CG", interval, 2)
DG_elt = TensorProductElement(DG_hori_elt, DG_vert_elt)
theta_elt = TensorProductElement(DG_hori_elt, Vt_vert_elt)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_elt))
DG1 = FunctionSpace(mesh, DG_elt)
new_field = Function(Vt_brok)
init_field = Function(Vt_brok)
DG1_field = Function(DG1)
# ------------------------------------------------------------------------ #
# Initial conditions
# ------------------------------------------------------------------------ #
_, z = SpatialCoordinate(mesh)
if profile == 'undershoot':
# A quadratic whose midpoint is lower than the top and bottom values
init_expr = (80. / 9.) * z**2 - 20. * z + 300.
elif profile == 'overshoot':
# A quadratic whose midpoint is higher than the top and bottom values
init_expr = (-80. / 9.) * z**2 + (100. / 3) * z + 300.
elif profile == 'linear':
# Linear profile which must be unchanged
init_expr = (20. / 3.) * z + 300.
else:
raise NotImplementedError
# Linear DG field has the same values at top and bottom as quadratic
DG_expr = (20. / 3.) * z + 300.
init_field.interpolate(init_expr)
DG1_field.interpolate(DG_expr)
# ------------------------------------------------------------------------ #
# Apply kernel
# ------------------------------------------------------------------------ #
kernel = kernels.LimitMidpoints(Vt_brok)
kernel.apply(new_field, DG1_field, init_field)
# ------------------------------------------------------------------------ #
# Check values
# ------------------------------------------------------------------------ #
tol = 1e-12
assert np.max(new_field.dat.data) <= np.max(init_field.dat.data) + tol, \
'LimitMidpoints kernel is giving an overshoot'
assert np.min(new_field.dat.data) >= np.min(init_field.dat.data) - tol, \
'LimitMidpoints kernel is giving an undershoot'
if profile == 'linear':
assert np.allclose(init_field.dat.data, new_field.dat.data), \
'For a profile with no maxima or minima, the LimitMidpoints ' + \
'kernel should leave the field unchanged'
def setup_3d_recovery(dirname):
L = 100.
H = 10.
W = 1.
deltax = L / 5.
deltay = W / 5.
deltaz = H / 5.
nlayers = int(H / deltaz)
ncolumnsx = int(L / deltax)
ncolumnsy = int(W / deltay)
m = RectangleMesh(ncolumnsx, ncolumnsy, L, W, quadrilateral=True)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
x, y, z = SpatialCoordinate(mesh)
# horizontal base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
u_hori = FiniteElement("RTCF", cell, 1)
w_hori = FiniteElement("DG", cell, 0)
# vertical base spaces
u_vert = FiniteElement("DG", interval, 0)
w_vert = FiniteElement("CG", interval, 1)
# build elements
u_element = HDiv(TensorProductElement(u_hori, u_vert))
w_element = HDiv(TensorProductElement(w_hori, w_vert))
theta_element = TensorProductElement(w_hori, w_vert)
v_element = u_element + w_element
# spaces
VDG0 = FunctionSpace(mesh, "DG", 0)
VCG1 = FunctionSpace(mesh, "CG", 1)
VDG1 = FunctionSpace(mesh, "DG", 1)
Vt = FunctionSpace(mesh, theta_element)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element))
Vu = FunctionSpace(mesh, v_element)
VuCG1 = VectorFunctionSpace(mesh, "CG", 1)
VuDG1 = VectorFunctionSpace(mesh, "DG", 1)
# set up initial conditions
np.random.seed(0)
expr = np.random.randn(
) + np.random.randn() * x + np.random.randn() * y + np.random.randn(
) * z + np.random.randn() * x * y + np.random.randn(
) * x * z + np.random.randn() * y * z + np.random.randn() * x * y * z
# our actual theta and rho and v
rho_CG1_true = Function(VCG1).interpolate(expr)
theta_CG1_true = Function(VCG1).interpolate(expr)
v_CG1_true = Function(VuCG1).interpolate(as_vector([expr, expr, expr]))
rho_Vt_true = Function(Vt).interpolate(expr)
# make the initial fields by projecting expressions into the lowest order spaces
rho_DG0 = Function(VDG0).interpolate(expr)
rho_CG1 = Function(VCG1)
theta_Vt = Function(Vt).interpolate(expr)
theta_CG1 = Function(VCG1)
v_Vu = Function(Vu).project(as_vector([expr, expr, expr]))
v_CG1 = Function(VuCG1)
rho_Vt = Function(Vt)
# make the recoverers and do the recovery
rho_recoverer = Recoverer(rho_DG0,
rho_CG1,
VDG=VDG1,
boundary_method=Boundary_Method.dynamics)
theta_recoverer = Recoverer(theta_Vt,
theta_CG1,
VDG=VDG1,
boundary_method=Boundary_Method.dynamics)
v_recoverer = Recoverer(v_Vu,
v_CG1,
VDG=VuDG1,
boundary_method=Boundary_Method.dynamics)
rho_Vt_recoverer = Recoverer(rho_DG0,
rho_Vt,
VDG=Vt_brok,
boundary_method=Boundary_Method.physics)
rho_recoverer.project()
theta_recoverer.project()
v_recoverer.project()
rho_Vt_recoverer.project()
rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true)
theta_diff = errornorm(theta_CG1, theta_CG1_true) / norm(theta_CG1_true)
v_diff = errornorm(v_CG1, v_CG1_true) / norm(v_CG1_true)
rho_Vt_diff = errornorm(rho_Vt, rho_Vt_true) / norm(rho_Vt_true)
return (rho_diff, theta_diff, v_diff, rho_Vt_diff)
def setup_sk(dirname):
nlayers = 10 # horizontal layers
columns = 30 # number of columns
L = 1.e5
m = PeriodicIntervalMesh(columns, L)
dt = 6.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# Set up points for output at the centre of the domain, edges and corners.
# The point at x=L*(13.0/30) is in the halo region for a two-way MPI decomposition
points_x = [0.0, L * (13.0 / 30), L / 2.0, L]
points_z = [0.0, H / 2.0, H]
points = np.array([p for p in itertools.product(points_x, points_z)])
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname + "/sk_nonlinear",
dumplist=['u'],
dumpfreq=5,
log_level=INFO,
point_data=[('rho', points), ('u', points)])
parameters = CompressibleParameters()
diagnostic_fields = [CourantNumber()]
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Thermodynamic constants required for setting initial conditions
# and reference profiles
g = parameters.g
N = parameters.N
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
x, z = SpatialCoordinate(mesh)
Tsurf = 300.
thetab = Tsurf * exp(N**2 * z / g)
theta_b = Function(Vt).interpolate(thetab)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b)
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
u0.project(as_vector([20.0, 0.0]))
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt)
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 2 * dt
def run_profliler(hybridization,
model_degree,
model_family,
mesh_degree,
cfl,
refinements,
layers,
debug,
rtol,
flexsolver=True,
stronger_smoother=False,
suppress_data_output=False):
nlayers = layers # Number of vertical layers
refinements = refinements # Number of horiz. cells = 20*(4^refinements)
hybrid = bool(hybridization)
# Set up problem parameters
parameters = CompressibleParameters()
a_ref = 6.37122e6 # Radius of the Earth (m)
X = 1.0 # Reduced-size Earth reduction factor
a = a_ref / X # Scaled radius of planet (m)
g = parameters.g # Acceleration due to gravity (m/s^2)
N = parameters.N # Brunt-Vaisala frequency (1/s)
p_0 = parameters.p_0 # Reference pressure (Pa, not hPa)
c_p = parameters.cp # SHC of dry air at const. pressure (J/kg/K)
R_d = parameters.R_d # Gas constant for dry air (J/kg/K)
kappa = parameters.kappa # R_d/c_p
T_eq = 300.0 # Isothermal atmospheric temperature (K)
p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator
u_0 = 20.0 # Maximum amplitude of the zonal wind (m/s)
d = 5000.0 # Width parameter for Theta'
lamda_c = 2.0 * np.pi / 3.0 # Longitudinal centerpoint of Theta'
phi_c = 0.0 # Lat. centerpoint of Theta' (equator)
deltaTheta = 1.0 # Maximum amplitude of Theta' (K)
L_z = 20000.0 # Vert. wave length of the Theta' perturb.
gamma = (1 - kappa) / kappa
cs = sqrt(c_p * T_eq / gamma) # Speed of sound in an air parcel
if model_family == "RTCF":
# Cubed-sphere mesh
m = CubedSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
elif model_family == "RT" or model_family == "BDFM":
m = IcosahedralSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
else:
raise ValueError("Unknown family: %s" % model_family)
cell_vs = interpolate(CellVolume(m), FunctionSpace(m, "DG", 0))
a_max = fmax(cell_vs)
dx_max = sqrt(a_max)
u_max = u_0
PETSc.Sys.Print("\nDetermining Dt from specified horizontal CFL: %s" % cfl)
dt = int(cfl * (dx_max / cs))
# Height position of the model top (m)
z_top = 1.0e4
deltaz = z_top / nlayers
vertical_cfl = dt * (cs / deltaz)
PETSc.Sys.Print("""
Problem parameters:\n
Profiling linear solver for the compressible Euler equations.\n
Speed of sound in compressible atmosphere: %s,\n
Hybridized compressible solver: %s,\n
Model degree: %s,\n
Model discretization: %s,\n
Mesh degree: %s,\n
Horizontal refinements: %s,\n
Vertical layers: %s,\n
Dx (max, m): %s,\n
Dz (m): %s,\n
Dt (s): %s,\n
horizontal CFL: %s,\n
vertical CFL: %s.
""" % (cs, hybrid, model_degree, model_family, mesh_degree, refinements,
nlayers, dx_max, deltaz, dt, cfl, vertical_cfl))
# Build volume mesh
mesh = ExtrudedMesh(m,
layers=nlayers,
layer_height=deltaz,
extrusion_type="radial")
x = SpatialCoordinate(mesh)
# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
W_Q1 = FunctionSpace(mesh, "CG", 1)
z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a
z = Function(W_Q1).interpolate(z_expr)
lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]))
lat = Function(W_Q1).interpolate(lat_expr)
lon = Function(W_Q1).interpolate(atan_2(x[1], x[0]))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=1, maxi=1)
dirname = 'meanflow_ref'
if hybrid:
dirname += '_hybridization'
# No output
output = OutputParameters(dumpfreq=3600,
dirname=dirname,
perturbation_fields=['theta', 'rho'],
dump_vtus=False,
dump_diagnostics=False,
checkpoint=False,
log_level='INFO')
diagnostics = Diagnostics(*fieldlist)
state = State(mesh,
vertical_degree=model_degree,
horizontal_degree=model_degree,
family=model_family,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# Initial conditions
u0 = state.fields.u
theta0 = state.fields.theta
rho0 = state.fields.rho
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
x = SpatialCoordinate(mesh)
# Random velocity field
CG2 = VectorFunctionSpace(mesh, "CG", 2)
urand = Function(CG2)
urand.dat.data[:] += np.random.randn(*urand.dat.data.shape)
u0.project(urand)
# Surface temperature
G = g**2 / (N**2 * c_p)
Ts_expr = G + (T_eq - G) * exp(-(u_max * N**2 / (4 * g * g)) * u_max *
(cos(2.0 * lat) - 1.0))
Ts = Function(W_Q1).interpolate(Ts_expr)
# Surface pressure
ps_expr = p_eq * exp((u_max / (4.0 * G * R_d)) * u_max *
(cos(2.0 * lat) - 1.0)) * (Ts / T_eq)**(1.0 / kappa)
ps = Function(W_Q1).interpolate(ps_expr)
# Background pressure
p_expr = ps * (1 + G / Ts * (exp(-N**2 * z / g) - 1))**(1.0 / kappa)
p = Function(W_Q1).interpolate(p_expr)
# Background temperature
Tb_expr = G * (1 - exp(N**2 * z / g)) + Ts * exp(N**2 * z / g)
Tb = Function(W_Q1).interpolate(Tb_expr)
# Background potential temperature
thetab_expr = Tb * (p_0 / p)**kappa
thetab = Function(W_Q1).interpolate(thetab_expr)
theta_b = Function(theta0.function_space()).interpolate(thetab)
rho_b = Function(rho0.function_space())
sin_tmp = sin(lat) * sin(phi_c)
cos_tmp = cos(lat) * cos(phi_c)
r = a * acos(sin_tmp + cos_tmp * cos(lon - lamda_c))
s = (d**2) / (d**2 + r**2)
theta_pert = deltaTheta * s * sin(2 * np.pi * z / L_z)
theta0.interpolate(theta_pert)
# Compute the balanced density
PETSc.Sys.Print("Computing balanced density field...\n")
# Use vert. hybridization preconditioner for initialization
pi_params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
'vert_hybridization': {
'ksp_type': 'gmres',
'pc_type': 'gamg',
'pc_gamg_sym_graph': True,
'ksp_rtol': 1e-12,
'ksp_atol': 1e-12,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_max_it': 3,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
if debug:
pi_params['vert_hybridization']['ksp_monitor_true_residual'] = None
compressible_hydrostatic_balance(state,
theta_b,
rho_b,
top=False,
pi_boundary=(p / p_0)**kappa,
solve_for_rho=False,
params=pi_params)
# Random potential temperature perturbation
theta0.assign(0.0)
theta0.dat.data[:] += np.random.randn(len(theta0.dat.data))
# Random density field
rho0.assign(0.0)
rho0.dat.data[:] += np.random.randn(len(rho0.dat.data))
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt, equation_form="advective")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn, subcycles=2)))
advected_fields.append(
("theta", SSPRK3(state, theta0, thetaeqn, subcycles=2)))
# Set up linear solver
if hybrid:
outer_solver_type = "Hybrid_SCPC"
PETSc.Sys.Print("""
Setting up hybridized solver on the traces.""")
if flexsolver:
inner_solver_type = "fgmres_gamg_gmres_smoother"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {
'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
else:
inner_solver_type = "fgmres_ml_richardson"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'ml',
'pc_mg_cycles': 1,
'pc_ml_maxNlevels': 25,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_richardson_scale': 0.8,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
if stronger_smoother:
inner_parameters['mg_levels']['ksp_max_it'] = 5
inner_solver_type += "_stronger"
if debug:
PETSc.Sys.Print("""Debugging on.""")
inner_parameters['ksp_monitor_true_residual'] = None
PETSc.Sys.Print("Inner solver: %s" % inner_solver_type)
# Use Firedrake static condensation interface
solver_parameters = {
'mat_type': 'matfree',
'pmat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
'pc_sc_eliminate_fields': '0, 1',
'condensed_field': inner_parameters
}
linear_solver = HybridizedCompressibleSolver(
state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
else:
outer_solver_type = "gmres_SchurPC"
PETSc.Sys.Print("""
Setting up GCR fieldsplit solver with Schur complement PC.""")
solver_parameters = {
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'fgmres',
'ksp_max_it': 100,
'ksp_rtol': rtol,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
},
'fieldsplit_1': {
'ksp_type': 'preonly',
'ksp_max_it': 30,
'ksp_monitor_true_residual': None,
'pc_type': 'hypre',
'pc_hypre_type': 'boomeramg',
'pc_hypre_boomeramg_max_iter': 1,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_coarsen_type': 'Falgout',
'pc_hypre_boomeramg_smooth_type': 'Euclid',
'pc_hypre_boomeramg_eu_bj': 1,
'pc_hypre_boomeramg_interptype': 'classical',
'pc_hypre_boomeramg_P_max': 0,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_strong_threshold': 0.25,
'pc_hypre_boomeramg_max_levels': 25,
'pc_hypre_boomeramg_no_CF': False
}
}
inner_solver_type = "hypre"
if debug:
solver_parameters['ksp_monitor_true_residual'] = None
linear_solver = CompressibleSolver(state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
param_info = ParameterInfo(dt=dt,
deltax=dx_max,
deltaz=deltaz,
horizontal_courant=cfl,
vertical_courant=vertical_cfl,
family=model_family,
model_degree=model_degree,
mesh_degree=mesh_degree,
solver_type=outer_solver_type,
inner_solver_type=inner_solver_type)
# Build profiler
profiler = Profiler(parameterinfo=param_info,
state=state,
advected_fields=advected_fields,
linear_solver=linear_solver,
forcing=compressible_forcing,
suppress_data_output=suppress_data_output)
PETSc.Sys.Print("Starting profiler...\n")
profiler.run(t=0, tmax=dt)
def setup_condens(dirname):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/condens",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Sum('water_v', 'water_c')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
Vt = theta0.function_space()
Vr = rho0.function_space()
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# declare tracer field and a background field
water_v0 = state.fields("water_v", Vt)
water_c0 = state.fields("water_c", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
# set up water_v
xc = 500.
zc = 350.
rc = 250.
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
w_expr = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r)))
# set up velocity field
u_max = 10.0
psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) *
sin(pi * x[1] / L))
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b)
rho0.interpolate(rho_b)
water_v0.interpolate(w_expr)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# set up advection schemes
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state,
Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("water_v", SSPRK3(state, water_v0, thetaeqn)))
advected_fields.append(("water_c", SSPRK3(state, water_c0, thetaeqn)))
physics_list = [Condensation(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 5.0
def run_dry_compressible(tmpdir):
dt = 6.0
tmax = 2 * dt
nlayers = 10 # horizontal layers
ncols = 10 # number of columns
Lx = 1000.0
Lz = 1000.0
m = PeriodicIntervalMesh(ncols, Lx)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz / nlayers)
output = OutputParameters(dirname=tmpdir + "/dry_compressible",
dumpfreq=2,
chkptfreq=2)
parameters = CompressibleParameters()
R_d = parameters.R_d
g = parameters.g
state = State(mesh, dt=dt, output=output, parameters=parameters)
eqn = CompressibleEulerEquations(state, "CG", 1)
# Initial conditions
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# Approximate hydrostatic balance
x, z = SpatialCoordinate(mesh)
T = Constant(300.0)
zH = R_d * T / g
p = Constant(100000.0) * exp(-z / zH)
theta0.interpolate(tde.theta(parameters, T, p))
rho0.interpolate(p / (R_d * T))
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Add perturbation
r = sqrt((x - Lx / 2)**2 + (z - Lz / 2)**2)
theta_pert = 1.0 * exp(-(r / (Lx / 5))**2)
theta0.interpolate(theta0 + theta_pert)
# Set up transport schemes
transported_fields = [
ImplicitMidpoint(state, "u"),
SSPRK3(state, "rho"),
SSPRK3(state, "theta")
]
# Set up linear solver for the timestepping scheme
linear_solver = CompressibleSolver(state, eqn)
# build time stepper
stepper = CrankNicolson(state,
eqn,
transported_fields,
linear_solver=linear_solver)
# Run
stepper.run(t=0, tmax=tmax)
# State for checking checkpoints
checkpoint_name = 'dry_compressible_chkpt'
new_path = join(abspath(dirname(__file__)), '..',
f'data/{checkpoint_name}')
check_output = OutputParameters(dirname=tmpdir + "/dry_compressible",
checkpoint_pickup_filename=new_path)
check_state = State(mesh,
dt=dt,
output=check_output,
parameters=parameters)
check_eqn = CompressibleEulerEquations(check_state, "CG", 1)
check_stepper = CrankNicolson(check_state, check_eqn, [])
check_stepper.run(t=0, tmax=0, pickup=True)
return state, check_state
def setup_saturated(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 3.
deltax = 400.
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# option to easily change between recovered and not if necessary
# default should be to use lowest order set of spaces
recovered = True
degree = 0 if recovered else 1
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + '/saturated_balance',
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['water_v'])
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = [Theta_e()]
state = State(mesh,
vertical_degree=degree,
horizontal_degree=degree,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water_v0 = state.fields("water_v", theta0.function_space())
water_c0 = state.fields("water_c", theta0.function_space())
moisture = ['water_v', 'water_c']
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
total_water = Constant(0.02)
theta_e = Function(Vt).interpolate(Tsurf)
water_t = Function(Vt).interpolate(total_water)
# Calculate hydrostatic Pi
saturated_hydrostatic_balance(state, theta_e, water_t)
water_c0.assign(water_t - water_v0)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho0), ('theta', theta0),
('water_v', water_v0)])
# Set up advection schemes
if recovered:
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, "DG", 1)
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu,
boundary_method=Boundary_Method.dynamics)
rho_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vr,
boundary_method=Boundary_Method.dynamics)
theta_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
ueqn = EmbeddedDGAdvection(state,
Vu,
equation_form="advective",
options=u_opts)
rhoeqn = EmbeddedDGAdvection(state,
Vr,
equation_form="continuity",
options=rho_opts)
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=theta_opts)
else:
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=EmbeddedDGOptions())
advected_fields = [('rho', SSPRK3(state, rho0, rhoeqn)),
('theta', SSPRK3(state, theta0, thetaeqn)),
('water_v', SSPRK3(state, water_v0, thetaeqn)),
('water_c', SSPRK3(state, water_c0, thetaeqn))]
if recovered:
advected_fields.append(('u', SSPRK3(state, u0, ueqn)))
else:
advected_fields.append(('u', ThetaMethod(state, u0, ueqn)))
linear_solver = CompressibleSolver(state, moisture=moisture)
# Set up forcing
if recovered:
compressible_forcing = CompressibleForcing(state,
moisture=moisture,
euler_poincare=False)
else:
compressible_forcing = CompressibleForcing(state, moisture=moisture)
# add physics
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state,
advected_fields,
linear_solver,
compressible_forcing,
physics_list=physics_list)
return stepper, tmax
T_eq = 300.0 # Isothermal atmospheric temperature (K)
p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator
u_max = 20.0 # Maximum amplitude of the zonal wind (m/s)
d = 5000.0 # Width parameter for Theta'
lamda_c = 2.0 * pi / 3.0 # Longitudinal centerpoint of Theta'
phi_c = 0.0 # Latitudinal centerpoint of Theta' (equator)
deltaTheta = 1.0 # Maximum amplitude of Theta' (K)
L_z = 20000.0 # Vertical wave length of the Theta' perturb.
# Cubed-sphere horizontal mesh
m = CubedSphereMesh(radius=a, refinement_level=refinements, degree=2)
# Build volume mesh
z_top = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m,
layers=nlayers,
layer_height=z_top / nlayers,
extrusion_type="radial")
x = SpatialCoordinate(mesh)
# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
W_Q1 = FunctionSpace(mesh, "CG", 1)
z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a
z = Function(W_Q1).interpolate(z_expr)
lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]))
lat = Function(W_Q1).interpolate(lat_expr)
lon = Function(W_Q1).interpolate(atan_2(x[1], x[0]))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
dirname = 'meanflow_ref_hybridization'
form, parameters=parameters)[0]
indices = IndexDict(
{idx: sympy.symbols(name)
for idx, name in index_names})
expr = expression(impero_kernel.tree,
impero_kernel.temporaries,
indices,
top=True)
p1 = sympy.symbols("p") + 1
'''Currently assume p+1 quad points in each direction.'''
return expr.subs([(i, p1) for i in indices.values()]).expand()
m = ExtrudedMesh(UnitSquareMesh(2, 2, quadrilateral=True), 2)
mass = form.mass(m, 6)
poisson = form.poisson(m, 6)
hyperelasticity = form.hyperelasticity(m, 6)
curl_curl = form.curl_curl(m, 6)
parameters = firedrake.parameters['form_compiler'].copy()
parameters['return_impero'] = True
parameters['mode'] = 'spectral'
for mode, action in (("assembly", False), ("action", True)):
print(mode)
print(" mass: ", complexity(mass, parameters, action))
print(" laplacian: ", complexity(poisson, parameters, action))
print(" hyperelasticity:", complexity(hyperelasticity, parameters,
def setup_tracer(dirname, hybridization):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname+"/tracer",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Difference('theta', 'tracer')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# declare tracer field and a background field
tracer0 = state.fields("tracer", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b,
solve_for_rho=True)
# set up perturbation to theta
xc = 500.
zc = 350.
rc = 250.
x = SpatialCoordinate(mesh)
r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2)
theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r)))
theta0.interpolate(theta_b + theta_pert)
rho0.interpolate(rho_b)
tracer0.interpolate(theta0)
state.initialise([('u', u0),
('rho', rho0),
('theta', theta0),
('tracer', tracer0)])
state.set_reference_profiles([('rho', rho_b),
('theta', theta_b)])
# set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("tracer", SSPRK3(state, tracer0, thetaeqn)))
# Set up linear solver
if hybridization:
linear_solver = HybridizedCompressibleSolver(state)
else:
linear_solver = CompressibleSolver(state)
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 100.0
def make_mesh(nx, ny):
return ExtrudedMesh(UnitIntervalMesh(nx), ny, layer_height=1.0 / ny)
def setup_saturated(dirname, recovered):
# set up grid and time stepping parameters
dt = 1.
tmax = 3.
deltax = 400.
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# option to easily change between recovered and not if necessary
# default should be to use lowest order set of spaces
degree = 0 if recovered else 1
output = OutputParameters(dirname=dirname + '/saturated_balance',
dumpfreq=1,
dumplist=['u'])
parameters = CompressibleParameters()
diagnostic_fields = [Theta_e()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
tracers = [WaterVapour(), CloudWater()]
if recovered:
u_transport_option = "vector_advection_form"
else:
u_transport_option = "vector_invariant_form"
eqns = CompressibleEulerEquations(state,
"CG",
degree,
u_transport_option=u_transport_option,
active_tracers=tracers)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water_v0 = state.fields("vapour_mixing_ratio")
water_c0 = state.fields("cloud_liquid_mixing_ratio")
moisture = ['vapour_mixing_ratio', 'cloud_liquid_mixing_ratio']
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
total_water = Constant(0.02)
theta_e = Function(Vt).interpolate(Tsurf)
water_t = Function(Vt).interpolate(total_water)
# Calculate hydrostatic exner
saturated_hydrostatic_balance(state, theta_e, water_t)
water_c0.assign(water_t - water_v0)
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up transport schemes
if recovered:
VDG1 = state.spaces("DG1_equispaced")
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element())
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu,
boundary_method=Boundary_Method.dynamics)
rho_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vr,
boundary_method=Boundary_Method.dynamics)
theta_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
wv_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
wc_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
else:
rho_opts = None
theta_opts = EmbeddedDGOptions()
wv_opts = EmbeddedDGOptions()
wc_opts = EmbeddedDGOptions()
transported_fields = [
SSPRK3(state, 'rho', options=rho_opts),
SSPRK3(state, 'theta', options=theta_opts),
SSPRK3(state, 'vapour_mixing_ratio', options=wv_opts),
SSPRK3(state, 'cloud_liquid_mixing_ratio', options=wc_opts)
]
if recovered:
transported_fields.append(SSPRK3(state, 'u', options=u_opts))
else:
transported_fields.append(ImplicitMidpoint(state, 'u'))
linear_solver = CompressibleSolver(state, eqns, moisture=moisture)
# add physics
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver,
physics_list=physics_list)
return stepper, tmax
def make_mesh(nx, ny, nz, quadrilateral=False):
return ExtrudedMesh(UnitSquareMesh(nx, ny,
quadrilateral=quadrilateral),
nz,
layer_height=1.0 / nz)
def setup_balance(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 5.
deltax = 400
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + '/dry_balance',
dumpfreq=10,
dumplist=['u'])
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = []
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
theta0.interpolate(Tsurf)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up advection schemes
ueqn = VectorInvariant(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=EmbeddedDGOptions())
advected_fields = [("u", ThetaMethod(state, u0, ueqn)),
("rho", SSPRK3(state, rho0, rhoeqn)),
("theta", SSPRK3(state, theta0, thetaeqn))]
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, tmax
def setup_recovered_space(dirname):
# declare grid shape, with length L and height H
L = 400.
H = 400.
nlayers = int(H / 20.)
ncolumns = int(L / 20.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/recovered_space_test",
dumpfreq=5,
dumplist=['u'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# declare initial fields
u0 = state.fields("u")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
VDG0 = FunctionSpace(mesh, "DG", 0)
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
# set up tracer field
tracer0 = state.fields("tracer", VDG0)
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# set up bubble
xc = 200.
zc = 200.
rc = 100.
tracer0.interpolate(
conditional(
sqrt((x - xc)**2.0) < rc,
conditional(
sqrt((z - zc)**2.0) < rc, Constant(0.2), Constant(0.0)),
Constant(0.0)))
# set up velocity field
u_max = Constant(10.0)
psi_expr = -u_max * z
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
state.initialise([('u', u0), ('tracer', tracer0)])
# set up advection schemes
recovered_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=VDG0,
boundary_method=Boundary_Method.dynamics)
tracereqn = EmbeddedDGAdvection(state,
VDG0,
equation_form="continuity",
options=recovered_opts)
# build advection dictionary
advected_fields = []
advected_fields.append(('tracer', SSPRK3(state, tracer0, tracereqn)))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 100.0
else:
tmax = 3600.
if '--hybridization' in sys.argv:
hybridization = True
else:
hybridization = False
nlayers = 50 # horizontal layers
columns = 50 # number of columns
L = 3.0e5
m = PeriodicIntervalMesh(columns, L)
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
dirname = 'sk_linear'
if hybridization:
dirname += '_hybridization'
output = OutputParameters(dirname=dirname,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
def setup_hori_limiters(dirname):
# declare grid shape
L = 400.
H = L
ncolumns = int(L / 10.)
nlayers = ncolumns
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
fieldlist = ['u']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/limiting_hori",
dumpfreq=5,
dumplist=['u'],
perturbation_fields=['theta0', 'theta1'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# make elements
# v is continuous in vertical, h is horizontal
cell = mesh._base_mesh.ufl_cell().cellname()
DG0_element = FiniteElement("DG", cell, 0)
CG1_element = FiniteElement("CG", cell, 1)
DG1_element = FiniteElement("DG", cell, 1)
CG2_element = FiniteElement("CG", cell, 2)
V0_element = TensorProductElement(DG0_element, CG1_element)
V1_element = TensorProductElement(DG1_element, CG2_element)
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
V0 = FunctionSpace(mesh, V0_element)
V1 = FunctionSpace(mesh, V1_element)
V0_brok = FunctionSpace(mesh, BrokenElement(V0.ufl_element()))
V0_spaces = (VDG1, VCG1, V0_brok)
# declare initial fields
u0 = state.fields("u")
theta0 = state.fields("theta0", V0)
theta1 = state.fields("theta1", V1)
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# Isentropic background state
Tsurf = 300.
thetab = Constant(Tsurf)
theta_b1 = Function(V1).interpolate(thetab)
theta_b0 = Function(V0).interpolate(thetab)
# set up bubble
xc = 200.
zc = 200.
rc = 100.
theta_expr = conditional(
sqrt((x - xc)**2.0) < rc,
conditional(sqrt((z - zc)**2.0) < rc, Constant(2.0), Constant(0.0)),
Constant(0.0))
theta_pert1 = Function(V1).interpolate(theta_expr)
theta_pert0 = Function(V0).interpolate(theta_expr)
# set up velocity field
u_max = Constant(10.0)
psi_expr = -u_max * z
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b0 + theta_pert0)
theta1.interpolate(theta_b1 + theta_pert1)
state.initialise([('u', u0), ('theta1', theta1), ('theta0', theta0)])
state.set_reference_profiles([('theta1', theta_b1), ('theta0', theta_b0)])
# set up advection schemes
thetaeqn1 = EmbeddedDGAdvection(state, V1, equation_form="advective")
thetaeqn0 = EmbeddedDGAdvection(state,
V0,
equation_form="advective",
recovered_spaces=V0_spaces)
# build advection dictionary
advected_fields = []
advected_fields.append(('u', NoAdvection(state, u0, None)))
advected_fields.append(('theta1',
SSPRK3(state,
theta1,
thetaeqn1,
limiter=ThetaLimiter(thetaeqn1))))
advected_fields.append(('theta0',
SSPRK3(state,
theta0,
thetaeqn0,
limiter=VertexBasedLimiter(VDG1))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 40.0