def interpolate(self, expression, subset=None):
"""Interpolate an expression onto this :class:`Function`.
:param expression: :class:`.Expression` or a UFL expression to interpolate
:returns: this :class:`Function` object"""
from firedrake import interpolation
return interpolation.interpolate(expression, self, subset=subset)
def initialize(self, obj):
if complex_mode:
raise NotImplementedError("HypreAMS preconditioner not yet implemented in complex mode")
Citations().register("Kolev2009")
A, P = obj.getOperators()
prefix = obj.getOptionsPrefix()
V = get_function_space(obj.getDM())
mesh = V.mesh()
family = str(V.ufl_element().family())
degree = V.ufl_element().degree()
if family != 'Nedelec 1st kind H(curl)' or degree != 1:
raise ValueError("Hypre AMS requires lowest order Nedelec elements! (not %s of degree %d)" % (family, degree))
P1 = FunctionSpace(mesh, "Lagrange", 1)
G = Interpolator(grad(TestFunction(P1)), V).callable().handle
pc = PETSc.PC().create(comm=obj.comm)
pc.incrementTabLevel(1, parent=obj)
pc.setOptionsPrefix(prefix + "hypre_ams_")
pc.setOperators(A, P)
pc.setType('hypre')
pc.setHYPREType('ams')
pc.setHYPREDiscreteGradient(G)
zero_beta = PETSc.Options(prefix).getBool("pc_hypre_ams_zero_beta_poisson", default=False)
if zero_beta:
pc.setHYPRESetBetaPoissonMatrix(None)
VectorP1 = VectorFunctionSpace(mesh, "Lagrange", 1)
pc.setCoordinates(interpolate(SpatialCoordinate(mesh), VectorP1).dat.data_ro.copy())
pc.setUp()
self.pc = pc
def interpolate(self, expression, subset=None):
"""Interpolate an expression onto this :class:`Function`.
:param expression: :class:`.Expression` or a UFL expression to interpolate
:returns: this :class:`Function` object"""
from firedrake import interpolation
return interpolation.interpolate(expression, self, subset=subset)
def assemble(self):
"""
Compute values and assemble interpolation matrix
Assembly function to compute the nonzero sparse matrix entries
and assemble the sparse matrix. Should only need to be called
once for each analysis and thus will return if ``is_assembled``
is ``True`` without re-assembling the matrix.
"""
if self.is_assembled:
return
mesh = self.function_space.ufl_domain()
W = VectorFunctionSpace(mesh, self.function_space.ufl_element())
X = interpolate(mesh.coordinates, W)
meshvals_local = np.array(X.dat.data_with_halos)
imin, imax = self.interp.getOwnershipRange()
# loop over all data points
for i in range(self.n_data):
cell = self.function_space.mesh().locate_cell(self.coords[i])
if not cell is None:
nodes = self.function_space.cell_node_list[cell]
points = meshvals_local[nodes]
interp_coords = interpolate_cell(self.coords[i], points)
for (node, val) in zip(nodes, interp_coords):
if node < self.n_mesh_local:
self.interp.setValue(imin + node, i, val)
self.interp.assemble()
self.is_assembled = True
def interpolate(self, expression, subset=None, ad_block_tag=None):
r"""Interpolate an expression onto this :class:`Function`.
:param expression: a UFL expression to interpolate
:param ad_block_tag: string for tagging the resulting block on the Pyadjoint tape
:returns: this :class:`Function` object"""
from firedrake import interpolation
return interpolation.interpolate(expression, self, subset=subset, ad_block_tag=ad_block_tag)
def sort_entities(self, dm, axis, dir, ndiv):
# compute
# [(pStart, (x, y, z)), (pEnd, (x, y, z))]
mesh = dm.getAttr("__firedrake_mesh__")
ele = mesh.coordinates.function_space().ufl_element()
V = mesh.coordinates.function_space()
if V.finat_element.entity_dofs(
) == V.finat_element.entity_closure_dofs():
# We're using DG or DQ for our coordinates, so we got
# a periodic mesh. We need to interpolate to CGk
# with access descriptor MAX to define a consistent opinion
# about where the vertices are.
CGkele = ele.reconstruct(family="Lagrange")
# Need to supply the actual mesh to the FunctionSpace constructor,
# not its weakref proxy (the variable `mesh`)
# as interpolation fails because they are not hashable
CGk = FunctionSpace(mesh.coordinates.function_space().mesh(),
CGkele)
coordinates = interpolate(mesh.coordinates, CGk, access=op2.MAX)
else:
coordinates = mesh.coordinates
select = partial(select_entity, dm=dm, exclude="pyop2_ghost")
entities = [(p, self.coords(dm, p, coordinates))
for p in filter(select, range(*dm.getChart()))]
minx = min(entities, key=lambda z: z[1][axis])[1][axis]
maxx = max(entities, key=lambda z: z[1][axis])[1][axis]
def keyfunc(z):
coords = tuple(z[1])
return (coords[axis], ) + tuple(coords[:axis] + coords[axis + 1:])
s = sorted(entities, key=keyfunc, reverse=(dir == -1))
divisions = numpy.linspace(minx, maxx, ndiv + 1)
(entities, coords) = zip(*s)
coords = [c[axis] for c in coords]
indices = numpy.searchsorted(coords[::dir], divisions)
out = []
for k in range(ndiv):
out.append(entities[indices[k]:indices[k + 1]])
out.append(entities[indices[-1]:])
return out
def _compute_G_vals(self):
"""
Compute nonzero values of the covariance matrix and the number of nonzero elements
This method pre-computes the values of the covariance function, storing those above the
cutoff. The nonzero elements for each row are stored in a dictionary and returned with
the maximum number found over all rows for this process. This is needed to allocate
memory for the sparse matrix and fill with the computed rows.
:returns: tuple holding a dictionary mapping row indices to a tuple containing the
matrix entries (as a numpy array) and the integer indices of the columns
represented by those matrix entries.
:rtype: tuple
"""
G_dict = {}
nnz = 0
int_basis = self._integrate_basis_functions()
mesh = self.function_space.ufl_domain()
W = VectorFunctionSpace(mesh, self.function_space.ufl_element())
X = interpolate(mesh.coordinates, W)
meshvals = X.vector().gather()
if X.dat.data.ndim == 2:
meshvals = np.reshape(meshvals, (-1, X.dat.data.shape[1]))
else:
meshvals = np.reshape(meshvals, (-1, 1))
assert meshvals.shape[
0] == self.nx, "error in gathering mesh coordinates"
for i in range(self.local_startind, self.local_endind):
row = (int_basis[i] * int_basis *
self.cov(meshvals[i], meshvals, self.sigma, self.l))[0]
row[i] += self.regularization
above_cutoff = (row / row[i] > self.cutoff)
G_dict[i] = (row[above_cutoff],
np.arange(0, self.nx,
dtype=PETSc.IntType)[above_cutoff])
new_nnz = int(np.sum(above_cutoff))
if new_nnz > nnz:
nnz = new_nnz
return G_dict, nnz
def initialize(self, obj):
A, P = obj.getOperators()
prefix = obj.getOptionsPrefix()
V = get_function_space(obj.getDM())
mesh = V.mesh()
family = str(V.ufl_element().family())
degree = V.ufl_element().degree()
if family != 'Raviart-Thomas' or degree != 1:
raise ValueError(
"Hypre ADS requires lowest order RT elements! (not %s of degree %d)"
% (family, degree))
P1 = FunctionSpace(mesh, "Lagrange", 1)
NC1 = FunctionSpace(mesh, "N1curl", 1)
# DiscreteGradient
G = Interpolator(grad(TestFunction(P1)), NC1).callable().handle
# DiscreteCurl
C = Interpolator(curl(TestFunction(NC1)), V).callable().handle
pc = PETSc.PC().create(comm=obj.comm)
pc.incrementTabLevel(1, parent=obj)
pc.setOptionsPrefix(prefix + "hypre_ads_")
pc.setOperators(A, P)
pc.setType('hypre')
pc.setHYPREType('ads')
pc.setHYPREDiscreteGradient(G)
pc.setHYPREDiscreteCurl(C)
V = VectorFunctionSpace(mesh, "Lagrange", 1)
linear_coordinates = interpolate(SpatialCoordinate(mesh),
V).dat.data_ro.copy()
pc.setCoordinates(linear_coordinates)
pc.setUp()
self.pc = pc
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 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 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)
