diff --git a/src/virtualship/instruments/adcp.py b/src/virtualship/instruments/adcp.py index 26c8122d..57d2e352 100644 --- a/src/virtualship/instruments/adcp.py +++ b/src/virtualship/instruments/adcp.py @@ -3,12 +3,11 @@ from typing import ClassVar import numpy as np -from parcels import ParticleFile, ParticleSet from virtualship.instruments.base import FetchSpec, Instrument 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 _write_underway_to_parquet, register_instrument # ===================================================== # SECTION: Dataclass @@ -22,14 +21,6 @@ class ADCP: name: ClassVar[str] = "ADCP" -# ===================================================== -# SECTION: non-sensor Particle Variables (non-sampling) -# ===================================================== - -# ADCP has no non-sensor variables, only sensor variables. -_ADCP_NONSENSOR_VARIABLES: list = [] - - # ===================================================== # SECTION: Kernels # ===================================================== @@ -41,7 +32,6 @@ def _sample_velocity(particles, fieldset): particles.z, particles.x, particles.y, - applyConversion=False, ) @@ -93,47 +83,37 @@ def simulate(self, measurements, out_path) -> None: 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 + # 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 + + # times in seconds since fieldset time origin, expanded across depth bins + 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 - - 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, - ) + times_full = np.repeat(times, NUM_BINS) + lons_full = np.repeat(lons, NUM_BINS) + lats_full = np.repeat(lats, NUM_BINS) + depths_full = np.tile(bins, len(times)) + + u, v = fieldset.UV.eval(t=times_full, z=depths_full, x=lons_full, y=lats_full) + + _write_underway_to_parquet( + dat_arrays=[u, v], + var_names=self.variables.keys(), + times_full=times_full, + lons_full=lons_full, + lats_full=lats_full, + depths_full=depths_full, + fieldset_time_origin=fieldset_starttime, + out_path=out_path, + ) diff --git a/src/virtualship/utils.py b/src/virtualship/utils.py index 9cb028f3..97220e1c 100644 --- a/src/virtualship/utils.py +++ b/src/virtualship/utils.py @@ -14,6 +14,8 @@ import copernicusmarine import numpy as np import parcels +import pyarrow as pa +import pyarrow.parquet as pq import pyproj import xarray as xr from parcels import FieldSet, Particle, Variable @@ -563,6 +565,90 @@ def _find_files_in_timerange( return [fname for _, fname in files_with_dates] +def _write_underway_to_parquet( + dat_arrays: list[np.ndarray], + var_names: list[str], + times_full: np.ndarray, + lons_full: np.ndarray, + lats_full: np.ndarray, + fieldset_time_origin: np.datetime64, + out_path: Path | str, + depths_full: np.ndarray | None = None, + depth_fill: float = np.nan, + 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 downstream workflows. + """ + breakpoint() + assert len(dat_arrays) == len(var_names), ( + "dat_arrays and var_names must have the same length" + ) + + n = len(times_full) + + 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()), + pa.field("dt", pa.float64()), + pa.field("state", pa.int32()), + ], + 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", # TODO: is the case as long as using Copernicus Marine data... + }, + ) + + 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}") + + if depths_full is None: + depths_full = np.full(n, depth_fill, dtype=np.float32) + + # build table with all data, including sampled variables + table = pa.table( + { + "t": pa.array(times_full.astype(np.float64)), + "z": pa.array(depths_full.astype(np.float32)) + if depths_full is not None + else pa.array(np.full(n, depth_fill, dtype=np.float32)), + "y": pa.array(lats_full.astype(np.float32)), + "x": pa.array(lons_full.astype(np.float32)), + "particle_id": pa.array( + np.arange(n, dtype=np.int64) + ), # doesn't necessarily correspond to an actual particle_id in the simulation, but required for Parcels ParticleFile schema + "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) + + def _compute_max_depths(measurements, fieldset) -> list[float]: """Compute the effective max depth for each measurement, capped by bathymetry.""" return [