diff --git a/profsea/components/core/base.py b/profsea/components/core/base.py index c946e29..9ec66bf 100644 --- a/profsea/components/core/base.py +++ b/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) diff --git a/profsea/components/spatial/fingerprint.py b/profsea/components/spatial/fingerprint.py index 4c281cf..1877e3f 100644 --- a/profsea/components/spatial/fingerprint.py +++ b/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 diff --git a/profsea/components/spatial/gia.py b/profsea/components/spatial/gia.py index 4ea424f..b929319 100644 --- a/profsea/components/spatial/gia.py +++ b/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: diff --git a/profsea/components/spatial/sterodynamic.py b/profsea/components/spatial/sterodynamic.py index 6023207..759a193 100644 --- a/profsea/components/spatial/sterodynamic.py +++ b/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 diff --git a/tests/components/test_fingerprint.py b/tests/components/test_fingerprint.py index de24b87..c8e1cb2 100644 --- a/tests/components/test_fingerprint.py +++ b/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): diff --git a/tests/components/test_gia.py b/tests/components/test_gia.py index adc1d4b..ef4079c 100644 --- a/tests/components/test_gia.py +++ b/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) diff --git a/tests/components/test_sterodynamic.py b/tests/components/test_sterodynamic.py index 5fe91aa..ad4c3a5 100644 --- a/tests/components/test_sterodynamic.py +++ b/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): diff --git a/tests/core/test_base.py b/tests/core/test_base.py index 6d687a1..b50bdf8 100644 --- a/tests/core/test_base.py +++ b/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."""