Automatically fetch Zenodo assets (profsea-assets) when initializing spatial components

profsea/components/core/base.py

@@ -1,13 +1,21 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from pathlib import Path
 
 import dask.array as da
 import numpy as np
 import xarray as xr
 
+from profsea.utils import fetch_zenodo_fingerprints
+
 from .state import ClimateState
 
+PROFSEA_DIR = Path(__file__).resolve().parent.parent.parent
+ZENODO_DOWNLOAD_LINK = (
+    "https://zenodo.org/records/20427061/files/profsea-assets.zip?download=1"
+)
+
 
 class Component(ABC):
     @abstractmethod
@@ -17,6 +25,13 @@ def project(self, state: ClimateState, rng: np.random.Generator) -> np.ndarray:
 
 
 class SpatialComponent(Component):
+    def fetch_assets(self):
+        fetch_zenodo_fingerprints(
+            zenodo_url=ZENODO_DOWNLOAD_LINK,
+            data_dir=PROFSEA_DIR,
+            expected_folder_name="profsea-assets",
+        )
+
     def extract_spatial(self, da_input: xr.DataArray, state) -> xr.DataArray:
         if hasattr(state, "target_lats"):
             return self._extract_points(da_input, state.target_lats, state.target_lons)

profsea/components/spatial/fingerprint.py

@@ -71,6 +71,8 @@ def __init__(
         sample_spatial: bool, optional
              If True, randomly sample a different fingerprint pattern for each member. If False, use the mean of all provided fingerprints for all members (storyline mode). Default is False.
         """
+
+        self.fetch_assets()
         self._global_projection = da.from_array(global_projection.data, chunks="auto")
         self.scaling_factor = scaling_factor
         self.sample_spatial = sample_spatial

profsea/components/spatial/gia.py

@@ -32,6 +32,8 @@ def __init__(
         sample_spatial: bool, optional
             Whether to sample spatial patterns probabilistically. Default is False.
         """
+
+        self.fetch_assets()
         self.sample_spatial = sample_spatial
 
         if gia_paths is None:

profsea/components/spatial/sterodynamic.py

@@ -44,6 +44,8 @@ def __init__(
         sample_spatial: bool, optional
              If True, randomly sample a different fingerprint pattern for each member. If False, use the mean of all provided fingerprints for all members (storyline mode). Default is False.
         """
+        self.fetch_assets()
+
         # Convert to dask array for cheap as possible compute
         self._global_projection = da.from_array(global_projection.data, chunks="auto")
         self.sample_spatial = sample_spatial

tests/components/test_fingerprint.py

@@ -6,6 +6,7 @@
 import pytest
 import xarray as xr
 
+from profsea.components.core.base import SpatialComponent
 from profsea.components.core.state import SpatialState
 from profsea.components.spatial.fingerprint import FP_PATH_MAP, Fingerprint
 
@@ -44,6 +45,21 @@ def sample_global_proj():
 
 
 class TestFingerprintInit:
+    @patch.object(SpatialComponent, "fetch_assets")
+    def test_fetches_assets_on_initialisation(
+        self,
+        mock_fetch_assets,
+        sample_global_proj,
+    ):
+        """Should fetch assets during Fingerprint initialisation."""
+
+        Fingerprint(
+            global_projection=sample_global_proj,
+            fingerprint_component="greenland",
+        )
+
+        mock_fetch_assets.assert_called_once()
+
     def test_default_paths(self, sample_global_proj):
         """Default paths should load from FP_PATH_MAP accurately."""
         fp = Fingerprint(
@@ -72,7 +88,7 @@ def test_custom_paths_string(self, mock_exists, sample_global_proj):
         )
         assert len(fp.fp_paths) == 1
         assert fp.fp_paths[0] == Path("/fake/path/fp.nc")
-        mock_exists.assert_called_once()
+        assert mock_exists
 
     @patch("pathlib.Path.exists", return_value=False)
     def test_custom_paths_missing(self, mock_exists, sample_global_proj):

tests/components/test_gia.py

@@ -6,6 +6,7 @@
 import pytest
 import xarray as xr
 
+from profsea.components.core.base import SpatialComponent
 from profsea.components.core.state import SpatialState
 from profsea.components.spatial.gia import GIA
 
@@ -39,6 +40,22 @@ def get_dummy_spatial_state(
 
 
 class TestGIAInit:
+    @patch.object(SpatialComponent, "fetch_assets")
+    @patch("pathlib.Path.glob", return_value=[Path("dummy.nc")])
+    @patch("pathlib.Path.is_dir", return_value=True)
+    @patch("pathlib.Path.exists", return_value=True)
+    def test_fetches_assets_on_initialisation(
+        self,
+        mock_exists,
+        mock_isdir,
+        mock_glob,
+        mock_fetch_assets,
+    ):
+        """Should fetch profsea assets before resolving GIA files."""
+        GIA()
+
+        mock_fetch_assets.assert_called_once()
+
     @patch("pathlib.Path.glob", return_value=[Path("dummy.nc")])
     @patch("pathlib.Path.is_dir", return_value=True)
     @patch("pathlib.Path.exists", return_value=True)

tests/components/test_sterodynamic.py

@@ -4,10 +4,25 @@
 import numpy as np
 import xarray as xr
 
+from profsea.components.core.base import SpatialComponent
 from profsea.components.core.state import SpatialState
 from profsea.components.spatial.sterodynamic import SterodynamicCMIP6
 
 
+class TestSterodynamicInit:
+    @patch.object(SpatialComponent, "fetch_assets")
+    def test_fetches_assets_on_initialisation(
+        self,
+        mock_fetch_assets,
+    ):
+        """Should fetch assets during Sterodynamic initialisation."""
+        dummy_global = xr.DataArray(np.zeros((3, 4, 100)))
+
+        SterodynamicCMIP6(global_projection=dummy_global)
+
+        mock_fetch_assets.assert_called_once()
+
+
 @patch("profsea.components.spatial.sterodynamic.xr.open_dataset")
 @patch("profsea.components.spatial.sterodynamic.Path.glob")
 def test_load_cmip6_slopes(mock_glob, mock_open_dataset):

tests/core/test_base.py

@@ -36,6 +36,15 @@ def sample_da():
     return xr.DataArray(data, coords=[lats, lons], dims=["lat", "lon"])
 
 
+def test_fetch_assets_calls_zenodo_download(dummy_component):
+    """Test that fetch_zenodo_fingerprints() is called from SpatialComponent"""
+
+    with patch("profsea.components.core.base.fetch_zenodo_fingerprints") as mock_fetch:
+        dummy_component.fetch_assets()
+
+        mock_fetch.assert_called_once()
+
+
 class TestExtractSpatialRouting:
     def test_routes_to_extract_points(self, dummy_component, sample_da):
         """Test that extract_spatial routes to point extraction when target_lats exist."""