text:
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE EmptyDataDecls             #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FunctionalDependencies     #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE Rank2Types                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE ViewPatterns               #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

{- |
Module      :  Numeric.LinearAlgebra.Static
Copyright   :  (c) Alberto Ruiz 2014
License     :  BSD3
Stability   :  experimental

Experimental interface with statically checked dimensions.

See code examples at http://dis.um.es/~alberto/hmatrix/static.html.

-}

module Numeric.LinearAlgebra.Static(
    -- * Vector
       ℝ, R,
    vec2, vec3, vec4, (&), (#), split, headTail,
    vector,
    linspace, range, dim,
    -- * Matrix
    L, Sq, build,
    row, col, (|||),(===), splitRows, splitCols,
    unrow, uncol,
    tr,
    eye,
    diag,
    blockAt,
    matrix,
    -- * Complex
    ℂ, C, M, Her, her, 𝑖,
    toComplex,
    fromComplex,
    complex,
    real,
    imag,
    sqMagnitude,
    magnitude,
    -- * Products
    (<>),(#>),(<.>),
    -- * Linear Systems
    linSolve, (<\>),
    -- * Factorizations
    svd, withCompactSVD, svdTall, svdFlat, Eigen(..),
    withNullspace, withOrth, qr, chol,
    -- * Norms
    Normed(..),
    -- * Random arrays
    Seed, RandDist(..),
    randomVector, rand, randn, gaussianSample, uniformSample,
    -- * Misc
    mean, meanCov,
    Disp(..), Domain(..),
    withVector, withMatrix, exactLength, exactDims,
    toRows, toColumns, withRows, withColumns,
    Sized(..), Diag(..), Sym, sym, mTm, unSym, (<·>)
) where

import           Control.Arrow               ((***))
import qualified Data.Bifunctor              as BF (first)
import           Data.Proxy                  (Proxy (..))
import           Data.Type.Equality          ((:~:) (Refl))
import           GHC.TypeLits
import           Internal.Static
import           Numeric.LinearAlgebra       hiding (C, Konst (..), R, build, chol, col,
                                              complex, diag, disp, dot, eig, eigSH,
                                              eigenvalues, eigenvaluesSH, fromComplex,
                                              fromList, gaussianSample, linspace, mTm,
                                              magnitude, matrix, meanCov, qr, rand, randn,
                                              randomVector, range, real, row, size, svd,
                                              sym, takeDiag, toColumns, toComplex, toRows,
                                              unSym, uniformSample, vector, ( #> ), (<.>),
                                              (<>), (<\>), (===), (|||))
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as LA
import           Text.Printf
#if MIN_VERSION_base(4,11,0)
import           Prelude                     hiding ((<>))
#endif

ud1 :: R n -> Vector ℝ
ud1 (R (Dim v)) = v

infixl 4 &
(&) :: forall n . KnownNat n
    => R n -> ℝ -> R (n+1)
u & x = u # (konst x :: R 1)

infixl 4 #
(#) :: forall n m . (KnownNat n, KnownNat m)
    => R n -> R m -> R (n+m)
(R u) # (R v) = R (vconcat u v)

vec2 :: ℝ -> ℝ -> R 2
vec2 a b = R (gvec2 a b)

vec3 :: ℝ -> ℝ -> ℝ -> R 3
vec3 a b c = R (gvec3 a b c)

vec4 :: ℝ -> ℝ -> ℝ -> ℝ -> R 4
vec4 a b c d = R (gvec4 a b c d)

vector :: KnownNat n => [ℝ] -> R n
vector = fromList

matrix :: (KnownNat m, KnownNat n) => [ℝ] -> L m n
matrix = fromList

linspace :: forall n . KnownNat n => (ℝ,ℝ) -> R n
linspace (a,b) = v
  where
    v = mkR (LA.linspace (size v) (a,b))

range :: forall n . KnownNat n => R n
range = v
  where
    v = mkR (LA.linspace d (1,fromIntegral d))
    d = size v

dim :: forall n . KnownNat n => R n
dim = v
  where
    v = mkR (scalar (fromIntegral $ size v))

--------------------------------------------------------------------------------

ud2 :: L m n -> Matrix ℝ
ud2 (L (Dim (Dim x))) = x

--------------------------------------------------------------------------------

diag :: KnownNat n => R n -> Sq n
diag = diagR 0

eye :: KnownNat n => Sq n
eye = diag 1

--------------------------------------------------------------------------------

blockAt :: forall m n . (KnownNat m, KnownNat n) =>  ℝ -> Int -> Int -> Matrix Float -> L m n
blockAt x r c a = res
  where
    z = scalar x
    z1 = LA.konst x (r,c)
    z2 = LA.konst x (max 0 (m'-(ra+r)), max 0 (n'-(ca+c)))
    ra = min (rows a) . max 0 $ m'-r
    ca = min (cols a) . max 0 $ n'-c
    sa = subMatrix (0,0) (ra, ca) a
    (m',n') = size res
    res = mkL $ fromBlocks [[z1,z,z],[z,sa,z],[z,z,z2]]

--------------------------------------------------------------------------------

row :: R n -> L 1 n
row = mkL . asRow . ud1

--col :: R n -> L n 1
col v = tr . row $ v

unrow :: L 1 n -> R n
unrow = mkR . head . LA.toRows . ud2

--uncol :: L n 1 -> R n
uncol v = unrow . tr $ v

infixl 2 ===
(===) :: (KnownNat r1, KnownNat r2, KnownNat c) => L r1 c -> L r2 c -> L (r1+r2) c
a === b = mkL (extract a LA.=== extract b)

infixl 3 |||
-- (|||) :: (KnownNat r, KnownNat c1, KnownNat c2) => L r c1 -> L r c2 -> L r (c1+c2)
a ||| b = tr (tr a === tr b)

type Sq n  = L n n
--type CSq n = CL n n

type GL = forall n m . (KnownNat n, KnownNat m) => L m n
type GSq = forall n . KnownNat n => Sq n

isKonst :: forall m n . (KnownNat m, KnownNat n) => L m n -> Maybe (ℝ,(Int,Int))
isKonst s@(unwrap -> x)
    | singleM x = Just (x `atIndex` (0,0), (size s))
    | otherwise = Nothing

isKonstC :: forall m n . (KnownNat m, KnownNat n) => M m n -> Maybe (ℂ,(Int,Int))
isKonstC s@(unwrap -> x)
    | singleM x = Just (x `atIndex` (0,0), (size s))
    | otherwise = Nothing

infixr 8 <>
(<>) :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => L m k -> L k n -> L m n
(<>) = mulR

infixr 8 #>
(#>) :: (KnownNat m, KnownNat n) => L m n -> R n -> R m
(#>) = appR

infixr 8 <·>
(<·>) :: KnownNat n => R n -> R n -> ℝ
(<·>) = dotR

infixr 8 <.>
(<.>) :: KnownNat n => R n -> R n -> ℝ
(<.>) = dotR

--------------------------------------------------------------------------------

class Diag m d | m -> d
  where
    takeDiag :: m -> d

instance KnownNat n => Diag (L n n) (R n)
  where
    takeDiag x = mkR (LA.takeDiag (extract x))

instance KnownNat n => Diag (M n n) (C n)
  where
    takeDiag x = mkC (LA.takeDiag (extract x))

--------------------------------------------------------------------------------

toComplex :: KnownNat n => (R n, R n) -> C n
toComplex (r,i) = mkC $ LA.toComplex (ud1 r, ud1 i)

fromComplex :: KnownNat n => C n -> (R n, R n)
fromComplex (C (Dim v)) = let (r,i) = LA.fromComplex v in (mkR r, mkR i)

complex :: KnownNat n => R n -> C n
complex r = mkC $ LA.toComplex (ud1 r, LA.konst 0 (size r))

real :: KnownNat n => C n -> R n
real = fst . fromComplex

imag :: KnownNat n => C n -> R n
imag = snd . fromComplex

sqMagnitude :: KnownNat n => C n -> R n
sqMagnitude c = let (r,i) = fromComplex c in r**2 + i**2

magnitude :: KnownNat n => C n -> R n
magnitude = sqrt . sqMagnitude

--------------------------------------------------------------------------------

linSolve :: (KnownNat m, KnownNat n) => L m m -> L m n -> Maybe (L m n)
linSolve (extract -> a) (extract -> b) = fmap mkL (LA.linearSolve a b)

(<\>) :: (KnownNat m, KnownNat n, KnownNat r) => L m n -> L m r -> L n r
(extract -> a) <\> (extract -> b) = mkL (a LA.<\> b)

svd :: (KnownNat m, KnownNat n) => L m n -> (L m m, R n, L n n)
svd (extract -> m) = (mkL u, mkR s', mkL v)
  where
    (u,s,v) = LA.svd m
    s' = vjoin [s, z]
    z = LA.konst 0 (max 0 (cols m - LA.size s))

svdTall :: (KnownNat m, KnownNat n, n <= m) => L m n -> (L m n, R n, L n n)
svdTall (extract -> m) = (mkL u, mkR s, mkL v)
  where
    (u,s,v) = LA.thinSVD m

svdFlat :: (KnownNat m, KnownNat n, m <= n) => L m n -> (L m m, R m, L n m)
svdFlat (extract -> m) = (mkL u, mkR s, mkL v)
  where
    (u,s,v) = LA.thinSVD m

--------------------------------------------------------------------------------

class Eigen m l v | m -> l, m -> v
  where
    eigensystem :: m -> (l,v)
    eigenvalues :: m -> l

newtype Sym n = Sym (Sq n) deriving Show

sym :: KnownNat n => Sq n -> Sym n
sym m = Sym $ (m + tr m)/2

mTm :: (KnownNat m, KnownNat n) => L m n -> Sym n
mTm x = Sym (tr x <> x)

unSym :: Sym n -> Sq n
unSym (Sym x) = x

𝑖 :: Sized ℂ s c => s
𝑖 = konst iC

newtype Her n = Her (M n n)

her :: KnownNat n => M n n -> Her n
her m = Her $ (m + LA.tr m)/2

instance (KnownNat n) => Disp (Sym n)
  where
    disp n (Sym x) = do
        let a = extract x
        let su = LA.dispf n a
        printf "Sym %d" (cols a) >> putStr (dropWhile (/='\n') $ su)

instance (KnownNat n) => Disp (Her n)
  where
    disp n (Her x) = do
        let a = extract x
        let su = LA.dispcf n a
        printf "Her %d" (cols a) >> putStr (dropWhile (/='\n') $ su)

instance KnownNat n => Eigen (Sym n) (R n) (L n n)
  where
    eigenvalues (Sym (extract -> m)) =  mkR . LA.eigenvaluesSH . LA.trustSym $ m
    eigensystem (Sym (extract -> m)) = (mkR l, mkL v)
      where
        (l,v) = LA.eigSH . LA.trustSym $ m

instance KnownNat n => Eigen (Sq n) (C n) (M n n)
  where
    eigenvalues (extract -> m) = mkC . LA.eigenvalues $ m
    eigensystem (extract -> m) = (mkC l, mkM v)
      where
        (l,v) = LA.eig m

chol :: KnownNat n => Sym n -> Sq n
chol (extract . unSym -> m) = mkL $ LA.chol $ LA.trustSym m

--------------------------------------------------------------------------------

withNullspace
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k . (KnownNat k) => L n k -> z)
    -> z
withNullspace (LA.nullspace . extract -> a) f =
    case someNatVal $ fromIntegral $ cols a of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL a :: L n k)

withOrth
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k. (KnownNat k) => L n k -> z)
    -> z
withOrth (LA.orth . extract -> a) f =
    case someNatVal $ fromIntegral $ cols a of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL a :: L n k)

withCompactSVD
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k . (KnownNat k) => (L m k, R k, L n k) -> z)
    -> z
withCompactSVD (LA.compactSVD . extract -> (u,s,v)) f =
    case someNatVal $ fromIntegral $ LA.size s of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL u :: L m k, mkR s :: R k, mkL v :: L n k)

--------------------------------------------------------------------------------

qr :: (KnownNat m, KnownNat n) => L m n -> (L m m, L m n)
qr (extract -> x) = (mkL q, mkL r)
  where
    (q,r) = LA.qr x

-- use qrRaw?

--------------------------------------------------------------------------------

split :: forall p n . (KnownNat p, KnownNat n, p<=n) => R n -> (R p, R (n-p))
split (extract -> v) = ( mkR (subVector 0 p' v) ,
                         mkR (subVector p' (LA.size v - p') v) )
  where
    p' = fromIntegral . natVal $ (undefined :: Proxy p) :: Int

headTail :: (KnownNat n, 1<=n) => R n -> (ℝ, R (n-1))
headTail = ((!0) . extract *** id) . split

splitRows :: forall p m n . (KnownNat p, KnownNat m, KnownNat n, p<=m) => L m n -> (L p n, L (m-p) n)
splitRows (extract -> x) = ( mkL (takeRows p' x) ,
                             mkL (dropRows p' x) )
  where
    p' = fromIntegral . natVal $ (undefined :: Proxy p) :: Int

splitCols :: forall p m n. (KnownNat p, KnownNat m, KnownNat n, KnownNat (n-p), p<=n) => L m n -> (L m p, L m (n-p))
splitCols = (tr *** tr) . splitRows . tr

toRows :: forall m n . (KnownNat m, KnownNat n) => L m n -> [R n]
toRows (LA.toRows . extract -> vs) = map mkR vs

withRows
    :: forall n z . KnownNat n
    => [R n]
    -> (forall m . KnownNat m => L m n -> z)
    -> z
withRows (LA.fromRows . map extract -> m) f =
    case someNatVal $ fromIntegral $ LA.rows m of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) -> f (mkL m :: L m n)

toColumns :: forall m n . (KnownNat m, KnownNat n) => L m n -> [R m]
toColumns (LA.toColumns . extract -> vs) = map mkR vs

withColumns
    :: forall m z . KnownNat m
    => [R m]
    -> (forall n . KnownNat n => L m n -> z)
    -> z
withColumns (LA.fromColumns . map extract -> m) f =
    case someNatVal $ fromIntegral $ LA.cols m of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy n)) -> f (mkL m :: L m n)

--------------------------------------------------------------------------------

build
  :: forall m n . (KnownNat n, KnownNat m)
    => (ℝ -> ℝ -> ℝ)
    -> L m n
build f = r
  where
    r = mkL $ LA.build (size r) f

--------------------------------------------------------------------------------

withVector
    :: forall z
     . Vector ℝ
    -> (forall n . (KnownNat n) => R n -> z)
    -> z
withVector v f =
    case someNatVal $ fromIntegral $ LA.size v of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) -> f (mkR v :: R m)

-- | Useful for constraining two dependently typed vectors to match each
-- other in length when they are unknown at compile-time.
exactLength
    :: forall n m . (KnownNat n, KnownNat m)
    => R m
    -> Maybe (R n)
exactLength v = do
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy m)
    return $ mkR (unwrap v)

withMatrix
    :: forall z
     . Matrix ℝ
    -> (forall m n . (KnownNat m, KnownNat n) => L m n -> z)
    -> z
withMatrix a f =
    case someNatVal $ fromIntegral $ rows a of
       Nothing -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) ->
           case someNatVal $ fromIntegral $ cols a of
               Nothing -> error "static/dynamic mismatch"
               Just (SomeNat (_ :: Proxy n)) ->
                  f (mkL a :: L m n)

-- | Useful for constraining two dependently typed matrices to match each
-- other in dimensions when they are unknown at compile-time.
exactDims
    :: forall n m j k . (KnownNat n, KnownNat m, KnownNat j, KnownNat k)
    => L m n
    -> Maybe (L j k)
exactDims m = do
    Refl <- sameNat (Proxy :: Proxy m) (Proxy :: Proxy j)
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy k)
    return $ mkL (unwrap m)

randomVector
    :: forall n . KnownNat n
    => Seed
    -> RandDist
    -> R n
randomVector s d = mkR (LA.randomVector s d
                          (fromInteger (natVal (Proxy :: Proxy n)))
                       )

rand
    :: forall m n . (KnownNat m, KnownNat n)
    => IO (L m n)
rand = mkL <$> LA.rand (fromInteger (natVal (Proxy :: Proxy m)))
                       (fromInteger (natVal (Proxy :: Proxy n)))

randn
    :: forall m n . (KnownNat m, KnownNat n)
    => IO (L m n)
randn = mkL <$> LA.randn (fromInteger (natVal (Proxy :: Proxy m)))
                         (fromInteger (natVal (Proxy :: Proxy n)))

gaussianSample
    :: forall m n . (KnownNat m, KnownNat n)
    => Seed
    -> R n
    -> Sym n
    -> L m n
gaussianSample s (extract -> mu) (Sym (extract -> sigma)) =
    mkL $ LA.gaussianSample s (fromInteger (natVal (Proxy :: Proxy m)))
                            mu (LA.trustSym sigma)

uniformSample
    :: forall m n . (KnownNat m, KnownNat n)
    => Seed
    -> R n    -- ^ minimums of each row
    -> R n    -- ^ maximums of each row
    -> L m n
uniformSample s (extract -> mins) (extract -> maxs) =
    mkL $ LA.uniformSample s (fromInteger (natVal (Proxy :: Proxy m)))
                           (zip (LA.toList mins) (LA.toList maxs))

meanCov
    :: forall m n . (KnownNat m, KnownNat n, 1 <= m)
    => L m n
    -> (R n, Sym n)
meanCov (extract -> vs) = mkR *** (Sym . mkL . LA.unSym) $ LA.meanCov vs

--------------------------------------------------------------------------------

class Domain field vec mat | mat -> vec field, vec -> mat field, field -> mat vec
  where
    mul :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => mat m k -> mat k n -> mat m n
    app :: forall m n . (KnownNat m, KnownNat n) => mat m n -> vec n -> vec m
    dot :: forall n . (KnownNat n) => vec n -> vec n -> field
    cross :: vec 3 -> vec 3 -> vec 3
    diagR ::  forall m n k . (KnownNat m, KnownNat n, KnownNat k) => field -> vec k -> mat m n
    dvmap :: forall n. KnownNat n => (field -> field) -> vec n -> vec n
    dmmap :: forall n m. (KnownNat m, KnownNat n) => (field -> field) -> mat n m -> mat n m
    outer :: forall n m. (KnownNat m, KnownNat n) => vec n -> vec m -> mat n m
    zipWithVector :: forall n. KnownNat n => (field -> field -> field) -> vec n -> vec n -> vec n
    det :: forall n. KnownNat n => mat n n -> field
    invlndet :: forall n. KnownNat n => mat n n -> (mat n n, (field, field))
    expm :: forall n. KnownNat n => mat n n -> mat n n
    sqrtm :: forall n. KnownNat n => mat n n -> mat n n
    inv :: forall n. KnownNat n => mat n n -> mat n n

instance Domain ℝ R L
  where
    mul = mulR
    app = appR
    dot = dotR
    cross = crossR
    diagR = diagRectR
    dvmap = mapR
    dmmap = mapL
    outer = outerR
    zipWithVector = zipWithR
    det = detL
    invlndet = invlndetL
    expm = expmL
    sqrtm = sqrtmL
    inv = invL

instance Domain ℂ C M
  where
    mul = mulC
    app = appC
    dot = dotC
    cross = crossC
    diagR = diagRectC
    dvmap = mapC
    dmmap = mapM'
    outer = outerC
    zipWithVector = zipWithC
    det = detM
    invlndet = invlndetM
    expm = expmM
    sqrtm = sqrtmM
    inv = invM

--------------------------------------------------------------------------------

mulR :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => L m k -> L k n -> L m n

mulR (isKonst -> Just (a,(_,k))) (isKonst -> Just (b,_)) = konst (a * b * fromIntegral k)

mulR (isDiag -> Just (0,a,_)) (isDiag -> Just (0,b,_)) = diagR 0 (mkR v :: R k)
  where
    v = a' * b'
    n = min (LA.size a) (LA.size b)
    a' = subVector 0 n a
    b' = subVector 0 n b

mulR (isDiag -> Just (0,a,_)) (extract -> b) = mkL (asColumn a * takeRows (LA.size a) b)

mulR (extract -> a) (isDiag -> Just (0,b,_)) = mkL (takeColumns (LA.size b) a * asRow b)

mulR a b = mkL (extract a LA.<> extract b)

appR :: (KnownNat m, KnownNat n) => L m n -> R n -> R m
appR (isDiag -> Just (0, w, _)) v = mkR (w * subVector 0 (LA.size w) (extract v))
appR m v                          = mkR (extract m LA.#> extract v)

dotR :: KnownNat n => R n -> R n -> ℝ
dotR (extract -> u) (extract -> v) = LA.dot u v

crossR :: R 3 -> R 3 -> R 3
crossR (extract -> x) (extract -> y) = vec3 z1 z2 z3
  where
    z1 = x!1*y!2-x!2*y!1
    z2 = x!2*y!0-x!0*y!2
    z3 = x!0*y!1-x!1*y!0

outerR :: (KnownNat m, KnownNat n) => R n -> R m -> L n m
outerR (extract -> x) (extract -> y) = mkL (LA.outer x y)

mapR :: KnownNat n => (ℝ -> ℝ) -> R n -> R n
mapR f (unwrap -> v) = mkR (LA.cmap f v)

zipWithR :: KnownNat n => (ℝ -> ℝ -> ℝ) -> R n -> R n -> R n
zipWithR f (extract -> x) (extract -> y) = mkR (LA.zipVectorWith f x y)

mapL :: (KnownNat n, KnownNat m) => (ℝ -> ℝ) -> L n m -> L n m
mapL f = overMatL' (LA.cmap f)

detL :: KnownNat n => Sq n -> ℝ
detL = LA.det . unwrap

invlndetL :: KnownNat n => Sq n -> (L n n, (ℝ, ℝ))
invlndetL = BF.first mkL . LA.invlndet . unwrap

expmL :: KnownNat n => Sq n -> Sq n
expmL = overMatL' LA.expm

sqrtmL :: KnownNat n => Sq n -> Sq n
sqrtmL = overMatL' LA.sqrtm

invL :: KnownNat n => Sq n -> Sq n
invL = overMatL' LA.inv

--------------------------------------------------------------------------------

mulC :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => M m k -> M k n -> M m n

mulC (isKonstC -> Just (a,(_,k))) (isKonstC -> Just (b,_)) = konst (a * b * fromIntegral k)

mulC (isDiagC -> Just (0,a,_)) (isDiagC -> Just (0,b,_)) = diagR 0 (mkC v :: C k)
  where
    v = a' * b'
    n = min (LA.size a) (LA.size b)
    a' = subVector 0 n a
    b' = subVector 0 n b

mulC (isDiagC -> Just (0,a,_)) (extract -> b) = mkM (asColumn a * takeRows (LA.size a) b)

mulC (extract -> a) (isDiagC -> Just (0,b,_)) = mkM (takeColumns (LA.size b) a * asRow b)

mulC a b = mkM (extract a LA.<> extract b)

appC :: (KnownNat m, KnownNat n) => M m n -> C n -> C m
appC (isDiagC -> Just (0, w, _)) v = mkC (w * subVector 0 (LA.size w) (extract v))
appC m v                           = mkC (extract m LA.#> extract v)

dotC :: KnownNat n => C n -> C n -> ℂ
dotC (extract -> u) (extract -> v) = LA.dot u v

crossC :: C 3 -> C 3 -> C 3
crossC (extract -> x) (extract -> y) = mkC (LA.fromList [z1, z2, z3])
  where
    z1 = x!1*y!2-x!2*y!1
    z2 = x!2*y!0-x!0*y!2
    z3 = x!0*y!1-x!1*y!0

outerC :: (KnownNat m, KnownNat n) => C n -> C m -> M n m
outerC (extract -> x) (extract -> y) = mkM (LA.outer x y)

mapC :: KnownNat n => (ℂ -> ℂ) -> C n -> C n
mapC f (unwrap -> v) = mkC (LA.cmap f v)

zipWithC :: KnownNat n => (ℂ -> ℂ -> ℂ) -> C n -> C n -> C n
zipWithC f (extract -> x) (extract -> y) = mkC (LA.zipVectorWith f x y)

mapM' :: (KnownNat n, KnownNat m) => (ℂ -> ℂ) -> M n m -> M n m
mapM' f = overMatM' (LA.cmap f)

detM :: KnownNat n => M n n -> ℂ
detM = LA.det . unwrap

invlndetM :: KnownNat n => M n n -> (M n n, (ℂ, ℂ))
invlndetM = BF.first mkM . LA.invlndet . unwrap

expmM :: KnownNat n => M n n -> M n n
expmM = overMatM' LA.expm

sqrtmM :: KnownNat n => M n n -> M n n
sqrtmM = overMatM' LA.sqrtm

invM :: KnownNat n => M n n -> M n n
invM = overMatM' LA.inv

--------------------------------------------------------------------------------

diagRectR :: forall m n k . (KnownNat m, KnownNat n, KnownNat k) => ℝ -> R k -> L m n
diagRectR x v
    | m' == 1 = mkL (LA.diagRect x ev m' n')
    | m'*n' > 0 = r
    | otherwise = matrix []
  where
    r = mkL (asRow (vjoin [scalar x, ev, zeros]))
    ev = extract v
    zeros = LA.konst x (max 0 ((min m' n') - LA.size ev))
    (m',n') = size r

diagRectC :: forall m n k . (KnownNat m, KnownNat n, KnownNat k) => ℂ -> C k -> M m n
diagRectC x v
    | m' == 1 = mkM (LA.diagRect x ev m' n')
    | m'*n' > 0 = r
    | otherwise = fromList []
  where
    r = mkM (asRow (vjoin [scalar x, ev, zeros]))
    ev = extract v
    zeros = LA.konst x (max 0 ((min m' n') - LA.size ev))
    (m',n') = size r

--------------------------------------------------------------------------------

mean :: (KnownNat n, 1<=n) => R n -> ℝ
mean v = v <·> (1/dim)

test :: (Bool, IO ())
test = (ok,info)
  where
    ok =   extract (eye :: Sq 5) == ident 5
           && (unwrap .unSym) (mTm sm :: Sym 3) == tr ((3><3)[1..]) LA.<> (3><3)[1..]
           && unwrap (tm :: L 3 5) == LA.matrix 5 [1..15]
           && thingS == thingD
           && precS == precD
           && withVector (LA.vector [1..15]) sumV == sumElements (LA.fromList [1..15])

    info = do
        print $ u
        print $ v
        print (eye :: Sq 3)
        print $ ((u & 5) + 1) <·> v
        print (tm :: L 2 5)
        print (tm <> sm :: L 2 3)
        print thingS
        print thingD
        print precS
        print precD
        print $ withVector (LA.vector [1..15]) sumV
        splittest

    sumV w = w <·> konst 1

    u = vec2 3 5

    𝕧 x = vector [x] :: R 1

    v = 𝕧 2 & 4 & 7

    tm :: GL
    tm = lmat 0 [1..]

    lmat :: forall m n . (KnownNat m, KnownNat n) => ℝ -> [ℝ] -> L m n
    lmat z xs = r
      where
        r = mkL . reshape n' . LA.fromList . take (m'*n') $ xs ++ repeat z
        (m',n') = size r

    sm :: GSq
    sm = lmat 0 [1..]

    thingS = (u & 1) <·> tr q #> q #> v
      where
        q = tm :: L 10 3

    thingD = vjoin [ud1 u, 1] LA.<.> tr m LA.#> m LA.#> ud1 v
      where
        m = LA.matrix 3 [1..30]

    precS = (1::Float) + (2::Float) * ((1 :: R 3) * (u & 6)) <·> konst 2 #> v
    precD = 1 + 2 * vjoin[ud1 u, 6] LA.<.> LA.konst 2 (LA.size (ud1 u) +1, LA.size (ud1 v)) LA.#> ud1 v

splittest
    = do
    let v = range :: R 7
        a = snd (split v) :: R 4
    print $ a
    print $ snd . headTail . snd . headTail $ v
    print $ first (vec3 1 2 3)
    print $ second (vec3 1 2 3)
    print $ third (vec3 1 2 3)
    print $ (snd $ splitRows eye :: L 4 6)
 where
    first v = fst . headTail $ v
    second v = first . snd . headTail $ v
    third v = first . snd . headTail . snd . headTail $ v

instance (KnownNat n', KnownNat m') => Testable (L n' m')
  where
    checkT _ = test

--------------------------------------------------------------------------------

instance KnownNat n => Normed (R n)
  where
    norm_0 v = norm_0 (extract v)
    norm_1 v = norm_1 (extract v)
    norm_2 v = norm_2 (extract v)
    norm_Inf v = norm_Inf (extract v)

instance (KnownNat m, KnownNat n) => Normed (L m n)
  where
    norm_0 m = norm_0 (extract m)
    norm_1 m = norm_1 (extract m)
    norm_2 m = norm_2 (extract m)
    norm_Inf m = norm_Inf (extract m)

mkSym f = Sym . f . unSym
mkSym2 f x y = Sym (f (unSym x) (unSym y))

instance KnownNat n =>  Num (Sym n)
  where
    (+) = mkSym2 (+)
    (*) = mkSym2 (*)
    (-) = mkSym2 (-)
    abs = mkSym abs
    signum = mkSym signum
    negate = mkSym negate
    fromInteger = Sym . fromInteger

instance KnownNat n => Fractional (Sym n)
  where
    fromRational = Sym . fromRational
    (/) = mkSym2 (/)

instance KnownNat n => Floating (Sym n)
  where
    sin   = mkSym sin
    cos   = mkSym cos
    tan   = mkSym tan
    asin  = mkSym asin
    acos  = mkSym acos
    atan  = mkSym atan
    sinh  = mkSym sinh
    cosh  = mkSym cosh
    tanh  = mkSym tanh
    asinh = mkSym asinh
    acosh = mkSym acosh
    atanh = mkSym atanh
    exp   = mkSym exp
    log   = mkSym log
    sqrt  = mkSym sqrt
    (**)  = mkSym2 (**)
    pi    = Sym pi

instance KnownNat n => Additive (Sym n) where
    add = (+)

instance KnownNat n => Transposable (Sym n) (Sym n) where
    tr  = id
    tr' = id

instance KnownNat n => Transposable (Her n) (Her n) where
    tr          = id
    tr' (Her m) = Her (tr' m)

max_stars_repo_name: schnecki/hmatrix-float

max_issues_repo_name: schnecki/hmatrix-float

max_forks_repo_name: schnecki/hmatrix-float


text:
{-# LANGUAGE DataKinds, RankNTypes, TypeFamilies #-}
module TestMnistFCNN
  ( testTrees, shortTestForCITrees, mnistTestCase2T, mnistTestCase2D
  ) where

import Prelude

import           Control.DeepSeq
import           Control.Monad (foldM, when)
import           Data.Array.Internal (valueOf)
import           Data.Coerce (coerce)
import           Data.List.Index (imap)
import           Data.Proxy (Proxy (Proxy))
import           Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime)
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat)
import qualified Numeric.LinearAlgebra as HM
import           System.IO (hFlush, hPutStrLn, stderr)
import           System.Random
import           Test.Tasty
import           Test.Tasty.HUnit hiding (assert)
import           Test.Tasty.QuickCheck hiding (label, shuffle)
import           Text.Printf

import HordeAd
import HordeAd.Core.OutdatedOptimizer
import HordeAd.Tool.MnistTools

import TestCommon

testTrees :: [TestTree]
testTrees = [ dumbMnistTests
            , bigMnistTests
            , vectorMnistTests
            , matrixMnistTests
            , fusedMnistTests
            ]

shortTestForCITrees :: [TestTree]
shortTestForCITrees = [ dumbMnistTests
                      , shortCIMnistTests
                      ]

sgdShow :: HasDelta r
        => r
        -> (a -> DualNumberVariables 'DModeGradient r -> DualMonadGradient r (DualNumber 'DModeGradient r))
        -> [a]  -- ^ training data
        -> Domain0 r  -- ^ initial parameters
        -> r
sgdShow gamma f trainData params0Init =
  let result =
        fst $ sgd gamma f trainData (params0Init, V.empty, V.empty, V.empty)
  in snd $ dReverse 1 (f $ head trainData) result

sgdTestCase :: String
            -> IO [a]
            -> (Int
                -> Int
                -> a
                -> DualNumberVariables 'DModeGradient Double
                -> DualMonadGradient Double
                                     (DualNumber 'DModeGradient Double))
            -> Double
            -> Double
            -> TestTree
sgdTestCase prefix trainDataIO trainWithLoss gamma expected =
  let widthHidden = 250
      widthHidden2 = 50
      nParams0 = fcnnMnistLen0 widthHidden widthHidden2
      vec = HM.randomVector 33 HM.Uniform nParams0 - HM.scalar 0.5
      name = prefix ++ " "
             ++ unwords [show widthHidden, show nParams0, show gamma]
  in testCase name $ do
       trainData <- trainDataIO
       sgdShow gamma (trainWithLoss widthHidden widthHidden2)
                     trainData vec
         @?= expected

mnistTestCase2
  :: String
  -> Int
  -> Int
  -> (Int
      -> Int
      -> MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2 prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
               gamma expected =
  let nParams0 = fcnnMnistLen0 widthHidden widthHidden2
      params0Init = HM.randomVector 44 HM.Uniform nParams0 - HM.scalar 0.5
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show gamma ]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: Domain0 Double -> (Int, [MnistData Double])
                    -> IO (Domain0 Double)
           runBatch !params0 (k, chunk) = do
             let f = trainWithLoss widthHidden widthHidden2
                 (!res, _, _, _) =
                   fst $ sgd gamma f chunk (params0, V.empty, V.empty, V.empty)
                 !trainScore = fcnnMnistTest0 (Proxy @Double)
                                         widthHidden widthHidden2 chunk res
                 !testScore  = fcnnMnistTest0 (Proxy @Double)
                                         widthHidden widthHidden2 testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int -> Domain0 Double -> IO (Domain0 Double)
           runEpoch n params0 | n > epochs = return params0
           runEpoch n params0 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params0 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 params0Init
       let testErrorFinal = 1 - fcnnMnistTest0 (Proxy @Double) widthHidden widthHidden2 testData res
       testErrorFinal @?= expected

mnistTestCase2V
  :: String
  -> Int
  -> Int
  -> (Int
      -> Int
      -> MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2V prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma expected =
  let (nParams0, nParams1, _, _) = fcnnMnistLen1 widthHidden widthHidden2
      params0Init = HM.randomVector 44 HM.Uniform nParams0 - HM.scalar 0.5
      params1Init = V.fromList $
        imap (\i nPV -> HM.randomVector (44 + nPV + i) HM.Uniform nPV
                        - HM.scalar 0.5)
             nParams1
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show (length nParams1)
                        , show (sum nParams1 + nParams0), show gamma ]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domain0 Double, Domain1 Double)
                    -> (Int, [MnistData Double])
                    -> IO (Domain0 Double, Domain1 Double)
           runBatch (!params0, !params1) (k, chunk) = do
             let f = trainWithLoss widthHidden widthHidden2
                 (resS, resV, _, _) =
                   fst $ sgd gamma f chunk (params0, params1, V.empty, V.empty)
                 res = (resS, resV)
                 !trainScore = fcnnMnistTest1
                                         widthHidden widthHidden2 chunk res
                 !testScore = fcnnMnistTest1
                                        widthHidden widthHidden2 testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int -> (Domain0 Double, Domain1 Double)
                    -> IO (Domain0 Double, Domain1 Double)
           runEpoch n params2 | n > epochs = return params2
           runEpoch n params2 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params2 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 (params0Init, params1Init)
       let testErrorFinal =
             1 - fcnnMnistTest1 widthHidden widthHidden2 testData res
       testErrorFinal @?= expected

fcnnMnistLossTanh :: DualMonad 'DModeGradient Double m
                => Int
                -> Int
                -> MnistData Double
                -> DualNumberVariables 'DModeGradient Double
                -> m (DualNumber 'DModeGradient Double)
fcnnMnistLossTanh widthHidden widthHidden2 (xs, targ) vec = do
  res <- fcnnMnist0 tanhAct softMaxAct widthHidden widthHidden2 xs vec
  lossCrossEntropy targ res

fcnnMnistLossRelu :: DualMonad 'DModeGradient Double m
                => Int
                -> Int
                -> MnistData Double
                -> DualNumberVariables 'DModeGradient Double
                -> m (DualNumber 'DModeGradient Double)
fcnnMnistLossRelu widthHidden widthHidden2 (xs, targ) vec = do
  res <- fcnnMnist0 reluAct softMaxAct widthHidden widthHidden2 xs vec
  lossCrossEntropy targ res

mnistTestCase2L
  :: String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2L prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma expected =
  let ((nParams0, nParams1, nParams2, _), totalParams, range, parameters0) =
        initializerFixed 44 0.5 (fcnnMnistLen2 widthHidden widthHidden2)
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: Domains Double
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double)
           runBatch (!params0, !params1, !params2, !paramsX) (k, chunk) = do
             let f = trainWithLoss
                 res = fst $ sgd gamma f chunk
                                 (params0, params1, params2, paramsX)
                 !trainScore = fcnnMnistTest2 @Double chunk res
                 !testScore = fcnnMnistTest2 @Double testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int
                    -> Domains Double
                    -> IO (Domains Double)
           runEpoch n params2 | n > epochs = return params2
           runEpoch n params2 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params2 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 parameters0
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2T
  :: Bool
  -> String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2T reallyWriteFile
                prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma expected =
  let ((nParams0, nParams1, nParams2, _), totalParams, range, !parameters0) =
        initializerFixed 44 0.5 (fcnnMnistLen2 widthHidden widthHidden2)
      name = prefix ++ " "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData0 <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       let !trainData = force $ shuffle (mkStdGen 6) trainData0
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domains Double, [(POSIXTime, Double)])
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runBatch ((!params0, !params1, !params2, !paramsX), !times)
                    (k, chunk) = do
             when (k `mod` 100 == 0) $ do
               hPutStrLn stderr $ printf "%s: %d " prefix k
               hFlush stderr
             let f = trainWithLoss
                 (!params0New, !value) =
                   sgd gamma f chunk (params0, params1, params2, paramsX)
             time <- getPOSIXTime
             return (params0New, (time, value) : times)
       let runEpoch :: Int
                    -> (Domains Double, [(POSIXTime, Double)])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runEpoch n params2times | n > epochs = return params2times
           runEpoch n (!params2, !times2) = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let !trainDataShuffled =
                   if n > 1
                   then shuffle (mkStdGen $ n + 5) trainData
                   else trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 1 trainDataShuffled
             res <- foldM runBatch (params2, times2) chunks
             runEpoch (succ n) res
       timeBefore <- getPOSIXTime
       (res, times) <- runEpoch 1 (parameters0, [])
       let ppTime (t, l) = init (show (t - timeBefore)) ++ " " ++ show l
       when reallyWriteFile $
         writeFile "walltimeLoss.txt" $ unlines $ map ppTime times
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2D
  :: Bool
  -> Int
  -> Bool
  -> String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2D reallyWriteFile miniBatchSize decay
                prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma0 expected =
  let np = fcnnMnistLen2 widthHidden widthHidden2
      ((nParams0, nParams1, nParams2, _), totalParams, range, !parameters0) =
        initializerFixed 44 0.5 np
      name = prefix ++ " "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma0, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData0 <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       let !trainData = force $ shuffle (mkStdGen 6) trainData0
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domains Double, [(POSIXTime, Double)])
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runBatch ((!params0, !params1, !params2, !paramsX), !times)
                    (k, chunk) = do
             when (k `mod` 100 == 0) $ do
               hPutStrLn stderr $ printf "%s: %d " prefix k
               hFlush stderr
             let f = trainWithLoss
                 gamma = if decay
                         then gamma0 * exp (- fromIntegral k * 1e-4)
                         else gamma0
                 (!params0New, !value) =
                   sgdBatchForward (33 + k * 7) miniBatchSize gamma f chunk
                                   (params0, params1, params2, paramsX) np
             time <- getPOSIXTime
             return (params0New, (time, value) : times)
       let runEpoch :: Int
                    -> (Domains Double, [(POSIXTime, Double)])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runEpoch n params2times | n > epochs = return params2times
           runEpoch n (!params2, !times2) = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let !trainDataShuffled =
                   if n > 1
                   then shuffle (mkStdGen $ n + 5) trainData
                   else trainData
                 chunks = take maxBatches
                          $ zip [1 ..]
                          $ chunksOf miniBatchSize trainDataShuffled
             res <- foldM runBatch (params2, times2) chunks
             runEpoch (succ n) res
       timeBefore <- getPOSIXTime
       (res, times) <- runEpoch 1 (parameters0, [])
       let ppTime (t, l) = init (show (t - timeBefore)) ++ " " ++ show l
       when reallyWriteFile $
         writeFile "walltimeLoss.txt" $ unlines $ map ppTime times
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2F
  :: Bool
  -> Int
  -> Bool
  -> String
  -> Int
  -> Int
  -> (MnistData Double
      -> DualNumberVariables 'DModeDerivative Double
      -> DualMonadForward Double (DualNumber 'DModeDerivative Double))
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
mnistTestCase2F reallyWriteFile miniBatchSize decay
                prefix epochs maxBatches trainWithLoss widthHidden widthHidden2
                gamma0 expected =
  let np = fcnnMnistLen2 widthHidden widthHidden2
      ((nParams0, nParams1, nParams2, _), totalParams, range, !parameters0) =
        initializerFixed 44 0.5 np
      name = prefix ++ " "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show totalParams, show gamma0, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData0 <- loadMnistData trainGlyphsPath trainLabelsPath
       testData <- loadMnistData testGlyphsPath testLabelsPath
       let !trainData = coerce $ force $ shuffle (mkStdGen 6) trainData0
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: (Domains Double, [(POSIXTime, Double)])
                    -> (Int, [MnistData Double])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runBatch ((!params0, !params1, !params2, !paramsX), !times)
                    (k, chunk) = do
             when (k `mod` 100 == 0) $ do
               hPutStrLn stderr $ printf "%s: %d " prefix k
               hFlush stderr
             let f = trainWithLoss
                 gamma = if decay
                         then gamma0 * exp (- fromIntegral k * 1e-4)
                         else gamma0
                 (!params0New, !value) =
                   sgdBatchFastForward (33 + k * 7) miniBatchSize gamma f chunk
                                       (params0, params1, params2, paramsX) np
             time <- getPOSIXTime
             return (params0New, (time, value) : times)
       let runEpoch :: Int
                    -> (Domains Double, [(POSIXTime, Double)])
                    -> IO (Domains Double, [(POSIXTime, Double)])
           runEpoch n params2times | n > epochs = return params2times
           runEpoch n (!params2, !times2) = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let !trainDataShuffled =
                   if n > 1
                   then shuffle (mkStdGen $ n + 5) trainData
                   else trainData
                 chunks = take maxBatches
                          $ zip [1 ..]
                          $ chunksOf miniBatchSize trainDataShuffled
             res <- foldM runBatch (params2, times2) chunks
             runEpoch (succ n) res
       timeBefore <- getPOSIXTime
       (res, times) <- runEpoch 1 (parameters0, [])
       let ppTime (t, l) = init (show (t - timeBefore)) ++ " " ++ show l
       when reallyWriteFile $
         writeFile "walltimeLoss.txt" $ unlines $ map ppTime times
       let testErrorFinal = 1 - fcnnMnistTest2 testData res
       testErrorFinal @?= expected

mnistTestCase2S
  :: forall widthHidden widthHidden2.
     (KnownNat widthHidden, KnownNat widthHidden2)
  => Proxy widthHidden -> Proxy widthHidden2
  -> String
  -> Int
  -> Int
  -> (forall d r m. DualMonad d r m
      => Proxy widthHidden -> Proxy widthHidden2
      -> MnistData r -> DualNumberVariables d r -> m (DualNumber d r))
  -> Double
  -> Double
  -> TestTree
mnistTestCase2S proxy proxy2
                prefix epochs maxBatches trainWithLoss gamma expected =
  let ((_, _, _, nParamsX), totalParams, range, parametersInit) =
        initializerFixed 44 0.5 (fcnnMnistLenS @widthHidden @widthHidden2)
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show (valueOf @widthHidden :: Int)
                        , show (valueOf @widthHidden2 :: Int)
                        , show nParamsX, show totalParams
                        , show gamma, show range ]
  in testCase name $ do
    hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
           prefix epochs maxBatches
    trainData <- loadMnistData trainGlyphsPath trainLabelsPath
    testData <- loadMnistData testGlyphsPath testLabelsPath
    let runBatch :: Domains Double
                 -> (Int, [MnistData Double])
                 -> IO (Domains Double)
        runBatch (!params0, !params1, !params2, !paramsX) (k, chunk) = do
          let f = trainWithLoss proxy proxy2
              res = fst $ sgd gamma f chunk
                              (params0, params1, params2, paramsX)
              !trainScore = fcnnMnistTestS @widthHidden @widthHidden2
                                          chunk res
              !testScore = fcnnMnistTestS @widthHidden @widthHidden2
                                         testData res
              !lenChunk = length chunk
          hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
          hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
          hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
          return res
    let runEpoch :: Int
                 -> Domains Double
                 -> IO (Domains Double)
        runEpoch n params2 | n > epochs = return params2
        runEpoch n params2 = do
          hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
          let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
              chunks = take maxBatches
                       $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
          !res <- foldM runBatch params2 chunks
          runEpoch (succ n) res
    res <- runEpoch 1 parametersInit
    let testErrorFinal = 1 - fcnnMnistTestS @widthHidden @widthHidden2
                                            testData res
    testErrorFinal @?= expected

dumbMnistTests :: TestTree
dumbMnistTests = testGroup "Dumb MNIST tests"
  [ testCase "1pretty-print in grey 3 2" $ do
      let (nParams0, lParams1, lParams2, _) = fcnnMnistLen2 4 3
          vParams1 = V.fromList lParams1
          vParams2 = V.fromList lParams2
          params0 = V.replicate nParams0 (1 :: Float)
          params1 = V.map (`V.replicate` 2) vParams1
          params2 = V.map (HM.konst 3) vParams2
          blackGlyph = V.replicate sizeMnistGlyph 4
          blackLabel = V.replicate sizeMnistLabel 5
          trainData = (blackGlyph, blackLabel)
          output = prettyPrintDf False (fcnnMnistLoss2 trainData)
                                 (params0, params1, params2, V.empty)
      -- printf "%s" output
      length output @?= 13348
  , testCase "2pretty-print in grey 3 2 fused" $ do
      let (nParams0, lParams1, lParams2, _) = fcnnMnistLen2 4 3
          vParams1 = V.fromList lParams1
          vParams2 = V.fromList lParams2
          params0 = V.replicate nParams0 (1 :: Float)
          params1 = V.map (`V.replicate` 2) vParams1
          params2 = V.map (HM.konst 3) vParams2
          blackGlyph = V.replicate sizeMnistGlyph 4
          blackLabel = V.replicate sizeMnistLabel 5
          trainData = (blackGlyph, blackLabel)
          output = prettyPrintDf True (fcnnMnistLossFused2 trainData)
                                 (params0, params1, params2, V.empty)
      --- printf "%s" output
      length output @?= 12431
  , testCase "3pretty-print on testset 3 2" $ do
      let (_, _, _, parameters0) = initializerFixed 44 0.5 (fcnnMnistLen2 4 3)
      testData <- loadMnistData testGlyphsPath testLabelsPath
      let trainDataItem = head testData
          output = prettyPrintDf True (fcnnMnistLoss2 trainDataItem) parameters0
      -- printf "%s" output
      length output @?= 16449
  , let blackGlyph = V.replicate sizeMnistGlyph 0
        blackLabel = V.replicate sizeMnistLabel 0
        trainData = replicate 10 (blackGlyph, blackLabel)
    in sgdTestCase "black"
         (return trainData) fcnnMnistLoss0 0.02 (-0.0)
  , let whiteGlyph = V.replicate sizeMnistGlyph 1
        whiteLabel = V.replicate sizeMnistLabel 1
        trainData = replicate 20 (whiteGlyph, whiteLabel)
    in sgdTestCase "white"
         (return trainData) fcnnMnistLoss0 0.02 23.02585095418536
  , let blackGlyph = V.replicate sizeMnistGlyph 0
        whiteLabel = V.replicate sizeMnistLabel 1
        trainData = replicate 50 (blackGlyph, whiteLabel)
    in sgdTestCase "black/white"
         (return trainData) fcnnMnistLoss0 0.02 23.025850929940457
  , let glyph = V.unfoldrExactN sizeMnistGlyph (uniformR (0, 1))
        label = V.unfoldrExactN sizeMnistLabel (uniformR (0, 1))
        trainData = map ((\g -> (glyph g, label g)) . mkStdGen) [1 .. 100]
    in sgdTestCase "random 100"
         (return trainData) fcnnMnistLoss0 0.02 11.089140063760212
  , sgdTestCase "first 100 trainset samples only"
      (take 100 <$> loadMnistData trainGlyphsPath trainLabelsPath)
      fcnnMnistLoss0 0.02 3.233123290489956
  , testCase "fcnnMnistTest0 on 0.1 params0 300 100 width 10k testset" $ do
      let nParams0 = fcnnMnistLen0 300 100
          params0 = V.replicate nParams0 0.1
      testData <- loadMnistData testGlyphsPath testLabelsPath
      (1 - fcnnMnistTest0 (Proxy @Double) 300 100 testData params0)
        @?= 0.902
  , testCase "fcnnMnistTest2VV on 0.1 params0 300 100 width 10k testset" $ do
      let (nParams0, nParams1, _, _) = fcnnMnistLen1 300 100
          params0 = V.replicate nParams0 0.1
          params1 = V.fromList $ map (`V.replicate` 0.1) nParams1
      testData <- loadMnistData testGlyphsPath testLabelsPath
      (1 - fcnnMnistTest1 300 100 testData (params0, params1))
        @?= 0.902
  , testCase "fcnnMnistTest2LL on 0.1 params0 300 100 width 10k testset" $ do
      let (nParams0, lParams1, lParams2, _) = fcnnMnistLen2 300 100
          vParams1 = V.fromList lParams1
          vParams2 = V.fromList lParams2
          params0 = V.replicate nParams0 0.1
          params1 = V.map (`V.replicate` 0.1) vParams1
          params2 = V.map (HM.konst 0.1) vParams2
      testData <- loadMnistData testGlyphsPath testLabelsPath
      (1 - fcnnMnistTest2 testData
                      (params0, params1, params2, V.empty))
        @?= 0.902
  , testProperty "Compare two forward derivatives and gradient for Mnist0" $
      \seed seedDs ->
      forAll (choose (1, 300)) $ \widthHidden ->
      forAll (choose (1, 100)) $ \widthHidden2 ->
      forAll (choose (0.01, 10)) $ \range ->
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = createRandomVector sizeMnistGlyph seed
            label = createRandomVector sizeMnistLabel seedDs
            mnistData :: MnistData Double
            mnistData = (glyph, label)
            nParams0 = fcnnMnistLen0 widthHidden widthHidden2
            paramShape = (nParams0, [], [], [])
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f :: forall d r m. (DualMonad d r m, r ~ Double)
              => DualNumberVariables d r -> m (DualNumber d r)
            f = fcnnMnistLoss0 widthHidden widthHidden2 mnistData
        in
            qcPropDom f parameters ds parametersPerturbation 1
  , testProperty "Compare two forward derivatives and gradient for Mnist1" $
      \seed seedDs ->
      forAll (choose (1, 2000)) $ \widthHidden ->
      forAll (choose (1, 5000)) $ \widthHidden2 ->
      forAll (choose (0.01, 0.5)) $ \range ->  -- large nn, so NaNs fast
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = createRandomVector sizeMnistGlyph seed
            label = createRandomVector sizeMnistLabel seedDs
            mnistData :: MnistData Double
            mnistData = (glyph, label)
            paramShape = fcnnMnistLen1 widthHidden widthHidden2
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f :: forall d r m. (DualMonad d r m, r ~ Double)
              => DualNumberVariables d r -> m (DualNumber d r)
            f = fcnnMnistLoss1 widthHidden widthHidden2 mnistData
        in
            qcPropDom f parameters ds parametersPerturbation 1
  , testProperty "Compare two forward derivatives and gradient for Mnist2" $
      \seed ->
      forAll (choose (0, sizeMnistLabel - 1)) $ \seedDs ->
      forAll (choose (1, 5000)) $ \widthHidden ->
      forAll (choose (1, 1000)) $ \widthHidden2 ->
      forAll (choose (0.01, 1)) $ \range ->
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = createRandomVector sizeMnistGlyph seed
            label = createRandomVector sizeMnistLabel seedDs
            labelOneHot = HM.konst 0 sizeMnistLabel V.// [(seedDs, 1)]
            mnistData, mnistDataOneHot :: MnistData Double
            mnistData = (glyph, label)
            mnistDataOneHot = (glyph, labelOneHot)
            paramShape = fcnnMnistLen2 widthHidden widthHidden2
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f, fOneHot, fFused
              :: forall d r m. (DualMonad d r m, r ~ Double)
                 => DualNumberVariables d r -> m (DualNumber d r)
            f = fcnnMnistLoss2 mnistData
            fOneHot = fcnnMnistLoss2 mnistDataOneHot
            fFused = fcnnMnistLossFused2 mnistDataOneHot
        in
            qcPropDom f       parameters ds parametersPerturbation 1 .&&.
            qcPropDom fOneHot parameters ds parametersPerturbation 1 .&&.
            qcPropDom fFused  parameters ds parametersPerturbation 1 .&&.
            cmpTwoSimple fOneHot fFused  parameters ds
  ]

bigMnistTests :: TestTree
bigMnistTests = testGroup "MNIST tests with a 2-hidden-layer nn"
  [ mnistTestCase2 "1 epoch, 1 batch" 1 1 fcnnMnistLoss0 300 100 0.02
                   0.1269
  , mnistTestCase2 "tanh: 1 epoch, 1 batch" 1 1 fcnnMnistLossTanh 300 100 0.02
                   0.6406000000000001
  , mnistTestCase2 "relu: 1 epoch, 1 batch" 1 1 fcnnMnistLossRelu 300 100 0.02
                   0.7248
  , mnistTestCase2 "1 epoch, 1 batch, wider" 1 1 fcnnMnistLoss0 500 150 0.02
                   0.1269
  , mnistTestCase2 "2 epochs, but only 1 batch" 2 1 fcnnMnistLoss0 300 100 0.02
                   9.809999999999997e-2
  , mnistTestCase2 "artificial 1 2 3 4 5" 1 2 fcnnMnistLoss0 3 4 5
                   0.8972
  , mnistTestCase2 "artificial 5 4 3 2 1" 5 4 fcnnMnistLoss0 3 2 1
                   0.8991
  ]

vectorMnistTests :: TestTree
vectorMnistTests = testGroup "MNIST VV tests with a 2-hidden-layer nn"
  [ mnistTestCase2V "1 epoch, 1 batch" 1 1 fcnnMnistLoss1 300 100 0.02
                    0.12960000000000005
  , mnistTestCase2V "1 epoch, 1 batch, wider" 1 1 fcnnMnistLoss1 500 150 0.02
                    0.13959999999999995
  , mnistTestCase2V "2 epochs, but only 1 batch" 2 1 fcnnMnistLoss1 300 100 0.02
                    0.10019999999999996
  , mnistTestCase2V "1 epoch, all batches" 1 99 fcnnMnistLoss1 300 100 0.02
                    5.389999999999995e-2
  , mnistTestCase2V "artificial 1 2 3 4 5" 1 2 fcnnMnistLoss1 3 4 5
                    0.8972
  , mnistTestCase2V "artificial 5 4 3 2 1" 5 4 fcnnMnistLoss1 3 2 1
                    0.7756000000000001
  ]

matrixMnistTests :: TestTree
matrixMnistTests = testGroup "MNIST LL tests with a 2-hidden-layer nn"
  [ mnistTestCase2L "1 epoch, 1 batch" 1 1 fcnnMnistLoss2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "1 epoch, 1 batch, wider" 1 1 fcnnMnistLoss2 500 150 0.02
                    0.15039999999999998
  , mnistTestCase2L "2 epochs, but only 1 batch" 2 1 fcnnMnistLoss2 300 100 0.02
                    8.879999999999999e-2
  , mnistTestCase2L "1 epoch, all batches" 1 99 fcnnMnistLoss2 300 100 0.02
                    5.1100000000000034e-2
  , mnistTestCase2L "artificial 1 2 3 4 5" 1 2 fcnnMnistLoss2 3 4 5
                    0.8972
  , mnistTestCase2T False
                    "artificial TL 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8865
  , mnistTestCase2D False 1 False
                    "artificial DL 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8991
  , mnistTestCase2F False 1 False
                    "artificial FL 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8991
--  , mnistTestCase2T True False
--                    "2 epochs, all batches, TL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    4.290000000000005e-2
--  , mnistTestCase2D True 1 False
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    0.9079
--  , mnistTestCase2D True 64 False
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    0.9261
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 0.02
--                    0.8993
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 2e-5
--                    0.9423
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 2e-4
--                    0.8714
--  , mnistTestCase2F True 64 True
--                    "2 epochs, all batches, FL, wider, to file"
--                    2 60000 fcnnMnistLoss2 500 150 2e-4
--                    0.8714
--  , mnistTestCase2D True 64 True
--                    "2 epochs, all batches, DL, wider, to file"
--                    2 60000 fcnnMnistLossFusedRelu2 1024 1024 2e-4
--                    0.902
--  , mnistTestCase2D False 64 True
--                    "2 epochs, all batches, 1024DL"
--                    2 60000 fcnnMnistLoss2 1024 1024 2e-4
--                    0.7465999999999999
--  , mnistTestCase2F False 64 True
--                    "2 epochs, all batches, 1024FL"
--                    2 60000 fcnnMnistLoss2 1024 1024 2e-4
--                    0.7465999999999999
  ]

fusedMnistTests :: TestTree
fusedMnistTests = testGroup "MNIST fused LL tests with a 2-hidden-layer nn"
  [ mnistTestCase2L "1 epoch, 1 batch" 1 1 fcnnMnistLossFused2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "1 epoch, 1 batch, wider" 1 1
                    fcnnMnistLossFused2 500 150 0.02
                    0.15039999999999998
  , mnistTestCase2L "2 epochs, but only 1 batch" 2 1
                    fcnnMnistLossFused2 300 100 0.02
                    8.879999999999999e-2
  , mnistTestCase2L "1 epoch, all batches" 1 99 fcnnMnistLossFused2 300 100 0.02
                    5.1100000000000034e-2
  , mnistTestCase2L "artificial 1 2 3 4 5" 1 2 fcnnMnistLossFused2 3 4 5
                    0.8972
  , mnistTestCase2L "artificial 5 4 3 2 1" 5 4 fcnnMnistLossFused2 3 2 1
                    0.7033
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 1 epoch, 1 batch" 1 1 fcnnMnistLossFusedS 0.02
                    0.1311
  , mnistTestCase2S (Proxy @500) (Proxy @150)
                    "S 1 epoch, 1 batch, wider" 1 1 fcnnMnistLossFusedS 0.02
                    0.12470000000000003
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 2 epochs, but only 1 batch" 2 1 fcnnMnistLossFusedS 0.02
                    9.630000000000005e-2
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 1 epoch, all batches" 1 99 fcnnMnistLossFusedS 0.02
                    5.620000000000003e-2
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "S artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "S artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedS 1
                    0.8246
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "SR 1 epoch, 1 batch" 1 1 fcnnMnistLossFusedReluS 0.02
                    0.7068
  , mnistTestCase2S (Proxy @500) (Proxy @150)
                    "SR 1 epoch, 1 batch, wider" 1 1
                    fcnnMnistLossFusedReluS 0.02
                    0.8874
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "SR 2 epochs, but 1 batch" 2 1 fcnnMnistLossFusedReluS 0.02
                    0.8352999999999999
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "SR 1 epoch, all batches" 1 99 fcnnMnistLossFusedReluS 0.02
                    0.6415
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "SR artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedReluS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "SR artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedReluS 1
                    0.8991
  ]

shortCIMnistTests :: TestTree
shortCIMnistTests = testGroup "Short CI MNIST tests"
  [ mnistTestCase2 "2 artificial 1 2 3 4 5" 1 2 fcnnMnistLoss0 3 4 5
                   0.8972
  , mnistTestCase2 "2 artificial 5 4 3 2 1" 5 4 fcnnMnistLoss0 3 2 1
                   0.8991
  , mnistTestCase2V "VV 1 epoch, 1 batch" 1 1 fcnnMnistLoss1 300 100 0.02
                    0.12960000000000005
  , mnistTestCase2V "VV artificial 1 2 3 4 5" 1 2 fcnnMnistLoss1 3 4 5
                    0.8972
  , mnistTestCase2V "VV artificial 5 4 3 2 1" 5 4 fcnnMnistLoss1 3 2 1
                    0.7756000000000001
  , mnistTestCase2L "LL 1 epoch, 1 batch" 1 1 fcnnMnistLoss2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "LL artificial 1 2 3 4 5" 1 2 fcnnMnistLoss2 3 4 5
                    0.8972
  , mnistTestCase2L "LL artificial 5 4 3 2 1" 5 4 fcnnMnistLoss2 3 2 1
                    0.8085
  , mnistTestCase2L "fused LL 1/1 batch" 1 1 fcnnMnistLossFused2 300 100 0.02
                    0.12339999999999995
  , mnistTestCase2L "fused LL artificial 1 2 3 4 5" 1 2
                    fcnnMnistLossFused2 3 4 5
                    0.8972
  , mnistTestCase2T False
                    "fused TL artificial 5 4 3 2 1" 5 4
                    fcnnMnistLossFused2 3 2 1
                    0.8865
  , mnistTestCase2D False 1 False
                    "fused DL artificial 5 4 3 2 1" 5 4
                    fcnnMnistLossFused2 3 2 1
                    0.8991
  , mnistTestCase2S (Proxy @300) (Proxy @100)
                    "S 1 epoch, 1 batch" 1 1 fcnnMnistLossFusedS 0.02
                    0.1311
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "S artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "S artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedS 1
                    0.8246
  , mnistTestCase2S (Proxy @3) (Proxy @4)
                    "SR artificial 1 2 3 4 5" 1 2 fcnnMnistLossFusedReluS 5
                    0.8972
  , mnistTestCase2S (Proxy @3) (Proxy @2)
                    "SR artificial 5 4 3 2 1" 5 4 fcnnMnistLossFusedReluS 1
                    0.8991
  ]

ext: hs

max_stars_repo_name: Mikolaj/horde-ad

max_issues_repo_name: Mikolaj/horde-ad

max_forks_repo_name: Mikolaj/horde-ad


text:
{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Vector
-- Copyright  : Copyright (c) , Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Immutable dense vectors.

module Numeric.LinearAlgebra.Vector (
    -- * Immutable vectors
    Vector,
    dim,

    -- * Vector construction
    fromList,
    zero,
    constant,

    -- * Accessing vectors
    at,
    unsafeAt,
    indices,
    elems,
    assocs,

    -- * Incremental vector updates
    update,
    unsafeUpdate,
    accum,
    unsafeAccum,

    -- * Derived vectors
    map,
    zipWith,
    unsafeZipWith,
    concat,

    -- * Vector views
    slice,
    unsafeSlice,
    splitAt,
    drop,
    take,

    -- * Vector properties
    sumAbs,
    norm2,
    whichMaxAbs,
    dot,
    unsafeDot,
    kronecker,

    -- * Vector math operations
    -- ** Num
    add,
    addWithScale,
    sub,
    scale,
    mul,
    negate,
    conjugate,
    abs,
    signum,

    -- ** Fractional
    div,
    recip,

    -- ** Floating
    sqrt,
    exp,
    log,
    pow,
    sin,
    cos,
    tan,
    asin,
    acos,
    atan,
    sinh,
    cosh,
    tanh,
    asinh,
    acosh,
    atanh,

    -- * Conversions between foreign pointers
    unsafeFromForeignPtr,
    unsafeToForeignPtr,

    -- * Mutable interface
    module Numeric.LinearAlgebra.Vector.ST,

    -- * Basic multivariate statistics
    module Numeric.LinearAlgebra.Vector.Statistics,

    ) where

import Prelude( Int, Double, ($), (*), return )
import Control.Monad.ST( runST, ST )
import Foreign( Storable )
import Foreign.BLAS( BLAS1, BLAS2 )
import Foreign.VMath( VNum, VFractional, VFloating )

import Numeric.LinearAlgebra.Vector.Base
import Numeric.LinearAlgebra.Vector.ST
import Numeric.LinearAlgebra.Vector.Statistics

infixr 8 `pow`
infixl 7 `div`
infixl 7 `mul`, `scale`, `kronecker`
infixl 6 `add`, `sub`

-- | Compute the sum of absolute values of entries in the vector.
sumAbs :: (BLAS1 e) => Vector e -> Double
sumAbs v = runST $ getSumAbs v
{-# INLINE sumAbs #-}

-- | Compute the 2-norm (Euclidean norm) of a vector.
norm2 :: (BLAS1 e) => Vector e -> Double
norm2 v = runST $ getNorm2 v
{-# INLINE norm2 #-}

-- | Get the index and norm of the element with absulte value.  Not valid
-- if any of the vector entries are @NaN@.  Raises an exception if the
-- vector has length @0@.
whichMaxAbs :: (BLAS1 e) => Vector e -> (Int, e)
whichMaxAbs v = runST $ getWhichMaxAbs v
{-# INLINE whichMaxAbs #-}

-- | Compute the dot product of two vectors.
dot :: (BLAS1 e) => Vector e -> Vector e -> e
dot v v' = runST $ getDot v v'
{-# INLINE dot #-}

unsafeDot :: (BLAS1 e) => Vector e -> Vector e -> e
unsafeDot v v' = runST $ unsafeGetDot v v'
{-# INLINE unsafeDot #-}

-- | Compute the kronecker product of two vectors.
kronecker :: (BLAS2 e) => Vector e -> Vector e -> Vector e
kronecker x y = create $ do
    z <- new_ (dim x * dim y)
    kroneckerTo z x y
    return z

-- | @add x y@ returns @x + y@.
add :: (VNum e) => Vector e -> Vector e -> Vector e
add = result2 addTo

-- | @sub x y@ returns @x - y@.
sub :: (VNum e) => Vector e -> Vector e -> Vector e
sub = result2 subTo

-- | @scale k x@ returns @k * x@.
scale :: (BLAS1 e) => e -> Vector e -> Vector e
scale k x = create $ do
    x' <- newCopy x
    scaleM_ k x'
    return x'

-- | @addWithScale alpha x y@ return @alpha * x + y@.
addWithScale :: (BLAS1 e) => e -> Vector e -> Vector e -> Vector e
addWithScale alpha x y = create $ do
    y' <- newCopy y
    addWithScaleM_ alpha x y'
    return y'

-- | @mul x y@ returns @x * y@.
mul :: (VNum e) => Vector e -> Vector e -> Vector e
mul = result2 mulTo

-- | @negate x@ returns @-x@.
negate :: (VNum e) => Vector e -> Vector e
negate = result negateTo

-- | @conjugate x@ returns @conjugate(x)@.
conjugate :: (VNum e) => Vector e -> Vector e
conjugate = result conjugateTo

-- | @abs x@ returns @abs(x)@.
abs :: (VNum e) => Vector e -> Vector e
abs = result absTo

-- | @signum x@ returns @signum(x)@.
signum :: (VNum e) => Vector e -> Vector e
signum = result signumTo

-- | @div x y@ returns @x / y@.
div :: (VFractional e) => Vector e -> Vector e -> Vector e
div = result2 divTo

-- | @recip x y@ returns @1 / x@.
recip :: (VFractional e) => Vector e -> Vector e
recip = result recipTo

-- | @sqrt x@ returns @sqrt(x)@.
sqrt :: (VFloating e) => Vector e -> Vector e
sqrt = result sqrtTo

-- | @exp x@ returns @exp(x)@.
exp :: (VFloating e) => Vector e -> Vector e
exp = result expTo

-- | @log x@ returns @log(x)@.
log :: (VFloating e) => Vector e -> Vector e
log = result logTo

-- | @pow x y@ returns @x ** y@.
pow :: (VFloating e) => Vector e -> Vector e -> Vector e
pow = result2 powTo

-- | @sin x@ returns @sin(x)@.
sin :: (VFloating e) => Vector e -> Vector e
sin = result sinTo

-- | @cos x@ returns @cos(x)@.
cos :: (VFloating e) => Vector e -> Vector e
cos = result cosTo

-- | @tan x@ returns @tan(x)@.
tan :: (VFloating e) => Vector e -> Vector e
tan = result tanTo

-- | @asin x@ returns @asin(x)@.
asin :: (VFloating e) => Vector e -> Vector e
asin = result asinTo

-- | @acos x@ returns @acos(x)@.
acos :: (VFloating e) => Vector e -> Vector e
acos = result acosTo

-- | @atan x@ returns @atan(x)@.
atan :: (VFloating e) => Vector e -> Vector e
atan = result atanTo

-- | @sinh x@ returns @sinh(x)@.
sinh :: (VFloating e) => Vector e -> Vector e
sinh = result sinhTo

-- | @cosh x@ returns @cosh(x)@.
cosh :: (VFloating e) => Vector e -> Vector e
cosh = result coshTo

-- | @tanh x@ returns @tanh(x)@.
tanh :: (VFloating e) => Vector e -> Vector e
tanh = result tanhTo

-- | @asinh x@ returns @asinh(x)@.
asinh :: (VFloating e) => Vector e -> Vector e
asinh = result asinhTo

-- | @acosh x@ returns @acosh(x)@.
acosh :: (VFloating e) => Vector e -> Vector e
acosh = result acoshTo

-- | @atanh x@ returns @atanh(x)@.
atanh :: (VFloating e) => Vector e -> Vector e
atanh = result atanhTo

result :: (Storable e, Storable f)
       => (forall s . STVector s f -> Vector e -> ST s a)
       -> Vector e
       -> Vector f
result f v = create $ newResult f v
{-# INLINE result #-}

result2 :: (Storable e, Storable f, Storable g)
        => (forall s . STVector s g -> Vector e -> Vector f -> ST s a)
        -> Vector e
        -> Vector f
        -> Vector g
result2 f v1 v2 = create $ newResult2 f v1 v2
{-# INLINE result2 #-}

newResult :: (RVector v, Storable e, Storable f)
          => (STVector s f -> v e -> ST s a)
          -> v e
          -> ST s (STVector s f)
newResult f v = do
    n <- getDim v
    z <- new_ n
    _ <- f z v
    return z
{-# INLINE newResult #-}

newResult2 :: (RVector v1, RVector v2, Storable e, Storable f, Storable g)
           => (STVector s g -> v1 e -> v2 f -> ST s a)
           -> v1 e
           -> v2 f
           -> ST s (STVector s g)
newResult2 f v1 v2 = do
    n <- getDim v1
    z <- new_ n
    _ <- f z v1 v2
    return z
{-# INLINE newResult2 #-}

max_stars_repo_name: patperry/hs-linear-algebra

max_issues_repo_name: cartazio/hs-cblas

max_forks_repo_name: cartazio/hs-cblas


text:
module Form where

import           Data.Complex
import           Control.Arrow
import           Settings

type Form = Double -> Complex Double

--f(t) = c * exp(n * i * 2pi * t)
form :: Complex Double -> Int -> Double -> Complex Double
form c n t = c * cis (fromIntegral n * 2 * pi * t)

-- exact coordinates
formsToCoordsExact :: [Form] -> [(Double, Double)]
formsToCoordsExact forms = map (realPart &&& imagPart) $ ps timeIntervals forms
  where
    ps []       _  = []
    ps (t : ts) fs = sum (map ($t) fs) : ps ts fs

-- [0, 0.001, ...]
timeIntervals :: [Double]
timeIntervals = ti amount [1]
  where
    ti 0 ds = ds
    ti i ds = ti (i - 1) (interval * i : ds)
    amount   = fromIntegral pointCount
    interval = 1 / (amount + 1)

formsToCoords :: [Form] -> [(Int, Int)]
formsToCoords fs = map (f *** f) $ formsToCoordsExact fs
    where f = (*) scaleFactor . round


text:
{-# LANGUAGE ViewPatterns, ScopedTypeVariables, TupleSections #-}

import qualified Text.Printf as TP (printf)
import qualified Data.List as L (intercalate)

import Data.Array.Repa (Z(Z), (:.)((:.)), (!), (-^), (+^), (*^), All(All), Any(Any), extent, size)
import qualified Data.Array.Repa as RP (fromFunction, fromListUnboxed, Array, U, D, DIM1, DIM2, DIM4, map)
import qualified Data.Array.Repa as RP (sumAllS, slice, delay, computeS, toList)
import qualified Data.Array.Repa.Repr.Vector as RP (V, fromListVector)
import Data.Complex (Complex((:+)), realPart, imagPart)
import Data.Bits (bit)

pdim1 :: RP.Array RP.U RP.DIM1 (Complex Double) -> String
pdim1 mat = pt ++ "\n" where
    (Z :. n) = extent mat
    pt = concat [ let (xr :+ xi) = mat ! (Z :. i) in TP.printf "%14.4f +I %8.4f" xr xi | i <- [0 .. n - 1] ]

pdim2 :: RP.Array RP.U RP.DIM2 (Complex Double) -> String
pdim2 mat = L.intercalate "\n" [ pt i | i <- [0 .. m - 1]] ++ "\n\n" where
    (Z :. m :. n) = extent mat
    pt i = concat [ let (xr :+ xi) = mat ! (Z :. i :. j) in TP.printf "%14.4f +I %8.4f" xr xi | j <- [0 .. n - 1] ]

identity :: RP.Array RP.U RP.DIM2 (Complex Double)
identity = RP.fromListUnboxed (Z :. 2 :. 2) [1, 0, 0, 1]

sigmaX :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaX = RP.fromListUnboxed (Z :. 2 :. 2) [0, 1, 1, 0]

sigmaY :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaY = RP.fromListUnboxed (Z :. 2 :. 2) [0, 0 :+ (-1), 0 :+ 1, 0]

sigmaZ :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaZ = RP.fromListUnboxed (Z :. 2 :. 2) [1, 0, 0, -1]

-- | Tensor product.
kron :: RP.Array RP.U RP.DIM2 (Complex Double)
     -> RP.Array RP.U RP.DIM2 (Complex Double)
     -> RP.Array RP.U RP.DIM2 (Complex Double)
kron mata matb = reduceBlock . RP.computeS $ RP.fromFunction (extent mata) iter where
    iter i = RP.computeS $ RP.map (* (mata ! i)) matb :: RP.Array RP.U RP.DIM2 (Complex Double)

-- | Heisenberg Hamiltonian
hamil :: RP.Array RP.U RP.DIM2 (Complex Double)
hamil = RP.computeS $ RP.map (* 0.25) (((kron sigmaX identity) `mmult` (kron identity sigmaX))
    +^ ((kron sigmaY identity) `mmult` (kron identity sigmaY))
    +^ ((kron sigmaZ identity) `mmult` (kron identity sigmaZ)))

-- | S^2 Operator
s2 :: RP.Array RP.U RP.DIM2 (Complex Double)
s2 = RP.computeS $ RP.map (* 0.25) (
    square (RP.computeS $ kron sigmaX identity +^ kron identity sigmaX)
    +^ square (RP.computeS $ kron sigmaY identity +^ kron identity sigmaY)
    +^ square (RP.computeS $ kron sigmaZ identity +^ kron identity sigmaZ))

-- | Square of a matrix
square :: RP.Array RP.U RP.DIM2 (Complex Double) -> RP.Array RP.U RP.DIM2 (Complex Double)
square mat = mmult mat mat

-- | Matrix Matrix Multiplication
mmult :: RP.Array RP.U RP.DIM2 (Complex Double)
      -> RP.Array RP.U RP.DIM2 (Complex Double)
      -> RP.Array RP.U RP.DIM2 (Complex Double)
mmult mata matb = RP.computeS $ RP.fromFunction (Z :. n :. m) iter where
    (Z :. n :. k) = extent mata
    (Z :. kp :. m) = extent matb
    iter :: RP.DIM2 -> Complex Double
    iter (Z :. i :. j) = RP.sumAllS $ RP.slice mata (Any :. i :. All) *^ RP.slice matb (Any :. All :. j)

-- | Matrix Vector Multiplication
mvmult :: RP.Array RP.U RP.DIM2 (Complex Double)
       -> RP.Array RP.U RP.DIM1 (Complex Double)
       -> RP.Array RP.U RP.DIM1 (Complex Double)
mvmult mat vec = RP.computeS $ RP.fromFunction (Z :. n) iter where
    (Z :. n :. m) = extent mat
    iter :: RP.DIM1 -> Complex Double
    iter (Z :. i) = RP.sumAllS $ RP.slice mat (Any :. i :. All) *^ vec

-- | Vector Vector Dot Product
dot :: RP.Array RP.U RP.DIM1 (Complex Double) -> RP.Array RP.U RP.DIM1 (Complex Double) -> Complex Double
dot veca vecb = RP.sumAllS $ veca *^ vecb

-- | Reduce blocked matrix to regular matrix.
reduceBlock :: RP.Array RP.V RP.DIM2 (RP.Array RP.U RP.DIM2 (Complex Double)) -- ^ matrix of blocks
            -> RP.Array RP.U RP.DIM2 (Complex Double)
reduceBlock mat = RP.computeS $ RP.fromFunction (Z :. nl :. nl) iter where
    f :: Int -> RP.Array RP.U RP.DIM2 (Complex Double) -> [(Int, Int)]
    f i block = let (Z :. _ :. n) = extent block in map (i, ) [0..n - 1]
    idx = let l = concat (zipWith f [0..] (RP.toList (RP.slice mat (Any :. (0 :: Int) :. All))))
        in RP.fromListUnboxed (Z :. (length l)) l
    (Z :. nl) = extent idx
    iter (Z :. i :. j) = let (ip, im) = idx ! (Z :. i); (jp, jm) = idx ! (Z :. j)
        in (mat ! (Z :. ip :. jp)) ! (Z :. im :. jm)

-- alpha = 0, beta = 1
-- |babb> = 1011 = 11
stateIndex :: String -> Int
stateIndex x = iter (length x - 1) x where
    iter n (x:xs) = (if n == 0 then 0 else iter (n - 1) xs) + (if x == 'a' then 0 else bit n)

-- | The kth eigenstate represented as n-component vector
eigenState :: Int -> Int -> RP.Array RP.U RP.DIM1 (Complex Double)
eigenState n k = RP.computeS $ RP.fromFunction (Z :. n) (\(Z :. i) -> if k == i then 1 else 0)

state_ab = eigenState (bit 2) $ stateIndex "ab"
state_ba = eigenState (bit 2) $ stateIndex "ba"
state_aa = eigenState (bit 2) $ stateIndex "aa"
state_bb = eigenState (bit 2) $ stateIndex "bb"
singlet = RP.computeS $ RP.map (* (2 ** (-0.5))) (state_ab -^ state_ba)
triplet = RP.computeS $ RP.map (* (2 ** (-0.5))) (state_ab +^ state_ba)

main = do
    putStrLn $ "[Singlet] (Si . Sj) S^2"
    print $ let state = singlet in dot state (hamil `mvmult` state)
    print $ let state = singlet in dot state (s2 `mvmult` state)
    putStrLn $ "[Triplet] (Si . Sj) S^2"
    print $ let state = triplet in dot state (hamil `mvmult` state)
    print $ let state = triplet in dot state (s2 `mvmult` state)
    putStrLn $ "[|aa>] (Si . Sj) S^2"
    print $ let state = state_aa in dot state (hamil `mvmult` state)
    print $ let state = state_aa in dot state (s2 `mvmult` state)

max_stars_repo_name: stczhc/notes

max_issues_repo_name: stczhc/notes

max_forks_repo_name: stczhc/notes


text:
{-# LANGUAGE FlexibleContexts, TypeFamilies #-}

module Neural.LayerSpawners where

import Control.Monad
import Control.Parallel.Strategies
import Numeric.LinearAlgebra

import Neural.FF
import Neural.Util
import Numeric.Extra

maxPooling :: (Int, Int) -> [Matrix Double] -> [Matrix Double]
maxPooling (m_pool, n_pool) = map (fromLists . (map (map maxElement)) . (toBlocksEvery m_pool n_pool))

-- propagates the gradient through the pooling.
-- the maximum element in each $1 block of $2 is set to its corresponding gradient, and all other elements are set to zero.
maxPoolingBP :: (Int, Int) -> [Matrix Double] -> [Matrix Double] -> [Matrix Double]
maxPoolingBP (m_pool, n_pool) = zipWith (\nb_k x_k -> fromBlocks $ zipWith (zipWith demax)
                                                     (toLists nb_k) (toBlocksEvery m_pool n_pool x_k))
   where demax nb__ x__ = accum (konst 0.0 (m_pool,n_pool) :: Matrix Double) (+) [(maxIndex x__, nb__)]

-- the default layer opts
layerOpts :: FFLayerOptions
layerOpts = FFLayerOptions {
  layerDepth = 1,
  layerActivationFn = map (cmap sigma),
  layerActivationFnDeriv = map (cmap sigmaPrime),
  pooling = (\_ -> id), -- default pooling and poolingBP don't affect input ∨ gradient
  poolingBP = (\_ nb _ -> nb),
  poolSize = (0,0)
  }

-- A fully connected layer is just one big slice that connects to all slices in the prev. layer -- there's only
-- one matrix in weights & biases each. For simplicity I'm using head to pull out these matrices.
-- all the matrices in the layer below are concatenated to fully connect the net .
-- For layer j:
-- z_j(y_i) = (w_j . y_i) + b_j
-- y_j(z_j) = sigma(z_j(y_i))
-- ∇b_j = δ_j = ((tr w_k) . δ_k) * dy_j/dz_j
-- ∇w_j = δ_j . (tr y_i)

fullyConnectedLayer :: (Int, Int) -> (Int, Int) -> FFLayerOptions -> FFLayerOptions -> IO (FFLayer)
fullyConnectedLayer (m_i,n_i) (m_j,n_j) layerOpts_i layerOpts_j = do
  let depth_i = layerDepth layerOpts_i
      yFnDeriv_i = layerActivationFnDeriv layerOpts_i
      yFn_j = layerActivationFn layerOpts_j

  w <- randn (m_j * n_j) (m_i * n_i * depth_i)
  b <- randn (m_j * n_j) 1

  return (FFLayer {
             weights = [w],
             biases = [b],

             zFn = (\x_j (w_j, b_j) ->
                     [reshape n_j $ flatten
                      $ (head w_j <> (asColumn $ fullyFlatten x_j)) + (head b_j)]),
             yFn = yFn_j,

             nablaW = (\nabla_b_j x_j -> [(flatten $ head nabla_b_j) <.||> (fullyFlatten x_j)]),
             nablaB = (\a _ -> a),
             backprop = (\nabla_b_j z_i w_j ->
                          head $ toBlocksEvery m_i n_i
                          $ (reshape ((cols $ head z_i) * depth_i) $ tr (head w_j) #||> (flatten $ head nabla_b_j))
                          * (head $ yFnDeriv_i [fromBlocks [z_i]]))
             })

-- TODO: build pooling into the conv layer, then delete pooling layer. Don't be an ass

-- None of k conv slices in the same layer interact -- leaving out k subscript for convenience
-- let (d,d) be the size of the kernel window
-- z_i[x,y] = w_i.(x_i[x..x+d,y..y+d]) + b
-- y_i = sigma(z_i)
--
--
-- NOTES:
-- bias, weights are the same for each neuron in a slice
-- weights and bias form a kernel which is applied across input, the output of the kernel forming
--     a new matrix of size (input - d, input - d)

convLayer :: (Int, Int) -> (Int, Int) -> FFLayerOptions -> FFLayerOptions -> IO (FFLayer)
convLayer (m_i,n_i) (m_j,n_j) layerOpts_i layerOpts_j = do
  let depth_i = layerDepth layerOpts_i
      depth_j = layerDepth layerOpts_j
      yFnDeriv_i = layerActivationFnDeriv layerOpts_i
      yFn_j = layerActivationFn layerOpts_j
      pooling_j = pooling layerOpts_j
      poolingBP_j = poolingBP layerOpts_j
      (m_pool,n_pool) = poolSize layerOpts_j

  w <- replicateM depth_j $ randn m_j n_j
  b <- replicateM depth_j $ randn (m_i - m_j + 1) (n_i - n_j + 1)

  return (FFLayer {
             weights = w,
             biases = b,

             zFn = (\x_j (w_j, b_j) ->
                     let zf w_ b_ x_ = (corr2 w_ (fliprl . flipud $ x_)) + b_
                     in pooling_j (m_pool,n_pool)
                        $ zipWith (\w_ b_ -> (sum $ map (zf w_ b_) x_j) / (fromIntegral depth_j)) w_j b_j),
             yFn = yFn_j,

             nablaW = (\nabla_b_j x_j ->
                        let nwf nabla_b_ x_ = corr2 nabla_b_ x_
                        in if depth_i == 1
                           then parMap rdeepseq (\nabla_b_ -> nwf nabla_b_ (head x_j)) nabla_b_j
                           else parZipWith rdeepseq nwf nabla_b_j x_j),
             nablaB = poolingBP_j (m_pool,n_pool),
             backprop = (\nabla_b_j z_i w_j ->
                          zipWith3 (\nabla_b_ zp_ w_ -> (conv2 nabla_b_ w_) * zp_) nabla_b_j (yFnDeriv_i z_i) w_j)
             })


text:
module FFT.DFTControl (
    dcdft,
    nesteddcdft,
    tfdft,
    nestedtfdft
) where
import Data.Complex
import FFT.Samples
import System.Environment
import Control.Parallel
import Strategies
import Criterion.Main
import FFT.Orig

-- twiddle factors
tw :: Int -> Int -> Complex Float
tw n k = cis (-2 * pi * fromIntegral k / fromIntegral n)

dcdft [] = []
dcdft xs = dc split threshold combine worker [0..n']
  where
    n = length xs
    n' = n-1
    worker = map workerInner
        where workerInner k = [sum (map (\j -> xs!!j * tw n (j*k)) [0..n'])]
    combine = concat
    threshold = (\x -> length x < (floor $ sqrt $ fromIntegral n'))
    split l = [front, back]
        where
            front = take p l
            back = drop p l
            p = length l `div` 2

nesteddcdft [] = []
nesteddcdft xs = dc split threshold combine worker [0..n']
  where
    n = length xs
    n' = n-1
    worker = map workerInner
        where workerInner k = [sum (dc split threshold combine (jWorker k) [0..n'])]
                where jWorker k = map (\j -> [xs!!j * tw n (j*k)])
    combine = concat
    threshold = (\x -> length x < (floor $ sqrt $ fromIntegral n'))
    split l = [front, back]
        where
            front = take p l
            back = drop p l
            p = length l `div` 2

tfdft [] = []
tfdft xs = taskfarm worker numWorkers [0..n']
  where
    n = length xs
    n' = n-1
    numWorkers = (floor $ sqrt $ fromIntegral n')
    worker k = [sum [ xs!!j * tw n (j*k) | j <- [0..n']]]

nestedtfdft [] = []
nestedtfdft xs = taskfarm worker numWorkers [0..n']
  where
    n = length xs
    n' = n-1
    numWorkers = (floor $ sqrt $ fromIntegral n')
    worker k = [foldr (+) (0 :+ 0) (concat (taskfarm (innerWorker k) numWorkers [0..n']))]
        where
            innerWorker k j = [xs!!j * tw n (j*k)]

max_stars_repo_name: BNJHope/parallel-fourier-transforms

max_issues_repo_name: BNJHope/parallel-fourier-transforms

max_forks_repo_name: BNJHope/parallel-fourier-transforms


text:
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Day13 where

import Control.Lens
import Control.Lens.TH

import qualified Data.ByteString as BS
import Data.Attoparsec.ByteString (Parser, (<?>))
import qualified Data.Attoparsec.ByteString as DAB
import qualified Data.Attoparsec.ByteString.Char8 as DABC8
import Data.Attoparsec.ByteString.Char8 (decimal, inClass, digit, endOfLine, char, anyChar, letter_ascii, notChar)

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import qualified Data.List as DL
import qualified Data.List.Split as DLS

import Debug.Trace

import Control.Arrow

import Data.Array.IArray (Array)
import qualified Data.Array.IArray as IA

import Data.Set (Set)
import qualified Data.Set as Set
import Data.Either (isRight)
import Data.Maybe (isJust, mapMaybe, catMaybes)
import Data.Bool (bool)

import Data.Vector (Vector, (!))
import qualified Data.Vector as Vec

import Data.Function (on)

import Control.Monad.State.Strict (State)
import qualified Control.Monad.State.Strict as State

import qualified Control.Monad as CM

import qualified Control.Foldl as L

import Data.Complex
import Data.Functor
import Data.Bifunctor (bimap)

-- Each input is the id of a bus line. Nothing means that the line, whatever its
-- id was, is out of service.
type Input = Maybe BusId

newtype BusId = BusId { getBusId :: Integer } deriving Show

-- Parases an unsigned decimal number
parseBusId :: Parser BusId
parseBusId = (BusId <$> decimal) <?> "BusId"

newtype Timestamp = Timestamp { getTimestamp :: Integer } deriving Show

parseTimestamp :: Parser Timestamp
parseTimestamp = (Timestamp <$> decimal) <?> "Timestamp"

-- Solve part 1
--
-- forall n m : int,
--
-- 	start + n = bus_id * m
--
-- n is the waiting time for the m-th departure of `bus_id`. This can be
-- rearranged as:
--
-- 	(1)	n = bus_id * m - start
--
-- `start` can be rewriten in terms of `bus_id`:
--
-- 	start = bus_id * q + r
--
-- where
--
-- 	q = start / bus_id
-- 	r = start % bus_id
--
-- with
--
-- 	r < bus_id
--
-- replacing in (1)
--
-- 	n = bus_id * (m - q) - r
--
-- So the smallest positive n happens when (m - q) == 1 and is thus equal to
--
-- 	n_min = bus_id - r
--
-- Returns Norhing in case there are no bus lines.
solver :: Timestamp -> [Input] -> Maybe Integer
solver start input = do
	-- Find the bus line with the earliest departure time.
	(busId, departureTime) <-
		L.fold ( L.minimumBy (compare `on` snd) ) $
		(id &&& mapper) <$> catMaybes input

	pure (getBusId busId * departureTime)
	where
	-- Compute the earliest departure time after `start` for a bus line.
	mapper :: BusId -> Integer
	mapper bid = getBusId bid - ( getTimestamp start `mod` getBusId bid )

-- Write equations for part 2
--
-- For each bus_id, n_i is the departure time for the m_i-th departure
--
-- n_1 = bus_id_1 * m_1
-- n_2 = bus_id_2 * m_2
-- ...
-- n_i = bus_id_i * m_i
--
-- Those departure times have to be consecutive so:
--
-- n_2 = n_1 + (ord_2 - ord_1)
-- ...
-- n_(i+1) = n_i + (ord_(i+1) - ord_i)
--
-- where ord_i is the order of the bus_id on the list. With ord_(i+1) > ord_i.
--
-- Applying this constrain we get:
--
-- n_1 = bus_id_1 * m_1
-- n_1 + (ord_2 - ord_1) = bus_id_2 * m_2
-- n_2 + (ord_3 - ord_2) = bus_id_3 * m_3
-- ...
-- n_i + (ord_(i+1) - ord_i) = bus_id_(i+1) * m_(i+1)
--
-- Leaving only n_1 we get:
--
-- n_1 = bus_id_1 * m_1
-- n_1 + (ord_2 - ord_1) = bus_id_2 * m_2
-- (n_1 + (ord_2 - ord_1)) + (ord_3 - ord_2) = bus_id_3 * m_3
-- ...
--
-- The last equation becomes
--
-- n_1 + (ord_3 - ord_1) = bus_id_3 * m_3
--
-- So all the equations become
--
-- n_1 = bus_id_1 * m_1
-- n_1 + (ord_2 - ord_1) = bus_id_2 * m_2
-- n_1 + (ord_3 - ord_1) = bus_id_3 * m_3
-- ...
-- n_1 + (ord_i - ord_1) = bus_id_i * m_i
--
-- Next I rewrite them like:
--
-- n_1 - bus_id_1 * m_1 = 0
-- n_1 + (ord_2 - ord_1) - bus_id_2 * m_2 = 0
-- ...
-- n_1 + (ord_i - ord_1) - bus_id_i * m_i = 0
--
-- Finally I replace n_1 in all the equations
--
-- bus_id_1 * m_1 + (ord_i - ord_1) - bus_id_i * m_i = 0
--
-- This equations can be stored using a 4-tuple
--
--	(bus_id_1, bus_id_i, (ord_i - ord_1), ord_i)
--
-- This functions just creates those 4-tuples to be further processed.
solver' input = case catMaybes (zipWith (\x -> fmap (x,)) [1..] input) of
	( (ord1, busId1) : rest) ->
		(busId1, busId1, 0, ord1) : fmap mapper rest
		where
		mapper (ordi, busIdi) = (busId1, busIdi, ordi - ord1, ordi)

-- Sample input
--
-- 939
-- 7,13,x,x,59,x,31,19
--
-- The x's represent lines that are out of service.
parseInput :: Parser (Timestamp, [Input])
parseInput =
	( (,) <$>
		(parseTimestamp <* endOfLine) <*>
		(DAB.choice
			[ Just <$> parseBusId
			, char 'x' $> Nothing ] `DAB.sepBy1'` char ',') ) <* endOfLine

runSolution :: FilePath -> IO ()
runSolution filePath = do
	putStrLn "**Day 13**"
	contents <- BS.readFile filePath
	let parseResult = DAB.parseOnly parseInput contents
	case parseResult of
		Left err -> putStrLn $ "Error :" ++ err
		Right (start, input) -> do

			putStrLn "Part 1"
			print $ solver start input

			putStrLn "\nPart 2"
			mapM_ (putStrLn . printEq) $ solver' input
	where
	-- Prints an equation to be solved.
	printEq (flid, bid, delta, ordi) =
		show (getBusId flid) ++ " * x[1] - " ++
		show (getBusId bid) ++ " * x[" ++ show ordi ++ "] + " ++
		show delta

-- Wrong
-- 1258

-- Part 1:
-- Answer = Just 410