def data2roms(data,grd,sparams,**kargs): ''' Interpolates data to roms 3D grid. The dict data must contain the prognostic variables temp, salt, u, v (3d) and ssh (zeta, 2d), as well as lon, lat (2d), depth (1d) and time/date info: date (data date), date0 (reference date) and time (difference between date and date0). The input data can be provided by load_data. Parameters ---------- data : dict with prognostic variables grd : ROMS netcdf grid file sparams : s-coordinates parameters, theta_s,theta_b, hc and NLevels **kargs: ij : axis for vertical interpolations (*i,j) ij_ind : list of i or j index for vertical interpolation, all by default (ij_ind=False) horizAux : if True, the original data horizontally interpolated is returned and can be used for next data2roms call with this same karg quiet : output messages flag (false by default) ''' ij='j' ij_ind=False horizAux=False quiet=False if 'ij' in kargs.keys(): ij = kargs['ij'] if 'ij_ind' in kargs.keys(): ij_ind = kargs['ij_ind'] if 'horizAux' in kargs.keys(): horizAux = kargs['horizAux'] if 'quiet' in kargs.keys(): quiet = kargs['quiet'] if not quiet: print 'using grid %s' % grd g=roms.Grid(grd) xr,yr,hr,mr=g.vars('r') xu,yu,hu,mu=g.vars('u') xv,yv,hv,mv=g.vars('v') ny,nx=hr.shape nz=sparams[3] NX=data['NX'] NY=data['NY'] NZ=data['NZ'] if not quiet: print 'calc s levels...' sshr=calc.griddata(data['lon'],data['lat'],data['ssh'],xr,yr,extrap=True) Zr = g.s_levels(sparams,sshr,hr,'r') Zu = g.s_levels(sparams,sshr,hr,'u') Zv = g.s_levels(sparams,sshr,hr,'v') # interp horiz: retHorizAux=horizAux is True if horizAux in (True,False): TEMP = np.ma.masked_all((NZ,ny,nx),data['temp'].dtype) SALT = np.ma.masked_all((NZ,ny,nx),data['salt'].dtype) U = np.ma.masked_all((NZ,ny,nx),data['u'].dtype) V = np.ma.masked_all((NZ,ny,nx),data['v'].dtype) if not quiet: print 'horizontal interpolation:' for i in range(NZ): if not quiet and i%10==0: print ' lev %d of %d' % (i,NZ) #import pylab #pylab.figure() #pylab.pcolormesh(data['lon'],data['lat'],data['temp'][i,...]) try: TEMP[i,...] = calc.griddata(data['lon'],data['lat'],data['temp'][i,...],xr,yr,extrap=True) except: pass try: SALT[i,...] = calc.griddata(data['lon'],data['lat'],data['salt'][i,...],xr,yr,extrap=True) except: pass try: U[i,...] = calc.griddata(data['lon'],data['lat'],data['u'][i,...],xr,yr,extrap=True) except: pass try: V[i,...] = calc.griddata(data['lon'],data['lat'],data['v'][i,...],xr,yr,extrap=True) except: pass # rotate U,V: if not quiet: print 'rotating U,V to grid angle' angle=g.use('angle') # rad U,V=calc.rot2d(U,V,angle) U=rt.rho2uvp3d(U,'u') V=rt.rho2uvp3d(V,'v') horizAux={} horizAux['TEMP'] = TEMP horizAux['SALT'] = SALT horizAux['U'] = U horizAux['V'] = V else: TEMP = horizAux['TEMP'] SALT = horizAux['SALT'] U = horizAux['U'] V = horizAux['V'] # interp vert: nxu=nx-1 nyv=ny-1 #> ----------------------------------------------------------------- useInd=not ij_ind is False if ij_ind is False: if ij=='j': ij_ind=range(ny) elif ij=='i': ij_ind=range(nx) else: try: iter(ij_ind) except: ij_ind=[ij_ind] if ij=='j': ny=nyv=len(ij_ind) elif ij=='i': nx=nxu=len(ij_ind) # -----------------------------------------------------------------< Temp = np.zeros((nz,ny ,nx ),data['temp'].dtype) Salt = np.zeros((nz,ny ,nx ),data['salt'].dtype) Uvel = np.zeros((nz,ny ,nxu),data['u'].dtype) Vvel = np.zeros((nz,nyv,nx ),data['v'].dtype) jslice=lambda x,ind: x[:,ind,:] islice=lambda x,ind: x[:,:,ind] ZZr = np.tile(data['depth'],(nx,ny,1)).T ZZu = np.tile(data['depth'],(nxu,ny,1)).T ZZv = np.tile(data['depth'],(nx,nyv,1)).T if not useInd is False: #>------------------------------------------ if ij=='j': slice=jslice sshr=sshr[ij_ind,:] hr =hr[ij_ind,:] elif ij=='i': slice=islice sshr=sshr[:,ij_ind] hr =hr[:,ij_ind] Zr,Zu,Zv,TEMP,SALT,U,V=[slice(k,ij_ind) for k in [Zr,Zu,Zv,TEMP,SALT,U,V]] # -----------------------------------------------------------------< if useInd: # then store distances for a possible bry file dtype=Temp.dtype distr=np.zeros((nz,ny, nx ),dtype) distu=np.zeros((nz,ny, nxu),dtype) distv=np.zeros((nz,nyv,nx ),dtype) if not quiet: print 'vertical interpolation:' if ij=='j': for j in range(ny): if not quiet and (ny<10 or (ny>=10 and j%10==0)): print ' j=%3d of %3d' % (j,ny) ind=ij_ind[j] dr=np.tile(calc.distance(xr[ind,:],yr[ind,:]),(nz,1)) du=np.tile(calc.distance(xu[ind,:],yu[ind,:]),(nz,1)) Dr=np.tile(calc.distance(xr[ind,:],yr[ind,:]),(NZ,1)) Du=np.tile(calc.distance(xu[ind,:],yu[ind,:]),(NZ,1)) if useInd: distr[:,j,:]=dr; distu[:,j,:]=du; Temp[:,j,:] = calc.griddata(Dr,ZZr[:,j,:],TEMP[:,j,:],dr,Zr[:,j,:],extrap=True) Salt[:,j,:] = calc.griddata(Dr,ZZr[:,j,:],SALT[:,j,:],dr,Zr[:,j,:],extrap=True) if 0 and j%10==0: print Dr.shape, ZZr[:,j,:].shape import pylab as pl pl.figure(1) pl.clf() pl.pcolormesh(Dr,ZZr[:,j,:],SALT[:,j,:]) pl.colorbar() clim=pl.gci().get_clim() pl.figure(2) pl.clf() pl.pcolormesh(dr,Zr[:,j,:],Salt[:,j,:]) pl.clim(clim) pl.colorbar() raw_input() Uvel[:,j,:] = calc.griddata(Du,ZZu[:,j,:],U[:,j,:], du,Zu[:,j,:],extrap=True) if j<Vvel.shape[1]: dv=np.tile(calc.distance(xv[ind,:],yv[ind,:]),(nz,1)) Dv=np.tile(calc.distance(xv[ind,:],yv[ind,:]),(NZ,1)) Vvel[:,j,:] = calc.griddata(Dv,ZZv[:,j,:],V[:,j,:], dv,Zv[:,j,:],extrap=True) if useInd: distv[:,j,:]=dv if np.any(np.isnan(Temp[:,j,:])): print 'found nan in temp',j if np.any(np.isnan(Salt[:,j,:])): print 'found nan in salt',j if np.any(np.isnan(Uvel[:,j,:])): print 'found nan in u',j if j<Vvel.shape[1] and np.any(np.isnan(Vvel[:,j,:])): print 'found nan in v',j elif ij=='i': for i in range(nx): if not quiet and (nx<10 or (nx>=10 and i%10==0)): print ' i=%3d of %3d' % (i,nx) ind=ij_ind[i] dr=np.tile(calc.distance(xr[:,ind],yr[:,ind]),(nz,1)) dv=np.tile(calc.distance(xv[:,ind],yv[:,ind]),(nz,1)) Dr=np.tile(calc.distance(xr[:,ind],yr[:,ind]),(NZ,1)) Dv=np.tile(calc.distance(xv[:,ind],yv[:,ind]),(NZ,1)) if useInd: distr[:,:,i]=dr; distv[:,:,i]=dv; Temp[:,:,i] = calc.griddata(Dr,ZZr[:,:,i],TEMP[:,:,i],dr,Zr[:,:,i],extrap=True) Salt[:,:,i] = calc.griddata(Dr,ZZr[:,:,i],SALT[:,:,i],dr,Zr[:,:,i],extrap=True) Vvel[:,:,i] = calc.griddata(Dv,ZZv[:,:,i],V[:,:,i], dv,Zv[:,:,i],extrap=True) if i<Uvel.shape[2]: du=np.tile(calc.distance(xu[:,ind],yu[:,ind]),(nz,1)) Du=np.tile(calc.distance(xu[:,ind],yu[:,ind]),(NZ,1)) Uvel[:,:,i] = calc.griddata(Du,ZZu[:,:,i],U[:,:,i], du,Zu[:,:,i],extrap=True) if useInd: distu[:,:,i]=du # uv bar: if not quiet: print 'calc uvbar' if useInd is False: ubar,vbar=rt.uvbar(Uvel,Vvel,sshr,hr,sparams) else: #>------------------------------------------------------------ sshu=calc.griddata(data['lon'],data['lat'],data['ssh'],xu,yu,extrap=True) sshv=calc.griddata(data['lon'],data['lat'],data['ssh'],xv,yv,extrap=True) if ij=='j': sshu=sshu[ij_ind,:] sshv=sshv[ij_ind,:] hu =hu[ij_ind,:] hv =hv[ij_ind,:] elif ij=='i': sshu=sshu[:,ij_ind] sshv=sshv[:,ij_ind] hu =hu[:,ij_ind] hv =hv[:,ij_ind] ubar=rt.barotropic(Uvel,sshu,hu,sparams) vbar=rt.barotropic(Vvel,sshv,hv,sparams) # -----------------------------------------------------------------< Vars=cb.odict() Vars['temp'] = Temp Vars['salt'] = Salt Vars['u'] = Uvel Vars['v'] = Vvel Vars['zeta'] = sshr Vars['ubar'] = ubar Vars['vbar'] = vbar Vars['date'] = data['date'] if not useInd is False: #>------------------------------------------ Vars['depth'] = Zr Vars['depthu'] = Zu Vars['depthv'] = Zv Vars['dist'] = distr Vars['distu'] = distu Vars['distv'] = distv # -----------------------------------------------------------------< if retHorizAux: return Vars, horizAux else: return Vars
def data2roms(data, grd, sparams, **kargs): ''' Interpolates data to roms 3D grid. The dict data must contain the prognostic variables temp, salt, u, v (3d) and ssh (zeta, 2d), as well as lon, lat (2d), depth (1d) and time/date info: date (data date), date0 (reference date) and time (difference between date and date0). The input data can be provided by load_data. Parameters ---------- data : dict with prognostic variables grd : ROMS netcdf grid file sparams : s-coordinates parameters, theta_s,theta_b, hc and NLevels **kargs: ij : axis for vertical interpolations (*i,j) ij_ind : list of i or j index for vertical interpolation, all by default (ij_ind=False) horizAux : if True, the original data horizontally interpolated is returned and can be used for next data2roms call with this same karg quiet : output messages flag (false by default) proj : projection - False, name or basemap proj - lcc by default if False, horizontal interpolations will use lonxlat instead of distances interp_opts: options for griddata rep_surf: repeat surface level (new upper level) ''' ij = kargs.get('ij', 'j') ij_ind = kargs.get('ij_ind', False) horizAux = kargs.get('horizAux', False) quiet = kargs.get('quiet', False) proj = kargs.get('proj', 'lcc') # lonxlat to distance before # horizontal interpolation interp_opts = kargs.get('interp_opts', {}) rep_surf = kargs.get('rep_surf', True) # create a surface upper level # before interpolation if not quiet: print 'using grid %s' % grd g = roms.Grid(grd) xr, yr, hr, mr = g.vars('r') xu, yu, hu, mu = g.vars('u') xv, yv, hv, mv = g.vars('v') ny, nx = hr.shape nz = sparams[3] if proj: print 'projecting coordinates...' if isinstance(proj, basestring): lonc = (xr.max() + xr.min()) / 2. latc = (yr.max() + yr.min()) / 2. from mpl_toolkits.basemap import Basemap proj = Basemap(projection=proj, width=1, height=1, resolution=None, lon_0=lonc, lat_0=latc, lat_1=latc) xr, yr = proj(xr, yr) xu, yu = proj(xu, yu) xv, yv = proj(xv, yv) dlon, dlat = proj(data['lon'], data['lat']) Rdz = 1 / 100. # distance to depth ratio (300km vs 3000m) distance = lambda x, y: np.append( 0., np.sqrt(np.diff(x)**2 + np.diff(y)**2).cumsum()) else: dlon, dlat = data['lon'], data['lat'] distance = calc.distance # needed for s_levels and for rep_surf! sshr = calc.griddata(dlon, dlat, data['ssh'], xr, yr, extrap=True, **interp_opts) # repeat surface: if rep_surf: # copy data cos dont want to change the original dataset: import copy data = copy.deepcopy(data) for vname in ['temp', 'salt', 'u', 'v', 'depth']: if data[vname].ndim == 1: # depth ! if np.ma.isMA(data[vname]): vstack = np.ma.hstack else: vstack = np.hstack else: if np.ma.isMA(data[vname]): vstack = np.ma.vstack else: vstack = np.vstack if data['depth'][0] > data['depth'][1]: # surf at ind 0 data[vname] = vstack((data[vname][0][np.newaxis], data[vname])) if vname == 'depth': data[vname][0] = sshr.max() else: data[vname] = vstack( (data[vname], data[vname][-1][np.newaxis])) if vname == 'depth': data[vname][-1] = sshr.max() data['NZ'] = data['NZ'] + 1 NX = data['NX'] NY = data['NY'] NZ = data['NZ'] if not quiet: print 'calc s levels...' Zr = g.s_levels(sparams, sshr, hr, 'r') Zu = g.s_levels(sparams, sshr, hr, 'u') Zv = g.s_levels(sparams, sshr, hr, 'v') # interp horiz: retHorizAux = horizAux is True if horizAux in (True, False): TEMP = np.ma.masked_all((NZ, ny, nx), data['temp'].dtype) SALT = np.ma.masked_all((NZ, ny, nx), data['salt'].dtype) U = np.ma.masked_all((NZ, ny, nx), data['u'].dtype) V = np.ma.masked_all((NZ, ny, nx), data['v'].dtype) if not quiet: print 'horizontal interpolation:' for i in range(NZ): if not quiet and i % 10 == 0: print ' lev %d of %d' % (i, NZ) #import pylab #pylab.figure() #pylab.pcolormesh(data['lon'],data['lat'],data['temp'][i,...]) try: TEMP[i, ...] = calc.griddata(dlon, dlat, data['temp'][i, ...], xr, yr, extrap=True, **interp_opts) except: pass try: SALT[i, ...] = calc.griddata(dlon, dlat, data['salt'][i, ...], xr, yr, extrap=True, **interp_opts) except: pass try: U[i, ...] = calc.griddata(dlon, dlat, data['u'][i, ...], xr, yr, extrap=True, **interp_opts) except: pass try: V[i, ...] = calc.griddata(dlon, dlat, data['v'][i, ...], xr, yr, extrap=True, **interp_opts) except: pass # rotate U,V: if not quiet: print 'rotating U,V to grid angle' angle = g.use('angle') # rad U, V = calc.rot2d(U, V, angle) U = rt.rho2uvp3d(U, 'u') V = rt.rho2uvp3d(V, 'v') horizAux = {} horizAux['TEMP'] = TEMP horizAux['SALT'] = SALT horizAux['U'] = U horizAux['V'] = V else: TEMP = horizAux['TEMP'] SALT = horizAux['SALT'] U = horizAux['U'] V = horizAux['V'] # interp vert: nxu = nx - 1 nyv = ny - 1 #> ----------------------------------------------------------------- useInd = not ij_ind is False if ij_ind is False: if ij == 'j': ij_ind = range(ny) elif ij == 'i': ij_ind = range(nx) else: try: iter(ij_ind) except: ij_ind = [ij_ind] if ij == 'j': ny = nyv = len(ij_ind) elif ij == 'i': nx = nxu = len(ij_ind) # -----------------------------------------------------------------< Temp = np.zeros((nz, ny, nx), data['temp'].dtype) Salt = np.zeros((nz, ny, nx), data['salt'].dtype) Uvel = np.zeros((nz, ny, nxu), data['u'].dtype) Vvel = np.zeros((nz, nyv, nx), data['v'].dtype) jslice = lambda x, ind: x[:, ind, :] islice = lambda x, ind: x[:, :, ind] ZZr = np.tile(data['depth'], (nx, ny, 1)).T ZZu = np.tile(data['depth'], (nxu, ny, 1)).T ZZv = np.tile(data['depth'], (nx, nyv, 1)).T if not useInd is False: #>------------------------------------------ if ij == 'j': slice = jslice sshr = sshr[ij_ind, :] hr = hr[ij_ind, :] elif ij == 'i': slice = islice sshr = sshr[:, ij_ind] hr = hr[:, ij_ind] Zr, Zu, Zv, TEMP, SALT, U, V = [ slice(k, ij_ind) for k in [Zr, Zu, Zv, TEMP, SALT, U, V] ] # -----------------------------------------------------------------< if useInd: # then store distances for a possible bry file dtype = Temp.dtype distr = np.zeros((nz, ny, nx), dtype) distu = np.zeros((nz, ny, nxu), dtype) distv = np.zeros((nz, nyv, nx), dtype) if not quiet: print 'vertical interpolation:' if ij == 'j': for j in range(ny): if not quiet and (ny < 10 or (ny >= 10 and j % 10 == 0)): print ' j=%3d of %3d' % (j, ny) ind = ij_ind[j] dr = np.tile(distance(xr[ind, :], yr[ind, :]), (nz, 1)) du = np.tile(distance(xu[ind, :], yu[ind, :]), (nz, 1)) Dr = np.tile(distance(xr[ind, :], yr[ind, :]), (NZ, 1)) Du = np.tile(distance(xu[ind, :], yu[ind, :]), (NZ, 1)) if useInd: distr[:, j, :] = dr distu[:, j, :] = du Temp[:, j, :] = calc.griddata(Rdz * Dr, ZZr[:, j, :], TEMP[:, j, :], Rdz * dr, Zr[:, j, :], extrap=True, **interp_opts) Salt[:, j, :] = calc.griddata(Rdz * Dr, ZZr[:, j, :], SALT[:, j, :], Rdz * dr, Zr[:, j, :], extrap=True, **interp_opts) if 0 and j % 10 == 0: print Dr.shape, ZZr[:, j, :].shape import pylab as pl pl.figure(1) pl.clf() pl.pcolormesh(Dr, ZZr[:, j, :], SALT[:, j, :]) pl.colorbar() clim = pl.gci().get_clim() pl.figure(2) pl.clf() pl.pcolormesh(dr, Zr[:, j, :], Salt[:, j, :]) pl.clim(clim) pl.colorbar() raw_input() Uvel[:, j, :] = calc.griddata(Rdz * Du, ZZu[:, j, :], U[:, j, :], Rdz * du, Zu[:, j, :], extrap=True, **interp_opts) if j < Vvel.shape[1]: dv = np.tile(distance(xv[ind, :], yv[ind, :]), (nz, 1)) Dv = np.tile(distance(xv[ind, :], yv[ind, :]), (NZ, 1)) Vvel[:, j, :] = calc.griddata(Rdz * Dv, ZZv[:, j, :], V[:, j, :], Rdz * dv, Zv[:, j, :], extrap=True, **interp_opts) if useInd: distv[:, j, :] = dv if np.any(np.isnan(Temp[:, j, :])): print 'found nan in temp', j if np.any(np.isnan(Salt[:, j, :])): print 'found nan in salt', j if np.any(np.isnan(Uvel[:, j, :])): print 'found nan in u', j if j < Vvel.shape[1] and np.any(np.isnan(Vvel[:, j, :])): print 'found nan in v', j elif ij == 'i': for i in range(nx): if not quiet and (nx < 10 or (nx >= 10 and i % 10 == 0)): print ' i=%3d of %3d' % (i, nx) ind = ij_ind[i] dr = np.tile(distance(xr[:, ind], yr[:, ind]), (nz, 1)) dv = np.tile(distance(xv[:, ind], yv[:, ind]), (nz, 1)) Dr = np.tile(distance(xr[:, ind], yr[:, ind]), (NZ, 1)) Dv = np.tile(distance(xv[:, ind], yv[:, ind]), (NZ, 1)) if useInd: distr[:, :, i] = dr distv[:, :, i] = dv Temp[:, :, i] = calc.griddata(Rdz * Dr, ZZr[:, :, i], TEMP[:, :, i], Rdz * dr, Zr[:, :, i], extrap=True, **interp_opts) Salt[:, :, i] = calc.griddata(Rdz * Dr, ZZr[:, :, i], SALT[:, :, i], Rdz * dr, Zr[:, :, i], extrap=True, **interp_opts) Vvel[:, :, i] = calc.griddata(Rdz * Dv, ZZv[:, :, i], V[:, :, i], Rdz * dv, Zv[:, :, i], extrap=True, **interp_opts) if i < Uvel.shape[2]: du = np.tile(distance(xu[:, ind], yu[:, ind]), (nz, 1)) Du = np.tile(distance(xu[:, ind], yu[:, ind]), (NZ, 1)) Uvel[:, :, i] = calc.griddata(Rdz * Du, ZZu[:, :, i], U[:, :, i], Rdz * du, Zu[:, :, i], extrap=True, **interp_opts) if useInd: distu[:, :, i] = du # uv bar: if not quiet: print 'calc uvbar' if useInd is False: ubar, vbar = rt.uvbar(Uvel, Vvel, sshr, hr, sparams) else: #>------------------------------------------------------------ sshu = calc.griddata(dlon, dlat, data['ssh'], xu, yu, extrap=True, **interp_opts) sshv = calc.griddata(dlon, dlat, data['ssh'], xv, yv, extrap=True, **interp_opts) if ij == 'j': sshu = sshu[ij_ind, :] sshv = sshv[ij_ind, :] hu = hu[ij_ind, :] hv = hv[ij_ind, :] elif ij == 'i': sshu = sshu[:, ij_ind] sshv = sshv[:, ij_ind] hu = hu[:, ij_ind] hv = hv[:, ij_ind] ubar = rt.barotropic(Uvel, sshu, hu, sparams) vbar = rt.barotropic(Vvel, sshv, hv, sparams) # -----------------------------------------------------------------< Vars = cb.odict() Vars['temp'] = Temp Vars['salt'] = Salt Vars['u'] = Uvel Vars['v'] = Vvel Vars['zeta'] = sshr Vars['ubar'] = ubar Vars['vbar'] = vbar Vars['date'] = data['date'] if not useInd is False: #>------------------------------------------ Vars['depth'] = Zr Vars['depthu'] = Zu Vars['depthv'] = Zv Vars['dist'] = distr Vars['distu'] = distu Vars['distv'] = distv # -----------------------------------------------------------------< if retHorizAux: return Vars, horizAux else: return Vars
def data2roms(data,grd,sparams,**kargs): ''' Interpolates data to roms 3D grid. The dict data must contain the prognostic variables temp, salt, u, v (3d) and ssh (zeta, 2d), as well as lon, lat (2d), depth (1d) and time/date info: date (data date), date0 (reference date) and time (difference between date and date0). The input data can be provided by load_data. Parameters ---------- data : dict with prognostic variables grd : ROMS netcdf grid file sparams : s-coordinates parameters, theta_s,theta_b, hc and NLevels **kargs: ij : axis for vertical interpolations (*i,j) ij_ind : list of i or j index for vertical interpolation, all by default (ij_ind=False) horizAux : if True, the original data horizontally interpolated is returned and can be used for next data2roms call with this same karg quiet : output messages flag (false by default) proj : projection - False, *auto* (use default grid projection) or basemap proj if False, horizontal interpolations will use lonxlat instead of distances interp_opts: options for griddata rep_surf: repeat surface level (new upper level) - If surface variables are present rep_surf must be True and for each variable (ex temp) will be used the surface variable (if present, ex ss_temp) instead of the upper layer. ''' ij = kargs.get('ij','j') ij_ind = kargs.get('ij_ind',False) horizAux = kargs.get('horizAux',False) quiet = kargs.get('quiet',False) proj = kargs.get('proj','auto') # lonxlat to distance before # horizontal interpolation interp_opts = kargs.get('interp_opts',{}) rep_surf = kargs.get('rep_surf',True) # create a surface upper level # before interpolation if not quiet: print('using grid %s' % grd) g=roms.Grid(grd) xr,yr,hr,mr=g.vars('r') xu,yu,hu,mu=g.vars('u') xv,yv,hv,mv=g.vars('v') ny,nx=hr.shape ny0,nx0=hr.shape nz=sparams[3] if proj=='auto': proj=g.get_projection() if proj: print('projecting coordinates...') xr,yr=proj(xr,yr) xu,yu=proj(xu,yu) xv,yv=proj(xv,yv) dlon,dlat=proj(data['lon'],data['lat']) Rdz=1/100. # distance to depth ratio (300km vs 3000m) distance=lambda x,y: np.append(0.,np.sqrt(np.diff(x)**2+np.diff(y)**2).cumsum()) else: dlon,dlat=data['lon'],data['lat'] Rdz=1. distance=calc.distance # needed for s_levels and for rep_surf! sshr=calc.griddata(dlon,dlat,data['ssh'],xr,yr,extrap=True,**interp_opts) # repeat surface: any_ssvar=False if rep_surf: # copy data cos dont want to change the original dataset: import copy data=copy.deepcopy(data) for vname in ['temp','salt','u','v','depth']: if data[vname].ndim==1: # depth ! if np.ma.isMA(data[vname]): vstack=np.ma.hstack else: vstack=np.hstack else: if np.ma.isMA(data[vname]): vstack=np.ma.vstack else: vstack=np.vstack ss_vname='ss_'+vname if ss_vname in data and not quiet: any_ssvar=True print('- using surface variable %s'%ss_vname) if data['depth'][0]>data['depth'][1]: # surf at ind 0 if ss_vname in data: data[vname]=vstack((data[ss_vname][np.newaxis],data[vname])) else: data[vname]=vstack((data[vname][0][np.newaxis],data[vname])) if vname=='depth': data[vname][0]=sshr.max() surf_ind=0 else: if ss_vname in data: data[vname]=vstack((data[vname],data[ss_vname][np.newaxis])) else: data[vname]=vstack((data[vname],data[vname][-1][np.newaxis])) if vname=='depth': data[vname][-1]=sshr.max() surf_ind=-1 data['NZ']=data['NZ']+1 NX=data['NX'] NY=data['NY'] NZ=data['NZ'] if not quiet: print('calc s levels...') Zr = g.s_levels(sparams,sshr,hr,'rr') Zu = g.s_levels(sparams,sshr,hr,'ur') Zv = g.s_levels(sparams,sshr,hr,'vr') # interp horiz: retHorizAux=horizAux is True if horizAux in (True,False): TEMP = np.ma.masked_all((NZ,ny,nx),data['temp'].dtype) SALT = np.ma.masked_all((NZ,ny,nx),data['salt'].dtype) U = np.ma.masked_all((NZ,ny,nx),data['u'].dtype) V = np.ma.masked_all((NZ,ny,nx),data['v'].dtype) if not quiet: print('horizontal interpolation:') for i in range(NZ): if not quiet and i%10==0: print(' lev %d of %d' % (i,NZ)) #import pylab #pylab.figure() #pylab.pcolormesh(data['lon'],data['lat'],data['temp'][i,...]) try: TEMP[i,...] = calc.griddata(dlon,dlat,data['temp'][i,...],xr,yr,extrap=True,**interp_opts) except: pass try: SALT[i,...] = calc.griddata(dlon,dlat,data['salt'][i,...],xr,yr,extrap=True,**interp_opts) except: pass try: U[i,...] = calc.griddata(dlon,dlat,data['u'][i,...],xr,yr,extrap=True,**interp_opts) except: pass try: V[i,...] = calc.griddata(dlon,dlat,data['v'][i,...],xr,yr,extrap=True,**interp_opts) except: pass # rotate U,V: if not quiet: print('rotating U,V to grid angle') U,V=calc.rot2d(U,V,g.angle) # g.angle in rad U=rt.rho2uvp3d(U,'u') V=rt.rho2uvp3d(V,'v') horizAux={} horizAux['TEMP'] = TEMP horizAux['SALT'] = SALT horizAux['U'] = U horizAux['V'] = V else: TEMP = horizAux['TEMP'] SALT = horizAux['SALT'] U = horizAux['U'] V = horizAux['V'] # interp vert: nxu=nx-1 nyv=ny-1 #> ----------------------------------------------------------------- useInd=not ij_ind is False if ij_ind is False: if ij=='j': ij_ind=range(ny) elif ij=='i': ij_ind=range(nx) else: try: iter(ij_ind) except: ij_ind=[ij_ind] if ij=='j': ny=nyv=len(ij_ind) elif ij=='i': nx=nxu=len(ij_ind) # -----------------------------------------------------------------< Temp = np.zeros((nz,ny ,nx ),data['temp'].dtype) Salt = np.zeros((nz,ny ,nx ),data['salt'].dtype) Uvel = np.zeros((nz,ny ,nxu),data['u'].dtype) Vvel = np.zeros((nz,nyv,nx ),data['v'].dtype) jslice=lambda x,ind: x[:,ind,:] islice=lambda x,ind: x[:,:,ind] if any_ssvar: ZZr = np.tile(data['depth'],(nx0,ny0,1)).T ZZu = np.tile(data['depth'],(nx0-1,ny0,1)).T ZZv = np.tile(data['depth'],(nx0,ny0-1,1)).T # replace 1st level by sea surface height! # instead of sshr.max(). This is at least slightly more correct, # and should be no problem for variables without surface counterpart. ZZr[surf_ind]=sshr ZZu[surf_ind]=(sshr[:,1:]+sshr[:,:-1])/2. ZZv[surf_ind]=(sshr[1:]+sshr[:-1])/2. else: ZZr = np.tile(data['depth'],(nx,ny,1)).T ZZu = np.tile(data['depth'],(nxu,ny,1)).T ZZv = np.tile(data['depth'],(nx,nyv,1)).T if not useInd is False: #>------------------------------------------ if ij=='j': slice=jslice sshr=sshr[ij_ind,:] hr =hr[ij_ind,:] elif ij=='i': slice=islice sshr=sshr[:,ij_ind] hr =hr[:,ij_ind] Zr,Zu,Zv,TEMP,SALT,U,V=[slice(k,ij_ind) for k in [Zr,Zu,Zv,TEMP,SALT,U,V]] if any_ssvar: ZZr,ZZu,ZZv=[slice(k,ij_ind) for k in [ZZr,ZZu,ZZv]] # -----------------------------------------------------------------< if useInd: # then store distances for a possible bry file dtype=Temp.dtype distr=np.zeros((nz,ny, nx ),dtype) distu=np.zeros((nz,ny, nxu),dtype) distv=np.zeros((nz,nyv,nx ),dtype) if not quiet: print('vertical interpolation:') if ij=='j': for j in range(ny): if not quiet and (ny<20 or (ny>=20 and j%20==0)): print(' j=%3d of %3d' % (j,ny)) ind=ij_ind[j] dr=np.tile(distance(xr[ind,:],yr[ind,:]),(nz,1)) du=np.tile(distance(xu[ind,:],yu[ind,:]),(nz,1)) Dr=np.tile(distance(xr[ind,:],yr[ind,:]),(NZ,1)) Du=np.tile(distance(xu[ind,:],yu[ind,:]),(NZ,1)) if useInd: distr[:,j,:]=dr; distu[:,j,:]=du; Temp[:,j,:] = calc.griddata(Rdz*Dr,ZZr[:,j,:],TEMP[:,j,:],Rdz*dr,Zr[:,j,:],extrap=True,**interp_opts) Salt[:,j,:] = calc.griddata(Rdz*Dr,ZZr[:,j,:],SALT[:,j,:],Rdz*dr,Zr[:,j,:],extrap=True,**interp_opts) if 0 and j%10==0: print(Dr.shape, ZZr[:,j,:].shape) import pylab as pl pl.figure(1) pl.clf() pl.pcolormesh(Dr,ZZr[:,j,:],TEMP[:,j,:]) pl.colorbar() clim=pl.gci().get_clim() pl.figure(2) pl.clf() pl.pcolormesh(dr,Zr[:,j,:],Temp[:,j,:]) pl.clim(clim) pl.colorbar() try: raw_input() except: input() # python 3 Uvel[:,j,:] = calc.griddata(Rdz*Du,ZZu[:,j,:],U[:,j,:], Rdz*du,Zu[:,j,:],extrap=True,**interp_opts) if j<Vvel.shape[1]: dv=np.tile(distance(xv[ind,:],yv[ind,:]),(nz,1)) Dv=np.tile(distance(xv[ind,:],yv[ind,:]),(NZ,1)) Vvel[:,j,:] = calc.griddata(Rdz*Dv,ZZv[:,j,:],V[:,j,:], Rdz*dv,Zv[:,j,:],extrap=True,**interp_opts) if useInd: distv[:,j,:]=dv if np.any(np.isnan(Temp[:,j,:])): print('found nan in temp',j) if np.any(np.isnan(Salt[:,j,:])): print('found nan in salt',j) if np.any(np.isnan(Uvel[:,j,:])): print('found nan in u',j) if j<Vvel.shape[1] and np.any(np.isnan(Vvel[:,j,:])): print('found nan in v',j) elif ij=='i': for i in range(nx): if not quiet and (nx<20 or (nx>=20 and i%20==0)): print(' i=%3d of %3d' % (i,nx)) ind=ij_ind[i] dr=np.tile(distance(xr[:,ind],yr[:,ind]),(nz,1)) dv=np.tile(distance(xv[:,ind],yv[:,ind]),(nz,1)) Dr=np.tile(distance(xr[:,ind],yr[:,ind]),(NZ,1)) Dv=np.tile(distance(xv[:,ind],yv[:,ind]),(NZ,1)) if useInd: distr[:,:,i]=dr; distv[:,:,i]=dv; Temp[:,:,i] = calc.griddata(Rdz*Dr,ZZr[:,:,i],TEMP[:,:,i],Rdz*dr,Zr[:,:,i],extrap=True,**interp_opts) Salt[:,:,i] = calc.griddata(Rdz*Dr,ZZr[:,:,i],SALT[:,:,i],Rdz*dr,Zr[:,:,i],extrap=True,**interp_opts) Vvel[:,:,i] = calc.griddata(Rdz*Dv,ZZv[:,:,i],V[:,:,i], Rdz*dv,Zv[:,:,i],extrap=True,**interp_opts) if i<Uvel.shape[2]: du=np.tile(distance(xu[:,ind],yu[:,ind]),(nz,1)) Du=np.tile(distance(xu[:,ind],yu[:,ind]),(NZ,1)) Uvel[:,:,i] = calc.griddata(Rdz*Du,ZZu[:,:,i],U[:,:,i], Rdz*du,Zu[:,:,i],extrap=True,**interp_opts) if useInd: distu[:,:,i]=du # uv bar: if not quiet: print('calc uvbar') if useInd is False: ubar,vbar=rt.uvbar(Uvel,Vvel,sshr,hr,sparams) else: #>------------------------------------------------------------ sshu=calc.griddata(dlon,dlat,data['ssh'],xu,yu,extrap=True,**interp_opts) sshv=calc.griddata(dlon,dlat,data['ssh'],xv,yv,extrap=True,**interp_opts) if ij=='j': sshu=sshu[ij_ind,:] sshv=sshv[ij_ind,:] hu =hu[ij_ind,:] hv =hv[ij_ind,:] elif ij=='i': sshu=sshu[:,ij_ind] sshv=sshv[:,ij_ind] hu =hu[:,ij_ind] hv =hv[:,ij_ind] ubar=rt.barotropic(Uvel,sshu,hu,sparams) vbar=rt.barotropic(Vvel,sshv,hv,sparams) # -----------------------------------------------------------------< #Vars=cb.odict() Vars=OrderedDict() Vars['temp'] = Temp Vars['salt'] = Salt Vars['u'] = Uvel Vars['v'] = Vvel Vars['zeta'] = sshr Vars['ubar'] = ubar Vars['vbar'] = vbar Vars['date'] = data['date'] if not useInd is False: #>------------------------------------------ Vars['depth'] = Zr Vars['depthu'] = Zu Vars['depthv'] = Zv Vars['dist'] = distr Vars['distu'] = distu Vars['distv'] = distv # -----------------------------------------------------------------< if retHorizAux: return Vars, horizAux else: return Vars
def data2roms(data, grd, sparams, **kargs): ''' Interpolates data to roms 3D grid. The dict data must contain the prognostic variables temp, salt, u, v (3d) and ssh (zeta, 2d), as well as lon, lat (2d), depth (1d) and time/date info: date (data date), date0 (reference date) and time (difference between date and date0). The input data can be provided by load_data. Parameters ---------- data : dict with prognostic variables grd : ROMS netcdf grid file sparams : s-coordinates parameters, theta_s,theta_b, hc and NLevels **kargs: ij : axis for vertical interpolations (*i,j) ij_ind : list of i or j index for vertical interpolation, all by default (ij_ind=False) horizAux : if True, the original data horizontally interpolated is returned and can be used for next data2roms call with this same karg quiet : output messages flag (false by default) ''' ij = 'j' ij_ind = False horizAux = False quiet = False if 'ij' in kargs.keys(): ij = kargs['ij'] if 'ij_ind' in kargs.keys(): ij_ind = kargs['ij_ind'] if 'horizAux' in kargs.keys(): horizAux = kargs['horizAux'] if 'quiet' in kargs.keys(): quiet = kargs['quiet'] if not quiet: print 'using grid %s' % grd g = roms.Grid(grd) xr, yr, hr, mr = g.vars('r') xu, yu, hu, mu = g.vars('u') xv, yv, hv, mv = g.vars('v') ny, nx = hr.shape nz = sparams[3] NX = data['NX'] NY = data['NY'] NZ = data['NZ'] if not quiet: print 'calc s levels...' sshr = calc.griddata(data['lon'], data['lat'], data['ssh'], xr, yr, extrap=True) Zr = g.s_levels(sparams, sshr, hr, 'r') Zu = g.s_levels(sparams, sshr, hr, 'u') Zv = g.s_levels(sparams, sshr, hr, 'v') # interp horiz: retHorizAux = horizAux is True if horizAux in (True, False): TEMP = np.ma.masked_all((NZ, ny, nx), data['temp'].dtype) SALT = np.ma.masked_all((NZ, ny, nx), data['salt'].dtype) U = np.ma.masked_all((NZ, ny, nx), data['u'].dtype) V = np.ma.masked_all((NZ, ny, nx), data['v'].dtype) if not quiet: print 'horizontal interpolation:' for i in range(NZ): if not quiet and i % 10 == 0: print ' lev %d of %d' % (i, NZ) #import pylab #pylab.figure() #pylab.pcolormesh(data['lon'],data['lat'],data['temp'][i,...]) try: TEMP[i, ...] = calc.griddata(data['lon'], data['lat'], data['temp'][i, ...], xr, yr, extrap=True) except: pass try: SALT[i, ...] = calc.griddata(data['lon'], data['lat'], data['salt'][i, ...], xr, yr, extrap=True) except: pass try: U[i, ...] = calc.griddata(data['lon'], data['lat'], data['u'][i, ...], xr, yr, extrap=True) except: pass try: V[i, ...] = calc.griddata(data['lon'], data['lat'], data['v'][i, ...], xr, yr, extrap=True) except: pass # rotate U,V: if not quiet: print 'rotating U,V to grid angle' angle = g.use('angle') # rad U, V = calc.rot2d(U, V, angle) U = rt.rho2uvp3d(U, 'u') V = rt.rho2uvp3d(V, 'v') horizAux = {} horizAux['TEMP'] = TEMP horizAux['SALT'] = SALT horizAux['U'] = U horizAux['V'] = V else: TEMP = horizAux['TEMP'] SALT = horizAux['SALT'] U = horizAux['U'] V = horizAux['V'] # interp vert: nxu = nx - 1 nyv = ny - 1 #> ----------------------------------------------------------------- useInd = not ij_ind is False if ij_ind is False: if ij == 'j': ij_ind = range(ny) elif ij == 'i': ij_ind = range(nx) else: try: iter(ij_ind) except: ij_ind = [ij_ind] if ij == 'j': ny = nyv = len(ij_ind) elif ij == 'i': nx = nxu = len(ij_ind) # -----------------------------------------------------------------< Temp = np.zeros((nz, ny, nx), data['temp'].dtype) Salt = np.zeros((nz, ny, nx), data['salt'].dtype) Uvel = np.zeros((nz, ny, nxu), data['u'].dtype) Vvel = np.zeros((nz, nyv, nx), data['v'].dtype) jslice = lambda x, ind: x[:, ind, :] islice = lambda x, ind: x[:, :, ind] ZZr = np.tile(data['depth'], (nx, ny, 1)).T ZZu = np.tile(data['depth'], (nxu, ny, 1)).T ZZv = np.tile(data['depth'], (nx, nyv, 1)).T if not useInd is False: #>------------------------------------------ if ij == 'j': slice = jslice sshr = sshr[ij_ind, :] hr = hr[ij_ind, :] elif ij == 'i': slice = islice sshr = sshr[:, ij_ind] hr = hr[:, ij_ind] Zr, Zu, Zv, TEMP, SALT, U, V = [ slice(k, ij_ind) for k in [Zr, Zu, Zv, TEMP, SALT, U, V] ] # -----------------------------------------------------------------< if useInd: # then store distances for a possible bry file dtype = Temp.dtype distr = np.zeros((nz, ny, nx), dtype) distu = np.zeros((nz, ny, nxu), dtype) distv = np.zeros((nz, nyv, nx), dtype) if not quiet: print 'vertical interpolation:' if ij == 'j': for j in range(ny): if not quiet and (ny < 10 or (ny >= 10 and j % 10 == 0)): print ' j=%3d of %3d' % (j, ny) ind = ij_ind[j] dr = np.tile(calc.distance(xr[ind, :], yr[ind, :]), (nz, 1)) du = np.tile(calc.distance(xu[ind, :], yu[ind, :]), (nz, 1)) Dr = np.tile(calc.distance(xr[ind, :], yr[ind, :]), (NZ, 1)) Du = np.tile(calc.distance(xu[ind, :], yu[ind, :]), (NZ, 1)) if useInd: distr[:, j, :] = dr distu[:, j, :] = du Temp[:, j, :] = calc.griddata(Dr, ZZr[:, j, :], TEMP[:, j, :], dr, Zr[:, j, :], extrap=True) Salt[:, j, :] = calc.griddata(Dr, ZZr[:, j, :], SALT[:, j, :], dr, Zr[:, j, :], extrap=True) if 0 and j % 10 == 0: print Dr.shape, ZZr[:, j, :].shape import pylab as pl pl.figure(1) pl.clf() pl.pcolormesh(Dr, ZZr[:, j, :], SALT[:, j, :]) pl.colorbar() clim = pl.gci().get_clim() pl.figure(2) pl.clf() pl.pcolormesh(dr, Zr[:, j, :], Salt[:, j, :]) pl.clim(clim) pl.colorbar() raw_input() Uvel[:, j, :] = calc.griddata(Du, ZZu[:, j, :], U[:, j, :], du, Zu[:, j, :], extrap=True) if j < Vvel.shape[1]: dv = np.tile(calc.distance(xv[ind, :], yv[ind, :]), (nz, 1)) Dv = np.tile(calc.distance(xv[ind, :], yv[ind, :]), (NZ, 1)) Vvel[:, j, :] = calc.griddata(Dv, ZZv[:, j, :], V[:, j, :], dv, Zv[:, j, :], extrap=True) if useInd: distv[:, j, :] = dv if np.any(np.isnan(Temp[:, j, :])): print 'found nan in temp', j if np.any(np.isnan(Salt[:, j, :])): print 'found nan in salt', j if np.any(np.isnan(Uvel[:, j, :])): print 'found nan in u', j if j < Vvel.shape[1] and np.any(np.isnan(Vvel[:, j, :])): print 'found nan in v', j elif ij == 'i': for i in range(nx): if not quiet and (nx < 10 or (nx >= 10 and i % 10 == 0)): print ' i=%3d of %3d' % (i, nx) ind = ij_ind[i] dr = np.tile(calc.distance(xr[:, ind], yr[:, ind]), (nz, 1)) dv = np.tile(calc.distance(xv[:, ind], yv[:, ind]), (nz, 1)) Dr = np.tile(calc.distance(xr[:, ind], yr[:, ind]), (NZ, 1)) Dv = np.tile(calc.distance(xv[:, ind], yv[:, ind]), (NZ, 1)) if useInd: distr[:, :, i] = dr distv[:, :, i] = dv Temp[:, :, i] = calc.griddata(Dr, ZZr[:, :, i], TEMP[:, :, i], dr, Zr[:, :, i], extrap=True) Salt[:, :, i] = calc.griddata(Dr, ZZr[:, :, i], SALT[:, :, i], dr, Zr[:, :, i], extrap=True) Vvel[:, :, i] = calc.griddata(Dv, ZZv[:, :, i], V[:, :, i], dv, Zv[:, :, i], extrap=True) if i < Uvel.shape[2]: du = np.tile(calc.distance(xu[:, ind], yu[:, ind]), (nz, 1)) Du = np.tile(calc.distance(xu[:, ind], yu[:, ind]), (NZ, 1)) Uvel[:, :, i] = calc.griddata(Du, ZZu[:, :, i], U[:, :, i], du, Zu[:, :, i], extrap=True) if useInd: distu[:, :, i] = du # uv bar: if not quiet: print 'calc uvbar' if useInd is False: ubar, vbar = rt.uvbar(Uvel, Vvel, sshr, hr, sparams) else: #>------------------------------------------------------------ sshu = calc.griddata(data['lon'], data['lat'], data['ssh'], xu, yu, extrap=True) sshv = calc.griddata(data['lon'], data['lat'], data['ssh'], xv, yv, extrap=True) if ij == 'j': sshu = sshu[ij_ind, :] sshv = sshv[ij_ind, :] hu = hu[ij_ind, :] hv = hv[ij_ind, :] elif ij == 'i': sshu = sshu[:, ij_ind] sshv = sshv[:, ij_ind] hu = hu[:, ij_ind] hv = hv[:, ij_ind] ubar = rt.barotropic(Uvel, sshu, hu, sparams) vbar = rt.barotropic(Vvel, sshv, hv, sparams) # -----------------------------------------------------------------< Vars = cb.odict() Vars['temp'] = Temp Vars['salt'] = Salt Vars['u'] = Uvel Vars['v'] = Vvel Vars['zeta'] = sshr Vars['ubar'] = ubar Vars['vbar'] = vbar Vars['date'] = data['date'] if not useInd is False: #>------------------------------------------ Vars['depth'] = Zr Vars['depthu'] = Zu Vars['depthv'] = Zv Vars['dist'] = distr Vars['distu'] = distu Vars['distv'] = distv # -----------------------------------------------------------------< if retHorizAux: return Vars, horizAux else: return Vars