Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 32 additions & 52 deletions src/virtualship/instruments/adcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
# =====================================================
Expand All @@ -41,7 +32,6 @@ def _sample_velocity(particles, fieldset):
particles.z,
particles.x,
particles.y,
applyConversion=False,
)


Expand Down Expand Up @@ -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,
)
86 changes: 86 additions & 0 deletions src/virtualship/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 [
Expand Down
Loading