class IncompressibleState(BaroclinicState):
def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
family="RT", z=None, k=None, Omega=None, mu=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=[],
on_sphere=False):
super(IncompressibleState, self).__init__(mesh=mesh,
vertical_degree=vertical_degree,
horizontal_degree=horizontal_degree,
family=family,
z=z, k=k, Omega=Omega, mu=mu,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
if parameters.geopotential:
raise RuntimeError("geopotential formulation is not compatible with incompressible Boussinesq")
def set_reference_profiles(self, b_ref):
"""
Initialise reference profiles
:arg b_ref: :class:`.Function` object, reference bouyancy
"""
self.bbar = Function(self.V[2])
self.bbar.project(b_ref)
class CompressibleState(BaroclinicState):
def set_reference_profiles(self, rho_ref, theta_ref):
"""
Initialise reference profiles
:arg rho_ref: :class:`.Function` object, initial rho
:arg theta_ref: :class:`.Function` object, initial theta
"""
self.rhobar = Function(self.V[1])
self.thetabar = Function(self.V[2])
self.rhobar.project(rho_ref)
self.thetabar.project(theta_ref)
class Diagnostics(object):
@staticmethod
def min(f_in):
fmin = op2.Global(1, [np.finfo(float).max], dtype=float)
if len(f_in.ufl_shape) > 0:
mesh = f_in.function_space().mesh()
V = FunctionSpace(mesh, "DG", 1)
f = Function(V).project(sqrt(inner(f_in, f_in)))
else:
f = f_in
op2.par_loop(
op2.Kernel(
"""void minify(double *a, double *b)
{
a[0] = a[0] > fabs(b[0]) ? fabs(b[0]) : a[0];
}""", "minify"), f.dof_dset.set, fmin(op2.MIN), f.dat(op2.READ))
return fmin.data[0]
@staticmethod
def max(f_in):
fmax = op2.Global(1, [np.finfo(float).min], dtype=float)
if len(f_in.ufl_shape) > 0:
mesh = f_in.function_space().mesh()
V = FunctionSpace(mesh, "DG", 1)
f = Function(V).project(sqrt(inner(f_in, f_in)))
else:
f = f_in
op2.par_loop(
op2.Kernel(
"""void maxify(double *a, double *b)
{
a[0] = a[0] < fabs(b[0]) ? fabs(b[0]) : a[0];
}""", "maxify"), f.dof_dset.set, fmax(op2.MAX), f.dat(op2.READ))
return fmax.data[0]
@staticmethod
def l2(f):
return sqrt(assemble(dot(f, f) * dx))
@staticmethod
def energy(h, u, g):
return assemble(0.5 * h * (inner(u, u) + g * h) * dx)
def max_courant_number(self, u, dt):
if not hasattr(self, "_area"):
V = FunctionSpace(u.function_space().mesh(), "DG", 0)
expr = TestFunction(V) * dx
self._area = assemble(expr)
self.Courant = Function(V)
self.Courant.project(sqrt(inner(u, u)) / sqrt(self._area) * dt)
return self.Courant.dat.data.max()
def _prepare_output(self, function, max_elem):
from firedrake import FunctionSpace, VectorFunctionSpace, \
TensorFunctionSpace, Function
from tsfc.finatinterface import create_element as create_finat_element
name = function.name()
# Need to project/interpolate?
# If space is not the max element, we must do so.
finat_elem = function.function_space().finat_element
if finat_elem == create_finat_element(max_elem):
return OFunction(array=get_array(function),
name=name,
function=function)
# OK, let's go and do it.
# Build appropriate space for output function.
shape = function.ufl_shape
if len(shape) == 0:
V = FunctionSpace(function.ufl_domain(), max_elem)
elif len(shape) == 1:
if numpy.prod(shape) > 3:
raise ValueError(
"Can't write vectors with more than 3 components")
V = VectorFunctionSpace(function.ufl_domain(),
max_elem,
dim=shape[0])
elif len(shape) == 2:
if numpy.prod(shape) > 9:
raise ValueError(
"Can't write tensors with more than 9 components")
V = TensorFunctionSpace(function.ufl_domain(),
max_elem,
shape=shape)
else:
raise ValueError("Unsupported shape %s" % (shape, ))
output = Function(V)
if self.project:
output.project(function)
else:
output.interpolate(function)
return OFunction(array=get_array(output), name=name, function=output)
class Advection(object, metaclass=ABCMeta):
"""
Base class for advection schemes.
:arg state: :class:`.State` object.
:arg field: field to be advected
:arg equation: :class:`.Equation` object, specifying the equation
that field satisfies
:arg solver_parameters: solver_parameters
:arg limiter: :class:`.Limiter` object.
:arg options: :class:`.AdvectionOptions` object
"""
def __init__(self, state, field, equation=None, *, solver_parameters=None,
limiter=None):
if equation is not None:
self.state = state
self.field = field
self.equation = equation
# get ubar from the equation class
self.ubar = self.equation.ubar
self.dt = self.state.timestepping.dt
# get default solver options if none passed in
if solver_parameters is None:
self.solver_parameters = equation.solver_parameters
else:
self.solver_parameters = solver_parameters
if logger.isEnabledFor(DEBUG):
self.solver_parameters["ksp_monitor_true_residual"] = True
self.limiter = limiter
if hasattr(equation, "options"):
self.discretisation_option = equation.options.name
self._setup(state, field, equation.options)
else:
self.discretisation_option = None
self.fs = field.function_space()
# setup required functions
self.dq = Function(self.fs)
self.q1 = Function(self.fs)
def _setup(self, state, field, options):
if options.name in ["embedded_dg", "recovered"]:
self.fs = options.embedding_space
self.xdg_in = Function(self.fs)
self.xdg_out = Function(self.fs)
self.x_projected = Function(field.function_space())
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
self.Projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
if options.name == "recovered":
# set up the necessary functions
self.x_in = Function(field.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)
def pre_apply(self, x_in, discretisation_option):
"""
Extra steps to advection if using an embedded method,
which might be either the plain embedded method or the
recovered space advection 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 advection 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 == "embedded_dg":
self.Projector.project()
elif discretisation_option == "recovered":
if self.limiter is not None:
self.x_brok_interpolator.interpolate()
self.x_out_projector.project()
else:
self.Projector.project()
x_out.assign(self.x_projected)
@abstractproperty
def lhs(self):
return self.equation.mass_term(self.equation.trial)
@abstractproperty
def rhs(self):
return self.equation.mass_term(self.q1) - self.dt*self.equation.advection_term(self.q1)
def update_ubar(self, xn, xnp1, alpha):
un = xn.split()[0]
unp1 = xnp1.split()[0]
self.ubar.assign(un + alpha*(unp1-un))
@cached_property
def solver(self):
# setup solver using lhs and rhs defined in derived class
problem = LinearVariationalProblem(self.lhs, self.rhs, self.dq)
solver_name = self.field.name()+self.equation.__class__.__name__+self.__class__.__name__
return LinearVariationalSolver(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 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
