From 4eed92f7c834076d1acaf890363b462c62e63548 Mon Sep 17 00:00:00 2001
From: Florian Hindenlang <florian.hindenlang@ipp.mpg.de>
Date: Mon, 6 Feb 2023 19:35:51 +0000
Subject: [PATCH] fix in grad btheta, fixes current computation

---
 notebooks/plot_mhd_equil.ipynb        | 104 ++++++++++++++++++++++----
 src/gvec_to_python/GVEC_functions.py  |   4 +-
 src/gvec_to_python/geometry/domain.py |  32 ++++----
 tests/test_gvec_equil.py              |  79 +++++++++++++++++++
 4 files changed, 184 insertions(+), 35 deletions(-)

diff --git a/notebooks/plot_mhd_equil.ipynb b/notebooks/plot_mhd_equil.ipynb
index 6b84748..70ff3c2 100644
--- a/notebooks/plot_mhd_equil.ipynb
+++ b/notebooks/plot_mhd_equil.ipynb
@@ -109,7 +109,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "gvec.mapping = 'pest' # use the mapping setter\n",
+    "gvec.mapping = 'gvec' # use the mapping setter\n",
     "\n",
     "n_planes = 4 \n",
     "s = np.linspace(0, 1, 20) # radial coordinate in [0, 1]\n",
@@ -129,7 +129,7 @@
     "fig = plt.figure(figsize=(13, np.ceil(n_planes/2) * 6.5))\n",
     "\n",
     "# poloidal plane grid\n",
-    "for n in range(v.size):\n",
+    "for n in range(a2.size):\n",
     "    if(gvec.mapping==\"wo_hmap\"):\n",
     "        rp = x[:, :, n].squeeze()\n",
     "        zp = y[:, :, n].squeeze()\n",
@@ -148,8 +148,8 @@
     "                ax.plot([rp[i, j], rp[i + 1, j]], [zp[i, j], zp[i + 1, j]], 'b', linewidth=.6)\n",
     "            if j < rp.shape[1] - 1:\n",
     "                ax.plot([rp[i, j], rp[i, j + 1]], [zp[i, j], zp[i, j + 1]], 'b', linewidth=.6)\n",
-    "    ax.set_xlabel('r')\n",
-    "    ax.set_ylabel('z')\n",
+    "    ax.set_xlabel('R')\n",
+    "    ax.set_ylabel('Z')\n",
     "    ax.axis('equal')\n",
     "    ax.set_title('Poloidal plane at $\\\\zeta$={0:4.3f} $\\pi$'.format(a2[n]/vfac*2/gvec.nfp))"
    ]
@@ -200,8 +200,8 @@
     "\n",
     "    ax = fig.add_subplot(int(np.ceil(n_planes/2)), 2, n + 1)\n",
     "    map = ax.contourf(rp, zp, detp, 30)\n",
-    "    ax.set_xlabel('r')\n",
-    "    ax.set_ylabel('z')\n",
+    "    ax.set_xlabel('R')\n",
+    "    ax.set_ylabel('Z')\n",
     "    ax.axis('equal')\n",
     "    ax.set_title('Jacobian determinant at $\\\\zeta$={0:4.3f} $\\pi$'.format(v[n]/vfac*2/gvec.nfp))\n",
     "    fig.colorbar(map, ax=ax, location='right')"
@@ -305,8 +305,8 @@
     "\n",
     "    ax = fig.add_subplot(int(np.ceil(n_planes/2)), 2, n + 1)\n",
     "    map = ax.contourf(rp, zp, pp, 30)\n",
-    "    ax.set_xlabel('r')\n",
-    "    ax.set_ylabel('z')\n",
+    "    ax.set_xlabel('R')\n",
+    "    ax.set_ylabel('Z')\n",
     "    ax.axis('equal')\n",
     "    ax.set_title('Pressure at $\\\\zeta$={0:4.3f} $\\pi$'.format(v[n]/vfac * 2 / gvec.nfp))\n",
     "    fig.colorbar(map, ax=ax, location='right')"
@@ -368,8 +368,8 @@
     "\n",
     "    ax = fig.add_subplot(int(np.ceil(n_planes/2)), 2, n + 1)\n",
     "    map = ax.contourf(rp, zp, ab, 30)\n",
-    "    ax.set_xlabel('r')\n",
-    "    ax.set_ylabel('z')\n",
+    "    ax.set_xlabel('R')\n",
+    "    ax.set_ylabel('Z')\n",
     "    ax.axis('equal')\n",
     "    ax.set_title('Magnetic field strength at $\\\\zeta$={0:4.3f} $\\pi$'.format(v[n]/vfac * 2 / gvec.nfp))\n",
     "    fig.colorbar(map, ax=ax, location='right')"
@@ -403,8 +403,8 @@
     "\n",
     "    ax = fig.add_subplot(int(np.ceil(n_planes/2)), 2, n + 1)\n",
     "    ax.quiver(rp, zp, drdu, dzdu, scale=20)\n",
-    "    ax.set_xlabel('r')\n",
-    "    ax.set_ylabel('z')\n",
+    "    ax.set_xlabel('R')\n",
+    "    ax.set_ylabel('Z')\n",
     "    ax.axis('equal')\n",
     "    ax.set_title('Theta derivative of mapping at $\\\\zeta$={0:4.3f} $\\pi$'.format(v[n] * 2 / gvec.domain.nfp))\n",
     "    #fig.colorbar(map, ax=ax, location='right')"
@@ -417,6 +417,45 @@
     "# Current density"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def bcart(s,th,ze):\n",
+    "    b,x=gvec.b_cart(s,th,ze)\n",
+    "    return b\n",
+    "\n",
+    "def jcart_FD(s,th,ze):\n",
+    "    eps=1.0e-8\n",
+    "    b_cart = bcart(s,th,ze)\n",
+    "    b_cart_s_eps   = bcart(s+eps,th,ze)\n",
+    "    b_cart_th_eps  = bcart(s,th+eps,ze)\n",
+    "    b_cart_ze_eps  = bcart(s,th,ze+eps)\n",
+    "    gradb=np.zeros((3,3))\n",
+    "    for d in [0,1,2]:\n",
+    "        b_ds =(b_cart_s_eps[d] -b_cart[d])/eps\n",
+    "        b_dth=(b_cart_th_eps[d]-b_cart[d])/eps\n",
+    "        b_dze=(b_cart_ze_eps[d]-b_cart[d])/eps\n",
+    "        gradb[d] = gvec.df_inv(s,th,ze).T@np.array([b_ds,b_dth,b_dze])\n",
+    "    j_cart_FD=np.zeros(3)\n",
+    "    for d in [0,1,2]:\n",
+    "        j_cart_FD[d] = gradb[(d+2)%3][(d+1)%3]-gradb[(d+1)%3][(d+2)%3]\n",
+    "    return j_cart_FD\n",
+    "\n",
+    "gvec.mapping = 'gvec'\n",
+    "s,th,ze = [0.49,0.5,0.3]\n",
+    "\n",
+    "mu0=4.0e-7*np.pi\n",
+    "j_ex ,xyz = gvec.j_cart(s, th, ze)\n",
+    "j_FD = jcart_FD(s,th,ze)/mu0\n",
+    "print('j_FD ',j_FD)\n",
+    "print('j_ex',np.array(j_ex))\n",
+    "print('j_ex-j_FD',np.array(j_ex)-j_FD)\n",
+    "assert(np.allclose(j_ex,j_FD,rtol=1e-6,atol=1e-6/mu0))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -456,8 +495,8 @@
     "\n",
     "    ax = fig.add_subplot(int(np.ceil(n_planes/2)), 2, n + 1)\n",
     "    map = ax.contourf(rp, zp, ab, 30)\n",
-    "    ax.set_xlabel('r')\n",
-    "    ax.set_ylabel('z')\n",
+    "    ax.set_xlabel('R')\n",
+    "    ax.set_ylabel('Z')\n",
     "    ax.axis('equal')\n",
     "    ax.set_title('Current (absolute value) at $\\\\zeta$={0:4.3f} $\\pi$'.format(v[n]/vfac * 2 / gvec.nfp))\n",
     "    fig.colorbar(map, ax=ax, location='right')"
@@ -498,7 +537,7 @@
     "\n",
     "print(ss.shape, uu.shape, vv.shape)\n",
     "\n",
-    "comp = 1\n",
+    "comp = 0\n",
     "\n",
     "for n in range(v.size):\n",
     "\n",
@@ -512,11 +551,44 @@
     "    map = ax.contourf(sp, up, ji, 30)\n",
     "    ax.set_xlabel('s')\n",
     "    ax.set_ylabel('u')\n",
-    "    ax.axis('equal')\n",
     "    ax.set_title('Current component {1} at $\\\\phi$={0:4.1f} $\\pi$'.format(v[n] * 2 / gvec.domain.nfp, comp))\n",
     "    fig.colorbar(map, ax=ax, location='right')"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure(figsize=(15, np.ceil(n_planes/2) * 6.5))\n",
+    "\n",
+    "ss, uu, vv = gvec.domain.prepare_args(s, u, v, sparse=False)\n",
+    "\n",
+    "print(ss.shape, uu.shape, vv.shape)\n",
+    "\n",
+    "comp = 0\n",
+    "\n",
+    "for n in range(v.size):\n",
+    "\n",
+    "    sp = ss[:, :, n].squeeze()\n",
+    "    up = uu[:, :, n].squeeze()\n",
+    "    vp = vv[:, :, n].squeeze()\n",
+    "\n",
+    "    if(comp==0):\n",
+    "        fbe = (p_s - (j[1] * b[2] - j[2] * b[1]))[:, :, n].squeeze()\n",
+    "    elif(comp==1):\n",
+    "        fbe = (j[2] * b[0] - j[0] * b[2])[:, :, n].squeeze()\n",
+    "    elif(comp==2):\n",
+    "        fbe = (j[0] * b[1] - j[1] * b[0])[:, :, n].squeeze()\n",
+    "    ax = fig.add_subplot(int(np.ceil(n_planes/2)), 2, n + 1)\n",
+    "    map = ax.contourf(sp, up, fbe, 30)\n",
+    "    ax.set_xlabel('s')\n",
+    "    ax.set_ylabel('u')\n",
+    "    ax.set_title('force balance error in comp {1} at $\\\\zeta$={0:4.1f} $\\pi$'.format(v[n]/vfac * 2 / gvec.nfp, comp))\n",
+    "    fig.colorbar(map, ax=ax, location='right')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/src/gvec_to_python/GVEC_functions.py b/src/gvec_to_python/GVEC_functions.py
index 801743c..dab5af4 100644
--- a/src/gvec_to_python/GVEC_functions.py
+++ b/src/gvec_to_python/GVEC_functions.py
@@ -581,11 +581,11 @@ class GVEC:
         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) 
+        term1 = - 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) 
+        term1 = - 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
 
diff --git a/src/gvec_to_python/geometry/domain.py b/src/gvec_to_python/geometry/domain.py
index 214ae1e..175505e 100644
--- a/src/gvec_to_python/geometry/domain.py
+++ b/src/gvec_to_python/geometry/domain.py
@@ -803,7 +803,7 @@ class GVEC_domain:
 
         J = np.array(((dX1ds, dX1dt, dX1dz),
                       (dX2ds, dX2dt, dX2dz),
-                      (zmat, zmat, omat)))
+                      ( zmat,  zmat,  omat)))
 
         return self.swap_J_axes(J)
 
@@ -823,8 +823,8 @@ class GVEC_domain:
         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)))
+                        (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')
@@ -834,8 +834,8 @@ class GVEC_domain:
         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)))
+                         (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')
@@ -843,8 +843,8 @@ class GVEC_domain:
         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)))
+                         (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)
 
@@ -875,17 +875,19 @@ class GVEC_domain:
             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]
+                tmp += [ddX[n][:, :, :, 0, 0] * dX2_dt + ddX[n][:, :, :, 1, 1] * dX1_ds
+                      - ddX[n][:, :, :, 1, 0] * dX1_dt - ddX[n][:, :, :, 0, 1] * 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]
+            #            X1ds/di * dX2_dt + dX1dt/di * dX1_ds 
+            #           -X2ds/di * dX1_dt - dX2dt/di * dX2_ds 
+            else:            
+                tmp += [ddX[n][0, 0] * dX2_dt + ddX[n][1, 1] * dX1_ds
+                      - ddX[n][1, 0] * dX1_dt - ddX[n][0, 1] * dX2_ds]
 
         return tmp[0], tmp[1], tmp[2]
 
@@ -1119,10 +1121,6 @@ class GVEC_domain:
                 s.ndim, v.ndim)
             if s.ndim == 1:
                 s, u, v = np.meshgrid(s, u, v, indexing='ij', sparse=sparse)
-            elif s.ndim == 3:
-                is_meshgrid=(np.all(np.minimum.reduce(s,axis=(1,2))==np.maximum.reduce(s,axis=(1,2))) and
-                             np.all(np.minimum.reduce(u,axis=(0,2))==np.maximum.reduce(u,axis=(0,2))) and
-                             np.all(np.minimum.reduce(v,axis=(0,1))==np.maximum.reduce(v,axis=(0,1))))
 
         return s, u, v
 
diff --git a/tests/test_gvec_equil.py b/tests/test_gvec_equil.py
index a0e1263..b112da0 100644
--- a/tests/test_gvec_equil.py
+++ b/tests/test_gvec_equil.py
@@ -212,6 +212,85 @@ def test_values(use_pyccel,gvec_file,unit_tor_domain):
         assert np.allclose(a_unit_pest[1], a_gvec[1])
         assert np.allclose(a_unit_pest[2], a_gvec[2])
 
+@pytest.mark.parametrize('use_pyccel,gvec_file', [(True ,"testcases/ellipstell/newBC_E4D6_M6N6/GVEC_ELLIPSTELL_E4D6_M6N6_State_0001_00200000")])   
+@pytest.mark.parametrize('mapping', ['gvec', 'pest', 'unit', 'unit_pest'])
+def test_derivs_with_FD(mapping,use_pyccel,gvec_file):
+    '''Compare exact derivatives of functions to their finite difference counterpart.'''
+
+    import os
+    import numpy as np
+    from gvec_to_python.reader.gvec_reader import create_GVEC_json
+    from gvec_to_python import GVEC
+    import gvec_to_python as _
+
+    def do_FD(func,s,th,ze,der=None):
+        '''  compute s,th,ze derivative of input function func'''
+        eps=1.0e-8
+        f0 = func(s,th,ze)
+        df_ds  = (func(s+eps,th,ze) - f0 )/ eps
+        df_dth = (func(s,th+eps,ze) - f0 )/ eps
+        df_dze = (func(s,th,ze+eps) - f0 )/ eps
+        if(der=="s"):
+            return df_ds
+        elif(der=="th"):
+            return df_dth
+        elif(der=="ze"):
+            return df_dze
+        else:
+            return df_ds,df_dth,df_dze
+
+    def bth(s,th,ze):
+        '''  computes B^theta=b^theta/det_df'''
+        return gvec._b_small_th(s, th, ze) / gvec.domain.det_df_gvec(s, th, ze)
+        
+    def bze(s,th,ze):
+        '''  computes B^zeta=b^zeta/det_df'''
+        return gvec._b_small_ze(s, th, ze) / gvec.domain.det_df_gvec(s, th, ze)
+
+    # give absolute paths to the files
+    pkg_path = _.__path__[0]
+    dat_file_in   = os.path.join( pkg_path, gvec_file+'.dat')
+    json_file_out = os.path.join( pkg_path, gvec_file+'.json')
+    create_GVEC_json(dat_file_in, json_file_out)
+
+    # main object 
+    gvec = GVEC(json_file_out, mapping=mapping, use_pyccel=use_pyccel)
+    print('gvec.mapping: ', gvec.mapping)
+
+    # test single-point float evaluation
+    s = .45
+    if 'unit' not in mapping:
+        th = 0.77*2*np.pi
+        ze = 0.33*2*np.pi
+    else:
+        th = .77
+        ze = .33
+
+
+    #check derivative of Jacobian of gvec mapping
+    dJf_FD = do_FD(gvec.domain.det_df_gvec,s,th,ze)
+    dJf_ex = [gvec.domain.dJ_gvec(s,th,ze,der=deriv) for deriv in  ["s","th","ze"]]
+    assert(np.allclose(dJf_FD,dJf_ex,rtol=1e-6,atol=1e-4))
+    
+    
+    # check derivative of metric tensor g of the gvec mapping
+    for deriv in ["s","th","ze"]:
+        dg_FD = do_FD(gvec.domain.g_gvec,s,th,ze,der=deriv)
+        dg_ex = gvec.domain.dg_gvec(s,th,ze,der=deriv)
+        assert(np.allclose(dg_FD,dg_ex,rtol=1e-6,atol=1e-4))
+    
+        
+    # check derivative of  B^theta 
+    dbth_FD = do_FD(bth,s,th,ze)
+    dbth_ex  = gvec._grad_b_theta(s, th, ze)
+    assert(np.allclose(dbth_FD,dbth_ex,rtol=1e-6,atol=1e-4))
+        
+    # check derivative of  B^zeta 
+    dbze_FD = do_FD(bze,s,th,ze)
+    dbze_ex  = gvec._grad_b_zeta(s, th, ze)
+    #print("diff dbze_FD-dbze_ex",dbze_FD-np.array(dbze_ex))
+    assert(np.allclose(dbze_FD,dbze_ex,rtol=1e-6,atol=1e-4))
+
 
 
 if __name__ == '__main__':
-- 
GitLab