diff --git a/src/virtualship/cli/_run.py b/src/virtualship/cli/_run.py index 18fa922d..a73b275d 100644 --- a/src/virtualship/cli/_run.py +++ b/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): diff --git a/src/virtualship/instruments/adcp.py b/src/virtualship/instruments/adcp.py index 0b9630d3..f7fc1800 100644 --- a/src/virtualship/instruments/adcp.py +++ b/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, ) diff --git a/src/virtualship/instruments/argo_float.py b/src/virtualship/instruments/argo_float.py index 5c987145..dc9d9453 100644 --- a/src/virtualship/instruments/argo_float.py +++ b/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__( diff --git a/src/virtualship/instruments/base.py b/src/virtualship/instruments/base.py index 30f4a57a..ff974996 100644 --- a/src/virtualship/instruments/base.py +++ b/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, diff --git a/src/virtualship/instruments/ctd.py b/src/virtualship/instruments/ctd.py index 50d641fd..66c40a86 100644 --- a/src/virtualship/instruments/ctd.py +++ b/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, ) diff --git a/src/virtualship/instruments/xbt.py b/src/virtualship/instruments/xbt.py index 80e2c484..a3fa33a4 100644 --- a/src/virtualship/instruments/xbt.py +++ b/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, ) diff --git a/src/virtualship/utils.py b/src/virtualship/utils.py index d142c580..aeca9082 100644 --- a/src/virtualship/utils.py +++ b/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 diff --git a/tests/instruments/test_base.py b/tests/instruments/test_base.py index 06194c6a..2b765269 100644 --- a/tests/instruments/test_base.py +++ b/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]), ):