From 07a5cf452a5316cf68e3e4c707646e99693fa260 Mon Sep 17 00:00:00 2001 From: Sahil Jhawar Date: Thu, 22 Oct 2026 22:26:37 +0200 Subject: [PATCH] fix solar wind variable handling when combining solar wind from multiple sources --- .../read_solar_wind_from_multiple_models.py | 6 ++- swvo/io/utils.py | 14 ++++++- ...st_read_solar_wind_from_multiple_models.py | 5 ++- tests/io/test_io_utils.py | 39 ++++++++++++++++++- 4 files changed, 60 insertions(+), 4 deletions(-) diff --git a/swvo/io/solar_wind/read_solar_wind_from_multiple_models.py b/swvo/io/solar_wind/read_solar_wind_from_multiple_models.py index fadfe40..9a67776 100644 --- a/swvo/io/solar_wind/read_solar_wind_from_multiple_models.py +++ b/swvo/io/solar_wind/read_solar_wind_from_multiple_models.py @@ -42,6 +42,8 @@ SWModel = DSCOVR | SWACE | SWOMNI | SWMIDL | SWIMAP | SWSWIFTEnsemble | SWENLIL_BKG | SWENLIL_CME ENLIL_MAX_LOOKBACK_DAYS = 5 + +SOLAR_WIND_COLUMNS = ["bavg", "bx_gsm", "by_gsm", "bz_gsm", "speed", "proton_density", "temperature", "pdyn"] logging.captureWarnings(True) @@ -190,7 +192,9 @@ def read_solar_wind_from_multiple_models( else: logger.info(f"{active_model.LABEL} data not available for future dates") - data_out = construct_updated_data_frame(data_out, data_one_model, active_model.LABEL) + data_out = construct_updated_data_frame( + data_out, data_one_model, active_model.LABEL, value_columns=SOLAR_WIND_COLUMNS + ) # A forecast model can stop short of end_time and pin the index there, which would leave # the next forecast model nowhere to write and end the loop before it is even read. diff --git a/swvo/io/utils.py b/swvo/io/utils.py index 64e8e02..ac0aa9e 100644 --- a/swvo/io/utils.py +++ b/swvo/io/utils.py @@ -82,6 +82,7 @@ def construct_updated_data_frame( data: list[pd.DataFrame] | pd.DataFrame, data_one_model: list[pd.DataFrame] | pd.DataFrame, model_label: str, + value_columns: list[str] | None = None, ) -> list[pd.DataFrame]: """ Construct an updated data frame providing the previous data frame and the data frame of the current model call. @@ -96,6 +97,9 @@ def construct_updated_data_frame( The data frame or list of data frames from the current model call. model_label : str The label of the model to add to the data frame. + value_columns : list[str] | None, optional + Columns that decide whether a row is empty. A row is filled from the new model when all of + these columns are NaN, regardless of the other columns. By default, all columns are checked. Returns ------- @@ -135,7 +139,15 @@ def construct_updated_data_frame( data_one_model[i].loc[data_one_model[i]["file_name"].isna(), "model"] = None if data[i].empty: data[i] = data_one_model[i] - empty_idx = data[i].index[data[i].isna().all(axis=1)] + if value_columns is not None: + # use value-columns + meta_cols = [c for c in ("model", "file_name") if c in data[i].columns] + check_cols = [c for c in value_columns if c in data[i].columns] + empty_idx = data[i].index[data[i][check_cols].isna().all(axis=1)] + # model and file_name are taken from the source that fills these rows + data[i].loc[empty_idx, meta_cols] = None + else: + empty_idx = data[i].index[data[i].isna().all(axis=1)] data[i].loc[empty_idx] = ( data[i].loc[empty_idx].combine_first(data_one_model[i].reindex(data[i].index).loc[empty_idx]) diff --git a/tests/io/solar_wind/test_read_solar_wind_from_multiple_models.py b/tests/io/solar_wind/test_read_solar_wind_from_multiple_models.py index 1640f7c..20973bb 100644 --- a/tests/io/solar_wind/test_read_solar_wind_from_multiple_models.py +++ b/tests/io/solar_wind/test_read_solar_wind_from_multiple_models.py @@ -123,7 +123,10 @@ def test_basic_historical_read(self, sample_times, expected_columns): assert isinstance(data, pd.DataFrame) assert all(col in data.columns for col in expected_columns) assert data.loc["2024-11-20 13:23:00+00:00"].model == "dscovr" - assert data.loc["2024-11-22 18:00:00+00:00"].model == "omni" + assert data.loc["2024-11-22 17:59:00+00:00"].model == "omni" + # OMNI has only a sym-h value at 18:00, which must not stop DSCOVR from filling the row + assert data.loc["2024-11-22 18:00:00+00:00"].model == "dscovr" + assert data.loc["2024-11-22 18:00:00+00:00"]["sym-h"] == -26.0 # no ace since dscovr and ace files are same in the test data and dscovr is before ace in the model_order assert not data["file_name"].isna().all() diff --git a/tests/io/test_io_utils.py b/tests/io/test_io_utils.py index 2043dbc..1b54e5a 100644 --- a/tests/io/test_io_utils.py +++ b/tests/io/test_io_utils.py @@ -8,7 +8,7 @@ import pandas as pd import pytest -from swvo.io.utils import any_nans, nan_percentage, start_of_next_month +from swvo.io.utils import any_nans, construct_updated_data_frame, nan_percentage, start_of_next_month def test_any_nans_single_dataframe_with_nans(): @@ -78,3 +78,40 @@ def test_start_of_next_month_keeps_timezone(): assert result == datetime(2025, 1, 1, tzinfo=timezone.utc) assert result.tzinfo == timezone.utc + + +def _omni_tail_and_ace_frames(): + index = pd.date_range("2026-09-30 00:00", periods=4, freq="1h", tz="UTC") + old = pd.DataFrame( + { + "speed": [400.0, None, None, None], + "bz_gsm": [1.0, None, None, None], + "sym-h": [-5.0, -6.0, -7.0, -8.0], + "file_name": ["omni.csv"] * 4, + "model": ["omni", None, None, None], + }, + index=index, + ) + new = pd.DataFrame({"speed": [410.0, 420.0, 430.0, 440.0], "bz_gsm": [2.0, 3.0, 4.0, 5.0]}, index=index) + new["file_name"] = "ace.csv" + return old, new + + +def test_construct_updated_data_frame_value_columns_fills_nan_tail(): + old, new = _omni_tail_and_ace_frames() + + result = construct_updated_data_frame(old, new, "ace", value_columns=["speed", "bz_gsm"])[0] + + assert result["speed"].tolist() == [400.0, 420.0, 430.0, 440.0] + assert result["model"].tolist() == ["omni", "ace", "ace", "ace"] + assert result["file_name"].tolist() == ["omni.csv", "ace.csv", "ace.csv", "ace.csv"] + assert result["sym-h"].tolist() == [-5.0, -6.0, -7.0, -8.0] + + +def test_construct_updated_data_frame_without_value_columns_keeps_rows_with_file_name(): + old, new = _omni_tail_and_ace_frames() + + result = construct_updated_data_frame(old, new, "ace")[0] + + assert result["speed"].isna().tolist() == [False, True, True, True] + assert result["file_name"].tolist() == ["omni.csv"] * 4