rainfallqc 0.2.1__tar.gz → 0.2.2__tar.gz

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rainfallqc
-Version: 0.2.1
+Version: 0.2.2
 Summary: Quality control for rainfall data
 Author-email: Tom Keel <tomkee@ceh.ac.uk>
 Maintainer-email: Tom Keel <tomkee@ceh.ac.uk>

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/pyproject.toml

@@ -68,7 +68,7 @@ homepage = "https://github.com/thomasjkeel/RainfallQC"
 target-version = ["py310", "py311", "py312", "py313"]
 
 [tool.bumpversion]
-current_version = "0.2.1"
+current_version = "0.2.2"
 commit = true
 message = "Bump version: {current_version} → {new_version}"
 commit_args = ""

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/__init__.py

@@ -2,9 +2,9 @@
 
 __author__ = """Tom Keel"""
 __email__ = "tomkee@ceh.ac.uk"
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
-from rainfallqc import utils
+from rainfallqc import core, utils
 from rainfallqc.checks import (
     comparison_checks,
     gauge_checks,
@@ -19,5 +19,6 @@ __all__ = [
     "gauge_checks",
     "neighbourhood_checks",
     "timeseries_checks",
+    "core",
     "utils",
 ]

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/checks/comparison_checks.py

@@ -11,9 +11,11 @@ import numpy as np
 import polars as pl
 import xarray as xr
 
+from rainfallqc.core.all_qc_checks import qc_check
 from rainfallqc.utils import data_readers, data_utils, neighbourhood_utils, stats
 
 
+@qc_check("check_annual_exceedance_etccdi_r99p", require_non_negative=True)
 def check_annual_exceedance_etccdi_r99p(
     data: pl.DataFrame, target_gauge_col: str, gauge_lat: int | float, gauge_lon: int | float
 ) -> list:
@@ -43,7 +45,9 @@ def check_annual_exceedance_etccdi_r99p(
     etccdi_r99p = data_readers.load_etccdi_data(etccdi_var="R99p")
 
     # 2. Get nearest local R99p value to the gauge coordinates
-    nearby_etccdi_r99p = neighbourhood_utils.get_nearest_etccdi_val_to_gauge(etccdi_r99p, gauge_lat, gauge_lon)
+    nearby_etccdi_r99p = neighbourhood_utils.get_nearest_non_nan_etccdi_val_to_gauge(
+        etccdi_r99p, etccdi_name="R99p", gauge_lat=gauge_lat, gauge_lon=gauge_lon
+    )
 
     # 3. Get sum of rainfall above the 99th percentile per year
     sum_rainfall_above_99percentile_per_year = get_sum_rainfall_above_percentile_per_year(
@@ -58,6 +62,7 @@ def check_annual_exceedance_etccdi_r99p(
     return exceedance_flags
 
 
+@qc_check("check_annual_exceedance_etccdi_prcptot", require_non_negative=True)
 def check_annual_exceedance_etccdi_prcptot(
     data: pl.DataFrame, target_gauge_col: str, gauge_lat: int | float, gauge_lon: int | float
 ) -> list:
@@ -87,7 +92,9 @@ def check_annual_exceedance_etccdi_prcptot(
     etccdi_prcptot = data_readers.load_etccdi_data(etccdi_var="PRCPTOT")
 
     # 2. Get nearest local PRCPTOT value to the gauge coordinates
-    nearby_etccdi_prcptot = neighbourhood_utils.get_nearest_etccdi_val_to_gauge(etccdi_prcptot, gauge_lat, gauge_lon)
+    nearby_etccdi_prcptot = neighbourhood_utils.get_nearest_non_nan_etccdi_val_to_gauge(
+        etccdi_prcptot, etccdi_name="PRCPTOT", gauge_lat=gauge_lat, gauge_lon=gauge_lon
+    )
 
     # 3. Get sum of rainfall above the 99th percentile per year
     sum_rainfall_above_99percentile_per_year = get_sum_rainfall_above_percentile_per_year(
@@ -102,6 +109,7 @@ def check_annual_exceedance_etccdi_prcptot(
     return exceedance_flags
 
 
+@qc_check("check_exceedance_of_rainfall_world_record", require_non_negative=True)
 def check_exceedance_of_rainfall_world_record(data: pl.DataFrame, target_gauge_col: str, time_res: str) -> pl.DataFrame:
     """
     Check exceedance of rainfall world record.
@@ -135,6 +143,7 @@ def check_exceedance_of_rainfall_world_record(data: pl.DataFrame, target_gauge_c
     return data_w_flags.select(["time", "world_record_check"])
 
 
+@qc_check("check_hourly_exceedance_etccdi_rx1day", require_non_negative=True)
 def check_hourly_exceedance_etccdi_rx1day(
     data: pl.DataFrame, target_gauge_col: str, gauge_lat: int | float, gauge_lon: int | float
 ) -> pl.DataFrame:
@@ -173,7 +182,9 @@ def check_hourly_exceedance_etccdi_rx1day(
     etccdi_rx1day = data_readers.load_etccdi_data(etccdi_var="Rx1day")
 
     # 3. Get nearest local Rx1day value to the gauge coordinates
-    nearby_etccdi_rx1day = neighbourhood_utils.get_nearest_etccdi_val_to_gauge(etccdi_rx1day, gauge_lat, gauge_lon)
+    nearby_etccdi_rx1day = neighbourhood_utils.get_nearest_non_nan_etccdi_val_to_gauge(
+        etccdi_rx1day, etccdi_name="Rx1day", gauge_lat=gauge_lat, gauge_lon=gauge_lon
+    )
 
     # 4. Get local maximum ETCCDI value
     max_nearby_etccdi_rx1day = np.max(nearby_etccdi_rx1day["Rx1day"])

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/checks/gauge_checks.py

@@ -10,9 +10,11 @@ Classes and functions ordered by appearance in IntenseQC framework.
 import polars as pl
 import scipy.stats
 
+from rainfallqc.core.all_qc_checks import qc_check
 from rainfallqc.utils import data_utils, stats
 
 
+@qc_check("check_years_where_nth_percentile_is_zero", require_non_negative=True)
 def check_years_where_nth_percentile_is_zero(data: pl.DataFrame, target_gauge_col: str, quantile: float) -> list:
     """
     Return years where the n-th percentiles is zero.
@@ -38,6 +40,7 @@ def check_years_where_nth_percentile_is_zero(data: pl.DataFrame, target_gauge_co
     return nth_perc.filter(pl.col(target_gauge_col) == 0)["time"].dt.year().to_list()
 
 
+@qc_check("check_years_where_annual_mean_k_top_rows_are_zero", require_non_negative=True)
 def check_years_where_annual_mean_k_top_rows_are_zero(data: pl.DataFrame, target_gauge_col: str, k: int) -> list:
     """
     Return year list where the annual mean top-K rows are zero.
@@ -63,6 +66,7 @@ def check_years_where_annual_mean_k_top_rows_are_zero(data: pl.DataFrame, target
     return data_top_k.filter(pl.col(target_gauge_col) == 0)["time"].dt.year().to_list()
 
 
+@qc_check("check_temporal_bias", require_non_negative=True)
 def check_temporal_bias(
     data: pl.DataFrame,
     target_gauge_col: str,
@@ -109,6 +113,7 @@ def check_temporal_bias(
     return int(p_val < p_threshold)
 
 
+@qc_check("check_intermittency", require_non_negative=True)
 def check_intermittency(
     data: pl.DataFrame, target_gauge_col: str, no_data_threshold: int = 2, annual_count_threshold: int = 5
 ) -> list:
@@ -135,20 +140,24 @@ def check_intermittency(
         List of years with intermittency issues.
 
     """
-    # 1. Replace missing values with NaN
+    # 1. Check data has consistent time step
+    data_utils.check_data_has_consistent_time_step(data)
+
+    # 2. Replace missing values with NaN
     data = data_utils.replace_missing_vals_with_nan(data, target_gauge_col)
 
-    # --- Step 2: Mark missing values ---
+    # 3. Mark missing values
     data = data.with_columns(pl.col(target_gauge_col).is_nan().alias("is_missing"))
 
-    # 3. Assign group numbers to consecutive missing values
+    # 4. Assign group numbers to consecutive missing values
     grouped = data.with_columns(
         pl.when(pl.col("is_missing"))
         .then((~pl.col("is_missing")).cum_sum())  # Only number missing stretches
         .otherwise(None)
         .alias("group")
     )
-    # 4. Count size of each missing group
+
+    # 5. Count size of each missing group
     group_counts = (
         grouped.filter(pl.col("is_missing"))
         .group_by("group")
@@ -156,21 +165,21 @@ def check_intermittency(
         .filter(pl.col("count") >= no_data_threshold)
     )
 
-    # 5. Keep only valid groups
+    # 6. Keep only valid groups
     valid_missing_groups = group_counts["group"].to_list()
     valid_missing = grouped.filter(pl.col("group").is_in(valid_missing_groups))
 
-    # 6. Get first and last time for each group
+    # 7. Get first and last time for each group
     first_last = (
         valid_missing.sort("time")
         .group_by("group")
         .agg([pl.first("time").alias("start_time"), pl.last("time").alias("end_time")])
     )
 
-    # 7. Add stable row index to full sorted data
+    # 8. Add stable row index to full sorted data
     df_sorted = data.sort("time").with_row_index("row_idx")
 
-    # 8. Lookup row indices of group bounds
+    # 9. Lookup row indices of group bounds
     first_last_with_idx = (
         first_last.join(df_sorted.select(["time", "row_idx"]), left_on="start_time", right_on="time")
         .rename({"row_idx": "start_idx"})
@@ -178,7 +187,7 @@ def check_intermittency(
         .rename({"row_idx": "end_idx"})
     )
 
-    # 9. Compute prev/next indices and fetch values
+    # 10. Compute prev/next indices and fetch values
     row_lookup = df_sorted.select([pl.col("row_idx"), pl.col(target_gauge_col)])
 
     first_last_with_vals = (
@@ -189,21 +198,22 @@ def check_intermittency(
         .join(row_lookup.rename({target_gauge_col: "next_val"}), left_on="next_idx", right_on="row_idx")
     )
 
-    # 10. Filter to groups bounded by zero
+    # 11. Filter to groups bounded by zero
     bounded_by_zero = first_last_with_vals.filter((pl.col("prev_val") == 0) & (pl.col("next_val") == 0))
 
-    # 11. Extract year from each bounded group's start time
+    # 12. Extract year from each bounded group's start time
     bounded_years = bounded_by_zero.select(pl.col("start_time").dt.year().alias("year"))
 
-    # 12. Count how many bounded groups occur per year
+    # 13. Count how many bounded groups occur per year
     year_counts = bounded_years.group_by("year").len()
 
-    # 13. Get years exceeding threshold
+    # 14. Get years exceeding threshold
     years_w_intermittency = year_counts.filter(pl.col("len") >= annual_count_threshold)["year"].to_list()
 
     return years_w_intermittency
 
 
+@qc_check("check_breakpoints", require_non_negative=True)
 def check_breakpoints(
     data: pl.DataFrame,
     target_gauge_col: str,
@@ -240,6 +250,7 @@ def check_breakpoints(
         return 0
 
 
+@qc_check("check_min_val_change", require_non_negative=True)
 def check_min_val_change(data: pl.DataFrame, target_gauge_col: str, expected_min_val: float) -> list:
     """
     Return years when the minimum recorded value changes.

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/checks/neighbourhood_checks.py

@@ -12,9 +12,11 @@ from typing import Iterable, List
 import numpy as np
 import polars as pl
 
+from rainfallqc.core.all_qc_checks import qc_check
 from rainfallqc.utils import data_readers, data_utils, neighbourhood_utils, stats
 
 
+@qc_check("check_wet_neighbours", require_non_negative=True)
 def check_wet_neighbours(
     neighbour_data: pl.DataFrame,
     target_gauge_col: str,
@@ -129,6 +131,7 @@ def check_wet_neighbours(
         return neighbour_data_w_wet_flags
 
 
+@qc_check("check_dry_neighbours", require_non_negative=True)
 def check_dry_neighbours(
     neighbour_data: pl.DataFrame,
     target_gauge_col: str,
@@ -261,6 +264,7 @@ def check_dry_neighbours(
         return neighbour_data_w_dry_flags
 
 
+@qc_check("check_monthly_neighbours", require_non_negative=True)
 def check_monthly_neighbours(
     monthly_neighbour_data: pl.DataFrame,
     target_gauge_col: str,
@@ -355,6 +359,7 @@ def check_monthly_neighbours(
     return monthly_neighbour_data_w_flags.select(["time", "majority_monthly_flag"])
 
 
+@qc_check("check_timing_offset", require_non_negative=True)
 def check_timing_offset(
     neighbour_data: pl.DataFrame,
     target_gauge_col: str,
@@ -436,6 +441,7 @@ def check_timing_offset(
     return offset_flag
 
 
+@qc_check("check_neighbour_affinity_index", require_non_negative=True)
 def check_neighbour_affinity_index(
     neighbour_data: pl.DataFrame, target_gauge_col: str, neighbouring_gauge_col: str
 ) -> float:
@@ -471,6 +477,7 @@ def check_neighbour_affinity_index(
     return stats.affinity_index(neighbour_data, binary_col="rain_not_minima")
 
 
+@qc_check("check_neighbour_correlation", require_non_negative=True)
 def check_neighbour_correlation(
     neighbour_data: pl.DataFrame, target_gauge_col: str, neighbouring_gauge_col: str
 ) -> float:
@@ -501,6 +508,7 @@ def check_neighbour_correlation(
     return stats.gauge_correlation(neighbour_data, target_col=target_gauge_col, other_col=neighbouring_gauge_col)
 
 
+@qc_check("check_daily_factor", require_non_negative=True)
 def check_daily_factor(
     neighbour_data: pl.DataFrame, target_gauge_col: str, neighbouring_gauge_col: str, averaging_method: str = "mean"
 ) -> float:
@@ -553,6 +561,7 @@ def check_daily_factor(
         raise ValueError(f"{averaging_method} not recognised, please use 'mean' or 'median'")
 
 
+@qc_check("check_monthly_factor", require_non_negative=True)
 def check_monthly_factor(
     neighbour_data: pl.DataFrame, target_gauge_col: str, neighbouring_gauge_col: str
 ) -> pl.DataFrame:

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/checks/timeseries_checks.py

@@ -11,11 +11,13 @@ import numpy as np
 import polars as pl
 import xarray as xr
 
+from rainfallqc.core.all_qc_checks import qc_check
 from rainfallqc.utils import data_readers, data_utils, neighbourhood_utils, spatial_utils, stats
 
 DAILY_DIVIDING_FACTOR = {"15m": 96, "hourly": 24, "daily": 1}
 
 
+@qc_check("check_dry_period_cdd", require_non_negative=True)
 def check_dry_period_cdd(
     data: pl.DataFrame, target_gauge_col: str, time_res: str, gauge_lat: int | float, gauge_lon: int | float
 ) -> pl.DataFrame:
@@ -53,7 +55,9 @@ def check_dry_period_cdd(
     etccdi_cdd_days = compute_dry_spell_days(etccdi_cdd)
 
     # 3. Get nearest local CDD value to the gauge coordinates
-    nearby_etccdi_cdd_days = neighbourhood_utils.get_nearest_etccdi_val_to_gauge(etccdi_cdd_days, gauge_lat, gauge_lon)
+    nearby_etccdi_cdd_days = neighbourhood_utils.get_nearest_non_nan_etccdi_val_to_gauge(
+        etccdi_cdd_days, etccdi_name="CDD", gauge_lat=gauge_lat, gauge_lon=gauge_lon
+    )
 
     # 4. Get local maximum CDD_days value
     max_etccdi_cdd_days = np.max(nearby_etccdi_cdd_days["CDD_days"])
@@ -74,6 +78,7 @@ def check_dry_period_cdd(
     return data_w_dry_spell_flags.select(["time", "dry_spell_flag"])
 
 
+@qc_check("check_daily_accumulations", require_non_negative=True)
 def check_daily_accumulations(
     data: pl.DataFrame,
     target_gauge_col: str,
@@ -141,6 +146,7 @@ def check_daily_accumulations(
     return data.select(["time", "daily_accumulation"])
 
 
+@qc_check("check_monthly_accumulations", require_non_negative=True)
 def check_monthly_accumulations(
     data: pl.DataFrame,
     target_gauge_col: str,
@@ -223,6 +229,7 @@ def check_monthly_accumulations(
     return gauge_data_monthly_accumulations.select(["time", "monthly_accumulation"])
 
 
+@qc_check("check_streaks", require_non_negative=True)
 def check_streaks(
     data: pl.DataFrame,
     target_gauge_col: str,
@@ -681,7 +688,9 @@ def get_local_etccdi_sdii_mean(gauge_lat: int | float, gauge_lon: int | float) -
     # 2. Compute spatial mean
     etccdi_sdii = spatial_utils.compute_spatial_mean_xr(etccdi_sdii, var_name="SDII")
     # 3. Get nearest local CDD value to the gauge coordinates
-    nearby_etccdi_sdii = neighbourhood_utils.get_nearest_etccdi_val_to_gauge(etccdi_sdii, gauge_lat, gauge_lon)
+    nearby_etccdi_sdii = neighbourhood_utils.get_nearest_non_nan_etccdi_val_to_gauge(
+        etccdi_sdii, etccdi_name="SDII", gauge_lat=gauge_lat, gauge_lon=gauge_lon
+    )
     # 4. Get local maximum CDD_days value
     nearby_etccdi_sdii_mean = np.max(nearby_etccdi_sdii["SDII_mean"])
     return nearby_etccdi_sdii_mean

rainfallqc-0.2.2/src/rainfallqc/core/all_qc_checks.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+"""Registry of all the QC checks in RainfallQC."""
+
+import functools
+import itertools
+
+import polars as pl
+
+from rainfallqc.utils import data_utils
+
+QC_CHECKS = {}
+
+
+def qc_check(name: str, require_non_negative: bool = False) -> callable:
+    """
+    Register a QC check and check for non-negative values if required.
+
+    Parameters
+    ----------
+    name :
+        Name of the QC check.
+    require_non_negative :
+        If True, check that the target gauge column has no negative values before running the QC check
+
+    Returns
+    -------
+    callable :
+        Decorator to register the QC check.
+
+    Raises
+    ------
+    ValueError :
+        If require_non_negative is True and the target gauge column contains negative values.
+
+    """
+
+    def decorator(func: callable) -> callable:
+        @functools.wraps(func)
+        def wrapper(df: pl.DataFrame, *args, **kwargs) -> list:
+            # Determine which column the user wants checked
+            columns_to_check = []
+            columns_to_check = get_columns_in_kwargs(
+                kwargs, kwarg_name="target_gauge_col", column_list=columns_to_check, name=name
+            )
+
+            for kwarg_name in ["neighbouring_gauge_col", "neighbouring_gauge_cols"]:
+                if kwarg_name in kwargs:
+                    columns_to_check = get_columns_in_kwargs(
+                        kwargs, kwarg_name=kwarg_name, column_list=columns_to_check, name=name
+                    )
+
+            # flatten column list for neighbouring_gauge_cols
+            columns_to_check = list(
+                itertools.chain.from_iterable(col if isinstance(col, list) else [col] for col in columns_to_check)
+            )
+
+            # Optional non-negative pre-check
+            for col in columns_to_check:
+                if require_non_negative and data_utils.check_for_negative_values(df, col):
+                    raise ValueError(f"{name} failed: column '{col}' contains negative values.")
+
+            # Run the actual QC check
+            return func(df, *args, **kwargs)
+
+        # Register for later use
+        QC_CHECKS[name] = wrapper
+        return wrapper
+
+    return decorator
+
+
+def get_columns_in_kwargs(kwargs: dict, kwarg_name: str, column_list: list, name: str) -> list:
+    """
+    Check that a column exists in the DataFrame.
+
+    Parameters
+    ----------
+    kwargs :
+        Dictionary of keyword arguments.
+    kwarg_name :
+        Name of the kwarg to check.
+    column_list :
+        List to append the column_name to if it exists.
+    name :
+        Name of the QC check (for error messages).
+
+    Raises
+    ------
+    ValueError :
+        If the column does not exist in the DataFrame.
+
+    """
+    col_name = kwargs.get(kwarg_name)  # user defined column name for that kwarg
+
+    if col_name is None:
+        raise ValueError(f"{name} requires '{kwarg_name}' to be provided.")
+    column_list.append(col_name)
+    return column_list

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/qc_frameworks/apply_qc_framework.py

@@ -57,7 +57,6 @@ def run_qc_framework(
             for k, v in combined_kwargs.items()
             if k in accepted_keys or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values())
         }
-
         qc_results[qc_method] = qc_func(data, **filtered_kwargs)
 
     return qc_results

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/utils/data_readers.py

@@ -693,7 +693,7 @@ class GPCCNetworkReader(GaugeNetworkReader):
         )
 
     def load_network_data(
-        self, data_paths: List[str] | np.ndarray[str], target_gauge_col: str, missing_val: int | float = -999
+        self, data_paths: List[str] | np.ndarray[str], target_gauge_col: str, missing_val: int | float = -999.9
     ) -> pl.DataFrame:
         """
         Load GPCC network data based on file paths.

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/utils/data_utils.py

@@ -107,7 +107,10 @@ def check_data_has_consistent_time_step(data: pl.DataFrame) -> None:
     unique_timesteps = get_data_timesteps(data)
     if unique_timesteps.len() != 1:
         timestep_strings = [format_timedelta_duration(td) for td in unique_timesteps]
-        raise ValueError(f"Data has a inconsistent time step. Data has following time steps: {timestep_strings}")
+        raise ValueError(
+            f"""Data has a inconsistent time step. Data has following time steps: {timestep_strings}.
+            One potential fix is to resample the data using: 'data.upsample(\"time\", every=\"[insert_timestep]\")'"""
+        )
 
 
 def check_data_is_monthly(data: pl.DataFrame) -> None:
@@ -175,6 +178,26 @@ def check_data_is_specific_time_res(data: pl.DataFrame, time_res: str | list) ->
         raise ValueError(f"Invalid time step. Expected one of {allowed_res}, but data has time-step(s): {time_step}")
 
 
+def check_for_negative_values(df: pl.DataFrame, target_gauge_col: str) -> bool:
+    """
+    Check if the target column contains any negative values.
+
+    Parameters
+    ----------
+    df :
+        DataFrame to check.
+    target_gauge_col :
+        Column to check for negative values.
+
+    Raises
+    ------
+    ValueError
+        If negative values are found in the target column.
+
+    """
+    return (df[target_gauge_col] < 0).any()
+
+
 def convert_datarray_seconds_to_days(series_seconds: xr.DataArray) -> np.ndarray:
     """
     Convert xarray series from seconds to days. For some reason the CDD data from ETCCDI is in seconds.

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/utils/neighbourhood_utils.py

@@ -8,6 +8,12 @@ import numpy as np
 import polars as pl
 import xarray as xr
 
+from rainfallqc.utils import spatial_utils
+
+STATION_ID_COL = "station_id"
+START_DATETIME_COL = "start_datetime"
+END_DATETIME_COL = "end_datetime"
+
 
 def get_target_neighbour_non_zero_minima(
     data: pl.DataFrame, target_col: str, other_col: str, default_minima: float = 0.1
@@ -121,6 +127,9 @@ def get_ids_of_n_nearest_overlapping_neighbouring_gauges(
     distance_threshold: int | float,
     n_closest: int,
     min_overlap_days: int,
+    station_id_col: str = STATION_ID_COL,
+    start_datetime_col: str = START_DATETIME_COL,
+    end_datetime_col: str = END_DATETIME_COL,
 ) -> set:
     """
     Get gauge IDs of nearest n time-overlapping neighbouring gauges.
@@ -137,6 +146,12 @@ def get_ids_of_n_nearest_overlapping_neighbouring_gauges(
         Number of closest neighbours.
     min_overlap_days :
         Minimum overlap between target and neighbouring gauges
+    station_id_col :
+        Column name for station ID in gauge_network_metadata (default 'station_id')
+    start_datetime_col  :
+        Column name for start datetime in gauge_network_metadata (default 'start_datetime')
+    end_datetime_col  :
+        Column name for end datetime in gauge_network_metadata (default 'end_datetime')
 
     Returns
     -------
@@ -145,11 +160,17 @@ def get_ids_of_n_nearest_overlapping_neighbouring_gauges(
 
     """
     # 1. Compute distances between neighbours and target
-    neighbour_distances_df = compute_km_distances_from_target_id(gauge_network_metadata, target_id=target_id)
+    neighbour_distances_df = compute_km_distances_from_target_id(
+        gauge_network_metadata, target_id=target_id, station_id_col=station_id_col
+    )
 
     # 2. Compute overlapping days between neighbours and target
     neighbour_overlap_days_df = compute_temporal_overlap_days_from_target_id(
-        gauge_network_metadata, target_id=target_id
+        gauge_network_metadata,
+        target_id=target_id,
+        station_id_col=station_id_col,
+        start_datetime_col=start_datetime_col,
+        end_datetime_col=end_datetime_col,
     )
 
     # 3. Subset n_closest based on distance threshold
@@ -163,8 +184,8 @@ def get_ids_of_n_nearest_overlapping_neighbouring_gauges(
     )
 
     # 5. get all ids of neighbouring gauges
-    neighbour_distances_ids = neighbour_distances_df["station_id"].to_list()
-    neighbour_overlap_ids = neighbour_overlap_days_df["station_id"].to_list()
+    neighbour_distances_ids = neighbour_distances_df[station_id_col].to_list()
+    neighbour_overlap_ids = neighbour_overlap_days_df[station_id_col].to_list()
 
     # 6. Select gauge IDs meeting both conditions
     all_neighbour_ids = set(neighbour_distances_ids).intersection(set(neighbour_overlap_ids))
@@ -202,7 +223,13 @@ def compute_temporal_overlap_days(
     return overlap_days
 
 
-def compute_temporal_overlap_days_from_target_id(gauge_network_metadata: pl.DataFrame, target_id: str) -> pl.DataFrame:
+def compute_temporal_overlap_days_from_target_id(
+    gauge_network_metadata: pl.DataFrame,
+    target_id: str,
+    station_id_col: str,
+    start_datetime_col: str,
+    end_datetime_col: str,
+) -> pl.DataFrame:
     """
     Compute overlap in days between target gauges and its neighbours.
 
@@ -214,6 +241,12 @@ def compute_temporal_overlap_days_from_target_id(gauge_network_metadata: pl.Data
         Metadata for gauge network. Each gauge must have 'longitude' and 'latitude'.
     target_id :
         Target gauge to compare against.
+    station_id_col :
+        Column name for station ID in gauge_network_metadata
+    start_datetime_col  :
+        Column name for start datetime in gauge_network_metadata
+    end_datetime_col  :
+        Column name for end datetime in gauge_network_metadata
 
     Returns
     -------
@@ -222,16 +255,16 @@ def compute_temporal_overlap_days_from_target_id(gauge_network_metadata: pl.Data
 
     """
     # 1. Get target station and start and end date
-    target_station = gauge_network_metadata.filter(pl.col("station_id") == target_id)
+    target_station = gauge_network_metadata.filter(pl.col(station_id_col) == target_id)
     start_1, end_1 = (
-        target_station["start_datetime"].item(),
-        target_station["end_datetime"].item(),
+        target_station[start_datetime_col].item(),
+        target_station[end_datetime_col].item(),
     )
 
     # 2. Compute overlap days between target station to other start and end date
     neighbour_overlap_days = {}
     for other_station_id, start_2, end_2 in gauge_network_metadata[
-        ["station_id", "start_datetime", "end_datetime"]
+        [station_id_col, start_datetime_col, end_datetime_col]
     ].rows():
         if target_id == other_station_id:
             continue
@@ -241,7 +274,7 @@ def compute_temporal_overlap_days_from_target_id(gauge_network_metadata: pl.Data
     # 3. Convert to pl.Dataframe
     neighbour_overlap_days_df = pl.DataFrame(
         {
-            "station_id": neighbour_overlap_days.keys(),
+            station_id_col: neighbour_overlap_days.keys(),
             "overlap_days": neighbour_overlap_days.values(),
         }
     )
@@ -272,7 +305,9 @@ def get_neighbours_with_min_overlap_days(
     return neighbour_overlap_days_df.filter(pl.col("overlap_days") >= min_overlap_days)
 
 
-def compute_km_distances_from_target_id(gauge_network_metadata: pl.DataFrame, target_id: str) -> pl.DataFrame:
+def compute_km_distances_from_target_id(
+    gauge_network_metadata: pl.DataFrame, target_id: str, station_id_col: str
+) -> pl.DataFrame:
     """
     Compute kilometre distances between gauges in network and target gauges.
 
@@ -282,6 +317,8 @@ def compute_km_distances_from_target_id(gauge_network_metadata: pl.DataFrame, ta
         Metadata for gauge network. Each gauge must have 'longitude' and 'latitude'.
     target_id :
         Target gauge to compare against.
+    station_id_col :
+        Column name for station ID in gauge_network_metadata
 
     Returns
     -------
@@ -290,7 +327,7 @@ def compute_km_distances_from_target_id(gauge_network_metadata: pl.DataFrame, ta
 
     """
     # 1. Get target station lat and lon
-    target_station = gauge_network_metadata.filter(pl.col("station_id") == target_id)
+    target_station = gauge_network_metadata.filter(pl.col(station_id_col) == target_id)
     target_latlon = (
         target_station["latitude"].item(),
         target_station["longitude"].item(),
@@ -299,7 +336,7 @@ def compute_km_distances_from_target_id(gauge_network_metadata: pl.DataFrame, ta
     # 2. Calculate lat/lon distances from the target gauge
     neighbour_distances = {}
     for other_station_id, other_lat, other_lon in gauge_network_metadata[
-        ["station_id", "latitude", "longitude"]
+        [station_id_col, "latitude", "longitude"]
     ].rows():
         neighbour_distances[other_station_id] = geopy.distance.geodesic(
             target_latlon, (other_lat, other_lon)
@@ -308,7 +345,7 @@ def compute_km_distances_from_target_id(gauge_network_metadata: pl.DataFrame, ta
     # 3. Convert to pl.Dataframe
     neighbour_distances_df = pl.DataFrame(
         {
-            "station_id": neighbour_distances.keys(),
+            station_id_col: neighbour_distances.keys(),
             "distance": neighbour_distances.values(),
         }
     )
@@ -358,31 +395,57 @@ def get_n_closest_neighbours(
         return sorted_close_neighbours.filter(pl.col("distance") <= nth_distance)
 
 
-def get_nearest_etccdi_val_to_gauge(
-    etccdi_data: xr.Dataset, gauge_lat: int | float, gauge_lon: int | float
+def get_nearest_non_nan_etccdi_val_to_gauge(
+    etccdi_data: xr.Dataset,
+    etccdi_name: str,
+    gauge_lat: int | float,
+    gauge_lon: int | float,
+    max_distance_km: int | float = 500,
 ) -> xr.Dataset:
     """
-    Get the value at the nearest ETCCDI grid cell to the gauge coordinates.
+    Get the value at the nearest non-nan ETCCDI grid cell to the gauge coordinates.
 
     Parameters
     ----------
     etccdi_data
         ETCCDI data with given variable to check
+    etccdi_name :
+        ETCCDI variable name to check
     gauge_lat :
         latitude of the rain gauge
     gauge_lon :
         longitude of the rain gauge
+    max_distance_km :
+        Maximum distance in km to search for a non-nan value (default 500 km)
 
     Returns
     -------
     nearby_etccdi_data :
-        ETCCDI data at the nearest grid cell
+        ETCCDI data at the nearest grid cell with non-nan values
 
     """
-    # 1. Get local data
-    nearby_etccdi_data = etccdi_data.sel(
-        lon=gauge_lon,
-        lat=gauge_lat,
-        method="nearest",
+    # 1. Stack lat/lon into a single dimension
+    stacked = etccdi_data.stack(points=("lat", "lon"))
+
+    # 2. Compute haversine distance to each point
+    dists = spatial_utils.haversine(stacked["lon"], stacked["lat"], gauge_lon, gauge_lat)
+
+    # 3. Sort by distance
+    sorted_idx = dists.argsort()
+
+    # Loop through points in ascending distance order
+    for idx in sorted_idx.values:
+        dist_km = dists.isel(points=idx).item()
+        if dist_km <= max_distance_km:
+            if not np.isnan(stacked[etccdi_name].isel(points=idx)).all():
+                sel_lon = stacked["lon"].isel(points=idx).item()
+                sel_lat = stacked["lat"].isel(points=idx).item()
+                return etccdi_data.sel(lon=sel_lon, lat=sel_lat)
+        else:
+            # Because sorted by distance, if we've exceeded max distance, we can break early
+            break
+
+    raise ValueError(
+        f"""No non-NaN point found within {max_distance_km} km of ({gauge_lat}, {gauge_lon}).
+        Assuming EPSG:4326 coordinates."""
     )
-    return nearby_etccdi_data

{rainfallqc-0.2.1 → rainfallqc-0.2.2}/src/rainfallqc/utils/spatial_utils.py

@@ -8,6 +8,8 @@ Classes and functions ordered alphabetically.
 import numpy as np
 import xarray as xr
 
+EARTH_RADIUS_KM = 6371.0  # Radius of the Earth in kilometers
+
 
 def compute_spatial_mean_xr(data: xr.Dataset, var_name: str) -> xr.Dataset:
     """
@@ -38,3 +40,33 @@ def compute_spatial_mean_xr(data: xr.Dataset, var_name: str) -> xr.Dataset:
         np.ma.mean(data_masked, axis=0),
     )  # axis 0 is time
     return data
+
+
+def haversine(lon1: float, lat1: float, lon2: float, lat2: float) -> float:
+    """
+    Great circle distance (km) between two points on Earth.
+
+    Parameters
+    ----------
+    lon1 : float
+        Longitude of point 1
+    lat1 : float
+        Latitude of point 1
+    lon2 : float
+        Longitude of point 2
+    lat2 : float
+        Latitude of point 2
+
+    Returns
+    -------
+    distance : float
+        Distance between the two points in km
+
+    """
+    R = EARTH_RADIUS_KM
+    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, float(lon2), float(lat2)])
+    dlon = lon2 - lon1
+    dlat = lat2 - lat1
+    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
+    c = 2 * np.arcsin(np.sqrt(a))
+    return R * c