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 __init__(self, state, equation_set, transport_schemes,
auxiliary_equations_and_schemes=None,
linear_solver=None,
diffusion_schemes=None,
physics_list=None, **kwargs):
self.maxk = kwargs.pop("maxk", 4)
self.maxi = kwargs.pop("maxi", 1)
self.alpha = kwargs.pop("alpha", 0.5)
if kwargs:
raise ValueError("unexpected kwargs: %s" % list(kwargs.keys()))
schemes = []
self.transport_schemes = []
self.active_transport = []
for scheme in transport_schemes:
schemes.append((scheme, transport))
assert scheme.field_name in equation_set.field_names
self.active_transport.append((scheme.field_name, scheme))
self.diffusion_schemes = []
if diffusion_schemes is not None:
for scheme in diffusion_schemes:
assert scheme.field_name in equation_set.field_names
schemes.append((scheme, diffusion))
self.diffusion_schemes.append((scheme.field_name, scheme))
problem = [(equation_set, tuple(schemes))]
if auxiliary_equations_and_schemes is not None:
problem.append(*auxiliary_equations_and_schemes)
self.auxiliary_schemes = [
(eqn.field_name, scheme)
for eqn, scheme in auxiliary_equations_and_schemes]
else:
self.auxiliary_schemes = []
self.tracers_to_copy = []
for name in equation_set.field_names:
# Extract time derivative for that prognostic
mass_form = equation_set.residual.label_map(
lambda t: (t.has_label(time_derivative) and t.get(prognostic) == name),
map_if_false=drop)
# Copy over field if the time derivative term has no linearisation
if not mass_form.terms[0].has_label(linearisation):
self.tracers_to_copy.append(name)
# TODO: why was this False? Should this be an argument?
apply_bcs = True
super().__init__(state, problem, apply_bcs, physics_list)
self.field_name = equation_set.field_name
W = equation_set.function_space
self.xrhs = Function(W)
self.dy = Function(W)
if linear_solver is None:
self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha)
else:
self.linear_solver = linear_solver
self.forcing = Forcing(equation_set, self.alpha)
self.bcs = equation_set.bcs
def test_update_multiple_attr(self):
"""
Test that updates propagate as expected.
"""
attrs = vars(self.prob).copy()
old_J = Function(self.V)
new_J = Function(self.V)
old_P = old_J**2 * dx
new_P = new_J**2 * dx
self.prob.J = old_J
self.prob.P = old_P
self.prob._update_func('J', new_J)
# No change to existing vars
for a, v in attrs.items():
self.assertEqual(v, getattr(self.prob, a))
# Change J
self.assertNotEqual(self.prob.J, old_J, 'New J is equal to old J')
self.assertEqual(self.prob.J, new_J, 'J is not correct')
# Change P
self.assertNotEqual(self.prob.P, old_P, 'New P is equal to old P')
self.assertEqual(self.prob.P, new_P, 'P is not correct')
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.u0 = Function(self.V)
self.u1 = Function(self.V)
self.h = Function(self.U)
self.a = Function(self.U)
self.p = TestFunction(self.V)
theta = conditions.theta
self.uh = (1-theta) * self.u0 + theta * self.u1
ep_dot = self.strain(grad(self.uh))
self.initial_condition((self.u0, conditions.ic['u']),(self.u1, self.u0),
(self.a, conditions.ic['a']),(self.h, conditions.ic['h']))
zeta = self.zeta(self.h, self.a, self.delta(self.uh))
eta = zeta * params.e ** -2
sigma = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - 0.5 * self.Ice_Strength(self.h,self.a) * Identity(2)
self.eqn = self.momentum_equation(self.h, self.u1, self.u0, self.p, sigma, params.rho, self.uh,
conditions.ocean_curr, params.rho_a, params.C_a, params.rho_w,
params.C_w, conditions.geo_wind, params.cor, self.timestep)
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
self.eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=self.uh, test=self.p)
self.bcs = DirichletBC(self.V, conditions.bc['u'], "on_boundary")
def get_trace_nullspace_vecs(forward, nullspace, V, V_d, TraceSpace):
"""Gets the nullspace vectors corresponding to the Schur complement
system for the multipliers.
:arg forward: A Slate expression denoting the forward elimination
operator.
:arg nullspace: A nullspace for the original mixed problem
:arg V: The original "unbroken" space.
:arg V_d: The broken space.
:arg TraceSpace: The space of approximate traces.
Returns: A list of vectors describing the nullspace of the multiplier
system
"""
from firedrake import project, assemble, Function
vecs = nullspace.getVecs()
tmp = Function(V)
tmp_b = Function(V_d)
tnsp_tmp = Function(TraceSpace)
forward_action = forward * tmp_b
new_vecs = []
for v in vecs:
with tmp.dat.vec as t:
v.copy(t)
project(tmp, tmp_b)
assemble(forward_action, tensor=tnsp_tmp)
with tnsp_tmp.dat.vec_ro as v:
new_vecs.append(v.copy())
return new_vecs
def remove_initial_w(u, Vv):
bc = DirichletBC(u.function_space()[0], 0.0, "bottom")
bc.apply(u)
uv = Function(Vv).project(u)
ustar = Function(u.function_space()).project(uv)
uin = Function(u.function_space()).assign(u - ustar)
u.assign(uin)
def test_1D_recovery(geometry, mesh, expr):
# horizontal base spaces
cell = mesh.ufl_cell().cellname()
# DG1
DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced")
DG1 = FunctionSpace(mesh, DG1_elt)
# spaces
DG0 = FunctionSpace(mesh, "DG", 0)
CG1 = FunctionSpace(mesh, "CG", 1)
# our actual theta and rho and v
rho_CG1_true = Function(CG1).interpolate(expr)
# make the initial fields by projecting expressions into the lowest order spaces
rho_DG0 = Function(DG0).interpolate(expr)
rho_CG1 = Function(CG1)
# make the recoverers and do the recovery
rho_recoverer = Recoverer(rho_DG0, rho_CG1, VDG=DG1, boundary_method=Boundary_Method.dynamics)
rho_recoverer.project()
rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true)
tolerance = 1e-7
error_message = ("""
Incorrect recovery for {variable} with {boundary} boundary method
on {geometry} 1D domain
""")
assert rho_diff < tolerance, error_message.format(variable='rho', boundary='dynamics', geometry=geometry)
def test_average(geometry, mesh):
vec_CG1 = VectorFunctionSpace(mesh, "CG", 1)
# We will fill DG_field with values, and average them to CG_field
weights = Function(vec_CG1)
true_values = Function(vec_CG1)
true_values = setup_values(geometry, true_values)
kernel = kernels.AverageWeightings(vec_CG1)
kernel.apply(weights)
tolerance = 1e-12
if geometry == "1D":
for i, (weight, true) in enumerate(
zip(weights.dat.data[:], true_values.dat.data[:])):
assert abs(weight - true
) < tolerance, "Weight not correct at position %i" % i
elif geometry == "2D":
for i, (weight, true) in enumerate(
zip(weights.dat.data[:], true_values.dat.data[:])):
for weight_j, true_j in zip(weight, true):
assert abs(
weight_j - true_j
) < tolerance, "Weight not correct at position %i" % i
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
def compressible_eady_initial_v(state, theta0, rho0, v):
f = state.parameters.f
cp = state.parameters.cp
# exner function
Vr = rho0.function_space()
Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0))
# get Pi gradient
Vu = state.spaces("HDiv")
g = TrialFunction(Vu)
wg = TestFunction(Vu)
n = FacetNormal(state.mesh)
a = inner(wg, g)*dx
L = -div(wg)*Pi*dx + inner(wg, n)*Pi*ds_tb
pgrad = Function(Vu)
solve(a == L, pgrad)
# get initial v
m = TrialFunction(Vr)
phi = TestFunction(Vr)
a = phi*f*m*dx
L = phi*cp*theta0*pgrad[0]*dx
solve(a == L, v)
return v
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 setup_values(geometry, DG0_field, weights):
x = SpatialCoordinate(weights.function_space().mesh())
coords_CG1 = Function(weights.function_space()).interpolate(x)
coords_DG0 = Function(DG0_field.function_space()).interpolate(x)
if geometry == "1D":
# Let us focus on the point at x = 1.0
# The test is if at CG_field[CG_index] we get the average of the corresponding DG_field values
CG_index = set_val_at_point(coords_CG1, 1.0)
set_val_at_point(coords_DG0, 0.5, DG0_field, 6.0)
set_val_at_point(coords_DG0, 1.5, DG0_field, -10.0)
set_val_at_point(coords_CG1, 1.0, weights, 2.0)
true_values = 0.5 * (6.0 - 10.0)
elif geometry == "2D":
# Let us focus on the point at (1,1)
# The test is if at CG_field[CG_index] we get the average of the corresponding DG_field values
# We do it for both components of the vector field
CG_index = set_val_at_point(coords_CG1, [1.0, 1.0])
set_val_at_point(coords_CG1, [1.0, 1.0], weights, [4.0, 4.0])
set_val_at_point(coords_DG0, [0.5, 0.5], DG0_field, [6.0, -3.0])
set_val_at_point(coords_DG0, [1.5, 0.5], DG0_field, [-7.0, -6.0])
set_val_at_point(coords_DG0, [0.5, 1.5], DG0_field, [0.0, 3.0])
set_val_at_point(coords_DG0, [1.5, 1.5], DG0_field, [-9.0, -1.0])
true_values = [
0.25 * (6.0 - 7.0 + 0.0 - 9.0), 0.25 * (-3.0 - 6.0 + 3.0 - 1.0)
]
return DG0_field, weights, true_values, CG_index
def setup(self, state):
if not self._initialised:
# check geometric dimension is 3D
if state.mesh.geometric_dimension() != 3:
raise ValueError(
'Spherical components only work when the geometric dimension is 3!'
)
space = FunctionSpace(state.mesh, "CG", 1)
super().setup(state, space=space)
V = VectorFunctionSpace(state.mesh, "CG", 1)
self.x, self.y, self.z = SpatialCoordinate(state.mesh)
self.x_hat = Function(V).interpolate(
Constant(as_vector([1.0, 0.0, 0.0])))
self.y_hat = Function(V).interpolate(
Constant(as_vector([0.0, 1.0, 0.0])))
self.z_hat = Function(V).interpolate(
Constant(as_vector([0.0, 0.0, 1.0])))
self.R = sqrt(self.x**2 + self.y**2) # distance from z axis
self.r = sqrt(self.x**2 + self.y**2 +
self.z**2) # distance from origin
self.f = state.fields(self.fname)
if np.prod(self.f.ufl_shape) != 3:
raise ValueError(
'Components can only be found of a vector function space in 3D.'
)
def _ad_convert_type(self, value, options=None):
from firedrake import Function, TrialFunction, TestFunction, assemble
options = {} if options is None else options
riesz_representation = options.get("riesz_representation", "l2")
if riesz_representation == "l2":
return Function(self.function_space(), val=value)
elif riesz_representation == "L2":
ret = Function(self.function_space())
u = TrialFunction(self.function_space())
v = TestFunction(self.function_space())
M = assemble(firedrake.inner(u, v) * firedrake.dx)
firedrake.solve(M, ret, value)
return ret
elif riesz_representation == "H1":
ret = Function(self.function_space())
u = TrialFunction(self.function_space())
v = TestFunction(self.function_space())
M = assemble(
firedrake.inner(u, v) * firedrake.dx +
firedrake.inner(firedrake.grad(u), firedrake.grad(v)) *
firedrake.dx)
firedrake.solve(M, ret, value)
return ret
elif callable(riesz_representation):
return riesz_representation(value)
else:
raise NotImplementedError("Unknown Riesz representation %s" %
riesz_representation)
def __init__(self, A, pyt_grad_op, pyt_op, actx, kappa):
"""
:arg kappa: The wave number
"""
self.actx = actx
self.k = kappa
self.pyt_op = pyt_op
self.pyt_grad_op = pyt_grad_op
self.A = A
self.meshmode_src_connection = meshmode_src_connection
# {{{ Create some functions needed for multing
self.x_fntn = Function(fspace)
# CG
self.potential_int = Function(fspace)
self.potential_int.dat.data[:] = 0.0
self.grad_potential_int = Function(vfspace)
self.grad_potential_int.dat.data[:] = 0.0
self.pyt_result = Function(fspace)
self.n = FacetNormal(mesh)
self.v = TestFunction(fspace)
# some meshmode ones
self.x_mm_fntn = self.meshmode_src_connection.discr.empty(
self.actx, dtype='c')
def test_physics_recovery_kernels(boundary):
m = IntervalMesh(3, 3)
mesh = ExtrudedMesh(m, layers=3, layer_height=1.0)
cell = m.ufl_cell().cellname()
hori_elt = FiniteElement("DG", cell, 0)
vert_elt = FiniteElement("CG", interval, 1)
theta_elt = TensorProductElement(hori_elt, vert_elt)
Vt = FunctionSpace(mesh, theta_elt)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_elt))
initial_field = Function(Vt)
true_field = Function(Vt_brok)
new_field = Function(Vt_brok)
initial_field, true_field = setup_values(boundary, initial_field,
true_field)
kernel = kernels.PhysicsRecoveryTop(
) if boundary == "top" else kernels.PhysicsRecoveryBottom()
kernel.apply(new_field, initial_field)
tolerance = 1e-12
index = 11 if boundary == "top" else 6
assert abs(true_field.dat.data[index] - new_field.dat.data[index]) < tolerance, \
"Value at %s from physics recovery is not correct" % boundary
def test_upper_and_lower_bound_vector(self):
"""
Test that setting lower and upper bounds work as expected for vector
functions.
"""
x = SpatialCoordinate(self.mesh)
self.prob.to_bound = as_tensor([
Function(self.V).interpolate(x[0]),
Function(self.V).interpolate(1 - x[0])
])
self.prob.test_func = self.prob.to_bound + as_tensor([1, 1])
self.assertAlmostEqual(self.prob.test_func[0]([0.1, 0.1, 0.1]), 1.1)
self.assertAlmostEqual(self.prob.test_func[1]([0.1, 0.1, 0.1]), 1.9)
self.assertAlmostEqual(self.prob.test_func[0]([0.6, 0.1, 0.1]), 1.6)
self.assertAlmostEqual(self.prob.test_func[1]([0.6, 0.1, 0.1]), 1.4)
self.prob.bound('to_bound', lower=0.45, upper=0.55)
self.assertAlmostEqual(self.prob.test_func[0]([0.1, 0.1, 0.1]), 1.45)
self.assertAlmostEqual(self.prob.test_func[1]([0.1, 0.1, 0.1]), 1.55)
self.assertAlmostEqual(self.prob.test_func[0]([0.6, 0.1, 0.1]), 1.55)
self.assertAlmostEqual(self.prob.test_func[1]([0.6, 0.1, 0.1]), 1.45)
def PeriodicIntervalMesh(ncells, length):
"""Generate a periodic mesh of an interval.
:arg ncells: The number of cells over the interval.
:arg length: The length the interval."""
m = CircleManifoldMesh(ncells)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=1)
old_coordinates = Function(m.coordinates)
new_coordinates = Function(coord_fs)
periodic_kernel = """double Y,pi;
Y = 0.5*(old_coords[0][1]-old_coords[1][1]);
pi=3.141592653589793;
for(int i=0;i<2;i++){
new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2;
if(new_coords[i][0]<0.) new_coords[i][0] += 1;
if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0;
new_coords[i][0] *= L;
}"""
periodic_kernel = periodic_kernel.replace('L', str(length))
par_loop(
periodic_kernel, dx, {
"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ)
})
m.coordinates = new_coordinates
return m
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 setup_solver(self, up_init=None):
""" Setup the solvers
"""
self.up0 = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up0, "/up")
chk_in.close()
self.u0, self.p0 = split(self.up0)
self.up = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up, "/up")
chk_in.close()
self.u1, self.p1 = split(self.up)
self.up.sub(0).rename("velocity")
self.up.sub(1).rename("pressure")
v, q = TestFunctions(self.W)
h = CellVolume(self.mesh)
u_norm = sqrt(dot(self.u0, self.u0))
if self.has_nullspace:
nullspace = MixedVectorSpaceBasis(
self.W,
[self.W.sub(0), VectorSpaceBasis(constant=True)])
else:
nullspace = None
tau = ((2.0 / self.dt)**2 + (2.0 * u_norm / h)**2 +
(4.0 * self.nu / h**2)**2)**(-0.5)
# temporal discretization
F = (1.0 / self.dt) * inner(self.u1 - self.u0, v) * dx
# weak form
F += (+inner(dot(self.u0, nabla_grad(self.u1)), v) * dx +
self.nu * inner(grad(self.u1), grad(v)) * dx -
(1.0 / self.rho) * self.p1 * div(v) * dx +
div(self.u1) * q * dx - inner(self.forcing, v) * dx)
# residual form
R = (+(1.0 / self.dt) * (self.u1 - self.u0) +
dot(self.u0, nabla_grad(self.u1)) - self.nu * div(grad(self.u1)) +
(1.0 / self.rho) * grad(self.p1) - self.forcing)
# GLS
F += tau * inner(
+dot(self.u0, nabla_grad(v)) - self.nu * div(grad(v)) +
(1.0 / self.rho) * grad(q), R) * dx
self.problem = NonlinearVariationalProblem(F, self.up, self.bcs)
self.solver = NonlinearVariationalSolver(
self.problem,
nullspace=nullspace,
solver_parameters=self.solver_parameters)
def setup(self, equation, uadv, *active_labels):
super().setup(equation, uadv, *active_labels)
self.k1 = Function(self.fs)
self.k2 = Function(self.fs)
self.k3 = Function(self.fs)
self.k4 = Function(self.fs)
def get_latlon_mesh(mesh):
"""Build 2D projected mesh of spherical mesh"""
crds_orig = mesh.coordinates
mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1)
crds_dg = Function(mesh_dg_fs)
crds_latlon = Function(mesh_dg_fs)
par_loop(
"""
for (int i=0; i<3; i++) {
for (int j=0; j<3; j++) {
dg[i][j] = cg[i][j];
}
}
""", dx, {
'dg': (crds_dg, WRITE),
'cg': (crds_orig, READ)
})
# lat-lon 'x' = atan2(y, x)
crds_latlon.dat.data[:, 0] = arctan2(crds_dg.dat.data[:, 1],
crds_dg.dat.data[:, 0])
# lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
crds_latlon.dat.data[:, 1] = (arcsin(
crds_dg.dat.data[:, 2] /
np_sqrt(crds_dg.dat.data[:, 0]**2 + crds_dg.dat.data[:, 1]**2 +
crds_dg.dat.data[:, 2]**2)))
crds_latlon.dat.data[:, 2] = 0.0
kernel = op2.Kernel(
"""
#define PI 3.141592653589793
#define TWO_PI 6.283185307179586
void splat_coords(double **coords) {
double diff0 = (coords[0][0] - coords[1][0]);
double diff1 = (coords[0][0] - coords[2][0]);
double diff2 = (coords[1][0] - coords[2][0]);
if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) {
const int sign0 = coords[0][0] < 0 ? -1 : 1;
const int sign1 = coords[1][0] < 0 ? -1 : 1;
const int sign2 = coords[2][0] < 0 ? -1 : 1;
if (sign0 < 0) {
coords[0][0] += TWO_PI;
}
if (sign1 < 0) {
coords[1][0] += TWO_PI;
}
if (sign2 < 0) {
coords[2][0] += TWO_PI;
}
}
}
""", "splat_coords")
op2.par_loop(kernel, crds_latlon.cell_set,
crds_latlon.dat(op2.RW, crds_latlon.cell_node_map()))
return Mesh(crds_latlon)
def remove_initial_w(u):
Vu = u.function_space()
Vv = FunctionSpace(Vu._ufl_domain, Vu.ufl_element()._elements[-1])
bc = DirichletBC(Vu[0], 0.0, "bottom")
bc.apply(u)
uv = Function(Vv).project(u)
ustar = Function(u.function_space()).project(uv)
uin = Function(u.function_space()).assign(u - ustar)
u.assign(uin)
def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None):
# get F
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
v = TrialFunction(Vv)
w = TestFunction(Vv)
if top:
bstring = "top"
else:
bstring = "bottom"
bcs = [DirichletBC(Vv, 0.0, bstring)]
a = inner(w, v) * dx
L = inner(state.k, w) * b0 * dx
F = Function(Vv)
solve(a == L, F, bcs=bcs)
# define mixed function space
VDG = state.spaces("DG")
WV = (Vv) * (VDG)
# get pprime
v, pprime = TrialFunctions(WV)
w, phi = TestFunctions(WV)
bcs = [DirichletBC(WV[0], zero(), bstring)]
a = (inner(w, v) + div(w) * pprime + div(v) * phi) * dx
L = phi * div(F) * dx
w1 = Function(WV)
if params is None:
params = {
'ksp_type': 'gmres',
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'pc_fieldsplit_schur_fact_type': 'full',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_1_ksp_type': 'preonly',
'fieldsplit_1_pc_type': 'gamg',
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 4,
'ksp_atol': 1.e-08,
'ksp_rtol': 1.e-08
}
solve(a == L, w1, bcs=bcs, solver_parameters=params)
v, pprime = w1.split()
p0.project(pprime)
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 create_sc_nullspace(P, V, V_facet, comm):
"""Gets the nullspace vectors corresponding to the Schur complement
system.
:arg P: The H1 operator from the ImplicitMatrixContext.
:arg V: The H1 finite element space.
:arg V_facet: The finite element space of H1 basis functions
restricted to the mesh skeleton.
Returns: A nullspace (if there is one) for the Schur-complement system.
"""
from firedrake import Function
nullspace = P.getNullSpace()
if nullspace.handle == 0:
# No nullspace
return None
vecs = nullspace.getVecs()
tmp = Function(V)
scsp_tmp = Function(V_facet)
new_vecs = []
# Transfer the trace bit (the nullspace vector restricted
# to facet nodes is the nullspace for the condensed system)
kernel = """
for (int i=0; i<%(d)d; ++i){
for (int j=0; j<%(s)d; ++j){
x_facet[i*%(s)d + j] = x_h[i*%(s)d + j];
}
}""" % {
"d": V_facet.finat_element.space_dimension(),
"s": np.prod(V_facet.shape)
}
for v in vecs:
with tmp.dat.vec_wo as t:
v.copy(t)
par_loop(kernel, ufl.dx, {
"x_facet": (scsp_tmp, WRITE),
"x_h": (tmp, READ)
})
# Map vecs to the facet space
with scsp_tmp.dat.vec_ro as v:
new_vecs.append(v.copy())
# Normalize
for v in new_vecs:
v.normalize()
sc_nullspace = PETSc.NullSpace().create(vectors=new_vecs, comm=comm)
return sc_nullspace
def test_checkpointing(tmpdir):
dirname_1 = str(tmpdir) + '/checkpointing_1'
dirname_2 = str(tmpdir) + '/checkpointing_2'
state_1, stepper_1, dt = setup_checkpointing(dirname_1)
state_2, stepper_2, dt = setup_checkpointing(dirname_2)
# ------------------------------------------------------------------------ #
# Run for 4 time steps and store values
# ------------------------------------------------------------------------ #
initialise_fields(state_1)
stepper_1.run(t=0.0, tmax=4 * dt)
# ------------------------------------------------------------------------ #
# Start again, run for 2 time steps, checkpoint and then run for 2 more
# ------------------------------------------------------------------------ #
initialise_fields(state_2)
stepper_2.run(t=0.0, tmax=2 * dt)
# Wipe fields, then pickup
state_2.fields('u').project(as_vector([-10.0, 0.0]))
state_2.fields('rho').interpolate(Constant(0.0))
state_2.fields('theta').interpolate(Constant(0.0))
stepper_2.run(t=2 * dt, tmax=4 * dt, pickup=True)
# ------------------------------------------------------------------------ #
# Compare fields against saved values for run without checkpointing
# ------------------------------------------------------------------------ #
# Pick up from both stored checkpoint files
# This is the best way to compare fields from different meshes
for field_name in ['u', 'rho', 'theta']:
with DumbCheckpoint(dirname_1 + '/chkpt', mode=FILE_READ) as chkpt:
field_1 = Function(state_1.fields(field_name).function_space(),
name=field_name)
chkpt.load(field_1)
# These are preserved in the comments for when we can use CheckpointFile
# mesh = chkpt.load_mesh(name='firedrake_default_extruded')
# field_1 = chkpt.load_function(mesh, name=field_name)
with DumbCheckpoint(dirname_2 + '/chkpt', mode=FILE_READ) as chkpt:
field_2 = Function(state_1.fields(field_name).function_space(),
name=field_name)
chkpt.load(field_2)
# These are preserved in the comments for when we can use CheckpointFile
# field_2 = chkpt.load_function(mesh, name=field_name)
error = norm(field_1 - field_2)
assert error < 1e-15, \
f'Checkpointed field {field_name} is not equal to non-checkpointed field'
def extend_function_to_3d(func, mesh_extruded):
"""
Returns a 3D view of a 2D :class:`Function` on the extruded domain.
The 3D function resides in V x R function space, where V is the function
space of the source function. The 3D function shares the data of the 2D
function.
"""
fs = func.function_space()
# assert fs.mesh().geometric_dimension() == 2, 'Function must be in 2D space'
ufl_elem = fs.ufl_element()
family = ufl_elem.family()
degree = ufl_elem.degree()
name = func.name()
if isinstance(ufl_elem, ufl.VectorElement):
# vector function space
fs_extended = get_functionspace(mesh_extruded,
family,
degree,
'R',
0,
dim=2,
vector=True)
else:
fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0)
func_extended = Function(fs_extended, name=name, val=func.dat._data)
func_extended.source = func
return func_extended
def setup(self, state):
if not self._initialised:
space = state.fields("theta").function_space()
broken_space = FunctionSpace(state.mesh, BrokenElement(space.ufl_element()))
boundary_method = Boundary_Method.physics if (state.vertical_degree == 0 and state.horizontal_degree == 0) else None
super().setup(state, space=space)
# now let's attach all of our fields
self.u = state.fields("u")
self.rho = state.fields("rho")
self.theta = state.fields("theta")
self.rho_averaged = Function(space)
self.recoverer = Recoverer(self.rho, self.rho_averaged, VDG=broken_space, boundary_method=boundary_method)
try:
self.r_v = state.fields("water_v")
except NotImplementedError:
self.r_v = Constant(0.0)
try:
self.r_c = state.fields("water_c")
except NotImplementedError:
self.r_c = Constant(0.0)
try:
self.rain = state.fields("rain")
except NotImplementedError:
self.rain = Constant(0.0)
# now let's store the most common expressions
self.pi = thermodynamics.pi(state.parameters, self.rho_averaged, self.theta)
self.T = thermodynamics.T(state.parameters, self.theta, self.pi, r_v=self.r_v)
self.p = thermodynamics.p(state.parameters, self.pi)
self.r_l = self.r_c + self.rain
self.r_t = self.r_v + self.r_c + self.rain
def _build_forcing_solvers(self):
super(EadyForcing, self)._build_forcing_solvers()
# b_forcing
dbdy = self.state.parameters.dbdy
Vb = self.state.spaces("HDiv_v")
F = TrialFunction(Vb)
gamma = TestFunction(Vb)
self.bF = Function(Vb)
u0, _, b0 = split(self.x0)
a = gamma * F * dx
L = -self.scaling * gamma * (dbdy *
inner(u0, as_vector([0., 1., 0.]))) * dx
b_forcing_problem = LinearVariationalProblem(a, L, self.bF)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.b_forcing_solver = LinearVariationalSolver(
b_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="BForcingSolver")
