/
example_radial_rotation.py
90 lines (61 loc) · 2.94 KB
/
example_radial_rotation.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
import math
from datetime import timedelta as delta
import numpy as np
import pytest
from parcels import AdvectionRK4
from parcels import FieldSet
from parcels import JITParticle
from parcels import ParticleSet
from parcels import ScipyParticle
ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
def radial_rotation_fieldset(xdim=200, ydim=200): # Define 2D flat, square fieldset for testing purposes.
lon = np.linspace(0, 60, xdim, dtype=np.float32)
lat = np.linspace(0, 60, ydim, dtype=np.float32)
x0 = 30. # Define the origin to be the centre of the Field.
y0 = 30.
U = np.zeros((ydim, xdim), dtype=np.float32)
V = np.zeros((ydim, xdim), dtype=np.float32)
T = delta(days=1)
omega = 2*np.pi/T.total_seconds() # Define the rotational period as 1 day.
for i in range(lon.size):
for j in range(lat.size):
r = np.sqrt((lon[i]-x0)**2 + (lat[j]-y0)**2) # Define radial displacement.
assert(r >= 0.)
assert(r <= np.sqrt(x0**2 + y0**2))
theta = math.atan2((lat[j]-y0), (lon[i]-x0)) # Define the polar angle.
assert(abs(theta) <= np.pi)
U[j, i] = r * math.sin(theta) * omega
V[j, i] = -r * math.cos(theta) * omega
data = {'U': U, 'V': V}
dimensions = {'lon': lon, 'lat': lat}
return FieldSet.from_data(data, dimensions, mesh='flat')
def true_values(age): # Calculate the expected values for particle 2 at the endtime.
x = 20*math.sin(2*np.pi*age/(24.*60.**2)) + 30.
y = 20*math.cos(2*np.pi*age/(24.*60.**2)) + 30.
return [x, y]
def rotation_example(fieldset, outfile, mode='jit', method=AdvectionRK4):
npart = 2 # Test two particles on the rotating fieldset.
pset = ParticleSet.from_line(fieldset, size=npart, pclass=ptype[mode],
start=(30., 30.),
finish=(30., 50.)) # One particle in centre, one on periphery of Field.
runtime = delta(hours=17)
dt = delta(minutes=5)
outputdt = delta(hours=1)
pset.execute(method, runtime=runtime, dt=dt, moviedt=None,
output_file=pset.ParticleFile(name=outfile, outputdt=outputdt))
return pset
@pytest.mark.parametrize('mode', ['scipy', 'jit'])
def test_rotation_example(mode, tmpdir):
fieldset = radial_rotation_fieldset()
outfile = tmpdir.join("RadialParticle")
pset = rotation_example(fieldset, outfile, mode=mode)
assert(pset[0].lon == 30. and pset[0].lat == 30.) # Particle at centre of Field remains stationary.
vals = true_values(pset[1].time)
assert(np.allclose(pset[1].lon, vals[0], 1e-5)) # Check advected values against calculated values.
assert(np.allclose(pset[1].lat, vals[1], 1e-5))
if __name__ == "__main__":
filename = 'radial_rotation'
fieldset = radial_rotation_fieldset()
fieldset.write(filename)
outfile = "RadialParticle"
rotation_example(fieldset, outfile)