interp_awap_by_mean.py

import xarray as xr
import numpy as np
import numba
import pandas as pd
from obspy.geodetics import kilometers2degrees


@numba.jit(parallel=True)
def _interp_over_lon_lat(data2d, lon_awap, lat_awap, lon_wrf_plus, lon_wrf_minus, lat_wrf_plus, lat_wrf_minus):
    """
    Private low-level function written using numba to speed-up the computation
    
    Parameters
    ----------
    data2d : 2darray
        Input observational data to interpolate    
    lon_awap : 1darray
        Longitudinal coordinates of data
    lat_awap : 1darray
        Latitudinal coordinates of data
    lon_wrf_plus : 
        Longitudinal coordinates of
    lon_wrf_minus : 
        Longitudinal coordinates of
    lat_wrf_plus : 
        Latitudinal coordinates of
    lat_wrf_minus : 
        Latitudinal coordinates of
                
    Returns
    -------
    data_interp : 2darray
        Data filtered on grid_out
    """
   
    ny, nx = np.shape(lat_wrf_plus)
    data2d_interp = np.zeros([ny, nx])

    for la in range(ny):
        for lo in range(nx):
            cond = ((lat_awap >= lat_wrf_minus[la,lo]) 
                    & (lat_awap <= lat_wrf_plus[la,lo]) 
                    & (lon_awap >= lon_wrf_minus[la,lo]) 
                    & (lon_awap <= lon_wrf_plus[la,lo]))
            if np.sum(cond) != 0:
                data2d_interp[la, lo] = np.mean(data2d[np.where(cond)])
            else:
                data2d_interp[la, lo] = np.nan
                    
    return data2d_interp
    

@numba.jit(parallel=True)
def _loop_over_time(data3d, lon_awap, lat_awap, lon_wrf_plus, lon_wrf_minus, lat_wrf_plus, lat_wrf_minus):
    
    ny, nx = np.shape(lat_wrf_plus)
    nt = len(data3d) 
    data3d_interp = np.zeros([nt, ny, nx])
    for t in range(nt):
        data3d_interp[t, :, :] = _interp_over_lon_lat(data3d[t, :, :], lon_awap, lat_awap, lon_wrf_plus, lon_wrf_minus, lat_wrf_plus, lat_wrf_minus)
    
    return data3d_interp


def interp_awap_by_mean(ds_awap, grid_wrf_file):
    """
    Perform an interpolation based on the mean of WRF grid cells
    
    Parameters
    ----------
    ds_awap : xr.Dataset (dims time, lat, lon)
        Input observational data to interpolate
    mask_awap: xr.DataArray
        Input observational mask    
    grid_wrf_file : xr.Dataset
        wrf_output grid on which interpolate the observations
            
    Returns
    -------
    data_interp : xr.DataArray
        Data interpolated on wrf grid
    """
    lat_min, lat_max = ds_awap.lat.min(), ds_awap.lat.max()
    lon_min, lon_max = ds_awap.lon.min(), ds_awap.lon.max()
    
    cond = (grid_wrf_file.XLONG_M >= lon_min) & (grid_wrf_file.XLONG_M <= lon_max) & (grid_wrf_file.XLAT_M >= lat_min) & (grid_wrf_file.XLAT_M <= lat_max)
    grid_file = grid_wrf_file.where(cond, drop=True).isel(Time=0)
    
    #ds = grid_file.isel(west_east=slice(1, None), south_north=slice(None, -1))
    ds = grid_file.isel(west_east=slice(1, None), south_north=slice(1, None))
    dx_wrf, dy_wrf = latlon2dydxWRF(ds['XLAT_M'], ds['XLONG_M'])
    
    lat_wrf = grid_file.XLAT_M[1:-1,1:-1].data
    lon_wrf = grid_file.XLONG_M[1:-1,1:-1].data
    dlon_wrf = kilometers2degrees(dx_wrf, radius=6371 * 1e3)
    dlat_wrf = kilometers2degrees(dy_wrf, radius=6371 * 1e3)      

    lon_wrf_plus = np.asarray(lon_wrf + dlon_wrf/2)
    lon_wrf_minus = np.asarray(lon_wrf - dlon_wrf/2)
    lat_wrf_plus = np.asarray(lat_wrf + dlat_wrf/2)
    lat_wrf_minus = np.asarray(lat_wrf - dlat_wrf/2)
    
    #create lat lon in 2d for awap
    lat2d, lon2d = xr.broadcast(ds_awap.lat, ds_awap.lon)
    
    lat_awap = np.asarray(lat2d)
    lon_awap = np.asarray(lon2d)
    time_awap = np.asarray(ds_awap.time)
    
    data3d = np.asarray(ds_awap.precip)
    
    res = _loop_over_time(data3d, lon_awap, lat_awap,
                          lon_wrf_plus, lon_wrf_minus,
                          lat_wrf_plus, lat_wrf_minus)
    
    data_interp = xr.DataArray(res, name='precip',
                               dims=('time', 'south_north', 'west_east'),
                               coords={'time':ds_awap.time,
                                       'lat':grid_file.XLAT_M[1:-1,1:-1],
                                       'lon':grid_file.XLONG_M[1:-1,1:-1]})
    
    return data_interp


def test_interp2d(ds_awap, grid_wrf_file):
    """
    Perform an interpolation based on the mean of WRF grid cells
    
    Parameters
    ----------
    ds_awap : xr.Dataset (2D field : lat, lon)
        Input observational data to interpolate
    mask_awap: xr.DataArray
        Input observational mask    
    grid_wrf_file : xr.Dataset
        wrf_output grid on which interpolate the observations
            
    Returns
    -------
    data_interp : xr.DataArray
        Data interpolated on wrf grid
    """
    lat_min, lat_max = ds_awap.lat.min(), ds_awap.lat.max()
    lon_min, lon_max = ds_awap.lon.min(), ds_awap.lon.max()
    
    cond = (grid_wrf_file.XLONG_M >= lon_min) & (grid_wrf_file.XLONG_M <= lon_max) & (grid_wrf_file.XLAT_M >= lat_min) & (grid_wrf_file.XLAT_M <= lat_max)
    grid_file = grid_wrf_file.where(cond, drop=True).isel(Time=0)
    
    ds = grid_file.isel(west_east=slice(1, None), south_north=slice(1, None))
    dx_wrf, dy_wrf = latlon2dydxWRF(ds['XLAT_M'], ds['XLONG_M'])
    
    lat_wrf = grid_file.XLAT_M[1:-1,1:-1].data
    lon_wrf = grid_file.XLONG_M[1:-1,1:-1].data
    dlon_wrf = kilometers2degrees(dx_wrf, radius=6371 * 1e3)
    dlat_wrf = kilometers2degrees(dy_wrf, radius=6371 * 1e3)

    lon_wrf_plus = np.asarray(lon_wrf + dlon_wrf/2)
    lon_wrf_minus = np.asarray(lon_wrf - dlon_wrf/2)
    lat_wrf_plus = np.asarray(lat_wrf + dlat_wrf/2)
    lat_wrf_minus = np.asarray(lat_wrf - dlat_wrf/2)
    
    #create 2d lat lon for awap
    lat2d, lon2d = xr.broadcast(ds_awap.lat, ds_awap.lon)
    
    lat_awap = np.asarray(lat2d)
    lon_awap = np.asarray(lon2d)
    
    data2d = np.asarray(ds_awap.precip)
    
    res = _interp_over_lon_lat(data2d, lon_awap, lat_awap,
                               lon_wrf_plus, lon_wrf_minus,
                               lat_wrf_plus, lat_wrf_minus)
    
    data2d_interp = xr.DataArray(res, name='precip', 
                                 dims=('south_north', 'west_east'), 
                                 coords={'time':ds_awap.time,
                                         'lat':grid_file.XLAT_M[1:-1,1:-1],
                                         'lon':grid_file.XLONG_M[1:-1,1:-1]})
    
    return data2d_interp


def latlon2dydxWRF(lat, lon, label='upper'):
    """
    Convert latitude and longitude arrays to elementary displacements in dy
    and dx of wrf_output files
    Parameters
    ----------
    lat : array-like
        Latitudinal spherical coordinates
    lon : array-like
        Longitudinal spherical coordinates
    dim : str
        Dimension along which the differentiation is performed, generally
        associated with the time dimension.
    label : {'upper', 'lower'}, optional
        The new coordinate in dimension dim will have the values of
        either the minuend\u2019s or subtrahend\u2019s coordinate for values \u2018upper\u2019
        and \u2018lower\u2019, respectively.
    Returns
    -------
    dy : array-like
        Zonal elementary displacement in cartesian coordinates
    dx : array-like
        Meridional elementary displacement in cartesian coordinates
    """
    EARTH_RADIUS = 6371 * 1e3
    dlat = lat.diff('south_north', label=label)
    dlon = lon.diff('west_east', label=label)
    dy = np.pi / 180. * EARTH_RADIUS * dlat
    # Need to slice the latitude data when there are duplicate values
    if label is 'upper':
        dx = (np.cos(np.pi / 180. * lat.isel(**{'west_east': slice(1, None)})) *
              np.pi / 180. * EARTH_RADIUS * dlon)
        dy = dy.isel(**{'west_east': slice(1, None)})
        dx = dx.isel(**{'south_north': slice(1, None)})
    elif label is 'lower':
        dx = (np.cos(np.pi / 180. * lat.isel(**{'west_east': slice(None, -1)})) *
              np.pi / 180. * EARTH_RADIUS * dlon)
        dy = dy.isel(**{'west_east': slice(None, -1)})
        dx = dx.isel(**{'south_north': slice(None, -1)})        
    return dy, dx