Use spline class directly in low-level code

pygyro/advection/Makefile

@@ -18,7 +18,7 @@ all: accelerated_advection_steps$(SO_EXT)
 	@rm -f .ACC.*
 	@touch $@
 
-DEPS := ../initialisation/$(NAME_PREFIX)initialiser_funcs.py ../splines/$(NAME_PREFIX)spline_eval_funcs.py ../splines/$(NAME_PREFIX)cubic_uniform_spline_eval_funcs.py
+DEPS := ../initialisation/$(NAME_PREFIX)initialiser_funcs.py ../splines/$(NAME_PREFIX)splines.py
 
 ifneq ($(ACC), numba)
 ifneq ($(ACC), pythran)

pygyro/advection/advection.py

@@ -283,10 +283,7 @@ def step(self, f: np.ndarray, cIdx: int, rIdx: int = 0):
             get_lagrange_vals(i, self._shifts[rIdx, cIdx],
                               self._LagrangeVals, self._points[0],
                               self._thetaShifts[rIdx, cIdx],
-                              self._thetaSpline.basis.knots,
-                              self._thetaSpline.basis.degree,
-                              self._thetaSpline.coeffs,
-                              self._thetaSpline.basis.cubic_uniform)
+                              self._thetaSpline)
 
         flux_advection(*self._nPoints, f,
                        self._lagrangeCoeffs[rIdx, cIdx],
@@ -366,12 +363,11 @@ def step(self, f: np.ndarray, dt: float, c: float, r: float):
         self._interpolator.compute_interpolant(f, self._spline)
 
         v_parallel_advection_eval_step(f, self._points-c*dt, r, self._points[0],
-                                       self._points[-1], self._spline.basis.knots,
-                                       self._spline.basis.degree, self._spline.coeffs,
+                                       self._points[-1], self._spline,
                                        self._constants.CN0, self._constants.kN0,
                                        self._constants.deltaRN0, self._constants.rp,
                                        self._constants.CTi, self._constants.kTi,
-                                       self._constants.deltaRTi, self._edgeType, self._spline.basis.cubic_uniform)
+                                       self._constants.deltaRTi, self._edgeType)
 
     def gridStep(self, grid: Grid, phi: Grid, parGrad: ParallelGradient, parGradVals: np.array, dt: float):
         for i, r in grid.getCoords(0):
@@ -474,41 +470,30 @@ def step(self, f: np.ndarray, dt: float, phi: Spline2D, v: float):
         assert f.shape == self._nPoints
         self._interpolator.compute_interpolant(f, self._spline)
 
-        phiBases = phi.basis
-        polBases = self._spline.basis
-
         if (self._explicit):
             poloidal_advection_step_expl(f, float(dt), v, self._points[1],
                                          self._points[0], self._drPhi_0,
                                          self._dqPhi_0, self._drPhi_k,
                                          self._dqPhi_k, self._endPts_k1_q,
                                          self._endPts_k1_r, self._endPts_k2_q,
-                                         self._endPts_k2_r, phiBases[0].knots,
-                                         phiBases[1].knots, phi.coeffs,
-                                         phiBases[0].degree, phiBases[1].degree,
-                                         polBases[0].knots, polBases[1].knots,
-                                         self._spline.coeffs, polBases[0].degree,
-                                         polBases[1].degree, self._constants.CN0,
+                                         self._endPts_k2_r, phi,
+                                         self._spline, self._constants.CN0,
                                          self._constants.kN0, self._constants.deltaRN0,
                                          self._constants.rp, self._constants.CTi,
                                          self._constants.kTi, self._constants.deltaRTi,
-                                         self._constants.B0, phiBases[0].cubic_uniform, self._nulEdge)
+                                         self._constants.B0, self._nulEdge)
         else:
             poloidal_advection_step_impl(f, float(dt), v, self._points[1],
                                          self._points[0], self._drPhi_0,
                                          self._dqPhi_0, self._drPhi_k,
                                          self._dqPhi_k, self._endPts_k1_q,
                                          self._endPts_k1_r, self._endPts_k2_q,
-                                         self._endPts_k2_r, phiBases[0].knots,
-                                         phiBases[1].knots, phi.coeffs,
-                                         phiBases[0].degree, phiBases[1].degree,
-                                         polBases[0].knots, polBases[1].knots,
-                                         self._spline.coeffs, polBases[0].degree,
-                                         polBases[1].degree, self._constants.CN0,
+                                         self._endPts_k2_r, phi,
+                                         self._spline, self._constants.CN0,
                                          self._constants.kN0, self._constants.deltaRN0,
                                          self._constants.rp, self._constants.CTi,
                                          self._constants.kTi, self._constants.deltaRTi,
-                                         self._constants.B0, self._TOL, phiBases[0].cubic_uniform, self._nulEdge)
+                                         self._constants.B0, self._TOL, self._nulEdge)
 
     def exact_step(self, f, endPts, v):
         assert f.shape == self._nPoints

pygyro/poisson/poisson_solver.py

@@ -5,7 +5,7 @@
 from numpy.polynomial.legendre import leggauss
 
 from ..model.grid import Grid
-from ..splines.splines import BSplines, Spline1D, make_knots
+from ..splines.splines import BSplines, Spline1D, Spline1DComplex, make_knots
 from ..splines.spline_interpolators import SplineInterpolator1D
 from ..initialisation import initialiser_funcs as init
 from .poisson_tools import get_perturbed_rho, get_rho
@@ -254,24 +254,32 @@ def __init__(self, degree: int, rspline: BSplines, nr: int, nTheta: int,
                 # Find the integral of the multiplication of these splines
                 # and their coefficients and save the value in the
                 # appropriate place for the matrix
+                rEval = np.empty_like(evalPts)
+                rDeriv = np.empty_like(evalPts)
+                splEval = np.empty_like(evalPts)
+                splDeriv = np.empty_like(evalPts)
+                self._rspline[s_j].eval_vector(evalPts, rEval)
+                spline.eval_vector(evalPts, splEval)
+                self._rspline[s_j].eval_vector(evalPts, rDeriv, der=1)
+                spline.eval_vector(evalPts, splDeriv, der=1)
                 massCoeffs[j][i] = np.sum(np.tile(self._weights, end-start) * multFactor *
-                                          rhoFactor(evalPts) * self._rspline[s_j].eval(evalPts) * spline.eval(evalPts) * evalPts)
+                                          rhoFactor(evalPts) * rEval * splEval * evalPts)
                 k2PhiPsiCoeffs[j][i] = np.sum(np.tile(self._weights, end-start) * multFactor *
-                                              ddThetaFactor(evalPts) * self._rspline[s_j].eval(evalPts) * spline.eval(evalPts) * evalPts)
+                                              ddThetaFactor(evalPts) * rEval * splEval * evalPts)
                 PhiPsiCoeffs[j][i] = np.sum(np.tile(self._weights, end-start) * multFactor *
-                                            rFactor(evalPts) * self._rspline[s_j].eval(evalPts) * spline.eval(evalPts) * evalPts)
+                                            rFactor(evalPts) * rEval * splEval * evalPts)
                 dPhidPsi = np.sum(np.tile(self._weights, end-start) * multFactor *
-                                  -ddrFactor(evalPts) * self._rspline[s_j].eval(evalPts, 1) * spline.eval(evalPts, 1) * evalPts)
+                                  -ddrFactor(evalPts) * rDeriv * splDeriv * evalPts)
                 dPhidPsiCoeffs[j][i] = dPhidPsi + \
                     np.sum(np.tile(self._weights, end-start) * multFactor *
-                           -ddrFactor(evalPts) * self._rspline[s_j].eval(evalPts, 1) * spline.eval(evalPts))
+                           -ddrFactor(evalPts) * rDeriv * splEval)
                 dPhidPsiCoeffs[self._rspline.degree*2-j][i] = dPhidPsi + \
                     np.sum(np.tile(self._weights, end-start) * multFactor *
-                           -ddrFactor(evalPts) * self._rspline[s_j].eval(evalPts) * spline.eval(evalPts, 1))
+                           -ddrFactor(evalPts) * rEval * splDeriv)
                 dPhiPsiCoeffs[j][i] = np.sum(np.tile(self._weights, end-start) * multFactor *
-                                             drFactor(evalPts) * self._rspline[s_j].eval(evalPts, 1) * spline.eval(evalPts) * evalPts)
+                                             drFactor(evalPts) * rDeriv * splEval * evalPts)
                 dPhiPsiCoeffs[self._rspline.degree*2-j][i] = np.sum(np.tile(self._weights, end-start) * multFactor *
-                                                                    drFactor(evalPts) * self._rspline[s_j].eval(evalPts) * spline.eval(evalPts, 1) * evalPts)
+                                                                    drFactor(evalPts) * rEval * splDeriv * evalPts)
 
         # Create the diagonal matrices
         # Diagonal matrices contain many 0 valued points so sparse
@@ -299,7 +307,7 @@ def __init__(self, degree: int, rspline: BSplines, nr: int, nTheta: int,
 
         # Create the tools required for the interpolation
         self._interpolator = SplineInterpolator1D(self._rspline, dtype=complex)
-        self._spline = Spline1D(self._rspline, np.complex128)
+        self._spline = Spline1DComplex(self._rspline)
         self._real_spline = Spline1D(self._rspline)
 
         self._realMem = np.empty(nr)

pygyro/splines/Makefile

@@ -8,7 +8,7 @@
 # Main targets
 #----------------------------------------------------------
 
-all: spline_eval_funcs$(SO_EXT) cubic_uniform_spline_eval_funcs$(SO_EXT)
+all: spline_eval_funcs$(SO_EXT) cubic_uniform_spline_eval_funcs$(SO_EXT) splines$(SO_EXT)
 
 .ACC.$(ACC):
 	@rm -f .ACC.*
@@ -28,6 +28,13 @@ else
 	$(TOOL) $< $(TOOL_FLAGS) -o $@
 endif
 
+splines$(SO_EXT): $(NAME_PREFIX)splines.py .ACC.$(ACC) spline_eval_funcs$(SO_EXT) cubic_uniform_spline_eval_funcs$(SO_EXT)
+ifneq ($(ACC), pythran)
+	$(TOOL) $< $(TOOL_FLAGS)
+else
+	$(TOOL) $< $(TOOL_FLAGS) -o $@
+endif
+
 
 clean:
 	rm -f *.o *.so *.mod .ACC.pycc .lock_acquisition.lock

pygyro/splines/cubic_uniform_spline_eval_funcs.py

@@ -1,7 +1,9 @@
-from typing import Final
+from typing import Final, TypeVar
 from pyccel.decorators import pure, stack_array
 from numpy import empty
 
+CoeffType = TypeVar('CoeffType', float, complex)
+
 
 @pure
 def cu_find_span(xmin: 'float', xmax: 'float', dx: 'float', x: 'float', ncells: 'int'):
@@ -132,7 +134,7 @@ def cu_basis_funs_1st_der(span: 'int', offset: 'float', dx: 'float', ders: 'floa
 
 @pure
 @stack_array('basis')
-def cu_eval_spline_1d_scalar(x: 'float', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[float[:]]', der: 'int') -> 'float':
+def cu_eval_spline_1d_scalar(x: 'float', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[CoeffType[:]]', der: 'int') -> 'CoeffType':
     """
     TODO
     """
@@ -146,15 +148,15 @@ def cu_eval_spline_1d_scalar(x: 'float', knots: 'Final[float[:]]', degree: 'int'
     elif (der == 1):
         cu_basis_funs_1st_der(span, offset, dx, basis)
 
-    y = 0.0
+    y = 0.0*coeffs[0]
     for j in range(4):
         y += coeffs[span-3+j]*basis[j]
     return y
 
 
 @pure
 @stack_array('basis')
-def cu_eval_spline_1d_vector(x: 'Final[float[:]]', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[float[:]]', y: 'float[:]', der: 'int' = 0):
+def cu_eval_spline_1d_vector(x: 'Final[float[:]]', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[CoeffType[:]]', y: 'CoeffType[:]', der: 'int' = 0):
     """
     TODO
     """
@@ -184,7 +186,7 @@ def cu_eval_spline_1d_vector(x: 'Final[float[:]]', knots: 'Final[float[:]]', deg
 @pure
 @stack_array('basis1', 'basis2', 'theCoeffs')
 def cu_eval_spline_2d_scalar(x: 'float', y: 'float', kts1: 'Final[float[:]]', deg1: 'int', kts2: 'Final[float[:]]', deg2: 'int',
-                             coeffs: 'Final[float[:,:]]', der1: 'int' = 0, der2: 'int' = 0) -> 'float':
+                             coeffs: 'Final[CoeffType[:,:]]', der1: 'int' = 0, der2: 'int' = 0) -> 'CoeffType':
     """
     TODO
     """
@@ -208,10 +210,10 @@ def cu_eval_spline_2d_scalar(x: 'float', y: 'float', kts1: 'Final[float[:]]', de
     elif (der2 == 1):
         cu_basis_funs_1st_der(span2, offset2, dy, basis2)
 
-    theCoeffs = empty((4, 4))
+    theCoeffs = empty((4, 4), dtype=type(coeffs[0, 0]))
     theCoeffs[:, :] = coeffs[span1-deg1:span1+1, span2-deg2:span2+1]
 
-    z = 0.0
+    z = 0.0*coeffs[0, 0]
     for i in range(4):
         theCoeffs[i, 0] = theCoeffs[i, 0]*basis2[0]
         for j in range(1, 4):
@@ -223,7 +225,7 @@ def cu_eval_spline_2d_scalar(x: 'float', y: 'float', kts1: 'Final[float[:]]', de
 @pure
 @stack_array('basis1', 'basis2', 'theCoeffs')
 def cu_eval_spline_2d_cross(X: 'Final[float[:]]', Y: 'Final[float[:]]', kts1: 'Final[float[:]]', deg1: 'int', kts2: 'Final[float[:]]', deg2: 'int',
-                            coeffs: 'Final[float[:,:]]', z: 'float[:,:]', der1: 'int' = 0, der2: 'int' = 0):
+                            coeffs: 'Final[CoeffType[:,:]]', z: 'CoeffType[:,:]', der1: 'int' = 0, der2: 'int' = 0):
     """
     TODO
     """
@@ -234,7 +236,7 @@ def cu_eval_spline_2d_cross(X: 'Final[float[:]]', Y: 'Final[float[:]]', kts1: 'F
 
     basis1 = empty(4)
     basis2 = empty(4)
-    theCoeffs = empty((4, 4))
+    theCoeffs = empty((4, 4), dtype=type(coeffs[0, 0]))
 
     if (der1 == 0 and der2 == 0):
         for i, x in enumerate(X):
@@ -315,7 +317,7 @@ def cu_eval_spline_2d_cross(X: 'Final[float[:]]', Y: 'Final[float[:]]', kts1: 'F
 @pure
 @stack_array('basis1', 'basis2', 'theCoeffs')
 def cu_eval_spline_2d_vector(x: 'float[:]', y: 'float[:]', kts1: 'float[:]', deg1: 'int', kts2: 'float[:]', deg2: 'int',
-                             coeffs: 'float[:,:]', z: 'float[:]', der1: 'int' = 0, der2: 'int' = 0):
+                             coeffs: 'CoeffType[:,:]', z: 'CoeffType[:]', der1: 'int' = 0, der2: 'int' = 0):
     """
     TODO
     """
@@ -326,7 +328,7 @@ def cu_eval_spline_2d_vector(x: 'float[:]', y: 'float[:]', kts1: 'float[:]', deg
 
     basis1 = empty(4)
     basis2 = empty(4)
-    theCoeffs = empty((4, 4))
+    theCoeffs = empty((4, 4), dtype=type(z[0]))
 
     if (der1 == 0):
         if (der2 == 0):

pygyro/splines/spline_eval_funcs.py

@@ -1,7 +1,9 @@
-from typing import Final
+from typing import Final, TypeVar
 from pyccel.decorators import pure, stack_array
 from numpy import empty
 
+CoeffType = TypeVar('CoeffType', float, complex)
+
 
 @pure
 def nu_find_span(knots: 'Final[float[:]]', degree: 'int', x: 'float') -> int:
@@ -167,7 +169,7 @@ def nu_basis_funs_1st_der(knots: 'Final[float[:]]', degree: 'int', x: 'float', s
 
 @pure
 @stack_array('basis')
-def nu_eval_spline_1d_scalar(x: 'float', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[float[:]]', der: 'int') -> 'float':
+def nu_eval_spline_1d_scalar(x: 'float', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[CoeffType[:]]', der: 'int') -> 'CoeffType':
     """
     TODO
     """
@@ -179,15 +181,15 @@ def nu_eval_spline_1d_scalar(x: 'float', knots: 'Final[float[:]]', degree: 'int'
     elif (der == 1):
         nu_basis_funs_1st_der(knots, degree, x, span, basis)
 
-    y = 0.0
+    y = 0.0*coeffs[0]
     for j in range(degree+1):
         y += coeffs[span-degree+j]*basis[j]
     return y
 
 
 @pure
 @stack_array('basis')
-def nu_eval_spline_1d_vector(x: 'Final[float[:]]', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[float[:]]', y: 'float[:]', der: 'int' = 0):
+def nu_eval_spline_1d_vector(x: 'Final[float[:]]', knots: 'Final[float[:]]', degree: 'int', coeffs: 'Final[CoeffType[:]]', y: 'CoeffType[:]', der: 'int' = 0):
     """
     TODO
     """
@@ -215,7 +217,7 @@ def nu_eval_spline_1d_vector(x: 'Final[float[:]]', knots: 'Final[float[:]]', deg
 @pure
 @stack_array('basis1', 'basis2', 'theCoeffs')
 def nu_eval_spline_2d_scalar(x: 'float', y: 'float', kts1: 'Final[float[:]]', deg1: 'int', kts2: 'Final[float[:]]', deg2: 'int',
-                             coeffs: 'Final[float[:,:]]', der1: 'int' = 0, der2: 'int' = 0) -> 'float':
+                             coeffs: 'Final[CoeffType[:,:]]', der1: 'int' = 0, der2: 'int' = 0) -> 'CoeffType':
     """
     TODO
     """
@@ -234,10 +236,10 @@ def nu_eval_spline_2d_scalar(x: 'float', y: 'float', kts1: 'Final[float[:]]', de
     elif (der2 == 1):
         nu_basis_funs_1st_der(kts2, deg2, y, span2, basis2)
 
-    theCoeffs = empty((deg1+1, deg2+1))
+    theCoeffs = empty((deg1+1, deg2+1), dtype=type(coeffs[0, 0]))
     theCoeffs[:, :] = coeffs[span1-deg1:span1+1, span2-deg2:span2+1]
 
-    z = 0.0
+    z = 0.0*coeffs[0, 0]
     for i in range(deg1+1):
         theCoeffs[i, 0] = theCoeffs[i, 0]*basis2[0]
         for j in range(1, deg2+1):
@@ -249,13 +251,13 @@ def nu_eval_spline_2d_scalar(x: 'float', y: 'float', kts1: 'Final[float[:]]', de
 @pure
 @stack_array('basis1', 'basis2', 'theCoeffs')
 def nu_eval_spline_2d_cross(X: 'Final[float[:]]', Y: 'Final[float[:]]', kts1: 'Final[float[:]]', deg1: 'int', kts2: 'Final[float[:]]', deg2: 'int',
-                            coeffs: 'Final[float[:,:]]', z: 'float[:,:]', der1: 'int' = 0, der2: 'int' = 0):
+                            coeffs: 'Final[CoeffType[:,:]]', z: 'CoeffType[:,:]', der1: 'int' = 0, der2: 'int' = 0):
     """
     TODO
     """
     basis1 = empty(deg1+1)
     basis2 = empty(deg2+1)
-    theCoeffs = empty((deg1+1, deg2+1))
+    theCoeffs = empty((deg1+1, deg2+1), dtype=type(z[0, 0]))
 
     if (der1 == 0 and der2 == 0):
         for i, x in enumerate(X):
@@ -337,13 +339,13 @@ def nu_eval_spline_2d_cross(X: 'Final[float[:]]', Y: 'Final[float[:]]', kts1: 'F
 @pure
 @stack_array('basis1', 'basis2', 'theCoeffs')
 def nu_eval_spline_2d_vector(x: 'float[:]', y: 'float[:]', kts1: 'float[:]', deg1: 'int', kts2: 'float[:]', deg2: 'int',
-                             coeffs: 'float[:,:]', z: 'float[:]', der1: 'int' = 0, der2: 'int' = 0):
+                             coeffs: 'CoeffType[:,:]', z: 'CoeffType[:]', der1: 'int' = 0, der2: 'int' = 0):
     """
     TODO
     """
     basis1 = empty(deg1+1)
     basis2 = empty(deg2+1)
-    theCoeffs = empty((deg1+1, deg2+1))
+    theCoeffs = empty((deg1+1, deg2+1), dtype=type(z[0]))
 
     if (der1 == 0):
         if (der2 == 0):