<jupyter_start><jupyter_text>ATM 651 Introduction of Atmospheric Dynamics: Project optionUniversity of Miami, RSMAS. 2019 Fall  2D Shallow water equation model This notebook is adapted from [Dr. <NAME>'s webpage.](https://personalpages.manchester.ac.uk/staff/paul.connolly/teaching/practicals/shallow_water_equations.html)This model integrates the **LINEARIZED** [shallow water equations](https://en.wikipedia.org/wiki/Shallow_water_equations) in conservative form in a channel using the [Lax-Wendroff method](https://en.wikipedia.org/wiki/Lax%E2%80%93Wendroff_method), a numerical method for the solution of hyperbolic partial differential equations, based on [finite differences](https://en.wikipedia.org/wiki/Finite_difference_method). It is second-order accurate in both space and time.<jupyter_code>import numpy as np
import sys
import os
from scipy.special import erfcinv as erfcinv
import tqdm as tqdm  # pause jupyter and pip install tqdm if necessary
import time
import scipy.io as sio
import lax_wendroff_Hsrc as lw
import warnings
warnings.simplefilter("ignore")<jupyter_output><empty_output><jupyter_text>Main codeThere are options for different simulations and flexible parameters<jupyter_code># SECTION 0: Definitions (normally don't modify this section)
# Possible orographies
FLAT=0;

# ------------------------------------------------------------------
# SECTION 1: Configuration
g    = 9.81;                # Acceleration due to gravity (m/s2)
f    = 1e-4;              # Coriolis parameter (s-1)

#f=0.;
beta = 1.6e-11;             # Meridional gradient of f (s-1m-1)
#beta=0.;
#beta=5e-10;

output_interval = 24.*3600.; # Time between outputs (s)
forecast_length = 30*24*3600.0; # Forecast length (s)

orography = FLAT
initially_geostrophic = True;   # Can be "True" or "False"
add_random_height_noise = False; # Can be "True" or "False"

nx=254; # Number of zonal gridpoints
ny=128;  # Number of meridional gridpoints

dx=100.0e3; # Zonal grid spacing (m)
dy=dx;      # Meridional grid spacing


# ------------------------------------------------------------------
# SECTION 2: Act on the configuration information

x=np.mgrid[0:nx]*dx; # Zonal distance coordinate (m)
y=np.mgrid[0:ny]*dy; # Meridional distance coordinate (m)
[Y,X] = np.meshgrid(y,x); # Create matrices of the coordinate variables

# Create the orography field "H"
if orography == FLAT:
    H = np.zeros((nx, ny));

# Create the initial height field (in meters)
height = 100.*np.ones((nx, ny))

# Specify the range of heights to plot in metres
plot_height_range = np.array([95., 105.]);

# external heat source adding to h equation
hsrc = H.copy()
hsrc[np.where(np.sqrt((X-np.mean(x))**2.+(Y-np.mean(y))**2.) <= 1000e3)] = -1E-4 # m/s

HEATING_OFF_TIME = 2*24*3600.0


# Calculate linear wave speed 
cgrav = np.sqrt(g*np.mean(height))
cgrav
dt = dx/cgrav /20. # factor of 5 or greater for safe time integration

# Set up grid and output arrays 
nt = int(np.fix(forecast_length/dt)+1); # Number of timesteps
timesteps_between_outputs = np.fix(output_interval/dt);
noutput = int(np.ceil(nt/timesteps_between_outputs)); # Number of output frames

dt
# Coriolis parameter as a matrix of values varying in y only
F = f+beta*(Y-np.mean(y));

# Initialize the wind to rest
u=np.zeros((nx, ny));
v=np.zeros((nx, ny));

# Define h as the depth of the fluid (whereas "height" is the height of
# the upper surface)
h = height - H;

# Initialize the 3D arrays where the output data will be stored
u_save = np.zeros((nx, ny, noutput));
v_save = np.zeros((nx, ny, noutput));
h_save = np.zeros((nx, ny, noutput));
t_save = np.zeros((noutput,1));

# Index to stored data
i_save = 0;



# ------------------------------------------------------------------
# SECTION 3: Main loop
for n in range(0,nt):
   # Every fixed number of timesteps we store the fields
    if np.mod(n,timesteps_between_outputs) == 0:
   
        max_u = np.sqrt(np.max(u[:]*u[:]+v[:]*v[:]));
      
        print("Time = %f hours (max %f); max(|u|) = %f"  
           % ((n)*dt/3600., forecast_length/3600./24., max_u) )
   
        u_save[:,:,i_save] = u;
        v_save[:,:,i_save] = v;
        h_save[:,:,i_save] = h;
        t_save[i_save] = (n)*dt;
        i_save = i_save+1;
  

   # Compute the accelerations
    u_accel = F[1:-1,1:-1]*v[1:-1,1:-1] \
              - (g/(2.*dx))*(H[2:,1:-1]-H[0:-2,1:-1]);
    v_accel = -F[1:-1,1:-1]*u[1:-1,1:-1] \
              - (g/(2.*dy))*(H[1:-1,2:]-H[1:-1,0:-2]);
    h_accel = hsrc * (n*dt < HEATING_OFF_TIME) # height source (heating, turned off at HEATING_OFF_TIME)

   # Call the Lax-Wendroff scheme to move forward one timestep
    (unew, vnew, h_new) = lw.lax_wendroff(dx, dy, dt, g, u, v, h, u_accel, v_accel, h_accel);

   # Update the wind and height fields, taking care to enforce 
   # boundary conditions 
    """
    u = unew([end 1:end 1],[1 1:end end]);
    v = vnew([end 1:end 1],[1 1:end end]);
    v(:,[1 end]) = 0;
    h(:,2:end-1) = h_new([end 1:end 1],:);
    """
    u[1:-1,1:-1] = unew[0:,0:];
    u[[-1,0],1:-1]  = unew[[0,-1],:]
    u[1:-1,[0,-1]]  = unew[:,[0,-1]]

    v[1:-1,1:-1] = vnew[0:,0:];
    v[[-1,0],0]  = vnew[[0,-1],0]
    v[1:-1,[0,-1]]  = vnew[:,[0,-1]]
   
    v[:,[0, -1]] = 0.;

    h[1:-1,1:-1] = h_new[0:,0:];
    h[[0,-1],1:-1]  = h_new[[-1,0],:]


print('Run completed successfully');<jupyter_output>Time = 0.000000 hours (max 30.000000); max(|u|) = 0.000000
Time = 23.990001 hours (max 30.000000); max(|u|) = 1.352118
Time = 47.980002 hours (max 30.000000); max(|u|) = 2.563657
Time = 71.970003 hours (max 30.000000); max(|u|) = 2.367241
Time = 95.960004 hours (max 30.000000); max(|u|) = 2.204398
Time = 119.950005 hours (max 30.000000); max(|u|) = 2.174828
Time = 143.940006 hours (max 30.000000); max(|u|) = 2.045164
Time = 167.930007 hours (max 30.000000); max(|u|) = 2.019507
Time = 191.920008 hours (max 30.000000); max(|u|) = 1.962617
Time = 215.910008 hours (max 30.000000); max(|u|) = 1.928198
Time = 239.900009 hours (max 30.000000); max(|u|) = 1.889617
Time = 263.890010 hours (max 30.000000); max(|u|) = 1.844571
Time = 287.880011 hours (max 30.000000); max(|u|) = 1.835131
Time = 311.870012 hours (max 30.000000); max(|u|) = 1.769892
Time = 335.860013 hours (max 30.000000); max(|u|) = 1.715859
Time = 359.850014 hours (max 30.000000); max(|u|) = 1.683258
Time = 383.840015 hours (max 3[...]<jupyter_text>plot and save figures<jupyter_code>import matplotlib.pyplot as plt
from matplotlib import rc

# create a folder in the current directory
os.system('rm -r fig');os.system('mkdir fig')

#f,(ax1, ax2) = plt.subplots(2, sharex=True, sharey=False)
f=plt.figure(figsize=(12,10))
ax1=plt.subplot(2,1,1)
ax2=plt.subplot(2,1,2)

#ax1.autoscale(enable=True, axis='y', tight=True)

# Axis units are thousands of kilometers (x and y are in metres)
x_1000km = x * 1.e-6
y_1000km = y * 1.e-6

# Set colormap to have 64 entries
ncol=64;

# Interval between arrows in the velocity vector plot
interval = 6;

# Set this to "True" to save each frame as a png file
plot_frames = True;

# Decide whether to show height in metres or km
if np.mean(plot_height_range) > 1000:
    height_scale = 0.001;
    height_title = 'Height (km)';
else:
    height_scale = 1;
    height_title = 'Height (m)';


print('Maximum orography height = %f m' % np.max(H[:]));
u = np.squeeze(u_save[:,:,0]);
vorticity = np.zeros(np.shape(u));

# Loop through the frames of the animation
for it in range(0,noutput):

    # Extract the height and velocity components for this frame
    h = np.squeeze(h_save[:,:,it]);
    u = np.squeeze(u_save[:,:,it]);
    v = np.squeeze(v_save[:,:,it]);

    # Compute the vorticity
    vorticity[1:-1,1:-1] = (1./dy)*(u[1:-1,0:-2]-u[1:-1,2:]) \
      + (1./dx)*(v[2:,1:-1]-v[0:-2,1:-1]);
  
# First plot the height field
    if it==0:

        # Plot the height field
        im=ax1.imshow(np.transpose(h+H), \
         extent=[np.min(x_1000km),np.max(x_1000km),np.min(y_1000km),np.max(y_1000km)], \
         cmap='jet')
        # Set other axes properties and plot a colorbar
        cb1=plt.colorbar(im,ax=ax1,shrink=0.5)
        cb1.set_label('height (m)')
        # Contour the terrain:
        cs=ax1.contour(x_1000km,y_1000km,np.transpose(H),levels=np.arange(0,100,10), colors='k')

        # Plot the velocity vectors
        Q = ax1.quiver(x_1000km[2::interval],y_1000km[2::interval], \
         np.transpose(u[2::interval,2::interval]), \
         np.transpose(v[2::interval,2::interval]), scale=1,scale_units='xy',pivot='mid')
        ax1.set_ylabel('Y distance (1000s of km)');
        ax1.set_title(height_title);
        tx1=ax1.text(0, np.max(y_1000km), 'Time = %.1f hours' % (t_save[it]/3600.));

        # Now plot the vorticity
        im2=ax2.imshow(np.transpose(vorticity), \
         extent=[np.min(x_1000km),np.max(x_1000km),np.min(y_1000km),np.max(y_1000km)], \
         cmap='jet',vmax=1e-5,vmin=-1e-5)
        # Set other axes properties and plot a colorbar
        cb2=plt.colorbar(im2,ax=ax2,shrink=0.5)
        cb2.set_label('vorticity (s$^{-1}$)')
        ax2.set_xlabel('X distance (1000s of km)');
        ax2.set_ylabel('Y distance (1000s of km)');
        ax2.set_title('Relative vorticity (s$^{-1}$)');
        tx2=ax2.text(0, np.max(y_1000km), 'Time = %.1f hours' % (t_save[it]/3600.));
        
        plt.tight_layout()

    else:
        # top plot:
        im.set_data(np.transpose(H+h))
        cs.set_array(np.transpose(h))
        Q.set_UVC(np.transpose(u[2::interval,2::interval]), \
               np.transpose(v[2::interval,2::interval]))
        tx1.set_text('Time = %.1f hours' % (t_save[it]/3600.));

        # bottom plot:
        im2.set_data(np.transpose(vorticity))
        tx2.set_text('Time = %.1f hours' % (t_save[it]/3600.));

    im.set_clim((plot_height_range));
    im2.set_clim((-1e-5,1e-5));
    ax1.axis((0., np.max(x_1000km), 0., np.max(y_1000km)));
    ax2.axis((0., np.max(x_1000km), 0., np.max(y_1000km)));
    
    # To make an animation we can save the frames as a 
    # sequence of images
    if plot_frames:
        plt.savefig('./fig/frame%03d.png' % it,format='png',dpi=200)

#    plt.pause(0.1)
!ls ./fig/
!open ./fig/frame???.png
!rm ./fig/frame???.png
plt.imshow(hsrc.transpose())
hsrc.max()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18864057773006734, 'fraction_numerical': 0.08138967502683712, 'mean_word_length': 3.2120838471023427, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 14, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '307954', 'license': 'MIT', 'n_tokens_mistral': 4458, 'n_tokens_neox': 3917, 'n_words': 924, 'path': 'Notebooks/SWEmodel-mass sink turned off after 2 days.ipynb', 'repo_name': 'kp374/ATM651_IntroAtmDynamics', 'size': '10254'}