Reduce peak ram

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "polars>=1.31.0",
     "pyarrow>=20.0.0",
     "hatanaka>=2.8.1",
+    "psutil>=7.2.2",
 ]
 
 [project.scripts]

src/prx/atmospheric_corrections.py

@@ -2,9 +2,9 @@
 import pandas as pd
 import georinex
 import logging
-
+import polars as pl
 from numpy.typing import NDArray
-
+from pathlib import Path
 from prx.util import deg_2_rad, ecef_2_geodetic, timedelta_2_weeks_and_seconds
 import prx.constants as constants
 
@@ -85,8 +85,10 @@ def compute_l1_iono_delay_klobuchar(
     return iono_correction_l1_m
 
 
-def add_iono_column(
-    flat_obs, rinex_3_ephemerides_files, approximate_receiver_ecef_position_m
+def compute_iono_column(
+    flat_obs: pl.DataFrame,
+    rinex_3_ephemerides_files: list[Path],
+    approximate_receiver_ecef_position_m: NDArray[np.float64],
 ):
     # create a dictionary containing the headers of the different NAV files.
     # The keys are the "YYYYDDD" (year and day of year) and are located at
@@ -95,62 +97,48 @@ def add_iono_column(
         file.name[12:19]: georinex.rinexheader(file)
         for file in rinex_3_ephemerides_files
     }
-
-    idx_all_days = []
-    iono_all_days = []
+    [latitude_user_rad, longitude_user_rad, __] = ecef_2_geodetic(
+        approximate_receiver_ecef_position_m
+    )
+    flat_obs = flat_obs.with_row_index()
     for file in rinex_3_ephemerides_files:
         # get year and doy from NAV filename
         year = int(file.name[12:16])
         doy = int(file.name[16:19])
-
-        # Selection criteria: time of emission belonging to the day of the current NAV file
-        mask = (
+        if "IONOSPHERIC CORR" not in nav_header_dict[f"{year:03d}" + f"{doy:03d}"]:
+            logging.warning(f"Missing iono model parameters for day {doy:03d}")
+            continue
+        # Assign iono model parameters to rows
+        # Assignment criteria: time of emission belonging to the day of the current NAV file
+        matching_rows = ((pl.col("time_of_emission_isagpst") >= pd.Timestamp(year=year, month=1, day=1) + pd.Timedelta(days=doy - 1)) &
+                         (pl.col("time_of_emission_isagpst") < pd.Timestamp(year=year, month=1, day=1) + pd.Timedelta(days=doy)) &
+                         pl.col("observation_type").str.starts_with("C"))
+        day_df = flat_obs.filter(matching_rows).select("index", "time_of_emission_isagpst", "elevation_rad", "azimuth_rad", "carrier_frequency_hz")
+        time_of_emission_weeksecond_isagpst = timedelta_2_weeks_and_seconds(
             (
-                flat_obs.time_of_emission_isagpst
-                >= pd.Timestamp(year=year, month=1, day=1) + pd.Timedelta(days=doy - 1)
-            )
-            & (
-                flat_obs.time_of_emission_isagpst
-                < pd.Timestamp(year=year, month=1, day=1) + pd.Timedelta(days=doy)
+                    day_df.select("time_of_emission_isagpst")
+                    - constants.system_time_scale_rinex_utc_epoch["GPST"]
             )
-            & (flat_obs.observation_type.str.startswith("C"))
-        )
-        mask_idx = mask.loc[mask].index
-        idx_all_days.append(mask_idx)
-        if "IONOSPHERIC CORR" in nav_header_dict[f"{year:03d}" + f"{doy:03d}"]:
-            logging.info(f"Computing iono delay for {year}-{doy:03d}")
-            time_of_emission_weeksecond_isagpst = timedelta_2_weeks_and_seconds(
-                flat_obs.loc[mask_idx, "time_of_emission_isagpst"]
-                - constants.system_time_scale_rinex_utc_epoch["GPST"]
-            )[1].to_numpy()
-            [latitude_user_rad, longitude_user_rad, __] = ecef_2_geodetic(
-                approximate_receiver_ecef_position_m
-            )
-            iono_all_days.append(
-                compute_l1_iono_delay_klobuchar(
-                    time_of_emission_weeksecond_isagpst,
-                    nav_header_dict[f"{year:03d}" + f"{doy:03d}"]["IONOSPHERIC CORR"][
-                        "GPSA"
-                    ],
-                    nav_header_dict[f"{year:03d}" + f"{doy:03d}"]["IONOSPHERIC CORR"][
-                        "GPSB"
-                    ],
-                    flat_obs.loc[mask_idx, "elevation_rad"],
-                    flat_obs.loc[mask_idx, "azimuth_rad"],
-                    latitude_user_rad,
-                    longitude_user_rad,
-                )
-                * (
-                    constants.carrier_frequencies_hz()["G"]["L1"][1] ** 2
-                    / flat_obs.loc[mask_idx, "carrier_frequency_hz"] ** 2
+            .to_series()
+            .cast(dtype=pl.Duration(time_unit="ns"))
+        )[1]
+        delay = compute_l1_iono_delay_klobuchar(
+            time_of_emission_weeksecond_isagpst.to_numpy().reshape((-1, 1)),
+            nav_header_dict[f"{year:03d}" + f"{doy:03d}"]["IONOSPHERIC CORR"]["GPSA"],
+            nav_header_dict[f"{year:03d}" + f"{doy:03d}"]["IONOSPHERIC CORR"]["GPSB"],
+            day_df.select(pl.col("elevation_rad")).to_numpy().reshape((-1, 1)),
+            day_df.select(pl.col("azimuth_rad")).to_numpy().reshape((-1, 1)),
+            latitude_user_rad,
+            longitude_user_rad,
+        ) * (
+                        constants.carrier_frequencies_hz()["G"]["L1"][1] ** 2
+        / day_df.select(pl.col("carrier_frequency_hz")).to_numpy().reshape((-1, 1)) ** 2
                 )
-            )
-        else:
-            logging.warning(f"Missing iono model parameters for day {doy:03d}")
-            iono_all_days.append(np.full(mask_idx.shape, np.nan))
-    delays = np.ones((len(flat_obs.index))) * np.nan
-    delays[np.concatenate(idx_all_days)] = np.concatenate(iono_all_days)
-    return delays
+        day_df = day_df.select(["index"]).with_columns(
+            pl.Series("iono_delay_m", delay.flatten())
+        )
+        flat_obs = flat_obs.join(day_df, on="index", how="left")
+    return flat_obs.drop("index")
 
 
 def compute_tropo_delay_saastamoinen(height, el, lat, humi=0.7):

src/prx/main.py

@@ -27,11 +27,10 @@
 @util.timeit
 def write_prx_file(
     prx_header: dict,
-    prx_records_pd: pd.DataFrame,
+    prx_records: pl.DataFrame,
     file_name_without_extension: Path,
 ):
     output_file = Path(f"{str(file_name_without_extension)}.csv")
-    prx_records = pl.from_pandas(prx_records_pd)
     prx_records = prx_records.with_columns(
         (pl.col("elevation_rad") * cDegPerRad).alias("sat_elevation_deg"),
         (pl.col("azimuth_rad") * cDegPerRad).alias("sat_azimuth_deg"),
@@ -216,7 +215,11 @@ def assign_carrier_frequencies(flat_obs):
     return flat_obs
 
 
+from line_profiler import profile
+
+
 @util.timeit
+@profile
 def build_records_levels_12(
     rinex_3_obs_file,
     rinex_3_ephemerides_files,
@@ -249,7 +252,7 @@ def build_records_levels_12(
     log.info("Computing times of emission in satellite time")
     per_sat = flat_obs.pivot(
         index=["time_of_reception_in_receiver_time", "satellite"],
-        columns=["observation_type"],
+        on=["observation_type"],
         values="observation_value",
     )
     per_sat = per_sat.with_columns(
@@ -341,15 +344,14 @@ def build_records_levels_12(
         day_query["query_time_isagpst"] = day_query["query_time_isagpst"].astype(
             "datetime64[ns]"
         )
-        sat_states_per_day.append(
-            pl.from_pandas(
-                rinex_evaluate.compute_parallel(
-                    file,
-                    day_query,
-                    joblib_backend=joblib_backend,
-                )
+        day_sat_states = pl.from_pandas(
+            rinex_evaluate.compute_parallel(
+                file,
+                day_query,
+                joblib_backend=joblib_backend,
             )
         )
+        sat_states_per_day.append(day_sat_states)
         if prx_level == 1:  # drop sat group delay
             sat_states_per_day[-1] = sat_states_per_day[-1].drop(["sat_code_bias_m"])
     sat_states = pl.concat(sat_states_per_day)
@@ -424,14 +426,13 @@ def build_records_levels_12(
 
     if prx_level == 2:
         # add iono correction
-        iono_delay = atmo.add_iono_column(
-            flat_obs.to_pandas(),
+        flat_obs = atmo.compute_iono_column(
+            flat_obs,
             rinex_3_ephemerides_files,
             approximate_receiver_ecef_position_m,
         )
-        flat_obs = flat_obs.with_columns(iono_delay_m=iono_delay)
 
-    return flat_obs.to_pandas()
+    return flat_obs
 
 
 def build_records_level_3(
@@ -467,7 +468,7 @@ def build_records_level_3(
     log.info("Computing times of emission in satellite time")
     per_sat = flat_obs.pivot(
         index=["time_of_reception_in_receiver_time", "satellite"],
-        columns=["observation_type"],
+        on=["observation_type"],
         values="observation_value",
     ).reset_index()
     per_sat["time_scale"] = (
@@ -618,28 +619,14 @@ def build_records_level_3(
     # set frequency slot to 1 for non-GLONASS satellites
     flat_obs.loc[flat_obs.satellite.str[0] != "R", "frequency_slot"] = int(1)
 
-    def assign_carrier_frequencies(flat_obs):
-        freq_dict = pd.json_normalize(carrier_frequencies_hz(), sep="_").to_dict(
-            orient="records"
-        )[0]
-        assignable = flat_obs.frequency_slot.notna()
-        keys = (
-            flat_obs.satellite[assignable].str[0]
-            + "_L"
-            + flat_obs["observation_type"][assignable].str[1]
-            + "_"
-            + flat_obs.frequency_slot[assignable].astype(int).astype(str)
-        )
-        flat_obs.loc[:, "carrier_frequency_hz"] = keys.map(freq_dict)
-        return flat_obs
 
     flat_obs = assign_carrier_frequencies(flat_obs).drop(columns=["frequency_slot"])
 
     # TODO: change to IONEX when implemented
     rnx3_nav_files = nav_file_discovery.discover_or_download_auxiliary_files(
         rinex_3_obs_file
     )["broadcast_ephemerides"]
-    iono_delay = atmo.add_iono_column(
+    iono_delay = atmo.compute_iono_column(
         flat_obs, rnx3_nav_files, approximate_receiver_ecef_position_m
     )
     flat_obs["iono_delay_m"] = iono_delay
@@ -700,13 +687,13 @@ def process(
             metadata["processing_start_time"] = t0
 
             # build record
-            records = build_records_level_3(
+            records = pl.from_pandas(build_records_level_3(
                 rinex_3_obs_file,
                 aux_files["sp3_orb"],
                 aux_files["atx"],
                 metadata["approximate_receiver_ecef_position_m"],
                 model_tropo,
-            )
+            ))
     metadata["processing_time"] = str(
         pd.Timestamp.now() - metadata["processing_start_time"]
     )