diff --git a/tests/instruments/test_adcp.py b/tests/instruments/test_adcp.py index fdbd50e3..d119f3e2 100644 --- a/tests/instruments/test_adcp.py +++ b/tests/instruments/test_adcp.py @@ -1,33 +1,119 @@ """Test the simulation of ADCP instruments.""" +import datetime + import numpy as np +import py +import xarray as xr from parcels import FieldSet from virtual_ship import Location, Spacetime from virtual_ship.instruments.adcp import simulate_adcp -def test_simulate_adcp() -> None: - MAX_DEPTH = -1000 - MIN_DEPTH = -5 - BIN_SIZE = 24 +def test_simulate_adcp(tmpdir: py.path.LocalPath) -> None: + # maximum depth the ADCP can measure + MAX_DEPTH = -1000 # -1000 + # minimum depth the ADCP can measure + MIN_DEPTH = -5 # -5 + # How many samples to take in the complete range between max_depth and min_depth. + NUM_BINS = 40 + + # 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)), + ] + + # expected observations at sample points + expected_obs = [ + { + "V": {"surface": 5, "max_depth": 6}, + "U": {"surface": 7, "max_depth": 8}, + "lat": sample_points[0].location.lat, + "lon": sample_points[0].location.lon, + "time": base_time + datetime.timedelta(seconds=0), + }, + { + "V": {"surface": 9, "max_depth": 10}, + "U": {"surface": 11, "max_depth": 12}, + "lat": sample_points[1].location.lat, + "lon": sample_points[1].location.lon, + "time": base_time + datetime.timedelta(seconds=1), + }, + ] + + # create fieldset based on the expected observations + # indices are time, depth, latitude, longitude + v = np.zeros((2, 2, 2, 2)) + v[0, 0, 0, 0] = expected_obs[0]["V"]["max_depth"] + v[0, 1, 0, 0] = expected_obs[0]["V"]["surface"] + v[1, 0, 1, 1] = expected_obs[1]["V"]["max_depth"] + v[1, 1, 1, 1] = expected_obs[1]["V"]["surface"] + + u = np.zeros((2, 2, 2, 2)) + u[0, 0, 0, 0] = expected_obs[0]["U"]["max_depth"] + u[0, 1, 0, 0] = expected_obs[0]["U"]["surface"] + 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( - {"U": 0, "V": 0}, { - "lon": 0, - "lat": 0, - "time": [np.datetime64("1950-01-01") + np.timedelta64(632160, "h")], + "V": v, + "U": u, + }, + { + "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"]), + ] + ), }, ) - sample_points = [Spacetime(Location(0, 0), 0)] + # perform simulation + out_path = tmpdir.join("out.zarr") simulate_adcp( fieldset=fieldset, - out_file_name="test", + out_path=out_path, max_depth=MAX_DEPTH, min_depth=MIN_DEPTH, - bin_size=BIN_SIZE, + num_bins=NUM_BINS, sample_points=sample_points, ) + + results = xr.open_zarr(out_path) + + # test if output is as expected + assert len(results.trajectory) == 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"), + ]: + 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 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() + 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_ship_underwater_st.py b/tests/instruments/test_ship_underwater_st.py index 784ff4cc..a0535c8d 100644 --- a/tests/instruments/test_ship_underwater_st.py +++ b/tests/instruments/test_ship_underwater_st.py @@ -18,50 +18,50 @@ def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None: # arbitrary time offset for the dummy fieldset base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d") - # variabes we are going to compare between expected and actual observations - variables = ["salinity", "temperature", "lat", "lon"] - # where to sample sample_points = [ - Spacetime(Location(3, 0), base_time + datetime.timedelta(seconds=0)), - Spacetime(Location(7, 0), 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": 1, - "temperature": 2, - "lat": 3, - "lon": 0, + "salinity": 5, + "temperature": 6, + "lat": sample_points[0].location.lat, + "lon": sample_points[0].location.lon, "time": base_time + datetime.timedelta(seconds=0), }, { - "salinity": 5, - "temperature": 6, - "lat": 7, - "lon": 0, + "salinity": 7, + "temperature": 8, + "lat": sample_points[1].location.lat, + "lon": sample_points[1].location.lon, "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"] + + 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( { - "U": np.zeros((2, 2)), - "V": np.zeros((2, 2)), - "salinity": [ - [expected_obs[0]["salinity"], 0], - [0, expected_obs[1]["salinity"]], - ], - "temperature": [ - [expected_obs[0]["temperature"], 0], - [0, expected_obs[1]["temperature"]], - ], + "V": np.zeros((2, 2, 2)), + "U": np.zeros((2, 2, 2)), + "salinity": salinity, + "temperature": temperature, }, { - "lon": 0, "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"]), @@ -71,6 +71,7 @@ def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None: }, ) + # perform simulation out_path = tmpdir.join("out.zarr") simulate_ship_underwater_st( @@ -80,21 +81,23 @@ def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None: sample_points=sample_points, ) + # test if output is as expected results = xr.open_zarr(out_path) - # below we assert if output makes sense - assert len(results.trajectory) == 1 # expect a singe trajectory + 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 # for every obs, check if the variables match the expected observations - for obs_i, exp in zip(results.sel(trajectory=traj).obs, expected_obs, strict=True): + for i, (obs_i, exp) in enumerate( + zip(results.sel(trajectory=traj).obs, expected_obs, strict=True) + ): obs = results.sel(trajectory=traj, obs=obs_i) - for var in variables: + for var in ["salinity", "temperature", "lat", "lon"]: obs_value = obs[var].values.item() exp_value = exp[var] assert np.isclose( obs_value, exp_value - ), f"Observation incorrect {obs_i=} {var=} {obs_value=} {exp_value=}." + ), f"Observation incorrect {i=} {var=} {obs_value=} {exp_value=}." diff --git a/virtual_ship/instruments/adcp.py b/virtual_ship/instruments/adcp.py index 2d7d032c..096facf1 100644 --- a/virtual_ship/instruments/adcp.py +++ b/virtual_ship/instruments/adcp.py @@ -1,11 +1,14 @@ """ADCP instrument.""" import numpy as np -from parcels import FieldSet, JITParticle, ParticleSet, Variable +import py +from parcels import FieldSet, ParticleSet, ScipyParticle, Variable from ..spacetime import Spacetime -_ADCPParticle = JITParticle.add_variables( +# we specifically use ScipyParticle because we have many small calls to execute +# there is some overhead with JITParticle and this ends up being significantly faster +_ADCPParticle = ScipyParticle.add_variables( [ Variable("U", dtype=np.float32, initial=np.nan), Variable("V", dtype=np.float32, initial=np.nan), @@ -21,25 +24,25 @@ def _sample_velocity(particle, fieldset, time): def simulate_adcp( fieldset: FieldSet, - out_file_name: str, + out_path: str | py.path.LocalPath, max_depth: float, min_depth: float, - bin_size: float, + num_bins: int, sample_points: list[Spacetime], ) -> None: """ Use parcels to simulate an ADCP in a fieldset. :param fieldset: The fieldset to simulate the ADCP in. - :param out_file_name: The file to write the results to. + :param out_path: The path to write the results to. :param max_depth: Maximum depth the ADCP can measure. :param min_depth: Minimum depth the ADCP can measure. - :param bin_size: How many samples to take in the complete range between max_depth and min_depth. + :param num_bins: How many samples to take in the complete range between max_depth and min_depth. :param sample_points: The places and times to sample at. """ sample_points.sort(key=lambda p: p.time) - bins = np.arange(max_depth, min_depth, bin_size) + bins = np.linspace(max_depth, min_depth, num_bins) num_particles = len(bins) particleset = ParticleSet.from_list( fieldset=fieldset, @@ -53,14 +56,22 @@ def simulate_adcp( ) # define output file for the simulation - out_file = particleset.ParticleFile( - name=out_file_name, - ) + # outputdt set to infinite as we just want to write at the end of every call to 'execute' + out_file = particleset.ParticleFile(name=out_path, outputdt=np.inf) for point in sample_points: particleset.lon_nextloop[:] = point.location.lon particleset.lat_nextloop[:] = point.location.lat - particleset.time_nextloop[:] = point.time + particleset.time_nextloop[:] = fieldset.time_origin.reltime( + np.datetime64(point.time) + ) - particleset.execute([_sample_velocity], dt=1, runtime=1, verbose_progress=False) - out_file.write(particleset, time=particleset[0].time) + # perform one step using the particleset + # dt and runtime are set so exactly one step is made. + particleset.execute( + [_sample_velocity], + dt=1, + runtime=1, + verbose_progress=False, + output_file=out_file, + ) diff --git a/virtual_ship/sailship.py b/virtual_ship/sailship.py index 928756c6..9c895912 100644 --- a/virtual_ship/sailship.py +++ b/virtual_ship/sailship.py @@ -203,10 +203,11 @@ def sailship(config: VirtualShipConfiguration): print("Simulating onboard ADCP.") simulate_adcp( fieldset=config.adcp_fieldset, - out_file_name=os.path.join("results", "adcp.zarr"), + out_path=os.path.join("results", "adcp.zarr"), max_depth=config.ADCP_settings["max_depth"], min_depth=-5, - bin_size=config.ADCP_settings["bin_size_m"], + num_bins=(-5 - config.ADCP_settings["max_depth"]) + // config.ADCP_settings["bin_size_m"], sample_points=adcps, )