Skip to content

Commit bad0828

Browse files
VeckoTheGeckoVeckoTheGecko
committed
Refactor constants
1 parent d187c00 commit bad0828

3 files changed

Lines changed: 49 additions & 44 deletions

File tree

‎parcels/_datasets/structured/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,7 @@
11
"""Structured datasets."""
2+
3+
_N = 30
4+
X = _N
5+
Y = 2 * _N
6+
Z = 3 * _N
7+
T = 13

‎parcels/_datasets/structured/generic.py‎

Lines changed: 39 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,14 @@
33
import numpy as np
44
import xarray as xr
55

6-
__all__ = ["N", "T", "datasets"]
6+
from . import T, X, Y, Z
77

8-
N = 30
9-
T = 13
8+
__all__ = ["T", "X", "Y", "Z", "datasets"]
109

1110

1211
def _rotated_curvilinear_grid():
13-
XG = np.arange(N)
14-
YG = np.arange(2 * N)
12+
XG = np.arange(X)
13+
YG = np.arange(Y)
1514
LON, LAT = np.meshgrid(XG, YG)
1615

1716
angle = -np.pi / 24
@@ -22,12 +21,12 @@ def _rotated_curvilinear_grid():
2221

2322
return xr.Dataset(
2423
{
25-
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
26-
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
27-
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
28-
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
29-
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
30-
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
24+
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
25+
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
26+
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
27+
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
28+
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
29+
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
3130
},
3231
coords={
3332
"XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -36,15 +35,15 @@ def _rotated_curvilinear_grid():
3635
"YC": (["YC"], YG + 0.5, {"axis": "Y"}),
3736
"ZG": (
3837
["ZG"],
39-
np.arange(3 * N),
38+
np.arange(Z),
4039
{"axis": "Z", "c_grid_axis_shift": -0.5},
4140
),
4241
"ZC": (
4342
["ZC"],
44-
np.arange(3 * N) + 0.5,
43+
np.arange(Z) + 0.5,
4544
{"axis": "Z"},
4645
),
47-
"depth": (["ZG"], np.arange(3 * N), {"axis": "Z"}),
46+
"depth": (["ZG"], np.arange(Z), {"axis": "Z"}),
4847
"time": (["time"], xr.date_range("2000", "2001", T), {"axis": "T"}),
4948
"lon": (
5049
["YG", "XG"],
@@ -75,8 +74,8 @@ def _polar_to_cartesian(r, theta):
7574
def _unrolled_cone_curvilinear_grid():
7675
# Not a great unrolled cone, but this is good enough for testing
7776
# you can use matplotlib pcolormesh to plot
78-
XG = np.arange(N)
79-
YG = np.arange(2 * N) * 0.25
77+
XG = np.arange(X)
78+
YG = np.arange(Y) * 0.25
8079

8180
pivot = -10, 0
8281
LON, LAT = np.meshgrid(XG, YG)
@@ -97,12 +96,12 @@ def _unrolled_cone_curvilinear_grid():
9796

9897
return xr.Dataset(
9998
{
100-
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
101-
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
102-
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
103-
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
104-
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
105-
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
99+
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
100+
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
101+
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
102+
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
103+
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
104+
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
106105
},
107106
coords={
108107
"XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -111,15 +110,15 @@ def _unrolled_cone_curvilinear_grid():
111110
"YC": (["YC"], YG + 0.5, {"axis": "Y"}),
112111
"ZG": (
113112
["ZG"],
114-
np.arange(3 * N),
113+
np.arange(Z),
115114
{"axis": "Z", "c_grid_axis_shift": -0.5},
116115
),
117116
"ZC": (
118117
["ZC"],
119-
np.arange(3 * N) + 0.5,
118+
np.arange(Z) + 0.5,
120119
{"axis": "Z"},
121120
),
122-
"depth": (["ZG"], np.arange(3 * N), {"axis": "Z"}),
121+
"depth": (["ZG"], np.arange(Z), {"axis": "Z"}),
123122
"time": (["time"], xr.date_range("2000", "2001", T), {"axis": "T"}),
124123
"lon": (
125124
["YG", "XG"],
@@ -139,43 +138,43 @@ def _unrolled_cone_curvilinear_grid():
139138
"2d_left_rotated": _rotated_curvilinear_grid(),
140139
"ds_2d_left": xr.Dataset(
141140
{
142-
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
143-
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
144-
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
145-
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
146-
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
147-
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
141+
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
142+
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
143+
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
144+
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
145+
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
146+
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
148147
},
149148
coords={
150149
"XG": (
151150
["XG"],
152-
2 * np.pi / N * np.arange(0, N),
151+
2 * np.pi / X * np.arange(0, X),
153152
{"axis": "X", "c_grid_axis_shift": -0.5},
154153
),
155-
"XC": (["XC"], 2 * np.pi / N * (np.arange(0, N) + 0.5), {"axis": "X"}),
154+
"XC": (["XC"], 2 * np.pi / X * (np.arange(0, X) + 0.5), {"axis": "X"}),
156155
"YG": (
157156
["YG"],
158-
2 * np.pi / (2 * N) * np.arange(0, 2 * N),
157+
2 * np.pi / (Y) * np.arange(0, Y),
159158
{"axis": "Y", "c_grid_axis_shift": -0.5},
160159
),
161160
"YC": (
162161
["YC"],
163-
2 * np.pi / (2 * N) * (np.arange(0, 2 * N) + 0.5),
162+
2 * np.pi / (Y) * (np.arange(0, Y) + 0.5),
164163
{"axis": "Y"},
165164
),
166165
"ZG": (
167166
["ZG"],
168-
np.arange(3 * N),
167+
np.arange(Z),
169168
{"axis": "Z", "c_grid_axis_shift": -0.5},
170169
),
171170
"ZC": (
172171
["ZC"],
173-
np.arange(3 * N) + 0.5,
172+
np.arange(Z) + 0.5,
174173
{"axis": "Z"},
175174
),
176-
"lon": (["XG"], 2 * np.pi / N * np.arange(0, N)),
177-
"lat": (["YG"], 2 * np.pi / (2 * N) * np.arange(0, 2 * N)),
178-
"depth": (["ZG"], np.arange(3 * N)),
175+
"lon": (["XG"], 2 * np.pi / X * np.arange(0, X)),
176+
"lat": (["YG"], 2 * np.pi / (Y) * np.arange(0, Y)),
177+
"depth": (["ZG"], np.arange(Z)),
179178
"time": (["time"], xr.date_range("2000", "2001", T), {"axis": "T"}),
180179
},
181180
),

‎tests/v4/test_gridadapter.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import pytest
55
from numpy.testing import assert_allclose
66

7-
from parcels._datasets.structured.generic import N, T, datasets
7+
from parcels._datasets.structured.generic import T, X, Y, Z, datasets
88
from parcels.grid import Grid as OldGrid
99
from parcels.tools.converters import TimeConverter
1010
from parcels.v4.grid import Grid as NewGrid
@@ -17,9 +17,9 @@
1717
GridTestCase(datasets["ds_2d_left"], "lat", datasets["ds_2d_left"].YG.values),
1818
GridTestCase(datasets["ds_2d_left"], "depth", datasets["ds_2d_left"].ZG.values),
1919
GridTestCase(datasets["ds_2d_left"], "time", datasets["ds_2d_left"].time.values),
20-
GridTestCase(datasets["ds_2d_left"], "xdim", N),
21-
GridTestCase(datasets["ds_2d_left"], "ydim", 2 * N),
22-
GridTestCase(datasets["ds_2d_left"], "zdim", 3 * N),
20+
GridTestCase(datasets["ds_2d_left"], "xdim", X),
21+
GridTestCase(datasets["ds_2d_left"], "ydim", Y),
22+
GridTestCase(datasets["ds_2d_left"], "zdim", Z),
2323
GridTestCase(datasets["ds_2d_left"], "tdim", T),
2424
GridTestCase(datasets["ds_2d_left"], "time_origin", TimeConverter(datasets["ds_2d_left"].time.values[0])),
2525
]

0 commit comments

Comments
 (0)