Fixed some issues with transfer classes (#422)

* Fixes for transfer classes

* Cleanup of `mesh_to_mesh`

* Bug fixes

* Fixed `fft_to_fft`

pySDC/helpers/transfer_helper.py

@@ -138,7 +138,7 @@ def restriction_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1):
 
 def interpolation_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1, equidist_nested=True):
     """
-    Function to contruct the restriction matrix in 1d using barycentric interpolation
+    Function to construct the restriction matrix in 1d using barycentric interpolation
 
     Args:
         fine_grid (np.ndarray): a one dimensional 1d array containing the nodes of the fine grid

pySDC/implementations/convergence_controller_classes/check_iteration_estimator.py

@@ -6,7 +6,7 @@
 class CheckIterationEstimatorNonMPI(ConvergenceController):
     def __init__(self, controller, params, description, **kwargs):
         """
-        Initalization routine
+        Initialization routine
 
         Args:
             controller (pySDC.Controller): The controller

pySDC/implementations/datatype_classes/mesh.py

@@ -199,7 +199,7 @@ def __new__(cls, init, *args, **kwargs):
     def __getattr__(self, name):
         if name in self.components:
             if self.shape[0] == len(self.components):
-                return self[self.components.index(name)]
+                return self[self.components.index(name)].view(mesh)
             else:
                 raise AttributeError(f'Cannot access {name!r} in {type(self)!r} because the shape is unexpected.')
         else:

pySDC/implementations/problem_classes/Brusselator.py

@@ -30,6 +30,7 @@ def __init__(self, alpha=0.1, **kwargs):
         shape = (2,) + (self.init[0])
         self.iU = 0
         self.iV = 1
+        self.ncomp = 2  # needed for transfer class
         self.init = (shape, self.comm, np.dtype('float'))
 
     def _eval_explicit_part(self, u, t, f_expl):

pySDC/implementations/problem_classes/GrayScott_MPIFFT.py

@@ -79,6 +79,7 @@ def __init__(self, Du=1.0, Dv=0.01, A=0.09, B=0.086, **kwargs):
         shape = (2,) + (self.init[0])
         self.iU = 0
         self.iV = 1
+        self.ncomp = 2  # needed for transfer class
         self.init = (shape, self.comm, self.xp.dtype('float'))
 
         self._makeAttributeAndRegister('Du', 'Dv', 'A', 'B', localVars=locals(), readOnly=True)

pySDC/implementations/sweeper_classes/imex_1st_order.py

@@ -42,16 +42,14 @@ def integrate(self):
             list of dtype_u: containing the integral as values
         """
 
-        # get current level and problem description
         L = self.level
+        P = L.prob
 
         me = []
-
         # integrate RHS over all collocation nodes
         for m in range(1, self.coll.num_nodes + 1):
-            me.append(L.dt * self.coll.Qmat[m, 1] * (L.f[1].impl + L.f[1].expl))
-            # new instance of dtype_u, initialize values with 0
-            for j in range(2, self.coll.num_nodes + 1):
+            me.append(P.dtype_u(P.init, val=0.0))
+            for j in range(1, self.coll.num_nodes + 1):
                 me[m - 1] += L.dt * self.coll.Qmat[m, j] * (L.f[j].impl + L.f[j].expl)
 
         return me

pySDC/implementations/transfer_classes/TransferMesh.py

@@ -4,15 +4,14 @@
 import pySDC.helpers.transfer_helper as th
 from pySDC.core.Errors import TransferError
 from pySDC.core.SpaceTransfer import space_transfer
-from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh, comp2_mesh
 
 
 class mesh_to_mesh(space_transfer):
     """
-    Custon base_transfer class, implements Transfer.py
+    Custom base_transfer class, implements Transfer.py
 
     This implementation can restrict and prolong between nd meshes with dirichlet-0 or periodic boundaries
-    via matrix-vector products
+    via matrix-vector products.
 
     Attributes:
         Rspace: spatial restriction matrix, dim. Nf x Nc
@@ -30,7 +29,7 @@ def __init__(self, fine_prob, coarse_prob, params):
         """
 
         # invoke super initialization
-        super(mesh_to_mesh, self).__init__(fine_prob, coarse_prob, params)
+        super().__init__(fine_prob, coarse_prob, params)
 
         if self.params.rorder % 2 != 0:
             raise TransferError('Need even order for restriction')
@@ -153,51 +152,31 @@ def restrict(self, F):
         Args:
             F: the fine level data (easier to access than via the fine attribute)
         """
-        if isinstance(F, mesh):
-            G = self.coarse_prob.dtype_u(self.coarse_prob.init)
-            if hasattr(self.fine_prob, 'ncomp'):
-                for i in range(self.fine_prob.ncomp):
-                    tmpF = F[..., i].flatten()
-                    tmpG = self.Rspace.dot(tmpF)
-                    G[..., i] = tmpG.reshape(self.coarse_prob.nvars)
-            else:
-                tmpF = F.flatten()
-                tmpG = self.Rspace.dot(tmpF)
-                G[:] = tmpG.reshape(self.coarse_prob.nvars)
-        elif isinstance(F, imex_mesh):
-            G = self.coarse_prob.dtype_f(self.coarse_prob.init)
-            if hasattr(self.fine_prob, 'ncomp'):
-                for i in range(self.fine_prob.ncomp):
-                    tmpF = F.impl[..., i].flatten()
-                    tmpG = self.Rspace.dot(tmpF)
-                    G.impl[..., i] = tmpG.reshape(self.coarse_prob.nvars)
-                    tmpF = F.expl[..., i].flatten()
-                    tmpG = self.Rspace.dot(tmpF)
-                    G.expl[..., i] = tmpG.reshape(self.coarse_prob.nvars)
-            else:
-                tmpF = F.impl.flatten()
-                tmpG = self.Rspace.dot(tmpF)
-                G.impl[:] = tmpG.reshape(self.coarse_prob.nvars)
-                tmpF = F.expl.flatten()
-                tmpG = self.Rspace.dot(tmpF)
-                G.expl[:] = tmpG.reshape(self.coarse_prob.nvars)
-        elif isinstance(F, comp2_mesh):
-            G = self.coarse_prob.dtype_f(self.coarse_prob.init)
+        G = type(F)(self.coarse_prob.init)
+
+        def _restrict(fine, coarse):
             if hasattr(self.fine_prob, 'ncomp'):
                 for i in range(self.fine_prob.ncomp):
-                    tmpF = F.comp1[..., i].flatten()
-                    tmpG = self.Rspace.dot(tmpF)
-                    G.comp1[..., i] = tmpG.reshape(self.coarse_prob.nvars)
-                    tmpF = F.comp2[..., i].flatten()
-                    tmpG = self.Rspace.dot(tmpF)
-                    G.comp2[..., i] = tmpG.reshape(self.coarse_prob.nvars)
+                    if fine.shape[-1] == self.fine_prob.ncomp:
+                        tmpF = fine[..., i].flatten()
+                        tmpG = self.Rspace.dot(tmpF)
+                        coarse[..., i] = tmpG.reshape(self.coarse_prob.nvars)
+                    elif fine.shape[0] == self.fine_prob.ncomp:
+                        tmpF = fine[i, ...].flatten()
+                        tmpG = self.Rspace.dot(tmpF)
+                        coarse[i, ...] = tmpG.reshape(self.coarse_prob.nvars)
+                    else:
+                        raise TransferError('Don\'t know how to restrict for this problem with multiple components')
             else:
-                tmpF = F.comp1.flatten()
-                tmpG = self.Rspace.dot(tmpF)
-                G.comp1[:] = tmpG.reshape(self.coarse_prob.nvars)
-                tmpF = F.comp2.flatten()
+                tmpF = fine.flatten()
                 tmpG = self.Rspace.dot(tmpF)
-                G.comp2[:] = tmpG.reshape(self.coarse_prob.nvars)
+                coarse[:] = tmpG.reshape(self.coarse_prob.nvars)
+
+        if hasattr(type(F), 'components'):
+            for comp in F.components:
+                _restrict(F.__getattr__(comp), G.__getattr__(comp))
+        elif type(F).__name__ == 'mesh':
+            _restrict(F, G)
         else:
             raise TransferError('Wrong data type for restriction, got %s' % type(F))
         return G
@@ -208,51 +187,32 @@ def prolong(self, G):
         Args:
             G: the coarse level data (easier to access than via the coarse attribute)
         """
-        if isinstance(G, mesh):
-            F = self.fine_prob.dtype_u(self.fine_prob.init)
-            if hasattr(self.fine_prob, 'ncomp'):
-                for i in range(self.fine_prob.ncomp):
-                    tmpG = G[..., i].flatten()
-                    tmpF = self.Pspace.dot(tmpG)
-                    F[..., i] = tmpF.reshape(self.fine_prob.nvars)
-            else:
-                tmpG = G.flatten()
-                tmpF = self.Pspace.dot(tmpG)
-                F[:] = tmpF.reshape(self.fine_prob.nvars)
-        elif isinstance(G, imex_mesh):
-            F = self.fine_prob.dtype_f(self.fine_prob.init)
-            if hasattr(self.fine_prob, 'ncomp'):
-                for i in range(self.fine_prob.ncomp):
-                    tmpG = G.impl[..., i].flatten()
-                    tmpF = self.Pspace.dot(tmpG)
-                    F.impl[..., i] = tmpF.reshape(self.fine_prob.nvars)
-                    tmpG = G.expl[..., i].flatten()
-                    tmpF = self.Rspace.dot(tmpG)
-                    F.expl[..., i] = tmpF.reshape(self.fine_prob.nvars)
-            else:
-                tmpG = G.impl.flatten()
-                tmpF = self.Pspace.dot(tmpG)
-                F.impl[:] = tmpF.reshape(self.fine_prob.nvars)
-                tmpG = G.expl.flatten()
-                tmpF = self.Pspace.dot(tmpG)
-                F.expl[:] = tmpF.reshape(self.fine_prob.nvars)
-        elif isinstance(G, comp2_mesh):
-            F = self.fine_prob.dtype_f(self.fine_prob.init)
+        F = type(G)(self.fine_prob.init)
+
+        def _prolong(coarse, fine):
             if hasattr(self.fine_prob, 'ncomp'):
                 for i in range(self.fine_prob.ncomp):
-                    tmpG = G.comp1[..., i].flatten()
-                    tmpF = self.Pspace.dot(tmpG)
-                    F.comp1[..., i] = tmpF.reshape(self.fine_prob.nvars)
-                    tmpG = G.comp2[..., i].flatten()
-                    tmpF = self.Rspace.dot(tmpG)
-                    F.comp2[..., i] = tmpF.reshape(self.fine_prob.nvars)
+                    if coarse.shape[-1] == self.fine_prob.ncomp:
+                        tmpG = coarse[..., i].flatten()
+                        tmpF = self.Pspace.dot(tmpG)
+                        fine[..., i] = tmpF.reshape(self.fine_prob.nvars)
+                    elif coarse.shape[0] == self.fine_prob.ncomp:
+                        tmpG = coarse[i, ...].flatten()
+                        tmpF = self.Pspace.dot(tmpG)
+                        fine[i, ...] = tmpF.reshape(self.fine_prob.nvars)
+                    else:
+                        raise TransferError('Don\'t know how to prolong for this problem with multiple components')
             else:
-                tmpG = G.comp1.flatten()
-                tmpF = self.Pspace.dot(tmpG)
-                F.comp1[:] = tmpF.reshape(self.fine_prob.nvars)
-                tmpG = G.comp2.flatten()
+                tmpG = coarse.flatten()
                 tmpF = self.Pspace.dot(tmpG)
-                F.comp2[:] = tmpF.reshape(self.fine_prob.nvars)
+                fine[:] = tmpF.reshape(self.fine_prob.nvars)
+            return fine
+
+        if hasattr(type(F), 'components'):
+            for comp in G.components:
+                _prolong(G.__getattr__(comp), F.__getattr__(comp))
+        elif type(G).__name__ == 'mesh':
+            F[:] = _prolong(G, F)
         else:
             raise TransferError('Wrong data type for prolongation, got %s' % type(G))
         return F

pySDC/implementations/transfer_classes/TransferMesh_FFT.py

@@ -2,7 +2,6 @@
 
 from pySDC.core.Errors import TransferError
 from pySDC.core.SpaceTransfer import space_transfer
-from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 
 
 class mesh_to_mesh_fft(space_transfer):
@@ -26,9 +25,9 @@ def __init__(self, fine_prob, coarse_prob, params):
             params: parameters for the transfer operators
         """
         # invoke super initialization
-        super(mesh_to_mesh_fft, self).__init__(fine_prob, coarse_prob, params)
+        super().__init__(fine_prob, coarse_prob, params)
 
-        self.ratio = int(self.fine_prob.params.nvars / self.coarse_prob.params.nvars)
+        self.ratio = int(self.fine_prob.nvars / self.coarse_prob.nvars)
 
     def restrict(self, F):
         """
@@ -37,11 +36,11 @@ def restrict(self, F):
         Args:
             F: the fine level data (easier to access than via the fine attribute)
         """
-        if isinstance(F, mesh):
-            G = mesh(self.coarse_prob.init, val=0.0)
+        G = type(F)(self.coarse_prob.init, val=0.0)
+
+        if type(F).__name__ == 'mesh':
             G[:] = F[:: self.ratio]
-        elif isinstance(F, imex_mesh):
-            G = imex_mesh(self.coarse_prob.init, val=0.0)
+        elif type(F).__name__ == 'imex_mesh':
             G.impl[:] = F.impl[:: self.ratio]
             G.expl[:] = F.expl[:: self.ratio]
         else:
@@ -55,28 +54,21 @@ def prolong(self, G):
         Args:
             G: the coarse level data (easier to access than via the coarse attribute)
         """
-        if isinstance(G, mesh):
-            F = mesh(self.fine_prob.init, val=0.0)
-            tmpG = np.fft.rfft(G)
-            tmpF = np.zeros(self.fine_prob.init[0] // 2 + 1, dtype=np.complex128)
-            halfG = int(self.coarse_prob.init[0] / 2)
-            tmpF[0:halfG] = tmpG[0:halfG]
-            tmpF[-1] = tmpG[-1]
-            F[:] = np.fft.irfft(tmpF) * self.ratio
-        elif isinstance(G, imex_mesh):
-            F = imex_mesh(G)
-            tmpG_impl = np.fft.rfft(G.impl)
-            tmpF_impl = np.zeros(self.fine_prob.init[0] // 2 + 1, dtype=np.complex128)
-            halfG = int(self.coarse_prob.init[0] / 2)
-            tmpF_impl[0:halfG] = tmpG_impl[0:halfG]
-            tmpF_impl[-1] = tmpG_impl[-1]
-            F.impl[:] = np.fft.irfft(tmpF_impl) * self.ratio
-            tmpG_expl = np.fft.rfft(G.expl)
-            tmpF_expl = np.zeros(self.fine_prob.init[0] // 2 + 1, dtype=np.complex128)
-            halfG = int(self.coarse_prob.init[0] / 2)
-            tmpF_expl[0:halfG] = tmpG_expl[0:halfG]
-            tmpF_expl[-1] = tmpG_expl[-1]
-            F.expl[:] = np.fft.irfft(tmpF_expl) * self.ratio
+        F = type(G)(self.fine_prob.init, val=0.0)
+
+        def _prolong(coarse):
+            coarse_hat = np.fft.rfft(coarse)
+            fine_hat = np.zeros(self.fine_prob.init[0] // 2 + 1, dtype=np.complex128)
+            half_idx = self.coarse_prob.init[0] // 2
+            fine_hat[0:half_idx] = coarse_hat[0:half_idx]
+            fine_hat[-1] = coarse_hat[-1]
+            return np.fft.irfft(fine_hat) * self.ratio
+
+        if type(G).__name__ == 'mesh':
+            F[:] = _prolong(G)
+        elif type(G).__name__ == 'imex_mesh':
+            F.impl[:] = _prolong(G.impl)
+            F.expl[:] = _prolong(G.expl)
         else:
             raise TransferError('Unknown data type, got %s' % type(G))
         return F