import xray
import numpy as np
import matplotlib.pyplot as plt
import atmos as atm
from atmos import split_timedim

filename = '/home/jennifer/datastore/atmos-tools/uwnd.mon.mean.nc'
ds = atm.ncload(filename)
lat = ds['lat'].values.astype(np.float64)
lon = ds['lon'].values.astype(np.float64)
u = ds['uwnd']
# Remove Jan-Jun of 2015
tlast = -7
u = u[:tlast]

unew = split_timedim(u.values, 12)
unew2 = split_timedim(u.values, 12, slowfast=False)

# Check that reshaping worked properly
m, y, k = 11, 20, 10
data1 = u[y*12 + m, k].values
data2 = unew[y, m, k]
data3 = unew2[m, y, k]
print(np.array_equal(data1,data2))
print(np.array_equal(data2, data3))

# Plot data to check
xi, yi = np.meshgrid(lon, lat)
plt.figure()
plt.subplot(211)
plt.pcolormesh(xi, yi, data1, cmap='jet')
plt.colorbar()
plt.subplot(212)
plt.pcolormesh(xi,yi, data2, cmap='jet')
plt.colorbar()

# ----------------------------------------------------------------------
# DataArray version
years = np.arange(1979, 2015)
months = np.arange(1, 13)
unew = split_timedim(u, 12, time0_name='year', time0_vals=years,
                     time1_name='month', time1_vals=months)
unew2 = split_timedim(u, 12, slowfast=False)

# Check that reshaping worked properly
m, y, k = 11, 20, 10
data1 = u[y*12 + m, k]
data2 = unew[y, m, k]
data3 = unew2[m, y, k]
print(np.array_equal(data1,data2))
print(np.array_equal(data2, data3))

# Plot data to check
plt.figure()
plt.subplot(211)
atm.pcolor_latlon(data1, cmap='jet')
plt.subplot(212)
atm.pcolor_latlon(data2, cmap='jet')

@@ -113,45 +113,201 @@
 im%0A%0A
-filename = '/home/jennifer/datastore/
+datadir = '/home/jwalker/eady/datastore/'%0A#datadir = '/home/jennifer/datastore/'%0A%0A# ----------------------------------------------------------------------%0A# NCEP2 Monthly%0Afilename = datadir + '
 atmo
@@ -362,106 +362,83 @@
 me)%0A
-lat
+u
  = ds%5B'
-lat'%5D.values.astype(np.float64)%0Alon = ds%5B'lon'%5D.values.astype(np.float64)%0Au = ds%5B'uwnd'%5D
+uwnd'%5D%0Alat = atm.get_coord(u, 'lat')%0Alon = atm.get_coord(u, 'lon')%0A
 %0A# R
@@ -1639,24 +1639,868 @@
 tlon(data2, cmap='jet')%0A
+%0A# ----------------------------------------------------------------------%0A# MERRA Daily%0Afilename = datadir + 'merra/daily/merra_u200_198601.nc'%0Ads = atm.ncload(filename)%0Au = ds%5B'U'%5D%0Alat = atm.get_coord(u, 'lat')%0Alon = atm.get_coord(u, 'lon')%0A%0A# Number of time points per day%0An = 8%0Adays = np.arange(1, 32)%0Ahrs = np.arange(0, 24, 3)%0Aunew = split_timedim(u, n, time0_name='days', time0_vals=days,%0A                     time1_name='hours', time1_vals=hrs)%0Aunew2 = split_timedim(u, n, slowfast=False)%0A%0A# Check that reshaping worked properly%0Ady, hr, k = 1, 3, 0%0Adata1 = u%5Bdy*n + hr, k%5D%0Adata2 = unew%5Bdy, hr, k%5D%0Adata3 = unew2%5Bhr, dy, k%5D%0Aprint(np.array_equal(data1,data2))%0Aprint(np.array_equal(data2, data3))%0A%0A# Plot data to check%0Aplt.figure()%0Aplt.subplot(211)%0Aatm.pcolor_latlon(data1, cmap='jet')%0Aplt.subplot(212)%0Aatm.pcolor_latlon(data2, cmap='jet')%0A%0A%0A

Test with MERRA daily and NCEP2 monthly