def getpsichi(self, ugrid, vgrid, ntrunc=None): """ compute streamfunction and velocity potential on grid given vector wind. @param ugrid: rank 2 or 3 numpy float32 array containing grid of zonal winds. Must have shape (nlat,nlon) or (nlat,nlon,nt), where nt is the number of grids to be transformed. If ugrid is rank 2, nt is assumed to be 1. @param vgrid: rank 2 or 3 numpy float32 array containing grid of meridional winds. Must have shape (nlat,nlon) or (nlat,nlon,nt), where nt is the number of grids to be transformed. Both ugrid and vgrid must have the same shape. @keyword ntrunc: optional spectral truncation limit. (default self.nlat-1) @return: C{B{psigrid, chigrid}} - rank 2 or 3 numpy float32 arrays of gridded streamfunction and velocity potential. Shapes are either (nlat,nlon) or (nlat,nlon,nt). """ # make sure ugrid,vgrid are rank 2 or 3 and same shape. idim = ugrid.ndim shapeu = ugrid.shape shapev = vgrid.shape if ntrunc is None: ntrunc = self.nlat-1 if shapeu != shapev: msg = 'getvrtdivspec input arrays must be same shape!' raise ValueError(msg) if len(shapeu) !=2 and len(shapeu) !=3: msg = 'getvrtdivspec needs rank two or three arrays!' raise ValueError(msg) if shapeu[0] != self.nlat or shapeu[1] != self.nlon: msg = 'getpsichi needs input arrays whose first two dimensions are si%d and %d, got %d and %d' % (self.nlat, self.nlon, ugrid.shape[0], ugrid.shape[1],) raise ValueError(msg) # check ntrunc. if ntrunc < 0 or ntrunc+1 > ugrid.shape[0]: msg = 'ntrunc must be between 0 and %d' % (ugrid.shape[0]-1,) raise ValueError(msg) # compute spectral coeffs of vort, div. vrtspec, divspec = self.getvrtdivspec(ugrid, vgrid, ntrunc) # number of grids to compute. if len(vrtspec.shape) == 1: nt = 1 vrtspec = numpy.reshape(vrtspec, ((ntrunc+1)*(ntrunc+2)/2,1)) divspec = numpy.reshape(divspec, ((ntrunc+1)*(ntrunc+2)/2,1)) else: nt = vrtspec.shape[1] # convert to spectral coeffs of psi, chi. psispec = _spherepack.invlap(vrtspec, self.rsphere) chispec = _spherepack.invlap(divspec, self.rsphere) # inverse transform to grid. psigrid = self.spectogrd(psispec) chigrid = self.spectogrd(chispec) if idim == 2: return numpy.squeeze(psigrid), numpy.squeeze(chigrid) else: return psigrid, chigrid
def getpsichi(self, ugrid, vgrid, ntrunc=None): """ compute streamfunction and velocity potential on grid given vector wind. @param ugrid: rank 2 or 3 numpy float32 array containing grid of zonal winds. Must have shape (nlat,nlon) or (nlat,nlon,nt), where nt is the number of grids to be transformed. If ugrid is rank 2, nt is assumed to be 1. @param vgrid: rank 2 or 3 numpy float32 array containing grid of meridional winds. Must have shape (nlat,nlon) or (nlat,nlon,nt), where nt is the number of grids to be transformed. Both ugrid and vgrid must have the same shape. @keyword ntrunc: optional spectral truncation limit. (default self.nlat-1) @return: C{B{psigrid, chigrid}} - rank 2 or 3 numpy float32 arrays of gridded streamfunction and velocity potential. Shapes are either (nlat,nlon) or (nlat,nlon,nt). """ # make sure ugrid,vgrid are rank 2 or 3 and same shape. idim = ugrid.ndim shapeu = ugrid.shape shapev = vgrid.shape if ntrunc is None: ntrunc = self.nlat - 1 if shapeu != shapev: msg = 'getvrtdivspec input arrays must be same shape!' raise ValueError(msg) if len(shapeu) != 2 and len(shapeu) != 3: msg = 'getvrtdivspec needs rank two or three arrays!' raise ValueError(msg) if shapeu[0] != self.nlat or shapeu[1] != self.nlon: msg = 'getpsichi needs input arrays whose first two dimensions are si%d and %d, got %d and %d' % ( self.nlat, self.nlon, ugrid.shape[0], ugrid.shape[1], ) raise ValueError(msg) # check ntrunc. if ntrunc < 0 or ntrunc + 1 > ugrid.shape[0]: msg = 'ntrunc must be between 0 and %d' % (ugrid.shape[0] - 1, ) raise ValueError(msg) # compute spectral coeffs of vort, div. vrtspec, divspec = self.getvrtdivspec(ugrid, vgrid, ntrunc) # number of grids to compute. if len(vrtspec.shape) == 1: nt = 1 vrtspec = numpy.reshape(vrtspec, ((ntrunc + 1) * (ntrunc + 2) // 2, 1)) divspec = numpy.reshape(divspec, ((ntrunc + 1) * (ntrunc + 2) // 2, 1)) else: nt = vrtspec.shape[1] # convert to spectral coeffs of psi, chi. psispec = _spherepack.invlap(vrtspec, self.rsphere) chispec = _spherepack.invlap(divspec, self.rsphere) # inverse transform to grid. psigrid = self.spectogrd(psispec) chigrid = self.spectogrd(chispec) if idim == 2: return numpy.squeeze(psigrid), numpy.squeeze(chigrid) else: return psigrid, chigrid
def create_sfn_velpot_vortdiv(reanal): ds_U = netCDF4.Dataset( 'reanalysis_clean/{0}.monthly.U.nc'.format(reanal)) ds_VOR = netCDF4.Dataset( 'reanalysis_clean/{0}.monthly.VORTICITY.nc'.format(reanal)) ds_DIV = netCDF4.Dataset( 'reanalysis_clean/{0}.monthly.DIVERGENCE.nc'.format(reanal)) U = ds_U.variables['U'] VOR = ds_VOR.variables['VORTICITY'] DIV = ds_DIV.variables['DIVERGENCE'] ntime, nlev, nlat, nlon = U.shape assert(U.shape == VOR.shape) assert(U.shape == DIV.shape) out_fh_1 = create_output_file(reanal, 'monthly', 'HORIZ_SFN', 'm2 s-1', "Atmospheric Horizontal Streamfunction", ds_U.variables['level'][:], ds_U.variables['latitude'][:], ds_U.variables['longitude'][:], noclobber=True, compress=True) out_fh_2 = create_output_file(reanal, 'monthly', 'HORIZ_VEL_POT', 'm2 s-1', "Atmospheric Horizontal Streamfunction", ds_U.variables['level'][:], ds_U.variables['latitude'][:], ds_U.variables['longitude'][ :], noclobber=True, compress=True) horiz_sfn = out_fh_1.variables['HORIZ_SFN'] vel_pot = out_fh_2.variables['HORIZ_VEL_POT'] time_1 = out_fh_1.variables['time'] time_2 = out_fh_2.variables['time'] lats = ds_U.variables['latitude'][:] lons = ds_U.variables['longitude'][:] # spharm requires grid to go from North to South # and from 0 to 360 positive # Some data can be from -180 to +180 but since a sphere is invariant under # rotations, that should be fine, as long as the direction is eastward. if lats[-1] > lats[0]: lats_increasing = True assert(np.allclose(np.linspace(-90, 90, nlat), lats, 1.E-5, 1.E-5)) else: lats_increasing = False assert(np.allclose(np.linspace(90, -90, nlat), lats, 1.E-5, 1.E-5)) if lons[-1] > lons[0]: lons_increasing = True else: lons_increasing = False # check if the grid is regular and increasing assert(lons[1]-lons[0] > 0.) assert(np.allclose(np.diff(lons), lons[1]-lons[0], 1.0E-4)) psi = np.zeros(U[0].shape, dtype=np.float32) chi = np.zeros(U[0].shape, dtype=np.float32) # this could be written in a much faster way, but it's fast enough for now s = spharm.Spharmt(nlon=nlon, nlat=nlat) for itime in range(ntime): print("{0}/{1}".format(itime+1, ntime)) time_1[itime] = ds_U.variables['time'][itime] time_2[itime] = ds_U.variables['time'][itime] vor_current = VOR[itime] div_current = DIV[itime] if type(vor_current) is np.ma.core.MaskedArray: vor_current = vor_current.filled(0.) if type(div_current) is np.ma.core.MaskedArray: div_current = div_current.filled(0.) if lats_increasing: vor_current = vor_current[:, ::-1, :] div_current = div_current[:, ::-1, :] if not lons_increasing: vor_current = vor_current[:, :, ::-1] div_current = div_current[:, :, ::-1] for ilev in range(nlev): print("\t{0}".format(ilev+1)) vor_spec = s.grdtospec(vor_current[ilev]) div_spec = s.grdtospec(div_current[ilev]) psi_spec = _spherepack.invlap(vor_spec, s.rsphere) chi_spec = _spherepack.invlap(div_spec, s.rsphere) psi_current = np.squeeze(s.spectogrd(psi_spec)) chi_current = np.squeeze(s.spectogrd(chi_spec)) print("\t\t{0:12.4g}\t{1:12.4g}".format( np.mean(psi_current), np.mean(chi_current))) psi[ilev] = psi_current chi[ilev] = chi_current assert(not (np.allclose(psi[ilev], chi[ilev], 1.0E-6, 1.0E-6))) if lats_increasing: psi = psi[:, ::-1, :] chi = chi[:, ::-1, :] if not lons_increasing: psi = psi[:, :, ::-1] chi = chi[:, :, ::-1] horiz_sfn[itime] = psi vel_pot[itime] = chi if itime % 10 == 0: out_fh_1.sync() out_fh_2.sync() out_fh_1.close() out_fh_2.close()
def create_sfn_velpot_vortdiv(reanal): ds_U = netCDF4.Dataset('reanalysis_clean/{0}.monthly.U.nc'.format(reanal)) ds_VOR = netCDF4.Dataset( 'reanalysis_clean/{0}.monthly.VORTICITY.nc'.format(reanal)) ds_DIV = netCDF4.Dataset( 'reanalysis_clean/{0}.monthly.DIVERGENCE.nc'.format(reanal)) U = ds_U.variables['U'] VOR = ds_VOR.variables['VORTICITY'] DIV = ds_DIV.variables['DIVERGENCE'] ntime, nlev, nlat, nlon = U.shape assert (U.shape == VOR.shape) assert (U.shape == DIV.shape) out_fh_1 = create_output_file(reanal, 'monthly', 'HORIZ_SFN', 'm2 s-1', "Atmospheric Horizontal Streamfunction", ds_U.variables['level'][:], ds_U.variables['latitude'][:], ds_U.variables['longitude'][:], noclobber=True, compress=True) out_fh_2 = create_output_file(reanal, 'monthly', 'HORIZ_VEL_POT', 'm2 s-1', "Atmospheric Horizontal Streamfunction", ds_U.variables['level'][:], ds_U.variables['latitude'][:], ds_U.variables['longitude'][:], noclobber=True, compress=True) horiz_sfn = out_fh_1.variables['HORIZ_SFN'] vel_pot = out_fh_2.variables['HORIZ_VEL_POT'] time_1 = out_fh_1.variables['time'] time_2 = out_fh_2.variables['time'] lats = ds_U.variables['latitude'][:] lons = ds_U.variables['longitude'][:] # spharm requires grid to go from North to South # and from 0 to 360 positive # Some data can be from -180 to +180 but since a sphere is invariant under # rotations, that should be fine, as long as the direction is eastward. if lats[-1] > lats[0]: lats_increasing = True assert (np.allclose(np.linspace(-90, 90, nlat), lats, 1.E-5, 1.E-5)) else: lats_increasing = False assert (np.allclose(np.linspace(90, -90, nlat), lats, 1.E-5, 1.E-5)) if lons[-1] > lons[0]: lons_increasing = True else: lons_increasing = False # check if the grid is regular and increasing assert (lons[1] - lons[0] > 0.) assert (np.allclose(np.diff(lons), lons[1] - lons[0], 1.0E-4)) psi = np.zeros(U[0].shape, dtype=np.float32) chi = np.zeros(U[0].shape, dtype=np.float32) # this could be written in a much faster way, but it's fast enough for now s = spharm.Spharmt(nlon=nlon, nlat=nlat) for itime in range(ntime): print("{0}/{1}".format(itime + 1, ntime)) time_1[itime] = ds_U.variables['time'][itime] time_2[itime] = ds_U.variables['time'][itime] vor_current = VOR[itime] div_current = DIV[itime] if type(vor_current) is np.ma.core.MaskedArray: vor_current = vor_current.filled(0.) if type(div_current) is np.ma.core.MaskedArray: div_current = div_current.filled(0.) if lats_increasing: vor_current = vor_current[:, ::-1, :] div_current = div_current[:, ::-1, :] if not lons_increasing: vor_current = vor_current[:, :, ::-1] div_current = div_current[:, :, ::-1] for ilev in range(nlev): print("\t{0}".format(ilev + 1)) vor_spec = s.grdtospec(vor_current[ilev]) div_spec = s.grdtospec(div_current[ilev]) psi_spec = _spherepack.invlap(vor_spec, s.rsphere) chi_spec = _spherepack.invlap(div_spec, s.rsphere) psi_current = np.squeeze(s.spectogrd(psi_spec)) chi_current = np.squeeze(s.spectogrd(chi_spec)) print("\t\t{0:12.4g}\t{1:12.4g}".format(np.mean(psi_current), np.mean(chi_current))) psi[ilev] = psi_current chi[ilev] = chi_current assert (not (np.allclose(psi[ilev], chi[ilev], 1.0E-6, 1.0E-6))) if lats_increasing: psi = psi[:, ::-1, :] chi = chi[:, ::-1, :] if not lons_increasing: psi = psi[:, :, ::-1] chi = chi[:, :, ::-1] horiz_sfn[itime] = psi vel_pot[itime] = chi if itime % 10 == 0: out_fh_1.sync() out_fh_2.sync() out_fh_1.close() out_fh_2.close()