Source code for firedrake.adjoint_utils.mesh

from functools import wraps
from pyadjoint.overloaded_type import OverloadedType
from pyadjoint.tape import no_annotations




class MeshGeometryMixin(OverloadedType):
    @staticmethod
    def _ad_annotate_init(init):
        @wraps(init)
        def wrapper(self, *args, **kwargs):
            from firedrake.mesh import ExtrudedMeshTopology
            from .blocks import MeshInputBlock, MeshOutputBlock

            OverloadedType.__init__(self, *args, **kwargs)
            init(self, *args, **kwargs)
            self._ad_coordinate_space = None

            # attach information to the mesh coordinates, this does not work for
            # meshes with multiple cell types
            if not isinstance(self.topology, ExtrudedMeshTopology) and len(self.topology.dm_cell_types) == 1:
                f = self._coordinates_function
                f.block_class = MeshInputBlock
                f._ad_floating_active = True
                f._ad_args = [self]

                f._ad_output_args = [self]
                f.output_block_class = MeshOutputBlock
                f._ad_outputs = [self]
        return wrapper

    @no_annotations
    def _ad_create_checkpoint(self):
        return self.coordinates.copy(deepcopy=True)

    @no_annotations
    def _ad_restore_at_checkpoint(self, checkpoint):
        self.coordinates.assign(checkpoint)
        return self

    def _ad_function_space(self):
        if self._ad_coordinate_space is None:
            self._ad_coordinate_space = self.coordinates.function_space().ufl_function_space()
        return self._ad_coordinate_space