Newer
Older
import iris
import iris.util
import iris.coords
import iris.analysis
import iris.exceptions
import numpy as np
import netCDF4
from earthdiagnostics import cdftools
from earthdiagnostics.box import Box
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, \
DiagnosticIntOption, DiagnosticDomainOption, \
DiagnosticBoolOption, DiagnosticBasinListOption, DiagnosticVariableOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile
import diagonals.regsum as regsum
from diagonals.mesh_helpers.nemo import Nemo
class RegionSum(Diagnostic):
"""
For 3D fields, a horizontal mean for each level is also given.
If a spatial window is specified, the mean value is computed only
in this window.
:original author: Javier Vegas-Regidor <javier.vegas@bsc.es>
:created: March 2017
:param data_manager: data management object
:type data_manager: DataManager
:param startdate: startdate
:type startdate: str
:param member: member number
:type member: int
:param chunk: chunk's number
:type chunk: int
:param variable: variable to average
:type variable: str
:param box: box used to restrict the vertical mean
:type box: Box
"""
alias = 'regsum'
"Diagnostic alias for the configuration file"
def __init__(self, data_manager, startdate, member, chunk, domain,
variable, grid_point, box, save3d, basins,
Diagnostic.__init__(self, data_manager)
self.startdate = startdate
self.member = member
self.chunk = chunk
self.domain = domain
self.box = box
self.save3d = save3d
self.grid = grid
self.declared = {}
self.lat_name = 'lat'
self.lon_name = 'lon'
if self._different_type(other):
return False
return self.startdate == other.startdate and \
self.member == other.member and self.chunk == other.chunk and \
self.box == other.box and self.variable == other.variable and \
self.grid_point == other.grid_point and self.grid == other.grid \
and self.basin == other.basin
return 'Region sum Startdate: {0.startdate} Member: {0.member}' \
'Chunk: {0.chunk} Variable: {0.variable} ' \
'Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d}' \
'Original grid: {0.grid} Basin: {0.basins}'.format(self)
def __hash__(self):
return hash(str(self))
@classmethod
def generate_jobs(cls, diags, options):
"""
Create a job for each chunk to compute the diagnostic
:param diags: Diagnostics manager class
:type diags: Diags
:param options: variable, minimum depth (level), maximum depth (level)
:type options: list[str]
:return:
"""
options_available = (DiagnosticDomainOption(),
DiagnosticVariableOption(diags.data_manager.config.var_manager),
DiagnosticOption('grid_point', 'T'),
DiagnosticBasinListOption('basins', Basins().Global),
DiagnosticIntOption('min_depth', -1),
DiagnosticIntOption('max_depth', -1),
DiagnosticIntOption('min_lat', -1),
DiagnosticIntOption('max_lat', -1),
DiagnosticIntOption('min_lon', -1),
DiagnosticIntOption('max_lon', -1),
DiagnosticBoolOption('save3D', True),
DiagnosticOption('grid', ''))
options = cls.process_options(options, options_available)
box = Box()
box.min_depth = options['min_depth']
box.max_depth = options['max_depth']
box.min_lat = options['min_lat']
box.max_lat = options['max_lat']
box.min_lon = options['min_lon']
box.max_lon = options['max_lon']
basins = options['basins']
if not basins:
Log.error('Basins not recognized')
return()
job_list = list()
for startdate, member, chunk in diags.config.experiment.get_chunk_list():
job_list.append(RegionSum(diags.data_manager, startdate, member, chunk,
options['domain'], options['variable'],
options['grid_point'].lower(), box,
options['save3D'], options['basins'],
options['grid']))
return job_list
def request_data(self):
self.variable_file = self.request_chunk(self.domain, self.variable,
self.startdate,
self.member, self.chunk,
grid=self.grid)
def declare_data_generated(self):
if self.box.min_depth == -1 and self.box.max_depth == -1:
box_save = None
else:
box_save = self.box
self._declare_var('sum', False, box_save)
self._declare_var('sum', True, box_save)
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
has_levels = self._fix_file_metadata()
data = self._load_data()
masks = {}
self.basins.sort()
for basin in self.basins:
masks[basin] = Utils.get_mask(basin)
mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc')
if self.box.min_lat is not -1 and self.box.max_lat is not -1 and self.box.min_lon is not -1 and self.box.max_lat is not -1:
name = '{0}_{1}'.format(Box.get_lat_str(self.box), Box.get_lon_str(self.box))
masks[name] = mesh.get_region_mask(self.box.min_lat,
self.box.max_lat,
self.box.min_lon,
self.box.max_lon)
if has_levels:
self._sum_3d_variable(data, mesh, masks)
else:
self._sum_2d_var(data, mesh, masks)
def _sum_2d_var(self, data, mesh, masks):
areacello = mesh.get_areacello(cell_point=self.grid_point)
varsum = regsum.compute_regsum_2D(data.data, masks, areacello)
self._save_result_2D('sum', varsum, data)
def _sum_3d_variable(self, data, mesh, masks):
areacello = mesh.get_areacello(cell_point=self.grid_point)
tmask = iris.load_cube('mesh_hgr.nc', '{0}mask'.format(self.grid_point))
e3 = self._try_load_cube(3)
e3 = self._rename_depth(e3)
e3.coord('depth').bounds = data.coord('depth').bounds
if self.box.min_depth is not -1 and self.box.max_depth is not -1:
depth_constraint = iris.Constraint(depth=lambda c: self.box.min_depth <= c <= self.box.max_depth)
e3 = e3.extract(depth_constraint)
data = data.extract(depth_constraint)
tmask = tmask.extract(depth_constraint)
volcello = areacello*e3.data.astype(np.float32)
varsum = regsum.compute_regsum_3D(data.data, masks, volcello,
tmask.data)
self._save_result_2D('sum', varsum, data)
if self.save3d:
varsum3d = regsum.compute_regsum_levels(data.data, masks,
volcello, tmask)
self._save_result_3D('sum', varsum3d, data)
def _try_load_cube(self, number):
try:
cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}'.format(number, self.grid_point))
except iris.exceptions.ConstraintMismatchError:
cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}_0'.format(number, self.grid_point))
cube = iris.util.squeeze(cube)
dims = len(cube.shape)
try:
cube.coord('i')
except iris.exceptions.CoordinateNotFoundError:
cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 1]), var_name='i'), dims - 1)
try:
cube.coord('j')
except iris.exceptions.CoordinateNotFoundError:
cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 2]), var_name='j'), dims - 2)
return cube
def _load_data(self):
coords = []
handler = Utils.open_cdf(self.variable_file.local_file)
for variable in handler.variables:
if variable in ('time', 'lev', 'lat', 'lon', 'latitude', 'longitude', 'leadtime', 'time_centered'):
coords.append(variable)
if variable == 'time_centered':
handler.variables[variable].standard_name = ''
handler.variables[self.variable].coordinates = ' '.join(coords)
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
data = iris.load_cube(self.variable_file.local_file)
return self._rename_depth(data)
def _rename_depth(self, data):
for coord_name in ('model_level_number', 'Vertical T levels', 'lev'):
if data.coords(coord_name):
coord = data.coord(coord_name)
coord.standard_name = 'depth'
coord.long_name = 'depth'
break
return data
def _fix_file_metadata(self):
handler = Utils.open_cdf(self.variable_file.local_file)
var = handler.variables[self.variable]
coordinates = ''
has_levels = False
for dimension in handler.variables.keys():
if dimension in ['time', 'lev', 'lat', 'latitude', 'lon', 'longitude', 'i', 'j']:
coordinates += ' {0}'.format(dimension)
if dimension == 'lev':
has_levels = True
var.coordinates = coordinates
handler.close()
return has_levels
def _declare_var(self, var, threed, box_save):
if threed:
if not self.save3d:
return False
final_name = '{1}3d{0}'.format(var, self.variable)
else:
final_name = '{1}{0}'.format(var, self.variable)
self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, final_name, self.startdate, self.member,
self.chunk, box=box_save, region=self.basins)
if threed:
if not self.save3d:
return False
original_name = '{0}_{1}'.format(var, self.variable)
final_name = '{1}3d{0}'.format(var, self.variable)
levels = ',lev'
else:
original_name = '{0}_3D{1}'.format(var, self.variable)
final_name = '{1}{0}'.format(var, self.variable)
levels = ''
temp2 = TempFile.get()
Utils.nco().ncks(
input=mean_file,
output=temp2,
options=('-v {0},{2.lat_name},{2.lon_name}{1}'.format(original_name, levels, self),)
)
var_handler = handler.variables[original_name]
if hasattr(var_handler, 'valid_min'):
del var_handler.valid_min
if hasattr(var_handler, 'valid_max'):
del var_handler.valid_max
handler.close()
self.declared[final_name].set_local_file(temp2, diagnostic=self, rename_var=original_name, region=self.basins)
if threed:
if not self.save3d:
return False
final_name = '{1}3d{0}'.format(var, self.variable)
else:
final_name = '{1}{0}'.format(var, self.variable)
self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, final_name, self.startdate, self.member,
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
self.chunk, box=box_save, region=self.basins, grid=self.grid)
def _save_result_2D(self, var, result, data):
final_name = '{1}{0}'.format(var, self.variable)
temp = TempFile.get()
handler_source = Utils.open_cdf(self.variable_file.local_file)
handler_temp = Utils.open_cdf(temp, 'w')
Utils.copy_variable(handler_source, handler_temp, 'time', True, True)
handler_temp.createDimension('region', len(result))
handler_temp.createDimension('region_length', 50)
var_region = handler_temp.createVariable('region', 'S1',
('region', 'region_length'))
var = handler_temp.createVariable('{1}{0}'.format(var, self.variable),
float, ('time', 'region',),)
var.units = '{0}'.format(data.units)
for i, basin in enumerate(result):
var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
var[..., i] = result[basin]
handler_temp.close()
self.declared[final_name].set_local_file(temp, diagnostic=self)
def _save_result_3D(self, var, result, data):
final_name = '{1}3d{0}'.format(var, self.variable)
temp = TempFile.get()
handler_source = Utils.open_cdf(self.variable_file.local_file)
handler_temp = Utils.open_cdf(temp, 'w')
Utils.copy_variable(handler_source, handler_temp, 'time', True, True)
handler_temp.createDimension('region', len(result))
handler_temp.createDimension('region_length', 50)
handler_temp.createDimension('lev', data.shape[1])
var_level = handler_temp.createVariable('lev', float, 'lev')
var_level[...] = data.coord('depth').points
var_level.units = 'm'
var_level.axis = 'Z'
var_level.positive = 'down'
var_level.long_name = 'ocean depth coordinate'
var_level.standard_name = 'depth'
var_region = handler_temp.createVariable('region', 'S1',
('region', 'region_length'))
var = handler_temp.createVariable('{1}3d{0}'.format(var, self.variable),
float, ('time', 'lev', 'region',),)
var.units = '{0}'.format(data.units)
for i, basin in enumerate(result):
var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
var[..., i] = result[basin]
handler_temp.close()
self.declared[final_name].set_local_file(temp, diagnostic=self)