Make interpolator work for vector valued base functions

kinnala · Aug 5, 2024 · fe3f474 · fe3f474
1 parent c6b3675
commit fe3f474
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/skfem/assembly/basis/cell_basis.py b/skfem/assembly/basis/cell_basis.py
@@ -190,15 +190,20 @@ def probes(self, x: ndarray):
                 for k in range(self.Nbfun)
             ]
         ).flatten()
+        #number of components of the current base functions
+        try:
+            comp=len(self.elem.lbasis(self.elem.dim*[0],0)[0])
+        except:    
+            comp=1
         return coo_matrix(
             (
                 phis,
                 (
-                    np.tile(np.arange(x.shape[1]), self.Nbfun),
-                    self.element_dofs[:, cells].flatten(),
+                    np.tile(np.arange(comp*x.shape[1]), self.Nbfun),
+                    self.element_dofs[:, np.tile(cells,comp)].flatten(),
                 ),
             ),
-            shape=(x.shape[1], self.N),
+            shape=(comp*x.shape[1], self.N),
         )
 
     def point_source(self, x: ndarray) -> ndarray: