MITgcmutils/llc/llc.py

34625de9eb Oliv*0001 from __future__ import print_function
e592bd33f2 Mart*0002 import sys
                0003 import numpy as np
                0004 import matplotlib.pyplot as plt
                0005 import matplotlib.tri as tri
                0006
                0007 def contourf(*arguments, **kwargs):
7621b5d564 Oliv*0008     """
e592bd33f2 Mart*0009     Create a contourf plot of a 2-D llc array (with tricontour).
                0010
7621b5d564 Oliv*0011     Call signatures::
                0012
                0013         contourf(X, Y, C, N, **kwargs)
                0014
                0015         contourf(X, Y, C, V, **kwargs)
                0016
                0017     Parameters
                0018     ----------
                0019     X : array-like
                0020         x coordinates of the grid points
                0021
                0022     Y : array-like
                0023         y coordinates of the grid points
                0024
                0025     C : array-like
                0026         array of color values.
e592bd33f2 Mart*0027
7621b5d564 Oliv*0028     N : int
                0029         number of levels
                0030
                0031     V : list of float
                0032         list of levels
                0033
                0034     kwargs
                0035         passed to tricontour.
e592bd33f2 Mart*0036
                0037     """
                0038
                0039     arglen = len(arguments)
                0040     h = []
                0041     if arglen >= 3:
                0042         data = np.copy(arguments[2].flatten())
                0043         x = arguments[0].flatten()
                0044         y = arguments[1].flatten()
                0045
7621b5d564 Oliv*0046         # Create the Triangulation;
                0047         # no triangles so Delaunay triangulation created.
e592bd33f2 Mart*0048         triang = tri.Triangulation(x, y)
                0049         ntri = triang.triangles.shape[0]
                0050
                0051         # Mask off unwanted triangles.
                0052         mask = np.where(data[triang.triangles].prod(axis=1)==0., 1, 0)
                0053         triang.set_mask(mask)
7621b5d564 Oliv*0054
e592bd33f2 Mart*0055         if arglen == 3:
                0056             h = plt.tricontourf(triang, data, **kwargs)
                0057         elif arglen == 4:
                0058             h = plt.tricontourf(triang, data, arguments[3], **kwargs)
                0059         else:
34625de9eb Oliv*0060             print("wrong number of arguments")
                0061             print("need at least 3 or 4 arguments")
e592bd33f2 Mart*0062             sys.exit(__doc__)
                0063
                0064         # show the triangles for debugging
                0065         #plt.triplot(triang, color='0.7')
                0066
                0067     else:
34625de9eb Oliv*0068         print("wrong number of arguments")
                0069         print("need at least x,y,fld")
e592bd33f2 Mart*0070         sys.exit(__doc__)
                0071
                0072     return h
                0073
                0074 def contour(*arguments, **kwargs):
7621b5d564 Oliv*0075     """
e592bd33f2 Mart*0076     Create a contour plot of a 2-D llc array (with tricontour).
                0077
7621b5d564 Oliv*0078     Call signatures::
                0079
                0080         contour(X, Y, C, N, **kwargs)
e592bd33f2 Mart*0081
7621b5d564 Oliv*0082         contour(X, Y, C, V, **kwargs)
                0083
                0084     Parameters
                0085     ----------
                0086     X : array-like
                0087         x coordinates of the grid points
                0088
                0089     Y : array-like
                0090         y coordinates of the grid points
                0091
                0092     C : array-like
                0093         array of color values.
                0094
                0095     N : int
                0096         number of levels
                0097
                0098     V : list of float
                0099         list of levels
                0100
                0101     kwargs
                0102         passed to tricontour.
e592bd33f2 Mart*0103
                0104     """
                0105
                0106     arglen = len(arguments)
                0107     h = []
                0108     if arglen >= 3:
                0109         data = arguments[2].flatten()
                0110         x = arguments[0].flatten()
                0111         y = arguments[1].flatten()
                0112
7621b5d564 Oliv*0113         # Create the Triangulation;
                0114         # no triangles so Delaunay triangulation created.
e592bd33f2 Mart*0115         triang = tri.Triangulation(x, y)
                0116         ntri = triang.triangles.shape[0]
                0117
                0118         # Mask off unwanted triangles.
                0119         mask = np.where(data[triang.triangles].prod(axis=1)==0., 1, 0)
                0120         triang.set_mask(mask)
7621b5d564 Oliv*0121
e592bd33f2 Mart*0122         if arglen == 3:
                0123             h = plt.tricontour(triang, data, **kwargs)
                0124         elif arglen == 4:
                0125             h = plt.tricontour(triang, data, arguments[3], **kwargs)
                0126         else:
34625de9eb Oliv*0127             print("wrong number of arguments")
                0128             print("need at least 3 or 4 arguments")
e592bd33f2 Mart*0129             sys.exit(__doc__)
7621b5d564 Oliv*0130
e592bd33f2 Mart*0131         # show the triangles for debugging
                0132         #plt.triplot(triang, color='0.7')
                0133
                0134     else:
34625de9eb Oliv*0135         print("wrong number of arguments")
                0136         print("need at least x,y,fld")
e592bd33f2 Mart*0137         sys.exit(__doc__)
                0138
                0139     return h
                0140
                0141 def flat(fld, **kwargs):
                0142     """convert mds data into global 2D field
                0143     only fields with 2 to 5 dimensions are allowed"""
                0144
                0145     ndims = len(fld.shape)
7621b5d564 Oliv*0146     if ndims == 2:
05d0413b02 Mart*0147         gfld = _flat2D(fld, **kwargs)
7621b5d564 Oliv*0148     elif ndims == 3:
                0149         gfld = [ _flat2D(fld[a,:,:], **kwargs)
                0150                  for a in range(fld.shape[0]) ]
                0151     elif ndims == 4:
                0152         gfld = [ [ _flat2D(fld[a,b,:,:], **kwargs)
                0153                    for b in range(fld.shape[1]) ]
e592bd33f2 Mart*0154                  for a in range(fld.shape[0]) ]
7621b5d564 Oliv*0155     elif ndims == 5:
                0156         gfld = [ [ [ _flat2D(fld[a,b,c,:,:], **kwargs)
e592bd33f2 Mart*0157                      for c in range(fld.shape[2]) ]
                0158                    for b in range(fld.shape[1]) ]
                0159                  for a in range(fld.shape[0]) ]
                0160     else:
34625de9eb Oliv*0161         print("wrong number of dimensions")
                0162         print("only 2 to 5 dimensions are allowed")
e592bd33f2 Mart*0163         sys.exit(__doc__)
                0164
                0165     gfld = np.array(gfld)
                0166
                0167     return gfld
                0168
                0169 def _flat2D(fld, center='Atlantic'):
                0170     """convert mds 2D data into global 2D field"""
                0171
                0172     nx = fld.shape[1]
                0173     ny = fld.shape[0]
27d195aed6 Oliv*0174     n = ny//nx//4
7621b5d564 Oliv*0175
e592bd33f2 Mart*0176     # eastern and western hemispheres
                0177     eastern=np.concatenate((fld[:n*nx,:],fld[n*nx:2*(n*nx)]),axis=1)
                0178     tmp    = fld[2*(n*nx)+nx:,        ::-1]
                0179     western=np.concatenate((tmp[2::n,:].transpose(),
                0180                             tmp[1::n,:].transpose(),
                0181                             tmp[0::n,:].transpose()))
                0182     # Arctic face is special
                0183     arctic  = fld[2*(n*nx):2*(n*nx)+nx,:]
27d195aed6 Oliv*0184     arctice = np.concatenate((np.triu(arctic[::-1,:nx//2].transpose()),
                0185                               np.zeros((nx//2,nx))),axis=1)
                0186     # arcticw = np.concatenate((arctic[:,nx:nx//2-1:-1].transpose(),
                0187     #                           np.zeros((nx//2,nx//2)),
                0188     #                           arctic[nx:nx//2-1:-1,nx//2-1::-1]),axis=1)
                0189     mskr = np.tri(nx//2)[::-1,:]
                0190     arcticw = np.concatenate((arctic[0:nx//2,nx:nx//2-1:-1].transpose(),
                0191                               arctic[nx//2:nx,nx:nx//2-1:-1].transpose()*mskr,
                0192                               np.triu(arctic[nx:nx//2-1:-1,nx:nx//2-1:-1]),
                0193                               arctic[nx:nx//2-1:-1,nx//2-1::-1]*mskr),axis=1)
e592bd33f2 Mart*0194     #
                0195     if center == 'Pacific':
                0196         gfld = np.concatenate( ( np.concatenate((eastern,arctice)),
                0197                                  np.concatenate((western,arcticw)) ), axis=1)
                0198     else:
                0199         gfld = np.concatenate( ( np.concatenate((western,arcticw)),
                0200                                  np.concatenate((eastern,arctice)) ), axis=1)
7621b5d564 Oliv*0201
e592bd33f2 Mart*0202     return gfld
                0203
                0204 def _mds2D(fld,center='Atlantic'):
                0205     """convert global 2D 'flat field' to mds 2D data"""
7621b5d564 Oliv*0206
e592bd33f2 Mart*0207     ni = fld.shape[-1]
                0208     nj = fld.shape[-2]
27d195aed6 Oliv*0209     nx = ni//4
e592bd33f2 Mart*0210     ny = nx*(3*4+1)
27d195aed6 Oliv*0211     n = ny//nx//4
e592bd33f2 Mart*0212
                0213     # arctic face
                0214     arcticw = fld[n*nx:,:nx]
                0215     arctice = fld[n*nx:,2*nx:3*nx]
                0216     arctic = np.concatenate((arctice,arcticw[::-1,::-1]),axis=0)
                0217
                0218     # eastern and western hemispheres
                0219     eastern=fld[:n*nx,2*nx:]
                0220     # this is tricky
                0221     western=fld[:n*nx,:2*nx]
                0222
                0223     mdsfld = np.concatenate((eastern[:,:nx],
                0224                              eastern[:,nx:],
                0225                              arctic[:,::-1].transpose(),
                0226                              western[::-1,:].transpose().reshape((2*n*nx,nx))),
                0227                              axis=0)
                0228     return mdsfld
                0229
                0230 def mds(fld,center='Atlantic'):
7621b5d564 Oliv*0231     """convert global 'flat' field into mds data;
e592bd33f2 Mart*0232     only fields with 2 to 5 dimensions are allowed"""
                0233
                0234     ndims = len(fld.shape)
7621b5d564 Oliv*0235     if ndims == 2:
05d0413b02 Mart*0236         mdsfld = _mds2D(fld, **kwargs)
7621b5d564 Oliv*0237     elif ndims == 3:
                0238         mdsfld = [ _mds2D(fld[a,:,:], **kwargs)
                0239                  for a in range(fld.shape[0]) ]
                0240     elif ndims == 4:
                0241         mdsfld = [ [ _mds2D(fld[a,b,:,:], **kwargs)
                0242                    for b in range(fld.shape[1]) ]
e592bd33f2 Mart*0243                  for a in range(fld.shape[0]) ]
7621b5d564 Oliv*0244     elif ndims == 5:
                0245         mdsfld = [ [ [ _mds2D(fld[a,b,c,:,:], **kwargs)
e592bd33f2 Mart*0246                      for c in range(fld.shape[2]) ]
                0247                    for b in range(fld.shape[1]) ]
                0248                  for a in range(fld.shape[0]) ]
                0249     else:
34625de9eb Oliv*0250         print("wrong number of dimensions")
                0251         print("only 2 to 5 dimensions are allowed")
e592bd33f2 Mart*0252         sys.exit(__doc__)
                0253
                0254     mdsfld = np.array(mdsfld)
7621b5d564 Oliv*0255
e592bd33f2 Mart*0256     return mdsfld
                0257
                0258 def faces(fld):
8c4d22700d Mart*0259     """convert mds multidimensional data into a list with 6 faces"""
7621b5d564 Oliv*0260
8c4d22700d Mart*0261     ndim = len(fld.shape)
                0262     if ndim == 2:
                0263         f = _faces2D(fld)
                0264     else:
                0265         # use list for dynamical memory allocation, because it is fast
                0266         if ndim == 3:
                0267             ff = []
                0268             nk = fld.shape[0]
                0269             for k in range(nk):
                0270                 fld2D = fld[k,:,:]
                0271                 f2D = _faces2D(fld2D)
                0272                 ff.append(f2D)
                0273         elif ndim == 4:
                0274             ff = []
                0275             nk = fld.shape[1]
                0276             nl = fld.shape[0]
                0277             for l in range(nl):
                0278                 for k in range(nk):
                0279                     fld2D = fld[l,k,:,:]
                0280                     f2D = _faces2D(fld2D)
                0281                     ff.append(f2D)
                0282         elif ndim == 5:
                0283             ff = []
                0284             nk = fld.shape[2]
                0285             nl = fld.shape[1]
                0286             nm = fld.shape[0]
                0287             for m in range(nm):
                0288                 for l in range(nl):
                0289                     for k in range(nk):
                0290                         fld2D = fld[m,l,k,:,:]
                0291                         f2D = _faces2D(fld2D)
                0292                         ff.append(f2D)
                0293         # permute list indices so that face index is the first
                0294         ff = np.transpose(ff)
                0295         f  = []
                0296         for listIndex in range(len(ff)):
                0297             # for each face turn list into array, glue together and reshape
                0298             nx = ff[listIndex][0].shape[-1]
                0299             ny = ff[listIndex][0].shape[-2]
                0300             if   ndim == 3: rshp =       (nk,ny,nx)
                0301             elif ndim == 4: rshp =    (nl,nk,ny,nx)
                0302             elif ndim == 5: rshp = (nm,nl,nk,ny,nx)
                0303             f.append(np.concatenate(np.array(ff[listIndex])).reshape(rshp))
                0304
                0305     return f
                0306
96a6fb1c14 Mart*0307 def faces2mds(ff):
                0308     """convert 6 faces to mds 2D data,
                0309     inverse opertation of llc.faces"""
7621b5d564 Oliv*0310
96a6fb1c14 Mart*0311     ndims = len(ff[0].shape)
                0312     shp = list(ff[0].shape)
                0313     shp[-2]=2*ff[0].shape[-2]
                0314     wd = np.concatenate( (ff[3],ff[4]),axis=-2 ).reshape(shp)
                0315     f  = np.concatenate( (ff[0],ff[1],ff[2],wd),axis=-2)
                0316
                0317     return f
                0318
8c4d22700d Mart*0319 def _faces2D(fld):
e592bd33f2 Mart*0320     """convert mds 2D data into a list with 6 faces"""
7621b5d564 Oliv*0321
e592bd33f2 Mart*0322     nx = fld.shape[-1]
                0323     ny = fld.shape[-2]
27d195aed6 Oliv*0324     n = ny//nx//4
7621b5d564 Oliv*0325
e592bd33f2 Mart*0326     # divide into faces
                0327     f = []
                0328     f.append(fld[:n*nx,:])
                0329     f.append(fld[n*nx:2*(n*nx),:])
                0330     # arctic face
                0331     f.append(fld[2*(n*nx):2*(n*nx)+nx,:])
96a6fb1c14 Mart*0332     # western hemisphere
                0333     wd = fld[2*(n*nx)+nx:,:].reshape(2*nx,n*nx)
                0334     f.append(wd[:nx,:])
                0335     f.append(wd[nx:,:])
e592bd33f2 Mart*0336     # pseudo-sixth face
                0337     f.append(np.zeros((nx,nx)))
                0338
                0339     return f
                0340
                0341
                0342 def _sqCoord(a):
7621b5d564 Oliv*0343     b = np.squeeze(a)
610fa23955 Mart*0344     return b
                0345
                0346 def _sqData(a):
                0347     b = np.copy(np.squeeze(a))
5e02a422c5 Mart*0348     b = np.ma.masked_where(b==0., b)
                0349     b = np.ma.masked_where(np.isnan(b), b)
e592bd33f2 Mart*0350     return b
                0351
                0352 def pcol(*arguments, **kwargs):
7621b5d564 Oliv*0353     """
e592bd33f2 Mart*0354     Create a pseudo-color plot of a 2-D llc array (with plt.pcolormesh).
                0355
7621b5d564 Oliv*0356     Call signatures::
                0357
                0358         pcol(X, Y, C, **kwargs)
e592bd33f2 Mart*0359
7621b5d564 Oliv*0360         pcol(X, Y, C, m, **kwargs)
                0361
                0362     Parameters
                0363     ----------
                0364     X : array-like
                0365         x coordinates of the grid point corners (G-points)
                0366
                0367     Y : array-like
                0368         y coordinates of the grid point corners (G-points)
                0369
                0370     C : array-like
                0371         array of color values.
                0372
                0373     m : Basemap instance, optional
                0374         map projection to use.
                0375         NOTE: currently not all projections work
                0376
                0377     kwargs
                0378         passed to plt.pcolormesh.
e592bd33f2 Mart*0379
                0380     """
                0381
                0382     arglen = len(arguments)
                0383     h = []
                0384     mapit = False
                0385     if arglen < 3:
34625de9eb Oliv*0386         print("wrong number of arguments")
                0387         print("need at least x,y,fld")
e592bd33f2 Mart*0388         sys.exit(__doc__)
                0389     elif arglen > 3:
                0390         mapit = True
                0391         m = arguments[3]
                0392
1dad64efb0 Mart*0393     if mapit:
                0394         # not all projections work, catch few of these here
7621b5d564 Oliv*0395         if ( (m.projection == 'hammer') |
                0396              (m.projection == 'robin')  |
                0397              (m.projection == 'moll')   |
04b905f0b9 Mart*0398              (m.projection == 'cea') ):
1dad64efb0 Mart*0399             sys.exit("selected projection '"+m.projection
                0400                      +"' is not supported")
                0401
7621b5d564 Oliv*0402         # these projections use simple code for the Arctic face;
1dad64efb0 Mart*0403         # all others require more complicted methods
                0404         stereographicProjection = (m.projection == 'npaeqd')  | \
                0405                                   (m.projection == 'spaeqd')  | \
                0406                                   (m.projection == 'nplaea')  | \
                0407                                   (m.projection == 'splaea')  | \
                0408                                   (m.projection == 'npstere') | \
                0409                                   (m.projection == 'spstere') | \
                0410                                   (m.projection == 'stere')
                0411     else:
                0412         stereographicProjection = False
                0413
7621b5d564 Oliv*0414
e592bd33f2 Mart*0415     xg = arguments[0]
                0416     yg = arguments[1]
                0417     data = arguments[2]
                0418
                0419     nx = data.shape[-1]
                0420     ny = data.shape[-2]
27d195aed6 Oliv*0421     n = ny//nx//4
7621b5d564 Oliv*0422
e592bd33f2 Mart*0423     # color range
                0424     cax = [data.min(),data.max()]
                0425     # overwrite if necessary
                0426     if 'vmin' in kwargs: cax[0] = kwargs.pop('vmin','')
                0427     if 'vmax' in kwargs: cax[1] = kwargs.pop('vmax','')
                0428     # divide into faces
                0429     f0 = []
                0430     f0.append(faces(xg))
                0431     f0.append(faces(yg))
                0432     f0.append(faces(data))
                0433     # fill holes in coordinate arrays
                0434 #    for t in [0,1,3,4]:
95edcfd937 Oliv*0435 #        inan = f0[2][t]==0 # _sqCoord(f0[2][t])==np.nan]
                0436 #        f0[0][t][inan]=np.nan
                0437 #        f0[1][t][inan]=np.nan
e592bd33f2 Mart*0438
                0439 #    for t in [0,1]:
                0440 #        for i in range(nx):
                0441 #            for j in range(n*nx):
                0442 #                if f0[0][t][j,i]==0:f0[0][t][100,i]
                0443 #                if f0[1][t][j,i]==0:f0[1][t][100,i]
                0444 #
                0445 #    for t in [3,4]:
                0446 #        for i in range(n*nx):
                0447 #            for j in range(nx):
                0448 #                if f0[0][t][j,i]==0:f0[0][t][j,239]
                0449 #                if f0[1][t][j,i]==0:f0[1][t][j,239]
7621b5d564 Oliv*0450
e592bd33f2 Mart*0451     # find the missing corners by interpolation
                0452     fo = []
                0453     fo.append( (f0[0][0][-1,0]+f0[0][2][-1,0]+f0[0][4][-1,0])/3. )
                0454     fo.append( (f0[1][2][-1,0]+f0[1][2][-1,0]+f0[1][4][-1,0])/3. )
95edcfd937 Oliv*0455     fo.append( np.nan )
e592bd33f2 Mart*0456     fe = []
                0457     fe.append( (f0[0][1][0,-1]+f0[0][3][0,-1])/2. )
                0458     fe.append( (f0[1][1][0,-1]+f0[1][3][0,-1])/2. )
95edcfd937 Oliv*0459     fe.append( np.nan )
616d857f4a Mart*0460     f = np.array(f0, dtype=object)
                0461     # fill some gaps at the face boundaries, but only for the coordinate arrays (k=0,1)
e592bd33f2 Mart*0462     for t in [0,2,4]:
27d195aed6 Oliv*0463         tp = 2*(t//2)
e592bd33f2 Mart*0464         tpp = tp
                0465         if tp==4: tpp = tp-6
616d857f4a Mart*0466         for k in [0,1]:
e592bd33f2 Mart*0467             tp = min(tp,3)
                0468             f[k][t] = np.concatenate((f0[k][t],f0[k][1+tp][:,:1]),axis=1)
                0469             if k==2: tmp = np.atleast_2d(np.append(f0[k][2+tpp][::-1,:1],fo[k]))
                0470             else:    tmp = np.atleast_2d(np.append(fo[k],f0[k][2+tpp][::-1,:1]))
                0471             f[k][t] = np.concatenate((f[k][t],tmp),axis=0)
                0472
                0473     for t in [1,3]:
27d195aed6 Oliv*0474         tp = 2*(t//2)
616d857f4a Mart*0475         for k in [0,1]:
e592bd33f2 Mart*0476             f[k][t] = np.concatenate((f0[k][t],f0[k][2+tp][:1,:]),axis=0)
                0477             if k==2: tmp = np.atleast_2d(np.append(f0[k][3+tp][:1,::-1],fe[k]))
                0478             else:    tmp = np.atleast_2d(np.append(fe[k],f0[k][3+tp][:1,::-1]))
                0479             f[k][t] = np.concatenate((f[k][t],tmp.transpose()),axis=1)
7621b5d564 Oliv*0480
29206b70d8 Mart*0481     # we do not really have a sixth face so we overwrite the southernmost row
                0482     # of face 4 and 5 by a hack:
                0483     for t in [3,4]:
                0484         f[0][t][:,-1] = f[0][t][:,-2]
                0485         f[1][t][:,-1] = -90. # degree = south pole
7621b5d564 Oliv*0486
e592bd33f2 Mart*0487     # make sure that only longitudes of one sign are on individual lateral faces
                0488     i0 = f[0][3]<0.
                0489     f[0][3][i0] = f[0][3][i0]+360.
                0490     # plot the lateral faces
                0491     ph = []
                0492     for t in [0,1,3,4]:
                0493         if mapit: x, y = m(_sqCoord(f[0][t]), _sqCoord(f[1][t]))
                0494         else:     x, y =   _sqCoord(f[0][t]), _sqCoord(f[1][t])
610fa23955 Mart*0495         ph.append(plt.pcolormesh(x,y,_sqData(f[2][t]), **kwargs))
e592bd33f2 Mart*0496     # plot more lateral faces to be able to select the longitude range later
7621b5d564 Oliv*0497     for t in [1,3,4]:
1dad64efb0 Mart*0498         f[0][t] = f[0][t]+ (-1)**t*360.
e592bd33f2 Mart*0499         if mapit: x, y = m(_sqCoord(f[0][t]), _sqCoord(f[1][t]))
                0500         else:     x, y =   _sqCoord(f[0][t]), _sqCoord(f[1][t])
610fa23955 Mart*0501         ph.append(plt.pcolormesh(x,y,_sqData(f[2][t]), **kwargs))
e592bd33f2 Mart*0502
7621b5d564 Oliv*0503     # Arctic face is special, because of the rotation of the grid by
1dad64efb0 Mart*0504     # rangle = 7deg (seems to be the default)
e592bd33f2 Mart*0505     t = 2
1dad64efb0 Mart*0506
                0507     if mapit & stereographicProjection:
                0508         x, y = m(_sqCoord(f[0][t]),_sqCoord(f[1][t]))
                0509         ph.append(plt.pcolormesh(x,y,_sqData(f[2][t]), **kwargs))
                0510     else:
                0511         rangle = 7.
                0512         # first half of Arctic tile
27d195aed6 Oliv*0513         nn = nx//2+1
1dad64efb0 Mart*0514         xx = np.copy(f[0][t][:nn,:])
                0515         yy = np.copy(f[1][t][:nn,:])
616d857f4a Mart*0516         # make sure that the data have one columns/rows fewer that the coordinates
                0517         zz = np.copy(f[2][t][:nn-1,:])
1dad64efb0 Mart*0518         xx = np.where(xx<rangle,xx+360,xx)
                0519         if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
                0520         else:     x, y =   _sqCoord(xx),_sqCoord(yy)
                0521         ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
                0522         # repeat for xx-360
                0523         xx = xx-360.
                0524         if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
                0525         else:     x, y =   _sqCoord(xx),_sqCoord(yy)
                0526         ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
                0527         # second half of Arctic tile
7621b5d564 Oliv*0528         nn = nx//2
1dad64efb0 Mart*0529         xx = np.copy(f[0][t][nn:,:])
                0530         yy = np.copy(f[1][t][nn:,:])
                0531         zz = np.copy(f[2][t][nn:,:])
                0532         #
                0533         if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
                0534         else:     x, y =   _sqCoord(xx),_sqCoord(yy)
                0535         ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
                0536         # repeat for xx+360
                0537         xx = xx + 360.
                0538         if mapit: x, y = m(_sqCoord(xx),_sqCoord(yy))
                0539         else:     x, y =   _sqCoord(xx),_sqCoord(yy)
                0540         ph.append(plt.pcolormesh(x,y,_sqData(zz), **kwargs))
e592bd33f2 Mart*0541
                0542     if not mapit:
                0543         plt.xlim([-170,190])
                0544         plt.ylim([-90,90])
                0545
                0546     for im in ph:
                0547         im.set_clim(cax[0],cax[1])
                0548
                0549     return ph
616d857f4a Mart*0550
                0551 def _getDims(u,v):
                0552     """
                0553     Retrieve dimensions of input fields and do some sanity checks
                0554     """
                0555
                0556     lenu = len(np.shape(u))
                0557     nt = 1
                0558     nk = 1
                0559     ntt = 1
                0560     nkk = 1
                0561     if lenu == 2:
                0562         nju, niu = np.shape(u)
                0563         njv, niv = np.shape(v)
                0564     elif lenu == 3:
                0565         nk, nju, niu = np.shape(u)
                0566         nkk,njv, niv = np.shape(v)
                0567     elif lenu == 4:
                0568         nt, nk, nju, niu = np.shape(u)
                0569         ntt,nkk,njv, niv = np.shape(v)
                0570     else:
                0571         raise ValueError('Can only handle 2 to 4 dimensions')
                0572
                0573     if nju!=njv or niu!=niv or nk!=nkk or nt!=ntt:
                0574         raise ValueError('input arrays must have the same dimensions.\n'
                0575                          + 'u.shape = (%i,%i,%i,%i)\n'%(nt,nk,nju,niu)
                0576                          + 'v.shape = (%i,%i,%i,%i)'%(ntt,nkk,njv,niv)
                0577                          )
                0578
                0579     if nju!=13*niu:
                0580         raise ValueError('nju=%i not equal 13*niu, niu=%i\n'%(nju,niu) +
                0581                          'This is not and NOT an llc grid.')
                0582
                0583     return nt, nk, nju, niu
                0584
                0585 def div(u, v, dxg=None, dyg=None, rac=None, hfw=None, hfs=None):
                0586     """
                0587     Compute divergence of vector field (U,V) on llc grid
                0588
                0589     Call signatures::
                0590
                0591        divergence = div(U, V, DXG, DYG, RAC, HFW, HFS)
                0592        divergence = div(U, V)
                0593        divergence = div(U, V, DXG, DYG)
                0594        divergence = div(U, V, DXG, DYG, RAC)
                0595        divergence = div(U, V, DXG, DYG, hfw=HFW, hfs=HFS)
                0596
                0597     Parameters
                0598     ----------
                0599     u   : array-like (timelevel,depthlevel,jpoint,ipoint)
                0600           x-component of vector field at u-point
                0601
                0602     v   : array-like (timelevel,depthlevel,jpoint,ipoint)
                0603           y-component of vector field at v-point
                0604
                0605     dxg : array-like (jpoint,ipoint), optional
                0606           grid spacing in x across v-point, defaults to one
                0607
                0608     dyg : array-like (jpoint,ipoint), optional
                0609           grid spacing in y across u-point, defaults to one
                0610
                0611     rac : array-like (jpoint,ipoint), optional
                0612           grid cell area, defaults to dxg*dyg
                0613
                0614     hfw : array-like (depthlevel,jpoint,ipoint), optional
                0615           hFac at u-point, defaults to one
                0616
                0617     hfs : array-like (depthlevel,jpoint,ipoint), optional
                0618           hFac at v-point, defaults to one
                0619     """
                0620
                0621     nt, nk, nj, ni =  _getDims(u,v)
                0622
                0623     ushape = u.shape
                0624
                0625     if dxg is None:
                0626         dxg = np.ones((nj,ni))
                0627
                0628     if dyg is None:
                0629         dyg = np.ones((nj,ni))
                0630
                0631     if rac is None:
                0632         rac = dxg*dyg
                0633
                0634     if hfw is None:
                0635         hfw = np.ones((nk,nj,ni))
                0636
                0637     if hfs is None:
                0638         hfs = np.ones((nk,nj,ni))
                0639
                0640     u   = u.reshape(nt,nk,nj,ni)
                0641     v   = v.reshape(nt,nk,nj,ni)
                0642     hfw = hfw.reshape(nk,nj,ni)
                0643     hfs = hfs.reshape(nk,nj,ni)
                0644
95edcfd937 Oliv*0645     recip_racf = faces(1./np.where(rac==0.,np.inf,rac))
616d857f4a Mart*0646     divergence = np.zeros(u.shape)
                0647     for t in range(nt):
                0648         for k in range(nk):
                0649             uflx = faces(u[t,k,:,:]*hfw[k,:,:]*dyg)
                0650             vflx = faces(v[t,k,:,:]*hfs[k,:,:]*dxg)
                0651             divf = faces(np.zeros((nj,ni)))
                0652             for iface in range(len(uflx)-1):
                0653                 du = np.roll(uflx[iface],-1,axis=-1)-uflx[iface]
                0654                 dv = np.roll(vflx[iface],-1,axis=-2)-vflx[iface]
                0655                 # now take care of the connectivity
                0656                 if iface==0:
                0657                     du[:,-1] = uflx[1][:,   0] - uflx[iface][:,-1]
                0658                     dv[-1,:] = uflx[2][::-1,0] - vflx[iface][-1,:]
                0659                 if iface==1:
                0660                     du[:,-1] = vflx[3][0,::-1] - uflx[iface][:,-1]
                0661                     dv[-1,:] = vflx[2][0,:]    - vflx[iface][-1,:]
                0662                 if iface==2:
                0663                     du[:,-1] = uflx[3][:,   0] - uflx[iface][:,-1]
                0664                     dv[-1,:] = uflx[4][::-1,0] - vflx[iface][-1,:]
                0665                 if iface==3:
                0666                     du[:,-1] = 0. # hack
                0667                     dv[-1,:] = vflx[4][0,:]    - vflx[iface][-1,:]
                0668                 if iface==4:
                0669                     du[:,-1] = 0. # hack
                0670                     dv[-1,:] = uflx[0][::-1,0] - vflx[iface][-1,:]
                0671                 # putting it all together
                0672                 divf[iface] = (du + dv)*recip_racf[iface]
                0673
                0674             divergence[t,k,:,:] = faces2mds(divf)
                0675
                0676     return divergence.reshape(ushape)
                0677
                0678
                0679 def uv2c(u,v):
                0680     """
                0681     Average vector component (u,v) to center points on llc grid
                0682
                0683     Call signatures::
                0684
                0685        uc,vc = uv2c(U,V)
                0686
                0687     Parameters
                0688     ----------
                0689     U   : array-like (timelevel,depthlevel,jpoint,ipoint)
                0690           x-component of vector field at u-point
                0691
                0692     V   : array-like (timelevel,depthlevel,jpoint,ipoint)
                0693           y-component of vector field at v-point
                0694
                0695     """
                0696
                0697     nt, nk, nj, ni =  _getDims(u,v)
                0698     ushape = u.shape
                0699
                0700     u   = u.reshape(nt,nk,nj,ni)
                0701     v   = v.reshape(nt,nk,nj,ni)
                0702
                0703     uc = np.zeros(u.shape)
                0704     vc = np.zeros(v.shape)
                0705     for t in range(nt):
                0706         for k in range(nk):
                0707             uf  = faces(u[t,k,:,:])
                0708             vf  = faces(v[t,k,:,:])
                0709             ucf = faces(np.zeros((nj,ni)))
                0710             vcf = faces(np.zeros((nj,ni)))
                0711             for iface in range(len(uf)-1):
                0712                 uk = np.roll(uf[iface],-1,axis=-1)+uf[iface]
                0713                 vk = np.roll(vf[iface],-1,axis=-2)+vf[iface]
                0714                 # now take care of the connectivity
                0715                 if iface==0:
                0716                     uk[:,-1] = uf[1][:,   0] + uf[iface][:,-1]
                0717                     vk[-1,:] = uf[2][::-1,0] + vf[iface][-1,:]
                0718                 if iface==1:
                0719                     uk[:,-1] = vf[3][0,::-1] + uf[iface][:,-1]
                0720                     vk[-1,:] = vf[2][0,:]    + vf[iface][-1,:]
                0721                 if iface==2:
                0722                     uk[:,-1] = uf[3][:,   0] + uf[iface][:,-1]
                0723                     vk[-1,:] = uf[4][::-1,0] + vf[iface][-1,:]
                0724                 if iface==3:
                0725                     uk[:,-1] = 0. # hack
                0726                     vk[-1,:] = vf[4][0,:]    + vf[iface][-1,:]
                0727                 if iface==4:
                0728                     uk[:,-1] = 0. # hack
                0729                     vk[-1,:] = uf[0][::-1,0] + vf[iface][-1,:]
                0730                 #
                0731                 ucf[iface] = 0.5*uk
                0732                 vcf[iface] = 0.5*vk
                0733
                0734             uc[t,k,:,:] = faces2mds(ucf)
                0735             vc[t,k,:,:] = faces2mds(vcf)
                0736
                0737     return uc.reshape(ushape), vc.reshape(ushape)
                0738
                0739 def grad(X, dxc=None, dyc=None, hfw=None, hfs=None):
                0740     """
                0741     Compute horizontal gradient of scalar field X on llc grid
                0742
                0743     Call signatures::
                0744
                0745        dXdx, dXdy = div(X, DXC, DYC, HFW, HFS)
                0746        dXdx, dXdy = div(X)
                0747        dXdx, dXdy = div(X, DXC, DYC)
                0748        dXdx, dXdy = div(X, hfw=HFW, hfs=HFS)
                0749
                0750     Parameters
                0751     ----------
                0752     X   : array-like (timelevel,depthlevel,jpoint,ipoint)
                0753           scalar field at c-point
                0754
                0755     dxc : array-like (jpoint,ipoint), optional
                0756           grid spacing in x across u-point, defaults to one
                0757
                0758     dyc : array-like (jpoint,ipoint), optional
                0759           grid spacing in y across v-point, defaults to one
                0760
                0761     hfw : array-like (depthlevel,jpoint,ipoint), optional
                0762           hFac at u-point, defaults to one
                0763
                0764     hfs : array-like (depthlevel,jpoint,ipoint), optional
                0765           hFac at v-point, defaults to one
                0766     """
                0767
                0768     nt, nk, nj, ni =  _getDims(X,X)
                0769     if dxc is None:
                0770         dxc = np.ones((nj,ni))
                0771
                0772     if dyc is None:
                0773         dyc = np.ones((nj,ni))
                0774
                0775     if hfw is None:
                0776         hfw = np.ones((nk,nj,ni))
                0777
                0778     if hfs is None:
                0779         hfs = np.ones((nk,nj,ni))
                0780
                0781     mskw=np.where(hfw>0.,1.,0.).reshape(nk,nj,ni)
                0782     msks=np.where(hfs>0.,1.,0.).reshape(nk,nj,ni)
                0783
                0784     xshape = X.shape
                0785
                0786     X   = X.reshape(nt,nk,nj,ni)
                0787
                0788     dXdx = np.zeros(X.shape)
                0789     dXdy = np.zeros(X.shape)
                0790
95edcfd937 Oliv*0791     rdxc = faces(1./np.where(dxc==0.,np.inf,dxc))
                0792     rdyc = faces(1./np.where(dyc==0.,np.inf,dyc))
616d857f4a Mart*0793     for t in range(nt):
                0794         for k in range(nk):
                0795             xf  = faces(X[t,k,:,:])
                0796             duf = faces(np.zeros((nj,ni)))
                0797             dvf = faces(np.zeros((nj,ni)))
                0798             for iface in range(len(xf)-1):
                0799                 du = (xf[iface] - np.roll(xf[iface],1,axis=-1))*rdxc[iface]
                0800                 dv = (xf[iface] - np.roll(xf[iface],1,axis=-2))*rdyc[iface]
                0801                 # now take care of the connectivity
                0802                 if iface==0:
                0803                     du[:,0] = (xf[0][:,0] - xf[4][-1,::-1])*rdxc[0][:,0]
                0804                     dv[0,:] = 0. # hack
                0805                 if iface==1:
                0806                     du[:,0] = (xf[1][:,0] - xf[0][:,   -1])*rdxc[1][:,0]
                0807                     dv[0,:] = 0. # hack
                0808                 if iface==2:
                0809                     du[:,0] = (xf[2][:,0] - xf[0][-1,::-1])*rdxc[2][:,0]
                0810                     dv[0,:] = (xf[2][0,:] - xf[1][-1,   :])*rdxc[2][0,:]
                0811                 if iface==3:
                0812                     du[:,0] = (xf[3][:,0] - xf[2][:   ,-1])*rdxc[3][:,0]
                0813                     dv[0,:] = (xf[3][0,:] - xf[1][::-1,-1])*rdyc[3][0,:]
                0814                 if iface==4:
                0815                     du[:,0] = (xf[4][:,0] - xf[2][-1,::-1])*rdxc[4][:,0]
                0816                     dv[0,:] = (xf[4][0,:] - xf[3][-1,   :])*rdyc[4][0,:]
                0817
                0818                 duf[iface]=du
                0819                 dvf[iface]=dv
                0820
                0821             dXdx[t,k,:,:] = faces2mds(duf)*mskw[k,:,:]
                0822             dXdy[t,k,:,:] = faces2mds(dvf)*msks[k,:,:]
                0823
                0824
                0825     return dXdx.reshape(xshape), dXdy.reshape(xshape)