<jupyter_start><jupyter_text>MOM6 input files for layered stacked shallow water<jupyter_code>import numpy as np
import netCDF4 as nc
import matplotlib.pyplot as plt
import xarray as xr
# First define grid
# The numbers are the same as in Constantinou and Hogg, 2019 GRL
#NI, NJ, NK = 800, 400, 2 # number of points in zonal, meridional and vertical levels
#LENLON, LENLAT = 3200., 1600. # in km 

NI, NJ, NK = 400, 800, 2 # number of points in zonal, meridional and vertical levels
LENLON, LENLAT = 1600., 3200. # in km 

MAXIMUM_DEPTH = 4000.   # in m 
WESTLON, SOUTHLAT = 0.0, 0.0

X = np.linspace(WESTLON + LENLON/(NI*2), WESTLON + LENLON - LENLON/(2*NI), NI)
Y = np.linspace(SOUTHLAT + LENLAT/(NJ*2), SOUTHLAT + LENLAT - LENLAT/(NJ*2), NJ) - LENLAT/2.

Xm, Ym = np.meshgrid(X,Y)
def flat_bottom():
    D = MAXIMUM_DEPTH*np.ones((NJ,NI))
    return D

def Navid_ridge_bottom(H0, sigma):
    D = MAXIMUM_DEPTH*np.ones((NJ,NI))
    
    #ridge = np.exp(-((Xm - 400.)/sigma)**2) + np.exp(-((Xm - 1000.)/sigma)**2) \
    #+ np.exp(-((Xm - 1600.)/sigma)**2) + np.exp(-((Xm - 2200.)/sigma)**2) \
    #+ np.exp(-((Xm - 2800.)/sigma)**2)
    
    ridge = np.exp(-((Xm - 300.)/sigma)**2) + np.exp(-((Xm - 800.)/sigma)**2) \
    + np.exp(-((Xm - 1300.)/sigma)**2)
    
    ridge = ridge*H0/np.max(ridge)
    
    D = D - ridge
    return D

def guassian_ridge_bottom(H0, sigma):
    D = MAXIMUM_DEPTH*np.ones((NJ,NI))
    ridge = np.exp(-((Xm - LENLON/2)/sigma)**2)
    ridge = ridge - np.min(ridge)
    ridge = ridge*H0/np.max(ridge)
    D = D - ridge
    return D

def sin_ridge_bottom(H0, n):
    D = MAXIMUM_DEPTH*np.ones((NJ,NI))
    ridge = np.sin(2*np.pi*n*Xm/LENLON - np.pi/2.)
    ridge = (H0)*ridge
    ridge[ridge < 0.] = 0. # negative values are removed
    D = D - ridge
    return D

def zero_velocity():
    u = np.zeros((NK,NJ,NI))
    v = np.zeros((NK,NJ,NI))
    return u,v

def wind_stress(tau_0):
    taux = tau_0*np.sin(np.pi*Ym/LENLAT)**2
    tauy = np.zeros((NJ,NI))
    return taux, tauy

def wind_stress_fix_width(tau_0,jet_width):
    no_jet = 1. - jet_width
    #taux = tau_0*np.sin(np.pi*(Ym - no_jet*LENLAT/2.)/(jet_width*LENLAT))**2
    taux = tau_0*np.cos(np.pi*Ym/(jet_width*LENLAT))**2
    taux[:int(np.round(no_jet*NJ/2.)),:] = 0.
    taux[int(np.round((jet_width + no_jet/2.)*NJ)):,:] = 0.
    tauy = np.zeros((NJ,NI))
    return taux, tauy

def wind_stress_fix_grad_fix_area(tau_0, base_tau, base_jet_width):
    tau_g = np.sqrt(base_tau*tau_0)
    amp = tau_g/base_tau
    jet_width = base_jet_width*amp 
    no_jet = 1. - jet_width
    taux = tau_g*np.cos(np.pi*Ym/(jet_width*LENLAT))**2
    taux[:int(np.round(no_jet*NJ/2.)),:] = 0.
    taux[int(np.round((jet_width + no_jet/2.)*NJ)):,:] = 0.
    tauy = np.zeros((NJ,NI))
    return taux, tauy
    
def wind_stress_var_width(tau_0, base_tau, base_jet_width):
    amp = tau_0/base_tau
    jet_width = base_jet_width*amp 
    no_jet = 1. - jet_width
    taux = tau_0*np.cos(np.pi*Ym/(jet_width*LENLAT))**2
    taux[:int(np.round(no_jet*NJ/2.)),:] = 0.
    taux[int(np.round((jet_width + no_jet/2.)*NJ)):,:] = 0.
    tauy = np.zeros((NJ,NI))
    return taux, tauy

def wind_stress_Navid(tau_0):
    taux = tau_0*np.sin(np.pi*(Ym - 0*LENLAT/4.)/(3*LENLAT/4))**2
    taux[int(np.round(3*NJ/4)):,:] = 0.
    tauy = np.zeros((NJ,NI))
    return taux, tauy

# Returns an array h[k,j,i] representing the layer thickness of the k-th layer at (lat,lon) 
# and an array eta[k,j,i] with the heights of each interface.
# Note that h has NK values in the vertical while eta has NK+1 values.
# k = 0 is the surface layer, with increasing k descending down the layers.
# h[i,j,k] >= ANGSTROM, which is the minimum layer thickness.
# input : D[i,j], a two-dimensional array of the depth of the ocean.
def layer_thickness(eta0, depth):
    
    eta = np.zeros((NK+1,NJ,NI))
    h = np.zeros((NK,NJ,NI))
    
    eta[NK,:,:] = - depth
    
    if(NK ==2):
        eta[NK-1,:,:] = eta0[1]
    
    for k in range(0,NK):
        h[k,:,:] = eta[k,:,:] - eta[k+1,:,:]
        
    return h, eta
# Generate Topography
ridge_h0 = 400.
sigma = 165./2
n_ridge = 5

H1, H2 = (1000., 3000.)

#depth = guassian_ridge_bottom(ridge_h0, sigma)
#depth = sin_ridge_bottom(ridge_h0, n_ridge)
depth = Navid_ridge_bottom(ridge_h0,sigma)

if(NK==2):
    eta0 = [0., -H1, -MAXIMUM_DEPTH]
else:
    eta0 = [0., -MAXIMUM_DEPTH]

h, eta = layer_thickness(eta0, depth)

# g_prime can be computed for a given rossby deformation radius value. This is used in the MOM_input file
# Or layer densities can supplied in MOM_input
Rd = 20000.
f = -10**-4
gp = Rd**2*MAXIMUM_DEPTH*f**2/(H1*H2)

rho_0 = 1027.35
rho_1 = gp*rho_0/9.8 + rho_0

# This is only for NK = 1 or 2. Modify the code for more layers.
if (NK == 2):
    density = [rho_0, rho_1]
else:
    density = rho_0*0.25 + rho_1*0.75

print(gp, density)


fig = plt.figure(figsize=[6,3])
plt.contourf(X,Y,depth)
plt.xlabel('x (km)')
plt.ylabel('y (km)')
plt.colorbar()
#plt.plot(X, H1 + H2 - depth[0,:])
#plt.ylim([0, 4000.])
# Save topographic and thickness data to nc file
ds = xr.Dataset(data_vars={'depth': (('lat', 'lon'), depth),
                           'h': (('level', 'lat', 'lon'), h),
                           'eta': (('interface','lat', 'lon'), eta)},
                coords={'lat': Y, 'lon': X, 'depth_list': eta0, 'level': np.atleast_1d(density)})

ds.lat.attrs['units'] = 'km'
ds.lon.attrs['units'] = 'km'
ds.depth_list.attrs['units'] = 'm'
ds.level.attrs['units'] = 'kg m-3'
ds.depth.attrs['units'] = 'm'
ds.h.attrs['units'] = 'm'
ds.eta.attrs['units'] = 'm'

ds.lon.attrs['standard_name'] = 'Longitude'
ds.lat.attrs['standard_name'] = 'Latitude'
ds.depth_list.attrs['standard_name'] = 'Mean Interface Height'
ds.level.attrs['standard_name'] = 'Mean Layer Density'
ds.depth.attrs['standard_name'] = 'Total Depth'
ds.h.attrs['standard_name'] = 'Layer Depth'
ds.eta.attrs['standard_name'] = 'Interface Height'

if(NK ==2):
    file = 'topography_2L.nc'
else:
    file = 'topography_1L.nc'    
ds.to_netcdf(path=file)
ds<jupyter_output><empty_output><jupyter_text>$\large \text{Wind Stress Profile}$We assume that the zonal wind stress is maximum at the center of the domain and follows the following functional form$$\tau = \tau_o\cos^2\left[\frac{\pi y}{\delta L_y}\right], -\delta L_y/2 < y<\delta L_y/2$$Here, $\delta$, which has value between 0 and 1, is the relative width of the wind stress profile. In the first set of experiments, only $\tau_o$ is varied. In the second set of experiments, $\delta$ is also varied in response to changes in $\tau_o$ to keep the maximum magnitude of $\frac{\partial \tau}{\partial y}$ across also experiements and $\int_{\frac{-\delta L_y}{2}}^{\frac{\delta L_y}{2}} \tau dy = \frac{\tau_o\delta L_y}{2}$ is the same as in correponding experiment for which only $\tau_o$ was changed. Let's assume that $\tau_b$ and $\delta_b$ are base values, $\tau_f$ and $\delta_b$ are new values in the experiment with fixed width and $\tau_g$ and $\delta_g$ are new values in the experiment with fixed meridional gradient. It is trivial to show that, in order to conserve maximum value in gradient, $$\frac{\tau_g}{\delta_g} = \frac{\tau_b}{\delta_b}$$For the area to be the same,$$\tau_g\delta_g = \tau_f\delta_b$$From these relations, we have$$\tau_g = \sqrt{\tau_b\tau_f}, \delta_g = \delta_b\frac{\tau_f}{\tau_b}$$<jupyter_code># Check wind profile

tau_b = 0.1 # base value
jet_width = 0.3 # base value
tau_f = 1.

taux1, tauy = wind_stress_fix_width(tau_f, jet_width)
taux2, tauy = wind_stress_fix_grad_fix_area(tau_f, tau_b, jet_width)
print('Fix = ', np.sum(taux1[:,0])*LENLAT/NJ, ' Var = ', np.sum(taux2[:,0])*LENLAT/NJ)

plt.plot(taux1[:,0],Y, label='Fix Width')
plt.plot(taux2[:,0],Y, label='Var Width')
plt.ylabel('y (km)')    
plt.xlabel('Wind Stress (N/m$^2$)')
plt.legend()
plt.grid()
# Save data to nc file
def save_tau(taux, tauy, X, Y, mag, case):
    
    ds = xr.Dataset(
            data_vars={'taux':    (('lat', 'lon'), taux),
                       'tauy': (('lat', 'lon'), tauy)},
        coords={'lat': Y, 'lon': X})

    ds.lat.attrs['units'] = 'km'
    ds.lon.attrs['units'] = 'km'
    ds.taux.attrs['units'] = 'N m-2'
    ds.tauy.attrs['units'] = 'N m-2'

    ds.lon.attrs['standard_name'] = 'Longitude'
    ds.lat.attrs['standard_name'] = 'Latitude'
    ds.taux.attrs['standard_name'] = 'Zonal Wind Stress'
    ds.tauy.attrs['standard_name'] = 'Meridional Wind Stress'

    file = 'wind_stress_' + case + '_tau_' + str(round(mag*100)) + '.nc'
    ds.to_netcdf(path=file)
# Wind Stress
tau_0 = 0.1

mag = np.array([0.01, 0.03, 0.05, 0.1, 0.3, 0.5, 1.])

jet_width = 0.3

fig = plt.figure(figsize=[14,5])

plt.subplot(1,3,1)

for i in range(0,len(mag)):
    taux, tauy = wind_stress_fix_width(mag[i], jet_width)
    
    print('Fix Width = ', np.sum(taux[:,0])*LENLAT/NJ)
    save_tau(taux, tauy, X, Y, mag[i], 'FW')
    
    plt.plot(taux[:,0],Y, label='Tau = '+str(mag[i]))
plt.ylabel('y (km)')    
plt.xlabel('Wind Stress (N/m$^2$)')
plt.title('Fixed Jet Width')
plt.legend()
plt.grid()

plt.subplot(1,3,2)
for i in range(0,len(mag)):
    taux, tauy = wind_stress_fix_grad_fix_area(mag[i], tau_0, jet_width)
    print('Var Width = ', np.sum(taux[:,0])*LENLAT/NJ)
    
    save_tau(taux, tauy, X, Y, mag[i], 'VW')
    
    plt.plot(taux[:,0],Y, label='Tau = '+str(mag[i]))
    
plt.ylabel('y (km)')    
plt.xlabel('Wind Stress (N/m$^2$)')
plt.title('Fixed Tau gradient')
plt.legend()
plt.grid()


plt.subplot(1,3,3)
for i in range(0,len(mag)):
    taux, tauy = wind_stress_fix_grad_fix_area(mag[i], tau_0, jet_width)
    plt.plot(np.diff(taux[:,0])/(LENLAT*1000./(NJ*2)), 0.5*(Y[0:-1]+Y[1:]), label='Tau = '+str(mag[i]))
plt.ylabel('y (km)')    
plt.xlabel('Wind Stress Gradient (N/m$^3$)')
plt.title('Fixed Tau gradient')
plt.legend()
plt.grid()<jupyter_output>Fix Width =  4.800000000000001
Fix Width =  14.4
Fix Width =  24.000000000000004
Fix Width =  48.00000000000001
Fix Width =  144.0
Fix Width =  240.0
Fix Width =  480.0
Var Width =  4.799998118197038
Var Width =  14.39999600287121
Var Width =  23.999997828200385
Var Width =  48.00000000000001
Var Width =  143.99999908114583
Var Width =  240.00000095430767
Var Width =  480.00000017760465
{'fraction_non_alphanumeric': 0.1825124840889063, 'fraction_numerical': 0.0485655537060609, 'mean_word_length': 3.5577867023650156, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '579224', 'license': 'MIT', 'n_tokens_mistral': 4710, 'n_tokens_neox': 4254, 'n_words': 966, 'path': 'input-files/input.ipynb', 'repo_name': 'ElizabethYankovsky/Topography-Eddy-MOM6', 'size': '10213'}