Source code for firedrake.projection

import abc
import ufl
import finat.ufl
from ufl.domain import extract_unique_domain

import firedrake
from firedrake.petsc import PETSc
from firedrake.utils import cached_property, complex_mode, SLATE_SUPPORTS_COMPLEX
from firedrake import functionspaceimpl
from firedrake import function
from firedrake.adjoint_utils import annotate_project
from finat import HDivTrace


__all__ = ['project', 'Projector']


def sanitise_input(v, V):
    if isinstance(v, function.Function):
        return v
    elif isinstance(v, ufl.classes.Expr):
        return v
    else:
        raise ValueError("Can't project from source object %r" % v)


def create_output(V, name=None):
    if isinstance(V, functionspaceimpl.WithGeometry):
        return function.Function(V, name=name)
    elif isinstance(V, function.Function):
        return V
    else:
        raise ValueError("Can't project into target object %r" % V)


def check_meshes(source, target):
    source_mesh = extract_unique_domain(source)
    target_mesh = extract_unique_domain(target)
    if source_mesh is None:
        source_mesh = target_mesh
    if target_mesh is None:
        raise ValueError("Target space must have a mesh")
    if source_mesh.ufl_cell() != target_mesh.ufl_cell():
        raise ValueError("Mismatching cells in source (%r) and target (%r) meshes" %
                         (source_mesh.ufl_cell(), target_mesh.ufl_cell()))
    return source_mesh, target_mesh


[docs] @PETSc.Log.EventDecorator() @annotate_project def project(v, V, bcs=None, solver_parameters=None, form_compiler_parameters=None, use_slate_for_inverse=True, name=None, ad_block_tag=None): """Project a UFL expression into a :class:`.FunctionSpace` It is possible to project onto the trace space 'DGT', but not onto other trace spaces e.g. into the restriction of CG onto the facets. :arg v: the :class:`ufl.core.expr.Expr` to project :arg V: the :class:`.FunctionSpace` or :class:`.Function` to project into :kwarg bcs: boundary conditions to apply in the projection :kwarg solver_parameters: parameters to pass to the solver used when projecting. :kwarg form_compiler_parameters: parameters to the form compiler :kwarg use_slate_for_inverse: compute mass inverse cell-wise using SLATE (ignored for non-DG function spaces). :kwarg name: name of the resulting :class:`.Function` :kwarg ad_block_tag: string for tagging the resulting block on the Pyadjoint tape If ``V`` is a :class:`.Function` then ``v`` is projected into ``V`` and ``V`` is returned. If `V` is a :class:`.FunctionSpace` then ``v`` is projected into a new :class:`.Function` and that :class:`.Function` is returned.""" val = Projector(v, V, bcs=bcs, solver_parameters=solver_parameters, form_compiler_parameters=form_compiler_parameters, use_slate_for_inverse=use_slate_for_inverse).project() val.rename(name) return val
class Assigner(object): def __init__(self, source, target): self.source = source self.target = target def project(self): self.target.assign(self.source) return self.target class ProjectorBase(object, metaclass=abc.ABCMeta): def __init__(self, source, target, bcs=None, solver_parameters=None, form_compiler_parameters=None, constant_jacobian=True, use_slate_for_inverse=True): if solver_parameters is None: solver_parameters = {} else: solver_parameters = solver_parameters.copy() solver_parameters.setdefault("ksp_type", "cg") solver_parameters.setdefault("ksp_rtol", 1e-8) mat_type = solver_parameters.get("mat_type", firedrake.parameters["default_matrix_type"]) if mat_type == "nest": solver_parameters.setdefault("pc_type", "fieldsplit") solver_parameters.setdefault("fieldsplit_pc_type", "bjacobi") solver_parameters.setdefault("fieldsplit_sub_pc_type", "icc") elif mat_type == "matfree": solver_parameters.setdefault("pc_type", "jacobi") else: solver_parameters.setdefault("pc_type", "bjacobi") solver_parameters.setdefault("sub_pc_type", "icc") self.source = source self.target = target self.solver_parameters = solver_parameters self.form_compiler_parameters = form_compiler_parameters self.bcs = bcs self.constant_jacobian = constant_jacobian try: element = self.target.function_space().finat_element is_dg = element.entity_dofs() == element.entity_closure_dofs() is_variable_layers = self.target.function_space().mesh().variable_layers except AttributeError: # Mixed space is_dg = False is_variable_layers = True self.use_slate_for_inverse = (use_slate_for_inverse and is_dg and not is_variable_layers and (not complex_mode or SLATE_SUPPORTS_COMPLEX)) @cached_property def A(self): u = firedrake.TrialFunction(self.target.function_space()) v = firedrake.TestFunction(self.target.function_space()) F = self.target.function_space() mixed = isinstance(F.ufl_element(), finat.ufl.MixedElement) if not mixed and isinstance(F.finat_element, HDivTrace): if F.extruded: a = (firedrake.inner(u, v)*firedrake.ds_t + firedrake.inner(u, v)*firedrake.ds_v + firedrake.inner(u, v)*firedrake.ds_b + firedrake.inner(u('+'), v('+'))*firedrake.dS_h + firedrake.inner(u('+'), v('+'))*firedrake.dS_v) else: a = (firedrake.inner(u, v)*firedrake.ds + firedrake.inner(u('+'), v('+'))*firedrake.dS) else: a = firedrake.inner(u, v)*firedrake.dx if self.use_slate_for_inverse: a = firedrake.Tensor(a).inv A = firedrake.assemble(a, bcs=self.bcs, mat_type=self.solver_parameters.get("mat_type"), form_compiler_parameters=self.form_compiler_parameters) return A @cached_property def solver(self): return firedrake.LinearSolver(self.A, solver_parameters=self.solver_parameters) @property def apply_massinv(self): if not self.constant_jacobian: firedrake.assemble(self.A.a, tensor=self.A, bcs=self.bcs, form_compiler_parameters=self.form_compiler_parameters) if self.use_slate_for_inverse: def solve(x, b): with x.dat.vec_wo as x_, b.dat.vec_ro as b_: self.A.petscmat.mult(b_, x_) return solve else: return self.solver.solve @cached_property def residual(self): return firedrake.Cofunction(self.target.function_space().dual()) @abc.abstractproperty def rhs(self): pass def project(self): self.apply_massinv(self.target, self.rhs) return self.target class BasicProjector(ProjectorBase): """ A basic projector projects a UFL expression into a function space and places the result in a function from that function space, allowing the solver to be reused. The difference to the :class:`.SupermeshProjector` is that both function spaces are defined on the same mesh. """ @cached_property def rhs_form(self): v = firedrake.TestFunction(self.target.function_space()) F = self.target.function_space() mixed = isinstance(F.ufl_element(), finat.ufl.MixedElement) if not mixed and isinstance(F.finat_element, HDivTrace): # Project onto a trace space by supplying the respective form on the facets. # The measures on the facets differ between extruded and non-extruded mesh. # FIXME The restrictions of cg onto the facets is also a trace space, # but we only cover DGT. if F.extruded: form = (firedrake.inner(self.source, v)*firedrake.ds_t + firedrake.inner(self.source, v)*firedrake.ds_v + firedrake.inner(self.source, v)*firedrake.ds_b + firedrake.inner(firedrake.avg(self.source), firedrake.avg(v))*firedrake.dS_h + firedrake.inner(firedrake.avg(self.source), firedrake.avg(v))*firedrake.dS_v) else: form = (firedrake.inner(self.source, v)*firedrake.ds + firedrake.inner(firedrake.avg(self.source), firedrake.avg(v))*firedrake.dS) else: form = firedrake.inner(self.source, v)*firedrake.dx return form @cached_property def assembler(self): from firedrake.assemble import get_assembler return get_assembler(self.rhs_form, form_compiler_parameters=self.form_compiler_parameters).assemble @property def rhs(self): self.assembler(tensor=self.residual) return self.residual class SupermeshProjector(ProjectorBase): @cached_property def mixed_mass(self): from firedrake.supermeshing import assemble_mixed_mass_matrix return assemble_mixed_mass_matrix(self.source.function_space(), self.target.function_space()) @property def rhs(self): with self.source.dat.vec_ro as u, self.residual.dat.vec_wo as v: self.mixed_mass.mult(u, v) return self.residual
[docs] @PETSc.Log.EventDecorator() def Projector(v, v_out, bcs=None, solver_parameters=None, form_compiler_parameters=None, constant_jacobian=True, use_slate_for_inverse=False): """ A projector projects a UFL expression into a function space and places the result in a function from that function space, allowing the solver to be reused. Projection reverts to an assign operation if ``v`` is a :class:`.Function` and belongs to the same function space as ``v_out``. It is possible to project onto the trace space 'DGT', but not onto other trace spaces e.g. into the restriction of CG onto the facets. :arg v: the :class:`ufl.core.expr.Expr` or :class:`.Function` to project :arg V: :class:`.Function` (or :class:`~.FunctionSpace`) to put the result in. :arg bcs: an optional set of :class:`.DirichletBC` objects to apply on the target function space. :arg solver_parameters: parameters to pass to the solver used when projecting. :arg constant_jacobian: Is the projection matrix constant between calls? Say False if you have moving meshes. :arg use_slate_for_inverse: compute mass inverse cell-wise using SLATE (only valid for DG function spaces). """ target = create_output(v_out) source = sanitise_input(v, target.function_space()) source_mesh, target_mesh = check_meshes(source, target) if source.ufl_shape != target.ufl_shape: raise ValueError("Shape mismatch between source %s and target %s in project" % (source.ufl_shape, target.ufl_shape)) if isinstance(v, function.Function) and not bcs and v.function_space() == target.function_space(): return Assigner(source, target) elif source_mesh == target_mesh: return BasicProjector(source, target, bcs=bcs, solver_parameters=solver_parameters, form_compiler_parameters=form_compiler_parameters, constant_jacobian=constant_jacobian, use_slate_for_inverse=use_slate_for_inverse) else: if bcs is not None: raise ValueError("Haven't implemented supermesh projection with boundary conditions yet, sorry!") if not isinstance(source, function.Function) or source.ufl_element().family() == "Real": raise NotImplementedError("Only for source Functions, not %s" % type(source)) return SupermeshProjector(source, target, bcs=bcs, solver_parameters=solver_parameters, form_compiler_parameters=form_compiler_parameters, constant_jacobian=constant_jacobian, use_slate_for_inverse=use_slate_for_inverse)