Merge pull request #336 from sblunt/fix-tau-tp

after_date can accept multiple values now

orbitize/basis.py

@@ -1140,7 +1140,7 @@ def tau_to_tp(tau, ref_epoch, period, after_date=None):
         tau (float or np.array): value of tau to convert
         ref_epoch (float or np.array): date (in days, typically MJD) that tau is defined relative to
         period (float or np.array): period (in years) that tau is noralized with
-        after_date (float): tp will be the first periastron after this date. If None, use ref_epoch.
+        after_date (float or np.array): tp will be the first periastron after this date. If None, use ref_epoch.
 
     Returns:
         float or np.array: corresponding t_p of the taus
@@ -1151,7 +1151,7 @@ def tau_to_tp(tau, ref_epoch, period, after_date=None):
 
     if after_date is not None:
         num_periods = (after_date - tp)/period_days
-        num_periods = int(np.ceil(num_periods))
+        num_periods = np.ceil(num_periods).astype(int)
 
         tp += num_periods * period_days
 

tests/test_conversions.py

@@ -22,6 +22,10 @@ def test_tau_tp_conversion():
     tp = basis.tau_to_tp(tau, ref_epoch, period, after_date=47000)
     assert tp == pytest.approx(51000 - 9 * 365.25, rel=1e-7)
 
+    tps = basis.tau_to_tp(tau, ref_epoch, period, after_date=np.array([47000, 51000]))
+    assert tps[0] == pytest.approx(51000 - 9 * 365.25, rel=1e-7)
+    assert tps[1] == pytest.approx(51000 + 365.25, rel=1e-7)
+
     tau3 = basis.tp_to_tau(tp, ref_epoch, period)
     assert tau == pytest.approx(tau3, rel=1e-7)