def setup(self, state):
if not self._initialised:
space = state.fields("theta").function_space()
broken_space = FunctionSpace(state.mesh, BrokenElement(space.ufl_element()))
boundary_method = Boundary_Method.physics if (state.vertical_degree == 0 and state.horizontal_degree == 0) else None
super().setup(state, space=space)
# now let's attach all of our fields
self.u = state.fields("u")
self.rho = state.fields("rho")
self.theta = state.fields("theta")
self.rho_averaged = Function(space)
self.recoverer = Recoverer(self.rho, self.rho_averaged, VDG=broken_space, boundary_method=boundary_method)
try:
self.r_v = state.fields("water_v")
except NotImplementedError:
self.r_v = Constant(0.0)
try:
self.r_c = state.fields("water_c")
except NotImplementedError:
self.r_c = Constant(0.0)
try:
self.rain = state.fields("rain")
except NotImplementedError:
self.rain = Constant(0.0)
# now let's store the most common expressions
self.pi = thermodynamics.pi(state.parameters, self.rho_averaged, self.theta)
self.T = thermodynamics.T(state.parameters, self.theta, self.pi, r_v=self.r_v)
self.p = thermodynamics.p(state.parameters, self.pi)
self.r_l = self.r_c + self.rain
self.r_t = self.r_v + self.r_c + self.rain
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 __init__(self,
state,
V,
ibp=IntegrateByParts.ONCE,
equation_form="advective",
vector_manifold=False,
solver_params=None,
outflow=False,
options=None):
if options is None:
raise ValueError(
"Must provide an instance of the AdvectionOptions class")
else:
self.options = options
if options.name == "embedded_dg" and options.embedding_space is None:
V_elt = BrokenElement(V.ufl_element())
options.embedding_space = FunctionSpace(state.mesh, V_elt)
super().__init__(state=state,
V=options.embedding_space,
ibp=ibp,
equation_form=equation_form,
vector_manifold=vector_manifold,
solver_params=solver_params,
outflow=outflow)
def __init__(self, space):
"""
Initialise limiter
:arg space: the space in which the transported variables lies.
It should be a form of the DG1xCG2 space.
"""
if not space.extruded:
raise ValueError(
'The Theta Limiter can only be used on an extruded mesh')
# check that horizontal degree is 1 and vertical degree is 2
sub_elements = space.ufl_element().sub_elements()
if (sub_elements[0].family() not in ['Discontinuous Lagrange', 'DQ']
or sub_elements[1].family() != 'Lagrange'
or space.ufl_element().degree() != (1, 2)):
raise ValueError(
'Theta Limiter should only be used with the DG1xCG2 space')
# Transport will happen in broken form of Vtheta
mesh = space.mesh()
self.Vt_brok = FunctionSpace(mesh, BrokenElement(space.ufl_element()))
# Create equispaced DG1 space needed for limiting
cell = mesh._base_mesh.ufl_cell().cellname()
DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced")
DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced")
CG2_vert_elt = FiniteElement("CG", interval, 2)
DG1_element = TensorProductElement(DG1_hori_elt, DG1_vert_elt)
Vt_element = TensorProductElement(DG1_hori_elt, CG2_vert_elt)
DG1_equispaced = FunctionSpace(mesh, DG1_element)
Vt_equispaced = FunctionSpace(mesh, Vt_element)
Vt_brok_equispaced = FunctionSpace(
mesh, BrokenElement(Vt_equispaced.ufl_element()))
self.vertex_limiter = VertexBasedLimiter(DG1_equispaced)
self.field_hat = Function(Vt_brok_equispaced)
self.field_old = Function(Vt_brok_equispaced)
self.field_DG1 = Function(DG1_equispaced)
self._limit_midpoints_kernel = LimitMidpoints(Vt_brok_equispaced)
def __init__(self,
state,
V,
ibp="once",
equation_form="advective",
vector_manifold=False,
Vdg=None,
solver_params=None,
recovered_spaces=None,
outflow=False):
# give equation the property V0, the space that the function should live in
# in the absence of Vdg, this is used to set up the space for advection
# to take place in
self.V0 = V
self.recovered = False
if recovered_spaces is not None:
# Vdg must be None to use recovered spaces
if Vdg is not None:
raise ValueError(
'The recovered_spaces option is incompatible with the Vdg option'
)
else:
# check that the list or tuple of spaces is the right length
if len(recovered_spaces) != 3:
raise ValueError(
'recovered_spaces must be a list or tuple containing three spaces'
)
self.space = recovered_spaces[
0] # the space in which advection happens
self.V_rec = recovered_spaces[
1] # the recovered continuous space
self.V_brok = recovered_spaces[2] # broken version of V0
self.recovered = True
elif Vdg is None:
# Create broken space, functions and projector
V_elt = BrokenElement(V.ufl_element())
self.space = FunctionSpace(state.mesh, V_elt)
else:
self.space = Vdg
super().__init__(state=state,
V=self.space,
ibp=ibp,
equation_form=equation_form,
vector_manifold=vector_manifold,
solver_params=solver_params,
outflow=outflow)
def setup(self, state):
if not self._initialised:
space = state.fields("theta").function_space()
broken_space = FunctionSpace(state.mesh,
BrokenElement(space.ufl_element()))
h_deg = space.ufl_element().degree()[0]
v_deg = space.ufl_element().degree()[1] - 1
boundary_method = Boundary_Method.physics if (
v_deg == 0 and h_deg == 0) else None
super().setup(state, space=space)
# now let's attach all of our fields
self.u = state.fields("u")
self.rho = state.fields("rho")
self.theta = state.fields("theta")
self.rho_averaged = Function(space)
self.recoverer = Recoverer(self.rho,
self.rho_averaged,
VDG=broken_space,
boundary_method=boundary_method)
try:
self.r_v = state.fields("vapour_mixing_ratio")
except NotImplementedError:
self.r_v = Constant(0.0)
try:
self.r_c = state.fields("cloud_liquid_mixing_ratio")
except NotImplementedError:
self.r_c = Constant(0.0)
try:
self.rain = state.fields("rain_mixing_ratio")
except NotImplementedError:
self.rain = Constant(0.0)
# now let's store the most common expressions
self.exner = thermodynamics.exner_pressure(state.parameters,
self.rho_averaged,
self.theta)
self.T = thermodynamics.T(state.parameters,
self.theta,
self.exner,
r_v=self.r_v)
self.p = thermodynamics.p(state.parameters, self.exner)
self.r_l = self.r_c + self.rain
self.r_t = self.r_v + self.r_c + self.rain
def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup):
# Make mesh and state using routine from conftest
setup = tracer_setup(tmpdir, geometry, degree=0)
state = setup.state
mesh = state.mesh
gdim = state.mesh.geometric_dimension()
# Spaces for recovery
Vu = state.spaces("HDiv", family=setup.family, degree=setup.degree)
if geometry == "slice":
VDG1 = state.spaces("DG1_equispaced")
Vec_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element(), name='Vec_DG1')
Vec_CG1 = VectorFunctionSpace(mesh, "CG", 1, name='Vec_CG1')
Vu_brok = FunctionSpace(mesh, BrokenElement(Vu.ufl_element()))
rec_opts = RecoveredOptions(embedding_space=Vec_DG1,
recovered_space=Vec_CG1,
broken_space=Vu_brok,
boundary_method=Boundary_Method.dynamics)
else:
raise NotImplementedError(
f'Recovered spaces for geometry {geometry} have not been implemented'
)
# Make equation
eqn = AdvectionEquation(state, Vu, "f", ufamily=setup.family, udegree=1)
# Initialise fields
f_init = as_vector([setup.f_init] * gdim)
state.fields("f").project(f_init)
state.fields("u").project(setup.uexpr)
transport_scheme = [(eqn, SSPRK3(state, options=rec_opts))]
f_end = as_vector([setup.f_end] * gdim)
# Run and check error
error = run(state, transport_scheme, setup.tmax, f_end)
assert error < setup.tol, \
'The transport error is greater than the permitted tolerance'
gamma = TestFunction(Vr)
rho_trial = TrialFunction(Vr)
a = gamma * rho_trial * dxp
L = gamma * (rho_b * theta_b / theta0) * dxp
rho_problem = LinearVariationalProblem(a, L, rho0)
rho_solver = LinearVariationalSolver(rho_problem)
rho_solver.solve()
physics_boundary_method = Boundary_Method.physics if recovered else None
# find perturbed water_v
w_v = Function(Vt)
phi = TestFunction(Vt)
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0, rho_averaged,
VDG=FunctionSpace(mesh, BrokenElement(Vt.ufl_element())),
boundary_method=physics_boundary_method)
rho_recoverer.project()
pi = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, pi)
T = thermodynamics.T(state.parameters, theta0, pi, r_v=w_v)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
w_functional = (phi * w_v * dxp - phi * w_sat * dxp)
w_problem = NonlinearVariationalProblem(w_functional, w_v)
w_solver = NonlinearVariationalSolver(w_problem)
w_solver.solve()
water_v0.assign(w_v)
water_c0.assign(water_t - water_v0)
def unsaturated_hydrostatic_balance(state,
theta_d,
H,
exner0=None,
top=False,
exner_boundary=Constant(1.0),
max_outer_solve_count=40,
max_inner_solve_count=20):
"""
Given vertical profiles for dry potential temperature
and relative humidity compute hydrostatically balanced
virtual potential temperature, dry density and water vapour profiles.
The general strategy is to split up the solving into two steps:
1) finding rho to balance the theta profile
2) finding theta_v and r_v to get back theta_d and H
We iteratively solve these steps until we (hopefully)
converge to a solution.
:arg state: The :class:`State` object.
:arg theta_d: The initial dry potential temperature profile.
:arg H: The relative humidity profile.
:arg exner0: Optional function to put exner pressure into.
:arg top: If True, set a boundary condition at the top, otherwise
it will be at the bottom.
:arg exner_boundary: The value of exner on the specified boundary.
:arg max_outer_solve_count: Max number of iterations for outer loop of balance solver.
:arg max_inner_solve_count: Max number of iterations for inner loop of balanace solver.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
mr_v0 = state.fields('vapour_mixing_ratio')
# Calculate hydrostatic exner pressure
Vt = theta0.function_space()
Vr = rho0.function_space()
R_d = state.parameters.R_d
R_v = state.parameters.R_v
epsilon = R_d / R_v
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
logger.warning(
"default quadrature degree most likely not sufficient for this degree element"
)
# apply first guesses
theta0.assign(theta_d * 1.01)
mr_v0.assign(0.01)
v_deg = Vr.ufl_element().degree()[1]
if v_deg == 0:
method = Boundary_Method.physics
else:
method = None
rho_h = Function(Vr)
rho_averaged = Function(Vt)
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_recoverer = Recoverer(rho0,
rho_averaged,
VDG=Vt_broken,
boundary_method=method)
w_h = Function(Vt)
delta = 1.0
# make expressions for determining mr_v0
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged,
theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, mr_v0)
r_v_expr = thermodynamics.r_v(state.parameters, H, T, p)
# make expressions to evaluate residual
exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged,
theta0)
p_ev = thermodynamics.p(state.parameters, exner_ev)
T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, mr_v0)
RH_ev = thermodynamics.RH(state.parameters, mr_v0, T_ev, p_ev)
RH = Function(Vt)
for i in range(max_outer_solve_count):
# solve for rho with theta_vd and w_v guesses
compressible_hydrostatic_balance(state,
theta0,
rho_h,
top=top,
exner_boundary=exner_boundary,
mr_t=mr_v0,
solve_for_rho=True)
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
# calculate averaged rho
rho_recoverer.project()
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
# now solve for r_v
for j in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
mr_v0.assign(mr_v0 * (1 - delta) + delta * w_h)
# compute theta_vd
theta0.assign(theta_d * (1 + mr_v0 / epsilon))
# test quality of solution by re-evaluating expression
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
if i == max_outer_solve_count:
raise RuntimeError(
'Hydrostatic balance solve has not converged within %i' % i,
'iterations')
if exner0 is not None:
exner = thermodynamics.exner_pressure(state.parameters, rho0, theta0)
exner0.interpolate(exner)
# do one extra solve for rho
compressible_hydrostatic_balance(state,
theta0,
rho0,
top=top,
exner_boundary=exner_boundary,
mr_t=mr_v0,
solve_for_rho=True)
def saturated_hydrostatic_balance(state,
theta_e,
mr_t,
exner0=None,
top=False,
exner_boundary=Constant(1.0),
max_outer_solve_count=40,
max_theta_solve_count=5,
max_inner_solve_count=3):
"""
Given a wet equivalent potential temperature, theta_e, and the total moisture
content, mr_t, compute a hydrostatically balance virtual potential temperature,
dry density and water vapour profile.
The general strategy is to split up the solving into two steps:
1) finding rho to balance the theta profile
2) finding theta_v and r_v to get back theta_e and saturation
We iteratively solve these steps until we (hopefully)
converge to a solution.
:arg state: The :class:`State` object.
:arg theta_e: The initial wet equivalent potential temperature profile.
:arg mr_t: The total water pseudo-mixing ratio profile.
:arg exner0: Optional function to put exner pressure into.
:arg top: If True, set a boundary condition at the top, otherwise
it will be at the bottom.
:arg exner_boundary: The value of exner on the specified boundary.
:arg max_outer_solve_count: Max number of outer iterations for balance solver.
:arg max_theta_solve_count: Max number of iterations for theta solver (middle part of solve).
:arg max_inner_solve_count: Max number of iterations on the inner most
loop for the water vapour solver.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
mr_v0 = state.fields('vapour_mixing_ratio')
# Calculate hydrostatic exner pressure
Vt = theta0.function_space()
Vr = rho0.function_space()
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
logger.warning(
"default quadrature degree most likely not sufficient for this degree element"
)
theta0.interpolate(theta_e)
mr_v0.interpolate(mr_t)
v_deg = Vr.ufl_element().degree()[1]
if v_deg == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
rho_h = Function(Vr)
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
w_h = Function(Vt)
theta_h = Function(Vt)
theta_e_test = Function(Vt)
delta = 0.8
# expressions for finding theta0 and mr_v0 from theta_e and mr_t
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged,
theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, mr_v0)
r_v_expr = thermodynamics.r_sat(state.parameters, T, p)
theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, mr_v0, mr_t)
for i in range(max_outer_solve_count):
# solve for rho with theta_vd and w_v guesses
compressible_hydrostatic_balance(state,
theta0,
rho_h,
top=top,
exner_boundary=exner_boundary,
mr_t=mr_t,
solve_for_rho=True)
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
theta_e_test.assign(theta_e_expr)
if errornorm(theta_e_test, theta_e) < 1e-8:
break
# calculate averaged rho
rho_recoverer.project()
# now solve for r_v
for j in range(max_theta_solve_count):
theta_h.interpolate(theta_e / theta_e_expr * theta0)
theta0.assign(theta0 * (1 - delta) + delta * theta_h)
# break when close enough
if errornorm(theta_e_test, theta_e) < 1e-6:
break
for k in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
mr_v0.assign(mr_v0 * (1 - delta) + delta * w_h)
# break when close enough
theta_e_test.assign(theta_e_expr)
if errornorm(theta_e_test, theta_e) < 1e-6:
break
if i == max_outer_solve_count:
raise RuntimeError(
'Hydrostatic balance solve has not converged within %i' % i,
'iterations')
if exner0 is not None:
exner = thermodynamics.exner(state.parameters, rho0, theta0)
exner0.interpolate(exner)
# do one extra solve for rho
compressible_hydrostatic_balance(state,
theta0,
rho0,
top=top,
exner_boundary=exner_boundary,
mr_t=mr_t,
solve_for_rho=True)
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
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 setup_limiters(dirname):
dt = 0.01
Ld = 1.
tmax = 0.2
rotations = 0.1
m = PeriodicIntervalMesh(20, Ld)
mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld / 20))
output = OutputParameters(
dirname=dirname,
dumpfreq=1,
dumplist=['u', 'chemical', 'moisture_higher', 'moisture_lower'])
parameters = CompressibleParameters()
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
fieldlist = [
'u', 'rho', 'theta', 'chemical', 'moisture_higher', 'moisture_lower'
]
diagnostic_fields = []
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
x, z = SpatialCoordinate(mesh)
Vr = state.spaces("DG")
Vt = state.spaces("HDiv_v")
Vpsi = FunctionSpace(mesh, "CG", 2)
cell = mesh._base_mesh.ufl_cell().cellname()
DG0_element = FiniteElement("DG", cell, 0)
CG1_element = FiniteElement("CG", interval, 1)
Vt0_element = TensorProductElement(DG0_element, CG1_element)
Vt0 = FunctionSpace(mesh, Vt0_element)
Vt0_brok = FunctionSpace(mesh, BrokenElement(Vt0_element))
VCG1 = FunctionSpace(mesh, "CG", 1)
u = state.fields("u", dump=True)
chemical = state.fields("chemical", Vr, dump=True)
moisture_higher = state.fields("moisture_higher", Vt, dump=True)
moisture_lower = state.fields("moisture_lower", Vt0, dump=True)
x_lower = 2 * Ld / 5
x_upper = 3 * Ld / 5
z_lower = 6 * Ld / 10
z_upper = 8 * Ld / 10
bubble_expr_1 = conditional(
x > x_lower,
conditional(
x < x_upper,
conditional(z > z_lower, conditional(z < z_upper, 1.0, 0.0), 0.0),
0.0), 0.0)
bubble_expr_2 = conditional(
x > z_lower,
conditional(
x < z_upper,
conditional(z > x_lower, conditional(z < x_upper, 1.0, 0.0), 0.0),
0.0), 0.0)
chemical.assign(1.0)
moisture_higher.assign(280.)
chem_pert_1 = Function(Vr).interpolate(bubble_expr_1)
chem_pert_2 = Function(Vr).interpolate(bubble_expr_2)
moist_h_pert_1 = Function(Vt).interpolate(bubble_expr_1)
moist_h_pert_2 = Function(Vt).interpolate(bubble_expr_2)
moist_l_pert_1 = Function(Vt0).interpolate(bubble_expr_1)
moist_l_pert_2 = Function(Vt0).interpolate(bubble_expr_2)
chemical.assign(chemical + chem_pert_1 + chem_pert_2)
moisture_higher.assign(moisture_higher + moist_h_pert_1 + moist_h_pert_2)
moisture_lower.assign(moisture_lower + moist_l_pert_1 + moist_l_pert_2)
# set up solid body rotation for advection
# we do this slightly complicated stream function to make the velocity 0 at edges
# thus we avoid odd effects at boundaries
xc = Ld / 2
zc = Ld / 2
r = sqrt((x - xc)**2 + (z - zc)**2)
omega = rotations * 2 * pi / tmax
r_out = 9 * Ld / 20
r_in = 2 * Ld / 5
A = omega * r_in / (2 * (r_in - r_out))
B = -omega * r_in * r_out / (r_in - r_out)
C = omega * r_in**2 * r_out / (r_in - r_out) / 2
psi_expr = conditional(
r < r_in, omega * r**2 / 2,
conditional(r < r_out, A * r**2 + B * r + C,
A * r_out**2 + B * r_out + C))
psi = Function(Vpsi).interpolate(psi_expr)
gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)])
u.project(gradperp(psi))
state.initialise([('u', u), ('chemical', chemical),
('moisture_higher', moisture_higher),
('moisture_lower', moisture_lower)])
# set up advection schemes
dg_opts = EmbeddedDGOptions()
recovered_opts = RecoveredOptions(embedding_space=Vr,
recovered_space=VCG1,
broken_space=Vt0_brok,
boundary_method=Boundary_Method.dynamics)
chemeqn = AdvectionEquation(state, Vr, equation_form="advective")
moisteqn_higher = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=dg_opts)
moisteqn_lower = EmbeddedDGAdvection(state,
Vt0,
equation_form="advective",
options=recovered_opts)
# build advection dictionary
advected_fields = []
advected_fields.append(('chemical',
SSPRK3(state,
chemical,
chemeqn,
limiter=VertexBasedLimiter(Vr))))
advected_fields.append(('moisture_higher',
SSPRK3(state,
moisture_higher,
moisteqn_higher,
limiter=ThetaLimiter(Vt))))
advected_fields.append(('moisture_lower',
SSPRK3(state,
moisture_lower,
moisteqn_lower,
limiter=VertexBasedLimiter(Vr))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, tmax
def __init__(self, space):
"""
Initialise limiter
:param space: the space in which theta lies.
It should be the DG1xCG2 space.
"""
self.Vt = FunctionSpace(space.mesh(),
BrokenElement(space.ufl_element()))
# check this is the right space, currently working for 2D and 3D extruded meshes
# check that horizontal degree is 1 and vertical degree is 2
if self.Vt.ufl_element().degree()[0] != 1 or \
self.Vt.ufl_element().degree()[1] != 2:
raise ValueError('This is not the right limiter for this space.')
if not self.Vt.extruded:
raise ValueError('This is not the right limiter for this space.')
if self.Vt.mesh().topological_dimension() == 2:
# check that continuity of the spaces is correct
# this will fail if the space does not use broken elements
if self.Vt.ufl_element()._element.sobolev_space()[0].name != 'L2' or \
self.Vt.ufl_element()._element.sobolev_space()[1].name != 'H1':
raise ValueError(
'This is not the right limiter for this space.')
elif self.Vt.mesh().topological_dimension() == 3:
# check that continuity of the spaces is correct
# this will fail if the space does not use broken elements
if self.Vt.ufl_element()._element.sobolev_space()[0].name != 'L2' or \
self.Vt.ufl_element()._element.sobolev_space()[2].name != 'H1':
raise ValueError(
'This is not the right limiter for this space.')
else:
raise ValueError('This is not the right limiter for this space.')
self.DG1 = FunctionSpace(self.Vt.mesh(), 'DG',
1) # space with only vertex DOFs
self.vertex_limiter = VertexBasedLimiter(self.DG1)
self.theta_hat = Function(
self.DG1) # theta function with only vertex DOFs
self.theta_old = Function(self.Vt)
self.w = Function(self.Vt)
self.result = Function(self.Vt)
shapes = {
'nDOFs': self.Vt.finat_element.space_dimension(),
'nDOFs_base': int(self.Vt.finat_element.space_dimension() / 3)
}
averager_domain = "{{[i]: 0 <= i < {nDOFs}}}".format(**shapes)
theta_domain = "{{[i,j]: 0 <= i < {nDOFs_base} and 0 <= j < 2}}".format(
**shapes)
average_instructions = ("""
for i
vo[i] = vo[i] + v[i] / w[i]
end
""")
weight_instructions = ("""
for i
w[i] = w[i] + 1.0
end
""")
copy_into_DG1_instrs = ("""
for i
for j
theta_hat[i*2+j] = theta[i*3+j]
end
end
""")
copy_from_DG1_instrs = ("""
<float64> max_value = 0.0
<float64> min_value = 0.0
for i
for j
theta[i*3+j] = theta_hat[i*2+j]
end
max_value = fmax(theta_hat[i*2], theta_hat[i*2+1])
min_value = fmin(theta_hat[i*2], theta_hat[i*2+1])
if theta_old[i*3+2] > max_value
theta[i*3+2] = 0.5 * (theta_hat[i*2] + theta_hat[i*2+1])
elif theta_old[i*3+2] < min_value
theta[i*3+2] = 0.5 * (theta_hat[i*2] + theta_hat[i*2+1])
else
theta[i*3+2] = theta_old[i*3+2]
end
end
""")
self._average_kernel = (averager_domain, average_instructions)
_weight_kernel = (averager_domain, weight_instructions)
self._copy_into_DG1_kernel = (theta_domain, copy_into_DG1_instrs)
self._copy_from_DG1_kernel = (theta_domain, copy_from_DG1_instrs)
par_loop(_weight_kernel,
dx, {"w": (self.w, INC)},
is_loopy_kernel=True)
def _setup_solver(self):
import numpy as np
state = self.state
dt = state.dt
beta_ = dt * self.alpha
cp = state.parameters.cp
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("theta")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = Vrho.ufl_element().degree()[0]
v_deg = Vrho.ufl_element().degree()[1]
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
self.xrhs = Function(self.equations.function_space)
u_in, rho_in, theta_in = split(self.xrhs)[0:3]
# Build the function space for "broken" u, rho, and pressure trace
M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace))
w, phi, dl = TestFunctions(M)
u, rho, l0 = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar,
thetabar)
exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar,
thetabar)
exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar,
thetabar)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k, u) * dot(k, grad(thetabar)) * beta + theta_in
# Only include theta' (rather than exner') in the vertical
# component of the gradient
# The exner prime term (here, bars are for mean and no bars are
# for linear perturbations)
exner = exnerbar_theta * theta + exnerbar_rho * rho
# Vertical projection
def V(u):
return k * inner(u, k)
# hydrostatic projection
h_project = lambda u: u - k * inner(u, k)
# Specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
dS_hp = dS_h(degree=(self.quadrature_degree))
ds_vp = ds_v(degree=(self.quadrature_degree))
ds_tbp = (ds_t(degree=(self.quadrature_degree)) +
ds_b(degree=(self.quadrature_degree)))
# Add effect of density of water upon theta
if self.moisture is not None:
water_t = Function(Vtheta).assign(0.0)
for water in self.moisture:
water_t += self.state.fields(water)
theta_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
_l0 = TrialFunction(Vtrace)
_dl = TestFunction(Vtrace)
a_tr = _dl('+') * _l0('+') * (
dS_vp + dS_hp) + _dl * _l0 * ds_vp + _dl * _l0 * ds_tbp
def L_tr(f):
return _dl('+') * avg(f) * (
dS_vp + dS_hp) + _dl * f * ds_vp + _dl * f * ds_tbp
cg_ilu_parameters = {
'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
# Project field averages into functions on the trace space
rhobar_avg = Function(Vtrace)
exnerbar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar),
exnerbar_avg)
rho_avg_solver = LinearVariationalSolver(
rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
exner_avg_solver = LinearVariationalSolver(
exner_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='exnerbar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectExnerbar"):
exner_avg_solver.solve()
# "broken" u, rho, and trace system
# NOTE: no ds_v integrals since equations are defined on
# a periodic (or sphere) base mesh.
if any([t.has_label(hydrostatic) for t in self.equations.residual]):
u_mass = inner(w, (h_project(u) - u_in)) * dx
else:
u_mass = inner(w, (u - u_in)) * dx
eqn = (
# momentum equation
u_mass - beta_cp * div(theta_w * V(w)) * exnerbar * dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp
+ beta_cp * jump(theta_w * V(w), n=n) * exnerbar_avg('+') * dS_hp +
beta_cp * dot(theta_w * V(w), n) * exnerbar_avg * ds_tbp -
beta_cp * div(thetabar_w * w) * exner * dxp
# trace terms appearing after integrating momentum equation
+ beta_cp * jump(thetabar_w * w, n=n) * l0('+') * (dS_vp + dS_hp) +
beta_cp * dot(thetabar_w * w, n) * l0 * (ds_tbp + ds_vp)
# mass continuity equation
+ (phi *
(rho - rho_in) - beta * inner(grad(phi), u) * rhobar) * dx +
beta * jump(phi * u, n=n) * rhobar_avg('+') * (dS_v + dS_h)
# term added because u.n=0 is enforced weakly via the traces
+ beta * phi * dot(u, n) * rhobar_avg * (ds_tb + ds_v)
# constraint equation to enforce continuity of the velocity
# through the interior facets and weakly impose the no-slip
# condition
+ dl('+') * jump(u, n=n) * (dS_vp + dS_hp) + dl * dot(u, n) *
(ds_tbp + ds_vp))
# contribution of the sponge term
if hasattr(self.equations, "mu"):
eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Function for the hybridized solutions
self.urhol0 = Function(M)
hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
hybridized_solver = LinearVariationalSolver(
hybridized_prb,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
self.hybridized_solver = hybridized_solver
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = {
"i": Vu.finat_element.space_dimension(),
"j": np.prod(Vu.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self._weight = Function(Vu)
par_loop(weight_kernel, dx, {"w": (self._weight, INC)})
# Averaging kernel
self._average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# HDiv-conforming velocity
self.u_hdiv = Function(Vu)
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
self.theta = Function(Vtheta)
theta_eqn = gamma * (theta - theta_in + dot(k, self.u_hdiv) *
dot(k, grad(thetabar)) * beta) * dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn), self.theta)
self.theta_solver = LinearVariationalSolver(
theta_problem,
solver_parameters=cg_ilu_parameters,
options_prefix='thetabacksubstitution')
# Store boundary conditions for the div-conforming velocity to apply
# post-solve
self.bcs = self.equations.bcs['u']
def setup(self, equation, uadv=None, apply_bcs=True, *active_labels):
self.residual = equation.residual
if self.field_name is not None:
self.idx = equation.field_names.index(self.field_name)
self.fs = self.state.fields(self.field_name).function_space()
self.residual = self.residual.label_map(
lambda t: t.get(prognostic) == self.field_name,
lambda t: Term(
split_form(t.form)[self.idx].form,
t.labels),
drop)
bcs = equation.bcs[self.field_name]
else:
self.field_name = equation.field_name
self.fs = equation.function_space
self.idx = None
if type(self.fs.ufl_element()) is MixedElement:
bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs]
else:
bcs = equation.bcs[self.field_name]
if len(active_labels) > 0:
self.residual = self.residual.label_map(
lambda t: any(t.has_label(time_derivative, *active_labels)),
map_if_false=drop)
options = self.options
# -------------------------------------------------------------------- #
# Routines relating to transport
# -------------------------------------------------------------------- #
if hasattr(self.options, 'ibp'):
self.replace_transport_term()
self.replace_transporting_velocity(uadv)
# -------------------------------------------------------------------- #
# Wrappers for embedded / recovery methods
# -------------------------------------------------------------------- #
if self.discretisation_option in ["embedded_dg", "recovered"]:
# construct the embedding space if not specified
if options.embedding_space is None:
V_elt = BrokenElement(self.fs.ufl_element())
self.fs = FunctionSpace(self.state.mesh, V_elt)
else:
self.fs = options.embedding_space
self.xdg_in = Function(self.fs)
self.xdg_out = Function(self.fs)
if self.idx is None:
self.x_projected = Function(equation.function_space)
else:
self.x_projected = Function(self.state.fields(self.field_name).function_space())
new_test = TestFunction(self.fs)
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
# -------------------------------------------------------------------- #
# Make boundary conditions
# -------------------------------------------------------------------- #
if not apply_bcs:
self.bcs = None
elif self.discretisation_option in ["embedded_dg", "recovered"]:
# Transfer boundary conditions onto test function space
self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs]
else:
self.bcs = bcs
# -------------------------------------------------------------------- #
# Modify test function for SUPG methods
# -------------------------------------------------------------------- #
if self.discretisation_option == "supg":
# construct tau, if it is not specified
dim = self.state.mesh.topological_dimension()
if options.tau is not None:
# if tau is provided, check that is has the right size
tau = options.tau
assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!"
else:
# create tuple of default values of size dim
default_vals = [options.default*self.dt]*dim
# check for directions is which the space is discontinuous
# so that we don't apply supg in that direction
if is_cg(self.fs):
vals = default_vals
else:
space = self.fs.ufl_element().sobolev_space()
if space.name in ["HDiv", "DirectionalH"]:
vals = [default_vals[i] if space[i].name == "H1"
else 0. for i in range(dim)]
else:
raise ValueError("I don't know what to do with space %s" % space)
tau = Constant(tuple([
tuple(
[vals[j] if i == j else 0. for i, v in enumerate(vals)]
) for j in range(dim)])
)
self.solver_parameters = {'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
test = TestFunction(self.fs)
new_test = test + dot(dot(uadv, tau), grad(test))
if self.discretisation_option is not None:
# replace the original test function with one defined on
# the embedding space, as this is the space where the
# the problem will be solved
self.residual = self.residual.label_map(
all_terms,
map_if_true=replace_test_function(new_test))
if self.discretisation_option == "embedded_dg":
if self.limiter is None:
self.x_out_projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
else:
self.x_out_projector = Recoverer(self.xdg_out, self.x_projected)
if self.discretisation_option == "recovered":
# set up the necessary functions
self.x_in = Function(self.state.fields(self.field_name).function_space())
x_rec = Function(options.recovered_space)
x_brok = Function(options.broken_space)
# set up interpolators and projectors
self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function
self.x_brok_projector = Projector(x_rec, x_brok) # function projected back
self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in)
if self.limiter is not None:
self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok)
self.x_out_projector = Recoverer(x_brok, self.x_projected)
else:
self.x_out_projector = Projector(self.xdg_out, self.x_projected)
# setup required functions
self.dq = Function(self.fs)
self.q1 = Function(self.fs)
def initialize(self, pc):
""" Set up the problem context. Takes the original
mixed problem and transforms it into the equivalent
hybrid-mixed system.
A KSP object is created for the Lagrange multipliers
on the top/bottom faces of the mesh cells.
"""
from firedrake import (FunctionSpace, Function, Constant,
FiniteElement, TensorProductElement,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, interval, MixedElement,
BrokenElement)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
from ufl.cell import TensorProductCell
# Extract PC context
prefix = pc.getOptionsPrefix() + "vert_hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
# Magically determine which spaces are vector and scalar valued
for i, Vi in enumerate(V):
# Vector-valued spaces will have a non-empty value_shape
if Vi.ufl_element().value_shape():
self.vidx = i
else:
self.pidx = i
Vv = V[self.vidx]
Vp = V[self.pidx]
# Create the space of approximate traces in the vertical.
# NOTE: Technically a hack since the resulting space is technically
# defined in cell interiors, however the degrees of freedom will only
# be geometrically defined on edges. Arguments will only be used in
# surface integrals
deg, _ = Vv.ufl_element().degree()
# Assumes a tensor product cell (quads, triangular-prisms, cubes)
if not isinstance(Vp.ufl_element().cell(), TensorProductCell):
raise NotImplementedError(
"Currently only implemented for tensor product discretizations"
)
# Only want the horizontal cell
cell, _ = Vp.ufl_element().cell()._cells
DG = FiniteElement("DG", cell, deg)
CG = FiniteElement("CG", interval, 1)
Vv_tr_element = TensorProductElement(DG, CG)
Vv_tr = FunctionSpace(mesh, Vv_tr_element)
# Break the spaces
broken_elements = MixedElement(
[BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up relevant functions
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(Vv_tr)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up transfer kernels to and from the broken velocity space
# NOTE: Since this snippet of code is used in a couple places in
# in Gusto, might be worth creating a utility function that is
# is importable and just called where needed.
shapes = {
"i": Vv.finat_element.space_dimension(),
"j": np.prod(Vv.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self.weight = Function(Vv)
par_loop(weight_kernel, dx, {"w": (self.weight, INC)})
# Averaging kernel
self.average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# Original mixed operator replaced with "broken" arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(Vv_tr)
n = FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# Again, assumes tensor product structure. Why use this if you
# don't have some form of vertical extrusion?
Kform = gammar('+') * jump(sigma, n=n) * dS_h
# Here we deal with boundary conditions
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError(
"Dirichlet conditions for scalar variable not supported. Use a weak bc."
)
if bc.function_space().index != self.vidx:
raise NotImplementedError(
"Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * dot(sigma, n)
measures = []
trace_subdomains = []
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ds_t, "bottom": ds_b}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand * measure
else:
trace_subdomains = ["top", "bottom"]
trace_bcs = [
DirichletBC(Vv_tr, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(Vv_tr)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("vert_trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vv_tr)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
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 __init__(self,
v_CG1,
v_DG1,
method=Boundary_Method.physics,
coords_to_adjust=None):
self.v_DG1 = v_DG1
self.v_CG1 = v_CG1
self.v_DG1_old = Function(v_DG1.function_space())
self.coords_to_adjust = coords_to_adjust
self.method = method
mesh = v_CG1.function_space().mesh()
VDG0 = FunctionSpace(mesh, "DG", 0)
VCG1 = FunctionSpace(mesh, "CG", 1)
VDG1 = FunctionSpace(mesh, "DG", 1)
self.num_ext = Function(VDG0)
# check function spaces of functions
if self.method == Boundary_Method.dynamics:
if v_CG1.function_space() != VCG1:
raise NotImplementedError(
"This boundary recovery method requires v1 to be in CG1.")
if v_DG1.function_space() != VDG1:
raise NotImplementedError(
"This boundary recovery method requires v_out to be in DG1."
)
# check whether mesh is valid
if mesh.topological_dimension() == 2:
# if mesh is extruded then we're fine, but if not needs to be quads
if not VDG0.extruded and mesh.ufl_cell().cellname(
) != 'quadrilateral':
raise NotImplementedError(
'For 2D meshes this recovery method requires that elements are quadrilaterals'
)
elif mesh.topological_dimension() == 3:
# assume that 3D mesh is extruded
if mesh._base_mesh.ufl_cell().cellname() != 'quadrilateral':
raise NotImplementedError(
'For 3D extruded meshes this recovery method requires a base mesh with quadrilateral elements'
)
elif mesh.topological_dimension() != 1:
raise NotImplementedError(
'This boundary recovery is implemented only on certain classes of mesh.'
)
if coords_to_adjust is None:
raise ValueError(
'Need coords_to_adjust field for dynamics boundary methods'
)
elif self.method == Boundary_Method.physics:
# check that mesh is valid -- must be an extruded mesh
if not VDG0.extruded:
raise NotImplementedError(
'The physics boundary method only works on extruded meshes'
)
# base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
w_hori = FiniteElement("DG", cell, 0)
w_vert = FiniteElement("CG", interval, 1)
# build element
theta_element = TensorProductElement(w_hori, w_vert)
# spaces
Vtheta = FunctionSpace(mesh, theta_element)
Vtheta_broken = FunctionSpace(mesh, BrokenElement(theta_element))
if v_CG1.function_space() != Vtheta:
raise ValueError(
"This boundary recovery method requires v_CG1 to be in DG0xCG1 TensorProductSpace."
)
if v_DG1.function_space() != Vtheta_broken:
raise ValueError(
"This boundary recovery method requires v_DG1 to be in the broken DG0xCG1 TensorProductSpace."
)
else:
raise ValueError(
"Boundary method should be a Boundary Method Enum object.")
VuDG1 = VectorFunctionSpace(VDG0.mesh(), "DG", 1)
x = SpatialCoordinate(VDG0.mesh())
self.interpolator = Interpolator(self.v_CG1, self.v_DG1)
if self.method == Boundary_Method.dynamics:
# STRATEGY
# obtain a coordinate field for all the nodes
VuDG1 = VectorFunctionSpace(mesh, "DG", 1)
self.act_coords = Function(VuDG1).project(x) # actual coordinates
self.eff_coords = Function(VuDG1).project(
x) # effective coordinates
shapes = {
"nDOFs":
self.v_DG1.function_space().finat_element.space_dimension(),
"dim":
np.prod(VuDG1.shape, dtype=int)
}
num_ext_domain = ("{{[i]: 0 <= i < {nDOFs}}}").format(**shapes)
num_ext_instructions = ("""
<float64> SUM_EXT = 0
for i
SUM_EXT = SUM_EXT + EXT_V1[i]
end
NUM_EXT[0] = SUM_EXT
""")
coords_domain = ("{{[i, j, k, ii, jj, kk, ll, mm, iii, kkk]: "
"0 <= i < {nDOFs} and "
"0 <= j < {nDOFs} and 0 <= k < {dim} and "
"0 <= ii < {nDOFs} and 0 <= jj < {nDOFs} and "
"0 <= kk < {dim} and 0 <= ll < {dim} and "
"0 <= mm < {dim} and 0 <= iii < {nDOFs} and "
"0 <= kkk < {dim}}}").format(**shapes)
coords_insts = (
"""
<float64> sum_V1_ext = 0
<int> index = 100
<float64> dist = 0.0
<float64> max_dist = 0.0
<float64> min_dist = 0.0
"""
# only do adjustment in cells with at least one DOF to adjust
"""
if NUM_EXT[0] > 0
"""
# find the maximum distance between DOFs in this cell, to serve as starting point for finding min distances
"""
for i
for j
dist = 0.0
for k
dist = dist + pow(ACT_COORDS[i,k] - ACT_COORDS[j,k], 2.0)
end
dist = pow(dist, 0.5) {{id=sqrt_max_dist, dep=*}}
max_dist = fmax(dist, max_dist) {{id=max_dist, dep=sqrt_max_dist}}
end
end
"""
# loop through cells and find which ones to adjust
"""
for ii
if EXT_V1[ii] > 0.5
"""
# find closest interior node
"""
min_dist = max_dist
index = 100
for jj
if EXT_V1[jj] < 0.5
dist = 0.0
for kk
dist = dist + pow(ACT_COORDS[ii,kk] - ACT_COORDS[jj,kk], 2)
end
dist = pow(dist, 0.5)
if dist <= min_dist
index = jj
end
min_dist = fmin(min_dist, dist)
for ll
EFF_COORDS[ii,ll] = 0.5 * (ACT_COORDS[ii,ll] + ACT_COORDS[index,ll])
end
end
end
else
"""
# for DOFs that aren't exterior, use the original coordinates
"""
for mm
EFF_COORDS[ii, mm] = ACT_COORDS[ii, mm]
end
end
end
else
"""
# for interior elements, just use the original coordinates
"""
for iii
for kkk
EFF_COORDS[iii, kkk] = ACT_COORDS[iii, kkk]
end
end
end
""").format(**shapes)
elimin_domain = (
"{{[i, ii_loop, jj_loop, kk, ll_loop, mm, iii_loop, kkk_loop, iiii, iiiii]: "
"0 <= i < {nDOFs} and 0 <= ii_loop < {nDOFs} and "
"ii_loop + 1 <= jj_loop < {nDOFs} and ii_loop <= kk < {nDOFs} and "
"ii_loop + 1 <= ll_loop < {nDOFs} and ii_loop <= mm < {nDOFs} + 1 and "
"0 <= iii_loop < {nDOFs} and {nDOFs} - iii_loop <= kkk_loop < {nDOFs} + 1 and "
"0 <= iiii < {nDOFs} and 0 <= iiiii < {nDOFs}}}").format(
**shapes)
elimin_insts = (
"""
<int> ii = 0
<int> jj = 0
<int> ll = 0
<int> iii = 0
<int> jjj = 0
<int> i_max = 0
<float64> A_max = 0.0
<float64> temp_f = 0.0
<float64> temp_A = 0.0
<float64> c = 0.0
<float64> f[{nDOFs}] = 0.0
<float64> a[{nDOFs}] = 0.0
<float64> A[{nDOFs},{nDOFs}] = 0.0
"""
# We are aiming to find the vector a that solves A*a = f, for matrix A and vector f.
# This is done by performing row operations (swapping and scaling) to obtain A in upper diagonal form.
# N.B. several for loops must be executed in numerical order (loopy does not necessarily do this).
# For these loops we must manually iterate the index.
"""
if NUM_EXT[0] > 0.0
"""
# only do Gaussian elimination for elements with effective coordinates
"""
for i
"""
# fill f with the original field values and A with the effective coordinate values
"""
f[i] = DG1_OLD[i]
A[i,0] = 1.0
A[i,1] = EFF_COORDS[i,0]
if {nDOFs} > 3
A[i,2] = EFF_COORDS[i,1]
A[i,3] = EFF_COORDS[i,0]*EFF_COORDS[i,1]
if {nDOFs} > 7
A[i,4] = EFF_COORDS[i,{dim}-1]
A[i,5] = EFF_COORDS[i,0]*EFF_COORDS[i,{dim}-1]
A[i,6] = EFF_COORDS[i,1]*EFF_COORDS[i,{dim}-1]
A[i,7] = EFF_COORDS[i,0]*EFF_COORDS[i,1]*EFF_COORDS[i,{dim}-1]
end
end
end
"""
# now loop through rows/columns of A
"""
for ii_loop
A_max = fabs(A[ii,ii])
i_max = ii
"""
# loop to find the largest value in the ith column
# set i_max as the index of the row with this largest value.
"""
jj = ii + 1
for jj_loop
if fabs(A[jj,ii]) > A_max
i_max = jj
end
A_max = fmax(A_max, fabs(A[jj,ii]))
jj = jj + 1
end
"""
# if the max value in the ith column isn't in the ith row, we must swap the rows
"""
if i_max != ii
"""
# swap the elements of f
"""
temp_f = f[ii] {{id=set_temp_f, dep=*}}
f[ii] = f[i_max] {{id=set_f_imax, dep=set_temp_f}}
f[i_max] = temp_f {{id=set_f_ii, dep=set_f_imax}}
"""
# swap the elements of A
# N.B. kk runs from ii to (nDOFs-1) as elements below diagonal should be 0
"""
for kk
temp_A = A[ii,kk] {{id=set_temp_A, dep=*}}
A[ii, kk] = A[i_max, kk] {{id=set_A_ii, dep=set_temp_A}}
A[i_max, kk] = temp_A {{id=set_A_imax, dep=set_A_ii}}
end
end
"""
# scale the rows below the ith row
"""
ll = ii + 1
for ll_loop
if ll > ii
"""
# find scaling factor
"""
c = - A[ll,ii] / A[ii,ii]
"""
# N.B. mm runs from ii to (nDOFs-1) as elements below diagonal should be 0
"""
for mm
A[ll, mm] = A[ll, mm] + c * A[ii,mm]
end
f[ll] = f[ll] + c * f[ii]
end
ll = ll + 1
end
ii = ii + 1
end
"""
# do back substitution of upper diagonal A to obtain a
"""
iii = 0
for iii_loop
"""
# jjj starts at the bottom row and works upwards
"""
jjj = {nDOFs} - iii - 1 {{id=assign_jjj, dep=*}}
a[jjj] = f[jjj] {{id=set_a, dep=assign_jjj}}
for kkk_loop
a[jjj] = a[jjj] - A[jjj,kkk_loop] * a[kkk_loop]
end
a[jjj] = a[jjj] / A[jjj,jjj]
iii = iii + 1
end
"""
# Having found a, this gives us the coefficients for the Taylor expansion with the actual coordinates.
"""
for iiii
if {nDOFs} == 2
DG1[iiii] = a[0] + a[1]*ACT_COORDS[iiii,0]
elif {nDOFs} == 4
DG1[iiii] = a[0] + a[1]*ACT_COORDS[iiii,0] + a[2]*ACT_COORDS[iiii,1] + a[3]*ACT_COORDS[iiii,0]*ACT_COORDS[iiii,1]
elif {nDOFs} == 8
DG1[iiii] = a[0] + a[1]*ACT_COORDS[iiii,0] + a[2]*ACT_COORDS[iiii,1] + a[3]*ACT_COORDS[iiii,0]*ACT_COORDS[iiii,1] + a[4]*ACT_COORDS[iiii,{dim}-1] + a[5]*ACT_COORDS[iiii,0]*ACT_COORDS[iiii,{dim}-1] + a[6]*ACT_COORDS[iiii,1]*ACT_COORDS[iiii,{dim}-1] + a[7]*ACT_COORDS[iiii,0]*ACT_COORDS[iiii,1]*ACT_COORDS[iiii,{dim}-1]
end
end
"""
# if element is not external, just use old field values.
"""
else
for iiiii
DG1[iiiii] = DG1_OLD[iiiii]
end
end
""").format(**shapes)
_num_ext_kernel = (num_ext_domain, num_ext_instructions)
_eff_coords_kernel = (coords_domain, coords_insts)
self._gaussian_elimination_kernel = (elimin_domain, elimin_insts)
# find number of external DOFs per cell
par_loop(_num_ext_kernel,
dx, {
"NUM_EXT": (self.num_ext, WRITE),
"EXT_V1": (self.coords_to_adjust, READ)
},
is_loopy_kernel=True)
# find effective coordinates
logger.warning(
'Finding effective coordinates for boundary recovery. This could give unexpected results for deformed meshes over very steep topography.'
)
par_loop(_eff_coords_kernel,
dx, {
"EFF_COORDS": (self.eff_coords, WRITE),
"ACT_COORDS": (self.act_coords, READ),
"NUM_EXT": (self.num_ext, READ),
"EXT_V1": (self.coords_to_adjust, READ)
},
is_loopy_kernel=True)
elif self.method == Boundary_Method.physics:
top_bottom_domain = ("{[i]: 0 <= i < 1}")
bottom_instructions = ("""
DG1[0] = 2 * CG1[0] - CG1[1]
DG1[1] = CG1[1]
""")
top_instructions = ("""
DG1[0] = CG1[0]
DG1[1] = -CG1[0] + 2 * CG1[1]
""")
self._bottom_kernel = (top_bottom_domain, bottom_instructions)
self._top_kernel = (top_bottom_domain, top_instructions)
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 test_3D_cartesian_recovery(geometry, element, mesh, expr):
family = "RTCF" if element == "quadrilateral" else "BDM"
# horizontal base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
u_hori = FiniteElement(family, 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
# DG1
DG1_hori = FiniteElement("DG", cell, 1, variant="equispaced")
DG1_vert = FiniteElement("DG", interval, 1, variant="equispaced")
DG1_elt = TensorProductElement(DG1_hori, DG1_vert)
DG1 = FunctionSpace(mesh, DG1_elt)
vec_DG1 = VectorFunctionSpace(mesh, DG1_elt)
# spaces
DG0 = FunctionSpace(mesh, "DG", 0)
CG1 = FunctionSpace(mesh, "CG", 1)
Vt = FunctionSpace(mesh, theta_element)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element))
Vu = FunctionSpace(mesh, v_element)
vec_CG1 = VectorFunctionSpace(mesh, "CG", 1)
# our actual theta and rho and v
rho_CG1_true = Function(CG1).interpolate(expr)
theta_CG1_true = Function(CG1).interpolate(expr)
v_CG1_true = Function(vec_CG1).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(DG0).interpolate(expr)
rho_CG1 = Function(CG1)
theta_Vt = Function(Vt).interpolate(expr)
theta_CG1 = Function(CG1)
v_Vu = Function(Vu).project(as_vector([expr, expr, expr]))
v_CG1 = Function(vec_CG1)
rho_Vt = Function(Vt)
# make the recoverers and do the recovery
rho_recoverer = Recoverer(rho_DG0, rho_CG1, VDG=DG1, boundary_method=Boundary_Method.dynamics)
theta_recoverer = Recoverer(theta_Vt, theta_CG1, VDG=DG1, boundary_method=Boundary_Method.dynamics)
v_recoverer = Recoverer(v_Vu, v_CG1, VDG=vec_DG1, 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)
tolerance = 1e-7
error_message = ("""
Incorrect recovery for {variable} with {boundary} boundary method
on {geometry} 3D Cartesian domain with {element} elements
""")
assert rho_diff < tolerance, error_message.format(variable='rho', boundary='dynamics',
geometry=geometry, element=element)
assert v_diff < tolerance, error_message.format(variable='v', boundary='dynamics',
geometry=geometry, element=element)
assert theta_diff < tolerance, error_message.format(variable='rho', boundary='dynamics',
geometry=geometry, element=element)
assert rho_Vt_diff < tolerance, error_message.format(variable='rho', boundary='physics',
geometry=geometry, element=element)
a = gamma * rho_trial * dxp
L = gamma * (rho_b * theta_b / theta0) * dxp
rho_problem = LinearVariationalProblem(a, L, rho0)
rho_solver = LinearVariationalSolver(rho_problem)
rho_solver.solve()
physics_boundary_method = None
# find perturbed water_v
w_v = Function(Vt)
phi = TestFunction(Vt)
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0,
rho_averaged,
VDG=FunctionSpace(mesh,
BrokenElement(Vt.ufl_element())),
boundary_method=physics_boundary_method)
rho_recoverer.project()
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, r_v=w_v)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
w_functional = (phi * w_v * dxp - phi * w_sat * dxp)
w_problem = NonlinearVariationalProblem(w_functional, w_v)
w_solver = NonlinearVariationalSolver(w_problem)
w_solver.solve()
water_v0.assign(w_v)
water_c0.assign(water_t - water_v0)
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
def __init__(self, state, moments=AdvectedMoments.M3, limit=True):
super().__init__(state)
# function spaces
Vt = state.fields('rain').function_space()
Vu = state.fields('u').function_space()
# declare properties of class
self.state = state
self.moments = moments
self.rain = state.fields('rain')
self.v = state.fields('rainfall_velocity', Vu)
self.limit = limit
if moments == AdvectedMoments.M0:
# all rain falls at terminal velocity
terminal_velocity = Constant(5) # in m/s
if state.mesh.geometric_dimension() == 2:
self.v.project(as_vector([0, -terminal_velocity]))
elif state.mesh.geometric_dimension() == 3:
self.v.project(as_vector([0, 0, -terminal_velocity]))
elif moments == AdvectedMoments.M3:
# this advects the third moment M3 of the raindrop
# distribution, which corresponds to the mean mass
rho = state.fields('rho')
rho_w = Constant(1000.0) # density of liquid water
# assume n(D) = n_0 * D^mu * exp(-Lambda*D)
# n_0 = N_r * Lambda^(1+mu) / gamma(1 + mu)
N_r = Constant(10**5) # number of rain droplets per m^3
mu = 0.0 # shape constant of droplet gamma distribution
# assume V(D) = a * D^b * exp(-f*D) * (rho_0 / rho)^g
# take f = 0
a = Constant(362.) # intercept for velocity distr. in log space
b = 0.65 # inverse scale parameter for velocity distr.
rho0 = Constant(1.22) # reference density in kg/m^3
g = Constant(0.5) # scaling of density correction
# we keep mu in the expressions even though mu = 0
threshold = Constant(10**-10) # only do rainfall for r > threshold
Lambda = (N_r * pi * rho_w * gamma(4 + mu) /
(6 * gamma(1 + mu) * rho * self.rain))**(1. / 3)
Lambda0 = (N_r * pi * rho_w * gamma(4 + mu) /
(6 * gamma(1 + mu) * rho * threshold))**(1. / 3)
v_expression = conditional(
self.rain > threshold,
(a * gamma(4 + b + mu) / (gamma(4 + mu) * Lambda**b) *
(rho0 / rho)**g),
(a * gamma(4 + b + mu) / (gamma(4 + mu) * Lambda0**b) *
(rho0 / rho)**g))
else:
raise NotImplementedError(
'Currently we only have implementations for zero and one moment schemes for rainfall. Valid options are AdvectedMoments.M0 and AdvectedMoments.M3'
)
if moments != AdvectedMoments.M0:
if state.mesh.geometric_dimension() == 2:
self.determine_v = Projector(as_vector([0, -v_expression]),
self.v)
elif state.mesh.geometric_dimension() == 3:
self.determine_v = Projector(as_vector([0, 0, -v_expression]),
self.v)
# determine whether to do recovered space advection scheme
# if horizontal and vertical degrees are 0 do recovered space
if state.horizontal_degree == 0 and state.vertical_degree == 0:
VDG1 = FunctionSpace(Vt.mesh(), "DG", 1)
VCG1 = FunctionSpace(Vt.mesh(), "CG", 1)
Vbrok = FunctionSpace(Vt.mesh(), BrokenElement(Vt.ufl_element()))
boundary_method = Boundary_Method.dynamics
advect_options = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vbrok,
boundary_method=boundary_method)
else:
advect_options = EmbeddedDGOptions()
# need to define advection equation before limiter (as it is needed for the ThetaLimiter)
advection_equation = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
outflow=True,
options=advect_options)
# decide which limiter to use
if self.limit:
if state.horizontal_degree == 0 and state.vertical_degree == 0:
limiter = VertexBasedLimiter(VDG1)
elif state.horizontal_degree == 1 and state.vertical_degree == 1:
limiter = ThetaLimiter(Vt)
else:
logger.warning(
"There is no limiter yet implemented for the spaces used. NoLimiter() is being used for the rainfall in this case."
)
limiter = NoLimiter()
else:
limiter = None
# sedimentation will happen using a full advection method
self.advection_method = SSPRK3(state,
self.rain,
advection_equation,
limiter=limiter)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from ufl.algorithms.map_integrands import map_integrand_dags
from firedrake import (FunctionSpace, TrialFunction, TrialFunctions,
TestFunction, Function, BrokenElement,
MixedElement, FacetNormal, Constant,
DirichletBC, Projector)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import ArgumentReplacer, split_form
# Extract the problem context
prefix = pc.getOptionsPrefix()
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
V = test.function_space()
if V.mesh().cell_set._extruded:
# TODO: Merge FIAT branch to support TPC trace elements
raise NotImplementedError("Not implemented on extruded meshes.")
# Break the function spaces and define fully discontinuous spaces
broken_elements = [BrokenElement(Vi.ufl_element()) for Vi in V]
elem = MixedElement(broken_elements)
V_d = FunctionSpace(V.mesh(), elem)
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
# Replace the problems arguments with arguments defined
# on the new discontinuous spaces
replacer = ArgumentReplacer(arg_map)
new_form = map_integrand_dags(replacer, context.a)
# Create the space of approximate traces.
# The vector function space will have a non-empty value_shape
W = next(v for v in V if bool(v.ufl_element().value_shape()))
if W.ufl_element().family() in ["Raviart-Thomas", "RTCF"]:
tdegree = W.ufl_element().degree() - 1
else:
tdegree = W.ufl_element().degree()
# NOTE: Once extruded is ready, we will need to be aware of this
# and construct the appropriate trace space for the HDiv element
TraceSpace = FunctionSpace(V.mesh(), "HDiv Trace", tdegree)
# NOTE: For extruded, we will need to add "on_top" and "on_bottom"
trace_conditions = [
DirichletBC(TraceSpace, Constant(0.0), "on_boundary")
]
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_rhs = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_rhs = Function(V)
# Create the symbolic Schur-reduction
Atilde = Tensor(new_form)
gammar = TestFunction(TraceSpace)
n = FacetNormal(V.mesh())
# Vector trial function will have a non-empty ufl_shape
sigma = next(f for f in TrialFunctions(V_d) if bool(f.ufl_shape))
# NOTE: Once extruded is ready, this will change slightly
# to include both horizontal and vertical interior facets
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_rhs,
tensor=self.schur_rhs,
form_compiler_parameters=context.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullsp = P.getNullSpace()
if nullsp.handle != 0:
new_vecs = get_trace_nullspace_vecs(K * Atilde.inv, nullsp, V, V_d,
TraceSpace)
tr_nullsp = PETSc.NullSpace().create(vectors=new_vecs,
comm=pc.comm)
Smat.setNullSpace(tr_nullsp)
# Set up the KSP for the system of Lagrange multipliers
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.setOptionsPrefix(prefix + "trace_")
ksp.setTolerances(rtol=1e-13)
ksp.setOperators(Smat)
ksp.setUp()
ksp.setFromOptions()
self.ksp = ksp
# Now we construct the local tensors for the reconstruction stage
# TODO: Add support for mixed tensors and these variables
# become unnecessary
split_forms = split_form(new_form)
A = Tensor(next(sf.form for sf in split_forms if sf.indices == (0, 0)))
B = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 0)))
C = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 1)))
trial = TrialFunction(
FunctionSpace(V.mesh(), BrokenElement(W.ufl_element())))
K_local = Tensor(gammar('+') * ufl.dot(trial, n) * ufl.dS)
# Split functions and reconstruct each bit separately
sigma_h, u_h = self.broken_solution.split()
g, f = self.broken_rhs.split()
# Pressure reconstruction
M = B * A.inv * B.T + C
u_sol = M.inv * f + M.inv * (
B * A.inv * K_local.T * self.trace_solution - B * A.inv * g)
self._assemble_pressure = create_assembly_callable(
u_sol, tensor=u_h, form_compiler_parameters=context.fc_params)
# Velocity reconstruction
sigma_sol = A.inv * g + A.inv * (B.T * u_h -
K_local.T * self.trace_solution)
self._assemble_velocity = create_assembly_callable(
sigma_sol,
tensor=sigma_h,
form_compiler_parameters=context.fc_params)
# Set up the projector for projecting the broken solution
# into the unbroken finite element spaces
# NOTE: Tolerance here matters!
sigma_b, _ = self.broken_solution.split()
sigma_u, _ = self.unbroken_solution.split()
self.projector = Projector(sigma_b,
sigma_u,
solver_parameters={
"ksp_type": "cg",
"ksp_rtol": 1e-13
})
def __init__(self, state):
super().__init__(state)
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.rain = state.fields('rain')
rho = state.fields('rho')
try:
water_c = state.fields('water_c')
water_l = self.rain + water_c
except NotImplementedError:
water_l = self.rain
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
if state.vertical_degree == 0 and state.horizontal_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
# define some parameters as attributes
dt = state.timestepping.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters,
self.theta,
Pi,
r_v=self.water_v)
p = thermodynamics.p(state.parameters, Pi)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# expression for ventilation factor
a = Constant(1.6)
b = Constant(124.9)
c = Constant(0.2046)
C = a + b * (rho_averaged * self.rain)**c
# make appropriate condensation rate
f = Constant(5.4e5)
g = Constant(2.55e6)
h = Constant(0.525)
dot_r_evap = (((1 - self.water_v / w_sat) * C *
(rho_averaged * self.rain)**h) /
(rho_averaged * (f + g / (p * w_sat))))
# make evap_rate function, needs to be the same for all updates in one time step
evap_rate = Function(Vt)
# adjust evap rate so negative rain doesn't occur
self.lim_evap_rate = Interpolator(
conditional(
dot_r_evap < 0, 0.0,
conditional(self.rain < 0.0, 0.0,
min_value(dot_r_evap, self.rain / dt))), evap_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt)
self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 - dt * evap_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
# find perturbed rho
gamma = TestFunction(Vr)
rho_trial = TrialFunction(Vr)
lhs = gamma * rho_trial * dx
rhs = gamma * (rho_b * theta_b / theta0) * dx
rho_problem = LinearVariationalProblem(lhs, rhs, rho0)
rho_solver = LinearVariationalSolver(rho_problem)
rho_solver.solve()
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up transport schemes
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)
Vu_brok = FunctionSpace(mesh, BrokenElement(Vu.ufl_element()))
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu_brok,
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)
def __init__(self, state, iterations=1):
super().__init__(state)
self.iterations = iterations
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.water_c = state.fields('water_c')
rho = state.fields('rho')
try:
rain = state.fields('rain')
water_l = self.water_c + rain
except NotImplementedError:
water_l = self.water_c
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
if state.vertical_degree == 0 and state.horizontal_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
# define some parameters as attributes
dt = state.timestepping.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters,
self.theta,
Pi,
r_v=self.water_v)
p = thermodynamics.p(state.parameters, Pi)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# make appropriate condensation rate
dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 +
((L_v**2.0 * w_sat) /
(cp * R_v * T**2.0)))))
# make cond_rate function, that needs to be the same for all updates in one time step
cond_rate = Function(Vt)
# adjust cond rate so negative concentrations don't occur
self.lim_cond_rate = Interpolator(
conditional(dot_r_cond < 0,
max_value(dot_r_cond, -self.water_c / dt),
min_value(dot_r_cond, self.water_v / dt)), cond_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt)
self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 + dt * cond_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
def __init__(self,
v_CG1,
v_DG1,
method=Boundary_Method.physics,
coords_to_adjust=None):
self.v_DG1 = v_DG1
self.v_CG1 = v_CG1
self.v_DG1_old = Function(v_DG1.function_space())
self.coords_to_adjust = coords_to_adjust
self.method = method
mesh = v_CG1.function_space().mesh()
VDG0 = FunctionSpace(mesh, "DG", 0)
VCG1 = FunctionSpace(mesh, "CG", 1)
if VDG0.extruded:
cell = mesh._base_mesh.ufl_cell().cellname()
DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced")
DG1_vert_elt = FiniteElement("DG",
interval,
1,
variant="equispaced")
DG1_element = TensorProductElement(DG1_hori_elt, DG1_vert_elt)
else:
cell = mesh.ufl_cell().cellname()
DG1_element = FiniteElement("DG", cell, 1, variant="equispaced")
VDG1 = FunctionSpace(mesh, DG1_element)
self.num_ext = Function(VDG0)
# check function spaces of functions
if self.method == Boundary_Method.dynamics:
if v_CG1.function_space() != VCG1:
raise NotImplementedError(
"This boundary recovery method requires v1 to be in CG1.")
if v_DG1.function_space() != VDG1:
raise NotImplementedError(
"This boundary recovery method requires v_out to be in DG1."
)
# check whether mesh is valid
if mesh.topological_dimension() == 2:
# if mesh is extruded then we're fine, but if not needs to be quads
if not VDG0.extruded and mesh.ufl_cell().cellname(
) != 'quadrilateral':
raise NotImplementedError(
'For 2D meshes this recovery method requires that elements are quadrilaterals'
)
elif mesh.topological_dimension() == 3:
# assume that 3D mesh is extruded
if mesh._base_mesh.ufl_cell().cellname() != 'quadrilateral':
raise NotImplementedError(
'For 3D extruded meshes this recovery method requires a base mesh with quadrilateral elements'
)
elif mesh.topological_dimension() != 1:
raise NotImplementedError(
'This boundary recovery is implemented only on certain classes of mesh.'
)
if coords_to_adjust is None:
raise ValueError(
'Need coords_to_adjust field for dynamics boundary methods'
)
elif self.method == Boundary_Method.physics:
# check that mesh is valid -- must be an extruded mesh
if not VDG0.extruded:
raise NotImplementedError(
'The physics boundary method only works on extruded meshes'
)
# base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
w_hori = FiniteElement("DG", cell, 0, variant="equispaced")
w_vert = FiniteElement("CG", interval, 1, variant="equispaced")
# build element
theta_element = TensorProductElement(w_hori, w_vert)
# spaces
Vtheta = FunctionSpace(mesh, theta_element)
Vtheta_broken = FunctionSpace(mesh, BrokenElement(theta_element))
if v_CG1.function_space() != Vtheta:
raise ValueError(
"This boundary recovery method requires v_CG1 to be in DG0xCG1 TensorProductSpace."
)
if v_DG1.function_space() != Vtheta_broken:
raise ValueError(
"This boundary recovery method requires v_DG1 to be in the broken DG0xCG1 TensorProductSpace."
)
else:
raise ValueError(
"Boundary method should be a Boundary Method Enum object.")
VuDG1 = VectorFunctionSpace(VDG0.mesh(), DG1_element)
x = SpatialCoordinate(VDG0.mesh())
self.interpolator = Interpolator(self.v_CG1, self.v_DG1)
if self.method == Boundary_Method.dynamics:
# STRATEGY
# obtain a coordinate field for all the nodes
self.act_coords = Function(VuDG1).project(x) # actual coordinates
self.eff_coords = Function(VuDG1).project(
x) # effective coordinates
self.output = Function(VDG1)
shapes = {
"nDOFs":
self.v_DG1.function_space().finat_element.space_dimension(),
"dim":
np.prod(VuDG1.shape, dtype=int)
}
num_ext_domain = ("{{[i]: 0 <= i < {nDOFs}}}").format(**shapes)
num_ext_instructions = ("""
<float64> SUM_EXT = 0
for i
SUM_EXT = SUM_EXT + EXT_V1[i]
end
NUM_EXT[0] = SUM_EXT
""")
coords_domain = ("{{[i, j, k, ii, jj, kk, ll, mm, iii, kkk]: "
"0 <= i < {nDOFs} and "
"0 <= j < {nDOFs} and 0 <= k < {dim} and "
"0 <= ii < {nDOFs} and 0 <= jj < {nDOFs} and "
"0 <= kk < {dim} and 0 <= ll < {dim} and "
"0 <= mm < {dim} and 0 <= iii < {nDOFs} and "
"0 <= kkk < {dim}}}").format(**shapes)
coords_insts = (
"""
<float64> sum_V1_ext = 0
<int> index = 100
<float64> dist = 0.0
<float64> max_dist = 0.0
<float64> min_dist = 0.0
"""
# only do adjustment in cells with at least one DOF to adjust
"""
if NUM_EXT[0] > 0
"""
# find the maximum distance between DOFs in this cell, to serve as starting point for finding min distances
"""
for i
for j
dist = 0.0
for k
dist = dist + pow(ACT_COORDS[i,k] - ACT_COORDS[j,k], 2.0)
end
dist = pow(dist, 0.5) {{id=sqrt_max_dist, dep=*}}
max_dist = fmax(dist, max_dist) {{id=max_dist, dep=sqrt_max_dist}}
end
end
"""
# loop through cells and find which ones to adjust
"""
for ii
if EXT_V1[ii] > 0.5
"""
# find closest interior node
"""
min_dist = max_dist
index = 100
for jj
if EXT_V1[jj] < 0.5
dist = 0.0
for kk
dist = dist + pow(ACT_COORDS[ii,kk] - ACT_COORDS[jj,kk], 2)
end
dist = pow(dist, 0.5)
if dist <= min_dist
index = jj
end
min_dist = fmin(min_dist, dist)
for ll
EFF_COORDS[ii,ll] = 0.5 * (ACT_COORDS[ii,ll] + ACT_COORDS[index,ll])
end
end
end
else
"""
# for DOFs that aren't exterior, use the original coordinates
"""
for mm
EFF_COORDS[ii, mm] = ACT_COORDS[ii, mm]
end
end
end
else
"""
# for interior elements, just use the original coordinates
"""
for iii
for kkk
EFF_COORDS[iii, kkk] = ACT_COORDS[iii, kkk]
end
end
end
""").format(**shapes)
_num_ext_kernel = (num_ext_domain, num_ext_instructions)
_eff_coords_kernel = (coords_domain, coords_insts)
self.gaussian_elimination_kernel = kernels.GaussianElimination(
VDG1)
# find number of external DOFs per cell
par_loop(_num_ext_kernel,
dx, {
"NUM_EXT": (self.num_ext, WRITE),
"EXT_V1": (self.coords_to_adjust, READ)
},
is_loopy_kernel=True)
# find effective coordinates
logger.warning(
'Finding effective coordinates for boundary recovery. This could give unexpected results for deformed meshes over very steep topography.'
)
par_loop(_eff_coords_kernel,
dx, {
"EFF_COORDS": (self.eff_coords, WRITE),
"ACT_COORDS": (self.act_coords, READ),
"NUM_EXT": (self.num_ext, READ),
"EXT_V1": (self.coords_to_adjust, READ)
},
is_loopy_kernel=True)
elif self.method == Boundary_Method.physics:
self.bottom_kernel = kernels.PhysicsRecoveryBottom()
self.top_kernel = kernels.PhysicsRecoveryTop()
def __init__(self,
v_CG1,
v_DG1,
method=Boundary_Method.physics,
eff_coords=None):
self.v_DG1 = v_DG1
self.v_CG1 = v_CG1
self.v_DG1_old = Function(v_DG1.function_space())
self.eff_coords = eff_coords
self.method = method
mesh = v_CG1.function_space().mesh()
DG0 = FunctionSpace(mesh, "DG", 0)
CG1 = FunctionSpace(mesh, "CG", 1)
if DG0.extruded:
cell = mesh._base_mesh.ufl_cell().cellname()
DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced")
DG1_vert_elt = FiniteElement("DG",
interval,
1,
variant="equispaced")
DG1_element = TensorProductElement(DG1_hori_elt, DG1_vert_elt)
else:
cell = mesh.ufl_cell().cellname()
DG1_element = FiniteElement("DG", cell, 1, variant="equispaced")
DG1 = FunctionSpace(mesh, DG1_element)
self.num_ext = find_domain_boundaries(mesh)
# check function spaces of functions
if self.method == Boundary_Method.dynamics:
if v_CG1.function_space() != CG1:
raise NotImplementedError(
"This boundary recovery method requires v1 to be in CG1.")
if v_DG1.function_space() != DG1:
raise NotImplementedError(
"This boundary recovery method requires v_out to be in DG1."
)
if eff_coords is None:
raise ValueError(
'Need eff_coords field for dynamics boundary methods')
elif self.method == Boundary_Method.physics:
# check that mesh is valid -- must be an extruded mesh
if not DG0.extruded:
raise NotImplementedError(
'The physics boundary method only works on extruded meshes'
)
# base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
w_hori = FiniteElement("DG", cell, 0, variant="equispaced")
w_vert = FiniteElement("CG", interval, 1, variant="equispaced")
# build element
theta_element = TensorProductElement(w_hori, w_vert)
# spaces
Vtheta = FunctionSpace(mesh, theta_element)
Vtheta_broken = FunctionSpace(mesh, BrokenElement(theta_element))
if v_CG1.function_space() != Vtheta:
raise ValueError(
"This boundary recovery method requires v_CG1 to be in DG0xCG1 TensorProductSpace."
)
if v_DG1.function_space() != Vtheta_broken:
raise ValueError(
"This boundary recovery method requires v_DG1 to be in the broken DG0xCG1 TensorProductSpace."
)
else:
raise ValueError(
"Boundary method should be a Boundary Method Enum object.")
vec_DG1 = VectorFunctionSpace(DG0.mesh(), DG1_element)
x = SpatialCoordinate(DG0.mesh())
self.interpolator = Interpolator(self.v_CG1, self.v_DG1)
if self.method == Boundary_Method.dynamics:
# STRATEGY
# obtain a coordinate field for all the nodes
self.act_coords = Function(vec_DG1).project(
x) # actual coordinates
self.eff_coords = eff_coords # effective coordinates
self.output = Function(DG1)
self.on_exterior = find_domain_boundaries(mesh)
self.gaussian_elimination_kernel = kernels.GaussianElimination(DG1)
elif self.method == Boundary_Method.physics:
self.bottom_kernel = kernels.PhysicsRecoveryBottom()
self.top_kernel = kernels.PhysicsRecoveryTop()
