def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG',
1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
self._min_max_loop = """
for(int i = 0; i < maxq.dofs; i++) {
maxq[i][0] = fmax(maxq[i][0],q[0][0]);
minq[i][0] = fmin(minq[i][0],q[0][0]);
}
"""
# Perform limiting loop
self._limit_kernel = """
def __init__(self, equation):
"""
Initialise limiter
:param space : equation, as we need the broken space attached to it
"""
self.Vt = equation.space
# check this is the right space, only currently working for 2D extruded mesh
if self.Vt.extruded and self.Vt.mesh().topological_dimension() == 2:
# check that horizontal degree is 1 and vertical degree is 2
if self.Vt.ufl_element().degree()[0] is not 1 or \
self.Vt.ufl_element().degree()[1] is not 2:
raise ValueError('This is not the right limiter for this space.')
# check that continuity of the spaces is correct
# this will fail if the space does not use broken elements
if self.Vt.ufl_element()._element.sobolev_space()[0].name is not 'L2' or \
self.Vt.ufl_element()._element.sobolev_space()[1].name is not 'H1':
raise ValueError('This is not the right limiter for this space.')
else:
logger.warning('This limiter may not work for the space you are using.')
self.Q1DG = FunctionSpace(self.Vt.mesh(), 'DG', 1) # space with only vertex DOFs
self.vertex_limiter = VertexBasedLimiter(self.Q1DG)
self.theta_hat = Function(self.Q1DG) # theta function with only vertex DOFs
self.w = Function(self.Vt)
self.result = Function(self.Vt)
par_loop(_weight_kernel, dx, {"weight": (self.w, INC)})
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG', 1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
self._min_max_loop = """
for(int i = 0; i < maxq.dofs; i++) {
maxq[i][0] = fmax(maxq[i][0],q[0][0]);
minq[i][0] = fmin(minq[i][0],q[0][0]);
}
"""
# Perform limiting loop
self._limit_kernel = """
def test_ForcingCovariance_init(fs):
"test init method of ForcingCovariance"
# note: tests only handle case of a single process
# need to think about testing more complicated cases
# simple example in 1D
sigma = np.log(1.)
l = np.log(0.1)
fc = ForcingCovariance(fs, sigma, l)
n = Function(fs).vector().size()
n_local = Function(fs).vector().local_size()
M = PETSc.Mat().create()
M.setSizes(((n_local, -1), (n_local, -1)))
M.setFromOptions()
M.setUp()
start, end = M.getOwnershipRange()
assert fc.nx == n
assert fc.nx_local == n_local
assert fc.function_space == fs
assert_allclose(fc.sigma, sigma)
assert_allclose(fc.l, l)
assert_allclose(fc.cutoff, 1.e-3)
assert_allclose(fc.regularization, 1.e-8)
assert fc.local_startind == start
assert fc.local_endind == end
assert not fc.is_assembled
def predict_mean(self, coords, scale_mean=True):
r"""
Compute the predictive mean
This method computes the predictive mean of data values at unmeasured locations. It returns
the vector of predicted sensor values on the root process as numpy array. It requires only a
small overhead above the computational work of finding the posterior mean (i.e. you get
the mean value at new sensor locations for "free" once you have solved the posterior).
The optional ``scale_mean`` argument determines if the solution is to be re-scaled
by the model discrepancy scaling factor. This value is by default ``True``.
To re-scale to match the FEM solution, pass ``scale_mean=False``.
:param coords: Spatial coordinates at which the mean will be predicted. Must be a
2D Numpy array (or a 1D array, which will assume the second axis has length
1)
:type coords: ndarray
:param scale_mean: Boolean indicating if the mean should be scaled by the model
discrepancy scaling factor. Optional, default is ``True``
:type scale_mean: bool
:returns: FEM prediction at specified sensor locations as a numpy array on the root process.
All other processes will have a numpy array of length 0.
:rtype: ndarray
"""
if not isinstance(bool(scale_mean), bool):
raise TypeError("scale_mean argument must be boolean-like")
coords = np.array(coords, dtype=np.float64)
if coords.ndim == 1:
coords = np.reshape(coords, (-1, 1))
assert coords.ndim == 2, "coords must be a 1d or 2d array"
assert coords.shape[1] == self.data.get_n_dim(
), "axis 1 of coords must be the same length as the FEM dimension"
if self.Cu is None:
self.solve_prior()
if self.params is None:
raise ValueError("must set parameter values to make predictions")
rho = np.exp(self.params[0])
if scale_mean:
scalefact = rho
else:
scalefact = 1.
x = Function(self.G.function_space)
self.solve_posterior(x)
im = InterpolationMatrix(self.G.function_space, coords)
mu = scalefact * im.interp_mesh_to_data(x.vector())
im.destroy()
return mu
def test_InterpolationMatrix(fs, coords):
"test InterpolationMatrix with multiple processes"
n_proc = COMM_WORLD.size
nd = len(coords)
vec = Function(fs).vector()
im = InterpolationMatrix(fs, coords)
if COMM_WORLD.rank == 0:
gathered_sizes = (nd, nd)
else:
gathered_sizes = (0, 0)
assert im.n_data == nd
assert im.n_data_local == nd // n_proc
assert im.n_mesh == vec.size()
assert im.n_mesh_local == vec.local_size()
assert_allclose(im.coords, coords)
assert im.function_space == fs
assert im.meshspace_vector.size() == vec.size()
assert im.meshspace_vector.local_size() == vec.local_size()
assert im.dataspace_distrib.getSizes() == (nd // n_proc, nd)
assert im.dataspace_gathered.getSizes() == gathered_sizes
assert im.interp.getSizes() == ((vec.local_size(), vec.size()),
(nd // n_proc, nd))
assert not im.is_assembled
def solve_forcing_covariance(G, ls, rhs):
"""
Solve the forcing covariance part of the stat FEM
This function performs the basic solve needed to determine the
prior covariance matrix in the stat-fem method. Two solves of
the FEM are required, in addition to a sparse matrix multiplication.
The arguments provide the sparse Forcing Covariance matrix,
the Firedrake Linear Solver object representing the FEM, and
the RHS to be solved for this particular iteration.
Note that this solve function temporarily turns off the BCs for
the stiffness matrix. This is because Dirichlet BCs will enforce
strong boundary conditions on the FEM solves, which is not desired
here.
:param G: Forcing covariance matrix to be used in the solve.
:type G: ForcingCovariance
:param ls: Firedrake Linear Solver to be used in the solve.
:type ls: Firedrake LinearSolver
:param rhs: RHS vector to be used in the solve
:type rhs: Firedrake Vector
:returns: Solution to :math:`A^{-1}GA^{-1}b` where :math:`A` is the FEM
stiffness matrix and :math:`b` is the RHS vector.
:rtype: Firedrake Vector
"""
if not isinstance(G, ForcingCovariance):
raise TypeError("G must be a ForcingCovariance object")
if not isinstance(ls, LinearSolver):
raise TypeError("ls must be a firedrake LinearSolver")
if not isinstance(rhs, Vector):
raise TypeError("rhs must be a firedrake vector")
# turn off BC application temporarily
bcs = ls.A.bcs
ls.A.bcs = None
rhs_working = rhs.copy()
x = Function(G.function_space).vector()
ls.solve(x, rhs_working)
G.mult(x, rhs_working)
ls.solve(x, rhs_working)
# turn BCs back on
ls.A.bcs = bcs
return x.copy()
def test_solve_forcing_covariance(comm, fs, A, b, fc, A_numpy, cov):
"test solve_forcing_covariance"
rhs = Function(fs).vector()
rhs.set_local(np.ones(fc.get_nx_local()))
ls = LinearSolver(A)
result = solve_forcing_covariance(fc, ls, rhs)
result_actual = result.gather()
result_expected = np.linalg.solve(A_numpy, np.ones(nx + 1))
result_expected = np.dot(cov, result_expected)
result_expected = np.linalg.solve(A_numpy, result_expected)
assert_allclose(result_expected, result_actual, atol=1.e-10)
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG', 1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.centroids_rhs = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
domain = "{[i]: 0 <= i < maxq.dofs}"
instructions = """
for i
maxq[i] = fmax(maxq[i], q[0])
minq[i] = fmin(minq[i], q[0])
end
"""
self._min_max_loop = (domain, instructions)
# Perform limiting loop
domain = "{[i, ii]: 0 <= i < q.dofs and 0 <= ii < q.dofs}"
instructions = """
<float64> alpha = 1
<float64> qavg = qbar[0, 0]
for i
<float64> _alpha1 = fmin(alpha, fmin(1, (qmax[i] - qavg)/(q[i] - qavg)))
<float64> _alpha2 = fmin(alpha, fmin(1, (qavg - qmin[i])/(qavg - q[i])))
alpha = if(q[i] > qavg, _alpha1, if(q[i] < qavg, _alpha2, alpha))
end
for ii
q[ii] = qavg + alpha * (q[ii] - qavg)
end
"""
self._limit_kernel = (domain, instructions)
def solve_prior(self):
r"""
Solve base (prior) FEM plus covariance interpolated to the data locations
This method solves the prior FEM and covariance interpolated to the sensor locations.
It does not require setting parameter values, as the model discrepancy does not
influence these results. The covariance is cached as it is expensive to compute
and is re-used in all other solves.
In addition to caching the results, the method returns solution as numpy arrays
on the root process (rank 0).
Note that unlike the solve done in the meshspace, this uses a return value rather than a
Firedrake/PETSc style interface to place the solution in a pre-allocated ``Function``.
This is because each process has a different array size, so would require correctly
pre-allocating arrays of different lengths on each process.
:returns: FEM prior mean and covariance (as a tuple of numpy arrays) on the root process.
Non-root processes return numpy arrays of shape ``(0,)`` (mean) and ``(0, 0)``
(covariance).
:rtype: tuple of ndarrays
"""
# form interpolated prior covariance across all ensemble processes
self.Cu = interp_covariance_to_data(self.im, self.G, self.solver,
self.im, self.ensemble_comm)
# solve base FEM (prior mean) and interpolate to data space on root
self.x = Function(self.G.function_space)
if self.ensemble_comm.rank == 0:
self.solver.solve(self.x, self.b)
self.mu = self.im.interp_mesh_to_data(self.x.vector())
else:
self.mu = np.zeros(0)
return self.mu, self.Cu
def test_ForcingCovariance_mult(fs, fc, cov):
"test the multiplication method of ForcingCovariance"
fc.assemble()
x = Function(fs).vector()
x.set_local(np.ones(x.local_size()))
y = Function(fs).vector()
fc.mult(x, y)
ygathered = y.gather()
assert_allclose(ygathered, np.dot(cov, np.ones(nx + 1)))
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
if utils.complex_mode:
raise ValueError(
"We haven't decided what limiting complex valued fields means. Please get in touch if you have need."
)
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG',
1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.centroids_rhs = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
domain = "{[i]: 0 <= i < maxq.dofs}"
instructions = """
for i
maxq[i] = fmax(maxq[i], q[0])
minq[i] = fmin(minq[i], q[0])
end
"""
self._min_max_loop = (domain, instructions)
# Perform limiting loop
domain = "{[i, ii]: 0 <= i < q.dofs and 0 <= ii < q.dofs}"
instructions = """
<float64> alpha = 1
<float64> qavg = qbar[0, 0]
for i
<float64> _alpha1 = fmin(alpha, fmin(1, (qmax[i] - qavg)/(q[i] - qavg)))
<float64> _alpha2 = fmin(alpha, fmin(1, (qavg - qmin[i])/(qavg - q[i])))
alpha = _alpha1 if q[i] > qavg else (_alpha2 if q[i] < qavg else alpha)
end
for ii
q[ii] = qavg + alpha * (q[ii] - qavg)
end
"""
self._limit_kernel = (domain, instructions)
def new_snes_ctx(pc, op, bcs, mat_type, fcp=None):
""" Create a new SNES contex for nested preconditioning
"""
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.function import Function
from firedrake.ufl_expr import action
from firedrake.dmhooks import get_appctx
from firedrake.solving_utils import _SNESContext
dm = pc.getDM()
old_appctx = get_appctx(dm).appctx
u = Function(op.arguments()[-1].function_space())
F = action(op, u)
nprob = NonlinearVariationalProblem(F,
u,
bcs=bcs,
J=op,
form_compiler_parameters=fcp)
nctx = _SNESContext(nprob, mat_type, mat_type, old_appctx)
return nctx
def func(self, *args, **kwargs):
if len(args) > 0 and isinstance(
args[0], FunctionSpace):
# Extract size from V
if 'size' in kwargs.keys():
raise TypeError(
"Cannot specify 'size' when generating a random function from 'V'"
)
V = args[0]
f = Function(V)
args = args[1:]
with f.dat.vec_wo as v:
kwargs['size'] = (v.local_size, )
v.array[:] = getattr(self, c_a)(*args,
**kwargs)
return f
else:
# forward to the original implementation
return getattr(super(_Wrapper, self),
c_a)(*args, **kwargs)
def test_solve_forcing_covariance_parallel(my_ensemble, comm, fs, A, b, fc,
A_numpy, cov):
"test that solve_forcing_covariance can be called independently from an ensemble process"
if my_ensemble.ensemble_comm.rank == 0:
rhs = Function(fs).vector()
rhs.set_local(np.ones(fc.get_nx_local()))
ls = LinearSolver(A)
result = solve_forcing_covariance(fc, ls, rhs)
result_actual = result.gather()
result_expected = np.linalg.solve(A_numpy, np.ones(nx + 1))
result_expected = np.dot(cov, result_expected)
result_expected = np.linalg.solve(A_numpy, result_expected)
assert_allclose(result_expected, result_actual, atol=1.e-10)
elif my_ensemble.ensemble_comm.rank == 1:
rhs = Function(fs).vector()
rhs.set_local(0.5 * np.ones(fc.get_nx_local()))
ls = LinearSolver(A)
result = solve_forcing_covariance(fc, ls, rhs)
result_actual = result.gather()
result_expected = np.linalg.solve(A_numpy, 0.5 * np.ones(nx + 1))
result_expected = np.dot(cov, result_expected)
result_expected = np.linalg.solve(A_numpy, result_expected)
assert_allclose(result_expected, result_actual, atol=1.e-10)
def __init__(self, spaces, name=None):
"""
:param spaces: a list (or tuple) of :class:`FunctionSpace`\s
The function space may be created as ::
V = MixedFunctionSpace(spaces)
``spaces`` may consist of multiple occurances of the same space: ::
P1 = FunctionSpace(mesh, "CG", 1)
P2v = VectorFunctionSpace(mesh, "Lagrange", 2)
ME = MixedFunctionSpace([P2v, P1, P1, P1])
"""
if self._initialized:
return
self._spaces = [
IndexedFunctionSpace(s, i, self)
for i, s in enumerate(flatten(spaces))
]
self._mesh = self._spaces[0].mesh()
self._ufl_element = ufl.MixedElement(
*[fs.ufl_element() for fs in self._spaces])
self.name = name or '_'.join(str(s.name) for s in self._spaces)
self.rank = 1
self._index = None
self._initialized = True
dm = PETSc.DMShell().create()
from firedrake.function import Function
with Function(self).dat.vec_ro as v:
dm.setGlobalVector(v.duplicate())
dm.setAttr('__fs__', weakref.ref(self))
dm.setCreateFieldDecomposition(self.create_field_decomp)
dm.setCreateSubDM(self.create_subdm)
self._dm = dm
self._ises = self.dof_dset.field_ises
self._subspaces = []
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG',
1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.centroids_rhs = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
domain = "{[i]: 0 <= i < maxq.dofs}"
instructions = """
for i
maxq[i] = fmax(maxq[i], q[0])
minq[i] = fmin(minq[i], q[0])
end
"""
self._min_max_loop = (domain, instructions)
# Perform limiting loop
domain = "{[i, ii]: 0 <= i < q.dofs and 0 <= ii < q.dofs}"
instructions = """
<float64> alpha = 1
<float64> qavg = qbar[0, 0]
for i
<float64> _alpha1 = fmin(alpha, fmin(1, (qmax[i] - qavg)/(q[i] - qavg)))
<float64> _alpha2 = fmin(alpha, fmin(1, (qavg - qmin[i])/(qavg - q[i])))
alpha = if(q[i] > qavg, _alpha1, if(q[i] < qavg, _alpha2, alpha))
end
for ii
q[ii] = qavg + alpha * (q[ii] - qavg)
end
"""
self._limit_kernel = (domain, instructions)
def test_InterpolationMatrix_interp_mesh_to_data(fs, coords, meshcoords):
"test method to interpolate from distributed mesh to data gathered at root"
# simple 1D test
nd = len(coords)
im = InterpolationMatrix(fs, coords)
im.assemble()
input_ordered = np.array([3., 2., 7., 4., 0., 0., 2., 1., 1., 1., 5.])
f = Function(fs).vector()
meshcoords_ordered = np.linspace(0., 1., 11)
with f.dat.vec as vec:
imin, imax = vec.getOwnershipRange()
for i in range(imin, imax):
vec.setValue(
i,
input_ordered[np.where(meshcoords_ordered == meshcoords[i])])
if COMM_WORLD.rank == 0:
expected = np.array([1., 0., 5.5, 3.25])
else:
expected = np.zeros(0)
out = im.interp_mesh_to_data(f)
assert_allclose(out, expected, atol=1.e-10)
# failure due to bad input sizes
mesh2 = UnitIntervalMesh(12)
V2 = FunctionSpace(mesh2, "CG", 1)
f2 = Function(V2).vector()
f2.set_local(np.ones(f2.local_size()))
with pytest.raises(AssertionError):
im.interp_mesh_to_data(f2)
im.destroy()
class ThetaLimiter(object):
"""
A vertex based limiter for fields in the DG1xCG2 space,
i.e. temperature variables. This acts like the vertex-based
limiter implemented in Firedrake, but in addition corrects
the central nodes to prevent new maxima or minima forming.
"""
def __init__(self, equation):
"""
Initialise limiter
:param space : equation, as we need the broken space attached to it
"""
self.Vt = equation.space
# check this is the right space, only currently working for 2D extruded mesh
if self.Vt.extruded and self.Vt.mesh().topological_dimension() == 2:
# check that horizontal degree is 1 and vertical degree is 2
if self.Vt.ufl_element().degree()[0] is not 1 or \
self.Vt.ufl_element().degree()[1] is not 2:
raise ValueError('This is not the right limiter for this space.')
# check that continuity of the spaces is correct
# this will fail if the space does not use broken elements
if self.Vt.ufl_element()._element.sobolev_space()[0].name is not 'L2' or \
self.Vt.ufl_element()._element.sobolev_space()[1].name is not 'H1':
raise ValueError('This is not the right limiter for this space.')
else:
logger.warning('This limiter may not work for the space you are using.')
self.Q1DG = FunctionSpace(self.Vt.mesh(), 'DG', 1) # space with only vertex DOFs
self.vertex_limiter = VertexBasedLimiter(self.Q1DG)
self.theta_hat = Function(self.Q1DG) # theta function with only vertex DOFs
self.w = Function(self.Vt)
self.result = Function(self.Vt)
par_loop(_weight_kernel, dx, {"weight": (self.w, INC)})
def copy_vertex_values(self, field):
"""
Copies the vertex values from temperature space to
Q1DG space which only has vertices.
"""
par_loop(_copy_into_Q1DG_loop, dx,
{"theta": (field, READ),
"theta_hat": (self.theta_hat, RW)})
def copy_vertex_values_back(self, field):
"""
Copies the vertex values back from the Q1DG space to
the original temperature space.
"""
par_loop(_copy_from_Q1DG_loop, dx,
{"theta": (field, RW),
"theta_hat": (self.theta_hat, READ)})
def check_midpoint_values(self, field):
"""
Checks the midpoint field values are less than the maximum
and more than the minimum values. Amends them to the average
if they are not.
"""
par_loop(_check_midpoint_values_loop, dx,
{"theta": (field, RW)})
def remap_to_embedded_space(self, field):
"""
Remap from DG space to embedded DG space.
"""
self.result.assign(0.)
par_loop(_average_kernel, dx, {"vrec": (self.result, INC),
"v_b": (field, READ),
"weight": (self.w, READ)})
field.assign(self.result)
def apply(self, field):
"""
The application of the limiter to the theta-space field.
"""
assert field.function_space() == self.Vt, \
'Given field does not belong to this objects function space'
self.copy_vertex_values(field)
self.vertex_limiter.apply(self.theta_hat)
self.copy_vertex_values_back(field)
self.check_midpoint_values(field)
self.remap_to_embedded_space(field)
def test_ForcingCovariance_get_nx_local(fs, fc):
"test the get_nx_local method of ForcingCovariance"
n_local = Function(fs).vector().local_size()
assert fc.get_nx_local() == n_local
class VertexBasedLimiter(Limiter):
"""
A vertex based limiter for P1DG fields.
This limiter implements the vertex-based limiting scheme described in
Dmitri Kuzmin, "A vertex-based hierarchical slope limiter for p-adaptive
discontinuous Galerkin methods". J. Comp. Appl. Maths (2010)
http://dx.doi.org/10.1016/j.cam.2009.05.028
"""
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG', 1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.centroids_rhs = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
domain = "{[i]: 0 <= i < maxq.dofs}"
instructions = """
for i
maxq[i] = fmax(maxq[i], q[0])
minq[i] = fmin(minq[i], q[0])
end
"""
self._min_max_loop = (domain, instructions)
# Perform limiting loop
domain = "{[i, ii]: 0 <= i < q.dofs and 0 <= ii < q.dofs}"
instructions = """
<float64> alpha = 1
<float64> qavg = qbar[0, 0]
for i
<float64> _alpha1 = fmin(alpha, fmin(1, (qmax[i] - qavg)/(q[i] - qavg)))
<float64> _alpha2 = fmin(alpha, fmin(1, (qavg - qmin[i])/(qavg - q[i])))
alpha = if(q[i] > qavg, _alpha1, if(q[i] < qavg, _alpha2, alpha))
end
for ii
q[ii] = qavg + alpha * (q[ii] - qavg)
end
"""
self._limit_kernel = (domain, instructions)
def _construct_centroid_solver(self):
"""
Constructs a linear problem for computing the centroids
:return: LinearSolver instance
"""
u = TrialFunction(self.P0)
v = TestFunction(self.P0)
a = assemble(u * v * dx)
return LinearSolver(a, solver_parameters={'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'})
def _update_centroids(self, field):
"""
Update centroid values
"""
assemble(TestFunction(self.P0) * field * dx, tensor=self.centroids_rhs)
self.centroid_solver.solve(self.centroids, self.centroids_rhs)
def compute_bounds(self, field):
"""
Only computes min and max bounds of neighbouring cells
"""
self._update_centroids(field)
self.max_field.assign(-1.0e10) # small number
self.min_field.assign(1.0e10) # big number
par_loop(self._min_max_loop,
dx,
{"maxq": (self.max_field, MAX),
"minq": (self.min_field, MIN),
"q": (self.centroids, READ)},
is_loopy_kernel=True)
def apply_limiter(self, field):
"""
Only applies limiting loop on the given field
"""
par_loop(self._limit_kernel, dx,
{"qbar": (self.centroids, READ),
"q": (field, RW),
"qmax": (self.max_field, READ),
"qmin": (self.min_field, READ)},
is_loopy_kernel=True)
def apply(self, field):
"""
Re-computes centroids and applies limiter to given field
"""
assert field.function_space() == self.P1DG, \
'Given field does not belong to this objects function space'
self.compute_bounds(field)
self.apply_limiter(field)
def from_meshmode(self, mm_field, out=None):
r"""
Transport meshmode field from :attr:`discr` into an
appropriate firedrake function space.
If *out* is *None*, values at any firedrake
nodes associated to NO meshmode nodes are zeroed out.
If *out* is supplied, values at nodes associated to NO meshmode nodes
are not modified.
:arg mm_field: Either
* A :class:`~meshmode.dof_array.DOFArray` representing
a field of shape *tuple()* on :attr:`discr`
* A :class:`numpy.ndarray` of dtype "object" with
entries of class :class:`~meshmode.dof_array.DOFArray`
representing a field of shape *mm_field.shape*
on :attr:`discr`
See :class:`~meshmode.dof.DOFArray` for further requirements.
The :attr:`group_nr` entry of each
:class:`~meshmode.dof_array.DOFArray`
must be of shape *(nelements, nunit_dofs)* and
the *element_dtype* must match that used for
:class:`firedrake.function.Function`\ s
:arg out: If *None* then ignored, otherwise a
:class:`firedrake.function.Function`
of the right function space for the transported data
to be stored in. The shape of its function space must
match the shape of *mm_field*
:return: a :class:`firedrake.function.Function` holding the transported
data (*out*, if *out* was not *None*)
"""
# All firedrake functions are the same dtype
dtype = self.firedrake_fspace().mesh().coordinates.dat.data.dtype
self._validate_field(mm_field, "mm_field", dtype=dtype)
# get the shape of mm_field
from meshmode.dof_array import DOFArray
if not isinstance(mm_field, DOFArray):
fspace_shape = mm_field.shape
else:
fspace_shape = ()
# make sure out is a firedrake function in an appropriate
# function space
if out is not None:
self._validate_function(out, "out", fspace_shape, dtype)
else:
from firedrake.function import Function
# Translate shape so that don't always get a TensorFunctionSpace,
# but instead get FunctionSpace or VectorFunctionSpace when
# reasonable
shape = fspace_shape
if shape == ():
shape = None
elif len(shape) == 1:
shape = shape[0]
# make a function filled with zeros
out = Function(self.firedrake_fspace(shape))
out.dat.data[:] = 0.0
out_data = out.dat.data
# Handle firedrake dropping dimensions
if len(out.dat.data.shape) != 1 + len(fspace_shape):
shape = (out.dat.data.shape[0],) + fspace_shape
out_data = out_data.reshape(shape)
def resample_and_reorder(fd_data, dof_array):
# pull data into numpy
dof_np = dof_array.array_context.to_numpy(dof_array[self.group_nr])
# resample the data and store in firedrake ordering
# store resampled data in firedrake ordering
fd_data[self.mm2fd_node_mapping] = \
np.einsum("ij,kj->ik", dof_np, self._resampling_mat_mm2fd)
# If scalar, just reorder and resample out
if fspace_shape == ():
resample_and_reorder(out_data, mm_field)
else:
# otherwise, have to grab each dofarray and the corresponding
# data from *function_data*
for multi_index in np.ndindex(fspace_shape):
# have to be careful to take view and not copy
index = (np.s_[:],) + multi_index
fd_data = out_data[index]
dof_array = mm_field[multi_index]
resample_and_reorder(fd_data, dof_array)
return out
def test_ForcingCovariance_mult_parallel(my_ensemble, fs, fc, cov):
"test that the multiplication method of ForcingCovariance can be called independently in an ensemble"
fc.assemble()
if my_ensemble.ensemble_comm.rank == 0:
x = Function(fs).vector()
x.set_local(np.ones(x.local_size()))
y = Function(fs).vector()
fc.mult(x, y)
ygathered = y.gather()
assert_allclose(ygathered, np.dot(cov, np.ones(nx + 1)))
elif my_ensemble.ensemble_comm.rank == 1:
x = Function(fs).vector()
x.set_local(0.5 * np.ones(x.local_size()))
y = Function(fs).vector()
fc.mult(x, y)
ygathered = y.gather()
assert_allclose(ygathered, np.dot(cov, 0.5 * np.ones(nx + 1)))
class SCPC(SCBase):
needs_python_pmat = True
"""A Slate-based python preconditioner implementation of
static condensation for problems with up to three fields.
"""
@timed_function("SCPCInit")
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(Smat)
c_ksp.setUp()
c_ksp.setFromOptions()
self.condensed_ksp = c_ksp
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A, self.residual,
self.solution,
elim_fields)
def condensed_system(self, A, rhs, elim_fields):
"""Forms the condensed linear system by eliminating
specified unknowns.
:arg A: A Slate Tensor containing the mixed bilinear form.
:arg rhs: A firedrake function for the right-hand side.
:arg elim_fields: An iterable of field indices to eliminate.
"""
from firedrake.slate.static_condensation.la_utils import condense_and_forward_eliminate
return condense_and_forward_eliminate(A, rhs, elim_fields)
def local_solver_calls(self, A, rhs, x, elim_fields):
"""Provides solver callbacks for inverting local operators
and reconstructing eliminated fields.
:arg A: A Slate Tensor containing the mixed bilinear form.
:arg rhs: A firedrake function for the right-hand side.
:arg x: A firedrake function for the solution.
:arg elim_fields: An iterable of eliminated field indices
to recover.
"""
from firedrake.slate.static_condensation.la_utils import backward_solve
from firedrake.assemble import create_assembly_callable
fields = x.split()
systems = backward_solve(A, rhs, x, reconstruct_fields=elim_fields)
local_solvers = []
for local_system in systems:
Ae = local_system.lhs
be = local_system.rhs
i, = local_system.field_idx
local_solve = Ae.solve(be, decomposition="PartialPivLU")
solve_call = create_assembly_callable(
local_solve,
tensor=fields[i],
form_compiler_parameters=self.cxt.fc_params)
local_solvers.append(solve_call)
return local_solvers
@timed_function("SCPCUpdate")
def update(self, pc):
"""Update by assembling into the KSP operator. No
need to reconstruct symbolic objects.
"""
self._assemble_S()
self.S.force_evaluation()
def forward_elimination(self, pc, x):
"""Perform the forward elimination of fields and
provide the reduced right-hand side for the condensed
system.
:arg pc: a Preconditioner instance.
:arg x: a PETSc vector containing the incoming right-hand side.
"""
with self.residual.dat.vec_wo as v:
x.copy(v)
# Now assemble residual for the reduced problem
self._assemble_Srhs()
def sc_solve(self, pc):
"""Solve the condensed linear system for the
condensed field.
:arg pc: a Preconditioner instance.
"""
with self.condensed_rhs.dat.vec_ro as rhs:
if self.condensed_ksp.getInitialGuessNonzero():
acc = self.solution.split()[self.c_field].dat.vec
else:
acc = self.solution.split()[self.c_field].dat.vec_wo
with acc as sol:
self.condensed_ksp.solve(rhs, sol)
def backward_substitution(self, pc, y):
"""Perform the backwards recovery of eliminated fields.
:arg pc: a Preconditioner instance.
:arg y: a PETSc vector for placing the resulting fields.
"""
# Recover eliminated unknowns
for local_solver_call in self.local_solvers:
local_solver_call()
with self.solution.dat.vec_ro as w:
w.copy(y)
def view(self, pc, viewer=None):
"""Viewer calls for the various configurable objects in this PC."""
viewer.printfASCII("Static condensation preconditioner\n")
viewer.printfASCII("KSP to solve the reduced system:\n")
self.condensed_ksp.view(viewer=viewer)
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(Smat)
c_ksp.setUp()
c_ksp.setFromOptions()
self.condensed_ksp = c_ksp
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A, self.residual,
self.solution,
elim_fields)
class HybridSCPC(PCBase):
"""A Slate-based python preconditioner implementation of
static condensation for three-field hybridized problems. This
applies to the mixed-hybrid methods, such as the RT-H and BDM-H
methods, as well as hybridized-DG discretizations like the LDG-H
method.
"""
needs_python_pmat = True
@timed_function("HybridSCInit")
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field hybridized system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system for the Lagrange
multipliers. The scalar and flux fields are reconstructed
locally.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "hybrid_sc_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
# Retrieve the mixed function space, which is expected to
# be of the form: W = (DG_k)^n \times DG_k \times DG_trace
W = self.cxt.a.arguments()[0].function_space()
if len(W) != 3:
raise RuntimeError("Expecting three function spaces.")
# Assert a specific ordering of the spaces
# TODO: Clean this up
assert W[2].ufl_element().family() == "HDiv Trace"
# Extract trace space
T = W[2]
# Need to duplicate a trace space which is NOT
# associated with a subspace of a mixed space.
Tr = FunctionSpace(T.mesh(), T.ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
assert bc.function_space() == T, (
"BCs should be imposing vanishing conditions on traces")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Tr)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Tr, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.r_lambda = Function(T)
self.residual = Function(W)
self.solution = Function(W)
# Perform symbolics only once
S_expr, r_lambda_expr, u_h_expr, q_h_expr = self._slate_expressions
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
Smat = self.S.petscmat
# Set up ksp for the trace problem
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.incrementTabLevel(1, parent=pc)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
self._assemble_Srhs = create_assembly_callable(
r_lambda_expr,
tensor=self.r_lambda,
form_compiler_parameters=self.cxt.fc_params)
q_h, u_h, lambda_h = self.solution.split()
# Assemble u_h using lambda_h
self._assemble_u = create_assembly_callable(
u_h_expr, tensor=u_h, form_compiler_parameters=self.cxt.fc_params)
# Recover q_h using both u_h and lambda_h
self._assemble_q = create_assembly_callable(
q_h_expr, tensor=q_h, form_compiler_parameters=self.cxt.fc_params)
@cached_property
def _slate_expressions(self):
"""Returns all the relevant Slate expressions
for the static condensation and local recovery
procedures.
"""
# This operator has the form:
# | A B C |
# | D E F |
# | G H J |
# NOTE: It is often the case that D = B.T,
# G = C.T, H = F.T, and J = 0, but we're not making
# that assumption here.
_O = Tensor(self.cxt.a)
O = _O.blocks
# Extract sub-block:
# | A B |
# | D E |
# which has block row indices (0, 1) and block
# column indices (0, 1) as well.
M = O[:2, :2]
# Extract sub-block:
# | C |
# | F |
# which has block row indices (0, 1) and block
# column indices (2,)
K = O[:2, 2]
# Extract sub-block:
# | G H |
# which has block row indices (2,) and block column
# indices (0, 1)
L = O[2, :2]
# And the final block J has block row-column
# indices (2, 2)
J = O[2, 2]
# Schur complement for traces
S = J - L * M.inv * K
# Create mixed function for residual computation.
# This projects the non-trace residual bits into
# the trace space:
# -L * M.inv * | v1 v2 |^T
_R = AssembledVector(self.residual)
R = _R.blocks
v1v2 = R[:2]
v3 = R[2]
r_lambda = v3 - L * M.inv * v1v2
# Reconstruction expressions
q_h, u_h, lambda_h = self.solution.split()
# Local tensors needed for reconstruction
A = O[0, 0]
B = O[0, 1]
C = O[0, 2]
D = O[1, 0]
E = O[1, 1]
F = O[1, 2]
Se = E - D * A.inv * B
Sf = F - D * A.inv * C
v1, v2, v3 = self.residual.split()
# Solve locally using Cholesky factorizations
# (Se and A are symmetric positive definite)
u_h_expr = Se.solve(AssembledVector(v2) -
D * A.inv * AssembledVector(v1) -
Sf * AssembledVector(lambda_h),
decomposition="LLT")
q_h_expr = A.solve(AssembledVector(v1) - B * AssembledVector(u_h) -
C * AssembledVector(lambda_h),
decomposition="LLT")
return (S, r_lambda, u_h_expr, q_h_expr)
@timed_function("HybridSCUpdate")
def update(self, pc):
"""Update by assembling into the KSP operator. No
need to reconstruct symbolic objects.
"""
self._assemble_S()
def apply(self, pc, x, y):
"""Solve the reduced system for the Lagrange multipliers.
The system is assembled using operators constructed from
the Slate expressions in the initialize method of this PC.
Recovery of the scalar and flux fields are assembled cell-wise
from Slate expressions describing the local problem.
"""
with self.residual.dat.vec_wo as v:
x.copy(v)
with timed_region("HybridSCRHS"):
# Now assemble residual for the reduced problem
self._assemble_Srhs()
with timed_region("HybridSCSolve"):
# Solve the system for the Lagrange multipliers
with self.r_lambda.dat.vec_ro as b:
if self.trace_ksp.getInitialGuessNonzero():
acc = self.solution.split()[2].dat.vec
else:
acc = self.solution.split()[2].dat.vec_wo
with acc as x_trace:
self.trace_ksp.solve(b, x_trace)
with timed_region("HybridSCReconstruct"):
# Recover u_h and q_h
self._assemble_u()
self._assemble_q()
with self.solution.dat.vec_ro as w:
w.copy(y)
def applyTranspose(self, pc, x, y):
"""Apply the transpose of the preconditioner."""
raise NotImplementedError("Transpose application is not implemented.")
def view(self, pc, viewer=None):
viewer.printfASCII("Hybridized trace preconditioner\n")
viewer.printfASCII("KSP to solve trace system:\n")
self.trace_ksp.view(viewer=viewer)
def __init__(self,
function_space,
sigma,
l,
cutoff=1.e-3,
regularization=1.e-8,
cov=sqexp):
"""
Create new forcing covariance
Creates a new ForcingCovariance object from a function space, parameters, and
covariance function. Required parameters are the function space and sigma and
correlation length parameters needed to compute the covariance matrix.
Note that this just initializes the object, and does not compute the matrix
entries or assemble the final PETSc matrix. This is done using the ``assemble``
method, though if you attempt to use an unassembled matrix assembly will
automatically be done. However the domain decomposition is done here to determine
the number of DOFs handled by each process.
"""
# need to investigate parallelization here, load balancing likely to be uneven
# if we just use the local ownership from the distributed matrix
# since each row has an uneven amount of work
# know that we have reduced bandwidth (though unclear if this translates to a low
# bandwidth of the assembled covariance matrix)
if not isinstance(function_space, WithGeometry):
raise TypeError(
"bad input type for function_space: must be a FunctionSpace")
self.function_space = function_space
self.comm = function_space.comm
# extract mesh and process local information
self.nx = Function(self.function_space).vector().size()
self.nx_local = Function(self.function_space).vector().local_size()
# set parameters and covariance
assert regularization >= 0., "regularization parameter must be non-negative"
self.sigma = sigma
self.l = l
self.cutoff = cutoff
self.regularization = regularization
self.cov = cov
# get local ownership information of distributed matrix
vtemp = PETSc.Vec().create(comm=self.comm)
vtemp.setSizes((self.nx_local, -1))
vtemp.setFromOptions()
vtemp.setUp()
self.local_startind, self.local_endind = vtemp.getOwnershipRange()
vtemp.destroy()
self.is_assembled = False
self.G = None
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field hybridized system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system for the Lagrange
multipliers. The scalar and flux fields are reconstructed
locally.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "hybrid_sc_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
# Retrieve the mixed function space, which is expected to
# be of the form: W = (DG_k)^n \times DG_k \times DG_trace
W = self.cxt.a.arguments()[0].function_space()
if len(W) != 3:
raise RuntimeError("Expecting three function spaces.")
# Assert a specific ordering of the spaces
# TODO: Clean this up
assert W[2].ufl_element().family() == "HDiv Trace"
# Extract trace space
T = W[2]
# Need to duplicate a trace space which is NOT
# associated with a subspace of a mixed space.
Tr = FunctionSpace(T.mesh(), T.ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
assert bc.function_space() == T, (
"BCs should be imposing vanishing conditions on traces")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Tr)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Tr, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.r_lambda = Function(T)
self.residual = Function(W)
self.solution = Function(W)
# Perform symbolics only once
S_expr, r_lambda_expr, u_h_expr, q_h_expr = self._slate_expressions
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
Smat = self.S.petscmat
# Set up ksp for the trace problem
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.incrementTabLevel(1, parent=pc)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
self._assemble_Srhs = create_assembly_callable(
r_lambda_expr,
tensor=self.r_lambda,
form_compiler_parameters=self.cxt.fc_params)
q_h, u_h, lambda_h = self.solution.split()
# Assemble u_h using lambda_h
self._assemble_u = create_assembly_callable(
u_h_expr, tensor=u_h, form_compiler_parameters=self.cxt.fc_params)
# Recover q_h using both u_h and lambda_h
self._assemble_q = create_assembly_callable(
q_h_expr, tensor=q_h, form_compiler_parameters=self.cxt.fc_params)
def __init__(self, mesh, element, name=None, dim=1, rank=0):
"""
:param mesh: :class:`Mesh` to build this space on
:param element: :class:`ufl.FiniteElementBase` to build this space from
:param name: user-defined name for this space
:param dim: vector space dimension of a :class:`.VectorFunctionSpace`
:param rank: rank of the space, not the value rank
"""
self._ufl_element = element
# Compute the FIAT version of the UFL element above
self.fiat_element = fiat_utils.fiat_from_ufl_element(element)
if isinstance(mesh, mesh_t.ExtrudedMesh):
# Set up some extrusion-specific things
# The bottom layer maps will come from element_dof_list
# dof_count is the total number of dofs in the extruded mesh
# Get the flattened version of the FIAT element
self.flattened_element = fiat_utils.FlattenedElement(
self.fiat_element)
# Compute the number of DoFs per dimension on top/bottom and sides
entity_dofs = self.fiat_element.entity_dofs()
top_dim = mesh._plex.getDimension()
self._xtr_hdofs = [
len(entity_dofs[(d, 0)][0]) for d in range(top_dim + 1)
]
self._xtr_vdofs = [
len(entity_dofs[(d, 1)][0]) for d in range(top_dim + 1)
]
# Compute the dofs per column
self.dofs_per_column = eutils.compute_extruded_dofs(
self.fiat_element, self.flattened_element.entity_dofs(),
mesh._layers)
# Compute the offset for the extrusion process
self.offset = eutils.compute_offset(
self.fiat_element.entity_dofs(),
self.flattened_element.entity_dofs(),
self.fiat_element.space_dimension())
# Compute the top and bottom masks to identify boundary dofs
#
# Sorting the keys of the closure entity dofs, the whole cell
# comes last [-1], before that the horizontal facet [-2], before
# that vertical facets [-3]. We need the horizontal facets here.
closure_dofs = self.fiat_element.entity_closure_dofs()
b_mask = closure_dofs[sorted(closure_dofs.keys())[-2]][0]
t_mask = closure_dofs[sorted(closure_dofs.keys())[-2]][1]
self.bt_masks = (b_mask, t_mask) # conversion to tuple
self.extruded = True
self._dofs_per_entity = self.dofs_per_column
else:
# If not extruded specific, set things to None/False, etc.
self.offset = None
self.bt_masks = None
self.dofs_per_column = np.zeros(1, np.int32)
self.extruded = False
entity_dofs = self.fiat_element.entity_dofs()
self._dofs_per_entity = [
len(entity[0]) for d, entity in entity_dofs.iteritems()
]
self.name = name
self._dim = dim
self._mesh = mesh
self._index = None
dm = PETSc.DMShell().create()
dm.setAttr('__fs__', weakref.ref(self))
dm.setPointSF(mesh._plex.getPointSF())
# Create the PetscSection mapping topological entities to DoFs
sec = mesh._plex.createSection([1],
self._dofs_per_entity,
perm=mesh._plex_renumbering)
dm.setDefaultSection(sec)
self._global_numbering = sec
self._dm = dm
self._ises = None
self._halo = halo.Halo(dm)
# Compute entity class offsets
self.dof_classes = [0, 0, 0, 0]
for d in range(mesh._plex.getDimension() + 1):
ndofs = self._dofs_per_entity[d]
for i in range(4):
self.dof_classes[i] += ndofs * mesh._entity_classes[d, i]
# Tell the DM about the layout of the global vector
from firedrake.function import Function
with Function(self).dat.vec_ro as v:
self._dm.setGlobalVector(v.duplicate())
self._node_count = self._global_numbering.getStorageSize()
self.cell_node_list = mesh.create_cell_node_list(
self._global_numbering, self.fiat_element)
if mesh._plex.getStratumSize("interior_facets", 1) > 0:
self.interior_facet_node_list = \
dmplex.get_facet_nodes(mesh.interior_facets.facet_cell,
self.cell_node_list)
else:
self.interior_facet_node_list = np.array([], dtype=np.int32)
if mesh._plex.getStratumSize("exterior_facets", 1) > 0:
self.exterior_facet_node_list = \
dmplex.get_facet_nodes(mesh.exterior_facets.facet_cell,
self.cell_node_list)
else:
self.exterior_facet_node_list = np.array([], dtype=np.int32)
# Note: this is the function space rank. The value rank may be different.
self.rank = rank
# Empty map caches. This is a sui generis cache
# implementation because of the need to support boundary
# conditions.
self._cell_node_map_cache = {}
self._exterior_facet_map_cache = {}
self._interior_facet_map_cache = {}
class VertexBasedLimiter(Limiter):
"""
A vertex based limiter for P1DG fields.
This limiter implements the vertex-based limiting scheme described in
Dmitri Kuzmin, "A vertex-based hierarchical slope limiter for p-adaptive
discontinuous Galerkin methods". J. Comp. Appl. Maths (2010)
http://dx.doi.org/10.1016/j.cam.2009.05.028
"""
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
if utils.complex_mode:
raise ValueError(
"We haven't decided what limiting complex valued fields means. Please get in touch if you have need."
)
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG',
1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.centroids_rhs = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
domain = "{[i]: 0 <= i < maxq.dofs}"
instructions = """
for i
maxq[i] = fmax(maxq[i], q[0])
minq[i] = fmin(minq[i], q[0])
end
"""
self._min_max_loop = (domain, instructions)
# Perform limiting loop
domain = "{[i, ii]: 0 <= i < q.dofs and 0 <= ii < q.dofs}"
instructions = """
<float64> alpha = 1
<float64> qavg = qbar[0, 0]
for i
<float64> _alpha1 = fmin(alpha, fmin(1, (qmax[i] - qavg)/(q[i] - qavg)))
<float64> _alpha2 = fmin(alpha, fmin(1, (qavg - qmin[i])/(qavg - q[i])))
alpha = _alpha1 if q[i] > qavg else (_alpha2 if q[i] < qavg else alpha)
end
for ii
q[ii] = qavg + alpha * (q[ii] - qavg)
end
"""
self._limit_kernel = (domain, instructions)
def _construct_centroid_solver(self):
"""
Constructs a linear problem for computing the centroids
:return: LinearSolver instance
"""
u = TrialFunction(self.P0)
v = TestFunction(self.P0)
a = assemble(inner(u, v) * dx)
return LinearSolver(a,
solver_parameters={
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
})
def _update_centroids(self, field):
"""
Update centroid values
"""
assemble(inner(field, TestFunction(self.P0)) * dx,
tensor=self.centroids_rhs)
self.centroid_solver.solve(self.centroids, self.centroids_rhs)
def compute_bounds(self, field):
"""
Only computes min and max bounds of neighbouring cells
"""
self._update_centroids(field)
self.max_field.assign(-1.0e10) # small number
self.min_field.assign(1.0e10) # big number
par_loop(self._min_max_loop,
dx, {
"maxq": (self.max_field, MAX),
"minq": (self.min_field, MIN),
"q": (self.centroids, READ)
},
is_loopy_kernel=True)
def apply_limiter(self, field):
"""
Only applies limiting loop on the given field
"""
par_loop(self._limit_kernel,
dx, {
"qbar": (self.centroids, READ),
"q": (field, RW),
"qmax": (self.max_field, READ),
"qmin": (self.min_field, READ)
},
is_loopy_kernel=True)
def apply(self, field):
"""
Re-computes centroids and applies limiter to given field
"""
assert field.function_space() == self.P1DG, \
'Given field does not belong to this objects function space'
self.compute_bounds(field)
self.apply_limiter(field)
class LinearSolver(object):
r"""
Class encapsulating all solves on the same FEM model
This class forms the base of all Stat FEM computations for a linear problem. It requires
the base FEM problem, the forcing covariance (represented by a ``ForcingCovariance`` object),
the sensor locations, data, and uncertanties (represented by a ``ObsData`` object),
priors on the model discrepancy hyperparameters (optional), and an ensemble MPI communicator
for parallelizing the covariance solves (optional).
Once these are set, the prior solves can be done and cached, which is generally the most
computationally expensive part of the modeling. The class also contains methods for
performing parameter estimation (a ``logposterior`` to compute the negative log posterior
or marginal likelihood if no priors are specified and its associated derivatives), and
prediction of sensor values and uncertainties at unmeasured locations.
:ivar solver: the base Firedrake FEM solver
:type solver: Firedrake LinearSolver
:ivar b: the FEM RHS vector
:type b: Firedrake Vector or Function
:ivar G: Forcing Covariance sparse matrix
:type G: ForcingCovariance
:ivar data: Sensor locations, observed values, and uncertainties
:type data: ObsData
:ivar priors: list of prior distributions on hyperparameters or all ``None`` if uninformative
priors are assumed
:type priors: list
:ivar ensemble_comm: Firedrake Ensemble communicator for parallelizing covariance solves
:type ensemble_comm: MPI Communicator
:ivar params: Current set of parameters (a numpy array of length 3) representing the
data/model scaling factor :math:`{\rho}`, model discrepancy covariance,
and model discrepancy correlation length. All parameters are on a logarithmic
scale to enforce positivity.
:type params: ndarray
:ivar im: interpolation matrix used to interpolate between FEM mesh and sensor data
:type im: InterpolationMatrix
:ivar x: Prior FEM solution on distributed FEM mesh
:type x: Firedrake Function
:ivar mu: Prior FEM solution interpolated to sensor locations on root process (other processes
have arrays of length 0)
:type mu: ndarray
:ivar Cu: Prior FEM covariance interpolated to sensor locations on root process (other processes
have arrays of shape ``(0, 0)``
:type Cu: ndarray
:ivar current_logpost: Current value of the negative log-posterior (or log likelihood if prior
is uninformative)
:type current_logpost: float
"""
def __init__(self,
A,
b,
G,
data,
*,
priors=[None, None, None],
ensemble_comm=COMM_SELF,
P=None,
solver_parameters=None,
nullspace=None,
transpose_nullspace=None,
near_nullspace=None,
options_prefix=None):
r"""
Create a new object encapsulating all solves on the same FEM model
Initialize a new object for a given FEM problem to perform the Stat FEM solves.
This class forms the base of all Stat FEM computations for a linear problem. It requires
the base FEM problem, the forcing covariance (represented by a ``ForcingCovariance`` object),
the sensor locations, data, and uncertanties (represented by a ``ObsData`` object),
priors on the model discrepancy hyperparameters (optional), and an ensemble MPI communicator
for parallelizing the covariance solves (optional).
:param A: the FEM stiffness matrix
:type A: Firedrake Matrix
:param b: the FEM RHS vector
:type b: Firedrake Vector or Function
:param G: Forcing Covariance sparse matrix
:type G: ForcingCovariance
:param data: Sensor locations, observed values, and uncertainties
:type data: ObsData
:param priors: list of prior distributions on hyperparameters or all ``None`` if uninformative
priors are assumed (optional)
:type priors: list
:param ensemble_comm: Firedrake Ensemble communicator for parallelizing covariance solves (optional)
:type ensemble_comm: MPI Communicator
:param P: an optional `MatrixBase` to construct any
preconditioner from; if none is supplied ``A`` is
used to construct the preconditioner.
:type P: MatrixBase or other derived Matrix type
:param parameters: (optional) dict of solver parameters
:type parameters: dict
:param nullspace: an optional `VectorSpaceBasis` (or
`MixedVectorSpaceBasis`) spanning the null space
of the operator.
:type nullspace: VectorSpaceBasis or MixedVectorSpaceBasis
:param transpose_nullspace: as for the nullspace, but used to
make the right hand side consistent.
:type transpose_nullspace: VectorSpaceBasis or MixedVectorSpaceBasis
:param near_nullspace: as for the nullspace, but used to set
the near nullpace.
:type near_nullspace: VectorSpaceBasis or MixedVectorSpaceBasis
:param options_prefix: an optional prefix used to distinguish
PETSc options. If not provided a unique
prefix will be created. Use this option
if you want to pass options to the solver
from the command line in addition to
through the ``solver_parameters`` dict.
:type options_prefix: str
:returns: new ``LinearSolver`` instance
:rtype: LinearSolver
"""
if not isinstance(A, MatrixBase):
raise TypeError("A must be a firedrake matrix")
if not isinstance(b, (Function, Vector)):
raise TypeError("b must be a firedrake function or vector")
if not isinstance(G, ForcingCovariance):
raise TypeError("G must be a forcing covariance")
if not isinstance(data, ObsData):
raise TypeError("data must be an ObsData type")
if not isinstance(priors, list):
raise TypeError("priors must be a list of prior objects or None")
if not len(priors) == 3:
raise ValueError(
"priors must be a list of prior objects or None of length 3")
for p in priors:
if not p is None:
raise TypeError(
"priors must be a list of prior objects or None")
if not isinstance(ensemble_comm, type(COMM_WORLD)):
raise TypeError(
"ensemble_comm must be an MPI communicator created with a firedrake Ensemble"
)
self.solver = fdLS(A,
P=P,
solver_parameters=solver_parameters,
nullspace=nullspace,
transpose_nullspace=transpose_nullspace,
near_nullspace=near_nullspace,
options_prefix=options_prefix)
self.b = b
self.G = G
self.data = data
self.ensemble_comm = ensemble_comm
self.priors = list(priors)
self.params = None
self.im = InterpolationMatrix(G.function_space, data.get_coords())
self.x = None
self.mu = None
self.Cu = None
self.current_logpost = None
def __del__(self):
r"""
Delete the LinearSolver object
When deleting a LinearSolver, one needs to deallocate the memory for the interpolation
matrix. No inputs or return values.
"""
self.im.destroy()
def set_params(self, params):
r"""
Sets parameter values
Checks and sets new values of the hyperparameters. New parameters must be a numpy
array of length 3. First parameter is the data/model scaling factor :math:`{\rho}`,
second parameter is the model discrepancy covariance, and the third parameter is
the model discrepancy correlation length. All parameters are assumed to be on a
logarithmic scale to enforce positivity.
:param params: New set of parameters (must be a numpy array of length 3)
:type params: ndarray
:returns: None
"""
params = np.array(params, dtype=np.float64)
assert params.shape == (
3, ), "bad shape for model discrepancy parameters"
self.params = params
def solve_prior(self):
r"""
Solve base (prior) FEM plus covariance interpolated to the data locations
This method solves the prior FEM and covariance interpolated to the sensor locations.
It does not require setting parameter values, as the model discrepancy does not
influence these results. The covariance is cached as it is expensive to compute
and is re-used in all other solves.
In addition to caching the results, the method returns solution as numpy arrays
on the root process (rank 0).
Note that unlike the solve done in the meshspace, this uses a return value rather than a
Firedrake/PETSc style interface to place the solution in a pre-allocated ``Function``.
This is because each process has a different array size, so would require correctly
pre-allocating arrays of different lengths on each process.
:returns: FEM prior mean and covariance (as a tuple of numpy arrays) on the root process.
Non-root processes return numpy arrays of shape ``(0,)`` (mean) and ``(0, 0)``
(covariance).
:rtype: tuple of ndarrays
"""
# form interpolated prior covariance across all ensemble processes
self.Cu = interp_covariance_to_data(self.im, self.G, self.solver,
self.im, self.ensemble_comm)
# solve base FEM (prior mean) and interpolate to data space on root
self.x = Function(self.G.function_space)
if self.ensemble_comm.rank == 0:
self.solver.solve(self.x, self.b)
self.mu = self.im.interp_mesh_to_data(self.x.vector())
else:
self.mu = np.zeros(0)
return self.mu, self.Cu
def solve_posterior(self, x, scale_mean=False):
r"""
Solve FEM posterior in mesh space
Solve for the FEM posterior conditioned on the data on the FEM mesh. The solution
is stored in the preallocated Firedrake ``Function``.
Note that if an ensemble communicator was used to parallelize the covariance solves,
the solution is only stored in the root of the ensemble communicator. The Firedrake
``Function`` on the other processes will not be modified.
The optional ``scale_mean`` argument determines if the solution is to be re-scaled
by the model discrepancy scaling factor. This value is by default ``False``.
To re-scale to match the data, pass ``scale_mean=True``.
:param x: Firedrake ``Function`` for holding the solution. This is modified in place
by the method.
:type x: Firedrake Function
:param scale_mean: Boolean indicating if the mean should be scaled by the model
discrepancy scaling factor. Optional, default is ``False``
:type scale_mean: bool
:returns: None
"""
if not isinstance(bool(scale_mean), bool):
raise TypeError("scale_mean argument must be boolean-like")
# create interpolation matrix if not cached
if self.Cu is None or self.x is None:
self.solve_prior()
if self.params is None:
raise ValueError("must set parameter values to solve posterior")
rho = np.exp(self.params[0])
if scale_mean:
scalefact = rho
else:
scalefact = 1.
# remaining solves are just done on ensemble root
if self.ensemble_comm.rank == 0:
if self.G.comm.rank == 0:
Ks = self.data.calc_K_plus_sigma(self.params[1:])
try:
LK = cho_factor(Ks)
except LinAlgError:
raise LinAlgError(
"Error attempting to compute the Cholesky factorization "
+ "of the model discrepancy")
tmp_dataspace_1 = cho_solve(LK, self.data.get_data())
else:
tmp_dataspace_1 = np.zeros(0)
# interpolate to dataspace
tmp_meshspace_1 = self.im.interp_data_to_mesh(tmp_dataspace_1)
# solve forcing covariance and interpolate to dataspace
tmp_meshspace_2 = solve_forcing_covariance(
self.G, self.solver,
tmp_meshspace_1)._scale(rho) + self.x.vector()
tmp_dataspace_1 = self.im.interp_mesh_to_data(tmp_meshspace_2)
if self.G.comm.rank == 0:
try:
L = cho_factor(Ks + rho**2 * self.Cu)
except LinAlgError:
raise LinAlgError(
"Error attempting to compute the Cholesky factorization "
+ "of the model discrepancy plus forcing covariance")
tmp_dataspace_2 = cho_solve(L, tmp_dataspace_1)
else:
tmp_dataspace_2 = np.zeros(0)
tmp_meshspace_1 = self.im.interp_data_to_mesh(tmp_dataspace_2)
tmp_meshspace_1 = solve_forcing_covariance(
self.G, self.solver, tmp_meshspace_1)._scale(rho**2)
x.assign(
(tmp_meshspace_2 - tmp_meshspace_1)._scale(scalefact).function)
def solve_posterior_covariance(self, scale_mean=False):
r"""
Solve posterior FEM and covariance interpolated to the data locations
This method solves the posterior FEM and covariance interpolated to the sensor
locations. The method returns solution as numpy arrays on the root process (rank 0).
Note that unlike the solve done in the meshspace, this uses a return value rather than a
Firedrake/PETSc style interface to place the solution in a pre-allocated ``Function``.
This is because each process has a different array size, so would require correctly
pre-allocating arrays of different lengths on each process.
The optional ``scale_mean`` argument determines if the solution is to be re-scaled
by the model discrepancy scaling factor. This value is by default ``False``.
To re-scale to match the data, pass ``scale_mean=True``.
:returns: FEM posterior mean and covariance (as a tuple of numpy arrays) on the root process.
Non-root processes return numpy arrays of shape ``(0,)`` (mean) and ``(0, 0)``
(covariance).
:param scale_mean: Boolean indicating if the mean should be scaled by the model
discrepancy scaling factor. Optional, default is ``False``
:type scale_mean: bool
:rtype: tuple of ndarrays
"""
if not isinstance(bool(scale_mean), bool):
raise TypeError("scale_mean argument must be boolean-like")
# create interpolation matrix if not cached
if self.mu is None or self.Cu is None:
self.solve_prior()
if self.params is None:
raise ValueError("must set parameter values to solve posterior")
rho = np.exp(self.params[0])
if scale_mean:
scalefact = rho
else:
scalefact = 1.
if self.ensemble_comm.rank == 0 and self.G.comm.rank == 0:
try:
Ks = self.data.calc_K_plus_sigma(self.params[1:])
LK = cho_factor(Ks)
LC = cho_factor(Ks + rho**2 * self.Cu)
except LinAlgError:
raise LinAlgError(
"Cholesky factorization of one of the covariance matrices failed"
)
# compute posterior mean
muy = rho * np.dot(self.Cu, cho_solve(
LK, self.data.get_data())) + self.mu
muy_tmp = rho**2 * np.dot(self.Cu, cho_solve(LC, muy))
muy = muy - muy_tmp
# compute posterior covariance
Cuy = self.Cu - rho**2 * np.dot(self.Cu, cho_solve(LC, self.Cu))
else:
muy = np.zeros(0)
Cuy = np.zeros((0, 0))
return scalefact * muy, Cuy
def solve_prior_generating(self):
r"""
Solve for the prior of the generating process
This method solves for the prior of the generating process before looking at the data.
The main computational cost is solving for the prior of the covariance, so if this is
cached from a previous solve this is a simple calculation.
:returns: FEM prior mean and covariance of the true generating process (as a tuple of
numpy arrays) on the root process. Non-root processes return numpy arrays of
shape ``(0,)`` (mean) and ``(0, 0)`` (covariance).
:rtype: tuple of ndarrays
"""
# create interpolation matrix if not cached
if self.mu is None or self.Cu is None:
self.solve_prior()
if self.params is None:
raise ValueError(
"must set parameter values to solve prior of generating process"
)
rho = np.exp(self.params[0])
if self.G.comm.rank == 0 and self.ensemble_comm.rank == 0:
m_eta = rho * self.mu
C_eta = rho**2 * self.Cu + self.data.calc_K(self.params[1:])
else:
m_eta = np.zeros(0)
C_eta = np.zeros((0, 0))
return m_eta, C_eta
def solve_posterior_generating(self):
r"""
Solve for the posterior of the generating process
This method solves for the posterior of the generating process before looking at the data.
The main computational cost is solving for the prior of the covariance, so if this is
cached from a previous solve this is a simple calculation.
:returns: FEM posterior mean and covariance of the true generating process (as a tuple of
numpy arrays) on the root process. Non-root processes return numpy arrays of
shape ``(0,)`` (mean) and ``(0, 0)`` (covariance).
:rtype: tuple of ndarrays
"""
# create interpolation matrix if not cached
m_eta, C_eta = self.solve_prior_generating()
if self.ensemble_comm.rank == 0 and self.G.comm.rank == 0:
try:
L = cho_factor(self.data.get_unc()**2 *
np.eye(self.data.get_n_obs()) + C_eta)
except LinAlgError:
raise LinAlgError(
"Cholesky factorization of the covariance matrix failed")
C_etay = cho_solve(L, self.data.get_unc()**2 * C_eta)
m_etay = cho_solve(
L,
np.dot(C_eta, self.data.get_data()) +
self.data.get_unc()**2 * m_eta)
else:
m_etay = np.zeros(0)
C_etay = np.zeros((0, 0))
return m_etay, C_etay
def predict_mean(self, coords, scale_mean=True):
r"""
Compute the predictive mean
This method computes the predictive mean of data values at unmeasured locations. It returns
the vector of predicted sensor values on the root process as numpy array. It requires only a
small overhead above the computational work of finding the posterior mean (i.e. you get
the mean value at new sensor locations for "free" once you have solved the posterior).
The optional ``scale_mean`` argument determines if the solution is to be re-scaled
by the model discrepancy scaling factor. This value is by default ``True``.
To re-scale to match the FEM solution, pass ``scale_mean=False``.
:param coords: Spatial coordinates at which the mean will be predicted. Must be a
2D Numpy array (or a 1D array, which will assume the second axis has length
1)
:type coords: ndarray
:param scale_mean: Boolean indicating if the mean should be scaled by the model
discrepancy scaling factor. Optional, default is ``True``
:type scale_mean: bool
:returns: FEM prediction at specified sensor locations as a numpy array on the root process.
All other processes will have a numpy array of length 0.
:rtype: ndarray
"""
if not isinstance(bool(scale_mean), bool):
raise TypeError("scale_mean argument must be boolean-like")
coords = np.array(coords, dtype=np.float64)
if coords.ndim == 1:
coords = np.reshape(coords, (-1, 1))
assert coords.ndim == 2, "coords must be a 1d or 2d array"
assert coords.shape[1] == self.data.get_n_dim(
), "axis 1 of coords must be the same length as the FEM dimension"
if self.Cu is None:
self.solve_prior()
if self.params is None:
raise ValueError("must set parameter values to make predictions")
rho = np.exp(self.params[0])
if scale_mean:
scalefact = rho
else:
scalefact = 1.
x = Function(self.G.function_space)
self.solve_posterior(x)
im = InterpolationMatrix(self.G.function_space, coords)
mu = scalefact * im.interp_mesh_to_data(x.vector())
im.destroy()
return mu
def predict_covariance(self, coords, unc):
r"""
Compute the predictive covariance
This method computes the predictive covariance of data values at unmeasured locations.
It returns the array of predicted sensor value covariances on the root process as numpy
array. Unlike the mean, the predictive covariance requires doing two additional sets of
covariance solves: one on the new sensor locations (to get the baseline covariance),
and one set of solves that interpolates between the predictive points and the original
sensor locations. This can be thought of as doing the covariance solves at the new points
to get a baseline uncertainty, and then the cross-solves determine if any of the sensor
data is close enough to the predictive points to reduce this uncertainty.
:param coords: Spatial coordinates at which the mean will be predicted. Must be a
2D Numpy array (or a 1D array, which will assume the second axis has length
1)
:type coords: ndarray
:param unc: Uncertainty for unmeasured sensor locations (i.e. the statistical error one would
expect if these measurements were made). Can be a single non-negative float,
or an array of non-negative floats with the same length as the first axis of
``coords``.
:type unc: float or ndarray
:returns: FEM predictive covariance at specified sensor locations as a numpy array on the
root process. All other processes will have a numpy array of shape ``(0, 0)``.
:rtype: ndarray
"""
coords = np.array(coords, dtype=np.float64)
if coords.ndim == 1:
coords = np.reshape(coords, (-1, 1))
assert coords.ndim == 2, "coords must be a 1d or 2d array"
assert coords.shape[1] == self.data.get_n_dim(
), "axis 1 of coords must be the same length as the FEM dimension"
if self.Cu is None:
self.solve_prior()
if self.params is None:
raise ValueError("must set parameter values to make predictions")
rho = np.exp(self.params[0])
im_coords = InterpolationMatrix(self.G.function_space, coords)
if coords.shape[0] > self.data.get_n_obs():
Cucd = interp_covariance_to_data(im_coords, self.G, self.solver,
self.im, self.ensemble_comm)
else:
Cucd = interp_covariance_to_data(self.im, self.G, self.solver,
im_coords, self.ensemble_comm).T
Cucc = interp_covariance_to_data(im_coords, self.G, self.solver,
im_coords, self.ensemble_comm)
if self.ensemble_comm.rank == 0 and self.G.comm.rank == 0:
try:
Ks = self.data.calc_K_plus_sigma(self.params[1:])
LC = cho_factor(Ks + rho**2 * self.Cu)
except LinAlgError:
raise LinAlgError(
"Cholesky factorization of one of the covariance matrices failed"
)
# compute predictive covariance
Cuy = Cucc - rho**2 * np.dot(Cucd, cho_solve(LC, Cucd.T))
Cuy = ObsData(coords, np.zeros(coords.shape[0]),
unc).calc_K_plus_sigma(
self.params[1:]) + rho**2 * Cuy
else:
Cuy = np.zeros((0, 0))
im_coords.destroy()
return Cuy
def logposterior(self, params):
r"""
Compute the negative log posterior for a particular set of parameters
Computes the negative log posterior (negative marginal log-likelihood minus any
prior log-probabilities). This is computed on the root process and then broadcast
to all processes.
The main computational expense is computing the prior mean and covariance, which only
needs to be done once and can be cached. This also requires computing the Cholesky
decomposition of the covariance plus model discrepancy.
New parameters must be a numpy array of length 3. First parameter is the data/model
scaling factor :math:`{\rho}`, second parameter is the model discrepancy covariance,
and the third parameter is the model discrepancy correlation length. All parameters
are assumed to be on a logarithmic scale to enforce positivity.
:param params: New set of parameters (must be a numpy array of length 3)
:type params: ndarray
:returns: negative log posterior
:rtype: float
"""
self.set_params(params)
rho = np.exp(self.params[0])
if self.Cu is None or self.mu is None:
self.solve_prior()
# compute log-likelihood on root process and broadcast
if COMM_WORLD.rank == 0:
KCu = rho**2 * self.Cu + self.data.calc_K_plus_sigma(
self.params[1:])
try:
L = cho_factor(KCu)
except LinAlgError:
raise LinAlgError(
"Error attempting to factorize the covariance matrix " +
"in model_loglikelihood")
invKCudata = cho_solve(L, self.data.get_data() - rho * self.mu)
log_posterior = 0.5 * (
self.data.get_n_obs() * np.log(2. * np.pi) +
2. * np.sum(np.log(np.diag(L[0]))) +
np.dot(self.data.get_data() - rho * self.mu, invKCudata))
for i in range(3):
if not self.priors[i] is None:
log_posterior -= self.priors[i].logp(self.params[i])
else:
log_posterior = None
log_posterior = COMM_WORLD.bcast(log_posterior, root=0)
assert not log_posterior is None, "error in broadcasting the log likelihood"
COMM_WORLD.barrier()
return log_posterior
def logpost_deriv(self, params):
r"""
Compute the gradient of the negative log posterior for a particular set of parameters
Computes the gradient of the negative log posterior (negative marginal log-likelihood
minus any prior log-probabilities). This is computed on the root process and then broadcast
to all processes.
The main computational expense is computing the prior mean and covariance, which only
needs to be done once and can be cached. This also requires computing the Cholesky
decomposition of the covariance plus model discrepancy.
New parameters must be a numpy array of length 3. First parameter is the data/model
scaling factor :math:`{\rho}`, second parameter is the model discrepancy covariance,
and the third parameter is the model discrepancy correlation length. All parameters
are assumed to be on a logarithmic scale to enforce positivity.
The returned log posterior gradient is a numpy array of length 3, with each component
corresponding to the derivative of each of the input parameters.
:param params: New set of parameters (must be a numpy array of length 3)
:type params: ndarray
:returns: gradient of the negative log posterior
:rtype: ndarray
"""
self.set_params(params)
rho = np.exp(self.params[0])
if self.Cu is None or self.mu is None:
self.solve_prior()
# compute log-likelihood on root process
if COMM_WORLD.rank == 0:
KCu = rho**2 * self.Cu + self.data.calc_K_plus_sigma(params[1:])
try:
L = cho_factor(KCu)
except LinAlgError:
raise LinAlgError(
"Error attempting to factorize the covariance matrix " +
"in model_loglikelihood")
invKCudata = cho_solve(L, self.data.get_data() - rho * self.mu)
K_deriv = self.data.calc_K_deriv(self.params[1:])
deriv = np.zeros(3)
deriv[0] = (
-rho * np.dot(self.mu, invKCudata) -
rho**2 * np.linalg.multi_dot([invKCudata, self.Cu, invKCudata])
+ rho**2 * np.trace(cho_solve(L, self.Cu)))
for i in range(0, 2):
deriv[i + 1] = -0.5 * (
np.linalg.multi_dot([invKCudata, K_deriv[i], invKCudata]) -
np.trace(cho_solve(L, K_deriv[i])))
for i in range(3):
if not self.priors[i] is None:
deriv[i] -= self.priors[i].dlogpdtheta(self.params[i])
else:
deriv = None
deriv = COMM_WORLD.bcast(deriv, root=0)
assert not deriv is None, "error in broadcasting the log likelihood derivative"
COMM_WORLD.barrier()
return deriv
class VertexBasedLimiter(Limiter):
"""
A vertex based limiter for P1DG fields.
This limiter implements the vertex-based limiting scheme described in
Dmitri Kuzmin, "A vertex-based hierarchical slope limiter for p-adaptive
discontinuous Galerkin methods". J. Comp. Appl. Maths (2010)
http://dx.doi.org/10.1016/j.cam.2009.05.028
"""
def __init__(self, space):
"""
Initialise limiter
:param space : FunctionSpace instance
"""
self.P1DG = space
self.P1CG = FunctionSpace(self.P1DG.mesh(), 'CG',
1) # for min/max limits
self.P0 = FunctionSpace(self.P1DG.mesh(), 'DG', 0) # for centroids
# Storage containers for cell means, max and mins
self.centroids = Function(self.P0)
self.centroids_rhs = Function(self.P0)
self.max_field = Function(self.P1CG)
self.min_field = Function(self.P1CG)
self.centroid_solver = self._construct_centroid_solver()
# Update min and max loop
self._min_max_loop = """
for(int i = 0; i < maxq.dofs; i++) {
maxq[i][0] = fmax(maxq[i][0],q[0][0]);
minq[i][0] = fmin(minq[i][0],q[0][0]);
}
"""
# Perform limiting loop
self._limit_kernel = """
double alpha = 1.0;
double qavg = qbar[0][0];
for (int i=0; i < q.dofs; i++) {
if (q[i][0] > qavg)
alpha = fmin(alpha, fmin(1, (qmax[i][0] - qavg)/(q[i][0] - qavg)));
else if (q[i][0] < qavg)
alpha = fmin(alpha, fmin(1, (qavg - qmin[i][0])/(qavg - q[i][0])));
}
for (int i=0; i<q.dofs; i++) {
q[i][0] = qavg + alpha*(q[i][0] - qavg);
}
"""
def _construct_centroid_solver(self):
"""
Constructs a linear problem for computing the centroids
:return: LinearSolver instance
"""
u = TrialFunction(self.P0)
v = TestFunction(self.P0)
a = assemble(u * v * dx)
return LinearSolver(a,
solver_parameters={
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
})
def _update_centroids(self, field):
"""
Update centroid values
"""
assemble(TestFunction(self.P0) * field * dx, tensor=self.centroids_rhs)
self.centroid_solver.solve(self.centroids, self.centroids_rhs)
def compute_bounds(self, field):
"""
Only computes min and max bounds of neighbouring cells
"""
self._update_centroids(field)
self.max_field.assign(-1.0e10) # small number
self.min_field.assign(1.0e10) # big number
par_loop(
self._min_max_loop, dx, {
"maxq": (self.max_field, MAX),
"minq": (self.min_field, MIN),
"q": (self.centroids, READ)
})
def apply_limiter(self, field):
"""
Only applies limiting loop on the given field
"""
par_loop(
self._limit_kernel, dx, {
"qbar": (self.centroids, READ),
"q": (field, RW),
"qmax": (self.max_field, READ),
"qmin": (self.min_field, READ)
})
def apply(self, field):
"""
Re-computes centroids and applies limiter to given field
"""
assert field.function_space() == self.P1DG, \
'Given field does not belong to this objects function space'
self.compute_bounds(field)
self.apply_limiter(field)
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
A, P = pc.getOperators()
self.cxt = A.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A_tensor = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A_tensor, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
Smat = self.S.petscmat
# If a different matrix is used for preconditioning,
# assemble this as well
if A != P:
self.cxt_pc = P.getPythonContext()
P_tensor = Tensor(self.cxt_pc.a)
P_reduced_sys = self.condensed_system(P_tensor,
self.residual,
elim_fields)
S_pc_expr = P_reduced_sys.lhs
self.S_pc_expr = S_pc_expr
# Allocate and set the condensed operator
self.S_pc = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S_pc = create_assembly_callable(
S_pc_expr,
tensor=self.S_pc,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S_pc()
Smat_pc = self.S_pc.petscmat
else:
self.S_pc_expr = S_expr
Smat_pc = Smat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Create a SNESContext for the DM associated with the trace problem
self._ctx_ref = self.new_snes_ctx(pc,
S_expr,
bcs,
mat_type,
self.cxt.fc_params,
options_prefix=prefix)
# Push new context onto the dm associated with the condensed problem
c_dm = Vc.dm
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
# Set the dm for the condensed solver
c_ksp.setDM(c_dm)
c_ksp.setDMActive(False)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(A=Smat, P=Smat_pc)
self.condensed_ksp = c_ksp
with dmhooks.add_hooks(c_dm, self,
appctx=self._ctx_ref,
save=False):
c_ksp.setFromOptions()
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A_tensor,
self.residual,
self.solution,
elim_fields)
class SCPC(SCBase):
needs_python_pmat = True
"""A Slate-based python preconditioner implementation of
static condensation for problems with up to three fields.
"""
@timed_function("SCPCInit")
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
A, P = pc.getOperators()
self.cxt = A.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A_tensor = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A_tensor, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
Smat = self.S.petscmat
# If a different matrix is used for preconditioning,
# assemble this as well
if A != P:
self.cxt_pc = P.getPythonContext()
P_tensor = Tensor(self.cxt_pc.a)
P_reduced_sys = self.condensed_system(P_tensor,
self.residual,
elim_fields)
S_pc_expr = P_reduced_sys.lhs
self.S_pc_expr = S_pc_expr
# Allocate and set the condensed operator
self.S_pc = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S_pc = create_assembly_callable(
S_pc_expr,
tensor=self.S_pc,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S_pc()
Smat_pc = self.S_pc.petscmat
else:
self.S_pc_expr = S_expr
Smat_pc = Smat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Create a SNESContext for the DM associated with the trace problem
self._ctx_ref = self.new_snes_ctx(pc,
S_expr,
bcs,
mat_type,
self.cxt.fc_params,
options_prefix=prefix)
# Push new context onto the dm associated with the condensed problem
c_dm = Vc.dm
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
# Set the dm for the condensed solver
c_ksp.setDM(c_dm)
c_ksp.setDMActive(False)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(A=Smat, P=Smat_pc)
self.condensed_ksp = c_ksp
with dmhooks.add_hooks(c_dm, self,
appctx=self._ctx_ref,
save=False):
c_ksp.setFromOptions()
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A_tensor,
self.residual,
self.solution,
elim_fields)
def condensed_system(self, A, rhs, elim_fields):
"""Forms the condensed linear system by eliminating
specified unknowns.
:arg A: A Slate Tensor containing the mixed bilinear form.
:arg rhs: A firedrake function for the right-hand side.
:arg elim_fields: An iterable of field indices to eliminate.
"""
from firedrake.slate.static_condensation.la_utils import condense_and_forward_eliminate
return condense_and_forward_eliminate(A, rhs, elim_fields)
def local_solver_calls(self, A, rhs, x, elim_fields):
"""Provides solver callbacks for inverting local operators
and reconstructing eliminated fields.
:arg A: A Slate Tensor containing the mixed bilinear form.
:arg rhs: A firedrake function for the right-hand side.
:arg x: A firedrake function for the solution.
:arg elim_fields: An iterable of eliminated field indices
to recover.
"""
from firedrake.slate.static_condensation.la_utils import backward_solve
from firedrake.assemble import create_assembly_callable
fields = x.split()
systems = backward_solve(A, rhs, x, reconstruct_fields=elim_fields)
local_solvers = []
for local_system in systems:
Ae = local_system.lhs
be = local_system.rhs
i, = local_system.field_idx
local_solve = Ae.solve(be, decomposition="PartialPivLU")
solve_call = create_assembly_callable(
local_solve,
tensor=fields[i],
form_compiler_parameters=self.cxt.fc_params)
local_solvers.append(solve_call)
return local_solvers
@timed_function("SCPCUpdate")
def update(self, pc):
"""Update by assembling into the KSP operator. No
need to reconstruct symbolic objects.
"""
self._assemble_S()
# Only reassemble if a preconditioning operator
# is provided for the condensed system
if hasattr(self, "S_pc"):
self._assemble_S_pc()
def forward_elimination(self, pc, x):
"""Perform the forward elimination of fields and
provide the reduced right-hand side for the condensed
system.
:arg pc: a Preconditioner instance.
:arg x: a PETSc vector containing the incoming right-hand side.
"""
with self.residual.dat.vec_wo as v:
x.copy(v)
# Now assemble residual for the reduced problem
self._assemble_Srhs()
def sc_solve(self, pc):
"""Solve the condensed linear system for the
condensed field.
:arg pc: a Preconditioner instance.
"""
dm = self.condensed_ksp.getDM()
with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
with self.condensed_rhs.dat.vec_ro as rhs:
if self.condensed_ksp.getInitialGuessNonzero():
acc = self.solution.split()[self.c_field].dat.vec
else:
acc = self.solution.split()[self.c_field].dat.vec_wo
with acc as sol:
self.condensed_ksp.solve(rhs, sol)
def backward_substitution(self, pc, y):
"""Perform the backwards recovery of eliminated fields.
:arg pc: a Preconditioner instance.
:arg y: a PETSc vector for placing the resulting fields.
"""
# Recover eliminated unknowns
for local_solver_call in self.local_solvers:
local_solver_call()
with self.solution.dat.vec_ro as w:
w.copy(y)
def view(self, pc, viewer=None):
"""Viewer calls for the various configurable objects in this PC."""
super(SCPC, self).view(pc, viewer)
if hasattr(self, "condensed_ksp"):
viewer.printfASCII("Solving linear system using static condensation.\n")
self.condensed_ksp.view(viewer=viewer)
viewer.printfASCII("Locally reconstructing unknowns.\n")
