Added equilibrium currentdensity callables

src/gvec_to_python/GVEC_functions.py

@@ -147,18 +147,22 @@ class GVEC:
             self.l = None
             self.dl = None
             self.dl_inv = None
+            self.det_dl = None
         elif mapping == 'pest':
             self.l = self.domain.Pmap
             self.dl = self.domain.dP
             self.dl_inv = self.domain.dP_inv
+            self.det_dl = self.domain.det_dP
         elif mapping == 'unit':
             self.l = self.domain.Lmap
             self.dl = self.domain.dL
             self.dl_inv = self.domain.dL_inv
+            self.det_dl = self.domain.det_dL
         elif mapping == 'unit_pest':
             self.l = self.domain.PLmap
             self.dl = self.domain.dPL
             self.dl_inv = self.domain.dPL_inv
+            self.det_dl = self.domain.det_dPL
 
         self._mapping = mapping
 
@@ -378,6 +382,74 @@ class GVEC:
         '''
         return self.transform(self.a1(s, a1, a2), s, a1, a2, kind='push_1form'), self.f(s, a1, a2)
 
+    def jv(self, s, a1, a2):
+        '''Contra-variant components (vector-field) of the equilibirum current curl B / mu0 in one of the coordinates "gvec", "pest", "unit" or unit_pest.
+        Coordinates can be chosen using the mapping setter.
+
+        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. 
+
+        Returns
+        -------
+            A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).
+        '''
+        return self.transform(self.j2(s, a1, a2), s, a1, a2, kind='2_to_v')
+
+    def j1(self, s, a1, a2):
+        '''Co-variant components (1-form) of the equilibirum current curl B / mu0 in one of the coordinates "gvec", "pest", "unit" or unit_pest.
+        Coordinates can be chosen using the mapping setter.
+
+        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. 
+
+        Returns
+        -------
+            A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).
+        '''
+        return self.transform(self.j2(s, a1, a2), s, a1, a2, kind='2_to_1')
+
+    def j2(self, s, a1, a2):
+        '''2-form components of the equilibirum current curl B / mu0 as in one of the coordinates "gvec", "pest", "unit" or unit_pest.
+        Coordinates can be chosen using the mapping setter.
+
+        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. 
+
+        Returns
+        -------
+            A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).
+        '''
+        if self.l is None:
+            return self._j2_gvec(s, a1, a2)
+        else:
+            return self.transform(self._j2_gvec(*self.l(s, a1, a2)), s, a1, a2, kind='pull_2form', full_map=False)
+
+    def j_cart(self, s, a1, a2):
+        '''Cartesian components of equilibirum current curl B / mu0 evaluated at one of the coordinates "gvec", "pest", "unit" or unit_pest.
+        Coordinates can be chosen using the mapping setter.
+
+        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. 
+
+        Returns
+        -------
+            A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).
+                Second entry are the Cartesian coordinates of evaluation points.
+        '''
+        return self.transform(self.jv(s, a1, a2), s, a1, a2, kind='push_vector'), self.f(s, a1, a2)
+
     #---------------------------------
     # Internal methods (gvec fields)
     #---------------------------------
@@ -395,10 +467,10 @@ class GVEC:
             A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).'''
 
         phi = self.profiles.profile(s, th, ze, name='phi')
-        dphi = self.profiles.profile_derivative(s, th, ze, name='phi')
+        dphi = self.profiles.profile(s, th, ze, name='phi', der='s')
         chi = self.profiles.profile(s, th, ze, name='chi')
 
-        return (- self.lambda_(s, th, ze) * dphi, phi, -chi)
+        return (- self._lambda(s, th, ze) * dphi, phi, -chi)
 
     def _bv_gvec(self, s, th, ze):
         '''Magnetic field as vector-field (contra-variant components) in gvec standard coordinates (map f).
@@ -413,27 +485,152 @@ class GVEC:
             A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).'''
 
         det_df = self.domain.det_df_gvec(s, th, ze)
-        dphi = self.profiles.profile_derivative(s, th, ze, name='phi')
-        dchi = self.profiles.profile_derivative(s, th, ze, name='chi')
 
-        return (0.*dphi, (dchi - dphi * self.lambda_ze(s, th, ze)) / det_df, dphi * (1. + self.lambda_th(s, th, ze)) / det_df)
+        return (0.*det_df, self._b_small_th(s, th, ze) / det_df, self._b_small_ze(s, th, ze) / det_df)
+
+    def _j2_gvec(self, s, th, ze):
+        '''Equilibirum current curl B / mu0, as 2-form in gvec standard coordinates (map f).
+
+        Parameters
+        ----------
+            s, th, ze : float or array-like
+                Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid).
+
+        Returns
+        -------
+            A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).'''
+
+        mu0 = 1.2566370621219e-6
+        b = np.ascontiguousarray(self._bv_gvec(s, th, ze))
+
+        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]))
+
+        g = self.domain.dg_gvec(s, th, ze, der=None)
+        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
+
+        j1 = (dgb_dt[2] - dgb_dz[1]) / mu0
+        j2 = (dgb_dz[0] - dgb_ds[2]) / mu0
+        j3 = (dgb_ds[1] - dgb_dt[0]) / mu0
+
+        return j1, j2, j3
+
+    def _grad_b_theta(self, s, th, ze):
+        '''Gradient w.r.t (s, theta, zeta) of the second component of _bv_gvec, returned as 3-tuple.
+        '''
+        det_df = self.domain.det_df_gvec(s, th, ze)
+        dphi = self.profiles.profile(s, th, ze, name='phi', der='s')
+        dchi = self.profiles.profile(s, th, ze, name='chi', der='s')
+        ddphi = self.profiles.profile(s, th, ze, name='phi', der='ss')
+        ddchi = self.profiles.profile(s, th, ze, name='chi', der='ss')
+
+        term1 = ddchi - ddphi * self._lambda_ze(s, th, ze) - dphi * self._lambda_sze(s, th, ze)
+        term2 = - self._b_small_th(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='s') / det_df
+        db_ds = (term1 + term2) / det_df
+
+        term1 = dchi - dphi * self._lambda_thze(s, th, ze) 
+        term2 = - self._b_small_th(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='th') / det_df
+        db_dt = (term1 + term2) / det_df
+
+        term1 = dchi - dphi * self._lambda_zeze(s, th, ze) 
+        term2 = - self._b_small_th(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='ze') / det_df
+        db_dz = (term1 + term2) / det_df
+
+        return db_ds, db_dt, db_dz 
+
+    def _grad_b_zeta(self, s, th, ze):
+        '''Gradient w.r.t (s, theta, zeta) of the third component of _bv_gvec, returned as 3-tuple.
+        '''
+
+        det_df = self.domain.det_df_gvec(s, th, ze)
+        dphi = self.profiles.profile(s, th, ze, name='phi', der='s')
+        ddphi = self.profiles.profile(s, th, ze, name='phi', der='ss')
+
+        term1 = ddphi * (1. + self._lambda_th(s, th, ze)) + dphi * self._lambda_sth(s, th, ze)
+        term2 = - self._b_small_ze(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='s') / det_df
+        db_ds = (term1 + term2) / det_df
+
+        term1 = dphi * self._lambda_thth(s, th, ze) 
+        term2 = - self._b_small_ze(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='th') / det_df
+        db_dt = (term1 + term2) / det_df
+
+        term1 = dphi * self._lambda_thze(s, th, ze) 
+        term2 = - self._b_small_ze(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='ze') / det_df
+        db_dz = (term1 + term2) / det_df
+
+        return db_ds, db_dt, db_dz 
+
+    def _b_small_th(self, s, th, ze):
+        '''B^theta * J'''
+
+        dphi = self.profiles.profile(s, th, ze, name='phi', der='s')
+        dchi = self.profiles.profile(s, th, ze, name='chi', der='s')
+
+        return dchi - dphi * self._lambda_ze(s, th, ze)
+
+    def _b_small_ze(self, s, th, ze):
+        '''B^zeta * J'''
+
+        dphi = self.profiles.profile(s, th, ze, name='phi', der='s')
+
+        return dphi * (1. + self._lambda_th(s, th, ze))
 
     # lambda function
-    def lambda_(self, s, th, ze):
+    def _lambda(self, s, th, ze):
         '''The lambda function from gvec.'''
         return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef)
 
-    def lambda_s(self, s, th, ze):
+    def _lambda_s(self, s, th, ze):
         '''s-derivative of the lambda function from gvec.'''
         return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='s')
 
-    def lambda_th(self, s, th, ze):
+    def _lambda_th(self, s, th, ze):
         '''theta-derivative of the lambda function from gvec.'''
         return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='th')
 
-    def lambda_ze(self, s, th, ze):
+    def _lambda_ze(self, s, th, ze):
         '''zeta-derivative of the lambda function from gvec.'''
         return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='ze')
+
+    def _lambda_ss(self, s, th, ze):
+        '''Second derivative of the lambda function from gvec.'''
+        return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='ss')
+
+    def _lambda_thth(self, s, th, ze):
+        '''Second derivative of the lambda function from gvec.'''
+        return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='thth')
+
+    def _lambda_zeze(self, s, th, ze):
+        '''Second derivative of the lambda function from gvec.'''
+        return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='zeze')
+
+    def _lambda_sth(self, s, th, ze):
+        '''Second derivative of the lambda function from gvec.'''
+        return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='sth')
+
+    def _lambda_sze(self, s, th, ze):
+        '''Second derivative of the lambda function from gvec.'''
+        return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='sze')
+
+    def _lambda_thze(self, s, th, ze):
+        '''Second derivative of the lambda function from gvec.'''
+        return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='thze')
     
     #---------------
     # Helper methods
@@ -444,6 +641,7 @@ class GVEC:
         
         pull_vector:    out = dm_inv * b    (dm is the Jacobian of m)
         pull_1form:     out = dm.T * b
+        pull_2form:     out = det_dm * dm_inv * b
         push_vector:    out = dm * b
         push_1form:     out = dm_inv.T * b
         v_to_1:         out = g * b         (g is the metric tensor of m)
@@ -497,6 +695,16 @@ class GVEC:
                     mat = self.dl(s, a1, a2)
                     transpose = True
 
+        elif kind == 'pull_2form':
+            if full_map:
+                raise NotImplementedError
+            else:
+                if self.mapping == 'gvec':
+                    return b
+                else:
+                    mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.dl_inv(s, a1, a2)) * self.det_dl(s, a1, a2))
+                    transpose = False
+
         # push to physical
         elif kind == 'push_vector':
             if full_map:

src/gvec_to_python/base/base.py

@@ -16,20 +16,20 @@ class SplineFourierBases:
     
     Parameters
     ----------
-        sgrid : array[float]
-            Cell boundaries in s-direction in [0, 1].
+    sgrid : array[float]
+        Cell boundaries in s-direction in [0, 1].
 
-        degree : int
-            Degree of spline space.
-            
-        sin_cos : int 
-            Type of Fourier basis: 1: use only sine, 2: use only cosine, 3: use both.
+    degree : int
+        Degree of spline space.
+        
+    sin_cos : int 
+        Type of Fourier basis: 1: use only sine, 2: use only cosine, 3: use both.
 
-        mn : list
-            mn-mode numbers, with NFP premultiplied into the n-modes.
+    mn : list
+        mn-mode numbers, with NFP premultiplied into the n-modes.
 
-        range_sin : array-like
-            Index range of sine modes in `mn` list."""
+    range_sin : array-like
+        Index range of sine modes in `mn` list."""
 
     def __init__(self, sgrid, degree, sin_cos, mn, range_sin):
 
@@ -66,22 +66,25 @@ class SplineFourierBases:
 
         Parameters
         ----------
-            s : float or meshgrid numpy.ndarray
-                Logical coordinate in radial direction.
+        s : float or meshgrid numpy.ndarray
+            Logical coordinate in radial direction.
 
-            th : float or meshgrid numpy.ndarray
-                Angle theta in Tokamak coordinate along poloidal direction.
+        th : float or meshgrid numpy.ndarray
+            Angle theta in Tokamak coordinate along poloidal direction.
 
-            ze : float or meshgrid numpy.ndarray
-                Angle zeta in Tokamak coordinate along toroidal direction.
+        ze : float or meshgrid numpy.ndarray
+            Angle zeta in Tokamak coordinate along toroidal direction.
 
-            coef : array_like
-                A list of B-spline control points (coefficients) for each (m, n) Fourier mode.
+        coef : array_like
+            A list of B-spline control points (coefficients) for each (m, n) Fourier mode.
+
+        der : str
+            Which derivative to evaluate: None, "s", "th", "ze", "ss", "thth", "zeze", "sth", "sze" or "thze".
 
         Returns
         -------
-            float or meshgrid numpy.ndarray
-                Evaluated coordinate.
+        float or meshgrid numpy.ndarray
+            Evaluated coordinate.
         """
 
         s, th, ze = GVEC_domain.prepare_args(s, th, ze)

src/gvec_to_python/equilibrium/profiles.py

@@ -17,17 +17,17 @@ class GVEC_profiles:
     
     Parameters
     ----------
-        sgrid : array[float]
-            Cell boundaries in s-direction in [0, 1].
-            
-        degree : int
-            Degree for spline interpolation of the profiles (spos_grev = sgrid - 1 + degree).
-            
-        spos_grev : array[float]
-            Greville points in s-direction in [0, 1].
-            
-        phi_grev, chi_grev, iota_grev, pres_grev : array[float]
-            Profile values at spos_grev of "phi" (toroidal flux), "chi" (poloidal flux) "iota" and "pressure".'''
+    sgrid : array[float]
+        Cell boundaries in s-direction in [0, 1].
+        
+    degree : int
+        Degree for spline interpolation of the profiles (spos_grev = sgrid - 1 + degree).
+        
+    spos_grev : array[float]
+        Greville points in s-direction in [0, 1].
+        
+    phi_grev, chi_grev, iota_grev, pres_grev : array[float]
+        Profile values at spos_grev of "phi" (toroidal flux), "chi" (poloidal flux) "iota" and "pressure".'''
 
     def __init__(self, sgrid, degree, spos_grev, phi_grev, chi_grev, iota_grev, pres_grev):
 
@@ -69,21 +69,24 @@ class GVEC_profiles:
         '''Spline coefficients of pressure profile for basis of spline_space .'''
         return self._pres_coef
 
-    def profile(self, *args, name='phi'):
-        '''Profile as a function of s in [0, 1]. 
+    def profile(self, *args, name='phi', der=None):
+        '''Profile (derivatives) as a function of s in [0, 1]. 
         
         Parameters
         ----------
-            *args : float or array-like
-                Arguments of the profile function. If one array argument is given, it must be 1d.
-                If three arguments are given (s, a1, a2), a meshgrid evaluation is performed.
-        
-            name : str
-                Profile identifier; must be one of "phi" (toroidal flux), "chi" (poloidal flux) "iota", or "pressure".
+        *args : float or array-like
+            Arguments of the profile function. If one array argument is given, it must be 1d.
+            If three arguments are given (s, a1, a2), a meshgrid evaluation is performed.
+    
+        name : str
+            Profile identifier; must be one of "phi" (toroidal flux), "chi" (poloidal flux) "iota", or "pressure".
+
+        der : str
+            Which derivative to evaluate: None, "s" or "ss".
                 
         Returns
         -------
-            A numpy array.'''
+        A numpy array.'''
 
         if name == 'phi':
             coef = self.phi_coef
@@ -96,54 +99,24 @@ class GVEC_profiles:
         else:
             ValueError(f'Profile name must be "phi", "chi", "iota"  or "pressure", is {name}.')
 
-        if len(args) == 1:
-            return self.spline_space.evaluate_N(args[0], coef)
-
-        elif len(args) == 3:
-            s, a1, a2 = GVEC_domain.prepare_args(args[0], args[1], args[2], sparse=False)
-            if isinstance(s, np.ndarray):
-                vals = self.spline_space.evaluate_N(s.flatten(), coef).reshape(s.shape)
-            else:
-                vals = self.spline_space.evaluate_N(s, coef)
-            return vals
-
-        else:
-            ValueError(f'Number of independent variables must be 1 or 3, is {len(args)}.')
-
-    def profile_derivative(self, *args, name='phi'):
-        '''Profile derivative as a function of s in [0, 1]. 
-        
-        Parameters
-        ----------
-            *args : float or array-like
-                Arguments of the profile derivative. If one array argument is given, it must be 1d.
-                If three arguments are given (s, a1, a2), a meshgrid evaluation is performed.
-        
-            name : str
-                Profile identifier; must be one of "phi" (toroidal flux), "chi" (poloidal flux) "iota", or "pressure".'''
-        
-        if name == 'phi':
-            coef = self.phi_coef
-        elif name == 'chi':
-            coef = self.chi_coef
-        elif name == 'iota':
-            coef = self.iota_coef
-        elif name == 'pressure':
-            coef = self.pres_coef
+        if der is None:
+            _fun = self.spline_space.evaluate_N
+        elif der == 's':
+            _fun = self.spline_space.evaluate_dN
+        elif der == 'ss':
+            _fun = self.spline_space.evaluate_ddN
         else:
-            ValueError(f'Profile name must be "phi", "chi", "iota"  or "pressure", is {name}.')
+            raise ValueError('Only up to second order derivatives implemented.')
 
         if len(args) == 1:
-            return self.spline_space.evaluate_dN(args[0], coef)
-
+            return _fun(args[0], coef)
         elif len(args) == 3:
             s, a1, a2 = GVEC_domain.prepare_args(args[0], args[1], args[2], sparse=False)
             if isinstance(s, np.ndarray):
-                vals = self.spline_space.evaluate_dN(s.flatten(), coef).reshape(s.shape)
+                vals = _fun(s.flatten(), coef).reshape(s.shape)
             else:
-                vals = self.spline_space.evaluate_dN(s, coef)
+                vals = _fun(s, coef)
             return vals
-            
         else:
             ValueError(f'Number of independent variables must be 1 or 3, is {len(args)}.')
 

src/gvec_to_python/geometry/domain.py

@@ -248,6 +248,80 @@ class GVEC_domain:
 
         return np.linalg.inv(self.g_gvec(s, th, ze))
 
+    def dg_gvec(self, s, th, ze, der=None):
+        '''Partial derivative of the metric tensor. der in  {"s", "th", "ze"} indicates which derivative to evaluate.'''
+
+        if der is None:
+            return self.g_gvec(s, th, ze)
+        elif der == 's':
+            comp = 0
+        elif der == 'th':
+            comp = 1
+        elif der == 'ze':
+            comp = 2
+        else:
+            raise ValueError(f'Derivative {der} not defined.')
+
+        dX = self.dX(s, th, ze)
+        ddX_di = self.ddX(s, th, ze)[comp]
+
+        gh = self.gh(*self.Xmap(s, th, ze))
+        dgh_dq1 = self.dgh_dq1(*self.Xmap(s, th, ze))
+
+        # meshgrid evaluation
+        if dX.ndim == 5:
+            dq1_di = dX[:, :, :, 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
+
+        return term1 + term2 + term3
+
+    def dJ_gvec(self, s, th, ze, der=None):
+        '''Partial derivative of Jacobian determinant. der in  {"s", "th", "ze"} indicates which derivative to evaluate.'''
+
+        if der is None:
+            return self.det_df_gvec(s, th, ze)
+        elif der == 's':
+            comp = 0
+        elif der == 'th':
+            comp = 1
+        elif der == 'ze':
+            comp = 2
+        else:
+            raise ValueError(f'Derivative {der} not defined.')
+
+        dX = self.dX(s, th, ze)
+
+        # meshgrid evaluation
+        if dX.ndim == 5:
+            dq1_di = dX[:, :, :, 0, comp]
+        # single point evaluation
+        else:
+            dq1_di = dX[0, comp]
+
+        q1, q2, ze2 = self.Xmap(s, th, ze)
+
+        return q1 * dq1_di * self.det_dX(s, th, ze) + self.det_dh(q1, q2, ze2) * self.dJ_X(s, th, ze)[comp]
+
     # gvec_pest coordinates
     def f_pest(self, s2, th2, ze2):
         '''The map f_pest:(s2, th2, ze2) -> (x, y, z), where (s2, th2, ze2) in [0, 1] x [0, 2*pi]^2 are straight-field-line coordinates (PEST)
@@ -709,6 +783,47 @@ class GVEC_domain:
 
         return self.swap_J_axes(J)
 
+    def ddX(self, s, th, ze):
+        '''Returns a 3-tuple for the three partial derivatives of dX w.r.t s, th and ze.'''
+        s, th, ze = self.prepare_args(s, th, ze)
+
+        # partial derivative dJ/ds
+        ddX1_dds = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='ss')
+        ddX1_dsdt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='sth')
+        ddX1_dsdz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='sze')
+
+        ddX2_dds = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='ss')
+        ddX2_dsdt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='sth')
+        ddX2_dsdz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='sze')
+
+        zmat = np.zeros_like(ddX1_dds)
+
+        J_s = np.array(((ddX1_dds, ddX1_dsdt, ddX1_dsdz),
+                      (ddX2_dds, ddX2_dsdt, ddX2_dsdz),
+                      (zmat, zmat, zmat)))
+
+        # partial derivative dJ/dth
+        ddX1_ddt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='thth')
+        ddX1_dtdz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='thze')
+
+        ddX2_ddt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='thth')
+        ddX2_dtdz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='thze')
+
+        J_th = np.array(((ddX1_dsdt, ddX1_ddt, ddX1_dtdz),
+                      (ddX2_dsdt, ddX2_ddt, ddX2_dtdz),
+                      (zmat, zmat, zmat)))
+
+        # partial derivative dJ/dze
+        ddX1_ddz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='zeze')
+
+        ddX2_ddz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='zeze')
+
+        J_ze = np.array(((ddX1_dsdz, ddX1_dtdz, ddX1_ddz),
+                      (ddX2_dsdz, ddX2_dtdz, ddX2_ddz),
+                      (zmat, zmat, zmat)))
+
+        return self.swap_J_axes(J_s), self.swap_J_axes(J_th), self.swap_J_axes(J_ze)
+
     def dX_inv(self, s, th, ze):
         '''Inverse Jacobian matrix of the Xmap.'''
         return np.linalg.inv(self.dX(s, th, ze))
@@ -717,6 +832,39 @@ class GVEC_domain:
         '''Jacobian determinant of the Xmap.'''
         return np.linalg.det(self.dX(s, th, ze))
 
+    def dJ_X(self, s, th, ze):
+        '''Returns a 3-tuple for the three partial derivatives of det_dX w.r.t s, th and ze.'''
+
+        dX = self.dX(s, th, ze)
+
+        # meshgrid evaluation
+        if dX.ndim == 5:
+            dX1_ds = dX[:, :, :, 0, 0]
+            dX2_ds = dX[:, :, :, 1, 0]
+            dX1_dt = dX[:, :, :, 0, 1]
+            dX2_dt = dX[:, :, :, 1, 1]
+        # single point evaluation
+        else:
+            dX1_ds = dX[0, 0]
+            dX2_ds = dX[1, 0]
+            dX1_dt = dX[0, 1]
+            dX2_dt = dX[1, 1]
+
+        ddX = self.ddX(s, th, ze)
+
+        tmp = []
+        for n in range(3):
+            # meshgrid evaluation
+            if dX.ndim == 5:
+                tmp += [ddX[0][:, :, :, 0, n] * dX2_dt + ddX[1][:, :, :, 1, n] * dX1_ds
+                    - ddX[0][:, :, :, 1, n] * dX1_dt + ddX[1][:, :, :, 0, n] * dX2_ds]
+            # single point evaluation
+            else:
+                tmp += [ddX[0][0, n] * dX2_dt + ddX[1][1, n] * dX1_ds
+                    - ddX[0][1, n] * dX1_dt + ddX[1][0, n] * dX2_ds]
+
+        return tmp[0], tmp[1], tmp[2]
+
     # hmap
     def hmap(self, q1, q2, ze):
         '''The map (x, y, z) = h(q1, q2, ze) which is a simple torus.
@@ -755,6 +903,29 @@ class GVEC_domain:
         '''Jacobian determinant of the hmap.'''
         return np.linalg.det(self.dh(q1, q2, ze))
 
+    def gh(self, q1, q2, ze):
+        '''Metric tensor of hmap.'''
+
+        zmat = np.zeros_like(q1)
+        omat = np.ones_like(q1)
+
+        J = np.array(((omat, zmat, zmat),
+                      (zmat, omat, zmat),
+                      (zmat, zmat, q1**2)))
+
+        return self.swap_J_axes(J)
+
+    def dgh_dq1(self, q1, q2, ze):
+        '''Partial derivative of gh w.r.t q1.'''
+
+        zmat = np.zeros_like(q1)
+
+        J = np.array(((zmat, zmat, zmat),
+                      (zmat, zmat, zmat),
+                      (zmat, zmat, 2*q1)))
+
+        return self.swap_J_axes(J)
+
     # Lmap
     def Lmap(self, s, u, v):
         '''The map (s, th, ze) = L(s, u, v) defined by th = 2*pi*u and ze = 2*pi*v/nfp with Jacobian determinant 4*pi^2/nfp.
@@ -798,6 +969,11 @@ class GVEC_domain:
 
         return self.swap_J_axes(J)
 
+    def det_dL(self, s, u, v):
+        '''Jacobian determinant of the Lmap.'''
+        s, u, v = self.prepare_args(s, u, v, sparse=False)
+        return np.ones_like(s) * 4 * np.pi**2 / self.nfp
+
     # pest map
     def Pmap(self, s2, th2, ze2):
         '''The map (s, th, ze) = P(s2, th2, ze2) giving the gvec coordinates (s, th, ze) in [0, 1] x [0, 2*pi]^2
@@ -895,6 +1071,10 @@ class GVEC_domain:
         '''Inverse Jacobian matrix of the composition P ° L.'''
         return np.linalg.inv(self.dPL(s, u, v))
 
+    def det_dPL(self, s, u, v):
+        '''Jacobian determinant of the composition P ° L.'''
+        return np.linalg.det(self.dPL(s, u, v))
+
     @staticmethod
     def prepare_args(s, u, v, sparse=True):
         '''Checks consistency, prepares arguments and performs meshgrid if necessary.'''

tests/test_gvec_domain.py

@@ -106,7 +106,7 @@ def test_values(use_pyccel,gvec_file):
     th = np.random.rand(10)*2*np.pi
     ze = np.array([np.pi])
     s, th, ze = GVEC_domain.prepare_args(s, th, ze)
-    th2 = th + gvec.lambda_(s, th, ze)
+    th2 = th + gvec._lambda(s, th, ze)
 
     # check Pmap
     assert np.allclose(s, gvec.domain.Pmap(s, th2, ze)[0])

tests/test_gvec_equil.py

@@ -41,6 +41,10 @@ def test_mhd_fields(mapping, use_pyccel,gvec_file):
     assert np.all([isinstance(gvec.a1(s, a1, a2)[j], float) for j in range(3)])
     assert np.all([isinstance(gvec.a2(s, a1, a2)[j], float) for j in range(3)])
     assert np.all([isinstance(gvec.a_cart(s, a1, a2)[i][j], float) for j in range(3) for i in range(2)])
+    assert np.all([isinstance(gvec.jv(s, a1, a2)[j], float) for j in range(3)])
+    assert np.all([isinstance(gvec.j1(s, a1, a2)[j], float) for j in range(3)])
+    assert np.all([isinstance(gvec.j2(s, a1, a2)[j], float) for j in range(3)])
+    assert np.all([isinstance(gvec.j_cart(s, a1, a2)[i][j], float) for j in range(3) for i in range(2)])
 
     # test single-point array evaluation
     s = np.array([s])
@@ -57,6 +61,10 @@ def test_mhd_fields(mapping, use_pyccel,gvec_file):
     assert np.all([gvec.a1(s, a1, a2)[j].shape == (1, 1, 1) for j in range(3)])
     assert np.all([gvec.a2(s, a1, a2)[j].shape == (1, 1, 1) for j in range(3)])
     assert np.all([gvec.a_cart(s, a1, a2)[i][j].shape == (1, 1, 1) for j in range(3) for i in range(2)])
+    assert np.all([gvec.jv(s, a1, a2)[j].shape == (1, 1, 1) for j in range(3)])
+    assert np.all([gvec.j1(s, a1, a2)[j].shape == (1, 1, 1) for j in range(3)])
+    assert np.all([gvec.j2(s, a1, a2)[j].shape == (1, 1, 1) for j in range(3)])
+    assert np.all([gvec.j_cart(s, a1, a2)[i][j].shape == (1, 1, 1) for j in range(3) for i in range(2)])
 
     # test 1d array evaluation
     s = np.linspace(.01, 1, 4) # mappings are singular at s=0
@@ -87,6 +95,14 @@ def test_mhd_fields(mapping, use_pyccel,gvec_file):
     assert np.all([tmp[j].shape == shp3 for j in range(3)])
     tmp = gvec.a_cart(s, a1, a2)
     assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)])
+    tmp = gvec.jv(s, a1, a2)
+    assert np.all([tmp[j].shape == shp3 for j in range(3)])
+    tmp = gvec.j1(s, a1, a2)
+    assert np.all([tmp[j].shape == shp3 for j in range(3)])
+    tmp = gvec.j2(s, a1, a2)
+    assert np.all([tmp[j].shape == shp3 for j in range(3)])
+    tmp = gvec.j_cart(s, a1, a2)
+    assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)])
 
     # test 3d array evaluation
     s, a1, a2 = np.meshgrid(s, a1, a2, indexing='ij')
@@ -109,6 +125,14 @@ def test_mhd_fields(mapping, use_pyccel,gvec_file):
     assert np.all([tmp[j].shape == shp3 for j in range(3)])
     tmp = gvec.a_cart(s, a1, a2)
     assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)])
+    tmp = gvec.jv(s, a1, a2)
+    assert np.all([tmp[j].shape == shp3 for j in range(3)])
+    tmp = gvec.j1(s, a1, a2)
+    assert np.all([tmp[j].shape == shp3 for j in range(3)])
+    tmp = gvec.j2(s, a1, a2)
+    assert np.all([tmp[j].shape == shp3 for j in range(3)])
+    tmp = gvec.j_cart(s, a1, a2)
+    assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)])
 
 
 if __name__ == '__main__':

tests/test_gvec_profiles.py

@@ -28,9 +28,11 @@ def test_profile(name,gvec_file):
     ze = np.pi/2
 
     assert isinstance(gvec.profiles.profile(s, name=name), float)
-    assert isinstance(gvec.profiles.profile_derivative(s, name=name), float)
+    assert isinstance(gvec.profiles.profile(s, name=name, der='s'), float)
+    assert isinstance(gvec.profiles.profile(s, name=name, der='ss'), float)
     assert gvec.profiles.profile(s, th, ze, name=name) == gvec.profiles.profile(s, name=name)
-    assert gvec.profiles.profile_derivative(s, th, ze, name=name) == gvec.profiles.profile_derivative(s, name=name)
+    assert gvec.profiles.profile(s, th, ze, name=name, der='s') == gvec.profiles.profile(s, name=name, der='s')
+    assert gvec.profiles.profile(s, th, ze, name=name, der='ss') == gvec.profiles.profile(s, name=name, der='ss')
 
     # test single-point array evaluation
     s = np.array([.5])
@@ -38,11 +40,14 @@ def test_profile(name,gvec_file):
     ze = np.array([np.pi/2])
 
     assert gvec.profiles.profile(s, name=name).shape == (1,)
-    assert gvec.profiles.profile_derivative(s, name=name).shape == (1,)
+    assert gvec.profiles.profile(s, name=name, der='s').shape == (1,)
+    assert gvec.profiles.profile(s, name=name, der='ss').shape == (1,)
     assert gvec.profiles.profile(s, th, ze, name=name).shape == (1, 1, 1)
-    assert gvec.profiles.profile_derivative(s, th, ze, name=name).shape == (1, 1, 1)
+    assert gvec.profiles.profile(s, th, ze, name=name, der='s').shape == (1, 1, 1)
+    assert gvec.profiles.profile(s, th, ze, name=name, der='ss').shape == (1, 1, 1)
     assert gvec.profiles.profile(s, th, ze, name=name)[0, 0, 0] == gvec.profiles.profile(s, name=name)[0]
-    assert gvec.profiles.profile_derivative(s, th, ze, name=name)[0, 0, 0] == gvec.profiles.profile_derivative(s, name=name)[0]
+    assert gvec.profiles.profile(s, th, ze, name=name, der='s')[0, 0, 0] == gvec.profiles.profile(s, name=name, der='s')[0]
+    assert gvec.profiles.profile(s, th, ze, name=name, der='ss')[0, 0, 0] == gvec.profiles.profile(s, name=name, der='ss')[0]
 
     # test 1d array evaluation
     s = np.linspace(0, 1, 3)
@@ -51,18 +56,22 @@ def test_profile(name,gvec_file):
     shp3 = (s.size, th.size, ze.size)
 
     assert gvec.profiles.profile(s, name=name).shape == (s.size,)
-    assert gvec.profiles.profile_derivative(s, name=name).shape == (s.size,)
+    assert gvec.profiles.profile(s, name=name, der='s').shape == (s.size,)
+    assert gvec.profiles.profile(s, name=name, der='ss').shape == (s.size,)
     assert gvec.profiles.profile(s, th, ze, name=name).shape == shp3
-    assert gvec.profiles.profile_derivative(s, th, ze, name=name).shape == shp3
+    assert gvec.profiles.profile(s, th, ze, name=name, der='s').shape == shp3
+    assert gvec.profiles.profile(s, th, ze, name=name, der='ss').shape == shp3
     for n in range(s.size):
         assert np.all(gvec.profiles.profile(s, th, ze, name=name)[n] == gvec.profiles.profile(s, name=name)[n])
-        assert np.all(gvec.profiles.profile_derivative(s, th, ze, name=name)[n] == gvec.profiles.profile_derivative(s, name=name)[n])
+        assert np.all(gvec.profiles.profile(s, th, ze, name=name, der='s')[n] == gvec.profiles.profile(s, name=name, der='s')[n])
+        assert np.all(gvec.profiles.profile(s, th, ze, name=name, der='ss')[n] == gvec.profiles.profile(s, name=name, der='ss')[n])
 
     # test 3d array evaluation
     s, th, ze = np.meshgrid(s, th, ze, indexing='ij')
 
     assert gvec.profiles.profile(s, th, ze, name=name).shape == s.shape
-    assert gvec.profiles.profile_derivative(s, th, ze, name=name).shape == s.shape
+    assert gvec.profiles.profile(s, th, ze, name=name, der='s').shape == s.shape
+    assert gvec.profiles.profile(s, th, ze, name=name, der='ss').shape == s.shape
 
 
 if __name__ == '__main__':