diff --git a/src/virtualship/instruments/adcp.py b/src/virtualship/instruments/adcp.py index 26c8122d..bd253230 100644 --- a/src/virtualship/instruments/adcp.py +++ b/src/virtualship/instruments/adcp.py @@ -1,48 +1,38 @@ from collections.abc import Callable -from dataclasses import dataclass from typing import ClassVar import numpy as np -from parcels import ParticleFile, ParticleSet +import parcels -from virtualship.instruments.base import FetchSpec, Instrument +from virtualship.instruments.base import ( + FetchSpec, + UnderwayCoordinates, + UnderwayInstrument, +) from virtualship.instruments.sensors import SensorType from virtualship.instruments.types import InstrumentType -from virtualship.utils import build_particle_class_from_sensors, register_instrument +from virtualship.utils import register_instrument # ===================================================== -# SECTION: Dataclass +# SECTION: Kernels # ===================================================== -@dataclass -class ADCP: - """ADCP configuration.""" - - name: ClassVar[str] = "ADCP" - - -# ===================================================== -# SECTION: non-sensor Particle Variables (non-sampling) -# ===================================================== - -# ADCP has no non-sensor variables, only sensor variables. -_ADCP_NONSENSOR_VARIABLES: list = [] +# N.B. underway 'kernels' are special cases, where the particleset is not needed, and the kernel is not passed to `pset.execute()` as would be done for a typical Parcels workflow. +# Instead, the 'kernel' function is used only once to evaluate the fieldset at given times, depths, lats, lons. -# ===================================================== -# SECTION: Kernels -# ===================================================== +def _sample_underway_velocity(fieldset: parcels.FieldSet, coords: UnderwayCoordinates): + # eval + u, v = fieldset.UV.eval( + t=coords.times, z=coords.depths, x=coords.lons, y=coords.lats + ) + # convert from degrees s-1 to metres s-1 + u = u * 1852 * 60 * np.cos(np.deg2rad(coords.lats)) + v = v * 1852 * 60 -def _sample_velocity(particles, fieldset): - particles.U, particles.V = fieldset.UV.eval( - particles.t, - particles.z, - particles.x, - particles.y, - applyConversion=False, - ) + return u, v # ===================================================== @@ -51,11 +41,11 @@ def _sample_velocity(particles, fieldset): @register_instrument(InstrumentType.ADCP) -class ADCPInstrument(Instrument): +class ADCPInstrument(UnderwayInstrument): """ADCP instrument class.""" sensor_kernels: ClassVar[dict[SensorType, Callable]] = { - SensorType.VELOCITY: _sample_velocity, + SensorType.VELOCITY: _sample_underway_velocity, } def __init__(self, expedition, from_data): @@ -74,9 +64,9 @@ def __init__(self, expedition, from_data): def simulate(self, measurements, out_path) -> None: """Simulate ADCP measurements.""" - config_max_depth = ( - self.expedition.instruments_config.adcp_config.max_depth_meter - ) + adcp_config = self.expedition.instruments_config.adcp_config + + config_max_depth = adcp_config.max_depth_meter if config_max_depth < -1600.0: print( @@ -87,53 +77,42 @@ def simulate(self, measurements, out_path) -> None: MAX_DEPTH = config_max_depth MIN_DEPTH = -5.0 - NUM_BINS = self.expedition.instruments_config.adcp_config.num_bins + NUM_BINS = adcp_config.num_bins measurements.sort(key=lambda p: p.time) fieldset = self.load_input_data() - # build dynamic particle class from the active sensors - adcp_config = self.expedition.instruments_config.adcp_config - _ADCPParticle = build_particle_class_from_sensors( - adcp_config.sensors, _ADCP_NONSENSOR_VARIABLES + # times in seconds since fieldset time origin, expanded across depth bins + fieldset_starttime = fieldset.time_interval.left + times = np.array( + [ + (np.datetime64(point.time) - fieldset_starttime) + / np.timedelta64(1, "s") + for point in measurements + ] ) + lons = np.array([point.location.lon for point in measurements]) + lats = np.array([point.location.lat for point in measurements]) bins = np.linspace(MAX_DEPTH, MIN_DEPTH, NUM_BINS) - num_particles = len(bins) - particleset = ParticleSet( - fieldset=fieldset, - pclass=_ADCPParticle, - x=np.full( - num_particles, 0.0 - ), # initial lat/lon are irrelevant and will be overruled later - y=np.full(num_particles, 0.0), - z=bins, - ) - - out_file = ParticleFile(path=out_path, outputdt=np.inf) - # build kernel list from active sensors only - sampling_kernels = [ - self.sensor_kernels[sc.sensor_type] - for sc in adcp_config.sensors - if sc.enabled and sc.sensor_type in self.sensor_kernels - ] - - # TODO: need to overhaul ADCP/underway instruments generally... don't think this Parcels API works anymore - # TODO: a good time to implement https://github.com/Parcels-code/virtualship/issues/231 + # full sampling coordinates + coords = UnderwayCoordinates( + times=np.repeat(times, NUM_BINS), + lons=np.repeat(lons, NUM_BINS), + lats=np.repeat(lats, NUM_BINS), + depths=np.tile(bins, len(times)), + ) - for point in measurements: - particleset.lon_nextloop[:] = point.location.lon - particleset.lat_nextloop[:] = point.location.lat - particleset.time_nextloop[:] = fieldset.time_origin.reltime( - np.datetime64(point.time) - ) + sampled = self._sample_underway( + config_sensors=adcp_config.sensors, fieldset=fieldset, coords=coords + ) - particleset.execute( - sampling_kernels, - dt=1, - runtime=1, - verbose_progress=self.verbose_progress, - output_file=out_file, - ) + self._to_parquet( + dat_arrays=sampled, + var_names=self.variables.keys(), + fieldset_time_origin=fieldset_starttime, + out_path=out_path, + coords=coords, + ) diff --git a/src/virtualship/instruments/base.py b/src/virtualship/instruments/base.py index 80deeb02..ce0f8324 100644 --- a/src/virtualship/instruments/base.py +++ b/src/virtualship/instruments/base.py @@ -2,15 +2,19 @@ import abc import collections +import inspect import tempfile from dataclasses import dataclass from datetime import timedelta from itertools import pairwise from pathlib import Path -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Literal import copernicusmarine +import numpy as np import parcels +import pyarrow as pa +import pyarrow.parquet as pq import xarray as xr from yaspin import yaspin @@ -50,8 +54,11 @@ class Instrument(abc.ABC): sensor_kernels: ClassVar[dict[SensorType, collections.abc.Callable]] def __init_subclass__(cls, **kwargs: object) -> None: - """Ensure subclasses define sensor_kernels as class attribute.""" + """Ensure non-abstract subclasses (i.e. final/concrete instrument classes) define sensor_kernels as a class attribute.""" super().__init_subclass__(**kwargs) + if inspect.isabstract(cls): + return + if "sensor_kernels" not in cls.__dict__: raise TypeError( f"Instrument subclass '{cls.__name__}' must define 'sensor_kernels' as a class attribute." @@ -132,20 +139,17 @@ def simulate( def execute(self, measurements: list, out_path: str | Path) -> None: """Run instrument simulation.""" - TMP = True # TODO: just for dev; remove before merging instrument_name = self.__class__.__name__.split("Instrument")[0] if not self.verbose_progress: - if TMP: - with yaspin( - text=f"Simulating {instrument_name} measurements... ", - side="right", - spinner=ship_spinner, - ) as spinner: - self.simulate(measurements, out_path) - spinner.ok("✅\n") - else: + with yaspin( + text=f"Simulating {instrument_name} measurements... ", + side="right", + spinner=ship_spinner, + ) as spinner: self.simulate(measurements, out_path) + spinner.ok("✅\n") + else: print(f"Simulating {instrument_name} measurements... ") self.simulate(measurements, out_path) @@ -279,3 +283,141 @@ def _via_tmp_ds(ds) -> xr.Dataset: def instrument_type(self) -> InstrumentType: """Return the InstrumentType for this instrument instance.""" return next(k for k, v in INSTRUMENT_CLASS_MAP.items() if type(self) is v) + + +@dataclass(frozen=True) +class UnderwayCoordinates: + """1D sampling location arrays for underway instruments.""" + + times: np.ndarray # seconds since origin + lons: np.ndarray + lats: np.ndarray + depths: np.ndarray + + def __post_init__(self): + """Validate that all arrays are 1D and have the same length.""" + shapes = { + "times": self.times.shape, + "lons": self.lons.shape, + "lats": self.lats.shape, + "depths": self.depths.shape, + } + + for name, shape in shapes.items(): + if len(shape) != 1: + raise ValueError(f"Array '{name}' must be 1D, but got shape {shape}.") + + n = len(self.times) + if not (len(self.lons) == len(self.lats) == len(self.depths) == n): + raise ValueError( + f"Array length mismatch in UnderwayCoordinates: " + f"times={len(self.times)}, lons={len(self.lons)}, " + f"lats={len(self.lats)}, depths={len(self.depths)}" + ) + + +class UnderwayInstrument(Instrument): + """Intermediate base class for underway instruments, which perform variable sampling without ParticleSets.""" + + def _sample_underway( + self, + config_sensors: list, + fieldset: parcels.FieldSet, + coords: UnderwayCoordinates, + ): + """Perform variable sampling for underway instruments and their active sensors.""" + sampling_kernels = [ + self.sensor_kernels[sc.sensor_type] + for sc in config_sensors + if sc.enabled and sc.sensor_type in self.sensor_kernels + ] # active sensors only + + sampled = [ + kernel(fieldset, coords) for kernel in sampling_kernels + ] # perform sampling + + # ensure that sampled is a flat list of arrays, even if some kernels return tuples/lists of arrays + # e.g. ADCP kernel returns (u, v) tuple of arrays, whilst UnderwaterST returns single array of temperature/salinity + sampled_flat = [ + arr + for item in sampled + for arr in (item if isinstance(item, (tuple, list)) else (item,)) + ] + + return sampled_flat + + @staticmethod + def _to_parquet( + dat_arrays: list[np.ndarray], + var_names: list[str], + fieldset_time_origin: np.datetime64, + out_path: Path | str, + coords: UnderwayCoordinates, + compression: Literal["zstd", "gzip", "snappy", "brotli", None] = "zstd", + ) -> None: + """ + Write underway instrument data to a Parquet file mirroring the Parcels v4 ParticleFile schema. + + Designed so that output files can be re-read back in with Parcels.read_particlefile for consistent downstream workflows with non-underway instruments. + """ + assert len(dat_arrays) == len(var_names), ( + "dat_arrays and var_names must have the same length" + ) + + n = len(coords.times) + + origin_str = str(fieldset_time_origin).replace("T", " ") + t_metadata = {"units": f"seconds since {origin_str}", "calendar": "standard"} + + # base schema mirroring Parcels ParticleFile schema, not yet with sampled variables + base_schema = pa.schema( + [ + pa.field("t", pa.float64(), metadata=t_metadata), + pa.field("z", pa.float32()), + pa.field("y", pa.float32()), + pa.field("x", pa.float32()), + pa.field("particle_id", pa.int64()), + ], + metadata={ + "feature_type": "trajectory", + "Conventions": "CF-1.6/CF-1.7", + "ncei_template_version": "NCEI_NetCDF_Trajectory_Template_v2.0", + "parcels_version": parcels.__version__, + "parcels_grid_mesh": "spherical", + }, + ) + + for var in var_names: + base_schema = base_schema.append( + pa.field(var, pa.float32()) + ) # add sampled variable to schema + + out_path = Path(out_path) + if out_path.suffix != ".parquet": + raise ValueError( + f"out_path must end in '.parquet', got {out_path.suffix!r}" + ) + + # build table with all data, including sampled variables + table = pa.table( + { + "t": pa.array(coords.times.astype(np.float64)), + "z": pa.array(coords.depths.astype(np.float32)) + if coords.depths is not None + else pa.array(np.full(n, np.nan, dtype=np.float32)), + "y": pa.array(coords.lats.astype(np.float32)), + "x": pa.array(coords.lons.astype(np.float32)), + "particle_id": pa.array( + np.zeros(n, dtype=np.int64) + ), # ship is a single 'particle' (here represented by a constant particle_id of 0) + "dt": pa.array(np.full(n, np.nan, dtype=np.float64)), + "state": pa.array(np.zeros(n, dtype=np.int32)), + **{ + var: pa.array(dat.astype(np.float32)) + for var, dat in zip(var_names, dat_arrays, strict=True) + }, + }, + schema=base_schema, + ) + + pq.write_table(table, out_path, compression=compression) diff --git a/src/virtualship/instruments/ctd.py b/src/virtualship/instruments/ctd.py index d6764130..5b8cba61 100644 --- a/src/virtualship/instruments/ctd.py +++ b/src/virtualship/instruments/ctd.py @@ -168,12 +168,8 @@ def simulate(self, measurements, out_path) -> None: fieldset = self.load_input_data() - # use first active field for time reference - _time_ref_key = next(iter(self.variables)) - _time_ref_field = getattr(fieldset, _time_ref_key) - - fieldset_starttime = _time_ref_field.data.time.isel(time=0).values - fieldset_endtime = _time_ref_field.data.time.isel(time=-1).values + fieldset_starttime = fieldset.time_interval.left + fieldset_endtime = fieldset.time_interval.right # deploy time for all ctds should be later than fieldset start time if not all( diff --git a/src/virtualship/instruments/ship_underwater_st.py b/src/virtualship/instruments/ship_underwater_st.py index 1dc7522a..4aff0af5 100644 --- a/src/virtualship/instruments/ship_underwater_st.py +++ b/src/virtualship/instruments/ship_underwater_st.py @@ -1,51 +1,40 @@ from collections.abc import Callable -from dataclasses import dataclass from typing import ClassVar import numpy as np -from parcels import ParticleFile, ParticleSet +import parcels -from virtualship.instruments.base import FetchSpec, Instrument +from virtualship.instruments.base import ( + FetchSpec, + UnderwayCoordinates, + UnderwayInstrument, +) from virtualship.instruments.sensors import SensorType from virtualship.instruments.types import InstrumentType from virtualship.utils import ( - build_particle_class_from_sensors, register_instrument, ) -# ===================================================== -# SECTION: Dataclass -# ===================================================== - - -@dataclass -class Underwater_ST: - """Underwater_ST configuration.""" - - name: ClassVar[str] = "Underwater_ST" - - -# ===================================================== -# SECTION: non-sensor Particle Variables (non-sampling) -# ===================================================== - -# Underwater ST has no non-sensor variables, only sensor variables. -_ST_NONSENSOR_VARIABLES: list = [] - - # ===================================================== # SECTION: Kernels # ===================================================== +# N.B. underway 'kernels' are special cases, where the particleset is not needed, and the kernel is not passed to `pset.execute()` as would be done for a typical Parcels workflow. +# Instead, the 'kernel' function is used only once to evaluate the fieldset at given times, depths, lats, lons. -# define function sampling Salinity -def _sample_salinity(particles, fieldset): - particles.S = fieldset.S[particles.t, particles.z, particles.y, particles.x] +def _sample_underway_salinity(fieldset: parcels.FieldSet, coords: UnderwayCoordinates): + return fieldset.S.eval( + t=coords.times, z=coords.depths, x=coords.lons, y=coords.lats + ) -# define function sampling Temperature -def _sample_temperature(particles, fieldset): - particles.T = fieldset.T[particles.t, particles.z, particles.y, particles.x] + +def _sample_underway_temperature( + fieldset: parcels.FieldSet, coords: UnderwayCoordinates +): + return fieldset.T.eval( + t=coords.times, z=coords.depths, x=coords.lons, y=coords.lats + ) # ===================================================== @@ -54,12 +43,12 @@ def _sample_temperature(particles, fieldset): @register_instrument(InstrumentType.UNDERWATER_ST) -class Underwater_STInstrument(Instrument): +class Underwater_STInstrument(UnderwayInstrument): """Underwater_ST instrument class.""" sensor_kernels: ClassVar[dict[SensorType, Callable]] = { - SensorType.TEMPERATURE: _sample_temperature, - SensorType.SALINITY: _sample_salinity, + SensorType.TEMPERATURE: _sample_underway_temperature, + SensorType.SALINITY: _sample_underway_salinity, } def __init__(self, expedition, from_data): @@ -80,49 +69,37 @@ def __init__(self, expedition, from_data): def simulate(self, measurements, out_path) -> None: """Simulate underway salinity and temperature measurements.""" + st_config = self.expedition.instruments_config.ship_underwater_st_config + DEPTH = -2.0 measurements.sort(key=lambda p: p.time) fieldset = self.load_input_data() - # build dynamic particle class from the active sensors - st_config = self.expedition.instruments_config.ship_underwater_st_config - _ShipSTParticle = build_particle_class_from_sensors( - st_config.sensors, _ST_NONSENSOR_VARIABLES + # sampling times and locations + fieldset_starttime = fieldset.time_interval.left + times = np.array( + [ + (np.datetime64(point.time) - fieldset_starttime) + / np.timedelta64(1, "s") + for point in measurements + ] + ) + lons = np.array([point.location.lon for point in measurements]) + lats = np.array([point.location.lat for point in measurements]) + coords = UnderwayCoordinates( + times, lons, lats, depths=np.full_like(times, DEPTH) ) - particleset = ParticleSet( - fieldset=fieldset, - pclass=_ShipSTParticle, - x=0.0, - y=0.0, - z=DEPTH, + sampled = self._sample_underway( + config_sensors=st_config.sensors, fieldset=fieldset, coords=coords ) - out_file = ParticleFile(path=out_path, outputdt=np.inf) - - # build kernel list from active sensors only - sampling_kernels = [ - self.sensor_kernels[sc.sensor_type] - for sc in st_config.sensors - if sc.enabled and sc.sensor_type in self.sensor_kernels - ] - - # TODO: need to overhaul UNDERWATER_ST/underway instruments generally... don't think this Parcels API works anymore - # TODO: a good time to implement https://github.com/Parcels-code/virtualship/issues/231 - - for point in measurements: - particleset.lon_nextloop[:] = point.location.lon - particleset.lat_nextloop[:] = point.location.lat - particleset.time_nextloop[:] = fieldset.time_origin.reltime( - np.datetime64(point.time) - ) - - particleset.execute( - sampling_kernels, - dt=1, - runtime=1, - verbose_progress=self.verbose_progress, - output_file=out_file, - ) + self._to_parquet( + dat_arrays=sampled, + var_names=self.variables.keys(), + fieldset_time_origin=fieldset_starttime, + out_path=out_path, + coords=coords, + ) diff --git a/src/virtualship/instruments/xbt.py b/src/virtualship/instruments/xbt.py index 7046fde1..9de23092 100644 --- a/src/virtualship/instruments/xbt.py +++ b/src/virtualship/instruments/xbt.py @@ -120,12 +120,8 @@ def simulate(self, measurements, out_path) -> None: fieldset = self.load_input_data() - # use first active field for time reference - _time_ref_key = next(iter(self.variables)) - _time_ref_field = getattr(fieldset, _time_ref_key) - - fieldset_starttime = _time_ref_field.data.time.isel(time=0).values - fieldset_endtime = _time_ref_field.data.time.isel(time=-1).values + fieldset_starttime = fieldset.time_interval.left + fieldset_endtime = fieldset.time_interval.right # deploy time for all xbts should be later than fieldset start time if not all( diff --git a/tests/instruments/test_adcp.py b/tests/instruments/test_adcp.py index 16604b67..ab13f99c 100644 --- a/tests/instruments/test_adcp.py +++ b/tests/instruments/test_adcp.py @@ -4,10 +4,11 @@ from typing import ClassVar import numpy as np +import parcels +import polars as pl import pydantic import pytest import xarray as xr -from parcels import FieldSet from virtualship.instruments.adcp import ADCPInstrument from virtualship.instruments.sensors import SensorType @@ -22,6 +23,7 @@ BASE_TIME = datetime.datetime.strptime( "1950-01-01", "%Y-%m-%d" ) # arbitrary time offset for the dummy fieldset +MIN_DEPTH = -5 MAX_DEPTH = -1000 NUM_BINS = 40 @@ -53,8 +55,6 @@ class schedule: def test_simulate_adcp(tmpdir, adcp_expedition) -> None: - MIN_DEPTH = -5 - # where to sample sample_points = [ Spacetime(Location(1, 2), BASE_TIME + datetime.timedelta(seconds=0)), @@ -93,53 +93,58 @@ def test_simulate_adcp(tmpdir, adcp_expedition) -> None: u[1, 0, 1, 1] = expected_obs[1]["U"]["max_depth"] u[1, 1, 1, 1] = expected_obs[1]["U"]["surface"] - fieldset = FieldSet.from_data( - { - "V": v, - "U": u, + # make ds + times = np.array([expected_obs[0]["time"], expected_obs[1]["time"]]) + lats = np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]) + lons = np.array([expected_obs[0]["lon"], expected_obs[1]["lon"]]) + + ds_fields = xr.Dataset( + data_vars={ + "U": (["time", "depth", "lat", "lon"], u, {"units": "m s-1"}), + "V": (["time", "depth", "lat", "lon"], v, {"units": "m s-1"}), }, - { - "lat": np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]), - "lon": np.array([expected_obs[0]["lon"], expected_obs[1]["lon"]]), - "depth": np.array([MAX_DEPTH, MIN_DEPTH]), - "time": np.array( - [ - np.datetime64(expected_obs[0]["time"]), - np.datetime64(expected_obs[1]["time"]), - ] - ), + coords={ + "time": ("time", times, {"axis": "T"}), + "depth": ("depth", [MAX_DEPTH, MIN_DEPTH], {"units": "m", "axis": "Z"}), + "lat": ("lat", lats, {"units": "degrees_north"}), + "lon": ("lon", lons, {"units": "degrees_east"}), }, ) + # to fieldset + fields = {"U": ds_fields["U"], "V": ds_fields["V"]} + ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields) + fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) + adcp_instrument = ADCPInstrument(adcp_expedition, from_data=None) - out_path = tmpdir.join("out.zarr") + out_path = tmpdir.join("out.parquet") adcp_instrument.load_input_data = lambda: fieldset adcp_instrument.simulate(sample_points, out_path) - results = xr.open_zarr(out_path) + results = parcels.read_particlefile(out_path) - # test if output is as expected - assert len(results.trajectory) == NUM_BINS + assert np.unique(results["z"].to_numpy()).size == NUM_BINS # for every obs, check if the variables match the expected observations # we only verify at the surface and max depth of the adcp, because in between is tricky - for traj, vert_loc in [ - (results.trajectory[0], "max_depth"), - (results.trajectory[-1], "surface"), + for df_depth, vert_loc in [ + (results.filter(pl.col("z") == MAX_DEPTH), "max_depth"), + (results.filter(pl.col("z") == MIN_DEPTH), "surface"), ]: - obs_all = results.sel(trajectory=traj).obs - assert len(obs_all) == len(sample_points) - for i, (obs_i, exp) in enumerate(zip(obs_all, expected_obs, strict=True)): - obs = results.sel(trajectory=traj, obs=obs_i) - for var in ["lat", "lon"]: - obs_value = obs[var].values.item() - exp_value = exp[var] + assert len(df_depth) == len(sample_points) + + for i, (obs_i, exp) in enumerate( + zip(df_depth.iter_rows(named=True), expected_obs, strict=True) + ): + for var in [("y", "lat"), ("x", "lon")]: + obs_value = obs_i[var[0]] + exp_value = exp[var[1]] assert np.isclose(obs_value, exp_value), ( f"Observation incorrect {vert_loc=} {obs_i=} {var=} {obs_value=} {exp_value=}." ) for var in ["V", "U"]: - obs_value = obs[var].values.item() + obs_value = obs_i[var] exp_value = exp[var][vert_loc] assert np.isclose(obs_value, exp_value), ( f"Observation incorrect {vert_loc=} {i=} {var=} {obs_value=} {exp_value=}." diff --git a/tests/instruments/test_base.py b/tests/instruments/test_base.py index 93a38e90..e14bb630 100644 --- a/tests/instruments/test_base.py +++ b/tests/instruments/test_base.py @@ -1,12 +1,27 @@ +from dataclasses import dataclass +from typing import ClassVar from unittest.mock import MagicMock, patch +import numpy as np +import parcels +import pyarrow.parquet as pq import pytest import xarray as xr -from virtualship.instruments.base import FetchSpec, Instrument +from virtualship.instruments.base import ( + FetchSpec, + Instrument, + UnderwayCoordinates, + UnderwayInstrument, +) +from virtualship.instruments.sensors import SensorType from virtualship.instruments.types import InstrumentType from virtualship.utils import get_instrument_class +# ============================================================================= +# Instrument base class testing +# ============================================================================= + def test_FetchSpec(): fetch_spec = FetchSpec() @@ -193,3 +208,245 @@ def test_instrument_subclass_without_sensor_kernels_error(): class ErrorInstrument(Instrument): def simulate(self, data_dir, measurements, out_path): pass + + +# ============================================================================= +# UnderwayInstrument intermediate class testing +# ============================================================================= + + +@dataclass +class DummySensorConfig: + """Mock sensor configuration.""" + + sensor_type: SensorType + enabled: bool = True + + +class ConcreteUnderwayInstrument(UnderwayInstrument): + """Concrete subclass of UnderwayInstrument for testing.""" + + sensor_kernels: ClassVar = { + SensorType.TEMPERATURE: lambda fieldset, coords: np.array( + [15.0, 16.0], dtype=np.float32 + ), + SensorType.SALINITY: lambda fieldset, coords: np.array( + [35.0, 35.1], dtype=np.float32 + ), + SensorType.VELOCITY: lambda fieldset, coords: ( + np.array([0.5, 0.6], dtype=np.float32), # U vector component + np.array([-0.1, -0.2], dtype=np.float32), # V vector component + ), + } + + def simulate(self, measurements, out_path) -> None: # noqa + pass + + +@pytest.fixture +def sample_underway_coords(): + """Fixture providing valid 1D UnderwayCoordinates.""" + return UnderwayCoordinates( + times=np.array([0.0, 3600.0]), + lons=np.array([-5.0, -5.1]), + lats=np.array([50.0, 50.1]), + depths=np.array([-2.0, -2.0]), + ) + + +@pytest.fixture +def dummy_underway_inst(): + """Bypass __init__ and requirements for expedition object etc. for testing.""" + return ConcreteUnderwayInstrument.__new__(ConcreteUnderwayInstrument) + + +def test_underway_coordinates_validation(): + """UnderwayCoordinates validates array lengths upon instantiation.""" + # valid coordinates work cleanly + coords = UnderwayCoordinates( + times=np.array([0.0, 1.0]), + lons=np.array([10.0, 11.0]), + lats=np.array([20.0, 21.0]), + depths=np.array([-1.0, -1.0]), + ) + assert len(coords.times) == 2 + + # mismatched array lengths raise ValueError + with pytest.raises(ValueError, match="Array length mismatch"): + UnderwayCoordinates( + times=np.array([0.0, 1.0]), + lons=np.array([10.0]), # length 1 vs 2 + lats=np.array([20.0, 21.0]), + depths=np.array([-1.0, -1.0]), + ) + + +def test_sample_underway_filters_and_flattens( + dummy_underway_inst, sample_underway_coords +): + """_sample_underway evaluates active sensors and flattens multi-output tuple kernels.""" + configs = [ + DummySensorConfig(SensorType.VELOCITY, enabled=True), + DummySensorConfig(SensorType.TEMPERATURE, enabled=True), + DummySensorConfig(SensorType.SALINITY, enabled=True), + ] + + sampled = dummy_underway_inst._sample_underway( + config_sensors=configs, + fieldset=None, + coords=sample_underway_coords, + ) + + assert len(sampled) == 4 # total flattened arrays (u, v, temp, sal) + np.testing.assert_array_equal(sampled[0], np.array([0.5, 0.6], dtype=np.float32)) + np.testing.assert_array_equal(sampled[1], np.array([-0.1, -0.2], dtype=np.float32)) + np.testing.assert_array_equal(sampled[2], np.array([15.0, 16.0], dtype=np.float32)) + np.testing.assert_array_equal(sampled[3], np.array([35.0, 35.1], dtype=np.float32)) + + +def test_to_parquet_writes_valid_file( + tmp_path, dummy_underway_inst, sample_underway_coords +): + """_to_parquet writes a valid Parquet table with expected schema metadata and data values.""" + out_path = tmp_path / "output.parquet" + dat_arrays = [ + np.array([15.0, 16.0], dtype=np.float32), + np.array([35.0, 35.1], dtype=np.float32), + ] + var_names = ["temp", "sal"] + origin = np.datetime64("2026-01-01T00:00:00") + + dummy_underway_inst._to_parquet( + dat_arrays=dat_arrays, + var_names=var_names, + fieldset_time_origin=origin, + out_path=out_path, + coords=sample_underway_coords, + ) + + assert out_path.exists() + + # verify parquet table, metadata, and columns + table = pq.read_table(out_path) + schema = table.schema + + assert table.column_names == [ + "t", + "z", + "y", + "x", + "particle_id", + "temp", + "sal", + ] + assert schema.metadata[b"feature_type"] == b"trajectory" + assert b"units" in schema.field("t").metadata + + np.testing.assert_array_equal( + table["x"].to_numpy(), np.array(sample_underway_coords.lons, dtype=np.float32) + ) + np.testing.assert_array_equal(table["temp"].to_numpy(), dat_arrays[0]) + + +def _create_underway_parquet( + out_path, + var_names, + dat_arrays=None, + origin=np.datetime64("2026-01-01T00:00:00"), # noqa +): + """Helper to generate an UnderwayInstrument parquet output file.""" + coords = UnderwayCoordinates( + times=np.array([0.0, 3600.0]), + lons=np.array([-5.0, -5.1]), + lats=np.array([50.0, 50.1]), + depths=np.array([-2.0, -2.0]), + ) + + if dat_arrays is None: + dat_arrays = [ + np.array([15.0, 16.0], dtype=np.float32), + np.array([35.0, 35.1], dtype=np.float32), + ] + + UnderwayInstrument._to_parquet( + dat_arrays=dat_arrays, + var_names=var_names, + fieldset_time_origin=origin, + out_path=out_path, + coords=coords, + ) + + +def dummy_sample_temperature(particles, fieldset): + particles.T = fieldset.T[particles.t, particles.z, particles.y, particles.x] + + +def test_parquet_openable_by_parcels_read_particlefile(tmp_path): + """Test that a parquet file written by _to_parquet can be read back by parcels.read_particlefile.""" + parquet_path = tmp_path / "test_particles.parquet" + _create_underway_parquet( + out_path=parquet_path, + var_names=["temp", "sal"], + dat_arrays=[ + np.array([15.0, 16.0], dtype=np.float32), + np.array([35.0, 35.1], dtype=np.float32), + ], + ) + + # read back and assert values + results = parcels.read_particlefile(parquet_path) + assert len(results) == 2 + assert np.isclose(results["temp"][0], 15.0) + assert np.isclose(results["sal"][1], 35.1) + + +def test_underway_schema_matches_parcels(tmp_path): + """Verify that underway instrument parquet output base schema matches Parcels' ParticleFile.""" + # minimal Parcels FieldSet + T = np.zeros((2, 1, 1)) + T[0, 0, 0], T[1, 0, 0] = 15.0, 16.0 + + t1 = np.datetime64("2024-01-01T00:00:00") + t2 = np.datetime64("2024-01-02T00:00:00") + ds_fields = xr.Dataset( + data_vars={"temperature": (["time", "lat", "lon"], T, {"units": "degC"})}, + coords={ + "time": ("time", [t1, t2], {"axis": "T"}), + "lat": ("lat", [0.0], {"units": "degrees_north"}), + "lon": ("lon", [0.0], {"units": "degrees_east"}), + }, + ) + + fields = {"T": ds_fields["temperature"]} + ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields) + fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) + + # parcels simualtion + SampleParticle = parcels.Particle.add_variable(parcels.Variable("T")) + + pset = parcels.ParticleSet( + fieldset=fieldset, pclass=SampleParticle, t=t1, y=[0.0], x=[0.0] + ) + + parcels_path = tmp_path / "parcels_particles.parquet" + parcels_output = parcels.ParticleFile(parcels_path, outputdt=3600.0) + pset.execute( + [dummy_sample_temperature], + runtime=np.timedelta64(60, "m"), + dt=np.timedelta64(60, "m"), + output_file=parcels_output, + ) + parcels_df = parcels.read_particlefile(parcels_path) + + # UnderwayInstrument output + underway_path = tmp_path / "underway_particles.parquet" + _create_underway_parquet( + out_path=underway_path, + var_names=["T"], + dat_arrays=[np.array([15.0, 16.0], dtype=np.float32)], + origin=np.datetime64("2024-01-01T00:00:00"), + ) + underway_df = parcels.read_particlefile(underway_path) + + # assert schemas match + assert parcels_df.schema == underway_df.schema diff --git a/tests/instruments/test_ship_underwater_st.py b/tests/instruments/test_ship_underwater_st.py index 016734a1..1bf1689c 100644 --- a/tests/instruments/test_ship_underwater_st.py +++ b/tests/instruments/test_ship_underwater_st.py @@ -4,10 +4,10 @@ from typing import ClassVar import numpy as np +import parcels import pydantic import pytest import xarray as xr -from parcels import FieldSet from virtualship.instruments.sensors import SensorType from virtualship.instruments.ship_underwater_st import Underwater_STInstrument @@ -54,87 +54,88 @@ class schedule: def test_simulate_ship_underwater_st(tmpdir, underwater_st_expedition) -> None: - # arbitrary time offset for the dummy fieldset - base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d") - # where to sample sample_points = [ - Spacetime(Location(1, 2), base_time + datetime.timedelta(seconds=0)), - Spacetime(Location(3, 4), base_time + datetime.timedelta(seconds=1)), + Spacetime(Location(1, 2), BASE_TIME + datetime.timedelta(seconds=0)), + Spacetime(Location(3, 4), BASE_TIME + datetime.timedelta(seconds=1)), ] # expected observations at sample points expected_obs = [ { - "salinity": 5, - "temperature": 6, + "S": 5, + "T": 6, "lat": sample_points[0].location.lat, "lon": sample_points[0].location.lon, - "time": base_time + datetime.timedelta(seconds=0), + "time": BASE_TIME + datetime.timedelta(seconds=0), }, { - "salinity": 7, - "temperature": 8, + "S": 7, + "T": 8, "lat": sample_points[1].location.lat, "lon": sample_points[1].location.lon, - "time": base_time + datetime.timedelta(seconds=1), + "time": BASE_TIME + datetime.timedelta(seconds=1), }, ] # create fieldset based on the expected observations # indices are time, latitude, longitude salinity = np.zeros((2, 2, 2)) - salinity[0, 0, 0] = expected_obs[0]["salinity"] - salinity[1, 1, 1] = expected_obs[1]["salinity"] + salinity[0, 0, 0] = expected_obs[0]["S"] + salinity[1, 1, 1] = expected_obs[1]["S"] temperature = np.zeros((2, 2, 2)) - temperature[0, 0, 0] = expected_obs[0]["temperature"] - temperature[1, 1, 1] = expected_obs[1]["temperature"] - - fieldset = FieldSet.from_data( - { - "V": np.zeros((2, 2, 2)), - "U": np.zeros((2, 2, 2)), - "S": salinity, - "T": temperature, + temperature[0, 0, 0] = expected_obs[0]["T"] + temperature[1, 1, 1] = expected_obs[1]["T"] + + # make ds + times = np.array([expected_obs[0]["time"], expected_obs[1]["time"]]) + lats = np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]) + lons = np.array([expected_obs[0]["lon"], expected_obs[1]["lon"]]) + + ds_fields = xr.Dataset( + data_vars={ + "T": (["time", "lat", "lon"], temperature, {"units": "degC"}), + "S": (["time", "lat", "lon"], salinity, {"units": "psu"}), }, - { - "lat": np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]), - "lon": np.array([expected_obs[0]["lon"], expected_obs[1]["lon"]]), - "time": np.array( - [ - np.datetime64(expected_obs[0]["time"]), - np.datetime64(expected_obs[1]["time"]), - ] - ), + coords={ + "time": ("time", times, {"axis": "T"}), + "lat": ("lat", lats, {"units": "degrees_north"}), + "lon": ("lon", lons, {"units": "degrees_east"}), }, ) - from_data = None + # to fieldset + fields = {"T": ds_fields["T"], "S": ds_fields["S"]} + ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields) + fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) - st_instrument = Underwater_STInstrument(underwater_st_expedition, from_data) - out_path = tmpdir.join("out.zarr") + st_instrument = Underwater_STInstrument(underwater_st_expedition, from_data=None) + out_path = tmpdir.join("out.parquet") st_instrument.load_input_data = lambda: fieldset - # The instrument expects measurements as sample_points st_instrument.simulate(sample_points, out_path) - # test if output is as expected - results = xr.open_zarr(out_path) + results = parcels.read_particlefile(out_path) - assert len(results.trajectory) == 1 # expect a single trajectory - traj = results.trajectory.item() - assert len(results.sel(trajectory=traj).obs) == len( - sample_points - ) # expect as many obs as sample points + # expect a single depth level + assert np.unique(results["z"].to_numpy()).size == 1 + + # expect as many obs as sample points (given the period is 5 minutes and the sample points are 1 second apart) + assert len(results) == len(sample_points) # for every obs, check if the variables match the expected observations for i, (obs_i, exp) in enumerate( - zip(results.sel(trajectory=traj).obs, expected_obs, strict=True) + zip(results.iter_rows(named=True), expected_obs, strict=True) ): - obs = results.sel(trajectory=traj, obs=obs_i) - for var in ["salinity", "temperature", "lat", "lon"]: - obs_value = obs[var].values.item() + for var in [("y", "lat"), ("x", "lon")]: + obs_value = obs_i[var[0]] + exp_value = exp[var[1]] + assert np.isclose(obs_value, exp_value), ( + f"Observation incorrect {obs_i=} {var=} {obs_value=} {exp_value=}." + ) + for var in ["T", "S"]: + obs_value = obs_i[var] exp_value = exp[var] assert np.isclose(obs_value, exp_value), ( f"Observation incorrect {i=} {var=} {obs_value=} {exp_value=}."