interpolatecdo.py 9.69 KB
Newer Older
import numpy as np

from earthdiagnostics.diagnostic import *
from earthdiagnostics.modelingrealm import ModelingRealm, ModelingRealms
from earthdiagnostics.utils import Utils, TempFile

class InterpolateCDO(Diagnostic):
    """
    3-dimensional conservative interpolation to the regular atmospheric grid.
    It can also be used for 2D (i,j) variables

    :original author: Javier Vegas-Regidor<javier.vegas@bsc.es>

    :created: October 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 variable: variable's name
    :type variable: str
    :param domain: variable's domain
    :type domain: ModelingRealm
    :param model_version: model version
    :type model_version: str
    """

    "Diagnostic alias for the configuration file"

    BILINEAR = 'bilinear'
    BICUBIC = 'bicubic'
    CONSERVATIVE = 'conservative'
    CONSERVATIVE2 = 'conservative2'

    METHODS = [BILINEAR, BICUBIC, CONSERVATIVE, CONSERVATIVE2]

    def __init__(self, data_manager, startdate, member, chunk, domain, variable, target_grid, model_version,
                 mask_oceans, original_grid, weights):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.variable = variable
        self.domain = domain
        self.model_version = model_version
        self.required_vars = [variable]
        self.generated_vars = [variable]
        self.tempTemplate = ''
        self.grid = target_grid
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
        self.mask_oceans = mask_oceans
        self.weights = weights

    def __eq__(self, other):
        return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \
            self.model_version == other.model_version and self.domain == other.domain and \
            self.variable == other.variable and self.mask_oceans == other.mask_oceans and self.grid == other.grid and \
            self.original_grid == other.original_grid
        return 'Interpolate with CDO Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \
               'Variable: {0.domain}:{0.variable} Target grid: {0.grid} Original grid: {0.original_grid} ' \
               'Mask ocean: {0.mask_oceans} Model: {0.model_version}'.format(self)
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
    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: target_grid, variable, domain=ocean
            :type options: list[str]
            :return:
            """
Javier Vegas-Regidor's avatar
Javier Vegas-Regidor committed
        options_available = (DiagnosticDomainOption(default_value=ModelingRealms.ocean),
                             DiagnosticVariableListOption(diags.data_manager.config.var_manager, 'variables'),
                             DiagnosticOption('target_grid', diags.config.experiment.atmos_grid.lower()),
                             DiagnosticChoiceOption('method', InterpolateCDO.METHODS, InterpolateCDO.BILINEAR),
                             DiagnosticOption('original_grid', ''),
                             DiagnosticBoolOption('weights_from_mask', True)
                             )
        options = cls.process_options(options, options_available)
        target_grid = cls._translate_ifs_grids_to_cdo_names(options['target_grid'])
        if not target_grid:
            raise Exception('Target grid not provided')
        job_list = list()
        weights = TempFile.get()
        method = options['method'].lower()
        cls._compute_weights(diags, method, options, target_grid, weights)
        for var in options['variables']:
            for startdate, member, chunk in diags.config.experiment.get_chunk_list():
                job_list.append(InterpolateCDO(diags.data_manager, startdate, member, chunk,
                                               options['domain'], var, target_grid,
                                               diags.config.experiment.model_version, options['mask_oceans'],
                                               options['original_grid'], weights))
        return job_list

    @classmethod
    def _compute_weights(cls, diags, method, options, target_grid, weights):
        if options['weights_from_mask']:
            temp = cls.get_sample_grid_file()
        else:
            startdate, member, chunk = diags.config.experiment.get_chunk_list()[0]
            temp = diags.data_manager.get_file(options['domain'], options['variable'], startdate, member, chunk,
                                               grid=options['original_grid'])
        if method == InterpolateCDO.BILINEAR:
            Utils.cdo.genbil(target_grid, input=temp, output=weights)
        elif method == InterpolateCDO.BICUBIC:
            Utils.cdo.genbic(target_grid, input=temp, output=weights)
        elif method == InterpolateCDO.CONSERVATIVE:
            Utils.cdo.genycon(target_grid, input=temp, output=weights)
        elif method == InterpolateCDO.CONSERVATIVE2:
            Utils.cdo.gencon2(target_grid, input=temp, output=weights)

    @classmethod
    def get_sample_grid_file(cls):
        temp = TempFile.get()
        Utils.nco.ncks(input='mask.nc', output=temp, options=('-O -v tmask,lat,lon,gphif,glamf',))
        handler = Utils.openCdf(temp)
        lon = handler.variables['lon']
        lon.units = "degrees_east"
        lon.long_name = "Longitude"
        lon.nav_model = "Default grid"
        lon.standard_name = "longitude"
        lon.short_name = "lon"
        lon.bounds = 'lon_bnds'

        lat = handler.variables['lat']
        lat.units = "degrees_north"
        lat.long_name = "Latitude"
        lat.nav_model = "Default grid"
        lat.standard_name = "latitude"
        lat.short_name = "lat"
        lat.bounds = 'lat_bnds'

        handler.createDimension('bounds', 4)

        lon_bnds = handler.createVariable('lon_bnds', lon.datatype, ('j', 'i', 'bounds'))
        corner_lat = handler.variables['glamf'][0, ...]
        lon_bnds[:, :, 0] = corner_lat
        lon_bnds[:, :, 1] = np.roll(corner_lat, 1, 0)
        lon_bnds[:, :, 2] = np.roll(corner_lat, -1, 1)
        lon_bnds[:, :, 3] = np.roll(lon_bnds[:, :, 1], -1, 1)

        lat_bnds = handler.createVariable('lat_bnds', lat.datatype, ('j', 'i', 'bounds'))
        corner_lat = handler.variables['gphif'][0, ...]
        lat_bnds[:, :, 0] = corner_lat
        lat_bnds[:, :, 1] = np.roll(corner_lat, 1, 0)
        lat_bnds[:, :, 2] = np.roll(corner_lat, 1, 1)
        lat_bnds[:, :, 3] = np.roll(lat_bnds[:, :, 1], 1, 1)
        lat_bnds[0, :, 1] = lat_bnds[1, 0, 1] - 1
        lat_bnds[0, :, 3] = lat_bnds[1, 0, 3] - 1

        tmask = handler.variables['tmask']
        tmask.coordinates = 'time lev lat lon'

        handler.close()

        Utils.nco.ncks(input=temp, output=temp, options=('-O -x -v gphif,glamf',))
    def _translate_ifs_grids_to_cdo_names(cls, target_grid):
        if target_grid.upper().startswith('T159L'):
            target_grid = 't106grid'
        if target_grid.upper().startswith('T255L'):
            target_grid = 't170grid'
        if target_grid.upper().startswith('T511L'):
            target_grid = 't340grid'
    def request_data(self):
        self.original = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk,
                                           grid=self.original_grid)

    def declare_data_generated(self):
        self.regridded = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk,
                                            grid=self.grid)

    def compute(self):
        """
        Runs the diagnostic
        """
        variable_file = TempFile.get()
        Utils.copy_file(self.original.local_file, variable_file)
        Utils.rename_variables(variable_file, {'jpib': 'i', 'jpjb': 'j', 'x': 'i', 'y': 'j',
                                               'time_counter': 'time', 't': 'time',
                                               'SSTK_ens0': 'tos', 'SSTK_ens1': 'tos', 'SSTK_ens2': 'tos',
                                               'nav_lat': 'lat', 'nav_lon': 'lon'},
                               must_exist=False, rename_dimension=True)
        handler = Utils.openCdf(variable_file)
        var = handler.variables[self.variable]
        coordinates = list()
        for dim in var.dimensions:
            if dim == 'i':
                if 'lat' in handler.variables:
                    coordinates.append('lat')
                else:
                    coordinates.append('latitude')
            elif dim == 'j':
                if 'lon' in handler.variables:
                    coordinates.append('lon')
                else:
                    coordinates.append('longitude')
            else:
                coordinates.append(dim)
        var.coordinates = ' '.join(coordinates)
        if self.mask_oceans:
            mask = Utils.get_mask(Basins().Global).astype(float)
            mask[mask == 0] = np.nan
            var[:] = mask * var[:]
        handler.close()

        temp = TempFile.get()
        Utils.cdo.remap(','.join((self.grid.split('_')[0], self.weights)), input=variable_file, output=temp)

        handler = Utils.openCdf(temp)
        handler.variables[self.variable].units = units
        handler.close()

        self.regridded.set_local_file(temp)