Merge pull request #610 from nickssl/master

Added sm2mlt

pyspedas/__init__.py

@@ -28,6 +28,9 @@
 from .cotrans.cotrans_get_coord import cotrans_get_coord
 from .cotrans.cotrans_set_coord import cotrans_set_coord
 from .cotrans.tvector_rotate import tvector_rotate
+from .cotrans.cart2spc import cart2spc
+from .cotrans.spc2cart import spc2cart
+from .cotrans.sm2mlt import sm2mlt
 
 from .mms import mms_load_mec, mms_load_fgm, mms_load_scm, mms_load_edi, \
     mms_load_edp, mms_load_eis, mms_load_feeps, \

pyspedas/cotrans/cart2spc.py

@@ -0,0 +1,51 @@
+"""
+Convert Cartesian coordinates to spherical coordinates.
+
+This function is similar to cart2spc.pro in IDL SPEDAS.
+"""
+
+import numpy as np
+
+
+def cart2spc(x, y, z):
+    """
+    Convert Cartesian coordinates to spherical coordinates.
+
+    Parameters:
+    -----------
+    x : float
+        X coordinate in Cartesian system.
+    y : float
+        Y coordinate in Cartesian system.
+    z : float
+        Z coordinate in Cartesian system.
+
+    Returns:
+    --------
+    r : float
+        Radial distance from the origin.
+    theta : float
+        Polar angle (angle from the positive z-axis to the point) in radians. It ranges from [0, pi].
+    phi : float
+        Azimuthal angle (angle in the xy-plane from the positive x-axis to the point) in radians. It ranges from [0, 2*pi].
+
+    Notes:
+    ------
+    - Uses the following equations for conversion:
+        r = sqrt(x^2 + y^2 + z^2)
+        theta = arccos(z/r)
+        phi = 2*pi - arccos(x/sqrt(x^2 + y^2)) if y < 0, otherwise arccos(x/sqrt(x^2 + y^2))
+    """
+
+    x, y, z = np.array(x), np.array(y), np.array(z)
+    r = np.sqrt(x**2 + y**2 + z**2)
+
+    # Calculate phi
+    maskn = np.where(y < 0, 1, 0)
+    maskp = np.where(y > 0, 1, 0)
+    phi = 2*np.pi - maskn * np.arccos(x/np.sqrt(x**2 + y**2)) + maskp * np.arccos(x/np.sqrt(x**2 + y**2))
+
+    # Calculate theta
+    theta = np.arccos(z/r)
+
+    return r, theta, phi

pyspedas/cotrans/sm2mlt.py

@@ -0,0 +1,44 @@
+'''
+Convert Solar Magnetic (SM) coordinates to Magnetic Local Time (MLT).
+
+This function is similar to sm2mlt.pro in IDL SPEDAS.
+'''
+import numpy as np
+from pyspedas import cart2spc
+
+
+def sm2mlt(x_sm, y_sm, z_sm):
+    """
+    Convert Solar Magnetic (SM) coordinates to Magnetic Local Time (MLT).
+
+    Parameters:
+    -----------
+    x_sm : float
+        X coordinate in Solar Magnetic system.
+    y_sm : float
+        Y coordinate in Solar Magnetic system.
+    z_sm : float
+        Z coordinate in Solar Magnetic system.
+
+    Returns:
+    --------
+    mlt_0 : float
+        Magnetic Local Time.
+
+    Notes:
+    ------
+    - First, the SM coordinates are converted to spherical coordinates.
+    - The azimuthal angle from the conversion (phi) is then used to calculate MLT.
+    - If the resulting MLT is greater than 24, it is corrected by subtracting 24.
+    """
+
+    # Convert to spherical coordinates
+    r, theta, phi = cart2spc(x_sm, y_sm, z_sm)
+
+    # Calculate initial MLT
+    mlt_0 = 12.0 + phi * 24.0 / (2.0 * np.pi)
+
+    # Fix MLTs greater than 24
+    mlt_0[mlt_0 > 24.0] -= 24.0
+
+    return mlt_0

pyspedas/cotrans/spc2cart.py

@@ -0,0 +1,43 @@
+'''
+Convert spherical coordinates (r, theta, phi) to Cartesian coordinates (x, y, z).
+
+This function is similar to spc2cart.pro in IDL SPEDAS.
+'''
+import numpy as np
+
+
+def spc2cart(r, theta, phi):
+    """
+    Convert spherical coordinates (r, theta, phi) to Cartesian coordinates (x, y, z).
+
+    Parameters:
+    -----------
+    r : float
+        Radial distance from the origin.
+    theta : float
+        Polar angle (angle from the positive z-axis to the point) in radians. It ranges from [0, pi].
+    phi : float
+        Azimuthal angle (angle in the xy-plane from the positive x-axis to the point) in radians. It ranges from [0, 2*pi].
+
+    Returns:
+    --------
+    x : float
+        X coordinate in Cartesian system.
+    y : float
+        Y coordinate in Cartesian system.
+    z : float
+        Z coordinate in Cartesian system.
+
+    Notes:
+    ------
+    - Uses the following equations for conversion:
+        x = r*sin(theta)*cos(phi)
+        y = r*sin(theta)*sin(phi)
+        z = r*cos(theta)
+    """
+    r, theta, phi = np.array(r), np.array(theta), np.array(phi)
+    x = r * np.sin(theta) * np.cos(phi)
+    y = r * np.sin(theta) * np.sin(phi)
+    z = r * np.cos(theta)
+
+    return x, y, z

pyspedas/cotrans/tests/cotrans.py

@@ -16,11 +16,12 @@
 from pyspedas.cotrans.cotrans import cotrans
 from pyspedas.cotrans.fac_matrix_make import fac_matrix_make
 from pytplot import get_data, store_data, del_data
-from pyspedas import cotrans_get_coord, cotrans_set_coord
+from pyspedas import cotrans_get_coord, cotrans_set_coord, sm2mlt
 
 
 class CotransTestCases(unittest.TestCase):
     """Tests for cotrans."""
+
     def test_fac_matrix_make(self):
         doesntexist = fac_matrix_make('doesnt_exist')
 
@@ -37,19 +38,19 @@ def test_get_set_coord_wrappers(self):
         self.assertTrue(setcoord)
         after = cotrans_get_coord('test_coord')
         self.assertTrue(after == 'GSE')
-        md = get_data('test_coord',metadata=True)
+        md = get_data('test_coord', metadata=True)
         md['data_att']['units'] = 'km'
         setcoord = cotrans_set_coord('test_coord', 'GSM')
         self.assertTrue(setcoord)
-        md_after = get_data('test_coord',metadata=True)
+        md_after = get_data('test_coord', metadata=True)
         after = cotrans_get_coord('test_coord')
         self.assertTrue(after == 'GSM')
         self.assertTrue(md_after['data_att']['units'] == 'km')
         setcoord = cotrans_set_coord('doesnt_exist', 'GSM')
 
     def test_get_set_coords(self):
         """ Test for pytplot.set_coords/get_coords """
-        from pytplot import set_coords,get_coords
+        from pytplot import set_coords, get_coords
 
         del_data()
         doesntexist = get_coords('test_coord')
@@ -62,19 +63,19 @@ def test_get_set_coords(self):
         self.assertTrue(setcoord)
         after = get_coords('test_coord')
         self.assertTrue(after == 'GSE')
-        md = get_data('test_coord',metadata=True)
+        md = get_data('test_coord', metadata=True)
         md['data_att']['units'] = 'km'
         setcoord = set_coords('test_coord', 'GSM')
         self.assertTrue(setcoord)
-        md_after = get_data('test_coord',metadata=True)
+        md_after = get_data('test_coord', metadata=True)
         after = get_coords('test_coord')
         self.assertTrue(after == 'GSM')
         self.assertTrue(md_after['data_att']['units'] == 'km')
         setcoord = set_coords('doesnt_exist', 'GSM')
 
     def test_get_set_units(self):
         """ Test for pytplot.set_coords/get_coords """
-        from pytplot import set_units,get_units, set_coords, get_coords
+        from pytplot import set_units, get_units, set_coords, get_coords
 
         del_data()
         doesntexist = get_units('test_units')
@@ -86,19 +87,19 @@ def test_get_set_units(self):
         self.assertTrue(setunits)
         after = get_units('test_units')
         self.assertTrue(after == 'Km')
-        set_coords('test_units','GEO')
+        set_coords('test_units', 'GEO')
         setunits = set_units('test_units', 'mm')
         self.assertTrue(setunits)
-        coords_after=get_coords('test_units')
-        units_after=get_units('test_units')
+        coords_after = get_coords('test_units')
+        units_after = get_units('test_units')
         self.assertTrue(coords_after == 'GEO')
         self.assertTrue(units_after == 'mm')
 
     def test_dsl2gse(self):
         """Test themis.cotrans.dsl2gse."""
         del_data()
         # Try with missing variables. It should exit without problems.
-        dsl2gse('tha_fgl_dsl','tha_fgl_gse')
+        dsl2gse('tha_fgl_dsl', 'tha_fgl_gse')
         # Now load the needed variables.
         time_range = ['2017-03-23 00:00:00', '2017-03-23 23:59:59']
         pyspedas.themis.state(probe='a', trange=time_range,
@@ -154,7 +155,7 @@ def test_cotrans_coord_mismatch(self):
         # Metadata coordinate system is GEI, but requesting GSM->GEO transform.  This should generate an error message
         # and return failure.
         result = cotrans(name_in=name_in, name_out=name_out,
-                coord_in="gsm", coord_out="geo")
+                         coord_in="gsm", coord_out="geo")
         self.assertTrue(result == 0)
 
     def test_cotrans_igrf(self):
@@ -235,7 +236,7 @@ def test_all_cotrans(self):
         count = 0
         # Test non-existent systems.
         result = cotrans(name_out=name1, time_in=t, data_in=d,
-                coord_in="badcoord", coord_out="gei")
+                         coord_in="badcoord", coord_out="gei")
         self.assertTrue(result == 0)
         result = cotrans(name_out=name1, time_in=t, data_in=d,
                          coord_in="gei", coord_out="badcoord")
@@ -269,6 +270,26 @@ def test_all_cotrans(self):
                 self.assertTrue(abs(dd1[1]-dd2[1]) <= 1e-6)
                 self.assertTrue(abs(dd1[2]-dd2[2]) <= 1e-6)
 
+    def test_mlt(self):
+        '''Test sm2mlt.
+
+        Data is from goes18, 2023-01-01. [0, 100, 600, 1000]
+        '''
+
+        # SM coordinates and MLT values from IDL
+        sm_idl = [[33199.660, 25815.376, 2995.6487], [18814.747, 37614.243, 2990.7136],
+                  [-40428.563, -11607.452, 2995.3227], [15798.372, -38974.509, 2995.0028]]
+        mlt_idl = [14.524527, 16.228377, 1.0679544, 7.4710163]
+
+        x = [item[0] for item in sm_idl]
+        y = [item[1] for item in sm_idl]
+        z = [item[2] for item in sm_idl]
+
+        mlt_python = sm2mlt(x, y, z)
+
+        for i in range(4):
+            self.assertTrue(abs(mlt_idl[i]-mlt_python[i]) <= 1e-6)
+
 
 if __name__ == '__main__':
     unittest.main()