Random Cofunctions (#4237)

Author: pbrubeck

firedrake/randomfunctiongen.py

@@ -3,15 +3,15 @@
 ========
 
 This module wraps `numpy.random <https://numpy.org/doc/stable/reference/random/index.html>`__,
-and enables users to generate a randomised :class:`.Function` from a :class:`.FunctionSpace`.
+and enables users to generate a randomised :class:`.Function` from a :class:`.WithGeometry` or :class:`.FiredrakeDualSpace`.
 This module inherits almost all attributes from `numpy.random <https://numpy.org/doc/stable/reference/random/index.html>`__ with the following changes:
 
 Generator
 ---------
 
 A :class:`.Generator` wraps `numpy.random.Generator <https://numpy.org/doc/stable/reference/random/generator.html>`__.
 :class:`.Generator` inherits almost all distribution methods from `numpy.random.Generator <https://numpy.org/doc/stable/reference/random/generator.html>`__,
-and they can be used to generate a randomised :class:`.Function` by passing a :class:`.FunctionSpace` as the first argument.
+and they can be used to generate a randomised :class:`.Function` by passing a :class:`.WithGeometry` or :class:`.FiredrakeDualSpace` as the first argument.
 
 Example:
 
@@ -106,7 +106,7 @@
 
 from firedrake.function import Function
 from pyop2.mpi import COMM_WORLD
-from ufl import FunctionSpace
+from ufl.functionspace import BaseFunctionSpace
 
 _deprecated_attributes = ['RandomGenerator', ]
 
@@ -280,7 +280,7 @@ def funcgen(c_a):
 
                     @add_doc_string(getattr(_Base, c_a).__doc__)
                     def func(self, *args, **kwargs):
-                        if len(args) > 0 and isinstance(args[0], FunctionSpace):
+                        if len(args) > 0 and isinstance(args[0], BaseFunctionSpace):
                             raise NotImplementedError("%s.%s does not take FunctionSpace as argument" % (module_attr, c_a))
                         else:
                             return getattr(super(_Wrapper, self), c_a)(*args, **kwargs)
@@ -294,7 +294,7 @@ def func(self, *args, **kwargs):
                 def funcgen(c_a):
                     @add_doc_string(getattr(_Base, c_a).__doc__)
                     def func(self, *args, **kwargs):
-                        if len(args) > 0 and isinstance(args[0], FunctionSpace):
+                        if len(args) > 0 and isinstance(args[0], BaseFunctionSpace):
                             # Extract size from V
                             if 'size' in kwargs.keys():
                                 raise TypeError("Cannot specify 'size' when generating a random function from 'V'")

tests/firedrake/adjoint/test_split_and_subfunctions.py

@@ -209,7 +209,6 @@ def test_writing_to_subfunctions():
         u.assign(kappa)
         usub *= 2
         J = assemble(inner(u, u) * dx)
-        print(f"{type(J) = }")
         rf = ReducedFunctional(J, Control(kappa), tape=tape)
     pause_annotation()
 

tests/firedrake/randomfunctiongen/test_randomfunction.py

@@ -129,3 +129,39 @@ def test_randomfunc_generator_spawn():
     assert all([child.bit_generator._seed_seq.entropy == parent.bit_generator._seed_seq.entropy for child in children])
     assert all([child.bit_generator._seed_seq.spawn_key == parent.bit_generator._seed_seq.spawn_key + (i, ) for i, child in enumerate(children)])
     assert all([child.bit_generator._seed_seq.pool_size == parent.bit_generator._seed_seq.pool_size for child in children])
+
+
+@pytest.mark.parametrize("brng", ['PCG64'])
+@pytest.mark.parametrize("meth_args", [('beta', (0.3, 0.5))])
+def test_random_cofunction(brng, meth_args):
+    meth, args = meth_args
+
+    mesh = UnitSquareMesh(10, 10)
+    V0 = VectorFunctionSpace(mesh, "CG", 1)
+    V1 = FunctionSpace(mesh, "CG", 1)
+    V = V0 * V1
+    V = V.dual()
+
+    seed = 123456789
+    # Original
+    bgen = getattr(randomgen, brng)(seed=seed)
+    if brng == 'PCG64':
+        state = bgen.state
+        state['state'] = {'state': seed, 'inc': V.comm.Get_rank()}
+        bgen.state = state
+    rg_base = randomgen.Generator(bgen)
+    # Firedrake wrapper
+    fgen = getattr(randomfunctiongen, brng)(seed=seed)
+    if brng == 'PCG64':
+        state = fgen.state
+        state['state'] = {'state': seed, 'inc': V.comm.Get_rank()}
+        fgen.state = state
+    rg_wrap = randomfunctiongen.Generator(fgen)
+    for i in range(1, 10):
+        f = getattr(rg_wrap, meth)(V, *args)
+        with f.dat.vec_ro as v:
+            if meth in ('rand', 'randn'):
+                assert np.allclose(getattr(rg_base, meth)(v.local_size), v.array[:])
+            else:
+                kwargs = {'size': (v.local_size, )}
+                assert np.allclose(getattr(rg_base, meth)(*args, **kwargs), v.array[:])