Fix creating coefficient matrix for multi expr kernels

sumpy/expansion/__init__.py

@@ -398,11 +398,27 @@ class LinearRecurrenceBasedExpansionTermsWrangler(ExpansionTermsWrangler):
             :math:`S = N_{[r, :]}^{-T} N^T = N_{[r, :]}^{-T} N^T`.
             """
             n = nullspace(pde_mat)
-            n = n[offset*iexpr:offset*(iexpr+1), :]
-            idx = self.get_reduced_coeffs(n)
-            n = n[:, :len(idx)]
-            s = solve_symbolic(n.T[:, idx], n.T)
-            stored_identifiers = [mis[i] for i in idx]
+
+            # Move the rows corresponding to this expression to the front
+            rearrange = list(range(offset*iexpr, offset*(iexpr+1)))
+            for i in range(nexpr*offset):
+                if i < offset*iexpr or i >= offset*(iexpr+1):
+                    rearrange.append(i)
+            n = n[rearrange, :]
+
+            # Get the indices of the reduced coefficients
+            idx_all_exprs = self.get_reduced_coeffs(n)
+
+            s = solve_symbolic(n.T[:, idx_all_exprs], n.T)
+
+            # Filter out coefficients belonging to this expression
+            indices = []
+            for idx in idx_all_exprs:
+                if idx >= offset*iexpr and idx < offset*(iexpr+1):
+                    indices.append(idx)
+            s = s[:len(indices), offset*iexpr:offset*(iexpr+1)]
+
+            stored_identifiers = [mis[i] for i in indices]
         else:
             s = np.eye(len(mis))
             stored_identifiers = mis

test/test_kernels.py

@@ -413,8 +413,6 @@ def test_translations(ctx_getter, knl, local_expn_class, mpole_expn_class):
            issubclass(local_expn_class, VolumeTaylorExpansionBase):
         # FIXME: Embarrassing--but we run out of memory for higher orders.
         orders = [2, 3]
-    elif isinstance(knl, StokesletKernel):
-        orders = [3, 4, 5]
     else:
         orders = [2, 3, 4]
     nboxes = centers.shape[-1]