Merge pull request #1090 from lsst/tickets/DM-46581

DM-46581: Speed up TimeConverter

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "deprecated >= 1.2",
     "pydantic >=2,<3.0",
     "pyarrow >= 0.16",
+    "numpy",
 ]
 
 dynamic = ["version"]

python/lsst/daf/butler/_timespan.py

@@ -335,14 +335,23 @@ def __str__(self) -> str:
                 tail = f"{self.end.tai.strftime(fmt)})"
         return head + tail
 
+    def __repr_astropy__(self, t: astropy.time.Time | None) -> str:
+        # Provide our own repr for astropy time.
+        # For JD times we want to use jd1 and jd2 to maintain precision.
+        if isinstance(t, astropy.time.Time):
+            if t.format == "jd":
+                return f"astropy.time.Time({t.jd1}, {t.jd2}, scale='{t.scale}', format='{t.format}')"
+            else:
+                return f"astropy.time.Time('{t}', scale='{t.scale}', format='{t.format}')"
+        return str(t)
+
     def __repr__(self) -> str:
         # astropy.time.Time doesn't have an eval-friendly __repr__, so we
         # simulate our own here to make Timespan's __repr__ eval-friendly.
         # Interestingly, enum.Enum has an eval-friendly __str__, but not an
         # eval-friendly __repr__.
-        tmpl = "astropy.time.Time('{t}', scale='{t.scale}', format='{t.format}')"
-        begin = tmpl.format(t=self.begin) if isinstance(self.begin, astropy.time.Time) else str(self.begin)
-        end = tmpl.format(t=self.end) if isinstance(self.end, astropy.time.Time) else str(self.end)
+        begin = self.__repr_astropy__(self.begin)
+        end = self.__repr_astropy__(self.end)
         return f"Timespan(begin={begin}, end={end})"
 
     def __eq__(self, other: Any) -> bool:

python/lsst/daf/butler/time_utils.py

@@ -33,7 +33,9 @@
 from typing import Any, ClassVar
 
 import astropy.time
+import astropy.time.formats
 import astropy.utils.exceptions
+import numpy
 import yaml
 
 # As of astropy 4.2, the erfa interface is shipped independently and
@@ -48,6 +50,43 @@
 _LOG = logging.getLogger(__name__)
 
 
+class _FastTimeUnixTai(astropy.time.formats.TimeUnixTai):
+    """Special astropy time format that skips some checks of the parameters.
+    This format is only used internally by TimeConverter so we can trust that
+    it passes correct values to astropy Time class.
+    """
+
+    # number of seconds in a day
+    _SEC_PER_DAY: ClassVar[int] = 24 * 3600
+
+    # Name of this format, it is registered in astropy formats registry.
+    name = "unix_tai_fast"
+
+    def _check_val_type(self, val1: Any, val2: Any) -> tuple:
+        # We trust everything that is passed to us.
+        return val1, val2
+
+    def set_jds(self, val1: numpy.ndarray, val2: numpy.ndarray | None) -> None:
+        # Epoch time format is TimeISO with scalar jd1/jd2 arrays, convert them
+        # to floats to speed things up.
+        epoch = self._epoch._time
+        jd1, jd2 = float(epoch._jd1), float(epoch._jd2)
+
+        assert val1.ndim == 0, "Expect scalar"
+        whole_days, seconds = divmod(float(val1), self._SEC_PER_DAY)
+        if val2 is not None:
+            assert val2.ndim == 0, "Expect scalar"
+            seconds += float(val2)
+
+        jd1 += whole_days
+        jd2 += seconds / self._SEC_PER_DAY
+        while jd2 > 0.5:
+            jd2 -= 1.0
+            jd1 += 1.0
+
+        self._jd1, self._jd2 = numpy.array(jd1), numpy.array(jd2)
+
+
 class TimeConverter(metaclass=Singleton):
     """A singleton for mapping TAI times to integer nanoseconds.
 
@@ -112,8 +151,10 @@ def astropy_to_nsec(self, astropy_time: astropy.time.Time) -> int:
         delta = value - self.epoch
         # Special care needed to preserve nanosecond precision.
         # Usually jd1 has no fractional part but just in case.
-        jd1, extra_jd2 = divmod(delta.jd1, 1)
-        value = int(jd1) * self._NSEC_PER_DAY + int(round((delta.jd2 + extra_jd2) * self._NSEC_PER_DAY))
+        # Can use "internal" ._time.jd1 interface because we know that both
+        # are TAI. This is a few percent faster than using .jd1.
+        jd1, extra_jd2 = divmod(delta._time.jd1, 1)
+        value = int(jd1) * self._NSEC_PER_DAY + int(round((delta._time.jd2 + extra_jd2) * self._NSEC_PER_DAY))
         return value
 
     def nsec_to_astropy(self, time_nsec: int) -> astropy.time.Time:
@@ -136,10 +177,14 @@ def nsec_to_astropy(self, time_nsec: int) -> astropy.time.Time:
         that the number falls in the supported range and can produce output
         time that is outside of that range.
         """
-        jd1, jd2 = divmod(time_nsec, self._NSEC_PER_DAY)
-        delta = astropy.time.TimeDelta(float(jd1), float(jd2) / self._NSEC_PER_DAY, format="jd", scale="tai")
-        value = self.epoch + delta
-        return value
+        jd1, jd2 = divmod(time_nsec, 1_000_000_000)
+        time = astropy.time.Time(
+            float(jd1), jd2 / 1_000_000_000, format="unix_tai_fast", scale="tai", precision=6
+        )
+        # Force the format to be something more obvious to external users.
+        # There is a small overhead doing this but it does avoid confusion.
+        time.format = "jd"
+        return time
 
     def times_equal(
         self, time1: astropy.time.Time, time2: astropy.time.Time, precision_nsec: float = 1.0

tests/test_time_utils.py

@@ -105,7 +105,7 @@ def test_round_trip(self):
                     delta2_sec = delta2.to_value("sec")
                     # absolute precision should be better than half
                     # nanosecond, but there are rounding errors too
-                    self.assertLess(abs(delta2_sec), 0.51e-9)
+                    self.assertLess(abs(delta2_sec), 0.511e-9)
 
     def test_times_equal(self):
         """Test for times_equal method"""

tests/test_timespan.py

@@ -129,7 +129,17 @@ def testStrings(self):
         """Test __str__ against expected values and __repr__ with eval
         round-tripping.
         """
-        for ts in self.timespans:
+        timespans = self.timespans
+
+        # Add timespan that includes Julian day high precision version.
+        timespans.append(
+            Timespan(
+                begin=astropy.time.Time(2458850.0, -0.49930555555555556, format="jd", scale="tai"),
+                end=astropy.time.Time(2458850.0, -0.4986111111111111, format="jd", scale="tai"),
+            )
+        )
+
+        for ts in timespans:
             # Uncomment the next line and run this test directly for the most
             # important test: human inspection.
             #    print(str(ts), repr(ts))