# coding=utf-8
"""Rotate two u v variables to align with latitude and longitude"""
import shutil
from bscearth.utils.log import Log
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticDomainOption, DiagnosticVariableOption
from earthdiagnostics.utils import Utils, TempFile
[docs]class Rotation(Diagnostic):
"""
Rotate two u v variables to align with latitude and longitude
:original author: Virginie Guemas <virginie.guemas@bsc.es>
:contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>
:created: September 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 domain: variable's domain
:type domain: Domain
"""
alias = 'rotate'
"Diagnostic alias for the configuration file"
def __init__(self, data_manager, startdate, member, chunk, domain, variableu, variablev, executable):
Diagnostic.__init__(self, data_manager)
self.startdate = startdate
self.member = member
self.chunk = chunk
self.variableu = variableu
self.variablev = variablev
self.domain = domain
self.executable = executable
self.tempTemplate = None
def __eq__(self, other):
if self._different_type(other):
return False
return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \
self.domain == other.domain and self.variableu == other.variableu and \
self.variablev == other.variablev and self.executable == other.executable
def __str__(self):
return 'Rotate variables Startdate: {0} Member: {1} Chunk: {2} Variables: {3}:{4} , ' \
'{3}:{5}'.format(self.startdate, self.member, self.chunk, self.domain, self.variableu,
self.variablev)
[docs] @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, zonal, value, domain=ocean
:type options: list[str]
:return:
"""
options_available = (DiagnosticDomainOption(),
DiagnosticVariableOption(diags.data_manager.config.var_manager, 'variableu'),
DiagnosticVariableOption(diags.data_manager.config.var_manager, 'variablev'),
DiagnosticOption('executable',
'/home/Earth/jvegas/pyCharm/cfutools/interpolation/rotateUVorca'))
options = cls.process_options(options, options_available)
job_list = list()
for startdate, member, chunk in diags.config.experiment.get_chunk_list():
job_list.append(Rotation(diags.data_manager, startdate, member, chunk,
options['domain'], options['variableu'], options['variablev'],
options['executable']))
return job_list
[docs] def request_data(self):
"""Request data required by the diagnostic"""
self.ufile = self.request_chunk(self.domain, self.variableu, self.startdate, self.member, self.chunk)
self.vfile = self.request_chunk(self.domain, self.variablev, self.startdate, self.member, self.chunk)
[docs] def declare_data_generated(self):
"""Declare data to be generated by the diagnostic"""
self.urotated_file = self.declare_chunk(self.domain, self.variableu, self.startdate, self.member, self.chunk,
grid='rotated')
self.vrotated_file = self.declare_chunk(self.domain, self.variablev, self.startdate, self.member, self.chunk,
grid='rotated')
[docs] def compute(self):
"""Run the diagnostic"""
handler = Utils.open_cdf(self.ufile.local_file)
if 'lev' in handler.dimensions:
self.num_levels = handler.dimensions['lev'].size
self.has_levels = True
else:
self.num_levels = 1
self.has_levels = False
handler.close()
for lev in range(0, self.num_levels):
self._rotate_level(lev)
urotated = self._merge_levels(self.variableu, 'u')
vrotated = self._merge_levels(self.variablev, 'v')
ufile_handler = Utils.open_cdf(self.ufile.local_file)
self._add_metadata_and_vars(ufile_handler, urotated, self.variableu)
ufile_handler.close()
self.urotated_file.set_local_file(urotated)
vfile_handler = Utils.open_cdf(self.vfile.local_file)
self._add_metadata_and_vars(vfile_handler, vrotated, self.variablev)
vfile_handler.close()
self.vrotated_file.set_local_file(urotated)
def _merge_levels(self, var, direction):
temp = TempFile.get()
if self.has_levels:
Utils.nco.ncecat(input=self._get_level_file(0, direction), output=temp,
options=("-n {0},2,1 -v '{1}'".format(self.num_levels, var),))
handler = Utils.open_cdf(temp)
if 'record' in handler.dimensions:
handler.renameDimension('record', 'lev')
handler.close()
Utils.nco.ncpdq(input=temp, output=temp, options=('-O -h -a time,lev',))
Utils.rename_variables(temp, {'x': 'i', 'y': 'j'}, must_exist=False)
else:
Utils.move_file(self._get_level_file(0, direction), temp)
return temp
def _rotate_level(self, lev):
ufile = self._extract_level(self.ufile.local_file, self.variableu, lev)
vfile = self._extract_level(self.vfile.local_file, self.variablev, lev)
namelist_file = self._create_namelist(ufile, self._get_level_file(lev, 'u'),
vfile, self._get_level_file(lev, 'v'))
Utils.execute_shell_command('{0} {1}'.format(self.executable, namelist_file), Log.INFO)
def _extract_level(self, input_file, var, level):
temp = TempFile.get()
if self.has_levels:
Utils.nco.ncks(input=input_file, output=temp, options=('-O -d lev,{0} -v {1},lat,lon'.format(level, var),))
Utils.nco.ncwa(input=temp, output=temp, options=('-O -h -a lev',))
else:
shutil.copy(input_file, temp)
return temp
def _create_namelist(self, ufile, urotated, vfile, vrotated):
namelist_file = TempFile.get(suffix='')
rotate_namelist = open(namelist_file, 'w')
rotate_namelist.write('&nam_rotUV\n')
rotate_namelist.write(' Ufilein = "{0}"\n'.format(ufile))
rotate_namelist.write(' Uvarin = "{0}"\n'.format(self.variableu))
rotate_namelist.write(' Vfilein = "{0}"\n'.format(vfile))
rotate_namelist.write(' Vvarin = "{0}"\n'.format(self.variablev))
rotate_namelist.write(' meshmask = "mask.nc"\n')
rotate_namelist.write(' Ufileout = "{0}"\n'.format(urotated))
rotate_namelist.write('Vfileout = "{0}"\n'.format(vrotated))
rotate_namelist.writelines("/\n")
rotate_namelist.close()
return namelist_file
def _add_metadata_and_vars(self, reference_file_handler, rotaded_file, var_name):
rotated_handler = Utils.open_cdf(rotaded_file)
self._copy_extra_variables(reference_file_handler, rotated_handler)
Utils.copy_attributes(rotated_handler.variables[var_name], reference_file_handler.variables[var_name],
('_FillValue',))
rotated_handler.close()
def _copy_extra_variables(self, reference_file_handler, rotated_handler):
for var in reference_file_handler.variables.keys():
if var not in rotated_handler.variables.keys() and var not in [self.variableu, self.variablev]:
Utils.copy_variable(reference_file_handler, rotated_handler, var, True, True)
def _get_level_file(self, lev, direction):
if not self.tempTemplate:
self.tempTemplate = TempFile.get(suffix='_01.nc')
# self.tempTemplate = 'temp_01.nc'
return self.tempTemplate.replace('_01.nc', '_{1}_{0:02d}.nc'.format(lev + 1, direction))