Fix bugs in big_K, big_M.

[numerical-analysis.git] / src / FEM / R1.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Finite Element Method (FEM) to solve the problem,
--
--     -(A(x)*u'(x))' + c(x)*u*(x) = f(x) on (a,b)
--
--   with one of the two types of boundary conditions,
--
--     * Dirichlet: u(a) = u(b) = 0
--
--     * A(a)u'(a) = alpha, A(b)u'(b) = beta
--
--   if c(x) = 0, then it is assumed that the boundary conditions are
--   Dirichlet.
--
--   The code creates a linear system,
--
--     (K + M)x = F
--
--   to be solved for @x@, where @K@ (\"big_K\") is the stiffness
--   matrix, @M@ (\"big_M\") is the mass matrix, and @F@ (\"big_F\")
--   is the load vector.
--
--   Warning: until PLU factorization is implemented, we can only
--   solve the resulting system if it's positive definite!
--
module FEM.R1
where

import qualified Algebra.Algebraic as Algebraic ( C )
import qualified Algebra.Field as Field ( C )
import qualified Algebra.RealField as RealField ( C )
import qualified Algebra.ToRational as ToRational ( C )
import Data.Vector.Fixed ( Arity, S )
import NumericPrelude
import qualified Prelude as P

import Integration.Gaussian ( gaussian )
import Linear.Matrix (
  Col,
  Mat(..),
  Row,
  (!!!),
  construct,
  element_sum2,
  fromList,
  ifoldl2,
  map2,
  nrows,
  rows2,
  set_idx,
  toList,
  transpose,
  zip2,
  zipwith2 )
import Linear.System ( solve_positive_definite )
import Piecewise ( Piecewise(..), from_intervals )
import Polynomials.Orthogonal ( legendre )

-- | Dirichlet boundary conditions. Since u(a)=u(b)=0 are fixed,
--   there's no additional information conveyed by this type.
data Dirichlet a = Dirichlet { domain_dirichlet :: Interval a }

-- | Neumann boundary conditions. @alpha@ specifies A(a)u'(b) and
--   @beta@ specifies A(b)u'(b).
data Neumann a =
  Neumann { domain_neumann :: Interval a,
            alpha :: a,
            beta :: a }

-- | Boundary conditions can be either Dirichlet or Neumann.
type BoundaryConditions a = Either (Dirichlet a) (Neumann a)

type Interval a = (a,a)

-- | The data for the PDE that we are attempting to solve.
data PDE a =
  PDE {
    -- | A(x)
    big_A   :: (a -> a),
    -- | c(x)
    c      :: (a -> a),
    -- | f(x)
    f      :: (a -> a),

    -- | The boundary conditions. The domain also specifies the
    --   boundary in R^1.
    bdy    :: BoundaryConditions a }



-- | Non-PDE parameters for the finite element method. The additional
--   type parameter @n@ should be a type-level representation of the
--   largest element in @max_degrees@. It needs to be known statically
--   for the dimensions of the pointer matrix. The parameter @l@ is
--   the number of global basis functions. It's equal to the number of
--   /internal/ mesh nodes (i.e. m-1), plus the sum of (p_i - 1) for
--   each p_i in max_degrees.
data Params m n l a =
  Params {
    -- | A partition of the domain.
    mesh :: Col m (Interval a),

    -- | A list of the highest-degree polynomial that we will use over
    --   each interval in our mesh.
    max_degrees :: Col m Int }


-- | The pointer matrix mapping local basis elements to global
--   ones. The (i,j)th entry of the matrix contains the index of the
--   global basis function corresponding to the jth local basis
--   function over element i.
--
--   TODO: This works for Dirichlet boundary conditions only.
--
--   Note that the number of columns in the matrix depends on the
--
--   Examples:
--
--   >>> import Data.Vector.Fixed ( N5, N6 )
--   >>> import Linear.Matrix ( Col5, fromList )
--   >>> import Naturals ( N19 )
--
--   >>> let p = fromList [[3],[3],[5],[4],[5]] :: Col5 Int
--   >>> let mesh = undefined :: Col5 (Int,Int)
--   >>> let params = Params mesh p :: Params N5 N5 N19 Int
--   >>> let row1 = [0,1,5,6,0,0] :: [Int]
--   >>> let row2 = [1,2,7,8,0,0] :: [Int]
--   >>> let row3 = [2,3,9,10,11,12] :: [Int]
--   >>> let row4 = [3,4,13,14,15,0] :: [Int]
--   >>> let row5 = [4,0,16,17,18,19] :: [Int]
--   >>> let expected = fromList [row1,row2,row3,row4,row5] :: Mat N5 N6 Int
--   >>> pointer params == expected
--   True
--
pointer :: (Arity m, Arity n, Arity l) => Params m n l a -> Mat m (S n) Int
pointer params =
  construct lambda
  where
    -- | The number of polynomials in the local basis for element i.
    numpolys :: Int -> Int
    numpolys i = ((max_degrees params) !!! (i,0)) + 1

    lambda i j

      -- The first two (linear) basis functions are easy to do.
      | j < 2 = if i + j >= (nrows $ mesh params) then 0 else i + j

      -- No local basis function exists for this j.
      | j >= (numpolys i) = 0

      -- The first row we handle as a special case.
      | i == 0 = (nrows $ mesh params) + j - 2

      -- The first entry for this row not corresponding to one of the
      -- linear polynomials (which we did in the first guard). This
      -- grabs the biggest number in the previous row and begins counting
      -- from there.
      | j == 2 = (lambda (i-1) (numpolys (i-1) - 1)) + 1

      -- If none of the other cases apply, we can compute our value by
      -- looking to the left in the matrix.
      | otherwise = (lambda i (j-1)) + 1



-- | An affine transform taking the interval [x1,x2] into [-1,1].
--
--   Examples:
--
--   >>> let phi = affine (-6,7)
--   >>> phi (-6)
--   -1.0
--   >>>  phi 7
--   1.0
--
affine :: Field.C a => (a,a) -> (a -> a)
affine (x1,x2) x = (fromInteger 2)*(x - x1)/(x2 - x1) - (fromInteger 1)

-- | The inverse of 'affine'. It should send [-1,1] into [x1,x2].
--
--   Examples:
--
--   >>> let phi = affine_inv (-6,7)
--   >>> phi (-1)
--   -6.0
--   >>> phi 1
--   7.0
--
affine_inv :: Field.C a => (a,a) -> (a -> a)
affine_inv (x1,x2) x =
  x*(x2 - x1)/two + (x1 + x2)/two
  where
    two = fromInteger 2


-- * Load vector

-- | Normalized integrals of orthogonal basis functions over
--   [-1,1]. The test case below comes from Sage where the
--   orthogonality of the polynomials' derivatives can easily be
--   tested.
--
--   Examples:
--
--   >>> import qualified Algebra.Absolute as Absolute ( abs )
--
--   >>> let expected = 2.99624907925257
--   >>> let actual = big_N 4 1.5 :: Double
--   >>> Absolute.abs (actual - expected) < 1e-12
--   True
--
big_N :: forall a. (Algebraic.C a, RealField.C a) => Integer -> a -> a
big_N k x
  | k < 0 = error "requested a negative basis function"
  | k == 0 = (one - x) / (fromInteger 2)
  | k == 1 = (one + x) / (fromInteger 2)
  | otherwise =
      coeff * ( legendre k x - legendre (k-2) x )
      where
        two = fromInteger 2
        four = fromInteger 4
        coeff = one / (sqrt (four*(fromInteger k) - two)) :: a


-- | A matrix containing 'big_N' functions indexed by their
--   element/number. Each row in the matrix represents a finite element
--   (i.e. an interval in the mesh). Within row @i@, column @j@ contains
--   the @j@th 'big_N' basis function.
--
--   Any given 'big_N' will probably wind up in this matrix multiple
--   times; the basis functions don't change depending on the
--   interval. Only the /number/ of basis functions does. Computing
--   them this way allows us to easily construct a lookup \"table\" of
--   the proper dimensions.
--
--   The second example below relies on the fact that @big_N 3@ and
--   @big_N 6@ expand to Legendre polynomials (2,4) and (5,7)
--   respectively and so should be orthogonal over [-1,1].
--
--   Examples:
--
--   >>> import Data.Vector.Fixed ( N5, N6 )
--   >>> import Integration.Gaussian ( gaussian )
--   >>> import Linear.Matrix ( Col5, fromList )
--   >>> import Naturals ( N19 )
--
--   >>> let p = fromList [[3],[3],[5],[4],[5]] :: Col5 Int
--   >>> let mesh = undefined :: Col5 (Double,Double)
--   >>> let params = Params mesh p :: Params N5 N5 N19 Double
--   >>> let big_ns = big_N_matrix :: Mat N5 N6 (Double -> Double)
--   >>> let n1 = big_ns !!! (1,0)
--   >>> let n4 = big_ns !!! (4,0)
--   >>> n1 1.5 == n4 1.5
--   True
--   >>> let n1 = big_ns !!! (1,3)
--   >>> let n2 = big_ns !!! (2,4)
--   >>> gaussian (\x -> (n1 x) * (n2 x)) < 1e-12
--   True
--
big_N_matrix :: (Arity m, Arity n, Algebraic.C a, RealField.C a)
             => Mat m n (a -> a)
big_N_matrix =
  construct lambda
  where
    lambda _ j x = big_N (toInteger j) x


-- | The matrix of (N_i * N_j) functions used in the integrand of
--   the mass matrices.
big_Ns_matrix :: (Arity m, Arity n, Algebraic.C a, RealField.C a)
               => Mat m n (a -> a)
big_Ns_matrix =
  construct lambda
  where
    lambda i j x = (big_N (toInteger i) x) * (big_N (toInteger j) x)


-- | Compute the global load vector F.
--
--   Examples:
--
--   >>> import Linear.Matrix ( Col7, frobenius_norm )
--   >>> import FEM.R1.Example1 ( pde', params' )
--
--   >>> let f1 = [0.0418]
--   >>> let f2 = [0.0805]
--   >>> let f3 = [0.1007]
--   >>> let f4 = [-0.0045]
--   >>> let f5 = [-0.0332]
--   >>> let f6 = [-0.0054]
--   >>> let f7 = [-0.0267]
--   >>> let expected = fromList [f1,f2,f3,f4,f5,f6,f7] :: Col7 Double
--   >>> let actual = big_F pde' params'
--   >>> frobenius_norm (actual - expected) < 1e-4
--   True
--
big_F :: forall m n l a.
         (Arity l, Arity m, Arity n,
          Algebraic.C a, RealField.C a, ToRational.C a)
      => PDE a
      -> Params m n l a
      -> Col l a
big_F pde params =
  ifoldl2 accum zero (big_N_matrix :: Mat m (S n) (a -> a))
  where
    accum :: Int -> Int -> Col l a -> (a -> a) -> Col l a
    accum i j prev_F this_N =
      prev_F + this_F
      where
        two = fromInteger 2
        (x1,x2) = (mesh params) !!! (i,0)
        q = affine_inv (x1,x2)
        integrand x = ((f pde) (q x)) * (this_N x)

        -- The pointer matrix numbers from 1 so subtract one here to
        -- get the right index.
        k = ((pointer params) !!! (i,j)) - 1
        integral = (gaussian integrand)*(x2 - x1) / two
        this_F = set_idx zero (k,0) integral


-- * Stiffness matrix

-- | Derivatives of the 'big_N's, that is, orthogonal basis functions
--   over [-1,1]. The test case below comes from Sage where the
--   orthogonality of the polynomials' derivatives can easily be
--   tested. The indices are shifted by one so that k=0 is the first
--   basis function.
--
--   Examples:
--
--   >>> import qualified Algebra.Absolute as Absolute ( abs )
--
--   >>> let expected = 11.5757525403319
--   >>> let actual = big_N' 3 1.5 :: Double
--   >>> Absolute.abs (actual - expected) < 1e-10
--   True
--
big_N' :: forall a. (Algebraic.C a, RealField.C a) => Integer -> a -> a
big_N' k x
  | k < 0 = error "requested a negative basis function"
  | k == 0 = negate ( one / (fromInteger 2))
  | k == 1 = one / (fromInteger 2)
  | otherwise = coeff * ( legendre k x )
      where
        two = fromInteger 2
        coeff = sqrt ((two*(fromInteger k) + one) / two) :: a


-- | The matrix of (N_i' * N_j') functions used in the integrand of
--   the stiffness matrix.
big_N's_matrix :: (Arity m, Arity n, Algebraic.C a, RealField.C a)
               => Mat m n (a -> a)
big_N's_matrix =
  construct lambda
  where
    lambda i j x = (big_N' (toInteger i) x) * (big_N' (toInteger j) x)


big_K_elem :: forall m n l a b.
         (Arity l, Arity m, Arity n,
          Algebraic.C a, RealField.C a, ToRational.C a)
      => PDE a
      -> Params m n l a
      -> Int
      -> Int
      -> Mat l l a
      -> b
      -> Mat l l a
big_K_elem pde params _ k cur_K _ =
  ifoldl2 accum cur_K (big_N's_matrix :: Mat (S n) (S n) (a -> a))
  where
    accum :: Int -> Int -> Mat l l a -> (a -> a) -> Mat l l a
    accum i j prev_K these_N's =
      prev_K + this_K
      where
        two = fromInteger 2
        (x1,x2) = (mesh params) !!! (k,0)
        q = affine_inv (x1,x2)
        integrand x = ((big_A pde) (q x)) * (these_N's x)
        -- The pointer matrix numbers from 1 so subtract one below to
        -- get the right index. The indices i,j have upper bounds
        -- dependent on the element k. Since we statically create the
        -- matrix of basis function derivatives, we have to check here
        -- whether or not i,j exceed the max index.
        row_idx = ((pointer params) !!! (k,i)) - 1
        col_idx = ((pointer params) !!! (k,j)) - 1
        integral = (two/(x2 - x1))* (gaussian integrand)
        this_K = set_idx zero (row_idx, col_idx) integral



-- | Compute the \"big K\" stiffness matrix. There are three
--   parameters needed for K, namely i,j,k so a fold over a matrix will
--   not do. This little gimmick simulates a three-index fold by doing a
--   two-index fold over a row of the proper dimensions.
--
--   Examples:
--
--   >>> import Linear.Matrix ( Mat7, frobenius_norm )
--   >>> import FEM.R1.Example1 ( pde', params' )
--
--   >>> let k1 = [6, -3, 0, 0, 0, 0, 0] :: [Double]
--   >>> let k2 = [-3, 10.5, -7.5, 0, 0, 0, 0] :: [Double]
--   >>> let k3 = [0, -7.5, 12.5, 0, 0, 0, 0] :: [Double]
--   >>> let k4 = [0, 0, 0, 6, 0, 0, 0] :: [Double]
--   >>> let k5 = [0, 0, 0, 0, 6, 0, 0] :: [Double]
--   >>> let k6 = [0, 0, 0, 0, 0, 6, 0] :: [Double]
--   >>> let k7 = [0, 0, 0, 0, 0, 0, 15] :: [Double]
--   >>> let expected = fromList [k1,k2,k3,k4,k5,k6,k7] :: Mat7 Double
--   >>> let actual = big_K pde' params'
--   >>> frobenius_norm (actual - expected) < 1e-10
--   True
--
big_K :: forall m n l a.
         (Arity l, Arity m, Arity n,
          Algebraic.C a, RealField.C a, ToRational.C a)
      => PDE a
      -> Params m n l a
      -> Mat l l a
big_K pde params =
  ifoldl2 (big_K_elem pde params) zero col_idxs
  where
    col_idxs = fromList [map fromInteger [0..]] :: Row m a


-- * Mass matrix

big_M_elem :: forall m n l a b.
         (Arity l, Arity m, Arity n,
          Algebraic.C a, RealField.C a, ToRational.C a)
      => PDE a
      -> Params m n l a
      -> Int
      -> Int
      -> Mat l l a
      -> b
      -> Mat l l a
big_M_elem pde params _ k cur_M _ =
  ifoldl2 accum cur_M (big_Ns_matrix :: Mat (S n) (S n) (a -> a))
  where
    accum :: Int -> Int -> Mat l l a -> (a -> a) -> Mat l l a
    accum i j prev_M these_Ns =
      prev_M + this_M
      where
        two = fromInteger 2
        (x1,x2) = (mesh params) !!! (k,0)
        q = affine_inv (x1,x2)
        integrand x = ((c pde) (q x)) * (these_Ns x)
        -- The pointer matrix numbers from 1 so subtract one here to
        -- get the right index.
        row_idx = ((pointer params) !!! (k,i)) - 1
        col_idx = ((pointer params) !!! (k,j)) - 1
        integral = (x2 - x1)*(gaussian integrand) / two
        this_M = set_idx zero (row_idx, col_idx) integral


-- | Compute the \"big M\" mass matrix. There are three
--   parameters needed for M, namely i,j,k so a fold over a matrix will
--   not do. This little gimmick simulates a three-index fold by doing a
--   two-index fold over a row of the proper dimensions.
--
--   Examples:
--
--   >>> import Linear.Matrix ( Mat7, frobenius_norm )
--   >>> import FEM.R1.Example1 ( pde', params' )
--
--   >>> let m1 = [0.0723,0.0266,0,-0.0135,-0.0305,0.0058,0] :: [Double]
--   >>> let m2 = [0.0266,0.0897,0.0149,0,-0.0345,-0.0109,-0.0179] :: [Double]
--   >>> let m3 = [0,0.0149,0.0809,0,0,0,-0.0185] :: [Double]
--   >>> let m4 = [-0.0135,0,0,0.0110,0,0,0] :: [Double]
--   >>> let m5 = [-0.0305,-0.0345,0,0,0.0319,0.0018,0] :: [Double]
--   >>> let m6 = [0.0058,-0.0109,0,0,0.0018,0.0076,0] :: [Double]
--   >>> let m7 = [0,-0.0179,-0.0185,0,0,0,0.0178] :: [Double]
--
--   >>> let expected = fromList [m1,m2,m3,m4,m5,m6,m7] :: Mat7 Double
--   >>> let actual = big_M pde' params'
--   >>> frobenius_norm (actual - expected) < 1e-3
--   True
--
big_M :: forall m n l a.
         (Arity l, Arity m, Arity n,
          Algebraic.C a, RealField.C a, ToRational.C a)
      => PDE a
      -> Params m n l a
      -> Mat l l a
big_M pde params =
  ifoldl2 (big_M_elem pde params) zero col_idxs
  where
    col_idxs = fromList [map fromInteger [0..]] :: Row m a



-- | Determine the coefficient vector @x@ from the system @(K + M)x = F@.
--
--   Examples:
--
--   >>> import Linear.Matrix ( Col7, frobenius_norm )
--   >>> import FEM.R1.Example1 ( pde', params' )
--
--   >>> let c1 = [0.02366220347687] :: [Double]
--   >>> let c2 = [0.03431630082636] :: [Double]
--   >>> let c3 = [0.02841800893264] :: [Double]
--   >>> let c4 = [-0.00069489654996] :: [Double]
--   >>> let c5 = [-0.00518637005151] :: [Double]
--   >>> let c6 = [-0.00085028505337] :: [Double]
--   >>> let c7 = [-0.00170478210110] :: [Double]
--   >>>
--   >>> let expected = fromList [c1,c2,c3,c4,c5,c6,c7] :: Col7 Double
--   >>> let actual = coefficients pde' params'
--   >>> frobenius_norm (actual - expected) < 1e-8
--   True
--
coefficients :: forall m n l a.
                (Arity m, Arity n, Arity l,
                 Algebraic.C a, Eq a, RealField.C a, ToRational.C a)
             => PDE a
             -> Params m n (S l) a
             -> Col (S l) a
coefficients pde params =
  solve_positive_definite matrix b
  where
    matrix = (big_K pde params) + (big_M pde params)
    b      = big_F pde params


solution :: forall m n l a.
            (Arity m, Arity n, Arity l,
            Algebraic.C a, Eq a, RealField.C a, ToRational.C a, Show a)
         => PDE a
         -> Params m n (S l) a
         -> Piecewise a
solution pde params =
  from_intervals $ map head $ toList $ solved_column
  where
    global_coeffs :: Col (S l) a
    global_coeffs = coefficients pde params

    ptr :: Mat m (S n) Int
    ptr = pointer params

    -- Each mesh element has an associated row in the pointer
    -- matrix. Stick them together.
    mesh_with_ptr_rows :: Col m (Interval a, Row (S n) Int)
    mesh_with_ptr_rows = zip2 (mesh params) (rows2 ptr)

    make_local_coeffs :: (Interval a, Row (S n) Int) -> Row (S n) a
    make_local_coeffs (_, ptr_row) =
      construct lambda
        where
          lambda _ j = if (ptr_row !!! (0,j)) == zero
                       then zero
                       else global_coeffs !!! ((ptr_row !!! (0,j)) - 1, 0)

    -- Create a column vector for each mesh element containing the global
    -- coefficients corresponding to that element.
    local_coeffs :: Col m (Row (S n) a)
    local_coeffs = map2 make_local_coeffs mesh_with_ptr_rows

    global_basis_functions :: Col (S n) (a -> a)
    global_basis_functions =
      construct lambda
        where lambda i _ = big_N (toInteger i)

    mesh_with_coeffs :: Col m (Interval a, Row (S n) a)
    mesh_with_coeffs = zip2 (mesh params) local_coeffs

    solved_column :: Col m (Interval a, (a -> a))
    solved_column = map2 solve_piece $ mesh_with_coeffs

    solve_piece :: (Interval a, Row (S n) a) -> (Interval a, (a -> a))
    solve_piece (interval, coeffs_row) = (interval, g)
      where
        coeffs_col = transpose coeffs_row

        g x = element_sum2 $ zipwith2 combine coeffs_col global_basis_functions
          where
            xi = (affine interval) x
            combine ci ni = ci*(ni xi)