Merge pull request #21 from BayaneMdW/models

add models and fix

space/coordinates/coordinates.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pandas as pd
 
+from ..smath import norm as mnorm
+
 
 def spherical_to_cartesian(r, theta, phi):
     x = r * np.cos(theta)
@@ -22,6 +24,14 @@ def cartesian_to_spherical(x, y, z):
     return r, theta, phi
 
 
+def cartesian_to_parabolic(x, y, z, xf):
+    r = np.sqrt((x - xf) ** 2 + y ** 2 + z ** 2)
+    s = np.sqrt((x - xf) + r)
+    t = np.sqrt(-(x - xf) + r)
+    p = np.arctan2(z, y)
+    return s, t, p
+
+
 def choice_coordinate_system(R, theta, phi, **kwargs):
     coord_sys = kwargs.get('coord_sys','cartesian')
     if coord_sys == 'cartesian':
@@ -44,51 +54,57 @@ def add_cst_radiuus(pos, cst,coord_sys):
     return pd.DataFrame({'X': x, 'Y': y, 'Z': z})
 
 
-def swi_base(omni_data):
-    norm_V = np.sqrt(omni_data.Vx ** 2 + (omni_data.Vy + 29.8) ** 2 + omni_data.Vz ** 2)
-    X = np.array([-np.array([vx, vy + 29.8, vz]) / V for vx, vy, vz, V in
-                  zip(omni_data.Vx.values, omni_data.Vy.values, omni_data.Vz.values, norm_V)])
-    B = np.array([np.array([bx, by, bz]) for bx, by, bz in zip(np.sign(omni_data.COA.values) * omni_data.Bx.values,
-                                                               np.sign(omni_data.COA.values) * omni_data.By.values,
-                                                               np.sign(omni_data.COA.values) * omni_data.Bz.values)])
-
-    Z = np.cross(X, B)
-
-    norm_cross = np.array(np.sqrt(Z[:, 0] ** 2 + Z[:, 1] ** 2 + Z[:, 2] ** 2))
-    Z = np.array([z / n for z, n in zip(Z, norm_cross)])
-    Y = np.cross(Z, X)
+def change_coordinate_system(x, y, z, x_uni, y_uni, z_uni):
+    X = x_uni[:,0]*x + x_uni[:,1]*y + x_uni[:,2]*z
+    Y = y_uni[:,0]*x + y_uni[:,1]*y + y_uni[:,2]*z
+    Z = z_uni[:,0]*x + z_uni[:,1]*y + z_uni[:,2]*z
     return X, Y, Z
 
+def gpim_base(vx, vy, vz, bx, by, bz):
+    V = mnorm(vx, vy, vz)
+    X = np.array([-vx / V, -vy / V, -vz / V])
+    B_scal_X = bx * X[0] + by * X[1] + bz * X[2]
+    Y = np.zeros_like(X)
+    sign = np.sign(B_scal_X)
+
+    if isinstance(bx, float) | isinstance(bx, int):
+        if bx == 0:
+            if by != 0:
+                sign = -np.sign(by)
+            else:
+                sign = np.sign(bz)
+    else:
+        sign[(bx == 0) & (by != 0)] = -np.sign(by[(bx == 0) & (by != 0)])
+        sign[(bx == 0) & (by == 0)] = np.sign(bz[(bx == 0) & (by == 0)])
 
-def to_swi(omni_data, msh_data, pos_msh):
-    X_swi, Y_swi, Z_swi = swi_base(omni_data)
-    data = msh_data.copy()
-    o_data = omni_data.copy()
-    pos = pos_msh.copy()
-
-
-
-    o_data['Vx'] = X_swi[:,0]*omni_data['Vx']+X_swi[:,1]*(omni_data['Vy']+29.8)+X_swi[:,2]*omni_data['Vz']
-    o_data['Vy'] = Y_swi[:,0]*omni_data['Vx']+Y_swi[:,1]*(omni_data['Vy']+29.8)+Y_swi[:,2]*omni_data['Vz']
-    o_data['Vz'] = Z_swi[:,0]*omni_data['Vx']+Z_swi[:,1]*(omni_data['Vy']+29.8)+Z_swi[:,2]*omni_data['Vz']
-
-    o_data['Bx'] = np.sign(omni_data['COA'])*(X_swi[:,0]*omni_data['Bx']+X_swi[:,1]*omni_data['By']+X_swi[:,2]*omni_data['Bz'])
-    o_data['By'] = np.sign(omni_data['COA'])*(Y_swi[:,0]*omni_data['Bx']+Y_swi[:,1]*omni_data['By']+Y_swi[:,2]*omni_data['Bz'])
-    o_data['Bz'] = np.sign(omni_data['COA'])*(Z_swi[:,0]*omni_data['Bx']+Z_swi[:,1]*omni_data['By']+Z_swi[:,2]*omni_data['Bz'])
-
-    data['Vx'] = X_swi[:,0]*msh_data['Vx']+X_swi[:,1]*msh_data['Vy']+X_swi[:,2]*msh_data['Vz']
-    data['Vy'] = Y_swi[:,0]*msh_data['Vx']+Y_swi[:,1]*msh_data['Vy']+Y_swi[:,2]*msh_data['Vz']
-    data['Vz'] = Z_swi[:,0]*msh_data['Vx']+Z_swi[:,1]*msh_data['Vy']+Z_swi[:,2]*msh_data['Vz']
-
-    data['Bx'] = np.sign(omni_data['COA'])*(X_swi[:,0]*msh_data['Bx']+X_swi[:,1]*msh_data['By']+X_swi[:,2]*msh_data['Bz'])
-    data['By'] = np.sign(omni_data['COA'])*(Y_swi[:,0]*msh_data['Bx']+Y_swi[:,1]*msh_data['By']+Y_swi[:,2]*msh_data['Bz'])
-    data['Bz'] = np.sign(omni_data['COA'])*(Z_swi[:,0]*msh_data['Bx']+Z_swi[:,1]*msh_data['By']+Z_swi[:,2]*msh_data['Bz'])
-
-
-    pos['X'] = X_swi[:,0]*pos_msh['X']+X_swi[:,1]*pos_msh['Y']+X_swi[:,2]*pos_msh['Z']
-    pos['Y'] = Y_swi[:,0]*pos_msh['X']+Y_swi[:,1]*pos_msh['Y']+Y_swi[:,2]*pos_msh['Z']
-    pos['Z'] = Z_swi[:,0]*pos_msh['X']+Z_swi[:,1]*pos_msh['Y']+Z_swi[:,2]*pos_msh['Z']
+    Y[1] = -sign * by + sign * B_scal_X * X[1]
+    Y[2] = -sign * bz + sign * B_scal_X * X[2]
 
+    Y_norm = mnorm(Y[0], Y[1], Y[2])
+    Y = np.array([Y[0] / Y_norm, Y[1] / Y_norm, Y[2] / Y_norm])
+    X, Y = X.T, Y.T
+    Z = np.cross(X, Y)
+    return X, Y, Z
 
-    return data,pos,o_data
 
+def swi_base(vx, vy, vz, bx, by, bz):
+    coa_zhang19 = np.arctan(bx / np.sqrt(by ** 2 + bz ** 2))
+    sign = np.sign(coa_zhang19)
+    if isinstance(coa_zhang19, float) or isinstance(coa_zhang19, int):
+        if coa_zhang19 == 0:
+            sign = np.sign(by)
+            if by == 0:
+                sign = np.sign(bz)
+    else:
+        sign[coa_zhang19 == 0] = np.sign(by[coa_zhang19 == 0])
+        sign[(coa_zhang19 == 0) & (by == 0)] = np.sign(bz[(coa_zhang19 == 0) & (by == 0)])
+
+    V = mnorm(vx, vy, vz)
+    X = np.array([-vx / V, -vy / V, -vz / V])
+    B = np.array([sign * bx, sign * by, sign * bz])
+    Z = np.cross(X.T, B.T).T
+    Z_norm = mnorm(Z[0], Z[1], Z[2])
+    Z = np.array([Z[0] / Z_norm, Z[1] / Z_norm, Z[2] / Z_norm])
+    X, Z = X.T, Z.T
+    Y = np.cross(Z, X)
+    return X, Y, Z