#! /usr/bin/env python3
#
from fenics import *

def step5 ( my_grid, my_degree ):

#*****************************************************************************80
#
## step5 solves a Poisson equation with diffusivity |grad(W)|.
#
#  Discussion:
#
#    Use the mixed DPG method to solve a Poisson equation
#    with a piecewise constant diffusivity function kappa().
#
#      - div kappa(x,y) grad u = f in Omega = [0,1]x[0,1]
#                            u = 0 on dOmega
#
#    Use:
#
#      kappa(x,y) = |grad w| = sqrt ( wx^2 + wy^2 )
#
#      f(x,y) = 2*pi*pi*sin(pi*x) * sin(pi*y)
#
#    Where:
#
#      w(x,y) is, for right now, an arbitrary scalar field, which we 
#      will construct as the solution of
#      - div grad w = f in Omega = [0,1]x[0,1]
#                 w = 0 on dOmega
#
#    Later, we will be working with an iterative procedure, starting with
#    u0, and producing u1, u2, ... and w(x,y) will be the value of u
#    from the previous iterative step.
#
#    To my surprise, it seems that the degree of the Q space and element
#    must be set to 1.  Using a value of 2 results in nan values.
#    JVB, 13 November 2018.
#
#    In
#      Qs = FunctionSpace ( my_mesh, "BDM", 1, restriction = "facet" ) 
#    I can replace the obscure 'restriction = "facet"' argument by
#      Qs = FunctionSpace ( my_mesh, "BDMF", 1 )
#    and then I can use BDMF for the Qe statement as well.
#    JVB, 29 November 2018
#
#    I think it possible that the reason I couldn't use degree 2 or higher for
#    the BDM element was that I was not specifying a boundary condition.
#    If I add a zero boundary condition, the code runs.
#    JVB, 08 January 2019.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    28 November 2018
#
#  Author:
#
#    John Burkardt
#
#  Parameters:
#
#    Input, integer MY_GRID, the number of (pairs of) elements to use in the 
#    X and Y directions.  Number of triangular elements is 2*my_grid*my_grid.
#
#    Input, integer MY_DEGREE, specifies approximation degrees.
#    Error indicators: MY_DEGREE+2
#    Primary variable: MY_DEGREE+1
#    Fluxes:           1
#
  import matplotlib.pyplot as plt
  import numpy as np
#
#  Report input.
#
  print ( '' )
  print ( '  Case of %d x %d mesh' % ( my_grid, my_grid ) )
  print ( '  Approximation degree %d:' % ( my_degree ) )
#
#  Mesh the unit square.
#
  my_mesh = UnitSquareMesh ( my_grid, my_grid, 'right/left' )
#
#  Plot the mesh.
#
  plot ( my_mesh, title = 'step5 Mesh' )
  filename = 'step5_mesh.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.close ( )
#-------------------------------------------------------------------------------
#  Set up and solve for a field W.
#  |grad(w)| will be used as the diffusivity for U.
#-------------------------------------------------------------------------------
  f_expr = Expression ( '2 * pi * pi * sin ( pi * x[0] ) * sin ( pi * x[1] )', degree = 10 )
  Ws = FunctionSpace ( my_mesh, 'CG',  my_degree + 1 )
  wt = TrialFunction ( Ws )
  vt = TestFunction ( Ws )
  awv = dot ( grad ( wt ), grad ( vt ) ) * dx
  fv = f_expr * vt * dx
  w_expr = Expression ( "0.0", degree = 1 )
  wbc = DirichletBC ( Ws, w_expr, DomainBoundary() )
  w = Function ( Ws )
  solve ( awv == fv, w, wbc )
#
#  Plot w.
#
  fig = plot ( w, title = 'step5 w' )
  plt.colorbar ( fig )
  filename = 'step5_w.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.close ( )
#
#  Plot |grad w|.
#  I am surprised that grad(w)**2 means wx^2 + wy^2, that is, that the
#  squared terms are automatically added.
#
  fig = plot ( sqrt ( grad ( w )**2 ), title = 'step5 |grad w|' )
  plt.colorbar ( fig )
  filename = 'step5_w_grad_norm.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.close ( )

#-------------------------------------------------------------------------------
#  Now we seek to solve
#    - div |grad(W)| grad U = F
#-------------------------------------------------------------------------------
#  Define the right hand side F(X,Y):
#
  f_expr = Expression ( '2 * pi * pi * sin ( pi * x[0] ) * sin ( pi * x[1] )', degree = 10 )
#
#  Set function spaces for Error estimator, Primal variable, interfacial flux.
#
  Es = FunctionSpace ( my_mesh, 'DG',  my_degree + 2 )
  Us = FunctionSpace ( my_mesh, 'CG',  my_degree + 1 )
  Qs = FunctionSpace ( my_mesh, 'BDMF', 1 )
#
#  Set elements for Error estimator, Primal variable, interfacial flux.
#
  Ee = FiniteElement ( 'DG',  triangle, my_degree + 2 )
  Ue = FiniteElement ( 'CG',  triangle, my_degree + 1 )
  Qe = FiniteElement ( 'BDMF', triangle, 1 )
#
#  Define the mixed element, and the corresponding function space.
#
  EUQe = MixedElement ( Ee, Ue, Qe )
  EUQs = FunctionSpace ( my_mesh, EUQe )
#
#  Extract the individual trial and test function factors from the space.
#  Here, "e" for example is a symbol for typical trial functions from Es.
#
  ( e, u, q ) = TrialFunctions ( EUQs )
  ( et, ut, qt ) = TestFunctions ( EUQs )
#
#  Compute the normal vectors.
#
  n = FacetNormal ( my_mesh )
#
#  Define the inner product for the error estimator space.
#
  yip = dot ( grad ( e ), grad ( et ) ) * dx + e * et * dx
#
#  Set up the saddle point problem:
#
#  b( (u,q), et ) = ( kappa * grad u, grad et ) - < q.n, et >
#  This is an equation for U and Q.
#
  b1 = dot ( sqrt ( grad ( w )**2 ) * grad ( u ), grad ( et ) ) * dx \
     - dot ( q ( '+' ), n ( '+' ) ) * ( et ( '+' ) - et ( '-' ) ) * dS \
     - dot ( q, n ) * et * ds
#
#  b( (ut,qt), e ) = ( kappa * grad ut, grad e ) - < qt.n, e >
#  This is an equation for E.
#
  b2 = dot ( sqrt ( grad ( w )**2 ) * grad ( ut ), grad ( e ) ) * dx \
     - dot ( qt ( '+' ), n ( '+' ) ) * ( e ( '+' ) - e ( '-' ) ) * dS \
     - dot ( qt, n ) * e * ds
#
#  Set the saddle point problem:
#
#    yip + b1 = F * et * dx
#    b2       = 0
#
  a = yip + b1 + b2
  b = f_expr * et * dx
#
#  Apply the Dirichlet boundary condition to the second component (U).
#
  U = Expression ( "0", degree = 10 )
  bc = DirichletBC ( EUQs.sub(1), U, DomainBoundary() )
#
#  Solve.
#  This "e, u, q" is DIFFERENT from the e, u, q above!
#  Here, "e" is the actual finite element solution function.
#  Specify a nonlinear equation to be solved.
#
  euq = Function ( EUQs ) 
  solve ( a == b, euq, bcs = bc )
  e, u, q = euq.split ( deepcopy = True )
#
#  Plot the solution u.
#
  fig = plot ( u, title = 'step5 u' )
  plt.colorbar ( fig )
  filename = 'step5_u.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.close ( )
#
#  Plot |grad u|.
#
  gradun = sqrt ( grad ( u ) ** 2 )
  fig = plot ( gradun, title = 'step5 |grad u|' )
  plt.colorbar ( fig )
  filename = 'step5_u_grad_norm.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.close ( )
#
#  Plot the error indicators.
#
  fig = plot ( e, title = 'step5 indicators' )
  plt.colorbar ( fig )
  filename = 'step5_indicators.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.close ( )
#
#  Terminate.
#
  return

def step5_test ( ):

#*****************************************************************************80
#
## step5_test tests step5.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    25 October 2018
#
#  Author:
#
#    John Burkardt
#
  import dolfin
  import platform
  import time

  print ( time.ctime ( time.time() ) )
  print ( '' )
  print ( 'step5_test:' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
  print ( '  FENICS version %s'% ( dolfin.__version__ ) )
  print ( '  p-Laplacian investigation.' )
  print ( '  Diffusivity is |grad(w)| where w is a given scalar field.' )
#
#  Report level = only warnings or higher.
#
  level = 30
  set_log_level ( level )

  my_grid = 16
  my_degree = 1
  step5 ( my_grid, my_degree )
#
#  Terminate.
#
  print ( '' )
  print ( 'step5_test:' );
  print ( '  Normal end of execution.' )
  print ( '' )
  print ( time.ctime ( time.time() ) )

if ( __name__ == '__main__' ):

  step5_test ( )