Implementation of the drift_bounce_orbit python wrapper

python/IRBEM/IRBEM.py

@@ -213,7 +213,7 @@ def drift_shell(self, X, maginput):
         # DEFINE OUTPUTS HERE        
         positType = (((ctypes.c_double * 3) * 1000) * 48)
         posit = positType()
-        npositType = (48 * ctypes.c_long)
+        npositType = (48 * ctypes.c_int)
         nposit = npositType()
         lm, lstar, bmin, xj = [ctypes.c_double() for i in range(4)]
         blocalType = ((ctypes.c_double * 1000) * 48)
@@ -230,14 +230,84 @@ 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):
-        raise NotImplementedError()
-        return
+    def drift_bounce_orbit(self, X, maginput, alpha=90, R0=1):
+        """
+        This function traces a full drift-bounce orbit for particles with a specified pitch
+        angle at the input location.  The output is a full array of positions of the
+        drift-bounce orbit, usefull for computing drift-bounce averages.
+        Key differences from `drif_shell`:
+         (1) only positions between mirror points are returned,
+         (2) 25 rather than 48 azimuths are returned,
+         (3) Lstar accuracy respects options(3) and options(4),
+         (4) a new parameter R0 is required which specifies the minimum radial distance allowed
+             along the drift path (usually R0=1, but use R0<1 in the drift loss cone),
+         (5) hmin and hmin_lon outputs provide the altitude and longitude (GDZ) of the minimum
+             altitude point along the orbit (among those traced, not just those returned). A set
+             of internal/external field can be selected.
+
+        Parameters
+        ----------
+        X: dict
+            A dictionary that specifies the input time and location. The `time` key can be a
+            ISO-formatted time string, or a `datetime.datetime` or `pd.TimeStamp` objects.
+            The three location keys: `x1`, `x2`, and `x3` specify the location in the `sysaxes`.
+        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
+        -------
+        dict
+            Contains keys Lm, lstar or Φ, blocal, bmin, bmirr, XJ, POSIT, Nposit, hmin, hmin_lon
+        """
+        # Prep the magnetic field model inputs and samping spacetime location.
+        self._prepMagInput(maginput)
+        iyear, idoy, ut, x1, x2, x3 = self._prepTimeLoc(X)
+        alpha = ctypes.c_double(alpha)
+        R0 = ctypes.c_double(R0)
+
+        # DEFINE OUTPUTS HERE
+        positType = (((ctypes.c_double * 3) * 1000) * 25)
+        posit = positType()
+        npositType = (25 * ctypes.c_int)
+        nposit = npositType()
+        lm, lstar, bmin, bmirr, xj, hmin, hmin_lon = [ctypes.c_double() for i in range(7)]
+        blocalType = ((ctypes.c_double * 1000) * 25)
+        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), \
+                ctypes.byref(x2), ctypes.byref(x3), ctypes.byref(alpha), ctypes.byref(self.maginput), \
+                ctypes.byref(R0), ctypes.byref(lm), ctypes.byref(lstar), ctypes.byref(blocal), \
+                ctypes.byref(bmin), ctypes.byref(bmirr), ctypes.byref(xj), ctypes.byref(posit), \
+                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':posit, 'Nposit':nposit, 'hmin':hmin.value,
+            'hmin_lon': hmin_lon.value}
+        return self.drift_bounce_orbit_output
 
     def find_mirror_point(self, X, maginput, alpha):
         """
@@ -1117,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):
         """