From 9d786b30ee526002a4037776545421fa737c659e Mon Sep 17 00:00:00 2001 From: sloosvel Date: Thu, 23 Jan 2020 16:29:24 +0100 Subject: [PATCH] Adapt psi to python --- earthdiagnostics/ocean/psi.py | 129 ++++++++++++++++++++++++++++++---- 1 file changed, 114 insertions(+), 15 deletions(-) diff --git a/earthdiagnostics/ocean/psi.py b/earthdiagnostics/ocean/psi.py index b9b32f02..fabfc5a6 100644 --- a/earthdiagnostics/ocean/psi.py +++ b/earthdiagnostics/ocean/psi.py @@ -1,10 +1,27 @@ # coding=utf-8 """Compute the barotropic stream function""" -from earthdiagnostics import cdftools -from earthdiagnostics.diagnostic import Diagnostic +import os +import datetime + +import numpy as np + +import netCDF4 + +import iris +import iris.analysis +import iris.coords +import iris.util +from bscearth.utils.log import Log + +from earthdiagnostics.constants import Basins +from earthdiagnostics.diagnostic import Diagnostic, \ + DiagnosticBasinListOption from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile +import diagonals.psi as psi +from diagonals.mesh_helpers.nemo import Nemo + class Psi(Diagnostic): """ @@ -31,11 +48,12 @@ class Psi(Diagnostic): vsftbarot = 'vsftbarot' - def __init__(self, data_manager, startdate, member, chunk): + def __init__(self, data_manager, startdate, member, chunk, masks): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk + self.masks = masks def __eq__(self, other): if self._different_type(other): @@ -43,7 +61,11 @@ class Psi(Diagnostic): return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk def __str__(self): - return 'PSI Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, self.chunk) + return 'PSI Startdate: {0} Member: {1} Chunk: {2} Basins: {3}'.format( + self.startdate, self.member, self.chunk, ','.join(str(basin) for basin in self.masks.keys())) + + def __hash__(self): + return hash(str(self)) @classmethod def generate_jobs(cls, diags, options): @@ -56,27 +78,104 @@ class Psi(Diagnostic): :type options: list[str] :return: """ - if len(options) > 1: - raise Exception('The PSI diagnostic has no options') + options_available = (DiagnosticBasinListOption('basins', + [Basins().Global]),) + options = cls.process_options(options, options_available) + + basins = options['basins'] + if not basins: + Log.error('Basins not recognized') + return () + + masks = {} + basins.sort() + for basin in basins: + masks[basin] = Utils.get_mask(basin) + job_list = list() for startdate, member, chunk in diags.config.experiment.get_chunk_list(): - job_list.append(Psi(diags.data_manager, startdate, member, chunk)) + job_list.append(Psi(diags.data_manager, startdate, member, chunk, + masks)) return job_list def request_data(self): """Request data required by the diagnostic""" - self.uo = self.request_chunk(ModelingRealms.ocean, 'uo', self.startdate, self.member, self.chunk) - self.vo = self.request_chunk(ModelingRealms.ocean, 'vo', self.startdate, self.member, self.chunk) + self.uo = self.request_chunk(ModelingRealms.ocean, 'uo', + self.startdate, self.member, self.chunk) + self.vo = self.request_chunk(ModelingRealms.ocean, 'vo', + self.startdate, self.member, self.chunk) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" - self.psi = self.declare_chunk(ModelingRealms.ocean, Psi.vsftbarot, self.startdate, self.member, self.chunk) + self.psi = self.declare_chunk(ModelingRealms.ocean, Psi.vsftbarot, + self.startdate, self.member, self.chunk) def compute(self): """Run the diagnostic""" + self._fix_coordinates_attribute(self.uo.local_file, 'uo') + self._fix_coordinates_attribute(self.vo.local_file, 'vo') + uo_cube = iris.load_cube(self.uo.local_file) + vo_cube = iris.load_cube(self.vo.local_file) + + uo = np.ma.filled(uo_cube.data, 0.0).astype(np.float32) + vo = np.ma.filled(vo_cube.data, 0.0).astype(np.float32) + + mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc') + e2u = mesh.get_j_length(cell_point='U') + e3u = mesh.get_k_length(cell_point='U') + e1v = mesh.get_i_length(cell_point='V') + e3v = mesh.get_k_length(cell_point='V') + + results = psi.compute(self.masks, e2u, e1v, e3u, e3v, uo, vo) + + self.save(results) + + def save(self, result): temp = TempFile.get() - cdftools.run('cdfpsi', input_file=[self.uo.local_file, self.vo.local_file], output_file=temp, - options='-mean -mask') - Utils.rename_variable(temp, 'sobarstf', Psi.vsftbarot) - Utils.setminmax(temp, Psi.vsftbarot) - self.psi.set_local_file(temp) + handler_source = Utils.open_cdf(self.uo.local_file) + handler_temp = Utils.open_cdf(temp, 'w') + lat_name = next(alias for alias in ('lat', 'latitude') + if alias in handler_source.variables.keys()) + lon_name = next(alias for alias in ('lon', 'longitude') + if alias in handler_source.variables.keys()) + Utils.copy_variable(handler_source, handler_temp, 'time', True, True) + Utils.copy_variable(handler_source, handler_temp, 'i', False, True) + Utils.copy_variable(handler_source, handler_temp, 'j', False, True) + + Utils.copy_variable(handler_source, handler_temp, lat_name, True, True) + Utils.copy_variable(handler_source, handler_temp, lon_name, 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('vsftbarot', float, + ('time', 'j', 'i', 'region')) + var.units = 'm3/s' + var.coordinates = ' '.join((lat_name, lon_name)) + var.missing_value = 1e20 + var.fill_value = 1e20 + var.valid_min = -300e6 + var.valid_max = 300e6 + var.long_name = 'Barotropic_Stream_Function' + + for i, basin in enumerate(result): + var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50) + result[basin].mask = self.masks[basin]<1 + var[..., i] = result[basin] + handler_temp.close() + self.psi.set_local_file(temp, diagnostic=self) + + def _fix_coordinates_attribute(self, filepath, var_name): + add_coordinates = { + 'time_centered', 'leadtime' + } + handler = Utils.open_cdf(filepath) + coordinates = handler.variables[var_name].coordinates.split() + handler.variables[var_name].coordinates = \ + ' '.join(set(coordinates) | add_coordinates) + handler.close() + + + + + -- GitLab