Skip to content

Commit cee0954

Browse files
erikvansebilleerikvansebille
committed
Splitting particleset tests into two files
One for the creation of particlesets, and one focussed on the pset.execute behaviour
1 parent 0dfb97f commit cee0954

2 files changed

Lines changed: 136 additions & 115 deletions

File tree

‎tests/v4/test_particleset.py‎

Lines changed: 0 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -7,20 +7,14 @@
77
import xarray as xr
88

99
from parcels import (
10-
AdvectionEE,
1110
Field,
1211
FieldSet,
1312
Particle,
1413
ParticleSet,
1514
ParticleSetWarning,
16-
StatusCode,
17-
UXPiecewiseConstantFace,
18-
VectorField,
1915
xgcm,
2016
)
2117
from parcels._datasets.structured.generic import datasets as datasets_structured
22-
from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
23-
from parcels.uxgrid import UxGrid
2418
from parcels.xgrid import XGrid
2519

2620

@@ -228,118 +222,9 @@ def test_pset_remove_index(fieldset, npart=100):
228222
assert not np.any(np.in1d(pset.trajectory, indices_to_remove))
229223

230224

231-
def test_pset_remove_particle_in_kernel(fieldset, npart=100):
232-
pset = ParticleSet(fieldset, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart))
233-
234-
def DeleteKernel(particle, fieldset, time): # pragma: no cover
235-
if particle.lon >= 0.4:
236-
particle.delete()
237-
238-
pset.execute(pset.Kernel(DeleteKernel), runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s"))
239-
assert pset.size == 40
240-
241-
242225
def test_pset_iterator(fieldset):
243226
npart = 10
244227
pset = ParticleSet(fieldset, lon=np.zeros(npart), lat=np.ones(npart))
245228
for i, particle in enumerate(pset):
246229
assert particle.trajectory == i
247230
assert i == npart - 1
248-
249-
250-
def test_pset_stop_simulation(fieldset):
251-
pset = ParticleSet(fieldset, lon=0, lat=0, pclass=Particle)
252-
253-
def Delete(particle, fieldset, time): # pragma: no cover
254-
if time >= fieldset.U.time[0].values + np.timedelta64(4, "s"):
255-
return StatusCode.StopExecution
256-
257-
pset.execute(Delete, dt=np.timedelta64(1, "s"), runtime=np.timedelta64(21, "s"))
258-
assert pset[0].time == fieldset.U.time[0].values + np.timedelta64(4, "s")
259-
260-
261-
@pytest.mark.parametrize("verbose_progress", [True, False])
262-
def test_uxstommelgyre_pset_execute(verbose_progress):
263-
ds = datasets_unstructured["stommel_gyre_delaunay"]
264-
grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"])
265-
U = Field(
266-
name="U",
267-
data=ds.U,
268-
grid=grid,
269-
mesh_type="spherical",
270-
interp_method=UXPiecewiseConstantFace,
271-
)
272-
V = Field(
273-
name="V",
274-
data=ds.V,
275-
grid=grid,
276-
mesh_type="spherical",
277-
interp_method=UXPiecewiseConstantFace,
278-
)
279-
P = Field(
280-
name="P",
281-
data=ds.p,
282-
grid=grid,
283-
mesh_type="spherical",
284-
interp_method=UXPiecewiseConstantFace,
285-
)
286-
UV = VectorField(name="UV", U=U, V=V)
287-
fieldset = FieldSet([UV, UV.U, UV.V, P])
288-
pset = ParticleSet(
289-
fieldset,
290-
lon=[30.0],
291-
lat=[5.0],
292-
depth=[50.0],
293-
time=[np.timedelta64(0, "s")],
294-
pclass=Particle,
295-
)
296-
pset.execute(
297-
runtime=np.timedelta64(10, "m"),
298-
dt=np.timedelta64(60, "s"),
299-
pyfunc=AdvectionEE,
300-
verbose_progress=verbose_progress,
301-
)
302-
303-
304-
@pytest.mark.xfail(reason="Output file not implemented yet")
305-
def test_uxstommelgyre_pset_execute_output():
306-
ds = datasets_unstructured["stommel_gyre_delaunay"]
307-
grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"])
308-
U = Field(
309-
name="U",
310-
data=ds.U,
311-
grid=grid,
312-
mesh_type="spherical",
313-
interp_method=UXPiecewiseConstantFace,
314-
)
315-
V = Field(
316-
name="V",
317-
data=ds.V,
318-
grid=grid,
319-
mesh_type="spherical",
320-
interp_method=UXPiecewiseConstantFace,
321-
)
322-
P = Field(
323-
name="P",
324-
data=ds.p,
325-
grid=grid,
326-
mesh_type="spherical",
327-
interp_method=UXPiecewiseConstantFace,
328-
)
329-
UV = VectorField(name="UV", U=U, V=V)
330-
fieldset = FieldSet([UV, UV.U, UV.V, P])
331-
pset = ParticleSet(
332-
fieldset,
333-
lon=[30.0],
334-
lat=[5.0],
335-
depth=[50.0],
336-
time=[0.0],
337-
pclass=Particle,
338-
)
339-
output_file = pset.ParticleFile(
340-
name="stommel_uxarray_particles.zarr", # the file name
341-
outputdt=timedelta(minutes=5), # the time step of the outputs
342-
)
343-
pset.execute(
344-
runtime=np.timedelta64(10, "m"), dt=np.timedelta64(60, "s"), pyfunc=AdvectionEE, output_file=output_file
345-
)

‎tests/v4/test_particleset_execute.py‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
import numpy as np
2+
import pytest
3+
4+
from parcels import (
5+
AdvectionEE,
6+
Field,
7+
FieldSet,
8+
Particle,
9+
ParticleSet,
10+
StatusCode,
11+
UXPiecewiseConstantFace,
12+
VectorField,
13+
xgcm,
14+
)
15+
from parcels._datasets.structured.generic import datasets as datasets_structured
16+
from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
17+
from parcels.uxgrid import UxGrid
18+
from parcels.xgrid import XGrid
19+
20+
21+
@pytest.fixture
22+
def fieldset() -> FieldSet:
23+
ds = datasets_structured["ds_2d_left"]
24+
grid = XGrid(xgcm.Grid(ds))
25+
U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
26+
V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
27+
return FieldSet([U, V])
28+
29+
30+
def test_pset_remove_particle_in_kernel(fieldset, npart=100):
31+
pset = ParticleSet(fieldset, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart))
32+
33+
def DeleteKernel(particle, fieldset, time): # pragma: no cover
34+
if particle.lon >= 0.4:
35+
particle.delete()
36+
37+
pset.execute(pset.Kernel(DeleteKernel), runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s"))
38+
assert pset.size == 40
39+
40+
41+
def test_pset_stop_simulation(fieldset):
42+
pset = ParticleSet(fieldset, lon=0, lat=0, pclass=Particle)
43+
44+
def Delete(particle, fieldset, time): # pragma: no cover
45+
if time >= fieldset.U.time[0].values + np.timedelta64(4, "s"):
46+
return StatusCode.StopExecution
47+
48+
pset.execute(Delete, dt=np.timedelta64(1, "s"), runtime=np.timedelta64(21, "s"))
49+
assert pset[0].time == fieldset.U.time[0].values + np.timedelta64(4, "s")
50+
51+
52+
@pytest.mark.parametrize("verbose_progress", [True, False])
53+
def test_uxstommelgyre_pset_execute(verbose_progress):
54+
ds = datasets_unstructured["stommel_gyre_delaunay"]
55+
grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"])
56+
U = Field(
57+
name="U",
58+
data=ds.U,
59+
grid=grid,
60+
mesh_type="spherical",
61+
interp_method=UXPiecewiseConstantFace,
62+
)
63+
V = Field(
64+
name="V",
65+
data=ds.V,
66+
grid=grid,
67+
mesh_type="spherical",
68+
interp_method=UXPiecewiseConstantFace,
69+
)
70+
P = Field(
71+
name="P",
72+
data=ds.p,
73+
grid=grid,
74+
mesh_type="spherical",
75+
interp_method=UXPiecewiseConstantFace,
76+
)
77+
UV = VectorField(name="UV", U=U, V=V)
78+
fieldset = FieldSet([UV, UV.U, UV.V, P])
79+
pset = ParticleSet(
80+
fieldset,
81+
lon=[30.0],
82+
lat=[5.0],
83+
depth=[50.0],
84+
time=[np.timedelta64(0, "s")],
85+
pclass=Particle,
86+
)
87+
pset.execute(
88+
runtime=np.timedelta64(10, "m"),
89+
dt=np.timedelta64(60, "s"),
90+
pyfunc=AdvectionEE,
91+
verbose_progress=verbose_progress,
92+
)
93+
94+
95+
@pytest.mark.xfail(reason="Output file not implemented yet")
96+
def test_uxstommelgyre_pset_execute_output():
97+
ds = datasets_unstructured["stommel_gyre_delaunay"]
98+
grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"])
99+
U = Field(
100+
name="U",
101+
data=ds.U,
102+
grid=grid,
103+
mesh_type="spherical",
104+
interp_method=UXPiecewiseConstantFace,
105+
)
106+
V = Field(
107+
name="V",
108+
data=ds.V,
109+
grid=grid,
110+
mesh_type="spherical",
111+
interp_method=UXPiecewiseConstantFace,
112+
)
113+
P = Field(
114+
name="P",
115+
data=ds.p,
116+
grid=grid,
117+
mesh_type="spherical",
118+
interp_method=UXPiecewiseConstantFace,
119+
)
120+
UV = VectorField(name="UV", U=U, V=V)
121+
fieldset = FieldSet([UV, UV.U, UV.V, P])
122+
pset = ParticleSet(
123+
fieldset,
124+
lon=[30.0],
125+
lat=[5.0],
126+
depth=[50.0],
127+
time=[0.0],
128+
pclass=Particle,
129+
)
130+
output_file = pset.ParticleFile(
131+
name="stommel_uxarray_particles.zarr", # the file name
132+
outputdt=np.timedelta64(5, "m"), # the time step of the outputs
133+
)
134+
pset.execute(
135+
runtime=np.timedelta64(10, "m"), dt=np.timedelta64(60, "s"), pyfunc=AdvectionEE, output_file=output_file
136+
)

0 commit comments

Comments
 (0)