"Simple" matrix math for Agda

Mat.agda

module Mat where

open import Function
open import Data.Empty
open import Data.Nat using (ℕ; zero; suc)
open import Data.Fin using (Fin; zero; suc; toℕ)
open import Data.Fin.Props using ()
open import Data.Product
open import Relation.Nullary
open import Relation.Binary
open import Algebra

Vec : ∀ {a} (A : Set a) → ℕ → Set a
Vec A x = (i : Fin x) → A

pureVec : ∀ {x} {a} {A : Set a} → A → Vec A x
pureVec x = const x

foldVec : ∀ {n} {a b} {A : Set a} {B : Set b} → (A → B → B) → B → Vec A n → B
foldVec {zero}  f z xs = z
foldVec {suc n} f z xs = f (xs zero) (foldVec f z (xs ∘ suc))


Mat : ∀ {a} (A : Set a) → ℕ → ℕ → Set a
Mat A x y = (i : Fin x) (j : Fin y) → A

diagonal : ∀ {x} {a} {A : Set a} → A → Vec A x → Mat A x x
diagonal 0# x zero zero = x zero
diagonal 0# x zero (suc j) = 0#
diagonal 0# x (suc i) zero = 0#
diagonal 0# x (suc i) (suc j) = diagonal 0# (x ∘ suc) i j -- Not the most efficient algorithm, but it works

transpose : ∀ {x y} {a} {A : Set a} → Mat A x y → Mat A y x
transpose = flip

premultiply : ∀ {x y z} {a b c} {A : Set a} {B : Set b} {C : Set c} → (Vec A y → Vec B y → C) → Mat A x y → Mat B y z → Mat C x z
premultiply f xs ys = λ i j → f (xs i) (transpose ys j)


module Matrix-Semiring {c ℓ} (S : Semiring c ℓ) where
  module Semi = Semiring S
  open Semi renaming (Carrier to C) hiding (zero)
  open import Relation.Binary.EqReasoning setoid

  infixl 7 _*′_
  infixl 6 _+′_
  infix  4 _≈′_

  _≈′_ : ∀ {x y} → Mat C x y → Mat C x y → Set ℓ
  f ≈′ g = ∀ i j → f i j ≈ g i j

  0#′ : ∀ {x} → Mat C x x
  0#′ i j = 0#

  1#′ : ∀ {x} → Mat C x x
  1#′ = diagonal 0# (pureVec 1#)

  _+′_ : ∀ {x y} → Mat C x y → Mat C x y → Mat C x y
  f +′ g = λ i j → f i j + g i j

  sum : ∀ {x} → Vec C x → C
  sum = foldVec _+_ 0#

  sum-cong : ∀ {x} {xs ys : Vec C x} → (∀ i → xs i ≈ ys i) → sum xs ≈ sum ys
  sum-cong {zero}  eq = refl
  sum-cong {suc x} eq = +-cong (eq zero) (sum-cong (eq ∘ suc))

  sum-0# : ∀ {x} {xs : Vec C x} → (∀ i → xs i ≈ 0#) → sum xs ≈ 0#
  sum-0# {zero}  eq = refl
  sum-0# {suc x} eq = trans (+-cong (eq zero) (sum-0# (eq ∘ suc))) (proj₁ +-identity 0#)

  -- Argh, give me a commutative monoid solver
  sum-+ : ∀ {x} (xs ys : Vec C x) → sum (λ i → xs i + ys i) ≈ sum xs + sum ys
  sum-+ {zero}  xs ys = sym (proj₁ +-identity 0#)
  sum-+ {suc x} xs ys =
    begin
      (xs zero + ys zero) + sum ((λ i → xs i + ys i) ∘ suc)
    ≈⟨ +-cong refl (sum-+ (xs ∘ suc) (ys ∘ suc)) ⟩
      xs zero + ys zero + (sum (xs ∘ suc) + sum (ys ∘ suc))
    ≈⟨ +-assoc _ _ _ ⟩
      xs zero + (ys zero + (sum (xs ∘ suc) + sum (ys ∘ suc)))
    ≈⟨ +-cong refl (trans (trans (sym (+-assoc _ _ _)) (+-cong (+-comm _ _) refl)) (+-assoc _ _ _)) ⟩
      xs zero + (sum (xs ∘ suc) + (ys zero + sum (ys ∘ suc)))
    ≈⟨ sym (+-assoc _ _ _) ⟩
      (xs zero + sum (xs ∘ suc)) + (ys zero + sum (ys ∘ suc))
    ∎

  sum-distribˡ : ∀ {n} (x : C) (ys : Vec C n) → x * sum ys ≈ sum (λ i → x * ys i)
  sum-distribˡ {zero}  x ys = proj₂ Semi.zero _
  sum-distribˡ {suc n} x ys = trans (proj₁ distrib _ _ _) (+-cong refl (sum-distribˡ x (ys ∘ suc)))

  sum-distribʳ : ∀ {n} (xs : Vec C n) (y : C) → sum xs * y ≈ sum (λ i → xs i * y)
  sum-distribʳ {zero}  xs y = proj₁ Semi.zero _
  sum-distribʳ {suc n} xs y = trans (proj₂ distrib _ _ _) (+-cong refl (sum-distribʳ (xs ∘ suc) y))

  sum-comm : ∀ {x y} (xs : Fin x → Fin y → C) → sum (λ i → sum (λ j → xs i j)) ≈ sum (λ j → sum (λ i → xs i j))
  sum-comm xs = {!!} -- ugh

  dot : ∀ {x} → Vec C x → Vec C x → C
  dot f g = sum (λ i → f i * g i)

  trace : ∀ {x} → Mat C x x → C
  trace m = sum (λ i → m i i)

  dot-1#ˡ : ∀ {x} (i : Fin x) (xs : Vec C x) → dot (1#′ i) xs ≈ xs i
  dot-1#ˡ zero    xs = trans (+-cong (proj₁ *-identity _) (sum-0# (λ i → proj₁ Semi.zero (xs (suc i))))) (proj₂ +-identity _)
  dot-1#ˡ (suc i) xs = trans (+-cong (proj₁ Semi.zero _) (dot-1#ˡ i (xs ∘ suc))) (proj₁ +-identity _)

  dot-1#ʳ : ∀ {x} (i : Fin x) (xs : Vec C x) → dot xs (transpose 1#′ i) ≈ xs i
  dot-1#ʳ zero    xs = trans (+-cong (proj₂ *-identity _) (sum-0# (λ i → proj₂ Semi.zero (xs (suc i))))) (proj₂ +-identity _)
  dot-1#ʳ (suc i) xs = trans (+-cong (proj₂ Semi.zero _) (dot-1#ʳ i (xs ∘ suc))) (proj₁ +-identity _)

  _*′_ : ∀ {x y z} → Mat C x y → Mat C y z → Mat C x z
  _*′_ = premultiply dot

  isEquivalence′ : ∀ {x y} → IsEquivalence (_≈′_ {x} {y})
  isEquivalence′ = record
    { refl = λ i j → refl
    ; sym = λ f i j → sym (f i j)
    ; trans = λ f g i j → trans (f i j) (g i j)
    }

  assoc′ : ∀ {n} (x y z : Mat C n n) (i j : Fin n) → sum (λ k → sum (λ l → x i l * y l k) * z k j) ≈ sum (λ k → x i k * sum (λ l → y k l * z l j))
  assoc′ x y z i j =
    begin
      sum (λ k → sum (λ l → x i l * y l k) * z k j)
    ≈⟨ sum-cong (λ k → sum-distribʳ (λ l → x i l * y l k) (z k j)) ⟩
      sum (λ k → sum (λ l → x i l * y l k  * z k j))
    ≈⟨ sum-cong (λ k → sum-cong (λ m → *-assoc (x i m) (y m k) (z k j))) ⟩
      sum (λ k → sum (λ l → x i l * (y l k * z k j)))
    ≈⟨ sym (sum-comm (λ k l → x i k * (y k l * z l j))) ⟩
      sum (λ k → sum (λ l → x i k * (y k l * z l j)))
    ≈⟨ sym (sum-cong (λ k → sum-distribˡ (x i k) (λ l → y k l * z l j))) ⟩
      sum (λ k → x i k * sum (λ l → y k l * z l j))
    ∎


  semiring : (n : ℕ) → Semiring c ℓ
  semiring n = record
    { Carrier = Mat C n n
    ; _≈_ = _≈′_
    ; _+_ = _+′_
    ; _*_ = _*′_
    ; 0# = 0#′
    ; 1# = 1#′
    ; isSemiring = record
      { isSemiringWithoutAnnihilatingZero = record 
        { +-isCommutativeMonoid = record
          { isSemigroup = record
            { isEquivalence = isEquivalence′
            ; assoc = λ f g h i j → +-assoc (f i j) (g i j) (h i j)
            ; ∙-cong = λ f g i j → +-cong (f i j) (g i j)
            }
          ; identityˡ = λ f i j → proj₁ +-identity (f i j)
          ; comm = λ f g i j → +-comm (f i j) (g i j)
          }
        ; *-isMonoid = record
          { isSemigroup = record
            { isEquivalence = isEquivalence′
            ; assoc = assoc′
            ; ∙-cong = λ f g i j → sum-cong (λ k → *-cong (f i k) (g k j))
            }
          ; identity = (λ f i j → dot-1#ˡ i (transpose f j)) , (λ f i j → dot-1#ʳ j (f i))
          }
        ; distrib = (λ f g h i j → trans (sum-cong (λ k → proj₁ distrib (f i k) (g k j) (h k j))) (sum-+ (λ k → f i k * g k j) (λ k → f i k * h k j)))
                  , (λ f g h i j → trans (sum-cong (λ k → proj₂ distrib (f k j) (g i k) (h i k))) (sum-+ (λ k → g i k * f k j) (λ k → h i k * f k j)))
        }
      ; zero = (λ f i j → sum-0# (λ k → proj₁ Semi.zero (f k j)))
             , (λ f i j → sum-0# (λ k → proj₂ Semi.zero (f i k)))
      }
    }

module Test where
  open import Data.Vec renaming (Vec to Vector)
  open import Data.Nat.Properties using (commutativeSemiring)
  open Matrix-Semiring (CommutativeSemiring.semiring commutativeSemiring) public

  Matrix : ∀ {a} (A : Set a) → ℕ → ℕ → Set a
  Matrix A m n = Vector (Vector A n) m

  fromMat : ∀ {m n} {a} {A : Set a} → Mat A m n → Matrix A m n
  fromMat m = tabulate (λ i → tabulate (λ j → m i j))

  toMat : ∀ {m n} {a} {A : Set a} → Matrix A m n → Mat A m n
  toMat xs = λ i j → lookup j (lookup i xs)

  test : Matrix ℕ 3 3
  test = (0 ∷ 1 ∷ 1 ∷ []) 
       ∷ (3 ∷ 0 ∷ 4 ∷ [])
       ∷ (1 ∷ 1 ∷ 0 ∷ []) 
       ∷ []

  test2 = fromMat (toMat test *′ toMat test *′ toMat test)

  test3 : Matrix ℕ 3 3
  test3 = (7  ∷ 8 ∷ 8  ∷ [])
        ∷ (24 ∷ 7 ∷ 32 ∷ [])
        ∷ (8  ∷ 8 ∷ 7  ∷ [])
        ∷ []

  test5 = fromMat (1#′ {3})

  test6 : Matrix ℕ 5 5
  test6 = fromMat (diagonal 0 toℕ)

open Test

module Matrix-Ring {c ℓ} (S : Ring c ℓ) where
  module Full = Ring S

  -′_ : ∀ {x y} → Mat Full.Carrier x y → Mat Full.Carrier x y
  -′ f = λ i j → Full.- f i j

  ring : (n : ℕ) → Ring c ℓ
  ring n = record
    { Carrier = Carrier
    ; _≈_ = _≈_
    ; _+_ = _+_
    ; _*_ = _*_
    ; -_ = -′_
    ; 0# = 0#
    ; 1# = 1#
    ; isRing = record
      { +-isAbelianGroup = record
        { isGroup = record
          { isMonoid = +-isMonoid
          ; inverse = map (λ q f i j → q (f i j)) (λ q f i j → q (f i j)) Full.-‿inverse
          ; ⁻¹-cong = λ f i j → Full.-‿cong (f i j)
          }
        ; comm = +-comm
        }
      ; *-isMonoid = *-isMonoid
      ; distrib = distrib
      }
    }
    where open Semiring (Matrix-Semiring.semiring Full.semiring n)

Now we just need to feed in a matrix of types using the semiring of types I made a while back!