diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 57e633fa..972187d6 100644 --- a/pyadjoint/adjfloat.py +++ b/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 diff --git a/pyadjoint/optimization/tao_solver.py b/pyadjoint/optimization/tao_solver.py index dce9ef00..15503c00 100644 --- a/pyadjoint/optimization/tao_solver.py +++ b/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,21 +38,19 @@ 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 @@ -59,6 +58,15 @@ def __init__(self, X, *, comm=None): 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) diff --git a/pyadjoint/overloaded_type.py b/pyadjoint/overloaded_type.py index a6a02f68..84cac36f 100644 --- a/pyadjoint/overloaded_type.py +++ b/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.