Support noise in the general PDE class

pde/fields/base.py

@@ -193,7 +193,8 @@ def _write_to_image(self, filename: str, **kwargs):
 
 
     @abstractmethod
-    def copy(self: TField, data=None, label: str = None) -> TField: pass
+    def copy(self: TField, data: OptionalArrayLike = None, label: str = None) \
+        -> TField: pass
 
 
     def assert_field_compatible(self, other: 'FieldBase',
@@ -794,8 +795,9 @@ def from_state(cls, attributes: Dict[str, Any],
         return cls(attributes.pop('grid'), data=data, **attributes)
 
 
-    def copy(self: TDataField, data: np.ndarray = None, label: str = None) \
-            -> TDataField:
+    def copy(self: TDataField,
+             data: OptionalArrayLike = None,
+             label: str = None) -> TDataField:
         """ return a copy of the data, but not of the grid
         
         Args:

pde/fields/collection.py

@@ -11,7 +11,7 @@
 
 import numpy as np
 
-from .base import FieldBase, DataFieldBase
+from .base import FieldBase, DataFieldBase, OptionalArrayLike
 from .scalar import ScalarField
 from ..grids.base import GridBase
 
@@ -22,7 +22,7 @@ class FieldCollection(FieldBase):
 
 
     def __init__(self, fields: Sequence[DataFieldBase],
-                 data=None,
+                 data: OptionalArrayLike = None,
                  copy_fields: bool = False,
                  label: Optional[str] = None):
         """ 
@@ -35,7 +35,8 @@ def __init__(self, fields: Sequence[DataFieldBase],
             copy_fields (bool):
                 Flag determining whether the individual fields given in `fields`
                 are copied. 
-            label (str): Label of the field collection
+            label (str):
+                Label of the field collection
             
         Warning:
             If `data` is given and `copy_fields == False`, the data in the
@@ -67,7 +68,7 @@ def __init__(self, fields: Sequence[DataFieldBase],
         for field in self.fields:
             if not isinstance(field, DataFieldBase):
                 raise RuntimeError('Individual fields must be of type '
-                                   'DataFieldBase. FieldCollections cannot be '
+                                   'DataFieldBase. Field collections cannot be '
                                    'nested')
             start = len(new_data)
             this_data = field._data_flat
@@ -78,15 +79,15 @@ def __init__(self, fields: Sequence[DataFieldBase],
         # combine into one data field
         data_shape = (dof,) + grid.shape
         if data is None:
-            data = np.array(new_data, dtype=np.double)
+            data_arr = np.array(new_data, dtype=np.double)
         else:
-            data = np.asarray(data, dtype=np.double)
-            if data.shape != data_shape:
-                data = np.array(np.broadcast_to(data, data_shape))
-        assert data.shape == data_shape
+            data_arr = np.asarray(data, dtype=np.double)
+            if data_arr.shape != data_shape:
+                data_arr = np.array(np.broadcast_to(data_arr, data_shape))
+        assert data_arr.shape == data_shape
 
         # initialize the class
-        super().__init__(grid, data, label=label)        
+        super().__init__(grid, data_arr, label=label)        
 
         # link the data of the original fields back to self._data if they were
         # not copied
@@ -307,14 +308,16 @@ def unserialize_attributes(cls, attributes: Dict[str, str]) \
         return results
 
 
-    def copy(self, data=None, label: str = None) -> 'FieldCollection':
+    def copy(self, data: OptionalArrayLike = None, label: str = None) \
+            -> 'FieldCollection':
         """ return a copy of the data, but not of the grid
         
         Args:
-            data (:class:`numpy.ndarray`, optional): Data values at the support
-                points of the grid that define the field. Note that the data is
-                not copied but used directly.
-            label (str, optional): Name of the copied field
+            data (:class:`numpy.ndarray`, optional):
+                Data values at the support points of the grid that define the
+                field. Note that the data is not copied but used directly.
+            label (str, optional):
+                Name of the copied field
         """
         if label is None:
             label = self.label

pde/pdes/base.py

@@ -11,7 +11,8 @@
 
 import numpy as np
 
-from ..fields.base import FieldBase
+from ..fields import FieldCollection
+from ..fields.base import FieldBase, OptionalArrayLike
 from ..trackers.base import TrackerCollectionDataType
 from ..tools.numba import jit
 
@@ -34,14 +35,20 @@ class PDEBase(metaclass=ABCMeta):
     explicit time dependence. """
 
 
-    def __init__(self, noise: float = 0):
+    def __init__(self, noise: OptionalArrayLike = 0):
         """
         Args:
-            noise (float):
+            noise (float or :class:`numpy.ndarray`):
                 Magnitude of the additive Gaussian white noise that is supported
                 by default. If set to zero, a deterministic partial differential
-                equation will be solved. If another noise structure is required
-                the respective methods need to be overwritten.
+                equation will be solved. Different noise magnitudes can be
+                supplied for each field in coupled PDEs.
+                
+        Note:
+            If more complicated noise structures are required, the methods
+            :meth:`PDEBase.noise_realization` and
+            :meth:`PDEBase._make_noise_realization_numba` need to be overwritten
+            for the `numpy` and `numba` backend, respectively.
         """
         self._logger = logging.getLogger(self.__class__.__name__)
         self.noise = noise
@@ -126,23 +133,43 @@ def evolution_rate_numpy(state_data, t: float):
         return rhs
 
 
-    def noise_realization(self, state: FieldBase, t: float = 0) -> FieldBase:
+    def noise_realization(self, state: FieldBase, t: float = 0,
+                          label: str = 'Noise realization') -> FieldBase:
         """ returns a realization for the noise
         
         Args:
             state (:class:`~pde.fields.ScalarField`):
                 The scalar field describing the concentration distribution
-            t (float): The current time point
+            t (float):
+                The current time point
+            label (str):
+                The label for the returned field
             
         Returns:
             :class:`~pde.fields.ScalarField`:
             Scalar field describing the evolution rate of the PDE 
         """
         if self.noise:
-            data = np.random.normal(scale=self.noise, size=state.data.shape)
-            return state.copy(data=data, label='Noise realization')
+            if np.isscalar(self.noise):
+                # a single noise value is given for all fields
+                data = np.random.normal(scale=self.noise, size=state.data.shape)
+                return state.copy(data=data, label=label)
+
+            elif isinstance(state, FieldCollection):
+                # different noise strengths, assuming one for each field
+                noise = np.broadcast_to(self.noise, len(state))
+                fields = [f.copy(data=np.random.normal(scale=n,
+                                                       size=f.data.shape))
+                          for f, n in zip(state, noise)]
+                return FieldCollection(fields, label=label)
+
+            else:
+                # different noise strengths, but a single field 
+                raise RuntimeError('Multiple noise strengths were given for '
+                                   f'the single field {state}')
+
         else:
-            return state.copy(data=0, label='Noise realization')
+            return state.copy(data=0, label=label)
 
 
     def _make_noise_realization_numba(self, state: FieldBase) -> Callable:            
@@ -157,18 +184,41 @@ def _make_noise_realization_numba(self, state: FieldBase) -> Callable:
             Function determining the right hand side of the PDE
         """
         if self.noise:        
-            noise_strength = float(self.noise)
             data_shape = state.data.shape
 
-            @jit
-            def noise_realization(state_data: np.ndarray, t: float):
-                """ compiled helper function returning a noise realization """ 
-                return noise_strength * np.random.randn(*data_shape)
+            if np.isscalar(self.noise):
+                # a single noise value is given for all fields
+                noise_strength = float(self.noise)
+
+                @jit
+                def noise_realization(state_data: np.ndarray, t: float):
+                    """ helper function returning a noise realization """ 
+                    return noise_strength * np.random.randn(*data_shape)
+
+            elif isinstance(state, FieldCollection):
+                # different noise strengths, assuming one for each field
+                noise_strengths = np.empty(data_shape[0])
+                noise_arr = np.broadcast_to(self.noise, len(state))
+                for i, noise in enumerate(noise_arr):
+                    noise_strengths[state._slices[i]] = noise
+
+                @jit
+                def noise_realization(state_data: np.ndarray, t: float):
+                    """ helper function returning a noise realization """ 
+                    out = np.random.randn(*data_shape)
+                    for i in range(data_shape[0]):
+                        out[i] *= noise_strengths[i]
+                    return out
+
+            else:
+                # different noise strengths, but a single field 
+                raise RuntimeError('Multiple noise strengths were given for '
+                                   f'the single field {state}')
 
         else:
             @jit
             def noise_realization(state_data: np.ndarray, t: float):
-                """ compiled helper function returning a noise realization """ 
+                """ helper function returning a noise realization """ 
                 return None
 
         return noise_realization  # type: ignore    

pde/pdes/pde.py

@@ -12,7 +12,7 @@
 
 from pde.pdes.base import PDEBase
 from pde.fields import FieldCollection
-from pde.fields.base import FieldBase, DataFieldBase
+from pde.fields.base import FieldBase, DataFieldBase, OptionalArrayLike
 from pde.grids.boundaries.axes import BoundariesData
 from pde.tools.numba import nb, jit
 from pde.tools.docstrings import fill_in_docstring
@@ -39,6 +39,7 @@ class PDE(PDEBase):
     @fill_in_docstring
     def __init__(self,
                  rhs: "OrderedDict[str, str]",
+                 noise: OptionalArrayLike = 0,
                  bc: BoundariesData = 'natural',
                  bc_ops: Dict[str, BoundariesData] = None):
         """
@@ -56,6 +57,11 @@ def __init__(self,
                 Note that operators need to be specified with their full name,
                 i.e., `laplace` for a scalar Laplacian and `vector_laplace` for
                 a Laplacian operating on a vector field.
+            noise (float or :class:`numpy.ndarray`):
+                Magnitude of additive Gaussian white noise. The default value of
+                zero implies deterministic partial differential equations will
+                be solved. Different noise magnitudes can be supplied for each
+                field in coupled PDEs.
             bc:
                 Boundary conditions for the operators used in the expression.
                 The conditions here are applied to all operators that do not
@@ -71,15 +77,14 @@ def __init__(self,
         Note:
             The order in which the fields are given in `rhs` defines the order
             in which they need to appear in the `state` variable when the
-            evolution rate is calculated. In Python version 3.6, the insertion
-            order of a normal dictionary was not guaranteed and might thus lead
-            to surprising results.
-                
+            evolution rate is calculated. Note that the insertion order of
+            `dict` is guaranteed since Python version 3.7, so a normal
+            dictionary can be used to define the equations.
         """
         from ..tools.expressions import ScalarExpression
         from sympy.core.function import AppliedUndef
 
-        super().__init__()
+        super().__init__(noise=noise)
 
         # validate input
         if not isinstance(rhs, OrderedDict):

pde/pdes/tests/test_pde.py

@@ -124,3 +124,20 @@ def make_op(state):
     np.testing.assert_allclose(field.data, res.data)
 
     del UnitGrid._operators['undefined']  # reset original state
+
+
+
+@pytest.mark.parametrize('backend', ['numpy', 'numba'])
+def test_pde_noise(backend):
+    """ test noise operator on PDE class """
+    grid = UnitGrid([64, 64])
+    state = FieldCollection([ScalarField(grid), ScalarField(grid)])
+
+    eq = PDE({'a': 0, 'b': 0}, noise=.5)
+    res = eq.solve(state, t_range=1, backend=backend, dt=1)
+    assert res.data.std() == pytest.approx(.5, rel=0.1)
+
+    eq = PDE({'a': 0, 'b': 0}, noise=[0.01, 2.])
+    res = eq.solve(state, t_range=1, backend=backend, dt=1)
+    assert res.data[0].std() == pytest.approx(0.01, rel=0.1)
+    assert res.data[1].std() == pytest.approx(2., rel=0.1)