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)