matrix.egi

--
-- Matrices
--

inductive pattern Matrix a:=
  | quadCons (Matrix a) (Matrix a) (Matrix a) (Matrix a)
  | matCons Integer Integer a (Matrix a) (Matrix a) (Matrix a) (Matrix a)

def matrix : Matcher (Matrix MathValue) :=
  matcher
    | quadCons $ $ $ $ as (mathValue, matrix, matrix, matrix) with
      | $tgt ->
        match tensorShape tgt as list integer with
          | $m :: $n :: _ ->
            [(tgt_1_1, tgt_1_(2, n), tgt_(2, m)_1, tgt_(2, m)_(2, n))]
          | _ -> []
    | matCons #$i #$j $ $ $ $ $ as (mathValue, matrix, matrix, matrix, matrix) with
      | $tgt ->
        let ns := tensorShape tgt
            m := nth 1 ns
            n := nth 2 ns
         in [ ( tgt_i_j
            , tgt_(1, i - 1)_(1, j - 1)
            , tgt_(1, i - 1)_(j + 1, n)
            , tgt_(i + 1, m)_(1, j - 1)
            , tgt_(i + 1, m)_(j + 1, n) ) ]
    | #$val as () with
      | $tgt -> if val = tgt then [()] else []
    | $ as (something) with
      | $tgt -> [tgt]

def M.tensorIndexValue (index: TensorIndex) : MathValue :=
  match index as tensorIndex with
    | (subIndex $x | supIndex $x | diagIndex $x | userIndex $x) -> x

def M.attachOrdinaryIndex {a}
      (value: Matrix a) (index: TensorIndex) : Matrix a :=
  let variance := tensorIndexVariance index
      x := M.tensorIndexValue index
   in match assert "ordinary matrix indices must be up or down"
                   (variance = "down" || variance = "up")
            as bool with
        | #True ->
            if variance = "down"
              then subrefs value [x]
              else suprefs value [x]

-- Preserve an ordinary matrix's explicit index contract after local
-- `withSymbols` indices have become anonymous result axes.
def M.copyOrdinaryIndices {a}
      (source: Matrix a) (value: Matrix a) : Matrix a :=
  let indices := tensorIndices source
   in match assert
              "matrix result must match its source shape and ordinary indices"
              (length (tensorShape source) = 2
               && tensorShape source = tensorShape value
               && (indices = [] || length indices = 2)
               && all
                    (\variance -> variance = "down" || variance = "up")
                    (tensorVariances source))
        as bool with
        | #True -> foldl M.attachOrdinaryIndex value indices

def trace {Ring a} (t: Matrix a) : a :=
  withSymbols [i]
    contractWith (+) t~i_i

def sym {Field a} (a: Matrix a) : Matrix a :=
  M.copyOrdinaryIndices a
    (withSymbols [i, j]
      ((a_i_j + a_j_i) / 2))

def antisym {Field a} (a: Matrix a) : Matrix a :=
  M.copyOrdinaryIndices a
    (withSymbols [i, j]
      ((a_i_j - a_j_i) / 2))

def M.inverse (m: Matrix MathValue) : Matrix MathValue :=
  let d := M.det m
   in generateTensor
        (\[i, j] ->
          match m as matrix with
          | matCons #j #i _ $A $B $C $D ->
            if isEven (i + j)
              then M.det (M.join A B C D) / d
              else - (M.det (M.join A B C D) / d))
        (tensorShape m)

def M.* (s: Matrix MathValue) (t: Matrix MathValue) : Matrix MathValue := 
  withSymbols [i, j, k] (s~i~j . t_j_k)

def M.*' (s: Matrix MathValue) (t: Matrix MathValue) : Matrix MathValue :=
  withSymbols [i, j, k] (s~i~j .' t_j_k)

-- Matrix-vector multiplication.  M.* indexes its right operand twice
-- (t_j_k), which requires rank 2; a vector operand contracts on the
-- single shared index instead.
def MV.* (s: Matrix MathValue) (t: Vector MathValue) : Vector MathValue :=
  withSymbols [i, j] (s~i~j . t_j)

def M.power (t: Matrix MathValue) (k: Integer) : Matrix MathValue := 
  foldl M.* t (take (k - 1) (repeat1 t))

def M.comm (m1: Matrix MathValue) (m2: Matrix MathValue) : Matrix MathValue := 
  withSymbols [i, j, k] m1~i~j . m2_j_k - m2~i~j . m1_j_k

def M.join (A: Matrix MathValue) (B: Matrix MathValue) (C: Matrix MathValue) (D: Matrix MathValue)
  : Matrix MathValue :=
  let ashape := tensorShape A
      bshape := tensorShape B
      cshape := tensorShape C
      dshape := tensorShape D
  in let a1 := nth 1 ashape
         a2 := nth 2 ashape
         b1 := nth 1 bshape
         b2 := nth 2 bshape
         c1 := nth 1 cshape
         c2 := nth 2 cshape
         d1 := nth 1 dshape
         d2 := nth 2 dshape
     in let m1 := max a1 b1
            m2 := max a2 c2
            n1 := max c1 d1
            n2 := max b2 d2
        in generateTensor
             (\match as list integer with
               | [$i & ?(<= a1), $j & ?(<= a2)] -> A_i_j
               | [$i & ?(<= m1), $j]            -> B_i_(j - a2)
               | [$i,            $j & ?(<= m2)] -> C_(i - a1)_j
               | [$i,            $j]            -> D_(i - m1)_(j - m2))
             [m1 + n1, m2 + n2]

--
-- Determinant
--
def M.determinant (m: Matrix MathValue) : MathValue :=
  match tensorShape m as list integer with
    | [#0, #0] -> 1
    | [$n, #n] ->
      let (es, os) := evenAndOddPermutations' n
       in sum (map (\e -> product (map2 (\i j -> m_i_j) (between 1 n) e)) es) -
            sum (map (\o -> product (map2 (\i j -> m_i_j) (between 1 n) o)) os)
    | _ -> undefined

def M.det (m: Matrix MathValue) : MathValue := M.determinant m