diff --git a/CHANGELOG.md b/CHANGELOG.md index a6a5e5b..b103560 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,11 @@ All notable changes to this project will be documented in this file. at higher temporal resolutions. There can be step changes of more than 20% in some situations when going from hour 02:59 to 03:00 when new indices get used again on the 3-hourly boundaries. +- **ADDED** `pymsis.use_space_weather_file()` function. + - The function points pymsis to a custom space weather + file instead of the one bundled in the install location. + - The same file path can be set at startup with the `PYMSIS_SPACE_WEATHER_FILE` + environment variable. - **CHANGED** `get_f107_ap()` returns arrays of the same shape as the input - Previously, when a scalar date was passed to the utility function, the `ap` values were 1d (7,) rather than of shape (1, 7) corresponding to diff --git a/docs/source/reference/index.rst b/docs/source/reference/index.rst index 07b0571..c600fdf 100644 --- a/docs/source/reference/index.rst +++ b/docs/source/reference/index.rst @@ -49,3 +49,4 @@ get the proper F10.7 and ap values to use for a specific date and time. utils.download_f107_ap utils.get_f107_ap + utils.use_space_weather_file diff --git a/pymsis/__init__.py b/pymsis/__init__.py index 1b92db7..f3379f5 100644 --- a/pymsis/__init__.py +++ b/pymsis/__init__.py @@ -3,8 +3,9 @@ import importlib.metadata from pymsis.msis import Variable, calculate +from pymsis.utils import use_space_weather_file __version__ = importlib.metadata.version("pymsis") -__all__ = ["Variable", "__version__", "calculate"] +__all__ = ["Variable", "__version__", "calculate", "use_space_weather_file"] diff --git a/pymsis/utils.py b/pymsis/utils.py index 2e55e3a..342e017 100644 --- a/pymsis/utils.py +++ b/pymsis/utils.py @@ -1,5 +1,6 @@ """Utilities for obtaining input datasets.""" +import os import urllib.request import warnings from io import BytesIO @@ -13,9 +14,51 @@ _DATA_FNAME: str = "SW-All.csv" _F107_AP_URL: str = f"https://celestrak.org/SpaceData/{_DATA_FNAME}" -_F107_AP_PATH: Path = Path(pymsis.__file__).parent / _DATA_FNAME +_F107_AP_DEFAULT_FILE: Path = Path(pymsis.__file__).parent / _DATA_FNAME _DATA: dict[str, npt.NDArray] | None = None +_F107_AP_FILE: Path = Path( + os.environ.get("PYMSIS_SPACE_WEATHER_FILE", _F107_AP_DEFAULT_FILE) +) + + +def use_space_weather_file(file: str | Path | None = None) -> None: + """ + Direct pymsis to use a custom file to retrieve the F10.7 and ap data from. + + The custom data file must follow the same (.csv) format as data retrieved from + Celestrak. The legacy (.txt) file format is not supported. + + By default, the data is retrieved from CelesTrak and stored inside the + installed package location. Retrieving the data from a custom file path may + be useful if you are retrieving the data from a different source, if you + have a centralized location for the data, or would like to use custom data. + + Alternatively, the file path can be set using the environment variable + ``PYMSIS_SPACE_WEATHER_FILE`` to achieve the same result and + to avoid setting the space weather file path programmatically on each run. + + Setting a default and running `download_f107_ap()` will not download to + this custom file path. + + Parameters + ---------- + file : str or Path or None + Path to the F10.7 and ap data file, retrieved from CelesTrak or elsewhere. + If set to None, the space weather file downloaded from CelesTrak (stored in the + installed package location) will be used. + """ + file = Path(file) if file is not None else _F107_AP_DEFAULT_FILE + if not file.is_file(): + raise FileNotFoundError( + f"Provided custom space weather file does not exist: {file}" + ) + + # update the global path and data variables + global _F107_AP_FILE, _DATA # noqa: PLW0603 + _F107_AP_FILE = Path(file) + _DATA = None + def download_f107_ap() -> None: """ @@ -26,6 +69,9 @@ def download_f107_ap() -> None: This routine can be called to update the data as well if you would like to use newer data since the last time you downloaded the file. + If `use_space_weather_file()` has been called to set a custom file path, the file + will still be downloaded to the default location and thus ignored by pymsis. + Notes ----- The default data provider is CelesTrak, which gets data from other @@ -43,14 +89,30 @@ def download_f107_ap() -> None: Space Weather, https://doi.org/10.1029/2020SW002641 """ warnings.warn(f"Downloading ap and F10.7 data from {_F107_AP_URL}") + if _F107_AP_DEFAULT_FILE != _F107_AP_FILE: + warnings.warn( + "A custom space weather file has been set, but the downloaded file " + "will be stored in the default location and ignored. Unset the " + "custom file path using `use_space_weather_file(None)` to use the " + "downloaded file." + ) req = urllib.request.urlopen(_F107_AP_URL) - with _F107_AP_PATH.open("wb") as f: + with _F107_AP_DEFAULT_FILE.open("wb") as f: f.write(req.read()) def _load_f107_ap_data() -> dict[str, npt.NDArray]: """Load data from disk, if it isn't present go out and download it first.""" - if not _F107_AP_PATH.exists(): + default_file_exists = _F107_AP_DEFAULT_FILE.is_file() + custom_file_used = _F107_AP_FILE != _F107_AP_DEFAULT_FILE + + if custom_file_used and not _F107_AP_FILE.is_file(): + raise FileNotFoundError( + "Custom space weather file has been set but does not exist: " + f"{_F107_AP_FILE}" + ) + + if not custom_file_used and not default_file_exists: download_f107_ap() dtype = { @@ -90,7 +152,7 @@ def _load_f107_ap_data() -> dict[str, npt.NDArray]: # Use a buffer to read in and load so we can quickly get rid of # the extra "PRD" lines at the end of the file (unknown length # so we can't just go back in line lengths) - with _F107_AP_PATH.open() as fin: + with _F107_AP_FILE.open() as fin: with BytesIO() as fout: for line in fin: if "PRM" in line: @@ -168,7 +230,8 @@ def _load_f107_ap_data() -> dict[str, npt.NDArray]: "f107a": f107a_data, "warn_data": warn_data, } - globals()["_DATA"] = data + global _DATA # noqa: PLW0603 + _DATA = data return data diff --git a/tests/conftest.py b/tests/conftest.py index bbe64d9..044bdb5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,8 +10,10 @@ def local_path(monkeypatch): # Update the data location to our test data test_file = Path(__file__).parent / "f107_ap_test_data.txt" # Monkeypatch the url and expected download location, so we aren't - # dependent on an internet connection. - monkeypatch.setattr(utils, "_F107_AP_PATH", test_file) + # dependent on an internet connection. Patch both the active path and the + # default path so tests exercise the "default location" branch by default. + monkeypatch.setattr(utils, "_F107_AP_FILE", test_file) + monkeypatch.setattr(utils, "_F107_AP_DEFAULT_FILE", test_file) return test_file diff --git a/tests/test_utils.py b/tests/test_utils.py index c70d388..a9ef39a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,3 +1,5 @@ +import importlib + import numpy as np import pytest from numpy.testing import assert_allclose, assert_array_equal @@ -7,7 +9,9 @@ def test_downloading(monkeypatch, tmp_path): tmp_file = tmp_path / "testfile.txt" - monkeypatch.setattr(utils, "_F107_AP_PATH", tmp_file) + # download_f107_ap() always writes to the default location, so patch both + monkeypatch.setattr(utils, "_F107_AP_FILE", tmp_file) + monkeypatch.setattr(utils, "_F107_AP_DEFAULT_FILE", tmp_file) # just be nice and make sure we are getting the monkeypatched local file assert utils._F107_AP_URL.startswith("file://") with pytest.warns(UserWarning, match="Downloading ap and F10.7"): @@ -15,7 +19,7 @@ def test_downloading(monkeypatch, tmp_path): assert tmp_file.exists() with ( open(tmp_file, "rb") as f_downloaded, - open(utils._F107_AP_PATH, "rb") as f_expected, + open(utils._F107_AP_FILE, "rb") as f_expected, ): assert f_downloaded.read() == f_expected.read() @@ -27,7 +31,10 @@ def test_loading_data(monkeypatch, tmp_path): # Make sure a download warning is emitted if the file doesn't exist tmp_file = tmp_path / "testfile.txt" with monkeypatch.context() as m: - m.setattr(utils, "_F107_AP_PATH", tmp_file) + # Point both at the (missing) default location so the auto-download + # branch is exercised rather than the custom-path branch. + m.setattr(utils, "_F107_AP_FILE", tmp_file) + m.setattr(utils, "_F107_AP_DEFAULT_FILE", tmp_file) assert not tmp_file.exists() with pytest.warns(UserWarning, match="Downloading ap and F10.7"): utils._load_f107_ap_data() @@ -167,3 +174,44 @@ def test_get_f107_ap_interpolate_exact_match(date): assert_allclose(f107_nointerp, f107_interp) assert_allclose(f107a_nointerp, f107a_interp) assert_array_equal(ap_nointerp, ap_interp) + + +def test_use_space_weather_file(monkeypatch, tmp_path): + custom_file = tmp_path / "custom_sw.csv" + custom_file.write_text("placeholder") + # Pretend some data was already cached so we can verify it gets reset + monkeypatch.setattr(utils, "_DATA", {"dummy": np.array([1])}) + + utils.use_space_weather_file(custom_file) + # Check if the path has been updated + assert utils._F107_AP_FILE == custom_file + # Check if the data has been reset. + assert utils._DATA is None + + # Setting a path that doesn't exist should raise + missing = tmp_path / "does_not_exist.csv" + with pytest.raises(FileNotFoundError, match="does not exist"): + utils.use_space_weather_file(missing) + + # downloading should raise a warning that the file is ignored. + # Redirect the download to a throwaway path so we don't overwrite shared test data. + monkeypatch.setattr(utils, "_F107_AP_DEFAULT_FILE", tmp_path / "downloaded.csv") + with pytest.warns(UserWarning, match="A custom space weather file has been set"): + utils.download_f107_ap() + + +def test_space_weather_env_variable(monkeypatch, tmp_path): + # check if setting space weather file via env variable works + custom_file = tmp_path / "custom_sw.csv" + custom_file.write_text("placeholder") + monkeypatch.setenv("PYMSIS_SPACE_WEATHER_FILE", str(custom_file)) + importlib.reload(utils) + + assert utils._F107_AP_FILE == custom_file + + # Loading the data must validate that the file exists. + missing = tmp_path / "does_not_exist.csv" + monkeypatch.setenv("PYMSIS_SPACE_WEATHER_FILE", str(missing)) + importlib.reload(utils) + with pytest.raises(FileNotFoundError, match="Custom space weather file"): + utils._load_f107_ap_data()