Performance Regresson 2

firedrake/adjoint/ensemble_reduced_functional.py

@@ -28,8 +28,6 @@ class EnsembleReducedFunctional(ReducedFunctional):
     operation is employed to sum the functionals and their gradients over an ensemble
     communicator.
 
-    If gather_functional is present, then all the values of J are communicated to all ensemble ranks, and passed in a list to gather_functional, which is a reduced functional that expects a list of that size of the relevant types.
-
     Parameters
     ----------
     J : pyadjoint.OverloadedType
@@ -43,9 +41,6 @@ class EnsembleReducedFunctional(ReducedFunctional):
     scatter_control : bool
         Whether scattering a control (or a list of controls) over the ensemble communicator
         ``Ensemble.ensemble comm``.
-    gather_functional : An instance of the :class:`pyadjoint.ReducedFunctional`.
-        that takes in all of the Js.
-
 
     See Also
     --------
@@ -58,35 +53,15 @@ class EnsembleReducedFunctional(ReducedFunctional):
     works, please refer to the `Firedrake manual
     <https://www.firedrakeproject.org/parallelism.html#id8>`_.
     """
-    def __init__(self, J, control, ensemble, scatter_control=True,
-                 gather_functional=None):
+    def __init__(self, J, control, ensemble, scatter_control=True):
         super(EnsembleReducedFunctional, self).__init__(J, control)
         self.ensemble = ensemble
         self.scatter_control = scatter_control
-        self.gather_functional = gather_functional
-
-    def _allgather_J(self, J):
-        if isinstance(J, float):
-            vals = self.ensemble.ensemble_comm.allgather(J)
-        elif isinstance(J, firedrake.Function):
-            #  allgather not implemented in ensemble.py
-            vals = []
-            for i in range(self.ensemble.ensemble_comm.size):
-                J0 = J.copy(deepcopy=True)
-                vals.append(self.ensemble.bcast(J0, root=i))
-        else:
-            raise NotImplementedError(f"Functionals of type {type(J).__name__} are not supported.")
-        return vals
 
     def __call__(self, values):
         local_functional = super(EnsembleReducedFunctional, self).__call__(values)
-        ensemble_comm = self.ensemble.ensemble_comm
-        if self.gather_functional:
-            controls_g = self._allgather_J(local_functional)
-            total_functional = self.gather_functional(controls_g)
-        # if gather_functional is None then we do a sum
-        elif isinstance(local_functional, float):
-            total_functional = ensemble_comm.allreduce(sendobj=local_functional, op=MPI.SUM)
+        if isinstance(local_functional, float):
+            total_functional = self.ensemble.ensemble_comm.allreduce(sendobj=local_functional, op=MPI.SUM)
         elif isinstance(local_functional, firedrake.Function):
             total_functional = type(local_functional)(local_functional.function_space())
             total_functional = self.ensemble.allreduce(local_functional, total_functional)
@@ -113,19 +88,10 @@ def derivative(self, adj_input=1.0, options=None):
         --------
         :meth:`~.ensemble.Ensemble.allreduce`, :meth:`pyadjoint.ReducedFunctional.derivative`.
         """
-
-        if self.gather_functional:
-            dJg_dmg = self.gather_functional.derivative(adj_input=adj_input,
-                                                        options=options)
-            i = self.ensemble.ensemble_comm.rank
-            adj_input = dJg_dmg[i]
-
         dJdm_local = super(EnsembleReducedFunctional, self).derivative(adj_input=adj_input, options=options)
-
+        dJdm_local = Enlist(dJdm_local)
+        dJdm_total = []
         if self.scatter_control:
-            dJdm_local = Enlist(dJdm_local)
-            dJdm_total = []
-
             for dJdm in dJdm_local:
                 if not isinstance(dJdm, (firedrake.Function, float)):
                     raise NotImplementedError("This type of gradient is not supported.")

firedrake/ensemble.py

@@ -151,9 +151,9 @@ def bcast(self, f, root=0):
         :raises ValueError: if function communicator mismatches the ensemble spatial communicator.
         """
         self._check_function(f)
+
         with f.dat.vec as vec:
             self._ensemble_comm.Bcast(vec.array, root=root)
-
         return f
 
     @PETSc.Log.EventDecorator()

tests/ensemble_reduced_functional/test_reduced_functional.py

@@ -40,53 +40,6 @@ def test_verification():
     assert taylor_test(rf, x, Function(R, val=0.1))
 
 
-@pytest.mark.parallel(nprocs=4)
-@pytest.mark.skipcomplex  # Taping for complex-valued 0-forms not yet done
-def test_verification_gather_functional_adjfloat():
-    ensemble = Ensemble(COMM_WORLD, 2)
-    rank = ensemble.ensemble_comm.rank
-    mesh = UnitSquareMesh(4, 4, comm=ensemble.comm)
-    R = FunctionSpace(mesh, "R", 0)
-    x = function.Function(R, val=rank+1)
-    J = assemble(x * x * dx(domain=mesh))
-    a = AdjFloat(1.0)
-    b = AdjFloat(1.0)
-    Jg_m = [Control(a), Control(b)]
-    Jg = ReducedFunctional(a**2 + b**2, Jg_m)
-    rf = EnsembleReducedFunctional(J, Control(x), ensemble,
-                                   scatter_control=False,
-                                   gather_functional=Jg)
-    ensemble_J = rf(x)
-    dJdm = rf.derivative()
-    assert_allclose(ensemble_J, 1.0**4+2.0**4, rtol=1e-12)
-    assert_allclose(dJdm.dat.data_ro, 4*(rank+1)**3, rtol=1e-12)
-    assert taylor_test(rf, x, Function(R, val=0.1))
-
-
-@pytest.mark.parallel(nprocs=4)
-@pytest.mark.skipcomplex  # Taping for complex-valued 0-forms not yet done
-def test_verification_gather_functional_Function():
-    ensemble = Ensemble(COMM_WORLD, 2)
-    rank = ensemble.ensemble_comm.rank
-    mesh = UnitSquareMesh(4, 4, comm=ensemble.comm)
-    R = FunctionSpace(mesh, "R", 0)
-    x = function.Function(R, val=rank+1)
-    J = Function(R).assign(x**2)
-    a = Function(R).assign(1.0)
-    b = Function(R).assign(1.0)
-    Jg_m = [Control(a), Control(b)]
-    Jg = assemble((a**2 + b**2)*dx)
-    Jghat = ReducedFunctional(Jg, Jg_m)
-    rf = EnsembleReducedFunctional(J, Control(x), ensemble,
-                                   scatter_control=False,
-                                   gather_functional=Jghat)
-    ensemble_J = rf(x)
-    dJdm = rf.derivative()
-    assert_allclose(ensemble_J, 1.0**4+2.0**4, rtol=1e-12)
-    assert_allclose(dJdm.dat.data_ro, 4*(rank+1)**3, rtol=1e-12)
-    assert taylor_test(rf, x, Function(R, val=0.1))
-
-
 @pytest.mark.parallel(nprocs=6)
 @pytest.mark.skipcomplex  # Taping for complex-valued 0-forms not yet done
 def test_minimise():