From bfa69da5e5ea12302068ad74f0b0eec8b71f0891 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 10 Nov 2025 11:22:19 +0000 Subject: [PATCH 01/22] White noise generators and AR covariance operators Co-authored-by: Jack Betteridge --- firedrake/adjoint/__init__.py | 2 + firedrake/adjoint/covariance_operator.py | 445 ++++++++++++++++++ .../adjoint/test_covariance_operator.py | 187 ++++++++ 3 files changed, 634 insertions(+) create mode 100644 firedrake/adjoint/covariance_operator.py create mode 100644 tests/firedrake/adjoint/test_covariance_operator.py diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index d3d28e6129..d00ad97193 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -38,6 +38,8 @@ from firedrake.adjoint.ufl_constraints import UFLInequalityConstraint, \ UFLEqualityConstraint # noqa F401 from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 +from firedrake.adjoint.covariance_operator import ( # noqa F401 + NoiseBackend, WhiteNoiseGenerator, GaussianCovarianceOperator, DiffusionFormulation) import numpy_adjoint # noqa F401 import firedrake.ufl_expr import types diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py new file mode 100644 index 0000000000..dfd67b59cc --- /dev/null +++ b/firedrake/adjoint/covariance_operator.py @@ -0,0 +1,445 @@ +from enum import Enum +from functools import cached_property +from textwrap import dedent +from petsctools import get_petscvariables +from loopy import generate_code_v2 +from pyop2 import op2 +from firedrake.tsfc_interface import compile_form +from firedrake import ( + grad, inner, avg, action, outer, replace, + assemble, CellSize, FacetNormal, + dx, ds, dS, sqrt, pi, Constant, + Function, Cofunction, RieszMap, + TrialFunction, TestFunction, + FunctionSpace, VectorFunctionSpace, + BrokenElement, VectorElement, + RandomGenerator, PCG64, + LinearVariationalProblem, + LinearVariationalSolver, + LinearSolver, +) + + +class DiffusionFormulation(Enum): + CG = 'CG' + DG = 'DG' + + +def diffusion_form(u, v, kappa, formulation=DiffusionFormulation.CG): + if formulation == DiffusionFormulation.CG: + return inner(u, v)*dx + inner(kappa*grad(u), grad(v))*dx + + elif formulation == DiffusionFormulation.DG: + mesh = v.function_space().mesh() + n = FacetNormal(mesh) + h = CellSize(mesh) + h_avg = 0.5*(h('+') + h('-')) + alpha_h = Constant(4.0)/h_avg + gamma_h = Constant(8.0)/h + return ( + inner(u, v)*dx + kappa*( + inner(grad(u), grad(v))*dx + - inner(avg(2*outer(u, n)), avg(grad(v)))*dS + - inner(avg(grad(u)), avg(2*outer(v, n)))*dS + + alpha_h*inner(avg(2*outer(u, n)), avg(2*outer(v, n)))*dS + - inner(outer(u, n), grad(v))*ds + - inner(grad(u), outer(v, n))*ds + + gamma_h*inner(u, v)*ds + ) + ) + + else: + raise ValueError("Unknown DiffusionFormulation {formulation}") + + +class CholeskyFactorisation: + def __init__(self, V, form=None): + self._V = V + + if form is None: + self.form = inner(TrialFunction(V), + TestFunction(V))*dx + else: + self.form = form + + self._wrk = Function(V) + + @property + def function_space(self): + return self._V + + @cached_property + def _assemble_action(self): + from firedrake.assemble import get_assembler + return get_assembler(action(self.form, self._wrk)).assemble + + def assemble_action(self, u, tensor=None): + self._wrk.assign(u) + return self._assemble_action(tensor=tensor) + + @cached_property + def solver(self): + return LinearSolver( + assemble(self.form, mat_type='aij'), + solver_parameters={ + "ksp_type": "preonly", + "pc_type": "cholesky", + "pc_factor_mat_ordering_type": "nd"}) + + @cached_property + def pc(self): + return self.solver.ksp.getPC() + + def apply(self, u): + u = self.assemble_action(u) + v = Cofunction(self.space.dual()) + with u.dat.vec_ro as u_v, v.dat.vec_wo as v_v: + self.pc.applySymmetricLeft(u_v, v_v) + return v + + def apply_transpose(self, u): + v = Function(self.function_space) + with u.dat.vec_ro as u_v, v.dat.vec_wo as v_v: + self.pc.applySymmetricRight(u_v, v_v) + v = self.assemble_action(v) + return v + + +class NoiseBackend(Enum): + PYOP2 = 'pyop2' + PETSC = 'petsc' + + +class NoiseBackendBase: + def __init__(self, V, rng=None): + self._V = V + self._rng = rng or RandomGenerator(PCG64()) + + def sample(self, *, rng=None, tensor=None): + raise NotImplementedError + + @cached_property + def broken_space(self): + element = self.function_space.ufl_element() + mesh = self.function_space.mesh().unique() + if isinstance(element, VectorElement): + dim = element.num_sub_elements + scalar_element = element.sub_elements[0] + broken_element = BrokenElement(scalar_element) + Vbroken = VectorFunctionSpace( + mesh, broken_element, dim=dim) + else: + Vbroken = FunctionSpace( + mesh, BrokenElement(element)) + return Vbroken + + @property + def function_space(self): + return self._V + + @property + def rng(self): + return self._rng + + @cached_property + def riesz_map(self): + return RieszMap(self.function_space, constant_jacobian=True) + + +class PetscNoiseBackend(NoiseBackendBase): + def __init__(self, V, rng=None): + super().__init__(V, rng=rng) + self.cholesky = CholeskyFactorisation(self.broken_space) + + def sample(self, *, rng=None, tensor=None, apply_riesz=False): + V = self.function_space + rng = rng or self.rng + + # z + z = rng.standard_normal(self.broken_space) + # C z + Cz = self.cholesky.apply_transpose(z) + # L C z + b = Cofunction(V.dual()).interpolate(Cz) + + if apply_riesz: + b = b.riesz_representation(self.riesz_map) + + if tensor: + tensor.assign(b) + else: + tensor = b + + return tensor + + +class PyOP2NoiseBackend(NoiseBackendBase): + def __init__(self, V, rng=None): + super().__init__(V, rng=rng) + + u = TrialFunction(V) + v = TestFunction(V) + mass = inner(u, v)*dx + + # Create mass expression, assemble and extract kernel + mass_ker, *stuff = compile_form(mass, "mass") + mass_code = generate_code_v2(mass_ker.kinfo.kernel.code).device_code() + mass_code = mass_code.replace( + "void " + mass_ker.kinfo.kernel.name, + "static void " + mass_ker.kinfo.kernel.name) + + # Add custom code for doing Cholesky + # decomposition and applying to broken vector + name = mass_ker.kinfo.kernel.name + blocksize = mass_ker.kinfo.kernel.code[name].args[0].shape[0] + + cholesky_code = dedent( + f"""\ + extern void dpotrf_(char *UPLO, + int *N, + double *A, + int *LDA, + int *INFO); + + extern void dgemv_(char *TRANS, + int *M, + int *N, + double *ALPHA, + double *A, + int *LDA, + double *X, + int *INCX, + double *BETA, + double *Y, + int *INCY); + + {mass_code} + + void apply_cholesky(double *__restrict__ z, + double *__restrict__ b, + double const *__restrict__ coords) + {{ + char uplo[1]; + int32_t N = {blocksize}, LDA = {blocksize}, INFO = 0; + int32_t i=0, j=0; + uplo[0] = 'u'; + double H[{blocksize}*{blocksize}] = {{{{ 0.0 }}}}; + + char trans[1]; + int32_t stride = 1; + double scale = 1.0; + double zero = 0.0; + + {mass_ker.kinfo.kernel.name}(H, coords); + + uplo[0] = 'u'; + dpotrf_(uplo, &N, H, &LDA, &INFO); + for (int i = 0; i < N; i++) + for (int j = 0; j < N; j++) + if (j>i) + H[i*N + j] = 0.0; + + trans[0] = 'T'; + dgemv_(trans, &N, &N, &scale, H, &LDA, z, &stride, &zero, b, &stride); + }} + """ + ) + + # Get the BLAS and LAPACK compiler parameters to compile the kernel + # TODO: Ask CW if this is the right comm to use. + comm = V.mesh()._comm + if comm.rank == 0: + petsc_variables = get_petscvariables() + BLASLAPACK_LIB = petsc_variables.get("BLASLAPACK_LIB", "") + BLASLAPACK_LIB = comm.bcast(BLASLAPACK_LIB, root=0) + BLASLAPACK_INCLUDE = petsc_variables.get("BLASLAPACK_INCLUDE", "") + BLASLAPACK_INCLUDE = comm.bcast(BLASLAPACK_INCLUDE, root=0) + else: + BLASLAPACK_LIB = comm.bcast(None, root=0) + BLASLAPACK_INCLUDE = comm.bcast(None, root=0) + + self.cholesky_kernel = op2.Kernel( + cholesky_code, "apply_cholesky", + include_dirs=BLASLAPACK_INCLUDE.split(), + ldargs=BLASLAPACK_LIB.split()) + + def sample(self, *, rng=None, tensor=None, apply_riesz=False): + rng = rng or self.rng + + z = rng.standard_normal(self.broken_space) + b = Cofunction(self.function_space.dual()) + + z_arg = z.dat(op2.READ, self.broken_space.cell_node_map()) + b_arg = b.dat(op2.INC, self.function_space.cell_node_map()) + + mesh = self.function_space.mesh() + coords = mesh.coordinates + c_arg = coords.dat(op2.READ, coords.cell_node_map()) + + op2.par_loop( + self.cholesky_kernel, + mesh.cell_set, + z_arg, b_arg, c_arg + ) + + if apply_riesz: + b = b.riesz_representation(self.riesz_map) + + if tensor: + tensor.assign(b) + else: + tensor = b + + return tensor + + +class WhiteNoiseGenerator: + r""" Generates a white noise sample + + :arg V: The :class: `firedrake.FunctionSpace` to construct a + white noise sample on + :arg backend: The :enum: `NoiseBackend` specifying how to calculate + and apply the mass matrix square root. + :arg rng: Initialised random number generator to use for obtaining + random numbers. Defaults to PCG64. + + Returns a :firedrake.Function: with + b ~ Normal(0, M) + where b is the dat.data of the function returned + and M is the mass matrix. + + For details see [Croci et al 2018]: + https://epubs.siam.org/doi/10.1137/18M1175239 + """ + + # TODO: Add Croci to citations manager + + def __init__(self, V, backend=NoiseBackend.PYOP2, rng=None): + if backend == NoiseBackend.PYOP2: + self.backend = PyOP2NoiseBackend(V, rng=rng) + elif backend == NoiseBackend.PETSC: + self.backend = PetscNoiseBackend(V, rng=rng) + else: + raise ValueError( + f"Unrecognised white noise generation backend {backend}") + + self.function_space = self.backend.function_space + self.rng = self.backend.rng + + def sample(self, *, rng=None, tensor=None, apply_riesz=False): + return self.backend.sample( + rng=rng, tensor=tensor, apply_riesz=apply_riesz) + + +class GaussianCovarianceOperator: + def __init__(self, V, L, sigma=1, m=2, rng=None, + bcs=None, form=DiffusionFormulation.CG, + solver_parameters=None, options_prefix=None): + + self.rng = rng or WhiteNoiseGenerator(V) + self.function_space = self.rng.function_space + + if sigma <= 0: + raise ValueError("Variance must be positive.") + if L < 0: + raise ValueError("Correlation lengthscale must be positive.") + if m < 0: + raise ValueError("Number of iterations must be positive.") + if (m % 2) != 0: + raise ValueError("Number of iterations must be even.") + + self.stddev = sigma + self.lengthscale = L + self.iterations = m + + if self.iterations > 0: + # Calculate diffusion operator parameters + self.kappa = Constant(L*L/(2*m)) + lambda_g = Constant(sqrt(2*pi)*L) + self.lamda = Constant(sigma*sqrt(lambda_g)) + + # setup diffusion solver + u, v = TrialFunction(V), TestFunction(V) + if isinstance(form, DiffusionFormulation): + a = diffusion_form(u, v, self.kappa, formulation=form) + else: + a = form + self._rhs = Function(V) + rhs = inner(self._rhs, v)*dx + self._u = Function(V) + + self.solver = LinearVariationalSolver( + LinearVariationalProblem(a, rhs, self._u, bcs=bcs, + constant_jacobian=True), + solver_parameters=solver_parameters, + options_prefix=options_prefix) + + # setup mass solver + M = inner(u, v)*dx + rhs = replace(a, {u: self._rhs}) + + self.mass_solver = LinearVariationalSolver( + LinearVariationalProblem(M, rhs, self._u, bcs=bcs, + constant_jacobian=True), + solver_parameters=solver_parameters) + + def sample(self, *, rng=None, tensor=None): + tensor = tensor or Function(self.function_space) + rng = rng or self.rng + w = rng.sample(apply_riesz=True) + + if self.iterations == 0: + return tensor.assign(self.lamda*w) + + self._u.assign(w) + for _ in range(self.iterations//2): + self._rhs.assign(self._u) + self.solver.solve() + + return tensor.assign(self.lamda*self._u) + + def norm(self, x): + + if self.iterations == 0: + sigma_x = self.stddev*x + return assemble(inner(sigma_x, sigma_x)*dx) + + lamda1 = 1/self.lamda + + self._u.assign(lamda1*x) + for k in range(self.iterations//2): + self._rhs.assign(self._u) + self.mass_solver.solve() + + return assemble(inner(self._u, self._u)*dx) + + def apply_inverse(self, x, *, tensor=None): + tensor = tensor or Function(self.function_space) + + if self.iterations == 0: + variance1 = 1/(self.stddev*self.stddev) + return tensor.assign(variance1*x) + + lamda1 = 1/self.lamda + self._u.assign(lamda1*x) + + for k in range(self.iterations): + self._rhs.assign(self._u) + self.mass_solver.solve() + + return tensor.assign(lamda1*self._u) + + def apply_action(self, x, *, tensor=None): + tensor = tensor or Function(self.function_space) + + if self.iterations == 0: + variance = self.stddev*self.stddev + return tensor.assign(variance*x) + + self._u.assign(self.lamda*x) + + for k in range(self.iterations): + self._rhs.assign(self._u) + self.solver.solve() + + return tensor.assign(self.lamda*self._u) diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py new file mode 100644 index 0000000000..aa0ad075b0 --- /dev/null +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -0,0 +1,187 @@ +import pytest +import numpy as np +from scipy.sparse import csr_matrix +from firedrake import * +from firedrake.adjoint import * + + +def petsc2numpy_vec(petsc_vec): + """Allgather a PETSc.Vec.""" + gvec = petsc_vec + gather, lvec = PETSc.Scatter().toAll(gvec) + gather(gvec, lvec, addv=PETSc.InsertMode.INSERT_VALUES) + return lvec.array_r.copy() + + +def petsc2numpy_mat(petsc_mat): + """Allgather a PETSc.Mat.""" + comm = petsc_mat.getComm() + local_mat = petsc_mat.getRedundantMatrix( + comm.size, PETSc.COMM_SELF) + return csr_matrix( + local_mat.getValuesCSR()[::-1], + shape=local_mat.getSize() + ).todense() + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) +@pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vec1", "vec2"]) +@pytest.mark.parametrize("family", ("CG", "DG")) +@pytest.mark.parametrize("mesh_type", ("interval", "square")) +@pytest.mark.parametrize("backend", ("pyop2", "petsc")) +def test_white_noise(family, degree, mesh_type, dim, backend): + """Test that white noise generator converges to a mass matrix covariance. + """ + if backend == "petsc" and COMM_WORLD.size > 1: + pytest.skip( + "petsc backend for noise generation not implemented in parallel.") + + nx = 10 + # Mesh dimension + if mesh_type == 'interval': + mesh = UnitIntervalMesh(nx) + elif mesh_type == 'square': + mesh = UnitSquareMesh(nx, nx) + elif mesh_type == 'cube': + mesh = UnitCubeMesh(nx, nx, nx) + + # Variable rank + if dim > 0: + V = VectorFunctionSpace(mesh, family, degree, dim=dim) + else: + V = FunctionSpace(mesh, family, degree) + + # Finite element white noise has mass matrix covariance + M = inner(TrialFunction(V), TestFunction(V))*dx + covmat = petsc2numpy_mat( + assemble(M, mat_type='aij').petscmat) + + rng = RandomGenerator(PCG64(seed=13)) + + generator = WhiteNoiseGenerator( + V, backend=NoiseBackend(backend), rng=rng) + + # Test convergence as sample size increases + nsamples = [50, 100, 200, 400, 800] + + samples = np.empty((V.dim(), nsamples[-1])) + for i in range(nsamples[-1]): + with generator.sample().dat.vec_ro as bv: + samples[:, i] = petsc2numpy_vec(bv) + + covariances = [np.cov(samples[:, :ns]) for ns in nsamples] + + # Covariance matrix should converge at a rate of sqrt(n) + errors = [np.linalg.norm(cov-covmat) for cov in covariances] + normalised_errors = [err*sqrt(n) for err, n in zip(errors, nsamples)] + normalised_errors /= normalised_errors[-1] + + # Loose tolerance because RNG + tol = 0.2 + assert (1 - tol) < np.max(normalised_errors) < (1 + tol) + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.parametrize("m", (0, 2, 4)) +@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) +@pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vector1", "vector2"]) +@pytest.mark.parametrize("family", ("CG", "DG")) +@pytest.mark.parametrize("mesh_type", ("interval", "square")) +@pytest.mark.parametrize("backend", ("pyop2", "petsc")) +def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): + """Test that correlated noise generator has the right covariance matrix. + """ + if backend == "petsc" and COMM_WORLD.size > 1: + pytest.skip( + "petsc backend for noise generation not implemented in parallel.") + + nx = 16 + if mesh_type == 'interval': + mesh = UnitIntervalMesh(nx) + x, = SpatialCoordinate(mesh) + wexpr = cos(2*pi*x) + elif mesh_type == 'square': + mesh = UnitSquareMesh(nx, nx) + x, y = SpatialCoordinate(mesh) + wexpr = cos(2*pi*x)*cos(4*pi*x) + elif mesh_type == 'cube': + mesh = UnitCubeMesh(nx, nx, nx) + x, y, z = SpatialCoordinate(mesh) + wexpr = cos(2*pi*x)*cos(4*pi*y)*cos(pi*z) + if dim > 0: + V = VectorFunctionSpace(mesh, family, degree, dim=dim) + wexpr = as_vector([-1**(j+1)*wexpr for j in range(dim)]) + else: + V = FunctionSpace(mesh, family, degree) + + rng = WhiteNoiseGenerator( + V, backend=NoiseBackend(backend), + rng=RandomGenerator(PCG64(seed=13))) + + L = 0.1 + sigma = 0.9 + + solver_parameters = { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps' + } + + B = GaussianCovarianceOperator( + V, L, sigma, m, rng=rng, + solver_parameters=solver_parameters, options_prefix="", + form=DiffusionFormulation(family)) + + w = Function(V).project(wexpr) + wcheck = B.apply_inverse(B.apply_action(w)) + + tol = 1e-8 + # Particularly sensitive tests?? + if mesh_type == 'square' and family == 'DG' and degree == 2 and m == 4: + tol = 5e-4 + + assert errornorm(w, wcheck) < tol + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.parametrize("m", (0, 2, 4)) +@pytest.mark.parametrize("family", ("CG", "DG")) +@pytest.mark.parametrize("backend", ("pyop2", "petsc")) +def test_covariance_adjoint_norm(m, family, backend): + """Test that correlated noise generator has the right covariance matrix. + """ + if backend == "petsc" and COMM_WORLD.size > 1: + pytest.skip( + "petsc backend for noise generation not implemented in parallel.") + nx = 20 + L = 0.2 + sigma = 0.1 + + mesh = UnitIntervalMesh(nx) + x, = SpatialCoordinate(mesh) + + V = FunctionSpace(mesh, family, 1) + + u = Function(V).project(sin(2*pi*x)) + v = Function(V).project(2 - 0.5*sin(6*pi*x)) + + B = GaussianCovarianceOperator( + V, L, sigma, m, + form=DiffusionFormulation(family)) + + continue_annotation() + with set_working_tape() as tape: + w = Function(V).project(u**4 + v) + J = B.norm(w) + Jhat = ReducedFunctional(J, Control(u), tape=tape) + pause_annotation() + + m = Function(V).project(sin(2*pi*(x+0.2))) + h = Function(V).project(sin(4*pi*(x-0.2))) + + taylor = taylor_to_dict(Jhat, m, h) + + assert min(taylor['R0']['Rate']) > 0.95, taylor['R0'] + assert min(taylor['R1']['Rate']) > 1.95, taylor['R1'] + assert min(taylor['R2']['Rate']) > 2.95, taylor['R2'] From f8d1eddbac322f5fa2884900bf2f60221aaabd25 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 21 Nov 2025 18:46:26 +0000 Subject: [PATCH 02/22] Type Enums as local members of WhiteNoiseGenerator and GaussianCovariance --- firedrake/adjoint/__init__.py | 2 +- firedrake/adjoint/covariance_operator.py | 95 ++++++++++--------- .../adjoint/test_covariance_operator.py | 86 +++++++++++++---- 3 files changed, 116 insertions(+), 67 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index d00ad97193..f620314d99 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -39,7 +39,7 @@ UFLEqualityConstraint # noqa F401 from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 - NoiseBackend, WhiteNoiseGenerator, GaussianCovarianceOperator, DiffusionFormulation) + WhiteNoiseGenerator, GaussianCovariance) import numpy_adjoint # noqa F401 import firedrake.ufl_expr import types diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index dfd67b59cc..0b2934c0a5 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -20,38 +20,6 @@ ) -class DiffusionFormulation(Enum): - CG = 'CG' - DG = 'DG' - - -def diffusion_form(u, v, kappa, formulation=DiffusionFormulation.CG): - if formulation == DiffusionFormulation.CG: - return inner(u, v)*dx + inner(kappa*grad(u), grad(v))*dx - - elif formulation == DiffusionFormulation.DG: - mesh = v.function_space().mesh() - n = FacetNormal(mesh) - h = CellSize(mesh) - h_avg = 0.5*(h('+') + h('-')) - alpha_h = Constant(4.0)/h_avg - gamma_h = Constant(8.0)/h - return ( - inner(u, v)*dx + kappa*( - inner(grad(u), grad(v))*dx - - inner(avg(2*outer(u, n)), avg(grad(v)))*dS - - inner(avg(grad(u)), avg(2*outer(v, n)))*dS - + alpha_h*inner(avg(2*outer(u, n)), avg(2*outer(v, n)))*dS - - inner(outer(u, n), grad(v))*ds - - inner(grad(u), outer(v, n))*ds - + gamma_h*inner(u, v)*ds - ) - ) - - else: - raise ValueError("Unknown DiffusionFormulation {formulation}") - - class CholeskyFactorisation: def __init__(self, V, form=None): self._V = V @@ -105,11 +73,6 @@ def apply_transpose(self, u): return v -class NoiseBackend(Enum): - PYOP2 = 'pyop2' - PETSC = 'petsc' - - class NoiseBackendBase: def __init__(self, V, rng=None): self._V = V @@ -296,9 +259,9 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): class WhiteNoiseGenerator: r""" Generates a white noise sample - :arg V: The :class: `firedrake.FunctionSpace` to construct a + :arg V: The :class:`firedrake.FunctionSpace` to construct a white noise sample on - :arg backend: The :enum: `NoiseBackend` specifying how to calculate + :arg backend: The :enum:`WhiteNoiseGenerator.Backend` specifying how to calculate and apply the mass matrix square root. :arg rng: Initialised random number generator to use for obtaining random numbers. Defaults to PCG64. @@ -311,13 +274,17 @@ class WhiteNoiseGenerator: For details see [Croci et al 2018]: https://epubs.siam.org/doi/10.1137/18M1175239 """ - # TODO: Add Croci to citations manager - def __init__(self, V, backend=NoiseBackend.PYOP2, rng=None): - if backend == NoiseBackend.PYOP2: + class Backend(Enum): + PYOP2 = 'pyop2' + PETSC = 'petsc' + + def __init__(self, V, backend=None, rng=None): + backend = backend or self.Backend.PYOP2 + if backend == self.Backend.PYOP2: self.backend = PyOP2NoiseBackend(V, rng=rng) - elif backend == NoiseBackend.PETSC: + elif backend == self.Backend.PETSC: self.backend = PetscNoiseBackend(V, rng=rng) else: raise ValueError( @@ -331,11 +298,17 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): rng=rng, tensor=tensor, apply_riesz=apply_riesz) -class GaussianCovarianceOperator: +class GaussianCovariance: + class DiffusionForm(Enum): + CG = 'CG' + IP = 'IP' + def __init__(self, V, L, sigma=1, m=2, rng=None, - bcs=None, form=DiffusionFormulation.CG, + bcs=None, form=None, solver_parameters=None, options_prefix=None): + form = form or self.DiffusionForm.CG + self.rng = rng or WhiteNoiseGenerator(V) self.function_space = self.rng.function_space @@ -360,7 +333,7 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, # setup diffusion solver u, v = TrialFunction(V), TestFunction(V) - if isinstance(form, DiffusionFormulation): + if isinstance(form, self.DiffusionForm): a = diffusion_form(u, v, self.kappa, formulation=form) else: a = form @@ -377,6 +350,8 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, # setup mass solver M = inner(u, v)*dx rhs = replace(a, {u: self._rhs}) + # rhs = a(self._rhs, v) + # rhs = action(a, self._rhs) self.mass_solver = LinearVariationalSolver( LinearVariationalProblem(M, rhs, self._u, bcs=bcs, @@ -399,7 +374,6 @@ def sample(self, *, rng=None, tensor=None): return tensor.assign(self.lamda*self._u) def norm(self, x): - if self.iterations == 0: sigma_x = self.stddev*x return assemble(inner(sigma_x, sigma_x)*dx) @@ -443,3 +417,30 @@ def apply_action(self, x, *, tensor=None): self.solver.solve() return tensor.assign(self.lamda*self._u) + + +def diffusion_form(u, v, kappa, formulation): + if formulation == GaussianCovariance.DiffusionForm.CG: + return inner(u, v)*dx + inner(kappa*grad(u), grad(v))*dx + + elif formulation == GaussianCovariance.DiffusionForm.IP: + mesh = v.function_space().mesh() + n = FacetNormal(mesh) + h = CellSize(mesh) + h_avg = 0.5*(h('+') + h('-')) + alpha_h = Constant(4.0)/h_avg + gamma_h = Constant(8.0)/h + return ( + inner(u, v)*dx + kappa*( + inner(grad(u), grad(v))*dx + - inner(avg(2*outer(u, n)), avg(grad(v)))*dS + - inner(avg(grad(u)), avg(2*outer(v, n)))*dS + + alpha_h*inner(avg(2*outer(u, n)), avg(2*outer(v, n)))*dS + - inner(outer(u, n), grad(v))*ds + - inner(grad(u), outer(v, n))*ds + + gamma_h*inner(u, v)*ds + ) + ) + + else: + raise ValueError("Unknown GaussianCovariance.DiffusionForm {formulation}") diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index aa0ad075b0..011b3ab302 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -60,7 +60,7 @@ def test_white_noise(family, degree, mesh_type, dim, backend): rng = RandomGenerator(PCG64(seed=13)) generator = WhiteNoiseGenerator( - V, backend=NoiseBackend(backend), rng=rng) + V, backend=WhiteNoiseGenerator.Backend(backend), rng=rng) # Test convergence as sample size increases nsamples = [50, 100, 200, 400, 800] @@ -96,17 +96,17 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): pytest.skip( "petsc backend for noise generation not implemented in parallel.") - nx = 16 + nx = 20 if mesh_type == 'interval': - mesh = UnitIntervalMesh(nx) + mesh = PeriodicUnitIntervalMesh(nx) x, = SpatialCoordinate(mesh) wexpr = cos(2*pi*x) elif mesh_type == 'square': - mesh = UnitSquareMesh(nx, nx) + mesh = PeriodicUnitSquareMesh(nx, nx) x, y = SpatialCoordinate(mesh) wexpr = cos(2*pi*x)*cos(4*pi*x) elif mesh_type == 'cube': - mesh = UnitCubeMesh(nx, nx, nx) + mesh = PeriodicUnitCubeMesh(nx, nx, nx) x, y, z = SpatialCoordinate(mesh) wexpr = cos(2*pi*x)*cos(4*pi*y)*cos(pi*z) if dim > 0: @@ -116,7 +116,7 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): V = FunctionSpace(mesh, family, degree) rng = WhiteNoiseGenerator( - V, backend=NoiseBackend(backend), + V, backend=WhiteNoiseGenerator.Backend(backend), rng=RandomGenerator(PCG64(seed=13))) L = 0.1 @@ -128,32 +128,75 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): 'pc_factor_mat_solver_type': 'mumps' } - B = GaussianCovarianceOperator( - V, L, sigma, m, rng=rng, - solver_parameters=solver_parameters, options_prefix="", - form=DiffusionFormulation(family)) + if family == 'CG': + form = GaussianCovariance.DiffusionForm.CG + elif family == 'DG': + form = GaussianCovariance.DiffusionForm.IP + else: + raise ValueError("Do not know which diffusion form to use for family {family}") + + B = GaussianCovariance( + V, L, sigma, m, rng=rng, form=form, + solver_parameters=solver_parameters, + options_prefix="") w = Function(V).project(wexpr) wcheck = B.apply_inverse(B.apply_action(w)) - tol = 1e-8 - # Particularly sensitive tests?? - if mesh_type == 'square' and family == 'DG' and degree == 2 and m == 4: - tol = 5e-4 + tol = 1e-7 assert errornorm(w, wcheck) < tol @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) -@pytest.mark.parametrize("family", ("CG", "DG")) +@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) @pytest.mark.parametrize("backend", ("pyop2", "petsc")) -def test_covariance_adjoint_norm(m, family, backend): +def test_covariance_inverse_action_hdiv(m, degree, backend): """Test that correlated noise generator has the right covariance matrix. """ if backend == "petsc" and COMM_WORLD.size > 1: pytest.skip( "petsc backend for noise generation not implemented in parallel.") + + nx = 20 + mesh = PeriodicUnitSquareMesh(nx, nx) + x, y = SpatialCoordinate(mesh) + wexpr = cos(2*pi*x)*cos(4*pi*x) + + V = FunctionSpace(mesh, "BDM", degree) + wexpr = as_vector([-1**(j+1)*wexpr for j in range(2)]) + + L = 0.1 + sigma = 0.9 + + solver_parameters = { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps' + } + + form = GaussianCovariance.DiffusionForm.IP + + B = GaussianCovariance( + V, L, sigma, m, form=form, + solver_parameters=solver_parameters, + options_prefix="") + + w = Function(V).project(wexpr) + wcheck = B.apply_inverse(B.apply_action(w)) + + tol = 1e-7 + + assert errornorm(w, wcheck) < tol + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.parametrize("m", (0, 2, 4)) +@pytest.mark.parametrize("family", ("CG", "DG")) +def test_covariance_adjoint_norm(m, family): + """Test that correlated noise generator has the right covariance matrix. + """ nx = 20 L = 0.2 sigma = 0.1 @@ -166,9 +209,14 @@ def test_covariance_adjoint_norm(m, family, backend): u = Function(V).project(sin(2*pi*x)) v = Function(V).project(2 - 0.5*sin(6*pi*x)) - B = GaussianCovarianceOperator( - V, L, sigma, m, - form=DiffusionFormulation(family)) + if family == 'CG': + form = GaussianCovariance.DiffusionForm.CG + elif family == 'DG': + form = GaussianCovariance.DiffusionForm.IP + else: + raise ValueError("Do not know which diffusion form to use for family {family}") + + B = GaussianCovariance(V, L, sigma, m, form=form) continue_annotation() with set_working_tape() as tape: From 373b56040a8e1d9a3985d30ae9cefa1c7f8792fd Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 2 Dec 2025 12:28:39 +0000 Subject: [PATCH 03/22] initial covariance mat and pc impls --- firedrake/adjoint/__init__.py | 4 +- firedrake/adjoint/covariance_operator.py | 232 ++++++++++++++++++-- firedrake/adjoint/transformed_functional.py | 80 +++++++ 3 files changed, 300 insertions(+), 16 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index a60acb2438..6cf817328e 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -39,7 +39,9 @@ UFLEqualityConstraint # noqa F401 from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 - WhiteNoiseGenerator, GaussianCovariance) + WhiteNoiseGenerator, GaussianCovariance, + CovarianceOperatorMat, CovarianceOperatorMatrix, + CovarianceOperatorPC) from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 import numpy_adjoint # noqa F401 import firedrake.ufl_expr diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 0b2934c0a5..033295ef08 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -1,14 +1,15 @@ from enum import Enum from functools import cached_property from textwrap import dedent -from petsctools import get_petscvariables +from scipy.special import factorial +from petsctools import get_petscvariables, PCBase from loopy import generate_code_v2 from pyop2 import op2 from firedrake.tsfc_interface import compile_form from firedrake import ( grad, inner, avg, action, outer, replace, assemble, CellSize, FacetNormal, - dx, ds, dS, sqrt, pi, Constant, + dx, ds, dS, sqrt, Constant, Function, Cofunction, RieszMap, TrialFunction, TestFunction, FunctionSpace, VectorFunctionSpace, @@ -17,6 +18,7 @@ LinearVariationalProblem, LinearVariationalSolver, LinearSolver, + PETSc ) @@ -298,19 +300,78 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): rng=rng, tensor=tensor, apply_riesz=apply_riesz) +# Auto-regressive function parameters + +def lengthscale_m(Lar: float, M: int): + """Daley-equivalent lengthscale of M-th order autoregressive function. + + Parameters + ---------- + Lar : + Target Daley correlation lengthscale. + M : + Order of autoregressive function. + + Returns + ------- + L : + Lengthscale parameter for autoregressive function. + """ + return Lar/sqrt(2*M - 3) + + +def lambda_m(Lar: float, M: int): + """Normalisation factor for autoregressive function. + + Parameters + ---------- + Lar : + Target Daley correlation lengthscale. + M : + Order of autoregressive function. + + Returns + ------- + lambda : + Normalisation coefficient for autoregressive correlation operator. + """ + L = lengthscale_m(Lar, M) + num = (2**(2*M - 1))*factorial(M - 1)**2 + den = factorial(2*M - 2) + return L*num/den + + +def kappa_m(Lar: float, M: int): + """Diffusion coefficient for autoregressive function. + + Parameters + ---------- + Lar : + Target Daley correlation lengthscale. + M : + Order of autoregressive function. + + Returns + ------- + kappa : + Diffusion coefficient for autoregressive covariance operator. + """ + return lengthscale_m(Lar, M)**2 + + class GaussianCovariance: class DiffusionForm(Enum): CG = 'CG' IP = 'IP' def __init__(self, V, L, sigma=1, m=2, rng=None, - bcs=None, form=None, + bcs=None, form=None, function_space=None, solver_parameters=None, options_prefix=None): form = form or self.DiffusionForm.CG self.rng = rng or WhiteNoiseGenerator(V) - self.function_space = self.rng.function_space + self.function_space = function_space or self.rng.function_space if sigma <= 0: raise ValueError("Variance must be positive.") @@ -327,9 +388,9 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, if self.iterations > 0: # Calculate diffusion operator parameters - self.kappa = Constant(L*L/(2*m)) - lambda_g = Constant(sqrt(2*pi)*L) - self.lamda = Constant(sigma*sqrt(lambda_g)) + self.kappa = Constant(kappa_m(L, m)) + self.lambda_m = Constant(lambda_m(L, m)) + self._weight = Constant(sigma*sqrt(self.lambda_m)) # setup diffusion solver u, v = TrialFunction(V), TestFunction(V) @@ -350,8 +411,6 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, # setup mass solver M = inner(u, v)*dx rhs = replace(a, {u: self._rhs}) - # rhs = a(self._rhs, v) - # rhs = action(a, self._rhs) self.mass_solver = LinearVariationalSolver( LinearVariationalProblem(M, rhs, self._u, bcs=bcs, @@ -364,21 +423,21 @@ def sample(self, *, rng=None, tensor=None): w = rng.sample(apply_riesz=True) if self.iterations == 0: - return tensor.assign(self.lamda*w) + return tensor.assign(self._weight*w) self._u.assign(w) for _ in range(self.iterations//2): self._rhs.assign(self._u) self.solver.solve() - return tensor.assign(self.lamda*self._u) + return tensor.assign(self._weight*self._u) def norm(self, x): if self.iterations == 0: sigma_x = self.stddev*x return assemble(inner(sigma_x, sigma_x)*dx) - lamda1 = 1/self.lamda + lamda1 = 1/self._weight self._u.assign(lamda1*x) for k in range(self.iterations//2): @@ -394,7 +453,7 @@ def apply_inverse(self, x, *, tensor=None): variance1 = 1/(self.stddev*self.stddev) return tensor.assign(variance1*x) - lamda1 = 1/self.lamda + lamda1 = 1/self._weight self._u.assign(lamda1*x) for k in range(self.iterations): @@ -410,13 +469,13 @@ def apply_action(self, x, *, tensor=None): variance = self.stddev*self.stddev return tensor.assign(variance*x) - self._u.assign(self.lamda*x) + self._u.assign(self._weight*x) for k in range(self.iterations): self._rhs.assign(self._u) self.solver.solve() - return tensor.assign(self.lamda*self._u) + return tensor.assign(self._weight*self._u) def diffusion_form(u, v, kappa, formulation): @@ -444,3 +503,146 @@ def diffusion_form(u, v, kappa, formulation): else: raise ValueError("Unknown GaussianCovariance.DiffusionForm {formulation}") + + +class CovarianceOperatorMat: + class Operation(Enum): + ACTION = 'action' + INVERSE = 'inverse' + + def __init__(self, covariance, operation=None): + operation = operation or self.Operation.ACTION + + V = covariance.function_space + self.function_space = V + self.comm = V.mesh().comm + self.covariance = covariance + self.operation = operation + + primal = Function(V) + dual = Function(V.dual()) + + if operation == self.Operation.ACTION: + self.x = dual + self.y = primal + self._mult_op = covariance.apply_action + elif operation == self.Operation.INVERSE: + self.x = primal + self.y = dual + self._mult_op = covariance.apply_inverse + else: + raise ValueError( + f"Unrecognised CovarianceOperatorMat operation {operation}") + + def mult(self, mat, x, y): + with self.x.dat.vec_wo as v: + x.copy(result=v) + + self._mult_op(self.x, tensor=self.y) + + with self.y.dat.vec_ro as v: + v.copy(result=y) + + def view(self, mat, viewer=None): + if viewer is None: + return + if viewer.getType() != PETSc.Viewer.Type.ASCII: + return + + viewer.printfASCII(f" firedrake covariance operator matrix: {type(self).__name__}\n") + viewer.printfASCII(f" Applying the {str(self.operation)} of the covariance operator {type(self.covariance).__name__}\n") + + if type(self.covariance) is GaussianCovariance: + viewer.printfASCII(" Autoregressive covariance operator with:\n") + viewer.printfASCII(f" order: {self.covariance.iterations}\n") + viewer.printfASCII(f" correlation lengthscale: {self.covariance.lengthscale}\n") + viewer.printfASCII(f" standard deviation: {self.covariance.stddev}\n") + + if self.operation == self.Operation.ACTION: + viewer.printfASCII(" Information for the diffusion solver for applying the action:\n") + self.covariance.solver.snes.ksp.view(viewer) + elif self.operation == self.Operation.INVERSE: + viewer.printfASCII(" Information for the mass solver for applying the inverse:\n") + self.covariance.mass_solver.snes.ksp.view(viewer) + + +def CovarianceOperatorMatrix(covariance, operation=None): + ctx = CovarianceOperatorMat(covariance, operation=operation) + + sizes = covariance.function_space.dof_dset.layout_vec.getSizes() + + mat = PETSc.Mat().createPython( + (sizes, sizes), ctx, comm=ctx.comm) + mat.setUp() + mat.assemble() + return mat + + +class CovarianceOperatorPC(PCBase): + """ + Precondition the inverse covariance operator: + P = B : V* -> V + """ + needs_python_pmat = True + + def initialize(self, pc): + A, P = pc.getOperators() + + use_amat_prefix = self.parent_prefix + "pc_use_amat" + self.use_amat = PETSc.Options().getBool(use_amat_prefix, False) + mat = (A if self.use_amat else P).getPythonContext() + if not isinstance(mat, CovarianceOperatorMat): + raise TypeError( + "CovarianceOperatorPC needs a CovarianceOperatorMat") + covariance = mat.covariance + + self.covariance = covariance + self.mat = mat + + V = covariance.function_space + primal = Function(V) + dual = Function(V.dual()) + + # PC does the opposite of the Mat + if mat.operation == CovarianceOperatorMat.Operation.ACTION: + self.operation = CovarianceOperatorMat.Operation.INVERSE + self.x = primal + self.y = dual + self._apply_op = covariance.apply_inverse + elif mat.operation == self.Operation.INVERSE: + self.operation = CovarianceOperatorMat.Operation.ACTION + self.x = dual + self.y = primal + self._apply_op = covariance.apply_action + + def apply(self, pc, x, y): + with self.x.dat.vec_wo as xvec: + x.copy(result=xvec) + + self._apply_op(self.x, tensor=self.y) + + with self.y.dat.vec_ro as yvec: + yvec.copy(result=y) + + def update(self, pc): + pass + + def view(self, pc, viewer=None): + if viewer is None: + return + if viewer.getType() != PETSc.Viewer.Type.ASCII: + return + + viewer.printfASCII(f" firedrake covariance operator preconditioner: {type(self).__name__}\n") + viewer.printfASCII(f" Applying the {str(self.operation)} of the covariance operator {type(self.covariance).__name__}\n") + + if self.use_amat: + viewer.printfASCII(" using Amat matrix\n") + + if type(self.covariance) is GaussianCovariance: + if self.operation == self.Operation.ACTION: + viewer.printfASCII(" Information for the diffusion solver for applying the action:\n") + self.covariance.solver.snes.ksp.view(viewer) + elif self.operation == self.Operation.INVERSE: + viewer.printfASCII(" Information for the mass solver for applying the inverse:\n") + self.covariance.mass_solver.snes.ksp.view(viewer) diff --git a/firedrake/adjoint/transformed_functional.py b/firedrake/adjoint/transformed_functional.py index d3eeb2f46b..392117bcb9 100644 --- a/firedrake/adjoint/transformed_functional.py +++ b/firedrake/adjoint/transformed_functional.py @@ -3,10 +3,13 @@ from numbers import Real from operator import itemgetter from typing import Optional, Union +from functools import cached_property import firedrake as fd from firedrake.adjoint import Control, ReducedFunctional, Tape from firedrake.functionspaceimpl import WithGeometry +from firedrake.ufl_expr import action +from firedrake.assemble import get_assembler import finat import pyadjoint from pyadjoint import no_annotations @@ -77,6 +80,33 @@ def _pc(self): return pc + @cached_property + def _M_action_assembler(self): + wrk = fd.Function(self.space) + M = fd.inner(fd.TrialFunction(self.space), + fd.TestFunction(self.space))*fd.dx + return wrk, get_assembler(action(M, wrk)).assemble + + def _M_action(self, u: fd.Function, + tensor: Optional[fd.Cofunction] = None) -> fd.Cofunction: + r"""Apply the action of the mass matrix. + + Parameters + ---------- + + u : + The :class:`~firedrake.function.Function` being acted on. + + Returns + ------- + + firedrake.cofunction.Cofunction + The result of the action of the mass matrix on ``u``. + """ + wrk, assembler = self._M_action_assembler + wrk.assign(u) + return assembler(tensor=tensor) + def C_inv_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Cofunction: r"""For the Cholesky factorization @@ -137,6 +167,56 @@ def C_T_inv_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Function: pc.applySymmetricRight(u_v_s, v_v_s) return v + def C_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Cofunction: + r"""For the Cholesky factorization + + ... math : + + M = C C^T, + + compute the action of :math:`C`. + + Parameters + ---------- + + u + Compute :math:`C \tilde{u}` where :math:`\tilde{u}` is the + vector of degrees of freedom for :math:`u`. + + Returns + ------- + + firedrake.cofunction.Cofunction + Has vector of degrees of freedom :math:`C \tilde{u}`. + """ + # ??? + return self.C_T_inv_action(self._M_action(u)) + + def C_T_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Cofunction: + r"""For the Cholesky factorization + + ... math : + + M = C C^T, + + compute the action of :math:`C^{T}`. + + Parameters + ---------- + + u + Compute :math:`C^{T} \tilde{u}` where :math:`\tilde{u}` is the + vector of degrees of freedom for :math:`u`. + + Returns + ------- + + firedrake.function.Function + Has vector of degrees of freedom :math:`C^{T} \tilde{u}`. + """ + # ??? + return self._M_action(self.C_inv_action(u)) + class L2RieszMap(fd.RieszMap): """An :math:`L^2` Riesz map. From 9d93a0e6a3f5e0f0faf83351dc002633283ca754 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Sun, 7 Dec 2025 17:29:57 +0000 Subject: [PATCH 04/22] covariance mat and pc tests --- firedrake/adjoint/__init__.py | 3 +- firedrake/adjoint/covariance_operator.py | 121 ++++++++++------- .../adjoint/test_covariance_operator.py | 125 +++++++++++++++--- 3 files changed, 181 insertions(+), 68 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index 6cf817328e..8ae7ad29c0 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -40,8 +40,7 @@ from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 WhiteNoiseGenerator, GaussianCovariance, - CovarianceOperatorMat, CovarianceOperatorMatrix, - CovarianceOperatorPC) + CovarianceMat, CovarianceMatCtx, CovariancePC) from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 import numpy_adjoint # noqa F401 import firedrake.ufl_expr diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 033295ef08..9666931d27 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -366,7 +366,8 @@ class DiffusionForm(Enum): def __init__(self, V, L, sigma=1, m=2, rng=None, bcs=None, form=None, function_space=None, - solver_parameters=None, options_prefix=None): + solver_parameters=None, options_prefix=None, + mass_parameters=None, mass_prefix=None): form = form or self.DiffusionForm.CG @@ -395,39 +396,45 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, # setup diffusion solver u, v = TrialFunction(V), TestFunction(V) if isinstance(form, self.DiffusionForm): - a = diffusion_form(u, v, self.kappa, formulation=form) + K = diffusion_form(u, v, self.kappa, formulation=form) else: - a = form - self._rhs = Function(V) - rhs = inner(self._rhs, v)*dx + K = form + + M = inner(u, v)*dx + self._u = Function(V) + self._b = Cofunction(V.dual()) + + self._Mrhs = replace(M, {u: self._u}) + self._Krhs = replace(K, {u: self._u}) self.solver = LinearVariationalSolver( - LinearVariationalProblem(a, rhs, self._u, bcs=bcs, + LinearVariationalProblem(K, self._b, self._u, bcs=bcs, constant_jacobian=True), solver_parameters=solver_parameters, options_prefix=options_prefix) - # setup mass solver - M = inner(u, v)*dx - rhs = replace(a, {u: self._rhs}) - self.mass_solver = LinearVariationalSolver( - LinearVariationalProblem(M, rhs, self._u, bcs=bcs, + LinearVariationalProblem(M, self._b, self._u, bcs=bcs, constant_jacobian=True), - solver_parameters=solver_parameters) + solver_parameters=mass_parameters, + options_prefix=mass_prefix) def sample(self, *, rng=None, tensor=None): tensor = tensor or Function(self.function_space) rng = rng or self.rng - w = rng.sample(apply_riesz=True) if self.iterations == 0: - return tensor.assign(self._weight*w) + w = rng.sample(apply_riesz=True) + return tensor.assign(self.stddev*w) + + w = rng.sample(apply_riesz=False) - self._u.assign(w) - for _ in range(self.iterations//2): - self._rhs.assign(self._u) + for i in range(self.iterations//2): + if i == 0: + self._b.assign(w) + else: + assemble(self._Mrhs, tensor=self._b) self.solver.solve() return tensor.assign(self._weight*self._u) @@ -438,41 +445,51 @@ def norm(self, x): return assemble(inner(sigma_x, sigma_x)*dx) lamda1 = 1/self._weight - self._u.assign(lamda1*x) - for k in range(self.iterations//2): - self._rhs.assign(self._u) + + for i in range(self.iterations//2): + assemble(self._Krhs, tensor=self._b) self.mass_solver.solve() return assemble(inner(self._u, self._u)*dx) def apply_inverse(self, x, *, tensor=None): - tensor = tensor or Function(self.function_space) + """B^{-1} : V -> V* + """ + tensor = tensor or Cofunction(self.function_space.dual()) if self.iterations == 0: + riesz_map = self.rng.backend.riesz_map + Cx = x.riesz_representation(riesz_map) variance1 = 1/(self.stddev*self.stddev) - return tensor.assign(variance1*x) + return tensor.assign(variance1*Cx) - lamda1 = 1/self._weight + lamda1 = Constant(1/self._weight) self._u.assign(lamda1*x) - for k in range(self.iterations): - self._rhs.assign(self._u) - self.mass_solver.solve() + for i in range(self.iterations): + assemble(self._Krhs, tensor=self._b) + if i != self.iterations - 1: + self.mass_solver.solve() - return tensor.assign(lamda1*self._u) + return tensor.assign(lamda1*self._b) def apply_action(self, x, *, tensor=None): + """B : V* -> V + """ tensor = tensor or Function(self.function_space) if self.iterations == 0: + riesz_map = self.rng.backend.riesz_map + Cx = x.riesz_representation(riesz_map) variance = self.stddev*self.stddev - return tensor.assign(variance*x) + return tensor.assign(variance*Cx) - self._u.assign(self._weight*x) - - for k in range(self.iterations): - self._rhs.assign(self._u) + for i in range(self.iterations): + if i == 0: + self._b.assign(self._weight*x) + else: + assemble(self._Mrhs, tensor=self._b) self.solver.solve() return tensor.assign(self._weight*self._u) @@ -505,7 +522,7 @@ def diffusion_form(u, v, kappa, formulation): raise ValueError("Unknown GaussianCovariance.DiffusionForm {formulation}") -class CovarianceOperatorMat: +class CovarianceMatCtx: class Operation(Enum): ACTION = 'action' INVERSE = 'inverse' @@ -532,7 +549,7 @@ def __init__(self, covariance, operation=None): self._mult_op = covariance.apply_inverse else: raise ValueError( - f"Unrecognised CovarianceOperatorMat operation {operation}") + f"Unrecognised CovarianceMat operation {operation}") def mult(self, mat, x, y): with self.x.dat.vec_wo as v: @@ -552,7 +569,7 @@ def view(self, mat, viewer=None): viewer.printfASCII(f" firedrake covariance operator matrix: {type(self).__name__}\n") viewer.printfASCII(f" Applying the {str(self.operation)} of the covariance operator {type(self.covariance).__name__}\n") - if type(self.covariance) is GaussianCovariance: + if (type(self.covariance) is GaussianCovariance) and (self.covariance.iterations > 0): viewer.printfASCII(" Autoregressive covariance operator with:\n") viewer.printfASCII(f" order: {self.covariance.iterations}\n") viewer.printfASCII(f" correlation lengthscale: {self.covariance.lengthscale}\n") @@ -560,14 +577,18 @@ def view(self, mat, viewer=None): if self.operation == self.Operation.ACTION: viewer.printfASCII(" Information for the diffusion solver for applying the action:\n") - self.covariance.solver.snes.ksp.view(viewer) + ksp = self.covariance.solver.snes.ksp elif self.operation == self.Operation.INVERSE: viewer.printfASCII(" Information for the mass solver for applying the inverse:\n") - self.covariance.mass_solver.snes.ksp.view(viewer) + ksp = self.covariance.mass_solver.snes.ksp + level = ksp.getTabLevel() + ksp.setTabLevel(mat.getTabLevel() + 1) + ksp.view(viewer) + ksp.setTabLevel(level) -def CovarianceOperatorMatrix(covariance, operation=None): - ctx = CovarianceOperatorMat(covariance, operation=operation) +def CovarianceMat(covariance, operation=None): + ctx = CovarianceMatCtx(covariance, operation=operation) sizes = covariance.function_space.dof_dset.layout_vec.getSizes() @@ -578,12 +599,13 @@ def CovarianceOperatorMatrix(covariance, operation=None): return mat -class CovarianceOperatorPC(PCBase): +class CovariancePC(PCBase): """ Precondition the inverse covariance operator: P = B : V* -> V """ needs_python_pmat = True + prefix = "covariance" def initialize(self, pc): A, P = pc.getOperators() @@ -591,9 +613,10 @@ def initialize(self, pc): use_amat_prefix = self.parent_prefix + "pc_use_amat" self.use_amat = PETSc.Options().getBool(use_amat_prefix, False) mat = (A if self.use_amat else P).getPythonContext() - if not isinstance(mat, CovarianceOperatorMat): + + if not isinstance(mat, CovarianceMatCtx): raise TypeError( - "CovarianceOperatorPC needs a CovarianceOperatorMat") + "CovariancePC needs a CovarianceMatCtx") covariance = mat.covariance self.covariance = covariance @@ -604,13 +627,13 @@ def initialize(self, pc): dual = Function(V.dual()) # PC does the opposite of the Mat - if mat.operation == CovarianceOperatorMat.Operation.ACTION: - self.operation = CovarianceOperatorMat.Operation.INVERSE + if mat.operation == CovarianceMatCtx.Operation.ACTION: + self.operation = CovarianceMatCtx.Operation.INVERSE self.x = primal self.y = dual self._apply_op = covariance.apply_inverse - elif mat.operation == self.Operation.INVERSE: - self.operation = CovarianceOperatorMat.Operation.ACTION + elif mat.operation == CovarianceMatCtx.Operation.INVERSE: + self.operation = CovarianceMatCtx.Operation.ACTION self.x = dual self.y = primal self._apply_op = covariance.apply_action @@ -639,10 +662,10 @@ def view(self, pc, viewer=None): if self.use_amat: viewer.printfASCII(" using Amat matrix\n") - if type(self.covariance) is GaussianCovariance: - if self.operation == self.Operation.ACTION: + if (type(self.covariance) is GaussianCovariance) and (self.covariance.iterations > 0): + if self.operation == CovarianceMatCtx.Operation.ACTION: viewer.printfASCII(" Information for the diffusion solver for applying the action:\n") self.covariance.solver.snes.ksp.view(viewer) - elif self.operation == self.Operation.INVERSE: + elif self.operation == CovarianceMatCtx.Operation.INVERSE: viewer.printfASCII(" Information for the mass solver for applying the inverse:\n") self.covariance.mass_solver.snes.ksp.view(viewer) diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index 011b3ab302..34ca3f05c5 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -1,6 +1,7 @@ import pytest import numpy as np from scipy.sparse import csr_matrix +import petsctools from firedrake import * from firedrake.adjoint import * @@ -88,13 +89,9 @@ def test_white_noise(family, degree, mesh_type, dim, backend): @pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vector1", "vector2"]) @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) -@pytest.mark.parametrize("backend", ("pyop2", "petsc")) -def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): +def test_covariance_inverse_action(m, family, degree, mesh_type, dim): """Test that correlated noise generator has the right covariance matrix. """ - if backend == "petsc" and COMM_WORLD.size > 1: - pytest.skip( - "petsc backend for noise generation not implemented in parallel.") nx = 20 if mesh_type == 'interval': @@ -104,7 +101,7 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): elif mesh_type == 'square': mesh = PeriodicUnitSquareMesh(nx, nx) x, y = SpatialCoordinate(mesh) - wexpr = cos(2*pi*x)*cos(4*pi*x) + wexpr = cos(2*pi*x)*cos(4*pi*y) elif mesh_type == 'cube': mesh = PeriodicUnitCubeMesh(nx, nx, nx) x, y, z = SpatialCoordinate(mesh) @@ -116,8 +113,7 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): V = FunctionSpace(mesh, family, degree) rng = WhiteNoiseGenerator( - V, backend=WhiteNoiseGenerator.Backend(backend), - rng=RandomGenerator(PCG64(seed=13))) + V, rng=RandomGenerator(PCG64(seed=13))) L = 0.1 sigma = 0.9 @@ -141,9 +137,9 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): options_prefix="") w = Function(V).project(wexpr) - wcheck = B.apply_inverse(B.apply_action(w)) + wcheck = B.apply_action(B.apply_inverse(w)) - tol = 1e-7 + tol = 1e-10 assert errornorm(w, wcheck) < tol @@ -151,13 +147,9 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim, backend): @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) -@pytest.mark.parametrize("backend", ("pyop2", "petsc")) -def test_covariance_inverse_action_hdiv(m, degree, backend): +def test_covariance_inverse_action_hdiv(m, degree): """Test that correlated noise generator has the right covariance matrix. """ - if backend == "petsc" and COMM_WORLD.size > 1: - pytest.skip( - "petsc backend for noise generation not implemented in parallel.") nx = 20 mesh = PeriodicUnitSquareMesh(nx, nx) @@ -184,9 +176,9 @@ def test_covariance_inverse_action_hdiv(m, degree, backend): options_prefix="") w = Function(V).project(wexpr) - wcheck = B.apply_inverse(B.apply_action(w)) + wcheck = B.apply_action(B.apply_inverse(w)) - tol = 1e-7 + tol = 1e-8 assert errornorm(w, wcheck) < tol @@ -233,3 +225,102 @@ def test_covariance_adjoint_norm(m, family): assert min(taylor['R0']['Rate']) > 0.95, taylor['R0'] assert min(taylor['R1']['Rate']) > 1.95, taylor['R1'] assert min(taylor['R2']['Rate']) > 2.95, taylor['R2'] + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.parametrize("m", (0, 2, 4)) +@pytest.mark.parametrize("family", ("CG", "DG")) +@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) +@pytest.mark.parametrize("operation", ("action", "inverse")) +def test_covariance_mat(m, family, degree, operation): + """Test that correlated noise generator has the right covariance matrix. + """ + nx = 20 + L = 0.2 + sigma = 0.9 + + mesh = UnitIntervalMesh(nx) + coords, = SpatialCoordinate(mesh) + + V = FunctionSpace(mesh, family, degree) + + if family == 'CG': + form = GaussianCovariance.DiffusionForm.CG + elif family == 'DG': + form = GaussianCovariance.DiffusionForm.IP + else: + raise ValueError("Do not know which diffusion form to use for family {family}") + + B = GaussianCovariance(V, L, sigma, m, form=form) + + operation = CovarianceMatCtx.Operation(operation) + + mat = CovarianceMat(B, operation=operation) + + expr = 2*pi*coords + + if operation == CovarianceMatCtx.Operation.ACTION: + x = Function(V).project(expr).riesz_representation() + y = Function(V) + xcheck = x.copy(deepcopy=True) + ycheck = y.copy(deepcopy=True) + + B.apply_action(xcheck, tensor=ycheck) + + elif operation == CovarianceMatCtx.Operation.INVERSE: + x = Function(V).project(expr) + y = Function(V.dual()) + xcheck = x.copy(deepcopy=True) + ycheck = y.copy(deepcopy=True) + + B.apply_inverse(xcheck, tensor=ycheck) + + with x.dat.vec as xv, y.dat.vec as yv: + mat.mult(xv, yv) + + # flip to primal space to calculate norms + if operation == CovarianceMatCtx.Operation.INVERSE: + y = y.riesz_representation() + ycheck = ycheck.riesz_representation() + + assert errornorm(ycheck, y)/norm(ycheck) < 1e-12 + + if operation == CovarianceMatCtx.Operation.INVERSE: + y = y.riesz_representation() + ycheck = ycheck.riesz_representation() + + ksp = PETSc.KSP().create() + ksp.setOperators(mat) + + # poorly conditioned cases + if (degree == 2) and (m == 4): + tol = 1e-6 + else: + tol = 1e-8 + + petsctools.set_from_options( + ksp, options_prefix="action", + parameters={ + 'ksp_monitor': None, + 'ksp_type': 'richardson', + 'ksp_max_it': 3, + 'ksp_rtol': tol, + 'pc_type': 'python', + 'pc_python_type': 'firedrake.adjoint.CovariancePC', + } + ) + x.zero() + + with x.dat.vec as xv, y.dat.vec as yv: + with petsctools.inserted_options(ksp): + ksp.solve(yv, xv) + + # CovarianceOperator operations should + # be exact inverses of each other. + assert ksp.its == 1 + + if operation == CovarianceMatCtx.Operation.ACTION: + x = x.riesz_representation() + xcheck = xcheck.riesz_representation() + + assert errornorm(xcheck, x)/norm(xcheck) < tol From 6dad6ded9aa553bbb0c2a349a9a9e0622233f495 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 12:14:55 +0000 Subject: [PATCH 05/22] CovarianceOperatorBase and covariance docstrings --- firedrake/adjoint/__init__.py | 4 +- firedrake/adjoint/covariance_operator.py | 462 +++++++++++++++--- .../adjoint/test_covariance_operator.py | 9 +- 3 files changed, 410 insertions(+), 65 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index 8ae7ad29c0..21f9a72681 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -39,8 +39,8 @@ UFLEqualityConstraint # noqa F401 from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 - WhiteNoiseGenerator, GaussianCovariance, - CovarianceMat, CovarianceMatCtx, CovariancePC) + WhiteNoiseGenerator, AutoregressiveCovariance, + CovarianceMat, CovariancePC) from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 import numpy_adjoint # noqa F401 import firedrake.ufl_expr diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 9666931d27..95f25d3bd1 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -1,3 +1,4 @@ +import abc from enum import Enum from functools import cached_property from textwrap import dedent @@ -261,17 +262,24 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): class WhiteNoiseGenerator: r""" Generates a white noise sample - :arg V: The :class:`firedrake.FunctionSpace` to construct a - white noise sample on - :arg backend: The :enum:`WhiteNoiseGenerator.Backend` specifying how to calculate + Parameters + ---------- + V : + The :class:`~firedrake.functionspace.FunctionSpace` to construct a + white noise sample on. + backend : + The ``WhiteNoiseGenerator.Backend`` specifying how to calculate and apply the mass matrix square root. - :arg rng: Initialised random number generator to use for obtaining + rng : + Initialised random number generator to use for obtaining random numbers. Defaults to PCG64. - Returns a :firedrake.Function: with - b ~ Normal(0, M) - where b is the dat.data of the function returned - and M is the mass matrix. + Returns + ------- + firedrake.function.Function : + with b ~ Normal(0, M) where b is the dat.data of the + :class:`~.firedrake.function.Function` returned and + M is the mass matrix. For details see [Croci et al 2018]: https://epubs.siam.org/doi/10.1137/18M1175239 @@ -279,6 +287,14 @@ class WhiteNoiseGenerator: # TODO: Add Croci to citations manager class Backend(Enum): + """ + The backend to implement applying the mass matrix square root. + + See Also + -------- + PyOP2NoiseBackend + PetscNoiseBackend + """ PYOP2 = 'pyop2' PETSC = 'petsc' @@ -302,65 +318,255 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): # Auto-regressive function parameters -def lengthscale_m(Lar: float, M: int): - """Daley-equivalent lengthscale of M-th order autoregressive function. +def lengthscale_m(Lar: float, m: int): + """Daley-equivalent lengthscale of m-th order autoregressive function. Parameters ---------- - Lar : - Target Daley correlation lengthscale. - M : - Order of autoregressive function. + Lar : + Target Daley correlation lengthscale. + m : + Order of autoregressive function. Returns ------- - L : - Lengthscale parameter for autoregressive function. + L : float + Lengthscale parameter for autoregressive function. """ - return Lar/sqrt(2*M - 3) + return Lar/sqrt(2*m - 3) -def lambda_m(Lar: float, M: int): +def lambda_m(Lar: float, m: int): """Normalisation factor for autoregressive function. Parameters ---------- - Lar : - Target Daley correlation lengthscale. - M : - Order of autoregressive function. + Lar : + Target Daley correlation lengthscale. + m : + Order of autoregressive function. Returns ------- - lambda : - Normalisation coefficient for autoregressive correlation operator. + lambda : float + Normalisation coefficient for autoregressive correlation operator. """ - L = lengthscale_m(Lar, M) - num = (2**(2*M - 1))*factorial(M - 1)**2 - den = factorial(2*M - 2) + L = lengthscale_m(Lar, m) + num = (2**(2*m - 1))*factorial(m - 1)**2 + den = factorial(2*m - 2) return L*num/den -def kappa_m(Lar: float, M: int): +def kappa_m(Lar: float, m: int): """Diffusion coefficient for autoregressive function. Parameters ---------- - Lar : - Target Daley correlation lengthscale. - M : - Order of autoregressive function. + Lar : + Target Daley correlation lengthscale. + m : + Order of autoregressive function. Returns ------- - kappa : - Diffusion coefficient for autoregressive covariance operator. + kappa : float + Diffusion coefficient for autoregressive covariance operator. """ - return lengthscale_m(Lar, M)**2 + return lengthscale_m(Lar, m)**2 + + +class CovarianceOperatorBase: + """ + Abstract base class for a covariance operator B where + + .. math:: + + B: V^{*} \\to V \quad \\text{and} \quad B^{-1}: V \\to V^{*} + + The covariance operators can be used to: + + - calculate weighted norms :math:`\|x\|_{B^{-1}} = x^{T}B^{-1}x` + to account for uncertainty in optimisation methods. + + - generate samples from the normal distribution :math:`\mathcal{N}(0, B)` + using :math:`w = B^{1/2}z` where :math:`z\sim\mathcal{N}(0, I)`. + + Inheriting classes must implement the following methods: + + - ``sample`` + + - ``apply_inverse`` + + - ``apply_action`` + + - ``rng`` + + - ``function_space`` + + They may optionally implement ``norm`` to provide a more + efficient implementation. + + See Also + -------- + WhiteNoiseGenerator + AutoregressiveCovariance + CovarianceMatCtx + CovarianceMat + CovariancePC + """ # noqa: W605 + + @abc.abstractmethod + def rng(self): + """:class:`~.WhiteNoiseGenerator` for generating samples. + """ + raise NotImplementedError + + @abc.abstractmethod + def function_space(self): + """The function space V that the covariance operator maps to. + """ + raise NotImplementedError + + @abc.abstractmethod + def sample(self, *, rng: WhiteNoiseGenerator | None = None, + tensor: Function | None = None): + """ + Sample from :math:`\mathcal{N}(0, B)` by correlating a + white noise sample: :math:`w = B^{1/2}z`. + + Parameters + ---------- + rng : + Generator for the white noise sample. + If not provided then self.rng will be used. + tensor : + Optional location to place the result into. + + Returns + ------- + firedrake.function.Function : + The sample. + """ # noqa: W605 + raise NotImplementedError + + def norm(self, x: Function): + """Return the weighted norm :math:`\|x\|_{B^{-1}} = x^{T}B^{-1}x`. + + Default implementation uses ``apply_inverse`` to first calculate + the :class:`~firedrake.cofunction.Cofunction` :math:`y = B^{-1}x`, + then returns :math:`y(x)`. + + Inheriting classes may provide more efficient specialisations. + + Parameters + ---------- + x : + The :class:`~firedrake.function.Function` to take the norm of. + + Returns + ------- + pyadjoint.AdjFloat : + The norm of ``x``. + """ + return self.apply_inverse(x)(x) + + @abc.abstractmethod + def apply_inverse(self, x: Function, *, + tensor: Cofunction | None = None): + """Return :math:`y = B^{-1}x` where B is the covariance operator. + :math:`B^{-1}: V \\to V^{*}`. + + Parameters + ---------- + x : + The :class:`~firedrake.function.Function` to apply the inverse to. + tensor : + Optional location to place the result into. + + Returns + ------- + firedrake.cofunction.Cofunction : + The result of :math:`B^{-1}x` + """ # noqa: W605 + raise NotImplementedError + + @abc.abstractmethod + def apply_action(self, x: Cofunction, *, + tensor: Function | None = None): + """Return :math:`y = Bx` where B is the covariance operator. + :math:`B: V^{*} \\to V`. + + Parameters + ---------- + x : + The :class:`~firedrake.cofunction.Cofunction` to apply + the action to. + tensor : + Optional location to place the result into. + + Returns + ------- + firedrake.function.Function : + The result of :math:`B^{-1}x` + """ # noqa: W605 + raise NotImplementedError + + +class AutoregressiveCovariance(CovarianceOperatorBase): + """ + An m-th order autoregressive covariance operator using an implicit diffusion operator. + + Covariance operator B with a kernel that is the ``m``-th autoregressive + function can be calculated using ``m`` Backward Euler steps of a + diffusion operator, where the diffusion coefficient is specified by + the desired correlation lengthscale. + + If :math:`M` is the mass matrix, :math:`K` is the matrix for a single + Backward Euler step, and :math:`\lambda` is a normalisation factor, then the + m-th order correlation operator (unit variance) is: + + .. math:: + + B: V^{*} \\to V = \lambda((K^{-1}M)^{m}M^{-1})\lambda + + B^{-1}: V \\to V^{*} = (1/\lambda)M(M^{-1}K)^{m}(1/\lambda) + + This formulation leads to an efficient implementations for :math:`B^{1/2}` + by taking only m/2 steps of the diffusion operator. This can be used + to calculate weighted norms and sample from :math:`\mathcal{N}(0,B)`. + + .. math:: + + \|x\|_{B^{-1}} = \|(M^{-1}K)^{m/2}(1/\lambda)x\|_{M} + + w = B^{1/2}z = \lambda M^{-1}(MK^{-1})^{m/2}(M^{1/2}z) + The white noise sample :math:`M^{1/2}z` is generated by a + :class:`.WhiteNoiseGenerator`. + + References + ---------- + Mirouze, I. and Weaver, A. T., 2010: "Representation of correlation + functions in variational assimilation using an implicit diffusion + operator". Q. J. R. Meteorol. Soc. 136: 1421–1443, July 2010 Part B, + https://doi.org/10.1002/qj.643 + + See Also + -------- + WhiteNoiseGenerator + CovarianceOperatorBase + CovarianceMat + CovariancePC + """ # noqa: W605 -class GaussianCovariance: class DiffusionForm(Enum): + """ + The diffusion operator formulation. + + See Also + -------- + diffusion_form + """ CG = 'CG' IP = 'IP' @@ -371,8 +577,8 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, form = form or self.DiffusionForm.CG - self.rng = rng or WhiteNoiseGenerator(V) - self.function_space = function_space or self.rng.function_space + self._rng = rng or WhiteNoiseGenerator(V) + self._function_space = function_space or self.rng.function_space if sigma <= 0: raise ValueError("Variance must be positive.") @@ -420,9 +626,15 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, solver_parameters=mass_parameters, options_prefix=mass_prefix) + def function_space(self): + return self._function_space + + def rng(self): + return self._rng + def sample(self, *, rng=None, tensor=None): - tensor = tensor or Function(self.function_space) - rng = rng or self.rng + tensor = tensor or Function(self.function_space()) + rng = rng or self.rng() if self.iterations == 0: w = rng.sample(apply_riesz=True) @@ -454,12 +666,10 @@ def norm(self, x): return assemble(inner(self._u, self._u)*dx) def apply_inverse(self, x, *, tensor=None): - """B^{-1} : V -> V* - """ - tensor = tensor or Cofunction(self.function_space.dual()) + tensor = tensor or Cofunction(self.function_space().dual()) if self.iterations == 0: - riesz_map = self.rng.backend.riesz_map + riesz_map = self.rng().backend.riesz_map Cx = x.riesz_representation(riesz_map) variance1 = 1/(self.stddev*self.stddev) return tensor.assign(variance1*Cx) @@ -475,12 +685,10 @@ def apply_inverse(self, x, *, tensor=None): return tensor.assign(lamda1*self._b) def apply_action(self, x, *, tensor=None): - """B : V* -> V - """ - tensor = tensor or Function(self.function_space) + tensor = tensor or Function(self.function_space()) if self.iterations == 0: - riesz_map = self.rng.backend.riesz_map + riesz_map = self.rng().backend.riesz_map Cx = x.riesz_representation(riesz_map) variance = self.stddev*self.stddev return tensor.assign(variance*Cx) @@ -495,11 +703,47 @@ def apply_action(self, x, *, tensor=None): return tensor.assign(self._weight*self._u) -def diffusion_form(u, v, kappa, formulation): - if formulation == GaussianCovariance.DiffusionForm.CG: +def diffusion_form(u, v, kappa: Constant | Function, + formulation: AutoregressiveCovariance.DiffusionForm): + """ + Convenience function for common diffusion forms. + + Currently provides: + + - Standard continuous Galerkin form. + + - Interior penalty method for discontinuous spaces. + + + Parameters + ---------- + u : + :func:`~firedrake.ufl_expr.TrialFunction` to construct diffusion form with. + v : + :func:`~firedrake.ufl_expr.TestFunction` to construct diffusion form with. + kappa : + The diffusion coefficient. + formulation : + The type of diffusion form. + + Returns + ------- + ufl.Form : + The diffusion form over u and v. + + Raises + ------ + ValueError + Unrecognised formulation. + + See Also + -------- + AutoregressiveCovariance + """ + if formulation == AutoregressiveCovariance.DiffusionForm.CG: return inner(u, v)*dx + inner(kappa*grad(u), grad(v))*dx - elif formulation == GaussianCovariance.DiffusionForm.IP: + elif formulation == AutoregressiveCovariance.DiffusionForm.IP: mesh = v.function_space().mesh() n = FacetNormal(mesh) h = CellSize(mesh) @@ -519,18 +763,42 @@ def diffusion_form(u, v, kappa, formulation): ) else: - raise ValueError("Unknown GaussianCovariance.DiffusionForm {formulation}") + raise ValueError("Unknown AutoregressiveCovariance.DiffusionForm {formulation}") class CovarianceMatCtx: + """ + A python Mat context for a covariance operator. + Can apply either the action or inverse of the covariance. + + .. math:: + + B: V^{*} \\to V + + B^{-1}: V \\to V^{*} + + Parameters + ---------- + covariance : + The covariance operator. + operation : CovarianceMatCtx.Operation + Whether the matrix applies the action or inverse of the covariance operator. + + See Also + -------- + CovarianceOperatorBase + AutoregressiveCovariance + CovarianceMat + CovariancePC + """ # noqa: W605 class Operation(Enum): ACTION = 'action' INVERSE = 'inverse' - def __init__(self, covariance, operation=None): + def __init__(self, covariance: CovarianceOperatorBase, operation=None): operation = operation or self.Operation.ACTION - V = covariance.function_space + V = covariance.function_space() self.function_space = V self.comm = V.mesh().comm self.covariance = covariance @@ -552,6 +820,20 @@ def __init__(self, covariance, operation=None): f"Unrecognised CovarianceMat operation {operation}") def mult(self, mat, x, y): + """Apply the action or inverse of the covariance operator + to x, putting the result in y. + + y is not guaranteed to be zero on entry. + + Parameters + ---------- + A : PETSc.Mat + The PETSc matrix that self is the python context of. + x : PETSc.Vec + The vector acted on by the matrix. + y : PETSc.Vec + The result of the matrix action. + """ with self.x.dat.vec_wo as v: x.copy(result=v) @@ -569,7 +851,7 @@ def view(self, mat, viewer=None): viewer.printfASCII(f" firedrake covariance operator matrix: {type(self).__name__}\n") viewer.printfASCII(f" Applying the {str(self.operation)} of the covariance operator {type(self.covariance).__name__}\n") - if (type(self.covariance) is GaussianCovariance) and (self.covariance.iterations > 0): + if (type(self.covariance) is AutoregressiveCovariance) and (self.covariance.iterations > 0): viewer.printfASCII(" Autoregressive covariance operator with:\n") viewer.printfASCII(f" order: {self.covariance.iterations}\n") viewer.printfASCII(f" correlation lengthscale: {self.covariance.lengthscale}\n") @@ -588,6 +870,36 @@ def view(self, mat, viewer=None): def CovarianceMat(covariance, operation=None): + """ + A Mat for a covariance operator. + Can apply either the action or inverse of the covariance. + This is a convenience function to create a PETSc.Mat with a :class:`.CovarianceMatCtx` Python context. + + .. math:: + + B: V^{*} \\to V + + B^{-1}: V \\to V^{*} + + Parameters + ---------- + covariance : + The covariance operator. + operation : CovarianceMatCtx.Operation + Whether the matrix applies the action or inverse of the covariance operator. + + Returns + ------- + PETSc.Mat : + The python type Mat with a CovarianceMatCtx context. + + See Also + -------- + CovarianceOperatorBase + AutoregressiveCovariance + CovarianceMatCtx + CovariancePC + """ # noqa: W605 ctx = CovarianceMatCtx(covariance, operation=operation) sizes = covariance.function_space.dof_dset.layout_vec.getSizes() @@ -601,9 +913,27 @@ def CovarianceMat(covariance, operation=None): class CovariancePC(PCBase): """ - Precondition the inverse covariance operator: - P = B : V* -> V - """ + A python PC context for a covariance operator. + Will apply either the action or inverse of the covariance, + whichever is the opposite of the Mat operator. + + .. math:: + + B: V^{*} \\to V + + B^{-1}: V \\to V^{*} + + Available options: + + * ``-pc_use_amat`` - use Amat to apply the covariance operator. + + See Also + -------- + CovarianceOperatorBase + AutoregressiveCovariance + CovarianceMatCtx + CovarianceMat + """ # noqa: W605 needs_python_pmat = True prefix = "covariance" @@ -639,6 +969,20 @@ def initialize(self, pc): self._apply_op = covariance.apply_action def apply(self, pc, x, y): + """Apply the action or inverse of the covariance operator + to x, putting the result in y. + + y is not guaranteed to be zero on entry. + + Parameters + ---------- + pc : PETSc.PC + The PETSc preconditioner that self is the python context of. + x : PETSc.Vec + The vector acted on by the pc. + y : PETSc.Vec + The result of the pc application. + """ with self.x.dat.vec_wo as xvec: x.copy(result=xvec) @@ -662,7 +1006,7 @@ def view(self, pc, viewer=None): if self.use_amat: viewer.printfASCII(" using Amat matrix\n") - if (type(self.covariance) is GaussianCovariance) and (self.covariance.iterations > 0): + if (type(self.covariance) is AutoregressiveCovariance) and (self.covariance.iterations > 0): if self.operation == CovarianceMatCtx.Operation.ACTION: viewer.printfASCII(" Information for the diffusion solver for applying the action:\n") self.covariance.solver.snes.ksp.view(viewer) diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index 34ca3f05c5..77f3e29583 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -90,7 +90,7 @@ def test_white_noise(family, degree, mesh_type, dim, backend): @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) def test_covariance_inverse_action(m, family, degree, mesh_type, dim): - """Test that correlated noise generator has the right covariance matrix. + """Test that covariance operator action and inverse are opposites. """ nx = 20 @@ -148,7 +148,8 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim): @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) def test_covariance_inverse_action_hdiv(m, degree): - """Test that correlated noise generator has the right covariance matrix. + """Test that covariance operator action and inverse are opposites + for hdiv spaces. """ nx = 20 @@ -187,7 +188,7 @@ def test_covariance_inverse_action_hdiv(m, degree): @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("family", ("CG", "DG")) def test_covariance_adjoint_norm(m, family): - """Test that correlated noise generator has the right covariance matrix. + """Test that covariance operators are properly taped. """ nx = 20 L = 0.2 @@ -233,7 +234,7 @@ def test_covariance_adjoint_norm(m, family): @pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) @pytest.mark.parametrize("operation", ("action", "inverse")) def test_covariance_mat(m, family, degree, operation): - """Test that correlated noise generator has the right covariance matrix. + """Test that covariance mat and pc apply correct and opposite actions. """ nx = 20 L = 0.2 From 4a936239db3ff7f4b8312f1e221098fa4d84e6af Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 12:22:27 +0000 Subject: [PATCH 06/22] lint --- firedrake/adjoint/covariance_operator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 95f25d3bd1..a5a9d965b2 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -467,7 +467,7 @@ def norm(self, x: Function): ------- pyadjoint.AdjFloat : The norm of ``x``. - """ + """ # noqa: W605 return self.apply_inverse(x)(x) @abc.abstractmethod From 8403d7ef67f012e19ccb49cfb02ba764583aee87 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 13:11:23 +0000 Subject: [PATCH 07/22] Croci2018 citation --- firedrake/adjoint/covariance_operator.py | 8 +++++--- firedrake/citations.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index a5a9d965b2..ef5d2613e6 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -3,7 +3,7 @@ from functools import cached_property from textwrap import dedent from scipy.special import factorial -from petsctools import get_petscvariables, PCBase +import petsctools from loopy import generate_code_v2 from pyop2 import op2 from firedrake.tsfc_interface import compile_form @@ -215,7 +215,7 @@ def __init__(self, V, rng=None): # TODO: Ask CW if this is the right comm to use. comm = V.mesh()._comm if comm.rank == 0: - petsc_variables = get_petscvariables() + petsc_variables = petsctools.get_petscvariables() BLASLAPACK_LIB = petsc_variables.get("BLASLAPACK_LIB", "") BLASLAPACK_LIB = comm.bcast(BLASLAPACK_LIB, root=0) BLASLAPACK_INCLUDE = petsc_variables.get("BLASLAPACK_INCLUDE", "") @@ -311,6 +311,8 @@ def __init__(self, V, backend=None, rng=None): self.function_space = self.backend.function_space self.rng = self.backend.rng + petsctools.cite("Croci2018") + def sample(self, *, rng=None, tensor=None, apply_riesz=False): return self.backend.sample( rng=rng, tensor=tensor, apply_riesz=apply_riesz) @@ -911,7 +913,7 @@ def CovarianceMat(covariance, operation=None): return mat -class CovariancePC(PCBase): +class CovariancePC(petsctools.PCBase): """ A python PC context for a covariance operator. Will apply either the action or inverse of the covariance, diff --git a/firedrake/citations.py b/firedrake/citations.py index b3e9210b47..6cfb7dd49a 100644 --- a/firedrake/citations.py +++ b/firedrake/citations.py @@ -320,3 +320,19 @@ url = {http://arxiv.org/abs/1410.5620} } """) + +petsctools.add_citation("Croci2018", """ +@article{Croci2018, + title={Efficient White Noise Sampling and Coupling for Multilevel Monte Carlo with Nonnested Meshes}, + volume={6}, + ISSN={2166-2525}, + DOI={10.1137/18M1175239}, + number={4}, + journal={SIAM/ASA Journal on Uncertainty Quantification}, + author={Croci, M. and Giles, M. B. and Rognes, M. E. and Farrell, P. E.}, + year={2018}, + month={jan}, + pages={1630–1655}, + language={en} +} +""") From e37ca6e33a30ce730d306e389a67258777102b31 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 17:47:40 +0000 Subject: [PATCH 08/22] white noise backend docstrings and covariance debug --- firedrake/adjoint/__init__.py | 2 +- firedrake/adjoint/covariance_operator.py | 235 +++++++++++++----- .../adjoint/test_covariance_operator.py | 49 ++-- 3 files changed, 196 insertions(+), 90 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index 21f9a72681..c6401fde0d 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -40,7 +40,7 @@ from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 WhiteNoiseGenerator, AutoregressiveCovariance, - CovarianceMat, CovariancePC) + CovarianceMatCtx, CovarianceMat, CovariancePC) from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 import numpy_adjoint # noqa F401 import firedrake.ufl_expr diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index ef5d2613e6..0c70b80263 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -8,7 +8,7 @@ from pyop2 import op2 from firedrake.tsfc_interface import compile_form from firedrake import ( - grad, inner, avg, action, outer, replace, + grad, inner, avg, action, outer, assemble, CellSize, FacetNormal, dx, ds, dS, sqrt, Constant, Function, Cofunction, RieszMap, @@ -77,15 +77,68 @@ def apply_transpose(self, u): class NoiseBackendBase: + r""" + A base class for implementations of a mass matrix square root action + for generating white noise samples. + + Inheriting classes implement the method from [Croci et al 2018](https://epubs.siam.org/doi/10.1137/18M1175239) + + Generating the samples on the function space :math:`V` requires the following steps: + + 1. On each element generate a white noise sample :math:`z_{e}\sim\mathcal{N}(0, I)` + over all DoFs in the element. Equivalantly, generate the sample on the + discontinuous superspace :math:`V_{d}^{*}` containing :math:`V^{*}`. + (i.e. ``Vd.ufl_element() = BrokenElement(V.ufl_element``). + + 2. Apply the Cholesky factor :math:`C_{e}` of the element-wise mass matrix :math:`M_{e}` + to the element-wise sample (:math:`M_{e}=C_{e}C_{e}^{T}`). + + 3. Assemble the element-wise samples :math:`z_{e}\in V_{d}^{*}` into the global + sample vector :math:`z\in V^{*}`. If :math:`L` is the interpolation operator + then :math:`z=Lz_{e}=LC_{e}z_{e}`. + + 4. Optionally apply a Riesz map to :math:`z` to return a sample in :math:`V`. + + See Also + -------- + PyOP2NoiseBackend + PetscNoiseBackend + WhiteNoiseGenerator + """ + def __init__(self, V, rng=None): self._V = V self._rng = rng or RandomGenerator(PCG64()) - def sample(self, *, rng=None, tensor=None): + @abc.abstractmethod + def sample(self, *, rng=None, tensor=None, apply_riesz=False): + """ + Generate a white noise sample. + + Parameters + ---------- + rng : + A ``RandomGenerator`` to use for sampling IID vectors. + If ``None`` then ``self.rng`` is used. + + tensor : + Optional location to place the result into. + + apply_riesz : + Whether to apply the L2 Riesz map to return a sample in :math:`V`. + + Returns + ------- + Function | Cofunction : + The white noise sample in :math:`V` + """ raise NotImplementedError @cached_property def broken_space(self): + """ + The discontinuous superspace containing :math:`V`, ``self.function_space``. + """ element = self.function_space.ufl_element() mesh = self.function_space.mesh().unique() if isinstance(element, VectorElement): @@ -101,18 +154,44 @@ def broken_space(self): @property def function_space(self): + """The function space that the noise will be generated on. + """ return self._V @property def rng(self): + """The ``RandomGenerator`` to generate the IID sample on the broken function space. + """ return self._rng @cached_property def riesz_map(self): + """A :class:`~firedrake.cofunction.RieszMap` to cache the solver + for :meth:`~firedrake.cofunction.Cofunction.riesz_representation`. + """ return RieszMap(self.function_space, constant_jacobian=True) + """ + Generate a white noise sample. + + Parameters + ---------- + rng : + A ``RandomGenerator`` to use for sampling IID vectors. + If ``None`` then ``self.rng`` is used. + + tensor : + Optional location to place the result into. + + apply_riesz : + Whether to apply an L2 Riesz map to the result to return + a sample in the primal space. + """ class PetscNoiseBackend(NoiseBackendBase): + """ + A PETSc based implementation of a mass matrix square root action for generating white noise. + """ def __init__(self, V, rng=None): super().__init__(V, rng=rng) self.cholesky = CholeskyFactorisation(self.broken_space) @@ -140,6 +219,9 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): class PyOP2NoiseBackend(NoiseBackendBase): + """ + A PyOP2 based implementation of a mass matrix square root for generating white noise. + """ def __init__(self, V, rng=None): super().__init__(V, rng=rng) @@ -260,7 +342,7 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): class WhiteNoiseGenerator: - r""" Generates a white noise sample + r"""Generate white noise samples. Parameters ---------- @@ -271,8 +353,7 @@ class WhiteNoiseGenerator: The ``WhiteNoiseGenerator.Backend`` specifying how to calculate and apply the mass matrix square root. rng : - Initialised random number generator to use for obtaining - random numbers. Defaults to PCG64. + Initialised random number generator to use for sampling IID vectors. Returns ------- @@ -281,10 +362,21 @@ class WhiteNoiseGenerator: :class:`~.firedrake.function.Function` returned and M is the mass matrix. - For details see [Croci et al 2018]: - https://epubs.siam.org/doi/10.1137/18M1175239 + References + ---------- + Croci, M. and Giles, M. B and Rognes, M. E. and Farrell, P. E., 2018: + "Efficient White Noise Sampling and Coupling for Multilevel Monte Carlo + with Nonnested Meshes". SIAM/ASA J. Uncertainty Quantification, Vol. 6, + No. 4, pp. 1630--1655. + https://doi.org/10.1137/18M1175239 + + See Also + -------- + NoiseBackendBase + PyOP2NoiseBackend + PetscNoiseBackend + CovarianceOperatorBase """ - # TODO: Add Croci to citations manager class Backend(Enum): """ @@ -313,7 +405,29 @@ def __init__(self, V, backend=None, rng=None): petsctools.cite("Croci2018") - def sample(self, *, rng=None, tensor=None, apply_riesz=False): + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): + """ + Generate a white noise sample. + + Parameters + ---------- + rng : + A ``RandomGenerator`` to use for sampling IID vectors. + If ``None`` then ``self.rng`` is used. + + tensor : + Optional location to place the result into. + + apply_riesz : + Whether to apply the L2 Riesz map to return a sample in :math:`V`. + + Returns + ------- + Function | Cofunction : + The white noise sample in :math:`V` + """ return self.backend.sample( rng=rng, tensor=tensor, apply_riesz=apply_riesz) @@ -378,12 +492,12 @@ def kappa_m(Lar: float, m: int): class CovarianceOperatorBase: - """ + r""" Abstract base class for a covariance operator B where .. math:: - B: V^{*} \\to V \quad \\text{and} \quad B^{-1}: V \\to V^{*} + B: V^{*} \to V \quad \text{and} \quad B^{-1}: V \to V^{*} The covariance operators can be used to: @@ -415,7 +529,7 @@ class CovarianceOperatorBase: CovarianceMatCtx CovarianceMat CovariancePC - """ # noqa: W605 + """ @abc.abstractmethod def rng(self): @@ -432,7 +546,7 @@ def function_space(self): @abc.abstractmethod def sample(self, *, rng: WhiteNoiseGenerator | None = None, tensor: Function | None = None): - """ + r""" Sample from :math:`\mathcal{N}(0, B)` by correlating a white noise sample: :math:`w = B^{1/2}z`. @@ -448,11 +562,11 @@ def sample(self, *, rng: WhiteNoiseGenerator | None = None, ------- firedrake.function.Function : The sample. - """ # noqa: W605 + """ raise NotImplementedError def norm(self, x: Function): - """Return the weighted norm :math:`\|x\|_{B^{-1}} = x^{T}B^{-1}x`. + r"""Return the weighted norm :math:`\|x\|_{B^{-1}} = x^{T}B^{-1}x`. Default implementation uses ``apply_inverse`` to first calculate the :class:`~firedrake.cofunction.Cofunction` :math:`y = B^{-1}x`, @@ -469,14 +583,14 @@ def norm(self, x: Function): ------- pyadjoint.AdjFloat : The norm of ``x``. - """ # noqa: W605 + """ return self.apply_inverse(x)(x) @abc.abstractmethod def apply_inverse(self, x: Function, *, tensor: Cofunction | None = None): - """Return :math:`y = B^{-1}x` where B is the covariance operator. - :math:`B^{-1}: V \\to V^{*}`. + r"""Return :math:`y = B^{-1}x` where B is the covariance operator. + :math:`B^{-1}: V \to V^{*}`. Parameters ---------- @@ -489,14 +603,14 @@ def apply_inverse(self, x: Function, *, ------- firedrake.cofunction.Cofunction : The result of :math:`B^{-1}x` - """ # noqa: W605 + """ raise NotImplementedError @abc.abstractmethod def apply_action(self, x: Cofunction, *, tensor: Function | None = None): - """Return :math:`y = Bx` where B is the covariance operator. - :math:`B: V^{*} \\to V`. + r"""Return :math:`y = Bx` where B is the covariance operator. + :math:`B: V^{*} \to V`. Parameters ---------- @@ -510,12 +624,12 @@ def apply_action(self, x: Cofunction, *, ------- firedrake.function.Function : The result of :math:`B^{-1}x` - """ # noqa: W605 + """ raise NotImplementedError class AutoregressiveCovariance(CovarianceOperatorBase): - """ + r""" An m-th order autoregressive covariance operator using an implicit diffusion operator. Covariance operator B with a kernel that is the ``m``-th autoregressive @@ -529,9 +643,9 @@ class AutoregressiveCovariance(CovarianceOperatorBase): .. math:: - B: V^{*} \\to V = \lambda((K^{-1}M)^{m}M^{-1})\lambda + B: V^{*} \to V = \lambda((K^{-1}M)^{m}M^{-1})\lambda - B^{-1}: V \\to V^{*} = (1/\lambda)M(M^{-1}K)^{m}(1/\lambda) + B^{-1}: V \to V^{*} = (1/\lambda)M(M^{-1}K)^{m}(1/\lambda) This formulation leads to an efficient implementations for :math:`B^{1/2}` by taking only m/2 steps of the diffusion operator. This can be used @@ -550,7 +664,7 @@ class AutoregressiveCovariance(CovarianceOperatorBase): ---------- Mirouze, I. and Weaver, A. T., 2010: "Representation of correlation functions in variational assimilation using an implicit diffusion - operator". Q. J. R. Meteorol. Soc. 136: 1421–1443, July 2010 Part B, + operator". Q. J. R. Meteorol. Soc. 136: 1421–1443, July 2010 Part B. https://doi.org/10.1002/qj.643 See Also @@ -559,7 +673,7 @@ class AutoregressiveCovariance(CovarianceOperatorBase): CovarianceOperatorBase CovarianceMat CovariancePC - """ # noqa: W605 + """ class DiffusionForm(Enum): """ @@ -580,7 +694,7 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, form = form or self.DiffusionForm.CG self._rng = rng or WhiteNoiseGenerator(V) - self._function_space = function_space or self.rng.function_space + self._function_space = function_space or self.rng().function_space if sigma <= 0: raise ValueError("Variance must be positive.") @@ -611,19 +725,19 @@ def __init__(self, V, L, sigma=1, m=2, rng=None, M = inner(u, v)*dx self._u = Function(V) - self._b = Cofunction(V.dual()) + self._urhs = Function(V) - self._Mrhs = replace(M, {u: self._u}) - self._Krhs = replace(K, {u: self._u}) + self._Mrhs = action(M, self._urhs) + self._Krhs = action(K, self._urhs) self.solver = LinearVariationalSolver( - LinearVariationalProblem(K, self._b, self._u, bcs=bcs, + LinearVariationalProblem(K, self._Mrhs, self._u, bcs=bcs, constant_jacobian=True), solver_parameters=solver_parameters, options_prefix=options_prefix) self.mass_solver = LinearVariationalSolver( - LinearVariationalProblem(M, self._b, self._u, bcs=bcs, + LinearVariationalProblem(M, self._Krhs, self._u, bcs=bcs, constant_jacobian=True), solver_parameters=mass_parameters, options_prefix=mass_prefix) @@ -642,13 +756,11 @@ def sample(self, *, rng=None, tensor=None): w = rng.sample(apply_riesz=True) return tensor.assign(self.stddev*w) - w = rng.sample(apply_riesz=False) + w = rng.sample(apply_riesz=True) + self._u.assign(w) for i in range(self.iterations//2): - if i == 0: - self._b.assign(w) - else: - assemble(self._Mrhs, tensor=self._b) + self._urhs.assign(self._u) self.solver.solve() return tensor.assign(self._weight*self._u) @@ -662,7 +774,7 @@ def norm(self, x): self._u.assign(lamda1*x) for i in range(self.iterations//2): - assemble(self._Krhs, tensor=self._b) + self._urhs.assign(self._u) self.mass_solver.solve() return assemble(inner(self._u, self._u)*dx) @@ -680,26 +792,27 @@ def apply_inverse(self, x, *, tensor=None): self._u.assign(lamda1*x) for i in range(self.iterations): - assemble(self._Krhs, tensor=self._b) + self._urhs.assign(self._u) if i != self.iterations - 1: self.mass_solver.solve() + b = assemble(self._Krhs) - return tensor.assign(lamda1*self._b) + return tensor.assign(lamda1*b) def apply_action(self, x, *, tensor=None): tensor = tensor or Function(self.function_space()) + riesz_map = self.rng().backend.riesz_map + Cx = x.riesz_representation(riesz_map) + if self.iterations == 0: - riesz_map = self.rng().backend.riesz_map - Cx = x.riesz_representation(riesz_map) variance = self.stddev*self.stddev return tensor.assign(variance*Cx) + self._u.assign(self._weight*Cx) + for i in range(self.iterations): - if i == 0: - self._b.assign(self._weight*x) - else: - assemble(self._Mrhs, tensor=self._b) + self._urhs.assign(self._u) self.solver.solve() return tensor.assign(self._weight*self._u) @@ -769,15 +882,15 @@ def diffusion_form(u, v, kappa: Constant | Function, class CovarianceMatCtx: - """ + r""" A python Mat context for a covariance operator. Can apply either the action or inverse of the covariance. .. math:: - B: V^{*} \\to V + B: V^{*} \to V - B^{-1}: V \\to V^{*} + B^{-1}: V \to V^{*} Parameters ---------- @@ -792,7 +905,7 @@ class CovarianceMatCtx: AutoregressiveCovariance CovarianceMat CovariancePC - """ # noqa: W605 + """ class Operation(Enum): ACTION = 'action' INVERSE = 'inverse' @@ -872,16 +985,16 @@ def view(self, mat, viewer=None): def CovarianceMat(covariance, operation=None): - """ + r""" A Mat for a covariance operator. Can apply either the action or inverse of the covariance. This is a convenience function to create a PETSc.Mat with a :class:`.CovarianceMatCtx` Python context. .. math:: - B: V^{*} \\to V + B: V^{*} \to V - B^{-1}: V \\to V^{*} + B^{-1}: V \to V^{*} Parameters ---------- @@ -901,10 +1014,10 @@ def CovarianceMat(covariance, operation=None): AutoregressiveCovariance CovarianceMatCtx CovariancePC - """ # noqa: W605 + """ ctx = CovarianceMatCtx(covariance, operation=operation) - sizes = covariance.function_space.dof_dset.layout_vec.getSizes() + sizes = covariance.function_space().dof_dset.layout_vec.getSizes() mat = PETSc.Mat().createPython( (sizes, sizes), ctx, comm=ctx.comm) @@ -914,16 +1027,16 @@ def CovarianceMat(covariance, operation=None): class CovariancePC(petsctools.PCBase): - """ + r""" A python PC context for a covariance operator. Will apply either the action or inverse of the covariance, whichever is the opposite of the Mat operator. .. math:: - B: V^{*} \\to V + B: V^{*} \to V - B^{-1}: V \\to V^{*} + B^{-1}: V \to V^{*} Available options: @@ -935,7 +1048,7 @@ class CovariancePC(petsctools.PCBase): AutoregressiveCovariance CovarianceMatCtx CovarianceMat - """ # noqa: W605 + """ needs_python_pmat = True prefix = "covariance" @@ -954,7 +1067,7 @@ def initialize(self, pc): self.covariance = covariance self.mat = mat - V = covariance.function_space + V = covariance.function_space() primal = Function(V) dual = Function(V.dual()) diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index 77f3e29583..6cfe40275a 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -85,11 +85,10 @@ def test_white_noise(family, degree, mesh_type, dim, backend): @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) -@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) @pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vector1", "vector2"]) @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) -def test_covariance_inverse_action(m, family, degree, mesh_type, dim): +def test_covariance_inverse_action(m, family, mesh_type, dim): """Test that covariance operator action and inverse are opposites. """ @@ -107,10 +106,10 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim): x, y, z = SpatialCoordinate(mesh) wexpr = cos(2*pi*x)*cos(4*pi*y)*cos(pi*z) if dim > 0: - V = VectorFunctionSpace(mesh, family, degree, dim=dim) + V = VectorFunctionSpace(mesh, family, 1, dim=dim) wexpr = as_vector([-1**(j+1)*wexpr for j in range(dim)]) else: - V = FunctionSpace(mesh, family, degree) + V = FunctionSpace(mesh, family, 1) rng = WhiteNoiseGenerator( V, rng=RandomGenerator(PCG64(seed=13))) @@ -125,13 +124,13 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim): } if family == 'CG': - form = GaussianCovariance.DiffusionForm.CG + form = AutoregressiveCovariance.DiffusionForm.CG elif family == 'DG': - form = GaussianCovariance.DiffusionForm.IP + form = AutoregressiveCovariance.DiffusionForm.IP else: raise ValueError("Do not know which diffusion form to use for family {family}") - B = GaussianCovariance( + B = AutoregressiveCovariance( V, L, sigma, m, rng=rng, form=form, solver_parameters=solver_parameters, options_prefix="") @@ -146,8 +145,7 @@ def test_covariance_inverse_action(m, family, degree, mesh_type, dim): @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) -@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) -def test_covariance_inverse_action_hdiv(m, degree): +def test_covariance_inverse_action_hdiv(m): """Test that covariance operator action and inverse are opposites for hdiv spaces. """ @@ -157,7 +155,7 @@ def test_covariance_inverse_action_hdiv(m, degree): x, y = SpatialCoordinate(mesh) wexpr = cos(2*pi*x)*cos(4*pi*x) - V = FunctionSpace(mesh, "BDM", degree) + V = FunctionSpace(mesh, "BDM", 1) wexpr = as_vector([-1**(j+1)*wexpr for j in range(2)]) L = 0.1 @@ -169,9 +167,9 @@ def test_covariance_inverse_action_hdiv(m, degree): 'pc_factor_mat_solver_type': 'mumps' } - form = GaussianCovariance.DiffusionForm.IP + form = AutoregressiveCovariance.DiffusionForm.IP - B = GaussianCovariance( + B = AutoregressiveCovariance( V, L, sigma, m, form=form, solver_parameters=solver_parameters, options_prefix="") @@ -203,13 +201,13 @@ def test_covariance_adjoint_norm(m, family): v = Function(V).project(2 - 0.5*sin(6*pi*x)) if family == 'CG': - form = GaussianCovariance.DiffusionForm.CG + form = AutoregressiveCovariance.DiffusionForm.CG elif family == 'DG': - form = GaussianCovariance.DiffusionForm.IP + form = AutoregressiveCovariance.DiffusionForm.IP else: raise ValueError("Do not know which diffusion form to use for family {family}") - B = GaussianCovariance(V, L, sigma, m, form=form) + B = AutoregressiveCovariance(V, L, sigma, m, form=form) continue_annotation() with set_working_tape() as tape: @@ -231,9 +229,8 @@ def test_covariance_adjoint_norm(m, family): @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("family", ("CG", "DG")) -@pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) @pytest.mark.parametrize("operation", ("action", "inverse")) -def test_covariance_mat(m, family, degree, operation): +def test_covariance_mat(m, family, operation): """Test that covariance mat and pc apply correct and opposite actions. """ nx = 20 @@ -243,16 +240,16 @@ def test_covariance_mat(m, family, degree, operation): mesh = UnitIntervalMesh(nx) coords, = SpatialCoordinate(mesh) - V = FunctionSpace(mesh, family, degree) + V = FunctionSpace(mesh, family, 1) if family == 'CG': - form = GaussianCovariance.DiffusionForm.CG + form = AutoregressiveCovariance.DiffusionForm.CG elif family == 'DG': - form = GaussianCovariance.DiffusionForm.IP + form = AutoregressiveCovariance.DiffusionForm.IP else: raise ValueError("Do not know which diffusion form to use for family {family}") - B = GaussianCovariance(V, L, sigma, m, form=form) + B = AutoregressiveCovariance(V, L, sigma, m, form=form) operation = CovarianceMatCtx.Operation(operation) @@ -293,18 +290,14 @@ def test_covariance_mat(m, family, degree, operation): ksp = PETSc.KSP().create() ksp.setOperators(mat) - # poorly conditioned cases - if (degree == 2) and (m == 4): - tol = 1e-6 - else: - tol = 1e-8 + tol = 1e-8 petsctools.set_from_options( - ksp, options_prefix="action", + ksp, options_prefix=str(operation), parameters={ 'ksp_monitor': None, 'ksp_type': 'richardson', - 'ksp_max_it': 3, + 'ksp_max_it': 2, 'ksp_rtol': tol, 'pc_type': 'python', 'pc_python_type': 'firedrake.adjoint.CovariancePC', From 05751908bc74c610b4cae0a90b9ef4cc2a4829b1 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 18:02:03 +0000 Subject: [PATCH 09/22] use L2Cholesky from transformedRF for white noise generation --- firedrake/adjoint/__init__.py | 2 +- firedrake/adjoint/covariance_operator.py | 62 ++-------------- firedrake/adjoint/transformed_functional.py | 80 --------------------- 3 files changed, 7 insertions(+), 137 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index c6401fde0d..b42e2a47a8 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -38,10 +38,10 @@ from firedrake.adjoint.ufl_constraints import UFLInequalityConstraint, \ UFLEqualityConstraint # noqa F401 from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional # noqa F401 +from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 WhiteNoiseGenerator, AutoregressiveCovariance, CovarianceMatCtx, CovarianceMat, CovariancePC) -from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 import numpy_adjoint # noqa F401 import firedrake.ufl_expr import types diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 0c70b80263..7bdebd39a7 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -7,6 +7,7 @@ from loopy import generate_code_v2 from pyop2 import op2 from firedrake.tsfc_interface import compile_form +from firedrake.adjoint.transformed_functional import L2Cholesky from firedrake import ( grad, inner, avg, action, outer, assemble, CellSize, FacetNormal, @@ -18,64 +19,10 @@ RandomGenerator, PCG64, LinearVariationalProblem, LinearVariationalSolver, - LinearSolver, PETSc ) -class CholeskyFactorisation: - def __init__(self, V, form=None): - self._V = V - - if form is None: - self.form = inner(TrialFunction(V), - TestFunction(V))*dx - else: - self.form = form - - self._wrk = Function(V) - - @property - def function_space(self): - return self._V - - @cached_property - def _assemble_action(self): - from firedrake.assemble import get_assembler - return get_assembler(action(self.form, self._wrk)).assemble - - def assemble_action(self, u, tensor=None): - self._wrk.assign(u) - return self._assemble_action(tensor=tensor) - - @cached_property - def solver(self): - return LinearSolver( - assemble(self.form, mat_type='aij'), - solver_parameters={ - "ksp_type": "preonly", - "pc_type": "cholesky", - "pc_factor_mat_ordering_type": "nd"}) - - @cached_property - def pc(self): - return self.solver.ksp.getPC() - - def apply(self, u): - u = self.assemble_action(u) - v = Cofunction(self.space.dual()) - with u.dat.vec_ro as u_v, v.dat.vec_wo as v_v: - self.pc.applySymmetricLeft(u_v, v_v) - return v - - def apply_transpose(self, u): - v = Function(self.function_space) - with u.dat.vec_ro as u_v, v.dat.vec_wo as v_v: - self.pc.applySymmetricRight(u_v, v_v) - v = self.assemble_action(v) - return v - - class NoiseBackendBase: r""" A base class for implementations of a mass matrix square root action @@ -194,7 +141,9 @@ class PetscNoiseBackend(NoiseBackendBase): """ def __init__(self, V, rng=None): super().__init__(V, rng=rng) - self.cholesky = CholeskyFactorisation(self.broken_space) + self.cholesky = L2Cholesky(self.broken_space) + self._zb = Function(self.broken_space) + self.M = inner(self._zb, TestFunction(self.broken_space))*dx def sample(self, *, rng=None, tensor=None, apply_riesz=False): V = self.function_space @@ -203,7 +152,8 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): # z z = rng.standard_normal(self.broken_space) # C z - Cz = self.cholesky.apply_transpose(z) + self._zb.assign(self.cholesky.C_T_inv_action(z)) + Cz = assemble(self.M) # L C z b = Cofunction(V.dual()).interpolate(Cz) diff --git a/firedrake/adjoint/transformed_functional.py b/firedrake/adjoint/transformed_functional.py index 392117bcb9..d3eeb2f46b 100644 --- a/firedrake/adjoint/transformed_functional.py +++ b/firedrake/adjoint/transformed_functional.py @@ -3,13 +3,10 @@ from numbers import Real from operator import itemgetter from typing import Optional, Union -from functools import cached_property import firedrake as fd from firedrake.adjoint import Control, ReducedFunctional, Tape from firedrake.functionspaceimpl import WithGeometry -from firedrake.ufl_expr import action -from firedrake.assemble import get_assembler import finat import pyadjoint from pyadjoint import no_annotations @@ -80,33 +77,6 @@ def _pc(self): return pc - @cached_property - def _M_action_assembler(self): - wrk = fd.Function(self.space) - M = fd.inner(fd.TrialFunction(self.space), - fd.TestFunction(self.space))*fd.dx - return wrk, get_assembler(action(M, wrk)).assemble - - def _M_action(self, u: fd.Function, - tensor: Optional[fd.Cofunction] = None) -> fd.Cofunction: - r"""Apply the action of the mass matrix. - - Parameters - ---------- - - u : - The :class:`~firedrake.function.Function` being acted on. - - Returns - ------- - - firedrake.cofunction.Cofunction - The result of the action of the mass matrix on ``u``. - """ - wrk, assembler = self._M_action_assembler - wrk.assign(u) - return assembler(tensor=tensor) - def C_inv_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Cofunction: r"""For the Cholesky factorization @@ -167,56 +137,6 @@ def C_T_inv_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Function: pc.applySymmetricRight(u_v_s, v_v_s) return v - def C_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Cofunction: - r"""For the Cholesky factorization - - ... math : - - M = C C^T, - - compute the action of :math:`C`. - - Parameters - ---------- - - u - Compute :math:`C \tilde{u}` where :math:`\tilde{u}` is the - vector of degrees of freedom for :math:`u`. - - Returns - ------- - - firedrake.cofunction.Cofunction - Has vector of degrees of freedom :math:`C \tilde{u}`. - """ - # ??? - return self.C_T_inv_action(self._M_action(u)) - - def C_T_action(self, u: Union[fd.Function, fd.Cofunction]) -> fd.Cofunction: - r"""For the Cholesky factorization - - ... math : - - M = C C^T, - - compute the action of :math:`C^{T}`. - - Parameters - ---------- - - u - Compute :math:`C^{T} \tilde{u}` where :math:`\tilde{u}` is the - vector of degrees of freedom for :math:`u`. - - Returns - ------- - - firedrake.function.Function - Has vector of degrees of freedom :math:`C^{T} \tilde{u}`. - """ - # ??? - return self._M_action(self.C_inv_action(u)) - class L2RieszMap(fd.RieszMap): """An :math:`L^2` Riesz map. From 777ce758738550f665b648752ae20896ca1b464b Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 19:07:12 +0000 Subject: [PATCH 10/22] covariance: remove unnecessary assembles --- firedrake/adjoint/covariance_operator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 7bdebd39a7..23468972f4 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -745,7 +745,7 @@ def apply_inverse(self, x, *, tensor=None): self._urhs.assign(self._u) if i != self.iterations - 1: self.mass_solver.solve() - b = assemble(self._Krhs) + b = assemble(self._Krhs) return tensor.assign(lamda1*b) From 0b09b6bc56e447f02a2b341eb3952f7cbe68fc00 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 19:45:42 +0000 Subject: [PATCH 11/22] covariance: type hints --- firedrake/adjoint/covariance_operator.py | 115 +++++++++++++++++------ 1 file changed, 85 insertions(+), 30 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 23468972f4..075d8de93a 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -1,6 +1,7 @@ import abc from enum import Enum from functools import cached_property +from typing import Iterable from textwrap import dedent from scipy.special import factorial import petsctools @@ -8,6 +9,8 @@ from pyop2 import op2 from firedrake.tsfc_interface import compile_form from firedrake.adjoint.transformed_functional import L2Cholesky +from firedrake.functionspaceimpl import WithGeometry +from firedrake.bcs import BCBase from firedrake import ( grad, inner, avg, action, outer, assemble, CellSize, FacetNormal, @@ -46,6 +49,13 @@ class NoiseBackendBase: 4. Optionally apply a Riesz map to :math:`z` to return a sample in :math:`V`. + Parameters + ---------- + V : + The :func:`~.firedrake.functionspace.FunctionSpace` to generate the samples in. + rng : + The ``RandomGenerator`` to generate the samples on the discontinuous superspace. + See Also -------- PyOP2NoiseBackend @@ -53,12 +63,14 @@ class NoiseBackendBase: WhiteNoiseGenerator """ - def __init__(self, V, rng=None): + def __init__(self, V: WithGeometry, rng=None): self._V = V self._rng = rng or RandomGenerator(PCG64()) @abc.abstractmethod - def sample(self, *, rng=None, tensor=None, apply_riesz=False): + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): """ Generate a white noise sample. @@ -139,13 +151,15 @@ class PetscNoiseBackend(NoiseBackendBase): """ A PETSc based implementation of a mass matrix square root action for generating white noise. """ - def __init__(self, V, rng=None): + def __init__(self, V: WithGeometry, rng=None): super().__init__(V, rng=rng) self.cholesky = L2Cholesky(self.broken_space) self._zb = Function(self.broken_space) self.M = inner(self._zb, TestFunction(self.broken_space))*dx - def sample(self, *, rng=None, tensor=None, apply_riesz=False): + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): V = self.function_space rng = rng or self.rng @@ -172,7 +186,7 @@ class PyOP2NoiseBackend(NoiseBackendBase): """ A PyOP2 based implementation of a mass matrix square root for generating white noise. """ - def __init__(self, V, rng=None): + def __init__(self, V: WithGeometry, rng=None): super().__init__(V, rng=rng) u = TrialFunction(V) @@ -261,7 +275,9 @@ def __init__(self, V, rng=None): include_dirs=BLASLAPACK_INCLUDE.split(), ldargs=BLASLAPACK_LIB.split()) - def sample(self, *, rng=None, tensor=None, apply_riesz=False): + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): rng = rng or self.rng z = rng.standard_normal(self.broken_space) @@ -294,24 +310,20 @@ def sample(self, *, rng=None, tensor=None, apply_riesz=False): class WhiteNoiseGenerator: r"""Generate white noise samples. + Generates samples :math:`w\in V^{*}` with + :math:`w\sim\mathcal{N}(0, M)`, where :math:`M` is + the mass matrix, or its Riesz representer in :math:`V`. + Parameters ---------- V : The :class:`~firedrake.functionspace.FunctionSpace` to construct a white noise sample on. - backend : - The ``WhiteNoiseGenerator.Backend`` specifying how to calculate - and apply the mass matrix square root. + backend : WhiteNoiseGenerator.Backend + The backend specifying how to calculate and apply the mass matrix square root. rng : Initialised random number generator to use for sampling IID vectors. - Returns - ------- - firedrake.function.Function : - with b ~ Normal(0, M) where b is the dat.data of the - :class:`~.firedrake.function.Function` returned and - M is the mass matrix. - References ---------- Croci, M. and Giles, M. B and Rognes, M. E. and Farrell, P. E., 2018: @@ -340,7 +352,7 @@ class Backend(Enum): PYOP2 = 'pyop2' PETSC = 'petsc' - def __init__(self, V, backend=None, rng=None): + def __init__(self, V: WithGeometry, backend=None, rng=None): backend = backend or self.Backend.PYOP2 if backend == self.Backend.PYOP2: self.backend = PyOP2NoiseBackend(V, rng=rng) @@ -376,7 +388,7 @@ def sample(self, *, rng=None, Returns ------- Function | Cofunction : - The white noise sample in :math:`V` + The white noise sample """ return self.backend.sample( rng=rng, tensor=tensor, apply_riesz=apply_riesz) @@ -504,7 +516,7 @@ def sample(self, *, rng: WhiteNoiseGenerator | None = None, ---------- rng : Generator for the white noise sample. - If not provided then self.rng will be used. + If not provided then ``self.rng`` will be used. tensor : Optional location to place the result into. @@ -610,6 +622,35 @@ class AutoregressiveCovariance(CovarianceOperatorBase): The white noise sample :math:`M^{1/2}z` is generated by a :class:`.WhiteNoiseGenerator`. + Parameters + ---------- + V : + The function space that the covariance operator maps into. + L : + The correlation lengthscale. + sigma : + The standard deviation. + m : + The number of diffusion operator steps. + Equal to the order of the autoregressive function kernel. + rng : + White noise generator to seed generating correlated samples. + form : AutoregressiveCovariance.DiffusionForm | ufl.Form | None + The diffusion formulation or form. If a ``DiffusionForm`` then + :func:`.diffusion_form` will be used to generate the diffusion + form. Otherwise assumed to be a ufl.Form on ``V``. + Defaults to ``AutoregressiveCovariance.DiffusionForm.CG``. + bcs : + Boundary conditions for the diffusion operator. + solver_parameters : + The PETSc options for the diffusion operator solver. + options_prefix : + The options prefix for the diffusion operator solver. + mass_parameters : + The PETSc options for the mass matrix solver. + mass_prefix : + The options prefix for the matrix matrix solver. + References ---------- Mirouze, I. and Weaver, A. T., 2010: "Representation of correlation @@ -623,6 +664,7 @@ class AutoregressiveCovariance(CovarianceOperatorBase): CovarianceOperatorBase CovarianceMat CovariancePC + diffusion_form """ class DiffusionForm(Enum): @@ -636,15 +678,19 @@ class DiffusionForm(Enum): CG = 'CG' IP = 'IP' - def __init__(self, V, L, sigma=1, m=2, rng=None, - bcs=None, form=None, function_space=None, - solver_parameters=None, options_prefix=None, - mass_parameters=None, mass_prefix=None): + def __init__(self, V: WithGeometry, L: float | Constant, + sigma: float | Constant = 1., m: int = 2, + rng: WhiteNoiseGenerator | None = None, form=None, + bcs: BCBase | Iterable[BCBase] | None = None, + solver_parameters: dict | None = None, + options_prefix: str | None = None, + mass_parameters: dict | None = None, + mass_prefix: str | None = None): form = form or self.DiffusionForm.CG self._rng = rng or WhiteNoiseGenerator(V) - self._function_space = function_space or self.rng().function_space + self._function_space = self.rng().function_space if sigma <= 0: raise ValueError("Variance must be positive.") @@ -698,7 +744,8 @@ def function_space(self): def rng(self): return self._rng - def sample(self, *, rng=None, tensor=None): + def sample(self, *, rng: WhiteNoiseGenerator | None = None, + tensor: Function | None = None): tensor = tensor or Function(self.function_space()) rng = rng or self.rng() @@ -715,7 +762,7 @@ def sample(self, *, rng=None, tensor=None): return tensor.assign(self._weight*self._u) - def norm(self, x): + def norm(self, x: Function): if self.iterations == 0: sigma_x = self.stddev*x return assemble(inner(sigma_x, sigma_x)*dx) @@ -729,7 +776,8 @@ def norm(self, x): return assemble(inner(self._u, self._u)*dx) - def apply_inverse(self, x, *, tensor=None): + def apply_inverse(self, x: Function, *, + tensor: Cofunction | None = None): tensor = tensor or Cofunction(self.function_space().dual()) if self.iterations == 0: @@ -749,7 +797,8 @@ def apply_inverse(self, x, *, tensor=None): return tensor.assign(lamda1*b) - def apply_action(self, x, *, tensor=None): + def apply_action(self, x: Cofunction, *, + tensor: Function | None = None): tensor = tensor or Function(self.function_space()) riesz_map = self.rng().backend.riesz_map @@ -848,6 +897,7 @@ class CovarianceMatCtx: The covariance operator. operation : CovarianceMatCtx.Operation Whether the matrix applies the action or inverse of the covariance operator. + Defaults to ``Operation.ACTION``. See Also -------- @@ -908,6 +958,8 @@ def mult(self, mat, x, y): v.copy(result=y) def view(self, mat, viewer=None): + """View object. Method usually called by PETSc with e.g. -ksp_view. + """ if viewer is None: return if viewer.getType() != PETSc.Viewer.Type.ASCII: @@ -934,7 +986,8 @@ def view(self, mat, viewer=None): ksp.setTabLevel(level) -def CovarianceMat(covariance, operation=None): +def CovarianceMat(covariance: CovarianceOperatorBase, + operation: CovarianceMatCtx.Operation | None = None): r""" A Mat for a covariance operator. Can apply either the action or inverse of the covariance. @@ -956,7 +1009,7 @@ def CovarianceMat(covariance, operation=None): Returns ------- PETSc.Mat : - The python type Mat with a CovarianceMatCtx context. + The python type Mat with a :class:`CovarianceMatCtx` context. See Also -------- @@ -1060,6 +1113,8 @@ def update(self, pc): pass def view(self, pc, viewer=None): + """View object. Method usually called by PETSc with e.g. -ksp_view. + """ if viewer is None: return if viewer.getType() != PETSc.Viewer.Type.ASCII: From 4eb226864388aa7d9477d1462653ef048be47a66 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 20:47:34 +0000 Subject: [PATCH 12/22] covariance norm bugfix --- firedrake/adjoint/covariance_operator.py | 2 +- tests/firedrake/adjoint/test_covariance_operator.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 075d8de93a..c008bf02b9 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -764,7 +764,7 @@ def sample(self, *, rng: WhiteNoiseGenerator | None = None, def norm(self, x: Function): if self.iterations == 0: - sigma_x = self.stddev*x + sigma_x = (1/self.stddev)*x return assemble(inner(sigma_x, sigma_x)*dx) lamda1 = 1/self._weight diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index 6cfe40275a..affcd70bdc 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -25,6 +25,7 @@ def petsc2numpy_mat(petsc_mat): ).todense() +@pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) @pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vec1", "vec2"]) @@ -83,6 +84,7 @@ def test_white_noise(family, degree, mesh_type, dim, backend): assert (1 - tol) < np.max(normalised_errors) < (1 + tol) +@pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vector1", "vector2"]) @@ -143,6 +145,7 @@ def test_covariance_inverse_action(m, family, mesh_type, dim): assert errornorm(w, wcheck) < tol +@pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) def test_covariance_inverse_action_hdiv(m): @@ -182,6 +185,7 @@ def test_covariance_inverse_action_hdiv(m): assert errornorm(w, wcheck) < tol +@pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("family", ("CG", "DG")) @@ -226,6 +230,7 @@ def test_covariance_adjoint_norm(m, family): assert min(taylor['R2']['Rate']) > 2.95, taylor['R2'] +@pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) @pytest.mark.parametrize("family", ("CG", "DG")) From 078b1cf8ebf6fce6ac0c1a1b5e9e6ebdceedd0fd Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 8 Dec 2025 21:13:10 +0000 Subject: [PATCH 13/22] Update tests/firedrake/adjoint/test_covariance_operator.py --- tests/firedrake/adjoint/test_covariance_operator.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index affcd70bdc..e4b66899ef 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -35,9 +35,6 @@ def petsc2numpy_mat(petsc_mat): def test_white_noise(family, degree, mesh_type, dim, backend): """Test that white noise generator converges to a mass matrix covariance. """ - if backend == "petsc" and COMM_WORLD.size > 1: - pytest.skip( - "petsc backend for noise generation not implemented in parallel.") nx = 10 # Mesh dimension From d25c0c3fa8fc1eecf2cc07312c3db8b780a5e29f Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 9 Dec 2025 08:41:48 +0000 Subject: [PATCH 14/22] VOM white noise generator --- firedrake/adjoint/covariance_operator.py | 68 ++++++++++- .../adjoint/test_covariance_operator.py | 107 ++++++++++++------ 2 files changed, 134 insertions(+), 41 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index c008bf02b9..7922449eed 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -22,6 +22,7 @@ RandomGenerator, PCG64, LinearVariationalProblem, LinearVariationalSolver, + VertexOnlyMeshTopology, PETSc ) @@ -147,9 +148,49 @@ def riesz_map(self): """ +class VOMNoiseBackend(NoiseBackendBase): + """ + A PETSc based implementation of a mass matrix square root action + for generating white noise on a vertex only mesh. + """ + def __init__(self, V: WithGeometry, rng=None): + super().__init__(V, rng=rng) + self.cholesky = L2Cholesky(V) + self._zb = Function(V) + self.M = inner(self._zb, TestFunction(V))*dx + + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): + rng = rng or self.rng + + # z + z = rng.standard_normal(self.broken_space) + # C z + self._zb.assign(self.cholesky.C_T_inv_action(z)) + Cz = assemble(self.M) + + # Usually we would interpolate to the unbroken space, + # but here we're on a VOM so everything is broken. + # L C z + # b = Cofunction(V.dual()).interpolate(Cz) + b = Cz + + if apply_riesz: + b = b.riesz_representation(self.riesz_map) + + if tensor: + tensor.assign(b) + else: + tensor = b + + return tensor + + class PetscNoiseBackend(NoiseBackendBase): """ - A PETSc based implementation of a mass matrix square root action for generating white noise. + A PETSc based implementation of a mass matrix square root action + for generating white noise. """ def __init__(self, V: WithGeometry, rng=None): super().__init__(V, rng=rng) @@ -184,7 +225,8 @@ def sample(self, *, rng=None, class PyOP2NoiseBackend(NoiseBackendBase): """ - A PyOP2 based implementation of a mass matrix square root for generating white noise. + A PyOP2 based implementation of a mass matrix square root + for generating white noise. """ def __init__(self, V: WithGeometry, rng=None): super().__init__(V, rng=rng) @@ -337,6 +379,7 @@ class WhiteNoiseGenerator: NoiseBackendBase PyOP2NoiseBackend PetscNoiseBackend + VOMNoiseBackend CovarianceOperatorBase """ @@ -348,16 +391,31 @@ class Backend(Enum): -------- PyOP2NoiseBackend PetscNoiseBackend + VOMNoiseBackend """ PYOP2 = 'pyop2' PETSC = 'petsc' + VOM = 'vom' def __init__(self, V: WithGeometry, backend=None, rng=None): - backend = backend or self.Backend.PYOP2 + # Not all backends are valid for VOM. + if isinstance(V.mesh().topology, VertexOnlyMeshTopology): + backend = self.Backend(backend or self.Backend.VOM) + if backend != self.Backend.VOM: + raise ValueError( + f"Cannot use white noise backend {backend} with a VertexOnlyMesh." + " Please use WhiteNoiseGenerator.Backend.VOM") + else: + backend = self.Backend(backend or self.Backend.PYOP2) + + backend = self.Backend(backend) + if backend == self.Backend.PYOP2: self.backend = PyOP2NoiseBackend(V, rng=rng) elif backend == self.Backend.PETSC: self.backend = PetscNoiseBackend(V, rng=rng) + elif backend == self.Backend.VOM: + self.backend = VOMNoiseBackend(V, rng=rng) else: raise ValueError( f"Unrecognised white noise generation backend {backend}") @@ -688,6 +746,8 @@ def __init__(self, V: WithGeometry, L: float | Constant, mass_prefix: str | None = None): form = form or self.DiffusionForm.CG + if isinstance(form, str): + form = self.DiffusionForm(form) self._rng = rng or WhiteNoiseGenerator(V) self._function_space = self.rng().function_space @@ -911,7 +971,7 @@ class Operation(Enum): INVERSE = 'inverse' def __init__(self, covariance: CovarianceOperatorBase, operation=None): - operation = operation or self.Operation.ACTION + operation = self.Operation(operation or self.Operation.ACTION) V = covariance.function_space() self.function_space = V diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index affcd70bdc..f41908f6e7 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -25,6 +25,11 @@ def petsc2numpy_mat(petsc_mat): ).todense() +@pytest.fixture +def rng(): + return RandomGenerator(PCG64(seed=13)) + + @pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) @@ -32,7 +37,7 @@ def petsc2numpy_mat(petsc_mat): @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) @pytest.mark.parametrize("backend", ("pyop2", "petsc")) -def test_white_noise(family, degree, mesh_type, dim, backend): +def test_white_noise(family, degree, mesh_type, dim, backend, rng): """Test that white noise generator converges to a mass matrix covariance. """ if backend == "petsc" and COMM_WORLD.size > 1: @@ -59,10 +64,61 @@ def test_white_noise(family, degree, mesh_type, dim, backend): covmat = petsc2numpy_mat( assemble(M, mat_type='aij').petscmat) - rng = RandomGenerator(PCG64(seed=13)) + generator = WhiteNoiseGenerator(V, backend=backend, rng=rng) + + # Test convergence as sample size increases + nsamples = [50, 100, 200, 400, 800] + + samples = np.empty((V.dim(), nsamples[-1])) + for i in range(nsamples[-1]): + with generator.sample().dat.vec_ro as bv: + samples[:, i] = petsc2numpy_vec(bv) + + covariances = [np.cov(samples[:, :ns]) for ns in nsamples] + + # Covariance matrix should converge at a rate of sqrt(n) + errors = [np.linalg.norm(cov-covmat) for cov in covariances] + normalised_errors = [err*sqrt(n) for err, n in zip(errors, nsamples)] + normalised_errors /= normalised_errors[-1] + + # Loose tolerance because RNG + tol = 0.2 + assert (1 - tol) < np.max(normalised_errors) < (1 + tol) + - generator = WhiteNoiseGenerator( - V, backend=WhiteNoiseGenerator.Backend(backend), rng=rng) +@pytest.mark.skipcomplex +@pytest.mark.parallel([1, 2]) +@pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vec1", "vec2"]) +@pytest.mark.parametrize("mesh_type", ("interval", "square")) +def test_vom_white_noise(dim, mesh_type, rng): + """Test that white noise generator converges to a mass matrix covariance. + """ + + nx = 10 + nv = 10 + np.random.seed(13) + # Mesh dimension + if mesh_type == 'interval': + mesh = UnitIntervalMesh(nx) + points = np.random.random_sample((nv, 1)) + elif mesh_type == 'square': + mesh = UnitSquareMesh(nx, nx) + points = np.random.random_sample((nv, 2)) + + vom = VertexOnlyMesh(mesh, points) + + # Variable rank + if dim > 0: + V = VectorFunctionSpace(vom, "DG", 0, dim=dim) + else: + V = FunctionSpace(vom, "DG", 0) + + # Finite element white noise has mass matrix covariance + M = inner(TrialFunction(V), TestFunction(V))*dx + covmat = petsc2numpy_mat( + assemble(M, mat_type='aij').petscmat) + + generator = WhiteNoiseGenerator(V, backend='vom', rng=rng) # Test convergence as sample size increases nsamples = [50, 100, 200, 400, 800] @@ -113,9 +169,6 @@ def test_covariance_inverse_action(m, family, mesh_type, dim): else: V = FunctionSpace(mesh, family, 1) - rng = WhiteNoiseGenerator( - V, rng=RandomGenerator(PCG64(seed=13))) - L = 0.1 sigma = 0.9 @@ -125,15 +178,10 @@ def test_covariance_inverse_action(m, family, mesh_type, dim): 'pc_factor_mat_solver_type': 'mumps' } - if family == 'CG': - form = AutoregressiveCovariance.DiffusionForm.CG - elif family == 'DG': - form = AutoregressiveCovariance.DiffusionForm.IP - else: - raise ValueError("Do not know which diffusion form to use for family {family}") + form = 'IP' if family == 'DG' else 'CG' B = AutoregressiveCovariance( - V, L, sigma, m, rng=rng, form=form, + V, L, sigma, m, form=form, solver_parameters=solver_parameters, options_prefix="") @@ -170,10 +218,8 @@ def test_covariance_inverse_action_hdiv(m): 'pc_factor_mat_solver_type': 'mumps' } - form = AutoregressiveCovariance.DiffusionForm.IP - B = AutoregressiveCovariance( - V, L, sigma, m, form=form, + V, L, sigma, m, form='IP', solver_parameters=solver_parameters, options_prefix="") @@ -204,13 +250,7 @@ def test_covariance_adjoint_norm(m, family): u = Function(V).project(sin(2*pi*x)) v = Function(V).project(2 - 0.5*sin(6*pi*x)) - if family == 'CG': - form = AutoregressiveCovariance.DiffusionForm.CG - elif family == 'DG': - form = AutoregressiveCovariance.DiffusionForm.IP - else: - raise ValueError("Do not know which diffusion form to use for family {family}") - + form = 'IP' if family == 'DG' else 'CG' B = AutoregressiveCovariance(V, L, sigma, m, form=form) continue_annotation() @@ -247,22 +287,15 @@ def test_covariance_mat(m, family, operation): V = FunctionSpace(mesh, family, 1) - if family == 'CG': - form = AutoregressiveCovariance.DiffusionForm.CG - elif family == 'DG': - form = AutoregressiveCovariance.DiffusionForm.IP - else: - raise ValueError("Do not know which diffusion form to use for family {family}") + form = 'IP' if family == 'DG' else 'CG' B = AutoregressiveCovariance(V, L, sigma, m, form=form) - operation = CovarianceMatCtx.Operation(operation) - mat = CovarianceMat(B, operation=operation) expr = 2*pi*coords - if operation == CovarianceMatCtx.Operation.ACTION: + if operation == 'action': x = Function(V).project(expr).riesz_representation() y = Function(V) xcheck = x.copy(deepcopy=True) @@ -270,7 +303,7 @@ def test_covariance_mat(m, family, operation): B.apply_action(xcheck, tensor=ycheck) - elif operation == CovarianceMatCtx.Operation.INVERSE: + elif operation == 'inverse': x = Function(V).project(expr) y = Function(V.dual()) xcheck = x.copy(deepcopy=True) @@ -282,13 +315,13 @@ def test_covariance_mat(m, family, operation): mat.mult(xv, yv) # flip to primal space to calculate norms - if operation == CovarianceMatCtx.Operation.INVERSE: + if operation == 'inverse': y = y.riesz_representation() ycheck = ycheck.riesz_representation() assert errornorm(ycheck, y)/norm(ycheck) < 1e-12 - if operation == CovarianceMatCtx.Operation.INVERSE: + if operation == 'inverse': y = y.riesz_representation() ycheck = ycheck.riesz_representation() @@ -318,7 +351,7 @@ def test_covariance_mat(m, family, operation): # be exact inverses of each other. assert ksp.its == 1 - if operation == CovarianceMatCtx.Operation.ACTION: + if operation == 'action': x = x.riesz_representation() xcheck = xcheck.riesz_representation() From 3a011756e6dd906dd56742622bb6b5c668fc6433 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 9 Dec 2025 09:40:56 +0000 Subject: [PATCH 15/22] remove white noise backend enum --- firedrake/adjoint/__init__.py | 1 + firedrake/adjoint/covariance_operator.py | 197 ++++++++---------- .../adjoint/test_covariance_operator.py | 13 +- 3 files changed, 94 insertions(+), 117 deletions(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index b534abc331..584cfb7727 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -34,6 +34,7 @@ from firedrake.adjoint.transformed_functional import L2RieszMap, L2TransformedFunctional # noqa: F401 from firedrake.adjoint.covariance_operator import ( # noqa F401 WhiteNoiseGenerator, AutoregressiveCovariance, + PyOP2NoiseBackend, PetscNoiseBackend, VOMNoiseBackend, CovarianceMatCtx, CovarianceMat, CovariancePC) import numpy_adjoint # noqa F401 import firedrake.ufl_expr diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 7922449eed..8987b56504 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -148,81 +148,6 @@ def riesz_map(self): """ -class VOMNoiseBackend(NoiseBackendBase): - """ - A PETSc based implementation of a mass matrix square root action - for generating white noise on a vertex only mesh. - """ - def __init__(self, V: WithGeometry, rng=None): - super().__init__(V, rng=rng) - self.cholesky = L2Cholesky(V) - self._zb = Function(V) - self.M = inner(self._zb, TestFunction(V))*dx - - def sample(self, *, rng=None, - tensor: Function | Cofunction | None = None, - apply_riesz: bool = False): - rng = rng or self.rng - - # z - z = rng.standard_normal(self.broken_space) - # C z - self._zb.assign(self.cholesky.C_T_inv_action(z)) - Cz = assemble(self.M) - - # Usually we would interpolate to the unbroken space, - # but here we're on a VOM so everything is broken. - # L C z - # b = Cofunction(V.dual()).interpolate(Cz) - b = Cz - - if apply_riesz: - b = b.riesz_representation(self.riesz_map) - - if tensor: - tensor.assign(b) - else: - tensor = b - - return tensor - - -class PetscNoiseBackend(NoiseBackendBase): - """ - A PETSc based implementation of a mass matrix square root action - for generating white noise. - """ - def __init__(self, V: WithGeometry, rng=None): - super().__init__(V, rng=rng) - self.cholesky = L2Cholesky(self.broken_space) - self._zb = Function(self.broken_space) - self.M = inner(self._zb, TestFunction(self.broken_space))*dx - - def sample(self, *, rng=None, - tensor: Function | Cofunction | None = None, - apply_riesz: bool = False): - V = self.function_space - rng = rng or self.rng - - # z - z = rng.standard_normal(self.broken_space) - # C z - self._zb.assign(self.cholesky.C_T_inv_action(z)) - Cz = assemble(self.M) - # L C z - b = Cofunction(V.dual()).interpolate(Cz) - - if apply_riesz: - b = b.riesz_representation(self.riesz_map) - - if tensor: - tensor.assign(b) - else: - tensor = b - - return tensor - - class PyOP2NoiseBackend(NoiseBackendBase): """ A PyOP2 based implementation of a mass matrix square root @@ -349,6 +274,81 @@ def sample(self, *, rng=None, return tensor +class PetscNoiseBackend(NoiseBackendBase): + """ + A PETSc based implementation of a mass matrix square root action + for generating white noise. + """ + def __init__(self, V: WithGeometry, rng=None): + super().__init__(V, rng=rng) + self.cholesky = L2Cholesky(self.broken_space) + self._zb = Function(self.broken_space) + self.M = inner(self._zb, TestFunction(self.broken_space))*dx + + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): + V = self.function_space + rng = rng or self.rng + + # z + z = rng.standard_normal(self.broken_space) + # C z + self._zb.assign(self.cholesky.C_T_inv_action(z)) + Cz = assemble(self.M) + # L C z + b = Cofunction(V.dual()).interpolate(Cz) + + if apply_riesz: + b = b.riesz_representation(self.riesz_map) + + if tensor: + tensor.assign(b) + else: + tensor = b + + return tensor + + +class VOMNoiseBackend(NoiseBackendBase): + """ + A PETSc based implementation of a mass matrix square root action + for generating white noise on a vertex only mesh. + """ + def __init__(self, V: WithGeometry, rng=None): + super().__init__(V, rng=rng) + self.cholesky = L2Cholesky(V) + self._zb = Function(V) + self.M = inner(self._zb, TestFunction(V))*dx + + def sample(self, *, rng=None, + tensor: Function | Cofunction | None = None, + apply_riesz: bool = False): + rng = rng or self.rng + + # z + z = rng.standard_normal(self.broken_space) + # C z + self._zb.assign(self.cholesky.C_T_inv_action(z)) + Cz = assemble(self.M) + + # Usually we would interpolate to the unbroken space, + # but here we're on a VOM so everything is broken. + # L C z + # b = Cofunction(V.dual()).interpolate(Cz) + b = Cz + + if apply_riesz: + b = b.riesz_representation(self.riesz_map) + + if tensor: + tensor.assign(b) + else: + tensor = b + + return tensor + + class WhiteNoiseGenerator: r"""Generate white noise samples. @@ -361,8 +361,8 @@ class WhiteNoiseGenerator: V : The :class:`~firedrake.functionspace.FunctionSpace` to construct a white noise sample on. - backend : WhiteNoiseGenerator.Backend - The backend specifying how to calculate and apply the mass matrix square root. + backend : + The backend to calculate and apply the mass matrix square root. rng : Initialised random number generator to use for sampling IID vectors. @@ -383,45 +383,22 @@ class WhiteNoiseGenerator: CovarianceOperatorBase """ - class Backend(Enum): - """ - The backend to implement applying the mass matrix square root. - - See Also - -------- - PyOP2NoiseBackend - PetscNoiseBackend - VOMNoiseBackend - """ - PYOP2 = 'pyop2' - PETSC = 'petsc' - VOM = 'vom' + def __init__(self, V: WithGeometry, + backend: NoiseBackendBase | None = None, rng=None): - def __init__(self, V: WithGeometry, backend=None, rng=None): # Not all backends are valid for VOM. if isinstance(V.mesh().topology, VertexOnlyMeshTopology): - backend = self.Backend(backend or self.Backend.VOM) - if backend != self.Backend.VOM: + backend = backend or VOMNoiseBackend(V, rng=rng) + if not isinstance(backend, VOMNoiseBackend): raise ValueError( - f"Cannot use white noise backend {backend} with a VertexOnlyMesh." - " Please use WhiteNoiseGenerator.Backend.VOM") + f"Cannot use white noise backend {type(backend).__name__}" + " with a VertexOnlyMesh. Please use a VOMNoiseBackend.") else: - backend = self.Backend(backend or self.Backend.PYOP2) - - backend = self.Backend(backend) - - if backend == self.Backend.PYOP2: - self.backend = PyOP2NoiseBackend(V, rng=rng) - elif backend == self.Backend.PETSC: - self.backend = PetscNoiseBackend(V, rng=rng) - elif backend == self.Backend.VOM: - self.backend = VOMNoiseBackend(V, rng=rng) - else: - raise ValueError( - f"Unrecognised white noise generation backend {backend}") + backend = backend or PyOP2NoiseBackend(V, rng=rng) - self.function_space = self.backend.function_space - self.rng = self.backend.rng + self.backend = backend + self.function_space = backend.function_space + self.rng = backend.rng petsctools.cite("Croci2018") diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index f41908f6e7..2dbfdcfa75 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -36,13 +36,10 @@ def rng(): @pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vec1", "vec2"]) @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) -@pytest.mark.parametrize("backend", ("pyop2", "petsc")) -def test_white_noise(family, degree, mesh_type, dim, backend, rng): +@pytest.mark.parametrize("backend_type", (PyOP2NoiseBackend, PetscNoiseBackend), ids=("pyop2", "petsc")) +def test_white_noise(family, degree, mesh_type, dim, backend_type, rng): """Test that white noise generator converges to a mass matrix covariance. """ - if backend == "petsc" and COMM_WORLD.size > 1: - pytest.skip( - "petsc backend for noise generation not implemented in parallel.") nx = 10 # Mesh dimension @@ -64,7 +61,8 @@ def test_white_noise(family, degree, mesh_type, dim, backend, rng): covmat = petsc2numpy_mat( assemble(M, mat_type='aij').petscmat) - generator = WhiteNoiseGenerator(V, backend=backend, rng=rng) + generator = WhiteNoiseGenerator( + V, backend=backend_type(V, rng=rng)) # Test convergence as sample size increases nsamples = [50, 100, 200, 400, 800] @@ -118,7 +116,8 @@ def test_vom_white_noise(dim, mesh_type, rng): covmat = petsc2numpy_mat( assemble(M, mat_type='aij').petscmat) - generator = WhiteNoiseGenerator(V, backend='vom', rng=rng) + backend = VOMNoiseBackend(V, rng) + generator = WhiteNoiseGenerator(V, backend=backend) # Test convergence as sample size increases nsamples = [50, 100, 200, 400, 800] From 1ef616b23589b65112dc8fd612f5195113f72856 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 9 Dec 2025 09:56:00 +0000 Subject: [PATCH 16/22] propogate seed through covariance operators and white noise generators --- firedrake/adjoint/covariance_operator.py | 40 +++++++++++++++--------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 8987b56504..b95225336e 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -56,6 +56,8 @@ class NoiseBackendBase: The :func:`~.firedrake.functionspace.FunctionSpace` to generate the samples in. rng : The ``RandomGenerator`` to generate the samples on the discontinuous superspace. + seed : + Seed for the ``RandomGenerator``. Ignored if ``rng`` is given. See Also -------- @@ -64,9 +66,10 @@ class NoiseBackendBase: WhiteNoiseGenerator """ - def __init__(self, V: WithGeometry, rng=None): + def __init__(self, V: WithGeometry, rng=None, + seed: int | None = None): self._V = V - self._rng = rng or RandomGenerator(PCG64()) + self._rng = rng or RandomGenerator(PCG64(seed=seed)) @abc.abstractmethod def sample(self, *, rng=None, @@ -153,8 +156,9 @@ class PyOP2NoiseBackend(NoiseBackendBase): A PyOP2 based implementation of a mass matrix square root for generating white noise. """ - def __init__(self, V: WithGeometry, rng=None): - super().__init__(V, rng=rng) + def __init__(self, V: WithGeometry, rng=None, + seed: int | None = None): + super().__init__(V, rng=rng, seed=seed) u = TrialFunction(V) v = TestFunction(V) @@ -279,8 +283,9 @@ class PetscNoiseBackend(NoiseBackendBase): A PETSc based implementation of a mass matrix square root action for generating white noise. """ - def __init__(self, V: WithGeometry, rng=None): - super().__init__(V, rng=rng) + def __init__(self, V: WithGeometry, rng=None, + seed: int | None = None): + super().__init__(V, rng=rng, seed=seed) self.cholesky = L2Cholesky(self.broken_space) self._zb = Function(self.broken_space) self.M = inner(self._zb, TestFunction(self.broken_space))*dx @@ -315,8 +320,9 @@ class VOMNoiseBackend(NoiseBackendBase): A PETSc based implementation of a mass matrix square root action for generating white noise on a vertex only mesh. """ - def __init__(self, V: WithGeometry, rng=None): - super().__init__(V, rng=rng) + def __init__(self, V: WithGeometry, rng=None, + seed: int | None = None): + super().__init__(V, rng=rng, seed=seed) self.cholesky = L2Cholesky(V) self._zb = Function(V) self.M = inner(self._zb, TestFunction(V))*dx @@ -365,6 +371,8 @@ class WhiteNoiseGenerator: The backend to calculate and apply the mass matrix square root. rng : Initialised random number generator to use for sampling IID vectors. + seed : + Seed for the ``RandomGenerator``. Ignored if ``rng`` is given. References ---------- @@ -384,17 +392,18 @@ class WhiteNoiseGenerator: """ def __init__(self, V: WithGeometry, - backend: NoiseBackendBase | None = None, rng=None): + backend: NoiseBackendBase | None = None, + rng=None, seed: int | None = None): # Not all backends are valid for VOM. if isinstance(V.mesh().topology, VertexOnlyMeshTopology): - backend = backend or VOMNoiseBackend(V, rng=rng) + backend = backend or VOMNoiseBackend(V, rng=rng, seed=seed) if not isinstance(backend, VOMNoiseBackend): raise ValueError( f"Cannot use white noise backend {type(backend).__name__}" " with a VertexOnlyMesh. Please use a VOMNoiseBackend.") else: - backend = backend or PyOP2NoiseBackend(V, rng=rng) + backend = backend or PyOP2NoiseBackend(V, rng=rng, seed=seed) self.backend = backend self.function_space = backend.function_space @@ -670,6 +679,8 @@ class AutoregressiveCovariance(CovarianceOperatorBase): Equal to the order of the autoregressive function kernel. rng : White noise generator to seed generating correlated samples. + seed : + Seed for the ``RandomGenerator``. Ignored if ``rng`` is given. form : AutoregressiveCovariance.DiffusionForm | ufl.Form | None The diffusion formulation or form. If a ``DiffusionForm`` then :func:`.diffusion_form` will be used to generate the diffusion @@ -715,7 +726,8 @@ class DiffusionForm(Enum): def __init__(self, V: WithGeometry, L: float | Constant, sigma: float | Constant = 1., m: int = 2, - rng: WhiteNoiseGenerator | None = None, form=None, + rng: WhiteNoiseGenerator | None = None, + seed: int | None = None, form=None, bcs: BCBase | Iterable[BCBase] | None = None, solver_parameters: dict | None = None, options_prefix: str | None = None, @@ -726,11 +738,9 @@ def __init__(self, V: WithGeometry, L: float | Constant, if isinstance(form, str): form = self.DiffusionForm(form) - self._rng = rng or WhiteNoiseGenerator(V) + self._rng = rng or WhiteNoiseGenerator(V, seed=seed) self._function_space = self.rng().function_space - if sigma <= 0: - raise ValueError("Variance must be positive.") if L < 0: raise ValueError("Correlation lengthscale must be positive.") if m < 0: From 2c2af19933641d3440d8e007080f1a3a7a424e46 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 9 Dec 2025 13:24:59 +0000 Subject: [PATCH 17/22] docstrings --- firedrake/adjoint/covariance_operator.py | 29 ++++++++++++++---------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index b95225336e..c21b04d241 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -55,9 +55,11 @@ class NoiseBackendBase: V : The :func:`~.firedrake.functionspace.FunctionSpace` to generate the samples in. rng : - The ``RandomGenerator`` to generate the samples on the discontinuous superspace. + The :mod:`RandomGenerator ` to generate the samples + on the discontinuous superspace. seed : - Seed for the ``RandomGenerator``. Ignored if ``rng`` is given. + Seed for the :mod:`RandomGenerator `. + Ignored if ``rng`` is given. See Also -------- @@ -81,8 +83,8 @@ def sample(self, *, rng=None, Parameters ---------- rng : - A ``RandomGenerator`` to use for sampling IID vectors. - If ``None`` then ``self.rng`` is used. + A :mod:`RandomGenerator ` to use for + sampling IID vectors. If ``None`` then ``self.rng`` is used. tensor : Optional location to place the result into. @@ -123,7 +125,8 @@ def function_space(self): @property def rng(self): - """The ``RandomGenerator`` to generate the IID sample on the broken function space. + """The :mod:`RandomGenerator ` to generate the + IID sample on the broken function space. """ return self._rng @@ -139,8 +142,8 @@ def riesz_map(self): Parameters ---------- rng : - A ``RandomGenerator`` to use for sampling IID vectors. - If ``None`` then ``self.rng`` is used. + A :mod:`RandomGenerator ` to use for + sampling IID vectors. If ``None`` then ``self.rng`` is used. tensor : Optional location to place the result into. @@ -372,7 +375,8 @@ class WhiteNoiseGenerator: rng : Initialised random number generator to use for sampling IID vectors. seed : - Seed for the ``RandomGenerator``. Ignored if ``rng`` is given. + Seed for the :mod:`RandomGenerator `. + Ignored if ``rng`` is given. References ---------- @@ -420,8 +424,8 @@ def sample(self, *, rng=None, Parameters ---------- rng : - A ``RandomGenerator`` to use for sampling IID vectors. - If ``None`` then ``self.rng`` is used. + A :mod:`RandomGenerator ` to use for + sampling IID vectors. If ``None`` then ``self.rng`` is used. tensor : Optional location to place the result into. @@ -680,7 +684,8 @@ class AutoregressiveCovariance(CovarianceOperatorBase): rng : White noise generator to seed generating correlated samples. seed : - Seed for the ``RandomGenerator``. Ignored if ``rng`` is given. + Seed for the :mod:`RandomGenerator `. + Ignored if ``rng`` is given. form : AutoregressiveCovariance.DiffusionForm | ufl.Form | None The diffusion formulation or form. If a ``DiffusionForm`` then :func:`.diffusion_form` will be used to generate the diffusion @@ -899,7 +904,7 @@ def diffusion_form(u, v, kappa: Constant | Function, See Also -------- - AutoregressiveCovariance + AutoregressiveCovariance.DiffusionForm """ if formulation == AutoregressiveCovariance.DiffusionForm.CG: return inner(u, v)*dx + inner(kappa*grad(u), grad(v))*dx From 755f728deb1a6ec9fe20ce5c290748b8e86019a0 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 9 Dec 2025 13:51:33 +0000 Subject: [PATCH 18/22] docstrings --- firedrake/adjoint/__init__.py | 2 +- firedrake/adjoint/covariance_operator.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py index 584cfb7727..9445626a12 100644 --- a/firedrake/adjoint/__init__.py +++ b/firedrake/adjoint/__init__.py @@ -35,7 +35,7 @@ from firedrake.adjoint.covariance_operator import ( # noqa F401 WhiteNoiseGenerator, AutoregressiveCovariance, PyOP2NoiseBackend, PetscNoiseBackend, VOMNoiseBackend, - CovarianceMatCtx, CovarianceMat, CovariancePC) + CovarianceMat, CovariancePC) import numpy_adjoint # noqa F401 import firedrake.ufl_expr import types diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index c21b04d241..2776474cfe 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -959,6 +959,16 @@ class CovarianceMatCtx: CovariancePC """ class Operation(Enum): + """ + The covariance operation to apply with this Mat. + + See Also + -------- + CovarianceOperatorBase + AutoregressiveCovariance + CovarianceMat + CovariancePC + """ ACTION = 'action' INVERSE = 'inverse' @@ -1068,6 +1078,7 @@ def CovarianceMat(covariance: CovarianceOperatorBase, CovarianceOperatorBase AutoregressiveCovariance CovarianceMatCtx + CovarianceMatCtx.Operation CovariancePC """ ctx = CovarianceMatCtx(covariance, operation=operation) @@ -1081,6 +1092,9 @@ def CovarianceMat(covariance: CovarianceOperatorBase, return mat +CovarianceMat.Operation = CovarianceMatCtx.Operation + + class CovariancePC(petsctools.PCBase): r""" A python PC context for a covariance operator. From c85414980a2334ed87896209f30a35d82fc87a85 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 9 Dec 2025 18:18:57 +0000 Subject: [PATCH 19/22] covariance review feedback --- firedrake/adjoint/covariance_operator.py | 91 +++++++++++------------- 1 file changed, 43 insertions(+), 48 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 2776474cfe..88a80f1184 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -65,6 +65,7 @@ class NoiseBackendBase: -------- PyOP2NoiseBackend PetscNoiseBackend + VOMNoiseBackend WhiteNoiseGenerator """ @@ -135,29 +136,18 @@ def riesz_map(self): """A :class:`~firedrake.cofunction.RieszMap` to cache the solver for :meth:`~firedrake.cofunction.Cofunction.riesz_representation`. """ - return RieszMap(self.function_space, constant_jacobian=True) - """ - Generate a white noise sample. - - Parameters - ---------- - rng : - A :mod:`RandomGenerator ` to use for - sampling IID vectors. If ``None`` then ``self.rng`` is used. - - tensor : - Optional location to place the result into. - - apply_riesz : - Whether to apply an L2 Riesz map to the result to return - a sample in the primal space. - """ + return RieszMap(self.function_space, "L2", constant_jacobian=True) class PyOP2NoiseBackend(NoiseBackendBase): """ A PyOP2 based implementation of a mass matrix square root for generating white noise. + + See Also + -------- + NoiseBackendBase + WhiteNoiseGenerator """ def __init__(self, V: WithGeometry, rng=None, seed: int | None = None): @@ -285,6 +275,11 @@ class PetscNoiseBackend(NoiseBackendBase): """ A PETSc based implementation of a mass matrix square root action for generating white noise. + + See Also + -------- + NoiseBackendBase + WhiteNoiseGenerator """ def __init__(self, V: WithGeometry, rng=None, seed: int | None = None): @@ -320,32 +315,26 @@ def sample(self, *, rng=None, class VOMNoiseBackend(NoiseBackendBase): """ - A PETSc based implementation of a mass matrix square root action - for generating white noise on a vertex only mesh. - """ - def __init__(self, V: WithGeometry, rng=None, - seed: int | None = None): - super().__init__(V, rng=rng, seed=seed) - self.cholesky = L2Cholesky(V) - self._zb = Function(V) - self.M = inner(self._zb, TestFunction(V))*dx + A mass matrix square root for generating white noise + on a vertex only mesh. + Notes + ----- + Computationally this is a no-op because the mass matrix + on a vertex only mesh is the identity, but we need a + consistent interface with other white noise backends. + + See Also + -------- + NoiseBackendBase + WhiteNoiseGenerator + """ def sample(self, *, rng=None, tensor: Function | Cofunction | None = None, apply_riesz: bool = False): rng = rng or self.rng - # z - z = rng.standard_normal(self.broken_space) - # C z - self._zb.assign(self.cholesky.C_T_inv_action(z)) - Cz = assemble(self.M) - - # Usually we would interpolate to the unbroken space, - # but here we're on a VOM so everything is broken. - # L C z - # b = Cofunction(V.dual()).interpolate(Cz) - b = Cz + b = rng.standard_normal(self.function_space) if apply_riesz: b = b.riesz_representation(self.riesz_map) @@ -519,17 +508,20 @@ class CovarianceOperatorBase: Inheriting classes must implement the following methods: + - ``rng`` + + - ``function_space`` + + Inheriting classes may implement the following methods (at least one + of this list must be implemented for this class to be useful): + - ``sample`` - ``apply_inverse`` - ``apply_action`` - - ``rng`` - - - ``function_space`` - - They may optionally implement ``norm`` to provide a more + They may optionally override ``norm`` to provide a more efficient implementation. See Also @@ -553,7 +545,6 @@ def function_space(self): """ raise NotImplementedError - @abc.abstractmethod def sample(self, *, rng: WhiteNoiseGenerator | None = None, tensor: Function | None = None): r""" @@ -573,7 +564,9 @@ def sample(self, *, rng: WhiteNoiseGenerator | None = None, firedrake.function.Function : The sample. """ - raise NotImplementedError + raise NotImplementedError( + "Need to implementation for sampling w~N(0, B), for" + " example by calculating w=B^{1/2}z with z~N(0, I)") def norm(self, x: Function): r"""Return the weighted norm :math:`\|x\|_{B^{-1}} = x^{T}B^{-1}x`. @@ -596,7 +589,6 @@ def norm(self, x: Function): """ return self.apply_inverse(x)(x) - @abc.abstractmethod def apply_inverse(self, x: Function, *, tensor: Cofunction | None = None): r"""Return :math:`y = B^{-1}x` where B is the covariance operator. @@ -614,9 +606,10 @@ def apply_inverse(self, x: Function, *, firedrake.cofunction.Cofunction : The result of :math:`B^{-1}x` """ - raise NotImplementedError + raise NotImplementedError( + "Inverse of B not implemented. You probably" + " also want to implement apply_action.") - @abc.abstractmethod def apply_action(self, x: Cofunction, *, tensor: Function | None = None): r"""Return :math:`y = Bx` where B is the covariance operator. @@ -635,7 +628,9 @@ def apply_action(self, x: Cofunction, *, firedrake.function.Function : The result of :math:`B^{-1}x` """ - raise NotImplementedError + raise NotImplementedError( + "Action of B not implemented. You probably" + " also want to implement apply_inverse.") class AutoregressiveCovariance(CovarianceOperatorBase): From 708822584fb5f2c1080db627464ac2c652587201 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 9 Dec 2025 20:15:59 +0000 Subject: [PATCH 20/22] fiat branch for BrokenElement fix --- firedrake/adjoint/covariance_operator.py | 15 +++------------ pyproject.toml | 2 +- .../adjoint/test_covariance_operator.py | 16 +++++++++------- 3 files changed, 13 insertions(+), 20 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 88a80f1184..41a8c20501 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -105,18 +105,9 @@ def broken_space(self): """ The discontinuous superspace containing :math:`V`, ``self.function_space``. """ - element = self.function_space.ufl_element() - mesh = self.function_space.mesh().unique() - if isinstance(element, VectorElement): - dim = element.num_sub_elements - scalar_element = element.sub_elements[0] - broken_element = BrokenElement(scalar_element) - Vbroken = VectorFunctionSpace( - mesh, broken_element, dim=dim) - else: - Vbroken = FunctionSpace( - mesh, BrokenElement(element)) - return Vbroken + return FunctionSpace( + self.function_space.mesh().unique(), + BrokenElement(self.function_space.ufl_element())) @property def function_space(self): diff --git a/pyproject.toml b/pyproject.toml index f91756ecd3..f6278048cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ # TODO RELEASE "fenics-ufl @ git+https://github.com/FEniCS/ufl.git@main", # TODO RELEASE - "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main", + "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@JHopeCollins/mixed-broken-element", "h5py>3.12.1", "immutabledict", "libsupermesh", diff --git a/tests/firedrake/adjoint/test_covariance_operator.py b/tests/firedrake/adjoint/test_covariance_operator.py index 2dbfdcfa75..40219d6d85 100644 --- a/tests/firedrake/adjoint/test_covariance_operator.py +++ b/tests/firedrake/adjoint/test_covariance_operator.py @@ -33,7 +33,7 @@ def rng(): @pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("degree", (1, 2), ids=["degree1", "degree2"]) -@pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vec1", "vec2"]) +@pytest.mark.parametrize("dim", (0, 2, (2, 2)), ids=["scalar", "vec2", "tensor22"]) @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) @pytest.mark.parametrize("backend_type", (PyOP2NoiseBackend, PetscNoiseBackend), ids=("pyop2", "petsc")) @@ -47,11 +47,11 @@ def test_white_noise(family, degree, mesh_type, dim, backend_type, rng): mesh = UnitIntervalMesh(nx) elif mesh_type == 'square': mesh = UnitSquareMesh(nx, nx) - elif mesh_type == 'cube': - mesh = UnitCubeMesh(nx, nx, nx) # Variable rank - if dim > 0: + if not isinstance(dim, int): + V = TensorFunctionSpace(mesh, family, degree, shape=dim) + elif dim > 0: V = VectorFunctionSpace(mesh, family, degree, dim=dim) else: V = FunctionSpace(mesh, family, degree) @@ -86,7 +86,7 @@ def test_white_noise(family, degree, mesh_type, dim, backend_type, rng): @pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) -@pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vec1", "vec2"]) +@pytest.mark.parametrize("dim", (0, 2, (2, 2)), ids=["scalar", "vec2", "tensor22"]) @pytest.mark.parametrize("mesh_type", ("interval", "square")) def test_vom_white_noise(dim, mesh_type, rng): """Test that white noise generator converges to a mass matrix covariance. @@ -106,7 +106,9 @@ def test_vom_white_noise(dim, mesh_type, rng): vom = VertexOnlyMesh(mesh, points) # Variable rank - if dim > 0: + if not isinstance(dim, int): + V = TensorFunctionSpace(vom, "DG", 0, shape=dim) + elif dim > 0: V = VectorFunctionSpace(vom, "DG", 0, dim=dim) else: V = FunctionSpace(vom, "DG", 0) @@ -142,7 +144,7 @@ def test_vom_white_noise(dim, mesh_type, rng): @pytest.mark.skipcomplex @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("m", (0, 2, 4)) -@pytest.mark.parametrize("dim", (0, 1, 2), ids=["scalar", "vector1", "vector2"]) +@pytest.mark.parametrize("dim", (0, 2), ids=["scalar", "vector2"]) @pytest.mark.parametrize("family", ("CG", "DG")) @pytest.mark.parametrize("mesh_type", ("interval", "square")) def test_covariance_inverse_action(m, family, mesh_type, dim): From 42a5d15e074a6645b6c51137b744d8482c165514 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 9 Dec 2025 20:18:00 +0000 Subject: [PATCH 21/22] lint --- firedrake/adjoint/covariance_operator.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index 41a8c20501..ba2f47e894 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -17,8 +17,7 @@ dx, ds, dS, sqrt, Constant, Function, Cofunction, RieszMap, TrialFunction, TestFunction, - FunctionSpace, VectorFunctionSpace, - BrokenElement, VectorElement, + FunctionSpace, BrokenElement, RandomGenerator, PCG64, LinearVariationalProblem, LinearVariationalSolver, From df622931478c422eee6b7bb234febf9d1f64d285 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 12 Dec 2025 10:40:28 +0000 Subject: [PATCH 22/22] covariance: use FunctionSpace.broken_space --- firedrake/adjoint/covariance_operator.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/firedrake/adjoint/covariance_operator.py b/firedrake/adjoint/covariance_operator.py index ba2f47e894..a41cd6d73d 100644 --- a/firedrake/adjoint/covariance_operator.py +++ b/firedrake/adjoint/covariance_operator.py @@ -17,7 +17,6 @@ dx, ds, dS, sqrt, Constant, Function, Cofunction, RieszMap, TrialFunction, TestFunction, - FunctionSpace, BrokenElement, RandomGenerator, PCG64, LinearVariationalProblem, LinearVariationalSolver, @@ -71,6 +70,7 @@ class NoiseBackendBase: def __init__(self, V: WithGeometry, rng=None, seed: int | None = None): self._V = V + self._Vb = V.broken_space() self._rng = rng or RandomGenerator(PCG64(seed=seed)) @abc.abstractmethod @@ -99,14 +99,12 @@ def sample(self, *, rng=None, """ raise NotImplementedError - @cached_property + @property def broken_space(self): """ The discontinuous superspace containing :math:`V`, ``self.function_space``. """ - return FunctionSpace( - self.function_space.mesh().unique(), - BrokenElement(self.function_space.ufl_element())) + return self._Vb @property def function_space(self):