separatrix bounds checking for flux/pn5 ode classes

BlackHolePerturbationToolkit · Nov 25, 2024 · e94aa6c · e94aa6c
1 parent 4e7eb23
commit e94aa6c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 12 deletions.
diff --git a/few/trajectory/ode/flux.py b/few/trajectory/ode/flux.py
@@ -64,9 +64,8 @@ def separatrix_buffer_dist(self):
         return 0.1
 
     def evaluate_rhs(self, p: float, e: float, x: float, *args) -> list[float]:
-        if 6 + 2*e > p:
-            return np.zeros(3)
-
+        if e < 0 or p < 6 + 2*e:
+            return [0., 0., 0., 0., 0., 0.,]
 
         # directly evaluate the numba kernel for speed
         Omega_phi, Omega_theta, Omega_r = _KerrGeoCoordinateFrequencies_kernel_inner(0., p, e, x)
@@ -138,14 +137,17 @@ def separatrix_buffer_dist(self):
 
 
     def evaluate_rhs(self, p: float, e: float, x: float, *args) -> list[float]:
+        if e < 0:
+             return [0., 0., 0., 0., 0., 0.,]
+
+        p_sep = _get_separatrix_kernel_inner(self.a, e, x)
+
+        if p < p_sep:
+             return [0., 0., 0., 0., 0., 0.,]
 
         # directly evaluate the numba kernel for speed
         Omega_phi, Omega_theta, Omega_r = _KerrGeoCoordinateFrequencies_kernel_inner(self.a, p, e, x)
 
-        p_sep = _get_separatrix_kernel_inner(self.a, e, x)
-        if e < 0 or p < p_sep:
-             return [0., 0., 0., 0., 0., 0.,]
-
         risco = _get_separatrix_kernel_inner(self.a, 0., x)
         u = _p_to_u(p, p_sep)
         w = e**0.5

diff --git a/few/trajectory/ode/pn5.py b/few/trajectory/ode/pn5.py
@@ -1,5 +1,5 @@
 from .base import ODEBase
-from ...utils.utility import _KerrGeoCoordinateFrequencies_kernel_inner
+from ...utils.utility import _KerrGeoCoordinateFrequencies_kernel_inner, _get_separatrix_kernel_inner
 from ...utils.pn_map import _Y_to_xI_kernel_inner
 
 from numba import njit
@@ -26,14 +26,17 @@ def evaluate_rhs(self, p: float, e: float, Y: float, *args) -> list[float]:
 
         if e < 0:
              return [0., 0., 0., 0., 0., 0.,]
-
+
+        # directly evaluate the numba kernels for speed
+        xI = _Y_to_xI_kernel_inner(self.a, p, e, Y)
+        p_sep = _get_separatrix_kernel_inner(self.a, e, xI)
+        if p < p_sep:
+             return [0., 0., 0., 0., 0., 0.,]
+
         pdot = dpdt8H_5PNe10(self.a, p, e, Y, 10, 10)
         edot = dedt8H_5PNe10(self.a, p, e, Y, 10, 8)
         Ydot = dYdt8H_5PNe10(self.a, p, e, Y, 7, 10)
 
-        # directly evaluate the numba kernel for speed
-        xI = _Y_to_xI_kernel_inner(self.a, p, e, Y)
-
         # directly evaluate the numba kernel for speed
         frequencies = _KerrGeoCoordinateFrequencies_kernel_inner(self.a, p, e, xI)