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

FFT-Matrices-v2.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
                       | n `gcd` k > 1 = let g = n `gcd` k
                                         in W P (n `div` g) (k `div` g)
                       | 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 { unMat :: [[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]

infixl 6 %*%

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)

zerom :: Int -> Mat
zerom bigN = diag $ replicate bigN Z


-- FOURIER BUILDING BLOCKS

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

-- Multiplied-up resized Fourier matrices.
fourierMultiple :: Int -> Int -> Mat
fourierMultiple = fourierMultiple' 1
fourierMultiple' :: Int -> Int -> Int -> Mat
fourierMultiple' sign np p =
  let n = np `div` p
      zs = replicate n Z
      zs0 = concat $ replicate (p-1) zs
      shift = map (take np . (zs ++))
      f = unMat $ fourier' sign n
  in Mat $ concat $ take p $ iterate shift $ map (++ zs0) f

-- Modulo permutation matrix.
perm :: Int -> Int -> Mat
perm np p =
  let n = np `div` p
      zz = (O P) : replicate (np-1) Z
      zs = take n $ iterate (take np . ((replicate p Z) ++)) zz
      shift = map (take np . (Z :))
  in Mat $ concat $ take p $ iterate shift zs

-- Generalised I + D matrix.
idMatrix :: Int -> Int -> Mat
idMatrix = idMatrix' 1
idMatrix' :: Int -> Int -> Int -> Mat
idMatrix' sign wfac split =
  let ns = wfac `div` split
      w = W P wfac sign
      i = unMat $ unit ns
      ds j = (O P) : [w^(i*j) | i <- [1..ns-1]]
      d r c = unMat $ diag $ map (w^(ns*r*c) *) (ds c)
      row r = L.foldl' (zipWith (++)) i $ map (d r) [1..split-1]
  in Mat $ concat $ map row [0..split-1]


tst :: Int -> Int -> Mat
tst np p = idMatrix np p %*% fourierMultiple np p %*% perm np p


-- 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}"]