Workaround PyGeopack issue

radiation-belts · Jan 23, 2025 · 20ff111 · 20ff111
1 parent d322fef
commit 20ff111
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 20 deletions.
diff --git a/docs/index.rst b/docs/index.rst
@@ -54,8 +54,8 @@ Below is code which calculates L* using the magnetic fields obtain from TS05 and
     time = datetime(2015, 10, 2)
     params = models.get_tsyganenko_params(time)
     
-    # Evaluate TS05 model on regular grid 
-    axis = np.arange(-10, 10, 0.50)
+    # Evaluate TS05 model on regular grid. This takes 10-15 sec.
+    axis = np.arange(-10, 10, 0.25)
     x, y, z = np.meshgrid(axis, axis, axis)
     model = models.get_tsyganenko(
         "TS05", params, time,
@@ -65,7 +65,7 @@ Below is code which calculates L* using the magnetic fields obtain from TS05 and
         inner_boundary=1
     )
 
-    # Calculate L* 
+    # Calculate L*. This takes 20-25 sec.
     result = invariants.calculate_LStar(
         model,
         starting_point=(-6.6, 0, 0),

diff --git a/rbinvariantslib/models.py b/rbinvariantslib/models.py
@@ -468,19 +468,34 @@ def get_tsyganenko(
     Bx[~mask] = np.nan
     By[~mask] = np.nan
     Bz[~mask] = np.nan
-
-    Bx[mask], By[mask], Bz[mask] = gp.ModelField(
-        x_re_sm[mask],
-        y_re_sm[mask],
-        z_re_sm[mask],
-        Date=gp_date,
-        ut=gp_ut,
-        Model=gp_model,
-        CoordIn='SM',
-        CoordOut='SM',
-        **params
-    )
-
+
+    Bx_masked = np.zeros(mask.sum())
+    By_masked = Bx_masked.copy()
+    Bz_masked = Bx_masked.copy()
+    x_masked = x_re_sm[mask]
+    y_masked = y_re_sm[mask]
+    z_masked = z_re_sm[mask]
+    batch_size = 10_000
+
+    for step in range(0, mask.sum(), batch_size):
+        i = step
+        j = step + batch_size
+        Bx_masked[i:j], By_masked[i:j], Bz_masked[i:j] = gp.ModelField(
+            x_masked[i:j],
+            y_masked[i:j],
+            z_masked[i:j],
+            Date=gp_date,
+            ut=gp_ut,
+            Model=gp_model,
+            CoordIn='SM',
+            CoordOut='SM',
+            **params
+        )
+
+    Bx[mask] = Bx_masked
+    By[mask] = By_masked
+    Bz[mask] = Bz_masked    
+
     # Convert from nT to Gauss
     Bx = nanoTesla2Gauss(Bx)
     By = nanoTesla2Gauss(By)

diff --git a/rbinvariantslib/tests/test_examples.py b/rbinvariantslib/tests/test_examples.py
@@ -11,7 +11,7 @@ def test_tsyganenko():
     params = models.get_tsyganenko_params(time)
 
     # Evaluate TS05 model on regular grid
-    axis = np.arange(-10, 10, 0.5)
+    axis = np.arange(-10, 10, 0.25)
     x, y, z = np.meshgrid(axis, axis, axis)
     model = models.get_tsyganenko(
         "TS05", params, time,
@@ -20,15 +20,15 @@ def test_tsyganenko():
         z_re_sm_grid=z,
         inner_boundary=1.5
     )
-
+    
     # Calculate L*
     result = invariants.calculate_LStar(
         model,
         starting_point=(-6.6, 0, 0),
         starting_pitch_angle=60
     )
-    
-    assert abs(result.LStar - 5.79) < .1
+
+    assert abs(result.LStar - 5.8258) < .1
 
 
 def test_swmf():