# coding=utf-8 """Point-wise Ocean Heat Content in a specified ocean thickness (J/m-2)""" import numpy as np from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticBasinListOption from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile from bscearth.utils.log import Log import diagonals from diagonals import ohc from diagonals.mesh_helpers.nemo import Nemo import iris import iris.coords import netCDF4 class HeatContentLayer(Diagnostic): """ Point-wise Ocean Heat Content in a specified ocean thickness (J/m-2) :original author: Isabel Andreu Burillo :contributor: Virginie Guemas :contributor: Eleftheria Exarchou :contributor: Javier Vegas-Regidor :created: June 2012 :last modified: June 2016 :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 box: box to use for the calculations :type box: Box """ alias = 'ohclayer' "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, box, areas, weight, layers, basins, data_convention): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.box = box self.areas = areas self.weight = weight self.layers = layers self.basins = basins self.data_convention = data_convention def __str__(self): return 'Heat content layer Startdate: {0} Member: {1} Chunk: {2} Box: {3}'.format(self.startdate, self.member, self.chunk, self.box) @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: minimum depth, maximum depth, basin=Global :type options: list[str] """ options_available = ( DiagnosticIntOption('min_depth'), DiagnosticIntOption('max_depth'), DiagnosticBasinListOption('basins', 'global') ) options = cls.process_options(options, options_available) box = Box(True) box.min_depth = options['min_depth'] box.max_depth = options['max_depth'] basins = options['basins'] if not basins: Log.error('Basins not recognized') return () mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') layers = ((box.min_depth, box.max_depth),) masks = {} basins.sort() for basin in basins: masks[basin] = Utils.get_mask(basin) areacello = mesh.get_areacello() areas = ohc.get_basin_area(areacello, masks) job_list = list() e3t = mesh.get_k_length() mask = mesh.get_landsea_mask() depth = mesh.get_depth(cell_point='W') weight = ohc.get_weights(layers, mask, e3t, depth) del mask, depth, e3t for startdate, member, chunk in diags.config.experiment.get_chunk_list(): job_list.append(HeatContentLayer( diags.data_manager, startdate, member, chunk, box, areas, weight, layers, basins, diags.config.data_convention )) return job_list @classmethod def _get_mesh_info(cls): handler = Utils.open_cdf('mesh_zgr.nc') # mask = Utils.get_mask(options['basin']) mask = handler.variables['tmask'][:] if 'e3t' in handler.variables: e3t = handler.variables['e3t'][:] elif 'e3t_0' in handler.variables: e3t = handler.variables['e3t_0'][:] else: raise Exception('e3t variable can not be found') if 'gdepw' in handler.variables: depth = handler.variables['gdepw'][:] elif 'gdepw_0' in handler.variables: depth = handler.variables['gdepw_0'][:] else: raise Exception('gdepw variable can not be found') e3t_3d = e3t.shape != depth.shape if e3t_3d: mask = e3t_3d * mask else: e3t = e3t[0, :] while len(depth.shape) < 4: depth = np.expand_dims(depth, -1) handler.close() return depth, e3t, mask def request_data(self): """Request data required by the diagnostic""" self.thetao = self.request_chunk(ModelingRealms.ocean, 'thetao', self.startdate, self.member, self.chunk) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.heatc = self.declare_chunk(ModelingRealms.ocean, 'heatc', self.startdate, self.member, self.chunk, box=self.box) self.heatcsum = self.declare_chunk(ModelingRealms.ocean, 'heatcsum', self.startdate, self.member, self.chunk, box=self.box) def compute(self): """Run the diagnostic""" thetao_file = TempFile.get() results = TempFile.get() results1D = TempFile.get() Utils.copy_file(self.thetao.local_file, thetao_file) handler = Utils.open_cdf(thetao_file) Utils.convert_units(handler.variables['thetao'], 'K') heatc_sl, heatc_sl1D = ohc.compute(self.layers, self.weight, handler.variables['thetao'][:], self.areas) handler.sync() handler.renameVariable('thetao', 'heatc_sl') results = TempFile.get() handler_results = Utils.open_cdf(results, 'w') lat_name = next(alias for alias in ('lat', 'latitude') if alias in handler.variables.keys()) lon_name = next(alias for alias in ('lon', 'longitude') if alias in handler.variables.keys()) Utils.copy_variable(handler, handler_results, 'time', True, True) Utils.copy_variable(handler, handler_results, 'i', True, True) Utils.copy_variable(handler, handler_results, 'j', True, True) Utils.copy_variable(handler, handler_results, lat_name, True, True) Utils.copy_variable(handler, handler_results, lon_name, True, True) var = handler_results.createVariable('heatc', float, ('time', 'j', 'i'), fill_value=1.e20) var.units = 'J m-2''' var.coordinates = ' '.join((lat_name, lon_name)) handler_results.sync() # temporary fix, needs to loop over layers handler_results.variables['heatc'][:] = heatc_sl[0] handler_results.close() results1D = TempFile.get() handler_results1D = Utils.open_cdf(results1D, 'w') Utils.copy_variable(handler, handler_results1D, 'time', True, True) handler_results1D.createDimension('region', len(self.basins)) handler_results1D.createDimension('region_length', 50) var_region = handler_results1D.createVariable( 'region', 'S1', ('region', 'region_length')) var_ohc1D = handler_results1D.createVariable( 'heatcsum', float, ('time', 'region',),) handler_results1D.sync() for i, basin in enumerate(self.basins): var_region[i, ...] = netCDF4.stringtoarr(basin.name, 50) var_ohc1D[..., i] = heatc_sl1D[i] handler_results1D.close() Utils.setminmax(results, 'heatc') self.heatc.set_local_file(results) self.heatcsum.set_local_file(results1D)