From c163496b6724ae6b4313319578928e94783eb3c1 Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Sat, 6 Mar 2021 23:51:04 +0000 Subject: [PATCH 1/8] Adds Gaussian random fields functionality --- firedrake/__init__.py | 1 + firedrake/randomfields.py | 511 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 512 insertions(+) create mode 100644 firedrake/randomfields.py diff --git a/firedrake/__init__.py b/firedrake/__init__.py index 0fc9aeeed6..9d148b5759 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -119,6 +119,7 @@ from firedrake.randomfunctiongen import * from firedrake.external_operators import * from firedrake.progress_bar import ProgressBar # noqa: F401 +from firedrake.randomfields import * from firedrake.logging import * # Set default log level diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py new file mode 100644 index 0000000000..3d333c599b --- /dev/null +++ b/firedrake/randomfields.py @@ -0,0 +1,511 @@ +import numpy as np + +from firedrake.assemble import vector_arg +from firedrake.constant import Constant +from firedrake.function import Function +from firedrake.functionspace import FunctionSpace, VectorFunctionSpace +from firedrake.logging import warning +from firedrake.petsc import get_petsc_variables +from firedrake.randomfunctiongen import PCG64, RandomGenerator +from firedrake.ufl_expr import TestFunction, TrialFunction +from firedrake.utility_meshes import * +from firedrake.tsfc_interface import compile_form +from firedrake.variational_solver import LinearVariationalProblem, LinearVariationalSolver + +from loopy import generate_code_v2 + +from inspect import signature + +from math import ceil, gamma + +from ufl import as_vector, BrokenElement, CellVolume, dx, inner, grad, MixedElement, SpatialCoordinate + +from pyop2 import op2 +from pyop2.mpi import COMM_WORLD + +_default_pcg = PCG64() + + +def WhiteNoise(V, rng=None): + r""" Generates a white noise sample + + :arg V: The :class: `firedrake.FunctionSpace` to construct a + white noise sample on + :arg rng: Initialised random number generator to use for obtaining + random numbers + + 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 Paper [Croci et al 2018]: + https://arxiv.org/abs/1803.04857v2 + """ + # If no random number generator provided make a new one + if rng is None: + pcg = _default_pcg + rng = RandomGenerator(pcg) + + # Create broken space for independent samples + mesh = V.mesh() + broken_elements = MixedElement([BrokenElement(Vi.ufl_element()) for Vi in V]) + Vbrok = FunctionSpace(mesh, broken_elements) + iid_normal = rng.normal(Vbrok, 0.0, 1.0) + wnoise = Function(V) + + # We also need cell volumes for correction + DG0 = FunctionSpace(mesh, 'DG', 0) + vol = Function(DG0) + vol.interpolate(CellVolume(mesh)) + + # Create mass expression, assemble and extract kernel + u = TrialFunction(V) + v = TestFunction(V) + mass = inner(u, v)*dx + mass_ker, *stuff = compile_form(mass, "mass", coffee=False) + 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" decomp and applying to broken vector + blocksize = mass_ker.kinfo.kernel.code.args[0].shape[0] + + cholesky_code = 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, + double const *__restrict__ volume) +{{ + 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 one = 1.0; + double scale = 1.0/volume[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 + if COMM_WORLD.rank == 0: + petsc_variables = get_petsc_variables() + BLASLAPACK_LIB = petsc_variables.get("BLASLAPACK_LIB", "") + BLASLAPACK_LIB = COMM_WORLD.bcast(BLASLAPACK_LIB, root=0) + BLASLAPACK_INCLUDE = petsc_variables.get("BLASLAPACK_INCLUDE", "") + BLASLAPACK_INCLUDE = COMM_WORLD.bcast(BLASLAPACK_INCLUDE, root=0) + else: + BLASLAPACK_LIB = COMM_WORLD.bcast(None, root=0) + BLASLAPACK_INCLUDE = COMM_WORLD.bcast(None, root=0) + + cholesky_kernel = op2.Kernel(cholesky_code, + "apply_cholesky", + include_dirs=BLASLAPACK_INCLUDE.split(), + ldargs=BLASLAPACK_LIB.split()) + + # Construct arguments for par loop + def get_map(x): + return x.cell_node_map() + i, _ = mass_ker.indices + + z_arg = vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) + b_arg = vector_arg(op2.INC, get_map, i, function=wnoise, V=V) + coords = mesh.coordinates + volumes = vector_arg(op2.READ, get_map, i, function=vol, V=DG0) + + op2.par_loop(cholesky_kernel, + mesh.cell_set, + z_arg, + b_arg, + coords.dat(op2.READ, get_map(coords)), + volumes) + + return wnoise + + +def PaddedMesh(constructor, *args, **kwargs): + + # Enumerate different types of supported meshes + supported_1D_constructors = [IntervalMesh, UnitIntervalMesh] + supported_2D_constructors = [RectangleMesh, SquareMesh, UnitSquareMesh] + supported_2D_immersed_constructors = [CylinderMesh] + supported_3D_constructors = [BoxMesh, CubeMesh, UnitCubeMesh] + supported_periodic_constructors = [PeriodicRectangleMesh, PeriodicSquareMesh, PeriodicUnitSquareMesh] + xyz = ['x', 'y', 'z'] + # Note that all of: + # PeriodicIntervalMesh, PeriodicUnitIntervalMesh, CircleManifoldMesh, + # IcosahedralSphereMesh, UnitIcosahedralSphereMesh, + # OctahedralSphereMesh, UnitOctahedralSphereMesh, + # CubedSphereMesh, UnitCubedSphereMesh, + # TorusMesh, + # PeriodicBoxMesh, PeriodicUnitCubeMesh + # have no boundary, so don't need an auxiliary mesh + supported_constructors = supported_1D_constructors \ + + supported_2D_constructors \ + + supported_2D_immersed_constructors \ + + supported_3D_constructors \ + + supported_periodic_constructors + assert constructor in supported_constructors + + def _pad(N, L, pad): + r""" Defines the number of additional cells, additional length + and coordinate shift in one direction + :arg N: Number of cells + :arg L: Length + :arg pad: Desired amount of padding + """ + h = L/N + extra = ceil(pad/h) + aN = N + 2*extra + aL = h*aN + shift = h*extra + return aN, aL, shift + + # Remove the pad keyword and construct the original mesh + pad = kwargs.pop('pad') + mesh = constructor(*args, **kwargs) + + # Gather the size and number of cells arguments into a dict + sig = signature(constructor) + known_args = [ + 'ncells', 'nx', 'ny', 'nz', 'nr', 'nl', + 'length_or_left', 'right', 'L', 'Lx', 'Ly', 'Lz', 'depth' + ] + constructor_keys = [var for var in sig.parameters if (var in known_args)] + arg_dict = {k: v for k, v in zip(constructor_keys, args)} + # To avoid duplicate keys move any known args out of the kwargs dict + for argument in known_args: + try: + arg_dict[argument] = kwargs.pop(argument) + except KeyError: + pass + + # Construct the padded mesh + if constructor in supported_1D_constructors: + # aux_constructor = IntervalMesh + aN, aL, shift = _pad(arg_dict.get('ncells'), + abs(arg_dict.get('right', 0) - arg_dict.get('length_or_left', 1)), + pad) + aux_mesh = IntervalMesh(aN, aL, **kwargs) + aux_mesh.coordinates.dat.data[:] = aux_mesh.coordinates.dat.data - shift + elif constructor in supported_2D_immersed_constructors: + # aux_constructor = CylinderMesh + direction = xyz.index(kwargs.get('longitudinal_direction', 'z')) + shift = [0, 0, 0] + aN, aL, s = _pad(arg_dict.get('nl'), + arg_dict.get('depth', 1), + pad) + + aux_mesh = CylinderMesh(arg_dict.get('nr'), aN, depth=aL, **kwargs) + shift[direction] = s + aux_mesh.coordinates.dat.data[:, :] = aux_mesh.coordinates.dat.data - shift + elif constructor in supported_periodic_constructors: + aux_constructor = PeriodicRectangleMesh + if kwargs.get('direction') == 'both': + # Has no boundary! + return mesh, None + else: + c = kwargs.get('direction') + direction = (xyz.index(c) + 1) % mesh.topological_dimension() + aN = [arg_dict.get('nx'), arg_dict.get('ny')] + aL = [ + arg_dict.get('Lx', arg_dict.get('L', 1)), + arg_dict.get('Ly', arg_dict.get('L', 1)) + ] + shift = [0, 0] + n, l, s = _pad(arg_dict.get('n' + c), + arg_dict.get('L' + c, arg_dict.get('L', 1)), + pad) + aN[direction] = n + aL[direction] = l + shift[direction] = s + aux_mesh = PeriodicRectangleMesh(*aN, *aL, **kwargs) + aux_mesh.coordinates.dat.data[:, :] = aux_mesh.coordinates.dat.data - shift + else: + # We can handle rectangles and boxes together + if mesh.topological_dimension() == 2: + aux_constructor = RectangleMesh + elif mesh.topological_dimension() == 3: + aux_constructor = BoxMesh + aN = [] + aL = [] + shift = [] + for _, c in zip(range(mesh.topological_dimension()), xyz): + n, l, s = _pad(arg_dict.get('n' + c), + arg_dict.get('L' + c, arg_dict.get('L', 1)), + pad) + aN.append(n) + aL.append(l) + shift.append(s) + aux_mesh = aux_constructor(*aN, *aL, **kwargs) + aux_mesh.coordinates.dat.data[:, :] = aux_mesh.coordinates.dat.data - shift + + return mesh, aux_mesh + + +def TrimFunction(f, mesh): + r"""Really don't like using this function, but can't think of better method + Cuts a function to the provided mesh + """ + V = f.ufl_function_space() + + W = FunctionSpace(mesh, V.ufl_element()) + g = Function(W) + + coordspace = VectorFunctionSpace(mesh, V.ufl_element()) + ho_coords = Function(coordspace) + ho_coords.interpolate(as_vector(SpatialCoordinate(mesh))) + + g.dat.data[:] = f.at(ho_coords.dat.data) + return g + + +class GaussianRF: + r"""The distribution is a spatially varying Gaussian Random Field + with Matern covariance, GRF(mu, C): + mu is the mean and C is the covariance kernel: + sigma**2 + C(x, y) = ------------------------- * (kappa * r)**nu * K_nu(kappa r) + 2**(nu - 1) * Gamma(nu) + where + Gamma is the gamma function, + K_nu is the modified Bessel function of the second kind + r = ||x-y||_2, + sqrt(8 * nu) + kappa = --------------, + lambda + sigma is the variance, + nu is the desired smoothness, + lambda is the correlation length + + A realisation of a GRF is generated by calling the sample method of this class + """ + def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, + rng=None, solver_parameters=None, V_aux=None): + r"""Creates a spatially varying Gaussian Random Field with Matern covariance + :arg V: The :class: `firedrake.FunctionSpace` to construct a + white noise sample on + :arg mu: The mean of the distribution + :arg sigma: The standard deviation of the distribution + :arg smoothness: The smoothness of the functions from the distribution, + sampled functions will have *at least* the smoothness specified here + :arg correlation_length: The length scale over which the sampled + function will vary + :arg rng: Initialised random number generator to use for obtaining + random numbers, used for seeding the distribution + :arg solver_parameters: Override the solver parameters used in constructing + the random field. These default to {'ksp_type': 'cg', 'pc_type': 'gamg'} + """ + # Check for sensible arguments + assert sigma > 0 + assert smoothness > 0 + assert correlation_length > 0 + + self.mesh = V.mesh() + if self.mesh.coordinates.exterior_facet_node_map().values.size != 0 and V_aux is None: + warning( + r"""The mesh that the provided function space is defined on +has a boundary, but no auxilliary mesh has been provided. To prevent boundary +pollution effects you should provide a function space defined on a mesh that +is at least the correlation length larger on all boundary edges to avoid +boundary pollution effects. + """) + self.trueV = V + self.V = V_aux if V_aux is not None else V + self.mu = mu + self.sigma = sigma + + # Set symbols to match + self.dim = V.mesh().topological_dimension() + self.nu = smoothness + self.lambd = correlation_length + self.k = ceil((self.nu + self.dim/2)/2) + + # Calculate additional parameters + self.kappa = np.sqrt(8*self.nu)/self.lambd + sigma_hat2 = gamma(self.nu)*self.nu**(self.dim/2) + sigma_hat2 /= gamma(self.nu + self.dim/2) + sigma_hat2 *= (2/np.pi)**(self.dim/2) + sigma_hat2 *= self.lambd**(-self.dim) + self.eta = self.sigma/np.sqrt(sigma_hat2) + + # Setup RNG if provided + self.rng = rng + + # Setup modified Helmholtz problem + u = TrialFunction(self.V) + v = TestFunction(self.V) + self.wnoise = Function(self.V) + a = (inner(u, v) + Constant(1/(self.kappa**2))*inner(grad(u), grad(v)))*dx + l = Constant(self.eta)*inner(self.wnoise, v)*dx + + # Solve problem once + self.u_h = Function(self.V) + if solver_parameters is None: + self.solver_param = {'ksp_type': 'cg', 'pc_type': 'gamg'} + else: + self.solver_param = solver_parameters + base_problem = LinearVariationalProblem(a, l, self.u_h) + self.base_solver = LinearVariationalSolver( + base_problem, + solver_parameters=self.solver_param + ) + + # For smoother solutions we must iterate this solve + if self.k > 1: + self.u_j = Function(self.V) + l_j = inner(self.u_j, v)*dx + problem = LinearVariationalProblem(a, l_j, self.u_h) + self.solver = LinearVariationalSolver( + problem, + solver_parameters=self.solver_param + ) + + def get_parameters(self): + r""" Returns a dict of parameters used to specify the distribution + """ + rf_parameters = {'mu': self.mu, + 'sigma': self.sigma, + 'smoothness': self.smoothness, + 'correlation_length': self.correlation_length, + 'rng': self.rng + } + # ~ print('dimension:', d,self. 'iterations:', k) + # ~ print('kappa:', kappa, 'sigma hat squared:', sigma_hat2, 'eta:', eta) + # ~ print('kappa**-2:', 1/(kappa**2)) + return rf_parameters + + def sample(self, rng=None): + r"""Returns a realisation of a GRF with parameters specified + :arg rng: Initialised random number generator to use for obtaining + random numbers, used for seeding the distribution, will override + the generator specified at initialisation for the given sample. + """ + # If no random number generator provided use default + if rng is None: + if self.rng is None: + pcg = _default_pcg + rng = RandomGenerator(pcg) + else: + rng = self.rng + + # Generate a new white noise sample + self.wnoise.assign(WhiteNoise(self.V, rng)) + + # Solve + self.base_solver.solve() + + # Iterate solve until required smoothness achieved + if self.k > 1: + self.u_j.assign(self.u_h) + for _ in range(self.k - 1): + self.solver.solve() + self.u_j.assign(self.u_h) + + # Shift the function by the mean and return + self.u_h.dat.data[:] = self.u_h.dat.data + self.mu + + # Remove boundary + if self.V is not self.trueV: + sample = TrimFunction(self.u_h, self.mesh) + else: + sample = self.u_h + return sample + + +class LogGaussianRF(GaussianRF): + r"""The distribution is a spatially varying Log Gaussian Random Field + with Matern covariance LGRF(mu, C): + mu is the mean and C is the covariance kernel of the + _logarithm_ of the distribution: + sigma**2 + C(x, y) = ------------------------- * (kappa * r)**nu * K_nu(kappa r) + 2**(nu - 1) * Gamma(nu) + where + Gamma is the gamma function, + K_nu is the modified Bessel function of the second kind + r = ||x-y||_2, + sqrt(8 * nu) + kappa = --------------, + lambda + sigma is the variance, + nu is the desired smoothness, + lambda is the correlation length + + A realisation of a LGRF is generated by calling the sample method of this class + """ + def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, + rng=None, solver_parameters=None, mean=None, std_dev=None, V_aux=None): + r'''Creates a spatially varying Log Gaussian Random Field with Matern covariance + :arg V: The :class: `firedrake.FunctionSpace` to construct a + white noise sample on + :arg mu: The mean of the _logarithm_ of the distribution + :arg sigma: The standard deviation of the _logarithm_ of the distribution + :arg smoothness: The smoothness of the functions from the distribution, + sampled functions will have *at least* the smoothness specified here + :arg correlation_length: The length scale over which the sampled + function will vary + :arg rng: Initialised random number generator to use for obtaining + random numbers, used for seeding the distribution + :arg solver_parameters: Override the solver parameters used in constructing + the random field. These default to {'ksp_type': 'cg', 'pc_type': 'gamg'} + It is also possible to specify the mean and standard deviation of + the distribution directly by setting: + :arg mean: The mean of the distribution + :arg std_dev: The standard deviation of the distribution + Both of the above parameters must be set to completely specify the + distribution. If set the parameters mu and sigma will be ignored. + ''' + if (mean is not None) and (std_dev is not None): + mu = np.log(mean**2/np.sqrt(std_dev**2 + mean**2)) + sigma = np.sqrt(np.log(1 + (std_dev/mean)**2)) + elif (mean is not None) and (std_dev is not None): + raise ValueError(r"""Specify either: + - Both `mu` and `sigma`, which the expected value and variance of the logarithm of the random field respectively. + - Both `mean` and `std_dev`, which specify the distribution of the random field respectively. + """) + super().__init__(V, mu, sigma, smoothness, correlation_length, rng=None, V_aux=None) + + def sample(self, rng=None): + r"""Returns a realisation of a LGRF with parameters specified + :arg rng: Initialised random number generator to use for obtaining + random numbers, used for seeding the distribution, will override + the generator specified at initialisation for the given sample. + """ + super().sample(rng) + self.u_h.dat.data[:] = np.exp(self.u_h.dat.data) + return self.u_h From e9c67c413393083fa98efeea65b07ece03acd402 Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Sun, 7 Mar 2021 15:36:15 +0000 Subject: [PATCH 2/8] Fixed some small mistakes, trim doesn't work in parallel --- firedrake/randomfields.py | 39 +++++++++++++++++++++++++-------------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index 3d333c599b..6e504a44ca 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -157,8 +157,14 @@ def get_map(x): return wnoise -def PaddedMesh(constructor, *args, **kwargs): - +def PaddedMesh(constructor, *args, pad=0.2, **kwargs): + r"""A mesh as specified and an additional mesh with padding on all + of its boundaries. + :arg constructor: The utility function used to construct a a mesh + :arg args: Arguments used to construct the mesh + :arg pad: Padding to add to all boundaries + :arg kwargs: Additional keyword arguments used for constructing the mesh + """ # Enumerate different types of supported meshes supported_1D_constructors = [IntervalMesh, UnitIntervalMesh] supported_2D_constructors = [RectangleMesh, SquareMesh, UnitSquareMesh] @@ -196,7 +202,7 @@ def _pad(N, L, pad): return aN, aL, shift # Remove the pad keyword and construct the original mesh - pad = kwargs.pop('pad') + # ~ pad = kwargs.pop('pad') mesh = constructor(*args, **kwargs) # Gather the size and number of cells arguments into a dict @@ -279,7 +285,7 @@ def _pad(N, L, pad): def TrimFunction(f, mesh): r"""Really don't like using this function, but can't think of better method - Cuts a function to the provided mesh + Cuts a function to the provided mesh, doesn't work in parallel """ V = f.ufl_function_space() @@ -290,7 +296,7 @@ def TrimFunction(f, mesh): ho_coords = Function(coordspace) ho_coords.interpolate(as_vector(SpatialCoordinate(mesh))) - g.dat.data[:] = f.at(ho_coords.dat.data) + g.dat.data[:] = f.at(ho_coords.dat.data, tolerance=1e-8) return g @@ -339,10 +345,12 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, if self.mesh.coordinates.exterior_facet_node_map().values.size != 0 and V_aux is None: warning( r"""The mesh that the provided function space is defined on -has a boundary, but no auxilliary mesh has been provided. To prevent boundary -pollution effects you should provide a function space defined on a mesh that -is at least the correlation length larger on all boundary edges to avoid -boundary pollution effects. +has a boundary, but no auxiliary function space has been provided. To +prevent boundary pollution effects you should provide a function space V_aux +defined on a mesh that is at least the correlation length larger on all +boundary edges to avoid boundary pollution effects. The utility function +PaddedMesh is provided to gennerate such meshes for some of Firedrake's +utility meshes. """) self.trueV = V self.V = V_aux if V_aux is not None else V @@ -361,7 +369,9 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, sigma_hat2 /= gamma(self.nu + self.dim/2) sigma_hat2 *= (2/np.pi)**(self.dim/2) sigma_hat2 *= self.lambd**(-self.dim) - self.eta = self.sigma/np.sqrt(sigma_hat2) + # The factor of 0.5 doesn't appear in the paper, but numerical + # experiments show sample variance is out by a factor 2 + self.eta = 0.5*self.sigma/np.sqrt(sigma_hat2) # Setup RNG if provided self.rng = rng @@ -498,7 +508,8 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, - Both `mu` and `sigma`, which the expected value and variance of the logarithm of the random field respectively. - Both `mean` and `std_dev`, which specify the distribution of the random field respectively. """) - super().__init__(V, mu, sigma, smoothness, correlation_length, rng=None, V_aux=None) + super().__init__(V, mu, sigma, smoothness, correlation_length, + rng=rng, solver_parameters=solver_parameters, V_aux=V_aux) def sample(self, rng=None): r"""Returns a realisation of a LGRF with parameters specified @@ -506,6 +517,6 @@ def sample(self, rng=None): random numbers, used for seeding the distribution, will override the generator specified at initialisation for the given sample. """ - super().sample(rng) - self.u_h.dat.data[:] = np.exp(self.u_h.dat.data) - return self.u_h + sample = super().sample(rng) + sample.dat.data[:] = np.exp(sample.dat.data) + return sample From aa39ac7435f19c86a6213d34712821dcbc6b310c Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Tue, 11 May 2021 17:32:32 +0100 Subject: [PATCH 3/8] Use private function _vector_arg from assemble --- firedrake/randomfields.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index 6e504a44ca..902de2940a 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -1,6 +1,6 @@ import numpy as np -from firedrake.assemble import vector_arg +from firedrake.assemble import _vector_arg from firedrake.constant import Constant from firedrake.function import Function from firedrake.functionspace import FunctionSpace, VectorFunctionSpace @@ -142,10 +142,10 @@ def get_map(x): return x.cell_node_map() i, _ = mass_ker.indices - z_arg = vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) - b_arg = vector_arg(op2.INC, get_map, i, function=wnoise, V=V) + z_arg = _vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) + b_arg = _vector_arg(op2.INC, get_map, i, function=wnoise, V=V) coords = mesh.coordinates - volumes = vector_arg(op2.READ, get_map, i, function=vol, V=DG0) + volumes = _vector_arg(op2.READ, get_map, i, function=vol, V=DG0) op2.par_loop(cholesky_kernel, mesh.cell_set, From e6a8868e0d3e2d59520e06e254f6fdd0bb94c3ad Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Tue, 8 Jun 2021 14:35:09 +0100 Subject: [PATCH 4/8] Fixing up mistakes in scaling --- firedrake/randomfields.py | 30 ++++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index 902de2940a..898358a789 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -1,15 +1,16 @@ import numpy as np -from firedrake.assemble import _vector_arg +from firedrake.assemble import assemble, _vector_arg from firedrake.constant import Constant from firedrake.function import Function from firedrake.functionspace import FunctionSpace, VectorFunctionSpace from firedrake.logging import warning from firedrake.petsc import get_petsc_variables from firedrake.randomfunctiongen import PCG64, RandomGenerator +from firedrake.solving import solve +from firedrake.tsfc_interface import compile_form from firedrake.ufl_expr import TestFunction, TrialFunction from firedrake.utility_meshes import * -from firedrake.tsfc_interface import compile_form from firedrake.variational_solver import LinearVariationalProblem, LinearVariationalSolver from loopy import generate_code_v2 @@ -68,7 +69,8 @@ def WhiteNoise(V, rng=None): "static void " + mass_ker.kinfo.kernel.name) # Add custom code for doing "Cholesky" decomp and applying to broken vector - blocksize = mass_ker.kinfo.kernel.code.args[0].shape[0] + name = mass_ker.kinfo.kernel.name + blocksize = mass_ker.kinfo.kernel.code[name].args[0].shape[0] cholesky_code = f""" extern void dpotrf_(char *UPLO, @@ -105,7 +107,8 @@ def WhiteNoise(V, rng=None): char trans[1]; int32_t stride = 1; //double one = 1.0; - double scale = 1.0/volume[0]; + double scale = 1.0;//volume[0]; + //printf("%g\\n", scale); double zero = 0.0; {mass_ker.kinfo.kernel.name}(H, coords); @@ -371,7 +374,7 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, sigma_hat2 *= self.lambd**(-self.dim) # The factor of 0.5 doesn't appear in the paper, but numerical # experiments show sample variance is out by a factor 2 - self.eta = 0.5*self.sigma/np.sqrt(sigma_hat2) + self.eta = self.sigma/np.sqrt(sigma_hat2) # Setup RNG if provided self.rng = rng @@ -381,7 +384,8 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, v = TestFunction(self.V) self.wnoise = Function(self.V) a = (inner(u, v) + Constant(1/(self.kappa**2))*inner(grad(u), grad(v)))*dx - l = Constant(self.eta)*inner(self.wnoise, v)*dx + self.A = assemble(a) + # ~ b = assemble(self.eta*self.wnoise*dx) # Solve problem once self.u_h = Function(self.V) @@ -389,11 +393,11 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, self.solver_param = {'ksp_type': 'cg', 'pc_type': 'gamg'} else: self.solver_param = solver_parameters - base_problem = LinearVariationalProblem(a, l, self.u_h) - self.base_solver = LinearVariationalSolver( - base_problem, - solver_parameters=self.solver_param - ) + # ~ base_problem = LinearVariationalProblem(a, l, self.u_h) + # ~ self.base_solver = LinearVariationalSolver( + # ~ base_problem, + # ~ solver_parameters=self.solver_param + # ~ ) # For smoother solutions we must iterate this solve if self.k > 1: @@ -435,9 +439,11 @@ def sample(self, rng=None): # Generate a new white noise sample self.wnoise.assign(WhiteNoise(self.V, rng)) + b = assemble(self.eta*self.wnoise) # Solve - self.base_solver.solve() + solve(self.A, self.u_h, b, solver_parameters=self.solver_param) + # ~ self.base_solver.solve() # Iterate solve until required smoothness achieved if self.k > 1: From 2653b02675c7dbb467260ea4e20c537a6a9c7ee7 Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Thu, 22 Jul 2021 18:28:53 +0100 Subject: [PATCH 5/8] Reuse solver object for effiecieny, remove commented out blocks --- firedrake/randomfields.py | 56 +++++++++++++++++---------------------- 1 file changed, 24 insertions(+), 32 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index 898358a789..716f33ee83 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -2,12 +2,13 @@ from firedrake.assemble import assemble, _vector_arg from firedrake.constant import Constant +from firedrake.dmhooks import add_hooks from firedrake.function import Function from firedrake.functionspace import FunctionSpace, VectorFunctionSpace +from firedrake.linear_solver import LinearSolver from firedrake.logging import warning from firedrake.petsc import get_petsc_variables from firedrake.randomfunctiongen import PCG64, RandomGenerator -from firedrake.solving import solve from firedrake.tsfc_interface import compile_form from firedrake.ufl_expr import TestFunction, TrialFunction from firedrake.utility_meshes import * @@ -19,7 +20,7 @@ from math import ceil, gamma -from ufl import as_vector, BrokenElement, CellVolume, dx, inner, grad, MixedElement, SpatialCoordinate +from ufl import as_vector, BrokenElement, dx, inner, grad, MixedElement, SpatialCoordinate from pyop2 import op2 from pyop2.mpi import COMM_WORLD @@ -27,13 +28,14 @@ _default_pcg = PCG64() -def WhiteNoise(V, rng=None): +def WhiteNoise(V, rng=None, scale=1.0): r""" Generates a white noise sample :arg V: The :class: `firedrake.FunctionSpace` to construct a white noise sample on :arg rng: Initialised random number generator to use for obtaining random numbers + :arg scale: Multiplicative scale factor for the white noise Returns a :firedrake.Function: with b ~ Normal(0, M) @@ -54,11 +56,6 @@ def WhiteNoise(V, rng=None): iid_normal = rng.normal(Vbrok, 0.0, 1.0) wnoise = Function(V) - # We also need cell volumes for correction - DG0 = FunctionSpace(mesh, 'DG', 0) - vol = Function(DG0) - vol.interpolate(CellVolume(mesh)) - # Create mass expression, assemble and extract kernel u = TrialFunction(V) v = TestFunction(V) @@ -95,8 +92,7 @@ def WhiteNoise(V, rng=None): void apply_cholesky(double *__restrict__ z, double *__restrict__ b, - double const *__restrict__ coords, - double const *__restrict__ volume) + double const *__restrict__ coords) {{ char uplo[1]; int32_t N = {blocksize}, LDA = {blocksize}, INFO = 0; @@ -106,9 +102,7 @@ def WhiteNoise(V, rng=None): char trans[1]; int32_t stride = 1; - //double one = 1.0; - double scale = 1.0;//volume[0]; - //printf("%g\\n", scale); + double scale = {scale}; double zero = 0.0; {mass_ker.kinfo.kernel.name}(H, coords); @@ -148,14 +142,13 @@ def get_map(x): z_arg = _vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) b_arg = _vector_arg(op2.INC, get_map, i, function=wnoise, V=V) coords = mesh.coordinates - volumes = _vector_arg(op2.READ, get_map, i, function=vol, V=DG0) op2.par_loop(cholesky_kernel, mesh.cell_set, z_arg, b_arg, - coords.dat(op2.READ, get_map(coords)), - volumes) + coords.dat(op2.READ, get_map(coords)) + ) return wnoise @@ -372,8 +365,6 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, sigma_hat2 /= gamma(self.nu + self.dim/2) sigma_hat2 *= (2/np.pi)**(self.dim/2) sigma_hat2 *= self.lambd**(-self.dim) - # The factor of 0.5 doesn't appear in the paper, but numerical - # experiments show sample variance is out by a factor 2 self.eta = self.sigma/np.sqrt(sigma_hat2) # Setup RNG if provided @@ -385,7 +376,6 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, self.wnoise = Function(self.V) a = (inner(u, v) + Constant(1/(self.kappa**2))*inner(grad(u), grad(v)))*dx self.A = assemble(a) - # ~ b = assemble(self.eta*self.wnoise*dx) # Solve problem once self.u_h = Function(self.V) @@ -393,11 +383,15 @@ def __init__(self, V, mu=0, sigma=1, smoothness=1, correlation_length=0.2, self.solver_param = {'ksp_type': 'cg', 'pc_type': 'gamg'} else: self.solver_param = solver_parameters - # ~ base_problem = LinearVariationalProblem(a, l, self.u_h) - # ~ self.base_solver = LinearVariationalSolver( - # ~ base_problem, - # ~ solver_parameters=self.solver_param - # ~ ) + + self.base_solver = LinearSolver( + self.A, + solver_parameters=self.solver_param + ) + # We need the appctx if we want to perform geometric multigrid + lvproblem = LinearVariationalProblem(a, v*dx, Function(self.V)) + lvsolver = LinearVariationalSolver(lvproblem) + self.lvs_ctx = lvsolver._ctx # For smoother solutions we must iterate this solve if self.k > 1: @@ -418,9 +412,6 @@ def get_parameters(self): 'correlation_length': self.correlation_length, 'rng': self.rng } - # ~ print('dimension:', d,self. 'iterations:', k) - # ~ print('kappa:', kappa, 'sigma hat squared:', sigma_hat2, 'eta:', eta) - # ~ print('kappa**-2:', 1/(kappa**2)) return rf_parameters def sample(self, rng=None): @@ -438,12 +429,13 @@ def sample(self, rng=None): rng = self.rng # Generate a new white noise sample - self.wnoise.assign(WhiteNoise(self.V, rng)) - b = assemble(self.eta*self.wnoise) + b = WhiteNoise(self.V, rng, scale=self.eta) - # Solve - solve(self.A, self.u_h, b, solver_parameters=self.solver_param) - # ~ self.base_solver.solve() + # Solve adding an appctx from the equivalent linear variational problem + ksp = self.base_solver.ksp + dm = ksp.getDM() + with add_hooks(dm, ksp, appctx=self.lvs_ctx, save=False): + self.base_solver.solve(self.u_h, b) # Iterate solve until required smoothness achieved if self.k > 1: From ac05cd23c83357664d00e199b9412fc06103161c Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Tue, 3 Aug 2021 15:11:11 +0100 Subject: [PATCH 6/8] Corrected parameters --- firedrake/randomfields.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index 716f33ee83..b7959c6db0 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -408,8 +408,8 @@ def get_parameters(self): """ rf_parameters = {'mu': self.mu, 'sigma': self.sigma, - 'smoothness': self.smoothness, - 'correlation_length': self.correlation_length, + 'smoothness': self.nu, + 'correlation_length': self.lambd, 'rng': self.rng } return rf_parameters From f31a813200eed744bfd73aecb6a573c369015e00 Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Mon, 14 Oct 2024 15:22:41 +0100 Subject: [PATCH 7/8] Update code for changes --- firedrake/randomfields.py | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index b7959c6db0..d4a8af8392 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -1,6 +1,6 @@ import numpy as np -from firedrake.assemble import assemble, _vector_arg +from firedrake.assemble import assemble from firedrake.constant import Constant from firedrake.dmhooks import add_hooks from firedrake.function import Function @@ -20,10 +20,14 @@ from math import ceil, gamma -from ufl import as_vector, BrokenElement, dx, inner, grad, MixedElement, SpatialCoordinate +from ufl import as_vector, dx, inner, grad, SpatialCoordinate +from finat.ufl import BrokenElement, MixedElement + +# ~ from ufl.finiteelement import MixedElement from pyop2 import op2 from pyop2.mpi import COMM_WORLD +from pyop2.parloop import DatParloopArg _default_pcg = PCG64() @@ -60,7 +64,7 @@ def WhiteNoise(V, rng=None, scale=1.0): u = TrialFunction(V) v = TestFunction(V) mass = inner(u, v)*dx - mass_ker, *stuff = compile_form(mass, "mass", coffee=False) + 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) @@ -135,19 +139,29 @@ def WhiteNoise(V, rng=None, scale=1.0): ldargs=BLASLAPACK_LIB.split()) # Construct arguments for par loop - def get_map(x): - return x.cell_node_map() + # ~ def get_map(x): + # ~ return x.cell_node_map() i, _ = mass_ker.indices - z_arg = _vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) - b_arg = _vector_arg(op2.INC, get_map, i, function=wnoise, V=V) + # ~ (access, get_map, i, *, function, V) + # ~ if i is None: + # ~ map_ = get_map(V) + # ~ return function.dat(access, map_) + # ~ else: + # ~ map_ = get_map(V[i]) + # ~ return function.dat[i](access, map_) + + # ~ z_arg = _vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) + z_arg = iid_normal.dat(op2.READ, Vbrok.cell_node_map()) + # ~ b_arg = _vector_arg(op2.INC, get_map, i, function=wnoise, V=V) + b_arg = wnoise.dat(op2.INC, V.cell_node_map()) coords = mesh.coordinates op2.par_loop(cholesky_kernel, mesh.cell_set, z_arg, b_arg, - coords.dat(op2.READ, get_map(coords)) + coords.dat(op2.READ, coords.cell_node_map()) ) return wnoise From 4d89be36a1074cd9cde3417e268af2b26359776a Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Tue, 15 Oct 2024 22:53:26 +0100 Subject: [PATCH 8/8] Lint --- firedrake/randomfields.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/firedrake/randomfields.py b/firedrake/randomfields.py index d4a8af8392..2a63968ea6 100644 --- a/firedrake/randomfields.py +++ b/firedrake/randomfields.py @@ -23,11 +23,8 @@ from ufl import as_vector, dx, inner, grad, SpatialCoordinate from finat.ufl import BrokenElement, MixedElement -# ~ from ufl.finiteelement import MixedElement - from pyop2 import op2 from pyop2.mpi import COMM_WORLD -from pyop2.parloop import DatParloopArg _default_pcg = PCG64() @@ -138,22 +135,9 @@ def WhiteNoise(V, rng=None, scale=1.0): include_dirs=BLASLAPACK_INCLUDE.split(), ldargs=BLASLAPACK_LIB.split()) - # Construct arguments for par loop - # ~ def get_map(x): - # ~ return x.cell_node_map() i, _ = mass_ker.indices - # ~ (access, get_map, i, *, function, V) - # ~ if i is None: - # ~ map_ = get_map(V) - # ~ return function.dat(access, map_) - # ~ else: - # ~ map_ = get_map(V[i]) - # ~ return function.dat[i](access, map_) - - # ~ z_arg = _vector_arg(op2.READ, get_map, i, function=iid_normal, V=Vbrok) z_arg = iid_normal.dat(op2.READ, Vbrok.cell_node_map()) - # ~ b_arg = _vector_arg(op2.INC, get_map, i, function=wnoise, V=V) b_arg = wnoise.dat(op2.INC, V.cell_node_map()) coords = mesh.coordinates