Source code for epygram.geometries.SpectralGeometry

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (c) Météo France (2014-)
# This software is governed by the CeCILL-C license under French law.
# http://www.cecill.info
"""
Contains the class that handles spectral parameters and spectral transforms.
"""

from __future__ import print_function, absolute_import, unicode_literals, division

import numpy
import os

from bronx.system import memory

from epygram import config, epygramError, epylog
from epygram.util import RecursiveObject

transforms_lib_init_env_kwargs = dict()


def nearest_greater_FFT992compliant_int(guess):
    """
    Returns the first integer *n* greater than *guess* that satisfies
    n = 2^(1+i) x 3^j x 5^k, with (i,j,k) being positive integers.
    """
    # these are defined for dimensions up to 10000 points at least
    M2 = 14
    M3 = 10
    M5 = 7
    fft_compliant_dims = numpy.zeros((M2, M3, M5))
    for i in range(M2):
        for j in range(M3):
            for k in range(M5):
                fft_compliant_dims[i, j, k] = 2 ** (1 + i) * 3 ** j * 5 ** k
    fft_compliant_dims = sorted(list(set(fft_compliant_dims.flatten())))
    for i in range(len(fft_compliant_dims)):
        if fft_compliant_dims[i] >= guess:
            result = int(fft_compliant_dims[i])
            break
    return result


[docs]def truncation_from_gridpoint_dims(dimensions, grid='linear', stretching_coef=1.): """ Compute truncation from gridpoint dimensions, according to the kind of **grid**. :param dimensions: dict containing dimensions, among: {'X':..., 'Y':...} for LAM grids, {'lat_number':..., 'max_lon_number':...} for Gauss grids :param grid: how to choose the truncation, among ('linear', 'quadratic', 'cubic') :param stretching_coef: stretching or dilatation coefficient for stretched Gauss grids (has an impact on the gridpoint/spectral dimensions) Formula taken from "Spectral transforms in the cycle 45 of ARPEGE/IFS", http://www.umr-cnrm.fr/gmapdoc/IMG/pdf/ykts45.pdf """ truncation = {} spfactor = {'linear':2, 'quadratic':3, 'cubic':4}[grid] if all([k in dimensions.keys() for k in ('X', 'Y')]): # LAM truncation['in_X'] = (dimensions['X'] - 1) // spfactor truncation['in_Y'] = (dimensions['Y'] - 1) // spfactor elif all([k in dimensions.keys() for k in ('lat_number', 'max_lon_number')]): # Gauss if stretching_coef > 1.: truncation['max'] = min(2 * dimensions['lat_number'] - 3, dimensions['max_lon_number'] - 1) // spfactor else: truncation['max'] = (dimensions['max_lon_number'] - 1) // spfactor return truncation
def _guess_compliant_lonlat(nlon): nlon = nearest_greater_FFT992compliant_int(nlon) if nlon % 4 != 0: nlat = nlon // 2 + 1 else: nlat = nlon // 2 return nlon, nlat def gridpoint_dims_from_truncation(truncation, grid='linear', stretching_coef=1.): """ Compute truncation from gridpoint dimensions, according to the kind of **grid**. :param truncation: dict containing truncation info, among: {'in_X':..., 'in_Y':...} for LAM grids, {'max':...} for Gauss grids :param grid: how to choose the truncation, among ('linear', 'quadratic', 'cubic') :param stretching_coef: stretching or dilatation coefficient for stretched Gauss grids (has an impact on the gridpoint/spectral dimensions) Formula taken from "Spectral transforms in the cycle 45 of ARPEGE/IFS", http://www.umr-cnrm.fr/gmapdoc/IMG/pdf/ykts45.pdf """ dimensions = {} spfactor = {'linear':2, 'quadratic':3, 'cubic':4}[grid] if all([k in truncation.keys() for k in ('in_X', 'in_Y')]): # LAM dimensions['X'] = spfactor * truncation['in_X'] + 1 dimensions['X'] = nearest_greater_FFT992compliant_int(dimensions['X']) dimensions['Y'] = spfactor * truncation['in_Y'] + 1 dimensions['Y'] = nearest_greater_FFT992compliant_int(dimensions['Y']) elif all([k in truncation.keys() for k in ('max',)]): # Gauss nlon = spfactor * truncation['max'] + 1 nlon, nlat = _guess_compliant_lonlat(nlon) if stretching_coef > 1. and truncation['max'] > min(2*nlat-3, nlon) // spfactor: nlon, nlat = _guess_compliant_lonlat(nlon + 1) epylog.warning("Grid will be sub-{}: nlon > 2*tronc + 1 because of stretching !".format(grid)) dimensions['max_lon_number'] = nlon dimensions['lat_number'] = nlat return dimensions def complete_gridpoint_dimensions(longitudes, latitudes, truncation, grid, stretching_coef): """Complete missing gridpoint dimensions from guessing values.""" if latitudes is None and longitudes is None: dims = gridpoint_dims_from_truncation({'max': truncation}, grid=grid, stretching_coef=stretching_coef) latitudes = dims['lat_number'] longitudes = dims['max_lon_number'] elif longitudes is None: longitudes = 2 * self.latitudes elif latitudes is None: if longitudes % 4 != 0: latitudes = longitudes // 2 + 1 else: latitudes = longitudes // 2 return longitudes, latitudes
[docs]class SpectralGeometry(RecursiveObject): """Handles the spectral geometry and transforms for a H2DField.""" # in order to avoid calling etrans_inq for large truncations _usual_SPdatasize_for_global_trunc = {'triangular':{1198:719400, 1798:1619100}} def __init__(self, space, truncation): """ :param space: Name of spectral space. (among ['legendre', 'bi-fourier', 'fourier']) :param truncation: Handles the spectral truncation parameters. """ self.add_attr_inlist('space', ['legendre', 'bi-fourier', 'fourier']) self.add_attr_dict('truncation') self.space = space self.truncation = truncation super(SpectralGeometry, self).__init__() if os.name == 'posix': meminfo = memory.LinuxMemInfo() else: raise NotImplementedError('MemInfo for os.name=={}'.format(os.name)) self._total_system_memory = meminfo.system_RAM(unit='MiB') @classmethod def _translib(cls): """ Accessor to the inner transforms library, not imported before the first call. """ if not hasattr(cls, '_transforms_lib'): from epygram.extra import ialsptrans4py cls._transforms_lib = ialsptrans4py if transforms_lib_init_env_kwargs.pop('trigger', False): ialsptrans4py.init_env(**transforms_lib_init_env_kwargs) return cls._transforms_lib def _prevent_swapping(self): if (self.space == 'legendre' and (float(self.needed_memory) / (1024 ** 2.) >= config.prevent_swapping_legendre * self._total_system_memory)): needed_mem_in_mb = float(self.needed_memory) / (1024 ** 2.) total_mem_in_mb = float(self._total_system_memory) raise epygramError('Legendre spectral transforms need {:.1f} MB \ memory, while only {:.1f} MB is available: \ SWAPPING prevented !'.format(needed_mem_in_mb, total_mem_in_mb)) def _prevent_limited_stack(self): if (self.space == 'legendre' and self.truncation['max'] > 1200): epylog.warning('Caution: large Legendre truncation may need very large stacksize !')
[docs] def legendre_known_spectraldata_size(self): """In order to avoid calling trans_inq for large truncations.""" if self.truncation['shape'] in self._usual_SPdatasize_for_global_trunc: return self._usual_SPdatasize_for_global_trunc.get( self.truncation['shape']).get(self.truncation['max'], None)
[docs] def trans_inq(self, gpdims): """ Wrapper to IAL's TRANS_INQ. :param dict gpdims: gridpoints dimensions """ self._prevent_swapping() self._prevent_limited_stack() return self._translib().w_trans_inq(gpdims['lat_number'], self.truncation['max'], len(gpdims['lon_number_by_lat']), numpy.array(gpdims['lon_number_by_lat']), config.KNUMMAXRESOL)
[docs] def etrans_inq(self, gpdims): """ Wrapper to IAL's ETRANS_INQ. :param dict gpdims: gridpoints dimensions """ self._prevent_swapping() self._prevent_limited_stack() return self._translib().w_etrans_inq(gpdims['X'], gpdims['Y'], gpdims['X_CIzone'], gpdims['Y_CIzone'], self.truncation['in_X'], self.truncation['in_Y'], config.KNUMMAXRESOL, gpdims['X_resolution'], gpdims['Y_resolution'])
@property def needed_memory(self): """Memory needed for transforms, in bytes.""" if self.space == 'legendre': return self.truncation['max'] ** 3 / 2 * 8 else: raise NotImplementedError('space:' + self.space)
[docs] def sp2gp(self, data, gpdims, spectral_coeff_order=config.spectral_coeff_order): """ Makes the transform of the spectral data contained in *data* (assumed this spectral geometry is that of *data*) to gridpoint space, defined by its dimensions contained in *gpdims*, and returns the gridpoint data. :param data: spectral data :param dict gpdims: gridpoints dimensions :param spectral_coeff_order: among 'model' or 'FA', cf. default and description in config.spectral_coeff_order Input and output data are both 1D. """ self._prevent_swapping() self._prevent_limited_stack() if self.space == 'bi-fourier': gpdata = self._translib().w_spec2gpt_lam(gpdims['X'], gpdims['Y'], gpdims['X_CIzone'], gpdims['Y_CIzone'], self.truncation['in_X'], self.truncation['in_Y'], config.KNUMMAXRESOL, len(data), False, # no derivatives spectral_coeff_order != 'model', gpdims['X_resolution'], gpdims['Y_resolution'], data)[0] elif self.space == 'legendre': gpdata = self._translib().w_spec2gpt_gauss(gpdims['lat_number'], self.truncation['max'], config.KNUMMAXRESOL, sum(gpdims['lon_number_by_lat']), len(gpdims['lon_number_by_lat']), numpy.array(gpdims['lon_number_by_lat']), len(data), False, # no derivatives spectral_coeff_order != 'model', data)[0] elif self.space == 'fourier': if self.truncation['in_Y'] > 1: gpdata = self._translib().w_spec2gpt_fft1d(len(data), self.truncation['in_Y'], data, gpdims['Y']) else: gpdata = numpy.ones(gpdims['Y']) * data[0] else: raise epygramError("unknown spectral space:" + self.space + ".") return gpdata
[docs] def gp2sp(self, data, gpdims, spectral_coeff_order=config.spectral_coeff_order): """ Makes the transform of the gridpoint data contained in *data* to the spectral space and truncation of this object, and returns the spectral data. :param data: gridpoint data :param dict gpdims: gridpoints dimensions :param spectral_coeff_order: among 'model' or 'FA', cf. default and description in config.spectral_coeff_order Input and output data are both 1D. """ self._prevent_swapping() self._prevent_limited_stack() if self.space == 'bi-fourier': SPdatasize = self.etrans_inq(gpdims)[1] spdata = self._translib().w_gpt2spec_lam(SPdatasize, gpdims['X'], gpdims['Y'], gpdims['X_CIzone'], gpdims['Y_CIzone'], self.truncation['in_X'], self.truncation['in_Y'], config.KNUMMAXRESOL, gpdims['X_resolution'], gpdims['Y_resolution'], spectral_coeff_order != 'model', data) elif self.space == 'legendre': SPdatasize = self.trans_inq(gpdims)[1] SPdatasize *= 2 # complex coefficients spdata = self._translib().w_gpt2spec_gauss(SPdatasize, gpdims['lat_number'], self.truncation['max'], config.KNUMMAXRESOL, len(gpdims['lon_number_by_lat']), numpy.array(gpdims['lon_number_by_lat']), len(data), spectral_coeff_order != 'model', data) elif self.space == 'fourier': # 1D case SPdatasize = self.etrans_inq(gpdims)[1] if self.truncation['in_Y'] <= 1: spdata = numpy.zeros(SPdatasize) spdata[0] = data[0] else: raise NotImplementedError("direct transform for 1D fourier" + " transform.") else: raise epygramError("unknown spectral space:" + self.space + ".") return spdata
[docs] def compute_xy_spderivatives(self, data, gpdims, spectral_coeff_order=config.spectral_coeff_order): """ Compute the derivatives of the spectral data contained in *data* in spectral space (assumed this spectral geometry is that of *data*) and return it in gridpoint space, defined by its dimensions contained in *gpdims*. :param data: spectral data :param dict gpdims: gridpoints dimensions :param spectral_coeff_order: among 'model' or 'FA', cf. default and description in config.spectral_coeff_order Returns: (dz/dx, dz/dy) Input and output data are both 1D. """ self._prevent_swapping() self._prevent_limited_stack() if self.space == 'bi-fourier': gpdata = self._translib().w_spec2gpt_lam(gpdims['X'], gpdims['Y'], gpdims['X_CIzone'], gpdims['Y_CIzone'], self.truncation['in_X'], self.truncation['in_Y'], config.KNUMMAXRESOL, len(data), True, # derivatives spectral_coeff_order != 'model', gpdims['X_resolution'], gpdims['Y_resolution'], data) gpdata = (gpdata[2], gpdata[1]) elif self.space == 'legendre': gpdata = self._translib().w_spec2gpt_gauss(gpdims['lat_number'], self.truncation['max'], config.KNUMMAXRESOL, sum(gpdims['lon_number_by_lat']), len(gpdims['lon_number_by_lat']), numpy.array(gpdims['lon_number_by_lat']), len(data), True, # derivatives spectral_coeff_order != 'model', data) gpdata = (gpdata[2], gpdata[1]) elif self.space == 'fourier': raise NotImplementedError('fourier(1D): not yet !') else: raise epygramError("unknown spectral space:" + self.space + ".") return gpdata