diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py index 6f65d53e0..3b5ecf2aa 100644 --- a/src/parcels/_core/particleset.py +++ b/src/parcels/_core/particleset.py @@ -44,7 +44,8 @@ class ParticleSet: y : List of initial y (latitude) values for particles z : - Optional list of initial z values for particles. Default is 0m + Optional list of initial z values for particles. Default is vertical grid position closest to the surface (z=0) + that covers all fields in the fieldset. If none of the fields in the fieldset have a vertical grid, z=0 is used. t : Optional list of initial t (time) values for particles. Default is fieldset.U.grid.time[0] repeatdt : datetime.timedelta or float, optional @@ -78,11 +79,15 @@ def __init__( particle_ids = np.arange(x.size) if z is None: - minz = 0 + minz = None for field in self.fieldset.fields.values(): - if field.grid.depth is not None: - minz = min(minz, field.grid.depth[0]) - z = np.ones(x.size) * minz + for depth in field.grid.depth: + if minz is None or np.abs(depth) < np.abs(minz): + minz = depth + if minz is not None: + z = np.ones(x.size) * minz + else: + z = np.zeros(x.size) else: z = np.array(z).flatten() assert x.size == y.size and x.size == z.size, "x, y, z don't all have the same lengths" diff --git a/tests/test_particleset.py b/tests/test_particleset.py index c8211baf9..4e3ce9dc5 100644 --- a/tests/test_particleset.py +++ b/tests/test_particleset.py @@ -7,11 +7,13 @@ import xarray as xr from parcels import ( + FieldSet, Particle, ParticleSet, ParticleSetWarning, Variable, ) +from parcels._datasets.structured.generated import simple_UV_dataset from tests.common_kernels import DoNothing from tests.utils import round_and_hash_float_array @@ -174,3 +176,38 @@ def test_pset_iterator(fieldset): for i, particle in enumerate(pset): assert particle.particle_id == i assert i == npart - 1 + + +@pytest.mark.parametrize( + "depths", + [ + pytest.param(np.linspace(1, 10, 10), id="all_depths_positive"), + pytest.param(np.linspace(-10, -1, 10), id="all_depths_negative"), + ], +) +def test_pset_default_z_is_in_domain(depths): + ds = simple_UV_dataset(dims=(1, len(depths), 10, 10), mesh="flat") + ds = ds.assign_coords(depth=depths) + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + + pset = ParticleSet(fieldset, x=[0], y=[0]) + expected_z = depths[np.argmin(np.abs(depths))] + assert np.isclose(pset.z[0], expected_z) + + +@pytest.mark.parametrize( + "depths", + [ + pytest.param(np.concatenate([np.linspace(-15, -1, 5), np.linspace(0, 2, 5)]), id="depths_include_zero"), + pytest.param(np.concatenate([np.linspace(-9, -3, 3), np.linspace(2, 8, 3)]), id="closest_depth_is_positive"), + pytest.param(np.concatenate([np.linspace(-8, -2, 3), np.linspace(3, 9, 3)]), id="closest_depth_is_negative"), + ], +) +def test_pset_default_z_closest_to_zero(depths): + ds = simple_UV_dataset(dims=(1, len(depths), 10, 10), mesh="flat") + ds = ds.assign_coords(depth=depths) + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + + pset = ParticleSet(fieldset, x=[0], y=[0]) + expected_z = depths[np.argmin(np.abs(depths))] + assert np.isclose(pset.z[0], expected_z)