text:
{-# LANGUAGE CPP
           , GADTs
           , Rank2Types
           , DataKinds
           , TypeFamilies
           , FlexibleContexts
           , UndecidableInstances
           , LambdaCase
           , MultiParamTypeClasses
           , OverloadedStrings
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs -fsimpl-tick-factor=1000 -fno-warn-orphans #-}
module Language.Hakaru.Runtime.LogFloatPrelude where

#if __GLASGOW_HASKELL__ < 710
import           Data.Functor                    ((<$>))
import           Control.Applicative             (Applicative(..))
#endif
import           Data.Foldable                   as F
import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWCD
import           Data.Number.Natural
import           Data.Number.LogFloat            hiding (sum, product)
import qualified Data.Number.LogFloat            as LF
import           Data.STRef
import qualified Data.Vector                     as V
import qualified Data.Vector.Unboxed             as U
import qualified Data.Vector.Generic             as G
import qualified Data.Vector.Generic.Mutable     as M
import           Control.Monad
import           Control.Monad.ST
import           Numeric.SpecFunctions           (logBeta)
import           Prelude                         hiding (init, sum, product, exp, log, (**), pi)
import qualified Prelude                         as P
import           Language.Hakaru.Runtime.CmdLine (Parseable(..), Measure(..), makeMeasure)

-- This Read instance really should be the logfloat package
instance Read LogFloat where
    readsPrec p s = [(logFloat x, r) | (x, r) <- readsPrec p s]

instance Parseable LogFloat where
  parse = return . read

type family MinBoxVec (v1 :: * -> *) (v2 :: * -> *) :: * -> *
type instance MinBoxVec V.Vector v        = V.Vector
type instance MinBoxVec v        V.Vector = V.Vector
type instance MinBoxVec U.Vector U.Vector = U.Vector

type family MayBoxVec a :: * -> *
type instance MayBoxVec ()           = U.Vector
type instance MayBoxVec Int          = U.Vector
type instance MayBoxVec Double       = U.Vector
type instance MayBoxVec LogFloat     = U.Vector
type instance MayBoxVec Bool         = U.Vector
type instance MayBoxVec (U.Vector a) = V.Vector
type instance MayBoxVec (V.Vector a) = V.Vector
type instance MayBoxVec (a,b)        = MinBoxVec (MayBoxVec a) (MayBoxVec b)

newtype instance U.MVector s LogFloat = MV_LogFloat (U.MVector s Double)
newtype instance U.Vector    LogFloat = V_LogFloat  (U.Vector    Double)

instance U.Unbox LogFloat

instance M.MVector U.MVector LogFloat where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
#if __GLASGOW_HASKELL__ > 710
  {-# INLINE basicInitialize #-}
#endif
  {-# INLINE basicUnsafeReplicate #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  {-# INLINE basicClear #-}
  {-# INLINE basicSet #-}
  {-# INLINE basicUnsafeCopy #-}
  {-# INLINE basicUnsafeGrow #-}
  basicLength (MV_LogFloat v) = M.basicLength v
  basicUnsafeSlice i n (MV_LogFloat v) = MV_LogFloat $ M.basicUnsafeSlice i n v
  basicOverlaps (MV_LogFloat v1) (MV_LogFloat v2) = M.basicOverlaps v1 v2
  basicUnsafeNew n = MV_LogFloat `liftM` M.basicUnsafeNew n
#if __GLASGOW_HASKELL__ > 710
  basicInitialize (MV_LogFloat v) = M.basicInitialize v
#endif
  basicUnsafeReplicate n x = MV_LogFloat `liftM` M.basicUnsafeReplicate n (logFromLogFloat x)
  basicUnsafeRead (MV_LogFloat v) i = logToLogFloat `liftM` M.basicUnsafeRead v i
  basicUnsafeWrite (MV_LogFloat v) i x = M.basicUnsafeWrite v i (logFromLogFloat x)
  basicClear (MV_LogFloat v) = M.basicClear v
  basicSet (MV_LogFloat v) x = M.basicSet v (logFromLogFloat x)
  basicUnsafeCopy (MV_LogFloat v1) (MV_LogFloat v2) = M.basicUnsafeCopy v1 v2
  basicUnsafeMove (MV_LogFloat v1) (MV_LogFloat v2) = M.basicUnsafeMove v1 v2
  basicUnsafeGrow (MV_LogFloat v) n = MV_LogFloat `liftM` M.basicUnsafeGrow v n

instance G.Vector U.Vector LogFloat where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw #-}
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicUnsafeIndexM #-}
  {-# INLINE elemseq #-}
  basicUnsafeFreeze (MV_LogFloat v) = V_LogFloat `liftM` G.basicUnsafeFreeze v
  basicUnsafeThaw (V_LogFloat v) = MV_LogFloat `liftM` G.basicUnsafeThaw v
  basicLength (V_LogFloat v) = G.basicLength v
  basicUnsafeSlice i n (V_LogFloat v) = V_LogFloat $ G.basicUnsafeSlice i n v
  basicUnsafeIndexM (V_LogFloat v) i
                = logToLogFloat `liftM` G.basicUnsafeIndexM v i
  basicUnsafeCopy (MV_LogFloat mv) (V_LogFloat v)
                = G.basicUnsafeCopy mv v
  elemseq _ x z = G.elemseq (undefined :: U.Vector a) (logFromLogFloat x) z

type Prob = LogFloat

lam :: (a -> b) -> a -> b
lam = id
{-# INLINE lam #-}

app :: (a -> b) -> a -> b
app f x = f x
{-# INLINE app #-}

let_ :: a -> (a -> b) -> b
let_ x f = let x1 = x in f x1
{-# INLINE let_ #-}

ann_ :: a -> b -> b
ann_ _ a = a
{-# INLINE ann_ #-}

exp :: Double -> Prob
exp = logToLogFloat
{-# INLINE exp #-}

log :: Prob -> Double
log = logFromLogFloat
{-# INLINE log #-}

betaFunc :: Prob -> Prob -> Prob
betaFunc a b = exp (logBeta (fromProb a) (fromProb b))

uniform :: Double -> Double -> Measure Double
uniform lo hi = makeMeasure $ MWC.uniformR (lo, hi)
{-# INLINE uniform #-}

normal :: Double -> Prob -> Measure Double
normal mu sd = makeMeasure $ MWCD.normal mu (fromProb sd)
{-# INLINE normal #-}

beta :: Prob -> Prob -> Measure Prob
beta a b = makeMeasure $ \g ->
  unsafeProb <$> MWCD.beta (fromProb a) (fromProb b) g
{-# INLINE beta #-}

gamma :: Prob -> Prob -> Measure Prob
gamma a b = makeMeasure $ \g ->
  unsafeProb <$> MWCD.gamma (fromProb a) (fromProb b) g
{-# INLINE gamma #-}

categorical :: MayBoxVec Prob Prob -> Measure Int
categorical a = makeMeasure $ MWCD.categorical (U.map prep a)
  where prep p = fromLogFloat (p / m)
        m      = G.maximum a
{-# INLINE categorical #-}

plate :: (G.Vector (MayBoxVec a) a) =>
         Int -> (Int -> Measure a) -> Measure (MayBoxVec a a)
plate n f = G.generateM (fromIntegral n) $ \x ->
             f (fromIntegral x)
{-# INLINE plate #-}

bucket :: Int -> Int -> (forall s. Reducer () s a) -> a
bucket b e r = runST
             $ case r of Reducer{init=initR,accum=accumR,done=doneR} -> do
                          s' <- initR ()
                          F.mapM_ (\i -> accumR () i s') [b .. e - 1]
                          doneR s'
{-# INLINE bucket #-}

data Reducer xs s a =
    forall cell.
    Reducer { init  :: xs -> ST s cell
            , accum :: xs -> Int -> cell -> ST s ()
            , done  :: cell -> ST s a
            }

r_fanout :: Reducer xs s a
         -> Reducer xs s b
         -> Reducer xs s (a,b)
r_fanout Reducer{init=initA,accum=accumA,done=doneA}
         Reducer{init=initB,accum=accumB,done=doneB} = Reducer
   { init  = \xs       -> liftM2 (,) (initA xs) (initB xs)
   , accum = \bs i (s1, s2) ->
             accumA bs i s1 >> accumB bs i s2
   , done  = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2)
   }
{-# INLINE r_fanout #-}

r_index :: (G.Vector (MayBoxVec a) a)
        => (xs -> Int)
        -> ((Int, xs) -> Int)
        -> Reducer (Int, xs) s a
        -> Reducer xs s (MayBoxVec a a)
r_index n f Reducer{init=initR,accum=accumR,done=doneR} = Reducer
   { init  = \xs -> V.generateM (n xs) (\b -> initR (b, xs))
   , accum = \bs i v ->
             let ov = f (i, bs) in
             accumR (ov,bs) i (v V.! ov)
   , done  = \v -> fmap G.convert (V.mapM doneR v)
   }
{-# INLINE r_index #-}

r_split :: ((Int, xs) -> Bool)
        -> Reducer xs s a
        -> Reducer xs s b
        -> Reducer xs s (a,b)
r_split b Reducer{init=initA,accum=accumA,done=doneA}
          Reducer{init=initB,accum=accumB,done=doneB} = Reducer
   { init  = \xs -> liftM2 (,) (initA xs) (initB xs)
   , accum = \bs i (s1, s2) ->
             if (b (i,bs)) then accumA bs i s1 else accumB bs i s2
   , done  = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2)
   }
{-# INLINE r_split #-}

r_add :: Num a => ((Int, xs) -> a) -> Reducer xs s a
r_add e = Reducer
   { init  = \_ -> newSTRef 0
   , accum = \bs i s ->
             modifySTRef' s (+ (e (i,bs)))
   , done  = readSTRef
   }
{-# INLINE r_add #-}

r_nop :: Reducer xs s ()
r_nop = Reducer
   { init  = \_ -> return ()
   , accum = \_ _ _ -> return ()
   , done  = \_ -> return ()
   }
{-# INLINE r_nop #-}

pair :: a -> b -> (a, b)
pair = (,)
{-# INLINE pair #-}

true, false :: Bool
true  = True
false = False

nothing :: Maybe a
nothing = Nothing

just :: a -> Maybe a
just = Just

left :: a -> Either a b
left = Left

right :: b -> Either a b
right = Right

unit :: ()
unit = ()

data Pattern = PVar | PWild
newtype Branch a b =
    Branch { extract :: a -> Maybe b }

ptrue, pfalse :: a -> Branch Bool a
ptrue  b = Branch { extract = extractBool True  b }
pfalse b = Branch { extract = extractBool False b }
{-# INLINE ptrue  #-}
{-# INLINE pfalse #-}

extractBool :: Bool -> a -> Bool -> Maybe a
extractBool b a p | p == b     = Just a
                  | otherwise  = Nothing
{-# INLINE extractBool #-}

pnothing :: b -> Branch (Maybe a) b
pnothing b = Branch { extract = \ma -> case ma of
                                         Nothing -> Just b
                                         Just _  -> Nothing }

pjust :: Pattern -> (a -> b) -> Branch (Maybe a) b
pjust PVar c = Branch { extract = \ma -> case ma of
                                           Nothing -> Nothing
                                           Just x  -> Just (c x) }
pjust _ _ = error "TODO: Runtime.Prelude{pjust}"

pleft :: Pattern -> (a -> c) -> Branch (Either a b) c
pleft PVar f = Branch { extract = \ma -> case ma of
                                           Right _ -> Nothing
                                           Left x -> Just (f x) }
pleft _ _ = error "TODO: Runtime.Prelude{pLeft}"

pright :: Pattern -> (b -> c) -> Branch (Either a b) c
pright PVar f = Branch { extract = \ma -> case ma of
                                            Left _ -> Nothing
                                            Right x -> Just (f x) }
pright _ _ = error "TODO: Runtime.Prelude{pRight}"

ppair :: Pattern -> Pattern -> (x -> y -> b) -> Branch (x,y) b
ppair PVar  PVar c = Branch { extract = (\(x,y) -> Just (c x y)) }
ppair _     _    _ = error "ppair: TODO"

uncase_ :: Maybe a -> a
uncase_ (Just a) = a
uncase_ Nothing  = error "case_: unable to match any branches"
{-# INLINE uncase_ #-}

case_ :: a -> [Branch a b] -> b
case_ e [c1]     = uncase_ (extract c1 e)
case_ e [c1, c2] = uncase_ (extract c1 e `mplus` extract c2 e)
case_ e bs_      = go bs_
  where go []     = error "case_: unable to match any branches"
        go (b:bs) = case extract b e of
                      Just b' -> b'
                      Nothing -> go bs
{-# INLINE case_ #-}

branch :: (c -> Branch a b) -> c -> Branch a b
branch pat body = pat body
{-# INLINE branch #-}

dirac :: a -> Measure a
dirac = return
{-# INLINE dirac #-}

pose :: Prob -> Measure a -> Measure a
pose _ a = a
{-# INLINE pose #-}

superpose :: [(Prob, Measure a)]
          -> Measure a
superpose pms = do
  i <- categorical (G.fromList $ map fst pms)
  snd (pms !! i)
{-# INLINE superpose #-}

reject :: Measure a
reject = Measure $ \_ -> return Nothing

nat_ :: Int -> Int
nat_ = id

int_ :: Int -> Int
int_ = id

unsafeNat :: Int -> Int
unsafeNat = id

nat2prob :: Int -> Prob
nat2prob = fromIntegral

fromInt  :: Int -> Double
fromInt  = fromIntegral

nat2int  :: Int -> Int
nat2int  = id

nat2real :: Int -> Double
nat2real = fromIntegral

fromProb :: Prob -> Double
fromProb = fromLogFloat

unsafeProb :: Double -> Prob
unsafeProb = logFloat

real_ :: Rational -> Double
real_ = fromRational

prob_ :: NonNegativeRational -> Prob
prob_ = fromRational . fromNonNegativeRational

infinity :: Double
infinity = 1/0

abs_ :: Num a => a -> a
abs_ = abs

(**) :: Prob -> Double -> Prob
(**) = pow
{-# INLINE (**) #-}

pi :: Prob
pi = unsafeProb P.pi
{-# INLINE pi #-}

thRootOf :: Int -> Prob -> Prob
thRootOf a b = b ** (recip $ fromIntegral a)
{-# INLINE thRootOf #-}

array
    :: (G.Vector (MayBoxVec a) a)
    => Int
    -> (Int -> a)
    -> MayBoxVec a a
array n f = G.generate (fromIntegral n) (f . fromIntegral)
{-# INLINE array #-}

arrayLit :: (G.Vector (MayBoxVec a) a) => [a] -> MayBoxVec a a
arrayLit = G.fromList
{-# INLINE arrayLit #-}

(!) :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int -> a
a ! b = a G.! (fromIntegral b)
{-# INLINE (!) #-}

size :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int
size v = fromIntegral (G.length v)
{-# INLINE size #-}

reduce
    :: (G.Vector (MayBoxVec a) a)
    => (a -> a -> a)
    -> a
    -> MayBoxVec a a
    -> a
reduce f n v = G.foldr f n v
{-# INLINE reduce #-}

class Num a => Num' a where
    product :: Int -> Int -> (Int -> a) -> a
    product a b f = F.foldl' (\x y -> x * f y) 1 [a .. b-1]
    {-# INLINE product #-}
    summate :: Int -> Int -> (Int -> a) -> a
    summate a b f = F.foldl' (\x y -> x + f y) 0 [a .. b-1]
    {-# INLINE summate #-}

instance Num' Int
instance Num' Double
instance Num' LogFloat where
    product a b f = LF.product (map f [a .. b-1])
    {-# INLINE product #-}
    summate a b f = LF.sum     (map f [a .. b-1])
    {-# INLINE summate #-}

run :: Show a
    => MWC.GenIO
    -> Measure a
    -> IO ()
run g k = unMeasure k g >>= \case
           Just a  -> print a
           Nothing -> return ()

iterateM_
    :: Monad m
    => (a -> m a)
    -> a
    -> m b
iterateM_ f = g
    where g x = f x >>= g

withPrint :: Show a => (a -> IO b) -> a -> IO b
withPrint f x = print x >> f x

max_stars_repo_name: vmchale/hakaru

max_issues_repo_name: vmchale/hakaru

max_forks_repo_name: vmchale/hakaru


text:
{-# LANGUAGE Strict
           , OverloadedStrings#-}
module Lib where

import qualified Data.HashTable.IO as H
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as VS
import System.IO
import Data.List
import Control.Monad
import System.Process
import System.Random
import Text.Read (readMaybe)

someFunc :: IO ()
someFunc = putStrLn "someFunc"

nanCheck :: Double -> Double
nanCheck x = if show x == "NaN" || show x == "Infinity" then 0 else x

st :: Double
st = 10

delta :: Double
delta = 1

chars :: T.Text
chars = "qwertyuiopasdfghjklzxcvbnmęóąśłżźćń\t1234567890 "

el :: Char -> T.Text -> Bool
el = T.any . (==)

preprocess :: T.Text -> [T.Text]
preprocess = map (T.take 20) . T.words . T.map check . T.toLower
             where check x = if x `el` chars then x else ' '

type Probs = H.CuckooHashTable T.Text Double

readWords :: Handle -> IO Probs
readWords handl = do
  dict <- H.newSized 2500000
  loop handl dict
  randomIO >>= H.insert dict "<bias>"
  return dict
  where
    loop :: Handle -> Probs -> IO ()
    loop hdl dic =
        hIsEOF hdl >>= \x -> if x
            then return ()
            else do
              line <- preprocess <$> TIO.hGetLine hdl
              let ins word =
                    if (T.length word > 4)
                    then H.lookup dic word >>=
                                    \mprob -> case mprob of
                                       Just _ -> do
                                         return ()
                                       Nothing -> randomIO >>= H.insert dic word
                    else return ()
              case line of
                _:dat -> mapM_ ins dat
                _     -> hPutStrLn stderr "Can't parse line:" >> hPrint stderr line
              loop hdl dic

saveProbs :: Probs -> IO ()
saveProbs dic = do
  H.mapM_ (\(word, v) -> putStrLn . unwords $ [T.unpack . T.strip $ word, show v]) dic

readProbs :: String -> IO Probs
readProbs filename = do
  dict <- H.newSized 2500000
  withFile filename ReadMode $ readFromFile dict
  return dict
  where
    readFromFile :: Probs -> Handle -> IO ()
    readFromFile dic hdl =
      hIsEOF hdl >>= \x -> if x
                              then return ()
                              else do
                                w <- hGetLine hdl
                                case words w of
                                  [p, pr] ->
                                    H.insert dic (T.pack p) (read pr)
                                  _ -> return ()
                                readFromFile dic hdl

linear :: Vector Double -> Vector Double -> Double
linear = (<.>)

rmse :: Vector Double -> Vector Double -> Double -> Double
rmse weights inp ex = let
  prob = linear weights inp
  in (prob - ex)^2

gradRMSE :: Vector Double -> Vector Double -> Double -> Double -> Vector Double
gradRMSE weights inp ex m = VS.map (\x -> x * (linear weights inp - ex)/m) inp

counts :: Eq a => a -> [a] -> Double
counts = ((fromIntegral. length).).filter.(==)

evalLine :: Probs -> [T.Text] -> IO Double
evalLine dict line = do
  let wek@(_:words) = "<bias>" : (filter ((>4). T.length) . nub $ line)
      inp   = VS.cons 1 (foldr (\c acc -> VS.cons (log . (+1) . counts c $ line) acc) VS.empty words)
  weights <- foldM (\vec c -> do
                       mval <-  H.lookup dict c
                       case mval of
                         Just v -> return $ VS.cons v vec
                         Nothing-> return $ VS.cons 0 vec) VS.empty wek
  return ((1900+).(100*).linear weights $ inp)

updateWeights :: Probs -> [T.Text] -> Double -> Double -> Double -> IO Double
updateWeights dict line ex step m = do
  let wek@(_:words) = "<bias>":(filter ((>4).T.length) . nub $ line)
      inp   = VS.cons 1 (foldr (\c acc -> VS.cons (log . (+1) . counts c $ line) acc) VS.empty words)
  weights <- foldM (\vec c -> do
                       mval <- H.lookup dict c
                       case mval of
                         Just v -> return $ VS.snoc vec v
                         Nothing-> return $ VS.snoc vec 0) VS.empty wek
  let grad       = gradRMSE weights inp ex m
      newWeights = weights - (VS.map (clip.nanCheck.(step *)) grad)
      newError   = rmse newWeights inp ex
  foldM_ (\acc x -> do
            let newW = VS.head acc
            unless (newW == 0) $ H.insert dict x newW
            return $ VS.tail acc) newWeights words
  return $ nanCheck newError
  where
    clip x = max x 3

checkError :: Probs -> [T.Text] -> Double -> IO Double
checkError dict line ex = do
  let wek@(_:words) = "<bias>":nub line
      inp   = VS.cons 1 . foldr (\c acc -> VS.cons (log . (+1) . counts c $ line) acc) VS.empty $ words
  weights <- foldM (\vec c -> do
                       mval <- H.lookup dict c
                       case mval of
                         Just v -> return $ VS.snoc vec v
                         Nothing-> return $ VS.snoc vec 0) VS.empty wek
  return $ rmse weights inp ex

train :: [(Double,[T.Text])] -> Probs -> Int -> IO ()
train corpus dict maxEp = do
  let len     = length corpus
      testset = take (len`div`10) corpus
      trainset= drop (len`div`10) corpus
      trlen   = length trainset
      testlen = length testset
      trainHelp :: Double -> Int -> Double -> IO ()
      trainHelp err epoch step = do
          let trainOnLine (year, dat) = updateWeights dict dat year step (fromIntegral trlen)
          mapM_ trainOnLine trainset
          newErr <- sqrt <$> foldM (\er (ex, line) -> (+(er/(fromIntegral testlen))) <$> checkError dict line ex) 0 testset
          let
            lossRate   = (err - newErr)
            newSteptmp = if lossRate > 0 then step*2 else step/100
            newStep    = if newSteptmp > 20 then 20 else newSteptmp
          hPutStrLn stderr $ "***********************"
          hPutStrLn stderr $ "  Old Error: " ++ show err
          hPutStrLn stderr $ "  New Error: " ++ show newErr
          hPutStrLn stderr $ "  LossRate: " ++ show lossRate
          unless ((abs lossRate <= delta)||(epoch == maxEp)) $ trainHelp newErr (1+epoch) newStep
  startErr <- sqrt  <$> foldM (\er (ex, line) -> (+(er/fromIntegral testlen)) <$> checkError dict line ex) 0 testset
  trainHelp startErr 1 st

loadCorpus :: String -> IO [(Double,[T.Text])]
loadCorpus filename = do
  withFile filename ReadMode (loop [])
  where loop corp hdl =
          hIsEOF hdl >>= \x -> if x
            then return corp
            else do
              line <- preprocess <$> TIO.hGetLine hdl
              newEntry <- case line of
                cl:dat -> case readMaybe . T.unpack $ cl of
                  Just year -> return $ Just ((year-1900)/100, dat)
                  Nothing   -> hPutStrLn stderr "Can't read as a number:" >> hPrint stderr cl >> return Nothing
                _      -> hPutStrLn stderr "Can't parse line:" >> hPrint stderr line >> return Nothing
              case newEntry of
                Nothing -> loop corp hdl
                Just x  -> loop (x:corp) hdl

eval :: Handle -> Probs -> IO ()
eval hdl dict =
  hIsEOF hdl >>= \x ->
    if x
    then return ()
    else do
      line <- preprocess <$> TIO.hGetLine hdl
      pred <- evalLine dict line
      print pred
      eval hdl dict

ext: hs

max_stars_repo_name: Antystenes/LinearRegression

max_issues_repo_name: Antystenes/LinearRegression

max_forks_repo_name: Antystenes/LinearRegression


text:
{-# LANGUAGE OverloadedLists #-}

module OBJReader where

import qualified Data.HashMap as HM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as V
import           Data.Vector.Storable ((!))
import           Control.Arrow
import           Control.Lens
import           Data.List (nub, nubBy)
import qualified Numeric.LinearAlgebra as L
import qualified Graphics.Rendering.OpenGL.GL as GL
import           Data.Function (on)
import Debug.Trace

import Physics.Collision

splitOn x =
  foldr (\c -> if c == x then ([]:) else over (ix 0) (c:)) [[]]

addP (a,b) (c,d) = (a+b,c+d)

getFace :: [String] -> [(Int, Int)]
getFace =
  map ((\[v,_,n] -> (read v - 1,read n - 1)) . splitOn '/') . tail

readVectors :: String -> [[String]] -> [VS.Vector Float]
readVectors vtype =
  map (VS.fromList.map read.tail) . filter ((==vtype) . head)

readPoints = readVectors "v"
readNormals = readVectors "vn"

readBoxFromObj :: FilePath -> IO Primitive
readBoxFromObj fn = do
  text <- map words . lines <$> readFile fn
  let p = V.fromList . readPoints $ text
      q = V.fromList . readNormals $  text
      f = map (V.unsafeIndex p *** V.unsafeIndex q >>> uncurry (flip Squad)) .
          nubBy ((==) `on` snd) .
          map (head . getFace) .
          filter ((=="f").head) $ text
  return $ Box p (V.fromList f)

readOBJ :: FilePath -> IO (VS.Vector Float, VS.Vector GL.BaseInstance,[VS.Vector Float])
readOBJ fn = do
  text <- map words . lines <$> readFile fn
  let
    verts    = readPoints text
    normals :: [VS.Vector Float]
    normals  = readNormals text
    faces    = map getFace . filter ((=="f").head) $ text
    vertData = nub . concat $ faces
    vertIxMap= HM.fromList . flip zip [0..] $ vertData
    --                 u v color
    addColor = (VS.++ [1,1,0.5,0.5,0.5])
    computeT :: [(Int,Int)] -> VS.Vector Float
    computeT [v0, v1, v2] =
      let dpos1 = verts !! fst v1 - verts !! fst v0
          -- dpos2 = verts !! fst v2 - verts !! fst v0
          tang  = dpos1
          -- bitan = dpos2
      in tang
    nullTans = HM.fromList . flip zip (repeat $ VS.replicate 3 0) $ vertData
    insertT l h = let p = computeT l
                  in foldr (\k -> HM.insertWith (+) k p) h l
    tangents = foldr insertT nullTans faces
    normalize :: (Int, Int) -> VS.Vector Float -> VS.Vector Float
    normalize (_,k) t =
      let n                 = normals !! k
          normalizedTangent = normalizeVec $ t - L.scale (L.dot n t) n
          normalizedBitangent = L.cross normalizedTangent n
      in -- n VS.++ t
        normalizedTangent VS.++ normalizedBitangent
    normalizedTangents = HM.fromList . map (fst &&& uncurry normalize) . HM.toList $ tangents
    faceProc p = (verts!!) *** (normals!!) >>> uncurry (VS.++) >>> addColor >>> flip (VS.++) (normalizedTangents HM.! p) $ p
    vBuffData = VS.concat $ map faceProc vertData
    vertices  = map ((verts!!).fst) vertData
    indices  = VS.fromList . map (vertIxMap HM.!) . concat $ faces
--      result   = VS.concat $ map faceProc faces
  return (vBuffData, indices, vertices)

normalizeVec :: VS.Vector Float -> VS.Vector Float
normalizeVec v = L.scale (realToFrac . (1/) $ L.norm_2 v) v

-- compute

max_stars_repo_name: Antystenes/CPG

max_issues_repo_name: Antystenes/CPG

max_forks_repo_name: Antystenes/CPG


text:
module TestHelpers where

import Control.Newtype.Generics (over)
import Data.Complex
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.Data.Particle
import qualified Numeric.LinearAlgebra as LA
import qualified Test.QuickCheck as QC
import Test.QuickCheck (Arbitrary(..), Gen, Positive(..))

getRank :: DensityMatrix -> Rank
getRank (DensityMatrix dm) =
  let (_, s, _) = LA.compactSVD dm
   in LA.size s

arbDimRank :: QC.Gen (Dim, Rank)
arbDimRank = do
  dim <- QC.suchThat arbitrary (> 1)
  Positive rank <- QC.suchThat arbitrary (\(Positive v) -> v <= dim)
  return (dim, rank)

arbDM :: Dim -> Rank -> QC.Gen DensityMatrix
arbDM dim rank = do
  let genXs =
        QC.vectorOf
          dim
          (QC.vectorOf rank (QC.arbitrary :: QC.Gen (Complex Double)))
  xs <- QC.suchThat genXs (\ll -> any (/= 0.0 :+ 0.0) (zipWith (!!) ll [0 ..]))
  let a = LA.fromLists xs
      h = a LA.<> LA.tr a
      hTr = LA.sumElements $ LA.takeDiag h
      dm = h / LA.scalar hTr
  return $ DensityMatrix dm

arbWDM :: Dim -> Rank -> QC.Gen WeighedDensityMatrix
arbWDM dim rank = do
  Positive w <- QC.arbitrary
  dm <- arbDM dim rank
  return $ WeighedDensityMatrix (w, dm)

arbNormedVecDim :: Int -> Gen (LA.Matrix (Complex Double))
arbNormedVecDim dim = do
  let genXs = QC.vectorOf dim (arbitrary :: Gen (Complex Double))
  xs <- QC.suchThat genXs (any (/= 0))
  let sv = LA.asColumn . LA.fromList $ xs
  return $ sv / LA.toComplex (LA.scalar (LA.norm_2 sv), 0)

arbitraryWeighed :: Arbitrary x => Gen (Weight, x)
arbitraryWeighed = do
  dim <- QC.suchThat QC.getSize (> 0)
  w <- getPositive <$> arbitrary
  dm <- QC.resize dim arbitrary
  return (w, dm)

arbParticles :: Dim -> Rank -> NumberOfParticles -> QC.Gen Particles
arbParticles dim rank num = do
  vdms <- V.replicateM num (arbWDM dim rank)
  let wr = V.foldl' (\acc (WeighedDensityMatrix (wi, _)) -> acc + wi) 0 vdms
      vdmsn = V.map (over WeighedDensityMatrix (\(w, dm) -> (w / wr, dm))) vdms
  QC.Positive w <- QC.arbitrary
  return $ Particles rank w num vdmsn

instance Arbitrary DensityMatrix where
  arbitrary = do
    dim <- QC.suchThat QC.getSize (> 0)
    arbDM dim dim

instance Arbitrary PureStateVector where
  arbitrary = do
    dim <- QC.suchThat QC.getSize (> 0)
    sv <- arbNormedVecDim dim
    return . PureStateVector $ sv

instance Arbitrary WeighedDensityMatrix where
  arbitrary = do
    (w, dm) <- arbitraryWeighed
    return . WeighedDensityMatrix $ (w, dm)

instance Arbitrary WeighedPureStateVector where
  arbitrary = do
    (w, sv) <- arbitraryWeighed
    return . WeighedPureStateVector $ (w, sv)

instance QC.Arbitrary Particles where
  arbitrary = do
    (dim, rank) <- arbDimRank
    num <- QC.suchThat QC.arbitrary (> 100)
    arbParticles dim rank num

class MEq a where
  infix 4 <==>
  (<==>) :: a -> a -> Bool

instance MEq DensityMatrix where
  DensityMatrix dm1 <==> DensityMatrix dm2 = LA.norm_Frob (dm1 - dm2) < 1e-6

instance MEq WeighedDensityMatrix where
  WeighedDensityMatrix (w1, DensityMatrix dm1) <==> WeighedDensityMatrix (w2, DensityMatrix dm2) =
    LA.norm_Frob (LA.scalar (w1 :+ 0) * dm1 - LA.scalar (w2 :+ 0) * dm2) < 1e-6

instance MEq PureStateVector where
  sv1 <==> sv2 = 1 - fidelity sv1 sv2 < 1e-6

instance MEq WeighedPureStateVector where
  WeighedPureStateVector (w1, sv1) <==> WeighedPureStateVector (w2, sv2) =
    abs (w1 - w2) <= 1e-6 * 0.5 * (w1 + w2) && sv1 <==> sv2

max_stars_repo_name: Belinsky-L-V/HABQT

max_issues_repo_name: Belinsky-L-V/HABQT

max_forks_repo_name: Belinsky-L-V/HABQT


text:
{-# LANGUAGE BangPatterns #-}
module STC.DFTArray where

import           Array.UnboxedArray         as AU
import           Control.Monad              as M
import           Data.Array.Repa            as R
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           DFT.Plan
import           FokkerPlanck.FourierSeries
import           Image.IO
import           STC.Utils
import           System.FilePath
import           Text.Printf
import           Utils.Parallel

data DFTArray =
  DFTArray !Int -- rows
           !Int -- cols
           ![Double] -- thetaFreqs
           ![Double] -- rFreqs
           ![VS.Vector (Complex Double)]

instance Show DFTArray where
  show (DFTArray rows cols thetaFreqs rFreqs vecs) = printf "DFTArray rows = %d cols = %d\nthetaFreqs: %s\nrFreqs: %s\n%s\n%d\n" rows cols (show thetaFreqs) (show rFreqs) (show . VS.length . L.head $ vecs) (L.length vecs) --(show . L.take 10 . VS.toList . L.head $ vecs)

{-# INLINE getDFTArrayVector #-}
getDFTArrayVector :: DFTArray -> [VS.Vector (Complex Double)]
getDFTArrayVector (DFTArray _ _ _ _ vecs) = vecs

{-# INLINE dftArrayToRepa #-}
dftArrayToRepa :: DFTArray -> R.Array U DIM4 (Complex Double)
dftArrayToRepa (DFTArray rows cols thetaFreqs rFreqs vecs) =
  fromUnboxed (Z :. (L.length rFreqs) :. (L.length thetaFreqs) :. cols :. rows) .
  VS.convert . VS.concat $vecs

{-# INLINE dftArrayToRepa1 #-}
dftArrayToRepa1 :: Int -> Int -> DFTArray -> R.Array U DIM4 (Complex Double)
dftArrayToRepa1 numThetaFreq numRFreq (DFTArray rows cols thetaFreqs rFreqs vecs) =
  fromUnboxed (Z :. numRFreq :. numThetaFreq :. cols :. rows) .
  VS.convert . VS.concat $vecs

{-# INLINE repaToDFTArray #-}
repaToDFTArray ::
     [Double] -> [Double] -> R.Array U DIM4 (Complex Double) -> DFTArray
repaToDFTArray thetaFreqs rFreqs arr =
  let (Z :. (!numRFreq) :. (!numThetaFreq) :. cols :. rows) = extent arr
  in DFTArray rows cols thetaFreqs rFreqs .
     L.map
       (\(!rf, !tf) ->
          VS.convert . toUnboxed . computeS . R.slice arr $
          (Z :. rf :. tf :. All :. All)) $
     (,) <$> [0 .. numRFreq - 1] <*> [0 .. numThetaFreq - 1]

{-# INLINE parMapDFTArray #-}
parMapDFTArray ::
     (VS.Vector (Complex Double) -> VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
parMapDFTArray f (DFTArray rows cols thetaFreqs rFreqs vecs) =
  DFTArray rows cols thetaFreqs rFreqs . parMap rdeepseq f $ vecs

{-# INLINE parZipWithDFTArray #-}
parZipWithDFTArray ::
     (VS.Vector (Complex Double) -> a -> VS.Vector (Complex Double))
  -> DFTArray
  -> [a]
  -> DFTArray
parZipWithDFTArray f (DFTArray rows cols thetaFreqs rFreqs vecs) =
  DFTArray rows cols thetaFreqs rFreqs . parZipWith rdeepseq f vecs

-- {-# INLINE plotDFTArrayPower #-}
plotDFTArrayPower :: FilePath -> Int -> Int -> DFTArray -> IO ()
plotDFTArrayPower !filePath !rows !cols =
  plotImageRepa filePath .
  ImageRepa 8 .
  fromUnboxed (Z :. (1 :: Int) :. cols :. rows) .
  VS.convert . VS.map sqrt .
  L.foldl1' (VS.zipWith (+)) .
  parMap rdeepseq (VS.map (\x -> (magnitude x) ** 2 )) . getDFTArrayVector

-- {-# INLINE plotDFTArrayThetaR #-}
plotDFTArrayThetaR :: FilePath -> Int -> Int -> [[Complex Double]] -> DFTArray -> IO ()
plotDFTArrayThetaR !filePath !rows !cols !thetaRHarmonics arr = do
  let vecs =
        computeFourierSeriesThetaR thetaRHarmonics .
        getDFTArrayVector $
        arr
      img =
        VS.convert . L.foldl1' (VS.zipWith (+)) . L.map (VS.map (magnitude)) $
        vecs
      m = VU.maximum img
  plotImageRepa filePath .
    ImageRepa 8 .
    computeS .
    reduceContrast 100 .
    fromUnboxed (Z :. (1 :: Int) :. cols :. rows) -- . VU.map (\x -> (x / m)** 1)
   $
    img
    -- VS.convert . L.foldl1' (VS.zipWith (+)) . L.map (VS.map magnitude) $
    -- vecs
  -- M.zipWithM_
  --   (\i vec ->
  --      plotImageRepa (dropExtension filePath L.++ (printf "_%d.png" i)) .
  --      ImageRepa 8 .
  --      fromUnboxed (Z :. (1 :: Int) :. cols :. rows) .
  --      VS.convert . VS.map magnitude $
  --      vec)
  --   [0 .. L.length vecs - 1]
  --   vecs

{-# INLINE plotDFTArrayThetaRMag #-}
plotDFTArrayThetaRMag :: FilePath -> Int -> Int -> [[Complex Double]] -> DFTArray -> IO ()
plotDFTArrayThetaRMag !filePath !rows !cols !thetaRHarmonics =
  plotImageRepa filePath .
  ImageRepa 8 .
  fromUnboxed (Z :. (1 :: Int) :. cols :. rows) .
  VS.convert .
  -- VS.map sqrt .
  L.foldl1' (VS.zipWith (max)) .
  L.map (VS.map (\x -> (magnitude x)^2)) .
  computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector

{-# INLINE sparseArrayToDFTArray #-}
sparseArrayToDFTArray ::
     Int
  -> Int
  -> [Double]
  -> [Double]
  -> [(Double, Double)]
  -> [R.Array U DIM2 (Complex Double)]
  -> DFTArray
sparseArrayToDFTArray rows cols thetaFreqs rFreqs xs sparseArray =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !numThetaFreq = L.length thetaFreqs
      !numRFreq = L.length rFreqs
  in repaToDFTArray thetaFreqs rFreqs .
     fromUnboxed (Z :. numRFreq :. numThetaFreq :. cols :. rows) .
     toUnboxedVector .
     AU.accum
       (+)
       0
       ((0, 0, minC, minR), (numRFreq - 1, numThetaFreq - 1, maxC, maxR)) .
     L.concat .
     L.zipWith
       (\(x, y) arr ->
          R.toList . R.traverse arr id $ \f idx@(Z :. i :. j) ->
            ((i, j, round x, round y), f idx))
       xs -- .
     -- L.map
     --   (\arr ->
     --      let !s =
     --            VU.maximum . toUnboxed . computeS . R.map (\x -> (magnitude x)) $
     --            arr
     --      in R.map (/ (s :+ 0)) arr)
  $
     sparseArray

{-# INLINE sparseArrayToDFTArray' #-}
sparseArrayToDFTArray' ::
     Int
  -> Int
  -> [Double]
  -> [Double]
  -> [(Double, Double)]
  -> [R.Array U DIM1 (Complex Double)]
  -> DFTArray
sparseArrayToDFTArray' rows cols thetaFreqs rFreqs xs sparseArray =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !numThetaFreq = L.length thetaFreqs
      arr =
        fromUnboxed (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows) .
        toUnboxedVector .
        AU.accum (+) 0 ((0, minC, minR), (numThetaFreq - 1, maxC, maxR)) .
        L.concat .
        L.zipWith
          (\(x, y) arr ->
             R.toList . R.traverse arr id $ \f idx@(Z :. j) ->
               ((j, round x, round y), f idx))
          xs -- .
        -- L.map
        --   (\arr ->
        --      let !s =
        --            VU.maximum .
        --            toUnboxed . computeS . R.map (\x -> (magnitude x)) $
        --            arr
        --      in R.map (/ (s :+ 0)) arr)
        $
        sparseArray
  in repaToDFTArray thetaFreqs rFreqs arr

{-# INLINE sparseArrayToDFTArrayG' #-}
sparseArrayToDFTArrayG' ::
     Int
  -> Double
  -> Int
  -> Int
  -> [Double]
  -> [Double]
  -> [(Double, Double)]
  -> [R.Array U DIM1 (Complex Double)]
  -> DFTArray
sparseArrayToDFTArrayG' n std rows cols thetaFreqs rFreqs xs sparseArray =
  let idx = [-n .. n]
      idx2D =
        L.filter
          (\(i, j) -> i ^ 2 + j ^ 2 < n ^ 2)
          [(i, j) | i <- idx, j <- idx]
      ys =
        L.map
          (\(x', y') ->
             L.map
               (\(i, j) ->
                  let x = round $ x' + fromIntegral i :: Int
                      y = round $ y' + fromIntegral j :: Int
                      v =
                        (exp $
                         ((fromIntegral x - x') ^ 2 + (fromIntegral y - y') ^ 2) /
                         (-2 * std ^ 2)) /
                        (std ^ 2 * 2 * pi)
                  in (x, y, v))
               idx2D)
          xs
      (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !numThetaFreq = L.length thetaFreqs
      arr =
        fromUnboxed (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows) .
        toUnboxedVector .
        AU.accum (+) 0 ((0, minC, minR), (numThetaFreq - 1, maxC, maxR)) .
        L.concat .
        L.zipWith
          (\zs arr ->
             L.concatMap
               (\(x, y, v) ->
                  R.toList . R.traverse arr id $ \f idx@(Z :. j) ->
                    ((j, x, y), (v :+ 0) * f idx))
               zs)
          ys $
        sparseArray
  in repaToDFTArray thetaFreqs rFreqs arr

{-# INLINE sparseArrayToDFTArray''' #-}
sparseArrayToDFTArray''' ::
     Int
  -> Int
  -> [Double]
  -> [Double]
  -> [(Double, Double)]
  -> [R.Array U DIM1 (Complex Double)]
  -> [(DFTArray,DFTArray)]
sparseArrayToDFTArray''' rows cols thetaFreqs rFreqs xs sparseArray =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !numThetaFreq = L.length thetaFreqs
      !pairs =
        L.zip xs -- .
        -- L.map
        --   (\arr ->
        --      let !s =
        --            sqrt . VU.sum .
        --            toUnboxed . computeS . R.map (\x -> (magnitude x)^2) $
        --            arr
        --      in R.map (/ (s :+ 0)) arr) $
        sparseArray
      computeArray pair =
        fromUnboxed (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows) .
        toUnboxedVector .
        AU.accum (+) 0 ((0, minC, minR), (numThetaFreq - 1, maxC, maxR)) .
        L.concatMap
          (\((x, y), arr) ->
             R.toList . R.traverse arr id $ \f idx@(Z :. j) ->
               ((j, round x, round y), f idx)) $
        pair
      func _ [] = []
      func zs (y:ys) =
        (computeArray [y], computeArray (zs L.++ ys)) : func (y : zs) ys
  in L.map
       (\(x, y) ->
          ( repaToDFTArray thetaFreqs rFreqs x
          , repaToDFTArray thetaFreqs rFreqs y)) .
     func [] $
     pairs

{-# INLINE fourierCoefficientsToR2Z2 #-}
fourierCoefficientsToR2Z2 ::
     [VS.Vector (Complex Double)] -> DFTArray -> [VU.Vector (Complex Double)]
fourierCoefficientsToR2Z2 !harmonics =
  parMap
    rdeepseq
    (\vec -> VU.fromList . L.map (VS.sum . VS.zipWith (*) vec) $ harmonics) .
  getDFTArrayVector

plotFourierCoefficientsPower ::
     FilePath -> Int -> Int -> [VS.Vector (Complex Double)] -> DFTArray -> IO ()
plotFourierCoefficientsPower !filePath !numRows !numCols !harmonics =
  plotImageRepa filePath .
  ImageRepa 8 .
  fromUnboxed (Z :. (1 :: Int) :. numCols :. numRows) .
  VS.convert .
  L.foldl1' (VU.zipWith (+)) .
  parMap rdeepseq (VU.map (\x -> (magnitude x) ** 2)) .
  fourierCoefficientsToR2Z2 harmonics

max_stars_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_issues_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_forks_repo_name: XinhuaZhang/Stochastic-Completion-Field


text:
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}

{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeOperators       #-}
module Test.Grenade.Layers.BatchNorm where

import           Control.Monad
import           Data.List                         (zipWith5)
import           Data.Proxy
import           GHC.TypeLits

import           Numeric.LinearAlgebra.Static      (L, R)
import qualified Numeric.LinearAlgebra.Static      as H

import           Hedgehog

import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix

import           Grenade.Core
import           Grenade.Layers.BatchNormalisation
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

batchnorm :: forall channels rows columns momentum.
  (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => Bool -> R channels -> R channels -> R channels -> R channels -> BatchNorm channels rows columns momentum
batchnorm training gamma beta mean var =
  let ε      = 0.00001
  in BatchNorm training (BatchNormParams gamma beta) mean var ε mkListStore

prop_batchnorm_train_behaves_as_reference :: Property
prop_batchnorm_train_behaves_as_reference = property $ do
  height   :: Int <- forAll $ choose 2 100
  width    :: Int <- forAll $ choose 2 100
  channels :: Int <- forAll $ choose 1 100

  case (someNatVal (fromIntegral height), someNatVal (fromIntegral width), someNatVal (fromIntegral channels), channels) of
    (Just (SomeNat (Proxy :: Proxy h)), Just (SomeNat (Proxy :: Proxy w)), _, 1) -> do
      inp :: S ('D2 h w) <- forAll genOfShape
      guard . not $ elementsEqual inp
      g :: R 1 <- forAll randomVector
      b :: R 1 <- forAll randomVector
      m :: R 1 <- forAll randomVector
      v :: R 1 <- forAll randomPositiveVector

      let layer   = batchnorm False g b m v :: BatchNorm 1 h w 90
          S2D out = snd $ runForwards layer inp :: S ('D2 h w)
          S2D ref = run2DBatchNorm layer inp :: S ('D2 h w)
      H.extract out === H.extract ref

    (Just (SomeNat (Proxy :: Proxy h)), Just (SomeNat (Proxy :: Proxy w)), Just (SomeNat (Proxy :: Proxy c)), _) -> do
      inp :: S ('D3 h w c) <- forAll genOfShape
      g :: R c <- forAll randomVector
      b :: R c <- forAll randomVector
      m :: R c <- forAll randomVector
      v :: R c <- forAll randomPositiveVector

      let layer   = batchnorm False g b m v :: BatchNorm c h w 90
          S3D out = snd $ runForwards layer inp :: S ('D3 h w c)
          S3D ref = run3DBatchNorm layer inp    :: S ('D3 h w c)
      H.extract out === H.extract ref

prop_batchnorm_1D_forward_same_as_torch :: Property
prop_batchnorm_1D_forward_same_as_torch = withTests 1 $ property $ do
    let g = H.fromList weight :: R 10
        b = H.fromList bias   :: R 10
        m = H.fromList mean   :: R 10
        v = H.fromList var    :: R 10

        bn = batchnorm False g b m v :: BatchNorm 10 4 4 90

        mat    = H.fromList . concat . concat $ input :: L 40 4
        x      = S3D mat :: S ('D3 4 4 10)
        y      = snd $ runForwards bn x :: S ('D3 4 4 10)

        mat' = H.fromList . concat . concat $ ref_out :: L 40 4

    assert $ allClose y (S3D mat')
  where
    weight = [ -0.9323,  1.0161,  0.1728,  0.3656, -0.6816,  0.0334,  0.8494, -0.6669, -0.1527, -0.7004 ]
    bias   = [  0.7582,  1.0068, -0.2532, -1.5240, -1.0370,  0.8442,  0.5867, -1.2567,  0.4283, -0.0001 ]
    mean   = [ -0.4507,  0.9090, -1.4717,  0.5009,  0.8931, -0.4792,  0.0432,  0.4649, -0.6547, -1.3197 ]
    var    = [  1.2303,  1.4775,  0.8372,  0.1644,  0.9392,  0.2103,  0.4951,  0.2482,  0.7559,  0.3686 ]

    input  = [ [ [ -0.70681204, -0.20616523, -0.33806887, -1.52378976   ]
               , [ 0.056113367, -0.51263034, -0.28884589, -2.64030218   ]
               , [ -1.19894597, -1.16714501, -0.19816216, -1.13361239   ]
               , [ -0.81997509, -1.05715847,  0.59198695,  0.51939314   ] ]
             , [ [ -0.18338945, -1.08975303,  0.30558434,  0.85780441   ]
               , [ -0.47586514,  0.16499641,  2.18205571, -0.11155529   ]
               , [ 1.090167402,  0.92460924,  0.42982020,  1.30098605   ]
               , [ 0.286766794, -1.90825951, -0.91737461, -1.11035680   ] ]
             , [ [ 1.042808533,  0.08287286, -0.92343962, -0.49747768   ]
               , [ -0.21943949,  0.61554014, -2.25771808, -0.04292159   ]
               , [ 1.290057424, -1.07323992, -1.00024509,  1.30155622   ]
               , [ 0.472014425, -0.96431374,  0.77593171, -1.19090688   ] ]
             , [ [ 0.993361895,  0.82586401, -1.64278686,  1.25544464   ]
               , [ 0.239656539, -0.81472164,  1.32814168,  0.78350490   ]
               , [ -0.16597847,  0.74175131, -1.29834091, -1.28858852   ]
               , [ 1.307537318,  0.55525642, -0.04312540,  0.24699424   ] ]
             , [ [ 0.391699581, -0.09803850, -0.41061267,  0.34999904   ]
               , [ -2.22257169,  0.43748092, -1.21343314,  0.39576068   ]
               , [ 0.003147978, -1.00396716,  1.27623140,  1.17001295   ]
               , [ -0.58247902, -0.15453417, -0.37016496,  0.04613848   ] ]
             , [ [ 0.521356827,  0.94643139,  1.11394095,  0.60162323   ]
               , [ -0.90214585, -0.75316292,  2.20823979, -1.63446676   ]
               , [ 0.668517357,  0.62832462,  0.31174039, -0.04457542   ]
               , [ -0.24607617,  0.12855675, -1.62831199, -0.23100854   ] ]
             , [ [ -0.43619379, -0.41219231,  0.07910434, -0.20312546   ]
               , [ 1.670419093, -0.26496240, -1.53759109,  1.00907373   ]
               , [ -1.04028647, -1.37309467, -0.79040497, -0.15661381   ]
               , [ 0.009049783, -0.05525103,  1.44492578,  0.44786781   ] ]
             , [ [ 1.431640263, -0.12869687,  1.25025844,  0.07864278   ]
               , [ -1.69032764, -0.07707843,  0.11284181, -0.00826502   ]
               , [ -0.92387816, -0.83121442,  0.42292186, -0.49128937   ]
               , [ -1.62631051,  0.98236626, -1.69256067, -0.66552013   ] ]
             , [ [ 0.154654814,  0.59295737,  0.48604089,  0.46829459   ]
               , [ 0.624001921,  2.11190581, -1.80008912,  0.26847255   ]
               , [ -0.36086676,  0.94211035,  0.19112136, -0.04113261   ]
               , [ -0.94438538, -0.38932472, -0.29867526,  0.34307864   ] ]
             , [ [ 1.016388653, -0.41974341, -0.94618958,  0.22629515   ]
               , [ -2.04437517, -1.14956784,  0.38054388,  0.82105201   ]
               , [ 0.054255251,  1.03682625,  0.29021424, -0.42736151   ]
               , [ -0.00021907,  0.98816186,  0.23878140, -0.17728853   ] ] ]

    ref_out = [ [ [ 0.9734674096107483, 0.5526634454727173, 0.6635311841964722, 1.660154104232788]
                , [0.3322128653526306, 0.8102537393569946, 0.6221582293510437, 2.5986058712005615]
                , [1.3871161937713623, 1.360386848449707, 0.5459367036819458, 1.3322019577026367]
                , [1.068583369255066, 1.267940878868103, -0.11819997429847717, -0.05718337371945381] ]
              , [ [0.0936361625790596, -0.66402268409729, 0.5023852586746216, 0.9640040397644043]
                , [-0.15085376799106598, 0.3848632574081421, 2.070988893508911, 0.15368467569351196]
                , [1.1582437753677368, 1.019848346710205, 0.606238067150116, 1.334473967552185]
                , [0.4866550862789154, -1.3482388257980347, -0.5199258923530579, -0.6812460422515869] ]
              , [ [0.22167539596557617, 0.04038754478096962, -0.14965875446796417, -0.06921406835317612]
                , [-0.016705399379134178, 0.14098398387432098, -0.4016427993774414, 0.016630740836262703]
                , [0.2683693766593933, -0.17794916033744812, -0.16416378319263458, 0.2705409526824951]
                , [0.11387854814529419, -0.15737800300121307, 0.1712745875120163, -0.20017105340957642] ]
              , [ [-1.0799674987792969, -1.230993390083313, -3.456873655319214, -0.8436583876609802]
                , [-1.7595523595809937, -2.7102415561676025, -0.7781105041503906, -1.2691868543624878]
                , [-2.1252965927124023, -1.30683434009552, -3.146301031112671, -3.137507677078247]
                , [-0.7966886162757874, -1.4749890565872192, -2.0145251750946045, -1.7529363632202148] ]
              , [ [-0.6843588948249817, -0.3399200439453125, -0.1200828030705452, -0.655030369758606]
                , [1.1542901992797852, -0.7165574431419373, 0.44455069303512573, -0.6872150897979736]
                , [-0.41108575463294983, 0.29723069071769714, -1.306460976600647, -1.2317562103271484]
                , [0.0007928922423161566, -0.3001859486103058, -0.14853018522262573, -0.4413214921951294] ]
              , [ [0.9170715808868408, 0.9480301737785339, 0.9602301120758057, 0.9229174852371216]
                , [0.8133963942527771, 0.8242470026016235, 1.0399290323257446, 0.760060727596283]
                , [0.9277894496917725, 0.9248621463775635, 0.9018049836158752, 0.8758541345596313]
                , [0.8611786365509033, 0.88846355676651, 0.7605089545249939, 0.862276017665863] ]
              , [ [0.007999604567885399, 0.036973003298044205, 0.6300419569015503, 0.28934815526008606]
                , [2.5509982109069824, 0.21470165252685547, -1.3215526342391968, 1.7526549100875854]
                , [-0.7212311029434204, -1.1229807138442993, -0.4195865988731384, 0.34549471735954285]
                , [0.5454756021499634, 0.4678548276424408, 2.278794050216675, 1.0751949548721313] ]
              , [ [-2.550779104232788, -0.4621109068393707, -2.307981491088867, -0.7396559119224548]
                , [1.628288984298706, -0.5312073826789856, -0.7854347229003906, -0.6233210563659668]
                , [0.6023192405700684, 0.4782795310020447, -1.2005081176757812, 0.023255640640854836]
                , [1.542595624923706, -1.9493807554244995, 1.631278157234192, 0.2564810514450073] ]
              , [ [0.28615128993988037, 0.20917125046253204, 0.22794923186302185, 0.2310660481452942]
                , [0.20371884107589722, -0.05760492384433746, 0.6294671297073364, 0.2661612331867218]
                , [0.3766934275627136, 0.14784877002239227, 0.2797465920448303, 0.32053783535957336]
                , [0.4791780710220337, 0.381691575050354, 0.3657706081867218, 0.25305798649787903] ]
              , [ [-2.6950571537017822, -1.0383073091506958, -0.4309888482093811, -1.7835899591445923]
                , [0.8358993530273438, -0.19636774063110352, -1.9615343809127808, -2.469712972640991]
                , [-1.5851213932037354, -2.7186343669891357, -1.8573282957077026, -1.029518961906433]
                , [-1.5222787857055664, -2.66249418258667, -1.7979943752288818, -1.3180080652236938] ] ]

tests :: IO Bool
tests = checkParallel $$(discover)

-- REFERENCE FUNCTIONS

run2DBatchNorm :: forall h w m.
                  (KnownNat h, KnownNat w, KnownNat m)
               => BatchNorm 1 h w m -> S ('D2 h w) ->  S ('D2 h w)
run2DBatchNorm (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S2D x)
  = let [m]    = vectorToList runningMean
        [v]    = vectorToList runningVar
        [g]    = vectorToList gamma
        [b]    = vectorToList beta
        std    = sqrt $ v + ε
        x_norm = H.dmmap (\a -> (a - m) / std) x
        out    = H.dmmap (\a -> g * a + b) x_norm
    in S2D out

run3DBatchNorm :: forall h w m c.
                  (KnownNat h, KnownNat w, KnownNat m, KnownNat c)
               => BatchNorm c h w m -> S ('D3 h w c) ->  S ('D3 h w c)
run3DBatchNorm (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) inp
  = let ms     = vectorToList runningMean
        vs     = vectorToList runningVar
        gs     = vectorToList gamma
        bs     = vectorToList beta

        cs     = splitChannels inp :: [S ('D2 h w)]

        f c g b m v = let gs' = listToVector [g] :: R 1
                          bs' = listToVector [b] :: R 1
                          ms' = listToVector [m] :: R 1
                          vs' = listToVector [v] :: R 1
                          bn' = BatchNorm False (BatchNormParams gs' bs') ms' vs' ε undefined :: BatchNorm 1 h w m
                      in  run2DBatchNorm bn' c
      in combineChannels $ zipWith5 f cs gs bs ms vs

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade


text:
module STCR2S1Image where

import           Control.Monad           as M
import           Control.Monad.Parallel  as MP
import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           Data.Vector.Unboxed     as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           Image.IO
import           STC.CompletionFieldR2S1
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:numTrailStr:maxTrailStr:initOriStr:initSpeedStr:writeFlagStr:numIterationStr:thresholdStr:histFilePath:cutoffStr:imagePath:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      initOri = read initOriStr :: Double
      initSpeed = read initSpeedStr :: Double
      writeFlag = read writeFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      threshold = read thresholdStr :: Double
      cutoff = read cutoffStr :: Int
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2S1Image"
  imgRepa@(ImageRepa _ img') <- readImageRepa imagePath False
  let (Z :. _ :. cols :. rows) = extent img'
      img =
        computeS . R.traverse img' id $ \f idx@(Z :. k :. i :. j) ->
          if (sqrt . fromIntegral $ (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2) >
             32
            then 0
            else f idx
  createDirectoryIfMissing True folderPath
  plotImageRepa (folderPath </> "input.png") . ImageRepa 8 $ img
  flag <- doesFileExist histFilePath
  arrG <-
    if flag
      then getNormalizedHistogramArr .
           mapHistogram
             (\x ->
                let y = x :: Int
                 in fromIntegral y) <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        computeS . R.map magnitude <$>
          solveMonteCarloR2S1
            numThread
            numTrail
            numPoint
            numPoint
            numOrientation
            sigma
            tao
            cutoff
            histFilePath
            (0, 0, 0, 0, initOri / 180 * pi, initSpeed)
  powerMethod
    emptyPlan
    folderPath
    arrG
    numIteration
    writeFlag
    ""
    threshold
    (R.traverse img (const (Z :. numOrientation :. cols :. rows)) $ \f (Z :. _ :. i :. j) ->
       if f (Z :. (0 :: Int) :. i :. j) > 0
         then 1
         else 0)
    (R.traverse img (const (Z :. numOrientation :. cols :. rows)) $ \f (Z :. _ :. i :. j) ->
       if f (Z :. (0 :: Int) :. i :. j) > 0
         then 1 / fromIntegral numOrientation
         else 0)


text:
{-# LANGUAGE Strict #-}
module STC.Bias where

import           Array.UnboxedArray       as AU
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Storable     as VS
import           Data.Vector.Unboxed      as VU
import Data.Vector.Generic as VG
import           DFT.Plan
import           Filter.Utils
import           Pinwheel.FourierSeries2D
import           STC.Point
import           STC.Utils
import           Utils.List
import FourierPinwheel.AsteriskGaussian
import Filter.Utils
import           Data.Array.Accelerate.LLVM.PTX as A
import qualified Data.Array.Accelerate          as A
import FourierMethod.FourierSeries2D
import Image.IO
import Utils.Array
import System.FilePath
import Utils.Parallel
import FourierPinwheel.GaussianEnvelopePinwheel

{-# INLINE computeBias #-}
computeBias ::
     (Num a, Unbox a, VG.Vector vector a) => Int -> Int -> [Point] -> vector a
computeBias rows cols =
  let (minR, maxR) = computeRange rows
      (minC, maxC) = computeRange cols
   in VG.convert .
      toUnboxedVector .
      AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
      L.map (\(Point x y theta scale) -> ((round x, round y), 1))

computeBiasPinwheelBasis ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> IO (VS.Vector (Complex Double))
computeBiasPinwheelBasis plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      envelope =
        computeUnboxedS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in dftExecute plan (DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]) .
     VS.convert .
     toUnboxed .
     computeS .
     makeFilter2D .
     R.zipWith (*) envelope .
     fromListUnboxed (Z :. numR2Freq :. numR2Freq) $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasPinwheelBasis1 ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> VU.Vector (Complex Double)
computeBiasPinwheelBasis1 plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      envelope =
        toUnboxed . computeS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in VU.zipWith (*) envelope . VU.fromList $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasGaussian ::
     DFTPlan
  -> Int
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double))
computeBiasGaussian plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      freqCenter = div numR2Freq 2
      envelope =
        fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - freqCenter) ^ 2 + (j - freqCenter) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in dftExecute plan (DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]) .
     VS.convert .
     toUnboxed .
     computeS .
     makeFilter2D .
     R.zipWith (*) envelope . fromListUnboxed (Z :. numR2Freq :. numR2Freq) $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasGaussian1 ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> VU.Vector (Complex Double)
computeBiasGaussian1 plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      freqCenter = div numR2Freq 2
      envelope =
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - freqCenter) ^ 2 + (j - freqCenter) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in VU.zipWith (*) envelope . VU.fromList $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasFourierPinwheel ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheel plan numR2Freqs thetaFreq alpha periodR2 periodEnv radius stdTheta points = do
  asteriskGaussianVec <-
    asteriskGaussianLowPass
      plan
      numR2Freqs
      thetaFreq
      alpha
      periodR2
      periodEnv
      stdTheta
      radius
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      numThetaFreq = 2 * thetaFreq + 1
      shiftVec =
        VS.concat . L.replicate numThetaFreq . VS.fromList $
        [ L.foldl'
          (\b (Point x y theta scale) ->
             b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
          0
          points
        | freqY <- r2Freqs
        , freqX <- r2Freqs
        ]
      biasVec = VS.zipWith (*) asteriskGaussianVec shiftVec
      -- centerFreq = div numR2Freqs 2
      -- gaussian2D =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
      --     let i = i' - centerFreq
      --         j = j' - centerFreq
      --      in exp
      --           (pi * fromIntegral (i ^ 2 + j ^ 2) /
      --            ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
      --         (2 * pi * stdR2 ^ 2) :+
      --         0
      -- zeroVec = VS.replicate (numR2Freqs ^ 2) 0
      -- shiftVec =
      --   VS.concat .
      --   L.map
      --     (\af ->
      --        if af == 0
      --          then VS.zipWith (*) gaussian2D . VS.fromList $
      --               [ L.foldl'
      --                 (\b (Point x y theta scale) ->
      --                    b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
      --                 0
      --                 points
      --               | freqY <- r2Freqs
      --               , freqX <- r2Freqs
      --               ]
      --          else zeroVec) $
      --   [-thetaFreq .. thetaFreq]
  biasVecF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [numThetaFreq, numR2Freqs, numR2Freqs] [0, 1, 2]) .
    VU.convert .
    toUnboxed .
    computeS .
    makeFilter3D .
    fromUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) . VS.convert $
    -- shiftVec
    biasVec
  return (biasVec, biasVecF)
  -- return (shiftVec, biasVecF)

computeBiasFourierPinwheelFull ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> (VU.Vector (Complex Double))
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelFull plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 asteriskGaussianVec points = do
  let -- asteriskGaussianVec =
      --   gaussianPinwheel
      --     numR2Freqs
      --     periodR2
      --     stdR2
      --     10
      --     thetaFreq
      --     rFreq
      --     periodEnv
      --     stdTheta
      --     stdR
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
      shiftArr =
        fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) .
        parMap
          rdeepseq
          (\(freqY, freqX) ->
             L.foldl'
               (\b (Point x y theta scale) ->
                  b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
               0
               points) $
        [(freqY, freqX) | freqY <- r2Freqs, freqX <- r2Freqs]
      asteriskGaussianArr =
        fromUnboxed
          (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs)
          asteriskGaussianVec
  biasArr <-
    computeUnboxedP . R.traverse2 asteriskGaussianArr shiftArr const $ \fA fS idx@(Z :. _ :. _ :. i :. j) ->
      fA idx * fS (Z :. i :. j)
  biasVecF <-
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreq, numThetaFreq, numR2Freqs, numR2Freqs]
         [0, 1, 2, 3]) .
    VU.convert . toUnboxed . computeS . makeFilter4D $
    biasArr
  return
    ( VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $ biasArr
    , biasVecF)

computeBiasFourierPinwheelKoffkaCorss ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> (VU.Vector (Complex Double))
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelKoffkaCorss plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 asteriskGaussianVec points
      -- asteriskGaussianVec =
      --   gaussianPinwheel1
      --     numR2Freqs
      --     periodR2
      --     stdR2
      --     10
      --     thetaFreq
      --     rFreq
      --     periodEnv
      --     stdTheta
      --     stdR
 = do
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      thetaFreqs = [fromIntegral (-thetaFreq) .. fromIntegral thetaFreq]
      numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
      halfPI = 0.5 * pi
      centerFreqR2 = div numR2Freqs 2
      shiftArr =
        fromFunction (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) $ \(Z :. tf :. yf :. xf) ->
          let freqTheta = fromIntegral (tf - thetaFreq)
              freqY = fromIntegral (yf - centerFreqR2)
              freqX = fromIntegral (xf - centerFreqR2)
           in L.foldl'
                (\b (Point x y theta scale) ->
                   let phi =
                         if abs x < abs y
                           then atan2 0 x
                           else if abs x > abs y
                                  then atan2 y 0
                                  else error
                                         "computeBiasFourierPinwheelKoffkaCorss: x == y"
                    in b +
                       cis (-(freqX * x + freqY * y) * 2 * pi / periodR2) *
                       (cis (-freqTheta * (phi + halfPI)) +
                        cis (-freqTheta * (phi - halfPI))))
                0
                points
      -- shiftArr =
      --   fromListUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) .
      --   parMap
      --     rdeepseq
      --     (\(thetaFreq, freqY, freqX) ->
      --        L.foldl'
      --          (\b (Point x y theta scale) ->
      --             let phi =
      --                   if abs x < abs y
      --                     then atan2 0 x
      --                     else if abs x > abs y
      --                            then atan2 y 0
      --                            else error
      --                                   "computeBiasFourierPinwheelKoffkaCorss: x == y"
      --              in b +
      --                 cis (-(freqX * x + freqY * y) * 2 * pi / periodR2) *
      --                 (cis (-thetaFreq * (phi + halfPI)) +
      --                  cis (-thetaFreq * (phi - halfPI))))
      --          0
      --          points) $
      --   [ (thetaFreq, freqY, freqX)
      --   | thetaFreq <- thetaFreqs
      --   , freqY <- r2Freqs
      --   , freqX <- r2Freqs
      --   ]
      asteriskGaussianArr =
        fromUnboxed
          (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs)
          asteriskGaussianVec
  biasArr <-
    computeUnboxedP . R.traverse2 asteriskGaussianArr shiftArr const $ \fA fS idx@(Z :. _ :. k :. i :. j) ->
      fA idx * fS (Z :. k :. i :. j)
  biasVecF <-
    (computeUnboxedP . makeFilter4D $ biasArr) >>=
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreq, numThetaFreq, numR2Freqs, numR2Freqs]
         [0, 1, 2, 3]) .
    VU.convert . toUnboxed
  biasVec <- VU.convert . toUnboxed <$> computeUnboxedP (makeFilter2D biasArr)
  return (biasVec, biasVecF)

computeBiasFourierPinwheelDiscrete ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelDiscrete plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 points = do
  let biasVec =
        VU.convert .
        toUnboxed .
        computeS . makeFilter2D . fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        computeBias numR2Freqs numR2Freqs points
      bias = VS.concat . L.replicate (2 * thetaFreq + 1) $ biasVec
      zeroVec = VS.replicate (VS.length biasVec) 0
  biasVecF <-
    dftExecute plan (DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]) biasVec
  let biasF =
        VS.concat .
        L.map
          (\angularFreq ->
             if angularFreq == 0
               then biasVecF
               else zeroVec) $
        [-thetaFreq .. thetaFreq]
  return (bias, biasF)

computeBiasFourierPinwheelDiscrete' ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelDiscrete' plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 points = do
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      shiftArr =
        fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) .
        parMap
          rdeepseq
          (\(freqY, freqX) ->
             L.foldl'
               (\b (Point x y theta scale) ->
                  b + cis (-(freqX * (fromIntegral (round x :: Int)) + freqY * (fromIntegral (round y :: Int))) * 2 * pi / periodR2))
               0
               points) $
        [(freqY, freqX) | freqY <- r2Freqs, freqX <- r2Freqs]
      shiftVec =
        VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $ shiftArr
      centerFreq = div numR2Freqs 2
      envelope =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
          if i == centerFreq && j == centerFreq
            then 0
            else let r =
                       sqrt . fromIntegral $
                       (i - centerFreq) ^ 2 + (j - centerFreq) ^ 2
                  in (1 / (r * periodR2)) :+ 0
      envelopeShiftVec =
        VU.convert . toUnboxed . computeS . makeFilter2D $
        shiftArr -- *^ (centerHollowArray numR2Freqs envelope)
      zeroVec = VS.replicate (VS.length shiftVec) 0
  biasVec <-
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1])
      envelopeShiftVec
  let bias = VS.concat . L.replicate (2 * thetaFreq + 1) $ biasVec
      biasF =
        VS.concat .
        L.map
          (\angularFreq ->
             if angularFreq == 0
               then shiftVec
               else zeroVec) $
        [-thetaFreq .. thetaFreq]
  return (bias, biasF)

-- computeBiasFourierPinwheel ::
--      DFTPlan
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> [Point]
--   -> FilePath
--   -> [PTX]
--   -> Int
--   -> Int
--   -> IO (VS.Vector (Complex Double))
-- computeBiasFourierPinwheel plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2 stdTheta points folderPath ptxs numPoints numBatch = do
--   let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
--       numThetaFreq = 2 * thetaFreq + 1
--       gaussianVec =
--         asteriskGaussian1 numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2
--       shiftCoef =
--         [ L.foldl'
--           (\b (Point x y theta scale) ->
--              b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
--           0
--           points
--         | freqY <- r2Freqs
--         , freqX <- r2Freqs
--         ]
--       shiftVec = VS.concat . L.replicate numThetaFreq . VS.fromList $ shiftCoef
--       biasVec = VS.zipWith (*) (VS.concat gaussianVec) shiftVec
--       biasMat =
--         A.transpose .
--         A.use .
--         A.fromList (A.Z A.:. numThetaFreq A.:. (numR2Freqs ^ 2)) . VS.toList $
--         biasVec
--   biasR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       numThetaFreq
--       periodR2
--       1
--       numBatch
--       biasMat
--   plotImageRepa (folderPath </> "AsteriskBias.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     biasR2
--   let centerFreq = div numR2Freqs 2
--       shiftArr = fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) shiftCoef
--       gaussian2D =
--         computeUnboxedS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
--           let i = i' - centerFreq
--               j = j' - centerFreq
--            in (exp
--                  ((pi * (fromIntegral $ i ^ 2 + j ^ 2)) /
--                   ((-1) * periodR2 ^ 2 * stdR2 ^ 2))) /
--               (2 * pi * stdR2 ^ 2) :+
--               0
--       gaussian2DMat =
--         A.use .
--         A.fromList (A.Z A.:. (numR2Freqs ^ 2) A.:. (1 :: Int)) .
--         R.toList . R.zipWith (*) shiftArr $
--         gaussian2D
--   gaussian2DR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       1
--       periodR2
--       1
--       numBatch
--       gaussian2DMat
--   plotImageRepa (folderPath </> "Gaussian2D.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     gaussian2DR2
--   let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
--       inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
--   gaussian2DF <-
--     dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
--     gaussian2D
--   asteriskGaussianF <- dftExecuteBatchP plan planID gaussianVec
--   outputVec <-
--     fmap VS.concat .
--     dftExecuteBatchP plan inversePlanID .
--     parMap rdeepseq (VS.zipWith (*) gaussian2DF) $
--     asteriskGaussianF
--   let outputMat =
--         A.transpose .
--         A.use .
--         A.fromList (A.Z A.:. numThetaFreq A.:. (numR2Freqs ^ 2)) . VS.toList $
--         outputVec
--   outputR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       numThetaFreq
--       periodR2
--       1
--       numBatch
--       outputMat
--   plotImageRepa (folderPath </> "AsteriskGaussian.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     outputR2
--   dftExecute
--     plan
--     (DFTPlanID DFT1DG [numThetaFreq, numR2Freqs, numR2Freqs] [0, 1, 2]) .
--     VU.convert .
--     toUnboxed .
--     computeS .
--     makeFilter3D .
--     fromUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) . VS.convert $
--     -- shiftVec
    -- VS.zipWith (*) outputVec shiftVec

-- zeroVec = VS.replicate (numR2Freqs ^ 2) 0
-- shiftVec =
--   VS.concat .
--   L.map
--     (\angularFreq ->
--        if angularFreq == 0
--          then VS.fromList $
--               [ L.foldl'
--                 (\b (Point x y _ _) ->
--                    b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
--                 0
--                 points
--               | freqY <- r2Freqs
--               , freqX <- r2Freqs
--               ]
--          else zeroVec) $
--   [-thetaFreq .. thetaFreq]


text:
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Recurrent.Layers.BasicRecurrent (
    BasicRecurrent (..)
  ) where

import           Data.Proxy
import           Data.Singletons.TypeLits hiding (natVal)

import           Numeric.LinearAlgebra.Static

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Recurrent.Core

data BasicRecurrent :: Nat -- Input layer size
                    -> Nat -- Output layer size
                    -> * where
  BasicRecurrent :: ( KnownNat input
                    , KnownNat output
                    , KnownNat matrixCols
                    , matrixCols ~ (input + output))
                 => !(R output)   -- Bias neuron weights
                 -> !(R output)   -- Bias neuron momentum
                 -> !(L output matrixCols) -- Activation
                 -> !(L output matrixCols) -- Momentum
                 -> BasicRecurrent input output

data BasicRecurrent' :: Nat -- Input layer size
                     -> Nat -- Output layer size
                     -> * where
  BasicRecurrent' :: ( KnownNat input
                     , KnownNat output
                     , KnownNat matrixCols
                     , matrixCols ~ (input + output))
                  => !(R output)   -- Bias neuron gradients
                  -> !(L output matrixCols)
                  -> BasicRecurrent' input output

instance Show (BasicRecurrent i o) where
  show BasicRecurrent {} = "BasicRecurrent"

instance (KnownNat i, KnownNat o, KnownNat (i + o)) => UpdateLayer (BasicRecurrent i o) where
  type Gradient (BasicRecurrent i o) = (BasicRecurrent' i o)

  runUpdate LearningParameters {..} (BasicRecurrent oldBias oldBiasMomentum oldActivations oldMomentum) (BasicRecurrent' biasGradient activationGradient) =
    let newBiasMomentum = konst learningMomentum * oldBiasMomentum - konst learningRate * biasGradient
        newBias         = oldBias + newBiasMomentum
        newMomentum     = konst learningMomentum * oldMomentum - konst learningRate * activationGradient
        regulariser     = konst (learningRegulariser * learningRate) * oldActivations
        newActivations  = oldActivations + newMomentum - regulariser
    in BasicRecurrent newBias newBiasMomentum newActivations newMomentum

instance (KnownNat i, KnownNat o, KnownNat x, KnownNat (x*o), x ~ (i+o)) => RandomLayer (BasicRecurrent i o) where
  createRandomWith m gen = do
    wB <- getRandomVector i o m gen
    wN <- getRandomMatrix i o m gen
    let bm = konst 0
        mm = konst 0
    return $ BasicRecurrent wB bm wN mm
      where i = natVal (Proxy :: Proxy i)
            o = natVal (Proxy :: Proxy o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentUpdateLayer (BasicRecurrent i o) where
  type RecurrentShape (BasicRecurrent i o) = S ('D1 o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentLayer (BasicRecurrent i o) ('D1 i) ('D1 o) where

  type RecTape (BasicRecurrent i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
  -- Do a matrix vector multiplication and return the result.
  runRecurrentForwards (BasicRecurrent wB _ wN _) (S1D lastOutput) (S1D thisInput) =
    let thisOutput = S1D $ wB + wN #> (thisInput # lastOutput)
    in ((S1D lastOutput, S1D thisInput), thisOutput, thisOutput)

  -- Run a backpropogation step for a full connected layer.
  runRecurrentBackwards (BasicRecurrent _ _ wN _) (S1D lastOutput, S1D thisInput) (S1D dRec) (S1D dEdy) =
    let biasGradient        = (dRec + dEdy)
        layerGrad           = (dRec + dEdy) `outer` (thisInput # lastOutput)
        -- calcluate derivatives for next step
        (backGrad, recGrad) = split $ tr wN #> (dRec + dEdy)
    in  (BasicRecurrent' biasGradient layerGrad, S1D recGrad, S1D backGrad)

-------------------- GNum instances --------------------

instance (KnownNat i,KnownNat o,KnownNat (i+o)) => GNum (BasicRecurrent i o) where
  n |* (BasicRecurrent wB mB a nM) = (BasicRecurrent (fromRational n * wB) mB a nM)
  (BasicRecurrent wB mB a nM) |+ (BasicRecurrent wB2 mB2 a2 nM2)  = (BasicRecurrent (0.5 * (wB + wB2)) (0.5* (mB+mB2)) (0.5 * (a+a2)) (0.5 * (nM + nM2)))
  gFromRational r = (BasicRecurrent (fromRational r) 0 (fromRational r) 0)

max_stars_repo_name: koenigmaximilian/grenade

max_issues_repo_name: koenigmaximilian/grenade

max_forks_repo_name: koenigmaximilian/grenade


text:
module STCR2Z1T0ImageSparseCode where

import           Control.Monad                as M
import           Control.Monad.Parallel       as MP
import           Data.Array.Repa              as R
import           Data.Binary                  (decodeFile)
import           Data.Complex
import           Data.Conduit
import           Data.Conduit.List            as CL
import           Data.List                    as L
import           Data.Maybe
import           Data.Vector.Generic          as VG
import           Data.Vector.Storable         as VS
import           DFT.Plan
import           Filter.Gaussian
import           Filter.Pinwheel
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           Image.Transform
import           PVPFile.IO
import           STC
import           STC.OrientationScaleAnalysis
import           System.Directory
import           System.Environment
import           System.FilePath
import           System.Random
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initialScaleStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:alphaStr:pinwheelFlagStr:imagePath:reconPath:numIterationStr:writeSourceFlagStr:thresholdStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      initialScale = read initialScaleStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      threshold = read thresholdStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z1T0ImageSparseCode"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  imgHeader <- readPVPHeader imagePath
  reconHeader <- readPVPHeader reconPath
  let m = min (nBands imgHeader) (nBands reconHeader)
  print m
  xs <-
    runConduitRes $
    pvpFileSource imagePath .| (CL.drop (m - 1) >> CL.map id) .| CL.head
  ys <-
    runConduitRes $
    pvpFileSource reconPath .| (CL.drop (m - 1) >> CL.map id) .| CL.head
  let img' = pvpOutputData2Array . fromJust $ xs
      recon = pvpOutputData2Array . fromJust $ ys
      (Z :. cols :. rows :. _) = extent img'
  print . extent $ img'
  plotImageRepa
    (folderPath </> "recon.png")
    (ImageRepa 8 . computeS . rotate3DR $ recon)
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z1T0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          sigma
          tao
          1
          theta0Freqs
          thetaFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length theta0Freqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z1T0 <-
    computeUnboxedP $
    computeR2Z1T0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      thetaFreqs
      theta0Freqs
  -- let arr3d =
  --       rotate3D . R.slice arrR2Z1T0 $
  --       (Z :. All :. (L.length theta0Freqs - 1) :. All :. All)
  -- MP.mapM_
  --   (\i ->
  --      plotImageRepaComplex
  --        (folderPath </> "GreensR2Z1T0_" L.++ show (i + 1) L.++ ".png") .
  --      ImageRepa 8 .
  --      computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr3d $
  --      (Z :. All :. All :. i))
  --   [0 .. (L.length thetaFreqs) - 1]
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  img'' <-
    fmap (fromUnboxed (Z :. cols :. rows)) .
    dropPixel (4 / 5) . toUnboxed . R.sumS $
    img'
  let s = R.sumAllS img'
      img = R.map (/ s) . R.sumS $ img'
      initialDistSource =
        R.traverse
          img
          (\(Z :. nx :. ny) -> (Z :. (L.length theta0Freqs) :. nx :. ny))
          (\f (Z :. _ :. i :. j) -> f (Z :. i :. j) :+ 0)
  undefined
  -- powerMethod1
  --   plan
  --   folderPath
  --   cols
  --   rows
  --   numOrientation
  --   thetaFreqs
  --   theta0Freqs
  --   arrR2Z1T0
  --   numIteration
  --   writeSourceFlag
  --   ""
  --   threshold
  --   initialDistSource


text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}

module NormalizeSpec where

import Commons
import Data.Array (listArray)
import Data.Complex (Complex (..))
import qualified Data.IntMap.Strict as IM
import qualified Data.Map as Map
import Data.Maybe (fromJust)
import HashedExpression.Internal.Expression
import HashedExpression.Internal.Normalize
import HashedExpression.Internal.Utils
import HashedExpression.Interp (eval, (~=))
import HashedExpression.Operation
import HashedExpression.Prettify
import HashedExpression.Value
import Test.Hspec
import Var
import Prelude hiding ((^))

reFT :: (DimensionType d) => Expression d R -> Expression d R
reFT = xRe . ft

imFT :: (DimensionType d) => Expression d R -> Expression d R
imFT = xIm . ft

spec :: Spec
spec = do
  describe "Normalize spec" $ do
    specify "normalize scalar one zero" $ do
      constant 0.0 *. constant 9.0 `shouldNormalizeTo` constant 0.0
      x * one `shouldNormalizeTo` x
      one * x `shouldNormalizeTo` x
      x * zero `shouldNormalizeTo` zero
      zero * x `shouldNormalizeTo` zero
      y * (x * zero) `shouldNormalizeTo` zero
      zero * (x * one) `shouldNormalizeTo` zero
      zero * x * one `shouldNormalizeTo` zero
      zero * (x * y) `shouldNormalizeTo` zero
      (x * y) * zero `shouldNormalizeTo` zero
      (x * zero) * one `shouldNormalizeTo` zero
      ((x * y) * one) `shouldNormalizeTo` (x * y)
      x * y * z * one `shouldNormalizeTo` x * y * z
      product [x, y, z, t, w, zero] `shouldNormalizeTo` zero
    specify "normalize log and exponential" $ do
      log (exp x) `shouldNormalizeTo` x
      exp (log x) `shouldNormalizeTo` x
    specify "complex related" $ do
      ((x +: y) * (z +: w)) `shouldNormalizeTo` ((x * z - y * w) +: (x * w + y * z))
      xRe (x +: y) `shouldNormalizeTo` x
      xIm (x +: y) `shouldNormalizeTo` y
      (x +: y) + (u +: v) `shouldNormalizeTo` (x + u) +: (y + v)
      s *. (x +: y) `shouldNormalizeTo` (s *. x) +: (s *. y)
      (x +: y) * (z +: w) `shouldNormalizeTo` (x * z - y * w) +: (x * w + y * z)
    specify "dot product" $ do
      x <.> zero `shouldNormalizeTo` zero
      zero <.> x `shouldNormalizeTo` zero
      ((s *. x) <.> y) `shouldNormalizeTo` (s *. (x <.> y))
      x <.> (s *. y) `shouldNormalizeTo` s *. (x <.> y)
    specify "distributivity" $ do
      x * (y + z) `shouldNormalizeTo` (x * y + x * z)
      (y + z) * x `shouldNormalizeTo` (x * y + x * z)
      (x *. (y + z)) `shouldNormalizeTo` (x *. y + x *. z)
      (x <.> (y + z)) `shouldNormalizeTo` ((x <.> y) + (x <.> z))
      (y + z) <.> x `shouldNormalizeTo` (y <.> x) + (z <.> x)
      x * sum [y, z, t, u, v] `shouldNormalizeTo` sum (map (x *) [y, z, t, u, v])
      sum [y, z, t, u, v] * x `shouldNormalizeTo` sum (map (x *) [y, z, t, u, v])
      x *. sum [y, z, t, u, v] `shouldNormalizeTo` sum (map (x *.) [y, z, t, u, v])
      x <.> sum [y, z, t, u, v] `shouldNormalizeTo` sum (map (x <.>) [y, z, t, u, v])
      sum [y, z, t, u, v] <.> x `shouldNormalizeTo` sum (map (<.> x) [y, z, t, u, v])
      product [a, b, c, sum [x, y, z]] `shouldNormalizeTo` sum (map (product . (: [a, b, c])) [x, y, z])
      (x + y) * (z + t) * a * b `shouldNormalizeTo` (a * b * x * z + a * b * x * t + a * b * y * z + a * b * y * t)
    specify "flatten sum and product" $ do
      product [x * y, product [z, t, w], one] `shouldNormalizeTo` product [x, y, z, t, w]
      sum [x + y, sum [z, t, w + s], zero] `shouldNormalizeTo` sum [x, y, z, t, w, s]
    specify "group constants together" $ do
      product [one, one, x, y, one, z] `shouldNormalizeTo` product [x, y, z]
      sum [one, one, x, y, one, z] `shouldNormalizeTo` sum [constant 3, x, y, z]
      product [constant 1, constant 2, x, y, constant 3, z] `shouldNormalizeTo` product [constant 6, x, y, z]
    specify "combine same terms" $ do
      sum [one *. x1, x1, x1, constant 3 *. y1, y1] `shouldNormalizeTo` sum [constant 3 *. x1, constant 4 *. y1]
      sum [constant (-1) *. x1, x1, constant 3 *. y1, y1, z1] `shouldNormalizeTo` sum [constant 4 *. y1, z1]
      x1 - x1 `shouldNormalizeTo` zero1
      sum [one *. x, x, x, constant 3 *. y, y] `shouldNormalizeTo` sum [constant 3 *. x, constant 4 *. y]
      sum [constant (-1) *. x, x, constant 3 *. y, y, z] `shouldNormalizeTo` sum [constant 4 *. y, z]
      x - x `shouldNormalizeTo` zero
    specify "scale rules" $ do
      x *. (y *. v) `shouldNormalizeTo` (x * y) *. v
    specify "negate rules" $ do
      negate (negate x) `shouldNormalizeTo` normalize x
      negate (negate (x + y)) `shouldNormalizeTo` (x + y)
      negate zero `shouldNormalizeTo` zero
  describe "Normalize spec higher dimension" $ do
    specify "normalize one d one zero" $ do
      x1 * one1 `shouldNormalizeTo` x1
      one1 * x1 `shouldNormalizeTo` x1
      x1 * zero1 `shouldNormalizeTo` zero1
      zero1 * x1 `shouldNormalizeTo` zero1
      y1 * (x1 * zero1) `shouldNormalizeTo` zero1
      zero1 * (x1 * one1) `shouldNormalizeTo` zero1
      zero1 * x1 * one1 `shouldNormalizeTo` zero1
      zero1 * (x1 * y1) `shouldNormalizeTo` zero1
      (x1 * y1) * zero1 `shouldNormalizeTo` zero1
      (x1 * zero1) * one1 `shouldNormalizeTo` zero1
      (x1 * y1) * one1 `shouldNormalizeTo` (x1 * y1)
      x1 * y1 * z1 * one1 `shouldNormalizeTo` x1 * y1 * z1
    specify "dot product higher dimension with scaling and point wise" $ do
      x1 <.> zero1 `shouldNormalizeTo` zero
      zero1 <.> x1 `shouldNormalizeTo` zero
      (s *. x1) <.> y1 `shouldNormalizeTo` s *. (x1 <.> y1)
      x1 <.> (s *. y1) `shouldNormalizeTo` s *. (x1 <.> y1)
      x1 * (y1 + z1) `shouldNormalizeTo` x1 * y1 + x1 * z1
      (y1 + z1) * x1 `shouldNormalizeTo` x1 * y1 + x1 * z1
      s *. (y1 + z1) `shouldNormalizeTo` s *. y1 + s *. z1
      x1 <.> (y1 + z1) `shouldNormalizeTo` (x1 <.> y1) + (x1 <.> z1)
      (y1 + z1) <.> x1 `shouldNormalizeTo` (y1 <.> x1) + (z1 <.> x1)
    specify "log and exp higher" $ do
      log (exp x1) `shouldNormalizeTo` x1
      exp (log x1) `shouldNormalizeTo` x1
      log (exp x2) `shouldNormalizeTo` x2
      exp (log x2) `shouldNormalizeTo` x2
    specify "rotate rules" $ do
      rotate 1 (rotate 2 x1) `shouldNormalizeTo` rotate 3 x1
      rotate 2 (rotate 3 x1) `shouldNormalizeTo` rotate 5 x1
      rotate (1, 1) (rotate (2, 3) x2)
        `shouldNormalizeTo` rotate (3, 4) x2
      rotate (0, 1) (rotate (2, -3) x2)
        `shouldNormalizeTo` rotate (2, -2) x2
      rotate (0, 0, 0) x3 `shouldNormalizeTo` x3
      rotate (2, 3, 4) (s *. y3) `shouldNormalizeTo` s
        *. rotate (2, 3, 4) y3
      rotate (2, 3) (sum [x2, y2, z2])
        `shouldNormalizeTo` sum (map (rotate (2, 3)) [x2, y2, z2])
      rotate (2, 3) (product [x2, y2, z2])
        `shouldNormalizeTo` product (map (rotate (2, 3)) [x2, y2, z2])
  describe "Normalize piecewise" $ specify "some piecewise rules" $ do
    piecewise [1] c [x, y] `shouldNormalizeTo` piecewise [1] c [x, zero] + piecewise [1] c [zero, y]
    piecewise [1] c [x + z, y] `shouldNormalizeTo` piecewise [1] c [x, zero] + piecewise [1] c [z, zero] + piecewise [1] c [zero, y]
    piecewise [1] c [x *. y, zero] `shouldNormalizeTo` x *. piecewise [1] c [y, zero]
    piecewise [1, 2] c [zero, x *. y, zero] `shouldNormalizeTo` x *. piecewise [1, 2] c [zero, y, zero]
    piecewise [4, 5] x1 [zero1, x1 * (s *. y1), zero1] `shouldNormalizeTo` s *. piecewise [4, 5] x1 [zero1, x1 * y1, zero1]
    piecewise [1] c1 [c1 +: y1, z1 +: t1] `shouldNormalizeTo` piecewise [1] c1 [c1, z1] +: piecewise [1] c1 [y1, t1]
  describe "Fourier transform" $ specify "some Ft rules" $ do
    reFT (reFT x1) + imFT (imFT x1) `shouldNormalizeTo` constant (fromIntegral defaultDim1D) *. x1
    reFT (reFT x2) + imFT (imFT x2) `shouldNormalizeTo` constant (fromIntegral (default1stDim2D * default2ndDim2D)) *. x2
    let x = variable1D @10 "x"
        y = variable1D @10 "y"
        z = variable1D @10 "z"
        t = variable1D @10 "t"
        valMap =
          Map.fromList
            [ ("x", V1D $ listArray (0, 9) [1 ..]),
              ("y", V1D $ listArray (0, 9) [2, 5 ..]),
              ("z", V1D $ listArray (0, 9) [3, 8 ..]),
              ("t", V1D $ listArray (0, 9) [0, -1 ..])
            ]
    let res1 = eval valMap $ reFT (reFT (x * y + z * t))
        res2 = eval valMap . normalize $ reFT (reFT (x * y + z * t))
    res1 `shouldApprox` res2
    let res1 = eval valMap $ imFT (imFT (x * y + z * t))
        res2 = eval valMap . normalize $ imFT (imFT (x * y + z * t))
    res1 `shouldApprox` res2
    let res1 = eval valMap $ reFT (reFT (x * y + z * t)) + imFT (imFT (x * y + z * t))
        res2 = eval valMap . normalize $ reFT (reFT (x * y + z * t)) + imFT (imFT (x * y + z * t))
    res1 `shouldApprox` res2
  describe "Unit tests" $ specify "normalize scalar one zero" $ do
    x `shouldNormalizeTo` x
    constant 1 / x `shouldNormalizeTo` x ^ (-1)
    x + x `shouldNormalizeTo` constant 2 *. x
    x - x `shouldNormalizeTo` constant 0
    x * x `shouldNormalizeTo` x ^ 2
    x / x `shouldNormalizeTo` constant 1
    x + y `shouldNormalizeTo` x + y
    x - y `shouldNormalizeTo` x - y
    x * y `shouldNormalizeTo` x * y
    x / y `shouldNormalizeTo` x * y ^ (-1)
    (x + y) * (x + y) `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x + y) ^ 2 `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x - y) ^ 2 `shouldNormalizeTo` (constant (-2.0) *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x - y) * (x - y) `shouldNormalizeTo` (constant (-2.0) *. (x * y)) + (x ^ 2) + (y ^ 2)
    (x * y) ^ 2 `shouldNormalizeTo` (x ^ 2) * (y ^ 2)
    (x * y) * (x * y) `shouldNormalizeTo` (x ^ 2) * (y ^ 2)
    (x / y) * (x / y) `shouldNormalizeTo` (y ^ (-2)) * (x ^ 2)
    (x / y) ^ 2 `shouldNormalizeTo` (y ^ (-2)) * (x ^ 2)
    (constant 1 / x) * (constant 1 / x) `shouldNormalizeTo` x ^ (-2)
    (constant 1 / x) ^ 2 `shouldNormalizeTo` x ^ (-2)
    x * x `shouldNormalizeTo` x ^ 2
    (x * x) * x `shouldNormalizeTo` x ^ 3
    constant 1 / x `shouldNormalizeTo` x ^ (-1)
    x * x / x `shouldNormalizeTo` normalize x
    x / x `shouldNormalizeTo` constant 1
    x / x / x `shouldNormalizeTo` x ^ (-1)
    x / x / y `shouldNormalizeTo` y ^ (-1)
    (x * y) ^ 3 `shouldNormalizeTo` (x ^ 3) * (y ^ 3)
    ((x * y) ^ 3) * x / y `shouldNormalizeTo` (y ^ 2) * (x ^ 4)
    ((x * y) ^ 3) * x + y `shouldNormalizeTo` y + ((y ^ 3) * (x ^ 4))
    ((x * y) ^ 3) * (x + y) `shouldNormalizeTo` ((y ^ 3) * (x ^ 4)) + ((x ^ 3) * (y ^ 4))
    x / (x ^ 3) `shouldNormalizeTo` x ^ (-2)
    (x ^ 2) ^ 2 `shouldNormalizeTo` x ^ 4
    (x + y) * (x + y) `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    ((x + y) * (x + y)) * (x + y) `shouldNormalizeTo` (constant 3.0 *. (y * (x ^ 2))) + (constant 3.0 *. (x * (y ^ 2))) + (x ^ 3) + (y ^ 3)
    ((x * y) ^ 3) * (x + y) ^ 2 `shouldNormalizeTo` ((y ^ 3) * (x ^ 5)) + ((x ^ 3) * (y ^ 5)) + (constant 2.0 *. ((x ^ 4) * (y ^ 4)))
    (x + y) ^ 2 `shouldNormalizeTo` (constant 2.0 *. (x * y)) + (x ^ 2) + (y ^ 2)
    x / (x ^ 2) `shouldNormalizeTo` x ^ (-1)
    x / (x * x) `shouldNormalizeTo` x ^ (-1)
    x / x * x `shouldNormalizeTo` normalize x
    x / (x ^ 3) `shouldNormalizeTo` x ^ (-2)
    (x * x) / y `shouldNormalizeTo` (y ^ (-1)) * (x ^ 2)
    (x ^ 2) ^ 3 `shouldNormalizeTo` x ^ 6
    ((x ^ 3) ^ 3) ^ 2 `shouldNormalizeTo` x ^ 18
    y * (y ^ 2) ^ 2 `shouldNormalizeTo` y ^ 5
    x / ((y ^ 2) ^ 3) `shouldNormalizeTo` x * y ^ (-6)
    ((x * y) ^ 2) / x `shouldNormalizeTo` x * y ^ 2
    ((x + y) ^ 2) * ((x - y) ^ 2) `shouldNormalizeTo` (constant (-2.0) *. ((x ^ 2) * (y ^ 2))) + (x ^ 4) + (y ^ 4)
    ((x + y) ^ 2) + ((x - y) ^ 2) `shouldNormalizeTo` (constant 2.0 *. (x ^ 2)) + (constant 2.0 *. (y ^ 2))
    ((x + y) ^ 2) - ((x - y) ^ 2) `shouldNormalizeTo` constant 4.0 *. (x * y)
    (x + y) * (x - y) `shouldNormalizeTo` (constant (-1.0) *. (y ^ 2)) + (x ^ 2)
    (x + y) + (x - y) `shouldNormalizeTo` constant 2.0 *. x
    (x + y) - (x - y) `shouldNormalizeTo` constant 2.0 *. y
    (x + y) / (x - y) `shouldNormalizeTo` (x * ((x + (constant (-1.0) *. y)) ^ (-1))) + (y * ((x + (constant (-1.0) *. y)) ^ (-1)))
    x <.> y `shouldNormalizeTo` x * y
    x <.> constant 1 `shouldNormalizeTo` normalize x
    x <.> constant (-1) `shouldNormalizeTo` constant (-1) * x
    x <.> constant 0 `shouldNormalizeTo` constant 0
    x <.> x `shouldNormalizeTo` normalize x ^ 2
    x <.> x * y `shouldNormalizeTo` y * (x ^ 2)
    x <.> x + x `shouldNormalizeTo` x + (x ^ 2)
    x <.> (x + x) `shouldNormalizeTo` constant 2.0 *. (x ^ 2)
    x <.> (constant 1 / x) `shouldNormalizeTo` constant 1
    x <.> (y / x) `shouldNormalizeTo` normalize y
    (x <.> (x + y)) ^ 2 `shouldNormalizeTo` ((x ^ 2) * (y ^ 2)) + (constant 2.0 *. (y * (x ^ 3))) + (x ^ 4)

max_stars_repo_name: Turboscient/HashedExpression

max_issues_repo_name: Turboscient/HashedExpression

max_forks_repo_name: Turboscient/HashedExpression