Add calculate_observing_night

Author: moeyensj

src/adam_core/observers/__init__.py

@@ -1,5 +1,6 @@
 # flake8: noqa: F401
 from .observers import Observers
 from .state import get_observer_state
+from .utils import calculate_observing_night
 
 __all__ = ["Observers", "get_observer_state"]

src/adam_core/observers/tests/test_utils.py

@@ -0,0 +1,68 @@
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from ...time import Timestamp
+from ..utils import calculate_observing_night
+
+
+def test_calculate_observing_night() -> None:
+    # Rubin Observatory is UTC-7
+    # Reasonable observating times would be +- 12 hours from local midnight
+    # or 7:00 to 17:00 UTC
+
+    # ZTF is UTC-8
+    # Reasonable observating times would be +- 12 hours from local midnight
+    # or 8:00 to 16:00 UTC
+
+    # M22 is UTC+2
+    # Reasonable observating times would be +- 12 hours from local midnight
+    # or 22:00 to 10:00 UTC
+
+    # 000 is UTC
+    # Reasonable observating times would be +- 12 hours from local midnight
+    # or 0:00 to 12:00 UTC
+
+    codes = pa.array([
+        "X05", "X05", "X05", "X05", "X05",
+        "I41", "I41", "I41", "I41", "I41",
+        "M22", "M22", "M22", "M22", "M22",
+        "000", "000", "000", "000", "000",
+    ])
+    times_utc = Timestamp.from_mjd(
+        [   
+            # Rubin Observatory
+            59000 + 7 / 24 - 12 / 24,
+            59000 + 7 / 24 - 6 / 24,
+            59000 + 7 / 24,
+            59000 + 7 / 24 + 6 / 24,
+            59000 + 7 / 24 + 12 / 24,
+            # ZTF
+            59000 + 8 / 24 - 12 / 24,
+            59000 + 8 / 24 - 4 / 24,
+            59000 + 8 / 24,
+            59000 + 8 / 24 + 4 / 24,
+            59000 + 8 / 24 + 12 / 24,
+            # M22
+            59000 - 2 / 24 - 12 / 24,
+            59000 - 2 / 24 - 6 / 24,
+            59000 - 2 / 24,
+            59000 - 2 / 24 + 6 / 24,
+            59000 - 2 / 24 + 12 / 24,
+            # 000
+            59000 - 12 / 24,
+            59000 - 6 / 24,
+            59000,
+            59000 + 6 / 24,
+            59000 + 12 / 24,
+        ],
+        scale="utc",
+    )
+
+    observing_night = calculate_observing_night(codes, times_utc)
+    desired = pa.array([
+        58999, 58999, 58999, 58999, 59000,
+        58999, 58999, 58999, 58999, 59000,
+        58999, 58999, 58999, 58999, 59000,
+        58999, 58999, 58999, 58999, 59000,
+    ])  
+    assert pc.all(pc.equal(observing_night, desired)).as_py()

src/adam_core/observers/utils.py

@@ -0,0 +1,57 @@
+from datetime import timedelta
+from zoneinfo import ZoneInfo
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+from astropy.time import Time
+
+from ..time import Timestamp
+from .observers import OBSERVATORY_PARALLAX_COEFFICIENTS
+
+
+def calculate_observing_night(codes: pa.Array, times: Timestamp) -> pa.Array:
+    """
+    Compute the observing night for a given set of observatory codes and times. An observing night is defined as the night
+    during which the observation is made, in the local time of the observatory +- 12 hours. The observing night is defined
+    as the integer MJD of the night in the local time of the observatory.
+
+    Parameters
+    ----------
+    codes : pyarrow.Array (N)
+        An array of observatory codes.
+    times : Timestamp (N)
+        An array of times.
+
+    Returns
+    -------
+    observing_night : pyarrow.Array (N)
+        An array of observing nights.
+    """
+    half_day = timedelta(hours=12)
+    observing_night = np.empty(len(times), dtype=np.int64)
+
+    for code in pc.unique(codes):
+
+        parallax_coefficients = OBSERVATORY_PARALLAX_COEFFICIENTS.select("code", code)
+        if len(parallax_coefficients) == 0:
+            raise ValueError(f"Unknown observatory code: {code}")
+
+        tz = ZoneInfo(parallax_coefficients.timezone()[0])
+
+        mask = pc.equal(codes, code)
+        times_code = times.apply_mask(mask)
+
+        observing_night_code = Time(
+            [
+                time + time.astimezone(tz).utcoffset() - half_day
+                for time in times_code.to_astropy().datetime
+            ],
+            format="datetime",
+            scale="utc",
+        )
+        observing_night_code = observing_night_code.mjd.astype(int)
+
+        observing_night[mask.to_numpy(zero_copy_only=False)] = observing_night_code
+
+    return pa.array(observing_night)