Vectorize row-wise applies

src/prx/precise_corrections/sp3/evaluate.py

@@ -7,6 +7,7 @@
 from scipy.interpolate import KroghInterpolator
 from prx import constants, util
 from prx.precise_corrections.antex import antex_processing as atx_processing
+from prx.util import timedelta_2_seconds
 
 log = logging.getLogger(__name__)
 
@@ -32,9 +33,7 @@ def cached_load(file_path: Path, file_hash: str):
             if "position" not in col and "velocity" not in col:
                 df.rename(columns={col: col.replace("_x", "")}, inplace=True)
         # Convert timestamps to seconds since GPST epoch
-        df["gpst_s"] = (df["time"] - constants.cGpstUtcEpoch).apply(
-            util.timedelta_2_seconds
-        )
+        df["gpst_s"] = timedelta_2_seconds(df["time"] - constants.cGpstUtcEpoch)
         df.drop(columns=["time"], inplace=True)
         df["sat_clock_offset_m"] = (
             constants.cGpsSpeedOfLight_mps * df["clock"]

src/prx/rinex_nav/evaluate.py

@@ -9,7 +9,7 @@
 import georinex
 from prx import util
 from prx import constants
-from prx.util import timeit, try_repair_with_gfzrnx
+from prx.util import timeit, try_repair_with_gfzrnx, timedelta_2_seconds
 
 log = logging.getLogger(__name__)
 
@@ -40,7 +40,9 @@ def cached_load(rinex_file_path: Path, file_hash: str):
     return cached_load(rinex_file, file_content_hash)
 
 
-def time_scale_integer_second_offset_wrt_gpst(time_scale, utc_gpst_leap_seconds=None):
+def time_scale_integer_second_offset_wrt_gpst(
+    time_scale: str, utc_gpst_leap_seconds: int = None
+):
     if time_scale in ["GPST", "SBAST", "QZSST", "IRNSST", "GST"]:
         return pd.Timedelta(seconds=0)
     if time_scale == "BDT":
@@ -552,22 +554,13 @@ def compute_gal_inav_fnav_indicators(df):
     return df
 
 
-def to_isagpst(time, timescale, gpst_utc_leapseconds):
-    if (isinstance(time, pd.Timedelta) or isinstance(time, pd.Series)) and isinstance(
-        timescale, str
-    ):
-        return time - time_scale_integer_second_offset_wrt_gpst(
-            timescale, gpst_utc_leapseconds
-        )
-    if isinstance(time, pd.Series) and isinstance(timescale, pd.Series):
-        return time - timescale.apply(
-            lambda element: time_scale_integer_second_offset_wrt_gpst(
-                element, gpst_utc_leapseconds
-            )
-        )
-
-    assert False, (
-        f"Unexpected types: time is {type(time)}, timescale is {type(timescale)}"
+def to_isagpst(
+    time: pd.Timedelta | pd.Series,
+    timescale: str,
+    gpst_utc_leapseconds: int | None,
+) -> pd.Timedelta | pd.Series:
+    return time - time_scale_integer_second_offset_wrt_gpst(
+        timescale, gpst_utc_leapseconds
     )
 
 
@@ -593,12 +586,12 @@ def select_ephemerides(df, query):
         direction="backward",
     )
     # Compute times w.r.t. orbit and clock reference times used by downstream computations
-    query["query_time_wrt_ephemeris_reference_time_s"] = (
+    query["query_time_wrt_ephemeris_reference_time_s"] = timedelta_2_seconds(
         query["query_time_isagpst"] - query["ephemeris_reference_time_isagpst"]
-    ).apply(util.timedelta_2_seconds)
-    query["query_time_wrt_clock_reference_time_s"] = (
+    )
+    query["query_time_wrt_clock_reference_time_s"] = timedelta_2_seconds(
         query["query_time_isagpst"] - query["clock_reference_time_isagpst"]
-    ).apply(util.timedelta_2_seconds)
+    )
     query["ephemeris_valid"] = (query["query_time_isagpst"] < query["validity_end"]) & (
         query["query_time_isagpst"] > query["validity_start"]
     )

src/prx/rinex_nav/test/test_evaluate.py

@@ -11,7 +11,7 @@
 from prx.precise_corrections.sp3 import evaluate as sp3_evaluate
 from prx.rinex_nav import evaluate as rinex_nav_evaluate
 from prx import constants, converters
-from prx.util import week_and_seconds_2_timedelta
+from prx.util import week_and_seconds_2_timedelta, disk_cache
 import shutil
 import pytest
 import itertools
@@ -108,6 +108,7 @@ def input_for_test_2023(tmp_path_factory):
 
 def test_parse_nav_file(input_for_test):
     path_to_rnx3_nav_file = converters.anything_to_rinex_3(input_for_test["nav"])
+    disk_cache.clear()
     df = parse_rinex_nav_file(path_to_rnx3_nav_file)
     assert not df.empty
 

src/prx/test/test_helpers.py → src/prx/test/test_util.py

@@ -6,8 +6,11 @@
 from pathlib import Path
 import shutil
 import os
-
+import polars as pl
+import datetime
+from prx.constants import cSecondsPerDay
 from prx.converters import anything_to_rinex_3
+from prx.util import timedelta_2_seconds
 
 log = logging.getLogger(__name__)
 
@@ -344,3 +347,29 @@ def test_rinex_type_based_on_first_line(input_for_test):
     assert util.is_rinex_3_nav_file(nav)
     assert not util.is_rinex_3_obs_file(nav)
     assert not util.is_rinex_3_nav_file(obs)
+
+
+def test_timedelta_2_seconds():
+    expected_timedelta_s = cSecondsPerDay + 1.23456789
+    assert np.isclose(
+        timedelta_2_seconds(pd.Timedelta(days=1, seconds=1.23456789)),
+        expected_timedelta_s,
+        atol=1e-9,
+    )
+    assert np.isclose(
+        timedelta_2_seconds(pd.Series([pd.Timedelta(days=1, seconds=1.23456789)])).iloc[
+            0
+        ],
+        expected_timedelta_s,
+        atol=1e-9,
+    )
+    assert np.isclose(
+        timedelta_2_seconds(
+            pl.Series(
+                values=[datetime.timedelta(days=1, seconds=1.23456789)],
+                dtype=pl.Duration(time_unit="ns"),
+            )
+        )[0],
+        expected_timedelta_s,
+        atol=1e-12,
+    )

src/prx/util.py

@@ -9,7 +9,7 @@
 from pathlib import Path
 import importlib.metadata as md
 import git
-
+import polars as pl
 import georinex
 import joblib
 import numpy as np
@@ -233,21 +233,20 @@ def week_and_seconds_2_timedelta(weeks, seconds):
     return pd.Timedelta(weeks * constants.cSecondsPerWeek + seconds, "seconds")
 
 
-def timedelta_2_seconds(time_delta: pd.Timedelta):
-    if pd.isnull(time_delta):
-        return np.nan
-    assert isinstance(time_delta, pd.Timedelta), (
-        "time_delta must be of type pd.Timedelta"
-    )
-    integer_seconds = np.float64(round(time_delta.total_seconds()))
-    fractional_seconds = (
-        np.float64(
-            timedelta_2_nanoseconds(time_delta)
-            - integer_seconds * constants.cNanoSecondsPerSecond
-        )
-        / constants.cNanoSecondsPerSecond
+def timedelta_2_seconds(
+    time_delta: pd.Timedelta | pd.Series | pl.Series,
+) -> float | pd.Series | pl.Series:
+    if isinstance(time_delta, pd.Timedelta):
+        return timedelta_2_seconds(
+            pl.Series([time_delta.value], dtype=pl.Duration(time_unit="ns"))
+        )[0]
+    if isinstance(time_delta, pd.Series):
+        return timedelta_2_seconds(pl.from_pandas(time_delta)).to_pandas()
+    assert isinstance(time_delta, pl.Series)
+    assert time_delta.dtype.time_unit == "ns"
+    return (
+        time_delta.dt.total_nanoseconds().cast(float) / constants.cNanoSecondsPerSecond
     )
-    return integer_seconds + fractional_seconds
 
 
 def timedelta_2_nanoseconds(time_delta: pd.Timedelta):