interpolate.py 6.09 KB
Newer Older
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
# coding=utf-8
import shutil

import os
from autosubmit.config.log import Log
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.utils import Utils, TempFile


class Interpolate(Diagnostic):
    """
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    3-dimensional conservative interpolation to the regular atmospheric grid.
    It can also be used for 2D (i,j) variables
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    :original author: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    :created: November 2012
    :last modified: June 2016
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    :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's name
    :type variable: str
    :param domain: variable's domain
    :type domain: str
    :param nemo_version: NEMO model version
    :type nemo_version: str
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    """
    def __init__(self, data_manager, startdate, member, chunk, variable, domain, nemo_version):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.variable = variable
        self.domain = domain
        self.nemo_version = nemo_version
        self.required_vars = [variable]
        self.generated_vars = [variable]
        self.tempTemplate = ''

Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    def __str__(self):
        return 'Interpolate Startdate: {0} Member: {1} Chunk: {2} ' \
               'Variable: {3}:{4}'.format(self.startdate, self.member, self.chunk, self.domain, self.variable)

    @classmethod
    def generate_jobs(cls, diags, options):
        """
            Creates a job for each chunk to compute the diagnostic

            :param diags: Diagnostics manager class
            :type diags: Diags
            :param options: variable, domain=ocean
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
            :type options: list[str]
        num_options = len(options) - 1
        if num_options < 1:
            raise Exception('You must specify the variable to average vertically')
        if num_options > 2:
            raise Exception('You must specify between 1 and 2 parameters for the interpolation diagnostic')
        variable = options[1]
        if num_options >= 2:
            domain = int(options[2])
        else:
            domain = 'ocean'

        job_list = list()
        for startdate, member, chunk in diags.exp_manager.get_chunk_list():
            job_list.append(Interpolate(diags.data_manager, startdate, member, chunk,
                                        variable, domain, diags.nemo_version))
        return job_list

    def compute(self):
        """
        Runs the diagnostic
        """
        variable_file = self.data_manager.get_file(self.domain, self.variable, self.startdate, self.member, self.chunk)
        cdo = Utils.cdo
        nco = Utils.nco
        handler = Utils.openCdf(variable_file)
        if 'lev' in handler.dimensions:
            num_levels = handler.dimensions['lev'].size
            has_levels = True
        else:
            num_levels = 1
            has_levels = False
        handler.close()
        for lev in range(0, num_levels):
            self._interpolate_level(lev, has_levels, variable_file)

        temp = TempFile.get()
        if has_levels:
            nco.ncrcat(input=self._get_level_file(0), output=temp,
                       options="-n {0},2,1 -v '{1}'".format(num_levels, self.variable))

        else:
            Utils.move_file(self._get_level_file(0),  temp)

        handler = Utils.openCdf(temp)
        handler.renameDimension('record', 'lev')
        handler.close()
        nco.ncpdq(input=temp, output=temp, options='-O -h -a time,lev')

        if has_levels:
            nco.ncks(input=variable_file, output=temp, options='-A -v lev')
            for lev in range(0, num_levels):
                os.remove(self._get_level_file(lev))
        temp2 = TempFile.get()
        cdo.setgrid('t106grid', input=temp, output=temp2)
        os.remove(temp)
        if self.nemo_version[6:9] == '025':
            cdo.invertlatdata(input=temp2, output=temp2)
        if not has_levels:
            nco.ncks(input=temp2, output=temp2, options='-O -v sic,lat,lon,time')

        self.data_manager.send_file(temp2, self.domain, self.variable, self.startdate, self.member, self.chunk,
                                    grid='regular')

    def _get_level_file(self, lev):
        if not self.tempTemplate:
            self.tempTemplate = TempFile.get(suffix='_01.nc')
            # self.tempTemplate = 'temp_01.nc'
        return self.tempTemplate.replace('_01.nc', '_{0:02d}.nc'.format(lev + 1))

    def _interpolate_level(self, lev, has_levels, input_file):
        nco = Utils.nco
        temp = TempFile.get()
        if has_levels:
            nco.ncks(input=input_file, output=temp, options='-O -d lev,{0} -v {1}'.format(lev, self.variable))
            nco.ncwa(input=temp, output=temp, options='-O -h -a lev')
        else:
            shutil.copy(input_file, temp)
        namelist_file = TempFile.get(suffix='')
        scrip_use_in = open(namelist_file, 'w')
        scrip_use_in.writelines("&remap_inputs\n")
        scrip_use_in.writelines("   remap_wgt   = '/esnas/autosubmit/con_files/"
                                "weigths/{0}/rmp_{0}_to_regular_lev{1}.nc'\n".format(self.nemo_version, lev + 1))
        scrip_use_in.writelines("   infile      = '{0}'\n".format(temp))
        scrip_use_in.writelines("   invertlat   = FALSE\n")
        scrip_use_in.writelines("   var         = '{0}'\n".format(self.variable))
        scrip_use_in.writelines("   fromregular = FALSE\n")
        scrip_use_in.writelines("   outfile     = '{0}'\n".format(temp))
        scrip_use_in.writelines("/\n")
        scrip_use_in.close()
        Utils.execute_shell_command('/home/Earth/jvegas/pyCharm/cfutools/interpolation/scrip_use '
                                    '{0}'.format(namelist_file), Log.NO_LOG)
        os.remove(namelist_file)
        nco.ncecat(input=temp, output=temp, options="-O -h")
        shutil.move(temp, self._get_level_file(lev))
        Log.debug("Level {0} ready", lev)