def overturning(u, v, dxg, dyg, drf):
ny, nx = dxg.shape
nz = drf.shape[0]
dydr = np.tile(drf, (1, ny, nx)) * np.tile(dyg.reshape((1, ny, nx)),
(nz, 1, 1))
dxdr = np.tile(drf, (1, ny, nx)) * np.tile(dxg.reshape((1, ny, nx)),
(nz, 1, 1))
uf = llc.faces(u * dydr)
vf = llc.faces(v * dxdr)
vf0 = np.sum(vf[0], axis=-1)
vf1 = np.sum(vf[1], axis=-1)
uf3 = np.roll(np.sum(-(uf[3])[:, :, ::-1], axis=-2), 1, axis=1)
uf4 = np.roll(np.sum(-(uf[4])[:, :, ::-1], axis=-2), 1, axis=1)
psi = -np.cumsum((vf0 + vf1 + uf3 + uf4)[::-1, :], axis=0)[::-1, :]
# psi = -(np.cumsum((np.sum(vf[0]*mskv[0],axis=-1)
# + np.sum(vf[1]*mskv[1],axis=-1)
# + np.sum(-(uf[3]*msku[3])[:,:,::-1],axis=-2)
# + np.sum(-(uf[4]*msku[4])[:,:,::-1],axis=-2)
# )[::-1,:],axis=0)
# )[::-1,:]
return psi
def overturning(udr,vdr,zlu,zlv,dxg,dyg):
nl,ny,nx=vdr.shape
ufx=llc.faces(udr*np.tile(dyg.reshape((1,ny,nx)),(nl,1,1)))
vfx=llc.faces(vdr*np.tile(dxg.reshape((1,ny,nx)),(nl,1,1)))
zzu = llc.faces(zlu)
zzv = llc.faces(zlv)
# get rid of Mediterrean Sea (does not help)
vfx[0][:,600:670,110:] = 0.
vfx[0][:,618:627,98:110] = 0.
zzv[0][:,600:670,110:] = 0.
zzv[0][:,618:627,98:110] = 0.
zzu[0][:,600:670,110:] = 0.
zzu[0][:,618:627,98:110] = 0.
# integration over longitude
vf = np.sum(vfx[0]+vfx[1],axis=-1)[:,1:] \
- np.sum(ufx[3]+ufx[4],axis=-2)[:,:0:-1]
# average layer thickness over lontitude
zl = - ( 0.5*(zzu[0]+zzu[1]).mean(axis=-1) \
+ 0.5*(zzv[3]+zzv[4]).mean(axis=-2)[:,::-1] ).cumsum(axis=0)
# add a column of zeros at the southern end
vf = np.hstack((np.zeros((nl,1)),vf))
# integration over depth
psi = np.cumsum(vf,axis=0)
return psi, zl
def atlantic_maskuv(msku, mskv):
# this is pure convenience
fmsku = llc.faces(msku)
fmskv = llc.faces(mskv)
# mask southern ocean first
js = 374
nz, ny, nx = fmsku[0].shape
fmsku[0][:, :js, :] = 0.
fmsku[4][:, :, ny - js:] = 0.
fmskv[0][:, :js, :] = 0.
fmskv[4][:, :, ny - js:] = 0.
# atlantic mask
iw = 176
# this rough and removes part of the norwegian sea
fmsku[0][:, :, iw:] = 0
# mediterranean sea
fmsku[0][:, 599:651, 96:] = 0
fmsku[0][:, 651:662, 125:] = 0
# baltic and north sea
fmsku[0][:, 703:756, 140:] = 0
fmsku[0][:, 756:758, 173:] = 0
fmsku[1][:] = 0.
fmsku[2][:] = 0.
fmsku[3][:] = 0.
# north pacific
fmsku[4][:, :90, :] = 0
# south pacific
fmsku[4][:, :174, 287:] = 0
fmsku[4][:, :133, 263:] = 0
# hudson bay etc
fmsku[4][:, :160, :121] = 0
fmsku[4][:, :189, 26:121] = 0
#
# this rough and removes part of the norwegian sea
fmskv[0][:, :, iw - 1:] = 0
# mediterranean sea
fmskv[0][:, 599:651, 96:] = 0
fmskv[0][:, 651:662, 125:] = 0
# baltic and north sea
fmskv[0][:, 703:756, 140:] = 0
fmskv[0][:, 756:758, 172:] = 0
fmskv[1][:] = 0.
fmskv[2][:] = 0.
fmskv[3][:] = 0.
# north pacific
fmskv[4][:, :90, :] = 0
# south pacific
fmskv[4][:, :174, 287:] = 0
fmskv[4][:, :133, 263:] = 0
fmskv[4][:, 90, 262] = 0.
fmskv[4][:, 99:102, 262] = 0.
# hudson bay etc
fmskv[4][:, :160, :121] = 0
fmskv[4][:, :189, 25:121] = 0
return llc.faces2mds(fmsku), llc.faces2mds(fmskv)
def barostream(u, v, dxg, dyg, drf):
ny, nx = dxg.shape
nz = drf.shape[0]
dydr = np.tile(drf, (1, ny, nx)) * np.tile(dyg.reshape((1, ny, nx)),
(nz, 1, 1))
dxdr = np.tile(drf, (1, ny, nx)) * np.tile(dxg.reshape((1, ny, nx)),
(nz, 1, 1))
ubf = llc.faces(np.sum(u * dydr, axis=-3))
vbf = llc.faces(np.sum(v * dxdr, axis=-3))
ub = np.concatenate((ubf[0], ubf[1], vbf[3][:, ::-1].transpose(),
vbf[4][:, ::-1].transpose()),
axis=-1)
psib = np.cumsum(np.ma.masked_array(ub, ub == 0.), axis=0)
return psib
def atlantic_mask(mskc):
# this is pure convenience
fmsk = llc.faces(mskc)
# mask southern ocean first
js = 374
nz, ny, nx = fmsk[0].shape
fmsk[0][:, :374, 172:] = 0.
fmsk[4][:, :188, 527:] = 0.
# atlantic mask
# this rough and removes part of the norwegian sea
fmsk[0][:, :, 180:] = 0
fmsk[0][:, :380, 175:] = 0
# mediterranean sea
fmsk[0][:, 599:651, 96:] = 0
fmsk[0][:, 651:662, 125:] = 0
# baltic and north sea
fmsk[0][:, 703:756, 140:] = 0
fmsk[0][:, 756:768, 172:] = 0
fmsk[1][:] = 0.
fmsk[2][:] = 0.
fmsk[3][:] = 0.
# north pacific
fmsk[4][:, :90, :] = 0
# south pacific
fmsk[4][:, :174, 287:] = 0
fmsk[4][:, :133, 263:] = 0
fmsk[4][:, :102, 262] = 0
# hudson bay etc
fmsk[4][:, :160, :121] = 0
fmsk[4][:, :189, 26:121] = 0
return llc.faces2mds(fmsk)
def overturningw(w,rac):
nz,ny,nx=w.shape
dxdy=np.tile(rac.reshape((1,ny,nx)),(nz,1,1))
wfx=llc.faces(w*dxdy)
# integration over longitude
wf = ( np.sum(wfx[0]+wfx[1],axis=-1)
+ np.sum(wfx[3]+wfx[4],axis=-2)[:,::-1] )
# integration over latitudes
psiw = np.cumsum(wf,axis=-1)
return psiw
def barostream(u, v, dxg, dyg, drf):
ny, nx = dxg.shape
nz = drf.shape[0]
dydr = np.tile(drf, (1, ny, nx)) * np.tile(dyg.reshape((1, ny, nx)),
(nz, 1, 1))
dxdr = np.tile(drf, (1, ny, nx)) * np.tile(dxg.reshape((1, ny, nx)),
(nz, 1, 1))
ubf = llc.faces(np.sum(u * dydr, axis=-3))
vbf = llc.faces(np.sum(v * dxdr, axis=-3))
ub0 = np.concatenate(
(ubf[0], ubf[1], np.rot90(vbf[3], k=1), np.rot90(vbf[4], k=1)),
axis=-1)
# shift arctic
va = np.roll(vbf[2], -1, axis=0)
va[-1, :] = 0
vb = np.hstack((np.rot90(-va, k=-1), np.zeros(va.shape),
np.rot90(vbf[2], k=1), np.zeros(va.shape)))[:nx // 2, :]
ub = np.concatenate((ub0, vb), axis=-2)
psib = np.cumsum(np.ma.masked_array(ub, ub == 0.), axis=0)
return psib
def move(runoff,moves,mask):#{{{2
'''move runoff from land- to ocean-point
runoff: runoff field (on MITgcm grid)
moves: list of source and destination points
mask: land-ocean-mask
'''
area = rdmds(args.gridpath+'RAC')
new_runoff = np.zeros_like(runoff)
for month in np.arange(0,12): # loop over months
# move mds to list of faces and resort faces
runoff_copy = np.copy(runoff[month,:,:])*area
runoff_f = llc.faces(runoff_copy)
runoff_f = resort_faces(runoff_f)
for i in np.arange(len(moves)): # loop over points to move
# source and destination coordinates for face, y and x
source_face=moves[i][0][0]
source_y = moves[i][0][1]
source_x = moves[i][0][2]
dest_face = moves[i][1][0]
dest_y = moves[i][1][1]
dest_x = moves[i][1][2]
# print '###########################################'
# print 'source_coord',source_face,source_y,source_x
# print 'dest_coord', dest_face,dest_y,dest_x
# print 'dest',runoff_f[dest_face][dest_y,dest_x]
# add runoff of source cell to destination cell
runoff_f[dest_face][dest_y,dest_x] = runoff_f[source_face][source_y,source_x] + runoff_f[dest_face][dest_y,dest_x]
# print 'source', runoff_f[source_face][source_y,source_x]
# print 'new',runoff_f[dest_face][dest_y,dest_x]
# resort faces back and move list to mds array
runoff_f = resort_back_faces(runoff_f)
runoff_m = llc.faces2mds(runoff_f)
runoff_m = runoff_m*mask # delete runoff on land
area[area == 0.0] = 1e-33
runoff_m = runoff_m/area
new_runoff[month,:,:] = runoff_m # write runoff of month into array
return new_runoff
def atlantic_mask(msku, mskv):
# this is pure convenience
fmsku = llc.faces(msku)
fmskv = llc.faces(mskv)
# mask southern ocean first
fmsku[0][:, :123, :] = 0.
fmsku[4][:, :, 148:] = 0.
fmskv[0][:, :123, :] = 0.
fmskv[4][:, :, 147:] = 0.
# atlantic mask
fmsku[0][:, :, 57:] = 0
fmsku[0][:, 203:221, 37:] = 0
fmsku[0][:, 208:210, 33:37] = 0
fmsku[1][:] = 0.
fmsku[2][:] = 0.
fmsku[3][:] = 0.
fmsku[4][:, :30, :] = 0
fmsku[4][:, :51, 91:] = 0
fmsku[4][:, :59, 123:] = 0
fmsku[4][:, :42, 85:] = 0
fmsku[4][:, :54, :40] = 0
#
fmskv[0][:, :, 57:] = 0
fmskv[0][:, 203:221, 37:] = 0
fmskv[0][:, 209, 33:37] = 0
fmskv[1][:] = 0.
fmskv[2][:] = 0.
fmskv[3][:] = 0.
fmskv[4][:, :55, :40] = 0
fmskv[4][:, :31, :] = 0
fmskv[4][:, :36, 84:] = 0
fmskv[4][:, :51, 92:] = 0
fmskv[4][:, :60, 123:] = 0
fmskv[4][:, :45, 86:] = 0
return llc.faces2mds(fmsku), llc.faces2mds(fmskv)
#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
from MITgcmutils import rdmds
from MITgcmutils import llc
from myutils import sq, pcol
d=rdmds('diag3Dm',[120550, 120554])
print d.shape
df=llc.faces(d[:,:,:,:,:])
print len(df)
#print len(df[0])
for k in range(len(df)): print k, df[k].shape
def list_of_movements(runoff,mask,iterations,bathy,extension=10): #{{{2
'''Make list of points on land to move to ocean,
list contains source (on land) and destination (in ocean) coordinate
runoff: runoff field (on MITgcm grid)
mask: land-ocean-mask
iterations: number of surrounding points in which to search for ocean-point
bathy: bathymetry field
extension: number of cells by which to extend current face
in order to search for ocean points on other faces
'''
# array where only points on land are non-zero
mask_inv = np.logical_not(mask)
on_land = np.sum(runoff,axis=0)*mask_inv
# reshape mds arrays to faces
land_f = llc.faces(on_land)
mask_f = llc.faces(mask)
bathy_f = llc.faces(bathy)
# rotate and relocate faces: move (differently shaped) Arctic face to the end of the list (4)
# and rotate faces 2 and 3 in order to match orientation of faces 0 and 1
mask_f = resort_faces(mask_f)
land_f = resort_faces(land_f)
bathy_f = resort_faces(bathy_f)
moves = []
# vector of order of faces, needed to find faces left and right of current face
order = np.concatenate([[4],np.tile(np.array([0,1,2,3]),3),[4]])
for face in np.arange(0,5): # loop over faces
k = 5+face # place in order vector
if face == 4:
k = -1
# extend current face, add (10) surrounding cells at all boundaries
mask_ext = embed_face(mask_f,face,order,k,extension)
bathy_ext = embed_face(bathy_f,face,order,k,extension)
land_ext = embed_face(land_f,face,order,k,extension)
# plt.figure(3)
# pcol(sq(mask_ext))
# plt.show()
# index of points lying on land
index_land = np.transpose(np.nonzero(land_f[face]))
index_land = index_land+extension # account for extention of face
for point in index_land: # loop over points on land
# search in surrounding cells for ocean-points with an increasing radius in each iteration
for m in np.arange(1,iterations+1):
if np.any(mask_ext[point[0]-m:point[0]+(m+1),point[1]-m:point[1]+(m+1)] == 1.): # any point in ocean (mask = 1)
ocean = np.transpose(np.nonzero(mask_ext[point[0]-m:point[0]+(m+1),point[1]-m:point[1]+(m+1)])) # relative indices of all ocean-points to current point
out = point+ocean-m # absolute position of ocean points
ocean_min = ocean[np.argmin(bathy_ext[out[:][:,0],out[:][:,1]])] # find ocean-point with lowest elevation
point = point-extension # account for extension of face (get coordinate without extension)
source = np.insert(point,0,face) # coordinate of point on land [face,y,x]
outlet=point+ocean_min-m # coordinate of ocean-point with lowest elevation [y,x]
# print 'out',outlet
# find coordinate if ocean-point is outside of current face #{{{3
if outlet[1] < 0: # left of face
if face == 4:
outlet[1] = mask_f[order[k-1]].shape[0] + outlet[1] # find correct x-coordinate
outlet = rotate_coordinate(outlet,1,np.rot90(mask_f[order[k-1]])) # find coordinate on back-rotated left face
else:
outlet[1] = mask_f[order[k-1]].shape[1] + outlet[1] # find correct x-coordinate
dest = np.insert(outlet,0,order[k-1]) # coordinate of ocean-point [face,y,x]
elif outlet[1] >= mask_f[face].shape[1]: # right of face
if face == 4:
outlet[1] = outlet[1] - mask_f[face].shape[1]
outlet = rotate_coordinate(outlet,3,np.rot90(mask_f[order[k-3]]))
dest = np.insert(outlet,0,order[k-3])
else:
outlet[1] = outlet[1] - mask_f[face].shape[1]
dest = np.insert(outlet,0,order[k+1])
elif outlet[0] < 0: # below face
outlet[0] = mask_f[order[k-4]].shape[0] + outlet[0] # find correct y-coordinate
dest = np.insert(outlet,0,order[k-4]) # coordinate of ocean-point [face,y,x]
elif outlet[0] >= mask_f[face].shape[0]: # above face
outlet[0] = outlet[0] - mask_f[face].shape[0]
if face == 0:
dest = np.insert(outlet,0,order[0])
elif face == 4:
outlet = rotate_coordinate(outlet,order[k-2],mask_f[order[k-2]])
dest = np.insert(outlet,0,order[k-2])
else:
outlet = rotate_coordinate(outlet,order[k],mask_f[order[0]]) # back-rotate coordinate on Arctic face
dest = np.insert(outlet,0,order[0]) #}}}3
else:
dest = np.insert(outlet,0,order[k]) # coordinate of ocean-point [face,y,x]
move = list([source,dest]) # write source and destionation point as list elements
# print move
break
if m == iterations:
point = point - extension
raise ValueError('ERROR: no wet point found for point {} on face {}. To solve problem the iteration could be increased'.format(point, face))
moves.append(move) # append source and destionation point to list
return moves #}}}2
def mit_div(u, v, hfw, hfs, dxg, dyg, rac):
"""Compute divergence of vector field
Usage: divergence = mit_div(u,v,hfw,hfs,dxg,dyg,rac):
"""
lenu = len(np.shape(u))
nt = 1
nk = 1
if lenu == 2:
nju, niu = np.shape(u)
njv, niv = np.shape(v)
elif lenu == 3:
nk, nju, niu = np.shape(u)
nk, njv, niv = np.shape(v)
elif lenu == 4:
nt, nk, nju, niu = np.shape(u)
nt, nk, njv, niv = np.shape(v)
else:
raise ValueError('Can only handle 2 to 4 dimensions')
# check for 2D arrays
# in this case nk is the time dimension and we need to have
# a separate index for hfw and hfs
if len(np.shape(hfw)) == 2 or np.shape(hfw)[0] != nk:
nt = nk
nk = 1
nj = np.min([nju, njv])
ni = np.min([niu, niv])
cubed_sphere = 0
if niu == 6 * nju:
print('shape(u)[-1]=6*shape(u)[-2]: assuming cubed sphere fields')
cubed_sphere = 1
llcap = 0
if nju == 13 * niu:
print('shape(u)[-2]=13*shape(u)[-1]: assuming lat-lon-cap fields')
llcap = 1
if cubed_sphere:
n = nju
div = np.zeros((nt, nk, 6, n, n), dtype=u.dtype)
u = u.reshape(nt, nk, nju, niu)
v = v.reshape(nt, nk, njv, niv)
hfw = hfw.reshape(nk, nju, niu)
hfs = hfs.reshape(nk, njv, niv)
for t in range(0, nt):
for k in range(0, nk):
uflx = u[t, k, :, :] * hfw[k, :, :] * dyg
vflx = v[t, k, :, :] * hfs[k, :, :] * dxg
uflx = np.transpose(uflx.reshape(n, 6, n), [1, 0, 2])
vflx = np.transpose(vflx.reshape(n, 6, n), [1, 0, 2])
for iface in range(0, 6):
ifp1 = np.remainder(iface + 1, 6)
ifp2 = np.remainder(iface + 2, 6)
if np.remainder(iface + 1, 2):
# odd face
utmp = np.concatenate(
(uflx[iface, :, 1:n], uflx[ifp1, :, 0].reshape(
n, 1)), 1)
tmp = uflx[ifp2, np.arange(n, 0, -1) - 1,
0].reshape(1, n)
vtmp = np.concatenate((vflx[iface, 1:n, :], tmp), 0)
else:
# even face
tmp = vflx[ifp2, 0,
np.arange(n, 0, -1) - 1].reshape(n, 1)
utmp = np.concatenate((uflx[iface, :, 1:n], tmp), 1)
vtmp = np.concatenate(
(vflx[iface, 1:n, :], vflx[ifp1, 0, :].reshape(
1, n)), 0)
ij = np.arange(0, n) + n * iface
# recip_rac = 1./(rac(ij,:).*hf(ij,:,k));
recip_rac = 1. / rac[:, ij]
du = utmp - uflx[iface, :, :]
dv = vtmp - vflx[iface, :, :]
div[t, k, iface, :, :] = (du + dv) * recip_rac
div = np.transpose(div, [0, 1, 3, 2, 4]).reshape(nt, nk, n, n * 6)
elif llcap:
u = u.reshape(nt, nk, nju, niu)
v = v.reshape(nt, nk, njv, niv)
hfw = hfw.reshape(nk, nju, niu)
hfs = hfs.reshape(nk, njv, niv)
div = np.zeros(u.shape)
from MITgcmutils.llc import faces, faces2mds
racf = faces(rac)
for t in range(0, nt):
for k in range(0, nk):
uflx = faces(u[t, k, :, :] * hfw[k, :, :] * dyg)
vflx = faces(v[t, k, :, :] * hfs[k, :, :] * dxg)
divf = faces(np.zeros((nju, niu)))
for iface in range(len(uflx) - 1):
du = np.roll(uflx[iface], -1, axis=-1) - uflx[iface]
dv = np.roll(vflx[iface], -1, axis=-2) - vflx[iface]
# now take care of the connectivity
if iface == 0:
du[:, -1] = uflx[1][:, 0] - uflx[iface][:, -1]
dv[-1, :] = uflx[2][::-1, 0] - vflx[iface][-1, :]
if iface == 1:
du[:, -1] = vflx[3][0, ::-1] - uflx[iface][:, -1]
dv[-1, :] = vflx[2][0, :] - vflx[iface][-1, :]
if iface == 2:
du[:, -1] = uflx[3][:, 0] - uflx[iface][:, -1]
dv[-1, :] = uflx[4][::-1, 0] - vflx[iface][-1, :]
if iface == 3:
du[:, -1] = 0. # hack
dv[-1, :] = vflx[4][0, :] - vflx[iface][-1, :]
if iface == 4:
du[:, -1] = 0. # hack
dv[-1, :] = uflx[0][::-1, 0] - vflx[iface][-1, :]
# putting it all together
divf[iface] = (du + dv) / racf[iface]
div[t, k, :, :] = faces2mds(divf)
else:
div = np.zeros((nt, nk, nj, ni), dtype=u.dtype)
u = u.reshape(nt, nk, nju, niu)
v = v.reshape(nt, nk, njv, niv)
hfw = hfw.reshape(nk, nju, niu)
hfs = hfs.reshape(nk, njv, niv)
for t in range(0, nt):
for k in range(0, nk):
uflx = u[t, k, :, :] * hfw[k, :, :] * dyg
vflx = v[t, k, :, :] * hfs[k, :, :] * dxg
if niu > niv:
du = uflx[:, 1:] - uflx[:, :-2]
else:
du = np.roll(uflx, 1, 1) - uflx
if njv > nju:
dv = vflx[1:, :] - vflx[:-2, :]
else:
dv = np.roll(vflx, 1, 1) - vflx
div[t, k, :, :] = (du + dv) / rac
# remove singleton dimensions again
return np.squeeze(div)
