/
example_nemo_curvilinear.py
124 lines (96 loc) · 4.48 KB
/
example_nemo_curvilinear.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
from argparse import ArgumentParser
from datetime import timedelta as delta
from glob import glob
from os import path
import numpy as np
import pytest
from parcels import AdvectionRK4, AdvectionAnalytical
from parcels import FieldSet
from parcels import JITParticle
from parcels import ParticleFile
from parcels import ParticleSet
from parcels import ScipyParticle
ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
advection = {'RK4': AdvectionRK4, 'AA': AdvectionAnalytical}
def run_nemo_curvilinear(mode, outfile, advtype='RK4'):
"""Function that shows how to read in curvilinear grids, in this case from NEMO"""
data_path = path.join(path.dirname(__file__), 'NemoCurvilinear_data/')
filenames = {'U': {'lon': data_path + 'mesh_mask.nc4',
'lat': data_path + 'mesh_mask.nc4',
'data': data_path + 'U_purely_zonal-ORCA025_grid_U.nc4'},
'V': {'lon': data_path + 'mesh_mask.nc4',
'lat': data_path + 'mesh_mask.nc4',
'data': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4'}}
variables = {'U': 'U', 'V': 'V'}
dimensions = {'lon': 'glamf', 'lat': 'gphif'}
chunksize = {'lat': ('y', 256), 'lon': ('x', 512)}
field_set = FieldSet.from_nemo(filenames, variables, dimensions, chunksize=chunksize)
assert field_set.U.chunksize == chunksize
# Now run particles as normal
npart = 20
lonp = 30 * np.ones(npart)
if advtype == 'RK4':
latp = np.linspace(-70, 88, npart)
runtime = delta(days=160)
else:
latp = np.linspace(-70, 70, npart)
runtime = delta(days=15)
def periodicBC(particle, fieldSet, time):
if particle.lon > 180:
particle.lon -= 360
pset = ParticleSet.from_list(field_set, ptype[mode], lon=lonp, lat=latp)
pfile = ParticleFile(outfile, pset, outputdt=delta(days=1))
kernels = pset.Kernel(advection[advtype]) + periodicBC
pset.execute(kernels, runtime=runtime, dt=delta(hours=6),
output_file=pfile)
assert np.allclose(pset.lat - latp, 0, atol=2e-2)
def make_plot(trajfile):
from netCDF4 import Dataset
import matplotlib.pyplot as plt
import cartopy
class ParticleData(object):
def __init__(self):
self.id = []
def load_particles_file(fname, varnames):
T = ParticleData()
pfile = Dataset(fname, 'r')
T.id = pfile.variables['trajectory'][:]
for v in varnames:
setattr(T, v, pfile.variables[v][:])
return T
T = load_particles_file(trajfile, ['lon', 'lat', 'time'])
plt.axes(projection=cartopy.crs.PlateCarree())
plt.scatter(T.lon, T.lat, c=T.time, s=10)
plt.show()
@pytest.mark.parametrize('mode', ['jit']) # Only testing jit as scipy is very slow
def test_nemo_curvilinear(mode, tmpdir):
outfile = tmpdir.join('nemo_particles')
run_nemo_curvilinear(mode, outfile)
def test_nemo_curvilinear_AA(tmpdir):
outfile = tmpdir.join('nemo_particlesAA')
run_nemo_curvilinear('scipy', outfile, 'AA')
def test_nemo_3D_samegrid():
data_path = path.join(path.dirname(__file__), 'NemoNorthSeaORCA025-N006_data/')
ufiles = sorted(glob(data_path + 'ORCA*U.nc'))
vfiles = sorted(glob(data_path + 'ORCA*V.nc'))
wfiles = sorted(glob(data_path + 'ORCA*W.nc'))
mesh_mask = data_path + 'coordinates.nc'
filenames = {'U': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': ufiles},
'V': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': vfiles},
'W': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': wfiles}}
variables = {'U': 'uo',
'V': 'vo',
'W': 'wo'}
dimensions = {'U': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'},
'V': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'},
'W': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'}}
fieldset = FieldSet.from_nemo(filenames, variables, dimensions)
assert fieldset.U.dataFiles is not fieldset.W.dataFiles
if __name__ == "__main__":
p = ArgumentParser(description="""Chose the mode using mode option""")
p.add_argument('--mode', choices=('scipy', 'jit'), nargs='?', default='jit',
help='Execution mode for performing computation')
args = p.parse_args()
outfile = "nemo_particles"
run_nemo_curvilinear(args.mode, outfile)
make_plot(outfile+'.nc')