diff --git a/README.md b/README.md index da49708..3d59aba 100755 --- a/README.md +++ b/README.md @@ -3,6 +3,7 @@ SPDX-FileCopyrightText: 2026 GFZ Helmholtz Centre for Geosciences SPDX-FileContributor: Matyas Szabo-Roberts SPDX-FileContributor: Ruggero Vasile SPDX-FileContributor: Sahil Jhawar +SPDX-FileContributor: Simon Mischel SPDX-License-Identifier: Apache-2.0 --> @@ -96,6 +97,21 @@ This package provides tools to read, process, and analyze several key solar and Each index can be accessed via these dedicated reader classes, which handle downloading and read methods. See the code in `swvo/io` or API documentation for details on each index's implementation. +## Expanded OMNI data access + +`OMNIHighRes` and `OMNILowRes` preserve their established default columns and +also support `variables="all"` or an explicit variable-name selection. The +high-resolution reader exposes all 42 one-minute and 45 five-minute OMNIWeb +fields; the hourly reader exposes all 54 non-time OMNI2 fields and accepts both +historic 55-word and current 57-word source records. + +Variable metadata can be inspected without making a network request through a +reader instance's `available_variables()` method or the shared cadence-aware +`swvo.io.omni.variables.available_variables()` utility. See the +[complete OMNI variable guide](https://swvo.readthedocs.io/en/latest/omni_variables.html) +for examples, variable tables, cadence restrictions, cache-upgrade behavior, +compatibility guarantees, and NASA source acknowledgements. + ## Installation To install the package, run the following command: diff --git a/docs/index.rst b/docs/index.rst index bda520f..3e34246 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,5 +1,6 @@ .. SPDX-FileCopyrightText: 2026 GFZ Helmholtz Centre for Geosciences .. SPDX-FileContributor: Sahil Jhawar +.. SPDX-FileContributor: Simon Mischel .. .. SPDX-License-Identifier: Apache-2.0 @@ -27,5 +28,6 @@ User Guide :maxdepth: 1 Examples + Complete OMNI variables Changelog Contributing Guide diff --git a/docs/omni_variables.rst b/docs/omni_variables.rst new file mode 100644 index 0000000..e0a908f --- /dev/null +++ b/docs/omni_variables.rst @@ -0,0 +1,243 @@ +.. SPDX-FileCopyrightText: 2026 GFZ Helmholtz Centre for Geosciences +.. SPDX-FileContributor: Simon Mischel +.. +.. SPDX-License-Identifier: Apache-2.0 + +Complete OMNI variable access +============================= + +SWVO provides two general-purpose OMNI readers: + +* :class:`swvo.io.omni.OMNIHighRes` retrieves NASA OMNIWeb data at one- or + five-minute cadence. +* :class:`swvo.io.omni.OMNILowRes` retrieves the fixed-width hourly OMNI2 + files published by NASA SPDF. + +Both readers retain their original default columns. Complete data is opt-in +through the ``variables`` argument, so existing applications and the focused +``SWOMNI``, ``SymhOMNI``, ``KpOMNI``, ``DSTOMNI``, and ``F107OMNI`` readers +continue to receive their established schemas. + +Selecting variables +------------------- + +``None`` returns the legacy columns, ``"all"`` returns every variable valid +for the selected product and cadence, and a string or iterable selects a +subset. Subset order is preserved, duplicate names are removed, and aliases +are normalized to the canonical output name. + +.. code-block:: python + + from datetime import datetime, timezone + from pathlib import Path + + from swvo.io.omni import OMNIHighRes, OMNILowRes + + start = datetime(2024, 5, 1, tzinfo=timezone.utc) + end = datetime(2024, 5, 2, tzinfo=timezone.utc) + + high_res = OMNIHighRes(data_dir=Path("./omni/high_res")) + complete_1min = high_res.read(start, end, download=True, variables="all") + selected_5min = high_res.read( + start, + end, + cadence_min=5, + download=True, + variables=["speed", "bz_gsm", "proton_flux_10_mev"], + ) + + low_res = OMNILowRes(data_dir=Path("./omni/low_res")) + hourly = low_res.read(start, end, download=True, variables=["kp", "dst", "ae"]) + +The complete machine-readable registries are available without downloading +data. Reader instances expose cadence-appropriate discovery methods, and the +shared utility provides the same metadata without constructing a reader: + +.. code-block:: python + + from swvo.io.omni.variables import available_variables + + one_minute_metadata = high_res.available_variables(cadence=1) + five_minute_metadata = high_res.available_variables(cadence=5) + hourly_metadata = low_res.available_variables() + + # Standalone discovery uses None for hourly data and 1 or 5 for OMNIWeb. + one_minute_metadata = available_variables(cadence=1) + five_minute_metadata = available_variables(cadence=5) + hourly_metadata = available_variables(cadence=None) + +High-resolution variables +------------------------- + +The one-minute product contains the first 42 rows below. The three integral +proton fluxes are only available in the five-minute product. Fill values and +request identifiers follow the `NASA high-resolution OMNI interface +`_. + +.. csv-table:: High-resolution OMNI registry + :header: "Canonical name", "NASA ID", "Cadence (min)", "Unit", "Fill value", "Description", "Aliases" + :widths: 22, 7, 10, 13, 11, 30, 14 + + "imf_spacecraft_id", "4", "1, 5", "", "99", "IMF source spacecraft identifier", "" + "plasma_spacecraft_id", "5", "1, 5", "", "99", "Solar-wind plasma source spacecraft identifier", "" + "imf_point_count", "6", "1, 5", "count", "999", "Fine-scale points in the IMF average", "" + "plasma_point_count", "7", "1, 5", "count", "999", "Fine-scale points in the plasma average", "" + "percent_interpolated", "8", "1, 5", "%", "999", "Percentage of interpolated values", "" + "timeshift", "9", "1, 5", "s", "999999", "Bow-shock-nose time shift", "timeshift_seconds" + "timeshift_rms", "10", "1, 5", "s", "999999", "RMS time shift", "timeshift_rms_seconds" + "minimum_variance_rms", "11", "1, 5", "", "99.99", "RMS minimum-variance vector", "" + "observation_gap", "12", "1, 5", "s", "999999", "Time between observations", "observation_gap_seconds" + "bavg", "13", "1, 5", "nT", "9999.99", "Average IMF magnitude", "" + "bx_gsm", "14", "1, 5", "nT", "9999.99", "IMF Bx in GSE/GSM coordinates", "bx_gse" + "by_gse", "15", "1, 5", "nT", "9999.99", "IMF By in GSE coordinates", "" + "bz_gse", "16", "1, 5", "nT", "9999.99", "IMF Bz in GSE coordinates", "" + "by_gsm", "17", "1, 5", "nT", "9999.99", "IMF By in GSM coordinates", "" + "bz_gsm", "18", "1, 5", "nT", "9999.99", "IMF Bz in GSM coordinates", "" + "sigma_bavg", "19", "1, 5", "nT", "9999.99", "RMS standard deviation of IMF magnitude", "" + "sigma_b", "20", "1, 5", "nT", "9999.99", "RMS standard deviation of the IMF vector", "" + "speed", "21", "1, 5", "km/s", "99999.9", "Solar-wind bulk speed", "" + "vx_gse", "22", "1, 5", "km/s", "99999.9", "Solar-wind Vx in GSE coordinates", "" + "vy_gse", "23", "1, 5", "km/s", "99999.9", "Solar-wind Vy in GSE coordinates", "" + "vz_gse", "24", "1, 5", "km/s", "99999.9", "Solar-wind Vz in GSE coordinates", "" + "proton_density", "25", "1, 5", "1/cm^3", "999.99", "Proton number density", "" + "temperature", "26", "1, 5", "K", "9999999", "Proton temperature", "" + "pdyn", "27", "1, 5", "nPa", "99.99", "Solar-wind flow pressure", "flow_pressure" + "electric_field", "28", "1, 5", "mV/m", "999.99", "Convective electric field Ey", "ey" + "plasma_beta", "29", "1, 5", "", "999.99", "Plasma beta", "" + "alfven_mach_number", "30", "1, 5", "", "999.9", "Alfven Mach number", "" + "spacecraft_x_gse", "31", "1, 5", "Re", "9999.99", "Spacecraft X position in GSE coordinates", "" + "spacecraft_y_gse", "32", "1, 5", "Re", "9999.99", "Spacecraft Y position in GSE coordinates", "" + "spacecraft_z_gse", "33", "1, 5", "Re", "9999.99", "Spacecraft Z position in GSE coordinates", "" + "bsn_x_gse", "34", "1, 5", "Re", "9999.99", "Bow-shock-nose X position in GSE coordinates", "" + "bsn_y_gse", "35", "1, 5", "Re", "9999.99", "Bow-shock-nose Y position in GSE coordinates", "" + "bsn_z_gse", "36", "1, 5", "Re", "9999.99", "Bow-shock-nose Z position in GSE coordinates", "" + "ae", "37", "1, 5", "nT", "99999", "Auroral electrojet AE index", "" + "al", "38", "1, 5", "nT", "99999", "Auroral electrojet AL index", "" + "au", "39", "1, 5", "nT", "99999", "Auroral electrojet AU index", "" + "sym-d", "40", "1, 5", "nT", "99999", "SYM/D geomagnetic index", "sym_d" + "sym-h", "41", "1, 5", "nT", "99999", "SYM/H geomagnetic index", "sym_h" + "asy-d", "42", "1, 5", "nT", "99999", "ASY/D geomagnetic index", "asy_d" + "asy-h", "43", "1, 5", "nT", "99999", "ASY/H geomagnetic index", "asy_h" + "pcn", "44", "1, 5", "", "999.99", "Northern polar-cap index", "pc" + "magnetosonic_mach_number", "45", "1, 5", "", "99.9", "Magnetosonic Mach number", "" + "proton_flux_10_mev", "46", "5", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 10 MeV", "" + "proton_flux_30_mev", "47", "5", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 30 MeV", "" + "proton_flux_60_mev", "48", "5", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 60 MeV", "" + +Hourly variables +---------------- + +The timestamp index replaces the raw year, day-of-year, and hour words. The +remaining 54 fields follow NASA's current 57-word `hourly OMNI2 record +description `_. Historic +55-word records are also accepted; ``lyman_alpha`` and +``proton_quasi_invariant`` are returned as ``NaN`` for those records. + +.. csv-table:: Hourly OMNI2 registry + :header: "Canonical name", "Unit", "Fill value", "Description", "Aliases" + :widths: 25, 16, 12, 37, 14 + + "bartels_rotation_number", "", "9999", "Bartels rotation number", "" + "imf_spacecraft_id", "", "99", "IMF source spacecraft identifier", "" + "plasma_spacecraft_id", "", "99", "Solar-wind plasma source spacecraft identifier", "" + "imf_point_count", "count", "999", "Fine-scale points in the IMF average", "" + "plasma_point_count", "count", "999", "Fine-scale points in the plasma average", "" + "bavg", "nT", "999.9", "Average IMF magnitude", "" + "magnitude_average_field_vector", "nT", "999.9", "Magnitude of the average IMF vector", "" + "latitude_average_field", "degree", "999.9", "Latitude angle of the average IMF vector", "" + "longitude_average_field", "degree", "999.9", "Longitude angle of the average IMF vector", "" + "bx_gse_gsm", "nT", "999.9", "IMF Bx in GSE/GSM coordinates", "bx_gse" + "by_gse", "nT", "999.9", "IMF By in GSE coordinates", "" + "bz_gse", "nT", "999.9", "IMF Bz in GSE coordinates", "" + "by_gsm", "nT", "999.9", "IMF By in GSM coordinates", "" + "bz_gsm", "nT", "999.9", "IMF Bz in GSM coordinates", "" + "sigma_bavg", "nT", "999.9", "RMS standard deviation of IMF magnitude", "" + "sigma_b", "nT", "999.9", "RMS standard deviation of the IMF vector", "" + "sigma_bx", "nT", "999.9", "RMS standard deviation of IMF Bx", "" + "sigma_by", "nT", "999.9", "RMS standard deviation of IMF By", "" + "sigma_bz", "nT", "999.9", "RMS standard deviation of IMF Bz", "" + "temperature", "K", "9999999", "Proton temperature", "" + "proton_density", "1/cm^3", "999.9", "Proton number density", "" + "speed", "km/s", "9999", "Solar-wind bulk speed", "" + "speed_longitude", "degree", "999.9", "Solar-wind flow longitude in GSE coordinates", "" + "speed_latitude", "degree", "999.9", "Solar-wind flow latitude in GSE coordinates", "" + "alpha_proton_ratio", "", "9.999", "Alpha-to-proton density ratio", "" + "pdyn", "nPa", "99.99", "Solar-wind flow pressure", "flow_pressure" + "sigma_temperature", "K", "9999999", "RMS standard deviation of proton temperature", "" + "sigma_proton_density", "1/cm^3", "999.9", "RMS standard deviation of proton density", "" + "sigma_speed", "km/s", "9999", "RMS standard deviation of solar-wind speed", "" + "sigma_speed_longitude", "degree", "999.9", "RMS standard deviation of flow longitude", "" + "sigma_speed_latitude", "degree", "999.9", "RMS standard deviation of flow latitude", "" + "sigma_alpha_proton_ratio", "", "9.999", "RMS standard deviation of alpha-to-proton ratio", "" + "electric_field", "mV/m", "999.99", "Convective electric field", "e" + "plasma_beta", "", "999.99", "Plasma beta", "" + "alfven_mach_number", "", "999.9", "Alfven Mach number", "" + "kp", "", "99", "Three-hour Kp index", "" + "sunspot_number", "", "999", "Daily international sunspot number", "" + "dst", "nT", "99999", "Dst geomagnetic index", "" + "ae", "nT", "9999", "Auroral electrojet AE index", "" + "proton_flux_1_mev", "1/(cm^2 s sr)", "999999.99", "Integral proton flux above 1 MeV", "" + "proton_flux_2_mev", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 2 MeV", "" + "proton_flux_4_mev", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 4 MeV", "" + "proton_flux_10_mev", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 10 MeV", "" + "proton_flux_30_mev", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 30 MeV", "" + "proton_flux_60_mev", "1/(cm^2 s sr)", "99999.99", "Integral proton flux above 60 MeV", "" + "magnetospheric_flux_flag", "", "", "Magnetospheric proton-flux contamination flag", "" + "ap", "nT", "999", "Three-hour ap index", "" + "f107", "10^-22 W/(m^2 Hz)", "999.9", "Daily F10.7 solar radio flux adjusted to 1 AU", "" + "pcn", "", "999.9", "Northern polar-cap index", "pc" + "al", "nT", "99999", "Auroral electrojet AL index", "" + "au", "nT", "99999", "Auroral electrojet AU index", "" + "magnetosonic_mach_number", "", "99.9", "Magnetosonic Mach number", "" + "lyman_alpha", "W/m^2", "0.999999", "Daily solar Lyman-alpha irradiance", "" + "proton_quasi_invariant", "", "9.9999", "Solar-wind proton quasi-invariant", "" + +Cache behavior and compatibility +-------------------------------- + +Processed filenames and directory layouts are unchanged. A newly downloaded +period contains the complete schema for its product and cadence. A legacy +cache containing only the former default columns remains valid for default +reads. + +When a selected variable is absent from an existing cache: + +* ``download=True`` downloads and atomically replaces only the incomplete + period with a complete processed file. +* ``download=False`` raises ``ValueError`` listing the missing columns and the + command needed to upgrade the cache. + +Temporary output is written beside the destination with a ``.tmp`` suffix and +is moved into place only after successful parsing. The hourly downloader uses +an operating-system temporary directory, so the workflow is safe on Linux and +macOS and does not depend on the current working directory. + +Missing values and provenance +----------------------------- + +NASA fill values are converted to ``NaN`` while processing. Kp retains SWVO's +existing conversion from OMNI's integer tenths to conventional thirds. The +``file_name`` column is not a NASA variable; it identifies the processed source +file and is null when every selected variable is missing for a row. + +Errors are raised before network access for unknown names, empty selections, +non-string selections, invalid or equal time ranges, and one-minute requests +for the five-minute-only proton fluxes. Source records with missing headers, +unexpected word counts, nonnumeric fields, or invalid day/hour/minute values +are rejected instead of being silently shifted or partially interpreted. + +An unreadable processed cache raises a path-specific ``ValueError`` when +``download=False``. With ``download=True`` SWVO attempts one atomic replacement; +if the replacement cannot be read, the error states that the cache remains +unreadable. OMNIWeb corrected-range responses are followed only when the +suggested end date shortens the request, preventing retry loops. + +Data source and acknowledgement +------------------------------- + +High-resolution records are retrieved from the `NASA GSFC/SPDF OMNIWeb +service `_. Hourly records are +retrieved from the `NASA SPDF low-resolution OMNI archive +`_. Publications using +these data should follow the acknowledgement and citation guidance in NASA's +OMNI documentation. diff --git a/swvo/io/dst/omni.py b/swvo/io/dst/omni.py index 3391ef5..0befc3d 100644 --- a/swvo/io/dst/omni.py +++ b/swvo/io/dst/omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -8,30 +9,25 @@ from __future__ import annotations -import logging -import warnings -from datetime import datetime, timedelta, timezone -from pathlib import Path +from datetime import datetime, timedelta -import numpy as np import pandas as pd from swvo.io.omni import OMNILowRes -from swvo.io.utils import enforce_utc_timezone - -logger = logging.getLogger(__name__) - -logging.captureWarnings(True) class DSTOMNI(OMNILowRes): """ - Class for reading F10.7 data from OMNI DST files. + Class for reading Dst data from OMNI hourly files. Inherits the `download_and_process`, other private methods and attributes from OMNILowRes. """ + _READ_TIME_PADDING = timedelta(hours=0.9999) + # data is downloaded along with OMNI data, check file name in parent class - def read(self, start_time: datetime, end_time: datetime, download: bool = False) -> pd.DataFrame: + def read( # ty: ignore[invalid-method-override] + self, start_time: datetime, end_time: datetime, download: bool = False + ) -> pd.DataFrame: """ Read OMNI DST data for the given time range. @@ -49,66 +45,7 @@ def read(self, start_time: datetime, end_time: datetime, download: bool = False) :class:`pandas.DataFrame` OMNI DST data. """ - if start_time > end_time: - msg = "start_time must be before end_time" - logger.error(msg) - raise ValueError(msg) - - start_time = enforce_utc_timezone(start_time) - end_time = enforce_utc_timezone(end_time) - - file_paths, _ = self._get_processed_file_list(start_time, end_time) - - t = pd.date_range( - datetime(start_time.year, start_time.month, start_time.day), - datetime(end_time.year, end_time.month, end_time.day, 23, 00, 00), - freq=timedelta(hours=1), - tz=timezone.utc, - ) - data_out = pd.DataFrame(index=t) - data_out["dst"] = np.array([np.nan] * len(t)) - data_out["file_name"] = np.array([None] * len(t)) - - for file_path in file_paths: - if not file_path.exists(): - if download: - self.download_and_process(start_time, end_time) - else: - warnings.warn(f"File {file_path} not found") - continue - - df_one_file = self._read_single_file(file_path) - data_out = df_one_file.combine_first(data_out) - - data_out = data_out.truncate( - before=start_time - timedelta(hours=0.9999), - after=end_time + timedelta(hours=0.9999), - ) + data_out = super().read(start_time, end_time, download=download, variables="dst") data_out.index.name = "t" - data_out.drop(columns=["timestamp", "t"], inplace=True, errors="ignore") return data_out - - def _read_single_file(self, file_path: Path) -> pd.DataFrame: - """Read yearly OMNI DST file to a DataFrame. - - Parameters - ---------- - file_path : Path - Path to the file. - - Returns - ------- - pd.DataFrame - Data from yearly OMNI DST file. - """ - df = pd.read_csv(file_path) - df.drop(columns=["kp", "f107"], inplace=True) - - df["t"] = pd.to_datetime(df["timestamp"], utc=True) - df.index = df["t"] - - df["file_name"] = file_path - df.loc[df["dst"].isna(), "file_name"] = None - - return df diff --git a/swvo/io/f10_7/omni.py b/swvo/io/f10_7/omni.py index d8338ec..48711eb 100644 --- a/swvo/io/f10_7/omni.py +++ b/swvo/io/f10_7/omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -8,19 +9,12 @@ from __future__ import annotations -import logging from datetime import datetime, timedelta import numpy as np import pandas as pd from swvo.io.omni import OMNILowRes -from swvo.io.utils import enforce_utc_timezone - -logger = logging.getLogger(__name__) - - -logging.captureWarnings(True) class F107OMNI(OMNILowRes): @@ -29,8 +23,12 @@ class F107OMNI(OMNILowRes): Inherits the :func:`download_and_process`, other private methods and attributes from :class:`OMNILowRes`. """ + _READ_TIME_PADDING = timedelta(hours=23.9999) + # data is downloaded along with OMNI data, check file name in parent class - def read(self, start_time: datetime, end_time: datetime, download: bool = False) -> pd.DataFrame: + def read( # ty: ignore[invalid-method-override] + self, start_time: datetime, end_time: datetime, download: bool = False + ) -> pd.DataFrame: """ Extract F10.7 data from OMNI Low Resolution files. @@ -49,15 +47,7 @@ def read(self, start_time: datetime, end_time: datetime, download: bool = False) F10.7 from OMNI Low Resolution data. """ - if start_time > end_time: - msg = "start_time must be before end_time" - logger.error(msg) - raise ValueError(msg) - - start_time = enforce_utc_timezone(start_time) - end_time = enforce_utc_timezone(end_time) - - data_out = super().read(start_time, end_time, download=download) + data_out = super().read(start_time, end_time, download=download, variables="f107") f107_df = pd.DataFrame(index=data_out.index) @@ -67,9 +57,5 @@ def read(self, start_time: datetime, end_time: datetime, download: bool = False) # we return it just every 24 hours f107_df = f107_df.drop(f107_df[data_out.index.hour % 24 != 0].index, axis=0) # ty: ignore[unresolved-attribute] f107_df = f107_df.replace(999.9, np.nan) - f107_df = f107_df.truncate( - before=start_time - timedelta(hours=23.9999), - after=end_time + timedelta(hours=23.9999), - ) - + f107_df.loc[f107_df["f107"].isna(), "file_name"] = None return f107_df # noqa: RET504 diff --git a/swvo/io/kp/omni.py b/swvo/io/kp/omni.py index 9f4ce1a..1d93f1f 100755 --- a/swvo/io/kp/omni.py +++ b/swvo/io/kp/omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -8,18 +9,11 @@ from __future__ import annotations -import logging from datetime import datetime, timedelta import pandas as pd from swvo.io.omni import OMNILowRes -from swvo.io.utils import enforce_utc_timezone - -logger = logging.getLogger(__name__) - - -logging.captureWarnings(True) class KpOMNI(OMNILowRes): @@ -28,7 +22,11 @@ class KpOMNI(OMNILowRes): Inherits the :func:`download_and_process`, other private methods and attributes from :class:`OMNILowRes`. """ - def read(self, start_time: datetime, end_time: datetime, download: bool = False) -> pd.DataFrame: + _READ_TIME_PADDING = timedelta(hours=2.9999) + + def read( # ty: ignore[invalid-method-override] + self, start_time: datetime, end_time: datetime, download: bool = False + ) -> pd.DataFrame: """ Extract Kp data from OMNI Low Resolution files. @@ -47,23 +45,11 @@ def read(self, start_time: datetime, end_time: datetime, download: bool = False) Kp data from OMNI Low Resolution data. """ - if start_time > end_time: - msg = "start_time must be before end_time" - logger.error(msg) - raise ValueError(msg) - - data_out = super().read(start_time, end_time, download=download) + data_out = super().read(start_time, end_time, download=download, variables="kp") kp_df = pd.DataFrame(index=data_out.index) - start_time = enforce_utc_timezone(start_time) - end_time = enforce_utc_timezone(end_time) kp_df["kp"] = data_out["kp"] kp_df["file_name"] = data_out["file_name"] # we return it just every 3 hours kp_df = kp_df.drop(kp_df[data_out.index.hour % 3 != 0].index, axis=0) # ty: ignore[unresolved-attribute] - kp_df = kp_df.truncate( - before=start_time - timedelta(hours=2.9999), - after=end_time + timedelta(hours=2.9999), - ) - return kp_df # noqa: RET504 diff --git a/swvo/io/omni/omni_high_res.py b/swvo/io/omni/omni_high_res.py index a9aa6c1..b862e6e 100644 --- a/swvo/io/omni/omni_high_res.py +++ b/swvo/io/omni/omni_high_res.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -9,6 +10,7 @@ import calendar import logging import re +from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta, timezone from typing import List, Optional, Tuple @@ -17,6 +19,12 @@ import requests from swvo.io.base import BaseIO +from swvo.io.omni.variables import ( + HIGH_RES_DEFAULT_VARIABLES, + HIGH_RES_VARIABLES, + resolve_variable_names, +) +from swvo.io.omni.variables import available_variables as get_available_variables from swvo.io.utils import enforce_utc_timezone logger = logging.getLogger(__name__) @@ -50,6 +58,44 @@ class OMNIHighRes(BaseIO): PARALLEL_DOWNLOAD_THRESHOLD = 10 MAX_PARALLEL_DOWNLOADS = 10 + def available_variables(self, cadence: int = 1) -> pd.DataFrame: + """Return metadata for variables available at a cadence. + + Parameters + ---------- + cadence : int, optional + OMNI cadence in minutes. Must be ``1`` or ``5``. + + Returns + ------- + pandas.DataFrame + One row per variable with its canonical name, NASA request ID, + description, unit, supported cadences, and accepted aliases. + + Raises + ------ + AssertionError + If ``cadence`` is not ``1`` or ``5``. + """ + + self._validate_cadence(cadence) + return get_available_variables(cadence) + + def _validate_cadence(self, cadence_min: int) -> None: + if cadence_min not in (1, 5): + raise AssertionError("Only 1 or 5 minute cadence can be chosen for high resolution omni data.") + + def _cache_contains(self, file_path, variable_names: Iterable[str]) -> bool: + """Check a processed file's schema without loading its data.""" + + if not file_path.exists(): + return False + try: + columns = set(pd.read_csv(file_path, nrows=0).columns) + except (OSError, UnicodeDecodeError, pd.errors.ParserError, pd.errors.EmptyDataError): + return False + return set(variable_names).issubset(columns) + def download_and_process( self, start_time: datetime, @@ -78,17 +124,28 @@ def download_and_process( ------ AssertionError Raises `AssertionError` if the cadence is not 1 or 5 minutes. + ValueError + If ``start_time`` is not before ``end_time``. """ - assert cadence_min == 1 or cadence_min == 5, ( - "Only 1 or 5 minute cadence can be chosen for high resolution omni data." + self._validate_cadence(cadence_min) + start_time = enforce_utc_timezone(start_time) + end_time = enforce_utc_timezone(end_time) + if start_time >= end_time: + raise ValueError("start_time must be before end_time") + + complete_schema = resolve_variable_names( + HIGH_RES_VARIABLES, + "all", + HIGH_RES_DEFAULT_VARIABLES, + cadence=cadence_min, ) file_paths, time_intervals = self._get_processed_file_list(start_time, end_time, cadence_min) download_tasks = [] for file_path, time_interval in zip(file_paths, time_intervals): - if file_path.exists() and not reprocess_files: + if not reprocess_files and self._cache_contains(file_path, complete_schema): continue download_tasks.append((file_path, time_interval)) @@ -122,15 +179,27 @@ def _download_and_process_single_file( tmp_path = file_path.with_suffix(file_path.suffix + ".tmp") try: + variable_names = resolve_variable_names( + HIGH_RES_VARIABLES, + "all", + HIGH_RES_DEFAULT_VARIABLES, + cadence=cadence_min, + ) data = self._get_data_from_omni( start=time_interval[0], end=time_interval[1], cadence=cadence_min, + variable_names=variable_names, ) logger.debug("Processing file ...") - processed_df = self._process_single_month(data, original_end=time_interval[1], cadence_min=cadence_min) + processed_df = self._process_single_month( + data, + original_end=time_interval[1], + cadence_min=cadence_min, + variable_names=variable_names, + ) # Do not save empty DataFrames — no data available for this interval if processed_df.empty: @@ -151,9 +220,9 @@ def read( end_time: datetime, cadence_min: int = 1, download: bool = False, + variables: str | Iterable[str] | None = None, ) -> pd.DataFrame: - """ - Read OMNI High Resolution data for the given time range. + """Read high-resolution OMNI data for a time range. Parameters ---------- @@ -165,29 +234,43 @@ def read( Cadence of the data in minutes, defaults to 1 download : bool, optional Download data on the go, defaults to False. + variables : str or iterable of str or None, optional + Variables to return. ``None`` preserves the legacy nine-column + schema, ``"all"`` returns every variable supported by the selected + cadence, and a name or iterable selects a subset. Normalized aliases + such as ``"sym_h"`` are accepted. Returns ------- :class:`pandas.DataFrame` - OMNI High Resolution data. + Selected OMNI variables plus ``file_name`` provenance. The index is + timezone-aware UTC. Raises ------ AssertionError Raises `AssertionError` if the cadence is not 1 or 5 minutes. + ValueError + If the time range is invalid, a variable is unknown or unavailable + at the selected cadence, or an existing partial or unreadable cache + cannot satisfy the request. """ - assert cadence_min == 1 or cadence_min == 5, ( - "Only 1 or 5 minute cadence can be chosen for high resolution omni data." + self._validate_cadence(cadence_min) + variable_names = resolve_variable_names( + HIGH_RES_VARIABLES, + variables, + HIGH_RES_DEFAULT_VARIABLES, + cadence=cadence_min, ) - if start_time > end_time: + start_time = enforce_utc_timezone(start_time) + end_time = enforce_utc_timezone(end_time) + + if start_time >= end_time: msg = "start_time must be before end_time" logger.error(msg) raise ValueError(msg) - start_time = enforce_utc_timezone(start_time) - end_time = enforce_utc_timezone(end_time) - if start_time < datetime(self.START_YEAR, 1, 1, tzinfo=timezone.utc): logger.warning( "Start date chosen falls behind the existing data. Moving start date to first" @@ -195,16 +278,17 @@ def read( ) start_time = datetime(self.START_YEAR, 1, 1, tzinfo=timezone.utc) - assert start_time < end_time + if start_time >= end_time: + raise ValueError(f"Requested time range ends before OMNI data begin in {self.START_YEAR}.") - file_paths, _ = self._get_processed_file_list(start_time, end_time, cadence_min) + file_paths, time_intervals = self._get_processed_file_list(start_time, end_time, cadence_min) dfs = [] - for file_path in file_paths: + for file_path, time_interval in zip(file_paths, time_intervals): if not file_path.exists(): if download: - self.download_and_process(start_time, end_time, cadence_min=cadence_min) + self._download_and_process_single_file(file_path, time_interval, cadence_min) else: logger.warning(f"File {file_path} not found") continue @@ -215,8 +299,40 @@ def read( logger.warning(f"File {file_path} not available after download attempt, skipping.") continue - dfs.append(self._read_single_file(file_path)) - + try: + df = self._read_single_file(file_path) + except (OSError, UnicodeDecodeError, ValueError, KeyError, pd.errors.ParserError) as error: + if not download: + raise ValueError(f"Cannot read processed OMNI file {file_path}: {error}") from error + logger.info(f"Replacing unreadable OMNI cache {file_path}") + self._download_and_process_single_file(file_path, time_interval, cadence_min) + try: + df = self._read_single_file(file_path) + except (OSError, UnicodeDecodeError, ValueError, KeyError, pd.errors.ParserError) as retry_error: + raise ValueError( + f"Processed OMNI file {file_path} remains unreadable after attempted cache upgrade: " + f"{retry_error}" + ) from retry_error + missing = [name for name in variable_names if name not in df.columns] + if missing and download: + logger.info(f"Upgrading partial OMNI cache {file_path} for variables: {', '.join(missing)}") + self._download_and_process_single_file(file_path, time_interval, cadence_min) + df = self._read_single_file(file_path) + missing = [name for name in variable_names if name not in df.columns] + if missing: + raise ValueError( + f"Processed OMNI file {file_path} does not contain: {', '.join(missing)}. " + "Call read(..., download=True) or download_and_process(..., reprocess_files=True) " + "to upgrade the cache." + ) + + selected = df.loc[:, variable_names].copy() + selected["file_name"] = file_path + selected.loc[selected[variable_names].isna().all(axis=1), "file_name"] = None + dfs.append(selected) + + if not dfs: + raise ValueError("No OMNI High Resolution files are available for the requested time range.") data_out = pd.concat(dfs, ignore_index=False) if not data_out.empty: @@ -305,7 +421,11 @@ def _get_processed_file_list( return file_paths, time_intervals def _process_single_month( - self, data: list[str], original_end: Optional[datetime] = None, cadence_min: int = 1 + self, + data: list[str], + original_end: Optional[datetime] = None, + cadence_min: int = 1, + variable_names: Iterable[str] | None = None, ) -> pd.DataFrame: """Process monthly OMNI High Resolution data to a DataFrame. @@ -316,6 +436,10 @@ def _process_single_month( original_end : datetime, optional The original requested end time. Used to build a NaN-filled DataFrame when no data is available (e.g. the interval is beyond the OMNI data range). + cadence_min : int, optional + Cadence of the requested records. + variable_names : iterable of str, optional + Canonical variables in the same order as the OMNIWeb request. Returns ------- @@ -324,12 +448,18 @@ def _process_single_month( up to ``original_end`` if no data is available, or an empty DataFrame if ``original_end`` is not provided. """ - columns = ["bavg", "bx_gsm", "by_gsm", "bz_gsm", "speed", "proton_density", "temperature", "pdyn", "sym-h"] + columns = resolve_variable_names( + HIGH_RES_VARIABLES, + list(variable_names) if variable_names is not None else None, + HIGH_RES_DEFAULT_VARIABLES, + cadence=cadence_min, + ) # Empty data list signals that no data is available for this interval if not data: if original_end is None: return pd.DataFrame(columns=columns) + original_end = enforce_utc_timezone(original_end) # Build a NaN-filled DataFrame with the correct timestamps up to original_end index = pd.date_range( start=original_end.replace(day=1, hour=0, minute=0, second=0, microsecond=0), @@ -339,44 +469,50 @@ def _process_single_month( ) return pd.DataFrame(pd.NA, index=index, columns=columns) - header_line = next(line for line in data if line.strip().startswith("YYYY")) - columns = header_line.split() + if not any(line.strip().startswith("YYYY") for line in data): + raise ValueError("OMNIWeb response does not contain the expected YYYY data header.") - data_lines = [line for line in data if line.strip().startswith(("19", "20"))] + # OMNIWeb wraps records in HTML and its parameter list contains lines + # such as ``19 Vx Velocity``. Require all four leading time fields so + # those labels (and year-prefixed footer text) cannot become data rows. + record_prefix = re.compile(r"^\s*\d{4}\s+\d{1,3}\s+\d{1,2}\s+\d{1,2}(?:\s|$)") + data_lines = [line for line in data if record_prefix.match(line)] if not data_lines: msg = "DataFrame is empty." logger.error(msg) raise ValueError(msg) - df = pd.DataFrame([line.split() for line in data_lines], columns=columns) + raw_columns = ["YYYY", "DOY", "HR", "MN", *columns] + rows = [line.split() for line in data_lines] + invalid_widths = sorted({len(row) for row in rows if len(row) != len(raw_columns)}) + if invalid_widths: + raise ValueError( + f"OMNIWeb returned record widths {invalid_widths}; expected {len(raw_columns)} values " + f"for {len(columns)} selected variables." + ) + + df = pd.DataFrame(rows, columns=raw_columns) df = df.apply(pd.to_numeric) - df["timestamp"] = df["YYYY"].map(str).apply(lambda x: x + "-01-01 ") + df["HR"].map(str).apply( - lambda x: x.zfill(2) - ) - df["timestamp"] += df["MN"].map(str).apply(lambda x: ":" + x.zfill(2) + ":00") - df["timestamp"] = pd.to_datetime(df["timestamp"]) - df["timestamp"] = df["timestamp"] + df["DOY"].apply(lambda x: timedelta(days=int(x) - 1)) + invalid_time = ~df["DOY"].between(1, 366) | ~df["HR"].between(0, 23) | ~df["MN"].between(0, 59) + if invalid_time.any(): + raise ValueError("OMNIWeb returned an invalid day-of-year, hour, or minute field.") + + year_and_day = df["YYYY"].astype(int).astype(str) + df["DOY"].astype(int).astype(str).str.zfill(3) + df["timestamp"] = pd.to_datetime(year_and_day, format="%Y%j", utc=True) + if not (df["timestamp"].dt.year == df["YYYY"]).all(): + raise ValueError("OMNIWeb returned an invalid day-of-year for its record year.") + df["timestamp"] += pd.to_timedelta(df["HR"], unit="h") + pd.to_timedelta(df["MN"], unit="m") df.drop(columns=["YYYY", "HR", "MN", "DOY"], inplace=True) df.set_index("timestamp", inplace=True) - maxes = { - "bavg": 9999.9, - "bx_gsm": 9999.9, - "by_gsm": 9999.9, - "bz_gsm": 9999.9, - "speed": 99999.8, - "proton_density": 999.8, - "temperature": 9999998.0, - "pdyn": 99.0, - "sym-h": 99999.0, - } - - df.columns = maxes.keys() - for col, max_val in maxes.items(): - df[col] = df[col].where(df[col] < max_val, other=pd.NA) + metadata = {variable.name: variable for variable in HIGH_RES_VARIABLES} + for column in columns: + fill_value = metadata[column].fill_value + if fill_value is not None: + df[column] = df[column].where(df[column] < fill_value, other=pd.NA) if df.empty: msg = "DataFrame is empty after processing the month." @@ -408,7 +544,13 @@ def _read_single_file(self, file_path) -> pd.DataFrame: return df - def _get_data_from_omni(self, start: datetime, end: datetime, cadence: int = 1) -> list: + def _get_data_from_omni( + self, + start: datetime, + end: datetime, + cadence: int = 1, + variable_names: Iterable[str] | None = None, + ) -> list[str]: """ Fetches data from NASA's OMNIWeb service. @@ -418,56 +560,68 @@ def _get_data_from_omni(self, start: datetime, end: datetime, cadence: int = 1) that no data is available for the requested interval. """ - payload = { + self._validate_cadence(cadence) + start = enforce_utc_timezone(start) + end = enforce_utc_timezone(end) + selected_names = resolve_variable_names( + HIGH_RES_VARIABLES, + list(variable_names) if variable_names is not None else None, + HIGH_RES_DEFAULT_VARIABLES, + cadence=cadence, + ) + metadata = {variable.name: variable for variable in HIGH_RES_VARIABLES} + payload: dict[str, str | list[str]] = { "activity": "retrieve", "start_date": start.strftime("%Y%m%d"), "end_date": end.strftime("%Y%m%d"), + "vars": [str(metadata[name].nasa_id) for name in selected_names], } - common_vars = {"vars": ["13", "14", "17", "18", "21", "25", "26", "27", "41"]} if cadence == 1: - params = {"res": "min", "spacecraft": "omni_min"} - payload.update(params) - payload.update(common_vars) - elif cadence == 5: - params = {"res": "5min", "spacecraft": "omni_5min"} - payload.update(params) - payload.update(common_vars) - + payload.update({"res": "min", "spacecraft": "omni_min"}) else: - msg = f"Invalid cadence: {cadence}. Only 1 or 5 minutes are supported." - logger.error(msg) - raise ValueError(msg) + payload.update({"res": "5min", "spacecraft": "omni_5min"}) logger.debug(f"Fetching data from {self.URL} with payload: {payload}") - response = requests.post(self.URL, data=payload) + response = requests.post(self.URL, data=payload, timeout=30) response.raise_for_status() data = response.text.splitlines() - if data and "Error" in data[0]: + correct_range_line = next((line for line in data if "correct range" in line.lower()), None) + has_error = any(re.search(r"\berror\b", line, flags=re.IGNORECASE) for line in data) + if correct_range_line is not None or has_error: logger.warning("Received an error response from the server.") - for line in data: - if "correct range" in line: - # Use regex to find the valid date range (e.g., YYYYMMDD - YYYYMMDD) - match = re.search(r"correct range: \d{8} - (\d{8})", line) - if match: - new_end_date_str = match.group(1) - new_end_date = datetime.strptime(new_end_date_str, "%Y%m%d").replace(tzinfo=timezone.utc) + if correct_range_line is not None: + # Use regex to find the valid date range (e.g., YYYYMMDD - YYYYMMDD) + match = re.search(r"correct range:\s*\d{8}\s*-\s*(\d{8})", correct_range_line, re.IGNORECASE) + if match: + new_end_date_str = match.group(1) + new_end_date = datetime.strptime(new_end_date_str, "%Y%m%d").replace(tzinfo=timezone.utc) + logger.warning( + f"Invalid date range detected. Found suggested end date: {new_end_date.strftime('%Y-%m-%d')}" + ) + + # If the suggested end date is before the start date, no data is available + # for this range — return an empty list to signal an empty DataFrame + if new_end_date < start: logger.warning( - f"Invalid date range detected. Found suggested end date: {new_end_date.strftime('%Y-%m-%d')}" + f"Suggested end date {new_end_date.strftime('%Y-%m-%d')} is before " + f"start date {start.strftime('%Y-%m-%d')}. No data available for this range." + ) + return [] + + if new_end_date >= end: + raise ValueError( + "OMNIWeb returned an invalid corrected end date that would not shorten the request." ) - # If the suggested end date is before the start date, no data is available - # for this range — return an empty list to signal an empty DataFrame - if new_end_date < start: - logger.warning( - f"Suggested end date {new_end_date.strftime('%Y-%m-%d')} is before " - f"start date {start.strftime('%Y-%m-%d')}. No data available for this range." - ) - return [] - - # Recursively call the function with the original start date and the new end date - return self._get_data_from_omni(start=start, end=new_end_date, cadence=cadence) + # Recursively call the function with the original start date and the new end date + return self._get_data_from_omni( + start=start, + end=new_end_date, + cadence=cadence, + variable_names=selected_names, + ) msg = f"An unspecified error occurred: {data}" logger.error(msg) raise ValueError(msg) diff --git a/swvo/io/omni/omni_low_res.py b/swvo/io/omni/omni_low_res.py index a13c21b..9f775f8 100755 --- a/swvo/io/omni/omni_low_res.py +++ b/swvo/io/omni/omni_low_res.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -8,9 +9,10 @@ import logging import warnings +from collections.abc import Iterable from datetime import datetime, timedelta, timezone from pathlib import Path -from shutil import rmtree +from tempfile import TemporaryDirectory from typing import List, Tuple import numpy as np @@ -18,6 +20,12 @@ import requests from swvo.io.base import BaseIO +from swvo.io.omni.variables import ( + LOW_RES_DEFAULT_VARIABLES, + LOW_RES_VARIABLES, + resolve_variable_names, +) +from swvo.io.omni.variables import available_variables as get_available_variables from swvo.io.utils import enforce_utc_timezone logger = logging.getLogger(__name__) @@ -44,63 +52,36 @@ class OMNILowRes(BaseIO): URL = "https://spdf.gsfc.nasa.gov/pub/data/omni/low_res_omni/" LABEL = "omni" - HEADER = [ - "year", - "day", - "hour", - "BarRotNumber", - "id_imf", - "id_sw", - "%points_imfavg", - "%points_plasmaavg", - "B_mag_avg", - "bavg", - "lat_angle_avg_field", - "lon_angle_avg_field", - "bx_gse_gsm", - "by_gse", - "bz_gse", - "by_gsm", - "bz_gsm", - "sigma_mod_B", - "sigma_B", - "sigma_Bx", - "sigma_By", - "sigma_Bz", - "temperature", - "proton_density", - "speed", - "speed_angle_lon", - "speed_angle_lat", - "alpha_proton_ratio", - "flow_pressure", - "sigma_T", - "sigma_N", - "sigma_V", - "sigma_phi_V", - "sigma_theta_V", - "sigma_alpha_proton_ratio", - "e", - "plasma_beta", - "alfven_mach_n", - "Kp", - "sunspot_n", - "dst", - "ae", - "p_flux_1", - "p_flux_2", - "p_flux_4", - "p_flux_10", - "p_flux_30", - "p_flux_60", - "flag", - "ap", - "f107", - "pc", - "al", - "au", - "magnetosonic_mach_n", - ] + TIME_COLUMNS = ("year", "day", "hour") + HEADER = [*TIME_COLUMNS, *(variable.name for variable in LOW_RES_VARIABLES)] + + # Focused hourly readers expose values at their natural cadence and may + # include the nearest value around a requested boundary. Keeping that + # policy here lets this parent remain the single owner of normalized times. + _READ_TIME_PADDING: timedelta | None = None + + def available_variables(self) -> pd.DataFrame: + """Return metadata for every hourly OMNI2 output variable. + + Returns + ------- + pandas.DataFrame + One row per hourly variable with its canonical name, description, + unit, fill value, and accepted aliases. + """ + + return get_available_variables() + + def _cache_contains(self, file_path: Path, variable_names: Iterable[str]) -> bool: + """Check a processed file's schema without loading its data.""" + + if not file_path.exists(): + return False + try: + columns = set(pd.read_csv(file_path, nrows=0).columns) + except (OSError, UnicodeDecodeError, pd.errors.ParserError, pd.errors.EmptyDataError): + return False + return set(variable_names).issubset(columns) def download_and_process(self, start_time: datetime, end_time: datetime, reprocess_files: bool = False) -> None: """Download and process OMNI Low Resolution data files. @@ -117,44 +98,47 @@ def download_and_process(self, start_time: datetime, end_time: datetime, reproce Returns ------- None + + Raises + ------ + ValueError + If ``start_time`` is not before ``end_time``. """ - temporary_dir = Path("./temp_omni_low_res_wget") - temporary_dir.mkdir(exist_ok=True, parents=True) + start_time = enforce_utc_timezone(start_time) + end_time = enforce_utc_timezone(end_time) + if start_time >= end_time: + raise ValueError("start_time must be before end_time") file_paths, time_intervals = self._get_processed_file_list(start_time, end_time) + complete_schema = [variable.name for variable in LOW_RES_VARIABLES] - for file_path, time_interval in zip(file_paths, time_intervals): - if file_path.exists() and not reprocess_files: - continue - - tmp_path = file_path.with_suffix(file_path.suffix + ".tmp") + with TemporaryDirectory(prefix="swvo-omni-low-res-") as temporary_dir_name: + temporary_dir = Path(temporary_dir_name) + for file_path, time_interval in zip(file_paths, time_intervals): + if not reprocess_files and self._cache_contains(file_path, complete_schema): + continue - try: - filename = "omni2_" + str(time_interval[0].year) + ".dat" + tmp_path = file_path.with_suffix(file_path.suffix + ".tmp") - if file_path.exists(): - if reprocess_files: - file_path.unlink() - else: - continue + try: + filename = f"omni2_{time_interval[0].year}.dat" - logger.debug(f"Downloading file {self.URL + filename} ...") + logger.debug(f"Downloading file {self.URL + filename} ...") - self._download(temporary_dir, filename) + self._download(temporary_dir, filename) - logger.debug("Processing file ...") + logger.debug("Processing file ...") - processed_df = self._process_single_file(temporary_dir / filename) - processed_df.to_csv(tmp_path, index=True, header=True) - tmp_path.replace(file_path) + processed_df = self._process_single_file(temporary_dir / filename) + processed_df.to_csv(tmp_path, index=True, header=True) + tmp_path.replace(file_path) - except Exception as e: - logger.error(f"Failed to process {file_path}: {e}") - if tmp_path.exists(): - tmp_path.unlink() - continue - rmtree(temporary_dir, ignore_errors=True) + except Exception as e: + logger.error(f"Failed to process {file_path}: {e}") + if tmp_path.exists(): + tmp_path.unlink() + continue def _download(self, temporary_dir: Path, filename: str): response = requests.get(self.URL + filename, timeout=10) @@ -190,14 +174,8 @@ def _get_processed_file_list(self, start_time: datetime, end_time: datetime) -> for year in range(start_year, end_year + 1): file_path = self.data_dir / f"OMNI_LOW_RES_{year}.csv" file_paths.append(file_path) - if year == start_year: - interval_start = datetime(year, 1, 1, 0, 0, 0) - else: - interval_start = datetime(year, 1, 1, 0, 0, 0) - if year == end_year: - interval_end = datetime(year, 12, 31, 23, 59, 59) - else: - interval_end = datetime(year, 12, 31, 23, 59, 59) + interval_start = datetime(year, 1, 1, 0, 0, 0, tzinfo=end_time.tzinfo) + interval_end = datetime(year, 12, 31, 23, 59, 59, tzinfo=end_time.tzinfo) time_intervals.append((interval_start, interval_end)) return file_paths, time_intervals @@ -216,24 +194,49 @@ def _process_single_file(self, file_path: Path) -> pd.DataFrame: Yearly OMNI Low Resolution data. """ - data = pd.read_csv(file_path, sep=r"\s+", names=self.HEADER) - data["timestamp"] = data["year"].map(str).apply(lambda x: x + "-01-01 ") - data["timestamp"] = data["year"].map(str).apply(lambda x: x + "-01-01 ") + data["hour"].map(str).apply( - lambda x: x.zfill(2) - ) - data["timestamp"] = pd.to_datetime(data["timestamp"]) - data["timestamp"] = data["timestamp"] + data["day"].apply(lambda x: timedelta(days=int(x) - 1)) + try: + data = pd.read_csv(file_path, sep=r"\s+", header=None) + except (pd.errors.ParserError, pd.errors.EmptyDataError) as error: + raise ValueError(f"Cannot parse OMNI2 source file {file_path}: {error}") from error + variable_count = data.shape[1] - len(self.TIME_COLUMNS) + supported_counts = {len(LOW_RES_VARIABLES) - 2, len(LOW_RES_VARIABLES)} + if variable_count not in supported_counts: + raise ValueError( + f"Unsupported OMNI2 record width {data.shape[1]}; expected 55 historic or 57 current fields." + ) + + present_variables = list(LOW_RES_VARIABLES[:variable_count]) + data.columns = [*self.TIME_COLUMNS, *(variable.name for variable in present_variables)] + try: + data = data.apply(pd.to_numeric) + except (TypeError, ValueError) as error: + raise ValueError(f"OMNI2 source file {file_path} contains nonnumeric fields: {error}") from error + + invalid_time = ~data["day"].between(1, 366) | ~data["hour"].between(0, 23) + if invalid_time.any(): + raise ValueError("OMNI2 returned an invalid day-of-year or hour field.") + + year_and_day = data["year"].astype(int).astype(str) + data["day"].astype(int).astype(str).str.zfill(3) + data["timestamp"] = pd.to_datetime(year_and_day, format="%Y%j", utc=True) + if not (data["timestamp"].dt.year == data["year"]).all(): + raise ValueError("OMNI2 returned an invalid day-of-year for its record year.") + data["timestamp"] += pd.to_timedelta(data["hour"], unit="h") data.set_index("timestamp", inplace=True) - mask = data["dst"] >= 99998 - data.loc[mask, "dst"] = np.nan - mask = data["Kp"] == 99 - data.loc[mask, "Kp"] = np.nan + for variable in present_variables: + if variable.fill_value is not None: + data[variable.name] = data[variable.name].where( + data[variable.name] < variable.fill_value, + other=pd.NA, + ) + + # NASA added Lyman-alpha and the proton quasi-invariant after the + # historic 55-word layout. They remain explicit NaN columns for old data. + for variable in LOW_RES_VARIABLES[variable_count:]: + data[variable.name] = pd.NA - df = pd.DataFrame(index=data.index) - df["dst"] = data["dst"] - df["kp"] = data["Kp"] / 10 - df["f107"] = data["f107"] + df = data.loc[:, [variable.name for variable in LOW_RES_VARIABLES]].copy() + df["kp"] = df["kp"] / 10 # change rounded numbers to be equal to 1/3 or 2/3 to be consistent with other Kp products df.loc[round(df["kp"] % 1, 2) == 0.7, "kp"] = round(df.loc[round(df["kp"] % 1, 2) == 0.7, "kp"]) - 1 / 3 @@ -241,9 +244,14 @@ def _process_single_file(self, file_path: Path) -> pd.DataFrame: return df - def read(self, start_time: datetime, end_time: datetime, download: bool = False) -> pd.DataFrame: - """ - Read OMNI Low Resolution data for the given time range. + def read( + self, + start_time: datetime, + end_time: datetime, + download: bool = False, + variables: str | Iterable[str] | None = None, + ) -> pd.DataFrame: + """Read hourly OMNI2 data for a time range. Parameters ---------- @@ -253,22 +261,38 @@ def read(self, start_time: datetime, end_time: datetime, download: bool = False) End time for the data to be read. download : bool, optional Download data on the go, defaults to False. + variables : str or iterable of str or None, optional + Variables to return. ``None`` preserves the legacy ``dst``, ``kp``, + and ``f107`` schema, ``"all"`` returns all 54 non-time fields, and + a name or iterable selects a subset. Returns ------- :class:`pandas.DataFrame` - OMNI Low Resolution data. + Selected OMNI variables plus ``file_name`` provenance. The index is + timezone-aware UTC. + + Raises + ------ + ValueError + If the time range is invalid, a variable is unknown, or an existing + partial or unreadable cache cannot satisfy the request. """ START_YEAR = 1963 + variable_names = resolve_variable_names( + LOW_RES_VARIABLES, + variables, + LOW_RES_DEFAULT_VARIABLES, + ) - if start_time > end_time: + start_time = enforce_utc_timezone(start_time) + end_time = enforce_utc_timezone(end_time) + + if start_time >= end_time: msg = "start_time must be before end_time" logger.error(msg) raise ValueError(msg) - start_time = enforce_utc_timezone(start_time) - end_time = enforce_utc_timezone(end_time) - if start_time < datetime(START_YEAR, 1, 1, tzinfo=timezone.utc): logger.warning( "Start date chosen falls behind the existing data. Moving start date to first" @@ -276,31 +300,65 @@ def read(self, start_time: datetime, end_time: datetime, download: bool = False) ) start_time = datetime(START_YEAR, 1, 1, tzinfo=timezone.utc) - assert start_time < end_time + if start_time >= end_time: + raise ValueError(f"Requested time range ends before OMNI data begin in {START_YEAR}.") - file_paths, _ = self._get_processed_file_list(start_time, end_time) + file_paths, time_intervals = self._get_processed_file_list(start_time, end_time) t = pd.date_range( datetime(start_time.year, start_time.month, start_time.day), datetime(end_time.year, end_time.month, end_time.day, 23, 00, 00), freq=timedelta(hours=1), tz=timezone.utc, ) - data_out = pd.DataFrame(index=t) - data_out["kp"] = np.array([np.nan] * len(t)) - data_out["dst"] = np.array([np.nan] * len(t)) - data_out["f107"] = np.array([np.nan] * len(t)) - data_out["file_name"] = np.array([None] * len(t)) + data_out = pd.DataFrame(np.nan, index=t, columns=variable_names) + data_out["file_name"] = None - for file_path in file_paths: + for file_path, time_interval in zip(file_paths, time_intervals): if not file_path.exists(): if download: - self.download_and_process(start_time, end_time) + self.download_and_process(time_interval[0], time_interval[1]) if not file_path.exists(): warnings.warn(f"File {file_path} not found") continue - df_one_file = self._read_single_file(file_path) - data_out = df_one_file.combine_first(data_out) + try: + df_one_file = self._read_single_file(file_path) + except (OSError, UnicodeDecodeError, ValueError, KeyError, pd.errors.ParserError) as error: + if not download: + raise ValueError(f"Cannot read processed OMNI file {file_path}: {error}") from error + logger.info(f"Replacing unreadable OMNI cache {file_path}") + self.download_and_process(time_interval[0], time_interval[1], reprocess_files=True) + try: + df_one_file = self._read_single_file(file_path) + except (OSError, UnicodeDecodeError, ValueError, KeyError, pd.errors.ParserError) as retry_error: + raise ValueError( + f"Processed OMNI file {file_path} remains unreadable after attempted cache upgrade: " + f"{retry_error}" + ) from retry_error + missing = [name for name in variable_names if name not in df_one_file.columns] + if missing and download: + logger.info(f"Upgrading partial OMNI cache {file_path} for variables: {', '.join(missing)}") + self.download_and_process(time_interval[0], time_interval[1], reprocess_files=True) + df_one_file = self._read_single_file(file_path) + missing = [name for name in variable_names if name not in df_one_file.columns] + if missing: + raise ValueError( + f"Processed OMNI file {file_path} does not contain: {', '.join(missing)}. " + "Call read(..., download=True) or download_and_process(..., reprocess_files=True) " + "to upgrade the cache." + ) + + selected = df_one_file.loc[:, variable_names].copy() + selected["file_name"] = file_path + selected.loc[selected[variable_names].isna().all(axis=1), "file_name"] = None + data_out = selected.combine_first(data_out) + + data_out = data_out.loc[:, [*variable_names, "file_name"]] + if self._READ_TIME_PADDING is not None: + data_out = data_out.truncate( + before=start_time - self._READ_TIME_PADDING, + after=end_time + self._READ_TIME_PADDING, + ) return data_out @@ -317,12 +375,7 @@ def _read_single_file(self, file_path: Path) -> pd.DataFrame: pd.DataFrame Data from yearly OMNI Low Resolution file. """ - df = pd.read_csv(file_path) - - df["t"] = pd.to_datetime(df["timestamp"], utc=True) - df.index = df["t"] - - df["file_name"] = file_path - df.loc[df["kp"].isna(), "file_name"] = None + df = pd.read_csv(file_path, index_col="timestamp") + df.index = pd.to_datetime(df.index, utc=True) return df diff --git a/swvo/io/omni/variables.py b/swvo/io/omni/variables.py new file mode 100644 index 0000000..2907cd9 --- /dev/null +++ b/swvo/io/omni/variables.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: 2026 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel +# +# SPDX-License-Identifier: Apache-2.0 + +"""Variable metadata shared by the OMNI readers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable, Sequence + +import pandas as pd + + +@dataclass(frozen=True) +class OMNIVariable: + """Description of one processed OMNI variable.""" + + name: str + description: str + unit: str + fill_value: float | int | None + nasa_id: int | None = None + cadences: tuple[int, ...] = () + aliases: tuple[str, ...] = () + + +HIGH_RES_VARIABLES: tuple[OMNIVariable, ...] = ( + OMNIVariable("imf_spacecraft_id", "IMF source spacecraft identifier", "", 99, 4, (1, 5)), + OMNIVariable("plasma_spacecraft_id", "Solar-wind plasma source spacecraft identifier", "", 99, 5, (1, 5)), + OMNIVariable("imf_point_count", "Fine-scale points in the IMF average", "count", 999, 6, (1, 5)), + OMNIVariable("plasma_point_count", "Fine-scale points in the plasma average", "count", 999, 7, (1, 5)), + OMNIVariable("percent_interpolated", "Percentage of interpolated values", "%", 999, 8, (1, 5)), + OMNIVariable("timeshift", "Bow-shock-nose time shift", "s", 999999, 9, (1, 5), ("timeshift_seconds",)), + OMNIVariable("timeshift_rms", "RMS time shift", "s", 999999, 10, (1, 5), ("timeshift_rms_seconds",)), + OMNIVariable("minimum_variance_rms", "RMS minimum-variance vector", "", 99.99, 11, (1, 5)), + OMNIVariable( + "observation_gap", + "Time between observations", + "s", + 999999, + 12, + (1, 5), + ("observation_gap_seconds",), + ), + OMNIVariable("bavg", "Average IMF magnitude", "nT", 9999.99, 13, (1, 5)), + OMNIVariable("bx_gsm", "IMF Bx in GSE/GSM coordinates", "nT", 9999.99, 14, (1, 5), ("bx_gse",)), + OMNIVariable("by_gse", "IMF By in GSE coordinates", "nT", 9999.99, 15, (1, 5)), + OMNIVariable("bz_gse", "IMF Bz in GSE coordinates", "nT", 9999.99, 16, (1, 5)), + OMNIVariable("by_gsm", "IMF By in GSM coordinates", "nT", 9999.99, 17, (1, 5)), + OMNIVariable("bz_gsm", "IMF Bz in GSM coordinates", "nT", 9999.99, 18, (1, 5)), + OMNIVariable("sigma_bavg", "RMS standard deviation of IMF magnitude", "nT", 9999.99, 19, (1, 5)), + OMNIVariable("sigma_b", "RMS standard deviation of the IMF vector", "nT", 9999.99, 20, (1, 5)), + OMNIVariable("speed", "Solar-wind bulk speed", "km/s", 99999.9, 21, (1, 5)), + OMNIVariable("vx_gse", "Solar-wind Vx in GSE coordinates", "km/s", 99999.9, 22, (1, 5)), + OMNIVariable("vy_gse", "Solar-wind Vy in GSE coordinates", "km/s", 99999.9, 23, (1, 5)), + OMNIVariable("vz_gse", "Solar-wind Vz in GSE coordinates", "km/s", 99999.9, 24, (1, 5)), + OMNIVariable("proton_density", "Proton number density", "1/cm^3", 999.99, 25, (1, 5)), + OMNIVariable("temperature", "Proton temperature", "K", 9999999.0, 26, (1, 5)), + OMNIVariable("pdyn", "Solar-wind flow pressure", "nPa", 99.99, 27, (1, 5), ("flow_pressure",)), + OMNIVariable("electric_field", "Convective electric field Ey", "mV/m", 999.99, 28, (1, 5), ("ey",)), + OMNIVariable("plasma_beta", "Plasma beta", "", 999.99, 29, (1, 5)), + OMNIVariable("alfven_mach_number", "Alfven Mach number", "", 999.9, 30, (1, 5)), + OMNIVariable("spacecraft_x_gse", "Spacecraft X position in GSE coordinates", "Re", 9999.99, 31, (1, 5)), + OMNIVariable("spacecraft_y_gse", "Spacecraft Y position in GSE coordinates", "Re", 9999.99, 32, (1, 5)), + OMNIVariable("spacecraft_z_gse", "Spacecraft Z position in GSE coordinates", "Re", 9999.99, 33, (1, 5)), + OMNIVariable("bsn_x_gse", "Bow-shock-nose X position in GSE coordinates", "Re", 9999.99, 34, (1, 5)), + OMNIVariable("bsn_y_gse", "Bow-shock-nose Y position in GSE coordinates", "Re", 9999.99, 35, (1, 5)), + OMNIVariable("bsn_z_gse", "Bow-shock-nose Z position in GSE coordinates", "Re", 9999.99, 36, (1, 5)), + OMNIVariable("ae", "Auroral electrojet AE index", "nT", 99999, 37, (1, 5)), + OMNIVariable("al", "Auroral electrojet AL index", "nT", 99999, 38, (1, 5)), + OMNIVariable("au", "Auroral electrojet AU index", "nT", 99999, 39, (1, 5)), + OMNIVariable("sym-d", "SYM/D geomagnetic index", "nT", 99999, 40, (1, 5), ("sym_d",)), + OMNIVariable("sym-h", "SYM/H geomagnetic index", "nT", 99999, 41, (1, 5), ("sym_h",)), + OMNIVariable("asy-d", "ASY/D geomagnetic index", "nT", 99999, 42, (1, 5), ("asy_d",)), + OMNIVariable("asy-h", "ASY/H geomagnetic index", "nT", 99999, 43, (1, 5), ("asy_h",)), + OMNIVariable("pcn", "Northern polar-cap index", "", 999.99, 44, (1, 5), ("pc",)), + OMNIVariable("magnetosonic_mach_number", "Magnetosonic Mach number", "", 99.9, 45, (1, 5)), + OMNIVariable("proton_flux_10_mev", "Integral proton flux above 10 MeV", "1/(cm^2 s sr)", 99999.99, 46, (5,)), + OMNIVariable("proton_flux_30_mev", "Integral proton flux above 30 MeV", "1/(cm^2 s sr)", 99999.99, 47, (5,)), + OMNIVariable("proton_flux_60_mev", "Integral proton flux above 60 MeV", "1/(cm^2 s sr)", 99999.99, 48, (5,)), +) + + +LOW_RES_VARIABLES: tuple[OMNIVariable, ...] = ( + OMNIVariable("bartels_rotation_number", "Bartels rotation number", "", 9999), + OMNIVariable("imf_spacecraft_id", "IMF source spacecraft identifier", "", 99), + OMNIVariable("plasma_spacecraft_id", "Solar-wind plasma source spacecraft identifier", "", 99), + OMNIVariable("imf_point_count", "Fine-scale points in the IMF average", "count", 999), + OMNIVariable("plasma_point_count", "Fine-scale points in the plasma average", "count", 999), + OMNIVariable("bavg", "Average IMF magnitude", "nT", 999.9), + OMNIVariable("magnitude_average_field_vector", "Magnitude of the average IMF vector", "nT", 999.9), + OMNIVariable("latitude_average_field", "Latitude angle of the average IMF vector", "degree", 999.9), + OMNIVariable("longitude_average_field", "Longitude angle of the average IMF vector", "degree", 999.9), + OMNIVariable("bx_gse_gsm", "IMF Bx in GSE/GSM coordinates", "nT", 999.9, aliases=("bx_gse",)), + OMNIVariable("by_gse", "IMF By in GSE coordinates", "nT", 999.9), + OMNIVariable("bz_gse", "IMF Bz in GSE coordinates", "nT", 999.9), + OMNIVariable("by_gsm", "IMF By in GSM coordinates", "nT", 999.9), + OMNIVariable("bz_gsm", "IMF Bz in GSM coordinates", "nT", 999.9), + OMNIVariable("sigma_bavg", "RMS standard deviation of IMF magnitude", "nT", 999.9), + OMNIVariable("sigma_b", "RMS standard deviation of the IMF vector", "nT", 999.9), + OMNIVariable("sigma_bx", "RMS standard deviation of IMF Bx", "nT", 999.9), + OMNIVariable("sigma_by", "RMS standard deviation of IMF By", "nT", 999.9), + OMNIVariable("sigma_bz", "RMS standard deviation of IMF Bz", "nT", 999.9), + OMNIVariable("temperature", "Proton temperature", "K", 9999999.0), + OMNIVariable("proton_density", "Proton number density", "1/cm^3", 999.9), + OMNIVariable("speed", "Solar-wind bulk speed", "km/s", 9999.0), + OMNIVariable("speed_longitude", "Solar-wind flow longitude in GSE coordinates", "degree", 999.9), + OMNIVariable("speed_latitude", "Solar-wind flow latitude in GSE coordinates", "degree", 999.9), + OMNIVariable("alpha_proton_ratio", "Alpha-to-proton density ratio", "", 9.999), + OMNIVariable("pdyn", "Solar-wind flow pressure", "nPa", 99.99, aliases=("flow_pressure",)), + OMNIVariable("sigma_temperature", "RMS standard deviation of proton temperature", "K", 9999999.0), + OMNIVariable("sigma_proton_density", "RMS standard deviation of proton density", "1/cm^3", 999.9), + OMNIVariable("sigma_speed", "RMS standard deviation of solar-wind speed", "km/s", 9999.0), + OMNIVariable("sigma_speed_longitude", "RMS standard deviation of flow longitude", "degree", 999.9), + OMNIVariable("sigma_speed_latitude", "RMS standard deviation of flow latitude", "degree", 999.9), + OMNIVariable("sigma_alpha_proton_ratio", "RMS standard deviation of alpha-to-proton ratio", "", 9.999), + OMNIVariable("electric_field", "Convective electric field", "mV/m", 999.99, aliases=("e",)), + OMNIVariable("plasma_beta", "Plasma beta", "", 999.99), + OMNIVariable("alfven_mach_number", "Alfven Mach number", "", 999.9), + OMNIVariable("kp", "Three-hour Kp index", "", 99), + OMNIVariable("sunspot_number", "Daily international sunspot number", "", 999), + OMNIVariable("dst", "Dst geomagnetic index", "nT", 99999), + OMNIVariable("ae", "Auroral electrojet AE index", "nT", 9999), + OMNIVariable("proton_flux_1_mev", "Integral proton flux above 1 MeV", "1/(cm^2 s sr)", 999999.99), + OMNIVariable("proton_flux_2_mev", "Integral proton flux above 2 MeV", "1/(cm^2 s sr)", 99999.99), + OMNIVariable("proton_flux_4_mev", "Integral proton flux above 4 MeV", "1/(cm^2 s sr)", 99999.99), + OMNIVariable("proton_flux_10_mev", "Integral proton flux above 10 MeV", "1/(cm^2 s sr)", 99999.99), + OMNIVariable("proton_flux_30_mev", "Integral proton flux above 30 MeV", "1/(cm^2 s sr)", 99999.99), + OMNIVariable("proton_flux_60_mev", "Integral proton flux above 60 MeV", "1/(cm^2 s sr)", 99999.99), + OMNIVariable("magnetospheric_flux_flag", "Magnetospheric proton-flux contamination flag", "", None), + OMNIVariable("ap", "Three-hour ap index", "nT", 999), + OMNIVariable("f107", "Daily F10.7 solar radio flux adjusted to 1 AU", "10^-22 W/(m^2 Hz)", 999.9), + OMNIVariable("pcn", "Northern polar-cap index", "", 999.9, aliases=("pc",)), + OMNIVariable("al", "Auroral electrojet AL index", "nT", 99999), + OMNIVariable("au", "Auroral electrojet AU index", "nT", 99999), + OMNIVariable("magnetosonic_mach_number", "Magnetosonic Mach number", "", 99.9), + OMNIVariable("lyman_alpha", "Daily solar Lyman-alpha irradiance", "W/m^2", 0.999999), + OMNIVariable("proton_quasi_invariant", "Solar-wind proton quasi-invariant", "", 9.9999), +) + + +HIGH_RES_DEFAULT_VARIABLES: tuple[str, ...] = ( + "bavg", + "bx_gsm", + "by_gsm", + "bz_gsm", + "speed", + "proton_density", + "temperature", + "pdyn", + "sym-h", +) + +LOW_RES_DEFAULT_VARIABLES: tuple[str, ...] = ("dst", "kp", "f107") + + +def available_variables(cadence: int | None = None) -> pd.DataFrame: + """Return metadata for the OMNI variables available at a cadence. + + Parameters + ---------- + cadence : int | None, optional + ``None`` selects the hourly OMNI2 registry. ``1`` and ``5`` select the + corresponding high-resolution OMNIWeb registry. + + Returns + ------- + pandas.DataFrame + A new metadata table in registry order. Hourly rows contain ``name``, + ``description``, ``unit``, ``fill_value``, and ``aliases``. High- + resolution rows additionally contain ``nasa_id`` and ``cadences``. + + Raises + ------ + ValueError + If ``cadence`` is not ``None``, ``1``, or ``5``. + + Examples + -------- + >>> available_variables(cadence=1) # one-minute OMNIWeb fields + >>> available_variables(cadence=5) # five-minute OMNIWeb fields + >>> available_variables() # hourly OMNI2 fields + """ + + if cadence not in (None, 1, 5): + raise ValueError("cadence must be None for hourly OMNI2 data, or 1 or 5 minutes for OMNIWeb data") + + if cadence is None: + return pd.DataFrame( + [ + { + "name": variable.name, + "description": variable.description, + "unit": variable.unit, + "fill_value": variable.fill_value, + "aliases": variable.aliases, + } + for variable in LOW_RES_VARIABLES + ] + ) + + return pd.DataFrame( + [ + { + "name": variable.name, + "nasa_id": variable.nasa_id, + "description": variable.description, + "unit": variable.unit, + "fill_value": variable.fill_value, + "cadences": variable.cadences, + "aliases": variable.aliases, + } + for variable in HIGH_RES_VARIABLES + if cadence in variable.cadences + ] + ) + + +def resolve_variable_names( + registry: Sequence[OMNIVariable], + variables: str | Iterable[str] | None, + default_names: Sequence[str], + cadence: int | None = None, +) -> list[str]: + """Resolve user input to unique canonical variable names. + + Aliases are normalized, duplicates are removed while preserving the first + occurrence, and ``"all"`` is restricted to the requested cadence. + + Raises + ------ + ValueError + If the selection is empty, not an iterable of strings, contains an + unknown name, or requests a variable unavailable at ``cadence``. + """ + + by_name = {variable.name: variable for variable in registry} + aliases = {alias: variable.name for variable in registry for alias in variable.aliases} + + if variables is None: + requested = list(default_names) + elif isinstance(variables, str): + requested = ( + [variable.name for variable in registry if cadence is None or cadence in variable.cadences] + if variables == "all" + else [variables] + ) + else: + try: + requested = list(variables) + except TypeError as error: + raise ValueError("variables must be a variable name or an iterable of variable names") from error + + if not requested: + raise ValueError("variables must contain at least one OMNI variable name") + if not all(isinstance(name, str) for name in requested): + raise ValueError("OMNI variable names must be strings") + + resolved: list[str] = [] + invalid: list[str] = [] + unavailable: list[str] = [] + for name in requested: + canonical_name = aliases.get(name, name) + variable = by_name.get(canonical_name) + if variable is None: + invalid.append(name) + continue + if cadence is not None and cadence not in variable.cadences: + unavailable.append(name) + continue + if canonical_name not in resolved: + resolved.append(canonical_name) + + if invalid: + available = ", ".join(variable.name for variable in registry if cadence is None or cadence in variable.cadences) + raise ValueError(f"Unknown OMNI variables: {', '.join(invalid)}. Available variables: {available}") + if unavailable: + raise ValueError(f"OMNI variables unavailable at {cadence}-minute cadence: {', '.join(unavailable)}") + + return resolved diff --git a/swvo/io/symh/omni.py b/swvo/io/symh/omni.py index f6b761c..01589a9 100644 --- a/swvo/io/symh/omni.py +++ b/swvo/io/symh/omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -8,17 +9,11 @@ from __future__ import annotations -import logging -from datetime import datetime, timedelta +from datetime import datetime import pandas as pd from swvo.io.omni import OMNIHighRes -from swvo.io.utils import enforce_utc_timezone - -logger = logging.getLogger(__name__) - -logging.captureWarnings(True) class SymhOMNI(OMNIHighRes): @@ -27,7 +22,7 @@ class SymhOMNI(OMNIHighRes): Inherits the `download_and_process`, other private methods and attributes from OMNIHighRes. """ - def read( + def read( # ty: ignore[invalid-method-override] self, start_time: datetime, end_time: datetime, @@ -54,23 +49,17 @@ def read( OMNI SYM-H data. """ - if start_time > end_time: - msg = "start_time must be before end_time" - logger.error(msg) - raise ValueError(msg) - - start_time = enforce_utc_timezone(start_time) - end_time = enforce_utc_timezone(end_time) - data_out = super().read(start_time, end_time, cadence_min=cadence_min, download=download) + data_out = super().read( + start_time, + end_time, + cadence_min=cadence_min, + download=download, + variables="sym-h", + ) symh_df = pd.DataFrame(index=data_out.index) symh_df["sym-h"] = data_out["sym-h"] symh_df["file_name"] = data_out["file_name"] - symh_df = symh_df.truncate( - before=start_time - timedelta(minutes=cadence_min - 0.0000001), - after=end_time + timedelta(minutes=cadence_min + 0.0000001), - ) - return symh_df diff --git a/tests/io/dst/test_dst_omni.py b/tests/io/dst/test_dst_omni.py index 1f985f5..7213c71 100644 --- a/tests/io/dst/test_dst_omni.py +++ b/tests/io/dst/test_dst_omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -11,6 +12,7 @@ import pytest from swvo.io.dst import DSTOMNI +from swvo.io.omni import OMNILowRes TEST_DIR = os.path.dirname(__file__) DATA_DIR = Path(os.path.join(TEST_DIR, "data/")) @@ -78,19 +80,19 @@ def test_read_without_download(self, dstomni): assert all(df["file_name"].isnull()) def test_read_with_download(self, dstomni, mock_dstomni_data, mocker): - mocker.patch("pathlib.Path.exists", return_value=False) - mocker.patch.object(dstomni, "_read_single_file", return_value=mock_dstomni_data) - mocker.patch.object(dstomni, "download_and_process") + parent_result = mock_dstomni_data.loc[:, ["dst", "file_name"]] + parent_read = mocker.patch.object(OMNILowRes, "read", return_value=parent_result) start_time = datetime(2020, 1, 1) end_time = datetime(2020, 12, 31) df = dstomni.read(start_time, end_time, download=True) - dstomni.download_and_process.assert_called_once() + parent_read.assert_called_once_with(start_time, end_time, download=True, variables="dst") assert not df.empty assert all(df["dst"] == 150.0) assert "dst" in df.columns + assert df.index.name == "t" assert all(idx.tzinfo is not None for idx in df.index) assert all(idx.tzinfo is timezone.utc for idx in df.index) diff --git a/tests/io/omni/test_derived_compatibility.py b/tests/io/omni/test_derived_compatibility.py new file mode 100644 index 0000000..9034f6f --- /dev/null +++ b/tests/io/omni/test_derived_compatibility.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: 2026 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel +# +# SPDX-License-Identifier: Apache-2.0 + +from datetime import datetime, timezone + +import pandas as pd +import pytest + +from swvo.io.dst import DSTOMNI +from swvo.io.f10_7 import F107OMNI +from swvo.io.kp import KpOMNI +from swvo.io.omni import OMNILowRes +from swvo.io.symh import SymhOMNI + + +@pytest.mark.parametrize("reader_class", [DSTOMNI, F107OMNI, KpOMNI, SymhOMNI]) +@pytest.mark.parametrize( + "start,end", + [ + (datetime(2020, 1, 1), datetime(2020, 1, 1, tzinfo=timezone.utc)), + (datetime(2020, 1, 2), datetime(2020, 1, 1, tzinfo=timezone.utc)), + ], +) +def test_focused_omni_readers_reject_equal_or_reversed_mixed_timezone_ranges( + tmp_path, + reader_class, + start, + end, +): + reader = reader_class(data_dir=tmp_path) + + with pytest.raises(ValueError, match="start_time must be before end_time"): + reader.read(start, end, download=False) + + +@pytest.mark.parametrize( + "reader_class,variable,expected", + [ + ( + DSTOMNI, + "dst", + pd.date_range("2020-01-02 01:00", "2020-01-02 05:00", freq="h", tz="UTC"), + ), + ( + KpOMNI, + "kp", + pd.date_range("2020-01-02 00:00", "2020-01-02 06:00", freq="3h", tz="UTC"), + ), + ( + F107OMNI, + "f107", + pd.date_range("2020-01-02 00:00", "2020-01-03 00:00", freq="24h", tz="UTC"), + ), + ], +) +def test_hourly_focused_readers_preserve_boundary_padding_in_parent( + tmp_path, + mocker, + reader_class, + variable, + expected, +): + source_path = tmp_path / "complete.csv" + source_index = pd.date_range("2020-01-01", "2020-01-04", freq="h", tz="UTC") + pd.DataFrame({"timestamp": source_index, variable: 1.0}).to_csv(source_path, index=False) + + reader = reader_class(data_dir=tmp_path) + mocker.patch.object( + reader, + "_get_processed_file_list", + return_value=([source_path], [(datetime(2020, 1, 1), datetime(2021, 1, 1))]), + ) + + result = reader.read( + datetime(2020, 1, 2, 1, 30), + datetime(2020, 1, 2, 4, 30, tzinfo=timezone.utc), + download=False, + ) + + pd.testing.assert_index_equal(result.index, expected, check_names=False) + assert result.columns.tolist() == [variable, "file_name"] + assert result[variable].eq(1.0).all() + assert result["file_name"].eq(source_path).all() + + +@pytest.mark.parametrize( + "reader_class,variable", + [(DSTOMNI, "dst"), (KpOMNI, "kp"), (F107OMNI, "f107")], +) +def test_hourly_focused_readers_delegate_raw_times_to_parent( + tmp_path, + mocker, + reader_class, + variable, +): + start = datetime(2020, 1, 1) + end = datetime(2020, 1, 2, tzinfo=timezone.utc) + parent_result = pd.DataFrame( + {variable: [1.0], "file_name": [tmp_path / "source.csv"]}, + index=pd.DatetimeIndex(["2020-01-01T00:00:00Z"]), + ) + parent_read = mocker.patch.object(OMNILowRes, "read", return_value=parent_result) + + reader_class(data_dir=tmp_path).read(start, end, download=True) + + parent_read.assert_called_once_with(start, end, download=True, variables=variable) diff --git a/tests/io/omni/test_omni_high_res.py b/tests/io/omni/test_omni_high_res.py index 518f7fa..6d984d4 100644 --- a/tests/io/omni/test_omni_high_res.py +++ b/tests/io/omni/test_omni_high_res.py @@ -1,18 +1,21 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 import os import shutil from concurrent.futures import Future -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch import numpy as np import pandas as pd import pytest +import requests from swvo.io.omni.omni_high_res import OMNIHighRes +from swvo.io.omni.variables import HIGH_RES_DEFAULT_VARIABLES, HIGH_RES_VARIABLES TEST_DIR = os.path.dirname(__file__) DATA_DIR = Path(os.path.join(TEST_DIR, "data/OMNI")) @@ -37,14 +40,23 @@ def test_initialization_without_env_var(self): with pytest.raises(ValueError): OMNIHighRes() - def test_download_and_process(self, omni_high_res): + def test_download_and_process(self, tmp_path, mocker): + omni_high_res = OMNIHighRes(data_dir=tmp_path) start_time = datetime(2020, 1, 1, tzinfo=timezone.utc) - end_time = datetime(2020, 12, 31, tzinfo=timezone.utc) - # download this file without mocking + end_time = datetime(2020, 1, 31, tzinfo=timezone.utc) + all_names = [variable.name for variable in HIGH_RES_VARIABLES if 1 in variable.cadences] + processed = pd.DataFrame( + {name: [1.0] for name in all_names}, + index=pd.DatetimeIndex(["2020-01-01"], tz="UTC", name="timestamp"), + ) + mocker.patch.object(omni_high_res, "_get_data_from_omni", return_value=[]) + mocker.patch.object(omni_high_res, "_process_single_month", return_value=processed) + omni_high_res.download_and_process(start_time, end_time) - for file in (DATA_DIR / "OMNI").glob("OMNI_HIGH_RES_1min_2020*.csv"): - assert file.exists() + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + assert output.exists() + assert all(name in pd.read_csv(output, nrows=0).columns for name in all_names) def test_read_without_download(self, omni_high_res, mocker): start_time = datetime(2021, 1, 1, tzinfo=timezone.utc) @@ -54,25 +66,23 @@ def test_read_without_download(self, omni_high_res, mocker): ): # value error is raised when no files are found hence no concatenation is possible omni_high_res.read(start_time, end_time, download=False) - def test_read_with_download(self, omni_high_res, mocker): + def test_read_with_download(self, tmp_path, mocker): + omni_high_res = OMNIHighRes(data_dir=tmp_path) start_time = datetime(2022, 1, 1, tzinfo=timezone.utc) end_time = datetime(2022, 2, 28, tzinfo=timezone.utc) + index = pd.DatetimeIndex(["2022-01-01"], tz="UTC", name="timestamp") - mocker.patch.object( - omni_high_res, - "download_and_process", - wraps=omni_high_res.download_and_process, - ) - mocker.patch.object( - omni_high_res, - "_read_single_file", - wraps=omni_high_res._read_single_file, - ) + def create_month(file_path, time_interval, cadence_min): + file_path.parent.mkdir(parents=True, exist_ok=True) + timestamp = pd.DatetimeIndex([time_interval[0]], name="timestamp") + pd.DataFrame({name: [1.0] for name in HIGH_RES_DEFAULT_VARIABLES}, index=timestamp).to_csv(file_path) - omni_high_res.read(start_time, end_time, download=True) - omni_high_res.download_and_process.assert_called_once_with(start_time, end_time, cadence_min=1) + download_one = mocker.patch.object(omni_high_res, "_download_and_process_single_file", side_effect=create_month) + result = omni_high_res.read(start_time, end_time, download=True) - assert omni_high_res._read_single_file.call_count == 2 + assert download_one.call_count == 2 + assert list(result.columns) == [*HIGH_RES_DEFAULT_VARIABLES, "file_name"] + assert index[0] in result.index def test_download_and_process_calls_get_data_per_month(self, omni_high_res, mocker): start_time = datetime(2023, 1, 1, tzinfo=timezone.utc) @@ -169,18 +179,29 @@ def test_start_year_behind(self, omni_high_res, mocker): mocker.patch.object(pd.DataFrame, "truncate", return_value=pd.DataFrame()) with patch("logging.Logger.warning") as mock_warning: - dfs = omni_high_res.read(start_time, end_time) + with pytest.raises(ValueError, match="No OMNI"): + omni_high_res.read(start_time, end_time) mock_warning.assert_any_call( "Start date chosen falls behind the existing data. Moving start date to first available mission files..." ) - assert len(dfs) == 0, "Expected dfs list to be empty since no files are found." - - def test_year_transition(self, omni_high_res): + def test_year_transition(self, tmp_path, mocker): + omni_high_res = OMNIHighRes(data_dir=tmp_path) start_time = datetime(2012, 12, 31, 23, 59, 0, tzinfo=timezone.utc) - end_time = datetime(2012, 12, 31, 23, 59, 59, tzinfo=timezone.utc) + def create_month(file_path, time_interval, cadence_min): + file_path.parent.mkdir(parents=True, exist_ok=True) + timestamp = ( + datetime(2012, 12, 31, 23, 59, tzinfo=timezone.utc) + if time_interval[0].month == 12 + else datetime(2013, 1, 1, tzinfo=timezone.utc) + ) + index = pd.DatetimeIndex([timestamp], name="timestamp") + pd.DataFrame({name: 1.0 for name in HIGH_RES_DEFAULT_VARIABLES}, index=index).to_csv(file_path) + + mocker.patch.object(omni_high_res, "_download_and_process_single_file", side_effect=create_month) + result_df = omni_high_res.read(start_time, end_time, download=True) assert result_df.index.min() == pd.Timestamp("2012-12-31 23:59:00+00:00") @@ -188,9 +209,11 @@ def test_year_transition(self, omni_high_res): def test_process_single_month_parses_data_correctly(self, omni_high_res): data = [ + "19 Vx Velocity,km/s", + "20 Vy Velocity, km/s", "YYYY DOY HR MN bavg bx_gsm by_gsm bz_gsm speed proton_density temperature pdyn sym-h", "2020 1 0 0 5.1 1.2 2.3 3.4 400 5.5 1000000 99 -15", - "2020 1 0 1 9999.9 9999.9 9999.9 9999.9 99999.8 999.8 9999998.0 99 99999.0", + "2020 1 0 1 9999.99 9999.99 9999.99 9999.99 99999.9 999.99 9999999.0 99.99 99999.0", ] df = omni_high_res._process_single_month(data) @@ -232,9 +255,452 @@ def test_process_single_month_handles_missing_data_lines(self, omni_high_res): def test_process_single_month_raises_on_missing_header(self, omni_high_res): data = ["2020 1 0 0 5.1 1.2 2.3 3.4 400 5.5 1000000 -15"] - with pytest.raises(StopIteration): + with pytest.raises(ValueError, match="expected YYYY data header"): omni_high_res._process_single_month(data) + def test_available_variables_cover_both_cadences(self, omni_high_res): + one_minute = omni_high_res.available_variables(cadence=1) + five_minute = omni_high_res.available_variables(cadence=5) + + assert len(one_minute) == 42 + assert len(five_minute) == 45 + assert one_minute["nasa_id"].tolist() == list(range(4, 46)) + assert five_minute["nasa_id"].tolist() == list(range(4, 49)) + assert {"name", "nasa_id", "description", "unit", "fill_value", "cadences", "aliases"} == set( + one_minute.columns + ) + assert dict(zip(five_minute["name"], five_minute["unit"])) == { + variable.name: variable.unit for variable in HIGH_RES_VARIABLES + } + assert one_minute.set_index("name").loc["bavg", "fill_value"] == 9999.99 + + def test_available_variables_delegates_to_shared_utility(self, omni_high_res, mocker): + expected = pd.DataFrame({"name": ["speed"]}) + utility = mocker.patch("swvo.io.omni.omni_high_res.get_available_variables", return_value=expected) + + result = omni_high_res.available_variables(cadence=5) + + assert result is expected + utility.assert_called_once_with(5) + + def test_available_variables_validates_cadence_before_delegating(self, omni_high_res, mocker): + utility = mocker.patch("swvo.io.omni.omni_high_res.get_available_variables") + + with pytest.raises(AssertionError, match="Only 1 or 5"): + omni_high_res.available_variables(cadence=60) + + utility.assert_not_called() + + def test_private_helpers_are_instance_methods(self): + assert not isinstance(OMNIHighRes.__dict__["available_variables"], classmethod) + assert not isinstance(OMNIHighRes.__dict__["_validate_cadence"], staticmethod) + assert not isinstance(OMNIHighRes.__dict__["_cache_contains"], staticmethod) + + @pytest.mark.parametrize("cadence,expected_count", [(1, 42), (5, 45)]) + def test_processes_complete_cadence_schema(self, omni_high_res, cadence, expected_count): + names = [variable.name for variable in HIGH_RES_VARIABLES if cadence in variable.cadences] + data = [ + "YYYY DOY HR MN " + " ".join(str(index) for index in range(1, expected_count + 1)), + "2020 1 0 0 " + " ".join("1" for _ in names), + ] + + result = omni_high_res._process_single_month(data, cadence_min=cadence, variable_names=names) + + assert list(result.columns) == names + assert len(result.columns) == expected_count + assert result.index.tz is not None + + def test_complete_request_uses_nasa_ids_in_registry_order(self, omni_high_res, mocker): + response = mocker.Mock() + response.text = "YYYY DOY HR MN\n" + response.raise_for_status = mocker.Mock() + post = mocker.patch("requests.post", return_value=response) + names = [variable.name for variable in HIGH_RES_VARIABLES if 5 in variable.cadences] + + omni_high_res._get_data_from_omni( + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + cadence=5, + variable_names=names, + ) + + payload = post.call_args.kwargs["data"] + assert payload["vars"] == [str(identifier) for identifier in range(4, 49)] + assert payload["res"] == "5min" + assert payload["spacecraft"] == "omni_5min" + assert post.call_args.kwargs["timeout"] == 30 + + @pytest.mark.parametrize("cadence", [1, 5]) + def test_fill_values_are_masked_for_every_cadence_field(self, omni_high_res, cadence): + variables = [variable for variable in HIGH_RES_VARIABLES if cadence in variable.cadences] + row = " ".join(str(variable.fill_value) for variable in variables) + data = ["YYYY DOY HR MN", f"2020 1 0 0 {row}"] + + result = omni_high_res._process_single_month( + data, + cadence_min=cadence, + variable_names=[variable.name for variable in variables], + ) + + assert result.isna().all().all() + + def test_selection_preserves_order_deduplicates_and_accepts_aliases(self, tmp_path): + omni_high_res = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + pd.DataFrame( + {"speed": [400.0], "sym-h": [-10.0]}, + index=pd.DatetimeIndex(["2020-01-01"], tz="UTC", name="timestamp"), + ).to_csv(output) + + result = omni_high_res.read( + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 1, 0, 1, tzinfo=timezone.utc), + variables=["sym_h", "speed", "sym-h"], + ) + + assert list(result.columns) == ["sym-h", "speed", "file_name"] + assert result["file_name"].notna().all() + + def test_rejects_unknown_and_cadence_incompatible_variables(self, omni_high_res): + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 2, tzinfo=timezone.utc) + + with pytest.raises(ValueError, match="Unknown OMNI variables"): + omni_high_res.read(start, end, variables="not_a_variable") + with pytest.raises(ValueError, match="unavailable at 1-minute cadence"): + omni_high_res.read(start, end, cadence_min=1, variables="proton_flux_10_mev") + with pytest.raises(ValueError, match="at least one"): + omni_high_res.read(start, end, variables=[]) + + def test_partial_cache_requires_download_or_is_upgraded(self, tmp_path, mocker): + omni_high_res = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + index = pd.DatetimeIndex(["2020-01-01"], tz="UTC", name="timestamp") + pd.DataFrame({"speed": [400.0]}, index=index).to_csv(output) + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 1, 0, 1, tzinfo=timezone.utc) + + with pytest.raises(ValueError, match="does not contain: ae"): + omni_high_res.read(start, end, variables="ae") + + def upgrade(file_path, _time_interval, _cadence_min): + pd.DataFrame({"speed": [400.0], "ae": [25.0]}, index=index).to_csv(file_path) + + upgrade_file = mocker.patch.object( + omni_high_res, + "_download_and_process_single_file", + side_effect=upgrade, + ) + result = omni_high_res.read(start, end, download=True, variables="ae") + + upgrade_file.assert_called_once() + assert result["ae"].iloc[0] == 25.0 + + def test_failed_processing_preserves_existing_cache_and_removes_temporary_file(self, tmp_path, mocker): + omni_high_res = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + output.write_text("original-cache\n") + mocker.patch.object(omni_high_res, "_get_data_from_omni", return_value=["invalid"]) + mocker.patch.object(omni_high_res, "_process_single_month", side_effect=ValueError("bad response")) + + omni_high_res._download_and_process_single_file( + output, + (datetime(2020, 1, 1), datetime(2020, 1, 31)), + 1, + ) + + assert output.read_text() == "original-cache\n" + assert not output.with_suffix(".csv.tmp").exists() + + @pytest.mark.parametrize("method", ["read", "download_and_process"]) + @pytest.mark.parametrize( + "start,end", + [ + (datetime(2020, 1, 1), datetime(2020, 1, 1)), + (datetime(2020, 1, 2), datetime(2020, 1, 1)), + ], + ) + def test_equal_and_reversed_ranges_raise_value_error(self, omni_high_res, method, start, end): + with pytest.raises(ValueError, match="start_time must be before end_time"): + getattr(omni_high_res, method)(start, end) + + def test_mixed_naive_and_aware_times_are_normalized_before_comparison(self, tmp_path): + reader = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + pd.DataFrame( + {name: [1.0] for name in HIGH_RES_DEFAULT_VARIABLES}, + index=pd.DatetimeIndex(["2020-01-01T00:00:00Z"], name="timestamp"), + ).to_csv(output) + + result = reader.read( + datetime(2020, 1, 1), + datetime(2020, 1, 1, 0, 1, tzinfo=timezone.utc), + ) + + assert str(result.index.tz) == "UTC" + + @pytest.mark.parametrize( + "doy,hour,minute,error", + [ + (0, 0, 0, "invalid day-of-year, hour, or minute"), + (367, 0, 0, "invalid day-of-year, hour, or minute"), + (1, 24, 0, "invalid day-of-year, hour, or minute"), + (1, 0, 60, "invalid day-of-year, hour, or minute"), + (366, 0, 0, "invalid day-of-year for its record year"), + ], + ) + def test_rejects_invalid_time_fields(self, omni_high_res, doy, hour, minute, error): + row = " ".join("1" for _ in HIGH_RES_DEFAULT_VARIABLES) + data = ["YYYY DOY HR MN", f"2021 {doy} {hour} {minute} {row}"] + + with pytest.raises(ValueError, match=error): + omni_high_res._process_single_month(data) + + def test_rejects_malformed_width_and_non_numeric_values(self, omni_high_res): + with pytest.raises(ValueError, match="record widths"): + omni_high_res._process_single_month(["YYYY DOY HR MN", "2020 1 0 0 1 2"]) + + values = ["1" for _ in HIGH_RES_DEFAULT_VARIABLES] + values[2] = "not-a-number" + with pytest.raises(ValueError, match="Unable to parse|string"): + omni_high_res._process_single_month(["YYYY DOY HR MN", f"2020 1 0 0 {' '.join(values)}"]) + + def test_missing_header_has_public_value_error(self, omni_high_res): + with pytest.raises(ValueError, match="expected YYYY data header"): + omni_high_res._process_single_month(["2020 1 0 0 " + " ".join("1" for _ in range(9))]) + + def test_empty_response_builds_utc_cadence_grid(self, omni_high_res): + result = omni_high_res._process_single_month( + [], + original_end=datetime(2020, 1, 1, 0, 10), + cadence_min=5, + ) + + assert len(result) == 3 + assert result.index.tolist() == list(pd.date_range("2020-01-01", periods=3, freq="5min", tz="UTC")) + assert result.isna().all().all() + + def test_empty_response_without_end_preserves_requested_schema(self, omni_high_res): + result = omni_high_res._process_single_month([], variable_names=["speed", "sym_h"]) + + assert result.empty + assert list(result.columns) == ["speed", "sym-h"] + + def test_values_below_fill_sentinel_are_preserved(self, omni_high_res): + variables = [variable for variable in HIGH_RES_VARIABLES if 1 in variable.cadences] + values = [float(variable.fill_value) - 0.01 for variable in variables] + data = ["YYYY DOY HR MN", "2020 1 0 0 " + " ".join(map(str, values))] + + result = omni_high_res._process_single_month( + data, + cadence_min=1, + variable_names=[variable.name for variable in variables], + ) + + assert result.notna().all().all() + + def test_provenance_depends_only_on_selected_fields(self, tmp_path): + reader = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + index = pd.DatetimeIndex(["2020-01-01T00:00Z", "2020-01-01T00:01Z"], name="timestamp") + pd.DataFrame({"speed": [np.nan, 400.0], "ae": [20.0, np.nan]}, index=index).to_csv(output) + + result = reader.read( + index[0].to_pydatetime(), (index[1] + timedelta(minutes=1)).to_pydatetime(), variables="speed" + ) + + assert pd.isna(result["file_name"].iloc[0]) + assert result["file_name"].iloc[1] == output + + def test_complete_cache_is_skipped_but_partial_and_reprocess_are_downloaded(self, tmp_path, mocker): + reader = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + all_names = [variable.name for variable in HIGH_RES_VARIABLES if 1 in variable.cadences] + pd.DataFrame(columns=all_names).to_csv(output) + download = mocker.patch.object(reader, "_download_and_process_single_file") + start = datetime(2020, 1, 1) + end = datetime(2020, 1, 31) + + reader.download_and_process(start, end) + download.assert_not_called() + + pd.DataFrame(columns=HIGH_RES_DEFAULT_VARIABLES).to_csv(output) + reader.download_and_process(start, end) + assert download.call_count == 1 + + download.reset_mock() + pd.DataFrame(columns=all_names).to_csv(output) + reader.download_and_process(start, end, reprocess_files=True) + assert download.call_count == 1 + + def test_corrupt_cache_is_reported_or_replaced(self, tmp_path, mocker): + reader = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + output.write_bytes(b"\xff\xfe\x00broken") + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 1, 0, 1, tzinfo=timezone.utc) + + with pytest.raises(ValueError, match="Cannot read processed OMNI file"): + reader.read(start, end, variables="speed") + + def replace_cache(file_path, _interval, _cadence): + pd.DataFrame( + {"speed": [401.0]}, + index=pd.DatetimeIndex([start], name="timestamp"), + ).to_csv(file_path) + + replacement = mocker.patch.object(reader, "_download_and_process_single_file", side_effect=replace_cache) + result = reader.read(start, end, download=True, variables="speed") + + replacement.assert_called_once() + assert result["speed"].iloc[0] == 401.0 + + def test_unreadable_cache_after_failed_upgrade_has_clear_error(self, tmp_path, mocker): + reader = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + output.write_bytes(b"\xff\xfe") + mocker.patch.object(reader, "_download_and_process_single_file") + + with pytest.raises(ValueError, match="remains unreadable after attempted cache upgrade"): + reader.read( + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 1, 0, 1, tzinfo=timezone.utc), + download=True, + variables="speed", + ) + + def test_html_wrapped_correct_range_retries_with_suggested_end(self, omni_high_res, mocker): + first = mocker.Mock() + first.text = "\n

Error

\nCorrect range: 20200101 - 20200102" + first.raise_for_status = mocker.Mock() + second = mocker.Mock() + second.text = "YYYY DOY HR MN" + second.raise_for_status = mocker.Mock() + post = mocker.patch("requests.post", side_effect=[first, second]) + + result = omni_high_res._get_data_from_omni( + datetime(2020, 1, 1), + datetime(2020, 1, 31), + ) + + assert result == ["YYYY DOY HR MN"] + assert post.call_count == 2 + assert post.call_args_list[1].kwargs["data"]["end_date"] == "20200102" + + def test_correct_range_before_start_returns_empty(self, omni_high_res, mocker): + response = mocker.Mock() + response.text = "\nError\ncorrect range: 20190101 - 20191231" + response.raise_for_status = mocker.Mock() + mocker.patch("requests.post", return_value=response) + + result = omni_high_res._get_data_from_omni(datetime(2020, 1, 1), datetime(2020, 1, 2)) + + assert result == [] + + def test_unspecified_html_error_and_http_error_propagate(self, omni_high_res, mocker): + response = mocker.Mock() + response.text = "\n

Error while processing request

" + response.raise_for_status = mocker.Mock() + mocker.patch("requests.post", return_value=response) + with pytest.raises(ValueError, match="unspecified error"): + omni_high_res._get_data_from_omni(datetime(2020, 1, 1), datetime(2020, 1, 2)) + + response.raise_for_status.side_effect = requests.HTTPError("503") + with pytest.raises(requests.HTTPError): + omni_high_res._get_data_from_omni(datetime(2020, 1, 1), datetime(2020, 1, 2)) + + def test_parallel_worker_exception_is_not_silenced(self, tmp_path, mocker): + reader = OMNIHighRes(data_dir=tmp_path) + mocker.patch.object(reader, "_download_and_process_single_file", side_effect=RuntimeError("worker failed")) + + with pytest.raises(RuntimeError, match="worker failed"): + reader.download_and_process(datetime(2020, 1, 1), datetime(2020, 12, 31)) + + def test_range_entirely_before_mission_start_has_clear_error(self, omni_high_res): + with pytest.raises(ValueError, match="ends before OMNI data begin in 1981"): + omni_high_res.read(datetime(1970, 1, 1), datetime(1971, 1, 1)) + + @pytest.mark.parametrize( + "cadence,end,expected_months", + [ + (1, datetime(2020, 1, 31, 23, 58, 59, tzinfo=timezone.utc), ["202001"]), + (1, datetime(2020, 1, 31, 23, 59, tzinfo=timezone.utc), ["202001", "202002"]), + (5, datetime(2020, 1, 31, 23, 55, tzinfo=timezone.utc), ["202001", "202002"]), + ], + ) + def test_month_boundary_neighbor_selection(self, tmp_path, cadence, end, expected_months): + reader = OMNIHighRes(data_dir=tmp_path) + + paths, intervals = reader._get_processed_file_list( + datetime(2020, 1, 1, tzinfo=timezone.utc), + end, + cadence, + ) + + assert [path.stem[-6:] for path in paths] == expected_months + assert all(interval[0].tzinfo is not None and interval[1].tzinfo is not None for interval in intervals) + + def test_corrected_range_that_does_not_shorten_request_is_rejected(self, omni_high_res, mocker): + response = mocker.Mock() + response.text = "\nError\ncorrect range: 20200101 - 20200131" + response.raise_for_status = mocker.Mock() + post = mocker.patch("requests.post", return_value=response) + + with pytest.raises(ValueError, match="would not shorten"): + omni_high_res._get_data_from_omni(datetime(2020, 1, 1), datetime(2020, 1, 31)) + + post.assert_called_once() + + def test_successful_reprocessing_replaces_cache_atomically(self, tmp_path, mocker): + reader = OMNIHighRes(data_dir=tmp_path) + output = tmp_path / "2020/OMNI_HIGH_RES_1min_202001.csv" + output.parent.mkdir(parents=True) + output.write_text("old-cache\n") + all_names = [variable.name for variable in HIGH_RES_VARIABLES if 1 in variable.cadences] + processed = pd.DataFrame( + {name: [1.0] for name in all_names}, + index=pd.DatetimeIndex(["2020-01-01T00:00Z"], name="timestamp"), + ) + mocker.patch.object(reader, "_get_data_from_omni", return_value=[]) + mocker.patch.object(reader, "_process_single_month", return_value=processed) + + reader._download_and_process_single_file( + output, + (datetime(2020, 1, 1, tzinfo=timezone.utc), datetime(2020, 1, 31, tzinfo=timezone.utc)), + 1, + ) + + assert pd.read_csv(output)["speed"].iloc[0] == 1.0 + assert not output.with_suffix(".csv.tmp").exists() + + @pytest.mark.skipif( + os.environ.get("SWVO_RUN_LIVE_TESTS") != "1", + reason="set SWVO_RUN_LIVE_TESTS=1 for the one-day NASA OMNIWeb smoke check", + ) + @pytest.mark.parametrize("cadence", [1, 5]) + def test_live_one_day_complete_omniweb_response(self, tmp_path, cadence): + omni_high_res = OMNIHighRes(data_dir=tmp_path) + names = [variable.name for variable in HIGH_RES_VARIABLES if cadence in variable.cadences] + raw = omni_high_res._get_data_from_omni( + datetime(2024, 1, 1, tzinfo=timezone.utc), + datetime(2024, 1, 1, tzinfo=timezone.utc), + cadence=cadence, + variable_names=names, + ) + + result = omni_high_res._process_single_month(raw, cadence_min=cadence, variable_names=names) + + assert list(result.columns) == names + assert not result.empty + def test_remove_processed_file(self): shutil.rmtree(Path(TEST_DIR) / "data/OMNI/2022", ignore_errors=True) shutil.rmtree(Path(TEST_DIR) / "data/OMNI/2023", ignore_errors=True) diff --git a/tests/io/omni/test_omni_low_res.py b/tests/io/omni/test_omni_low_res.py index 0f731e4..ecb5668 100644 --- a/tests/io/omni/test_omni_low_res.py +++ b/tests/io/omni/test_omni_low_res.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 import os import warnings -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch @@ -12,6 +13,7 @@ import pytest from swvo.io.omni.omni_low_res import OMNILowRes +from swvo.io.omni.variables import LOW_RES_DEFAULT_VARIABLES, LOW_RES_VARIABLES TEST_DIR = os.path.dirname(__file__) DATA_DIR = Path(os.path.join(TEST_DIR, "data/")) @@ -89,6 +91,9 @@ def test_process_single_file(self, omni_low_res): assert isinstance(df, pd.DataFrame) assert all(column in df.columns for column in ["kp", "dst", "f107"]) + assert len(df.columns) == 54 + assert df["lyman_alpha"].isna().all() + assert df["proton_quasi_invariant"].isna().all() assert len(df) > 0 def test_read_single_file(self, omni_low_res): @@ -124,5 +129,321 @@ def test_start_year_behind(self, omni_low_res, mocker): assert all(df["dst"].isna()) assert all(df["file_name"].isnull()) + def test_available_variables_describe_all_non_time_fields(self, omni_low_res): + variables = omni_low_res.available_variables() + + assert len(variables) == 54 + assert variables["name"].tolist() == [variable.name for variable in LOW_RES_VARIABLES] + assert {"name", "description", "unit", "fill_value", "aliases"} == set(variables.columns) + assert dict(zip(variables["name"], variables["unit"])) == { + variable.name: variable.unit for variable in LOW_RES_VARIABLES + } + + def test_available_variables_delegates_to_shared_utility(self, omni_low_res, mocker): + expected = pd.DataFrame({"name": ["dst"]}) + utility = mocker.patch("swvo.io.omni.omni_low_res.get_available_variables", return_value=expected) + + result = omni_low_res.available_variables() + + assert result is expected + utility.assert_called_once_with() + + def test_private_cache_helper_is_an_instance_method(self): + assert not isinstance(OMNILowRes.__dict__["available_variables"], classmethod) + assert not isinstance(OMNILowRes.__dict__["_cache_contains"], staticmethod) + + def test_fill_values_are_masked_for_every_hourly_field(self, omni_low_res, tmp_path): + values = ["0" if variable.fill_value is None else str(variable.fill_value) for variable in LOW_RES_VARIABLES] + raw_file = tmp_path / "omni2_fill_values.dat" + raw_file.write_text(" ".join(["2020", "1", "0", *values]) + "\n") + + result = omni_low_res._process_single_file(raw_file) + + assert result.drop(columns="magnetospheric_flux_flag").isna().all().all() + assert result["magnetospheric_flux_flag"].iloc[0] == 0 + + def test_processes_current_57_word_record(self, omni_low_res, tmp_path): + historic_line = (Path(TEST_DIR) / "data/omni2_2020.dat").read_text().splitlines()[0] + current_file = tmp_path / "omni2_current.dat" + current_file.write_text(f"{historic_line} 0.005123 0.4567\n") + + result = omni_low_res._process_single_file(current_file) + + assert len(result.columns) == 54 + assert result["lyman_alpha"].iloc[0] == pytest.approx(0.005123) + assert result["proton_quasi_invariant"].iloc[0] == pytest.approx(0.4567) + + def test_all_and_subset_selection_preserve_public_schema(self, tmp_path): + omni_low_res = OMNILowRes(data_dir=tmp_path) + processed = omni_low_res._process_single_file(Path(TEST_DIR) / "data/omni2_2020.dat") + output = tmp_path / "OMNI_LOW_RES_2020.csv" + processed.to_csv(output, index=True) + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 1, 1, tzinfo=timezone.utc) + + default = omni_low_res.read(start, end) + all_variables = omni_low_res.read(start, end, variables="all") + subset = omni_low_res.read(start, end, variables=["pc", "speed", "pcn"]) + + assert list(default.columns) == [*LOW_RES_DEFAULT_VARIABLES, "file_name"] + assert list(all_variables.columns) == [*[variable.name for variable in LOW_RES_VARIABLES], "file_name"] + assert list(subset.columns) == ["pcn", "speed", "file_name"] + assert all_variables.index.tz is not None + assert subset["file_name"].notna().any() + + def test_rejects_unknown_and_empty_variable_selection(self, omni_low_res): + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 2, tzinfo=timezone.utc) + + with pytest.raises(ValueError, match="Unknown OMNI variables"): + omni_low_res.read(start, end, variables="not_a_variable") + with pytest.raises(ValueError, match="at least one"): + omni_low_res.read(start, end, variables=[]) + + def test_partial_cache_requires_download_or_is_upgraded(self, tmp_path, mocker): + omni_low_res = OMNILowRes(data_dir=tmp_path) + output = tmp_path / "OMNI_LOW_RES_2020.csv" + index = pd.DatetimeIndex(["2020-01-01"], tz="UTC", name="timestamp") + pd.DataFrame({"dst": [-5.0], "kp": [1.0], "f107": [70.0]}, index=index).to_csv(output) + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 1, 1, tzinfo=timezone.utc) + + with pytest.raises(ValueError, match="does not contain: speed"): + omni_low_res.read(start, end, variables="speed") + + def upgrade(_start, _end, reprocess_files=False): + assert reprocess_files + pd.DataFrame({"speed": [400.0]}, index=index).to_csv(output) + + upgrade_cache = mocker.patch.object(omni_low_res, "download_and_process", side_effect=upgrade) + result = omni_low_res.read(start, end, download=True, variables="speed") + + upgrade_cache.assert_called_once() + assert result["speed"].iloc[0] == 400.0 + + def test_failed_processing_preserves_existing_cache_and_removes_temporary_file(self, tmp_path, mocker): + omni_low_res = OMNILowRes(data_dir=tmp_path) + output = tmp_path / "OMNI_LOW_RES_2020.csv" + output.write_text("original-cache\n") + mocker.patch.object(omni_low_res, "_download") + mocker.patch.object(omni_low_res, "_process_single_file", side_effect=ValueError("bad response")) + + omni_low_res.download_and_process( + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 12, 30, tzinfo=timezone.utc), + reprocess_files=True, + ) + + assert output.read_text() == "original-cache\n" + assert not output.with_suffix(".csv.tmp").exists() + + @pytest.mark.parametrize("method", ["read", "download_and_process"]) + @pytest.mark.parametrize( + "start,end", + [ + (datetime(2020, 1, 1), datetime(2020, 1, 1)), + (datetime(2020, 1, 2), datetime(2020, 1, 1)), + ], + ) + def test_equal_and_reversed_ranges_raise_value_error(self, omni_low_res, method, start, end): + with pytest.raises(ValueError, match="start_time must be before end_time"): + getattr(omni_low_res, method)(start, end) + + def test_mixed_naive_and_aware_times_are_normalized_before_comparison(self, tmp_path): + reader = OMNILowRes(data_dir=tmp_path) + + result = reader.read( + datetime(2020, 1, 1), + datetime(2020, 1, 1, 1, tzinfo=timezone.utc), + variables="speed", + ) + + assert str(result.index.tz) == "UTC" + assert result["speed"].isna().all() + + @pytest.mark.parametrize( + "word_count", + [54, 56, 58], + ) + def test_rejects_unsupported_record_widths(self, omni_low_res, tmp_path, word_count): + raw_file = tmp_path / f"omni2_{word_count}.dat" + raw_file.write_text(" ".join(["2020", "1", "0", *(["1"] * (word_count - 3))]) + "\n") + + with pytest.raises(ValueError, match="expected 55 historic or 57 current"): + omni_low_res._process_single_file(raw_file) + + def test_empty_and_mixed_width_sources_have_clear_errors(self, omni_low_res, tmp_path): + empty = tmp_path / "empty.dat" + empty.write_text("") + with pytest.raises(ValueError, match="Cannot parse OMNI2 source file"): + omni_low_res._process_single_file(empty) + + historic = (Path(TEST_DIR) / "data/omni2_2020.dat").read_text().splitlines()[0] + mixed = tmp_path / "mixed.dat" + mixed.write_text(f"{historic}\n{historic} 0.1 0.2\n") + with pytest.raises(ValueError, match="Cannot parse OMNI2 source file"): + omni_low_res._process_single_file(mixed) + + @pytest.mark.parametrize( + "day,hour,error", + [ + (0, 0, "invalid day-of-year or hour"), + (367, 0, "invalid day-of-year or hour"), + (1, 24, "invalid day-of-year or hour"), + (366, 0, "invalid day-of-year for its record year"), + ], + ) + def test_rejects_invalid_hourly_time_fields(self, omni_low_res, tmp_path, day, hour, error): + values = ["1" for _ in range(len(LOW_RES_VARIABLES) - 2)] + raw_file = tmp_path / "invalid_time.dat" + raw_file.write_text(" ".join(["2021", str(day), str(hour), *values]) + "\n") + + with pytest.raises(ValueError, match=error): + omni_low_res._process_single_file(raw_file) + + def test_kp_integer_tenths_are_converted_to_conventional_thirds(self, omni_low_res, tmp_path): + kp_index = 3 + next(index for index, variable in enumerate(LOW_RES_VARIABLES) if variable.name == "kp") + rows = [] + for hour, kp in enumerate([0, 3, 7, 10, 13]): + words = ["2020", "1", str(hour), *(["1"] * (len(LOW_RES_VARIABLES) - 2))] + words[kp_index] = str(kp) + rows.append(" ".join(words)) + raw_file = tmp_path / "kp.dat" + raw_file.write_text("\n".join(rows) + "\n") + + result = omni_low_res._process_single_file(raw_file) + + assert result["kp"].tolist() == pytest.approx([0, 1 / 3, 2 / 3, 1, 4 / 3]) + + def test_provenance_depends_only_on_selected_fields(self, tmp_path): + reader = OMNILowRes(data_dir=tmp_path) + output = tmp_path / "OMNI_LOW_RES_2020.csv" + index = pd.DatetimeIndex(["2020-01-01T00:00Z", "2020-01-01T01:00Z"], name="timestamp") + pd.DataFrame({"speed": [float("nan"), 400.0], "dst": [-10.0, float("nan")]}, index=index).to_csv(output) + + result = reader.read( + index[0].to_pydatetime(), + (index[1] + timedelta(hours=1)).to_pydatetime(), + variables="speed", + ) + + assert pd.isna(result.loc[index[0], "file_name"]) + assert result.loc[index[1], "file_name"] == output + + def test_complete_cache_is_skipped_but_partial_and_reprocess_are_downloaded(self, tmp_path, mocker): + reader = OMNILowRes(data_dir=tmp_path) + output = tmp_path / "OMNI_LOW_RES_2020.csv" + all_names = [variable.name for variable in LOW_RES_VARIABLES] + pd.DataFrame(columns=all_names).to_csv(output) + download = mocker.patch.object(reader, "_download") + process = mocker.patch.object(reader, "_process_single_file", return_value=pd.DataFrame(columns=all_names)) + start = datetime(2020, 1, 1) + end = datetime(2020, 12, 30) + + reader.download_and_process(start, end) + download.assert_not_called() + process.assert_not_called() + + pd.DataFrame(columns=LOW_RES_DEFAULT_VARIABLES).to_csv(output) + reader.download_and_process(start, end) + assert download.call_count == 1 + assert process.call_count == 1 + + download.reset_mock() + process.reset_mock() + pd.DataFrame(columns=all_names).to_csv(output) + reader.download_and_process(start, end, reprocess_files=True) + assert download.call_count == 1 + assert process.call_count == 1 + + def test_downloader_uses_and_cleans_os_temporary_directory(self, tmp_path, mocker): + reader = OMNILowRes(data_dir=tmp_path) + observed_directories = [] + + def download(temporary_dir, filename): + observed_directories.append(temporary_dir) + assert temporary_dir.is_dir() + (temporary_dir / filename).write_text("raw") + + mocker.patch.object(reader, "_download", side_effect=download) + mocker.patch.object( + reader, + "_process_single_file", + return_value=pd.DataFrame( + {name: [1.0] for name in (variable.name for variable in LOW_RES_VARIABLES)}, + index=pd.DatetimeIndex(["2020-01-01T00:00Z"], name="timestamp"), + ), + ) + + reader.download_and_process(datetime(2020, 1, 1), datetime(2020, 12, 30)) + + assert len(observed_directories) == 1 + assert not observed_directories[0].exists() + + def test_corrupt_cache_is_reported_or_replaced(self, tmp_path, mocker): + reader = OMNILowRes(data_dir=tmp_path) + output = tmp_path / "OMNI_LOW_RES_2020.csv" + output.write_bytes(b"\xff\xfe\x00broken") + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 1, 1, tzinfo=timezone.utc) + + with pytest.raises(ValueError, match="Cannot read processed OMNI file"): + reader.read(start, end, variables="speed") + + def replace_cache(_start, _end, reprocess_files=False): + assert reprocess_files + pd.DataFrame( + {"speed": [402.0]}, + index=pd.DatetimeIndex([start], name="timestamp"), + ).to_csv(output) + + replacement = mocker.patch.object(reader, "download_and_process", side_effect=replace_cache) + result = reader.read(start, end, download=True, variables="speed") + + replacement.assert_called_once() + assert result.loc[pd.Timestamp(start), "speed"] == 402.0 + + def test_unreadable_cache_after_failed_upgrade_has_clear_error(self, tmp_path, mocker): + reader = OMNILowRes(data_dir=tmp_path) + output = tmp_path / "OMNI_LOW_RES_2020.csv" + output.write_bytes(b"\xff\xfe") + mocker.patch.object(reader, "download_and_process") + + with pytest.raises(ValueError, match="remains unreadable after attempted cache upgrade"): + reader.read( + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 1, 1, tzinfo=timezone.utc), + download=True, + variables="speed", + ) + + def test_year_file_selection_includes_only_required_boundary_neighbor(self, tmp_path): + reader = OMNILowRes(data_dir=tmp_path) + + regular, _ = reader._get_processed_file_list( + datetime(2020, 6, 1, tzinfo=timezone.utc), + datetime(2020, 12, 31, 20, 59, tzinfo=timezone.utc), + ) + boundary, _ = reader._get_processed_file_list( + datetime(2020, 6, 1, tzinfo=timezone.utc), + datetime(2020, 12, 31, 21, 0, tzinfo=timezone.utc), + ) + + assert [path.name for path in regular] == ["OMNI_LOW_RES_2020.csv"] + assert [path.name for path in boundary] == ["OMNI_LOW_RES_2020.csv", "OMNI_LOW_RES_2021.csv"] + + def test_range_entirely_before_mission_start_has_clear_error(self, omni_low_res): + with pytest.raises(ValueError, match="ends before OMNI data begin in 1963"): + omni_low_res.read(datetime(1950, 1, 1), datetime(1951, 1, 1)) + + def test_nonnumeric_hourly_field_has_clear_error(self, omni_low_res, tmp_path): + words = ["2020", "1", "0", *(["1"] * (len(LOW_RES_VARIABLES) - 2))] + words[10] = "not-a-number" + raw_file = tmp_path / "nonnumeric.dat" + raw_file.write_text(" ".join(words) + "\n") + + with pytest.raises(ValueError, match="contains nonnumeric fields"): + omni_low_res._process_single_file(raw_file) + def test_remove_processed_file(self): os.remove(Path(TEST_DIR) / "data/OMNI_LOW_RES_2020.csv") diff --git a/tests/io/omni/test_variables.py b/tests/io/omni/test_variables.py new file mode 100644 index 0000000..c19c8ee --- /dev/null +++ b/tests/io/omni/test_variables.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: 2026 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from swvo.io.omni.variables import ( + HIGH_RES_DEFAULT_VARIABLES, + HIGH_RES_VARIABLES, + LOW_RES_DEFAULT_VARIABLES, + LOW_RES_VARIABLES, + available_variables, + resolve_variable_names, +) + + +@pytest.mark.parametrize( + "registry,expected_count", + [(HIGH_RES_VARIABLES, 45), (LOW_RES_VARIABLES, 54)], +) +def test_registry_names_and_aliases_are_unambiguous(registry, expected_count): + names = [variable.name for variable in registry] + aliases = [alias for variable in registry for alias in variable.aliases] + + assert len(registry) == expected_count + assert len(names) == len(set(names)) + assert len(aliases) == len(set(aliases)) + assert set(names).isdisjoint(aliases) + assert all(variable.description for variable in registry) + assert all(variable.unit is not None for variable in registry) + + +def test_high_resolution_registry_ids_and_cadences_are_complete(): + assert [variable.nasa_id for variable in HIGH_RES_VARIABLES] == list(range(4, 49)) + assert sum(1 in variable.cadences for variable in HIGH_RES_VARIABLES) == 42 + assert sum(5 in variable.cadences for variable in HIGH_RES_VARIABLES) == 45 + assert all(set(variable.cadences).issubset({1, 5}) for variable in HIGH_RES_VARIABLES) + + +@pytest.mark.parametrize( + "cadence,expected_count,expected_names", + [ + (None, 54, [variable.name for variable in LOW_RES_VARIABLES]), + (1, 42, [variable.name for variable in HIGH_RES_VARIABLES if 1 in variable.cadences]), + (5, 45, [variable.name for variable in HIGH_RES_VARIABLES]), + ], +) +def test_available_variables_resolves_cadence_in_registry_order(cadence, expected_count, expected_names): + metadata = available_variables(cadence) + + assert len(metadata) == expected_count + assert metadata["name"].tolist() == expected_names + + +def test_available_variables_preserves_cadence_specific_metadata_schemas(): + hourly = available_variables() + one_minute = available_variables(1) + five_minute = available_variables(5) + + assert list(hourly.columns) == ["name", "description", "unit", "fill_value", "aliases"] + assert list(one_minute.columns) == [ + "name", + "nasa_id", + "description", + "unit", + "fill_value", + "cadences", + "aliases", + ] + assert list(five_minute.columns) == list(one_minute.columns) + assert one_minute["nasa_id"].tolist() == list(range(4, 46)) + assert five_minute["nasa_id"].tolist() == list(range(4, 49)) + + +def test_available_variables_returns_independent_dataframes(): + first = available_variables(1) + second = available_variables(1) + + first.loc[0, "name"] = "modified" + + assert second.loc[0, "name"] == HIGH_RES_VARIABLES[0].name + + +@pytest.mark.parametrize("cadence", [-1, 0, 2, 60, "1", object()]) +def test_available_variables_rejects_unsupported_cadences(cadence): + with pytest.raises(ValueError, match="None.*1 or 5"): + available_variables(cadence) + + +@pytest.mark.parametrize( + "registry,defaults,cadence", + [ + (HIGH_RES_VARIABLES, HIGH_RES_DEFAULT_VARIABLES, 1), + (LOW_RES_VARIABLES, LOW_RES_DEFAULT_VARIABLES, None), + ], +) +def test_default_selection_is_stable_and_returns_a_new_list(registry, defaults, cadence): + first = resolve_variable_names(registry, None, defaults, cadence) + second = resolve_variable_names(registry, None, defaults, cadence) + + assert first == list(defaults) + assert second == list(defaults) + assert first is not second + + +def test_all_selection_respects_cadence_and_registry_order(): + one_minute = resolve_variable_names(HIGH_RES_VARIABLES, "all", HIGH_RES_DEFAULT_VARIABLES, 1) + five_minute = resolve_variable_names(HIGH_RES_VARIABLES, "all", HIGH_RES_DEFAULT_VARIABLES, 5) + + assert one_minute == [variable.name for variable in HIGH_RES_VARIABLES if 1 in variable.cadences] + assert five_minute == [variable.name for variable in HIGH_RES_VARIABLES] + + +def test_generator_aliases_and_duplicates_preserve_first_canonical_occurrence(): + requested = (name for name in ["sym_h", "speed", "sym-h", "flow_pressure", "pdyn"]) + + result = resolve_variable_names(HIGH_RES_VARIABLES, requested, HIGH_RES_DEFAULT_VARIABLES, 1) + + assert result == ["sym-h", "speed", "pdyn"] + + +@pytest.mark.parametrize("selection", [[], (), iter(())]) +def test_empty_iterables_are_rejected(selection): + with pytest.raises(ValueError, match="at least one"): + resolve_variable_names(LOW_RES_VARIABLES, selection, LOW_RES_DEFAULT_VARIABLES) + + +@pytest.mark.parametrize("selection", [42, object(), ["speed", 3], [None]]) +def test_non_string_or_non_iterable_selections_are_rejected_clearly(selection): + with pytest.raises(ValueError, match="variable name|names must be strings"): + resolve_variable_names(HIGH_RES_VARIABLES, selection, HIGH_RES_DEFAULT_VARIABLES, 1) + + +def test_unknown_error_lists_only_cadence_valid_names(): + with pytest.raises(ValueError) as error: + resolve_variable_names(HIGH_RES_VARIABLES, "unknown", HIGH_RES_DEFAULT_VARIABLES, 1) + + message = str(error.value) + assert "bavg" in message + assert "proton_flux_10_mev" not in message + + +def test_cadence_error_preserves_requested_alias_in_message(): + with pytest.raises(ValueError, match="proton_flux_10_mev"): + resolve_variable_names(HIGH_RES_VARIABLES, ["speed", "proton_flux_10_mev"], HIGH_RES_DEFAULT_VARIABLES, 1) diff --git a/tests/io/solar_wind/test_omni.py b/tests/io/solar_wind/test_omni.py index 6ad5a8e..760f347 100644 --- a/tests/io/solar_wind/test_omni.py +++ b/tests/io/solar_wind/test_omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -11,6 +12,7 @@ import pandas as pd import pytest +from swvo.io.omni import OMNIHighRes from swvo.io.solar_wind import SWOMNI TEST_DIR = os.path.dirname(__file__) @@ -37,13 +39,7 @@ def test_initialization_without_env_var(self): SWOMNI() def test_download_and_process(self, swomni): - start_time = datetime(2020, 1, 1, tzinfo=timezone.utc) - end_time = datetime(2020, 12, 31, tzinfo=timezone.utc) - # download this file without mocking - swomni.download_and_process(start_time, end_time) - - for file in (DATA_DIR / "OMNI").glob("OMNI_HIGH_RES_1min_2020*.csv"): - assert file.exists() + assert swomni.download_and_process.__func__ is OMNIHighRes.download_and_process def test_read_without_download(self, swomni, mocker): start_time = datetime(2021, 1, 1, tzinfo=timezone.utc) @@ -81,13 +77,12 @@ def test_start_year_behind(self, swomni, mocker): mocker.patch.object(pd.DataFrame, "truncate", return_value=pd.DataFrame()) with patch("logging.Logger.warning") as mock_warning: - dfs = swomni.read(start_time, end_time) + with pytest.raises(ValueError, match="No OMNI"): + swomni.read(start_time, end_time) mock_warning.assert_any_call( "Start date chosen falls behind the existing data. Moving start date to first available mission files..." ) - assert len(dfs) == 0, "Expected dfs list to be empty since no files are found." - def test_remove_processed_file(self): shutil.rmtree(Path(TEST_DIR) / "data/OMNI/2020", ignore_errors=True) shutil.rmtree(Path(TEST_DIR) / "data/OMNI/2022", ignore_errors=True) diff --git a/tests/io/symh/test_symh_omni.py b/tests/io/symh/test_symh_omni.py index 7a8ac92..b49d101 100644 --- a/tests/io/symh/test_symh_omni.py +++ b/tests/io/symh/test_symh_omni.py @@ -1,4 +1,5 @@ # SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences +# SPDX-FileContributor: Simon Mischel # # SPDX-License-Identifier: Apache-2.0 @@ -10,6 +11,7 @@ import pandas as pd import pytest +from swvo.io.omni import OMNIHighRes from swvo.io.symh import SymhOMNI TEST_DIR = os.path.dirname(__file__) @@ -50,12 +52,7 @@ def test_initialization_without_env_var(self): SymhOMNI() def test_download_and_process(self, symhomni): - start_time = datetime(2020, 1, 1, tzinfo=timezone.utc) - end_time = datetime(2020, 12, 31, tzinfo=timezone.utc) - # download this file without mocking - symhomni.download_and_process(start_time, end_time) - for file in (DATA_DIR / "OMNI" / "2020").glob("OMNI_HIGH_RES_1min_2020*.csv"): - assert file.exists() + assert symhomni.download_and_process.__func__ is OMNIHighRes.download_and_process def test_read_without_download(self, symhomni): start_time = datetime(2021, 1, 1, tzinfo=timezone.utc) @@ -65,11 +62,22 @@ def test_read_without_download(self, symhomni): ): # value error is raised when no files are found hence no concatenation is possible symhomni.read(start_time, end_time, download=False) - def test_read_with_download(self, symhomni): - start_time = datetime(2021, 12, 1, tzinfo=timezone.utc) + def test_read_with_download(self, symhomni, mocker): + start_time = datetime(2021, 12, 1) end_time = datetime(2022, 1, 31, tzinfo=timezone.utc) + index = pd.date_range(start="2021-12-01", end="2022-01-31", freq="min", tz="UTC") + parent_result = pd.DataFrame({"sym-h": -15.0, "file_name": "some_file"}, index=index) + parent_read = mocker.patch.object(OMNIHighRes, "read", return_value=parent_result) + df = symhomni.read(start_time, end_time, download=True) + parent_read.assert_called_once_with( + start_time, + end_time, + cadence_min=1, + download=True, + variables="sym-h", + ) assert isinstance(df, pd.DataFrame) assert len(df) > 0 assert "sym-h" in df.columns