diff --git a/src/virtualship/instruments/argo_float.py b/src/virtualship/instruments/argo_float.py index 96b8e4f0..029a0aec 100644 --- a/src/virtualship/instruments/argo_float.py +++ b/src/virtualship/instruments/argo_float.py @@ -324,6 +324,11 @@ def simulate(self, measurements, out_path) -> None: drift_days=[argo.drift_days for argo in measurements], ) + # add initial conditions to sampling variables + self._sample_initial( + argo_float_particleset, fieldset, argo_float_config.sensors + ) + # define output file for the simulation out_file = ParticleFile( path=out_path, diff --git a/src/virtualship/instruments/base.py b/src/virtualship/instruments/base.py index ce0f8324..a9643718 100644 --- a/src/virtualship/instruments/base.py +++ b/src/virtualship/instruments/base.py @@ -279,6 +279,28 @@ def _via_tmp_ds(ds) -> xr.Dataset: del ds return xr.open_dataset(tmp_fpath) + @staticmethod + def _sample_initial( + pset: parcels.ParticleSet, + fieldset: parcels.FieldSet, + sensors_config: object, + ) -> parcels.ParticleSet: + """Perform initial Field sampling with ParticleSet.""" + for sensor in sensors_config: + if not sensor.enabled: + raise ValueError( + f"Attempted to initialise sensor '{sensor.sensor_type}' but it is not enabled in the expedition configuration." + ) + + fs_key = sensor.meta.fs_key + field = getattr(fieldset, fs_key) + particle_vars = [pv.name for pv in sensor.meta.particle_vars] + + for var in particle_vars: + setattr(pset, var, field[pset]) + + return pset + @property def instrument_type(self) -> InstrumentType: """Return the InstrumentType for this instrument instance.""" diff --git a/src/virtualship/instruments/ctd.py b/src/virtualship/instruments/ctd.py index 5b8cba61..1291890d 100644 --- a/src/virtualship/instruments/ctd.py +++ b/src/virtualship/instruments/ctd.py @@ -209,6 +209,9 @@ def simulate(self, measurements, out_path) -> None: winch_speed=[WINCH_SPEED for _ in measurements], ) + # add initial conditions to sampling variables + self._sample_initial(ctd_particleset, fieldset, ctd_config.sensors) + # define output file for the simulation out_file = ParticleFile(path=out_path, outputdt=OUTPUT_DT) diff --git a/src/virtualship/instruments/drifter.py b/src/virtualship/instruments/drifter.py index 3b52dc71..fe8b3738 100644 --- a/src/virtualship/instruments/drifter.py +++ b/src/virtualship/instruments/drifter.py @@ -154,6 +154,9 @@ def simulate(self, measurements, out_path) -> None: ], ) + # add initial conditions to sampling variables + self._sample_initial(drifter_particleset, fieldset, drifter_config.sensors) + # define output file for the simulation out_file = ParticleFile( path=out_path, diff --git a/src/virtualship/instruments/xbt.py b/src/virtualship/instruments/xbt.py index 9de23092..c59e6996 100644 --- a/src/virtualship/instruments/xbt.py +++ b/src/virtualship/instruments/xbt.py @@ -164,6 +164,9 @@ def simulate(self, measurements, out_path) -> None: fall_speed=[xbt.fall_speed for xbt in measurements], ) + # add initial conditions to sampling variables + self._sample_initial(xbt_particleset, fieldset, xbt_config.sensors) + out_file = ParticleFile(path=out_path, outputdt=OUTPUT_DT) # build kernel list from active sensors only diff --git a/tests/instruments/test_base.py b/tests/instruments/test_base.py index e14bb630..aa725fb4 100644 --- a/tests/instruments/test_base.py +++ b/tests/instruments/test_base.py @@ -16,8 +16,48 @@ ) from virtualship.instruments.sensors import SensorType from virtualship.instruments.types import InstrumentType +from virtualship.models.expedition import SensorConfig from virtualship.utils import get_instrument_class +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture() +def fieldset(): + """Minimal Parcels FieldSet containing a temperature field.""" + 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) + return parcels.FieldSet.from_sgrid_conventions(ds_fset) + + +@pytest.fixture() +def pset(fieldset): + """Minimal ParticleSet initialized with a custom Particle class and the fieldset fixture.""" + SampleParticle = parcels.Particle.add_variable(parcels.Variable("temperature")) + t1 = np.datetime64("2024-01-01T00:00:00") + + return parcels.ParticleSet( + fieldset=fieldset, pclass=SampleParticle, t=t1, y=[0.0], x=[0.0] + ) + + # ============================================================================= # Instrument base class testing # ============================================================================= @@ -210,6 +250,24 @@ def simulate(self, data_dir, measurements, out_path): pass +def test_instrument_samples_initial_conditions(fieldset, pset): + """_sample_initial adds initial conditions to particles.""" + psetT_preinit = pset.temperature.copy() # before sampling initial conditions + + sensor_config = SensorConfig(sensor_type=SensorType.TEMPERATURE, enabled=True) + pset = Instrument._sample_initial(pset, fieldset, [sensor_config]) + + psetT_postinit = pset.temperature # once initialised + + assert not np.array_equal(psetT_preinit, psetT_postinit), ( + "Initial conditions were not added." + ) + + assert np.allclose(psetT_postinit, [15.0]), ( + "Initial conditions do not match expected values." + ) + + # ============================================================================= # UnderwayInstrument intermediate class testing # ============================================================================= @@ -378,7 +436,9 @@ def _create_underway_parquet( def dummy_sample_temperature(particles, fieldset): - particles.T = fieldset.T[particles.t, particles.z, particles.y, particles.x] + particles.temperature = fieldset.T[ + particles.t, particles.z, particles.y, particles.x + ] def test_parquet_openable_by_parcels_read_particlefile(tmp_path): @@ -400,34 +460,8 @@ def test_parquet_openable_by_parcels_read_particlefile(tmp_path): assert np.isclose(results["sal"][1], 35.1) -def test_underway_schema_matches_parcels(tmp_path): +def test_underway_schema_matches_parcels(tmp_path, pset): """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(