predict(): accept one obscode per time (mixed-observatory prediction)

src/layup/predict.py

@@ -410,8 +410,11 @@ def predict(
     ----------
     data : numpy structured array
         The data to be processed.
-    obscode : str
-        The observer code.
+    obscode : str or sequence of str
+        The observer code. A single string is used for every time. A sequence of
+        length ``len(times)`` gives one observatory per time -- a mixed-observatory
+        prediction in one call (each time is light-corrected against its own
+        station).
     times : list
         The times for the predictions, in jd_tdb.
     primary_id_column_name : str
@@ -436,7 +439,22 @@ def predict(
 
     times_et = np.array([spice.str2et(f"jd {t} tdb") for t in times], dtype="<f8")
 
-    obs_data = np.array([(obscode, t) for t in times_et], dtype=[("stn", "<U10"), ("et", "<f8")])
+    # obscode may be a single code -- used for every time (the original behavior) --
+    # or one code per time, which predicts a mixed-observatory sequence in a single
+    # call: each time is light-corrected against its own station, and predict_sequence
+    # integrates each orbit once across the whole (sorted) set. A plain string
+    # reproduces the single-observatory result exactly.
+    if isinstance(obscode, str):
+        obscodes = [obscode] * len(times_et)
+    else:
+        obscodes = [str(o) for o in obscode]
+        if len(obscodes) != len(times_et):
+            raise ValueError(
+                f"predict: obscode has {len(obscodes)} entries but there are {len(times_et)} "
+                "times; pass a single obscode or exactly one per time."
+            )
+
+    obs_data = np.array(list(zip(obscodes, times_et)), dtype=[("stn", "<U10"), ("et", "<f8")])
 
     obs_pos_vel = layup_observatory.obscodes_to_barycentric(obs_data)
 

tests/layup/test_predict.py

@@ -125,6 +125,49 @@ def __init__(self, g=None):
     assert len(predictions) == n_uniq_ids * len(times)
 
 
+def test_predict_mixed_and_single_obscode():
+    """obscode may be a single string (used for all times) or one per time. A
+    string must be byte-identical to a same-code-per-time list (so existing
+    single-observatory callers are unaffected), a per-time list light-corrects
+    each time against its own station, and a length mismatch raises."""
+
+    class FakeCliArgs:
+        def __init__(self):
+            self.onsky_data = False
+
+    data = CSVDataReader(
+        get_test_filepath("fit_result_file_example.csv"), "csv", primary_id_column_name="provID"
+    ).read_rows()
+    epoch0 = float(data["epochMJD_TDB"][0]) + 2400000.5
+    times = [epoch0 - 100.0, epoch0 + 20.0, epoch0 + 250.0]
+    kw = dict(num_workers=1, cache_dir=None, primary_id_column_name="provID", args=FakeCliArgs())
+
+    # (1) a plain string == a same-code-per-time list, exactly (backward compatible)
+    a = predict(data, "X05", times, **kw)
+    b = predict(data, ["X05"] * len(times), times, **kw)
+    for c in a.dtype.names:
+        if a[c].dtype.kind in "fc":
+            np.testing.assert_array_equal(np.nan_to_num(a[c]), np.nan_to_num(b[c]))
+        else:
+            np.testing.assert_array_equal(a[c], b[c])
+
+    # (2) mixed obscodes: each time matches a single-obscode prediction at that station
+    obscodes = ["X05", "500", "X05"]
+    mixed = predict(data, obscodes, times, **kw)
+    by_epoch = {}
+    for r in mixed:
+        by_epoch.setdefault(round(float(r["epoch_JD_TDB"]), 6), {})[str(r["provID"])] = r
+    for oc, t in zip(obscodes, times):
+        for r in predict(data, oc, [t], **kw):
+            m = by_epoch[round(float(t), 6)][str(r["provID"])]
+            assert float(m["ra_deg"]) == pytest.approx(float(r["ra_deg"]), abs=1e-11)
+            assert float(m["dec_deg"]) == pytest.approx(float(r["dec_deg"]), abs=1e-11)
+
+    # (3) one-obscode-per-time is required when a sequence is given
+    with pytest.raises(ValueError):
+        predict(data, ["X05", "500"], times, **kw)
+
+
 def test_predict_sequence_marches_equivalently():
     """predict_sequence integrates each orbit once across the sorted set of times
     and interpolates at each (ASSIST integrate_or_interpolate), rather than
@@ -139,7 +182,6 @@ def __init__(self):
     data = CSVDataReader(
         get_test_filepath("fit_result_file_example.csv"), "csv", primary_id_column_name="provID"
     ).read_rows()
-
     epoch0 = float(data["epochMJD_TDB"][0]) + 2400000.5
     # times both before and after the epoch, exercising the forward + backward passes
     times = [epoch0 - 200.0, epoch0 - 40.0, epoch0 + 30.0, epoch0 + 300.0]