Removed matmul_5d; use numpy.swapaxes with @ instead.

src/gvec_to_python/geometry/domain.py

 import numpy as np
 from scipy.optimize import newton
 
-from gvec_to_python.geometry.kernels import matmul_5d_kernel, matvec_5d_kernel, transpose_5d_kernel
+from gvec_to_python.geometry.kernels import matvec_5d_kernel, transpose_5d_kernel
 
 
 class GVEC_domain:
@@ -59,7 +59,7 @@ class GVEC_domain:
 
     Evaluation points can be float or array-like (1d or 3d).
 
-    There are also three @staticmethods prepare_args, swap_J_axes and matmul_5d.
+    There are also three @staticmethods prepare_args, swap_J_axes.
 
     Parameters
     ----------
@@ -202,14 +202,7 @@ class GVEC_domain:
         tmp1 = self.dX(s, th, ze, flat_eval=flat_eval)  # Jacobian of the Xmap
         tmp2 = self.dh(*self.Xmap(s, th, ze, flat_eval=flat_eval))  # Jacobian of the hmap
 
-        # meshgrid evaluation
-        # if tmp1.ndim == 5:
-        #     out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel)
-        # # single-point evaluation
-        # else:
-        out = tmp2 @ tmp1
-
-        return out
+        return tmp2 @ tmp1
 
     def det_df_gvec(self, s, th, ze, flat_eval=False):
         '''The Jacobian determinant of f_gvec.
@@ -264,14 +257,8 @@ class GVEC_domain:
 
         tmp = np.ascontiguousarray(self.df_gvec(s, th, ze, flat_eval=flat_eval))  # Jacobian of f_gvec
         tmpT = tmp.swapaxes(-1, -2)
-        # # meshgrid evaluation
-        # if tmp.ndim == 5:
-        #     out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel)
-        # # single-point evaluation
-        # else:
-        out = tmpT @ tmp
 
-        return out
+        return tmpT @ tmp
 
     def g_gvec_inv(self, s, th, ze, flat_eval=False):
         '''Inverse metric tensor of f_gvec.
@@ -326,30 +313,20 @@ class GVEC_domain:
         gh = self.gh(*self.Xmap(s, th, ze, flat_eval=flat_eval))
         dgh_dq1 = self.dgh_dq1(*self.Xmap(s, th, ze, flat_eval=flat_eval))
 
-        # meshgrid evaluation
-        if dX.ndim == 5:
-            dq1_di = dX[:, :, :, 0, comp]
+        tmp = self.swap_J_back(dX)
+        dq1_di = tmp[0, comp]
         dgh_di = self.swap_J_axes(self.swap_J_back(dgh_dq1) * dq1_di)
     
-            tmp1 = self.matmul_5d(gh, dX, use_pyccel=self.use_pyccel)
-            tmp2 = self.matmul_5d(dgh_di, dX, use_pyccel=self.use_pyccel)
-            tmp3 = self.matmul_5d(gh, ddX_di, use_pyccel=self.use_pyccel)
-
-            term1 = self.matmul_5d(ddX_di, tmp1, transpose_a=True, use_pyccel=self.use_pyccel)
-            term2 = self.matmul_5d(dX, tmp2, transpose_a=True, use_pyccel=self.use_pyccel)
-            term3 = self.matmul_5d(dX, tmp3, transpose_a=True, use_pyccel=self.use_pyccel)
-        # single point evaluation
-        else:
-            dq1_di = dX[0, comp]
-            dgh_di = dgh_dq1 * dq1_di
-        
         tmp1 = gh @ dX    
         tmp2 = dgh_di @ dX
         tmp3 = gh @ ddX_di
 
-            term1 = ddX_di.T @ tmp1
-            term2 = dX.T @ tmp2
-            term3 = dX.T @ tmp3
+        ddX_diT = ddX_di.swapaxes(-1, -2)
+        dXT = dX.swapaxes(-1, -2)
+        
+        term1 = ddX_diT @ tmp1
+        term2 = dXT @ tmp2
+        term3 = dXT @ tmp3
 
         return term1 + term2 + term3
 
@@ -412,7 +389,6 @@ class GVEC_domain:
         -------
         A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).
         '''
-        print(f'pest: {flat_eval = }')
         return self.f_gvec(*self.Pmap(s2, th2, ze2, flat_eval=flat_eval), flat_eval=flat_eval)
 
     def df_pest(self, s2, th2, ze2, flat_eval=False):
@@ -430,20 +406,10 @@ class GVEC_domain:
         -------
         A 2d numpy array (single-point evaluation) or a 5d numpy array, 
         with the last two indices referring to the Jacobian entries at each point.'''
-        print(f'df_pest: {flat_eval = }')
         tmp1 = self.dP(s2, th2, ze2, flat_eval=flat_eval)  # Jacobian of the Pmap
-        print(f'{tmp1.shape = }')
         tmp2 = self.df_gvec(*self.Pmap(s2, th2, ze2, flat_eval=flat_eval), flat_eval=flat_eval)  # Jacobian of f_gvec
-        print(f'{tmp2.shape = }')
 
-        # # meshgrid evaluation
-        # if tmp1.ndim == 5:
-        #     out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel)
-        # # single-point evaluation
-        # else:
-        out = tmp2 @ tmp1
-
-        return out
+        return tmp2 @ tmp1
 
     def det_df_pest(self, s2, th2, ze2, flat_eval=False):
         '''The Jacobian determinant of f_pest.
@@ -459,7 +425,6 @@ class GVEC_domain:
         Returns
         -------
         A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation).'''
-        print(f'det_df_pest: {flat_eval = }')
         tmp1 = self.det_dP(s2, th2, ze2, flat_eval=flat_eval)
         tmp2 = self.det_df_gvec(*self.Pmap(s2, th2, ze2, flat_eval=flat_eval), flat_eval=flat_eval)
         return tmp1 * tmp2
@@ -478,7 +443,6 @@ class GVEC_domain:
         Returns
         -------
         A 2d numpy array (single-point evaluation) or a 5d numpy array.'''
-        print(f'df_pest_inv: {flat_eval = }')
         return np.linalg.inv(self.df_pest(s2, th2, ze2, flat_eval=flat_eval))
 
     def g_pest(self, s2, th2, ze2, flat_eval=False):
@@ -495,20 +459,11 @@ class GVEC_domain:
         Returns
         -------
         A 2d numpy array (single-point evaluation) or a 5d numpy array.'''
-        print(f'g_pest: {flat_eval = }')
+
         tmp = np.ascontiguousarray(self.df_pest(s2, th2, ze2, flat_eval=flat_eval))  # Jacobian of f_gvec
-        print(f'{tmp.shape = }')
         tmpT = tmp.swapaxes(-1, -2)
-        print(f'{tmpT.shape = }')
-        # # meshgrid evaluation
-        # if tmp.ndim == 5:
-        #     out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel)
-        # # single-point evaluation
-        # else:
-        out = tmpT @ tmp
-        print(f'{out.shape = }')
         
-        return out
+        return tmpT @ tmp
 
     def g_pest_inv(self, s2, th2, ze2, flat_eval=False):
         '''Inverse metric tensor of f_pest.
@@ -524,7 +479,6 @@ class GVEC_domain:
         Returns
         -------
         A 2d numpy array (single-point evaluation) or a 5d numpy array.'''
-        print(f'g_pest_inv: {flat_eval = }')
         return np.linalg.inv(self.g_pest(s2, th2, ze2, flat_eval=flat_eval))
 
     # unit cube coordinates
@@ -564,14 +518,7 @@ class GVEC_domain:
         tmp1 = self.dL(s, u, v, flat_eval=flat_eval)  # Jacobian of the Lmap
         tmp2 = self.df_gvec(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval)  # Jacobian of f_gvec
 
-        # meshgrid evaluation
-        if tmp1.ndim == 5:
-            out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel)
-        # single-point evaluation
-        else:
-            out = tmp2 @ tmp1
-
-        return out
+        return tmp2 @ tmp1
 
     def det_df_unit(self, s, u, v, flat_eval=False):
         '''The Jacobian determinant of f_unit.
@@ -623,15 +570,9 @@ class GVEC_domain:
         A 2d numpy array (single-point evaluation) or a 5d numpy array.'''
 
         tmp = np.ascontiguousarray(self.df_unit(s, u, v, flat_eval=flat_eval))  # Jacobian of f_unit
+        tmpT = tmp.swapaxes(-1, -2)
         
-        # meshgrid evaluation
-        if tmp.ndim == 5:
-            out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel)
-        # single-point evaluation
-        else:
-            out = tmp.T @ tmp
-
-        return out
+        return tmpT @ tmp
 
     def g_unit_inv(self, s, u, v, flat_eval=False):
         '''The inverse metric tensor of f_unit.
@@ -688,14 +629,7 @@ class GVEC_domain:
         tmp1 = self.dL(s, u, v, flat_eval=flat_eval)  # Jacobian of the Lmap
         tmp2 = self.df_pest(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval)  # Jacobian of f_pest
 
-        # meshgrid evaluation
-        if tmp1.ndim == 5:
-            out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel)
-        # single-point evaluation
-        else:
-            out = tmp2 @ tmp1
-
-        return out
+        return tmp2 @ tmp1
 
     def det_df_unit_pest(self, s, u, v, flat_eval=False):
         '''The Jacobian determinant of f_unit_pest.
@@ -747,15 +681,9 @@ class GVEC_domain:
         A 2d numpy array (single-point evaluation) or a 5d numpy array.'''
 
         tmp = np.ascontiguousarray(self.df_unit_pest(s, u, v, flat_eval=flat_eval))  # Jacobian of f_unit_pest
+        tmpT = tmp.swapaxes(-1, -2)
 
-        # meshgrid evaluation
-        if tmp.ndim == 5:
-            out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel)
-        # single-point evaluation
-        else:
-            out = tmp.T @ tmp
-
-        return out
+        return tmpT @ tmp
 
     def g_unit_pest_inv(self, s, u, v, flat_eval=False):
         '''The inverse metric tensor of f_unit_pest.
@@ -860,15 +788,9 @@ class GVEC_domain:
         A 2d numpy array (single-point evaluation) or a 5d numpy array.'''
 
         tmp = np.ascontiguousarray(self.df_wo_hmap(s, th, ze, flat_eval=flat_eval))  # Jacobian of f_wo_hmap
+        tmpT = tmp.swapaxes(-1, -2)
 
-        # meshgrid evaluation
-        if tmp.ndim == 5:
-            out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel)
-        # single-point evaluation
-        else:
-            out = tmp.T @ tmp
-
-        return out
+        return tmpT @ tmp
 
     def g_wo_hmap_inv(self, s, th, ze, flat_eval=False):
         '''Inverse metric tensor of f_wo_hmap.
@@ -1167,16 +1089,10 @@ class GVEC_domain:
         '''Jacobian of the Pmap.'''
 
         s, th, ze = self.Pmap(s2, th2, ze2, flat_eval=flat_eval)
-        print(f'{s.shape = }')
-        print(f'{th.shape = }')
-        print(f'{ze.shape = }')
 
         dlambda_ds = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='s', flat_eval=flat_eval)
         dlambda_dth = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='th', flat_eval=flat_eval)
         dlambda_dze = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='ze', flat_eval=flat_eval)
-        print(f'{dlambda_ds.shape = }')
-        print(f'{dlambda_dth.shape = }')
-        print(f'{dlambda_dze.shape = }')
 
         dth_ds2 = - dlambda_ds / (1. + dlambda_dth)
         dth_dth2 = 1. / (1. + dlambda_dth)
@@ -1189,8 +1105,6 @@ class GVEC_domain:
                       (dth_ds2, dth_dth2, dth_dze2),
                       (zmat, zmat, omat)))
 
-        print(f'{J.shape = }')
-
         return self.swap_J_axes(J)
 
     def dP_inv(self, s2, th2, ze2, flat_eval=False):
@@ -1222,14 +1136,7 @@ class GVEC_domain:
         tmp1 = self.dL(s, u, v)  # Jacobian of the Lmap
         tmp2 = self.dP(*self.Lmap(s, u, v))  # Jacobian of the Pmap
 
-        # meshgrid evaluation
-        if tmp1.ndim == 5:
-            out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel)
-        # single-point evaluation
-        else:
-            out = tmp2 @ tmp1
-
-        return out
+        return tmp2 @ tmp1
 
     def dPL_inv(self, s, u, v):
         '''Inverse Jacobian matrix of the composition P ° L.'''
@@ -1355,53 +1262,6 @@ class GVEC_domain:
 
         return out
 
-    @staticmethod
-    def matmul_5d(a, b, transpose_a=False, use_pyccel=False):
-        '''Performs matrix multiplication a*b in the last two indices of two 5-dim arrays.
-        If transpose_a=False and a is of shape (s0, s1, s2, l, m), then b must be of shape (s0, s1, s2, m, n).
-        
-        Parameters
-        ----------
-            a, b : array[float]
-                5-dimensional matrices to multiply.
-                
-            transpose_a : bool
-                Whether to transpose a.
-
-            use_pyccel : bool
-                Whether to use an accelerated kernel for the multiplication.
-                
-        Returns
-        -------
-            A 5-dim C-contiguous array with the result.'''
-        
-        a_shp = a.shape
-
-        assert isinstance(a, np.ndarray)
-        assert isinstance(b, np.ndarray)
-        assert a.ndim == 5
-        assert b.ndim == 5
-        assert a.shape[:3] == b.shape[:3]
-
-        if transpose_a:
-            tmp = GVEC_domain.transpose_5d(a, use_pyccel=use_pyccel)
-        else:
-            tmp = np.ascontiguousarray(a)
-
-        assert tmp.shape[4] == b.shape[3]
-
-        out = np.ascontiguousarray(np.empty(shape=(*a_shp[:3], tmp.shape[3], b.shape[4])))
-
-        if use_pyccel:
-            matmul_5d_kernel(tmp, np.ascontiguousarray(b), out)
-        else:
-            for i in range(a_shp[0]):
-                for j in range(a_shp[1]):
-                    for k in range(a_shp[2]):
-                        out[i, j, k] = tmp[i, j, k] @ b[i, j, k]
-
-        return out
-
     @staticmethod
     def matvec_5d(a, b, transpose_a=False, use_pyccel=False):
         '''Performs matrix-vector multiplication a*b with the last two indices of the 5-dim array a.

src/gvec_to_python/geometry/kernels.py

-def matmul_5d_kernel(a: 'float[:,:,:,:,:]', b: 'float[:,:,:,:,:]', out: 'float[:,:,:,:,:]'):
-    '''Performs matrix multiplication in the last two indices of 5-dim arrays.
-    If a is shape shape (s0, s1, s2, l, m) then b must be shape shape (s0, s1, s2, m, n).'''
-
-    out[:] = 0.
-
-    for i in range(out.shape[0]):
-        for j in range(out.shape[1]):
-            for k in range(out.shape[2]):
-                for l in range(out.shape[3]):
-                    for n in range(out.shape[4]):
-                        for m in range(a.shape[4]):
-                            out[i, j, k, l, n] += a[i, j, k, l, m] * b[i, j, k, m, n]
-
 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).

tests/test_gvec_domain.py

@@ -90,9 +90,6 @@ def test_output_formats(mapping, use_pyccel,gvec_file):
         a1 = np.linspace(0, 1, Np, endpoint=False)
         a2 = np.linspace(0, 1, Np, endpoint=False)
         
-    print(f'{s.shape = }')
-    print(f'{a1.shape = }')
-    print(f'{a2.shape = }')
     shp3 = s.shape
     shp5 = (s.size, 3, 3)