def build(self):
"""
Build the gaussian function.
Raises:
AttributeError: If required properties are not defined.
Returns:
Function: firedrake Constant set to the given value.
"""
for k in self.properties:
if self._props[k] is None:
raise AttributeError('"{}" has not been defined.'.format(k))
mean = self._props['mean']
sd = self._props['sd']
scale = self._props['scale']
xs = SpatialCoordinate(self.mesh)
if not isinstance(mean, list):
mean = [mean] * len(xs)
if not isinstance(sd, list):
sd = [sd] * len(xs)
gaussian = Function(self.V)
components = [exp(-(x - m)**2 / 2 / s**2)
for x, m, s in zip(xs, mean, sd)]
product = components[0]
for c in components[1:]:
product *= c
gaussian.interpolate(scale * product)
return gaussian
def find_domain_boundaries(mesh):
"""
Makes a scalar DG0 function whose values are 0. everywhere except for in
cells on the boundary of the domain, where the values are 1.0.
This allows boundary cells to be identified easily.
:arg mesh: the mesh.
"""
DG0 = FunctionSpace(mesh, "DG", 0)
CG1 = FunctionSpace(mesh, "CG", 1)
on_exterior_DG0 = Function(DG0)
on_exterior_CG1 = Function(CG1)
# we get values in CG1 initially as DG0 will not work for triangular elements
bc_codes = ['on_boundary', 'top', 'bottom']
bcs = [DirichletBC(CG1, Constant(1.0), bc_code) for bc_code in bc_codes]
for bc in bcs:
try:
bc.apply(on_exterior_CG1)
except ValueError:
pass
on_exterior_DG0.interpolate(on_exterior_CG1)
return on_exterior_DG0
class EmbeddedDGAdvection(DGAdvection):
def __init__(self, state, V, Vdg=None, continuity=False):
if Vdg is None:
Vdg_elt = BrokenElement(V.ufl_element())
Vdg = FunctionSpace(state.mesh, Vdg_elt)
super(EmbeddedDGAdvection, self).__init__(state, Vdg, continuity)
self.xdg_in = Function(Vdg)
self.xdg_out = Function(Vdg)
self.x_projected = Function(V)
pparameters = {'ksp_type':'cg',
'pc_type':'bjacobi',
'sub_pc_type':'ilu'}
self.Projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=pparameters)
def apply(self, x_in, x_out):
self.xdg_in.interpolate(x_in)
super(EmbeddedDGAdvection, self).apply(self.xdg_in, self.xdg_out)
self.Projector.project()
x_out.assign(self.x_projected)
def _bezier_plot(function, axes, **kwargs):
"""Plot a 1D function on a function space with order no more than 4 using
Bezier curves within each cell
:arg function: 1D :class:`~.Function` to plot
:arg axes: :class:`Axes <matplotlib.axes.Axes>` for plotting
:arg kwargs: additional key work arguments to plot
:return: matplotlib :class:`PathPatch <matplotlib.patches.PathPatch>`
"""
deg = function.function_space().ufl_element().degree()
mesh = function.function_space().mesh()
if deg == 0:
V = FunctionSpace(mesh, "DG", 1)
func = Function(V).interpolate(function)
return _bezier_plot(func, axes, **kwargs)
y_vals = _bezier_calculate_points(function)
x = SpatialCoordinate(mesh)
coords = Function(FunctionSpace(mesh, 'DG', deg))
coords.interpolate(x[0])
x_vals = _bezier_calculate_points(coords)
vals = np.dstack((x_vals, y_vals))
codes = {1: [Path.MOVETO, Path.LINETO],
2: [Path.MOVETO, Path.CURVE3, Path.CURVE3],
3: [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4]}
vertices = vals.reshape(-1, 2)
path = Path(vertices, np.tile(codes[deg],
function.function_space().cell_node_list.shape[0]))
kwargs["facecolor"] = kwargs.pop("facecolor", "none")
kwargs["linewidth"] = kwargs.pop("linewidth", 2.)
patch = matplotlib.patches.PathPatch(path, **kwargs)
axes.add_patch(patch)
return patch
def test_average(geometry, mesh):
cell = mesh.ufl_cell().cellname()
DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced")
vec_DG1 = VectorFunctionSpace(mesh, DG1_elt)
vec_DG0 = VectorFunctionSpace(mesh, "DG", 0)
vec_CG1 = VectorFunctionSpace(mesh, "CG", 1)
# We will fill DG1_field with values, and average them to CG_field
# First need to put the values into DG0 and then interpolate
DG0_field = Function(vec_DG0)
DG1_field = Function(vec_DG1)
CG_field = Function(vec_CG1)
weights = Function(vec_CG1)
DG0_field, weights, true_values, CG_index = setup_values(
geometry, DG0_field, weights)
DG1_field.interpolate(DG0_field)
kernel = kernels.Average(vec_CG1)
kernel.apply(CG_field, weights, DG1_field)
tolerance = 1e-12
if geometry == "1D":
assert abs(CG_field.dat.data[CG_index] - true_values) < tolerance
elif geometry == "2D":
assert abs(CG_field.dat.data[CG_index][0] - true_values[0]) < tolerance
assert abs(CG_field.dat.data[CG_index][1] - true_values[1]) < tolerance
def run_tracer(setup):
# Get initial conditions from shared config
state = setup.state
mesh = state.mesh
dt = state.dt
output = state.output
x = SpatialCoordinate(state.mesh)
H = 0.1
parameters = ShallowWaterParameters(H=H)
Omega = parameters.Omega
g = parameters.g
umax = setup.umax
R = setup.radius
fexpr = 2 * Omega * x[2] / R
# Need to create a new state containing parameters
state = State(mesh, dt=dt, output=output, parameters=parameters)
# Equations
eqns = LinearShallowWaterEquations(state,
setup.family,
setup.degree,
fexpr=fexpr)
tracer_eqn = AdvectionEquation(state, state.spaces("DG"), "tracer")
# Specify initial prognostic fields
u0 = state.fields("u")
D0 = state.fields("D")
tracer0 = state.fields("tracer", D0.function_space())
tracer_end = Function(D0.function_space())
# Expressions for initial fields corresponding to Williamson 2 test case
Dexpr = H - ((R * Omega * umax) * (x[2] * x[2] / (R * R))) / g
u0.project(setup.uexpr)
D0.interpolate(Dexpr)
tracer0.interpolate(setup.f_init)
tracer_end.interpolate(setup.f_end)
# set up transport schemes
transport_schemes = [ForwardEuler(state, "D")]
# Set up tracer transport
tracer_transport = [(tracer_eqn, SSPRK3(state))]
# build time stepper
stepper = CrankNicolson(state,
eqns,
transport_schemes,
auxiliary_equations_and_schemes=tracer_transport)
stepper.run(t=0, tmax=setup.tmax)
error = norm(state.fields("tracer") - tracer_end) / norm(tracer_end)
return error
def b(mesh, fs):
v = TestFunction(fs)
f = Function(fs)
x = SpatialCoordinate(mesh)
f.interpolate((4.*pi*pi)*sin(x[0]*pi*2))
L = f * v * dx
b = assemble(L)
return b
def bezier_plot(function, axes=None, **kwargs):
"""Plot a 1D function on a function space with order no more than 4 using
Bezier curve within each cell, return a matplotlib axes
:arg function: 1D function for plotting
:arg axes: Axes for plotting, if None, a new one will be created
:arg kwargs: additional key work arguments to plot
"""
try:
import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches
except ImportError:
raise RuntimeError("Matplotlib not importable, is it installed?")
deg = function.function_space().ufl_element().degree()
mesh = function.function_space().mesh()
if deg == 0:
V = FunctionSpace(mesh, "DG", 1)
func = Function(V).interpolate(function)
return bezier_plot(func, axes, **kwargs)
y_vals = _bezier_calculate_points(function)
x = SpatialCoordinate(mesh)
coords = Function(FunctionSpace(mesh, 'DG', deg))
coords.interpolate(x[0])
x_vals = _bezier_calculate_points(coords)
vals = np.dstack((x_vals, y_vals))
if axes is None:
figure = plt.figure()
axes = figure.add_subplot(111)
codes = {
1: [Path.MOVETO, Path.LINETO],
2: [Path.MOVETO, Path.CURVE3, Path.CURVE3],
3: [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4]
}
vertices = vals.reshape(-1, 2)
path = Path(
vertices,
np.tile(codes[deg],
function.function_space().cell_node_list.shape[0]))
patch = patches.PathPatch(path, facecolor='none', lw=2)
axes.add_patch(patch)
axes.plot(**kwargs)
return axes
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)
def bezier_plot(function, axes=None, **kwargs):
"""Plot a 1D function on a function space with order no more than 4 using
Bezier curve within each cell, return a matplotlib axes
:arg function: 1D function for plotting
:arg axes: Axes for plotting, if None, a new one will be created
:arg kwargs: additional key work arguments to plot
"""
try:
import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches
except ImportError:
raise RuntimeError("Matplotlib not importable, is it installed?")
deg = function.function_space().ufl_element().degree()
mesh = function.function_space().mesh()
if deg == 0:
V = FunctionSpace(mesh, "DG", 1)
func = Function(V).interpolate(function)
return bezier_plot(func, axes, **kwargs)
y_vals = _bezier_calculate_points(function)
x = SpatialCoordinate(mesh)
coords = Function(FunctionSpace(mesh, 'DG', deg))
coords.interpolate(x[0])
x_vals = _bezier_calculate_points(coords)
vals = np.dstack((x_vals, y_vals))
if axes is None:
figure = plt.figure()
axes = figure.add_subplot(111)
codes = {1: [Path.MOVETO, Path.LINETO],
2: [Path.MOVETO, Path.CURVE3, Path.CURVE3],
3: [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4]}
vertices = vals.reshape(-1, 2)
path = Path(vertices, np.tile(codes[deg],
function.function_space().cell_node_list.shape[0]))
patch = patches.PathPatch(path, facecolor='none', lw=2)
axes.add_patch(patch)
axes.plot(**kwargs)
return axes
def correct_eff_coords(eff_coords):
"""
Correct the effective coordinates calculated by simply averaging
which will not be correct at periodic boundaries.
:arg eff_coords: the effective coordinates in vec_DG1 space.
"""
mesh = eff_coords.function_space().mesh()
vec_CG1 = VectorFunctionSpace(mesh, "CG", 1)
if vec_CG1.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")
vec_DG1 = VectorFunctionSpace(mesh, DG1_element)
x = SpatialCoordinate(mesh)
if eff_coords.function_space() != vec_DG1:
raise ValueError('eff_coords needs to be in the vector DG1 space')
# obtain different coords in DG1
DG1_coords = Function(vec_DG1).interpolate(x)
CG1_coords_from_DG1 = Function(vec_CG1)
averager = Averager(DG1_coords, CG1_coords_from_DG1)
averager.project()
DG1_coords_from_averaged_CG1 = Function(vec_DG1).interpolate(
CG1_coords_from_DG1)
DG1_coords_diff = Function(vec_DG1).interpolate(
DG1_coords - DG1_coords_from_averaged_CG1)
# interpolate coordinates, adjusting those different coordinates
adjusted_coords = Function(vec_DG1)
adjusted_coords.interpolate(eff_coords + DG1_coords_diff)
return adjusted_coords
class SqueezedDQ1Filter:
"""
Filter that acts on squashed quads (wedges?)
In those cells in an extruded mesh that have a vertical facet that is not an external
boundary facet or an interior facet (i.e. they are not in dS_v or ds_v), the DQ1 values
of nodes on that facet are averaged, and all set to that average value. In 2D, assuming
that facet has been squeezed to length 0, this results in a DQ1 solution in the resulting
triangle that is linear.
"""
def __init__(self, dq1_space):
v = TestFunction(dq1_space)
# nodes on squeezed facets are not included in dS_v or ds_v
# and all other nodes are (for DQ1), so this step gives nonzero for these other nodes
self.marker = assemble(avg(v) * dS_v + v * ds_v)
# flip this: 1 for squeezed nodes, and 0 for all others
self.marker.assign(conditional(self.marker > 1e-12, 0, 1))
self.P0 = FunctionSpace(dq1_space.mesh(), "DG", 0)
self.u0 = Function(self.P0, name='averaged squeezed values')
self.u1 = Function(dq1_space, name='aux. squeezed values')
def apply(self, u):
self.u1.interpolate(self.marker * u)
self.u0.interpolate(2 * self.u1)
self.u1.interpolate(self.u0)
u.assign(self.marker * self.u1 + (1 - self.marker) * u)
def regularization_form(r):
mesh = UnitSquareMesh(2 ** r, 2 ** r)
x = SpatialCoordinate(mesh)
S = VectorFunctionSpace(mesh, "CG", 1)
beta = 4.0
reg_solver = RegularizationSolver(S, mesh, beta=beta, gamma=0.0, dx=dx)
# Exact solution with free Neumann boundary conditions for this domain
u_exact = Function(S)
u_exact_component = cos(x[0] * pi * 2) * cos(x[1] * pi * 2)
u_exact.interpolate(as_vector((u_exact_component, u_exact_component)))
f = Function(S)
theta = TestFunction(S)
f_component = (1 + beta * 8 * pi * pi) * u_exact_component
f.interpolate(as_vector((f_component, f_component)))
rhs_form = inner(f, theta) * dx
velocity = Function(S)
rhs = assemble(rhs_form)
reg_solver.solve(velocity, rhs)
File("solution_vel_unitsquare.pvd").write(velocity)
return norm(project(u_exact - velocity, S))
def test_cond_evap(tmpdir, process):
dirname = str(tmpdir)
state, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(
dirname, process)
water_v = state.fields('vapour_mixing_ratio')
water_c = state.fields('cloud_liquid_mixing_ratio')
theta_vd = state.fields('theta')
theta_d = Function(theta_vd.function_space())
theta_d.interpolate(
theta_vd / (1 + water_v * state.parameters.R_v / state.parameters.R_d))
# Check that water vapour is approximately equal to saturation amount
assert norm(water_v - mv_true) / norm(mv_true) < 0.01, \
f'Final vapour field is incorrect for {process}'
# Check that cloud has been created / removed
denom = norm(mc_true) if process == "condensation" else norm(mc_init)
assert norm(water_c - mc_true) / denom < 0.1, \
f'Final cloud field is incorrect for {process}'
# Check that theta pertubation has correct sign and size
assert norm(theta_d - theta_d_true) / norm(theta_d_true) < 0.01, \
f'Latent heating is incorrect for {process}'
# Check that total moisture conserved
filename = path.join(dirname, "cond_evap/diagnostics.nc")
data = Dataset(filename, "r")
water = data.groups["vapour_mixing_ratio_plus_cloud_liquid_mixing_ratio"]
total = water.variables["total"]
water_t_0 = total[0]
water_t_T = total[-1]
assert abs(water_t_0 - water_t_T) / water_t_0 < 1e-12, \
f'Total amount of water should be conserved by {process}'
def test_2D_dirichlet_regions():
# Can't do mesh refinement because MeshHierarchy ruins the domain tags
mesh = Mesh("./2D_mesh.msh")
dim = mesh.geometric_dimension()
x = SpatialCoordinate(mesh)
S = VectorFunctionSpace(mesh, "CG", 1)
beta = 1.0
reg_solver = RegularizationSolver(
S, mesh, beta=beta, gamma=0.0, dx=dx, design_domain=0
)
# Exact solution with free Neumann boundary conditions for this domain
u_exact = Function(S)
u_exact_component = (-cos(x[0] * pi * 2) + 1) * (-cos(x[1] * pi * 2) + 1)
u_exact.interpolate(
as_vector(tuple(u_exact_component for _ in range(dim)))
)
f = Function(S)
f_component = (
-beta
* (
4.0 * pi * pi * cos(2 * pi * x[0]) * (-cos(2 * pi * x[1]) + 1)
+ 4.0 * pi * pi * cos(2 * pi * x[1]) * (-cos(2 * pi * x[0]) + 1)
)
+ u_exact_component
)
f.interpolate(as_vector(tuple(f_component for _ in range(dim))))
theta = TestFunction(S)
rhs_form = inner(f, theta) * dx
velocity = Function(S)
rhs = assemble(rhs_form)
reg_solver.solve(velocity, rhs)
assert norm(project(domainify(u_exact, x) - velocity, S)) < 1e-3
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 ThetaLimiter(object):
"""
A vertex based limiter for fields in the DG1xCG2 space, i.e. temperature
variables in the next-to-lowest order set of spaces. 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
: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 apply(self, field):
"""
The application of the limiter to the theta-space field.
"""
assert field.function_space() == self.Vt_brok, \
"Given field does not belong to this object's function space"
# Obtain field in equispaced DG space and save original field
self.field_old.interpolate(field)
self.field_DG1.interpolate(field)
# Use vertex based limiter on DG1 field
self.vertex_limiter.apply(self.field_DG1)
# Limit midpoints in fully equispaced Vt space
self._limit_midpoints_kernel.apply(self.field_hat, self.field_DG1,
self.field_old)
# Return to original space
field.interpolate(self.field_hat)
def test_3D_dirichlet_regions():
# Can't do mesh refinement because MeshHierarchy ruins the domain tags
mesh = Mesh("./3D_mesh.msh")
dim = mesh.geometric_dimension()
x = SpatialCoordinate(mesh)
solver_parameters_amg = {
"ksp_type": "cg",
"ksp_converged_reason": None,
"ksp_rtol": 1e-7,
"pc_type": "hypre",
"pc_hypre_type": "boomeramg",
"pc_hypre_boomeramg_max_iter": 5,
"pc_hypre_boomeramg_coarsen_type": "PMIS",
"pc_hypre_boomeramg_agg_nl": 2,
"pc_hypre_boomeramg_strong_threshold": 0.95,
"pc_hypre_boomeramg_interp_type": "ext+i",
"pc_hypre_boomeramg_P_max": 2,
"pc_hypre_boomeramg_relax_type_all": "sequential-Gauss-Seidel",
"pc_hypre_boomeramg_grid_sweeps_all": 1,
"pc_hypre_boomeramg_truncfactor": 0.3,
"pc_hypre_boomeramg_max_levels": 6,
}
S = VectorFunctionSpace(mesh, "CG", 1)
beta = 1.0
reg_solver = RegularizationSolver(
S,
mesh,
beta=beta,
gamma=0.0,
dx=dx,
design_domain=0,
solver_parameters=solver_parameters_amg,
)
# Exact solution with free Neumann boundary conditions for this domain
u_exact = Function(S)
u_exact_component = (
(-cos(x[0] * pi * 2) + 1)
* (-cos(x[1] * pi * 2) + 1)
* (-cos(x[2] * pi * 2) + 1)
)
u_exact.interpolate(
as_vector(tuple(u_exact_component for _ in range(dim)))
)
f = Function(S)
f_component = (
-beta
* (
4.0
* pi
* pi
* cos(2 * pi * x[0])
* (-cos(2 * pi * x[1]) + 1)
* (-cos(2 * pi * x[2]) + 1)
+ 4.0
* pi
* pi
* cos(2 * pi * x[1])
* (-cos(2 * pi * x[0]) + 1)
* (-cos(2 * pi * x[2]) + 1)
+ 4.0
* pi
* pi
* cos(2 * pi * x[2])
* (-cos(2 * pi * x[1]) + 1)
* (-cos(2 * pi * x[0]) + 1)
)
+ u_exact_component
)
f.interpolate(as_vector(tuple(f_component for _ in range(dim))))
theta = TestFunction(S)
rhs_form = inner(f, theta) * dx
velocity = Function(S)
rhs = assemble(rhs_form)
reg_solver.solve(velocity, rhs)
error = norm(
project(
domainify(u_exact, x) - velocity,
S,
solver_parameters=solver_parameters_amg,
)
)
assert error < 5e-2
if COMM_WORLD.Get_rank() == 0:
print("Simulation for Williamson2 test case, model u-ad_D-ad", "\n",
"| Discr details: Poisson implicit, no EP, TM=0.5", "\n", "| dt", dt,
"| tmax", tmax, " | refinement level", ref_level, "\n", "| maxk",
maxk, "| number of dumps (exlcuding t=0):", dumpnumber,
"| nc, h5 dump frequency", nc_h5_dumpfreq, "\n", "| pickup", pickup,
"| Diagnostics:", tuple(nc_diag.keys()))
print("Starting Initial condition, and function setup at", ctime())
R, Omega, = 6371220., 7.292e-5
mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2)
mesh.init_cell_orientations(SpatialCoordinate(mesh))
x = SpatialCoordinate(mesh)
f = Function(FunctionSpace(mesh, "CG", 1))
f.interpolate(2 * Omega * x[2] / R)
g, H = 9.810616, 5960.
u_0 = 2 * pi * R / (12 * 24 * 60 * 60.)
uexpr = u_0 * as_vector([-x[1], x[0], 0.0]) / R
Dexpr = H - (R * Omega * u_0 + u_0**2 / 2.) * x[2]**2 / (g * R**2)
bexpr = Constant(0.)
# Build function spaces
degree = 2
family = ("DG", "BDM", "CG")
W0 = FunctionSpace(mesh, family[0], degree - 1, family[0])
W1_elt = FiniteElement(family[1], triangle, degree)
W1 = FunctionSpace(mesh, W1_elt, name="HDiv")
W2 = FunctionSpace(mesh, family[2], degree + 1)
M = MixedFunctionSpace((W1, W0))
if COMM_WORLD.Get_rank() == 0:
print("Simulation for Galewsky test case, model u-ad_flux", "\n",
"| Discr details: Poisson implicit, no EP, TM=0.5", "\n", "| dt", dt,
"| tmax", tmax, " | refinement level", ref_level, "\n", "| maxk",
maxk, "| dfr field, nc_h5:", field_dumpfreq, nc_h5_dumpfreq, "\n",
"| pickup", pickup, "| Diagnostics:", tuple(nc_diag.keys()))
print("Starting Initial condition, and function setup at", ctime())
R, Omega = 6371220., 7.292e-5
mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2)
mesh.init_cell_orientations(SpatialCoordinate(mesh))
x = SpatialCoordinate(mesh)
f = Function(FunctionSpace(mesh, "CG", 1))
f.interpolate(2 * Omega * x[2] / R)
g, H = 9.810616, 5960.
def latlon_coords(mesh):
"""Compute latitude-longitude coordinates given Cartesian ones"""
x0, y0, z0 = SpatialCoordinate(mesh)
unsafe = z0 / sqrt(x0 * x0 + y0 * y0 + z0 * z0)
safe = Min(Max(unsafe, -1.0), 1.0) # avoid silly roundoff errors
theta = asin(safe) # latitude
lamda = atan_2(y0, x0) # longitude
return theta, lamda
theta, lamda = latlon_coords(mesh)
except ImportError:
makeplots = False
# Set up base FEM, which solves Poisson's equation on a square mesh
nx = 101
mesh = UnitSquareMesh(nx - 1, nx - 1)
V = FunctionSpace(mesh, "CG", 1)
u = TrialFunction(V)
v = TestFunction(V)
f = Function(V)
x = SpatialCoordinate(mesh)
f.interpolate((8 * pi * pi) * sin(x[0] * pi * 2) * sin(x[1] * pi * 2))
a = (dot(grad(v), grad(u))) * dx
L = f * v * dx
bc = DirichletBC(V, 0., "on_boundary")
A = assemble(a, bcs=bc)
b = assemble(L)
u = Function(V)
solve(A, u, b)
# Create some fake data that is systematically different from the FEM solution.
def __init__(self,
mesh: object,
kappa: float,
comm_space: MPI.Comm,
mu: float = 5.,
*args,
**kwargs):
"""
Constructor
:param mesh: spatial domain
:param kappa: diffusion coefficient
:param mu: penalty weighting function
"""
super(Diffusion2D, self).__init__(*args, **kwargs)
# Spatial domain and function space
self.mesh = mesh
V = FunctionSpace(self.mesh, "DG", 1)
self.function_space = V
self.comm_space = comm_space
# Placeholder for time step - will be updated in the update method
self.dt = Constant(0.)
# Things we need for the form
gamma = TestFunction(V)
phi = TrialFunction(V)
self.f = Function(V)
n = FacetNormal(mesh)
# Set up the rhs and bilinear form of the equation
a = (inner(gamma, phi) * dx + self.dt *
(inner(grad(gamma),
grad(phi) * kappa) * dx -
inner(2 * avg(outer(phi, n)), avg(grad(gamma) * kappa)) * dS -
inner(avg(grad(phi) * kappa), 2 * avg(outer(gamma, n))) * dS +
mu * inner(2 * avg(outer(phi, n)),
2 * avg(outer(gamma, n) * kappa)) * dS))
rhs = inner(gamma, self.f) * dx
# Function to hold the solution
self.soln = Function(V)
# Setup problem and solver
prob = LinearVariationalProblem(a, rhs, self.soln)
self.solver = NonlinearVariationalSolver(prob)
# Set the data structure for any user-defined time point
self.vector_template = VectorDiffusion2D(size=len(self.function_space),
comm_space=self.comm_space)
# Set initial condition:
# Setting up a Gaussian blob in the centre of the domain.
self.vector_t_start = VectorDiffusion2D(size=len(self.function_space),
comm_space=self.comm_space)
x = SpatialCoordinate(self.mesh)
initial_tracer = exp(-((x[0] - 5)**2 + (x[1] - 5)**2))
tmp = Function(self.function_space)
tmp.interpolate(initial_tracer)
self.vector_t_start.set_values(np.copy(tmp.dat.data))
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 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)
max_outer_solve_count = 20
max_inner_solve_count = 10
PETSc.Sys.Print("Starting rho solver loop...\n")
for i in range(max_outer_solve_count):
# calculate averaged rho
rho_recoverer.project()
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
# first solve for r_v
for j in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
water_v0.assign(water_v0 * (1 - delta) + delta * w_h)
# compute theta_vd
theta0.assign(theta_d * (1 + water_v0 / epsilon))
# test quality of solution by re-evaluating expression
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
# now solve for rho with theta_vd and w_v guesses
rho_solver.solve()
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
x, = SpatialCoordinate(mesh)
Ld = simulation_parameters['Ld'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.dt = simulation_parameters['dt'][-1]
self.num_Xis = simulation_parameters['num_Xis'][-1]
self.Xi_family = simulation_parameters['Xi_family'][-1]
self.dXi = prognostic_variables.dXi
self.dWs = [Constant(0.0) for dw in range(self.num_Xis)]
self.dW_nums = prognostic_variables.dW_nums
self.Xi_functions = []
self.nXi_updates = simulation_parameters['nXi_updates'][-1]
self.smooth_t = simulation_parameters['smooth_t'][-1]
self.fixed_dW = simulation_parameters['fixed_dW'][-1]
if self.smooth_t is not None and self.nXi_updates > 1:
raise ValueError('Prescribing forcing and including multiple Xi updates are not compatible.')
if self.smooth_t is not None or self.fixed_dW is not None:
print('WARNING: Remember to change sigma to sigma * sqrt(dt) with the prescribed forcing option or the fixed_dW option.')
seed = simulation_parameters['seed'][-1]
np.random.seed(seed)
# make sure sigma is a Constant
if self.num_Xis != 0:
if isinstance(simulation_parameters['sigma'][-1], Constant):
self.sigma = simulation_parameters['sigma'][-1]
else:
self.sigma = Constant(simulation_parameters['sigma'][-1])
else:
self.sigma = Constant(0.0)
self.pure_xi_list = prognostic_variables.pure_xi_list
self.pure_xi_x_list = prognostic_variables.pure_xi_x_list
self.pure_xi_xx_list = prognostic_variables.pure_xi_xx_list
self.pure_xi_xxx_list = prognostic_variables.pure_xi_xxx_list
self.pure_xi_xxxx_list = prognostic_variables.pure_xi_xxxx_list
for xi in range(self.num_Xis):
self.pure_xi_list.append(Function(self.dXi.function_space()))
self.pure_xi_x_list.append(Function(self.dXi.function_space()))
self.pure_xi_xx_list.append(Function(self.dXi.function_space()))
self.pure_xi_xxx_list.append(Function(self.dXi.function_space()))
self.pure_xi_xxxx_list.append(Function(self.dXi.function_space()))
if self.Xi_family == 'sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin(2*(n+1)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos(2*(n+1)*pi*x/Ld))
elif self.Xi_family == 'double_sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin(4*(n+1)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos(4*(n+1)*pi*x/Ld))
elif self.Xi_family == 'high_freq_sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin((2*(n+1)+10)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos((2*(n+1)+10)*pi*x/Ld))
elif self.Xi_family == 'gaussians':
for n in range(self.num_Xis):
self.Xi_functions.append(self.sigma * 0.5*self.num_Xis*exp(-((x-Ld*(n+1)/(self.num_Xis +1.0))/2.)**2))
elif self.Xi_family == 'quadratic':
if self.num_Xis > 1:
raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(32/(Ld*Ld)*conditional(x > Ld/4,
conditional(x > 3*Ld/8,
conditional(x > 5*Ld/8,
conditional(x < 3*Ld/4,
self.sigma * (x - 3*Ld/4)**2,
0.0),
(x-Ld/2)**2+Ld**2/32),
(x-Ld/4)**2),
0.0))
elif self.Xi_family == 'proper_peak':
if self.num_Xis > 1:
raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(self.sigma * 0.5*2/(exp(x-Ld/2)+exp(-x+Ld/2)))
elif self.Xi_family == 'constant':
if self.num_Xis > 1:
raise NotImplementedError('Constant Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(self.sigma * (sin(0*pi*x/Ld)+1))
else:
raise NotImplementedError('Xi_family %s not implemented' % self.Xi_family)
# make lists of functions for xi_x, xi_xx and xi_xxx
if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']:
self.dXi_x = prognostic_variables.dXi_x
self.dXi_xx = prognostic_variables.dXi_xx
self.Xi_x_functions = []
self.Xi_xx_functions = []
for Xi_expr in self.Xi_functions:
Xi_x_function = Function(self.dXi_x.function_space())
Xi_xx_function = Function(self.dXi_xx.function_space())
phi_x = TestFunction(self.dXi_x.function_space())
phi_xx = TestFunction(self.dXi_xx.function_space())
Xi_x_eqn = phi_x * Xi_x_function * dx + phi_x.dx(0) * Xi_expr * dx
Xi_xx_eqn = phi_xx * Xi_xx_function * dx + phi_xx.dx(0) * Xi_x_function * dx
Xi_x_problem = NonlinearVariationalProblem(Xi_x_eqn, Xi_x_function)
Xi_xx_problem = NonlinearVariationalProblem(Xi_xx_eqn, Xi_xx_function)
Xi_x_solver = NonlinearVariationalSolver(Xi_x_problem)
Xi_xx_solver = NonlinearVariationalSolver(Xi_xx_problem)
# for some reason these solvers don't work for constant Xi functions
# so just manually make the derivatives be zero
if self.Xi_family == 'constant':
Xi_x_function.interpolate(0.0*x)
Xi_xx_function.interpolate(0.0*x)
else:
Xi_x_solver.solve()
Xi_xx_solver.solve()
self.Xi_x_functions.append(Xi_x_function)
self.Xi_xx_functions.append(Xi_xx_function)
# now make a master xi
Xi_expr = 0.0*x
for dW, Xi_function, pure_xi, pure_xi_x, pure_xi_xx, pure_xi_xxx, pure_xi_xxxx in zip(self.dWs, self.Xi_functions, self.pure_xi_list, self.pure_xi_x_list, self.pure_xi_xx_list, self.pure_xi_xxx_list, self.pure_xi_xxxx_list):
Xi_expr += dW * Xi_function
if self.scheme in ['upwind', 'LASCH']:
pure_xi.interpolate(as_vector([Xi_function]))
pure_xi_x.project(as_vector([Xi_function.dx(0)]))
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
xixx_scalar = Function(CG1)
xixx_eqn = psi * xixx_scalar * dx + psi.dx(0) * Xi_function.dx(0) * dx
prob = NonlinearVariationalProblem(xixx_eqn, xixx_scalar)
solver = NonlinearVariationalSolver(prob)
solver.solve()
pure_xi_xx.interpolate(as_vector([xixx_scalar]))
else:
pure_xi.interpolate(Xi_function)
# I guess we can't take the gradient of constants
if self.Xi_family != 'constant':
pure_xi_x.project(Xi_function.dx(0))
pure_xi_xx.project(pure_xi_x.dx(0))
pure_xi_xxx.project(pure_xi_xx.dx(0))
pure_xi_xxxx.project(pure_xi_xxx.dx(0))
if self.scheme in ['upwind', 'LASCH']:
self.dXi_interpolator = Interpolator(as_vector([Xi_expr]), self.dXi)
else:
self.dXi_interpolator = Interpolator(Xi_expr, self.dXi)
if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']:
# initialise blank expressions
Xi_x_expr = 0.0*x
Xi_xx_expr = 0.0*x
# make full expressions by adding all dW * Xi_xs
for dW, Xi_x_function, Xi_xx_function in zip(self.dWs, self.Xi_x_functions, self.Xi_xx_functions):
Xi_x_expr += dW * Xi_x_function
Xi_xx_expr += dW * Xi_xx_function
self.dXi_x_interpolator = Interpolator(Xi_x_expr, self.dXi_x)
self.dXi_xx_interpolator = Interpolator(Xi_xx_expr, self.dXi_xx)
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