diff --git a/src/virtualship/instruments/base.py b/src/virtualship/instruments/base.py index f07bc0cf..18c175d3 100644 --- a/src/virtualship/instruments/base.py +++ b/src/virtualship/instruments/base.py @@ -2,6 +2,7 @@ import abc import collections +import tempfile from dataclasses import dataclass from datetime import timedelta from itertools import pairwise @@ -192,6 +193,9 @@ def _generate_fieldset(self) -> parcels.FieldSet: Create and combine FieldSets for each variable, supporting both local and Copernicus Marine data sources. N.B. Per variable avoids issues when using copernicusmarine and creating directly one FieldSet of ds's sourced from different Copernicus Marine product IDs (which can also have different temporal resolutions), which is often the case for BGC variables. + + Includes an intermediate step of writing to tmp files, as per https://github.com/Parcels-code/parcels-benchmarks/pull/49 + TODO: the need for this step may be removed as Parcels x copernicusmarine integration improves, tracked in https://github.com/Parcels-code/Parcels/issues/2756 and xref'd in VirtualShip #357 (https://github.com/Parcels-code/virtualship/issues/357) """ fieldsets_list = [] keys = list(self.variables.keys()) @@ -227,15 +231,12 @@ def _generate_fieldset(self) -> parcels.FieldSet: ) field_var_name = var - # # negate depth and reindex (to suit Parcels XGrid strictly increasing depth convention) - # ds["depth"] = -ds["depth"] - # ds = ds.reindex(depth=ds["depth"][::-1]) - - # TODO: update when decision on handling of nans/0s in v4 is made (i.e. https://github.com/Parcels-code/Parcels/issues/2393) + # TODO: to be removed when Parcels #2746 is merged (i.e. https://github.com/Parcels-code/Parcels/pull/2746) ds = ds.fillna(0) fields = {key: ds[field_var_name]} ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields) + ds_fset = self._via_tmp_ds(ds_fset) fs = parcels.FieldSet.from_sgrid_conventions(ds_fset) @@ -257,3 +258,12 @@ def _generate_fieldset(self) -> parcels.FieldSet: base_fieldset.add_field(uv) return base_fieldset + + @staticmethod + def _via_tmp_ds(ds) -> xr.Dataset: + """Create and re-load a temporary local dataset.""" + tmpdir = tempfile.TemporaryDirectory() + tmp_fpath = Path(tmpdir.name).joinpath("tmp.nc") + ds.to_netcdf(tmp_fpath) + del ds + return xr.open_dataset(tmp_fpath) diff --git a/src/virtualship/instruments/drifter.py b/src/virtualship/instruments/drifter.py index 6319831e..3b52dc71 100644 --- a/src/virtualship/instruments/drifter.py +++ b/src/virtualship/instruments/drifter.py @@ -89,7 +89,7 @@ def __init__(self, expedition, from_data): **sensor_variables, } # advection variables (U and V) are always required for drifter simulation; sensor variables come from config fetch_spec = FetchSpec( - latlon_buffer=30.0, # TODO: generous buffer to limit tmp file size download, can potentially be removed in the future as and when Parcels streaming performance improves (see #358) + latlon_buffer=30.0, # TODO: generous buffer to reduce tmp file footprint, can potentially be removed in the future as/when Parcels streaming performance improves (see #358) time_buffer=expedition.instruments_config.drifter_config.lifetime.total_seconds() / (24 * 3600), # [days] depth_min=abs( diff --git a/tests/instruments/test_base.py b/tests/instruments/test_base.py index 1b090cb9..93a38e90 100644 --- a/tests/instruments/test_base.py +++ b/tests/instruments/test_base.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock, patch import pytest +import xarray as xr from virtualship.instruments.base import FetchSpec, Instrument from virtualship.instruments.types import InstrumentType @@ -89,25 +90,43 @@ def test_execute_calls_simulate(monkeypatch): dummy.simulate.assert_called_once() -def test_get_spec_value_buffer_and_limit(): +def test_fetch_spec_applied_to_instrument(): + """FetchSpec values are correctly stored on the instrument.""" mock_waypoint = MagicMock() mock_waypoint.location.latitude = 1.0 mock_waypoint.location.longitude = 2.0 mock_schedule = MagicMock() mock_schedule.waypoints = [mock_waypoint] + fetch_spec = FetchSpec(latlon_buffer=5.0, depth_min=10.0) dummy = DummyInstrument( expedition=MagicMock(schedule=mock_schedule), variables={"A": "a"}, add_bathymetry=False, allow_time_extrapolation=False, verbose_progress=False, - spacetime_buffer_size={"latlon": 5.0}, - limit_spec={"depth_min": 10.0}, + fetch_spec=fetch_spec, from_data=None, ) - assert dummy._get_spec_value("buffer", "latlon", 0.0) == 5.0 - assert dummy._get_spec_value("limit", "depth_min", None) == 10.0 - assert dummy._get_spec_value("buffer", "missing", 42) == 42 + assert dummy.fetch_spec.latlon_buffer == 5.0 + assert dummy.fetch_spec.depth_min == 10.0 + # unset values use dataclass defaults + assert dummy.fetch_spec.time_buffer == 0.0 + assert dummy.fetch_spec.depth_max is None + + +def test_via_tmp_ds_roundtrip(): + """_via_tmp_ds writes to a tmp file and re-opens it.""" + ds = xr.Dataset( + {"temperature": (["x", "y"], [[1.0, 2.0], [3.0, 4.0]])}, + coords={"x": [0, 1], "y": [10, 20]}, + ) + result = Instrument._via_tmp_ds(ds) + + assert isinstance(result, xr.Dataset) + assert "temperature" in result + assert ( + result is not ds + ) # result is new object loaded from tmp file, not the original def test_generate_fieldset_combines_fields(monkeypatch):