diff --git a/pyproject.toml b/pyproject.toml index d5d2921..48683c4 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/prx/atmospheric_corrections.py b/src/prx/atmospheric_corrections.py index 727ffdf..ea59303 100644 --- a/src/prx/atmospheric_corrections.py +++ b/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): diff --git a/src/prx/main.py b/src/prx/main.py index 28e3a3d..6f8b2f3 100644 --- a/src/prx/main.py +++ b/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,20 +619,6 @@ 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"]) @@ -639,7 +626,7 @@ def assign_carrier_frequencies(flat_obs): 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"] ) diff --git a/src/prx/rinex_nav/evaluate.py b/src/prx/rinex_nav/evaluate.py index 11b4be8..a7e4d3f 100644 --- a/src/prx/rinex_nav/evaluate.py +++ b/src/prx/rinex_nav/evaluate.py @@ -1,9 +1,11 @@ import logging from functools import lru_cache +import polars as pl import pandas as pd import numpy as np from pathlib import Path + import scipy from joblib import Parallel, delayed import georinex @@ -527,6 +529,7 @@ def compute_ephemeris_and_clock_offset_reference_times(group): ) df = df.reset_index(drop=True) df = compute_gal_inav_fnav_indicators(df) + df["frequency_slot"] = int(1) if "R" in df.constellation.unique(): df["frequency_slot"] = df.FreqNum.where(df.sv.str[0] == "R", 1).astype(int) df.attrs["ionospheric_corr_GPS"] = nav_ds.ionospheric_corr_GPS @@ -537,21 +540,22 @@ def compute_gal_inav_fnav_indicators(df): """ Based on RINEX 3.05, section A8 """ - df["fnav_or_inav"] = "" is_gal = df.sv.str[0] == "E" - df.loc[is_gal, "fnav_or_inav_indicator"] = np.bitwise_and( + df["fnav_or_inav_int"] = -1 + df.loc[is_gal, "fnav_or_inav_indicator_int"] = np.bitwise_and( df[is_gal].DataSrc.astype(np.uint).to_numpy(), 0b111 ) # We expect only the following navigation message types for Galileo: - indicators = set(df[is_gal].fnav_or_inav_indicator.unique()) + indicators = set(df.loc[is_gal, "fnav_or_inav_indicator_int"].unique()) assert len(indicators.intersection({1, 2, 4, 5})) == len(indicators), ( f"Unexpected Galileo navigation message type: {indicators}" ) - df.loc[is_gal & (df.fnav_or_inav_indicator == 1), "fnav_or_inav"] = "inav" - df.loc[is_gal & (df.fnav_or_inav_indicator == 2), "fnav_or_inav"] = "fnav" - df.loc[is_gal & (df.fnav_or_inav_indicator == 4), "fnav_or_inav"] = "inav" - df.loc[is_gal & (df.fnav_or_inav_indicator == 5), "fnav_or_inav"] = "inav" - return df + df["fnav_or_inav"] = "" + df.loc[is_gal & (df["fnav_or_inav_indicator_int"] == 1), "fnav_or_inav"] = "inav" + df.loc[is_gal & (df["fnav_or_inav_indicator_int"] == 2), "fnav_or_inav"] = "fnav" + df.loc[is_gal & (df["fnav_or_inav_indicator_int"] == 4), "fnav_or_inav"] = "inav" + df.loc[is_gal & (df["fnav_or_inav_indicator_int"] == 5), "fnav_or_inav"] = "inav" + return df.drop(columns=["fnav_or_inav_indicator_int"]) def to_isagpst( @@ -565,40 +569,48 @@ def to_isagpst( @timeit -def select_ephemerides(df, query): - df = df[df.ephemeris_reference_time_isagpst.notna()] - query = query.sort_values(by="query_time_isagpst") - df = df.sort_values(by="ephemeris_reference_time_isagpst") +def select_ephemerides(ephemerides: pl.DataFrame, query: pl.DataFrame): + ephemerides = ephemerides.drop_nulls(subset=["ephemeris_reference_time_isagpst"]) # Add fnav/inav indicator to query for to select the FNAV ephemeris for E5b signals, and INAV for other signals - query["fnav_or_inav"] = "" - query.loc[ - (query.sv.str[0] == "E") & (query.signal.str[1] == "5"), "fnav_or_inav" - ] = "fnav" - query.loc[ - (query.sv.str[0] == "E") & (query.signal.str[1] != "5"), "fnav_or_inav" - ] = "inav" - query = pd.merge_asof( - query, - df, + query = query.with_columns( + pl.when( + (pl.col("sv").str.starts_with("E")) + & (pl.col("signal").str.slice(1, 1) == "5") + ) + .then(pl.lit("fnav")) + .when( + (pl.col("sv").str.starts_with("E")) + & (pl.col("signal").str.slice(1, 1) != "5") + ) + .then(pl.lit("inav")) + .otherwise(pl.lit("")) + .alias("fnav_or_inav") + ) + query = query.sort("query_time_isagpst").join_asof( + ephemerides.sort("ephemeris_reference_time_isagpst"), left_on="query_time_isagpst", right_on="ephemeris_reference_time_isagpst", by=["sv", "fnav_or_inav"], - direction="backward", + strategy="backward", ) # Compute times w.r.t. orbit and clock reference times used by downstream computations - query["query_time_wrt_ephemeris_reference_time_s"] = timedelta_2_seconds( - query["query_time_isagpst"] - query["ephemeris_reference_time_isagpst"] - ) - query["query_time_wrt_clock_reference_time_s"] = timedelta_2_seconds( - query["query_time_isagpst"] - query["clock_reference_time_isagpst"] - ) - query["ephemeris_valid"] = (query["query_time_isagpst"] < query["validity_end"]) & ( - query["query_time_isagpst"] > query["validity_start"] + query = query.with_columns( + query_time_wrt_ephemeris_reference_time_s=timedelta_2_seconds( + query["query_time_isagpst"] - query["ephemeris_reference_time_isagpst"] + ), + query_time_wrt_clock_reference_time_s=timedelta_2_seconds( + query["query_time_isagpst"] - query["clock_reference_time_isagpst"] + ), + ephemeris_valid=( + (query["query_time_isagpst"] < query["validity_end"]) + & (query["query_time_isagpst"] > query["validity_start"]) + ).cast(bool), ) + return query -def extract_health_flag_from_query(query): +def extract_health_flag_from_query(query: pl.DataFrame) -> pl.DataFrame: """ Extracts the health flag for each row of a query from a `query` DataFrame containing ephemeris data. @@ -615,27 +627,31 @@ def extract_health_flag_from_query(query): "C" : "SatH1" """ - query = query.copy() - query["constellation"] = query["sv"].str[0] - - query["health_flag"] = query["health"] - if "C" in query["constellation"].unique(): - query.loc[query.constellation == "C", "health_flag"] = query.loc[ - query.constellation == "C", "SatH1" - ] - - return query["health_flag"] + query = query.with_columns( + pl.when(pl.col("sv").str.starts_with("C")) + .then(pl.col("SatH1")) + .otherwise(pl.col("health")) + .alias("health_flag") + ) + return query def compute_clock_offsets(df): - df["sat_clock_offset_m"] = constants.cGpsSpeedOfLight_mps * ( - df["SVclockBias"] - + df["SVclockDrift"] * df["query_time_wrt_clock_reference_time_s"] - + df["SVclockDriftRate"] * df["query_time_wrt_clock_reference_time_s"] ** 2 - ) - df["sat_clock_drift_mps"] = constants.cGpsSpeedOfLight_mps * ( - df["SVclockDrift"] - + 2 * df["SVclockDriftRate"] * df["query_time_wrt_clock_reference_time_s"] + df = df.with_columns( + sat_clock_offset_m=constants.cGpsSpeedOfLight_mps + * ( + pl.col("SVclockBias") + + pl.col("SVclockDrift") * pl.col("query_time_wrt_clock_reference_time_s") + + pl.col("SVclockDriftRate") + * pl.col("query_time_wrt_clock_reference_time_s") ** 2 + ), + sat_clock_drift_mps=constants.cGpsSpeedOfLight_mps + * ( + pl.col("SVclockDrift") + + 2 + * pl.col("SVclockDriftRate") + * pl.col("query_time_wrt_clock_reference_time_s") + ), ) return df @@ -652,57 +668,64 @@ def compute_parallel( # split dataframe into `n_chunks` smaller dataframes n_chunks = min(len(per_signal_query.index), 4) chunks = np.array_split(per_signal_query, n_chunks) - processed_chunks = parallel( - delayed(compute)( - rinex_nav_file_path, chunk, is_query_corrected_by_sat_clock_offset + if joblib_backend == "sequential": + processed_chunks = [ + compute(rinex_nav_file_path, chunk, is_query_corrected_by_sat_clock_offset) + for chunk in chunks + ] + else: + processed_chunks = parallel( + delayed(compute)( + rinex_nav_file_path, chunk, is_query_corrected_by_sat_clock_offset + ) + for chunk in chunks ) - for chunk in chunks - ) - return pd.concat(processed_chunks) + result = pd.concat(processed_chunks) + result["frequency_slot"] = result["frequency_slot"].astype(float) + return result def compute( - rinex_nav_file_path, per_signal_query, is_query_corrected_by_sat_clock_offset=False + rinex_nav_file_path, + per_signal_query_pd, + is_query_corrected_by_sat_clock_offset=False, ): + per_signal_query = pl.from_pandas(per_signal_query_pd) query_columns = per_signal_query.columns - # per_signal_query is a pd.DataFrame with the following columns - # - time_of_reception_in_receiver_time - # - observation_value - # - signal - # - sv - # - query_time_isagpst rinex_nav_file_path = Path(rinex_nav_file_path) - ephemerides = parse_rinex_nav_file(rinex_nav_file_path) + ephemerides = pl.from_pandas(parse_rinex_nav_file(rinex_nav_file_path)) # Group delays and clock offsets can be signal-specific, so we need to match ephemerides to code signals, # not only to satellites # Example: Galileo transmits E5a clock and group delay parameters in the F/NAV message, but parameters for other # signals in the I/NAV message - per_signal_query = select_ephemerides(ephemerides, per_signal_query) + per_signal_query = select_ephemerides( + ephemerides=ephemerides, query=per_signal_query + ) # compute satellite clock bias if is_query_corrected_by_sat_clock_offset: per_signal_query = compute_clock_offsets(per_signal_query) else: # compute satellite clock offset iteratively - t = per_signal_query.query_time_wrt_clock_reference_time_s + t = per_signal_query["query_time_wrt_clock_reference_time_s"] for _ in range(2): per_signal_query = compute_clock_offsets(per_signal_query) - per_signal_query.query_time_wrt_clock_reference_time_s = ( - t - per_signal_query.sat_clock_offset_m / constants.cGpsSpeedOfLight_mps + per_signal_query = per_signal_query.with_columns( + query_time_wrt_clock_reference_time_s=( + t - pl.col("sat_clock_offset_m") / constants.cGpsSpeedOfLight_mps + ) ) # Apply sat clock correction to the query time for satellite position computation - per_signal_query.query_time_wrt_ephemeris_reference_time_s -= ( - per_signal_query.sat_clock_offset_m / constants.cGpsSpeedOfLight_mps + per_signal_query = per_signal_query.with_columns( + query_time_wrt_ephemeris_reference_time_s=pl.col( + "query_time_wrt_ephemeris_reference_time_s" + ) + - pl.col("sat_clock_offset_m") / constants.cGpsSpeedOfLight_mps ) # Compute orbital states for each (satellite,ephemeris) pair only once: - per_sat_eph_query = ( - per_signal_query.groupby(["sv", "query_time_isagpst", "ephemeris_hash"]) - .first() - .reset_index() - ) - per_sat_eph_query = per_sat_eph_query.drop( - columns=["sat_clock_offset_m", "sat_clock_drift_mps"] - ) + per_sat_eph_query = per_signal_query.unique( + subset=["sv", "query_time_isagpst", "ephemeris_hash"] + ).drop(["sat_clock_offset_m", "sat_clock_drift_mps"]) def evaluate_orbit(sub_df): orbit_type = sub_df["orbit_type"].iloc[0] @@ -719,11 +742,12 @@ def evaluate_orbit(sub_df): sub_df[["x_m", "y_m", "z_m", "dx_mps", "dy_mps", "dz_mps"]] = np.nan return sub_df - per_sat_eph_query = per_sat_eph_query.groupby("orbit_type")[ - per_sat_eph_query.columns - ].apply(evaluate_orbit) - per_sat_eph_query = per_sat_eph_query.reset_index(drop=True) - per_sat_eph_query["health_flag"] = extract_health_flag_from_query(per_sat_eph_query) + per_sat_eph_query = pl.from_pandas( + per_sat_eph_query.to_pandas() + .groupby("orbit_type")[per_sat_eph_query.columns] + .apply(evaluate_orbit) + ) + per_sat_eph_query = extract_health_flag_from_query(per_sat_eph_query) columns_to_keep = [ "sv", "sat_pos_x_m", @@ -737,28 +761,38 @@ def evaluate_orbit(sub_df): "health_flag", "relativistic_clock_effect_m", ] - per_sat_eph_query = per_sat_eph_query[columns_to_keep] + per_sat_eph_query = per_sat_eph_query.select(columns_to_keep) # Merge the computed satellite states into the larger signal-specific query dataframe - per_signal_query = per_signal_query.merge( - per_sat_eph_query, on=["sv", "query_time_isagpst", "ephemeris_hash"] + per_signal_query = per_signal_query.join( + per_sat_eph_query, + on=["sv", "query_time_isagpst", "ephemeris_hash"], + how="inner", ) columns_to_keep = [ "sat_clock_offset_m", "sat_clock_drift_mps", ] + columns_to_keep - per_signal_query = compute_total_group_delays(per_signal_query) + per_signal_query = pl.from_pandas( + compute_total_group_delays(per_signal_query.to_pandas()) + ) if "signal" in per_signal_query.columns: columns_to_keep = ["signal", "sat_code_bias_m"] + columns_to_keep - columns_to_keep.append("frequency_slot") + columns_to_keep += ["frequency_slot"] computed_columns_to_keep = [ col for col in columns_to_keep if col not in query_columns ] - per_signal_query.loc[ - ~per_signal_query.ephemeris_valid, computed_columns_to_keep - ] = np.nan - per_signal_query = per_signal_query[columns_to_keep].reset_index(drop=True) - return per_signal_query + per_signal_query = per_signal_query.with_columns( + [ + pl.when(~pl.col("ephemeris_valid")) + .then(None) + .otherwise(pl.col(col)) + .alias(col) + for col in computed_columns_to_keep + ] + ) + per_signal_query = per_signal_query.select(columns_to_keep) + return per_signal_query.sort(by=["query_time_isagpst", "sv"]).to_pandas() def compute_total_group_delays( diff --git a/src/prx/rinex_nav/test/test_evaluate.py b/src/prx/rinex_nav/test/test_evaluate.py index 9283af4..ca61b0d 100644 --- a/src/prx/rinex_nav/test/test_evaluate.py +++ b/src/prx/rinex_nav/test/test_evaluate.py @@ -15,6 +15,7 @@ import shutil import pytest import itertools +import polars as pl # The following thresholds are the achieved maximum difference between broadcast and # MGEX precise orbit and clock solutions seen in this test. @@ -747,26 +748,34 @@ def test_select_ephemerides(): "TransTime": [9, 9, 100, 999], } ) - ephemerides = set_time_of_validity(ephemerides) - query = pd.DataFrame( - { - "sv": ["E01", "G01", "G01"], - "query_time_isagpst": [ - pd.Timedelta("100s"), - pd.Timedelta("50s"), - pd.Timedelta("90s"), - ], - "signal": ["C5X", "C1C", "C1C"], - } + ephemerides = pl.from_pandas(set_time_of_validity(ephemerides)) + query = pl.from_pandas( + pd.DataFrame( + { + "sv": ["E01", "G01", "G01"], + "query_time_isagpst": [ + pd.Timedelta("100s"), + pd.Timedelta("50s"), + pd.Timedelta("90s"), + ], + "signal": ["C5X", "C1C", "C1C"], + } + ) ) query_with_ephemerides = select_ephemerides(ephemerides, query) - query_with_ephemerides = query_with_ephemerides.sort_values( + query_with_ephemerides = query_with_ephemerides.sort( by=["sv", "query_time_isagpst"] - ).reset_index(drop=True) - assert query_with_ephemerides.query_time_isagpst.equals( - pd.Series([pd.Timedelta("100s"), pd.Timedelta("50s"), pd.Timedelta("90s")]) ) - assert query_with_ephemerides.ephemeris_hash.equals(pd.Series([1, 2, 2])) + assert query_with_ephemerides["query_time_isagpst"].equals( + pl.Series( + [ + pd.Timedelta("100s").value, + pd.Timedelta("50s").value, + pd.Timedelta("90s").value, + ] + ) + ) + assert query_with_ephemerides["ephemeris_hash"].equals(pl.Series([1, 2, 2])) def test_compute_health_flag(input_for_test_2): diff --git a/src/prx/test/benchmark.py b/src/prx/test/benchmark.py index a50b7e0..e8997e7 100644 --- a/src/prx/test/benchmark.py +++ b/src/prx/test/benchmark.py @@ -35,7 +35,9 @@ def run_case(case: dict, ram: bool, warm_parser_cache: bool) -> pd.DataFrame: p = cProfile.Profile() p.enable() # cProfile does not profile subprocesses with joblib's loky backend, use threading + t0 = pd.Timestamp.now() process(observation_file_path=obs_file, joblib_backend="threading") + t_exec = pd.Timestamp.now() - t0 p.disable() # RAM @@ -46,7 +48,7 @@ def run_case(case: dict, ram: bool, warm_parser_cache: bool) -> pd.DataFrame: memray_output.unlink(missing_ok=True) if not warm_parser_cache: disk_cache.clear() - with memray.Tracker(memray_output, follow_fork=True): + with memray.Tracker(memray_output, follow_fork=True, native_traces=True) as _: # Use multithreading here, memray does not track memory allocations in child processes with # joblib's "loky" backend. This likely makes prx slower, but we only care about memory allocation here. process( @@ -55,7 +57,9 @@ def run_case(case: dict, ram: bool, warm_parser_cache: bool) -> pd.DataFrame: reader = memray.FileReader(memray_output) metadata = reader.metadata peak_ram_mb = metadata.peak_memory / 1024 / 1024 - + logger.info( + f"Processed {obs_file.name} in {t_exec}: {case['epochs'] / t_exec.seconds} epochs/s, peak RAM [Mb]: {peak_ram_mb}" + ) # Run time stats_file = Path("benchmark_prx.prof").resolve() p.dump_stats(stats_file) @@ -67,9 +71,6 @@ def run_case(case: dict, ram: bool, warm_parser_cache: bool) -> pd.DataFrame: .sort_values(by="tottime", ascending=False) .reset_index(drop=True) ) - logger.info( - f"Processed {obs_file.name} in {df.iloc[0, :]['tottime']} seconds: {case['epochs'] / df.iloc[0, :]['tottime']} epochs/s, peak RAM [Mb]: {peak_ram_mb}" - ) df = df[["func", "tottime"]] df["function"] = df["func"].apply(lambda x: getattr(x, "co_name", None)) df["file"] = df["func"].apply(lambda x: getattr(x, "co_filename", None)) diff --git a/src/prx/test/benchmark_copying.py b/src/prx/test/benchmark_copying.py new file mode 100644 index 0000000..8e340d0 --- /dev/null +++ b/src/prx/test/benchmark_copying.py @@ -0,0 +1,46 @@ +import time +import tracemalloc + +import numpy as np +import pandas as pd +import polars as pl +from string import ascii_uppercase + + +def wraps_polars(df: pd.DataFrame): + pl_df = pl.from_pandas(df) + pl_df = pl_df.select(pl.all() * 2) + return pl_df.to_pandas( + split_blocks=True, + self_destruct=True, + ) + + +def main(): + array_bytes = 1e9 + n_columns = 10 + np_datatype = np.float64 + n_rows = int((array_bytes / n_columns) / 8) + time.sleep(1) + + pandas_df = pd.DataFrame( + { + column_name: np.random.random(n_rows).astype(np_datatype) + for column_name in ascii_uppercase[:n_columns] + } + ) + time.sleep(1) + + pandas_df = wraps_polars(pandas_df) + time.sleep(1) + + # input(f"Press Enter to exit process ({os.getpid()}) ...") + + +if __name__ == "__main__": + tracemalloc.start() + main() + time.sleep(1) + current, peak = tracemalloc.get_traced_memory() + print(f"tracemalloc current : {current / 10**6} MB") + print(f"tracemalloc peak: {peak / 10**6} MB") diff --git a/src/prx/test/test_main.py b/src/prx/test/test_main.py index 55e5d30..6f1e7b2 100644 --- a/src/prx/test/test_main.py +++ b/src/prx/test/test_main.py @@ -178,7 +178,9 @@ def test_prx_command_line_call(input_for_test_tlse): def test_prx_function_call(input_for_test_tlse): test_file = input_for_test_tlse - main.process(observation_file_path=test_file, prx_level=2) + main.process( + observation_file_path=test_file, prx_level=2, joblib_backend="sequential" + ) expected_prx_file = Path(str(test_file).replace("crx.gz", "csv")) assert expected_prx_file.exists() df = pd.read_csv(expected_prx_file, comment="#") diff --git a/src/prx/test/test_util.py b/src/prx/test/test_util.py index 0acddca..b22d248 100644 --- a/src/prx/test/test_util.py +++ b/src/prx/test/test_util.py @@ -233,7 +233,9 @@ def test_timedelta_2_weeks_and_seconds(): seconds_of_week_expected = [280800, 281400, 317400, 302400, np.nan] np.testing.assert_array_equal(week_computed, week_expected) - np.testing.assert_array_equal(seconds_of_week_computed, seconds_of_week_expected) + np.testing.assert_allclose( + seconds_of_week_computed, seconds_of_week_expected, atol=1e-15, rtol=0 + ) # We also expect the function to work for Series of timestamps week_series_computed, seconds_of_week_series_computed = ( diff --git a/src/prx/util.py b/src/prx/util.py index 06f3b5e..a21a9f0 100644 --- a/src/prx/util.py +++ b/src/prx/util.py @@ -219,14 +219,21 @@ def timestamp_to_mid_day(ts): ) -def timedelta_2_weeks_and_seconds(time_delta: pd.Timedelta | pd.Series): +def timedelta_2_weeks_and_seconds(time_delta: pd.Timedelta | pd.Series | pl.Series): + if time_delta is pd.NaT: + return np.nan, np.nan if isinstance(time_delta, pd.Timedelta): - wn_series, tow_series = timedelta_2_weeks_and_seconds(pd.Series([time_delta])) - return wn_series.iloc[0], tow_series.iloc[0] - in_nanoseconds = time_delta / pd.Timedelta(1, "ns") - weeks = np.floor(in_nanoseconds / constants.cNanoSecondsPerWeek) - week_nanoseconds = in_nanoseconds - weeks * constants.cNanoSecondsPerWeek - return weeks, week_nanoseconds.astype(np.float64) / constants.cNanoSecondsPerSecond + w, s = timedelta_2_weeks_and_seconds( + pl.Series([time_delta.value], dtype=pl.Duration(time_unit="ns")) + ) + return w[0], s[0] + if isinstance(time_delta, pd.Series): + w, s = timedelta_2_weeks_and_seconds(pl.from_pandas(time_delta)) + return w, s + seconds = timedelta_2_seconds(time_delta) + weeks = (seconds / constants.cSecondsPerWeek).floor() + week_seconds = seconds - weeks * constants.cSecondsPerWeek + return weeks, week_seconds def week_and_seconds_2_timedelta(weeks, seconds): @@ -244,8 +251,13 @@ def timedelta_2_seconds( return timedelta_2_seconds(pl.from_pandas(time_delta)).to_pandas() assert isinstance(time_delta, pl.Series) assert time_delta.dtype.time_unit == "ns" + # Don't lose float resolution unnecessarily here + nanoseconds = time_delta.dt.total_nanoseconds() + integer_seconds = (nanoseconds / constants.cNanoSecondsPerSecond).floor().cast(int) return ( - time_delta.dt.total_nanoseconds().cast(float) / constants.cNanoSecondsPerSecond + integer_seconds + + (nanoseconds - integer_seconds * constants.cNanoSecondsPerSecond).cast(float) + / constants.cNanoSecondsPerSecond ) diff --git a/uv.lock b/uv.lock index 88a208b..629b343 100644 --- a/uv.lock +++ b/uv.lock @@ -1226,6 +1226,7 @@ dependencies = [ { name = "pandas" }, { name = "plotly" }, { name = "polars" }, + { name = "psutil" }, { name = "pyarrow" }, { name = "requests" }, { name = "scipy" }, @@ -1260,6 +1261,7 @@ requires-dist = [ { name = "pandas", specifier = ">=2.2.3,<3.0.0" }, { name = "plotly", specifier = ">=5.24.1,<7.0.0" }, { name = "polars", specifier = ">=1.31.0" }, + { name = "psutil", specifier = ">=7.2.2" }, { name = "pyarrow", specifier = ">=20.0.0" }, { name = "requests", specifier = ">=2.32.3,<3.0.0" }, { name = "scipy", specifier = ">=1.14.1,<2.0.0" }, @@ -1278,6 +1280,28 @@ dev = [ { name = "viztracer", specifier = ">=1.1.1" }, ] +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + [[package]] name = "pyarrow" version = "21.0.0"