Br/add more field lines

tofu/geom/_core.py

@@ -3074,37 +3074,21 @@ def _get_phithetaproj_dist(
 
         return dist, indStruct
 
-    def plot_phithetaproj_dist(
-        self,
-        refpt=None,
-        ntheta=None,
-        nphi=None,
-        theta=None,
-        phi=None,
-        cmap=None,
-        ax=None,
-        fs=None,
-        tit=None,
-        wintit=None,
-        draw=None,
-    ):
-        dist, indStruct = self._get_phithetaproj_dist(
-            refpt=refpt, ntheta=ntheta, nphi=nphi, theta=theta, phi=phi
-        )
-        return _plot.Config_phithetaproj_dist(
-            self,
-            refpt,
-            dist,
-            indStruct,
-            cmap=cmap,
-            ax=ax,
-            fs=fs,
-            tit=tit,
-            wintit=wintit,
-            draw=draw,
-        )
+    def plot_phithetaproj_dist(self, refpt=None, ntheta=None, nphi=None,
+                               theta=None, phi=None, cmap=None, invertx=None,
+                               ax=None, fs=None, tit=None, wintit=None,
+                               draw=None):
+        dist, indStruct = self._get_phithetaproj_dist(refpt=refpt,
+                                                      ntheta=ntheta, nphi=nphi,
+                                                      theta=theta, phi=phi)
+        return _plot.Config_phithetaproj_dist(self, refpt, dist, indStruct,
+                                              cmap=cmap, ax=ax, fs=fs,
+                                              tit=tit, wintit=wintit,
+                                              invertx=invertx, draw=draw)
+
 
     def isInside(self, pts, In="(X,Y,Z)", log="any"):
+
         """ Return a 2D array of bool
 
         Equivalent to applying isInside to each Struct

tofu/geom/_plot.py

@@ -692,13 +692,15 @@ def Plot_Impact_3DPoly(T, Leg="", ax=None, Ang=_def.TorPAng,
 
 def Config_phithetaproj_dist(config, refpt, dist, indStruct,
                              distonly=False,
-                             cmap=None, vmin=None, vmax=None,
+                             cmap=None, vmin=None, vmax=None, invertx=None,
                              ax=None, fs=None, cbck=(0.8,0.8,0.8,0.8),
                              tit=None, wintit=None, legend=None, draw=None):
     if cmap is None:
         cmap = 'touch'
     lS = config.lStruct
     indsu = np.unique(indStruct)
+    if invertx is None:
+        invertx = True
 
     # set extent
     ratio = refpt[0] / np.nanmin(dist)
@@ -745,6 +747,8 @@ def Config_phithetaproj_dist(config, refpt, dist, indStruct,
             # labels.append( '%s_%s'%(lS[ii].Id.Cls, lS[ii].Id.Name) )
         # dax['cross'][0].legend(handles, labels, frameon=False,
                                # bbox_to_anchor=(1.01,1.), loc=2, borderaxespad=0.)
+    if invertx is True:
+        dax['dist'][0].invert_xaxis()
 
     if draw:
         fig.canvas.draw()

tofu/mag/equimap.py

@@ -20,19 +20,19 @@
     import imas_west
 except ImportError as err:
     pass
-#    print(err)
+    print(err)
 try:
     import pywed as pw
 except ImportError as err:
     pass
-#    print(err)
+    print(err)
 
 # Project modules
 try:
     import mag_ripple as mr
 except ImportError as err:
     pass
-#    print(err)
+    print(err)
 
 __all__ = ['get']
 

tofu/mag/magFieldLines.py

@@ -160,7 +160,7 @@ def __init__(self, shot, run=0, occ=0, user='imas_public', machine='west'):
                            bounds_error=False)
 
     def trace_mline(self, init_state, time, direction='FWD', \
-                    length_line=None, stp=None):
+                    length_line=None, stp=None, ripple=True):
         '''
         Traces the field line given a starting point.
         Integration step defined by stp and maximum length of the field line
@@ -174,6 +174,7 @@ def trace_mline(self, init_state, time, direction='FWD', \
             - time : array of times requested (list or numpy array)
             - direction : direction of integration 'FWD' forward or 'REV' reverse
               (string)
+            - ripple : take into account magnetic ripple or not (boolean)
 
         output:
             returns a dictionary containing:
@@ -209,7 +210,8 @@ def trace_mline(self, init_state, time, direction='FWD', \
 
         # !!!!!!!!!!!!!!!!!!!!!
         # HARD CODED CORRECTION
-        bt_intp_t *= -1
+        if (np.nanmean(bt_intp_t) > 0):
+            bt_intp_t *= -1
         # !!!!!!!!!!!!!!!!!!!!!
 
         # Interpolate current
@@ -236,7 +238,7 @@ def trace_mline(self, init_state, time, direction='FWD', \
 
             for jj in range(init_state.shape[1]):
                 out_prime = self.integrate_solve_ivp(init_state[:, jj], \
-                                                     direction, s, ds)
+                                                     direction, s, ds, ripple)
                 out_prime['init_point'] = init_state[:, jj]
                 out_prime['time']       = self.ar_time[ii]
                 # Return list of dict
@@ -246,7 +248,7 @@ def trace_mline(self, init_state, time, direction='FWD', \
 
         return outMagLine
 
-    def integrate_solve_ivp(self, init_state, direction, s, ds):
+    def integrate_solve_ivp(self, init_state, direction, s, ds, ripple=True):
         '''
         Integration step defined by stp and maximum length of the field line
         defined by s.
@@ -258,6 +260,7 @@ def integrate_solve_ivp(self, init_state, direction, s, ds):
               (string)
             - s : field line maximum length
             - ds : vector of steps
+            - ripple : boolean, take into account magnetic ripple or not
 
         output:
             returns a dictionary containing:
@@ -268,14 +271,26 @@ def integrate_solve_ivp(self, init_state, direction, s, ds):
             - y  : y coordinate (np.array)
             - cp : collision point with the wall (list)
         '''
-        if (direction=='FWD'):
-            sol = spode.solve_ivp(self.mfld3dcylfwd, [0, s], init_state,
-                                  method='RK23', t_eval=ds,
-                                  events=self.hit_wall_circ)
-        elif (direction=='REV'):
-            sol = spode.solve_ivp(self.mfld3dcylrev, [0, s], init_state,
-                                  method='RK23', t_eval=ds,
-                                  events=self.hit_wall_circ)
+        if (ripple):
+            if direction == 'FWD':
+                sol = spode.solve_ivp(self.mfld3dcylfwd, [0, s], init_state,
+                                      method='RK23', t_eval=ds,
+                                      events=self.hit_wall_circ)
+            elif direction == 'REV':
+                sol = spode.solve_ivp(self.mfld3dcylrev, [0, s], init_state,
+                                      method='RK23', t_eval=ds,
+                                      events=self.hit_wall_circ)
+        else:
+            if direction == 'FWD':
+                sol = spode.solve_ivp(self.mfld3dcylfwd_no_ripple,
+                                      [0, s], init_state,
+                                      method='RK23', t_eval=ds,
+                                      events=self.hit_wall_circ)
+            elif direction == 'REV':
+                sol = spode.solve_ivp(self.mfld3dcylrev_no_ripple,
+                                      [0, s], init_state,
+                                      method='RK23', t_eval=ds,
+                                      events=self.hit_wall_circ)
         sgf = sol.t
         rgf = sol.y[0]
         pgf = sol.y[1]
@@ -321,9 +336,43 @@ def mfld3dcylrev(self, s, state):
         '''
         R, P, Z = state
         Br, Bt, Bz = self.b_field_interp(R, P, Z)
-        #Br=self.fBp_r(R,Z)[0]+self.fBt_r(P,self.ftheta(R,Z),R)
-        #Bt=self.fBt_t(P,self.ftheta(R,Z),R)
-        #Bz=self.fBp_z(R,Z)[0]
+        # Br=self.fBp_r(R,Z)[0]+self.fBt_r(P,self.ftheta(R,Z),R)
+        # Bt=self.fBt_t(P,self.ftheta(R,Z),R)
+        # Bz=self.fBp_z(R,Z)[0]
+        B = np.sqrt(Br*Br + Bz*Bz + Bt*Bt)
+        d_R = -Br/B
+        d_P = -Bt/B*1/R
+        d_Z = -Bz/B
+        return [d_R, d_P, d_Z]
+
+    def mfld3dcylfwd_no_ripple(self, s, state):
+        '''
+        Returns the right end side of the field line system of equations in
+        cyclindrical (R, Phi, Z) coord.
+        This is the case for forward integration.
+        '''
+        R, P, Z = state
+        Br, Bt, Bz = self.b_field_interp_no_ripple(R, P, Z)
+        # Br=self.fBp_r(R,Z)[0]+self.fBt_r(P,self.ftheta(R,Z),R)
+        # Bt=self.fBt_t(P,self.ftheta(R,Z),R)
+        # Bz=self.fBp_z(R,Z)[0]
+        B = np.sqrt(Br*Br + Bz*Bz + Bt*Bt)
+        d_R = Br/B
+        d_P = Bt/B*1/R
+        d_Z = Bz/B
+        return [d_R, d_P, d_Z]
+
+    def mfld3dcylrev_no_ripple(self, s, state):
+        '''
+        Returns the right end side of the field line system of equations in
+        cyclindrical (R, Phi, Z) coord.
+        This is the case for backward integration.
+        '''
+        R, P, Z = state
+        Br, Bt, Bz = self.b_field_interp_no_ripple(R, P, Z)
+        # Br=self.fBp_r(R,Z)[0]+self.fBt_r(P,self.ftheta(R,Z),R)
+        # Bt=self.fBt_t(P,self.ftheta(R,Z),R)
+        # Bz=self.fBp_z(R,Z)[0]
         B = np.sqrt(Br*Br + Bz*Bz + Bt*Bt)
         d_R = -Br/B
         d_P = -Bt/B*1/R
@@ -340,14 +389,17 @@ def hit_wall_circ(self, s, state):
             mat file
         '''
         R, P, Z = state
-        #if np.abs(Z)<=0.5 and R>3.01 and np.deg2rad(89.5)<=P<=np.deg2rad(90.5):
-        #    return 0
-        #if np.abs(Z)<=0.5 and R>3.01 and np.deg2rad(179.5)<=P<=np.deg2rad(180.5):
-        #    print('Got 2')
-        #    return 0
-        #elif np.abs(Z)<=0.5 and R>3.01 and np.deg2rad(269.5)<=P<=np.deg2rad(270.5):
-        #    print('Got 3')
-        #    return 0
+        # if (np.abs(Z) <= 0.5 and R > 3.01
+        #     and np.deg2rad(89.5) <= P <= np.deg2rad(90.5)):
+        #     return 0
+        # if (np.abs(Z) <= 0.5 and R > 3.01
+        #     and np.deg2rad(179.5) <= P <= np.deg2rad(180.5)):
+        #     print('Got 2')
+        #     return 0
+        # elif (np.abs(Z) <= 0.5 and R > 3.01
+        #       and np.deg2rad(269.5) <= P <= np.deg2rad(270.5)):
+        #     print('Got 3')
+        #     return 0
         if self.wall_ck==False:
             Rc=2.460
             Zc=0.0
@@ -362,10 +414,9 @@ def hit_wall_circ(self, s, state):
 
     hit_wall_circ.terminal=True
 
-    def b_field_interp(self, R, Phi, Z, no_ripple=False):
+    def b_field_interp(self, R, Phi, Z):
         ''' Linear interpolation of B vector components at R, Phi, Z positions
         '''
-
         interp_points = np.vstack((R, Z)).transpose()
 
         br_intp = self.br_lin_intp.__call__(interp_points)
@@ -384,6 +435,17 @@ def b_field_interp(self, R, Phi, Z, no_ripple=False):
 
         return br_intp[0], bt_intp[0], bz_intp[0]
 
+    def b_field_interp_no_ripple(self, R, Phi, Z):
+        ''' Linear interpolation of B vector components at R, Phi, Z positions
+        '''
+        interp_points = np.vstack((R, Z)).transpose()
+
+        br_intp = self.br_lin_intp.__call__(interp_points)
+        bt_intp = self.bt_lin_intp.__call__(interp_points)
+        bz_intp = self.bz_lin_intp.__call__(interp_points)
+
+        return br_intp[0], bt_intp[0], bz_intp[0]
+
     def plot_trace(self, trace):
         '''
         Plots a summary of the magnetic field line trace.
@@ -511,3 +573,27 @@ def plot_trace_3D(self, trace):
         ax.set_xlabel('x [m]'); ax.set_ylabel('y [m]'); ax.set_zlabel('z [m]')
         plt.title('Magnetic field line trace')
         plt.show()
+
+    def plot_poincare(self, trace_plt, tor_angle=0):
+        '''
+        Plots Poincare plot at given toroidal angle
+
+        input:
+            - trace_plt : the magneticl field line data for several points
+                          one single time (dictionary)
+        '''
+
+        fig, ax = plt.subplots()
+
+        for ii in range(len(trace_plt)):
+            loc_cross = trace_plt[ii]['p'] % (tor_angle + 2*np.pi)
+            mask_cross = np.r_[False, loc_cross[1:]
+                               < loc_cross[:-1]] & np.r_[loc_cross[:-1]
+                                                         < loc_cross[1:],
+                                                         False]
+            ax.plot(trace_plt[ii]['r'][mask_cross],
+                    trace_plt[ii]['z'][mask_cross], 'o', markersize=3)
+
+        ax.set_aspect('equal')
+        ax.set_xlabel('R [m]')
+        ax.set_ylabel('Z [m]')

tofu/mag/mag_ripple/__init__.py

@@ -3,3 +3,20 @@
 '''
     Magnetic ripple
 '''
+
+import warnings
+import traceback
+
+try:
+    try:
+        from tofu.mag.mag_ripple import *
+    except Exception:
+        from .mag_ripple import *
+    del warnings, traceback
+except Exception as err:
+    msg = str(traceback.format_exc())
+    msg += "\n\n    => the optional sub-package is not usable\n"
+    warnings.warn(msg)
+    del msg, err
+
+__all__ = ['mag_ripple']