Br/add more field lines

tofu/mag/test_magFieldLines.py

@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
+# Also if needed: retab
+'''
+    TEST magFieldLines
+'''
+from __future__ import (unicode_literals, absolute_import,  \
+                        print_function, division)
+import doctest
+import os
+import sys
+
+#print('path 1 =', sys.path)
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+#print('path 2 =', sys.path)
+
+#import mag
+from mag import magFieldLines
+
+doctest.testmod(magFieldLines, verbose=True)

tofu/utils.py

@@ -623,7 +623,7 @@ def load_from_imas(shot=None, run=None, user=None, tokamak=None, version=None,
                    occ=None, indch=None, indDescription=None, equilibrium=None,
                    dsig=None, data=None, X=None, t0=None, dextra=None,
                    plot=True, plot_sig=None, plot_X=None, sharex=False,
-                   bck=True, indch_auto=True, t=None, init=None):
+                   bck=True, indch_auto=True, t=None, init=None, dR_sep=None):
     # -------------------
     # import imas2tofu
     try:
@@ -677,11 +677,33 @@ def load_from_imas(shot=None, run=None, user=None, tokamak=None, version=None,
         assert shot.size == 1
         import tofu.mag as tfm
         plot = True
+
+        multi = imas2tofu.MultiIDSLoader(shot=shot[0], run=run, user=user,
+                                        tokamak=tokamak, version=version,
+                                        ids='equilibrium')
+        equi = multi.to_Plasma2D()
+
         if t is None:
-            t = np.r_[38]
+            # Get time in the middle of equilibrium time interval
+            t = equi.ddata['equilibrium.t']['data'][int(0.5*equi.ddata['equilibrium.t']['data'].size)]
+            #t = np.r_[38]
         t = np.atleast_1d(t).ravel()
-        if init is None:
-            init = [[2.9],[0.],[0.]]
+
+        #init = [[2.9, 2.9, 2.9, 2.9], [0., np.pi/2, np.pi, 3.*np.pi/2], [0., 0., 0., 0.]]
+        #init = [[2.9], [0.], [0.]]
+        equi_ind_t     = np.abs(t - equi.ddata['equilibrium.t']['data']).argmin()
+        equi_ind_r_ext = np.argmax(equi.ddata['equilibrium.sep']['data'][equi_ind_t][0])
+        equi_r_ext     = equi.ddata['equilibrium.sep']['data'][equi_ind_t][0][equi_ind_r_ext]
+        equi_z_r_ext   = equi.ddata['equilibrium.sep']['data'][equi_ind_t][1][equi_ind_r_ext]
+
+        nbr_init = 10
+        if dR_sep is not None:
+            r_init = [equi_r_ext + dR_sep]*nbr_init
+        else:
+            r_init = [equi_r_ext]*nbr_init
+        phi_init = [ii*2.*np.pi/nbr_init for ii in range(nbr_init)]
+        z_init   = [equi_z_r_ext]*nbr_init
+        init_plt = [r_init, phi_init, z_init]
 
         if False:
             multi = imas2tofu.MultiIDSLoader(shot=shot[0], run=run, user=user,
@@ -690,20 +712,57 @@ def load_from_imas(shot=None, run=None, user=None, tokamak=None, version=None,
             config = multi.to_Config(plot=False)
         else:
             import tofu.geom as tfg
-            config = tfg.utils.create_config('B2')
+            config = tfg.utils.create_config('B3')
         if config.nStruct > 1:
             config.set_colors_random()
-        trace = tfm.MagFieldLines(int(shot[0])).trace_mline(init, t,
+        trace = tfm.MagFieldLines(int(shot[0])).trace_mline(init_plt, t,
                                                        direction='FWD',
-                                                       length_line=None,
+                                                       length_line=25,
+                                                       stp=None)
+        trace_rev = tfm.MagFieldLines(int(shot[0])).trace_mline(init_plt, t,
+                                                       direction='REV',
+                                                       length_line=25,
                                                        stp=None)
+
         refpt = np.r_[2.4,0.]
         dax = config.plot_phithetaproj_dist(refpt)
+
+        if init is not None:
+            trace_init = tfm.MagFieldLines(int(shot[0])).trace_mline(init, t,
+                                                           direction='FWD',
+                                                           length_line=25,
+                                                           stp=None)
+            trace_init_rev = tfm.MagFieldLines(int(shot[0])).trace_mline(init, t,
+                                                           direction='REV',
+                                                           length_line=25,
+                                                           stp=None)
+            trace_init[0] = trace_init[0] + trace_init_rev[0]
+
+            for kk in range(0,len(trace_init[0])):
+                phi_init = np.arctan2(np.sin(trace_init[0][kk]['p']),
+                                      np.cos(trace_init[0][kk]['p']))
+                theta_init = np.arctan2(trace_init[0][kk]['z']-refpt[1],
+                                        trace_init[0][kk]['r']-refpt[0])
+                indnan = ((np.abs(np.diff(phi_init)) > np.pi)
+                          | (np.abs(np.diff(theta_init)) > np.pi)).nonzero()[0] + 1
+                dax['dist'][0].plot(np.insert(phi_init, indnan, np.nan),
+                                    np.insert(theta_init, indnan, np.nan),
+                                    linewidth=3, color='red')
+                dax['cross'][0].plot(trace_init[0][kk]['r'], trace_init[0][kk]['z'],
+                                     linewidth=3, color='red')
+                x = trace_init[0][kk]['r']*np.cos(trace_init[0][kk]['p'])
+                y = trace_init[0][kk]['r']*np.sin(trace_init[0][kk]['p'])
+                dax['hor'][0].plot(x, y, linewidth=3, color='red')
+
         for ii in range(0,len(trace)):
+            # Concatenate trace lists
+            trace[ii] = trace[ii] + trace_rev[ii]
             for jj in range(0,len(trace[ii])):
                 lab = r't = %s s'%str(t[ii])
-                phi = np.arctan2(np.sin(trace[ii][jj]['p']), np.cos(trace[ii][jj]['p']))
-                theta = np.arctan2(trace[ii][jj]['z']-refpt[1], trace[ii][jj]['r']-refpt[0])
+                phi = np.arctan2(np.sin(trace[ii][jj]['p']),
+                                 np.cos(trace[ii][jj]['p']))
+                theta = np.arctan2(trace[ii][jj]['z']-refpt[1],
+                                   trace[ii][jj]['r']-refpt[0])
                 # insert nans for clean periodicity
                 indnan = ((np.abs(np.diff(phi)) > np.pi)
                           | (np.abs(np.diff(theta)) > np.pi)).nonzero()[0] + 1
@@ -715,6 +774,7 @@ def load_from_imas(shot=None, run=None, user=None, tokamak=None, version=None,
                 x = trace[ii][jj]['r']*np.cos(trace[ii][jj]['p'])
                 y = trace[ii][jj]['r']*np.sin(trace[ii][jj]['p'])
                 dax['hor'][0].plot(x, y, label=lab)
+        dax['t'][0].figure.suptitle('Shot {0}, t = {1:6.3f} s'.format(shot[0], t[0]))
         return dax
 
 

tofuplot.py

@@ -49,7 +49,7 @@
 _VERSION = '3'
 _LIDS_DIAG = MultiIDSLoader._lidsdiag
 _LIDS_PLASMA = tf.imas2tofu.MultiIDSLoader._lidsplasma
-_LIDS = _LIDS_DIAG + _LIDS_PLASMA
+_LIDS = _LIDS_DIAG + _LIDS_PLASMA + ['magfieldlines']
 _T0 = 'IGNITRON'
 _SHAREX = False
 _BCK = True
@@ -75,7 +75,7 @@ def call_tfloadimas(shot=None, run=_RUN, user=_USER,
                     tokamak=_TOKAMAK, version=_VERSION,
                     ids=None, quantity=None, X=None, t0=_T0,
                     sharex=_SHAREX, indch=None, indch_auto=None,
-                    background=_BCK):
+                    background=_BCK, t=None, dR_sep=None, init=None):
 
     lidspla = [ids_ for ids_ in ids if ids_ in _LIDS_PLASMA]
     if t0.lower() == 'none':
@@ -85,7 +85,8 @@ def call_tfloadimas(shot=None, run=_RUN, user=_USER,
                       tokamak=tokamak, version=version,
                       ids=ids, indch=indch, indch_auto=indch_auto,
                       plot_sig=quantity, plot_X=X,
-                      t0=t0, plot=True, sharex=sharex, bck=background)
+                      t0=t0, plot=True, sharex=sharex, bck=background,
+                      t=t, dR_sep=dR_sep, init=init)
 
     plt.show(block=True)
 
@@ -125,7 +126,7 @@ def _str2bool(v):
     msg = 'username of the DB where the datafile is located'
     parser.add_argument('-u','--user',help=msg, required=False, default=_USER)
     msg = 'tokamak name of the DB where the datafile is located'
-    parser.add_argument('-t','--tokamak',help=msg, required=False,
+    parser.add_argument('-tok','--tokamak',help=msg, required=False,
                         default=_TOKAMAK)
     parser.add_argument('-r','--run',help='run number',
                         required=False, type=int, default=_RUN)
@@ -143,6 +144,14 @@ def _str2bool(v):
                         nargs='+', default=None)
     parser.add_argument('-t0', '--t0', type=str, required=False,
                         help='Reference time event setting t = 0', default=_T0)
+    parser.add_argument('-t', '--t', type=float, required=False,
+                        help='Input time when needed')
+    parser.add_argument('-dR_sep', '--dR_sep', type=float, required=False,
+                        help='Distance to separatrix from r_ext to plot'
+                            +'magneticfield lines')
+    parser.add_argument('-init', '--init', type=float, required=False, nargs=3,
+                        help='Initial point from where trace magnetic field'
+                            +'line')
     parser.add_argument('-ich', '--indch', type=int, required=False,
                         help='indices of channels to be loaded',
                         nargs='+', default=None)