1
1
from datetime import timedelta
2
+ from pathlib import Path
3
+ from typing import Literal
2
4
3
5
import numpy as np
4
6
import parcels
7
9
ptype = {"scipy" : parcels .ScipyParticle , "jit" : parcels .JITParticle }
8
10
9
11
10
- def run_mitgcm_zonally_reentrant (mode ):
12
+ def run_mitgcm_zonally_reentrant (mode : Literal [ "scipy" , "jit" ], path : Path ):
11
13
"""Function that shows how to load MITgcm data in a zonally periodic domain."""
12
14
data_folder = parcels .download_example_dataset ("MITgcm_example_data" )
13
15
filenames = {
@@ -41,7 +43,7 @@ def periodicBC(particle, fieldset, time):
41
43
size = 10 ,
42
44
)
43
45
pfile = parcels .ParticleFile (
44
- "MIT_particles_" + str (mode ) + ".zarr" ,
46
+ str (path ) ,
45
47
pset ,
46
48
outputdt = timedelta (days = 1 ),
47
49
chunks = (len (pset ), 1 ),
@@ -52,12 +54,15 @@ def periodicBC(particle, fieldset, time):
52
54
)
53
55
54
56
55
- def test_mitgcm_output_compare ():
56
- run_mitgcm_zonally_reentrant ( "scipy" )
57
- run_mitgcm_zonally_reentrant ( "jit" )
57
+ def test_mitgcm_output_compare (tmpdir ):
58
+ def get_path ( mode : Literal [ "scipy" , "jit" ]) -> Path :
59
+ return tmpdir / f"MIT_particles_ { mode } .zarr"
58
60
59
- ds_jit = xr .open_zarr ("MIT_particles_jit.zarr" )
60
- ds_scipy = xr .open_zarr ("MIT_particles_scipy.zarr" )
61
+ for mode in ["scipy" , "jit" ]:
62
+ run_mitgcm_zonally_reentrant (mode , get_path (mode ))
63
+
64
+ ds_jit = xr .open_zarr (get_path ("jit" ))
65
+ ds_scipy = xr .open_zarr (get_path ("scipy" ))
61
66
62
67
np .testing .assert_allclose (ds_jit .lat .data , ds_scipy .lat .data )
63
68
np .testing .assert_allclose (ds_jit .lon .data , ds_scipy .lon .data )
0 commit comments