diff --git a/docs/source/conf.py b/docs/source/conf.py index 64959e17..d934de2f 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -2,6 +2,7 @@ extensions = ["autoapi.extension", "nbsphinx", + "numpydoc", "sphinx.ext.intersphinx", "sphinx_rtd_theme"] @@ -23,6 +24,8 @@ nbsphinx_execute = "auto" +numpydoc_validation_checks = {"all", "GL08"} + html_theme = "sphinx_rtd_theme" html_theme_options = {"display_version": False} diff --git a/tests/firedrake/test_block_system.py b/tests/firedrake/test_block_system.py index 06abb826..ef8ee7c6 100644 --- a/tests/firedrake/test_block_system.py +++ b/tests/firedrake/test_block_system.py @@ -4,10 +4,11 @@ from tlm_adjoint.firedrake.block_system import ( BlockMatrix, BlockNullspace, ConstantNullspace, DirichletBCNullspace, Eigensolver, LinearSolver as _BlockLinearSolver, MatrixFreeMatrix, - MatrixFunctionSolver, UnityNullspace, form_matrix) + MatrixFunctionSolver, UnityNullspace, WhiteNoiseSampler, form_matrix) from .test_base import * +import itertools import mpi4py.MPI as MPI # noqa: N817 import numbers import numpy as np @@ -604,3 +605,41 @@ def test_M_root(setup_test, test_leaks): error = var_copy(m_u_ref) var_axpy(error, -1.0, M_u) assert var_linf_norm(error) < 1.0e-15 + + +@pytest.mark.firedrake +@pytest.mark.skipif(SLEPc is None, reason="SLEPc not available") +@pytest.mark.skipif(DEFAULT_COMM.size > 1, reason="serial only") +@pytest.mark.skipif(complex_mode, reason="real only") +@pytest.mark.parametrize("precondition", [False, True]) +@seed_test +def test_white_noise_sampler(setup_test, test_leaks, + precondition): + mesh = UnitIntervalMesh(3) + space = FunctionSpace(mesh, "Lagrange", 3) + test = TestFunction(space) + trial = TrialFunction(space) + + M = assemble(inner(trial, test) * dx) + V_ref = np.zeros((space_global_size(space), space_global_size(space)), + dtype=space_dtype(space)) + for i, j in itertools.product(range(space_global_size(space)), + range(space_global_size(space))): + V_ref[i, j] = M.petscmat[i, j] + + rng = np.random.default_rng( + np.random.SeedSequence(entropy=np.random.get_state()[1][0])) + sampler = WhiteNoiseSampler( + space, rng, precondition=precondition, + solver_parameters={"mfn_tol": 1.0e-10}) + + V = np.zeros((space_global_size(space), space_global_size(space)), + dtype=space_dtype(space)) + N = 1000 + for _ in range(N): + X = var_get_values(sampler.sample()) + V += np.outer(X, X) + V /= N + error = abs(V - V_ref).max() + print(f"{error=}") + assert abs(error) < 0.015 diff --git a/tlm_adjoint/block_system.py b/tlm_adjoint/block_system.py index f916fcaa..81c3f4fa 100644 --- a/tlm_adjoint/block_system.py +++ b/tlm_adjoint/block_system.py @@ -1305,9 +1305,8 @@ def solve(self, u, v): u_petsc.to_petsc(u) v_petsc.to_petsc(v) - self.mfn.solve(u_petsc.vec, v_petsc.vec) - if self.mfn.getConvergedReason() <= 0: - raise RuntimeError("Convergence failure") + self.mfn.solve(u_petsc.vec, v_petsc.vec) + if self.mfn.getConvergedReason() <= 0: + raise RuntimeError("Convergence failure") - with paused_space_type_checking(): v_petsc.from_petsc(v) diff --git a/tlm_adjoint/firedrake/block_system.py b/tlm_adjoint/firedrake/block_system.py index 84a8207c..c54db940 100644 --- a/tlm_adjoint/firedrake/block_system.py +++ b/tlm_adjoint/firedrake/block_system.py @@ -1,8 +1,12 @@ """Firedrake specific extensions to :mod:`tlm_adjoint.block_system`. """ -from .backend import TestFunction, backend_assemble, backend_DirichletBC -from ..interface import packed, space_eq, var_axpy, var_inner, var_new +from .backend import ( + TestFunction, TrialFunction, backend_assemble, backend_DirichletBC, dx, + inner) +from ..interface import ( + packed, space_dtype, space_eq, var_axpy, var_copy, var_dtype, + var_get_values, var_inner, var_local_size, var_new, var_set_values) from ..block_system import ( BlockMatrix as _BlockMatrix, BlockNullspace, Eigensolver, @@ -12,6 +16,7 @@ from .backend_interface import assemble, matrix_multiply from .variables import Constant, Function +import numpy as np import ufl __all__ = \ @@ -34,7 +39,9 @@ "LinearSolver", "Eigensolver", - "MatrixFunctionSolver" + "MatrixFunctionSolver", + + "WhiteNoiseSampler" ] @@ -238,3 +245,117 @@ def __init__(self, A, *args, **kwargs): if isinstance(A, ufl.classes.Form): A = form_matrix(A) super().__init__(A, *args, **kwargs) + + +class WhiteNoiseSampler: + r"""White noise sampling. + + Utility class for drawing independent spatial white noise samples. + Generates a sample using + + .. math:: + + X = \Xi^{-T} \sqrt{ \Xi^T M \Xi } Z, + + where + + - :math:`M` is the mass matrix. + - :math:`\Xi` is a preconditioner. + - :math:`Z` is a vector whose elements are independent standard + Gaussian samples. + + The matrix square root is computed using SLEPc. + + Parameters + ---------- + + space : :class:`firedrake.functionspaceimpl.WithGeometry` + The function space. + rng : :class:`numpy.random._generator.Generator` + Pseudorandom number generator. + precondition : :class:`bool` + If `True` then :math:`\Xi` is set equal to the inverse of the + (principal) square root of the diagonal of :math:`M`. Otherwise it is + set equal to the identity. + M : :class:`firedrake.matrix.Matrix` + Mass matrix. Constructed by finite element assembly if not supplied. + solver_parameters : :class:`Mapping` + Solver parameters. + + Attributes + ---------- + + space : :class:`firedrake.functionspaceimpl.WithGeometry` + The function space. + rng : :class:`numpy.random._generator.Generator` + Pseudorandom number generator. + """ + + def __init__(self, space, rng, *, precondition=True, M=None, + solver_parameters=None): + if solver_parameters is None: + solver_parameters = {} + else: + solver_parameters = dict(solver_parameters) + if solver_parameters.get("mfn_type", "krylov") != "krylov": + raise ValueError("Invalid mfn_type") + if solver_parameters.get("fn_type", "sqrt") != "sqrt": + raise ValueError("Invalid fn_type") + solver_parameters.update({"mfn_type": "krylov", + "fn_type": "sqrt"}) + + if not issubclass(space_dtype(space), np.floating): + raise ValueError("Real space required") + if M is None: + test = TestFunction(space) + trial = TrialFunction(space) + M = assemble(inner(trial, test) * dx) + + if precondition: + M_diag = M.petscmat.getDiagonal() + pc = np.sqrt(M_diag.getArray(True)) + else: + pc = None + + def mult(x, y): + if pc is not None: + x = var_copy(x) + var_set_values(x, var_get_values(x) / pc) + matrix_multiply(M, x, tensor=y) + if pc is not None: + var_set_values(y, var_get_values(y) / pc) + + self._space = space + self._M = M + self._rng = rng + self._pc = pc + self._mfn = MatrixFunctionSolver( + MatrixFreeMatrix(space, space, mult), + solver_parameters=solver_parameters) + + @property + def space(self): + return self._space + + @property + def rng(self): + return self._rng + + def sample(self): + """Generate a new sample. + + Returns + ------- + X : :class:`firedrake.function.Function` or \ + :class:`firedrake.cofunction.Cofunction` + The sample. + """ + + Z = Function(self.space) + var_set_values( + Z, self.rng.standard_normal(var_local_size(Z), dtype=var_dtype(Z))) + X = Function(self.space) + self._mfn.solve(Z, X) + if self._pc is not None: + var_set_values(X, var_get_values(X) * self._pc) + return X