ENH: Add ability to point to a different Space Weather File location

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

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

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"]

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
 
 

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
 
 

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,15 +9,17 @@
 
 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"):
         utils.download_f107_ap()
     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()