-
Notifications
You must be signed in to change notification settings - Fork 34
/
grd.py
303 lines (253 loc) · 11.9 KB
/
grd.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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
import logging
from datetime import datetime
import numpy as np
import xarray as xr
from zarr.convenience import consolidate_metadata
import IOverticalGrid
import xesmf as xe
import ESMF
__author__ = 'Trond Kristiansen'
__email__ = 'me@trondkristiansen.com'
__created__ = datetime(2008, 12, 9)
__modified__ = datetime(2021, 3, 23)
__version__ = "1.5"
__status__ = "Development"
"""
This class creates an object based on the input file structure. Currently the class takes
two types of structures as input: SODA or ROMS
"""
class Grd:
def __init__(self, grdtype, confM2R):
"""
The object is initialised and created through the __init__ method
As an example of how to use, these lines return a grid object called grdTEST:
=> import grd
=> grdTEST = grd.grdClass("grdfilename","ROMS")
"""
self.type = grdtype
self.grdName = confM2R.outgrid_name
self.realm = confM2R.realm
logging.info(f"[M2R_grd] => Creating init for grid object {confM2R.outgrid_name}")
logging.info(f"[M2R_grd] => Initialized GRD object for grid type {self.type}\n")
def create_object(self, confM2R, grd_filename):
"""
This method creates a new object by reading the grd input file. All
dimensions (eta, xi, lon, lat etc.) are defined here and used throughout these scripts.
Also, the depth matrix is calculated in this function by calling IOverticalGrid.py (ROMS grid only). For
input model depths, the depth array is a one dimensional. If input data has 2 or 3 dimensions, this
has to be accounted for throughout the model2roms package as one dimension is currently only supported.
Object types currently supported: STATION, GLORYS, SODA3
Trond Kristiansen, 18.11.2009, edited 03.01.2017, 23.03.2021
"""
if self.type == 'FORCINGDATA' and confM2R.use_zarr:
mapper = confM2R.fs.get_mapper(grd_filename)
consolidate_metadata(mapper)
ds = xr.open_zarr(mapper, chunks={"time": 100}, consolidated=False).isel(time=0)
if confM2R.subset_indata:
ds = ds.sel(longitude=slice(confM2R.subset[2], confM2R.subset[3]),
latitude=slice(confM2R.subset[0], confM2R.subset[1]))
else:
ds = xr.open_dataset(grd_filename)
if self.type == 'FORCINGDATA':
logging.info(f"[M2R_grd] ---> Assuming {confM2R.grd_type} grid type for {self.type}")
logging.info(f"[M2R_grd] ---> Using dimension names {confM2R.lon_name} and {confM2R.lat_name} and {confM2R.depth_name}")
self.lon = ds[str(confM2R.lon_name)][:]
self.lat = ds[str(confM2R.lat_name)][:]
self.h = ds[str(confM2R.depth_name)][:]
self.nlevels = len(self.h)
self.fillval = -9.99e+33
self.hc = None
if self.lon.ndim == 1:
self.lon, self.lat = np.meshgrid(self.lon, self.lat)
# Create grid for ESMF interpolation
if confM2R.use_zarr:
self.esmfgrid,lon_shape, dim_names = xe.frontend.ds_to_ESMFgrid(ds, need_bounds=True, periodic=False, append=None)
else:
self.esmfgrid = ESMF.Grid(filename=grd_filename, filetype=ESMF.FileFormat.GRIDSPEC,
is_sphere=True, coord_names=[str(confM2R.lon_name), str(confM2R.lat_name)],
add_mask=False)
if confM2R.ocean_indata_type == 'SODA3':
self.fillval = -1.e+20
if confM2R.ocean_indata_type == 'SODA3_5DAY':
self.fillval = -1.e+20
if confM2R.ocean_indata_type == 'GLORYS':
self.fillval = 9.96921e+36
if confM2R.ocean_indata_type == 'NORESM':
# self.h = ds["depth"][:]
self.h = np.asarray([0, 5, 10, 15, 20, 25, 30, 40, 50, 62.5, 75, 87.5, 100, 112.5, 125,
137.5, 150, 175, 200, 225, 250, 275, 300, 350, 400, 450, 500, 550, 600,
650, 700, 750, 800, 850, 900, 950, 1000, 1050, 1100, 1150, 1200, 1250,
1300, 1350, 1400, 1450, 1500, 1625, 1750, 1875, 2000, 2250, 2500, 2750,
3000, 3250, 3500, 3750, 4000, 4250, 4500, 4750, 5000, 5250, 5500, 5750,
6000, 6250, 6500, 6750])
self.fillval = 32768
self.nlevels = len(self.h)
IOverticalGrid.get_z_levels(self)
if self.type == 'STATION':
self.lon = ds[confM2R.lon_name][:]
self.lat = ds[confM2R.lat_name][:]
self.h = ds[confM2R.depth_name][:]
self.time = ds[confM2R.time_name][:]
self.Lp = 1
self.Mp = 1
self.fillval = -9.99e+33
if self.type in ['ROMS']:
self.write_clim = True
self.write_bry = True
self.write_init = True
self.write_stations = False
self.lonname = 'lon_rho'
self.latname = 'lat_rho'
"""
Set initTime to 1 if you dont want the first time-step to be
the initial field (no ubar and vbar if time=0)
"""
self.inittime = 0
self.ocean_time = 0
self.NT = 2
self.tracer = self.NT
self.message = None # Used to store the date for printing to screen (IOwrite.py)
self.time = 0
self.reftime = 0
self.grdtype = 'regular'
self.mask_rho = ds["mask_rho"][:, :]
self.lon_rho = ds["lon_rho"][:, :]
self.lat_rho = ds["lat_rho"][:, :]
self.h = ds["h"][:, :]
masked_h = np.where(self.h > 0, self.h, self.h.max())
self.hmin = masked_h.min()
if "Vtransform" in ds.variables:
self.vtransform = ds["Vtransform"].values
else:
self.vtransform = confM2R.vtransform
if "s_rho" in ds.variables:
self.s_rho = ds["s_rho"].values
self.nlevels = len(self.s_rho)
else:
self.nlevels = confM2R.nlevels
if "Vstretching" in ds.variables:
self.vstretching = ds["Vstretching"].values
if "theta_s" in ds.variables:
self.theta_s = ds["theta_s"].values
else:
self.theta_s = confM2R.theta_s
if "theta_b" in ds.variables:
self.theta_b = ds["theta_b"].values
else:
self.theta_b = confM2R.theta_b
if "Tcline" in ds.variables:
self.tcline = ds["Tcline"].values
else:
self.tcline = confM2R.tcline
if "hc" in ds.variables:
self.hc = ds["hc"].values
else:
self.hc = confM2R.hc
if self.vtransform == 1:
self.hc = min(self.hmin, self.tcline)
self.hc = self.tcline
if self.tcline > self.hmin:
print('Vertical transformation parameters are not defined correctly in either gridid.txt '
'or in the history files: \n Tc\
line = %d and hmin = %d. \n You need to make sure that '
'tcline <= hmin when using transformation 1.' % (
self.tcline, self.hmin))
else:
self.hc = self.tcline
zeta = None
if zeta is None:
self.zeta = np.zeros(self.h.shape)
else:
self.zeta = zeta
# for findvar in ds:
# if findvar=="hraw":
# self.hraw = ds["hraw"][:,:,:]
self.lon_u = ds["lon_u"][:, :]
self.lat_u = ds["lat_u"][:, :]
self.mask_u = ds["mask_u"][:, :]
for findvar in ds:
if findvar == "lon_vert":
self.lon_vert = ds["lon_vert"][:, :]
self.lat_vert = ds["lat_vert"][:, :]
for findvar in ds:
if findvar == "x_rho":
self.x_rho = ds["x_rho"][:, :]
self.y_rho = ds["y_rho"][:, :]
for findvar in ds:
if findvar == "x_u":
self.x_u = ds["x_u"][:, :]
self.y_u = ds["y_u"][:, :]
for findvar in ds:
if findvar == "x_v":
self.x_v = ds["x_v"][:, :]
self.y_v = ds["y_v"][:, :]
for findvar in ds:
if findvar == "x_psi":
self.x_psi = ds["x_psi"][:, :]
self.y_psi = ds["y_psi"][:, :]
for findvar in ds:
if findvar == "x_vert":
self.x_vert = ds["x_vert"][:, :]
self.y_vert = ds["y_vert"][:, :]
for findvar in ds:
if findvar == "xl":
self.xl = ds["xl"]
self.el = ds["el"]
for findvar in ds:
if findvar == "dmde":
self.dmde = ds["dmde"][:, :]
self.dndx = ds["dndx"][:, :]
self.lon_v = ds["lon_v"][:, :]
self.lat_v = ds["lat_v"][:, :]
self.mask_v = ds["mask_v"][:, :]
# self.spherical = ds["spherical"][:]
self.lon_psi = self.lon_u[:-1, :]
self.lat_psi = self.lat_v[:, :-1]
self.mask_psi = self.mask_v[:, :-1]
# self.f = ds["f"][:, :]
self.angle = ds["angle"][:, :]
self.pm = ds["pm"][:, :]
self.invpm = 1.0 / np.asarray(ds["pm"][:, :])
self.pn = ds["pn"][:, :]
self.invpn = 1.0 / np.asarray(ds["pn"][:, :])
self.Lp = len(self.lat_rho[1, :])
self.Mp = len(self.lat_rho[:, 1])
self.fillval = -9.99e33
self.eta_rho = self.Mp
self.eta_u = self.Mp
self.eta_v = self.Mp - 1
self.eta_psi = self.Mp - 1
self.xi_rho = self.Lp
self.xi_u = self.Lp - 1
self.xi_v = self.Lp
self.xi_psi = self.Lp - 1
# Boolean to check if we need to initialize the CLIM file before writing
self.ioClimInitialized = False
self.ioInitInitialized = False
if self.lon_rho.ndim == 1:
self.lon_rho, self.lat_rho = np.meshgrid(self.lon_rho, self.lat_rho)
self.lon_u, self.lat_u = np.meshgrid(self.lon_u, self.lat_u)
self.lon_v, self.lat_v = np.meshgrid(self.lon_v, self.lat_v)
# Setup the vertical coordinate system
IOverticalGrid.calculateVgrid(self)
self.esmfgrid_u = ESMF.Grid(filename=grd_filename, filetype=ESMF.FileFormat.GRIDSPEC,
coord_names=['lon_u', 'lat_u'],
is_sphere=True,
add_mask=False)
self.esmfgrid_v = ESMF.Grid(filename=grd_filename, filetype=ESMF.FileFormat.GRIDSPEC,
is_sphere=True,
coord_names=['lon_v', 'lat_v'],
add_mask=False)
self.esmfgrid = ESMF.Grid(filename=grd_filename, filetype=ESMF.FileFormat.GRIDSPEC,
is_sphere=True,
coord_names=[self.lonname, self.latname],
add_mask=False)
def getdims(self):
if self.type in ["ROMS"]:
self.Lp = len(self.lat_rho[1, :])
self.Mp = len(self.lat_rho[:, 1])
if self.type in ['FORCINGDATA']:
self.Lp = len(self.lat[1, :])
self.Mp = len(self.lat[:, 1])
self.M = self.Mp - 1
self.L = self.Lp - 1