diff --git a/adam_core/time/tests/test_time.py b/adam_core/time/tests/test_time.py index 83ef4571..a3be0d50 100644 --- a/adam_core/time/tests/test_time.py +++ b/adam_core/time/tests/test_time.py @@ -1,6 +1,8 @@ import astropy.time import astropy.units +import numpy as np import numpy.testing as npt +import pyarrow as pa import pyarrow.compute as pc import pytest import quivr as qv @@ -87,6 +89,51 @@ def test_with_nulls(self): have = self.with_nulls.mjd() assert have.to_pylist() == [0, 10000.5, 20000.999999999999, None, None, None] + def test_from_list_ints(self): + data = [0, 10000, 20000] + have = Timestamp.from_mjd(data) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 0, 0] + + def test_from_list_floats(self): + data = [0.0, 10000.5, 20000.75] + have = Timestamp.from_mjd(data) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 43200_000_000_000, 64800_000_000_000] + + def test_from_numpy_ints(self): + data = np.array([0, 10000, 20000], dtype=np.int64) + have = Timestamp.from_mjd(data) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 0, 0] + + def test_from_numpy_floats(self): + data = np.array([0.0, 10000.5, 20000.75], dtype=np.float64) + have = Timestamp.from_mjd(data) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 43200_000_000_000, 64800_000_000_000] + + def test_from_pyarrow_ints(self): + data = pa.array([0, 10000, 20000], type=pa.int64()) + have = Timestamp.from_mjd(data) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 0, 0] + + def test_from_pyarrow_floats(self): + data = pa.array([0.0, 10000.5, 20000.75], type=pa.float64()) + have = Timestamp.from_mjd(data) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 43200_000_000_000, 64800_000_000_000] + + def test_from_pyarrow_table_float(self): + """ + This test case exercises chunked arrays, which have slightly different behavior + """ + data = pa.table({"mjd": [0.0, 10000.5, 20000.75]}) + have = Timestamp.from_mjd(data["mjd"]) + assert have.days.to_pylist() == [0, 10000, 20000] + assert have.nanos.to_pylist() == [0, 43200_000_000_000, 64800_000_000_000] + class TestAstropyTime: ts = Timestamp.from_kwargs( diff --git a/adam_core/time/time.py b/adam_core/time/time.py index 6cf04e63..8475210e 100644 --- a/adam_core/time/time.py +++ b/adam_core/time/time.py @@ -71,6 +71,7 @@ def from_iso8601( def from_mjd(cls, mjd: pa.lib.DoubleArray, scale: str = "tai") -> Timestamp: days = pc.floor(mjd) fractional_days = pc.subtract(mjd, days) + days = pc.cast(days, pa.int64()) nanos = pc.cast(pc.round(pc.multiply(fractional_days, 86400 * 1e9)), pa.int64()) return cls.from_kwargs(days=days, nanos=nanos, scale=scale)