Improve performance: memory footprint and speed

src/virtualship/cli/_run.py

@@ -179,17 +179,20 @@ def _run(
             attr = MeasurementsToSimulate.get_attr_for_instrumenttype(itype)
             measurements = getattr(schedule_results.measurements_to_simulate, attr)
 
-            # initialise instrument, execute simulation within context manager
-            with instrument_class(
+            # initialise instrument
+            instrument = instrument_class(
                 expedition=expedition,
                 from_data=Path(from_data) if from_data is not None else None,
-            ) as instrument:
-                instrument.execute(
-                    measurements=measurements,
-                    out_path=expedition_dir.joinpath(
-                        RESULTS, f"{itype.name.lower()}.parquet"
-                    ),
-                )
+            )
+
+            # execute simulation
+            instrument.execute(
+                measurements=measurements,
+                out_path=expedition_dir.joinpath(
+                    RESULTS, f"{itype.name.lower()}.parquet"
+                ),
+            )
+
         except Exception as e:
             # clean up if unexpected error occurs
             if os.path.exists(problems_dir):

src/virtualship/instruments/adcp.py

@@ -51,13 +51,17 @@ class ADCPInstrument(UnderwayInstrument):
     def __init__(self, expedition, from_data):
         """Initialize ADCPInstrument."""
         variables = expedition.instruments_config.adcp_config.active_variables()
+        fetch_spec = FetchSpec(
+            depth_min=0,  # ensures copernicusmarine fetches properly
+            depth_max=expedition.instruments_config.adcp_config.max_depth_meter,
+        )
 
         super().__init__(
             expedition,
             variables,
             add_bathymetry=False,
             verbose_progress=False,
-            fetch_spec=FetchSpec(),
+            fetch_spec=fetch_spec,
             from_data=from_data,
         )
 

src/virtualship/instruments/argo_float.py

@@ -166,6 +166,10 @@ def _argo_sample_temperature(particles, fieldset):
     depth_mask = particles.z < particles.min_depth  # still ascending
     mask = np.logical_and(phase_mask, depth_mask)
     particles.temperature[~mask] = np.nan  # no measurement outside the ascent
+    if not np.any(mask):
+        # TODO: tmp fix avoiding IndexError in Parcels' ChunkCachedArray vectorized indexing when sampling with an empty ParticleSet (Parcels issue: #2906)
+        # TODO: can be removed when fixed upstream in Parcels
+        return
     sampling_particles = particles[mask]
     sampling_particles.temperature = fieldset.T[sampling_particles]
 
@@ -176,6 +180,10 @@ def _argo_sample_salinity(particles, fieldset):
     depth_mask = particles.z < particles.min_depth  # still ascending
     mask = np.logical_and(phase_mask, depth_mask)
     particles.salinity[~mask] = np.nan  # no measurement outside the ascent
+    if not np.any(mask):
+        # TODO: tmp fix avoiding IndexError in Parcels' ChunkCachedArray vectorized indexing when sampling with an empty ParticleSet (Parcels issue: #2906)
+        # TODO: can be removed when fixed upstream in Parcels
+        return
     sampling_particles = particles[mask]
     sampling_particles.salinity = fieldset.S[sampling_particles]
 
@@ -250,6 +258,8 @@ def __init__(self, expedition, from_data):
             latlon_buffer=9.0,  # [degrees]
             time_buffer=expedition.instruments_config.argo_float_config.lifetime.total_seconds()
             / (24 * 3600),  # [days]
+            depth_min=expedition.instruments_config.argo_float_config.min_depth_meter,
+            depth_max=expedition.instruments_config.argo_float_config.max_depth_meter,
         )
 
         super().__init__(

src/virtualship/instruments/base.py

@@ -4,7 +4,6 @@
 import collections
 import inspect
 import itertools
-import tempfile
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -23,6 +22,7 @@
 from virtualship.utils import (
     COPERNICUSMARINE_PHYS_VARIABLES,
     INSTRUMENT_CLASS_MAP,
+    MAX_CACHE_BYTES,
     _find_files_in_timerange,
     _find_nc_file_with_variable,
     _get_bathy_data,
@@ -119,7 +119,6 @@ def __init__(
         self.add_bathymetry = add_bathymetry
         self.verbose_progress = verbose_progress
         self.fetch_spec = fetch_spec if fetch_spec is not None else FetchSpec()
-        self._tmp_dirs: list[tempfile.TemporaryDirectory] = []
 
         # filter to waypoints relevant to this instrument
         wps_in_use = self.expedition.schedule._get_wps_in_use()
@@ -137,26 +136,6 @@ def __init__(
         # spatio-temporal bounding box of all relevant waypoints
         self.bounds = SpatialBounds.from_waypoints(relevant_waypoints)
 
-    def close(self):
-        """Explicitly cleanup all tmp dirs."""
-        tmp_dirs = getattr(self, "_tmp_dirs", None)
-        if not tmp_dirs:
-            return
-        for tmp_dir in tmp_dirs:
-            try:
-                tmp_dir.cleanup()
-            except Exception:
-                pass  # i.e. best effort clean up
-        self._tmp_dirs = []
-
-    def __enter__(self):
-        """Enter the context manager."""
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Exit context manager, ensuring resource cleanup."""
-        self.close()
-
     def load_input_data(self) -> parcels.FieldSet:
         """Load and return the input data as a FieldSet for the instrument."""
         try:
@@ -227,11 +206,10 @@ def _generate_fieldset(self) -> parcels.FieldSet:
         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)
         """
         combined_fieldset = None
-        keys = list(self.variables.keys())
         time_buffer = self.fetch_spec.time_buffer
+        is_underway = self.instrument_type.is_underway
 
-        for key in keys:
-            var = self.variables[key]
+        for key, var in self.variables.items():
             physical = var in COPERNICUSMARINE_PHYS_VARIABLES
 
             if self.from_data is not None:  # load from local data
@@ -260,16 +238,15 @@ def _generate_fieldset(self) -> parcels.FieldSet:
             fields = {key: ds[field_var_name]}
             ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)
 
-            # streaming data performance is improved by writing to a temporary file, unnecessary for local data
-            if self.from_data is None:
-                ds_fset = self._via_tmp_ds(ds_fset)
+            # operations only necessary for non-underway instruments
+            if not is_underway:
+                fs = parcels.FieldSet.from_sgrid_conventions(ds_fset)
 
-            fs = parcels.FieldSet.from_sgrid_conventions(ds_fset)
+                # to ChunkCachedArrays for better Dask/memory management
+                fs = fs.to_chunk_cached_arrays(max_cache_bytes=MAX_CACHE_BYTES)
 
-            # non-underway instruments to windowed arrays, just in case any ds is Dask backed
-            # underway instruments should not to converted to windowed arrays, as they use one direct fieldset.eval() call which could cause a big memory usage if the fieldset is windowed
-            if not self.instrument_type.is_underway:
-                fs = fs.to_windowed_arrays()
+            else:
+                fs = parcels.FieldSet.from_sgrid_conventions(ds_fset)
 
             combined_fieldset = combined_fieldset + fs if combined_fieldset else fs
 
@@ -374,32 +351,6 @@ def _get_local_ds(self, files: list[Path]) -> xr.Dataset:
 
         return ds
 
-    def _via_tmp_ds(self, ds: xr.Dataset) -> xr.Dataset:
-        """Create and re-load a temporary local dataset without loading everything into RAM, using local Zarr store for improved performance and concurrent chunk writing."""
-        tmp_dir = tempfile.TemporaryDirectory()
-        self._tmp_dirs.append(tmp_dir)
-        tmp_store = Path(tmp_dir.name) / f"tmp_{id(ds)}.zarr"
-
-        ds_to_write = ds.copy()
-        for variable in ds_to_write.variables.values():
-            variable.encoding = {}
-
-        ds_to_write = ds_to_write.chunk(
-            {dim: size for dim, size in ds_to_write.sizes.items()}
-        )
-
-        ds_to_write.to_zarr(
-            tmp_store,
-            mode="w",
-            consolidated=False,
-        )
-
-        loaded_ds = xr.open_zarr(
-            tmp_store, chunks=None, consolidated=False
-        )  # chunks=None to avoid Dask backed
-
-        return loaded_ds
-
     @staticmethod
     def _sample_initial(
         pset: parcels.ParticleSet,

src/virtualship/instruments/ctd.py

@@ -134,13 +134,17 @@ class CTDInstrument(Instrument):
     def __init__(self, expedition, from_data):
         """Initialize CTDInstrument."""
         variables = expedition.instruments_config.ctd_config.active_variables()
+        fetch_spec = FetchSpec(
+            depth_min=expedition.instruments_config.ctd_config.min_depth_meter,
+            depth_max=expedition.instruments_config.ctd_config.max_depth_meter,
+        )
 
         super().__init__(
             expedition,
             variables,
             add_bathymetry=True,
             verbose_progress=False,
-            fetch_spec=FetchSpec(),
+            fetch_spec=fetch_spec,
             from_data=from_data,
         )
 

src/virtualship/instruments/xbt.py

@@ -89,13 +89,17 @@ class XBTInstrument(Instrument):
     def __init__(self, expedition, from_data):
         """Initialize XBTInstrument."""
         variables = expedition.instruments_config.xbt_config.active_variables()
+        fetch_spec = FetchSpec(
+            depth_min=expedition.instruments_config.xbt_config.min_depth_meter,
+            depth_max=expedition.instruments_config.xbt_config.max_depth_meter,
+        )
 
         super().__init__(
             expedition,
             variables,
             add_bathymetry=True,
             verbose_progress=False,
-            fetch_spec=FetchSpec(),
+            fetch_spec=fetch_spec,
             from_data=from_data,
         )
 

src/virtualship/utils.py

@@ -43,7 +43,7 @@
 # projection used to sail between waypoints
 PROJECTION = pyproj.Geod(ellps="WGS84")
 
-# caching for problems module
+# problems module
 CACHE = "cache"
 EXPEDITION_IDENTIFIER = "id_latest.txt"
 PROBLEMS_ENCOUNTERED = "problems_encountered_" + "{expedition_id}"
@@ -53,6 +53,9 @@
 EXPEDITION_ORIGINAL = "expedition_original.yaml"
 EXPEDITION_LATEST = "expedition_latest.yaml"
 
+# Parcels cacheing
+MAX_CACHE_BYTES = 300_000_000  # [Bytes per variable]
+
 
 # =====================================================
 # SECTION: Copernicus Marine Service constants

tests/instruments/test_base.py

@@ -167,8 +167,8 @@ def __init__(self, **fields):
             setattr(self, name, value)
         self.fields = {}
 
-    def to_windowed_arrays(self):
-        """Mimic FieldSet.to_windowed_arrays."""
+    def to_chunk_cached_arrays(self, **kwargs):
+        """Mimic FieldSet.to_chunk_cached_arrays."""
         return self
 
 
@@ -195,7 +195,6 @@ def test_load_input_data(mock_expedition):
             return_value="dummy_product_id",
         ),
         patch("copernicusmarine.open_dataset"),
-        patch.object(dummy, "_via_tmp_ds", side_effect=lambda ds: ds),
         patch("parcels.convert.copernicusmarine_to_sgrid"),
         patch(
             "parcels.FieldSet.from_sgrid_conventions", return_value=fake_fieldset
@@ -269,50 +268,6 @@ def test_fetch_spec_applied_to_instrument(mock_expedition):
     assert dummy.fetch_spec.depth_max is None
 
 
-def test_via_tmp_ds_roundtrip(mock_expedition):
-    """_via_tmp_ds writes to a tmp file and re-opens it."""
-    with DummyInstrument(
-        expedition=mock_expedition,
-        variables={"A": "a"},
-        add_bathymetry=False,
-        verbose_progress=False,
-        from_data=None,
-    ) as dummy:
-        ds = xr.Dataset(
-            {"temperature": (["x", "y"], [[1.0, 2.0], [3.0, 4.0]])},
-            coords={"x": [0, 1], "y": [10, 20]},
-        )
-        result = dummy._via_tmp_ds(ds)
-
-        assert isinstance(result, xr.Dataset)
-        assert "temperature" in result
-        assert result is not ds
-
-        result.close()
-        ds.close()
-
-
-def test_instrument_context_manager(mock_expedition):
-    """Test context manager cleanup of temporary directories."""
-    with DummyInstrument(
-        expedition=mock_expedition,
-        variables={"A": "a"},
-        add_bathymetry=False,
-        verbose_progress=False,
-        from_data=None,
-    ) as dummy:
-        ds = xr.Dataset(
-            {"temperature": (["x", "y"], [[1.0, 2.0], [3.0, 4.0]])},
-            coords={"x": [0, 1], "y": [10, 20]},
-        )
-        result = dummy._via_tmp_ds(ds)
-        assert len(dummy._tmp_dirs) == 1
-        result.close()
-        ds.close()
-
-    assert len(dummy._tmp_dirs) == 0
-
-
 def test_generate_fieldset_combines_fields(mock_expedition):
     dummy = DummyInstrument(
         expedition=mock_expedition,
@@ -325,12 +280,11 @@ def test_generate_fieldset_combines_fields(mock_expedition):
     fs_A = MagicMock()
     fs_B = MagicMock()
 
-    fs_A.to_windowed_arrays.return_value = fs_A
-    fs_B.to_windowed_arrays.return_value = fs_B
+    fs_A.to_chunk_cached_arrays.return_value = fs_A
+    fs_B.to_chunk_cached_arrays.return_value = fs_B
 
     with (
         patch.object(dummy, "_get_copernicus_ds"),
-        patch.object(dummy, "_via_tmp_ds"),
         patch("parcels.convert.copernicusmarine_to_sgrid"),
         patch("parcels.FieldSet.from_sgrid_conventions", side_effect=[fs_A, fs_B]),
     ):