Removed matvec_5d and all related kernels; added new methods...

Removed matvec_5d and all related kernels; added new methods `prepare_batch_vector` and `finalize_batch_vector` to use numpy routines instead.

src/gvec_to_python/GVEC_functions.py

@@ -212,17 +212,17 @@ class GVEC:
 
         Parameters
         ----------
-            s, a1, a2 : float or array-like
-                Coordinates in logical space. 
-                If list or np.array, must be either 1d or 3d (from meshgrid).
-                s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. 
+        s, a1, a2 : float or array-like
+            Coordinates in logical space. 
+            If list or np.array, must be either 1d or 3d (from meshgrid).
+            s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. 
 
-            flat_eval : bool
-                Whether to do flat (marker) evaluation when three array arguments are given. 
+        flat_eval : bool
+            Whether to do flat (marker) evaluation when three array arguments are given. 
 
         Returns
         -------
-            A float (single-point evaluation) or 3d numpy array (meshgrid evaluation).
+        A float (single-point evaluation) or 3d numpy array (meshgrid evaluation).
         '''
         return self.profiles.profile(s, a1, a2, name='pressure', flat_eval=flat_eval)
 
@@ -543,28 +543,29 @@ class GVEC:
 
         mu0 = 1.2566370621219e-6
         b = np.ascontiguousarray(self._bv_gvec(s, th, ze))
+        b = GVEC_domain.prepare_batch_vector(b)
 
         grad_bt = self._grad_b_theta(s, th, ze)
         grad_bz = self._grad_b_zeta(s, th, ze)
+        
         db_ds = np.ascontiguousarray((0.*grad_bt[0], grad_bt[0], grad_bz[0]))
         db_dt = np.ascontiguousarray((0.*grad_bt[1], grad_bt[1], grad_bz[1]))
         db_dz = np.ascontiguousarray((0.*grad_bt[2], grad_bt[2], grad_bz[2]))
+        db_ds = GVEC_domain.prepare_batch_vector(db_ds)
+        db_dt = GVEC_domain.prepare_batch_vector(db_dt)
+        db_dz = GVEC_domain.prepare_batch_vector(db_dz)
 
         g = self.domain.g_gvec(s, th, ze)
         dg_ds = self.domain.dg_gvec(s, th, ze, der='s')
         dg_dt = self.domain.dg_gvec(s, th, ze, der='th')
         dg_dz = self.domain.dg_gvec(s, th, ze, der='ze')
-
-        # meshgrid evaluation
-        if g.ndim == 5:
-            dgb_ds = GVEC_domain.matvec_5d(dg_ds, b, use_pyccel=self.use_pyccel) + GVEC_domain.matvec_5d(g, db_ds, use_pyccel=self.use_pyccel)
-            dgb_dt = GVEC_domain.matvec_5d(dg_dt, b, use_pyccel=self.use_pyccel) + GVEC_domain.matvec_5d(g, db_dt, use_pyccel=self.use_pyccel)
-            dgb_dz = GVEC_domain.matvec_5d(dg_dz, b, use_pyccel=self.use_pyccel) + GVEC_domain.matvec_5d(g, db_dz, use_pyccel=self.use_pyccel)
-        # single point evaluation
-        else:
-            dgb_ds = dg_ds @ b + g @ db_ds
-            dgb_dt = dg_dt @ b + g @ db_dt
-            dgb_dz = dg_dz @ b + g @ db_dz
+            
+        dgb_ds = dg_ds @ b + g @ db_ds
+        dgb_dt = dg_dt @ b + g @ db_dt
+        dgb_dz = dg_dz @ b + g @ db_dz
+        dgb_ds = GVEC_domain.finalize_batch_vector(dgb_ds)
+        dgb_dt = GVEC_domain.finalize_batch_vector(dgb_dt)
+        dgb_dz = GVEC_domain.finalize_batch_vector(dgb_dz)
 
         j1 = (dgb_dt[2] - dgb_dz[1]) / mu0
         j2 = (dgb_dz[0] - dgb_ds[2]) / mu0
@@ -806,15 +807,13 @@ class GVEC:
 
         else:
             raise NotImplementedError      
-
-        # meshgrid evaluation
-        if mat.ndim == 5:
-            out = GVEC_domain.matvec_5d(mat, np.ascontiguousarray(b), transpose_a=transpose, use_pyccel=self.use_pyccel)
-        # single-point evaluation
+                
+        b = GVEC_domain.prepare_batch_vector(np.ascontiguousarray(b))
+        if transpose:
+            matT = mat.swapaxes(-1, -2)
+            out = matT @ b
         else:
-            if transpose:
-                out = mat.T @ b
-            else:
-                out = mat @ b
+            out = mat @ b
+        out = GVEC_domain.finalize_batch_vector(out)
 
         return out[0], out[1], out[2]

src/gvec_to_python/Makefile

@@ -21,7 +21,6 @@ SOURCES := $(BK).py $(BEV1).py
 
 OUTPUTS := $(SOURCES:.py=$(SO_EXT))
 
-
 #--------------------------------------
 # PYCCELIZE
 #--------------------------------------
@@ -35,7 +34,6 @@ $(BK)$(SO_EXT) : $(BK).py
 $(BEV1)$(SO_EXT) : $(BEV1).py $(BK)$(SO_EXT)
 	pyccel $(FLAGS) $< 
 
-
 #--------------------------------------
 # CLEAN UP
 #--------------------------------------

src/gvec_to_python/geometry/domain.py

@@ -1229,6 +1229,47 @@ class GVEC_domain:
             J = np.moveaxis(J, -1, 0)
             J = np.moveaxis(J, -1, 0)
         return J
+    
+    @staticmethod
+    def prepare_batch_vector(v):
+        """Swap axes of a batch of vectors and add singelton dimension, 
+        such that it is compatible with numpy's matmul (@).
+
+        Parameters
+        ----------
+            v : numpy.ndarray of shape (3) or (3, ...)
+                A batch of vectors.
+
+        Returns
+        -------
+            numpy.ndarray of shape (3) or (..., 3, 1)
+                A batch of vectors.
+        """
+
+        if v.ndim > 1 and v.shape[0] == 3:
+            v = np.moveaxis(v, 0, -1)
+            v = v.reshape(*v.shape, 1)
+        return v
+
+    @staticmethod
+    def finalize_batch_vector(v):
+        """Swap axes of a batch of vectors and remove singelton dimension.
+
+        Parameters
+        ----------
+            v : numpy.ndarray of shape (3) or (..., 3, 1)
+                A batch of vectors.
+
+        Returns
+        -------
+            numpy.ndarray of shape (3) or (..., 3, 1)
+                A batch of vectors.
+        """
+
+        if v.ndim > 1 and v.shape[-2:] == (3, 1):
+            v = np.squeeze(v, axis=-1)
+            v = np.moveaxis(v, -1, 0)
+        return v
 
     @staticmethod
     def transpose_5d(a, use_pyccel=False):

src/gvec_to_python/geometry/kernels.py

-def matvec_5d_kernel(a: 'float[:,:,:,:,:]', b: 'float[:,:,:,:]', out: 'float[:,:,:,:]'):
-    '''Performs matrix-vector multiplication in the last two indices of the 5-dim array a.
-    a must have shape (s0, s1, s2, n, m) and b must have "swapped axes" with shape (m, s0, s1, s2).
-    Hence the last index of a contracts with the first index of b.'''
-
-    out[:] = 0.
-
-    for i in range(out.shape[1]):
-        for j in range(out.shape[2]):
-            for k in range(out.shape[3]):
-                for l in range(out.shape[0]):
-                    for m in range(b.shape[0]):
-                        out[l, i, j, k] += a[i, j, k, l, m] * b[m, i, j, k]
-
-def transpose_5d_kernel(a: 'float[:,:,:,:,:]', out: 'float[:,:,:,:,:]'):
-    '''Matrix transpose in the last two indices of a 5-dim array.
-    If a is shape shape (., ., ., m, n) then out must be shape shape (., ., ., n, m).'''
-
-    for i in range(out.shape[0]):
-        for j in range(out.shape[1]):
-            for k in range(out.shape[2]):
-                for n in range(out.shape[3]):
-                    for m in range(out.shape[4]):
-                        out[i, j, k, n, m] = a[i, j, k, m, n]
\ No newline at end of file