grids for Diana

example.py

from acf.raster import grid_info

#----------------------------------------------#

fnames = os.listdir(force_dir)

lons = np.empty(len(fnames))
lats = np.empty(len(fnames))
for i, fname in enumerate(fnames):
    pre, lat, lon = fname.split('_')
    lons[i] = float(lon)
    lats[i] = float(lat)
    
target, (y, x) = grid_info(lons, lats)

target.mask.plot()

#----------------------------------------------#

variables = ['prcp', 'tmax', 'tmin']

mask_shape = target['mask'].shape
var_shape = (len(time), ) + mask_shape
coords = dict(time=time, **target['mask'].coords)

for var in variables:
    target[var] = xray.DataArray(np.full(var_shape, np.nan), coords=coords,
                                 attrs=attrs[var],
                                 dims=('time', 'latitude', 'longitude'))

for i, fname in enumerate(fnames):
    df = pd.read_table(os.path.join(force_dir, fname),
                       engine='python', sep=' ',
                       names=['prcp', 'tmax', 'tmin', 'wind'])
    
    for var in variables:
        target[var].values[:, y[i], x[i]] = df[var].values
        
#----------------------------------------------#

grid_info.py

import numpy as np
import xray
from scipy.spatial import cKDTree


def latlon2yx(plats, plons, glats, glons):
    '''find y x coordinates'''

    if glons.ndim == 1 or glats.ndim == 1:
        glons, glats = np.meshgrid(glons, glats)

    combined = np.dstack(([glats.ravel(), glons.ravel()]))[0]
    points = list(np.vstack((np.array(plats), np.array(plons))).transpose())

    mytree = cKDTree(combined)
    dist, indexes = mytree.query(points, k=1)
    y, x = np.unravel_index(np.array(indexes), glons.shape)
    return y, x
    

def grid_info(lons, lats, decimals=4):
    '''return a xray.Dataset describing the grid shape/mask of a domain from 1D
    numpy Arrays of lons/lats'''
    grid = xray.Dataset()

    # get unique lats and lons
    lon = np.sort(np.unique(lons.round(decimals=decimals)))
    lat = np.sort(np.unique(lats.round(decimals=decimals)))

    # coords
    lon_da = xray.DataArray(lon, dims=('longitude', ), name='longitude',
                            attrs={'long_name': 'longitude coordinate'})
    lat_da = xray.DataArray(lat, dims=('latitude', ), name='latitude',
                            attrs={'long_name': 'latitude coordinate'})

    shape = (len(lat), len(lon))
    y, x = latlon2yx(lats, lons, lat, lon)

    # mask
    mask = np.full(shape, np.nan)
    mask[y, x] = 1
    grid['mask'] = xray.DataArray(mask, dims=('latitude', 'longitude'),
                                  name='mask',
                                  coords={'latitude': lat_da,
                                          'longitude': lon_da},
                                  attrs={'long_name': 'domain mask',
                                         'comment': '0 indicates grid cell '
                                         'is not active'})

    return grid, (y, x)