returnall=False):

# Remove All NaN Points

ntime,nlat,nlon = sla.shape

slars = sla.reshape(ntime,nlat*nlon)

okdata,knan,okpts = proc.find_nan(slars,0)

npts = okdata.shape[1]

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

# Calculate Correlation and Covariance Matrices

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

srho = np.corrcoef(okdata.T,okdata.T)

scov = np.cov(okdata.T,okdata.T)

srho = srho[:npts,:npts]

scov = scov[:npts,:npts]

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

# Calculate Distance Matrix

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

lonmesh,latmesh = np.meshgrid(lon,lat)

coords = np.vstack([lonmesh.flatten(),latmesh.flatten()]).T

coords = coords[okpts,:]

coords1 = coords.copy()

coords2 = np.zeros(coords1.shape)

coords2[:,0] = np.radians(coords1[:,1]) # First point is latitude

coords2[:,1] = np.radians(coords1[:,0]) # Second Point is Longitude

sdist = haversine_distances(coords2,coords2) * 6371

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

# Combine the Matrices

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

a_fac = np.sqrt(-distthres/(2*np.log(0.5))) # Calcuate so exp=0.5 when distance is 3000km

expterm = np.exp(-sdist/(2*a_fac**2))

distance_matrix = 1-expterm*srho

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

# Do Clustering (scipy)

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

cdist = squareform(distance_matrix,checks=False)

linked = linkage(cdist,'weighted')

clusterout = fcluster(linked, nclusters,criterion='maxclust')

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

# Calculate the uncertainty

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

uncertout = np.zeros(clusterout.shape)

for i in range(len(clusterout)):

covpt = scov[i,:] #

cid = clusterout[i] #

covin = covpt[np.where(clusterout==cid)]

covout = covpt[np.where(clusterout!=cid)]

uncertout[i] = np.mean(covin)/np.mean(covout)

# Apply rules from Thompson and Merrifield (Do this later)

# if uncert > 2, set to 2

# if uncert <0.5, set to 0

#uncertout[uncertout>2] = 2

#uncertout[uncertout<0.5] = 0

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

# Replace into full array

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

clustered = np.zeros(nlat*nlon)*np.nan

clustered[okpts] = clusterout

clustered = clustered.reshape(nlat,nlon)

cluster_count = []

for i in range(nclusters):

cid = i+1

cnt = (clustered==cid).sum()

cluster_count.append(cnt)

print("Found %i points in cluster %i" % (cnt,cid))

uncert = np.zeros(nlat*nlon)*np.nan

uncert[okpts] = uncertout

uncert = uncert.reshape(nlat,nlon)

if returnall:

return clustered,uncert,cluster_count,srho,scov,sdist,distance_matrix

# Set some defaults

ucolors = ('Blues','Purples','Greys','Blues','Reds','Oranges','Greens')

proj = ccrs.PlateCarree(central_longitude=180)

cmap = cm.get_cmap("jet",nclusters)

fig,ax = plt.subplots(1,1,subplot_kw={'projection':proj})

ax = slutil.add_coast_grid(ax)

gl = ax.gridlines(ccrs.PlateCarree(central_longitude=0),draw_labels=True,

linewidth=2, color='gray', alpha=0.5, linestyle="dotted",lw=0.75)

gl.xlabels_top = False

gl.ylabels_right = False

pcm = ax.pcolormesh(lon5,lat5,clustered,cmap=cmap,transform=ccrs.PlateCarree())#,cmap='Accent')#@,cmap='Accent')

plt.colorbar(pcm,ax=ax,orientation='horizontal')

ax.set_title("Clustering Results \n nclusters=%i %s" % (nclusters,expname))

plt.savefig("%sCluster_results_n%i_%s.png"%(outfigpath,nclusters,expname),dpi=200,transparent=True)

# Plot raw uncertainty

fig,ax = plt.subplots(1,1,subplot_kw={'projection':proj})

ax = slutil.add_coast_grid(ax)

pcm = plt.pcolormesh(lon5,lat5,uncert,cmap='copper',transform=ccrs.PlateCarree())

ax.set_title(r"Uncertainty $(<\sigma^{2}_{out,x}>/<\sigma^{2}_{in,x}>)$")

fig.colorbar(pcm,ax=ax,fraction=0.02)

plt.savefig(outfigpath+"Uncertainty.png",dpi=200)

# Apply Thompson and Merrifield thresholds

uncert[uncert>2] = 2

uncert[uncert<0.5]=0

# Plot Cluster Uncertainty

fig1,ax1 = plt.subplots(1,1,subplot_kw={'projection':proj})

ax1 = slutil.add_coast_grid(ax1)

for i in range(nclusters):

cid = i+1

if (i+1) > len(ucolors):

ci=i%len(ucolors)

else:

ci=i

cuncert = uncert.copy()

cuncert[clustered!=cid] *= np.nan

ax1.pcolormesh(lon5,lat5,cuncert,vmin=0,vmax=2,cmap=ucolors[ci],transform=ccrs.PlateCarree())

#fig.colorbar(pcm,ax=ax)

ax1.set_title("Clustering Output (nclusters=%i) %s "% (nclusters,expname))

plt.savefig(outfigpath+"Cluster_with_Shaded_uncertainties_%s.png" % expname,dpi=200)

ntime,nlat,nlon = sla.shape

slain = sla.copy()

# Preallocate

allclusters = []

alluncert = []

allcount = []

rempts = np.zeros((nlat*nlon))*np.nan

# Loop

flag = True

it = 0

while flag or it < maxiter:

expname = "iteration%02i" % (it+1)

print("Iteration %i ========================="%it)

# Perform Clustering

clustered,uncert,cluster_count = cluster_ssh(slain,lat,lon,nclusters,distthres=distthres)

# Visualize Results

fig,ax,fig1,ax1 = plot_results(clustered,uncert,expname,lat,lon,outfigpath)

# Save results

allclusters.append(clustered)

alluncert.append(uncert)

allcount.append(cluster_count)

# Check cluster counts

for i in range(nclusters):

cid = i+1

flag = False

if cluster_count[i] < minpts:

flag = True # Set flag to continue running

print("\tCluster %i (count=%i) will be removed" % (cid,cluster_count[i]))

clusteredrs = clustered.reshape(nlat*nlon)

slainrs = slain.reshape(ntime,nlat*nlon)

slainrs[:,clusteredrs==cid] = np.nan # Assign NaN Values

rempts[clusteredrs==cid] = it # Record iteration of removal

slain = slainrs.reshape(ntime,nlat,nlon)

it += 1

print("COMPLETE after %i iterations"%it)

rempts = rempts.reshape(nlat,nlon)