diff --git a/tests/instruments/test_argo.py b/tests/instruments/test_argo.py deleted file mode 100644 index 1b6493ea..00000000 --- a/tests/instruments/test_argo.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Test the simulation of Argo floats.""" - -from datetime import timedelta - -import numpy as np -from parcels import FieldSet - -from virtual_ship import Location, Spacetime -from virtual_ship.instruments.argo_float import ArgoFloat, simulate_argo_floats - - -def test_simulate_argo_floats() -> None: - DRIFT_DEPTH = -1000 - MAX_DEPTH = -2000 - VERTICAL_SPEED = -0.10 - CYCLE_DAYS = 10 - DRIFT_DAYS = 9 - - fieldset = FieldSet.from_data( - {"U": 0, "V": 0, "T": 0, "S": 0}, - { - "lon": 0, - "lat": 0, - "time": [np.datetime64("1950-01-01") + np.timedelta64(632160, "h")], - }, - ) - - min_depth = -fieldset.U.depth[0] - - argo_floats = [ - ArgoFloat( - spacetime=Spacetime(location=Location(latitude=0, longitude=0), time=0), - min_depth=min_depth, - max_depth=MAX_DEPTH, - drift_depth=DRIFT_DEPTH, - vertical_speed=VERTICAL_SPEED, - cycle_days=CYCLE_DAYS, - drift_days=DRIFT_DAYS, - ) - ] - - simulate_argo_floats( - argo_floats=argo_floats, - fieldset=fieldset, - out_file_name="test", - outputdt=timedelta(minutes=5), - ) diff --git a/tests/instruments/test_argo_float.py b/tests/instruments/test_argo_float.py new file mode 100644 index 00000000..b25e80e7 --- /dev/null +++ b/tests/instruments/test_argo_float.py @@ -0,0 +1,74 @@ +"""Test the simulation of Argo floats.""" + +from datetime import datetime, timedelta + +import numpy as np +import py +import xarray as xr +from parcels import FieldSet + +from virtual_ship import Location, Spacetime +from virtual_ship.instruments.argo_float import ArgoFloat, simulate_argo_floats + + +def test_simulate_argo_floats(tmpdir: py.path.LocalPath) -> None: + # arbitrary time offset for the dummy fieldset + base_time = datetime.strptime("1950-01-01", "%Y-%m-%d") + + DRIFT_DEPTH = -1000 + MAX_DEPTH = -2000 + VERTICAL_SPEED = -0.10 + CYCLE_DAYS = 10 + DRIFT_DAYS = 9 + + CONST_TEMPERATURE = 1.0 # constant temperature in fieldset + CONST_SALINITY = 1.0 # constant salinity in fieldset + + v = np.full((2, 2, 2), 1.0) + u = np.full((2, 2, 2), 1.0) + t = np.full((2, 2, 2), CONST_TEMPERATURE) + s = np.full((2, 2, 2), CONST_SALINITY) + + fieldset = FieldSet.from_data( + {"V": v, "U": u, "T": t, "S": s}, + { + "lon": np.array([0.0, 10.0]), + "lat": np.array([0.0, 10.0]), + "time": [ + np.datetime64(base_time + timedelta(seconds=0)), + np.datetime64(base_time + timedelta(hours=4)), + ], + }, + ) + + # argo floats to deploy + argo_floats = [ + ArgoFloat( + spacetime=Spacetime(location=Location(latitude=0, longitude=0), time=0), + min_depth=0.0, + max_depth=MAX_DEPTH, + drift_depth=DRIFT_DEPTH, + vertical_speed=VERTICAL_SPEED, + cycle_days=CYCLE_DAYS, + drift_days=DRIFT_DAYS, + ) + ] + + # perform simulation + out_path = tmpdir.join("out.zarr") + + simulate_argo_floats( + fieldset=fieldset, + out_path=out_path, + argo_floats=argo_floats, + outputdt=timedelta(minutes=5), + endtime=None, + ) + + # test if output is as expected + results = xr.open_zarr(out_path) + + # check the following variables are in the dataset + assert len(results.trajectory) == len(argo_floats) + for var in ["lon", "lat", "z", "temperature", "salinity"]: + assert var in results, f"Results don't contain {var}" diff --git a/virtual_ship/instruments/argo_float.py b/virtual_ship/instruments/argo_float.py index cd014671..031c8309 100644 --- a/virtual_ship/instruments/argo_float.py +++ b/virtual_ship/instruments/argo_float.py @@ -2,9 +2,10 @@ import math from dataclasses import dataclass -from datetime import timedelta +from datetime import datetime, timedelta import numpy as np +import py from parcels import ( AdvectionRK4, FieldSet, @@ -50,7 +51,7 @@ class ArgoFloat: def _argo_float_vertical_movement(particle, fieldset, time): if particle.cycle_phase == 0: # Phase 0: Sinking with vertical_speed until depth is drift_depth - particle_ddepth += ( # noqa See comment above about particle_* variables. + particle_ddepth += ( # noqa Parcels defines particle_* variables, which code checkers cannot know. particle.vertical_speed * particle.dt ) if particle.depth + particle_ddepth <= particle.drift_depth: @@ -115,31 +116,31 @@ def _check_error(particle, fieldset, time): def simulate_argo_floats( - argo_floats: list[ArgoFloat], fieldset: FieldSet, - out_file_name: str, + out_path: str | py.path.LocalPath, + argo_floats: list[ArgoFloat], outputdt: timedelta, + endtime: datetime | None, ) -> None: """ Use parcels to simulate a set of Argo floats in a fieldset. - :param argo_floats: A list of Argo floats to simulate. :param fieldset: The fieldset to simulate the Argo floats in. - :param out_file_name: The file to write the results to. + :param out_path: The path to write the results to. + :param argo_floats: A list of Argo floats to simulate. :param outputdt: Interval which dictates the update frequency of file output during simulation + :param endtime: Stop at this time, or if None, continue until the end of the fieldset. """ - lon = [argo.spacetime.location.lon for argo in argo_floats] - lat = [argo.spacetime.location.lat for argo in argo_floats] - time = [argo.spacetime.time for argo in argo_floats] + DT = 10.0 # dt of Argo float simulation integrator # define parcel particles argo_float_particleset = ParticleSet( fieldset=fieldset, pclass=_ArgoParticle, - lon=lon, - lat=lat, + lat=[argo.spacetime.location.lat for argo in argo_floats], + lon=[argo.spacetime.location.lon for argo in argo_floats], depth=[argo.min_depth for argo in argo_floats], - time=time, + time=[argo.spacetime.time for argo in argo_floats], min_depth=[argo.min_depth for argo in argo_floats], max_depth=[argo.max_depth for argo in argo_floats], drift_depth=[argo.drift_depth for argo in argo_floats], @@ -149,14 +150,17 @@ def simulate_argo_floats( ) # define output file for the simulation - out_file = argo_float_particleset.ParticleFile( - name=out_file_name, - outputdt=outputdt, - chunks=(1, 500), - ) + out_file = argo_float_particleset.ParticleFile(name=out_path, outputdt=outputdt) - # get time when the fieldset ends + # get earliest between fieldset end time and provide end time fieldset_endtime = fieldset.time_origin.fulltime(fieldset.U.grid.time_full[-1]) + if endtime is None: + actual_endtime = fieldset_endtime + elif endtime > fieldset_endtime: + print("WARN: Requested end time later than fieldset end time.") + actual_endtime = fieldset_endtime + else: + actual_endtime = np.timedelta64(endtime) # execute simulation argo_float_particleset.execute( @@ -166,7 +170,8 @@ def simulate_argo_floats( _keep_at_surface, _check_error, ], - endtime=fieldset_endtime, - dt=outputdt, + endtime=actual_endtime, + dt=DT, output_file=out_file, + verbose_progress=False, ) diff --git a/virtual_ship/sailship.py b/virtual_ship/sailship.py index d9dc4a85..b0eb3d6c 100644 --- a/virtual_ship/sailship.py +++ b/virtual_ship/sailship.py @@ -239,10 +239,11 @@ def sailship(config: VirtualShipConfiguration): print("Simulating argo floats") simulate_argo_floats( + out_path=os.path.join("results", "argo_floats.zarr"), argo_floats=argo_floats, fieldset=config.argo_float_fieldset, - out_file_name=os.path.join("results", "argo_floats.zarr"), outputdt=timedelta(minutes=5), + endtime=None, ) # convert CTD data to CSV