From 70d6188daef304ddeaef3ec14f54a244b307e2dc Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 28 Aug 2025 22:22:12 +0500 Subject: [PATCH 01/25] Added rf_eigensolver to init. --- firedrake/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/firedrake/__init__.py b/firedrake/__init__.py index cb57a882f8..f23dc56bca 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -83,6 +83,7 @@ def init_petsc(): from firedrake.utility_meshes import * from firedrake.variational_solver import * from firedrake.eigensolver import * +from firedrake.rf_eigensolver import * from firedrake.ensemble import * from firedrake.randomfunctiongen import * from firedrake.external_operators import * From 091e5ff647487552f6c34321d49123afc651f1c2 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 28 Aug 2025 22:52:15 +0500 Subject: [PATCH 02/25] Restricted functional context. --- firedrake/restricted_functional_ctx.py | 216 +++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 firedrake/restricted_functional_ctx.py diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py new file mode 100644 index 0000000000..6ccabde970 --- /dev/null +++ b/firedrake/restricted_functional_ctx.py @@ -0,0 +1,216 @@ +from contextlib import contextmanager +from functools import wraps +import itertools +from enum import Enum +from numbers import Complex +import warnings +from petsctools import * +from firedrake.adjoint import PETScVecInterface, check_rf_action, TLMAction, \ + AdjointAction, HessianAction, get_valid_comm, \ + flatten_parameters, ReducedFunctionalMatCtx +from firedrake import Function, Cofunction, TestFunction, TrialFunction, solve, inner, dx, assemble +from firedrake.__future__ import interpolate + +import numpy as np + +try: + import petsc4py.PETSc as PETSc +except ModuleNotFoundError: + PETSc = None + + +def new_restricted_control_variable(reduced_functional, function_space, dual=False): + """Return new variables suitable for storing a control value or its dual + by interpolating into the space over which the 'ReducedFunctional' is + defined. + + Args: + ---------- + reduced_functional (ReducedFunctional): The `ReducedFunctional` whose + controls are to be copied. + dual (bool): whether to return a dual type. If False then a primal type is returned. + + Returns: + ---------- + tuple[OverloadedType]: New variables suitable for storing a control value. + """ + + return tuple(Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) + for control in reduced_functional.controls) + +def interpolate_vars(variables, function_space): + """ + Interpolates primal/dual variables to restricted/unrestricted function spaces. + + Args: + ---------- + variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), + List(firedrake.Function), List(firedrake.Cofunction), + firedrake.Function, firedrake.Cofunction): + Variables that are to be interpolated into unrestricted/restricted primal/dual spaces. + restricted_space (Optional[FunctionSpace, RestrictedFunctionSpace]): + The function space where `PETSc.Vec` objects will live. + + Returns: + ---------- + The same variable/variables but interpolated into the necessary function space. + """ + if isinstance(variables, (tuple, list)): + if isinstance(variables[0], Function): + return tuple(Function(function_space).interpolate(v) for v in variables) + else: + return tuple(Cofunction(function_space.dual()).interpolate(v) for v in variables) + elif isinstance(variables, Function): + return Function(function_space).interpolate(variables) + else: + return Cofunction(function_space.dual()).interpolate(variables) + + +class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): + """ + `PETSc.Mat.Python context to apply to the operator representing linearisation of a `ReducedFunctional`. + Optional restriction of the control space to a provided `FunctionSpace`. + + Args: + ---------- + rf (ReducedFunctional): + Defines the forward model, and used to compute operator actions. + action (RFAction): + Whether to apply the TLM, adjoint, or Hessian action. + apply_riesz (bool): + Whether to apply the Riesz map before returning the + result of the action to `PETSc`. + appctx (Optional[dict]): + User provided context. + comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): + Communicator that the `ReducedFunctional` is defined over. + restricted_space (Optional[FunctionSpace]): + If provided, the control space will be restricted to the passed space. + """ + + def __init__(self, rf, action=HessianAction, *, + apply_riesz=False, appctx=None, comm=PETSc.COMM_WORLD, + restricted_space=None): + if restricted_space is None: + raise ValueError("restricted_space must be provided for RestrictedReducedFunctionalMatCtx") + + super().__init__(rf, action=action, apply_riesz=apply_riesz, + appctx=appctx, comm=comm) + + self.function_space = rf.controls[0].control.function_space() + self.restricted_space = restricted_space or self.function_space + + # Build restricted interfaces + self.restricted_control_interface = ( + PETScVecInterface([Function(self.restricted_space).interpolate(c.control) + for c in rf.controls]) + if restricted_space else self.control_interface + ) + + if action in (AdjointAction, TLMAction): + self.restricted_functional_interface = PETScVecInterface(Function(restricted_space).interpolate(rf.functional), comm=comm) if restricted_space else self.functional_interface + + # Override x/y interfaces and variables based on the action + if action == HessianAction: + self.xresinterface = self.restricted_control_interface + self.yresinterface = self.restricted_control_interface + self.xres = new_restricted_control_variable(rf, self.restricted_space) + self.yres = new_restricted_control_variable(rf, self.restricted_space) + + elif action == AdjointAction: + self.xresinterface = self.restricted_functional_interface + self.yresinterface = self.restricted_control_interface + self.xres = Function(self.restricted_space).interpolate(rf.functional)._ad_copy()._ad_init_zero(dual=True) + self.yres = new_restricted_control_variable(rf, self.restricted_space) + + elif action == TLMAction: + self.xresinterface = self.restricted_control_interface + self.yresinterface = self.restricted_functional_interface + self.xres = new_restricted_control_variable(rf, self.restricted_space) + self.yres = Function(self.restricted_space).interpolate(rf.functional)._ad_copy()._ad_init_zero(dual=True) + + def mult(self, A, x, y): + """ + Compute `y = Ax` and store in `y`. + Ax is represented as the forward action of implicit matrix A. + + Args: + ---------- + A (ReducedFunctional): + Implicit matrix for which Ax is defined. + x (PETSc.Vec): + `PETSc.Vec` which is converted to an `OverloadedType` object. + + """ + self.xresinterface.from_petsc(x, self.xres) + interpolated_x = interpolate_vars(self.xres, self.function_space) + out = self.mult_impl(A, interpolated_x) + interpolated_y = interpolate_vars(out, self.restricted_space) + self.yresinterface.to_petsc(y, interpolated_y) + if self._shift != 0: + y.axpy(self._shift, x) + + def multTranspose(self, A, x, y): + """ + Compute `y = A^Tx` and store in `y`. + A^Tx is represented as the action of the transpose of implicit matrix A. + + Args: + ---------- + A (ReducedFunctional): + Implicit matrix for which Ax is defined. + x (PETSc.Vec): + `PETSc.Vec` which is converted to an `OverloadedType` object. + + """ + self.yresinterface.from_petsc(x, self.yres) + interpolated_x = interpolate_vars(self.yres, self.function_space) + out = self.mult_impl_transpose(A, interpolated_x) + interpolated_y = interpolate_vars(out, self.restricted_space) + self.xresinterface.to_petsc(y, interpolated_y) + if self._shift != 0: + y.axpy(self._shift, x) + + +def RestrictedReducedFunctionalMat(rf, action=HessianAction, *, apply_riesz=False, appctx=None, comm=None, restricted_space=None): + """ + `PETSc.Mat` to apply the action of a linearisation of a `ReducedFunctional`. + + If V is the control space and U is the functional space, each action has the following map: + Jhat : V -> U + TLM : V -> U + Adjoint : U* -> V* + Hessian : V x U* -> V* | V -> V* + + Args: + ---------- + rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. + action (RFAction): Whether to apply the TLM, adjoint, or Hessian action. + apply_riesz (bool): Whether to apply the Riesz map before returning the + result of the action to `PETSc`. + appctx (Optional[dict]): User provided context. + comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the `ReducedFunctional` is defined over. + restricted_space (Optional[FunctionSpace]): If provided, the control space will be restricted to the passed space. + + Returns + ---------- + mat (PETSc.Mat): + The `PETSc.Mat` whose action and transpose action are defined by the context. + """ + ctx = RestrictedReducedFunctionalMatCtx( + rf, action, appctx=appctx, apply_riesz=apply_riesz, comm=comm, restricted_space=restricted_space) + + ncol = ctx.xresinterface.n + Ncol = ctx.xresinterface.N + + nrow = ctx.yresinterface.n + Nrow = ctx.yresinterface.N + + mat = PETSc.Mat().createPython( + ((nrow, Nrow), (ncol, Ncol)), + ctx, comm=ctx.control_interface.comm) + if action == HessianAction: + mat.setOption(PETSc.Mat.Option.SYMMETRIC, True) + mat.setUp() + mat.assemble() + return mat From 82588032c63780c5b706d38b338fad5592942acd Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 28 Aug 2025 22:52:41 +0500 Subject: [PATCH 03/25] RF eigensolver abstraction. --- firedrake/rf_eigensolver.py | 400 ++++++++++++++++++++++++++++++++++++ 1 file changed, 400 insertions(+) create mode 100644 firedrake/rf_eigensolver.py diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py new file mode 100644 index 0000000000..f7b2db1301 --- /dev/null +++ b/firedrake/rf_eigensolver.py @@ -0,0 +1,400 @@ +"""Specify and solve finite element eigenproblems.""" +from petsctools import OptionsManager, flatten_parameters +from firedrake.assemble import assemble +from firedrake.bcs import extract_subdomain_ids, restricted_function_space, DirichletBC +from firedrake.function import Function +from firedrake.ufl_expr import TrialFunction, TestFunction +from firedrake import utils +from firedrake.restricted_functional_ctx import RestrictedReducedFunctionalMat +from firedrake.exceptions import ConvergenceError +from ufl import replace, inner, dx, Argument, Form +from firedrake.petsc import PETSc +from petsc4py import PETSc +from firedrake.adjoint import ReducedFunctional, ReducedFunctionalMat, TLMAction, AdjointAction +try: + from slepc4py import SLEPc +except ImportError: + SLEPc = None +__all__ = ["RFEigenproblem", + "RFEigensolver"] + + +class RFEigenproblem: + """Generalised linear eigenvalue problem. + + The problem has the form, find `u`, `λ` such that:: + + A(u, v) = λM(u, v) ∀ v ∈ V + + Parameters + ---------- + rf : ReducedFunctional + The bilinear operator A. + M : ufl.Form + The mass form M(u, v), defaults to inner(u, v) * dx if is None and identity=False. + If identity is True, M is treated as identity operator. + bcs : DirichletBC or list of DirichletBC + The Dirichlet boundary conditions. + bc_shift: float + The value to shift the boundary condition eigenvalues by. This value + will be ignored if restrict==True. + restrict: bool + If True, replace the function spaces of u and v with their restricted + version. The output space remains unchanged. + identity: bool + If True, M is replaced by identity matrix. Differentiate between mass + matrix being identity and inner product operator. + action: (Optional[TLMAction, AdjointAction]): + Defines forward action of implicit matrix operator. + apply_riez: bool + If True, when defining A's adjoint action, we can output into the dual + space. + action: str + Determine whether forward action is TLM or adjoint. + appctx (Optional[dict]): + Context forwarded to the linearisation. + comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): + Communicator that the rf is defined over. + + + + Notes + ----- + If restrict==True, the arguments of A and M will be replaced, such that + their function space is replaced by the equivalent RestrictedFunctionSpace + class. This avoids the computation of eigenvalues associated with the + Dirichlet boundary conditions. This in turn prevents convergence failures, + and allows only the non-boundary eigenvalues to be returned. The + eigenvectors will be in the original, non-restricted space. + + If restrict==False and Dirichlet boundary conditions are supplied, then + these conditions will result in the eigenproblem having a nullspace spanned + by the basis functions with support on the boundary. To facilitate + solution, this is shifted by the specified amount. It is the user's + responsibility to ensure that the shift is not close to an actual + eigenvalue of the system. + + Also, we assume operator that we linearise maps from V to V. + """ + def __init__(self, rf, M=None, bcs=None, bc_shift=0.0, restrict=True, + identity=False, apply_riesz=True, action="tlm", appctx=None, comm=None): + if not SLEPc: + raise ImportError( + "Unable to import SLEPc, eigenvalue computation not possible " + "(see https://www.firedrakeproject.org/install.html#slepc)" + ) + + if not isinstance(rf, ReducedFunctional) and not isinstance(M, ReducedFunctional): + raise TypeError( + "At least A or M in the GEP has to be a ReducedFunctional instance. Use 'LinearEigenproblem' otherwise." + ) + + self.output_space = self.output_space(rf) + self.bc_shift = bc_shift + self.restrict = restrict and bcs + self.restricted_space = None + self.bcs = bcs + self.apply_riesz = apply_riesz + self.appctx = appctx + self.comm = comm + self.action = self.match_forward_action(action) + + if not M and not identity: + M = self.generate_inner_form() + + if self.restrict: + union_bcs = self.get_bcs_union(rf) if isinstance(rf, ReducedFunctional) else bcs + V_res = restricted_function_space(self.output_space, extract_subdomain_ids(union_bcs)) + self.restricted_space = V_res + self.bcs = [bc.reconstruct(V=V_res, indices=bc._indices) for bc in union_bcs] + self.M = self.restrict_obj(M, V_res) + self.A = self.restrict_obj(rf, V_res) + else: + self.bcs = bcs + self.M = None if identity else self.as_petsc_mat(M) + self.A = self.as_petsc_mat(rf) + + def restrict_obj(self, obj, V_res): + """ + Substitutes variational formulation with functions from restricted space if necessary. + Otherwise, tries to convert obj to `SLEPc`-compatible operator. + + Parameters + ---------- + obj : Optional[ReducedFunctional, ufl.Form] + `Firedrake` operator converted to a 'PETSc.Mat'. + V_res : RestrictedFunctionSpace + Function space to which restriction occurs.. + + Returns + ------- + Optional[None, PETSc.Mat]: + `SLEPc`-compatible operator representing obj. + """ + if isinstance(obj, Form): + args = obj.arguments() + v, u = args + u_res = TrialFunction(V_res) + v_res = TestFunction(V_res) + obj = replace(obj, {u: u_res, v: v_res}) + return self.as_petsc_mat(obj) if obj else None + + def as_petsc_mat(self, obj): + """ + Convert a `ufl.Form` or a `ReducedFunctional` to a `PETSc.Mat` operator. + Forward action of implicit matrix determined by action. + + Parameters + ---------- + obj : Optional[ReducedFunctional, ufl.Form] + `Firedrake` operator converted to a 'PETSc.Mat'. + + Returns + ------- + `PETSc.Mat` object that represents the operator obj. + """ + if self.restricted_space: + if isinstance(obj, ReducedFunctional): + return RestrictedReducedFunctionalMat(obj, action=self.action, apply_riesz=self.apply_riesz, appctx=self.appctx, + comm=self.comm, restricted_space=self.restricted_space) + if isinstance(obj, ReducedFunctional): + return ReducedFunctionalMat(obj, action=self.action, apply_riesz=self.apply_riesz, appctx=self.appctx, comm=self.comm) + return assemble( + obj, bcs=self.bcs, + weight=self.bc_shift and 1./self.bc_shift + ).petscmat + + def match_forward_action(self, action): + """ + Match forward action of matrix using passed string. + + Parameters + ---------- + action : str + String representing the forward action of the matrix. 'tlm' or 'adjoint'. + + Returns + ------- + Optional[TLMAction, AdjointAction]: + Action corresponding to the string. + """ + match action: + case "tlm": + return TLMAction + case "adjoint": + return AdjointAction + + def output_space(self, rf): + """ + Determines the function space that operator outputs into. + + Parameters + ---------- + rf : Optional[ReducedFunctional, ufl.Form] + Operator whose output space we want to find. + + Returns + ------- + firedrake.FunctionSpace: + Function space into which operator outputs to. + """ + if isinstance(rf, ReducedFunctional): + return rf._controls[0].function_space() + return rf.arguments()[0].function_space() + + def generate_inner_form(self): + """ + Generates the mass matrix in the `L2`-norm. + + Parameters + ---------- + space : FunctionSpace + Function space over which we define the mass matrix. + + Returns + ------- + ufl.Form: + `ufl.Form` corresponding to the `L2`-norm. + """ + trial = TrialFunction(self.output_space) + test = TestFunction(self.output_space) + mass_form = inner(trial, test) * dx + return mass_form + + def get_bcs_union(self, rf): + """ + Finds the union off all the boundary conditions. This is done so as to remove the + nodes related to these boundary conditions from the variational formulation. + + Parameters + ---------- + rf : ReducedFunctional + Provides `pyadjoint.Tape` over which collection of the + union of the boundary conditions occurs. + + Returns + ------- + List[DirichletBC]]: + List of all unique `DirichletBC`. + """ + tape = rf.tape + common_dofs = None + all_bcs = [] + + for block in tape._blocks: + if hasattr(block, 'bcs'): + bcs = block.bcs + print([type(bc) for bc in bcs]) + if isinstance(bcs, DirichletBC): + all_bcs.append(bcs) + elif isinstance(bcs, (list, tuple)): + all_bcs.extend([b for b in bcs if isinstance(b, DirichletBC)]) + + return list(set(all_bcs)) + + @utils.cached_property + def dm(self): + r"""Return the dm associated with the output space.""" + if self.restrict: + return self.restricted_space.dm + else: + return self.output_space.dm + + +class RFEigensolver(OptionsManager): + r"""Solve a LinearEigenproblem. + + Parameters + ---------- + problem : LinearEigenproblem + The eigenproblem to solve. + n_evals : int + The number of eigenvalues to compute. + options_prefix : str + The options prefix to use for the eigensolver. + solver_parameters : dict + PETSc options for the eigenvalue problem. + ncv : int + Maximum dimension of the subspace to be used by the solver. See + `SLEPc.EPS.setDimensions`. + mpd : int + Maximum dimension allowed for the projected problem. See + `SLEPc.EPS.setDimensions`. + + Notes + ----- + + Users will typically wish to set solver parameters specifying the symmetry + of the eigenproblem and which eigenvalues to search for first. + + The former is set using the options available for `EPSSetProblemType + `__. + + For example if the bilinear form is symmetric (Hermitian in complex mode), + one would add this entry to `solver_options`:: + + "eps_gen_hermitian": None + + As always when specifying PETSc options, `None` indicates that the option + in question is a flag and hence doesn't take an argument. + + The eigenvalues to search for first are specified using the options for + `EPSSetWhichEigenPairs `__. + + For example, to look for the eigenvalues with largest real part, one + would add this entry to `solver_options`:: + + "eps_largest_real": None + """ + """ + Preconditioners won't work because of MATSHELL. + Only iterative solvers work as well (refer to PETSc docs). + Spectral transform changed with bcs (default) because sinvert wont work because of MATSHELL. + If info about A and M, could pass HEP to problem_type parameter. If M identity, NHEP. If not, GNHEP. + + + """ + + DEFAULT_EPS_PARAMETERS = {"eps_type": "krylovschur", + "eps_tol": 1e-10, + "eps_target": 0.0} + + def __init__(self, problem, n_evals, *, options_prefix=None, + solver_parameters=None, ncv=None, mpd=None): + + self.es = SLEPc.EPS().create(comm=problem.dm.comm) + self._problem = problem + self.n_evals = n_evals + self.ncv = ncv + self.mpd = mpd + solver_parameters = flatten_parameters(solver_parameters or {}) + for key in self.DEFAULT_EPS_PARAMETERS: + value = self.DEFAULT_EPS_PARAMETERS[key] + solver_parameters.setdefault(key, value) + super().__init__(solver_parameters, options_prefix) + self.set_from_options(self.es) + + def solve(self): + r"""Solve the eigenproblem. + + Returns + ------- + int + The number of Eigenvalues found. + """ + A = self._problem.A + M = self._problem.M + if not isinstance(A, PETSc.Mat): + raise TypeError("A must be a PETSc matrix, got %s" % A.__class__) + if not isinstance(M, PETSc.Mat) and (M is not None): + raise TypeError("M must be a PETSc matrix or None (identity), got %s" % M.__class__) + + self.es.setDimensions(nev=self.n_evals, ncv=self.ncv, mpd=self.mpd) + self.es.setOperators(A, M) + with self.inserted_options(): + self.es.solve() + nconv = self.es.getConverged() + if nconv == 0: + raise ConvergenceError("Did not converge any eigenvalues.") + return nconv + + def check_es_convergence(self): + r"""Check the convergence of the eigenvalue problem.""" + r = self.es.getConvergedReason() + try: + reason = SLEPc.EPS.ConvergedReasons[r] + except KeyError: + reason = ("unknown reason (petsc4py enum incomplete?), " + "try with -eps_converged_reason") + if r < 0: + raise ConvergenceError( + r"""Eigenproblem failed to converge after %d iterations. + Reason: + %s""" % (self.es.getIterationNumber(), reason) + ) + + def eigenvalue(self, i): + r"""Return the i-th eigenvalue of the solved problem.""" + return self.es.getEigenvalue(i) + + def eigenfunction(self, i): + r"""Return the i-th eigenfunction of the solved problem. + + Returns + ------- + (Function, Function) + The real and imaginary parts of the eigenfunction. + """ + if self._problem.restrict: + eigenmodes_real = Function(self._problem.restricted_space) + eigenmodes_imag = Function(self._problem.restricted_space) + else: + eigenmodes_real = Function(self._problem.output_space) # fn of V + eigenmodes_imag = Function(self._problem.output_space) + with eigenmodes_real.dat.vec_wo as vr: + with eigenmodes_imag.dat.vec_wo as vi: + self.es.getEigenvector(i, vr, vi) # gets the i-th eigenvector + if self._problem.restrict: + eigenmodes_real = Function(self._problem.output_space).interpolate(eigenmodes_real) + eigenmodes_imag = Function(self._problem.output_space).interpolate(eigenmodes_imag) + return eigenmodes_real, eigenmodes_imag # returns Firedrake fns From 7567ee4ac77d838c1c2f7e9b8f7f7483334a3ea7 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 28 Aug 2025 22:52:52 +0500 Subject: [PATCH 04/25] RF eigensolver tests. --- .../firedrake/adjoint/test_rf_eigensolver.py | 161 ++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 tests/firedrake/adjoint/test_rf_eigensolver.py diff --git a/tests/firedrake/adjoint/test_rf_eigensolver.py b/tests/firedrake/adjoint/test_rf_eigensolver.py new file mode 100644 index 0000000000..37b48d85fe --- /dev/null +++ b/tests/firedrake/adjoint/test_rf_eigensolver.py @@ -0,0 +1,161 @@ + +import pytest +import numpy as np +from firedrake import * +from firedrake.adjoint import * +from firedrake.bcs import DirichletBC +from firedrake.rf_eigensolver import * +from firedrake.adjoint import PETScVecInterface, AdjointAction +from firedrake.restricted_functional_ctx import new_restricted_control_variable, interpolate_vars + +def generate_expressions(): + """ + Returns a list of three expressions + """ + mesh = UnitSquareMesh(50, 50, quadrilateral=False) + x, y = SpatialCoordinate(mesh) + return [2 * x, exp(-x - y), sin(pi * x * y)] + +@pytest.mark.parametrize("expression", generate_expressions()) +def test_compare_eigensolvers_helmholtz(expression): + """Test TLM vs assembled form for Helmholtz eigenproblem.""" + n = 50 + num_eigenvalues = 10 + + mesh = UnitSquareMesh(n, n, quadrilateral=False) + x, y = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "CG", 1) + bc = DirichletBC(V, 0.0, "on_boundary") + + u = TrialFunction(V) + v = TestFunction(V) + a = (inner(grad(u), grad(v)) + u * v) * dx + + # Form based Ax = cMx + problem_form = LinearEigenproblem(a, bcs=bc, restrict=True) + solver_form = LinearEigensolver(problem_form, n_evals=num_eigenvalues) + solver_form.solve() + nconv = solver_form.es.getConverged() + + evals_form = [solver_form.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + + # ReducedFunctional-based, Jx = 1/cx + continue_annotation() + f = Function(V) + f.interpolate(2 * x) + control = Control(f) + u_1 = Function(V) + v = TestFunction(V) + F = (inner(grad(u_1), grad(v)) + u_1 * v - f * v) * dx + solve(F == 0, u_1, bcs=bc) + J = ReducedFunctional(u_1, controls=[control]) + pause_annotation() + + problem_rf = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=True, identity=True) + solver_rf = RFEigensolver(problem_rf, n_evals=num_eigenvalues) + solver_rf.es.setTarget(1) + solver_rf.solve() + nconv = solver_rf.es.getConverged() + + evals_rf = [1/solver_rf.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + + for ev1, ev2 in zip(evals_form, evals_rf): + assert abs(ev1 - ev2) < 1e-6 + + set_working_tape(Tape()) + +@pytest.mark.parametrize("restrict", [False]) +@pytest.mark.parametrize("riesz", ["L2"]) +def test_apply_riesz_param(restrict, riesz): + """Test implicit matrix for Helmholtz eigenproblem.""" + n = 50 + mesh = UnitSquareMesh(n, n, quadrilateral=False) + x, y = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "CG", 1) + bc = DirichletBC(V, 0.0, "on_boundary") + + # Tape the solve + continue_annotation() + f = Function(V) + f.interpolate(2 * x * y) + control = Control(f) + u_1 = Function(V) + v = TestFunction(V) + F = (inner(grad(u_1), grad(v)) + u_1 * v - f * v) * dx + solve(F == 0, u_1, bcs=bc) + J = ReducedFunctional(u_1, controls=[control]) + pause_annotation() + + #Check adjoint action + problem_rf_true = RFEigenproblem(J, bcs=bc, apply_riesz=True, restrict=restrict, action="adjoint") + problem_rf_false = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=restrict, action="adjoint") + + mat_true = problem_rf_true.A + mat_false = problem_rf_false.A + + x = mat_true.createVecRight() + Ax = mat_true.createVecLeft() + Ax2 = mat_true.createVecLeft() + x.setRandom() + + mat_true.mult(x, Ax) + mat_false.mult(x, Ax2) + + interface_space = problem_rf_true.restricted_space if restrict else V + + interface_Ax = PETScVecInterface([Function(interface_space).interpolate(c.control) + for c in J.controls]) + func_representation = new_restricted_control_variable(J, interface_space, dual=False) + func_representation_dual = new_restricted_control_variable(J, interface_space, dual=True) + interface_Ax.from_petsc(Ax, func_representation) + interface_Ax.from_petsc(Ax2, func_representation_dual) + + if restrict: + #Intepolate from restricted to unrestricted + func_representation = interpolate_vars(func_representation, V) + func_representation_dual= interpolate_vars(func_representation_dual, V) + + set_working_tape(Tape()) + assert errornorm(func_representation[0], func_representation_dual[0].riesz_representation(riesz_map=riesz)) < 1e-6 + +def test_compare_tlm_adjoint(): + """Test TLM and adjoint give same eigenvalues for Helmholtz eigenproblem.""" + n = 50 + num_eigenvalues = 10 + + mesh = UnitSquareMesh(n, n, quadrilateral=False) + x, y = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "CG", 1) + bc = DirichletBC(V, 0.0, "on_boundary") + + # ReducedFunctional-based, Jx = 1/cx + continue_annotation() + f = Function(V) + f.interpolate(2 * x) + control = Control(f) + u_1 = Function(V) + v = TestFunction(V) + F = (inner(grad(u_1), grad(v)) + u_1 * v - f * v) * dx + solve(F == 0, u_1, bcs=bc) + J = ReducedFunctional(u_1, controls=[control]) + pause_annotation() + + problem_rf_a = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=True, identity=True, action="adjoint") + solver_rf_a = RFEigensolver(problem_rf_a, n_evals=num_eigenvalues) + solver_rf_a.es.setTarget(1) + solver_rf_a.solve() + nconv = solver_rf_a.es.getConverged() + + evals_rf_a = [solver_rf_a.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + + problem_rf_t = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=True, identity=True, action="tlm") + solver_rf_t = RFEigensolver(problem_rf_t, n_evals=num_eigenvalues) + solver_rf_t.es.setTarget(1) + solver_rf_t.solve() + nconv = solver_rf_t.es.getConverged() + + evals_rf_t = [solver_rf_t.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + for ev1, ev2 in zip(evals_rf_a, evals_rf_t): + print(abs(ev1 - ev2)) + assert abs(ev1 - ev2) < 1e-6 + set_working_tape(Tape()) From dabb809120f9ce9f7f61c5d21bb7c40a9381e3a9 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 28 Aug 2025 22:55:39 +0500 Subject: [PATCH 05/25] Init so necessary new pyadjoint objects added. --- firedrake/adjoint/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index d3d28e6129..5d86591111 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -27,6 +27,10 @@ continue_disk_checkpointing, stop_disk_checkpointing, \ checkpointable_mesh # noqa F401 from firedrake.adjoint_utils import get_solve_blocks # noqa F401 +from pyadjoint.optimization.tao_solver import ReducedFunctionalMat, ReducedFunctionalMatCtx, \ + HessianAction, TLMAction, AdjointAction, check_rf_action, \ + get_valid_comm, flatten_parameters, \ + ReducedFunctionalMatCtx, PETScVecInterface from pyadjoint.verification import taylor_test, taylor_to_dict # noqa F401 from pyadjoint.drivers import compute_gradient, compute_derivative, compute_hessian # noqa F401 From 8ab669ab7b504a1cb6d29b7711e19628b9f0ba10 Mon Sep 17 00:00:00 2001 From: KP Date: Fri, 29 Aug 2025 01:03:39 +0500 Subject: [PATCH 06/25] Corrected docstring in RRFC. --- firedrake/restricted_functional_ctx.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py index 6ccabde970..db3d1e5700 100644 --- a/firedrake/restricted_functional_ctx.py +++ b/firedrake/restricted_functional_ctx.py @@ -136,10 +136,12 @@ def mult(self, A, x, y): Args: ---------- - A (ReducedFunctional): + A (PETSc.Mat): Implicit matrix for which Ax is defined. x (PETSc.Vec): - `PETSc.Vec` which is converted to an `OverloadedType` object. + `PETSc.Vec` to which operator is applied to. + y (PETSc.Vec): + `PETSc.Vec` which is the result of the action of this operator. """ self.xresinterface.from_petsc(x, self.xres) @@ -157,10 +159,12 @@ def multTranspose(self, A, x, y): Args: ---------- - A (ReducedFunctional): + A (PETSc.Mat): Implicit matrix for which Ax is defined. x (PETSc.Vec): - `PETSc.Vec` which is converted to an `OverloadedType` object. + `PETSc.Vec` to which transpose of operator is applied to. + y (PETSc.Vec): + `PETSc.Vec` which is the result of the action of this operator. """ self.yresinterface.from_petsc(x, self.yres) From 5a9e76ff88fe6da5c8494b7a44ab31ae58b7478e Mon Sep 17 00:00:00 2001 From: KP Date: Fri, 29 Aug 2025 17:42:18 +0300 Subject: [PATCH 07/25] Remove wrong comment. --- firedrake/restricted_functional_ctx.py | 1 - 1 file changed, 1 deletion(-) diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py index db3d1e5700..25ef4a2c2e 100644 --- a/firedrake/restricted_functional_ctx.py +++ b/firedrake/restricted_functional_ctx.py @@ -110,7 +110,6 @@ def __init__(self, rf, action=HessianAction, *, if action in (AdjointAction, TLMAction): self.restricted_functional_interface = PETScVecInterface(Function(restricted_space).interpolate(rf.functional), comm=comm) if restricted_space else self.functional_interface - # Override x/y interfaces and variables based on the action if action == HessianAction: self.xresinterface = self.restricted_control_interface self.yresinterface = self.restricted_control_interface From ffcdea378930b1f0fcd790f7b69ba434602d90ce Mon Sep 17 00:00:00 2001 From: KP Date: Sun, 31 Aug 2025 18:58:22 +0300 Subject: [PATCH 08/25] Linted. Except wants me to remove things from init. --- firedrake/adjoint/__init__.py | 10 +- firedrake/restricted_functional_ctx.py | 115 ++++++---- firedrake/rf_eigensolver.py | 209 ++++++++++-------- .../firedrake/adjoint/test_rf_eigensolver.py | 75 +++++-- 4 files changed, 245 insertions(+), 164 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index 5d86591111..e28d8c5c31 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -27,10 +27,12 @@ continue_disk_checkpointing, stop_disk_checkpointing, \ checkpointable_mesh # noqa F401 from firedrake.adjoint_utils import get_solve_blocks # noqa F401 -from pyadjoint.optimization.tao_solver import ReducedFunctionalMat, ReducedFunctionalMatCtx, \ - HessianAction, TLMAction, AdjointAction, check_rf_action, \ - get_valid_comm, flatten_parameters, \ - ReducedFunctionalMatCtx, PETScVecInterface +from pyadjoint.optimization.tao_solver import ( + ReducedFunctionalMat, ReducedFunctionalMatCtx, HessianAction, + TLMAction, AdjointAction, check_rf_action, get_valid_comm, + flatten_parameters, PETScVecInterface +) + from pyadjoint.verification import taylor_test, taylor_to_dict # noqa F401 from pyadjoint.drivers import compute_gradient, compute_derivative, compute_hessian # noqa F401 diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py index 25ef4a2c2e..766f9729f6 100644 --- a/firedrake/restricted_functional_ctx.py +++ b/firedrake/restricted_functional_ctx.py @@ -1,17 +1,11 @@ -from contextlib import contextmanager -from functools import wraps -import itertools -from enum import Enum -from numbers import Complex -import warnings -from petsctools import * -from firedrake.adjoint import PETScVecInterface, check_rf_action, TLMAction, \ - AdjointAction, HessianAction, get_valid_comm, \ - flatten_parameters, ReducedFunctionalMatCtx -from firedrake import Function, Cofunction, TestFunction, TrialFunction, solve, inner, dx, assemble -from firedrake.__future__ import interpolate - -import numpy as np +from firedrake.adjoint import ( + PETScVecInterface, + TLMAction, + AdjointAction, + HessianAction, + ReducedFunctionalMatCtx, +) +from firedrake import Function, Cofunction try: import petsc4py.PETSc as PETSc @@ -34,18 +28,20 @@ def new_restricted_control_variable(reduced_functional, function_space, dual=Fal ---------- tuple[OverloadedType]: New variables suitable for storing a control value. """ + return tuple( + Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) + for control in reduced_functional.controls + ) - return tuple(Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) - for control in reduced_functional.controls) def interpolate_vars(variables, function_space): """ Interpolates primal/dual variables to restricted/unrestricted function spaces. - + Args: ---------- - variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), - List(firedrake.Function), List(firedrake.Cofunction), + variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), + List(firedrake.Function), List(firedrake.Cofunction), firedrake.Function, firedrake.Cofunction): Variables that are to be interpolated into unrestricted/restricted primal/dual spaces. restricted_space (Optional[FunctionSpace, RestrictedFunctionSpace]): @@ -59,7 +55,9 @@ def interpolate_vars(variables, function_space): if isinstance(variables[0], Function): return tuple(Function(function_space).interpolate(v) for v in variables) else: - return tuple(Cofunction(function_space.dual()).interpolate(v) for v in variables) + return tuple( + Cofunction(function_space.dual()).interpolate(v) for v in variables + ) elif isinstance(variables, Function): return Function(function_space).interpolate(variables) else: @@ -70,17 +68,17 @@ class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): """ `PETSc.Mat.Python context to apply to the operator representing linearisation of a `ReducedFunctional`. Optional restriction of the control space to a provided `FunctionSpace`. - + Args: ---------- rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. action (RFAction): Whether to apply the TLM, adjoint, or Hessian action. - apply_riesz (bool): + apply_riesz (bool): Whether to apply the Riesz map before returning the result of the action to `PETSc`. - appctx (Optional[dict]): + appctx (Optional[dict]): User provided context. comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the `ReducedFunctional` is defined over. @@ -88,27 +86,41 @@ class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): If provided, the control space will be restricted to the passed space. """ - def __init__(self, rf, action=HessianAction, *, - apply_riesz=False, appctx=None, comm=PETSc.COMM_WORLD, - restricted_space=None): + def __init__( + self, + rf, + action=HessianAction, + *, + apply_riesz=False, + appctx=None, + comm=PETSc.COMM_WORLD, + restricted_space=None, + ): if restricted_space is None: - raise ValueError("restricted_space must be provided for RestrictedReducedFunctionalMatCtx") + raise ValueError( + "restricted_space must be provided for RestrictedReducedFunctionalMatCtx" + ) - super().__init__(rf, action=action, apply_riesz=apply_riesz, - appctx=appctx, comm=comm) + super().__init__(rf, action=action, apply_riesz=apply_riesz, appctx=appctx, comm=comm) self.function_space = rf.controls[0].control.function_space() self.restricted_space = restricted_space or self.function_space # Build restricted interfaces self.restricted_control_interface = ( - PETScVecInterface([Function(self.restricted_space).interpolate(c.control) - for c in rf.controls]) - if restricted_space else self.control_interface + PETScVecInterface( + [Function(self.restricted_space).interpolate(c.control) for c in rf.controls] + ) + if restricted_space + else self.control_interface ) if action in (AdjointAction, TLMAction): - self.restricted_functional_interface = PETScVecInterface(Function(restricted_space).interpolate(rf.functional), comm=comm) if restricted_space else self.functional_interface + self.restricted_functional_interface = ( + PETScVecInterface(Function(restricted_space).interpolate(rf.functional), comm=comm) + if restricted_space + else self.functional_interface + ) if action == HessianAction: self.xresinterface = self.restricted_control_interface @@ -119,14 +131,24 @@ def __init__(self, rf, action=HessianAction, *, elif action == AdjointAction: self.xresinterface = self.restricted_functional_interface self.yresinterface = self.restricted_control_interface - self.xres = Function(self.restricted_space).interpolate(rf.functional)._ad_copy()._ad_init_zero(dual=True) + self.xres = ( + Function(self.restricted_space) + .interpolate(rf.functional) + ._ad_copy() + ._ad_init_zero(dual=True) + ) self.yres = new_restricted_control_variable(rf, self.restricted_space) elif action == TLMAction: self.xresinterface = self.restricted_control_interface self.yresinterface = self.restricted_functional_interface self.xres = new_restricted_control_variable(rf, self.restricted_space) - self.yres = Function(self.restricted_space).interpolate(rf.functional)._ad_copy()._ad_init_zero(dual=True) + self.yres = ( + Function(self.restricted_space) + .interpolate(rf.functional) + ._ad_copy() + ._ad_init_zero(dual=True) + ) def mult(self, A, x, y): """ @@ -141,7 +163,6 @@ def mult(self, A, x, y): `PETSc.Vec` to which operator is applied to. y (PETSc.Vec): `PETSc.Vec` which is the result of the action of this operator. - """ self.xresinterface.from_petsc(x, self.xres) interpolated_x = interpolate_vars(self.xres, self.function_space) @@ -164,7 +185,6 @@ def multTranspose(self, A, x, y): `PETSc.Vec` to which transpose of operator is applied to. y (PETSc.Vec): `PETSc.Vec` which is the result of the action of this operator. - """ self.yresinterface.from_petsc(x, self.yres) interpolated_x = interpolate_vars(self.yres, self.function_space) @@ -174,8 +194,16 @@ def multTranspose(self, A, x, y): if self._shift != 0: y.axpy(self._shift, x) - -def RestrictedReducedFunctionalMat(rf, action=HessianAction, *, apply_riesz=False, appctx=None, comm=None, restricted_space=None): + +def RestrictedReducedFunctionalMat( + rf, + action=HessianAction, + *, + apply_riesz=False, + appctx=None, + comm=None, + restricted_space=None, +): """ `PETSc.Mat` to apply the action of a linearisation of a `ReducedFunctional`. @@ -194,14 +222,15 @@ def RestrictedReducedFunctionalMat(rf, action=HessianAction, *, apply_riesz=Fals appctx (Optional[dict]): User provided context. comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the `ReducedFunctional` is defined over. restricted_space (Optional[FunctionSpace]): If provided, the control space will be restricted to the passed space. - + Returns ---------- mat (PETSc.Mat): The `PETSc.Mat` whose action and transpose action are defined by the context. """ ctx = RestrictedReducedFunctionalMatCtx( - rf, action, appctx=appctx, apply_riesz=apply_riesz, comm=comm, restricted_space=restricted_space) + rf, action, appctx=appctx, apply_riesz=apply_riesz, comm=comm, restricted_space=restricted_space + ) ncol = ctx.xresinterface.n Ncol = ctx.xresinterface.N @@ -211,7 +240,9 @@ def RestrictedReducedFunctionalMat(rf, action=HessianAction, *, apply_riesz=Fals mat = PETSc.Mat().createPython( ((nrow, Nrow), (ncol, Ncol)), - ctx, comm=ctx.control_interface.comm) + ctx, + comm=ctx.control_interface.comm, + ) if action == HessianAction: mat.setOption(PETSc.Mat.Option.SYMMETRIC, True) mat.setUp() diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index f7b2db1301..4e6198d05a 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -7,16 +7,16 @@ from firedrake import utils from firedrake.restricted_functional_ctx import RestrictedReducedFunctionalMat from firedrake.exceptions import ConvergenceError -from ufl import replace, inner, dx, Argument, Form -from firedrake.petsc import PETSc +from ufl import replace, inner, dx, Form from petsc4py import PETSc from firedrake.adjoint import ReducedFunctional, ReducedFunctionalMat, TLMAction, AdjointAction + try: from slepc4py import SLEPc except ImportError: SLEPc = None -__all__ = ["RFEigenproblem", - "RFEigensolver"] + +__all__ = ["RFEigenproblem", "RFEigensolver"] class RFEigenproblem: @@ -31,64 +31,73 @@ class RFEigenproblem: rf : ReducedFunctional The bilinear operator A. M : ufl.Form - The mass form M(u, v), defaults to inner(u, v) * dx if is None and identity=False. + The mass form M(u, v), defaults to inner(u, v) * dx if None and identity=False. If identity is True, M is treated as identity operator. bcs : DirichletBC or list of DirichletBC The Dirichlet boundary conditions. - bc_shift: float - The value to shift the boundary condition eigenvalues by. This value - will be ignored if restrict==True. - restrict: bool + bc_shift : float + The value to shift the boundary condition eigenvalues by. Ignored if restrict == True. + restrict : bool If True, replace the function spaces of u and v with their restricted version. The output space remains unchanged. - identity: bool + identity : bool If True, M is replaced by identity matrix. Differentiate between mass matrix being identity and inner product operator. - action: (Optional[TLMAction, AdjointAction]): + action : (Optional[TLMAction, AdjointAction]) Defines forward action of implicit matrix operator. - apply_riez: bool - If True, when defining A's adjoint action, we can output into the dual - space. - action: str + apply_riez : bool + If True, when defining A's adjoint action, we can output into the dual space. + action : str Determine whether forward action is TLM or adjoint. - appctx (Optional[dict]): + appctx : Optional[dict] Context forwarded to the linearisation. - comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): + comm : Optional[petsc4py.PETSc.Comm, mpi4py.MPI.Comm] Communicator that the rf is defined over. - - Notes ----- - If restrict==True, the arguments of A and M will be replaced, such that + If restrict == True, the arguments of A and M will be replaced, such that their function space is replaced by the equivalent RestrictedFunctionSpace class. This avoids the computation of eigenvalues associated with the Dirichlet boundary conditions. This in turn prevents convergence failures, and allows only the non-boundary eigenvalues to be returned. The eigenvectors will be in the original, non-restricted space. - If restrict==False and Dirichlet boundary conditions are supplied, then + If restrict == False and Dirichlet boundary conditions are supplied, then these conditions will result in the eigenproblem having a nullspace spanned by the basis functions with support on the boundary. To facilitate solution, this is shifted by the specified amount. It is the user's responsibility to ensure that the shift is not close to an actual eigenvalue of the system. - Also, we assume operator that we linearise maps from V to V. + Also, we assume the operator we linearise maps from V to V. """ - def __init__(self, rf, M=None, bcs=None, bc_shift=0.0, restrict=True, - identity=False, apply_riesz=True, action="tlm", appctx=None, comm=None): + + def __init__( + self, + rf, + M=None, + bcs=None, + bc_shift: float = 0.0, + restrict: bool = True, + identity: bool = False, + apply_riesz: bool = True, + action: str = "tlm", + appctx=None, + comm=None, + ): if not SLEPc: raise ImportError( "Unable to import SLEPc, eigenvalue computation not possible " "(see https://www.firedrakeproject.org/install.html#slepc)" ) - + if not isinstance(rf, ReducedFunctional) and not isinstance(M, ReducedFunctional): raise TypeError( - "At least A or M in the GEP has to be a ReducedFunctional instance. Use 'LinearEigenproblem' otherwise." + "At least A or M in the GEP has to be a ReducedFunctional instance. " + "Use 'LinearEigenproblem' otherwise." ) - + self.output_space = self.output_space(rf) self.bc_shift = bc_shift self.restrict = restrict and bcs @@ -113,58 +122,68 @@ def __init__(self, rf, M=None, bcs=None, bc_shift=0.0, restrict=True, self.bcs = bcs self.M = None if identity else self.as_petsc_mat(M) self.A = self.as_petsc_mat(rf) - + def restrict_obj(self, obj, V_res): """ - Substitutes variational formulation with functions from restricted space if necessary. - Otherwise, tries to convert obj to `SLEPc`-compatible operator. + Substitute variational formulation with functions from restricted space if necessary. + Otherwise, try to convert obj to `SLEPc`-compatible operator. Parameters ---------- obj : Optional[ReducedFunctional, ufl.Form] `Firedrake` operator converted to a 'PETSc.Mat'. V_res : RestrictedFunctionSpace - Function space to which restriction occurs.. - + Function space to which restriction occurs. + Returns ------- - Optional[None, PETSc.Mat]: + Optional[None, PETSc.Mat] `SLEPc`-compatible operator representing obj. """ if isinstance(obj, Form): args = obj.arguments() v, u = args u_res = TrialFunction(V_res) - v_res = TestFunction(V_res) + v_res = TestFunction(V_res) obj = replace(obj, {u: u_res, v: v_res}) return self.as_petsc_mat(obj) if obj else None def as_petsc_mat(self, obj): """ - Convert a `ufl.Form` or a `ReducedFunctional` to a `PETSc.Mat` operator. + Convert a `ufl.Form` or a `ReducedFunctional` to a `PETSc.Mat` operator. Forward action of implicit matrix determined by action. Parameters ---------- obj : Optional[ReducedFunctional, ufl.Form] `Firedrake` operator converted to a 'PETSc.Mat'. - + Returns ------- - `PETSc.Mat` object that represents the operator obj. + PETSc.Mat + Object that represents the operator obj. """ if self.restricted_space: if isinstance(obj, ReducedFunctional): - return RestrictedReducedFunctionalMat(obj, action=self.action, apply_riesz=self.apply_riesz, appctx=self.appctx, - comm=self.comm, restricted_space=self.restricted_space) + return RestrictedReducedFunctionalMat( + obj, + action=self.action, + apply_riesz=self.apply_riesz, + appctx=self.appctx, + comm=self.comm, + restricted_space=self.restricted_space, + ) if isinstance(obj, ReducedFunctional): - return ReducedFunctionalMat(obj, action=self.action, apply_riesz=self.apply_riesz, appctx=self.appctx, comm=self.comm) + return ReducedFunctionalMat( + obj, action=self.action, apply_riesz=self.apply_riesz, appctx=self.appctx, comm=self.comm + ) return assemble( - obj, bcs=self.bcs, - weight=self.bc_shift and 1./self.bc_shift - ).petscmat - - def match_forward_action(self, action): + obj, + bcs=self.bcs, + weight=self.bc_shift and 1.0 / self.bc_shift, + ).petscmat + + def match_forward_action(self, action: str): """ Match forward action of matrix using passed string. @@ -172,11 +191,11 @@ def match_forward_action(self, action): ---------- action : str String representing the forward action of the matrix. 'tlm' or 'adjoint'. - + Returns ------- - Optional[TLMAction, AdjointAction]: - Action corresponding to the string. + Optional[TLMAction, AdjointAction] + Action corresponding to the string. """ match action: case "tlm": @@ -186,17 +205,17 @@ def match_forward_action(self, action): def output_space(self, rf): """ - Determines the function space that operator outputs into. + Determine the function space that the operator outputs into. Parameters ---------- rf : Optional[ReducedFunctional, ufl.Form] - Operator whose output space we want to find. - + Operator whose output space we want to find. + Returns ------- - firedrake.FunctionSpace: - Function space into which operator outputs to. + firedrake.FunctionSpace + Function space into which operator outputs to. """ if isinstance(rf, ReducedFunctional): return rf._controls[0].function_space() @@ -204,45 +223,39 @@ def output_space(self, rf): def generate_inner_form(self): """ - Generates the mass matrix in the `L2`-norm. + Generate the mass matrix in the `L2`-norm. - Parameters - ---------- - space : FunctionSpace - Function space over which we define the mass matrix. - Returns ------- - ufl.Form: - `ufl.Form` corresponding to the `L2`-norm. + ufl.Form + `ufl.Form` corresponding to the `L2`-norm. """ trial = TrialFunction(self.output_space) test = TestFunction(self.output_space) mass_form = inner(trial, test) * dx return mass_form - + def get_bcs_union(self, rf): """ - Finds the union off all the boundary conditions. This is done so as to remove the - nodes related to these boundary conditions from the variational formulation. + Find the union of all the boundary conditions to remove the + nodes related to these boundary conditions from the variational formulation. Parameters ---------- rf : ReducedFunctional - Provides `pyadjoint.Tape` over which collection of the + Provides `pyadjoint.Tape` over which collection of the union of the boundary conditions occurs. - + Returns ------- - List[DirichletBC]]: + list[DirichletBC] List of all unique `DirichletBC`. """ tape = rf.tape - common_dofs = None all_bcs = [] for block in tape._blocks: - if hasattr(block, 'bcs'): + if hasattr(block, "bcs"): bcs = block.bcs print([type(bc) for bc in bcs]) if isinstance(bcs, DirichletBC): @@ -259,7 +272,7 @@ def dm(self): return self.restricted_space.dm else: return self.output_space.dm - + class RFEigensolver(OptionsManager): r"""Solve a LinearEigenproblem. @@ -283,7 +296,6 @@ class RFEigensolver(OptionsManager): Notes ----- - Users will typically wish to set solver parameters specifying the symmetry of the eigenproblem and which eigenvalues to search for first. @@ -306,22 +318,28 @@ class RFEigensolver(OptionsManager): "eps_largest_real": None """ - """ - Preconditioners won't work because of MATSHELL. - Only iterative solvers work as well (refer to PETSc docs). - Spectral transform changed with bcs (default) because sinvert wont work because of MATSHELL. - If info about A and M, could pass HEP to problem_type parameter. If M identity, NHEP. If not, GNHEP. - - - """ - - DEFAULT_EPS_PARAMETERS = {"eps_type": "krylovschur", - "eps_tol": 1e-10, - "eps_target": 0.0} - - def __init__(self, problem, n_evals, *, options_prefix=None, - solver_parameters=None, ncv=None, mpd=None): + # Preconditioners won't work because of MATSHELL. + # Only iterative solvers work as well (refer to PETSc docs). + # Spectral transform changed with bcs (default) because sinvert won't work because of MATSHELL. + # If info about A and M, could pass HEP to problem_type parameter. If M identity, NHEP. If not, GNHEP. + + DEFAULT_EPS_PARAMETERS = { + "eps_type": "krylovschur", + "eps_tol": 1e-10, + "eps_target": 0.0, + } + + def __init__( + self, + problem, + n_evals, + *, + options_prefix=None, + solver_parameters=None, + ncv=None, + mpd=None, + ): self.es = SLEPc.EPS().create(comm=problem.dm.comm) self._problem = problem self.n_evals = n_evals @@ -340,14 +358,14 @@ def solve(self): Returns ------- int - The number of Eigenvalues found. + The number of eigenvalues found. """ A = self._problem.A M = self._problem.M if not isinstance(A, PETSc.Mat): - raise TypeError("A must be a PETSc matrix, got %s" % A.__class__) + raise TypeError(f"A must be a PETSc matrix, got {A.__class__}") if not isinstance(M, PETSc.Mat) and (M is not None): - raise TypeError("M must be a PETSc matrix or None (identity), got %s" % M.__class__) + raise TypeError(f"M must be a PETSc matrix or None (identity), got {M.__class__}") self.es.setDimensions(nev=self.n_evals, ncv=self.ncv, mpd=self.mpd) self.es.setOperators(A, M) @@ -357,7 +375,7 @@ def solve(self): if nconv == 0: raise ConvergenceError("Did not converge any eigenvalues.") return nconv - + def check_es_convergence(self): r"""Check the convergence of the eigenvalue problem.""" r = self.es.getConvergedReason() @@ -368,9 +386,8 @@ def check_es_convergence(self): "try with -eps_converged_reason") if r < 0: raise ConvergenceError( - r"""Eigenproblem failed to converge after %d iterations. - Reason: - %s""" % (self.es.getIterationNumber(), reason) + f"Eigenproblem failed to converge after {self.es.getIterationNumber()} iterations.\n" + f"Reason:\n{reason}" ) def eigenvalue(self, i): @@ -389,12 +406,12 @@ def eigenfunction(self, i): eigenmodes_real = Function(self._problem.restricted_space) eigenmodes_imag = Function(self._problem.restricted_space) else: - eigenmodes_real = Function(self._problem.output_space) # fn of V + eigenmodes_real = Function(self._problem.output_space) eigenmodes_imag = Function(self._problem.output_space) with eigenmodes_real.dat.vec_wo as vr: with eigenmodes_imag.dat.vec_wo as vi: - self.es.getEigenvector(i, vr, vi) # gets the i-th eigenvector + self.es.getEigenvector(i, vr, vi) if self._problem.restrict: eigenmodes_real = Function(self._problem.output_space).interpolate(eigenmodes_real) eigenmodes_imag = Function(self._problem.output_space).interpolate(eigenmodes_imag) - return eigenmodes_real, eigenmodes_imag # returns Firedrake fns + return eigenmodes_real, eigenmodes_imag diff --git a/tests/firedrake/adjoint/test_rf_eigensolver.py b/tests/firedrake/adjoint/test_rf_eigensolver.py index 37b48d85fe..ec93766113 100644 --- a/tests/firedrake/adjoint/test_rf_eigensolver.py +++ b/tests/firedrake/adjoint/test_rf_eigensolver.py @@ -1,21 +1,24 @@ - import pytest -import numpy as np from firedrake import * from firedrake.adjoint import * from firedrake.bcs import DirichletBC from firedrake.rf_eigensolver import * -from firedrake.adjoint import PETScVecInterface, AdjointAction -from firedrake.restricted_functional_ctx import new_restricted_control_variable, interpolate_vars +from firedrake.adjoint import PETScVecInterface +from firedrake.restricted_functional_ctx import ( + new_restricted_control_variable, + interpolate_vars, +) + def generate_expressions(): """ - Returns a list of three expressions + Returns a list of three expressions. """ mesh = UnitSquareMesh(50, 50, quadrilateral=False) x, y = SpatialCoordinate(mesh) return [2 * x, exp(-x - y), sin(pi * x * y)] + @pytest.mark.parametrize("expression", generate_expressions()) def test_compare_eigensolvers_helmholtz(expression): """Test TLM vs assembled form for Helmholtz eigenproblem.""" @@ -37,7 +40,9 @@ def test_compare_eigensolvers_helmholtz(expression): solver_form.solve() nconv = solver_form.es.getConverged() - evals_form = [solver_form.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + evals_form = [ + solver_form.eigenvalue(i) for i in range(0, min(num_eigenvalues, nconv)) + ] # ReducedFunctional-based, Jx = 1/cx continue_annotation() @@ -57,13 +62,16 @@ def test_compare_eigensolvers_helmholtz(expression): solver_rf.solve() nconv = solver_rf.es.getConverged() - evals_rf = [1/solver_rf.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + evals_rf = [ + 1 / solver_rf.eigenvalue(i) for i in range(0, min(num_eigenvalues, nconv)) + ] for ev1, ev2 in zip(evals_form, evals_rf): assert abs(ev1 - ev2) < 1e-6 - + set_working_tape(Tape()) + @pytest.mark.parametrize("restrict", [False]) @pytest.mark.parametrize("riesz", ["L2"]) def test_apply_riesz_param(restrict, riesz): @@ -86,9 +94,13 @@ def test_apply_riesz_param(restrict, riesz): J = ReducedFunctional(u_1, controls=[control]) pause_annotation() - #Check adjoint action - problem_rf_true = RFEigenproblem(J, bcs=bc, apply_riesz=True, restrict=restrict, action="adjoint") - problem_rf_false = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=restrict, action="adjoint") + # Check adjoint action + problem_rf_true = RFEigenproblem( + J, bcs=bc, apply_riesz=True, restrict=restrict, action="adjoint" + ) + problem_rf_false = RFEigenproblem( + J, bcs=bc, apply_riesz=False, restrict=restrict, action="adjoint" + ) mat_true = problem_rf_true.A mat_false = problem_rf_false.A @@ -103,20 +115,30 @@ def test_apply_riesz_param(restrict, riesz): interface_space = problem_rf_true.restricted_space if restrict else V - interface_Ax = PETScVecInterface([Function(interface_space).interpolate(c.control) - for c in J.controls]) + interface_Ax = PETScVecInterface( + [Function(interface_space).interpolate(c.control) for c in J.controls] + ) func_representation = new_restricted_control_variable(J, interface_space, dual=False) - func_representation_dual = new_restricted_control_variable(J, interface_space, dual=True) + func_representation_dual = new_restricted_control_variable( + J, interface_space, dual=True + ) interface_Ax.from_petsc(Ax, func_representation) interface_Ax.from_petsc(Ax2, func_representation_dual) if restrict: - #Intepolate from restricted to unrestricted + # Interpolate from restricted to unrestricted func_representation = interpolate_vars(func_representation, V) - func_representation_dual= interpolate_vars(func_representation_dual, V) - + func_representation_dual = interpolate_vars(func_representation_dual, V) + set_working_tape(Tape()) - assert errornorm(func_representation[0], func_representation_dual[0].riesz_representation(riesz_map=riesz)) < 1e-6 + assert ( + errornorm( + func_representation[0], + func_representation_dual[0].riesz_representation(riesz_map=riesz), + ) + < 1e-6 + ) + def test_compare_tlm_adjoint(): """Test TLM and adjoint give same eigenvalues for Helmholtz eigenproblem.""" @@ -140,22 +162,31 @@ def test_compare_tlm_adjoint(): J = ReducedFunctional(u_1, controls=[control]) pause_annotation() - problem_rf_a = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=True, identity=True, action="adjoint") + problem_rf_a = RFEigenproblem( + J, bcs=bc, apply_riesz=False, restrict=True, identity=True, action="adjoint" + ) solver_rf_a = RFEigensolver(problem_rf_a, n_evals=num_eigenvalues) solver_rf_a.es.setTarget(1) solver_rf_a.solve() nconv = solver_rf_a.es.getConverged() - evals_rf_a = [solver_rf_a.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + evals_rf_a = [ + solver_rf_a.eigenvalue(i) for i in range(0, min(num_eigenvalues, nconv)) + ] - problem_rf_t = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=True, identity=True, action="tlm") + problem_rf_t = RFEigenproblem( + J, bcs=bc, apply_riesz=False, restrict=True, identity=True, action="tlm" + ) solver_rf_t = RFEigensolver(problem_rf_t, n_evals=num_eigenvalues) solver_rf_t.es.setTarget(1) solver_rf_t.solve() nconv = solver_rf_t.es.getConverged() - evals_rf_t = [solver_rf_t.eigenvalue(i) for i in range(0,min(num_eigenvalues, nconv))] + evals_rf_t = [ + solver_rf_t.eigenvalue(i) for i in range(0, min(num_eigenvalues, nconv)) + ] for ev1, ev2 in zip(evals_rf_a, evals_rf_t): print(abs(ev1 - ev2)) assert abs(ev1 - ev2) < 1e-6 + set_working_tape(Tape()) From 949b9d78499ffc9dce9d4bcd8aed92457f42f32a Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 09:11:15 +0300 Subject: [PATCH 09/25] Numpydoc followed. RFEigenproblem docstring made clearer. --- firedrake/restricted_functional_ctx.py | 22 +++++++++++----------- firedrake/rf_eigensolver.py | 10 ++++++++-- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py index 766f9729f6..69d893d179 100644 --- a/firedrake/restricted_functional_ctx.py +++ b/firedrake/restricted_functional_ctx.py @@ -18,14 +18,14 @@ def new_restricted_control_variable(reduced_functional, function_space, dual=Fal by interpolating into the space over which the 'ReducedFunctional' is defined. - Args: + Parameters ---------- reduced_functional (ReducedFunctional): The `ReducedFunctional` whose controls are to be copied. dual (bool): whether to return a dual type. If False then a primal type is returned. - Returns: - ---------- + Returns + ------- tuple[OverloadedType]: New variables suitable for storing a control value. """ return tuple( @@ -38,7 +38,7 @@ def interpolate_vars(variables, function_space): """ Interpolates primal/dual variables to restricted/unrestricted function spaces. - Args: + Parameters ---------- variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), List(firedrake.Function), List(firedrake.Cofunction), @@ -47,8 +47,8 @@ def interpolate_vars(variables, function_space): restricted_space (Optional[FunctionSpace, RestrictedFunctionSpace]): The function space where `PETSc.Vec` objects will live. - Returns: - ---------- + Returns + ------- The same variable/variables but interpolated into the necessary function space. """ if isinstance(variables, (tuple, list)): @@ -69,7 +69,7 @@ class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): `PETSc.Mat.Python context to apply to the operator representing linearisation of a `ReducedFunctional`. Optional restriction of the control space to a provided `FunctionSpace`. - Args: + Parameters ---------- rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. @@ -155,7 +155,7 @@ def mult(self, A, x, y): Compute `y = Ax` and store in `y`. Ax is represented as the forward action of implicit matrix A. - Args: + Parameters ---------- A (PETSc.Mat): Implicit matrix for which Ax is defined. @@ -177,7 +177,7 @@ def multTranspose(self, A, x, y): Compute `y = A^Tx` and store in `y`. A^Tx is represented as the action of the transpose of implicit matrix A. - Args: + Parameters ---------- A (PETSc.Mat): Implicit matrix for which Ax is defined. @@ -213,7 +213,7 @@ def RestrictedReducedFunctionalMat( Adjoint : U* -> V* Hessian : V x U* -> V* | V -> V* - Args: + Parameters ---------- rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. action (RFAction): Whether to apply the TLM, adjoint, or Hessian action. @@ -224,7 +224,7 @@ def RestrictedReducedFunctionalMat( restricted_space (Optional[FunctionSpace]): If provided, the control space will be restricted to the passed space. Returns - ---------- + ------- mat (PETSc.Mat): The `PETSc.Mat` whose action and transpose action are defined by the context. """ diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 4e6198d05a..e8edf0a73d 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -20,11 +20,17 @@ class RFEigenproblem: - """Generalised linear eigenvalue problem. + """Generalised eigenvalue problem. - The problem has the form, find `u`, `λ` such that:: + The linear problem has the form, find `u`, `λ` such that:: A(u, v) = λM(u, v) ∀ v ∈ V + + The new problem that can be formed has the form, find `v`, `k` such that:: + + dJ_m(v) = kv + + where for J: V -> V, dJ_m : V -> V is the linearisation of J around the state in the taped run. Parameters ---------- From 433d065464d082fd32a221d5285cfcccee829d18 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 09:26:38 +0300 Subject: [PATCH 10/25] Tidy up comments. --- firedrake/rf_eigensolver.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index e8edf0a73d..fa16883e09 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -312,6 +312,11 @@ class RFEigensolver(OptionsManager): one would add this entry to `solver_options`:: "eps_gen_hermitian": None + + Note that for eigenproblems arising from taped linearised operators, matrices in the + formulation are implicit (`PETSc.Mat.PYTHON` type), which restricts the available + options to those that do not use the explicit matrix entries. For instance, spectral + transforms won't work and solvers are restricted. As always when specifying PETSc options, `None` indicates that the option in question is a flag and hence doesn't take an argument. @@ -325,11 +330,6 @@ class RFEigensolver(OptionsManager): "eps_largest_real": None """ - # Preconditioners won't work because of MATSHELL. - # Only iterative solvers work as well (refer to PETSc docs). - # Spectral transform changed with bcs (default) because sinvert won't work because of MATSHELL. - # If info about A and M, could pass HEP to problem_type parameter. If M identity, NHEP. If not, GNHEP. - DEFAULT_EPS_PARAMETERS = { "eps_type": "krylovschur", "eps_tol": 1e-10, From 7c365cf2d221962a39380ed7ee1458d70781ecd5 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 09:31:09 +0300 Subject: [PATCH 11/25] Stop inheriting OptionsManager. --- firedrake/rf_eigensolver.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index fa16883e09..289ef01ab0 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -280,7 +280,7 @@ def dm(self): return self.output_space.dm -class RFEigensolver(OptionsManager): +class RFEigensolver: r"""Solve a LinearEigenproblem. Parameters @@ -355,8 +355,8 @@ def __init__( for key in self.DEFAULT_EPS_PARAMETERS: value = self.DEFAULT_EPS_PARAMETERS[key] solver_parameters.setdefault(key, value) - super().__init__(solver_parameters, options_prefix) - self.set_from_options(self.es) + self.options_manager = OptionsManager(solver_parameters, options_prefix) + self.options_manager.set_from_options(self.es) def solve(self): r"""Solve the eigenproblem. @@ -375,7 +375,7 @@ def solve(self): self.es.setDimensions(nev=self.n_evals, ncv=self.ncv, mpd=self.mpd) self.es.setOperators(A, M) - with self.inserted_options(): + with self.options_manager.inserted_options(): self.es.solve() nconv = self.es.getConverged() if nconv == 0: From 61f1208d12353dbd8966adb57d7ed7076f53d7a1 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 10:14:50 +0300 Subject: [PATCH 12/25] Type-hint where needed. Added RFAction to __init__. --- firedrake/restricted_functional_ctx.py | 30 ++++++++++++++++---------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py index 69d893d179..d3d5e3f400 100644 --- a/firedrake/restricted_functional_ctx.py +++ b/firedrake/restricted_functional_ctx.py @@ -1,19 +1,27 @@ from firedrake.adjoint import ( + ReducedFunctional, PETScVecInterface, TLMAction, AdjointAction, HessianAction, + RFAction, ReducedFunctionalMatCtx, ) from firedrake import Function, Cofunction +from typing import Optional try: import petsc4py.PETSc as PETSc + import mpi4py.MPI as MPI except ModuleNotFoundError: PETSc = None +OneOrManyFunction = Function | list[Function] | tuple[Function, ...] +OneOrManyCofunction = Cofunction | list[Cofunction] | tuple[Cofunction, ...] +VarsInterpolate = OneOrManyFunction | OneOrManyCofunction +Comms = PETSc.Comm | MPI.Comm -def new_restricted_control_variable(reduced_functional, function_space, dual=False): +def new_restricted_control_variable(reduced_functional:ReducedFunctional, function_space, dual=False): """Return new variables suitable for storing a control value or its dual by interpolating into the space over which the 'ReducedFunctional' is defined. @@ -34,7 +42,7 @@ def new_restricted_control_variable(reduced_functional, function_space, dual=Fal ) -def interpolate_vars(variables, function_space): +def interpolate_vars(variables:VarsInterpolate, function_space): """ Interpolates primal/dual variables to restricted/unrestricted function spaces. @@ -88,17 +96,17 @@ class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): def __init__( self, - rf, - action=HessianAction, + rf: ReducedFunctional, + action: RFAction = HessianAction, *, apply_riesz=False, - appctx=None, - comm=PETSc.COMM_WORLD, + appctx: Optional[dict] = None, + comm: Comms = PETSc.COMM_WORLD, restricted_space=None, ): if restricted_space is None: raise ValueError( - "restricted_space must be provided for RestrictedReducedFunctionalMatCtx" + "restricted_space must be provided for RestrictedReducedFunctionalMatCtx." ) super().__init__(rf, action=action, apply_riesz=apply_riesz, appctx=appctx, comm=comm) @@ -196,12 +204,12 @@ def multTranspose(self, A, x, y): def RestrictedReducedFunctionalMat( - rf, - action=HessianAction, + rf: ReducedFunctional, + action: RFAction = HessianAction, *, apply_riesz=False, - appctx=None, - comm=None, + appctx: Optional[dict] = None, + comm: Comms = None, restricted_space=None, ): """ From 21c1983e0bf49dec0cd34263af17436d32cc8d9c Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 10:15:15 +0300 Subject: [PATCH 13/25] Now added RFAction to init. --- firedrake/adjoint/__init__.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index e28d8c5c31..8803e72af8 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -29,11 +29,10 @@ from firedrake.adjoint_utils import get_solve_blocks # noqa F401 from pyadjoint.optimization.tao_solver import ( ReducedFunctionalMat, ReducedFunctionalMatCtx, HessianAction, - TLMAction, AdjointAction, check_rf_action, get_valid_comm, - flatten_parameters, PETScVecInterface + TLMAction, AdjointAction, RFAction, check_rf_action, + flatten_parameters, PETScVecInterface, get_valid_comm ) - from pyadjoint.verification import taylor_test, taylor_to_dict # noqa F401 from pyadjoint.drivers import compute_gradient, compute_derivative, compute_hessian # noqa F401 from pyadjoint.adjfloat import AdjFloat # noqa F401 From c2f80abd380df8290a3e43aad2b398f4cca66e86 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 10:29:06 +0300 Subject: [PATCH 14/25] Combined restricted_functional_ctx and rf_eigensolver. --- firedrake/restricted_functional_ctx.py | 258 ------------------------ firedrake/rf_eigensolver.py | 259 ++++++++++++++++++++++++- 2 files changed, 256 insertions(+), 261 deletions(-) delete mode 100644 firedrake/restricted_functional_ctx.py diff --git a/firedrake/restricted_functional_ctx.py b/firedrake/restricted_functional_ctx.py deleted file mode 100644 index d3d5e3f400..0000000000 --- a/firedrake/restricted_functional_ctx.py +++ /dev/null @@ -1,258 +0,0 @@ -from firedrake.adjoint import ( - ReducedFunctional, - PETScVecInterface, - TLMAction, - AdjointAction, - HessianAction, - RFAction, - ReducedFunctionalMatCtx, -) -from firedrake import Function, Cofunction -from typing import Optional - -try: - import petsc4py.PETSc as PETSc - import mpi4py.MPI as MPI -except ModuleNotFoundError: - PETSc = None - -OneOrManyFunction = Function | list[Function] | tuple[Function, ...] -OneOrManyCofunction = Cofunction | list[Cofunction] | tuple[Cofunction, ...] -VarsInterpolate = OneOrManyFunction | OneOrManyCofunction -Comms = PETSc.Comm | MPI.Comm - -def new_restricted_control_variable(reduced_functional:ReducedFunctional, function_space, dual=False): - """Return new variables suitable for storing a control value or its dual - by interpolating into the space over which the 'ReducedFunctional' is - defined. - - Parameters - ---------- - reduced_functional (ReducedFunctional): The `ReducedFunctional` whose - controls are to be copied. - dual (bool): whether to return a dual type. If False then a primal type is returned. - - Returns - ------- - tuple[OverloadedType]: New variables suitable for storing a control value. - """ - return tuple( - Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) - for control in reduced_functional.controls - ) - - -def interpolate_vars(variables:VarsInterpolate, function_space): - """ - Interpolates primal/dual variables to restricted/unrestricted function spaces. - - Parameters - ---------- - variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), - List(firedrake.Function), List(firedrake.Cofunction), - firedrake.Function, firedrake.Cofunction): - Variables that are to be interpolated into unrestricted/restricted primal/dual spaces. - restricted_space (Optional[FunctionSpace, RestrictedFunctionSpace]): - The function space where `PETSc.Vec` objects will live. - - Returns - ------- - The same variable/variables but interpolated into the necessary function space. - """ - if isinstance(variables, (tuple, list)): - if isinstance(variables[0], Function): - return tuple(Function(function_space).interpolate(v) for v in variables) - else: - return tuple( - Cofunction(function_space.dual()).interpolate(v) for v in variables - ) - elif isinstance(variables, Function): - return Function(function_space).interpolate(variables) - else: - return Cofunction(function_space.dual()).interpolate(variables) - - -class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): - """ - `PETSc.Mat.Python context to apply to the operator representing linearisation of a `ReducedFunctional`. - Optional restriction of the control space to a provided `FunctionSpace`. - - Parameters - ---------- - rf (ReducedFunctional): - Defines the forward model, and used to compute operator actions. - action (RFAction): - Whether to apply the TLM, adjoint, or Hessian action. - apply_riesz (bool): - Whether to apply the Riesz map before returning the - result of the action to `PETSc`. - appctx (Optional[dict]): - User provided context. - comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): - Communicator that the `ReducedFunctional` is defined over. - restricted_space (Optional[FunctionSpace]): - If provided, the control space will be restricted to the passed space. - """ - - def __init__( - self, - rf: ReducedFunctional, - action: RFAction = HessianAction, - *, - apply_riesz=False, - appctx: Optional[dict] = None, - comm: Comms = PETSc.COMM_WORLD, - restricted_space=None, - ): - if restricted_space is None: - raise ValueError( - "restricted_space must be provided for RestrictedReducedFunctionalMatCtx." - ) - - super().__init__(rf, action=action, apply_riesz=apply_riesz, appctx=appctx, comm=comm) - - self.function_space = rf.controls[0].control.function_space() - self.restricted_space = restricted_space or self.function_space - - # Build restricted interfaces - self.restricted_control_interface = ( - PETScVecInterface( - [Function(self.restricted_space).interpolate(c.control) for c in rf.controls] - ) - if restricted_space - else self.control_interface - ) - - if action in (AdjointAction, TLMAction): - self.restricted_functional_interface = ( - PETScVecInterface(Function(restricted_space).interpolate(rf.functional), comm=comm) - if restricted_space - else self.functional_interface - ) - - if action == HessianAction: - self.xresinterface = self.restricted_control_interface - self.yresinterface = self.restricted_control_interface - self.xres = new_restricted_control_variable(rf, self.restricted_space) - self.yres = new_restricted_control_variable(rf, self.restricted_space) - - elif action == AdjointAction: - self.xresinterface = self.restricted_functional_interface - self.yresinterface = self.restricted_control_interface - self.xres = ( - Function(self.restricted_space) - .interpolate(rf.functional) - ._ad_copy() - ._ad_init_zero(dual=True) - ) - self.yres = new_restricted_control_variable(rf, self.restricted_space) - - elif action == TLMAction: - self.xresinterface = self.restricted_control_interface - self.yresinterface = self.restricted_functional_interface - self.xres = new_restricted_control_variable(rf, self.restricted_space) - self.yres = ( - Function(self.restricted_space) - .interpolate(rf.functional) - ._ad_copy() - ._ad_init_zero(dual=True) - ) - - def mult(self, A, x, y): - """ - Compute `y = Ax` and store in `y`. - Ax is represented as the forward action of implicit matrix A. - - Parameters - ---------- - A (PETSc.Mat): - Implicit matrix for which Ax is defined. - x (PETSc.Vec): - `PETSc.Vec` to which operator is applied to. - y (PETSc.Vec): - `PETSc.Vec` which is the result of the action of this operator. - """ - self.xresinterface.from_petsc(x, self.xres) - interpolated_x = interpolate_vars(self.xres, self.function_space) - out = self.mult_impl(A, interpolated_x) - interpolated_y = interpolate_vars(out, self.restricted_space) - self.yresinterface.to_petsc(y, interpolated_y) - if self._shift != 0: - y.axpy(self._shift, x) - - def multTranspose(self, A, x, y): - """ - Compute `y = A^Tx` and store in `y`. - A^Tx is represented as the action of the transpose of implicit matrix A. - - Parameters - ---------- - A (PETSc.Mat): - Implicit matrix for which Ax is defined. - x (PETSc.Vec): - `PETSc.Vec` to which transpose of operator is applied to. - y (PETSc.Vec): - `PETSc.Vec` which is the result of the action of this operator. - """ - self.yresinterface.from_petsc(x, self.yres) - interpolated_x = interpolate_vars(self.yres, self.function_space) - out = self.mult_impl_transpose(A, interpolated_x) - interpolated_y = interpolate_vars(out, self.restricted_space) - self.xresinterface.to_petsc(y, interpolated_y) - if self._shift != 0: - y.axpy(self._shift, x) - - -def RestrictedReducedFunctionalMat( - rf: ReducedFunctional, - action: RFAction = HessianAction, - *, - apply_riesz=False, - appctx: Optional[dict] = None, - comm: Comms = None, - restricted_space=None, -): - """ - `PETSc.Mat` to apply the action of a linearisation of a `ReducedFunctional`. - - If V is the control space and U is the functional space, each action has the following map: - Jhat : V -> U - TLM : V -> U - Adjoint : U* -> V* - Hessian : V x U* -> V* | V -> V* - - Parameters - ---------- - rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. - action (RFAction): Whether to apply the TLM, adjoint, or Hessian action. - apply_riesz (bool): Whether to apply the Riesz map before returning the - result of the action to `PETSc`. - appctx (Optional[dict]): User provided context. - comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the `ReducedFunctional` is defined over. - restricted_space (Optional[FunctionSpace]): If provided, the control space will be restricted to the passed space. - - Returns - ------- - mat (PETSc.Mat): - The `PETSc.Mat` whose action and transpose action are defined by the context. - """ - ctx = RestrictedReducedFunctionalMatCtx( - rf, action, appctx=appctx, apply_riesz=apply_riesz, comm=comm, restricted_space=restricted_space - ) - - ncol = ctx.xresinterface.n - Ncol = ctx.xresinterface.N - - nrow = ctx.yresinterface.n - Nrow = ctx.yresinterface.N - - mat = PETSc.Mat().createPython( - ((nrow, Nrow), (ncol, Ncol)), - ctx, - comm=ctx.control_interface.comm, - ) - if action == HessianAction: - mat.setOption(PETSc.Mat.Option.SYMMETRIC, True) - mat.setUp() - mat.assemble() - return mat diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 289ef01ab0..0d215e66d9 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -2,20 +2,36 @@ from petsctools import OptionsManager, flatten_parameters from firedrake.assemble import assemble from firedrake.bcs import extract_subdomain_ids, restricted_function_space, DirichletBC -from firedrake.function import Function +from firedrake.function import Function, Cofunction from firedrake.ufl_expr import TrialFunction, TestFunction from firedrake import utils from firedrake.restricted_functional_ctx import RestrictedReducedFunctionalMat from firedrake.exceptions import ConvergenceError from ufl import replace, inner, dx, Form from petsc4py import PETSc -from firedrake.adjoint import ReducedFunctional, ReducedFunctionalMat, TLMAction, AdjointAction +from firedrake.adjoint import ( + ReducedFunctional, + PETScVecInterface, + TLMAction, + AdjointAction, + HessianAction, + RFAction, + ReducedFunctionalMat, + ReducedFunctionalMatCtx, +) +from typing import Optional try: + import petsc4py.PETSc as PETSc from slepc4py import SLEPc -except ImportError: +except ModuleNotFoundError: + PETSc = None SLEPc = None +OneOrManyFunction = Function | list[Function] | tuple[Function, ...] +OneOrManyCofunction = Cofunction | list[Cofunction] | tuple[Cofunction, ...] +VarsInterpolate = OneOrManyFunction | OneOrManyCofunction + __all__ = ["RFEigenproblem", "RFEigensolver"] @@ -421,3 +437,240 @@ def eigenfunction(self, i): eigenmodes_real = Function(self._problem.output_space).interpolate(eigenmodes_real) eigenmodes_imag = Function(self._problem.output_space).interpolate(eigenmodes_imag) return eigenmodes_real, eigenmodes_imag + + +def new_restricted_control_variable(reduced_functional:ReducedFunctional, function_space, dual=False): + """Return new variables suitable for storing a control value or its dual + by interpolating into the space over which the 'ReducedFunctional' is + defined. + + Parameters + ---------- + reduced_functional (ReducedFunctional): The `ReducedFunctional` whose + controls are to be copied. + dual (bool): whether to return a dual type. If False then a primal type is returned. + + Returns + ------- + tuple[OverloadedType]: New variables suitable for storing a control value. + """ + return tuple( + Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) + for control in reduced_functional.controls + ) + + +def interpolate_vars(variables:VarsInterpolate, function_space): + """ + Interpolates primal/dual variables to restricted/unrestricted function spaces. + + Parameters + ---------- + variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), + List(firedrake.Function), List(firedrake.Cofunction), + firedrake.Function, firedrake.Cofunction): + Variables that are to be interpolated into unrestricted/restricted primal/dual spaces. + restricted_space (Optional[FunctionSpace, RestrictedFunctionSpace]): + The function space where `PETSc.Vec` objects will live. + + Returns + ------- + The same variable/variables but interpolated into the necessary function space. + """ + if isinstance(variables, (tuple, list)): + if isinstance(variables[0], Function): + return tuple(Function(function_space).interpolate(v) for v in variables) + else: + return tuple( + Cofunction(function_space.dual()).interpolate(v) for v in variables + ) + elif isinstance(variables, Function): + return Function(function_space).interpolate(variables) + else: + return Cofunction(function_space.dual()).interpolate(variables) + + +class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): + """ + `PETSc.Mat.Python context to apply to the operator representing linearisation of a `ReducedFunctional`. + Optional restriction of the control space to a provided `FunctionSpace`. + + Parameters + ---------- + rf (ReducedFunctional): + Defines the forward model, and used to compute operator actions. + action (RFAction): + Whether to apply the TLM, adjoint, or Hessian action. + apply_riesz (bool): + Whether to apply the Riesz map before returning the + result of the action to `PETSc`. + appctx (Optional[dict]): + User provided context. + comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): + Communicator that the `ReducedFunctional` is defined over. + restricted_space (Optional[FunctionSpace]): + If provided, the control space will be restricted to the passed space. + """ + + def __init__( + self, + rf: ReducedFunctional, + action: RFAction = HessianAction, + *, + apply_riesz=False, + appctx: Optional[dict] = None, + comm=PETSc.COMM_WORLD, + restricted_space=None, + ): + if restricted_space is None: + raise ValueError( + "restricted_space must be provided for RestrictedReducedFunctionalMatCtx." + ) + + super().__init__(rf, action=action, apply_riesz=apply_riesz, appctx=appctx, comm=comm) + + self.function_space = rf.controls[0].control.function_space() + self.restricted_space = restricted_space or self.function_space + + # Build restricted interfaces + self.restricted_control_interface = ( + PETScVecInterface( + [Function(self.restricted_space).interpolate(c.control) for c in rf.controls] + ) + if restricted_space + else self.control_interface + ) + + if action in (AdjointAction, TLMAction): + self.restricted_functional_interface = ( + PETScVecInterface(Function(restricted_space).interpolate(rf.functional), comm=comm) + if restricted_space + else self.functional_interface + ) + + if action == HessianAction: + self.xresinterface = self.restricted_control_interface + self.yresinterface = self.restricted_control_interface + self.xres = new_restricted_control_variable(rf, self.restricted_space) + self.yres = new_restricted_control_variable(rf, self.restricted_space) + + elif action == AdjointAction: + self.xresinterface = self.restricted_functional_interface + self.yresinterface = self.restricted_control_interface + self.xres = ( + Function(self.restricted_space) + .interpolate(rf.functional) + ._ad_copy() + ._ad_init_zero(dual=True) + ) + self.yres = new_restricted_control_variable(rf, self.restricted_space) + + elif action == TLMAction: + self.xresinterface = self.restricted_control_interface + self.yresinterface = self.restricted_functional_interface + self.xres = new_restricted_control_variable(rf, self.restricted_space) + self.yres = ( + Function(self.restricted_space) + .interpolate(rf.functional) + ._ad_copy() + ._ad_init_zero(dual=True) + ) + + def mult(self, A, x, y): + """ + Compute `y = Ax` and store in `y`. + Ax is represented as the forward action of implicit matrix A. + + Parameters + ---------- + A (PETSc.Mat): + Implicit matrix for which Ax is defined. + x (PETSc.Vec): + `PETSc.Vec` to which operator is applied to. + y (PETSc.Vec): + `PETSc.Vec` which is the result of the action of this operator. + """ + self.xresinterface.from_petsc(x, self.xres) + interpolated_x = interpolate_vars(self.xres, self.function_space) + out = self.mult_impl(A, interpolated_x) + interpolated_y = interpolate_vars(out, self.restricted_space) + self.yresinterface.to_petsc(y, interpolated_y) + if self._shift != 0: + y.axpy(self._shift, x) + + def multTranspose(self, A, x, y): + """ + Compute `y = A^Tx` and store in `y`. + A^Tx is represented as the action of the transpose of implicit matrix A. + + Parameters + ---------- + A (PETSc.Mat): + Implicit matrix for which Ax is defined. + x (PETSc.Vec): + `PETSc.Vec` to which transpose of operator is applied to. + y (PETSc.Vec): + `PETSc.Vec` which is the result of the action of this operator. + """ + self.yresinterface.from_petsc(x, self.yres) + interpolated_x = interpolate_vars(self.yres, self.function_space) + out = self.mult_impl_transpose(A, interpolated_x) + interpolated_y = interpolate_vars(out, self.restricted_space) + self.xresinterface.to_petsc(y, interpolated_y) + if self._shift != 0: + y.axpy(self._shift, x) + + +def RestrictedReducedFunctionalMat( + rf: ReducedFunctional, + action: RFAction = HessianAction, + *, + apply_riesz=False, + appctx: Optional[dict] = None, + comm=None, + restricted_space=None, +): + """ + `PETSc.Mat` to apply the action of a linearisation of a `ReducedFunctional`. + + If V is the control space and U is the functional space, each action has the following map: + Jhat : V -> U + TLM : V -> U + Adjoint : U* -> V* + Hessian : V x U* -> V* | V -> V* + + Parameters + ---------- + rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. + action (RFAction): Whether to apply the TLM, adjoint, or Hessian action. + apply_riesz (bool): Whether to apply the Riesz map before returning the + result of the action to `PETSc`. + appctx (Optional[dict]): User provided context. + comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the `ReducedFunctional` is defined over. + restricted_space (Optional[FunctionSpace]): If provided, the control space will be restricted to the passed space. + + Returns + ------- + mat (PETSc.Mat): + The `PETSc.Mat` whose action and transpose action are defined by the context. + """ + ctx = RestrictedReducedFunctionalMatCtx( + rf, action, appctx=appctx, apply_riesz=apply_riesz, comm=comm, restricted_space=restricted_space + ) + + ncol = ctx.xresinterface.n + Ncol = ctx.xresinterface.N + + nrow = ctx.yresinterface.n + Nrow = ctx.yresinterface.N + + mat = PETSc.Mat().createPython( + ((nrow, Nrow), (ncol, Ncol)), + ctx, + comm=ctx.control_interface.comm, + ) + if action == HessianAction: + mat.setOption(PETSc.Mat.Option.SYMMETRIC, True) + mat.setUp() + mat.assemble() + return mat From 9343a04249c64adf48315db9ec32291914902d2b Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 10:29:40 +0300 Subject: [PATCH 15/25] Removed the old import. --- firedrake/rf_eigensolver.py | 1 - 1 file changed, 1 deletion(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 0d215e66d9..f39a0a27f8 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -5,7 +5,6 @@ from firedrake.function import Function, Cofunction from firedrake.ufl_expr import TrialFunction, TestFunction from firedrake import utils -from firedrake.restricted_functional_ctx import RestrictedReducedFunctionalMat from firedrake.exceptions import ConvergenceError from ufl import replace, inner, dx, Form from petsc4py import PETSc From 30a91bf8894a3eac0d17f3f9104d6e2298f6ebb4 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 12:06:30 +0300 Subject: [PATCH 16/25] "=" changed to is. --- firedrake/rf_eigensolver.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index f39a0a27f8..6a4c7f3577 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -547,13 +547,13 @@ def __init__( else self.functional_interface ) - if action == HessianAction: + if action is HessianAction: self.xresinterface = self.restricted_control_interface self.yresinterface = self.restricted_control_interface self.xres = new_restricted_control_variable(rf, self.restricted_space) self.yres = new_restricted_control_variable(rf, self.restricted_space) - elif action == AdjointAction: + elif action is AdjointAction: self.xresinterface = self.restricted_functional_interface self.yresinterface = self.restricted_control_interface self.xres = ( @@ -564,7 +564,7 @@ def __init__( ) self.yres = new_restricted_control_variable(rf, self.restricted_space) - elif action == TLMAction: + elif action is TLMAction: self.xresinterface = self.restricted_control_interface self.yresinterface = self.restricted_functional_interface self.xres = new_restricted_control_variable(rf, self.restricted_space) @@ -668,7 +668,7 @@ def RestrictedReducedFunctionalMat( ctx, comm=ctx.control_interface.comm, ) - if action == HessianAction: + if action is HessianAction: mat.setOption(PETSc.Mat.Option.SYMMETRIC, True) mat.setUp() mat.assemble() From ae46cdf218e9662238f9d2fd454864f08106bf13 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 12:55:08 +0300 Subject: [PATCH 17/25] Type-hinted necessary things. --- firedrake/rf_eigensolver.py | 91 ++++++++++++++++--------------------- 1 file changed, 38 insertions(+), 53 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 6a4c7f3577..ce8a40badf 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -6,6 +6,7 @@ from firedrake.ufl_expr import TrialFunction, TestFunction from firedrake import utils from firedrake.exceptions import ConvergenceError +from firedrake.functionspaceimpl import WithGeometry from ufl import replace, inner, dx, Form from petsc4py import PETSc from firedrake.adjoint import ( @@ -27,11 +28,21 @@ PETSc = None SLEPc = None +try: + import mpi4py.MPI as MPI +except ModuleNotFoundError: + MPI = None + + OneOrManyFunction = Function | list[Function] | tuple[Function, ...] OneOrManyCofunction = Cofunction | list[Cofunction] | tuple[Cofunction, ...] VarsInterpolate = OneOrManyFunction | OneOrManyCofunction +Comms = PETSc.Comm | MPI.Comm | None +AppCtxType = dict | None +MassType = ReducedFunctional | Form | None +BcType = DirichletBC | list[DirichletBC] -__all__ = ["RFEigenproblem", "RFEigensolver"] +__all__ = ["RFEigenproblem", "RFEigensolver", "RestrictedReducedFunctionalMat"] class RFEigenproblem: @@ -96,16 +107,16 @@ class RFEigenproblem: def __init__( self, - rf, - M=None, - bcs=None, + rf: ReducedFunctional, + M: MassType = None, + bcs: BcType = None, bc_shift: float = 0.0, restrict: bool = True, identity: bool = False, apply_riesz: bool = True, action: str = "tlm", - appctx=None, - comm=None, + appctx: AppCtxType = None, + comm: Comms = None, ): if not SLEPc: raise ImportError( @@ -144,7 +155,7 @@ def __init__( self.M = None if identity else self.as_petsc_mat(M) self.A = self.as_petsc_mat(rf) - def restrict_obj(self, obj, V_res): + def restrict_obj(self, obj: MassType, V_res: WithGeometry): """ Substitute variational formulation with functions from restricted space if necessary. Otherwise, try to convert obj to `SLEPc`-compatible operator. @@ -169,7 +180,7 @@ def restrict_obj(self, obj, V_res): obj = replace(obj, {u: u_res, v: v_res}) return self.as_petsc_mat(obj) if obj else None - def as_petsc_mat(self, obj): + def as_petsc_mat(self, obj: MassType): """ Convert a `ufl.Form` or a `ReducedFunctional` to a `PETSc.Mat` operator. Forward action of implicit matrix determined by action. @@ -224,7 +235,7 @@ def match_forward_action(self, action: str): case "adjoint": return AdjointAction - def output_space(self, rf): + def output_space(self, rf: MassType): """ Determine the function space that the operator outputs into. @@ -256,7 +267,7 @@ def generate_inner_form(self): mass_form = inner(trial, test) * dx return mass_form - def get_bcs_union(self, rf): + def get_bcs_union(self, rf: ReducedFunctional): """ Find the union of all the boundary conditions to remove the nodes related to these boundary conditions from the variational formulation. @@ -353,13 +364,13 @@ class RFEigensolver: def __init__( self, - problem, - n_evals, + problem: RFEigenproblem, + n_evals: int, *, - options_prefix=None, - solver_parameters=None, - ncv=None, - mpd=None, + options_prefix: str = None, + solver_parameters: dict | None = None, + ncv: int = None, + mpd: int = None, ): self.es = SLEPc.EPS().create(comm=problem.dm.comm) self._problem = problem @@ -411,11 +422,11 @@ def check_es_convergence(self): f"Reason:\n{reason}" ) - def eigenvalue(self, i): + def eigenvalue(self, i: int): r"""Return the i-th eigenvalue of the solved problem.""" return self.es.getEigenvalue(i) - def eigenfunction(self, i): + def eigenfunction(self, i: int): r"""Return the i-th eigenfunction of the solved problem. Returns @@ -438,7 +449,7 @@ def eigenfunction(self, i): return eigenmodes_real, eigenmodes_imag -def new_restricted_control_variable(reduced_functional:ReducedFunctional, function_space, dual=False): +def new_restricted_control_variable(reduced_functional: ReducedFunctional, function_space: WithGeometry, dual: bool=False): """Return new variables suitable for storing a control value or its dual by interpolating into the space over which the 'ReducedFunctional' is defined. @@ -459,7 +470,7 @@ def new_restricted_control_variable(reduced_functional:ReducedFunctional, functi ) -def interpolate_vars(variables:VarsInterpolate, function_space): +def interpolate_vars(variables: VarsInterpolate, function_space: WithGeometry): """ Interpolates primal/dual variables to restricted/unrestricted function spaces. @@ -516,10 +527,10 @@ def __init__( rf: ReducedFunctional, action: RFAction = HessianAction, *, - apply_riesz=False, - appctx: Optional[dict] = None, - comm=PETSc.COMM_WORLD, - restricted_space=None, + apply_riesz: bool = False, + appctx: AppCtxType = None, + comm: Comms = PETSc.COMM_WORLD, + restricted_space: Optional[WithGeometry] = None, ): if restricted_space is None: raise ValueError( @@ -576,19 +587,6 @@ def __init__( ) def mult(self, A, x, y): - """ - Compute `y = Ax` and store in `y`. - Ax is represented as the forward action of implicit matrix A. - - Parameters - ---------- - A (PETSc.Mat): - Implicit matrix for which Ax is defined. - x (PETSc.Vec): - `PETSc.Vec` to which operator is applied to. - y (PETSc.Vec): - `PETSc.Vec` which is the result of the action of this operator. - """ self.xresinterface.from_petsc(x, self.xres) interpolated_x = interpolate_vars(self.xres, self.function_space) out = self.mult_impl(A, interpolated_x) @@ -598,19 +596,6 @@ def mult(self, A, x, y): y.axpy(self._shift, x) def multTranspose(self, A, x, y): - """ - Compute `y = A^Tx` and store in `y`. - A^Tx is represented as the action of the transpose of implicit matrix A. - - Parameters - ---------- - A (PETSc.Mat): - Implicit matrix for which Ax is defined. - x (PETSc.Vec): - `PETSc.Vec` to which transpose of operator is applied to. - y (PETSc.Vec): - `PETSc.Vec` which is the result of the action of this operator. - """ self.yresinterface.from_petsc(x, self.yres) interpolated_x = interpolate_vars(self.yres, self.function_space) out = self.mult_impl_transpose(A, interpolated_x) @@ -624,10 +609,10 @@ def RestrictedReducedFunctionalMat( rf: ReducedFunctional, action: RFAction = HessianAction, *, - apply_riesz=False, + apply_riesz: bool = False, appctx: Optional[dict] = None, - comm=None, - restricted_space=None, + comm: Comms = None, + restricted_space: WithGeometry | None = None, ): """ `PETSc.Mat` to apply the action of a linearisation of a `ReducedFunctional`. From d610892ff094ad65d868791b81c3db6b2678dfb2 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 13:03:27 +0300 Subject: [PATCH 18/25] No double type-hinting. --- firedrake/rf_eigensolver.py | 49 ++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 25 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index ce8a40badf..1a73f0d34f 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -162,9 +162,9 @@ def restrict_obj(self, obj: MassType, V_res: WithGeometry): Parameters ---------- - obj : Optional[ReducedFunctional, ufl.Form] + obj : `Firedrake` operator converted to a 'PETSc.Mat'. - V_res : RestrictedFunctionSpace + V_res : Function space to which restriction occurs. Returns @@ -187,7 +187,7 @@ def as_petsc_mat(self, obj: MassType): Parameters ---------- - obj : Optional[ReducedFunctional, ufl.Form] + obj : `Firedrake` operator converted to a 'PETSc.Mat'. Returns @@ -221,7 +221,7 @@ def match_forward_action(self, action: str): Parameters ---------- - action : str + action : String representing the forward action of the matrix. 'tlm' or 'adjoint'. Returns @@ -241,12 +241,12 @@ def output_space(self, rf: MassType): Parameters ---------- - rf : Optional[ReducedFunctional, ufl.Form] + rf : Operator whose output space we want to find. Returns ------- - firedrake.FunctionSpace + firedrake.functionspaceimpl.WithGeometry Function space into which operator outputs to. """ if isinstance(rf, ReducedFunctional): @@ -274,7 +274,7 @@ def get_bcs_union(self, rf: ReducedFunctional): Parameters ---------- - rf : ReducedFunctional + rf : Provides `pyadjoint.Tape` over which collection of the union of the boundary conditions occurs. @@ -456,9 +456,10 @@ def new_restricted_control_variable(reduced_functional: ReducedFunctional, funct Parameters ---------- - reduced_functional (ReducedFunctional): The `ReducedFunctional` whose + reduced_functional: The `ReducedFunctional` whose controls are to be copied. - dual (bool): whether to return a dual type. If False then a primal type is returned. + function_space: Function space to which we restrict variables to. + dual: whether to return a dual type. If False then a primal type is returned. Returns ------- @@ -476,11 +477,9 @@ def interpolate_vars(variables: VarsInterpolate, function_space: WithGeometry): Parameters ---------- - variables (Optional[tuple(firedrake.Function), tuple(firedrake.Cofunction), - List(firedrake.Function), List(firedrake.Cofunction), - firedrake.Function, firedrake.Cofunction): + variables: Variables that are to be interpolated into unrestricted/restricted primal/dual spaces. - restricted_space (Optional[FunctionSpace, RestrictedFunctionSpace]): + restricted_space: The function space where `PETSc.Vec` objects will live. Returns @@ -507,18 +506,18 @@ class RestrictedReducedFunctionalMatCtx(ReducedFunctionalMatCtx): Parameters ---------- - rf (ReducedFunctional): + rf: Defines the forward model, and used to compute operator actions. - action (RFAction): + action: Whether to apply the TLM, adjoint, or Hessian action. - apply_riesz (bool): + apply_riesz: Whether to apply the Riesz map before returning the result of the action to `PETSc`. - appctx (Optional[dict]): + appctx: User provided context. - comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): + comm: Communicator that the `ReducedFunctional` is defined over. - restricted_space (Optional[FunctionSpace]): + restricted_space: If provided, the control space will be restricted to the passed space. """ @@ -625,13 +624,13 @@ def RestrictedReducedFunctionalMat( Parameters ---------- - rf (ReducedFunctional): Defines the forward model, and used to compute operator actions. - action (RFAction): Whether to apply the TLM, adjoint, or Hessian action. - apply_riesz (bool): Whether to apply the Riesz map before returning the + rf: Defines the forward model, and used to compute operator actions. + action: Whether to apply the TLM, adjoint, or Hessian action. + apply_riesz: Whether to apply the Riesz map before returning the result of the action to `PETSc`. - appctx (Optional[dict]): User provided context. - comm (Optional[petsc4py.PETSc.Comm,mpi4py.MPI.Comm]): Communicator that the `ReducedFunctional` is defined over. - restricted_space (Optional[FunctionSpace]): If provided, the control space will be restricted to the passed space. + appctx: User provided context. + comm: Communicator that the `ReducedFunctional` is defined over. + restricted_space: If provided, the control space will be restricted to the passed space. Returns ------- From 9ed9f14d4e852c5005c6c9934d6f3d10b8036c49 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 13:07:54 +0300 Subject: [PATCH 19/25] Clearer docstring in RFEigensolver. --- firedrake/rf_eigensolver.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 1a73f0d34f..a9775bc561 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -307,22 +307,26 @@ def dm(self): class RFEigensolver: - r"""Solve a LinearEigenproblem. + r"""Solves an RFEigenproblem. + This is a generalisation of the linear eigenproblem which is + additionally able to cope with eigenproblems arising from + linearisations of operators represented by sequences of + Firedrake operations. Parameters ---------- - problem : LinearEigenproblem + problem : The eigenproblem to solve. - n_evals : int + n_evals : The number of eigenvalues to compute. - options_prefix : str + options_prefix : The options prefix to use for the eigensolver. - solver_parameters : dict + solver_parameters : PETSc options for the eigenvalue problem. - ncv : int + ncv : Maximum dimension of the subspace to be used by the solver. See `SLEPc.EPS.setDimensions`. - mpd : int + mpd : Maximum dimension allowed for the projected problem. See `SLEPc.EPS.setDimensions`. From 58a40ad3975f555abcbc94b7989c2ad9dff6649b Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 13:14:56 +0300 Subject: [PATCH 20/25] Updated doc for more clarity. --- firedrake/rf_eigensolver.py | 34 ++++++++++++++++------------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index a9775bc561..7b1522dc2a 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -46,44 +46,42 @@ class RFEigenproblem: - """Generalised eigenvalue problem. - - The linear problem has the form, find `u`, `λ` such that:: - - A(u, v) = λM(u, v) ∀ v ∈ V - - The new problem that can be formed has the form, find `v`, `k` such that:: + """ + The new problem that can be formed has the form, find `u`, `λ` such that:: - dJ_m(v) = kv + dJ_m(u) = λu where for J: V -> V, dJ_m : V -> V is the linearisation of J around the state in the taped run. + + Additionally, the class has the capabilities of solving the generalised eigenvalue problem + defined in `LinearEigenproblem`. Parameters ---------- - rf : ReducedFunctional + rf : The bilinear operator A. - M : ufl.Form + M : The mass form M(u, v), defaults to inner(u, v) * dx if None and identity=False. If identity is True, M is treated as identity operator. - bcs : DirichletBC or list of DirichletBC + bcs : The Dirichlet boundary conditions. - bc_shift : float + bc_shift : The value to shift the boundary condition eigenvalues by. Ignored if restrict == True. - restrict : bool + restrict : If True, replace the function spaces of u and v with their restricted version. The output space remains unchanged. - identity : bool + identity : If True, M is replaced by identity matrix. Differentiate between mass matrix being identity and inner product operator. action : (Optional[TLMAction, AdjointAction]) Defines forward action of implicit matrix operator. - apply_riez : bool + apply_riez : If True, when defining A's adjoint action, we can output into the dual space. - action : str + action : Determine whether forward action is TLM or adjoint. - appctx : Optional[dict] + appctx : Context forwarded to the linearisation. - comm : Optional[petsc4py.PETSc.Comm, mpi4py.MPI.Comm] + comm : Communicator that the rf is defined over. Notes From 4578bc2da6d564d6aa4f92833604634e7d8163f4 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 19:12:31 +0300 Subject: [PATCH 21/25] Removed action and added options. --- firedrake/rf_eigensolver.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 7b1522dc2a..195675399b 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -73,12 +73,10 @@ class RFEigenproblem: identity : If True, M is replaced by identity matrix. Differentiate between mass matrix being identity and inner product operator. - action : (Optional[TLMAction, AdjointAction]) - Defines forward action of implicit matrix operator. apply_riez : If True, when defining A's adjoint action, we can output into the dual space. - action : - Determine whether forward action is TLM or adjoint. + action : (Union["tlm", "adjoint"]) + Defines forward action of implicit matrix operator. appctx : Context forwarded to the linearisation. comm : From 88062bb96810295484306697df2abedee442d788 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 21:13:12 +0300 Subject: [PATCH 22/25] Ensure argument and docstring align. --- firedrake/rf_eigensolver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 195675399b..3660326eb6 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -479,7 +479,7 @@ def interpolate_vars(variables: VarsInterpolate, function_space: WithGeometry): ---------- variables: Variables that are to be interpolated into unrestricted/restricted primal/dual spaces. - restricted_space: + function_space: The function space where `PETSc.Vec` objects will live. Returns From 2b4b8990c0f4ef560ea19c35b7ac1783ed267f49 Mon Sep 17 00:00:00 2001 From: Kostya <109463618+Shawniz@users.noreply.github.com> Date: Thu, 4 Sep 2025 19:21:41 +0100 Subject: [PATCH 23/25] Initialise zero primal/dual space objects without interpolation. Co-authored-by: Josh Hope-Collins --- firedrake/rf_eigensolver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 195675399b..da2a7dd7da 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -466,7 +466,7 @@ def new_restricted_control_variable(reduced_functional: ReducedFunctional, funct tuple[OverloadedType]: New variables suitable for storing a control value. """ return tuple( - Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) + Function(function_space.dual() if dual else function_space) for control in reduced_functional.controls ) From ec9a7315b20b52741f36e7e770d8d71d8cf05829 Mon Sep 17 00:00:00 2001 From: KP Date: Thu, 4 Sep 2025 21:29:45 +0300 Subject: [PATCH 24/25] Set operators in init so don't have to be reset in solve. Also init zero more effective. --- firedrake/rf_eigensolver.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/firedrake/rf_eigensolver.py b/firedrake/rf_eigensolver.py index 3660326eb6..b5e9e9d0a6 100644 --- a/firedrake/rf_eigensolver.py +++ b/firedrake/rf_eigensolver.py @@ -384,14 +384,6 @@ def __init__( self.options_manager = OptionsManager(solver_parameters, options_prefix) self.options_manager.set_from_options(self.es) - def solve(self): - r"""Solve the eigenproblem. - - Returns - ------- - int - The number of eigenvalues found. - """ A = self._problem.A M = self._problem.M if not isinstance(A, PETSc.Mat): @@ -401,6 +393,15 @@ def solve(self): self.es.setDimensions(nev=self.n_evals, ncv=self.ncv, mpd=self.mpd) self.es.setOperators(A, M) + + def solve(self): + r"""Solve the eigenproblem. + + Returns + ------- + int + The number of eigenvalues found. + """ with self.options_manager.inserted_options(): self.es.solve() nconv = self.es.getConverged() @@ -466,7 +467,7 @@ def new_restricted_control_variable(reduced_functional: ReducedFunctional, funct tuple[OverloadedType]: New variables suitable for storing a control value. """ return tuple( - Function(function_space).interpolate(control.control)._ad_init_zero(dual=dual) + Function(function_space.dual() if dual else function_space) for control in reduced_functional.controls ) From f63c1de67bf5aa4143c04ba863b21a16e42183a1 Mon Sep 17 00:00:00 2001 From: Kostya <109463618+Shawniz@users.noreply.github.com> Date: Thu, 4 Sep 2025 19:36:52 +0100 Subject: [PATCH 25/25] Using context manager when working with tape. Co-authored-by: Josh Hope-Collins --- tests/firedrake/adjoint/test_rf_eigensolver.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/firedrake/adjoint/test_rf_eigensolver.py b/tests/firedrake/adjoint/test_rf_eigensolver.py index ec93766113..fe4e42df1b 100644 --- a/tests/firedrake/adjoint/test_rf_eigensolver.py +++ b/tests/firedrake/adjoint/test_rf_eigensolver.py @@ -46,14 +46,15 @@ def test_compare_eigensolvers_helmholtz(expression): # ReducedFunctional-based, Jx = 1/cx continue_annotation() - f = Function(V) - f.interpolate(2 * x) - control = Control(f) - u_1 = Function(V) - v = TestFunction(V) - F = (inner(grad(u_1), grad(v)) + u_1 * v - f * v) * dx - solve(F == 0, u_1, bcs=bc) - J = ReducedFunctional(u_1, controls=[control]) + with set_working_tape() as tape: + f = Function(V) + f.interpolate(2 * x) + control = Control(f) + u_1 = Function(V) + v = TestFunction(V) + F = (inner(grad(u_1), grad(v)) + u_1 * v - f * v) * dx + solve(F == 0, u_1, bcs=bc) + J = ReducedFunctional(u_1, controls=[control], tape=tape) pause_annotation() problem_rf = RFEigenproblem(J, bcs=bc, apply_riesz=False, restrict=True, identity=True)