title: Shallow water model using Python
body: <p>I've spent an ungodly amount of hours debugging my Python code to simulate the shallow water model. I cannot seem to reproduce the "Rossby" wave. The code can run but does not produce correct results. The code is below and should run by just using <code>python</code> or <code>python3</code>. A pop-up window will display the results. Below is the image of the incorrect results. The model should be lighter in color on the left of the "tower" due to rotation. Right now, it looks diagonal, my guess is that the <code>rotU</code> and <code>rotV</code> terms are causing the <code>dUdT</code> and <code>dVdT</code> to update the the results but not in the correct grids. 
  <a href="https://i.stack.imgur.com/Cbmq5.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/Cbmq5.png" alt="enter image description here"></a></p>
  
  <pre><code>import numpy
  import matplotlib.pyplot as plt
  import matplotlib.ticker as tkr
  import math
  
  ncol = 10        # grid size (number of cells)
  nrow = ncol
  
  nSlices = 10000   # maximum number of frames to show in the plot
  ntAnim = 1      # number of time steps for each frame
  
  horizontalWrap = True # determines whether the flow wraps around, connecting
                         # the left and right-hand sides of the grid, or whether
                         # there's a wall there. 
  rotationScheme = "PlusMinus"
  
  initialPerturbation = "Tower"
  textOutput = False
  plotOutput = True
  arrowScale = 30
  
  dT = 600 # seconds
  G = 9.8e-4 # m/s2, hacked (low-G) to make it run faster
  HBackground = 4000 # meters
  
  dX = 10.E3 # meters, small enough to respond quickly.  This is a very small ocean
  # on a very small, low-G planet.  
  flowConst = G  # 1/s2
  dragConst = 1.E-6  # about 10 days decay time
  
  rotConst = []
  for irow in range(0,nrow):
      rotationScheme is "PlusMinus":
          rotConst.append(-3.5e-5 * (1. - 0.8 * (irow - (nrow-1)/2) / nrow)) # rot 50% +-
  
  itGlobal = 0
  
  U = numpy.zeros((nrow, ncol+1))
  V = numpy.zeros((nrow+1, ncol))
  H = numpy.zeros((nrow, ncol+1))
  dUdT = numpy.zeros((nrow, ncol))
  dVdT = numpy.zeros((nrow, ncol))
  dHdT = numpy.zeros((nrow, ncol))
  dHdX = numpy.zeros((nrow, ncol+1))
  dHdY = numpy.zeros((nrow, ncol))
  dUdX = numpy.zeros((nrow, ncol))
  dVdY = numpy.zeros((nrow, ncol))
  rotV = numpy.zeros((nrow,ncol)) # interpolated to u locations
  rotU = numpy.zeros((nrow,ncol)) #              to v
  
  
  midCell = int(ncol/2)
  if initialPerturbation is "Tower":
      H[midCell,midCell] = 1
  
  
  ###############################################################
  def animStep():    
      global stepDump, itGlobal
  
      for time in range(0,ntAnim):
          #### Longitudinal derivative ##########################
          # Calculate dHdX
          for ix in range(0, nrow-1):
              for iy in range(0, ncol-1):
                  # Calculate the slope for X
                  dHdX[ix,iy+1] = (H[ix,iy+1] - H[ix,iy]) / dX    
          # Update the boundary cells
          if horizontalWrap is True: # Wrapping around
              U[:,ncol] = U[:,0]
              H[:,ncol] = H[:,0]
          else: # Bounded by walls on the left and right
              U[:,0] = 0 # U at the left-hand side is zero
              U[:,ncol-1] = 0 # U at the right-had side is zero
          # Calculate dUdX
          for ix in range(0, nrow):
              for iy in range(0, ncol):     
                  # Calculate the difference in U
                  dUdX[ix,iy] = (U[ix,iy+1] - U[ix,iy]) / dX
          ########################################################
          #### Latitudinal derivative ############################
          # Calculate dHdY
          dHdY[0,:] = 0 # The top boundary gradient dHdY set to zero
          for ix in range(0, nrow-1): # NOTE: the top row is zero
              for iy in range(0, ncol):
                  # Calculate the slope for Y
                  dHdY[ix+1,iy] = (H[ix+1,iy] - H[ix, iy]) / dX
          # Calculate dVdY
          V[0,:] = 0 # North wall is zero
          V[nrow,:] = 0 # South wall is zero
          for ix in range(0, nrow):
              for iy in range(0, ncol):
                  # Calculate the difference in V
                  dVdY[ix,iy] = (V[ix+1,iy] - V[ix,iy]) / dX
          #########################################################
          #### Rotational terms ###################################
  
              for ix in range(0, nrow):
                  for iy in range(0, ncol):
                      rotU[ix,iy] = rotConst[ix] * U[ix,iy] # Interpolated to U
                      rotV[ix,iy] = rotConst[ix] * V[ix,iy] # Interpolated to V
  
  
  
          ##########################################################  
          #### Time derivatives ####################################
          ## dUdT
          for ix in range(0, nrow):
              for iy in range(0, ncol):
                  dUdT[ix,iy] = (rotV[ix,iy]) - (flowConst * dHdX[ix,iy]) - (dragConst * U[ix,iy]) + windU[ix]
          ## dVdT
          for ix in range(0, nrow):
              for iy in range(0, ncol):       
                  dVdT[ix,iy] = (-rotU[ix,iy]) - (flowConst * dHdY[ix,iy]) - (dragConst * V[ix,iy])
          ## dHdT  
          for ix in range(0, nrow):
              for iy in range(0, ncol):
                  dHdT[ix,iy] = -(dUdX[ix,iy] + dVdY[ix,iy]) * (HBackground / dX)   
  
          # Step Forward One Time Step
          for ix in range(0,nrow):
              for iy in range(0,ncol):
                  U[ix,iy] = U[ix,iy] + (dUdT[ix,iy] * dT)
          for ix in range(0,nrow):
              for iy in range(0,ncol):
                  V[ix,iy] = V[ix,iy] + (dVdT[ix,iy] * dT)
          for ix in range(0,nrow):
              for iy in range(0,ncol):
                  H[ix,iy] = H[ix,iy] + (dHdT[ix,iy] * dT)
          ###########################################################        
          #### Maintain the ghost cells #############################
          # North wall velocity zero
          V[0,:] = 0
          # Horizontal wrapping
          if horizontalWrap is True:
              U[:,ncol] = U[:,0]
              H[:,ncol] = H[:,0]
          else:
              U[:,0] = 0
              U[:,ncol] = 0 
      itGlobal = itGlobal + ntAnim
  ###################################################################
  
  def firstFrame():
      global fig, ax, hPlot
      fig, ax = plt.subplots()
      ax.set_title("H")   
      hh = H[:,0:ncol]
      loc = tkr.IndexLocator(base=1, offset=1)
      ax.xaxis.set_major_locator(loc)
      ax.yaxis.set_major_locator(loc)
      grid = ax.grid(which='major', axis='both', linestyle='-')
      hPlot = ax.imshow(hh, interpolation='nearest', clim=(-0.5,0.5))   
      plotArrows()
      plt.show(block=False)
  
  
  def plotArrows():
      global quiv, quiv2
      xx = []
      yy = []
      uu = []
      vv = []
      for irow in range( 0, nrow ):
          for icol in range( 0, ncol ):
              xx.append(icol - 0.5)
              yy.append(irow )
              uu.append( U[irow,icol] * arrowScale )
              vv.append( 0 )
      quiv = ax.quiver( xx, yy, uu, vv, color='white', scale=1)
      for irow in range( 0, nrow ):
          for icol in range( 0, ncol ):
              xx.append(icol)
              yy.append(irow - 0.5)
              uu.append( 0 )
              vv.append( -V[irow,icol] * arrowScale )
      quiv2 = ax.quiver( xx, yy, uu, vv, color='white', scale=1)
  
  def updateFrame():
      global fig, ax, hPlot, quiv, quiv2
      hh = H[:,0:ncol]
      hPlot.set_array(hh)
      quiv.remove()    
      quiv2.remove()
      plotArrows()
      fig.canvas.draw()
      plt.pause(0.01)
      print("Time: ", math.floor( itGlobal * dT / 86400.*10)/10, "days")
      print("H: ", H[1,1])
  
  def textDump():
      #print("time step ", itGlobal)    
      #print("H", H)
      #print("rotV" )
      #print( rotV)
      #print("V" )
      #print( V)
      #print("dHdX" )
      #print( dHdX)
      #print("dHdY" )
      #print( dHdY)
      #print("dVdY" )
      #print( dVdY)
      #print("dHdT" )
      #print( dHdT)
      #print("dUdT" )
      #print( dUdT)
      #print("dVdT" )
      #print( dVdT)
      #print("inter_u")
      #print(inter_u)
      #print("inter_u1")
      #print(inter_u1)
      #print("inter_v")
      #print(inter_v)
      #print("rotU" )
      #print( rotU)
      #print("U" )
      #print( U)
      #print("dUdX" )
      #print( dUdX)
      #print("rotConst")
      #print(rotConst)
  
  
  if textOutput is True:
      textDump()
  if plotOutput is True:
      firstFrame()
  for i_anim_step in range(0,nSlices):
      animStep()
      if textOutput is True:
          textDump()
      if plotOutput is True:
          updateFrame()
  </code></pre>