#! /usr/bin/env python
#
def hw31 ( ):
#
## HW31 does homework problem 3.1
#
  print ""
  print "HW31"
  print "  Python version"
  print "  Solve the PDE with 6, 11, and 21 nodes."
  print "  Compute and print the L2 and H1 errors."
  print ""
  print "  Nodes       L2              H1"
  print ""
  node_num = 6

  for i in range ( 0, 5 ):
    l2, h1 = fem1d ( node_num )
    print "  %2d  %14f  %14f" % ( node_num, l2, h1 )
    node_num = 2 * node_num - 1

  return

def fem1d ( node_num ):
#
## FEM1D solves the PDE with a given number of nodes.
#
#  Modifications:
#
#    This code started out as a copy of the "fem1d.py" code.
#    We needed to make a function HW31 that called FEM1D.
#    We needed to make NODE_NUM an input quantity.
#    We needed to delete lots of print out.
#    We needed to compute the L2 and H1 error norms.
#    This means we had to add a function for the derivative of the
#    exact solution.
#
#  Location:
#
#    http://people.sc.fsu.edu/~jburkardt/classes/fem_2014/hw31/hw31.py
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    01 October 2014
#
#  Author:
#
#    John Burkardt
#
  from math import sqrt
  import matplotlib.pyplot as plt
  import numpy as np
  import scipy.linalg as la
#
#  NODE_NUM counts the nodes.  N counts the elements.
#
  n = node_num - 1


#
#  Define the mesh, N+1 points between A and B.
#  These will be X[0] through X[N].
#
  a = 0.0
  b = 1.0
# n = 5
  x = np.linspace ( a, b, n + 1 )
#
#  Set a 3 point quadrature rule on the reference interval [0,1].
#
  ng = 3

  xg = np.array ( ( \
    0.112701665379258311482073460022, \
    0.5, \
    0.887298334620741688517926539978 ) )

  wg = np.array ( ( \
    5.0 / 18.0, \
    8.0 / 18.0, \
    5.0 / 18.0 ) )
#
#  Compute the system matrix A and right hand side RHS.
#
  A = np.zeros ( ( n + 1, n + 1 ) )
  rhs = np.zeros ( n + 1 )
#
#  Look at element E: (0, 1, 2, ..., N-1).
#
  for e in range ( 0, n ):

    xl = x[e]
    xr = x[e+1]
#
#  Consider quadrature point Q: (0, 1, 2 ) in element E.
#
    for q in range ( 0, ng ):
#
#  Map XG and WG from [0,1] to
#      XQ and QQ in [XL,XR].
#
      xq = xl + xg[q] * ( xr - xl )
      wq = wg[q] * ( xr - xl )
#
#  Consider the I-th test function PHI(I,X) and its derivative PHI'(I,X).
#
      for i_local in range ( 0, 2 ):
        i = i_local + e

        if ( i_local == 0 ):
          phii = ( xq - xr ) / ( xl - xr )
          phiip = 1.0 / ( xl - xr )
        else:
          phii = ( xq - xl ) / ( xr - xl )
          phiip = 1.0 / ( xr - xl )

        rhs[i] = rhs[i] + wq * phii * rhs_fn ( xq )
#
#  Consider the J-th basis function PHI(J,X) and its derivative PHI'(J,X).
#  (It turns out we don't need PHI for this particular problem, only PHI')
#
        for j_local in range ( 0, 2 ):
          j = j_local + e

          if ( j_local == 0 ):
            phijp = 1.0 / ( xl - xr )
          else:
            phijp = 1.0 / ( xr - xl )

          A[i][j] = A[i][j] + wq * phiip * phijp
#
#  Modify the linear system to enforce the left boundary condition.
#
  A[0,0] = 1.0
  A[0,1:n+1] = 0.0
  rhs[0] = exact_fn ( x[0] )
#
#  Modify the linear system to enforce the right boundary condition.
#
  A[n,n] = 1.0
  A[n,0:n] = 0.0
  rhs[n] = exact_fn ( x[n] )
#
#  Solve the linear system.
#
  c = la.solve ( A, rhs )
#
#  Compute the L2 and H1 error norms.
#
  l2 = 0.0
  h1 = 0.0

  for e in range ( 0, n ):

    l = e
    r = e + 1

    xl = x[l]
    xr = x[r]

    for q in range ( 0, ng ):

      xq = xl + xg[q] * ( xr - xl )
      wq = wg[q] * ( xr - xl )

      phil = ( xr - xq ) / ( xr - xl )
      phir = ( xq - xl ) / ( xr - xl )

      phipl = -1 / ( xr - xl )
      phipr = +1 / ( xr - xl )

      ufem = c[l] * phil + c[r] * phir
      uexact = exact_fn ( xq )

      l2 = l2 + wq * ( ufem - uexact ) ** 2

      upfem = c[l] * phipl + c[r] * phipr
      upexact = exactp_fn ( xq )

      h1 = h1 + wq * ( upfem - upexact ) ** 2

  l2 = sqrt ( l2 )
  h1 = sqrt ( h1 )
#
#  Terminate.
#
  return l2, h1
#
#  That is the end of the main program.
#  Now we list some helper functions.
#
def exact_fn ( x ):
#
## EXACT_FN evaluates the exact solution.
#
  from math import exp
  value = x * ( 1 - x ) * exp ( x )
  return value

def exactp_fn ( x ):
#
## EXACTP_FN evaluates the derivative of the exact solution.
#
  from math import exp
  value = ( 1 - x - x * x ) * exp ( x )
  return value

def rhs_fn ( x ):
#
## RHS_FN evaluates the right hand side.
#
  from math import exp
  value = x * ( x + 3 ) * exp ( x )
  return value
#
#  If this script is called directly, then run it as a program.
#
if ( __name__ == '__main__' ):
  hw31 ( )