From 2ba3d6c9e71b927606ab6747c5e09b901bc9b6ba Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 16 Apr 2026 13:39:41 +0100 Subject: [PATCH 1/3] do not import petsc until we really need it --- pyadjoint/optimization/tao_solver.py | 75 ++++++++++++++++++---------- 1 file changed, 48 insertions(+), 27 deletions(-) diff --git a/pyadjoint/optimization/tao_solver.py b/pyadjoint/optimization/tao_solver.py index 000d0138..a279e4b5 100644 --- a/pyadjoint/optimization/tao_solver.py +++ b/pyadjoint/optimization/tao_solver.py @@ -1,5 +1,6 @@ from enum import Enum from numbers import Complex +from functools import cached_property import numpy as np @@ -9,14 +10,33 @@ from .optimization_solver import OptimizationSolver -try: - import petsc4py.PETSc as PETSc -except ModuleNotFoundError: - PETSc = None -try: - import petsctools -except ModuleNotFoundError: - petsctools = None +def import_petsctools(): + """Return the petsctools module. + + Delay importing petsctools so that we can put TAOSolver + in the top-level pyadjoint namespace without requiring + petsctools or petsc4py to be installed. + """ + try: + import petsctools + return petsctools + except ModuleNotFoundError as err: + raise RuntimeError("petsctools not available") from err + + +def import_petsc(): + """Return the petsc4py.PETSc module. + + Delay importing petsc4py so that we can put TAOSolver + in the top-level pyadjoint namespace without requiring + petsctools or petsc4py to be installed. + """ + try: + from petsc4py import PETSc + return PETSc + except ModuleNotFoundError as err: + raise RuntimeError("PETSc not available") from err + __all__ = [ "TAOConvergenceError", @@ -37,10 +57,7 @@ class PETScVecInterface: """ def __init__(self, x, *, comm=None): - if PETSc is None: - raise RuntimeError("PETSc not available") - if petsctools is None: - raise RuntimeError("petsctools not available") + PETSc = import_petsc() x = Enlist(x) comm = valid_comm(comm) @@ -80,6 +97,7 @@ def new_petsc(self): Returns: petsc4py.PETSc.Vec: The new :class:`petsc4py.PETSc.Vec`. """ + PETSc = import_petsc() vec = PETSc.Vec().create(comm=self.comm) vec.setSizes((self.n, self.N)) @@ -152,6 +170,7 @@ def valid_comm(comm): petsc4py.PETSc.COMM_WORLD if `comm is None`, otherwise `comm.tompi4py()`. """ if comm is None: + PETSc = import_petsc() comm = PETSc.COMM_WORLD if hasattr(comm, "tompi4py"): comm = comm.tompi4py() @@ -435,6 +454,7 @@ def ReducedFunctionalMat(rf, action=RFOperation.HESSIAN, *, apply_riesz=False, a be reevaluated at every call to `mult`. comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the rf is defined over. """ + PETSc = import_petsc() if action == RFOperation.HESSIAN: ctx = ReducedFunctionalHessianMat( rf, appctx=appctx, apply_riesz=apply_riesz, @@ -511,6 +531,7 @@ def RieszMapMat(controls, symmetric=True, comm=None): symmetric (bool): Whether the Riesz map attached to the Control is symmetric. comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the controls are defined over. """ + PETSc = import_petsc() ctx = RieszMapMatCtx(controls, comm=comm) n = ctx.vec_interface.n @@ -623,16 +644,6 @@ class TAOConvergenceError(Exception): """ -if PETSc is None: - _tao_reasons = {} -else: - # Same approach as in _make_reasons in firedrake/solving_utils.py, - # Firedrake master branch 57e21cc8ebdb044c1d8423b48f3dbf70975d5548 - _tao_reasons = {getattr(PETSc.TAO.Reason, key): key - for key in dir(PETSc.TAO.Reason) - if not key.startswith("_")} - - class TAOSolver(OptimizationSolver): """Use TAO to solve an optimization problem. @@ -648,10 +659,8 @@ class TAOSolver(OptimizationSolver): def __init__(self, problem, parameters, *, options_prefix=None, appctx=None, Pmat=None, comm=None): - if PETSc is None: - raise RuntimeError("PETSc not available") - if petsctools is None: - raise RuntimeError("petsctools not available") + PETSc = import_petsc() + petsctools = import_petsctools() if not isinstance(problem, MinimizationProblem): raise TypeError("MinimizationProblem required") @@ -795,12 +804,24 @@ def x(self): return self._x + @cached_property + def _tao_reasons(self): + """Dictionary of TAO convergence reason int codes -> python objects + """ + PETSc = import_petsc() + # Same approach as in _make_reasons in firedrake/solving_utils.py, + # Firedrake master branch 57e21cc8ebdb044c1d8423b48f3dbf70975d5548 + return {getattr(PETSc.TAO.Reason, key): key + for key in dir(PETSc.TAO.Reason) + if not key.startswith("_")} + def solve(self): """Solve the optimization problem. Returns: OverloadedType or Sequence[OverloadedType]: The solution. """ + petsctools = import_petsctools() controls = self.tao_objective.reduced_functional.controls m = tuple(control.tape_value()._ad_copy() for control in controls) @@ -814,7 +835,7 @@ def solve(self): # Using the same format as Firedrake linear solver errors raise TAOConvergenceError( f"TAOSolver failed to converge after {self.tao.getIterationNumber()} iterations " - f"with reason: {_tao_reasons[self.tao.getConvergedReason()]}") + f"with reason: {self._tao_reasons[self.tao.getConvergedReason()]}") if isinstance(controls, Enlist): return controls.delist(m) else: From 717088add1546dc153195e2879ae48905415122f Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 17 Apr 2026 09:40:53 +0100 Subject: [PATCH 2/3] import petsctools and PETSc inside each function in TAO solver --- pyadjoint/optimization/tao_solver.py | 46 ++++++---------------------- 1 file changed, 9 insertions(+), 37 deletions(-) diff --git a/pyadjoint/optimization/tao_solver.py b/pyadjoint/optimization/tao_solver.py index a279e4b5..a341dbea 100644 --- a/pyadjoint/optimization/tao_solver.py +++ b/pyadjoint/optimization/tao_solver.py @@ -10,34 +10,6 @@ from .optimization_solver import OptimizationSolver -def import_petsctools(): - """Return the petsctools module. - - Delay importing petsctools so that we can put TAOSolver - in the top-level pyadjoint namespace without requiring - petsctools or petsc4py to be installed. - """ - try: - import petsctools - return petsctools - except ModuleNotFoundError as err: - raise RuntimeError("petsctools not available") from err - - -def import_petsc(): - """Return the petsc4py.PETSc module. - - Delay importing petsc4py so that we can put TAOSolver - in the top-level pyadjoint namespace without requiring - petsctools or petsc4py to be installed. - """ - try: - from petsc4py import PETSc - return PETSc - except ModuleNotFoundError as err: - raise RuntimeError("PETSc not available") from err - - __all__ = [ "TAOConvergenceError", "TAOSolver" @@ -57,7 +29,7 @@ class PETScVecInterface: """ def __init__(self, x, *, comm=None): - PETSc = import_petsc() + from petsc4py import PETSc x = Enlist(x) comm = valid_comm(comm) @@ -97,7 +69,7 @@ def new_petsc(self): Returns: petsc4py.PETSc.Vec: The new :class:`petsc4py.PETSc.Vec`. """ - PETSc = import_petsc() + from petsc4py import PETSc vec = PETSc.Vec().create(comm=self.comm) vec.setSizes((self.n, self.N)) @@ -170,7 +142,7 @@ def valid_comm(comm): petsc4py.PETSc.COMM_WORLD if `comm is None`, otherwise `comm.tompi4py()`. """ if comm is None: - PETSc = import_petsc() + from petsc4py import PETSc comm = PETSc.COMM_WORLD if hasattr(comm, "tompi4py"): comm = comm.tompi4py() @@ -454,7 +426,7 @@ def ReducedFunctionalMat(rf, action=RFOperation.HESSIAN, *, apply_riesz=False, a be reevaluated at every call to `mult`. comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the rf is defined over. """ - PETSc = import_petsc() + from petsc4py import PETSc if action == RFOperation.HESSIAN: ctx = ReducedFunctionalHessianMat( rf, appctx=appctx, apply_riesz=apply_riesz, @@ -531,7 +503,7 @@ def RieszMapMat(controls, symmetric=True, comm=None): symmetric (bool): Whether the Riesz map attached to the Control is symmetric. comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the controls are defined over. """ - PETSc = import_petsc() + from petsc4py import PETSc ctx = RieszMapMatCtx(controls, comm=comm) n = ctx.vec_interface.n @@ -659,8 +631,8 @@ class TAOSolver(OptimizationSolver): def __init__(self, problem, parameters, *, options_prefix=None, appctx=None, Pmat=None, comm=None): - PETSc = import_petsc() - petsctools = import_petsctools() + from petsc4py import PETSc + import petsctools if not isinstance(problem, MinimizationProblem): raise TypeError("MinimizationProblem required") @@ -808,7 +780,7 @@ def x(self): def _tao_reasons(self): """Dictionary of TAO convergence reason int codes -> python objects """ - PETSc = import_petsc() + from petsc4py import PETSc # Same approach as in _make_reasons in firedrake/solving_utils.py, # Firedrake master branch 57e21cc8ebdb044c1d8423b48f3dbf70975d5548 return {getattr(PETSc.TAO.Reason, key): key @@ -821,7 +793,7 @@ def solve(self): Returns: OverloadedType or Sequence[OverloadedType]: The solution. """ - petsctools = import_petsctools() + import petsctools controls = self.tao_objective.reduced_functional.controls m = tuple(control.tape_value()._ad_copy() for control in controls) From 92fd6398a6256843cda5d5122bcce13d293aa319 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 17 Apr 2026 09:49:26 +0100 Subject: [PATCH 3/3] bump minimum petsctools version for the OptionsManager free functions --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9b4ac2ca..53e80f78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,7 @@ visualisation = [ ] tao = [ "petsc4py", - "petsctools>2025.0" + "petsctools>2025.3" ]