Added tests for drift_bounce_orbit and drift_shell in python test suite.

python/IRBEM/IRBEM.py

@@ -230,9 +230,13 @@ def drift_shell(self, X, maginput):
                 ctypes.byref(nposit))
         # Format the output into a dictionary, and convert ctypes arrays into
         # native Python format.
+        posit = np.array(posit)
+        nposit = np.array(nposit)
+        for i,n in enumerate(nposit):
+            posit[i,n:,:] = np.nan
         self.drift_shell_output = {'Lm':lm.value, 'blocal':np.array(blocal),
             'bmin':bmin.value, 'lstar':lstar.value, 'xj':xj.value, 
-            'POSIT':np.array(posit), 'Nposit':np.array(nposit)} 
+            'POSIT':posit, 'Nposit':nposit} 
         return self.drift_shell_output
 
     def drift_bounce_orbit(self, X, maginput, alpha=90, R0=1):
@@ -259,6 +263,11 @@ def drift_bounce_orbit(self, X, maginput, alpha=90, R0=1):
         maginput: dict
             The magnetic field input dictionary. See the online documentation for the valid
             keys and the corresponding models.
+        alpha: float
+            The local pitch angle.
+        R0: float
+            The radius, in units of RE, of the reference surface (i.e. altitude) between which
+            the line is traced.
 
         Returns
         -------
@@ -281,7 +290,6 @@ def drift_bounce_orbit(self, X, maginput, alpha=90, R0=1):
         blocal = blocalType()
 
         if self.TMI: print("Running IRBEM-LIB drift_bounce_orbit")
-
         self._irbem_obj.drift_bounce_orbit2_1_(ctypes.byref(self.kext), ctypes.byref(self.options),\
                 ctypes.byref(self.sysaxes), ctypes.byref(iyear),\
                 ctypes.byref(idoy), ctypes.byref(ut), ctypes.byref(x1), \
@@ -291,9 +299,13 @@ def drift_bounce_orbit(self, X, maginput, alpha=90, R0=1):
                 ctypes.byref(nposit), ctypes.byref(hmin), ctypes.byref(hmin_lon))
         # Format the output into a dictionary, and convert ctypes arrays into
         # native Python format.
+        posit = np.array(posit)
+        nposit = np.array(nposit)
+        for i,n in enumerate(nposit):
+            posit[i,n:,:] = np.nan
         self.drift_bounce_orbit_output = {'Lm':lm.value, 'blocal':np.array(blocal),
             'bmin':bmin.value, 'bmirr':bmirr.value, 'lstar':lstar.value, 'xj':xj.value,
-            'POSIT':np.array(posit), 'Nposit':np.array(nposit), 'hmin':hmin.value,
+            'POSIT':posit, 'Nposit':nposit, 'hmin':hmin.value,
             'hmin_lon': hmin_lon.value}
         return self.drift_bounce_orbit_output
 
@@ -1175,4 +1187,4 @@ def _load_shared_object(path=None):
 """
 beta = lambda Ek, Erest = 511: np.sqrt(1-((Ek/Erest)+1)**(-2)) # Ek,a dErest must be in keV
 gamma = lambda Ek, Erest = 511:np.sqrt(1-beta(Ek, Erest = 511)**2)**(-1/2)
-vparalel = lambda Ek, Bm, B, Erest = 511:c*beta(Ek, Erest)*np.sqrt(1 - np.abs(B/Bm))
+vparalel = lambda Ek, Bm, B, Erest = 511:c*beta(Ek, Erest)*np.sqrt(1 - np.abs(B/Bm))

python/IRBEM/test_IRBEM.py

@@ -23,6 +23,7 @@ class TestIRBEM(unittest.TestCase):
     """
     def setUp(self):
         self.model = IRBEM.MagFields(options=[0,0,0,0,0], verbose=False, kext='T89')
+        self.dipol_model = IRBEM.MagFields(options=[0,0,5,0,5], verbose=False, kext=0)
         # Create four sets of model inputs: one set of input values, 
         # array of input values, one input value with a string 
         # datetime object, and array of inputs including a 
@@ -170,6 +171,40 @@ def test_get_mlt(self):
         self.model.get_mlt(input_dict)
         self.assertAlmostEqual(self.model.get_mlt_output, true_MLT)
         return
+
+    def test_drift_shell(self):
+        """
+        Tests the drift_shell IRBEM function.
+        """
+        res = self.dipol_model.drift_shell(self.X, self.maginput)
+        Lm = res['Lm']
+        self.assertAlmostEqual(Lm, 4.326679, 2)
+        L_posit = self._compute_dipole_L_shell(res['POSIT'])
+        self.assertLessEqual(np.nanmax(np.abs(L_posit-Lm))/Lm, 1e-2)
+
+    def test_drift_bounce_orbit(self):
+        """
+        Tests the drift_bounce_orbit IRBEM function.
+        """
+        res = self.dipol_model.drift_bounce_orbit(self.X, self.maginput)
+        Lm = res['Lm']
+        self.assertAlmostEqual(Lm, 4.326679, 2)
+        alt = self.X['x1']
+        self.assertLessEqual((res['hmin']-alt)/alt, 1e-2)
+        L_posit = self._compute_dipole_L_shell(res['POSIT'])
+        self.assertLessEqual(np.nanmax(np.abs(L_posit-Lm))/Lm, 1e-2)
+
+    @staticmethod
+    def _compute_dipole_L_shell(posit):
+        """
+        Compute the dipole L-shell for each point in coordinate list.
+        """
+        x = posit[...,0]
+        y = posit[...,1]
+        z = posit[...,2]
+        r = np.sqrt(x**2+y**2+z**2)
+        theta = np.arctan2(np.sqrt(x**2+y**2),z)
+        return r/np.sin(theta)**2
 
     def assertAlmostEqualDict(self, A, B):
         """