diff --git a/tofu/mag/test_magFieldLines.py b/tofu/mag/test_magFieldLines.py new file mode 100644 index 000000000..a6205e5da --- /dev/null +++ b/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) diff --git a/tofu/utils.py b/tofu/utils.py index 7679c6864..a859142f2 100644 --- a/tofu/utils.py +++ b/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 diff --git a/tofuplot.py b/tofuplot.py index f8e3a2fb6..1aad802c3 100755 --- a/tofuplot.py +++ b/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)