From 598e79d2c1c3cfa502f4204c9cf22e462f5d596f Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 6 Sep 2019 16:44:40 +0200 Subject: [PATCH] Adapt to last version of diagonals --- earthdiagnostics/ocean/heatcontentlayer.py | 142 ++++++++++++--------- 1 file changed, 79 insertions(+), 63 deletions(-) diff --git a/earthdiagnostics/ocean/heatcontentlayer.py b/earthdiagnostics/ocean/heatcontentlayer.py index 023909a0..bf0cc580 100644 --- a/earthdiagnostics/ocean/heatcontentlayer.py +++ b/earthdiagnostics/ocean/heatcontentlayer.py @@ -4,10 +4,19 @@ import numpy as np from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins -from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticBasinOption +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): """ @@ -36,15 +45,17 @@ class HeatContentLayer(Diagnostic): alias = 'ohclayer' "Diagnostic alias for the configuration file" - def __init__(self, data_manager, startdate, member, chunk, box, weight, min_level, max_level, data_convention): + 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.min_level = min_level - self.max_level = max_level + self.layers = layers + self.basins = basins self.required_vars = ['so', 'mlotst'] self.generated_vars = ['scvertsum'] self.data_convention = data_convention @@ -65,62 +76,47 @@ class HeatContentLayer(Diagnostic): """ options_available = (DiagnosticIntOption('min_depth'), DiagnosticIntOption('max_depth'), - DiagnosticBasinOption('basin', Basins().Global)) + DiagnosticBasinListOption('basins', + 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() - max_level, min_level, weight = cls._compute_weights(box) + 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, - weight, min_level, max_level, + areas, weight, layers, basins, diags.config.data_convention )) return job_list - @classmethod - def _compute_weights(cls, box): - depth, e3t, mask = cls._get_mesh_info() - - def calculate_weight(e3t_point, depth_point, mask_point): - """Calculate the weight for each cell""" - if not mask_point: - return 0 - top = depth_point - bottom = top + e3t_point - if bottom < box.min_depth or top > box.max_depth: - return 0 - else: - if top < box.min_depth: - top = box.min_depth - if bottom > box.max_depth: - bottom = box.max_depth - - return (bottom - top) * 1020 * 4000 - - calc = np.vectorize(calculate_weight, otypes='f') - weight = calc(e3t, depth, mask) - max_level, min_level = cls._get_used_levels(weight) - weight = weight[:, min_level:max_level, :] - return max_level, min_level, weight - - @classmethod - def _get_used_levels(cls, weight): - # Now we will reduce to the levels with any weight != 0 to avoid loading too much data on memory - levels = weight.shape[1] - min_level = 0 - while min_level < levels and not weight[:, min_level, :].any(): - min_level += 1 - max_level = min_level - while max_level < (levels - 1) and weight[:, max_level + 1, :].any(): - max_level += 1 - return max_level, min_level - @classmethod def _get_mesh_info(cls): handler = Utils.open_cdf('mesh_zgr.nc') @@ -150,42 +146,62 @@ class HeatContentLayer(Diagnostic): def request_data(self): """Request data required by the diagnostic""" - self.thetao = self.request_chunk(ModelingRealms.ocean, 'thetao', self.startdate, self.member, self.chunk) + 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.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""" - nco = Utils.nco() + 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 = np.sum(handler.variables['thetao'][:, self.min_level:self.max_level, :] * self.weight, 1) + heatc_sl, heatc_sl1D = ohc.compute(self.layers, self.weight, + handler.variables['thetao'][:], + self.areas) handler.sync() handler.renameVariable('thetao', 'heatc_sl') - handler.close() - nco.ncks( - input=thetao_file, - output=results, - options=( - '-O -v {0.lon_name},{0.lat_name},time'.format(self.data_convention), - ) - ) - Utils.rename_variables(results, {'x': 'i', 'y': 'j'}, False) - handler_results = Utils.open_cdf(results) - var = handler_results.createVariable('heatc', float, ('time', 'j', 'i'), fill_value=1.e20) + results = TempFile.get() + handler_results = Utils.open_cdf(results, 'w') + 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.rename_variables(results, {'x': 'i', 'y': 'j'}, False) + var = handler_results.createVariable('heatc', float, + ('time', 'j', 'i'), + fill_value=1.e20) var.units = 'J m-2' handler_results.sync() - handler_results.variables['heatc'][:] = heatc_sl + handler_results.variables['heatc'][:] = heatc_sl[0] # temporary fix, needs to loop over layers 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) -- GitLab