# http://stackoverflow.com/questions/3684484/peak-detection-in-a-2d-array/3689710#3689710

"""

Takes an array and detects the troughs using the

local minimum filter.Returns a boolean mask of the troughs

(i.e. 1 when the pixel's value is the neighborhood

minimum, 0 otherwise)

Parameters:

--------------------

arr: numpy array where the land is masked with 0

detected_minima: numpy array

"""

# define an connected neighborhood

neighborhood = morphology.generate_binary_structure(2,2)

neighborhood = morphology.binary_dilation(neighborhood,iterations = 20)

# apply the local minimum filter; all locations of minimum value

# in their neighborhood are set to 1

# filter multiple times to get just one point per cyclone(A.O.)

tmp = filters.minimum_filter(arr, footprint=neighborhood)

tmp = filters.minimum_filter(tmp, footprint=neighborhood)

local_min = (filters.minimum_filter(tmp, footprint=neighborhood)==arr)

background = (arr==0)

#

# a little technicality: we must erode the background in order to

# successfully subtract it from local_min, otherwise a line will

# appear along the background border (artifact of the local minimum filter)

eroded_background = morphology.binary_erosion(

background, structure=neighborhood, border_value=1)

#

# we obtain the final mask, containing only peaks,

# by removing the background from the local_min mask

detected_minima = local_min - eroded_background

"""buffer_coast()

Returns an array that has been buffered to remove

ocean data close to coasts.

Parameters:

--------------------

pdata: numpy array of pressure data

buf: 2d buffer shape. (5,5) minimum for continuous coast outline

edgedif: expected difference between water and land values

mask: numpy array. 0 at coast, 1 away from coast

"""

pdata[np.isnan(pdata)] = 0

if len(mask.shape) <= 1:

print("buffer_coast: loading default mask")

mask = np.load('/glade/scratch/aordonez/landmask_stereo.npy')

elif len(mask.shape) > 2:

mask = mask[0,:,:]

if mask.shape != pdata.shape:

print("buffer_coast: mask must be same shape as data")

return

bi = morphology.generate_binary_structure(2,2)

mask = morphology.binary_dilation(mask,\

structure = bi, iterations = buf)

newmask = 1-mask.astype(int)

newmask = newmask.astype(float)

newmask[newmask < 1] = np.nan

pdata = pdata * newmask

"""

buffer_points(data,buf = (5,5))

returns an array of buffered data points

data: binary array, where 1 is the value to buffer

buf: the number of iterations for dilation

"""

bi = morphology.generate_binary_structure(2,2)

mask = morphology.binary_dilation(data,\

structure = bi, iterations = buf)

"""

find_cyclone_center

Returns a matrix (time x lon x lat). Cells with

a "1" indicate a low pressure center; cells equal "0"

otherwise.

For a pixel to be counted as a low, it must meet these criteria:

There must be a local minima in the sea level pressure (SLP),

a local maxima in the laplacian of SLP, greater than 15% ice

cover, and the SLP must be between the bounds 'pmin' and 'pmax'.

Grid cells near the coast are not included due to noise from

the stereo regridding and rotation done in get_boxes()

Parameters:

--------------------

psl: numpy array of sea level pressure. land areas masked with 0

icefrac: numpy array of sea ice concentration on atmosphere grid,

max = 1

pmax: numeric, maximum allowed value of central pressure

pmin: numeric, minimum allowed value of central pressure

lows: numpy array

"""

time,rows,cols = psl.shape

lows = np.zeros((psl.shape))

landmask = np.load('/glade/scratch/aordonez/landmask_stereo.npy')

# find the lows

for n in range(0,psl.shape[0]):

# smooth the pressure field?

psl1 = filters.gaussian_filter(psl[n,:,:],3)

#psl1 = buffer_coast(psl1,buf=1,mask=landmask)

#psl1 = psl[n,:,:]

lap = filters.laplace(psl1)

#lap = laplacian[n,:,:]

#lap = filters.gaussian_filter(lap,3)

#lap = buffer_coast(lap,buf=1,mask=landmask)

lapmax = detect_local_minima(lap*-1.)

# include cells immediately surrounding maxes

lapmax = buffer_points(lapmax,buf=1)

#ptmp = buffer_coast(psl[n,:,:],buf = 1, mask = landmask)

low_n = detect_local_minima(psl1)

lows[n,:,:] = np.select([(low_n == True) &

(icefrac[n,:,:] > 0.15) &

(psl[n,:, :] <= pmax) &

(psl[n,:,:] >= pmin) &

(lapmax ==1)],[low_n])

"""

find_cyclone_center

Returns a matrix (time x lon x lat). Cells with

a "1" indicate a low pressure center; cells equal "0"

otherwise.

For a pixel to be counted as a high, it must meet these criteria:

There must be a local maxima in the sea level pressure (SLP),

a local minima in the laplacian of SLP, greater than 15% ice

cover, and the SLP must be between the bounds 'pmin' and 'pmax'.

Grid cells near the coast are not included due to noise from

the stereo regridding and rotation done in get_boxes()

Parameters:

--------------------

psl: numpy array of sea level pressure. land areas masked with 0

icefrac: numpy array of sea ice concentration on atmosphere grid,

max = 1

pmax: numeric, maximum allowed value of central pressure

pmin: numeric, minimum allowed value of central pressure

lows: numpy array

"""

time,rows,cols = psl.shape

lows = np.zeros((psl.shape))

landmask = np.load('/glade/scratch/aordonez/landmask_stereo.npy')

# find the lows

for n in range(0,psl.shape[0]):

lap = filters.laplace(psl[n,:,:])

lapmax = detect_local_minima(lap)

# include cells immediately surrounding mins

lapmax = buffer_points(lapmax,buf=2)

ptmp = buffer_coast(psl[n,:,:],buf = 1, mask = landmask)

low_n = detect_local_minima(ptmp * -1.)

lows[n,:,:] = np.select([(low_n == True) &

(icefrac[n,:,:] > 0.15) &

#(lat > 66) &

(psl[n,:, :] <= pmax) &

(psl[n,:,:] >= pmin) &

(lapmax ==1) &

(coast == 0)],[low_n])

"""ID_parked_lows

Returns the locations of low pressure systems

which stay in one location over the coarse of many days.

Parameters:

-------------------

lows: a binary array where '1' is the location of a low

lat: an array of latitudes for the lows grid

lon: an array of longitudes for the lows grid

"""

"""

for t in times:

buffer around each low

find buffers that overlap

delete these for now; just include moving storms

"""

lows_copy = np.copy(lows)

time = lows.shape[0]

for day in range(1,time):

bi = morphology.generate_binary_structure(2,2)

buffered_lows_new = morphology.binary_dilation(lows[day,:,:],structure = bi,iterations = 1).astype(int)

buffered_lows_old = morphology.binary_dilation(lows[day-1,:,:],structure = bi,iterations = 1).astype(int)

storms, _ = ndimage.label(buffered_lows_new)

# id which storms are overlapping

lowsum = buffered_lows_new + buffered_lows_old

k = np.where(lowsum == 2)

label_list = np.unique(storms[k])

# delete those storms

l = np.select([lows[day,:,:] == 1],[storms])

lnew = np.in1d(l,label_list).astype(int)

lnew = np.reshape(lnew,(lows.shape[1],lows.shape[2]))

l[lnew == 1] = 0

l[l > 0] = 1

lows_copy[day,:,:] = l

"""

find_cyclone_center

Returns a matrix (time x lon x lat). Cells with

a "1" indicate a low pressure center; cells equal "0"

otherwise.

For a pixel to be counted as a low, it must meet these criteria:

There must be a local minima in the sea level pressure (SLP),

a local maxima in the laplacian of SLP, greater than 15% ice

cover, and the SLP must be between the bounds 'pmin' and 'pmax'.

Grid cells near the coast are not included due to noise from

the stereo regridding and rotation done in get_boxes()

Parameters:

--------------------

psl: numpy array of sea level pressure. land areas masked with 0

icefrac: numpy array of sea ice concentration on atmosphere grid,

max = 1

pmax: numeric, maximum allowed value of central pressure

pmin: numeric, minimum allowed value of central pressure

lows: numpy array

"""

time,rows,cols = psl.shape

lows = np.zeros((psl.shape))

# find the lows

for n in range(0,psl.shape[0]):

lap = filters.laplace(psl[n,:,:])

lapmax = detect_local_minima(lap*-1.)

ptmp = buffer_coast(psl[n,:,:],buf = (5,5))

low_n = detect_local_minima(ptmp)

lows[n,:,:] = np.select([(low_n == True) &

(icefrac[n,:,:] > 0.15) &

#(icefrac[n,:,:] < 0.70) &

(lat < -55) &

(psl[n,:, :] <= pmax) &

(psl[n,:,:] >= pmin) &

(lapmax ==1) &

(coast == 0)],[low_n])

"""

box = get_boxes(lows, data, size)

Clips a square of length(2 x size) + 1 around each low

pressure center in lows and returns an array with all the

boxes.

Parameters:

--------------------

lows: binary matrix where 1 = low pressure center

data: numpy array, land masked with 0

size: numeric, half the length of the 2D subset box

edgedif: numeric, roughly the difference

in value between data and land grid cells

box: numpy array of data around low pressure centers

"""

lon[lon < 0.] = lon[lon < 0.] + 360.

long_size = ((size *2) + 1)

mylow = np.where(lows == 1)

nlows = mylow[0].shape[0]

data_box = np.zeros((nlows,long_size,long_size))

lat_box = np.zeros(data_box.shape)

lon_box = np.zeros(data_box.shape)

(tmax, ymax, xmax) = data.shape

if len(landmask.shape) == 3:

landmask = landmask[0,:,:]

# get lon where north is up

lon0 = lon[0,(int(xmax/2))-1]

count = 0

indlist = np.zeros((nlows))

for ind in range(0,nlows):

time = mylow[0][ind]

lowrow = mylow[1][ind]

lowcol = mylow[2][ind]

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

# rotation to north

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

mylon = lon[lowrow,lowcol]

low_mask = np.zeros((ymax,xmax))

low_mask[lowrow,lowcol] = 1

if lon0 < mylon:

deg = mylon - lon0

elif lon0 >= mylon:

deg = (360 + mylon) - lon0

low_rotated = interpolation.rotate(low_mask, deg, order = 2)

# because of interpolation, lows != 1

ynew,xnew = np.where(low_rotated == low_rotated.max())

if len(ynew.shape) > 1:

print("get_boxes: problem with rotation: too many indices for max")

print("get_boxes: exiting script")

#return

data_rotated = interpolation.rotate(data[time,:,:], deg, order =2)

# try to ignore data outside map

data_rotated[data_rotated == 0.0] = np.nan

# take out noisy grid cells near coast

landmask_rot = interpolation.rotate(landmask, deg, order = 2)

landmask_rot[landmask_rot < 0.5] = 0

landmask_rot[landmask_rot >= 0.5] = 1

data_rotated = buffer_coast(data_rotated, buf = 1, mask = landmask_rot)

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

# extracting box

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

y1 = int(ynew - size)

y2 = int(ynew + size + 1)

x1 = int(xnew - size)

x2 = int(xnew + size + 1)

if (y1 < 0) | (x1 < 0) | (y2 > ymax) | (x2 > xmax):

# too close to edge of map

continue

else:

data_box[count,:,:] = data_rotated[y1:y2,x1:x2]

#data_box[count,:,:] = data[ind,y1:y2,x1:x2]

#lat_box[count,:,:] = lat[y1:y2,x1:x2]

#lon_box[count,:,:] = lon[y1:y2,x1:x2]

indlist[count] = ind

count += 1

"""

box = get_boxes_no_rotation(lows, data, size)

Clips a square of length(2 x size) + 1 around each low

pressure center in lows and returns an array with all the

boxes. Like get_boxes, but does not rotate the box

relative to north.

Parameters:

--------------------

lows: binary matrix where 1 = low pressure center

data: numpy array, land masked with 0

size: numeric, half the length of the 2D subset box

edgedif: numeric, roughly the difference

in value between data and land grid cells

box: numpy array of data around low pressure centers

"""

lon[lon < 0.] = lon[lon < 0.] + 360.

long_size = ((size *2) + 1)

mylow = np.where(lows == 1)

nlows = mylow[0].shape[0]

data_box = np.zeros((nlows,long_size,long_size))

lat_box = np.zeros(data_box.shape)

lon_box = np.zeros(data_box.shape)

(tmax, ymax, xmax) = data.shape

# get lon where north is up

lon0 = lon[0,(xmax/2)-1]

count = 0

indlist = np.zeros((nlows))

for ind in range(0,nlows):

time = mylow[0][ind]

lowrow = mylow[1][ind]

lowcol = mylow[2][ind]

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

# buffer out coast

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

# take out noisy grid cells near coast

coast = buffer_coast(data[time,:,:], buf = 1, edgedif = edgedif)

data_buffered = data[time,:,:] * coast

#ynew,xnew = lowrow,lowcol

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

# extracting box

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

y1 = lowrow - size

y2 = lowrow + size + 1

x1 = lowcol - size

x2 = lowcol + size + 1

if (y1 < 0) | (x1 < 0) | (y2 > ymax) | (x2 > xmax):

# too close to edge of map

continue

else:

data_box[count,:,:] = data_buffered[y1:y2,x1:x2]

#data_box[count,:,:] = data[ind,y1:y2,x1:x2]

lat_box[count,:,:] = lat[y1:y2,x1:x2]

lon_box[count,:,:] = lon[y1:y2,x1:x2]

indlist[count] = ind

count += 1

# regrid to conformal conic

# equations from https://en.wikipedia.org/wiki/Lambert_conformal_conic_projection

row,col = lat.shape

lat = lat * (np.pi / 180.)

lon = lon * (np.pi / 180.)

lon_ref = lon_ref * (np.pi / 180.)

lat_ref = lat_ref * (np.pi / 180.)

lat_stnd1 = np.complex(lat_stnd1 * (np.pi / 180.))

lat_stnd2 = np.complex(lat_stnd2 * (np.pi / 180.))

#lon_ref = lon[0,col/2]

#lat_ref = lat[row/2,col/2]

#lat_stnd2 = np.complex(lat[row/2,col/2] - 0.01)

#lat_stnd1 = np.complex(lat[row/2,col/2] + 0.01)

n_top = np.log(np.cos(lat_stnd1) * 1./np.cos(lat_stnd2))

n_bottom = np.log(np.tan(0.25 * np.pi + 0.5 * lat_stnd2) *

1./np.tan(0.25 * np.pi + 0.5 * lat_stnd1))

n = n_top / n_bottom

F = (np.cos(lat_stnd1) * np.power(np.tan(0.25 * np.pi + 0.5 * lat_stnd1),n)) / n

rho = F * 1./np.tan(0.25 * np.pi + 0.5 * lat)**n

rho_0 = F * 1./np.tan(0.25 * np.pi + 0.5 * lat_ref)**n

x = rho * np.sin(n * (lon - lon_ref))

y = rho_0 - rho * np.cos(n * (lon - lon_ref))

"""like get_boxes, but clips from a

small, conic-projected area

Since 3-d data is all on planes, cannot use griddata

on 3-d array all at once

lows: binary array of low pressure center locations

data: dictionary of 3-d data. One of the entries must be 'psl'

lat: array of latitude data

lon: array of longitude data

types: char dictionary indicating if data grid is 'atm' or 'ice'

dataregrid: dictionary containing clipped, aligned data for compositing

"""

mylow = np.where(lows == 1)

nlows = mylow[0].shape[0]

indlist = np.zeros((nlows,1))

# area of interest is small region at center of regrid:

xnew,ynew = np.meshgrid(np.arange(-0.2,0.205,0.005),np.arange(-0.2,0.205,0.005))

dataregrid = {}

for item in data.keys():

dataregrid[item] = np.zeros((nlows,xnew.shape[0],xnew.shape[1]))

count = 0

for ind in range(0,nlows):

if ind % 10 == 0:

print(ind)

time = mylow[0][ind]

lowrow = mylow[1][ind]

lowcol = mylow[2][ind]

lattest = latlist['psl'][lowrow,lowcol]

lontest = lonlist['psl'][lowrow,lowcol]

if lattest < -80:

continue

x,y = regrid_to_conic(latlist['psl'],lonlist['psl'],lattest,lontest,lattest+5,lattest - 5)

xi,yi = regrid_to_conic(latlist['ice'],lonlist['ice'],lattest,lontest,lattest+5,lattest - 5)

# first, get pressure info and do extra quality control:

d = data['psl'][time,:,:]

k = np.where(np.isnan(x) == False)

ki = np.where(np.isnan(xi) == False)

s = interpolate.griddata((x[k],y[k]),d[k],(xnew,ynew),method = 'linear')

# draw box around center and compare mean at low with mean around low

sslice = s[s.shape[0]/2-10:s.shape[0]/2+10,s.shape[1]/2-10:s.shape[1]/2+10]

if s[s.shape[0]/2,s.shape[1]/2] < np.nanmean(sslice):

dataregrid['psl'][count,:,:] = s

indlist[count] = ind

# low ok, get the rest of the variables:

for item in data.keys():

if item != 'psl':

d = data[item][time,:,:]

if types[item] == 'atm':

s = interpolate.griddata((x[k],y[k]),

d[k],(xnew,ynew),method = 'linear')

elif types[item] == 'ice':

s = interpolate.griddata((xi[ki],yi[ki]),

d[ki],(xnew,ynew),

method = 'linear')

dataregrid[item][count,:,:] = s

count += 1

# since we eliminated some of the lows, trim data to new low count:

for item in data.keys():

dataregrid[item] = dataregrid[item][0:count,:,:]

"""plot_lows_on_map

Function for quickly assessing the find_cyclone_center

results. Makes a plot of sea level pressure with

identified low pressure centers marked

Parameters:

--------------------

lows: 3D numpy array

psl: 3D numpy array

time (optional): numeric

"""

lowsmap = lows[time,:,:]

pslmap = psl[time,:,:]

k = np.where(lowsmap == 1)

f,ax = plt.subplots(1,1)

ax.pcolormesh(pslmap, vmin = 90000, vmax = 104000)

ax.scatter(k[1], k[0],color = 'k')

boxplot = box[time,:,:]

f,axs = plt.subplots(1,1)

h = axs.pcolormesh(boxplot,vmin = 90000, vmax = 104000)

f.colorbar(h,ax = axs)

f,axs = plt.subplots(1,1)

h = axs.pcolormesh(np.nanmean(data,axis = 0),vmin = cmin, vmax = cmax)

f.colorbar(h,ax = axs)

n = len(wgts)

half = n

datama = np.zeros(data.shape)

for ind in range(half,data.shape[0]-half):

n0 = ind - half

nf = ind

datama[ind,:,:] = np.average(data[n0:nf,:,:],axis = 0,weights = wgts)

datama = data- datama

datama[0:half,:,:] = 0

datama[(datama.shape[0]-half):,:,:] = 0

"""Finds the trend over the previous ndays

number of days at each gridcell

"""

s = data.shape

data = np.reshape(data,(s[0],s[1]*s[2]))

datatrend = np.zeros(data.shape)

if b:

datab = np.zeros(data.shape)

for ind in range(ndays,data.shape[0]-ndays):

n0 = ind - ndays

nf = ind

tmp = np.polyfit(range(0,ndays),data[n0:nf,:],1)

datatrend[ind,:] = tmp[0,:]

if b:

datab[ind,:] = tmp[1,:]

datatrend = np.reshape(datatrend,s)

datatrend[0:ndays,:,:] = 0

if b:

return datatrend, np.reshape(datab,s)

else:

"""Uses the linear trend over the past

ndays number of days to predict what

the value of data is on a given day

day: day at which prediction is desired

"""

m,b = get_nday_trend(data,ndays,b = True)

predict = np.zeros(data.shape)

predict = m * (ndays+day) + b

"""Pulls out a timeseries only containing days in

March, April, and May. Input timeseries may not contain

partial years of data.

"""

nyrs = data.shape[0] / 365

mam = len(range(59,151))

data_mam = np.zeros((nyrs*mam,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_mam[0:mam,:,:] = data[59:151,:,:]

last_ind = mam

else:

data_mam[last_ind:last_ind + mam,:,:] = data[59+(yr*365):151+(yr*365),:]

last_ind = last_ind + mam

"""Pulls out a timeseries only containing days in

June, July, and August

"""

nyrs = int(data.shape[0] / 365)

jja = len(range(151,243))

data_jja = np.zeros((nyrs*jja,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_jja[0:jja,:,:] = data[151:243,:,:]

last_ind = jja

else:

data_jja[last_ind:last_ind + jja,:,:] = data[151+(yr*365):243+(yr*365),:,:]

last_ind = last_ind + jja

"""Pulls out a timeseries only containing days in

June, July

"""

nyrs = int(data.shape[0] / 365 )

jja = len(range(151,212))

data_jja = np.zeros((nyrs*jja,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_jja[0:jja,:,:] = data[151:212,:,:]

last_ind = jja

else:

data_jja[last_ind:last_ind + jja,:,:] = data[151+(yr*365):212+(yr*365),:,:]

last_ind = last_ind + jja

"""Pulls out a timeseries only containing days in

June, July, and August

"""

nyrs = data.shape[0] / 365

jja = 30

data_jja = np.zeros((nyrs*jja,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_jja[0:jja,:,:] = data[151:181,:,:]

last_ind = jja

else:

data_jja[last_ind:last_ind + jja,:,:] = data[151+(yr*365):181+(yr*365),:,:]

last_ind = last_ind + jja

"""Pulls out a timeseries only containing days in

June, July, and August

"""

nyrs = data.shape[0] / 365

jja = 31

data_jja = np.zeros((nyrs*jja,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_jja[0:jja,:,:] = data[212:243,:,:]

last_ind = jja

else:

data_jja[last_ind:last_ind + jja,:,:] = data[212+(yr*365):243+(yr*365),:,:]

last_ind = last_ind + jja

"""Pulls out a timeseries only containing days in

June, July, and August

"""

nyrs = data.shape[0] / 365

ndays = 30

data_sep = np.zeros((nyrs*ndays,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_sep[0:ndays,:,:] = data[243:243+ndays,:,:]

last_ind = ndays

else:

data_sep[last_ind:last_ind + ndays,:,:] = data[243+(yr*365):243+ndays+(yr*365),:,:]

last_ind = last_ind + ndays

"""Pulls out a timeseries only containing days in

August, September, and October

"""

nyrs = data.shape[0] / 365

son = len(range(212,304))

data_son = np.zeros((nyrs*son,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_son[0:son,:,:] = data[212:304,:,:]

last_ind = son

else:

data_son[last_ind:last_ind + son,:,:] = data[212+(yr*365):304+(yr*365),:,:]

last_ind = last_ind + son

"""Pulls out a timeseries only containing days in

September, October, and November

"""

nyrs = data.shape[0] / 365

son = len(range(243,334))

data_son = np.zeros((nyrs*son,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_son[0:son,:,:] = data[243:334,:,:]

last_ind = son

else:

data_son[last_ind:last_ind + son,:,:] = data[243+(yr*365):334+(yr*365),:,:]

last_ind = last_ind + son

"""Pulls out a timeseries only containing days in

December, January, and February

"""

nyrs = int(data.shape[0] / 365)

jf = len(range(0,59))

d = len(range(334,365))

data_djf = np.zeros((nyrs*(d+jf),data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_djf[0:59,:,:] = data[0:59,:,:]

data_djf[59:59+d,:,:] =data[334:365,:,:]

last_ind = 59+d

else:

data_djf[last_ind:last_ind + jf,:,:] = data[0+(yr*365):59+(yr*365),:]

data_djf[last_ind + jf:last_ind + d + jf,:,:] = data[334+(yr*365):365+(yr*365),:]

last_ind = last_ind + d + jf

"""Pulls out a timeseries only containing days in

January, and February

"""

nyrs = int(data.shape[0] / 365 )

jf = len(range(0,59))

data_djf = np.zeros((nyrs*(jf),data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_djf[0:59,:,:] = data[0:59,:,:]

last_ind = jf

else:

data_djf[last_ind:last_ind + jf,:,:] = data[0+(yr*365):59+(yr*365),:]

last_ind = last_ind + jf

nyrs = data.shape[0] / 365

month_len = [31,28,31,30,31,30,31,31,30,31,30,31]

mlen = month_len[month]

start_day = [0,31,59,90,120,151,181,212,243,273,304,334]

start = start_day[month]

data_month = np.zeros((nyrs*mlen,data.shape[1],data.shape[2]))

for yr in range(0,nyrs):

if yr == 0:

data_month[0:mlen,:,:] = data[start:start+mlen,:,:]

last_ind = mlen

else:

span = yr*365

data_month[last_ind:last_ind+mlen,:,:] = data[start+span:start+mlen+span,:,:]

last_ind = last_ind + mlen

"""Calculates the daily anomalies relative to values

over the ice pack where concentration >= 15%

"""

if np.max(icedata) > 1:

icedata[icedata < 15] = 0.

icedata[icedata > 0] = 1.

else:

icedata[icedata < 0.15] = 0.

icedata[icedata > 0] = 1.

tarea = np.tile(tarea,(icedata.shape[0],1,1))

# compute area weights over sea ice

wgt_total = np.nansum(np.nansum((tarea * icedata),axis =2), axis = 1)

wgt_total = np.tile(wgt_total,(vardata.shape[1],vardata.shape[2],1))

wgt_total = np.transpose(wgt_total,(2,0,1))

wgt_total[icedata != 1] = 0.

wgt = tarea * icedata / wgt_total

# get time series of mean values

varmean = np.nansum(np.nansum(wgt * vardata, axis = 2),axis = 1)

varmean_overall = np.nanmean(varmean,axis = 0)

#print varmean[0]

varmean = np.tile(varmean, (vardata.shape[1],vardata.shape[2],1))

varmean = np.transpose(varmean,(2,0,1))

vardata = vardata - varmean

titles,sname,Pmean,Tmean,Umean,Vmean):

"""Makes a composite plot of the input data field with

temperature and pressure contours and wind streamlines.

Saves the plot.

"""

# figure out rows and columns

if len(varlist) < 3:

row_n = 1

col_n = len(varlist)

else:

row_n = 2

col_n = int(np.ceil(len(varlist)/2.))

# create plot

f,axs = plt.subplots(nrows = row_n, ncols = col_n, figsize=(row_n * 3., col_n * 2.5))

axs = np.reshape(axs,(row_n*col_n,1))

Tlevels = np.arange(200,310,5)

for ind in range(0,len(varlist)):

cmax = np.max([abs(np.nanmax(varlist[ind])),abs(np.nanmin(varlist[ind]))])

ax1 = axs[ind,0].pcolormesh(X1,Y1,varlist[ind],cmap='PuOr',norm=colors.SymLogNorm(linthresh=0.1,linscale=0.1,vmin=-1*cmax,vmax=cmax))

c1 = axs[ind,0].contour(X5,Y5,Pmean,colors='k',linewidths = 1)

c2 = axs[ind,0].contour(X5,Y5,Tmean,Tlevels,colors='r',linewidths = 1)

#axs[ind,0].streamplot(X5,Y5,Umean,Vmean,linewidth = 1)

plot_spaced_quivers(axs[ind,0],X5,Y5,Umean,Vmean,spacing = 5);

axs[ind,0].clabel(c1, inline=1, fontsize=8)

axs[ind,0].clabel(c2, inline=1, fontsize=8)

axs[ind,0].set_title(titles[ind], fontsize = 10)

cb1 = f.colorbar(ax1, ax = axs[ind,0], ticks = [-1.0*cmax,-0.5*cmax,-0.2*cmax,0.,0.2*cmax,0.5*cmax,1.0*cmax])

cb1.ax.tick_params(labelsize = 8)

#plt.setp(axs[ind,0].get_xticklabels(), fontsize = 8)

#plt.setp(axs[ind,0].get_yticklabels(), fontsize = 8)

axs[ind,0].get_xaxis().set_ticks([])

axs[ind,0].get_yaxis().set_ticks([])

savename = "cyclone_composite_" + sname + "_" + res + "_c_" + storm_set + "_" + season + ".png"