#! /usr/bin/env python3
#
def shallow_water_1d_test ( ):

#*****************************************************************************80
#
## shallow_water_1d_test() tests shallow_water_1d().
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    09 December 2016
#
#  Author:
#
#    John Burkardt
#
#  Reference:
#
#    Cleve Moler,
#    "The Shallow Water Equations",
#    Experiments with MATLAB.
#
  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'shallow_water_1d_test:' )
  print ( '  Compute a solution of the discrete shallow water equations' )
  print ( '  over a 1-dimensional domain.' )
#
#  Set parameters.
#
  nx = 41
  nt = 100
  x_length = 1.0
  t_length = 0.2
  g = 9.8
#
#  Compute H and UH.
#
  h_array, uh_array, x, t = shallow_water_1d ( nx, nt, x_length, t_length, g )

  x_min = min ( x )
  x_max = max ( x )
  
  h_min = 0.0
  h_max = np.amax ( h_array )
 
  uh_max = np.amax ( uh_array )
  uh_min = np.amin ( uh_array )
#
#  Animation of H.
#
  for it in range ( 0, nt + 1 ):

    plt.axis ( [ x_min, x_max, h_min, h_max ] )
    plt.fill_between ( x, 0, h_array[:,it] )
    title_string = ( 'H(T), Step %3d, Time = %f' % ( it, t[it] ) )
    plt.title ( title_string )
    plt.xlabel ( 'X' )
    plt.ylabel ( 'H(X,T)' )
    plt.show ( block = False )
    plt.close ( )
#
#  Animation of UH.
#
  for it in range ( 0, nt + 1 ):

    plt.axis ( [ x_min, x_max, h_min, h_max ] )
    plt.fill_between ( x, 0, uh_array[:,it] )
    title_string = ( 'UH(T), Step %3d, Time = %f' % ( it, t[it] ) )
    plt.title ( title_string )
    plt.xlabel ( 'X' )
    plt.ylabel ( 'UH(X,T)' )
    plt.show ( block = False )
    plt.close ( )
#
#  Terminate.
#
  print ( '' )
  print ( 'shallow_water_1d_test:' )
  print ( '  Normal end of execution.' )
  return

def shallow_water_1d ( nx, nt, x_length, t_length, g ):

#*****************************************************************************80
#
## shallow_water_1d() approximates the 1D shallow water equations.
#
#  Discussion:
#
#    This code can be considered a 1D version of Cleve Moler's shallow
#    water equation solver.
#
#    The version of the shallow water equations being solved here is in
#    conservative form, and omits the Coriolis force.  The state variables
#    are H (the height) and UH (the mass velocity).
#
#    The equations have the form
#
#      dH/dt + d UH/dx = 0
#
#      d UH/dt + d ( U^2 H + 1/2 g H^2 )/dx = 0
#
#    Here U is the ordinary velocity, U = UH/H, and g is the gravitational
#    acceleration.
#
#    The initial conditions are used to specify ( H, UH ) at an equally
#    spaced set of points, and then the Lax-Wendroff method is used to advance
#    the solution through a number of equally spaced points in time, with 
#    boundary conditions supplying the first and last spatial values.
#
#
#    Some input values will result in an unstable calculation that
#    quickly blows up.  This is related to the Courant-Friedrichs-Lewy
#    condition, which requires that DT be small enough, relative to DX and
#    the velocity, that information cannot cross an entire cell.
#
#    A "reasonable" set of input quantities is
#
#      shallow_water_1d ( 41, 100, 1.0, 0.2, 9.8 )
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    09 December 2016
#
#  Author:
#
#    John Burkardt
#
#  Reference:
#
#    Cleve Moler,
#    "The Shallow Water Equations",
#    Experiments with MATLAB.
#
#  Input:
#
#    integer NX, the number of spatial nodes.
#
#    integer NT, the number of times steps.
#
#    real X_LENGTH, the length of the region.
#
#    real T_LENGTH, the time extent.
#
#    real G, the gravity constant.  G = 9.8 meters per second^2.
#
#  Output:
#
#    real H_ARRAY(NX,NT+1), the height for all space and time points.
#
#    real UH_ARRAY(NX,NT+1), the mass velocity for all space and time points.
#
#    real X(NX), the X coordinates.
#
#    real T(NT+1), the T coordinates.
#
  import numpy as np
  import platform

  print ( '' )
  print ( 'shallow_water_1d:' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
#
#  Allocate vectors.
#
  h = np.zeros ( nx )
  uh = np.zeros ( nx )
  hm = np.zeros ( nx - 1 )
  uhm = np.zeros ( nx - 1 )
  x = np.zeros ( nx )
  t = np.zeros ( nt + 1 )
  h_array = np.zeros ( [ nx, nt + 1 ] )
  uh_array = np.zeros ( [ nx, nt + 1 ] )
#
#  Define the locations of the nodes and time steps and the spacing.
#
  x = np.linspace ( 0, x_length, nx )
  t = np.linspace ( 0, t_length, nt + 1 )

  dx = x_length / float ( nx - 1 )
  dt = t_length / float ( nt )
#
#  Apply the initial conditions.
#
  h, uh = initial_conditions ( nx, nt, h, uh, x )
#
#  Apply the boundary conditions.
#
  h, uh = boundary_conditions ( nx, nt, h, uh, t[0] )
#
#  Store the first time step into H_ARRAY and UH_ARRAY.
#
  h_array[0:nx,0] = h[0:nx]
  uh_array[0:nx,0] = uh[0:nx]
#
#  Take NT more time steps.
#
  for it in range ( 1, nt + 1 ):
#
#  Take a half time step, estimating H and UH at the NX-1 spatial midpoints.
#
    hm[0:nx-1] = ( h[0:nx-1] + h[1:nx] ) / 2.0 \
      - ( dt / 2.0 ) * ( uh[1:nx] - uh[0:nx-1] ) / dx

    uhm[0:nx-1] = ( uh[0:nx-1] + uh[1:nx] ) / 2.0 \
      - ( dt / 2.0 ) * ( \
        uh[1:nx] ** 2    / h[1:nx]   + 0.5 * g * h[1:nx] ** 2 \
      - uh[0:nx-1] ** 2  / h[0:nx-1] - 0.5 * g * h[0:nx-1] ** 2 ) / dx
#
#  Take a full time step, evaluating the derivative at the half time step,
#  to estimate the solution at the NX-2 nodes.
#
    h[1:nx-1] = h[1:nx-1] \
      - dt * ( uhm[1:nx-1] - uhm[0:nx-2] ) / dx

    uh[1:nx-1] = uh[1:nx-1] \
      - dt * ( \
        uhm[1:nx-1] ** 2  / hm[1:nx-1] + 0.5 * g * hm[1:nx-1] ** 2 \
      - uhm[0:nx-2] ** 2  / hm[0:nx-2] - 0.5 * g * hm[0:nx-2] ** 2 ) / dx
#
#  Update the boundary conditions.
#
    h, uh = boundary_conditions ( nx, nt, h, uh, t[it] )
#
#  Copy data into the big arrays.
#
    h_array[0:nx,it] = h[0:nx]
    uh_array[0:nx,it] = uh[0:nx]
#
#  Write data.
#
  filename_x = 'sw1d_x.txt'
  filename_t = 'sw1d_t.txt'
  filename_h = 'sw1d_h.txt'
  filename_uh = 'sw1d_uh.txt'

  r8vec_write ( filename_x, nx, x )
  r8vec_write ( filename_t, nt + 1, t )
  r8mat_write ( filename_h, nx, nt + 1, h_array )
  r8mat_write ( filename_uh, nx, nt + 1, uh_array )

  print ( '' )
  print ( '  X  values saved in file "%s"' % ( filename_x ) )
  print ( '  T  values saved in file "%s"' % ( filename_t ) )
  print ( '  H  values saved in file "%s"' % ( filename_h ) )
  print ( '  UH values saved in file "%s"' % ( filename_uh ) )
#
#  Terminate.
#
  print ( '' )
  print ( 'shallow_water_1d:' )
  print ( '  Normal end of execution.' )

  return h_array, uh_array, x, t 

def boundary_conditions ( nx, nt, h, uh, t ):

#*****************************************************************************80
#
## boundary_conditions() sets the boundary conditions.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    09 December 2016
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer NX, the number of spatial nodes.
#
#    integer NT, the number of times steps.
#
#    real H(NX), the height for all space.
#
#    real UH(NX), the mass velocity for all space.
#
#    real T, the current time.
#
#  Output:
#
#    real H(NX), the height, with H(1) and H(NX) adjusted for
#    boundary conditions.
#
#    real UH(NX), the mass velocity, with UH(1) and UH(NX)
#    adjusted for boundary conditions.
#
  bc = 1
#
#  Periodic boundary conditions on H and UH.
#
  if ( bc == 1 ):
    h[0] = h[nx-2]
    h[nx-1] = h[1]
    uh[0] = uh[nx-2]
    uh[nx-1] = uh[1]
#
#  Free boundary conditions on H and UH.
#
  elif ( bc == 2 ):
    h[0] = h[1]
    h[nx-1] = h[nx-2]
    uh[0] = uh[1]
    uh[nx-1] = uh[nx-2]
#
#  Reflective boundary conditions on UH, free boundary conditions on H.
#
  elif ( bc == 3 ):
    h[0] = h[1]
    h[nx-1] = h[nx-2]
    uh[0] = - uh[1]
    uh[nx-1] = - uh[nx-2]

  return h, uh

def initial_conditions ( nx, nt, h, uh, x ):

#*****************************************************************************80
#
## initial_conditions() sets the initial conditions.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    09 December 2016
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer NX, the number of spatial nodes.
#
#    integer NT, the number of times steps.
#
#    real H(NX,1), an array to hold the height.
#
#    real UH(NX,1), an array to hold the mass velocity.
#
#    real X(NX,1), the coordinates of the nodes.
#
#  Output:
#
#    real H(NX,1), the initial height for all space.
#
#    real UH(NX,1), the initial mass velocity for all space.
#
  import numpy as np

  h = 2.0 + np.sin ( 2.0 * np.pi * x )
  uh = np.zeros ( nx )

  return h, uh

def r8mat_write ( filename, m, n, a ):

#*****************************************************************************80
#
## r8mat_write() writes an R8MAT to a file.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    12 October 2014
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    string FILENAME, the name of the output file.
#
#    integer M, the number of rows in A.
#
#    integer N, the number of columns in A.
#
#    real A(M,N), the matrix.
#
  output = open ( filename, 'w' )

  for i in range ( 0, m ):
    for j in range ( 0, n ):
      s = '  %g' % ( a[i,j] )
      output.write ( s )
    output.write ( '\n' )

  output.close ( )

  return

def r8mat_write_test ( ):

#*****************************************************************************80
#
## r8mat_write_test() tests r8mat_write().
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    12 October 2014
#
#  Author:
#
#    John Burkardt
#
  import numpy as np
  import platform

  print ( '' )
  print ( 'r8mat_write_test:' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
  print ( '  r8mat_write writes an R8MAT to a file.' )

  filename = 'r8mat_write_test.txt'
  m = 5
  n = 3
  a = np.array ( (  \
    ( 1.1, 1.2, 1.3 ), \
    ( 2.1, 2.2, 2.3 ), \
    ( 3.1, 3.2, 3.3 ), \
    ( 4.1, 4.2, 4.3 ), \
    ( 5.1, 5.2, 5.3 ) ) )
  r8mat_write ( filename, m, n, a )

  print ( '' )
  print ( '  Created file "%s".' % ( filename ) )
#
#  Terminate.
#
  print ( '' )
  print ( 'r8mat_write_test:' )
  print ( '  Normal end of execution.' )
  return

def r8vec_write ( filename, n, a ):

#*****************************************************************************80
#
## r8vec_write() writes an R8VEC to a file.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    06 November 2014
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    string FILENAME, the name of the output file.
#
#    integer N, the number of entries in A.
#
#    real A(N), the matrix.
#
  output = open ( filename, 'w' )

  for i in range ( 0, n ):
    s = '  %g\n' % ( a[i] )
    output.write ( s )

  output.close ( )

  return

def r8vec_write_test ( ):

#*****************************************************************************80
#
## r8vec_write_test() tests r8vec_write().
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    06 November 2014
#
#  Author:
#
#    John Burkardt
#
  import numpy as np
  import platform

  print ( '' )
  print ( 'r8vec_write_test:' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
  print ( '  r8vec_write writes an R8VEC to a file.' )

  filename = 'r8vec_write_test.txt'
  n = 5
  a = np.array ( ( 1.1, 2.2, 3.3, 4.4, 5.5 ) )
  r8vec_write ( filename, n, a )

  print ( '' )
  print ( '  Created file "%s".' % ( filename ) )
#
#  Terminate.
#
  print ( '' )
  print ( 'r8vec_write_test:' )
  print ( '  Normal end of execution.' )
  return
  
def timestamp ( ):

#*****************************************************************************80
#
## timestamp() prints the date as a timestamp.
#
#  Licensing:
#
#    This code is distributed under the MIT license. 
#
#  Modified:
#
#    06 April 2013
#
#  Author:
#
#    John Burkardt
#
  import time

  t = time.time ( )
  print ( time.ctime ( t ) )

  return None

if ( __name__ == '__main__' ):
  timestamp ( )
  shallow_water_1d_test ( )
  timestamp ( )