def spht(field, ntrunc=None, gridtype='gaussian'): """Transform a field on lat-lon grid to spherical harmonics. Currently only works for: - fields defined on the lat-lon points, not latb-lonb. - triangular truncation. Returns an xarray.DataArray """ nlat, nlon = len(field.lat), len(field.lon) grid = spharm.Spharmt(nlon, nlat, gridtype=gridtype) if ntrunc is None: ntrunc = nlat - 1 other_dims = [d for d in field.dims if d not in ('lat', 'lon')] # need the field in N-S, E-W form, (lat, lon, other coords) vfield = (field.sel(lat=sorted(cbbt.lat, reverse=True), lon=sorted(cbbt.lon)).transpose( 'lat', 'lon', *other_dims)) # calculate the spectral coefficients, given the truncation level coeffs = grid.grdtospec(vfield.values, ntrunc) m, n = spharm.getspecindx(ntrunc) # put the coefficients into a grid, half of which will be empty # due to the triangular trunctation cc = np.zeros((ntrunc + 1, ntrunc + 1) + tuple(coeffs.shape[1:]), dtype=np.complex128) cc[m, n] = coeffs coords = [('m', np.arange(0, ntrunc + 1)), ('n', np.arange(0, ntrunc + 1))] for d in other_dims: coords.append((d, field.coords[d])) return xr.DataArray(data=cc, coords=coords, name=field.name)
def sph_filter(field, l_cut, gridtype='gaussian'): """Remove all waves above a specific spherical wavenumber l_cut""" nlat, nlon = len(field.lat), len(field.lon) grid = spharm.Spharmt(nlon, nlat, gridtype=gridtype) other_dims = [d for d in field.dims if d not in ('lat', 'lon')] # need the field in N-S, E-W form, (lat, lon, other coords) vfield = (field.sel(lat=sorted(cbbt.lat, reverse=True), lon=sorted(cbbt.lon)).transpose( 'lat', 'lon', *other_dims)) ntrunc = nlat - 1 m, n = spharm.getspecindx(ntrunc) l = m + n # transform to spherical modes, eliminate high wavenumbers and transform back coeffs = grid.grdtospec(vfield.values, ntrunc) mask = l > l_cut coeffs[mask] = 0 values = grid.spectogrd(coeffs) # copy the original field and comply to the same coordinate ordering nfield = vfield.copy() nfield.values = values return nfield.transpose(*field.dims)
def __init__(self, rsphere=EARTH_RADIUS, omega=EARTH_OMEGA, resolution=2.5, ntrunc=None, legfunc="stored"): """Grid constructor. - The radius and angular velocity of the planet are given by `rsphere` in m and `omega` in 1/s, respectively (the default values correspond to those of the Earth). - The grid resolution is uniform and specified with the `resolution` parameter in degrees. - Uses spherical harmonics for some operations utilizing the `spharm`/`pyspharm` package. By default (`ntrunc=None`, `legfunc="stored"`), the spherical harmonics are truncated after (number of latitudes - 1) functions and precomputed. Consult the documentation of `spharm.Spharmt` for more information on these parameters. """ # Planet parameters (radius, angular velocity) self.rsphere = rsphere self.omega = omega # Setup longitude and latitude grid. Make sure longitude 0° is not # repeated and that both latitudes 90° and -90° exist. Latitudes start # at the North pole, which is the convention used by pyspharm. self.nlon = int(360. / resolution) self.nlat = int(180. / resolution) + 1 if self.nlat % 2 == 0: raise ValueError( "Number of latitudes must be odd but is {}".format(self.nlat)) self.lons = np.linspace(0., 360., self.nlon, endpoint=False) self.lats = np.linspace(90., -90., self.nlat, endpoint=True) self.lon, self.lat = np.meshgrid(self.lons, self.lats) # Grid spacing in degrees self.dlon = resolution self.dlat = -resolution # Spherical coordinate grid (for use with trigonometric functions) self.lams = np.deg2rad(self.lons) self.phis = np.deg2rad(self.lats) self.lam, self.phi = np.meshgrid(self.lams, self.phis) # Grid spacing in radians self.dlam = np.deg2rad(self.dlon) self.dphi = np.deg2rad(self.dlat) # Precompute Coriolis field self.fcor = self.coriolis(self.lat) # Spherical harmonic transform object self._spharm = spharm.Spharmt(self.nlon, self.nlat, rsphere=rsphere, gridtype="regular", legfunc=legfunc) self._ntrunc = (self.nlat - 1) if ntrunc is None else ntrunc _, self.specindxn = spharm.getspecindx(self._ntrunc) # Eigenvalues of the horizontal Laplacian for each spherical harmonic. # Force use of complex64 datatype (= 2 * float32) because spharm will # cast to float32 components anyway but the multiplication with the # python scalars results in float64 values. self.laplacian_eigenvalues = (self.specindxn * (1. + self.specindxn) / rsphere / rsphere).astype( np.complex64, casting="same_kind")
def test_case(): """ Runs an example case: extratropical zonal jets with superimposed sinusoidal NH vorticity perturbations and a gaussian vorticity tendency forcing. """ from time import time start = time() # 1) LET'S CREATE SOME INITIAL CONDITIONS lon_list = list(np.arange(0, 360.1, 2.5)) lat_list = list(np.arange(-87.5, 88, 2.5))[::-1] lamb, theta = np.meshgrid([x * np.pi / 180. for x in lon_list], [x * np.pi / 180. for x in lat_list]) # Mean state: zonal extratropical jets ubar = 25 * np.cos(theta) - 30 * np.cos(theta)**3 + 300 * np.sin( theta)**2 * np.cos(theta)**6 vbar = np.zeros(np.shape(ubar)) # Initial perturbation: sinusoidal vorticity perturbations A = 1.5 * 8e-5 # vorticity perturbation amplitude m = 4 # zonal wavenumber theta0 = np.deg2rad(45) # center lat = 45 N thetaW = np.deg2rad(15) vort_pert = 0.5 * A * np.cos(theta) * np.exp(-( (theta - theta0) / thetaW)**2) * np.cos(m * lamb) # Get U' and V' from this vorticity perturbation s = spharm.Spharmt(len(lon_list), len(lat_list), gridtype='regular', legfunc='computed', rsphere=NL.Re) uprime, vprime = s.getuv(s.grdtospec(vort_pert), np.zeros(np.shape(s.grdtospec(vort_pert)))) # Full initial conditions dictionary: ics = { 'u_bar': ubar, 'v_bar': vbar, 'u_prime': uprime, 'v_prime': vprime, 'lons': lon_list, 'lats': lat_list, 'start_time': datetime(2017, 1, 1, 0) } # 2) LET'S ALSO FEED IN A GAUSSIAN NH RWS FORCING amplitude = 10e-10 # s^-2 forcing = np.zeros(np.shape(ubar)) x, y = np.meshgrid(np.linspace(-1, 1, 10), np.linspace(-1, 1, 10)) d = np.sqrt(x * x + y * y) sigma, mu = 0.5, 0.0 g = np.exp(-((d - mu)**2 / (2.0 * sigma**2))) # GAUSSIAN CURVE lat_i = np.where(np.array(lat_list) == 35.)[0][0] lon_i = np.where(np.array(lon_list) == 160.)[0][0] forcing[lat_i:lat_i + 10, lon_i:lon_i + 10] = g * amplitude # 3) INTEGRATE! model = Model(ics, forcing=forcing) model.integrate() print('TOTAL INTEGRATION TIME: {:.02f} minutes'.format( (time() - start) / 60.))
def super_rotation(linearized=False, idate=None, ntrunc=42): """ Creates ICs for an idealized case: "zonal flow corresponding to a super-rotation of the atmosphere with a maximum value of ~15.4 m/s on the equator." (Sardeshmukh and Hoskins 1988) Args ---- linearized : bool True if this is a linearized model. idate : datetime Forecast initialization date ntrunc : int Triangular trunction (e.g., 42 for T42). Returns ------- ics : dict Model initial state (a dictionary of DataArrays) """ # Set the init. date if not provided if idate is None: idate = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0) # Create a gaussian lat/lon grid if ntrunc not in TRUNC_NLATS.keys(): raise ValueError( 'Truncation T{} is not in the dictionary TRUNC_NLATS'.format( ntrunc)) nlat = TRUNC_NLATS[ntrunc] lons, lats = gaussian_latlon_grid(nlat) theta = lats * np.pi / 180. lamb = lons * np.pi / 180. # Mean state: zonal extratropical jets ubar = Re * Omega * np.cos(theta) / 30. vbar = np.zeros(ubar.shape) # Get the mean state vorticity from ubar and vbar s = spharm.Spharmt(lamb.shape[1], lamb.shape[0], gridtype='gaussian', rsphere=Re, legfunc='computed') vortb_spec, _ = s.getvrtdivspec(ubar, vbar) vort_bar = s.spectogrd(vortb_spec) vort_prime = np.zeros(vort_bar.shape) # Generate the state return _generate_state(idate, lats, lons, vort_bar, vort_prime, linearized=linearized)
def __init__(self, ics, forcing=None): """ Initializes the model. Requires: ics -----> Dictionary of linearized fields and space/time dimensions keys: u_bar, v_bar, u_prime, v_prime, lats, lons, start_time forcing -> a 2D array (same shape as model fields) containing a vorticity tendency [s^-2] to be imposed at each integration time step """ # 1) STORE SPACE/TIME VARIABLES (DIMENSIONS) # Get the latitudes and longitudes (as lists) self.lats = ics['lats'] self.lons = ics['lons'] self.start_time = ics['start_time'] # datetime self.curtime = self.start_time # 2) GENERATE/STORE NONDIVERGENT INITIAL STATE # Set up the spherical harmonic transform object self.s = spharm.Spharmt(self.nlons(), self.nlats(), rsphere=NL.Re, gridtype='regular', legfunc='computed') # Truncation for the spherical transformation if NL.M is None: self.ntrunc = self.nlats() else: self.ntrunc = NL.M # Use the object to get the initial conditions # First convert to vorticity using spharm object vortb_spec, div_spec = self.s.getvrtdivspec(ics['u_bar'], ics['v_bar']) vortp_spec, div_spec = self.s.getvrtdivspec(ics['u_prime'], ics['v_prime']) div_spec = np.zeros( vortb_spec.shape) # Only want NON-DIVERGENT part of wind # Re-convert this to u-v winds to get the non-divergent component # of the wind field self.ub, self.vb = self.s.getuv(vortb_spec, div_spec) # MEAN WINDS self.up, self.vp = self.s.getuv(vortp_spec, div_spec) # PERTURBATION WINDS # Use these winds to get the streamfunction (psi) and # velocity potential (chi) self.psib, chi = self.s.getpsichi(self.ub, self.vb) # MEAN STREAMFUNCTION self.psip, chi = self.s.getpsichi( self.up, self.vp) # PERTURBATION STREAMFUNCTION # Convert the spectral vorticity to grid self.vort_bar = self.s.spectogrd(vortb_spec) # MEAN RELATIVE VORTICITY self.vortp = self.s.spectogrd( vortp_spec) # PERTURBATION RELATIVE VORTICITY # 3) STORE A COUPLE MORE VARIABLES # Map projections for plotting self.bmaps = create_basemaps(self.lons, self.lats) # Get the vorticity tendency forcing (if any) for integration self.forcing = forcing
def make_spharmt(self): try: self.spharmt = self.vectorwind.s except AttributeError: self.spharmt = spharm.Spharmt(self.n_lon, self.n_lat, rsphere=self._rsphere, gridtype=self._gridtype, legfunc=self._legfunc)
def regrid(lon_in,lat_in,data,nlon_o,nlat_o): ''' Given data(nlat,nlon), regrid the data into (nlat_o,nlon_o) using the spharm.regrid routine If the input data is not on a complete sphere, copy the input onto a complete sphere before regridding Input: lon_in - the longitude the input data is on lat_in - the latitude the input data is on data - numpy.ndarray nlon_o - integer nlat_o - integer Return: numpy.ndarray of shape (nlat_o,nlon_o) ''' if not _SPHARM_INSTALLED: raise _import_error assert data.ndim <= 3 # determine if the domain is a complete sphere dlon = numpy.diff(lon_in[0:2]) AllLon = dlon*(len(lon_in)+1) > 360. dlat = numpy.diff(lat_in[0:2]) AllLat = dlat*(len(lat_in)+1) > 180. sliceobj = [slice(None),]*data.ndim if AllLon and AllLat: # The data covers the whole sphere inputdata = data nlat_i = len(lat_in) nlon_i = len(lon_in) lat = lat_in lon = lon_in else: inputdata,lon,lat = regional2global(data,lon_in,lat_in) nlat_i = len(lat) nlon_i = len(lon) gridin = spharm.Spharmt(nlon_i,nlat_i) gridout = spharm.Spharmt(nlon_o,nlat_o) return spharm.regrid(gridin,gridout,inputdata)
def rws_from_tropical_divergence(state, center=(0., 145.), amp=6e-6, width=12): # Get desired state variables lats = state['latitude'].values lons = state['longitude'].values vort_bar = state['base_atmosphere_relative_vorticity'].values s = spharm.Spharmt(lats.shape[1], lons.shape[0], gridtype='regular', rsphere=Re) vortb_spec = s.grdtospec(vort_bar) ubar, vbar = s.getuv(vortb_spec, np.zeros(vortb_spec.shape)) divergence = gaussian_blob_2d(lats, lons, center, width, amp) # Calculate the Rossby Wave Source # Term 1 zetabar_spec, _ = s.getvrtdivspec(ubar, vbar) zetabar = s.spectogrd(zetabar_spec) + 2 * Omega * np.sin(np.deg2rad(lats)) term1 = -zetabar * divergence # Term 2 uchi, vchi = s.getuv(np.zeros(zetabar_spec.shape), s.grdtospec(divergence)) dzeta_dx, dzeta_dy = s.getgrad(s.grdtospec(zetabar)) term2 = - uchi * dzeta_dx - vchi * dzeta_dy rws = term1 + term2 return rws, lats, lons
def compute_del4vort(vort, ntrunc): # Compute del^4(vorticity) with spherical harmonics # Approximating del^4 as: d4_dx4 + d4_dy4 + 2 * (d2_dx2 * d2_dy2) s = spharm.Spharmt(vort.shape[1], vort.shape[0], rsphere=6378100., gridtype='gaussian', legfunc='computed') vspec = s.grdtospec(vort, ntrunc=ntrunc) # First order dvort_dx, dvort_dy = s.getgrad(vspec) # Second order d2vort_dx2, _ = s.getgrad(s.grdtospec(dvort_dx, ntrunc=ntrunc)) _, d2vort_dy2 = s.getgrad(s.grdtospec(dvort_dy, ntrunc=ntrunc)) # Fourth order d4vort_dx4, _ = s.getgrad( s.grdtospec(s.getgrad(s.grdtospec(d2vort_dx2, ntrunc=ntrunc))[0], ntrunc=ntrunc)) _, d4vort_dy4 = s.getgrad( s.grdtospec(s.getgrad(s.grdtospec(d2vort_dy2, ntrunc=ntrunc))[1], ntrunc=ntrunc)) # Put it all together to approximate del^4 del4vort = d4vort_dx4 + d4vort_dy4 + (2 * d2vort_dx2 * d2vort_dy2) return del4vort
def __init__(self, earth, ice, grid=None, topo=None): """ Compute glacial isostatic adjustment on a globe. Paramaters ---------- earth : <giapy.code.earth_tools.earthSpherical.SphericalEarth> ice : <giapy.code.icehistory.IceHistory / PersistentIceHistory> grid : <giapy.code.map_tools.GridObject> topo : numpy.ndarray Methods ------- performConvolution """ self.earth = earth self.ice = ice self.nlat, self.nlon = ice.shape # The grid used is a cylindrical projection (equispaced lat/lon grid # unless otherwise specified) if grid is not None: self.grid = grid else: self.grid = GridObject(mapparam={'projection': 'cyl'}, shape=ice.shape) self.topo = topo # Precompute and store harmonic transform coefficients, for # computational efficiency, but at a memory cost. self.harmTrans = spharm.Spharmt(self.nlon, self.nlat, legfunc='stored')
grbs = pygrib.open('../sampledata/spherical_pressure_level.grib1') g = grbs[1] fld = g.values # ECMWF normalizes the spherical harmonic coeffs differently than NCEP. # (m=0,n=0 is global mean, instead of sqrt(2)/2 times global mean) fld = 2. * fld / np.sqrt(2.) fldr = fld[0::2] fldi = fld[1::2] fld = np.zeros(fldr.shape, 'F') fld.real = fldr fld.imag = fldi nlons = 360 nlats = 181 s = spharm.Spharmt(nlons, nlats) data = s.spectogrd(fld) lons = (360. / nlons) * np.arange(nlons) lats = 90. - (180. / (nlats - 1)) * np.arange(nlats) lons, lats = np.meshgrid(lons, lats) # stack grids side-by-side (in longitiudinal direction), so # any range of longitudes (between -360 and 360) may be plotted on a world map. lons = np.concatenate((lons - 360, lons), 1) lats = np.concatenate((lats, lats), 1) data = np.concatenate((data, data), 1) # setup miller cylindrical map projection. m = Basemap(llcrnrlon=-180.,llcrnrlat=-90,urcrnrlon=180.,urcrnrlat=90.,\ resolution='l',area_thresh=10000.,projection='mill') x, y = m(lons, lats) CS = m.contourf(x, y, data, 15, cmap=plt.cm.jet) ax = plt.gca()
import cartopy.crs as ccrs import cartopy.util import spharm import numpy as np import cartopy import cartopy.crs as ccrs import cartopy.util import spharm import numpy as np field = np.random.normal(0, 1, (64, 128)) x = spharm.Spharmt(128, 64, rsphere=6.4e6, gridtype='gaussian', legfunc='computed') spec = x.grdtospec(field) field_2 = x.spectogrd(spec) #plt.figure(figsize=(10,10)) fig, ax = plt.subplots(2, 1, figsize=(10, 10)) ax[0].imshow(field) ax[1].imshow(field_2) #plt.colorbar() plt.figure(figsize=(10, 10)) plt.imshow(field_2)
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 test_case(): """ Runs an example case: extratropical zonal jets with superimposed sinusoidal NH vorticity perturbations and a gaussian vorticity tendency forcing. """ from time import time start = time() # 1) LET'S CREATE SOME INITIAL CONDITIONS lons = np.arange(0, 360.1, 2.5) lats = np.arange(-87.5, 88, 2.5)[::-1] lamb, theta = np.meshgrid(lons * np.pi / 180., lats * np.pi / 180.) # Mean state: zonal extratropical jets ubar = NL.mag * np.cos(theta) - 30 * np.cos(theta)**3 + 300 * np.sin( theta)**2 * np.cos(theta)**6 vbar = np.zeros(np.shape(ubar)) # Initial perturbation: sinusoidal vorticity perturbations theta0 = np.deg2rad(NL.pert_center_lat) # center lat = 45 N thetaW = np.deg2rad(NL.pert_width) #15 vort_pert = 0.5 * NL.A * np.cos(theta) * np.exp(-( (theta - theta0) / thetaW)**2) * np.cos(NL.m * lamb) # Get U' and V' from this vorticity perturbation s = spharm.Spharmt(len(lons), len(lats), gridtype='regular', legfunc='computed', rsphere=NL.Re) uprime, vprime = s.getuv(s.grdtospec(vort_pert), np.zeros(np.shape(s.grdtospec(vort_pert)))) # Full initial conditions dictionary: ics = { 'u_bar': ubar, 'v_bar': vbar, 'u_prime': uprime, 'v_prime': vprime, 'lons': lons, 'lats': lats, 'start_time': datetime(2017, 1, 1, 0) } # 2) LET'S ALSO FEED IN A GAUSSIAN NH RWS FORCING (CAN BE USED TO CREATE SYNTEHTIC HURRICANES) amplitude = NL.forcing_amp # s^-2 forcing = np.zeros(np.shape(ubar)) x, y = np.meshgrid(np.linspace(-1, 1, 10), np.linspace(-1, 1, 10)) d = np.sqrt(x * x + y * y) sigma, mu = 0.5, 0.0 g = np.exp(-((d - mu)**2 / (2.0 * sigma**2))) # GAUSSIAN CURVE source_lat = NL.forcing_lat source_lon = NL.forcing_lon # The location of the forcing lat_i = np.where(lats == source_lat)[0][0] lon_i = np.where(lons == source_lon)[0][0] if NL.use_forcing == True: forcing[lat_i:lat_i + 10, lon_i:lon_i + 10] = g * amplitude else: forcing[:, :] = 0 # 3) INTEGRATE! model = Model(ics, forcing=forcing) model.integrate() print('TOTAL INTEGRATION TIME: {:.02f} minutes'.format( (time() - start) / 60.)) plt.plot((NL.dt * np.arange(len(model.tot_ke))) / 3600., (1 - model.tot_ke / model.expected_ke) * 100) #plt.plot(np.arange(len(model.tot_ke)), model.tot_ke, 'o-') plt.title('Model Kinetic Energy Error vs. Time Step') plt.xlabel("Model Time [hr]") plt.ylabel(u'Model Kinetic Energy Error [%]') print("Expected Kinetic Energy [m^2/s^2]:", model.expected_ke) plt.savefig(NL.figdir + '/model_ke_ts.png', bbox_inches='tight')
def from_u_and_v_winds(lats, lons, ubar, vbar, uprime=None, vprime=None, interp=True, linearized=False, ntrunc=None, idate=None): """ Creates ICs from numpy arrays describing the intial wind field. Args ---- lats : numpy array 1D array of global latitudes (in degrees) lons : numpy array 1D array of global longitudes (in degrees) ubar : numpy array 2D array (nlat, nlon) of mean state zonal winds vbar : numpy array 2D array (nlat, nlon) of mean state meridional winds uprime : numpy array 2D array (nlat, nlon) of perturbation zonal winds vprime : numpy array 2D array (nlat, nlon) of perturbation meridional winds interp : bool If True, fields will be interpolated to a gaussian grid. linearized : bool True if this is a linearized model. ntrunc : int Triangular trunction (e.g., 42 for T42). idate : datetime Foreast initialization date Returns ------- ics : dict Model initial state (a dictionary of DataArrays) """ # Set the init. date if not provided if idate is None: idate = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0) # If only ubar and vbar are provided, then forecast will be nonlinear if uprime is None or vprime is None: uprime = np.zeros(ubar.shape) vprime = np.zeros(vbar.shape) linearized = False # Interpolate to a gaussian grid, if necessary gridlons, gridlats = np.meshgrid(lons, lats) if interp: if ntrunc not in TRUNC_NLATS.keys(): raise ValueError( 'Truncation T{} is not in the dictionary TRUNC_NLATS'.format( ntrunc)) nlat = TRUNC_NLATS[ntrunc] lats, lons, ubar = interp_to_gaussian(gridlats, gridlons, ubar, nlat=nlat, return_latlon=True) vbar = interp_to_gaussian(gridlats, gridlons, vbar, nlat=nlat) if linearized: uprime = interp_to_gaussian(gridlats, gridlons, uprime, nlat=nlat) vprime = interp_to_gaussian(gridlats, gridlons, vprime, nlat=nlat) else: uprime = np.zeros(ubar.shape) vprime = np.zeros(vbar.shape) else: lons, lats = gridlons, gridlats # Get the mean state & perturbation vorticity from the winds s = spharm.Spharmt(lats.shape[1], lats.shape[0], gridtype='gaussian', rsphere=get_constant('planetary_radius', 'm'), legfunc='computed') vortb_spec, _ = s.getvrtdivspec(ubar, vbar, ntrunc=ntrunc) vort_bar = s.spectogrd(vortb_spec) vortp_spec, _ = s.getvrtdivspec(uprime, vprime, ntrunc=ntrunc) vort_prime = s.spectogrd(vortp_spec) # Generate the state return _generate_state(idate, lats, lons, vort_bar, vort_prime, linearized=linearized)
def integrate(init_cond, bmaps, ntimes=480): """ Function that integrates the barotropic model using spherical harmonics Input: init_cond : dictionary of intial conditions containing u and v initial conditions, the latitude and longitudes describing the grids, and a starting time bmaps : a dictionary of global and regional basemaps and projected coordinates ntimes : Number of timesteps to integrate """ # Get the initial u and v wind fields u = init_cond['u_in'] v = init_cond['v_in'] ntrunc = len(init_cond['lats']) start_time = init_cond['start_time'] # Create a radian grid lat_list_r = [x * np.pi / 180. for x in init_cond['lats']] lon_list_r = [x * np.pi / 180. for x in init_cond['lons']] # Meshgrid lons, lats = np.meshgrid(init_cond['lons'], init_cond['lats']) lamb, theta = np.meshgrid(lon_list_r, lat_list_r) dlamb = np.gradient(lamb)[1] dtheta = np.gradient(theta)[0] # Here is the Coriolis parameter f = 2 * 7.292E-5 * np.sin(theta) # Set up the spherical harmonic transform object s = spharm.Spharmt(len(init_cond['lons']), len(init_cond['lats']), rsphere=Re, gridtype='regular', legfunc='computed') # Use the object to plot the initial conditions # First convert to vorticity using spharm object vort_spec, div_spec = s.getvrtdivspec(u, v) div_spec = np.zeros( vort_spec.shape) # Only want non-divergent part of wind # Re-convert this to u-v winds to get the non-divergent component # of the wind field u, v = s.getuv(vort_spec, div_spec) # Use these winds to get the streamfunction (psi) and # velocity potential (chi) psi, chi = s.getpsichi(u, v) # Convert the spectral vorticity to grid vort_now = s.spectogrd(vort_spec) # Plot Initial Conditions curtime = start_time plot_figures(0, curtime, u, v, vort_now, psi, bmaps) # Now loop through the timesteps for n in xrange(ntimes): # Compute spectral vorticity from u and v wind vort_spec, div_spec = s.getvrtdivspec(u, v) # Now get the actual vorticity vort_now = s.spectogrd(vort_spec) div = np.zeros( vort_now.shape) # Divergence is zero in barotropic vorticity # Here we actually compute vorticity tendency # Compute tendency with beta as only forcing vort_tend_rough = -2. * 7.292E-5/(Re**2) * d_dlamb(psi,dlamb) -\ Jacobian(psi,vort_now,theta,dtheta,dlamb) # Apply hyperdiffusion if requested for smoothing if use_hyperdiffusion: vort_tend = add_hyperdiffusion(s, vort_now, vort_tend_rough, ntrunc) else: vort_tend = vort_tend_rough if n == 0: # First step just do forward difference # Vorticity at next time is just vort + vort_tend * dt vort_next = vort_now + vort_tend[:, :, 0] * dt else: # Otherwise do leapfrog vort_next = vort_prev + vort_tend[:, :, 0] * 2 * dt # Invert this new vort to get the new psi (or rather, uv winds) # First go back to spectral space vort_spec = s.grdtospec(vort_next) div_spec = s.grdtospec(div) # Now use the spharm methods to get new u and v grid u, v = s.getuv(vort_spec, div_spec) psi, chi = s.getpsichi(u, v) #raw_input() # Change vort_now to vort_prev # and if not first step add Robert filter # (from Held barotropic model) # to dampen out crazy modes r = 0.2 if n == 0: vort_prev = vort_now else: vort_prev = (1. - 2. * r) * vort_now + r * (vort_next + vort_prev) cur_fhour = (n + 1) * dt / 3600. curtime = start_time + timedelta(hours=cur_fhour) # Output every six hours if cur_fhour % output_freq == 0 and plot_output: # Go from psi to geopotential #phi = divide(psi * f,9.81) print("Plotting hour", cur_fhour) plot_figures(cur_fhour, curtime, u, v, vort_next, psi, bmaps)
def sinusoidal_perts_on_zonal_jet(linearized=False, idate=None, ntrunc=42, amp=8e-5, m=4, theta0=45., theta_w=15.): """ Creates ICs for an idealized case: extratropical zonal jets with superimposed sinusoidal NH vorticity perturbations. Taken from the GFDL model documentation (default case) Args ---- linearized : bool True if this is a linearized model. idate : datetime Forecast initialization date ntrunc : int Triangular trunction (e.g., 42 for T42). amp : float Vorticity perturbation amplitude [s^-1]. m : int Vorticity perturbation zonal wavenumber. theta0 : float Center latitude [degrees] for the vorticity perturbations. theta_w : float Halfwidth [degrees lat/lon] of the vorticity perturbations. Returns ------- ics : dict Model initial state (a dictionary of DataArrays) """ # Set the init. date if not provided if idate is None: idate = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0) # Create a gaussian lat/lon grid if ntrunc not in TRUNC_NLATS.keys(): raise ValueError( 'Truncation T{} is not in the dictionary TRUNC_NLATS'.format( ntrunc)) nlat = TRUNC_NLATS[ntrunc] lons, lats = gaussian_latlon_grid(nlat) theta = np.deg2rad(lats) lamb = np.deg2rad(lons) # Mean state: zonal extratropical jets ubar = 25 * np.cos(theta) - 30 * np.cos(theta)**3 + \ 300 * np.sin(theta)**2 * np.cos(theta)**6 vbar = np.zeros(np.shape(ubar)) # Get the mean state vorticity from ubar and vbar s = spharm.Spharmt(lamb.shape[1], lamb.shape[0], gridtype='gaussian', rsphere=Re, legfunc='computed') vortb_spec, _ = s.getvrtdivspec(ubar, vbar, ntrunc=42) vort_bar = s.spectogrd(vortb_spec) # Initial perturbation: sinusoidal vorticity perturbations theta0 = np.deg2rad(theta0) # center lat --> radians theta_w = np.deg2rad(theta_w) # halfwidth ---> radians vort_prime = 0.5 * amp * np.cos(theta) * np.exp(-((theta-theta0)/theta_w)**2) * \ np.cos(m*lamb) vort_prime = s.spectogrd(s.grdtospec(vort_prime, ntrunc=42)) # Generate the state return _generate_state(idate, lats, lons, vort_bar, vort_prime, linearized=linearized)
def __init__(self, nx, ny): self.spharmt = spharm.Spharmt(nx, ny)