├── README.md ├── barotropic_spectral.py ├── hyperdiffusion.py └── namelist.py /README.md: -------------------------------------------------------------------------------- 1 | Barotropic-Python 2 | ================= 3 | ----------------------------------------------------------------- 4 | 5 | A simple barotropic model written in Python using the ``spharm`` package for spherical harmonics. 6 | 7 | Currently set up to use __linearized__ initial conditions 8 | 9 | Written by Luke Madaus (5/2012) - University of Washington 10 | 11 | Restructured by Nick Weber (10/2017) - University of Washington 12 | 13 | ----------------------------------------------------------------- 14 | 15 | __**Requires**__: 16 | 17 | - PySPHARM -- Python interface to NCAR SPHEREPACK library: 18 | https://code.google.com/p/pyspharm/ 19 | - netCDF4 -- Python interface to netCDF4 library 20 | - numpy, datetime, matplotlib 21 | 22 | 23 | Based on the Held-Suarez Barotropic model, including hyperdiffusion. 24 | A brief description of their model may be found at: 25 | http://data1.gfdl.noaa.gov/~arl/pubrel/m/atm_dycores/src/atmos_spectral_barotropic/barotropic.pdf 26 | 27 | The basic premise of the code is to take upper-level u and v winds from any dataset 28 | (forecast or analysis), extract the non-divergent component, compute vorticity, and 29 | advect along this vorticity using the barotropic vorticity equation. As this model 30 | uses "real" atmospheric data (which is not barotropic), the results are rarely stable 31 | beyond ~5 days of forecasting. 32 | 33 | ----------------------------------------------------------------- 34 | 35 | __**Contents**__: 36 | 37 | - **``barotropic_spectral.py``** -- contains the ``Model`` class, which handles the initialization, 38 | integration, plotting, and I/O for the barotropic model 39 | - **``namelist.py``** -- functions as a traditional model namelist, containing the various 40 | configuration parameters for the barotropic model 41 | - **``hyperdiffusion.py``** -- contains functions for applying hyperdiffusion to the vorticity 42 | tendecy equation (helps prevent the model from blowing up) 43 | 44 | -------------------------------------------------------------------------------- /barotropic_spectral.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import numpy as np 3 | from datetime import datetime, timedelta 4 | import matplotlib.pyplot as plt 5 | from mpl_toolkits.basemap import Basemap 6 | import spharm 7 | import os 8 | from hyperdiffusion import del4_filter, apply_des_filter 9 | import namelist as NL # <---- IMPORTANT! Namelist containing constants and other model parameters 10 | 11 | 12 | class Model: 13 | """ 14 | Class for storing/plotting flow fields for a homogeneous (constant density), 15 | non-divergent, and incompressible fluid on a sphere and integrating them forward 16 | with the barotropic vorticity equation. 17 | """ 18 | 19 | def __init__(self, ics, forcing=None): 20 | """ 21 | Initializes the model. 22 | 23 | Requires: 24 | ics -----> Dictionary of linearized fields and space/time dimensions 25 | keys: u_bar, v_bar, u_prime, v_prime, lats, lons, start_time 26 | forcing -> a 2D array (same shape as model fields) containing a 27 | vorticity tendency [s^-2] to be imposed at each integration time step 28 | """ 29 | # 1) STORE SPACE/TIME VARIABLES (DIMENSIONS) 30 | # Get the latitudes and longitudes (as lists) 31 | self.lats = ics['lats'] 32 | self.lons = ics['lons'] 33 | self.start_time = ics['start_time'] # datetime 34 | self.curtime = self.start_time 35 | 36 | 37 | # 2) GENERATE/STORE NONDIVERGENT INITIAL STATE 38 | # Set up the spherical harmonic transform object 39 | self.s = spharm.Spharmt(self.nlons(), self.nlats(), rsphere=NL.Re, 40 | gridtype='regular', legfunc='computed') 41 | # Truncation for the spherical transformation 42 | if NL.M is None: 43 | self.ntrunc = self.nlats() 44 | else: 45 | self.ntrunc = NL.M 46 | # Use the object to get the initial conditions 47 | # First convert to vorticity using spharm object 48 | vortb_spec, div_spec = self.s.getvrtdivspec(ics['u_bar'], ics['v_bar']) 49 | vortp_spec, div_spec = self.s.getvrtdivspec(ics['u_prime'], ics['v_prime']) 50 | div_spec = np.zeros(vortb_spec.shape) # Only want NON-DIVERGENT part of wind 51 | # Re-convert this to u-v winds to get the non-divergent component 52 | # of the wind field 53 | self.ub, self.vb = self.s.getuv(vortb_spec, div_spec) # MEAN WINDS 54 | self.up, self.vp = self.s.getuv(vortp_spec, div_spec) # PERTURBATION WINDS 55 | # Use these winds to get the streamfunction (psi) and 56 | # velocity potential (chi) 57 | self.psib,chi = self.s.getpsichi(self.ub, self.vb) # MEAN STREAMFUNCTION 58 | self.psip,chi = self.s.getpsichi(self.up, self.vp) # PERTURBATION STREAMFUNCTION 59 | # Convert the spectral vorticity to grid 60 | self.vort_bar = self.s.spectogrd(vortb_spec) # MEAN RELATIVE VORTICITY 61 | self.vortp = self.s.spectogrd(vortp_spec) # PERTURBATION RELATIVE VORTICITY 62 | 63 | 64 | # 3) STORE A COUPLE MORE VARIABLES 65 | # Map projections for plotting 66 | self.bmaps = create_basemaps(self.lons, self.lats) 67 | # Get the vorticity tendency forcing (if any) for integration 68 | self.forcing = forcing 69 | 70 | 71 | #==== Some simple dimensional functions ========================================== 72 | def nlons(self): 73 | return len(self.lons) 74 | def nlats(self): 75 | return len(self.lats) 76 | 77 | 78 | 79 | #==== Primary function: model integrator ========================================= 80 | def integrate(self): 81 | """ 82 | Integrates the barotropic model using spherical harmonics. 83 | Simulation configuration is set in namelist.py 84 | """ 85 | # Create a radian grid 86 | lat_list_r = [x * np.pi/180. for x in self.lats] 87 | lon_list_r = [x * np.pi/180. for x in self.lons] 88 | 89 | # Meshgrid 90 | lons,lats = np.meshgrid(self.lons, self.lats) 91 | lamb, theta = np.meshgrid(lon_list_r, lat_list_r) 92 | 93 | # Need these for derivatives later 94 | dlamb = np.gradient(lamb)[1] 95 | dtheta = np.gradient(theta)[0] 96 | 97 | 98 | # Plot Initial Conditions 99 | if NL.plot_freq != 0: 100 | self.plot_figures(0) 101 | 102 | 103 | # Now loop through the timesteps 104 | for n in range(NL.ntimes): 105 | 106 | # Here we actually compute vorticity tendency 107 | # Compute tendency with beta as only forcing 108 | vort_tend = -2. * NL.omega/(NL.Re**2) * d_dlamb(self.psip + self.psib, dlamb) - \ 109 | Jacobian(self.psip+self.psib, self.vortp+self.vort_bar, theta, dtheta, dlamb) 110 | 111 | 112 | # Apply hyperdiffusion if requested for smoothing 113 | if NL.diff_opt==1: 114 | vort_tend -= del4_filter(self.vortp, self.lats, self.lons) 115 | elif NL.diff_opt==2: 116 | vort_tend = apply_des_filter(self.s, self.vortp, vort_tend, self.ntrunc, 117 | t = (n+1) * NL.dt / 3600.).squeeze() 118 | 119 | 120 | # Now add any imposed vorticity tendency forcing 121 | if self.forcing is not None: 122 | vort_tend += self.forcing 123 | 124 | if n == 0: 125 | # First step just do forward difference 126 | # Vorticity at next time is just vort + vort_tend * dt 127 | vortp_next = self.vortp + vort_tend * NL.dt 128 | else: 129 | # Otherwise do leapfrog 130 | vortp_next = vortp_prev + vort_tend * 2 * NL.dt 131 | 132 | 133 | # Invert this new vort to get the new psi (or rather, uv winds) 134 | # First go back to spectral space 135 | vortp_spec = self.s.grdtospec(vortp_next) 136 | div_spec = np.zeros(np.shape(vortp_spec)) # Divergence is zero in barotropic vorticity 137 | 138 | # Now use the spharm methods to update the u and v grid 139 | self.up, self.vp = self.s.getuv(vortp_spec, div_spec) 140 | self.psip, chi = self.s.getpsichi(self.up, self.vp) 141 | 142 | # Change vort_now to vort_prev 143 | # and if not first step, add Robert filter to dampen out crazy modes 144 | if n == 0: 145 | vortp_prev = self.vortp 146 | else: 147 | vortp_prev = (1.-2.*NL.r)*self.vortp + NL.r*(vortp_next + vortp_prev) 148 | 149 | # Update the vorticity 150 | self.vortp = self.s.spectogrd(vortp_spec) 151 | 152 | # Update the current time 153 | cur_fhour = (n+1) * NL.dt / 3600. 154 | self.curtime = self.start_time + timedelta(hours = cur_fhour) 155 | 156 | # Make figure(s) every hours 157 | if NL.plot_freq!=0 and cur_fhour % NL.plot_freq == 0: 158 | # Go from psi to geopotential 159 | print("Plotting hour", cur_fhour) 160 | self.plot_figures(int(cur_fhour)) 161 | 162 | 163 | #==== Plotting utilities ========================================================= 164 | def plot_figures(self, n, winds='total', vorts='total', psis='pert', showforcing=True, 165 | vortlevs=np.array([-10,-8,-6,-4,-2,2,4,6,8,10])*1e-5, 166 | windlevs=np.arange(20,61,4), hgtlevs=np.linspace(-500,500,26), 167 | forcelevs=np.array([-15,-12,-9,-6,-3,3,6,9,12,15])*1e-10): 168 | """ 169 | Make global and regional plots of the flow. 170 | 171 | Requires: 172 | n ------------------> timestep number 173 | winds, vorts, psis -> are we plotting the 'mean', 'pert', or 'total' field? 174 | showforcing --------> if True, contour the vorticity tendency forcing 175 | *levs --------------> contour/fill levels for the respective variables 176 | """ 177 | # What wind component(s) are we plotting? 178 | if winds=='pert': u = self.up; v = self.vp 179 | elif winds=='mean': u = self.ub; v = self.vb 180 | else: u = self.up+self.ub; v = self.vp+self.vb 181 | # What vorticity component(s) are we plotting? 182 | if vorts=='pert': vort = self.vortp 183 | elif vorts=='mean': vort = self.vort_bar 184 | else: vort = self.vortp + self.vort_bar 185 | # What streamfunction component(s) are we plotting? 186 | if psis=='pert': psi = self.psip 187 | elif psis=='mean': psi = self.psib 188 | else: psi = self.psip + self.psib 189 | 190 | # MAKE GLOBAL ZETA & WIND BARB MAP 191 | fig, ax = plt.subplots(figsize=(10,8)) 192 | fig.subplots_adjust(bottom=0.2, left=0.05, right=0.95) 193 | 194 | xx, yy = self.bmaps['global_x'], self.bmaps['global_y'] 195 | cs = ax.contourf(xx, yy, vort, vortlevs, cmap=plt.cm.RdBu_r, extend='both') 196 | self.bmaps['global'].drawcoastlines() 197 | ax.quiver(xx[::2,::2], yy[::2,::2], u[::2,::2], v[::2,::2]) 198 | # Plot the forcing 199 | if showforcing and self.forcing is not None: 200 | ax.contour(xx, yy, self.forcing, forcelevs, linewidths=2, colors='darkorchid') 201 | ax.set_title('relative vorticity [s$^{-1}$] and winds [m s$^{-1}$] at %03d hours' % n) 202 | # Colorbar 203 | cax = fig.add_axes([0.05, 0.12, 0.9, 0.03]) 204 | plt.colorbar(cs, cax=cax, orientation='horizontal') 205 | # Save figure 206 | if not os.path.isdir(NL.figdir+'/global'): os.makedirs(NL.figdir+'/global') 207 | plt.savefig('{}/global/zeta_wnd_{:03d}.png'.format(NL.figdir,n), bbox_inches='tight') 208 | plt.close() 209 | 210 | # MAKE REGIONAL HEIGHT & WIND SPEED MAP 211 | phi = np.divide(psi * NL.omega, NL.g) 212 | fig, ax = plt.subplots(figsize=(10,6)) 213 | fig.subplots_adjust(bottom=0.2, left=0.05, right=0.95) 214 | xx, yy = self.bmaps['regional_x'], self.bmaps['regional_y'] 215 | # Calculate wind speed 216 | wspd = np.sqrt(u**2 + v**2) 217 | cs = ax.contourf(xx, yy, wspd, windlevs, cmap=plt.cm.viridis, extend='max') 218 | self.bmaps['regional'].drawcoastlines() 219 | self.bmaps['regional'].drawcountries() 220 | self.bmaps['regional'].drawstates() 221 | hgtconts = ax.contour(xx, yy, phi, hgtlevs, colors='k') 222 | # Plot the forcing 223 | if showforcing and self.forcing is not None: 224 | ax.contour(xx, yy, self.forcing, forcelevs, linewidths=2, colors='darkorchid') 225 | ax.set_title('geopotential height [m] and wind speed [m s$^{-1}$] at %03d hours' % n) 226 | # Colorbar 227 | cax = fig.add_axes([0.05, 0.12, 0.9, 0.03]) 228 | plt.colorbar(cs, cax=cax, orientation='horizontal') 229 | # Save figure 230 | if not os.path.isdir(NL.figdir+'/regional'): os.makedirs(NL.figdir+'/regional') 231 | plt.savefig('{}/regional/hgt_wspd_{:03d}.png'.format(NL.figdir,n), bbox_inches='tight') 232 | plt.close() 233 | 234 | 235 | ########################################################################################################### 236 | ##### Other Utilities ##################################################################################### 237 | ########################################################################################################### 238 | 239 | 240 | def create_basemaps(lons,lats): 241 | """ Setup global and regional basemaps for eventual plotting """ 242 | print("Creating basemaps for plotting") 243 | 244 | long, latg = np.meshgrid(lons,lats) 245 | 246 | # Set up a global map 247 | bmap_globe = Basemap(projection='merc',llcrnrlat=-70, urcrnrlat=70, 248 | llcrnrlon=0,urcrnrlon=360,lat_ts=20,resolution='c') 249 | xg,yg = bmap_globe(long,latg) 250 | 251 | # Set up a regional map (currently Pacific and N. America) 252 | bmap_reg = Basemap(projection='merc',llcrnrlat=0,urcrnrlat=65,llcrnrlon=80, 253 | urcrnrlon=290, lat_ts=20,resolution='l') 254 | xr,yr = bmap_reg(long,latg) 255 | 256 | return {'global' : bmap_globe, 257 | 'global_x' : xg, 258 | 'global_y' : yg, 259 | 'regional' : bmap_reg, 260 | 'regional_x' : xr, 261 | 'regional_y' : yr, 262 | } 263 | 264 | def d_dlamb(field,dlamb): 265 | """ Finds a finite-difference approximation to gradient in 266 | the lambda (longitude) direction""" 267 | out = np.divide(np.gradient(field)[1],dlamb) 268 | return out 269 | 270 | def d_dtheta(field,dtheta): 271 | """ Finds a finite-difference approximation to gradient in 272 | the theta (latitude) direction """ 273 | out = np.divide(np.gradient(field)[0],dtheta) 274 | return out 275 | 276 | def Jacobian(A,B,theta,dtheta,dlamb): 277 | """ Returns the Jacobian of two fields in spherical coordinates """ 278 | term1 = d_dlamb(A,dlamb) * d_dtheta(B,dtheta) 279 | term2 = d_dlamb(B,dlamb) * d_dtheta(A,dtheta) 280 | return 1./(NL.Re**2 * np.cos(theta)) * (term1 - term2) 281 | 282 | ########################################################################################################### 283 | 284 | def test_case(): 285 | """ 286 | Runs an example case: extratropical zonal jets with superimposed sinusoidal NH vorticity 287 | perturbations and a gaussian vorticity tendency forcing. 288 | """ 289 | from time import time 290 | start = time() 291 | 292 | # 1) LET'S CREATE SOME INITIAL CONDITIONS 293 | lons = np.arange(0, 360.1, 2.5) 294 | lats = np.arange(-87.5, 88, 2.5)[::-1] 295 | lamb, theta = np.meshgrid(lons * np.pi/180., lats * np.pi/180.) 296 | # Mean state: zonal extratropical jets 297 | ubar = 25 * np.cos(theta) - 30 * np.cos(theta)**3 + 300 * np.sin(theta)**2 * np.cos(theta)**6 298 | vbar = np.zeros(np.shape(ubar)) 299 | # Initial perturbation: sinusoidal vorticity perturbations 300 | A = 1.5 * 8e-5 # vorticity perturbation amplitude 301 | m = 4 # zonal wavenumber 302 | theta0 = np.deg2rad(45) # center lat = 45 N 303 | thetaW = np.deg2rad(15) 304 | vort_pert = 0.5*A*np.cos(theta)*np.exp(-((theta-theta0)/thetaW)**2)*np.cos(m*lamb) 305 | # Get U' and V' from this vorticity perturbation 306 | s = spharm.Spharmt(len(lons), len(lats), gridtype='regular', legfunc='computed', rsphere=NL.Re) 307 | uprime, vprime = s.getuv(s.grdtospec(vort_pert), np.zeros(np.shape(s.grdtospec(vort_pert)))) 308 | # Full initial conditions dictionary: 309 | ics = {'u_bar' : ubar, 310 | 'v_bar' : vbar, 311 | 'u_prime': uprime, 312 | 'v_prime': vprime, 313 | 'lons' : lons, 314 | 'lats' : lats, 315 | 'start_time' : datetime(2017,1,1,0)} 316 | 317 | # 2) LET'S ALSO FEED IN A GAUSSIAN NH RWS FORCING 318 | amplitude = 10e-10 # s^-2 319 | forcing = np.zeros(np.shape(ubar)) 320 | x, y = np.meshgrid(np.linspace(-1,1,10), np.linspace(-1,1,10)) 321 | d = np.sqrt(x*x+y*y) 322 | sigma, mu = 0.5, 0.0 323 | g = np.exp(-( (d-mu)**2 / ( 2.0 * sigma**2 ) ) ) # GAUSSIAN CURVE 324 | lat_i = np.where(lats==35.)[0][0] 325 | lon_i = np.where(lons==160.)[0][0] 326 | forcing[lat_i:lat_i+10, lon_i:lon_i+10] = g*amplitude 327 | 328 | # 3) INTEGRATE! 329 | model = Model(ics, forcing=forcing) 330 | model.integrate() 331 | print('TOTAL INTEGRATION TIME: {:.02f} minutes'.format((time()-start)/60.)) 332 | 333 | if __name__ == '__main__': 334 | test_case() -------------------------------------------------------------------------------- /hyperdiffusion.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | """ 4 | Module for handling vorticity tendency diffusion in a barotropic model. 5 | """ 6 | 7 | import numpy as np 8 | import namelist as NL 9 | 10 | #==================================================================================== 11 | #==== N. Weber's new hyperdiffusion scheme ========================================== 12 | #==================================================================================== 13 | 14 | def del4_filter(zeta, lats, lons): 15 | """ 16 | Calculates/returns the diffusion term of the barotropic vorticity equation, 17 | which is subtracted from the overall vorticity tendency. 18 | 19 | Diffusion is calculated with the following formula: 20 | D = k * del^4(vorticity) 21 | [default namelist k value (2.338e16) taken from Sardeshmukh and Hoskins 1988] 22 | 23 | Requires: 24 | data --> 2D numpy array of data (e.g., global vorticity field) 25 | shape(data) = (nlats, nlons) 26 | lats --> 1D array/list of the corresponding latitudes (in degrees) 27 | lons --> 1D array/list of the corresponding longitudes (in degrees) 28 | 29 | Returns: 30 | 2-dimensional numpy array (same shape as ) 31 | """ 32 | del4vort = del4(zeta, lats, lons) 33 | return NL.k * del4vort 34 | 35 | 36 | def del4(data, lats, lons): 37 | """ 38 | Computes the del^4 operator on 2D global data in x-y space. 39 | Uses the given lats/lons to compute spatial derivatives in meters. 40 | 41 | Requires: 42 | data --> 2D numpy array of data (e.g., global vorticity field) 43 | shape(data) = (nlats, nlons) 44 | lats --> 1D array/list of the corresponding latitudes (in degrees) 45 | lons --> 1D array/list of the corresponding longitudes (in degrees) 46 | 47 | Returns: 48 | 2-dimensional numpy array (same shape as ) 49 | """ 50 | # Get the data resolution (in degrees) -- assumes uniform lat/lon spacing 51 | res = lons[1] - lons[0] 52 | _, la = np.meshgrid(lons * np.pi / 180., lats * np.pi / 180.) 53 | 54 | # Use to calculate the distance (in meters) between each point 55 | dy = 111000. * res 56 | dx = np.cos(la[:,1:]) * 111000. * res 57 | 58 | # Calculate 2nd and 4th derivatives in x and y directions 59 | d2data_dy2 = second_derivative(data, dy, axis=0) 60 | d2data_dx2 = second_derivative(data, dx, axis=1) 61 | d4data_dy4 = fourth_derivative(data, dy, axis=0) 62 | d4data_dx4 = fourth_derivative(data, dx, axis=1) 63 | 64 | # Use the above to calculate/return del^4(data) 65 | return d4data_dy4 + d4data_dx4 + (2 * d2data_dy2 * d2data_dx2) 66 | 67 | 68 | def second_derivative(data, delta, axis=0): 69 | """ 70 | Computes the second derivative of Nd-array along the desired axis. 71 | 72 | Requires: 73 | data ---> N-dimensional numpy array 74 | delta --> float or 1-dimensional array/list (same length as desired axis) 75 | indicating the distance between data points 76 | axis ---> desired axis to take the derivative along 77 | 78 | Returns: 79 | N-dimensional numpy array (same shape as ) 80 | """ 81 | n = len(data.shape) 82 | 83 | # If is not a list/array (i.e., if the mesh is uniform), 84 | # create a list of deltas that is the same length as the desired axis 85 | deltashape = list(data.shape) 86 | deltashape[axis] -= 1 87 | if type(delta) in [float, int, np.float64]: 88 | delta = np.ones(deltashape) * delta 89 | elif type(delta) in [list, np.ndarray]: 90 | shp = list(data.shape) 91 | shp[axis] -= 1 92 | if np.shape(delta) != tuple(shp): 93 | print(np.shape(delta), tuple(shp)) 94 | raise ValueError('input should has invalid shape') 95 | else: 96 | raise ValueError('input should be value or array') 97 | 98 | # create slice objects --- initially all are [:, :, ..., :] 99 | slice0 = [slice(None)] * n 100 | slice1 = [slice(None)] * n 101 | slice2 = [slice(None)] * n 102 | delta_slice0 = [slice(None)] * n 103 | delta_slice1 = [slice(None)] * n 104 | 105 | # First handle centered case 106 | slice0[axis] = slice(None, -2) 107 | slice1[axis] = slice(1, -1) 108 | slice2[axis] = slice(2, None) 109 | delta_slice0[axis] = slice(None, -1) 110 | delta_slice1[axis] = slice(1, None) 111 | combined_delta = delta[delta_slice0] + delta[delta_slice1] 112 | center = 2 * (data[slice0] / (combined_delta * delta[delta_slice0]) - 113 | data[slice1] / (delta[delta_slice0] * delta[delta_slice1]) + 114 | data[slice2] / (combined_delta * delta[delta_slice1])) 115 | 116 | # Fill the left boundary (pad it with the edge value) 117 | slice0[axis] = slice(None,1) 118 | left = center[slice0].repeat(1, axis=axis) 119 | 120 | # Fill the right boundary (pad it with the edge value) 121 | slice0[axis] = slice(-1, None) 122 | right = center[slice0].repeat(1, axis=axis) 123 | 124 | return np.concatenate((left, center, right), axis=axis) 125 | 126 | 127 | 128 | 129 | def fourth_derivative(data, delta, axis=0): 130 | """ 131 | Computes the fourth derivative of Nd-array along the desired axis. 132 | 133 | Requires: 134 | data ---> N-dimensional numpy array 135 | delta --> float or 1-dimensional array/list (same length as desired axis) 136 | indicating the distance between data points 137 | axis ---> desired axis to take the derivative along 138 | 139 | Returns: 140 | N-dimensional numpy array (same shape as ) 141 | """ 142 | n = len(data.shape) 143 | 144 | # If is not a list (i.e., if the mesh is uniform), create a list 145 | # of deltas that is the same length as the desired axis 146 | deltashape = list(data.shape) 147 | deltashape[axis] -= 1 148 | if type(delta) in [float, int, np.float64]: 149 | delta = np.ones(deltashape) * delta 150 | elif type(delta) in [list, np.ndarray]: 151 | shp = list(data.shape) 152 | shp[axis] -= 1 153 | if np.shape(delta) != tuple(shp): 154 | print(np.shape(delta), tuple(shp)) 155 | raise ValueError('input should has invalid shape') 156 | raise ValueError('input should has invalid shape') 157 | else: 158 | raise ValueError('input should be value or array') 159 | 160 | 161 | # create slice objects --- initially all are [:, :, ..., :] 162 | slice0 = [slice(None)] * n 163 | slice1 = [slice(None)] * n 164 | slice2 = [slice(None)] * n 165 | slice3 = [slice(None)] * n 166 | slice4 = [slice(None)] * n 167 | delta_slice0 = [slice(None)] * n 168 | delta_slice1 = [slice(None)] * n 169 | delta_slice2 = [slice(None)] * n 170 | delta_slice3 = [slice(None)] * n 171 | 172 | 173 | # First handle centered case 174 | slice0[axis] = slice(None, -4) 175 | slice1[axis] = slice(1, -3) 176 | slice2[axis] = slice(2, -2) 177 | slice3[axis] = slice(3, -1) 178 | slice4[axis] = slice(4, None) 179 | delta_slice0[axis] = slice(None, -3) 180 | delta_slice1[axis] = slice(1, -2) 181 | delta_slice2[axis] = slice(2, -1) 182 | delta_slice3[axis] = slice(3, None) 183 | center = f4(data[slice0], data[slice1], data[slice2], data[slice3], data[slice4], 184 | delta[delta_slice0], delta[delta_slice1], delta[delta_slice2], delta[delta_slice3]) 185 | 186 | # Fill the left boundary (pad it with the edge value) 187 | slice0[axis] = slice(None,1) 188 | left = center[slice0].repeat(2, axis=axis) 189 | 190 | # Fill the right boundary (pad it with the edge value) 191 | slice0[axis] = slice(-1, None) 192 | right = center[slice0].repeat(2, axis=axis) 193 | 194 | return np.concatenate((left, center, right), axis=axis) 195 | 196 | 197 | 198 | 199 | def f4(f0, f1, f2, f3, f4, d0, d1, d2, d3): 200 | """ 201 | Computes the fourth derivative with the approximation: 202 | f^4(x) = [f(x-2) - 4*f(x-1) + 6*f(x) - 4*f(x+1) + f(x+2)] / hx^4 203 | (modified for non-uniform grid spacing) 204 | 205 | Requires: 206 | f0,f1,f2,f3,f4 -> values at the staggered locations (numerator) 207 | d0,d1,d2,d3 ----> distances (deltas) between the staggered locations (denominator) 208 | """ 209 | d01 = d0 + d1 210 | d02 = d0 + d1 + d2 211 | d03 = d0 + d1 + d2 + d3 212 | d12 = d1 + d2 213 | d13 = d1 + d2 + d3 214 | d23 = d2 + d3 215 | return 24 * (f0/(d0*d01*d02*d03) - f1/(d0*d1*d12*d13) + f2/(d01*d1*d2*d23) - 216 | f3/(d02*d12*d2*d3) + f4/(d03*d13*d23*d3)) 217 | 218 | 219 | 220 | #==================================================================================== 221 | #==== L. Madaus's original hyperdiffusion scheme ==================================== 222 | #==================================================================================== 223 | 224 | def apply_des_filter(s, cur_vort, vort_tend, ntrunc, t=0): 225 | """ Add spectral hyperdiffusion and return a new 226 | vort_tend """ 227 | # Convert to spectral grids 228 | print('vorticity:', np.shape(cur_vort)) 229 | vort_spec = s.grdtospec(cur_vort) 230 | print('spectral vorticity:', np.shape(vort_spec)) 231 | vort_tend_spec = s.grdtospec(vort_tend) 232 | total_length = vort_spec.shape[0] 233 | 234 | # Reshape to 2-d array 235 | vort_spec = np.reshape(vort_spec,(ntrunc,-1)) 236 | print('reshaped:', np.shape(vort_spec)) 237 | vort_tend_spec = np.reshape(vort_tend_spec,(ntrunc,-1)) 238 | new_vort_tend_spec = np.array(vort_tend_spec,dtype=np.complex) 239 | 240 | DES = compute_dampening_eddy_sponge(vort_tend_spec.shape) 241 | 242 | num = vort_tend_spec - DES * vort_spec 243 | den = np.ones(np.shape(DES), dtype=np.complex) + DES * np.complex(NL.dt,0.) 244 | new_vort_tend_spec[:,:] = num / den 245 | 246 | 247 | # Reshape the new vorticity tendency and convert back to grid 248 | new_vort_tend_spec = np.reshape(new_vort_tend_spec, (total_length,-1)) 249 | 250 | new_vort_tend = s.spectogrd(new_vort_tend_spec) 251 | 252 | return new_vort_tend 253 | 254 | 255 | def compute_dampening_eddy_sponge(fieldshape): 256 | """ Computes the eddy sponge by getting the eigenvalues 257 | of the Laplacian for each spectral coefficient and 258 | multiplying them by a dampening factor nu 259 | (specified in the namelist) 260 | From Held and Suarez 261 | """ 262 | 263 | # Need some arrays 264 | m_vals = np.arange(fieldshape[0]) 265 | n_vals = np.arange(fieldshape[1]) 266 | 267 | spherical_wave = np.zeros(fieldshape) 268 | 269 | fourier_wave = m_vals * NL.fourier_inc 270 | for n in n_vals: 271 | spherical_wave[:,n] = fourier_wave + n 272 | 273 | # Now for the laplacian 274 | eigen_laplacian = np.divide(np.multiply(spherical_wave,np.add(spherical_wave,1.)),NL.Re**2) 275 | 276 | # Dampening Eddy Sponge values 277 | DES = np.multiply(eigen_laplacian, NL.nu) 278 | DES_cpx = np.array(DES, dtype=np.complex) 279 | 280 | return DES_cpx -------------------------------------------------------------------------------- /namelist.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """ 3 | This module contains the variables used to run barotropic_spectral.py 4 | """ 5 | import os 6 | 7 | # Integration options 8 | dt = 900. # Timestep (seconds) 9 | ntimes = 960 # Number of time steps to integrate 10 | plot_freq = 6 # Frequency of output plots in hours (if 0, no plots are made) 11 | M = None # Truncation (if None, defaults to # latitudes) 12 | r = 0.2 # Coefficient for Robert Filter 13 | 14 | 15 | # I/O parameters 16 | figdir = os.path.join(os.getcwd(), 'figures') # Figure directory 17 | 18 | 19 | # Diffusion parameters 20 | diff_opt = 1 # Hyperdiffusion option (0 = none, 1 = del^4, 2 = DES) 21 | k = 2.338e16 # Diffusion coefficient for del^4 hyperdiffusion (diff_opt=1) 22 | nu = 1E-4 # Dampening coefficient for DES hyperdiffusion (diff_opt=2) 23 | fourier_inc = 1 # Fourier increment for computing dampening eddy sponge (diff_opt=2) 24 | 25 | 26 | # Constants 27 | Re = 6378100. # Radius of earth (m) 28 | omega = 7.292E-5 # Earth's angular momentum (s^-1) 29 | g = 9.81 # Gravitational acceleration (m s^-2) --------------------------------------------------------------------------------