Make vectorization consistent across all time conversion functions (#…

…1324)

* Make vectorization consistent across all time conversion functions
Fix bug causing iteration over a 0-d array when numpy.vectorize returns a 0-d array rather than a scalar

* Add test coverage ensuring that met_to_datetime64 returns a scalar when a scalar is input
Modify met_to_datetime64 to return a scalar when a scalar is input

imap_processing/spice/time.py

@@ -23,6 +23,36 @@
 TTJ2000_EPOCH = np.datetime64("2000-01-01T11:58:55.816", "ns")
 
 
+@typing.no_type_check
+def _vectorize(pyfunc: typing.Callable, **vectorize_kwargs) -> typing.Callable:
+    """
+    Convert 0-D arrays from numpy.vectorize to scalars.
+
+    For details on numpy.vectorize, see
+    https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html
+
+    Parameters
+    ----------
+    pyfunc : callable
+        A python function or method.
+    **vectorize_kwargs :
+        Keyword arguments to pass to numpy.vectorize.
+
+    Returns
+    -------
+    out : callable
+        A vectorized function.
+    """
+    vectorized_func = np.vectorize(pyfunc, **vectorize_kwargs)
+
+    def wrapper(*args, **kwargs):  # numpydoc ignore=GL08
+        # Calling the vectorized function with [()] will convert 0-D arrays
+        # to scalars
+        return vectorized_func(*args, **kwargs)[()]
+
+    return wrapper
+
+
 def met_to_sclkticks(met: npt.ArrayLike) -> npt.NDArray[float]:
     """
     Convert Mission Elapsed Time (MET) to floating point spacecraft clock ticks.
@@ -93,8 +123,8 @@ def ttj2000ns_to_et(tt_ns: npt.ArrayLike) -> npt.NDArray[float]:
         Number of seconds since the J2000 epoch in the TDB timescale.
     """
     tt_seconds = np.asarray(tt_ns, dtype=np.float64) / 1e9
-    vectorized_unitim = np.vectorize(
-        spiceypy.unitim, [float], excluded=["insys", "outsys"]
+    vectorized_unitim = _vectorize(
+        spiceypy.unitim, otypes=[float], excluded=["insys", "outsys"]
     )
     return vectorized_unitim(tt_seconds, "TT", "ET")
 
@@ -141,9 +171,7 @@ def met_to_datetime64(
     numpy.ndarray[str]
         The mission elapsed time converted to UTC string.
     """
-    if isinstance(met, typing.Iterable):
-        return np.asarray([np.datetime64(utc) for utc in met_to_utc(met)])
-    return np.datetime64(met_to_utc(met))
+    return np.array(met_to_utc(met), dtype=np.datetime64)[()]
 
 
 @typing.no_type_check
@@ -168,16 +196,14 @@ def _sct2e_wrapper(
     ephemeris_time: np.ndarray
         Ephemeris time, seconds past J2000.
     """
-    if isinstance(sclk_ticks, Collection):
-        return np.array([spiceypy.sct2e(IMAP_SC_ID, s) for s in sclk_ticks])
-    else:
-        return spiceypy.sct2e(IMAP_SC_ID, sclk_ticks)
+    vectorized_sct2e = _vectorize(spiceypy.sct2e, otypes=[float], excluded=[0])
+    return vectorized_sct2e(IMAP_SC_ID, sclk_ticks)
 
 
 @typing.no_type_check
 @ensure_spice
 def sct_to_ttj2000s(
-    sclk_ticks: Union[float, Collection[float]],
+    sclk_ticks: Union[float, Iterable[float]],
 ) -> Union[float, np.ndarray]:
     """
     Convert encoded spacecraft clock "ticks" to terrestrial time (TT).
@@ -190,24 +216,21 @@ def sct_to_ttj2000s(
 
     Parameters
     ----------
-    sclk_ticks : Union[float, Collection[float]]
+    sclk_ticks : Union[float, Iterable[float]]
         Input sclk ticks value(s) to be converted to ephemeris time.
 
     Returns
     -------
     terrestrial_time: np.ndarray[float]
         Terrestrial time, seconds past J2000.
     """
-    if isinstance(sclk_ticks, Collection):
-        return np.array(
-            [
-                spiceypy.unitim(spiceypy.sct2e(IMAP_SC_ID, s), "ET", "TT")
-                for s in sclk_ticks
-            ]
-        )
-    else:
+
+    def conversion(sclk_ticks):  # numpydoc ignore=GL08
         return spiceypy.unitim(spiceypy.sct2e(IMAP_SC_ID, sclk_ticks), "ET", "TT")
 
+    vectorized_func = _vectorize(conversion, otypes=[float])
+    return vectorized_func(sclk_ticks)
+
 
 @typing.no_type_check
 @ensure_spice
@@ -231,10 +254,8 @@ def str_to_et(
     ephemeris_time: np.ndarray
         Ephemeris time, seconds past J2000.
     """
-    if isinstance(time_str, str):
-        return spiceypy.str2et(time_str)
-    else:
-        return np.array([spiceypy.str2et(t) for t in time_str])
+    vectorized_str2et = _vectorize(spiceypy.str2et, otypes=[float])
+    return vectorized_str2et(time_str)
 
 
 @typing.no_type_check

imap_processing/tests/spice/test_time.py

@@ -55,6 +55,9 @@ def test_ttj2000ns_to_et(furnish_time_kernels):
     epoch = int(spiceypy.unitim(et, "ET", "TT") * 1e9)
     j2000s = ttj2000ns_to_et(epoch)
     assert j2000s == et
+    # Test for bug when spiceypy tries to iterate over 0-d array returned by
+    # np.vectorize for the scalar case
+    assert not spiceypy.support_types.is_iterable(et)
     # Test array input
     ets = np.arange(et, et + 10000, 100)
     epoch = np.array([spiceypy.unitim(et, "ET", "TT") * 1e9 for et in ets]).astype(
@@ -109,11 +112,17 @@ def test_met_to_datetime64(furnish_time_kernels, utc):
         et_arr = spiceypy.str2et(utc)
         sclk_ticks = np.array([spiceypy.sce2c(IMAP_SC_ID, et) for et in et_arr])
     else:
-        expected_dt64 = np.asarray(np.datetime64(utc))
+        expected_dt64 = np.datetime64(utc)
         et = spiceypy.str2et(utc)
         sclk_ticks = spiceypy.sce2c(IMAP_SC_ID, et)
     met = sclk_ticks * TICK_DURATION
     dt64 = met_to_datetime64(met)
+
+    if isinstance(utc, list):
+        assert isinstance(dt64, np.ndarray)
+        assert dt64.dtype.name == "datetime64[ns]"
+    else:
+        assert isinstance(dt64, np.datetime64)
     np.testing.assert_array_equal(
         dt64.astype("datetime64[us]"), expected_dt64.astype("datetime64[us]")
     )
@@ -146,6 +155,7 @@ def test_str_to_et(furnish_time_kernels):
     expected_et = 553333629.1837274
     actual_et = str_to_et(utc)
     assert expected_et == actual_et
+    assert isinstance(actual_et, float)
 
     # Test list input
     list_of_utc = [