Toy algebra system code for Fourier matrix decomposition (http://www.skybluetrades.net/blog/posts/2013/11/15/data-analysis-fft-2.html)

FFT-Matrices-v1.hs

{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
import qualified Data.List as L

-- MATRIX ENTRIES

data Sign = P | M deriving (Eq, Show)

data Entry = Z | O Sign | W Sign Int Int deriving Eq

negS :: Sign -> Sign
negS P = M
negS M = P

multS :: Sign -> Sign -> Sign
multS s1 s2 | s1 == s2  = P
            | otherwise = M

instance Num Entry where
  abs _ = error "Absolute value not appropriate"
  signum _ = error "Sign not appropriate"
  negate (O s) = O (negS s)
  negate (W s n p) = W (negS s) n p
  Z + x = x
  x + Z = x
  _ + _ = error "Adding two non-zero terms!"
  x - y = x + negate y
  Z * x = Z
  x * Z = Z
  (O s1) * (O s2) = O (multS s1 s2)
  (O s1) * (W s2 n p) = W (multS s1 s2) n p
  (W s1 n p) * (O s2) = W (multS s1 s2) n p
  (W s1 n1 p1) * (W s2 n2 p2)
    | n1 == n2 = W (multS s1 s2) n1 ((p1+p2) `mod` n1)
    | gcd n1 n2 == 1 = error "Multiplying different omegas"
    | otherwise = if n1 > n2
                  then W (multS s1 s2) n1 ((p1+p2*(n1`div`n2))`mod`n1)
                  else W (multS s1 s2) n2 ((p2+p1*(n2`div`n1))`mod`n2)
  fromInteger n
    | n == 0 = Z
    | n == 1 = O P
    | n == -1 = O M
    | otherwise = error "Invalid integer"

normalise :: Entry -> Entry
normalise = fix norm
  where norm (W s _ 0) = O s
        norm (W M n k) | even n = W P n ((k + n `div` 2) `mod` n)
                       | otherwise = W M n (k `mod` n)
        norm (W s 2 k) | odd k = O s * O M
                       | otherwise = O s
        norm (W P n k) | even n && even k = W P (n `div` 2) (k `div` 2)
                       | n == 2 * k = O M
                       | otherwise = W P n (k `mod` n)
        norm x = x
        fix f x = let xs = iterate f x
                  in fst $ head $ dropWhile (\(x,y) -> x/=y) $ zip xs (tail xs)


-- MATRICES

newtype Mat = Mat [[Entry]] deriving Eq

normMat :: Mat -> Mat
normMat (Mat rs) = Mat $ map (map normalise) rs

transpose :: Mat -> Mat
transpose (Mat xs) = Mat (L.transpose xs)

dot :: [Entry] -> [Entry] -> Entry
dot v1 v2 = L.sum $ L.zipWith (*) v1 v2

(%*%) :: Mat -> Mat -> Mat
(Mat m1) %*% (Mat m2) =
  normMat $ Mat [[r `dot` c | c <- L.transpose m2] | r <- m1]

diag :: [Entry] -> Mat
diag ds = let bigN = length ds
              zs = replicate (bigN - 1) Z
              zsplits = map (flip splitAt zs) [0..bigN-1]
              rows = zipWith (\d (b, a) -> b ++ [d] ++ a) ds zsplits
          in Mat rows

unit :: Int -> Mat
unit bigN = diag $ replicate bigN (O P)


-- FOURIER BUILDING BLOCKS

-- Basic Fourier matrix.
fourier :: Int -> Mat
fourier bigN =
  normMat $ Mat [[(W P bigN 1)^(n*k) | k <- [0..bigN-1]] | n <- [0..bigN-1]]

-- Doubled-up half-size Fourier matrices.
fourierDouble :: Int -> Mat
fourierDouble n2 =
  let n = n2 `div` 2
      zs = replicate n Z
      Mat f = fourier n
  in Mat $ (map (++ zs) f) ++ (map (zs ++) f)

-- Even/odd permutation matrix.
perm :: Int -> Mat
perm n2 = let n = n2 `div` 2
              evens = ((O P) : replicate (n2-1) Z) :
                      map (take n2 . (Z :) . (Z :)) evens
              odds = map (take n2 . (Z :)) evens
          in Mat $ take n evens ++ take n odds

-- I + D matrix.
idMatrix :: Int -> Mat
idMatrix n2 = let n = n2 `div` 2
                  Mat i = unit n
                  ds = (O P) : [W P n2 i | i <- [1..n-1]]
                  Mat d = diag ds
                  Mat nd = diag $ map negate ds
              in Mat $ (zipWith (++) i d) ++ (zipWith (++) i nd)



-- OUTPUT HELPERS

instance Show Entry where
  show Z     = "  0  "
  show (O s) = if s == P then "  1  " else " -1  "
  show (W s n k) = (if s == P then " " else "-") ++
                   "W" ++ show n ++ "^" ++ show k

instance Show Mat where
  show (Mat rs) = unlines $ map (L.intercalate " " . map show) rs


class ToLaTeX a where
  toLaTeX :: a -> String

instance ToLaTeX Entry where
  toLaTeX Z = "0"
  toLaTeX (O P) = "1"
  toLaTeX (O M) = "-1"
  toLaTeX (W s n k) = let sl = if s == M then "-" else ""
                          pow = if k == 1 then "" else "^{" ++ show k ++ "}"
                      in sl ++ "\\omega_{" ++ show n ++ "}" ++ pow

instance ToLaTeX Mat where
  toLaTeX (Mat rs) = unlines $
                     ["\\begin{pmatrix}"] ++
                     map ((++ "\\\\") . L.intercalate " & " . map toLaTeX) rs
                     ++ ["\\end{pmatrix}"]


-- GENERAL UTILITIES

-- Doubled-up any matrix.
doubleUp :: Mat -> Mat
doubleUp (Mat m) =
  let n = length m
      zs = replicate n Z
  in Mat $ (map (++ zs) m) ++ (map (zs ++) m)