Merge pull request #1 from Ig-dolci/dolci/test_concatenatevec_tao

Dolci/test concatenatevec tao

pyadjoint/adjfloat.py

@@ -128,6 +128,12 @@ def _ad_str(self):
         """Return the string of the taped value of this variable."""
         return str(self.block_variable.saved_output)
 
+    def _ad_to_petsc(self):
+        NotImplementedError("It requires more thought to return a PETSc Vec from `AdjFloat`")
+
+    def _ad_from_petsc(self, vec):
+        NotImplementedError("It requires more thought.")
+
 
 _min = min
 _max = max

pyadjoint/optimization/tao_solver.py

@@ -1,4 +1,6 @@
 from functools import cached_property
+from sys import maxsize
+from enum import Enum
 import numbers
 import weakref
 
@@ -7,7 +9,6 @@
 from .optimization_solver import OptimizationSolver
 
 
-import numpy as np
 try:
     import petsc4py.PETSc as PETSc
 except ModuleNotFoundError:
@@ -37,28 +38,35 @@ def __init__(self, X, *, comm=None):
             raise RuntimeError("PETSc not available")
 
         X = Enlist(X)
+        vecs, indices = self._get_vecs_indexes(X)
+        _, isis = PETSc.Vec().concatenate(vecs)
+        self._concatenate_vecs = vecs
+        self._isis_concatenated_vecs = [i for i in isis]
         if comm is None:
             comm = PETSc.COMM_WORLD
         if hasattr(comm, "tompi4py"):
             comm = comm.tompi4py()
-
-        indices = []
         n = 0
         N = 0
-        for x in X:
-            y = x._ad_copy()
-            y_a = np.zeros(y._ad_dim(), dtype=PETSc.ScalarType)
-            _, x_n = y._ad_assign_numpy(y, y_a, offset=0)
-            del y, y_a
-            indices.append((n, n + x_n))
-            n += x_n
+        for x, i in zip(X, indices):
+            indices.append((n, n + i))
+            n += i
             N += x._ad_dim()
 
         self._comm = comm
         self._indices = tuple(indices)
         self._n = n
         self._N = N
 
+    def _get_vecs_indexes(self, X):
+        indices = []
+        vecs = []
+        for x in X:
+            vec = x._ad_to_petsc()
+            indices.append(vec.getLocalSize())
+            vecs.append(vec)
+        return vecs, indices
+
     @property
     def comm(self):
         """Communicator.
@@ -70,7 +78,6 @@ def comm(self):
     def indices(self):
         """Local index ranges for variables.
         """
-
         return self._indices
 
     @property
@@ -111,38 +118,25 @@ def from_petsc(self, y, X):
         """
 
         X = Enlist(X)
-        y_a = y.getArray(True)
-
-        if y_a.shape != (self.n,):
-            raise ValueError("Invalid shape")
-        if len(X) != len(self.indices):
+        if len(X) != len(self._isis_concatenated_vecs):
             raise ValueError("Invalid length")
-
-        for (i0, i1), x in zip(self.indices, X):
-            _, x_i1 = x._ad_assign_numpy(x, y_a, offset=i0)
-            if i1 != x_i1:
-                raise ValueError("Invalid index")
+        for iset, x in zip(self._isis_concatenated_vecs, X):
+            x._ad_from_petsc(y.getSubVector(iset))
 
     def to_petsc(self, x, Y):
         """Copy data from variables to a :class:`petsc4py.PETSc.Vec`.
 
         Args:
             x (petsc4py.PETSc.Vec): The output :class:`petsc4py.PETSc.Vec`.
-            Y (numbers.Complex, OverloadedType or Sequence[OverloadedType]):
+            Y (OverloadedType or Sequence[OverloadedType]):
                 Values for input variables.
         """
 
         Y = Enlist(Y)
-        if len(Y) != len(self.indices):
-            raise ValueError("Invalid length")
-
-        x_a = np.zeros(self.n, dtype=PETSc.ScalarType)
-        for (i0, i1), y in zip(self.indices, Y):
-            if isinstance(y, numbers.Complex):
-                x_a[i0:i1] = y
-            else:
-                x_a[i0:i1] = y._ad_to_list(y)
-        x.setArray(x_a)
+        for iset, y in zip(self._isis_concatenated_vecs, Y):
+            v_i = x.getSubVector(iset)
+            y._ad_to_petsc().copy(result=v_i)
+            x.restoreSubVector(iset, v_i)
 
 
 class PETScOptions:
@@ -187,6 +181,11 @@ def update(self, d):
             self[key] = value
 
 
+class BoundType(Enum):
+    LOWER = 0
+    UPPER = 1
+
+
 class TAOObjective:
     def __init__(self, rf):
         self.reduced_functional = rf
@@ -265,7 +264,6 @@ def __init__(self, problem, parameters, *, comm=None, convert_options=None):
             comm=comm)
         n, N = vec_interface.n, vec_interface.N
         to_petsc, from_petsc = vec_interface.to_petsc, vec_interface.from_petsc
-
         tao = PETSc.TAO().create(comm=comm)
 
         def objective_gradient(tao, x, g):
@@ -325,13 +323,36 @@ def mult(self, A, x, y):
         tao.setGradientNorm(M_inv_matrix)
 
         if problem.bounds is not None:
-            x_lb = vec_interface.new_petsc()
-            x_ub = vec_interface.new_petsc()
-            to_petsc(x_lb, tuple(np.finfo(PETSc.ScalarType).min if lb is None else lb for lb, _ in problem.bounds))
-            to_petsc(x_ub, tuple(np.finfo(PETSc.ScalarType).max if ub is None else ub for _, ub in problem.bounds))
-            tao.setVariableBounds(x_lb, x_ub)
-            x_lb.destroy()
-            x_ub.destroy()
+            if not all(len(bound) == 2 for bound in problem.bounds):
+                raise ValueError("Each bound should be a tuple of length 2 (lb, ub)")
+            if not len(problem.bounds) == len(problem.reduced_functional.controls):
+                raise ValueError(
+                    "bounds should be of length number of controls of the ReducedFunctional"
+                )
+
+            new_concatenate_vecs_lb = vec_interface.new_petsc()
+            new_concatenate_vecs_ub = vec_interface.new_petsc()
+            lbs = []
+            ubs = []
+            for bound, control in zip(
+                problem.bounds, taoobjective.reduced_functional.controls
+            ):
+                lb, ub = bound
+                lb = self._prepared_bound(
+                    control, lb, bound_type=BoundType.LOWER
+                )
+                ub = self._prepared_bound(
+                    control, ub, bound_type=BoundType.UPPER
+                )
+
+                lbs.append(lb)
+                ubs.append(ub)
+            to_petsc(new_concatenate_vecs_lb, lbs)
+            to_petsc(new_concatenate_vecs_ub, ubs)
+
+            tao.setVariableBounds(
+                new_concatenate_vecs_lb, new_concatenate_vecs_ub
+            )
 
         options = PETScOptions(f"_pyadjoint__{tao.name:s}_")
         options.update(parameters)
@@ -392,15 +413,33 @@ def finalize_callback(*args):
             tao, H_matrix, M_inv_matrix, B_0_matrix_pc, B_0_matrix, x)
         finalize.atexit = False
 
+    def _prepared_bound(self, control, bound, bound_type=BoundType.UPPER):
+        if bound is None:
+            bound = control.control._ad_copy()
+            if bound_type == BoundType.LOWER:
+                bound._ad_assign(-maxsize)
+            else:
+                bound._ad_assign(maxsize)
+        elif isinstance(bound, numbers.Number):
+            bound_num = bound
+            bound = control.control._ad_copy()
+            bound._ad_assign(bound_num)
+        elif not isinstance(bound, type(control)):
+            raise TypeError(
+                "This bound {bound} should be None, a float, or a %s." % type(control)
+            )
+        return bound
+
     def solve(self):
         """Solve the optimisation problem.
 
         Returns:
             OverloadedType or tuple[OverloadedType]: The solution.
         """
-
-        M = tuple(m.tape_value()._ad_copy()
-                  for m in self.taoobjective.reduced_functional.controls)
+        M = tuple(
+            m.tape_value()._ad_copy()
+            for m in self.taoobjective.reduced_functional.controls
+        )
         self.vec_interface.to_petsc(self.x, M)
         self.tao.solve()
         self.vec_interface.from_petsc(self.x, M)

pyadjoint/overloaded_type.py

@@ -142,6 +142,19 @@ def _ad_restore_at_checkpoint(self, checkpoint):
         """
         raise NotImplementedError
 
+    def _ad_assign(self, other):
+        """This method must be overridden.
+
+        The method should implement a routine for assigning other to `self`.
+
+        Note:
+            This method should not be employed for immutable types.
+
+        Args:
+            other (obj): The object which should be assigned to `self`.
+        """
+        raise NotImplementedError
+
     def _ad_mul(self, other):
         """This method must be overridden.
 
@@ -240,6 +253,31 @@ def _ad_will_add_as_output(self):
         """
         return True
 
+    def _ad_to_petsc(self):
+        """This method must be overridden.
+
+        This method should implement a routine to return a new instance that is
+        a copy of the PETSc Vec associated with the overloaded object.
+
+        Returns:
+            PETSc.Vec: A new instance that is a copy of the PETSc Vec of the
+                overloaded object.
+        """
+        raise NotImplementedError
+
+    def _ad_from_petsc(self, vec):
+        """This method must be overridden.
+
+        The method should implement a routine to update the PETSc Vec associated
+        with the overloaded object from the given PETSc Vec `vec`.
+
+        Args:
+            vec (PETSc.Vec): The PETSc Vec from which to update the `self`
+                associated PETSc Vec.
+
+        """
+        raise NotImplementedError
+
     @staticmethod
     def _ad_assign_numpy(dst, src, offset):
         """This method must be overridden.