currents2grib

# -*- coding: utf-8 -*-
"""
Created on Wed Sep 07 18:46:15 2016

@author: Компьютер
"""
##################################
#
#  !!!!!MAKE IT FOR EACH LEVEL SEPARETELY!!!!
#
###################################

from netCDF4 import Dataset
import numpy as np
import pyresample
import glob
import os
#import time
from pylab import *

def get_data_from_txt(fname):
    ff =open(fname, 'r')
    lines = ff.readlines()
    lats = []
    lons = []
    uu = []
    vv = []
    for line in lines:
        try:
            lat, lon, u, v = line.split()[0], line.split()[1], line.split()[2], \
                             line.split()[3]
            lats.append(float(lat))
            lons.append(float(lon))
            uu.append(float(u))
            vv.append(float(v))
        except:
            pass
    ff.close()
    return lats, lons, uu, vv

def interpolate_data(uu_grid, vv_grid):
    np.set_printoptions(formatter={'float': '{: .3f}'.format})
    # Initial model grid
    orig_def = pyresample.geometry.SwathDefinition(lons=ilon, lats=ilat)
    # Target regular grid
    targ_def = pyresample.geometry.SwathDefinition(lons=lons_out, lats=lats_out)

    # Interpolate U
    uu_interpolated = pyresample.kd_tree.resample_nearest(orig_def, uu_grid, \
                           targ_def, radius_of_influence=16000,\
                            fill_value=0.) #, sigmas=250000,)
    # Interpolate V
    vv_interpolated = pyresample.kd_tree.resample_nearest(orig_def, vv_grid, \
                           targ_def, radius_of_influence=16000, \
                            fill_value=0.) #, sigmas=250000,)
    return uu_interpolated, vv_interpolated

basedir = 'currents'

# Open initial data
# Get date and lat lon
ffiles = glob.glob('currents\\*.dat')
idate = os.path.basename(ffiles[0]).split('_')[1].split('.')[0]
lats, lons, uu, vv = get_data_from_txt(ffiles[0])
lats, lons, uu, vv = np.array(lats), np.array(lons), np.array(uu), np.array(vv)

# Reshape lon, lat, U, V
ilon, ilat       = lons[:].reshape(9,113), lats[:].reshape(9,113)

# Make regular lon lat grid
ilats_out = np.arange(min(lats),max(lats),0.125,dtype='float16')
ilons_out = np.arange(min(lons),max(lons),0.125,dtype='float16')
lons_out, lats_out = np.meshgrid(ilons_out, ilats_out)

# the output array to write will be nlats x nlons
nlats = len(ilats_out)
nlons = len(ilons_out)

# Create NETCDF4
#dataset = Dataset('AARI_SEA_CURRENTS_%s.nc' % idate,'w',format='NETCDF4_CLASSIC')
dataset = Dataset('AARI_SEA_CURRENTS_%s.nc' % idate,'w',format='NETCDF3_CLASSIC')
latitudes =   dataset.createDimension('latitude', nlats)
longitudes =   dataset.createDimension('longitude', nlons)
#level = dataset.createDimension('level', 4)
time =  dataset.createDimension('time', None)

# Create coordinate variables for 4-dimensions
latitudes =  dataset.createVariable('latitude', 'f4', ('latitude',))
longitudes = dataset.createVariable('longitude','f4', ('longitude',))
#levels =     dataset.createVariable('level', np.int32, ('level',))
times =      dataset.createVariable('time', np.float64, ('time',))

# Create the actual 4-d variable
u = dataset.createVariable('u', 'f4',('time','latitude','longitude'))#,zlib=True)#,least_significant_digit=7)
v = dataset.createVariable('v', 'f4',('time','latitude','longitude'))#,zlib=True)#,least_significant_digit=7)

# Global Attributes
dataset.description = 'Sea currents from AARI-IOCM model'
import time
dataset.history = 'Created ' + time.ctime(time.time())
dataset.source = 'AARI-IOCM'

# Variable Attributes
latitudes.units = 'degrees_north'
longitudes.units = 'degrees_east'

#levels.units = 'meters'
u.units = 'm/s'
v.units = 'm/s'

times.units = 'hours since 0001-01-01 00:00:00'
times.calendar = 'gregorian'

latitudes[:] = ilats_out
longitudes[:] = ilons_out

nlats = len(dataset.dimensions['latitude'])
nlons = len(dataset.dimensions['longitude'])

# Levels
ll = []
for ifile in ffiles:
    lll = os.path.basename(ifile)[1:3]
    ll.append(lll)

uniq_ll = set(ll)
uu_ll   = list(uniq_ll)
m       = np.array(uu_ll)
m.sort()

# Fill LEVEL values
#levels[:] =  m

# Get number of Times
t_num = len(glob.glob('%s\\w%02d*' % (basedir,int(m[0]))))

# Go through levels
#level_ch = 0
#for ilevel in levels:
#    print '\nLEVEL: %s\n' % ilevel
ilevel = m[0]

ffiles = glob.glob('%s\\w%02d*' % (basedir,int(ilevel)))
#print ffiles
time_ch = 0


for ifile in ffiles:
    #print 'level_ch: %s' % level_ch
    print 'time_ch: %s' % time_ch

    print ifile

    lats, lons, uu, vv = get_data_from_txt(ifile)
    lats, lons, uu, vv = np.array(lats), np.array(lons), np.array(uu), np.array(vv)

    uu_grid, vv_grid = uu.reshape(9,113), vv.reshape(9,113)
    uu_interpolated, vv_interpolated = interpolate_data(uu_grid, vv_grid)
    #plt.clf()
    #imshow(uu_interpolted)
    #plt.show()

    #print uu_interpolated[10:20,4:50] #, vv_interpolated
    #print np.min(uu_interpolated[~np.isnan(uu_interpolated)])
    #print np.min(vv_interpolated[~np.isnan(vv_interpolated)])

    u[time_ch,:] = uu_interpolated
    v[time_ch,:] = vv_interpolated

    time_ch = time_ch + 1
    #level_ch = level_ch + 1


# Fill in times.
from datetime import datetime, timedelta
from netCDF4 import num2date, date2num
dates = []

for n in range(u.shape[0]):
    dates.append(datetime(2016, 4, 4) + n * timedelta(hours=6))

times[:] = date2num(dates, units = times.units, calendar = times.calendar)
dates = num2date(times[:], units=times.units, calendar=times.calendar)

dataset.close()