Ensemble.sequential context manager (#4964)

Co-authored-by: David A. Ham <david.ham@imperial.ac.uk>

Author: JHopeCollins

firedrake/ensemble/ensemble.py

@@ -1,6 +1,8 @@
 from functools import wraps
 import weakref
+from contextlib import contextmanager
 from itertools import zip_longest
+from types import SimpleNamespace
 
 from firedrake.petsc import PETSc
 from firedrake.function import Function
@@ -584,3 +586,93 @@ def isendrecv(self, fsend: Function | Cofunction, dest: int, sendtag: int = 0,
         requests.extend([self._ensemble_comm.Irecv(dat.data, source=source, tag=recvtag)
                          for dat in frecv.dat])
         return requests
+
+    @contextmanager
+    def sequential(self, *, synchronise: bool = False, reverse: bool = False, **kwargs):
+        """
+        Context manager for executing code on each ensemble
+        member consecutively (ordered by increasing
+        :attr:`~.Ensemble.ensemble_rank`).
+
+        Any data in ``kwargs`` will be made available in the returned
+        context and will be communicated forward after each ensemble
+        member exits. :class:`.Function` or :class:`.Cofunction`
+        ``kwargs`` will be sent with the corresponding Ensemble methods.
+
+        For example:
+
+        .. code-block:: python3
+
+            with ensemble.sequential(index=0) as ctx:
+                print(ensemble.ensemble_rank, ctx.index)
+                ctx.index += 2
+
+        Would print:
+
+        .. code-block::
+
+            0 0
+            1 2
+            2 4
+            3 6
+            ...
+
+        If ``reverse is True`` then the ensemble ranks will be looped through
+        in decreasing order i.e. ``ensemble_rank == (ensemble_size - 1)`` will
+        run first, then ``ensemble_rank == (ensemble_size - 2)`` etc.
+
+        Parameters
+        ----------
+        synchronise :
+            If True then MPI_Barrier will be called on the ``global_comm``
+            at the beginning and end of this method.
+
+        reverse :
+            If True then will iterate through spatial comms in order of
+            decreasing ``ensemble_rank``.
+
+        kwargs :
+            Data to be passed forward by each rank and made available
+            in the returned ``ctx``.
+        """
+        rank = self.ensemble_rank
+        if reverse:  # send backwards
+            src = rank + 1
+            dst = rank - 1
+            first_rank = (rank == self.ensemble_size - 1)
+            last_rank = (rank == 0)
+        else:  # send forwards
+            src = rank - 1
+            dst = rank + 1
+            first_rank = (rank == 0)
+            last_rank = (rank == self.ensemble_size - 1)
+
+        if synchronise:
+            self.global_comm.Barrier()
+
+        if not first_rank:
+            for i, (k, v) in enumerate(kwargs.items()):
+                if isinstance(v, (Function, Cofunction)):
+                    # Functions are sent in-place, everything else is pickled
+                    recv_args = [kwargs[k]]
+                else:
+                    recv_args = []
+                kwargs[k] = self.recv(*recv_args, source=src, tag=rank+i*100)
+
+        ctx = SimpleNamespace(**kwargs)
+        yield ctx
+
+        if not last_rank:
+            for i, v in enumerate((getattr(ctx, k)
+                                   for k in kwargs.keys())):
+                try:
+                    self.send(v, dest=dst, tag=dst+i*100)
+                except Exception as error:
+                    raise TypeError(
+                        "Failed to send object of type {type(v)__name__}. kwargs for"
+                        " Ensemble.sequential must be Functions, Cofunctions,"
+                        " or acceptable arguments to mpi4py.MPI.Comm.send."
+                    ) from error
+
+        if synchronise:
+            self.global_comm.Barrier()

tests/firedrake/ensemble/test_ensemble.py

@@ -380,3 +380,43 @@ def test_ensemble_solvers(ensemble, W, urank, urank_sum):
     ensemble.allreduce(u_separate, usum)
 
     parallel_assert(errornorm(u_combined, usum) < 1e-8)
+
+
+@pytest.mark.parallel(nprocs=6)
+@pytest.mark.parametrize("direction", ["forward", "reverse"])
+def test_ensemble_sequential(ensemble, direction):
+    """
+    Test that the sequential context manager sends forward
+    the correct values after each rank has executed, for both
+    intrinsic types (float) and Firedrake types (Function).
+    """
+
+    rank = ensemble.ensemble_rank
+    mesh = UnitIntervalMesh(1, comm=ensemble.comm)
+    R = FunctionSpace(mesh, "R", 0)
+
+    reverse = direction == "reverse"
+
+    idx_i = 0
+    idx_f = Function(R).zero()
+    two = Function(R).assign(2)
+
+    with ensemble.sequential(reverse=reverse, idx_i=idx_i, idx_f=idx_f) as ctx:
+        recv_i = float(ctx.idx_i)
+        recv_f = float(ctx.idx_f)
+
+        ctx.idx_i += 2
+        ctx.idx_f += two
+
+    if reverse:
+        expected = 2*(ensemble.ensemble_size - 1 - rank)
+    else:
+        expected = 2*rank
+
+    parallel_assert(
+        recv_i == expected,
+        msg=f"Failed to send int properly. Expecting {expected} but received {recv_i}")
+
+    parallel_assert(
+        abs(float(recv_f)-expected) < 1e-12,
+        msg=f"Failed to send Function properly. Expecting {expected} but received {float(recv_f)}")