# Source code for MITgcmutils.llc

from __future__ import print_function
import sys
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri

[docs]def contourf(*arguments, **kwargs):
"""
Create a contourf plot of a 2-D llc array (with tricontour).

Call signatures::

contourf(X, Y, C, N, **kwargs)

contourf(X, Y, C, V, **kwargs)

Parameters
----------
X : array-like
x coordinates of the grid points

Y : array-like
y coordinates of the grid points

C : array-like
array of color values.

N : int
number of levels

V : list of float
list of levels

kwargs
passed to tricontour.

"""

arglen = len(arguments)
h = []
if arglen >= 3:
data = np.copy(arguments[2].flatten())
x = arguments[0].flatten()
y = arguments[1].flatten()

# Create the Triangulation;
# no triangles so Delaunay triangulation created.
triang = tri.Triangulation(x, y)
ntri = triang.triangles.shape[0]

if arglen == 3:
h = plt.tricontourf(triang, data, **kwargs)
elif arglen == 4:
h = plt.tricontourf(triang, data, arguments[3], **kwargs)
else:
print("wrong number of arguments")
print("need at least 3 or 4 arguments")
sys.exit(__doc__)

# show the triangles for debugging
#plt.triplot(triang, color='0.7')

else:
print("wrong number of arguments")
print("need at least x,y,fld")
sys.exit(__doc__)

return h

[docs]def contour(*arguments, **kwargs):
"""
Create a contour plot of a 2-D llc array (with tricontour).

Call signatures::

contour(X, Y, C, N, **kwargs)

contour(X, Y, C, V, **kwargs)

Parameters
----------
X : array-like
x coordinates of the grid points

Y : array-like
y coordinates of the grid points

C : array-like
array of color values.

N : int
number of levels

V : list of float
list of levels

kwargs
passed to tricontour.

"""

arglen = len(arguments)
h = []
if arglen >= 3:
data = arguments[2].flatten()
x = arguments[0].flatten()
y = arguments[1].flatten()

# Create the Triangulation;
# no triangles so Delaunay triangulation created.
triang = tri.Triangulation(x, y)
ntri = triang.triangles.shape[0]

if arglen == 3:
h = plt.tricontour(triang, data, **kwargs)
elif arglen == 4:
h = plt.tricontour(triang, data, arguments[3], **kwargs)
else:
print("wrong number of arguments")
print("need at least 3 or 4 arguments")
sys.exit(__doc__)

# show the triangles for debugging
#plt.triplot(triang, color='0.7')

else:
print("wrong number of arguments")
print("need at least x,y,fld")
sys.exit(__doc__)

return h

[docs]def flat(fld, **kwargs):
"""convert mds data into global 2D field
only fields with 2 to 5 dimensions are allowed"""

ndims = len(fld.shape)
if ndims == 2:
gfld = _flat2D(fld, **kwargs)
elif ndims == 3:
gfld = [ _flat2D(fld[a,:,:], **kwargs)
for a in range(fld.shape[0]) ]
elif ndims == 4:
gfld = [ [ _flat2D(fld[a,b,:,:], **kwargs)
for b in range(fld.shape[1]) ]
for a in range(fld.shape[0]) ]
elif ndims == 5:
gfld = [ [ [ _flat2D(fld[a,b,c,:,:], **kwargs)
for c in range(fld.shape[2]) ]
for b in range(fld.shape[1]) ]
for a in range(fld.shape[0]) ]
else:
print("wrong number of dimensions")
print("only 2 to 5 dimensions are allowed")
sys.exit(__doc__)

gfld = np.array(gfld)

return gfld

def _flat2D(fld, center='Atlantic'):
"""convert mds 2D data into global 2D field"""

nx = fld.shape[1]
ny = fld.shape[0]
n = ny//nx//4

# eastern and western hemispheres
eastern=np.concatenate((fld[:n*nx,:],fld[n*nx:2*(n*nx)]),axis=1)
tmp    = fld[2*(n*nx)+nx:,        ::-1]
western=np.concatenate((tmp[2::n,:].transpose(),
tmp[1::n,:].transpose(),
tmp[0::n,:].transpose()))
# Arctic face is special
arctic  = fld[2*(n*nx):2*(n*nx)+nx,:]
arctice = np.concatenate((np.triu(arctic[::-1,:nx//2].transpose()),
np.zeros((nx//2,nx))),axis=1)
# arcticw = np.concatenate((arctic[:,nx:nx//2-1:-1].transpose(),
#                           np.zeros((nx//2,nx//2)),
#                           arctic[nx:nx//2-1:-1,nx//2-1::-1]),axis=1)
mskr = np.tri(nx//2)[::-1,:]
arcticw = np.concatenate((arctic[0:nx//2,nx:nx//2-1:-1].transpose(),
arctic[nx//2:nx,nx:nx//2-1:-1].transpose()*mskr,
np.triu(arctic[nx:nx//2-1:-1,nx:nx//2-1:-1]),
arctic[nx:nx//2-1:-1,nx//2-1::-1]*mskr),axis=1)
#
if center == 'Pacific':
gfld = np.concatenate( ( np.concatenate((eastern,arctice)),
np.concatenate((western,arcticw)) ), axis=1)
else:
gfld = np.concatenate( ( np.concatenate((western,arcticw)),
np.concatenate((eastern,arctice)) ), axis=1)

return gfld

def _mds2D(fld,center='Atlantic'):
"""convert global 2D 'flat field' to mds 2D data"""

ni = fld.shape[-1]
nj = fld.shape[-2]
nx = ni//4
ny = nx*(3*4+1)
n = ny//nx//4

# arctic face
arcticw = fld[n*nx:,:nx]
arctice = fld[n*nx:,2*nx:3*nx]
arctic = np.concatenate((arctice,arcticw[::-1,::-1]),axis=0)

# eastern and western hemispheres
eastern=fld[:n*nx,2*nx:]
# this is tricky
western=fld[:n*nx,:2*nx]

mdsfld = np.concatenate((eastern[:,:nx],
eastern[:,nx:],
arctic[:,::-1].transpose(),
western[::-1,:].transpose().reshape((2*n*nx,nx))),
axis=0)
return mdsfld

[docs]def mds(fld,center='Atlantic'):
"""convert global 'flat' field into mds data;
only fields with 2 to 5 dimensions are allowed"""

ndims = len(fld.shape)
if ndims == 2:
mdsfld = _mds2D(fld, **kwargs)
elif ndims == 3:
mdsfld = [ _mds2D(fld[a,:,:], **kwargs)
for a in range(fld.shape[0]) ]
elif ndims == 4:
mdsfld = [ [ _mds2D(fld[a,b,:,:], **kwargs)
for b in range(fld.shape[1]) ]
for a in range(fld.shape[0]) ]
elif ndims == 5:
mdsfld = [ [ [ _mds2D(fld[a,b,c,:,:], **kwargs)
for c in range(fld.shape[2]) ]
for b in range(fld.shape[1]) ]
for a in range(fld.shape[0]) ]
else:
print("wrong number of dimensions")
print("only 2 to 5 dimensions are allowed")
sys.exit(__doc__)

mdsfld = np.array(mdsfld)

return mdsfld

[docs]def faces(fld):
"""convert mds multidimensional data into a list with 6 faces"""

ndim = len(fld.shape)
if ndim == 2:
f = _faces2D(fld)
else:
# use list for dynamical memory allocation, because it is fast
if ndim == 3:
ff = []
nk = fld.shape[0]
for k in range(nk):
fld2D = fld[k,:,:]
f2D = _faces2D(fld2D)
ff.append(f2D)
elif ndim == 4:
ff = []
nk = fld.shape[1]
nl = fld.shape[0]
for l in range(nl):
for k in range(nk):
fld2D = fld[l,k,:,:]
f2D = _faces2D(fld2D)
ff.append(f2D)
elif ndim == 5:
ff = []
nk = fld.shape[2]
nl = fld.shape[1]
nm = fld.shape[0]
for m in range(nm):
for l in range(nl):
for k in range(nk):
fld2D = fld[m,l,k,:,:]
f2D = _faces2D(fld2D)
ff.append(f2D)
# permute list indices so that face index is the first
ff = np.transpose(ff)
f  = []
for listIndex in range(len(ff)):
# for each face turn list into array, glue together and reshape
nx = ff[listIndex][0].shape[-1]
ny = ff[listIndex][0].shape[-2]
if   ndim == 3: rshp =       (nk,ny,nx)
elif ndim == 4: rshp =    (nl,nk,ny,nx)
elif ndim == 5: rshp = (nm,nl,nk,ny,nx)
f.append(np.concatenate(np.array(ff[listIndex])).reshape(rshp))

return f

[docs]def faces2mds(ff):
"""convert 6 faces to mds 2D data,
inverse opertation of llc.faces"""

ndims = len(ff[0].shape)
shp = list(ff[0].shape)
shp[-2]=2*ff[0].shape[-2]
wd = np.concatenate( (ff[3],ff[4]),axis=-2 ).reshape(shp)
f  = np.concatenate( (ff[0],ff[1],ff[2],wd),axis=-2)

return f

def _faces2D(fld):
"""convert mds 2D data into a list with 6 faces"""

nx = fld.shape[-1]
ny = fld.shape[-2]
n = ny//nx//4

# divide into faces
f = []
f.append(fld[:n*nx,:])
f.append(fld[n*nx:2*(n*nx),:])
# arctic face
f.append(fld[2*(n*nx):2*(n*nx)+nx,:])
# western hemisphere
wd = fld[2*(n*nx)+nx:,:].reshape(2*nx,n*nx)
f.append(wd[:nx,:])
f.append(wd[nx:,:])
# pseudo-sixth face
f.append(np.zeros((nx,nx)))

return f

def _sqCoord(a):
b = np.squeeze(a)
return b

def _sqData(a):
b = np.copy(np.squeeze(a))
return b

[docs]def pcol(*arguments, **kwargs):
"""
Create a pseudo-color plot of a 2-D llc array (with plt.pcolormesh).

Call signatures::

pcol(X, Y, C, **kwargs)

pcol(X, Y, C, m, **kwargs)

Parameters
----------
X : array-like
x coordinates of the grid point corners (G-points)

Y : array-like
y coordinates of the grid point corners (G-points)

C : array-like
array of color values.

m : Basemap instance, optional
map projection to use.
NOTE: currently not all projections work

kwargs
passed to plt.pcolormesh.

"""

arglen = len(arguments)
h = []
mapit = False
if arglen < 3:
print("wrong number of arguments")
print("need at least x,y,fld")
sys.exit(__doc__)
elif arglen > 3:
mapit = True
m = arguments[3]

if mapit:
# not all projections work, catch few of these here
if ( (m.projection == 'hammer') |
(m.projection == 'robin')  |
(m.projection == 'moll')   |
(m.projection == 'cea') ):
sys.exit("selected projection '"+m.projection
+"' is not supported")

# these projections use simple code for the Arctic face;
# all others require more complicted methods
stereographicProjection = (m.projection == 'npaeqd')  | \
(m.projection == 'spaeqd')  | \
(m.projection == 'nplaea')  | \
(m.projection == 'splaea')  | \
(m.projection == 'npstere') | \
(m.projection == 'spstere') | \
(m.projection == 'stere')
else:
stereographicProjection = False

xg = arguments[0]
yg = arguments[1]
data = arguments[2]

nx = data.shape[-1]
ny = data.shape[-2]
n = ny//nx//4

# color range
cax = [data.min(),data.max()]
# overwrite if necessary
if 'vmin' in kwargs: cax[0] = kwargs.pop('vmin','')
if 'vmax' in kwargs: cax[1] = kwargs.pop('vmax','')
# divide into faces
f0 = []
f0.append(faces(xg))
f0.append(faces(yg))
f0.append(faces(data))
# fill holes in coordinate arrays
#    for t in [0,1,3,4]:
#        inan = f0[2][t]==0 # _sqCoord(f0[2][t])==np.NaN]
#        f0[0][t][inan]=np.NaN
#        f0[1][t][inan]=np.NaN

#    for t in [0,1]:
#        for i in range(nx):
#            for j in range(n*nx):
#                if f0[0][t][j,i]==0:f0[0][t][100,i]
#                if f0[1][t][j,i]==0:f0[1][t][100,i]
#
#    for t in [3,4]:
#        for i in range(n*nx):
#            for j in range(nx):
#                if f0[0][t][j,i]==0:f0[0][t][j,239]
#                if f0[1][t][j,i]==0:f0[1][t][j,239]

# find the missing corners by interpolation
fo = []
fo.append( (f0[0][0][-1,0]+f0[0][2][-1,0]+f0[0][4][-1,0])/3. )
fo.append( (f0[1][2][-1,0]+f0[1][2][-1,0]+f0[1][4][-1,0])/3. )
fo.append( np.NaN )
fe = []
fe.append( (f0[0][1][0,-1]+f0[0][3][0,-1])/2. )
fe.append( (f0[1][1][0,-1]+f0[1][3][0,-1])/2. )
fe.append( np.NaN )
f  = np.copy(f0)
# fill some gaps at the face boundaries
for t in [0,2,4]:
tp = 2*(t//2)
tpp = tp
if tp==4: tpp = tp-6
for k in [0,1,2]:
tp = min(tp,3)
f[k][t] = np.concatenate((f0[k][t],f0[k][1+tp][:,:1]),axis=1)
if k==2: tmp = np.atleast_2d(np.append(f0[k][2+tpp][::-1,:1],fo[k]))
else:    tmp = np.atleast_2d(np.append(fo[k],f0[k][2+tpp][::-1,:1]))
f[k][t] = np.concatenate((f[k][t],tmp),axis=0)

for t in [1,3]:
tp = 2*(t//2)
for k in [0,1,2]:
f[k][t] = np.concatenate((f0[k][t],f0[k][2+tp][:1,:]),axis=0)
if k==2: tmp = np.atleast_2d(np.append(f0[k][3+tp][:1,::-1],fe[k]))
else:    tmp = np.atleast_2d(np.append(fe[k],f0[k][3+tp][:1,::-1]))
f[k][t] = np.concatenate((f[k][t],tmp.transpose()),axis=1)

# we do not really have a sixth face so we overwrite the southernmost row
# of face 4 and 5 by a hack:
for t in [3,4]:
f[0][t][:,-1] = f[0][t][:,-2]
f[1][t][:,-1] = -90. # degree = south pole

# make sure that only longitudes of one sign are on individual lateral faces
i0 = f[0][3]<0.
f[0][3][i0] = f[0][3][i0]+360.
# plot the lateral faces
ph = []
for t in [0,1,3,4]:
if mapit: x, y = m(_sqCoord(f[0][t]), _sqCoord(f[1][t]))
else:     x, y =   _sqCoord(f[0][t]), _sqCoord(f[1][t])
ph.append(plt.pcolormesh(x,y,_sqData(f[2][t]), **kwargs))
# plot more lateral faces to be able to select the longitude range later
for t in [1,3,4]:
f[0][t] = f[0][t]+ (-1)**t*360.
if mapit: x, y = m(_sqCoord(f[0][t]), _sqCoord(f[1][t]))
else:     x, y =   _sqCoord(f[0][t]), _sqCoord(f[1][t])
ph.append(plt.pcolormesh(x,y,_sqData(f[2][t]), **kwargs))

# Arctic face is special, because of the rotation of the grid by
# rangle = 7deg (seems to be the default)
t = 2

if mapit & stereographicProjection:
x, y = m(_sqCoord(f[0][t]),_sqCoord(f[1][t]))
ph.append(plt.pcolormesh(x,y,_sqData(f[2][t]), **kwargs))
else:
rangle = 7.
# first half of Arctic tile
nn = nx//2+1
xx = np.copy(f[0][t][:nn,:])
yy = np.copy(f[1][t][:nn,:])
zz = np.copy(f[2][t][:nn,:])
xx = np.where(xx<rangle,xx+360,xx)
if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
else:     x, y =   _sqCoord(xx),_sqCoord(yy)
ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
# repeat for xx-360
xx = xx-360.
if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
else:     x, y =   _sqCoord(xx),_sqCoord(yy)
ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
# second half of Arctic tile
nn = nx//2
xx = np.copy(f[0][t][nn:,:])
yy = np.copy(f[1][t][nn:,:])
zz = np.copy(f[2][t][nn:,:])
#
if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
else:     x, y =   _sqCoord(xx),_sqCoord(yy)
ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
# repeat for xx+360
xx = xx + 360.
if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
else:     x, y =   _sqCoord(xx),_sqCoord(yy)
ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))

if not mapit:
plt.xlim([-170,190])
plt.ylim([-90,90])

for im in ph:
im.set_clim(cax[0],cax[1])

return ph