Generic multi-component mesh

pySDC/implementations/datatype_classes/mesh.py

@@ -80,7 +80,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, out=None, **kwargs):
             else:
                 args.append(input_)
 
-        results = super(mesh, self).__array_ufunc__(ufunc, method, *args, **kwargs).view(mesh)
+        results = super(mesh, self).__array_ufunc__(ufunc, method, *args, **kwargs).view(type(self))
         if type(self) == type(results):
             results._comm = comm
         return results
@@ -149,74 +149,66 @@ def bcast(self, root=None, comm=None):
         return self
 
 
-class imex_mesh(object):
+class MultiComponentMesh(mesh):
+    r"""
+    Generic mesh with multiple components.
+
+    To make a specific multi-component mesh, derive from this class and list the components as strings in the class
+    attribute ``components``. An example:
+
+    ```
+    class imex_mesh(MultiComponentMesh):
+        components = ['impl', 'expl']
+    ```
+
+    Instantiating such a mesh will expand the mesh along an added first dimension for each component and allow access
+    to the components with ``.``. Continuing the above example:
+
+    ```
+    init = ((100,), None, numpy.dtype('d'))
+    f = imex_mesh(init)
+    f.shape  # (2, 100)
+    f.expl.shape  # (100,)
+    ```
+
+    Note that the components are not attributes of the mesh: ``"expl" in dir(f)`` will return False! Rather, the
+    components are handled in ``__getattr__``. This function is called if an attribute is not found and returns a view
+    on to the component if appropriate. Importantly, this means that you cannot name a component like something that
+    is already an attribute of ``mesh`` or ``numpy.ndarray`` because this will not result in calls to ``__getattr__``.
+
+    There are a couple more things to keep in mind:
+     - Because a ``MultiComponentMesh`` is just a ``numpy.ndarray`` with one more dimension, all components must have
+       the same shape.
+     - You can use the entire ``MultiComponentMesh`` like a ``numpy.ndarray`` in operations that accept arrays, but make
+       sure that you really want to apply the same operation on all components if you do.
+     - If you omit the assignment operator ``[:]`` during assignment, you will not change the mesh at all. Omitting this
+       leads to all kinds of trouble throughout the code. But here you really cannot get away without.
     """
-    RHS data type for meshes with implicit and explicit components
 
-    This data type can be used to have RHS with 2 components (here implicit and explicit)
+    components = []
 
-    Attributes:
-        impl (mesh.mesh): implicit part
-        expl (mesh.mesh): explicit part
-    """
-
-    def __init__(self, init, val=0.0):
-        """
-        Initialization routine
-
-        Args:
-            init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested)
-                  or another imex_mesh object
-            val (float): an initial number (default: 0.0)
-        Raises:
-            DataError: if init is none of the types above
-        """
-
-        if isinstance(init, type(self)):
-            self.impl = mesh(init.impl)
-            self.expl = mesh(init.expl)
-        elif (
-            isinstance(init, tuple)
-            and (init[1] is None or isinstance(init[1], MPI.Intracomm))
-            and isinstance(init[2], np.dtype)
-        ):
-            self.impl = mesh(init, val=val)
-            self.expl = mesh(init, val=val)
-        # something is wrong, if none of the ones above hit
+    def __new__(cls, init, *args, **kwargs):
+        if isinstance(init, tuple):
+            shape = (init[0],) if type(init[0]) is int else init[0]
+            obj = super().__new__(cls, ((len(cls.components), *shape), *init[1:]), *args, **kwargs)
         else:
-            raise DataError('something went wrong during %s initialization' % type(self))
+            obj = super().__new__(cls, init, *args, **kwargs)
 
+        return obj
 
-class comp2_mesh(object):
-    """
-    RHS data type for meshes with 2 components
+    def __getattr__(self, name):
+        if name in self.components:
+            if self.shape[0] == len(self.components):
+                return self[self.components.index(name)]
+            else:
+                raise AttributeError(f'Cannot access {name!r} in {type(self)!r} because the shape is unexpected.')
+        else:
+            raise AttributeError(f"{type(self)!r} does not have attribute {name!r}!")
 
-    Attributes:
-        comp1 (mesh.mesh): first part
-        comp2 (mesh.mesh): second part
-    """
 
-    def __init__(self, init, val=0.0):
-        """
-        Initialization routine
+class imex_mesh(MultiComponentMesh):
+    components = ['impl', 'expl']
 
-        Args:
-            init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested)
-                  or another comp2_mesh object
-        Raises:
-            DataError: if init is none of the types above
-        """
 
-        if isinstance(init, type(self)):
-            self.comp1 = mesh(init.comp1)
-            self.comp2 = mesh(init.comp2)
-        elif (
-            isinstance(init, tuple)
-            and (init[1] is None or isinstance(init[1], MPI.Intracomm))
-            and isinstance(init[2], np.dtype)
-        ):
-            self.comp1 = mesh(init, val=val)
-            self.comp2 = mesh(init, val=val)
-        # something is wrong, if none of the ones above hit
-        else:
-            raise DataError('something went wrong during %s initialization' % type(self))
+class comp2_mesh(MultiComponentMesh):
+    components = ['comp1', 'comp2']

pySDC/tests/test_datatypes/test_multicomponent_mesh.py

@@ -0,0 +1,64 @@
+import pytest
+
+
+@pytest.mark.base
+@pytest.mark.parametrize('shape', [1, (3,), (2, 4)])
+def test_MultiComponentMesh(shape):
+    from pySDC.implementations.datatype_classes.mesh import MultiComponentMesh
+    import numpy as np
+
+    class TestMesh(MultiComponentMesh):
+        components = ['a', 'b']
+
+    # instantiate meshes
+    init = (shape, None, np.dtype('D'))
+    A = TestMesh(init)
+    B = TestMesh(A)
+
+    # fill part of the meshes with values
+    a = np.random.random(shape)
+    b = np.random.random(shape)
+    zero = np.zeros_like(a)
+    A.a[:] = a
+    B.a[:] = b
+
+    # check that the meshes have been prepared appropriately
+    for M, m in zip([A, B], [a, b]):
+        assert M.shape == (len(TestMesh.components),) + ((shape,) if type(shape) is int else shape)
+        assert np.allclose(M.a, m)
+        assert np.allclose(M.b, zero)
+        assert np.shares_memory(M, M.a)
+        assert np.shares_memory(M, M.b)
+        assert not np.shares_memory(M.a, m)
+
+    # check that various computations give the desired results
+    assert np.allclose(A.a + B.a, a + b)
+    assert np.allclose((A + B).a, a + b)
+    assert np.allclose((A + B).b, zero)
+
+    C = A - B
+    assert np.allclose(C.a, a - b)
+    assert np.allclose(C.b, zero)
+    assert not np.shares_memory(A, C)
+    assert not np.shares_memory(B, C)
+
+    D = np.exp(A)
+    assert type(D) == TestMesh
+    assert np.allclose(D.a, np.exp(a))
+    assert np.allclose(D.b, zero + 1)
+    assert not np.shares_memory(A, D)
+
+    B *= A
+    assert np.allclose(B.a, a * b)
+    assert np.allclose(A.a, a)
+    assert np.allclose(B.b, zero)
+    assert np.allclose(A.b, zero)
+    assert not np.shares_memory(A, B)
+
+    A /= 10.0
+    assert np.allclose(A.a, a / 10)
+    assert np.allclose(A.b, zero)
+
+
+if __name__ == '__main__':
+    test_MultiComponentMesh(1)