EnsembleBlockDiagonalMat and EnsembleBJacobiPC

firedrake/ensemble/__init__.py

@@ -1,3 +1,5 @@
-from firedrake.ensemble.ensemble import *  # noqa: F401
-from firedrake.ensemble.ensemble_function import *  # noqa: F401
-from firedrake.ensemble.ensemble_functionspace import *  # noqa: F401
+from firedrake.ensemble.ensemble import Ensemble  # noqa: F401
+from firedrake.ensemble.ensemble_function import EnsembleFunction, EnsembleCofunction  # noqa: F401
+from firedrake.ensemble.ensemble_functionspace import EnsembleFunctionSpace, EnsembleDualSpace  # noqa: F401
+from firedrake.ensemble.ensemble_mat import EnsembleBlockDiagonalMatCtx, EnsembleBlockDiagonalMat  # noqa: F401
+from firedrake.ensemble.ensemble_pc import EnsembleBJacobiPC  # noqa: F401

firedrake/ensemble/ensemble.py

@@ -7,8 +7,6 @@
 from firedrake.cofunction import Cofunction
 from pyop2.mpi import MPI, internal_comm
 
-__all__ = ("Ensemble", )
-
 
 def _ensemble_mpi_dispatch(func):
     """

firedrake/ensemble/ensemble_function.py

@@ -9,35 +9,36 @@
 from firedrake.function import Function
 from firedrake.norms import norm
 
-__all__ = ("EnsembleFunction", "EnsembleCofunction")
-
 
 class EnsembleFunctionBase(EnsembleFunctionMixin):
     """
     A mixed (co)function defined on a :class:`~.ensemble.Ensemble`.
     The subcomponents are distributed over the ensemble members, and
-    are specified locally in a :class:`~firedrake.EnsembleFunctionSpace`.
+    are specified locally in an
+    :class:`~.ensemble_functionspace.EnsembleFunctionSpace`.
 
     Parameters
     ----------
 
-    function_space : `~ensemble_functionspace.EnsembleFunctionSpace`.
+    function_space :
         The function space of the (co)function.
 
     Notes
     -----
-    Passing an `EnsembleDualSpace` to `EnsembleFunction`
-    will return an instance of :class:`~firedrake.EnsembleCofunction`.
+    Passing an :class:`~.ensemble_functionspace.EnsembleDualSpace` to
+    :class:`EnsembleFunction` will return an instance of :class:`EnsembleCofunction`.
 
     This class does not carry UFL symbolic information, unlike a
     :class:`~firedrake.function.Function`. UFL expressions can only be defined
     locally on each ensemble member using a `~firedrake.function.Function`
-    from `EnsembleFunction.subfunctions`.
+    from ``EnsembleFunction.subfunctions``.
 
     See Also
     --------
-    - Primal ensemble objects: :class:`~ensemble_functionspace.EnsembleFunctionSpace` and :class:`~firedrake.EnsembleFunction`.
-    - Dual ensemble objects: :class:`~firedrake.EnsembleDualSpace` and :class:`~firedrake.EnsembleCofunction`.
+    .ensemble_functionspace.EnsembleFunctionSpace
+    .ensemble_functionspace.EnsembleDualSpace
+    EnsembleFunction
+    EnsembleCofunction
     """
 
     @PETSc.Log.EventDecorator()
@@ -134,7 +135,7 @@ def assign(self, other, subsets=None):
             The value to assign from.
 
         subsets : Collection[Optional[:class:`pyop2.types.set.Subset`]]
-            One subset for each local :class:`firedrake.functionFunction`.
+            One subset for each local :class:`firedrake.function.Function`.
             None elements will be ignored. The values of each local function
             will only be assigned on the nodes on the corresponding subset.
         """
@@ -215,7 +216,7 @@ def __rmul__(self, other):
     @contextmanager
     def vec(self):
         """
-        Context manager for the global :class:`petsc4py.PETSc.Vec` with
+        Context manager for the global ``PETSc.Vec`` with
         read/write access.
 
         It is invalid to access the ``Vec`` outside of a context manager.
@@ -234,7 +235,7 @@ def vec(self):
     @contextmanager
     def vec_ro(self):
         """
-        Context manager for the global :class:`petsc4py.PETSc.Vec` with
+        Context manager for the global ``PETSc.Vec`` with
         read only access.
 
         It is invalid to access the ``Vec`` outside of a context manager.
@@ -250,7 +251,7 @@ def vec_ro(self):
     @contextmanager
     def vec_wo(self):
         """
-        Context manager for the global :class:`petsc4py.PETSc.Vec` with
+        Context manager for the global ``PETSc.Vec`` with
         write only access.
 
         It is invalid to access the ``Vec`` outside of a context manager.
@@ -272,7 +273,7 @@ class EnsembleFunction(EnsembleFunctionBase):
     """
     A mixed Function defined on a :class:`~.ensemble.Ensemble`.
     The subcomponents are distributed over the ensemble members, and
-    are specified locally in a :class:`~firedrake.ensemble.ensemble_functionspace.EnsembleFunctionSpace`.
+    are specified locally in an :class:`~firedrake.ensemble.ensemble_functionspace.EnsembleFunctionSpace`.
 
     Parameters
     ----------

firedrake/ensemble/ensemble_functionspace.py

@@ -7,8 +7,6 @@
 from firedrake.ensemble.ensemble import Ensemble
 from firedrake.functionspace import MixedFunctionSpace
 
-__all__ = ("EnsembleFunctionSpace", "EnsembleDualSpace")
-
 
 def _is_primal_or_dual(local_spaces, ensemble):
     """
@@ -83,13 +81,16 @@ class EnsembleFunctionSpaceBase:
     will return an instance of :class:`EnsembleDualSpace`.
 
     This class does not carry UFL symbolic information, unlike a
-    :class:`~firedrake.functionspaceimpl.FunctionSpace`. UFL expressions can only be defined locally
-    on each ensemble member using a :class:`~firedrake.functionspaceimpl.FunctionSpace` from
+    :class:`~firedrake.functionspace.FunctionSpace`. UFL expressions can only be defined locally
+    on each ensemble member using a :class:`~firedrake.functionspace.FunctionSpace` from
     `EnsembleFunctionSpace.local_spaces`.
 
-    See also:
-    - Primal ensemble objects: :class:`EnsembleFunctionSpace` and :class:`~firedrake.ensemble.ensemble_function.EnsembleFunction`.
-    - Dual ensemble objects: :class:`EnsembleDualSpace` and :class:`~firedrake.ensemble.ensemble_function.EnsembleCofunction`.
+    See Also
+    --------
+    EnsembleFunctionSpace
+    EnsembleDualSpace
+    .ensemble_function.EnsembleFunction
+    .ensemble_function.EnsembleCofunction
     """
     def __init__(self, local_spaces: Collection, ensemble: Ensemble):
         meshes = set(V.mesh().unique() for V in local_spaces)
@@ -143,7 +144,7 @@ def local_spaces(self):
         return self._local_spaces
 
     def mesh(self):
-        """The :class:`~firedrake.Mesh` on the local ensemble.comm.
+        """The :class:`~firedrake.mesh.Mesh` on the local ensemble.comm.
         """
         return self._mesh
 
@@ -156,7 +157,7 @@ def dual(self):
 
     @cached_property
     def nlocal_spaces(self):
-        """The total number of subspaces across all ensemble ranks.
+        """The number of subspaces on this ensemble rank.
         """
         return len(self.local_spaces)
 
@@ -184,6 +185,12 @@ def nglobal_dofs(self):
         """
         return self.ensemble_comm.allreduce(self.nlocal_comm_dofs)
 
+    @cached_property
+    def global_spaces_offset(self):
+        """Index of the first local subspace in the global mixed space.
+        """
+        return self.ensemble.ensemble_comm.exscan(self.nlocal_spaces) or 0
+
     def _component_indices(self, i):
         """
         Return the indices into the local mixed function storage
@@ -194,15 +201,19 @@ def _component_indices(self, i):
 
     def create_vec(self):
         """Return a PETSc Vec on the ``Ensemble.global_comm`` with the same layout
-        as a :class:`~firedrake.ensemble.ensemble_functionspace.EnsembleFunction`
-        or :class:`~firedrake.ensemble.ensemble_functionspace.EnsembleCofunction`
+        as an :class:`~firedrake.ensemble.ensemble_function.EnsembleFunction`
+        or :class:`~firedrake.ensemble.ensemble_function.EnsembleCofunction`
         in this function space.
         """
         vec = PETSc.Vec().create(comm=self.global_comm)
-        vec.setSizes((self.nlocal_dofs, self.nglobal_dofs))
+        vec.setSizes((self.nlocal_rank_dofs, self.nglobal_dofs))
         vec.setUp()
         return vec
 
+    @cached_property
+    def layout_vec(self):
+        return self.create_vec()
+
     def __eq__(self, other):
         if not isinstance(other, type(self)):
             local_eq = False