def argument(self, o):
from ufl import split
from firedrake import MixedFunctionSpace, FunctionSpace
V = o.function_space()
if len(V) == 1:
# Not on a mixed space, just return ourselves.
return o
if o in self._arg_cache:
return self._arg_cache[o]
V_is = V.split()
indices = self.blocks[o.number()]
try:
indices = tuple(indices)
nidx = len(indices)
except TypeError:
# Only one index provided.
indices = (indices, )
nidx = 1
if nidx == 1:
W = V_is[indices[0]]
W = FunctionSpace(W.mesh(), W.ufl_element())
a = (Argument(W, o.number(), part=o.part()), )
else:
W = MixedFunctionSpace([V_is[i] for i in indices])
a = split(Argument(W, o.number(), part=o.part()))
args = []
for i in range(len(V_is)):
if i in indices:
c = indices.index(i)
a_ = a[c]
if len(a_.ufl_shape) == 0:
args += [a_]
else:
args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
else:
args += [
Zero()
for j in numpy.ndindex(V_is[i].ufl_element().value_shape())
]
return self._arg_cache.setdefault(o, as_vector(args))
def argument(self, o):
from ufl import split
from firedrake import MixedFunctionSpace, FunctionSpace
V = o.function_space()
if len(V) == 1:
# Not on a mixed space, just return ourselves.
return o
if o in self._arg_cache:
return self._arg_cache[o]
V_is = V.split()
indices = self.blocks[o.number()]
try:
indices = tuple(indices)
nidx = len(indices)
except TypeError:
# Only one index provided.
indices = (indices, )
nidx = 1
if nidx == 1:
W = V_is[indices[0]]
W = FunctionSpace(W.mesh(), W.ufl_element())
a = (Argument(W, o.number(), part=o.part()), )
else:
W = MixedFunctionSpace([V_is[i] for i in indices])
a = split(Argument(W, o.number(), part=o.part()))
args = []
for i in range(len(V_is)):
if i in indices:
c = indices.index(i)
a_ = a[c]
if len(a_.ufl_shape) == 0:
args += [a_]
else:
args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
else:
args += [Zero()
for j in numpy.ndindex(
V_is[i].ufl_element().value_shape())]
return self._arg_cache.setdefault(o, as_vector(args))
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 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
class TimeDiscretisation(object, metaclass=ABCMeta):
"""
Base class for time discretisation schemes.
:arg state: :class:`.State` object.
:arg field: field to be evolved
:arg equation: :class:`.Equation` object, specifying the equation
that field satisfies
:arg solver_parameters: solver_parameters
:arg limiter: :class:`.Limiter` object.
:arg options: :class:`.DiscretisationOptions` object
"""
def __init__(self, state, field_name=None, solver_parameters=None,
limiter=None, options=None):
self.state = state
self.field_name = field_name
self.dt = self.state.dt
self.limiter = limiter
self.options = options
if options is not None:
self.discretisation_option = options.name
else:
self.discretisation_option = None
# get default solver options if none passed in
if solver_parameters is None:
self.solver_parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
else:
self.solver_parameters = solver_parameters
if logger.isEnabledFor(DEBUG):
self.solver_parameters["ksp_monitor_true_residual"] = None
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 pre_apply(self, x_in, discretisation_option):
"""
Extra steps to discretisation if using an embedded method,
which might be either the plain embedded method or the
recovered space scheme.
:arg x_in: the input set of prognostic fields.
:arg discretisation option: string specifying which scheme to use.
"""
if discretisation_option == "embedded_dg":
try:
self.xdg_in.interpolate(x_in)
except NotImplementedError:
self.xdg_in.project(x_in)
elif discretisation_option == "recovered":
self.x_in.assign(x_in)
self.x_rec_projector.project()
self.x_brok_projector.project()
self.xdg_interpolator.interpolate()
def post_apply(self, x_out, discretisation_option):
"""
The projection steps, returning a field to its original space
for an embedded DG scheme. For the case of the
recovered scheme, there are two options dependent on whether
the scheme is limited or not.
:arg x_out: the outgoing field.
:arg discretisation_option: string specifying which option to use.
"""
if discretisation_option == "recovered" and self.limiter is not None:
self.x_brok_interpolator.interpolate()
self.x_out_projector.project()
x_out.assign(self.x_projected)
@abstractproperty
def lhs(self):
l = self.residual.label_map(
lambda t: t.has_label(time_derivative),
map_if_true=replace_subject(self.dq, self.idx),
map_if_false=drop)
return l.form
@abstractproperty
def rhs(self):
r = self.residual.label_map(
all_terms,
map_if_true=replace_subject(self.q1, self.idx))
r = r.label_map(
lambda t: t.has_label(time_derivative),
map_if_false=lambda t: -self.dt*t)
return r.form
def replace_transport_term(self):
"""
This routine allows the default transport term to be replaced with a
different one, specified through the transport options.
This is necessary because when the prognostic equations are declared,
the whole transport
"""
# Extract transport term of equation
old_transport_term_list = self.residual.label_map(
lambda t: t.has_label(transport), map_if_false=drop)
# If there are more transport terms, extract only the one for this variable
if len(old_transport_term_list.terms) > 1:
raise NotImplementedError('Cannot replace transport terms when there are more than one')
# Then we should only have one transport term
old_transport_term = old_transport_term_list.terms[0]
# If the transport term has an ibp label, then it could be replaced
if old_transport_term.has_label(ibp_label) and hasattr(self.options, 'ibp'):
# Do the options specify a different ibp to the old transport term?
if old_transport_term.labels['ibp'] != self.options.ibp:
# Set up a new transport term
field = self.state.fields(self.field_name)
test = TestFunction(self.fs)
# Set up new transport term (depending on the type of transport equation)
if old_transport_term.labels['transport'] == TransportEquationType.advective:
new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp)
elif old_transport_term.labels['transport'] == TransportEquationType.conservative:
new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp)
else:
raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}')
# Finally, drop the old transport term and add the new one
self.residual = self.residual.label_map(
lambda t: t.has_label(transport), map_if_true=drop)
self.residual += subject(new_transport_term, field)
def replace_transporting_velocity(self, uadv):
# replace the transporting velocity in any terms that contain it
if any([t.has_label(transporting_velocity) for t in self.residual]):
assert uadv is not None
if uadv == "prognostic":
self.residual = self.residual.label_map(
lambda t: t.has_label(transporting_velocity),
map_if_true=lambda t: Term(ufl.replace(
t.form, {t.get(transporting_velocity): split(t.get(subject))[0]}), t.labels)
)
else:
self.residual = self.residual.label_map(
lambda t: t.has_label(transporting_velocity),
map_if_true=lambda t: Term(ufl.replace(
t.form, {t.get(transporting_velocity): uadv}), t.labels)
)
self.residual = transporting_velocity.update_value(self.residual, uadv)
@cached_property
def solver(self):
# setup solver using lhs and rhs defined in derived class
problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.dq, bcs=self.bcs)
solver_name = self.field_name+self.__class__.__name__
return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
@abstractmethod
def apply(self, x_in, x_out):
"""
Function takes x as input, computes L(x) as defined by the equation,
and returns x_out as output.
:arg x: :class:`.Function` object, the input Function.
:arg x_out: :class:`.Function` object, the output Function.
"""
pass
class ThetaLimiter(object):
"""
A vertex based limiter for fields in the DG1xCG2 space,
i.e. temperature variables. This acts like the vertex-based
limiter implemented in Firedrake, but in addition corrects
the central nodes to prevent new maxima or minima forming.
"""
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 copy_vertex_values(self, field):
"""
Copies the vertex values from temperature space to
DG1 space which only has vertices.
"""
par_loop(self._copy_into_DG1_kernel,
dx, {
"theta": (field, READ),
"theta_hat": (self.theta_hat, WRITE)
},
is_loopy_kernel=True)
def copy_vertex_values_back(self, field):
"""
Copies the vertex values back from the DG1 space to
the original temperature space, and checks that the
midpoint values are within the minimum and maximum
at the adjacent vertices.
If outside of the minimum and maximum, correct the values
to be the average.
"""
par_loop(self._copy_from_DG1_kernel,
dx, {
"theta": (field, WRITE),
"theta_hat": (self.theta_hat, READ),
"theta_old": (self.theta_old, READ)
},
is_loopy_kernel=True)
def apply(self, field):
"""
The application of the limiter to the theta-space field.
"""
assert field.function_space() == self.Vt, \
'Given field does not belong to this objects function space'
self.theta_old.assign(field)
self.copy_vertex_values(field)
self.vertex_limiter.apply(self.theta_hat)
self.copy_vertex_values_back(field)
def setup_limiters(dirname, direction, grid_params, ic_params):
# declare grid shape
L, H, ncolumns, nlayers = grid_params
# 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)
output = OutputParameters(dirname=dirname,
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
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()))
# declare initial fields
u0 = state.fields("u")
theta0 = state.fields("theta0", V0)
theta1 = state.fields("theta1", V1)
Tsurf, xc, zc, rc, u_max = ic_params
# Isentropic background state
thetab = Constant(Tsurf)
theta_b1 = Function(V1).interpolate(thetab)
theta_b0 = Function(V0).interpolate(thetab)
# set up bubble
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)
if direction == "horizontal":
Vpsi = FunctionSpace(mesh, "CG", 2)
psi_expr = -u_max * z
psi0 = Function(Vpsi).interpolate(psi_expr)
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
u0.project(gradperp(psi0))
elif direction == "vertical":
u0.project(as_vector([0, -u_max]))
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
dg_opts = EmbeddedDGOptions()
recovered_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=V0_brok)
thetaeqn1 = EmbeddedDGAdvection(state,
V1,
equation_form="advective",
options=dg_opts)
thetaeqn0 = EmbeddedDGAdvection(state,
V0,
equation_form="advective",
options=recovered_opts)
# build advection dictionary
advected_fields = []
advected_fields.append(
('theta1', SSPRK3(state, theta1, thetaeqn1, limiter=ThetaLimiter(V1))))
advected_fields.append(('theta0',
SSPRK3(state,
theta0,
thetaeqn0,
limiter=VertexBasedLimiter(VDG1))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 40.0
def NavierStokesBrinkmannForm(
W: fd.FunctionSpace,
w: fd.Function,
nu,
phi: Union[fd.Function, Product] = None,
brinkmann_penalty: fd.Constant = None,
brinkmann_min=0.0,
design_domain=None,
hs: Callable = hs,
beta_gls=0.9,
) -> ufl.form:
"""Returns the Galerkin Least Squares formulation for the Navier-Stokes problem with a Brinkmann term
Args:
W (fd.FunctionSpace): [description]
w (fd.Function): [description]
phi (fd.Function): [description]
nu ([type]): [description]
brinkmann_penalty ([type], optional): [description]. Defaults to None.
design_domain ([type], optional): Region where the level set is defined. Defaults to None.
Returns:
ufl.form: Nonlinear form
"""
mesh = w.ufl_domain()
W_elem = W.ufl_element()
assert isinstance(W_elem, fd.MixedElement)
if brinkmann_penalty:
assert isinstance(brinkmann_penalty, fd.Constant)
assert W_elem.num_sub_elements() == 2
for W_sub_elem in W_elem.sub_elements():
assert W_sub_elem.family() == "Lagrange"
assert W_sub_elem.degree() == 1
assert isinstance(W_elem.sub_elements()[0], fd.VectorElement)
v, q = fd.TestFunctions(W)
u, p = fd.split(w)
# Main NS form
F = (nu * inner(grad(u), grad(v)) * dx + inner(dot(grad(u), u), v) * dx -
p * div(v) * dx + div(u) * q * dx)
# Brinkmann terms for design
def add_measures(list_dd, **kwargs):
return sum((dx(dd, kwargs) for dd in list_dd[1::]), dx(list_dd[0]))
def alpha(phi):
return brinkmann_penalty * hs(phi) + fd.Constant(brinkmann_min)
if brinkmann_penalty and phi is not None:
if design_domain is not None:
dx_brinkmann = partial(add_measures, Enlist(design_domain))
else:
dx_brinkmann = dx
F = F + alpha(phi) * inner(u, v) * dx_brinkmann()
# GLS stabilization
R_U = dot(u, grad(u)) - nu * div(grad(u)) + grad(p)
if isinstance(beta_gls, (float, int)):
beta_gls = fd.Constant(beta_gls)
h = fd.CellSize(mesh)
tau_gls = beta_gls * ((4.0 * dot(u, u) / h**2) + 9.0 *
(4.0 * nu / h**2)**2)**(-0.5)
theta_U = dot(u, grad(v)) - nu * div(grad(v)) + grad(q)
F = F + tau_gls * inner(R_U, theta_U) * dx()
if brinkmann_penalty and phi is not None:
tau_gls_alpha = beta_gls * ((4.0 * dot(u, u) / h**2) + 9.0 *
(4.0 * nu / h**2)**2 +
(alpha(phi) / 1.0)**2)**(-0.5)
R_U_alpha = R_U + alpha(phi) * u
theta_alpha = theta_U + alpha(phi) * v
F = F + tau_gls_alpha * inner(R_U_alpha, theta_alpha) * dx_brinkmann()
if (design_domain is not None
): # Substract this domain from the original integral
F = F - tau_gls * inner(R_U, theta_U) * dx_brinkmann()
return F
