Skip to content

Commit cace230

Browse files
committed
Assign NaN to results of action-angle approximations for unbound orbits in the orbit interface; fixes #383
1 parent 6a74587 commit cace230

File tree

3 files changed

+824
-19
lines changed

3 files changed

+824
-19
lines changed

galpy/actionAngle/actionAngleStaeckel.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,13 @@ def _evaluate(self, *args, **kwargs):
214214
else:
215215
# Set up the actionAngleStaeckelSingle object
216216
aASingle = actionAngleStaeckelSingle(
217-
R[0], vR[0], vT[0], z[0], vz[0], pot=self._pot, delta=delta
217+
R[0],
218+
vR[0],
219+
vT[0],
220+
z[0],
221+
vz[0],
222+
pot=self._pot,
223+
delta=delta[0] if hasattr(delta, "__len__") else delta,
218224
)
219225
return (
220226
numpy.atleast_1d(aASingle.JR(**copy.copy(kwargs))),

galpy/orbit/Orbits.py

+90-18
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
toPlanarPotential,
3939
)
4040
from ..potential.DissipativeForce import _isDissipative
41+
from ..potential.plotEscapecurve import _INF
4142
from ..potential.Potential import _check_c
4243
from ..util import conversion, coords, galpyWarning, galpyWarningVerbose, plot
4344
from ..util._optional_deps import (
@@ -2601,6 +2602,27 @@ def _setupaA(self, pot=None, type="staeckel", **kwargs):
26012602
self._aA = actionAngle.actionAngleSpherical(pot=self._aAPot, **kwargs)
26022603
return None
26032604

2605+
def _unbound_indx_and_aAkwargs(self):
2606+
"""Internal function to compute the index of unbound orbits"""
2607+
if self.dim() == 2:
2608+
indx = self.E(
2609+
pot=self._aAPot, use_physical=False, dontreshape=True
2610+
) < evaluateplanarPotentials(
2611+
toPlanarPotential(self._aAPot), _INF, use_physical=False
2612+
)
2613+
elif self.dim() == 3:
2614+
indx = self.E(
2615+
pot=self._aAPot, use_physical=False, dontreshape=True
2616+
) < evaluatePotentials(self._aAPot, _INF, 0.0, use_physical=False)
2617+
if hasattr(self._aA, "_delta"):
2618+
if hasattr(self._aA._delta, "__len__"):
2619+
aAkwargs = {"delta": self._aA._delta[indx]}
2620+
else:
2621+
aAkwargs = {"delta": self._aA._delta}
2622+
else:
2623+
aAkwargs = {}
2624+
return indx, aAkwargs
2625+
26042626
def _setup_EccZmaxRperiRap(self, pot=None, **kwargs):
26052627
"""Internal function to compute e,zmax,rperi,rap and cache it for reuse"""
26062628
self._setupaA(pot=pot, **kwargs)
@@ -2619,18 +2641,32 @@ def _setup_EccZmaxRperiRap(self, pot=None, **kwargs):
26192641
tz = numpy.zeros(self.size)
26202642
tvz = numpy.zeros(self.size)
26212643
# self.dim() == 1 error caught by _setupaA
2644+
# Exclude unbound orbits (aAkwargs deals with delta processing)
2645+
indx, aAkwargs = self._unbound_indx_and_aAkwargs()
26222646
(
26232647
self._aA_ecc,
26242648
self._aA_zmax,
26252649
self._aA_rperi,
26262650
self._aA_rap,
2651+
) = (
2652+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2653+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2654+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2655+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2656+
)
2657+
(
2658+
self._aA_ecc[indx],
2659+
self._aA_zmax[indx],
2660+
self._aA_rperi[indx],
2661+
self._aA_rap[indx],
26272662
) = self._aA.EccZmaxRperiRap(
2628-
self.R(use_physical=False, dontreshape=True),
2629-
self.vR(use_physical=False, dontreshape=True),
2630-
self.vT(use_physical=False, dontreshape=True),
2631-
tz,
2632-
tvz,
2663+
self.R(use_physical=False, dontreshape=True)[indx],
2664+
self.vR(use_physical=False, dontreshape=True)[indx],
2665+
self.vT(use_physical=False, dontreshape=True)[indx],
2666+
tz[indx],
2667+
tvz[indx],
26332668
use_physical=False,
2669+
**aAkwargs,
26342670
)
26352671
return None
26362672

@@ -2652,6 +2688,8 @@ def _setup_actionsFreqsAngles(self, pot=None, **kwargs):
26522688
tz = numpy.zeros(self.size)
26532689
tvz = numpy.zeros(self.size)
26542690
# self.dim() == 1 error caught by _setupaA
2691+
# Exclude unbound orbits (aAkwargs deals with delta processing)
2692+
indx, aAkwargs = self._unbound_indx_and_aAkwargs()
26552693
(
26562694
self._aA_jr,
26572695
self._aA_jp,
@@ -2662,14 +2700,36 @@ def _setup_actionsFreqsAngles(self, pot=None, **kwargs):
26622700
self._aA_wr,
26632701
self._aA_wp,
26642702
self._aA_wz,
2703+
) = (
2704+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2705+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2706+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2707+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2708+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2709+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2710+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2711+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2712+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2713+
)
2714+
(
2715+
self._aA_jr[indx],
2716+
self._aA_jp[indx],
2717+
self._aA_jz[indx],
2718+
self._aA_Or[indx],
2719+
self._aA_Op[indx],
2720+
self._aA_Oz[indx],
2721+
self._aA_wr[indx],
2722+
self._aA_wp[indx],
2723+
self._aA_wz[indx],
26652724
) = self._aA.actionsFreqsAngles(
2666-
self.R(use_physical=False, dontreshape=True),
2667-
self.vR(use_physical=False, dontreshape=True),
2668-
self.vT(use_physical=False, dontreshape=True),
2669-
tz,
2670-
tvz,
2671-
self.phi(use_physical=False, dontreshape=True),
2725+
self.R(use_physical=False, dontreshape=True)[indx],
2726+
self.vR(use_physical=False, dontreshape=True)[indx],
2727+
self.vT(use_physical=False, dontreshape=True)[indx],
2728+
tz[indx],
2729+
tvz[indx],
2730+
self.phi(use_physical=False, dontreshape=True)[indx],
26722731
use_physical=False,
2732+
**aAkwargs,
26732733
)
26742734
return None
26752735

@@ -2694,14 +2754,26 @@ def _setup_actions(self, pot=None, **kwargs):
26942754
# tz = numpy.zeros(self.size)
26952755
# tvz = numpy.zeros(self.size)
26962756
# self.dim() == 1 error caught by _setupaA
2697-
self._aA_jr, self._aA_jp, self._aA_jz = self._aA(
2698-
self.R(use_physical=False, dontreshape=True),
2699-
self.vR(use_physical=False, dontreshape=True),
2700-
self.vT(use_physical=False, dontreshape=True),
2701-
tz,
2702-
tvz,
2703-
self.phi(use_physical=False, dontreshape=True),
2757+
# Exclude unbound orbits
2758+
indx, aAkwargs = self._unbound_indx_and_aAkwargs()
2759+
(
2760+
self._aA_jr,
2761+
self._aA_jp,
2762+
self._aA_jz,
2763+
) = (
2764+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2765+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2766+
numpy.zeros_like(self.R(use_physical=False, dontreshape=True)) + numpy.nan,
2767+
)
2768+
self._aA_jr[indx], self._aA_jp[indx], self._aA_jz[indx] = self._aA(
2769+
self.R(use_physical=False, dontreshape=True)[indx],
2770+
self.vR(use_physical=False, dontreshape=True)[indx],
2771+
self.vT(use_physical=False, dontreshape=True)[indx],
2772+
tz[indx],
2773+
tvz[indx],
2774+
self.phi(use_physical=False, dontreshape=True)[indx],
27042775
use_physical=False,
2776+
**aAkwargs,
27052777
)
27062778
return None
27072779

0 commit comments

Comments
 (0)