kl_geotiff

from netCDF4 import Dataset
import numpy as np
import sys
import matplotlib.pyplot as plt
import datetime
import os
import matplotlib.ticker as ticker
import glob
import pyresample

def kl_2_ll(ffile2):
    ffile = open(ffile2,'r')
    lines = ffile.readlines()
    lon_00 = np.zeros(shape=(81, 77))
    lat_00 = np.zeros(shape=(81, 77))
    for iline in lines:
        #float(ilat), float(ilon), int(stolb), int(stroka) = iline.split()
        ilat = float(iline.split()[0])
        ilon = float(iline.split()[1])
        stolb = int(iline.split()[2])
        stroka = int(iline.split()[3])
        lon_00[stroka-1,stolb-1] = ilon
        lat_00[stroka-1,stolb-1] = ilat
    ffile.close()
    return lon_00, lat_00

def resample_kl_ice_grid(fname, lons, lats, data):
    # Reproject to Klyachkin ice model grid
    lon_kl, lat_kl = kl_2_ll('grid/grid_kl.txt')

    # Setka dannyh kotorye nuzhno sproecirovat
    orig_def = pyresample.geometry.SwathDefinition(lons=lons, lats=lats)

    # Klyachkin rep
    targ_def_kl = pyresample.geometry.SwathDefinition(lons=lon_kl, lats=lat_kl)


    # Klyachkin
    nn_kl = pyresample.kd_tree.resample_gauss(orig_def, data, \
                           targ_def_kl, radius_of_influence=100000, neighbours=1,\
                           sigmas=250000, fill_value=np.nan)

    np.savetxt("%s.txt" % fname, nn_kl, fmt = "%8.3f")
    return nn_kl

def calc_deformations(dx, dy, normalization=False, normalization_time=None, cell_size=10.):
    ''' Calculate deformation fields from U and V '''
    m_div = np.empty((dx.shape[0], dx.shape[1],))
    m_div[:] = np.NAN
    m_curl = np.empty((dx.shape[0], dx.shape[1],))
    m_curl[:] = np.NAN
    m_shear = np.empty((dx.shape[0], dx.shape[1],))
    m_shear[:] = np.NAN
    m_tdef = np.empty((dx.shape[0], dx.shape[1],))
    m_tdef[:] = np.NAN

    # Invert meridional component
    dy = dy * (-1)

    # Calculate magnitude
    mag = np.hypot(dx, dy)

    # Normilize u and v to 1 hour
    if not normalization:
        pass
    else:
        # Normilize to cell size
        dx = dx / cell_size
        dy = dy / cell_size

        # Normilize to time
        dx = dx / normalization_time
        dy = dy / normalization_time

    # Test
    plt.clf()
    plt.imshow(m_div)

    for i in range(1, dx.shape[0] - 1):
        for j in range(1, dx.shape[1] - 1):
            # div
            if (np.isnan(dx[i, j + 1]) == False and np.isnan(dx[i, j - 1]) == False
                    and np.isnan(dy[i - 1, j]) == False and np.isnan(dy[i + 1, j]) == False
                    and (np.isnan(dx[i, j]) == False or np.isnan(dy[i, j]) == False)):
                # m_div[i,j] = 0.5*((u_int[i,j + 1] - u_int[i,j - 1])  + (v_int[i + 1,j] - v_int[i - 1,j]))/m_cell_size
                m_div[i, j] = 0.5 * ((dx[i, j + 1] - dx[i, j - 1])
                                                      + (dy[i - 1, j] - dy[i + 1, j]))
                #print m_div[i,j]
            # curl
            if (np.isnan(dy[i, j + 1]) == False and np.isnan(dy[i, j - 1]) == False and
                    np.isnan(dx[i - 1, j]) == False and np.isnan(dx[i + 1, j]) == False
                    and (np.isnan(dx[i, j]) == False or np.isnan(dy[i, j]) == False)):
                m_curl[i, j] = 0.5 * (dy[i, j + 1] - dy[i, j - 1]
                                      - dx[i - 1, j] + dx[i + 1, j]) / cell_size


                # m_curl[i,j] = 0.5*(v_int[i,j - 1] - v_int[i,j + 1]  - u_int[i - 1,j] + u_int[i + 1,j])/m_cell_size

            # shear
            if (np.isnan(dy[i + 1, j]) == False and np.isnan(dy[i - 1, j]) == False and
                    np.isnan(dx[i, j - 1]) == False and np.isnan(dx[i, j + 1]) == False and \
                    np.isnan(dy[i, j - 1]) == False and np.isnan(dy[i, j + 1]) == False and \
                    np.isnan(dx[i + 1, j]) == False and np.isnan(dx[i - 1, j]) == False and \
                    (np.isnan(dx[i, j]) == False or np.isnan(dy[i, j]) == False)):
                dc_dc = 0.5 * (dy[i + 1, j] - dy[i - 1, j])
                dr_dr = 0.5 * (dx[i, j - 1] - dx[i, j + 1])
                dc_dr = 0.5 * (dy[i, j - 1] - dy[i, j + 1])
                dr_dc = 0.5 * (dx[i + 1, j] - dx[i - 1, j])
                m_shear[i, j] = np.sqrt(
                    (dc_dc - dr_dr) * (dc_dc - dr_dr) + (dc_dr - dr_dc) * (dc_dr - dr_dc)) / cell_size

                '''
                # Den
                dc_dc = 0.5*(v_int[i + 1,j] - v_int[i - 1,j])
                dr_dr = 0.5*(u_int[i,j + 1] - u_int[i,j - 1])
                dc_dr = 0.5*(v_int[i,j + 1] - v_int[i,j - 1])
                dr_dc = 0.5*(u_int[i + 1,j] - u_int[i - 1,j])

                m_shear[i,j] = np.sqrt((dc_dc -dr_dr) * (dc_dc -dr_dr) + (dc_dr - dr_dc) * (dc_dr - dr_dc))/m_cell_size
                '''

            # deformation
            if (np.isnan(m_shear[i, j]) == False and np.isnan(m_div[i, j]) == False):
                m_tdef[i, j] = np.hypot(m_shear[i, j], m_div[i, j])

    # Invert dy back
    dy = dy * (-1)

    # data = np.vstack((np.ravel(xx_int), np.ravel(yy_int), np.ravel(m_div), np.ravel(u_int), np.ravel(v_int))).T
    divergence = m_div

    # TODO: Test
    '''
    plt.imshow(divergence, cmap='RdBu', vmin=-0.06, vmax=0.06, interpolation='None')
    # Plot u and v values inside cells (for testing porposes)
    for ii in range(dx.shape[1]):
        for jj in range(dx.shape[0]):
            try:
                if not np.isnan(dx[ii,jj]):
                    plt.text(jj, ii,
                             'dx:%.2f\ndy:%.2f\n%.2f\nij:(%s,%s)' %
                             (dx[ii,jj], dy[ii,jj], divergence[ii,jj], ii, jj),
                             horizontalalignment='center',
                             verticalalignment='center', fontsize=0.05, color='k')
            except:
                pass

    plt.savefig('test.png', bbox_inches='tight', dpi=1000)
    '''

    curl = m_curl
    shear = m_shear
    total_deform = m_tdef
    return mag, divergence, curl, shear, total_deform

def fmt(x, pos):
    a, b = '{:.0e}'.format(x).split('e')
    b = int(b)
    return r'${} \times 10^{{{}}}$'.format(a, b)

def plot_div(fname):
    xxc = ff.variables['xc'][:]
    yyc = ff.variables['yc'][:]

    xc, yc = np.meshgrid(xxc, yyc)

    dX = ff.variables['dX'][0,:,:]
    dY = ff.variables['dY'][0,:,:]

    # Make nan where
    dX[dX==9.96921e+36] = np.nan
    dY[dY==9.96921e+36] = np.nan

    # Lats Lons
    lats = ff.variables['lat'][:]
    lons = ff.variables['lon'][:]

    # Calculate time delta and define norm time
    ddd1 = os.path.basename(fname).split('_')[5].split('-')[0]
    dt1_str = '%s-%s-%s %s:%s:%s' % (ddd1[0:4], ddd1[4:6], ddd1[6:8], ddd1[8:10], ddd1[10:12],
                                     ddd1[12:14])
    dt1 = datetime.datetime.strptime(dt1_str, '%Y-%m-%d %H:%M:%S')

    ddd2 = os.path.basename(fname).split('_')[5].split('-')[1]
    dt2_str = '%s-%s-%s %s:%s:%s' % (ddd2[0:4], ddd2[4:6], ddd2[6:8], ddd2[8:10], ddd2[10:12],
                                     ddd2[12:14])
    dt2 = datetime.datetime.strptime(dt2_str, '%Y-%m-%d %H:%M:%S')

    # Time for normalization
    time_delta = (dt2 - dt1).total_seconds()
    norm_times = 12 * 3600 # normilize to 12 hours
    normalization_time = time_delta / norm_times

    #print 'Date1: %s' % dt1
    #print 'Date2: %s' % dt2
    #print 'Time delta: %s, Normalization time: %s' % (time_delta, normalization_time)

    # Caclulation
    mag, divergence, curl, shear, total_deform = calc_deformations(dX, dY, normalization=True,
                                                                   normalization_time = normalization_time,
                                                                   cell_size=cell_size)

    # Calculate min and max
    minima = np.nanpercentile(divergence, 2)
    maxima = np.nanpercentile(divergence, 98)

    mmax = maxima

    if minima > mmax:
        mmax = minima
    mmin = -mmax

    # Divergence max and min
    #print 'Div min: %s   Div max: %s' % (np.nanmin(divergence), np.nanmax(divergence))
    #print 'Prec. Div min: %s   Div max: %s' % (mmin, mmax)
    #print ('Minima: %s   Maxima: %s') % (minima, maxima)

    plt.clf()

    # Make plot with vertical (default) colorbar
    fig, ax = plt.subplots()

    div_plot = plt.imshow(divergence, cmap='RdBu', vmin=-0.06, vmax=0.06, interpolation='None')
    ax.set_title(u'Скорость дивергенции дрейфа льда\n%s - %s' % (dt1_str, dt2_str))

    plt.quiver(range(dX.shape[1]), range(dX.shape[0]), dX, (-1) * dY, facecolor='#000000',
               angles='xy', scale_units='xy', alpha=.5)
               #, linewidth=0.01,
               #width=0.0012, headwidth=4.5, scale=0.2)
    '''
    # Testing
    # Plot u and v values inside cells (for testing porposes)
    for ii in range(dX.shape[1]):
        for jj in range(dX.shape[0]):
            try:
                if not np.isnan(dX[ii,jj]):
                    plt.text(jj, ii,
                             'dx:%.2f\ndy:%.2f\n%.2f' % (dX[ii,jj], dY[ii,jj], divergence[ii,jj]),
                             horizontalalignment='center',
                             verticalalignment='center', fontsize=0.1, color='k')
            except:
                pass
    '''


    #ticks=[mmin, 0, mmax]
    cbar = fig.colorbar(div_plot,  format=ticker.FuncFormatter(fmt))
    #cbar.ax.set_yticklabels(['< %.3f' % mmin, '0', '> %.3f' % mmax])
    cbar.ax.set_ylabel(u'1/12 часов', rotation=90)

    plt.savefig('div_plot/div_%s_%s.png' % (ddd1, ddd2), bbox_inches='tight', dpi=150)

    return divergence


########################
# CONSTANTS
########################

cell_size = 10.182338

########################
# END CONSTANTS
########################

#fname = '../drift/ice_drift_polstereo_sarbased_north_20170304020917-20170305020024_7156.nc'

ffiles = glob.glob('../drift/*201706*.nc')
num_of_files = len(ffiles) - 1

for i, fname in enumerate(ffiles):
    print '%s of %s' % (i, num_of_files)
    ff = Dataset(fname, 'r')
    # Lats Lons
    lats = ff.variables['lat'][:]
    lons = ff.variables['lon'][:]
    lats[lats==9.96921e+36] = np.nan
    lons[lons==9.96921e+36] = np.nan

    lon_min = np.nanmin(lons)
    lon_max = np.nanmax(lons)
    lat_min = np.nanmin(lats)
    lat_max = np.nanmax(lats)

    # If lats and lons inside bbox of Kara sea
    if lon_min >= 54. and lon_max <= 84. and lat_min >= 67.1 and lat_max <= 77.3:
        div = plot_div(fname)
        nn_kl = resample_kl_ice_grid('grid/aari_ice_model_grid/sar_div_%s' %
                                     os.path.basename(fname).split('_')[5], lons, lats, div)
        #plot_mag(fname)