text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

module NQS.Internal.Types
  ( -- * Numeric types
    ℂ
  , ℝ

  -- , Estimate(..)
  -- , NormalError(..)
  -- , EnergyMeasurement(..)
  , RealOf
  -- , SRMeasurement(..)
  {-
  , withRbm
  , withRbmPure
  , withMRbm
  , withMRbmPure
  -}
    -- * Vectors and matrices
  , DenseVector
  , DenseMatrix
  , MDenseVector
  , MDenseMatrix
  , slice
  , withForeignPtrPrim
  , newVectorAligned
  , newDenseVector
  , newDenseMatrix

  , unsafeRow
  , unsafeColumn
  , unsafeWriteVector

  , Mutable(..)
  , Variant(..)
  , Orientation(..)
  , Transpose(..)
  , MatUpLo(..)
  , Sing(..)

  , CheckValid(..)
  , isValidVector
  , isValidMatrix
  , badVectorInfo
  , HasBuffer(..)
  , HasStride(..)
  , HasDim(..)

  , orientationOf
  , ToNative(..)

  , DenseWorkspace(..)
  , MCConfig(..)
  , defaultMCConfig
  , CGConfig(..)
  , SRConfig(..)

  , HasMc(..)
  , HasCg(..)
  , HasSteps(..)
  , HasRuns(..)
  , HasRate(..)
  , HasRestarts(..)
  , HasRegulariser(..)
  , HasMaxIter(..)
  , HasMagnetisation(..)

  , mapVectorM
  , mapMatrixM
  , zipWithVectorM
  , zipWithMatrixM
  -- , asTuple
  -- , fromTuple
  , FreezeThaw(..)

  , HasMean(..)
  , HasVar(..)
  -- , HasEnergy(..)
  -- , HasForce(..)
  -- , HasDerivatives(..)
  ) where

import Foreign.Storable
import Foreign.ForeignPtr

import Control.DeepSeq
import Control.Exception (assert)
import Control.Monad ((>=>), unless)
import Control.Monad.Primitive
import Control.Monad.ST.Strict
import Data.Complex
import Debug.Trace
import GHC.Generics (Generic)
import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes)

import Data.Semigroup ((<>))
import Lens.Micro
import Lens.Micro.TH

import System.IO.Unsafe
import Data.Coerce
import Data.Bits

import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.ForeignPtr.Unsafe

import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import Data.Vector.Storable (Vector, MVector(..))

import qualified Data.Vector.Generic as GV
import qualified Data.Vector as Boxed
import Data.Aeson
import qualified Data.Aeson.Types as Aeson

import Data.Singletons

import qualified Numerical.HBLAS.MatrixTypes   as HBLAS
import           Numerical.HBLAS.MatrixTypes    ( Variant(..)
                                                , Orientation(..)
                                                , Transpose(..)
                                                , MatUpLo(..)
                                                )
import           Numerical.HBLAS.BLAS.FFI       ( CBLAS_UPLOT(..) )

import Data.Kind
import Data.Proxy

deriving instance Generic Variant
deriving instance NFData Variant

deriving instance Generic Orientation
deriving instance NFData Orientation

class ToNative a b where
  encode :: a -> b

data family Mutable (v :: * -> *) :: * -> * -> *

-- | After some prototyping, I've concluded (possibly incorrectly) that a
-- single-precision 'Float' is enough for our purposes.
--
-- [This answer](https://stackoverflow.com/a/40538415) also mentions that 32-bit
-- floats should be enought for most neural network applications.
type ℝ = Float
type ℂ = Complex ℝ

data instance Sing (orient :: Orientation) where
  SRow :: Sing 'Row
  SColumn :: Sing 'Column

instance SingI 'Row where sing = SRow
instance SingI 'Column where sing = SColumn

instance SingKind Orientation where
  type Demote Orientation = Orientation
  fromSing SRow = Row
  fromSing SColumn = Column
  toSing x = case x of
    Row -> SomeSing SRow
    Column -> SomeSing SColumn

data DenseVector (variant :: Variant) (a :: *) = DenseVector
  { _denseVectorDim    :: {-# UNPACK #-}!Int
  , _denseVectorStride :: {-# UNPACK #-}!Int
  , _denseVectorBuffer :: {-# UNPACK #-}!(Vector a)
  } deriving (Show, Generic, NFData)

data DenseMatrix (orientation :: Orientation) (a :: *) = DenseMatrix
  { _denseMatrixRows   :: {-# UNPACK #-}!Int
  , _denseMatrixCols   :: {-# UNPACK #-}!Int
  , _denseMatrixStride :: {-# UNPACK #-}!Int
  , _denseMatrixBuffer :: {-# UNPACK #-}!(Vector a)
  } deriving (Show, Generic, NFData)

data instance Mutable (DenseVector variant) s a = MDenseVector
  { _mDenseVectorDim :: {-# UNPACK #-}!Int
  , _mDenseVectorStride :: {-# UNPACK #-}!Int
  , _mDenseVectorBuffer :: {-# UNPACK #-}!(MVector s a)
  } deriving (Generic, NFData)

type MDenseVector variant = Mutable (DenseVector variant)

data instance Mutable (DenseMatrix orientation) s a = MDenseMatrix
  { _mDenseMatrixRows :: {-# UNPACK #-}!Int
  , _mDenseMatrixCols :: {-# UNPACK #-}!Int
  , _mDenseMatrixStride :: {-# UNPACK #-}!Int
  , _mDenseMatrixBuffer :: {-# UNPACK #-}!(MVector s a)
  } deriving (Generic, NFData)

type MDenseMatrix orientation = Mutable (DenseMatrix orientation)

makeFields ''DenseVector
makeFields ''DenseMatrix

unsafeIndexVector :: Storable a => DenseVector 'Direct a -> Int -> a
{-# INLINE unsafeIndexVector #-}
unsafeIndexVector !(DenseVector _ stride buff) !i = V.unsafeIndex buff (i * stride)

unsafeIndexMatrix :: forall a orient. (Storable a, SingI orient)
                  => DenseMatrix orient a -> (Int, Int) -> a
{-# INLINE unsafeIndexMatrix #-}
unsafeIndexMatrix !(DenseMatrix _ _ stride buff) !(r, c) = V.unsafeIndex buff i
  where i = case (sing :: Sing orient) of
              SRow -> r * stride + c
              SColumn -> r + c * stride

unsafeReadVector :: (Storable a, PrimMonad m)
                 => MDenseVector 'Direct (PrimState m) a -> Int -> m a
{-# INLINE unsafeReadVector #-}
unsafeReadVector !(MDenseVector _ stride buff) !i = MV.unsafeRead buff (i * stride)

unsafeReadMatrix :: forall a m orient. (Storable a, PrimMonad m, SingI orient)
                 => MDenseMatrix orient (PrimState m) a -> (Int, Int) -> m a
{-# INLINE unsafeReadMatrix #-}
unsafeReadMatrix !(MDenseMatrix _ _ stride buff) !(r, c) = MV.unsafeRead buff i
  where i = case (sing :: Sing orient) of
              SRow -> r * stride + c
              SColumn -> r + c * stride

unsafeWriteVector :: (Storable a, PrimMonad m)
                  => MDenseVector 'Direct (PrimState m) a -> Int -> a -> m ()
{-# INLINE unsafeWriteVector #-}
unsafeWriteVector !(MDenseVector _ stride buff) !i !x = MV.unsafeWrite buff (i * stride) x

unsafeWriteMatrix :: forall a m orient. (Storable a, PrimMonad m, SingI orient)
                  => MDenseMatrix orient (PrimState m) a -> (Int, Int) -> a -> m ()
{-# INLINE unsafeWriteMatrix #-}
unsafeWriteMatrix !(MDenseMatrix _ _ stride buff) !(r, c) !x = MV.unsafeWrite buff i x
  where i = case (sing :: Sing orient) of
              SRow -> r * stride + c
              SColumn -> r + c * stride

-- | Returns a row of a matrix.
unsafeRow
  :: forall orient s a
   . (Storable a, SingI orient)
  => Int
  -> MDenseMatrix orient s a
  -> MDenseVector 'Direct s a
{-# INLINE unsafeRow #-}
unsafeRow !i !(MDenseMatrix rows cols stride buff) = assert False $
  assert (i < rows) $ case (sing :: Sing orient) of
    SRow    -> MDenseVector cols 1 (MV.slice (i * stride) cols buff)
    SColumn -> MDenseVector cols stride (MV.slice i (cols * stride) buff)

-- | Returns a column of a matrix.
unsafeColumn
  :: forall orient s a
   . (Storable a, SingI orient)
  => Int
  -> MDenseMatrix orient s a
  -> MDenseVector 'Direct s a
{-# INLINE unsafeColumn #-}
unsafeColumn !i !(MDenseMatrix rows cols stride buff) =
  assert (i < cols) $ case (sing :: Sing orient) of
    SRow    -> if rows == 1
                 then MDenseVector 1 1 (MV.slice i 1 buff)
                 else MDenseVector rows stride (MV.slice i ((rows - 1)* stride + 1) buff)
    SColumn -> MDenseVector rows 1 (MV.slice (i * stride) rows buff)

data DenseWorkspace s a = DenseWorkspace
  { _denseWorkspaceForce       :: !(MDenseVector 'Direct s a) -- ^ Force
  , _denseWorkspaceDerivatives :: !(MDenseMatrix 'Row s a)    -- ^ Derivatives
  , _denseWorkspaceDelta       :: !(MDenseVector 'Direct s a) -- ^ Old delta
  } deriving (Generic, NFData)

-- | Configuration for Monte-Carlo sampling.
data MCConfig =
  MCConfig
    { _mCConfigSteps :: {-# UNPACK #-}!(Int, Int, Int)
        -- ^ A range of steps for a single run. It is very similar to Python's
        -- [range](). @(low, high, step)@ corresponds to [low, low + step, low +
        -- 2 * step, ..., high).
        --
        -- /Note:/ only positive @step@s are supported.
    , _mCConfigThreads :: {-# UNPACK #-}!(Int, Int, Int)
    , _mCConfigRuns :: {-# UNPACK #-}!Int
        -- ^ Number of Monte-Carlo runs to perform.
        --
        -- /Note:/ for optimal work scheduling, make that the total number of
        -- threads is divisible by the number of runs.
    , _mCConfigFlips :: {-# UNPACK #-}!Int
        -- ^ Number of spin-flips to do at each step.
    , _mCConfigRestarts :: {-# UNPACK #-}!Int
        -- ^ Allowed number of restarts.
        --
        -- A restart happens when the function computing local energy notices
        -- that a particular spin-flip results in a spin configuration with
        -- significantly higher probability. Monte-Carlo sampler is then reset
        -- and this new spin configuration is used as the initial one.
    , _mCConfigMagnetisation :: !(Maybe Int)
        -- ^ Specifies the magnetisation over which to sample.
    }

makeFields ''MCConfig

defaultMCConfig :: MCConfig
defaultMCConfig = MCConfig (1000, 11000, 1) (4, 1, 1) 4 2 5 Nothing

-- | Configuration for the Conjugate Gradient solver.
data CGConfig a =
  CGConfig { _cGConfigMaxIter :: {-# UNPACK #-}!Int
               -- ^ Maximum number of iterations.
           , _cGConfigTol :: !a
               -- ^ Tolerance.
           }

makeFields ''CGConfig

-- | Configuration for the Stochastic Reconfiguration algorithm.
data SRConfig a =
  SRConfig
    { _sRConfigMaxIter :: !Int
        -- ^ Maximum number of iterations to perform.
    , _sRConfigRegulariser :: !(Maybe (Int -> ℂ))
        -- ^ Regulariser for the S matrix. If it is @'Just' f@,
        -- then @λ = f i@ where @i@ is the currect iteration is used to
        -- regularise the matrix S according to @S <- S + λ1@.
    , _sRConfigRate :: !(Int -> ℂ)
        -- ^ Learning rate as a function of the iteration.
    , _sRConfigCg :: !(CGConfig ℝ)
        -- ^ Configuration for the Monte-Carlo sampling.
    , _sRConfigMc :: !MCConfig
        -- ^ Configuration for the Conjugate Gradient solver.
    }

-- | Returns a slice of the vector.
slice
  :: forall a s
   . Storable a
  => Int -- ^ Start index
  -> Int -- ^ Length
  -> MDenseVector 'Direct s a -- ^ Source vector
  -> MDenseVector 'Direct s a -- ^ Slice
slice !i !n !(MDenseVector size stride buff) =
  MDenseVector n stride (MV.slice (stride * i) (stride * n) buff)

makeFields ''SRConfig

type family RealOf a :: *

type instance RealOf Float = Float
type instance RealOf Double = Double
type instance RealOf (Complex a) = a

mallocVectorAligned :: forall a. Storable a => Int -> Int -> IO (ForeignPtr a)
mallocVectorAligned n alignment =
  mallocPlainForeignPtrAlignedBytes (n * sizeOf (undefined :: a)) alignment

newVectorAligned
  :: (Storable a, PrimMonad m) => Int -> Int -> m (MVector (PrimState m) a)
newVectorAligned n alignment =
  unsafePrimToPrim $! MVector n <$> mallocVectorAligned n alignment

newDenseVector
  :: (Storable a, PrimMonad m) => Int -> m (MDenseVector 'Direct (PrimState m) a)
newDenseVector n = MDenseVector n 1 <$> newVectorAligned n 64

-- | Default alignment used when allocating new vectors and matrices.
defaultAlignment :: Int
defaultAlignment = 64

roundUpTo :: Int -> Int -> Int
{-# INLINE roundUpTo #-}
roundUpTo !alignment !n = assert (isValidAlignment alignment && n >= 0) $
  (n + alignment - 1) .&. complement (alignment - 1)
  where isValidAlignment !x = x > 0 && (x .&. (x - 1) == 0)

newDenseMatrix
  :: forall orient a m. (Storable a, SingI orient, PrimMonad m)
  => Int
  -> Int
  -> m (MDenseMatrix orient (PrimState m) a)
newDenseMatrix rows cols = case (sing :: Sing orient) of
  SRow ->
    let ldim = assert (defaultAlignment `mod` sizeOf (undefined :: a) == 0) $
                roundUpTo (defaultAlignment `div` sizeOf (undefined :: a)) cols
    in  MDenseMatrix rows cols ldim
          <$> newVectorAligned (rows * ldim) defaultAlignment
  SColumn ->
    let ldim = roundUpTo defaultAlignment rows
    in  MDenseMatrix rows cols ldim
          <$> newVectorAligned (ldim * cols) defaultAlignment

class HasOrientation s a | s -> a where
  orientationOf :: s -> a

instance SingI orient
  => HasOrientation (DenseMatrix orient a) Orientation where
    orientationOf _ = fromSing (sing :: Sing orient)

instance SingI orient
  => HasOrientation (Mutable (DenseMatrix orient) s a) Orientation where
    orientationOf _ = fromSing (sing :: Sing orient)

isValidMatrix :: Orientation -- ^ Row- vs. column-major layout
              -> Int -- ^ Number of rows
              -> Int -- ^ Number of columns
              -> Int -- ^ Stride
              -> Int -- ^ Buffer size
              -> Bool
isValidMatrix orient rows cols i size =
  i >= 0 && (rows == 0 && cols == 0 || rows > 0 && cols > 0 && validAccesses orient)
  where validAccesses Row    = (rows - 1) * i < size && cols <= i
        validAccesses Column = (cols - 1) * i < size && rows <= i

-- | Constructs a nice message describing the problem in the BLAS vector.
badMatrixInfo :: String -- ^ Function name
              -> String -- ^ Argument name
              -> Int -- ^ Number of rows
              -> Int -- ^ Number of columns
              -> Int -- ^ Stride
              -> Int -- ^ Size of the underlying buffer
              -> String -- ^ Error message
badMatrixInfo funcName argName rows cols i size
  | rows < 0 = preamble <> " has a negative number of rows: " <> show rows <> "."
  | cols < 0 = preamble <> " has a negative number of columns: " <> show cols <> "."
  | i < 0 = preamble <> " has a negative stride: " <> show i <> "."
  | otherwise = preamble <> " has invalid range of accesses: #rows = " <> show rows <>
    ", #cols = " <> show cols <> ", stride = " <> show i <> ", bufferSize = " <> show size <> "."
  where preamble = funcName <> ": " <> argName

-- | Returns whether strides and dimensions are consistent.
isValidVector :: Int -- ^ Logical dimension
              -> Int -- ^ Stride
              -> Int -- ^ Buffer size
              -> Bool
isValidVector n i size = i >= 0 && (n == 0 || n > 0 && (n - 1) * i < size)

-- | Constructs a nice message describing the problem in the BLAS vector.
badVectorInfo :: String -- ^ Function name
              -> String -- ^ Argument name
              -> Int -- ^ Logical vector dimension
              -> Int -- ^ Vector stride
              -> Int -- ^ Size of the underlying buffer
              -> String -- ^ Error message
badVectorInfo funcName argName n i size
  | n < 0 = preamble <> " has a negative logical dimension: " <> show n <> "."
  | i < 0 = preamble <> " has a negative stride: " <> show i <> "."
  | otherwise = preamble <> " has invalid range of accesses: dim = " <> show n <>
                ", stride = " <> show i <> ", bufferSize = " <> show size <> "."
  where preamble = funcName <> ": " <> argName

class CheckValid a where
  assertValid :: String -> String -> a -> b -> b

instance Storable a => CheckValid (DenseVector variant a) where
  assertValid funcName argName !(DenseVector dim stride (V.length -> size))
    | isValidVector dim stride size = id
    | otherwise = error $! badVectorInfo funcName argName dim stride size

instance Storable a => CheckValid (Mutable (DenseVector variant) s a) where
  assertValid funcName argName !(MDenseVector dim stride (MV.length -> size))
    | isValidVector dim stride size = id
    | otherwise = error $! badVectorInfo funcName argName dim stride size

instance (Storable a, SingI orient) => CheckValid (DenseMatrix orient a) where
  assertValid funcName argName !m@(DenseMatrix rows cols stride (V.length -> size))
    | isValidMatrix (orientationOf m) rows cols stride size = id
    | otherwise = error $! badMatrixInfo funcName argName rows cols stride size

instance (Storable a, SingI orient) => CheckValid (Mutable (DenseMatrix orient) s a) where
  assertValid funcName argName !m@(MDenseMatrix rows cols stride (MV.length -> size))
    | isValidMatrix (orientationOf m) rows cols stride size = id
    | otherwise = error $! badMatrixInfo funcName argName rows cols stride size

instance Storable a => HasBuffer (MDenseVector variant s a) (MVector s a) where
  buffer inj (MDenseVector dim stride buf) = MDenseVector dim stride <$> inj buf

instance Storable a => HasBuffer (MDenseMatrix orientation s a) (MVector s a) where
  buffer inj (MDenseMatrix rows cols stride buf) = MDenseMatrix rows cols stride <$> inj buf

instance Storable a => HasStride (MDenseVector variant s a) Int where
  stride inj (MDenseVector dim stride buf) = (\x -> MDenseVector dim x buf) <$> inj stride

instance Storable a => HasStride (MDenseMatrix orientation s a) Int where
  stride inj (MDenseMatrix rows cols stride buf) = (\x -> MDenseMatrix rows cols x buf) <$> inj stride

instance Storable a => HasDim (MDenseVector variant s a) Int where
  dim inj (MDenseVector dim stride buf) = (\x -> MDenseVector x stride buf) <$> inj dim

instance Storable a => HasDim (MDenseMatrix orientation s a) (Int, Int) where
  dim inj (MDenseMatrix rows cols stride buf) =
    (\(x, y) -> MDenseMatrix x y stride buf) <$> inj (rows, cols)

instance Storable a => HasDim (DenseMatrix orientation a) (Int, Int) where
  dim inj (DenseMatrix rows cols stride buf) =
    (\(x, y) -> DenseMatrix x y stride buf) <$> inj (rows, cols)

class FreezeThaw (v :: * -> *) a where
  unsafeFreeze :: PrimMonad m => (Mutable v) (PrimState m) a -> m (v a)
  unsafeThaw :: PrimMonad m => v a -> m ((Mutable v) (PrimState m) a)

instance Storable a => FreezeThaw (DenseVector variant) a where
  unsafeFreeze (MDenseVector dim stride mv) = DenseVector dim stride <$> V.unsafeFreeze mv
  unsafeThaw (DenseVector dim stride mv) = MDenseVector dim stride <$> V.unsafeThaw mv

instance Storable a => FreezeThaw (DenseMatrix orientation) a where
  unsafeFreeze (MDenseMatrix rows cols stride mv) = DenseMatrix rows cols stride <$> V.unsafeFreeze mv
  unsafeThaw (DenseMatrix rows cols stride v) = MDenseMatrix rows cols stride <$> V.unsafeThaw v

touchForeignPtrPrim :: PrimMonad m => ForeignPtr a -> m ()
{-# NOINLINE touchForeignPtrPrim #-}
touchForeignPtrPrim fp = unsafeIOToPrim $! touchForeignPtr fp

withForeignPtrPrim :: PrimMonad m => ForeignPtr a -> (Ptr a -> m b) -> m b
{-# INLINE withForeignPtrPrim #-}
withForeignPtrPrim p func = do r <- func (unsafeForeignPtrToPtr p)
                               touchForeignPtrPrim p
                               return r

mapVectorM ::
     (PrimMonad m, Storable a, Storable b)
  => (a -> m b)
  -> MDenseVector 'Direct (PrimState m) a
  -> MDenseVector 'Direct (PrimState m) b
  -> m ()
mapVectorM f x y =
  assertValid "NQS.Rbm.zipWithVectorM" "x" x $
  assertValid "NQS.Rbm.zipWithVectorM" "x" y $
  assert (x ^. dim == y ^. dim) $ go 0
  where
    n = x ^. dim
    go !i
      | i == n = return ()
      | otherwise = do
        xi <- unsafeReadVector x i
        yi <- f xi
        unsafeWriteVector y i yi
        go (i + 1)

zipWithVectorM ::
     (PrimMonad m, Storable a, Storable b, Storable c)
  => (a -> b -> m c)
  -> MDenseVector 'Direct (PrimState m) a
  -> MDenseVector 'Direct (PrimState m) b
  -> MDenseVector 'Direct (PrimState m) c
  -> m ()
zipWithVectorM f x y z =
  assertValid "NQS.Rbm.zipWithVectorM" "x" x $
  assertValid "NQS.Rbm.zipWithVectorM" "x" y $
  assertValid "NQS.Rbm.zipWithVectorM" "x" z $
  assert (y ^. dim == n && z ^. dim == n) $ go 0
  where
    n = x ^. dim
    go !i
      | i == n = return ()
      | otherwise = do
        xi <- unsafeReadVector x i
        yi <- unsafeReadVector y i
        zi <- f xi yi
        unsafeWriteVector z i zi
        go (i + 1)

mapMatrixM ::
     forall m a b orientX orientY.
     (PrimMonad m, Storable a, Storable b, SingI orientX, SingI orientY)
  => (a -> m b)
  -> MDenseMatrix orientX (PrimState m) a
  -> MDenseMatrix orientY (PrimState m) b
  -> m ()
mapMatrixM f x y =
  assertValid "NQS.Rbm.zipWithMatrixM" "x" x $
  assertValid "NQS.Rbm.zipWithMatrixM" "y" y $
  assert (x ^. dim == y ^. dim) $
  case (sing :: Sing orientY) of
    SRow -> stepperRow 0 0
    SColumn -> stepperColumn 0 0
  where
    !(n, m) = x ^. dim
    stepperRow !i !j
      | j < m && i < n = go i j >> stepperRow i (j + 1)
      | i < n = stepperRow (i + 1) 0
      | otherwise = return ()
    stepperColumn !i !j
      | i < n && j < m = go i j >> stepperColumn (i + 1) j
      | j < m = stepperColumn 0 (j + 1)
      | otherwise = return ()
    go !i !j = do
      xij <- unsafeReadMatrix x (i, j)
      yij <- f xij
      unsafeWriteMatrix y (i, j) yij

zipWithMatrixM :: forall m a b c orientX orientY orientZ.
     (PrimMonad m, Storable a, Storable b, Storable c, SingI orientX, SingI orientY, SingI orientZ)
  => (a -> b -> m c)
  -> MDenseMatrix orientX (PrimState m) a
  -> MDenseMatrix orientY (PrimState m) b
  -> MDenseMatrix orientZ (PrimState m) c
  -> m ()
zipWithMatrixM f x y z =
  assertValid "NQS.Rbm.zipWithMatrixM" "x" x $
  assertValid "NQS.Rbm.zipWithMatrixM" "y" y $
  assertValid "NQS.Rbm.zipWithMatrixM" "z" z $
  assert (y ^. dim == (n, m) && z ^. dim == (n, m)) $
    case (sing :: Sing orientZ) of
      SRow -> stepperRow 0 0
      SColumn -> stepperColumn 0 0
  where
    !(n, m) = x ^. dim
    stepperRow !i !j
      | j < m && i < n = go i j >> stepperRow i (j + 1)
      | i < n = stepperRow (i + 1) 0
      | otherwise = return ()
    stepperColumn !i !j
      | i < n && j < m = go i j >> stepperColumn (i + 1) j
      | j < m = stepperColumn 0 (j + 1)
      | otherwise = return ()
    go !i !j = do
      xij <- unsafeReadMatrix x (i, j)
      yij <- unsafeReadMatrix y (i, j)
      zij <- f xij yij
      unsafeWriteMatrix z (i, j) zij

{-
withRbm :: PrimMonad m => Rbm a -> (Ptr (RbmCore a) -> m b) -> m b
withRbm (Rbm p) func = withForeignPtrPrim p func

withRbmPure :: Rbm a -> (Ptr (RbmCore a) -> b) -> b
withRbmPure (Rbm p) func = unsafePerformIO $! withForeignPtrPrim p (return . func)

withMRbm :: PrimMonad m => MRbm (PrimState m) a -> (Ptr (RbmCore a) -> m b) -> m b
withMRbm (MRbm p) func = withForeignPtrPrim p func

withMRbmPure :: MRbm s a -> (Ptr (RbmCore a) -> b) -> b
withMRbmPure (MRbm p) func = unsafePerformIO $! withForeignPtr p (return . func)
-}

data Estimate e a = Estimate
  { estPoint :: !a
  , estError :: !(e a)
  } deriving (Generic, NFData)

newtype NormalError a = NormalError a
  deriving (Eq, Ord, Read, Show, Generic, NFData)

data EnergyMeasurement a = EnergyMeasurement
  { _energyMeasurementMean :: !a
  , _energyMeasurementVar :: !a
  } deriving (Generic)

makeFields ''EnergyMeasurement

-- | Simple for loop. Counts from /start/ to /end/-1.
for :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
for !start !end f = assert (start <= end) $ loop start
 where
  loop !i | i == end  = return ()
          | otherwise = f i >> loop (i + 1)
{-# INLINE for #-}

for' :: Monad m => a -> Int -> (a -> Int -> m a) -> m a
for' !x0 !n f = assert (n >= 0) $ loop x0 0
 where
  loop !x !i | i == n    = return x
             | otherwise = f x i >>= \x' -> loop x' (i + 1)
{-# INLINE for' #-}

generateVectorM
  :: forall a m
   . (Monad m, Storable a)
  => Int
  -> (Int -> m a)
  -> m (DenseVector 'Direct a)
{-# NOINLINE generateVectorM #-}
generateVectorM !n f
  | n < 0     = error $! "generateVectorM: Invalid dimension: " <> show n
  | n == 0    = return $! DenseVector 0 0 V.empty
  | otherwise = go (runST $ newDenseVector n >>= unsafeFreeze)
 where
  go :: DenseVector 'Direct a -> m (DenseVector 'Direct a)
  go !v' = for' v' n $ \v i -> f i >>= \x -> return (set v i x)
  set !v !i !x = runST $ do
    !v' <- unsafeThaw v
    unsafeWriteVector v' i x
    unsafeFreeze v'

generateMatrixM
  :: forall orient m a. (Monad m, Storable a, SingI orient)
  => Int
  -> Int
  -> (Int -> Int -> m a)
  -> m (DenseMatrix orient a)
{-# NOINLINE generateMatrixM #-}
generateMatrixM !rows !cols f
  | rows < 0 || cols < 0 = error $! "generateMatrixM: Invalid dimensions: " <> show (rows, cols)
  | rows == 0 || cols == 0 = return $! DenseMatrix 0 0 0 V.empty
  | otherwise = go (runST $ newDenseMatrix rows cols >>= unsafeFreeze)
 where
  go :: DenseMatrix orient a -> m (DenseMatrix orient a)
  go !m'' = case (sing :: Sing orient) of
    SRow -> for' m'' rows $ \m' i -> for' m' cols $ \m j ->
              f i j >>= \x -> return (set m i j x)
    SColumn -> for' m'' cols $ \m' j -> for' m' rows $ \m i ->
                 f i j >>= \x -> return (set m i j x)
  set !m !i !j !x = runST $ do
    !m' <- unsafeThaw m
    unsafeWriteMatrix m' (i, j) x
    unsafeFreeze m'

instance (Storable a, FromJSON a) => FromJSON (DenseVector 'Direct a) where
  parseJSON = withArray "vector elements (i.e. Array)" $ \v -> do
    generateVectorM (GV.length v) $ \i -> parseJSON (v `GV.unsafeIndex` i)

instance ToJSON a => ToJSON (Complex a) where
  toJSON (x :+ y) = toJSON [x, y]
  toEncoding (x :+ y) = toEncoding [x, y]

instance FromJSON a => FromJSON (Complex a) where
  parseJSON = withArray "Complex" $ \v ->
    (:+) <$> parseJSON (v GV.! 0)
         <*> parseJSON (v GV.! 1)

instance (Storable a, FromJSON a, SingI orient) => FromJSON (DenseMatrix orient a) where
  parseJSON = withArray "Matrix rows" $ \matrix -> do
    (n, m) <- checkDimensions matrix
    generateMatrixM n m $ \i j ->
      flip (withArray "a matrix row") (matrix `GV.unsafeIndex` i) $ \row ->
        parseJSON (row `GV.unsafeIndex` j)
    where
      checkDimensions x@(GV.length -> n)
        | n == 0 = return (0, 0)
        | otherwise = do
          !m <- withArray "first row (i.e. Array)" (return . GV.length) (GV.unsafeHead x)
          for 0 n $ \i ->
            flip (withArray "a matrix row (i.e. Array)") (x `GV.unsafeIndex` i) $ \row ->
              unless (GV.length row == m) $ fail $! "Matrix row #" <> show i <> " has wrong dimension."
          return (n, m)

vectorAsVector :: Storable a => DenseVector 'Direct a -> Boxed.Vector a
vectorAsVector !(DenseVector n stride buff) =
  Boxed.generate n (\i -> buff `V.unsafeIndex` (i * stride))

-- | Returns a row of a matrix.
unsafeRowAsVector
  :: forall orient a
   . (Storable a, SingI orient)
  => Int
  -> DenseMatrix orient a
  -> Boxed.Vector a
unsafeRowAsVector !i !(DenseMatrix rows cols stride buff) =
  assert (i < rows) $ case (sing :: Sing orient) of
    SRow    -> GV.generate cols (\j -> buff `V.unsafeIndex` (i * stride + j))
    SColumn -> GV.generate cols (\j -> buff `V.unsafeIndex` (i + j * stride))

instance (Storable a, ToJSON a) => ToJSON (DenseVector 'Direct a) where
  toEncoding = toEncoding . vectorAsVector

instance (Storable a, ToJSON a, SingI orient) => ToJSON (DenseMatrix orient a) where
  toEncoding x = toEncoding $ Boxed.generate (x ^. dim . _1) (\i -> unsafeRowAsVector i x)

max_stars_repo_name: twesterhout/tcm-swarm

max_issues_repo_name: twesterhout/tcm-swarm

max_forks_repo_name: twesterhout/tcm-swarm


text:
module Main where

import Control.Monad
import Data.Complex
import Data.Int
import Graphics.UI.GLUT
import qualified Streamly as St
import qualified Streamly.Prelude as S

--import qualified Streamly.Prelude as S
--import qualified Control.Monad.Parallel as MP
iterations = 400

x // y = fromIntegral x / fromIntegral y

-- Divides [a] into [[a], [a], ...] with each sublist of length n,
-- except the last sublist which has length <= n.
chunkify n [] = []
chunkify n xs =
  let (xs', rest) = splitAt n xs
  in xs' : chunkify n rest

-- Converts a coordinate in screen space to a vertex.
pix2vert (Size w h) (x, y) =
  Vertex2 ((3 // w * fromIntegral x) - 2.0) ((2 // h * fromIntegral y) - 1.0)

-- List of all of the vertices that represent screen pixels.
vertices :: IO [Vertex2 GLfloat]
vertices =
  get windowSize >>= \(Size w h) ->
    return $ [pix2vert (Size w h) (x, y) | x <- [0 .. w - 1], y <- [0 .. h - 1]]

-- Gets the color for a number of iterations.
color3 r g b = Color3 r g b

getcolor :: Int -> Color3 Float
getcolor iter
  | iter == iterations = color3 0 0 0
  | otherwise = color3 (amt * 0.5) amt (amt * 0.5)
  where
    amt = iter // iterations

-- Returns the number of iterations <= the maximum iterations of the
-- Mandelbrot set at the given vertex.
mandel (Vertex2 r i) =
  length . takeWhile (\z -> magnitude z <= 2) . take iterations $
  iterate (\z -> z ^ 2 + (r :+ i)) 0

-- plots one point.
drawVert v = do
  color . getcolor $ mandel v
  vertex v

-- draws all the vertices in slices (to update the display while drawing).
display' :: [[Vertex2 GLfloat]] -> IO ()
display' chunks =
  let foo = S.each chunks :: St.ParallelT IO ([Vertex2 GLfloat]) -- O(n)
      bar =
        (\vs -> do
           renderPrimitive Points $ mapM_ drawVert vs
           flush -- ::[Vertex2 GLfloat] -> IO ()
         )
  in do S.mapM_ bar foo
        displayCallback $= display

-- draws the whole fractal
display :: IO ()
display = do
  clear [ColorBuffer]
  displayCallback $=
    (do v <- vertices
        let c = chunkify 256 v
        let (x:y:z:_) = chunkify 3 c
                                    --display' c
        St.runStreamT $
          St.liftIO (display' x) St.<|> St.liftIO (display' y) St.<|>
          St.liftIO (display' z))
  get currentWindow >>= postRedisplay

main = do
  getArgsAndInitialize
  initialDisplayMode $= [SingleBuffered, RGBMode]
  initialWindowSize $= Size 1200 1024
  initialWindowPosition $= Position 100 100
  createWindow "Mandelbrot"
  clearColor $= Color4 0 0 0 0
  matrixMode $= Projection
  loadIdentity
  ortho (-2) 1 (-1) 1 (-1) 1
  displayCallback $= display
  mainLoop
-- import Control.Parallel.Strategies
-- type Point = (Float, Float)
-- -- z = x + i y ; p = u + iv
-- -- next p z = z^2 + p -- order of application reversed for currying
-- next :: Point -> Point -> Point
-- next (x,y) (u,v) = (x * x - y * y + u, 2 * x * y + v)
-- mandelbrot :: Point -> [Point]
-- mandelbrot p = iterate (next p) (0,0)
-- julia :: Point -> Point -> [Point]
-- julia c = iterate (next c)
-- fairlyClose :: Point -> Bool
-- fairlyClose (u,v) = (u * u + v * v) < 100
-- inMandelbrotSet :: Point -> Bool
-- inMandelbrotSet p = all fairlyClose (mandelbrot p)
-- chooseColor :: [a] -> [Point] -> a
-- chooseColor palette = (palette !!) . length . take n . takeWhile fairlyClose
--                        where n = length palette - 1
-- type Image a = Point -> a
-- fracImage :: (Point -> [Point]) -> [a] -> Image a
-- fracImage fractal palette = chooseColor palette . fractal
-- type Grid a = [[a]]
-- grid :: Int -> Int -> Point -> Point -> Grid Point
-- grid c r (xmin, ymin) (xmax, ymax)
--   = [[(x,y) | x <- for c xmin xmax] | y <- for r ymin ymax]
-- for :: Int -> Float -> Float -> [Float]
-- for n min max = take n [min, min + delta ..]
--                 where delta = (max - min) / fromIntegral (n - 1)
-- pmap = parMap rseq
-- sample :: Grid Point -> Image a -> Grid a
-- sample points image = map (map image) points
-- draw ::
--   Grid Point
--   -> (Point -> [Point]) -- mandelbrot
--   -> [a] -- [Char]
--   -> (Grid a -> b) -- the charRender IO function
--   -> b
-- draw points fractal palette render
--   = render (sample points (fracImage fractal palette))
-- charPalette :: [Char]
-- charPalette = " ,.‘\"~:;o-!|?/<>X+={^O#%&@8*$"
-- --putStr . unlines
-- charRender :: Grid Char -> IO ()
-- charRender x = return ()--putStr . unlines
-- main :: IO ()
-- main = draw points mandelbrot charPalette charRender
--           where points = grid (1790 * 300) (1370 * 300) (-2.25,-1.5) (0.75,1.5)
-- figure2 = draw points(julia (0.32,0.043)) charPalette charRender
--        where points = grid 79 37 (-1.5,-1.5) (1.5,1.5)

ext: hs

max_stars_repo_name: Abhiroop/streamlytest

max_issues_repo_name: Abhiroop/streamlytest

max_forks_repo_name: Abhiroop/streamlytest


text:
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
module Lib where

import           Config
import           Control.Monad.IO.Class    (liftIO)
import           Control.Monad.Logger      (runNoLoggingT)
import           Control.Monad.Trans
import           Data.Aeson                hiding (json)
import           Data.Complex
import           Data.IORef
import           Data.Monoid
import           Data.Text                 (pack)
import qualified Data.Text                 as T
import           Data.Time.Clock           (getCurrentTime)
import           Database.Persist.Sqlite   (SqlBackend, createSqlitePool,
                                            runMigration, runSqlPool)
import           Db                        (createToken, findToken, migrateAll,
                                            removeOldTokens, runSQL)
import           Eval                      (eval)
import           GHC.Generics              (Generic)
import           Network.HTTP.Types.Status
import           Text.Read                 (readMaybe)
import           Text.StringRandom         (stringRandomIO)
import           Web.Spock
import           Web.Spock.Config
import Services

start :: Config -> IO ()
start config@(Config db port) =
    do pool <- runNoLoggingT $ createSqlitePool db 5
       runNoLoggingT $ runSqlPool (runMigration migrateAll) pool
       spockCfg <- defaultSpockCfg () (PCPool pool) config
       runSpock port (spock spockCfg app)

app :: ApiApp ()
app =
    do post "login" postLogin
       post "compute" postCompute


text:
module STCLocalEigenVectorEndPoint where

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

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:histFileName:numIterationStr:writeSourceFlagStr:cutoffRadiusStr:patchNormFlagStr:patchNormSizeStr:shape2DStr:useFFTWWisdomFlagStr:fftwWisdomFileName:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      cutoffRadius = read cutoffRadiusStr :: Int
      patchNormFlag = read patchNormFlagStr :: Bool
      patchNormSize = read patchNormSizeStr :: Int
      shape2D@(Points _ minDist _) = read shape2DStr :: Points Shape2D
      useFFTWWisdomFlag = read useFFTWWisdomFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCLocalEigenVectorEndPoint"
      histFilePath = folderPath </> histFileName
      fftwWisdomFilePath = folderPath </> fftwWisdomFileName
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          thetaFreqs
          thetaFreqs
          scaleFreqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scaleFreqs
             , L.length thetaFreqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  let localEigenVector =
        computeLocalEigenVector
          (ParallelParams numThread 0)
          (pinwheelHollow 2)
          (cutoff cutoffRadius radialArr)
          numPoint
          numPoint
          maxScale
          thetaFreqs
          scaleFreqs
      pinwheelArr =
        -- computeS .
        R.traverse2
          (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
          (fromListUnboxed (Z :. (L.length scaleFreqs)) scaleFreqs)
          (\(Z :. numThetaFreq) (Z :. numScaleFreq) ->
             (Z :. numThetaFreq :. numScaleFreq :. numPoint :. numPoint)) $ \ft fs (Z :. k :. l :. i :. j) ->
          let r =
                sqrt . fromIntegral $
                (i - center numPoint) ^ 2 + (j - center numPoint) ^ 2
           in if r < fromIntegral cutoffRadius
                then 0
                else pinwheelHollow
                       1
                       (ft (Z :. k))
                       (fs (Z :. l))
                       maxScale
                       0
                       (i - center numPoint)
                       (j - center numPoint)
      -- localEigenVector = (computeS $
      --                     R.traverse2
      --                       localEigenVector'
      --                       (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
      --                       const $ \f fFreq idx@(Z :. k :. _ :. _ :. _) ->
      --                       f idx * exp (0 :+ fFreq (Z :. k) * pi))
      pointSet = makeShape2D shape2D
      shapeArr =
        getShape2DRepaArray
          numPoint
          numPoint
          (L.map
             (\(x, y) ->
                (x + fromIntegral numPoint / 2, y + fromIntegral numPoint / 2))
             pointSet)
      xs =
        L.map (\(x, y) -> R2S1RPPoint (x, y, 0, 1)) . getShape2DIndexList $
        pointSet
      bias = computeBiasR2T0S0 numPoint numPoint thetaFreqs scaleFreqs xs
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (pinwheelHollow 4)
      (cutoff cutoffRadius radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      thetaFreqs
      scaleFreqs
  plan' <-
    make4DPlan emptyPlan useFFTWWisdomFlag fftwWisdomFilePath localEigenVector
  plan <-
    makeR2Z2T0S0Plan plan' useFFTWWisdomFlag fftwWisdomFilePath arrR2Z2T0S0
  plotImageRepa (folderPath </> "Shape.png") . ImageRepa 8 $ shapeArr
  filterF <- dft4D plan . computeS . makeFilter2D $ localEigenVector
  initialVec <-
    R.zipWith (*) bias <$>
    convolve4D
      plan
      filterF
      (computeInitialDistributionR2T0S0'
         numPoint
         numPoint
         thetaFreqs
         scaleFreqs
         maxScale
         xs)
  -- eigenVec <-
  --   powerMethodR2Z2
  --     plan
  --     folderPath
  --     numPoint
  --     numPoint
  --     numOrientation
  --     thetaFreqs
  --     numScale
  --     scaleFreqs
  --     maxScale
  --     pinwheelArr
  --     -- localEigenVector
  --     PowerMethodGlobal
  --     numIteration
  --     writeSourceFlag
  --     (printf
  --        "_%d_%d_%d_%d_%d_%d_%.2f_%.2f"
  --        numPoint
  --        (round thetaFreq :: Int)
  --        (round scaleFreq :: Int)
  --        (round maxScale :: Int)
  --        (round tao :: Int)
  --        cutoffRadius
  --        thetaSigma
  --        scaleSigma)
  --     bias
  --     initialVec
  --     -- bias
  --     -- (R.traverse bias id $ \f idx@(Z :. k :. l :. _ :. _) ->
  --     --    if k == div (L.length thetaFreqs) 2 && l == div (L.length scaleFreqs) 2
  --     --      then f idx
  --     --      else 0)
  pathNormMethod <-
    if patchNormFlag
      then do
        let points =
              createIndex2D .
              L.map (\(i, j) -> (i + div numPoint 2, j + div numPoint 2)) .
              getShape2DIndexList $
              pointSet
            ys =
              pointCluster
                (connectionMatrixP
                   (ParallelParams numThread 1)
                   (minDist + 1)
                   points) $
              points
        M.zipWithM_
          (\i ->
             plotImageRepa (folderPath </> (printf "Cluster%03d.png" i)) .
             ImageRepa 8)
          [1 :: Int ..] .
          cluster2Array numPoint numPoint $
          ys
        return . PowerMethodConnection $ ys
      else return PowerMethodGlobal
  sourceR2Z2 <-
    powerMethodR2Z2T0S0Reversal
      plan
      folderPath
      numPoint
      numPoint
      numOrientation
      thetaFreqs
      thetaFreqs
      numScale
      scaleFreqs
      scaleFreqs
      maxScale
      arrR2Z2T0S0
      pathNormMethod
      0
      writeSourceFlag
      (printf
         "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_Final"
         numPoint
         (round thetaFreq :: Int)
         (round scaleFreq :: Int)
         (round maxScale :: Int)
         (round tao :: Int)
         cutoffRadius
         thetaSigma
         scaleSigma)
      0.0
      0.0
      bias
      -- (computeInitialEigenVectorR2T0S0
      --    numPoint
      --    numPoint
      --    thetaFreqs
      --    scaleFreqs
      --    thetaFreqs
      --    scaleFreqs
      --    xs)
      (R.traverse
         -- eigenVec
         initialVec
         (\(Z :. numThetaFreq :. numScaleFreq :. x :. y) ->
            (Z :. numThetaFreq :. numScaleFreq :. numThetaFreq :. numScaleFreq :.
             x :.
             y)) $ \f idx@(Z :. tf :. sf :. t0f :. s0f :. i :. j) ->
         f (Z :. tf :. sf :. i :. j))
  powerMethodR2Z2
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    numScale
    scaleFreqs
    maxScale
    localEigenVector
    PowerMethodGlobal
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_TEST"
       numPoint
       (round thetaFreq :: Int)
       (round scaleFreq :: Int)
       (round maxScale :: Int)
       (round tao :: Int)
       cutoffRadius
       thetaSigma
       scaleSigma)
    bias
    (R.zipWith (*) bias sourceR2Z2)

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:
{-# language DerivingVia #-}
{-# language FlexibleInstances #-}

module Geometry.Point
  ( APoint(..), Point (..), Label (..)
  , mkPoint, mkLabel
  ) where

import Data.Complex
import Data.Bool
import Numeric
import Geometry.Base

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

-- | The transparent class fot point-like objects.
class APoint p where
  toPoint :: p -> Point
  asPoint :: Point -> p

instance APoint Cmp where
  toPoint = Point
  asPoint (Point p) = p

instance APoint XY where
  toPoint = Point . cmp
  asPoint (Point p) = xy p

instance APoint p => APoint (Maybe p) where
  toPoint = maybe (asCmp 0) toPoint
  asPoint = pure . asPoint

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

-- | The type representing a point.
newtype Point = Point Cmp

-- | The general point constructor.
mkPoint :: Affine a => a -> Point
mkPoint = Point . cmp

instance APoint Point where
  toPoint = id
  asPoint = id

instance Show Point where
  show p = concat ["<Point (", sx, " ", sy, ")>"]
    where sx = show $ getX p
          sy = show $ getY p

instance Eq Point where
  p1 == p2 = cmp p1 ~= cmp p2

instance Metric Point where
  dist p1 p2 = cmp p1 `dist` cmp p2
  dist2 p1 p2 = cmp p1 `dist2` cmp p2

instance Trans Point where
  transform t (Point p) = Point $ transformCmp t p

instance Affine Point where
  cmp (Point p) = p
  asCmp = Point

instance Figure Point where
  isTrivial _ = False
  refPoint = cmp
  box = pointBox

--------------------------------------------------------------------------------
-- | The type representing a label.
newtype Label = Label Cmp
  deriving ( Eq
           , Metric
           , Trans
           , Affine
           , Figure
           , APoint) via Point

-- | The general label constructor (the label test is set by @label@ decorator).
mkLabel :: Affine a => a -> Label
mkLabel = Label . cmp

instance Show Label where
  show p = concat ["<Label (", sx, " ", sy, ")>"]
    where sx = show $ getX p
          sy = show $ getY p

max_stars_repo_name: MethaHardworker/geometry

max_issues_repo_name: MethaHardworker/geometry

max_forks_repo_name: MethaHardworker/geometry


text:
{-# LANGUAGE QuasiQuotes #-}

{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

-- | C code generation of primitive Feldspar expressions

module Feldspar.Primitive.Backend.C where

import Data.Complex

import Data.Constraint (Dict (..))
import Data.Proxy

import Language.C.Quote.C
import qualified Language.C.Syntax as C

import Language.C.Monad
import Language.Embedded.Backend.C

import Language.Syntactic

import Feldspar.Primitive.Representation

-- Note: This module assumes a 32-bit target. For example the C function `abs`
-- is used up to 32-bits, and `labs` is used above that.

viewLitPrim :: ASTF (Primitive :&: PrimTypeRep) a -> Maybe a
viewLitPrim (Sym (Lit a :&: _)) = Just a
viewLitPrim (Sym (Pi :&: _))    = Just pi
viewLitPrim _ = Nothing

instance CompTypeClass PrimType'
  where
    compType _ (_ :: proxy a) = case primTypeRep :: PrimTypeRep a of
      BoolT   -> addInclude "<stdbool.h>" >> return [cty| typename bool     |]
      Int8T   -> addInclude "<stdint.h>"  >> return [cty| typename int8_t   |]
      Int16T  -> addInclude "<stdint.h>"  >> return [cty| typename int16_t  |]
      Int32T  -> addInclude "<stdint.h>"  >> return [cty| typename int32_t  |]
      Int64T  -> addInclude "<stdint.h>"  >> return [cty| typename int64_t  |]
      Word8T  -> addInclude "<stdint.h>"  >> return [cty| typename uint8_t  |]
      Word16T -> addInclude "<stdint.h>"  >> return [cty| typename uint16_t |]
      Word32T -> addInclude "<stdint.h>"  >> return [cty| typename uint32_t |]
      Word64T -> addInclude "<stdint.h>"  >> return [cty| typename uint64_t |]
      FloatT  -> return [cty| float |]
      DoubleT -> return [cty| double |]
      ComplexFloatT  -> addInclude "<tgmath.h>" >> return [cty| float  _Complex |]
      ComplexDoubleT -> addInclude "<tgmath.h>" >> return [cty| double _Complex |]

    compLit _ a = case primTypeOf a of
      BoolT   -> do addInclude "<stdbool.h>"
                    return $ if a then [cexp| true |] else [cexp| false |]
      Int8T   -> return [cexp| $a |]
      Int16T  -> return [cexp| $a |]
      Int32T  -> return [cexp| $a |]
      Int64T  -> return [cexp| $a |]
      Word8T  -> return [cexp| $a |]
      Word16T -> return [cexp| $a |]
      Word32T -> return [cexp| $a |]
      Word64T -> return [cexp| $a |]
      FloatT  -> return [cexp| $a |]
      DoubleT -> return [cexp| $a |]
      ComplexFloatT  -> return $ compComplexLit a
      ComplexDoubleT -> return $ compComplexLit a

-- | Compile a complex literal
compComplexLit :: (Eq a, Num a, ToExp a) => Complex a -> C.Exp
compComplexLit (r :+ 0) = [cexp| $r |]
compComplexLit (0 :+ i) = [cexp| $i * I |]
compComplexLit (r :+ i) = [cexp| $r + $i * I |]

addTagMacro :: MonadC m => m ()
addTagMacro = addGlobal [cedecl|$esc:("#define TAG(tag,exp) (exp)")|]

-- | Compile a unary operator
compUnOp :: MonadC m => C.UnOp -> ASTF PrimDomain a -> m C.Exp
compUnOp op a = do
    a' <- compPrim $ Prim a
    return $ C.UnOp op a' mempty

-- | Compile a binary operator
compBinOp :: MonadC m =>
    C.BinOp -> ASTF PrimDomain a -> ASTF PrimDomain b -> m C.Exp
compBinOp op a b = do
    a' <- compPrim $ Prim a
    b' <- compPrim $ Prim b
    return $ C.BinOp op a' b' mempty

-- | Compile a function call
compFun :: MonadC m => String -> Args (AST PrimDomain) sig -> m C.Exp
compFun fun args = do
    as <- sequence $ listArgs (compPrim . Prim) args
    return [cexp| $id:fun($args:as) |]

-- | Compile a call to 'abs'
compAbs :: MonadC m => PrimTypeRep a -> ASTF PrimDomain a -> m C.Exp
compAbs t a = case viewPrimTypeRep t of
    PrimTypeBool     -> error "compAbs: type BoolT not supported"
    PrimTypeIntWord (IntType _)  -> addInclude "<stdlib.h>" >> compFun "abs" (a :* Nil)
    PrimTypeIntWord (WordType _) -> compPrim $ Prim a
    _ -> addInclude "<tgmath.h>" >> compFun "fabs" (a :* Nil)
      -- Floating and complex types

complexSign_def = [cedecl|
double _Complex feld_complexSign(double _Complex c) {
    double z = cabs(c);
    if (z == 0) {
        return 0;
    } else {
        return (creal(c)/z + I*(cimag(c)/z));
    }
}
|]

complexSignf_def = [cedecl|
float _Complex feld_complexSignf(float _Complex c) {
    float z = cabsf(c);
    if (z == 0) {
        return 0;
    } else {
        return (crealf(c)/z + I*(cimagf(c)/z));
    }
}
|]

-- | Compile a call to 'signum'
compSign :: MonadC m => PrimTypeRep a -> ASTF PrimDomain a -> m C.Exp
compSign t a = case viewPrimTypeRep t of
    PrimTypeBool -> error "compSign: type BoolT not supported"
    PrimTypeIntWord (WordType _) -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", $a' > 0) |]
    PrimTypeIntWord (IntType _) -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", ($a' > 0) - ($a' < 0)) |]
    PrimTypeFloating FloatType -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", (float) (($a' > 0) - ($a' < 0))) |]
    PrimTypeFloating DoubleType -> do
        addTagMacro
        a' <- compPrim $ Prim a
        return [cexp| TAG("signum", (double) (($a' > 0) - ($a' < 0))) |]
    PrimTypeComplex ComplexDoubleType -> do
        addInclude "<tgmath.h>"
        addGlobal complexSign_def
        a' <- compPrim $ Prim a
        return [cexp| feld_complexSign($a') |]
    PrimTypeComplex ComplexFloatType -> do
        addInclude "<complex.h>"
        addGlobal complexSignf_def
        a' <- compPrim $ Prim a
        return [cexp| feld_complexSignf($a') |]
  -- TODO The floating point cases give `sign (-0.0) = 0.0`, which is (slightly)
  -- wrong. They should return -0.0. I don't know whether it's correct for other
  -- strange values.

-- | Compile a type casted primitive expression
compCast :: MonadC m => PrimTypeRep a -> ASTF PrimDomain b -> m C.Exp
compCast t a = compPrim (Prim a) >>= compCastExp t

-- | Applies type cast on a compiled expression
compCastExp :: MonadC m => PrimTypeRep a -> C.Exp -> m C.Exp
compCastExp t a
    | Dict <- witPrimType t = do
        t' <- compType (Proxy :: Proxy PrimType') t
        return [cexp|($ty:t') $a|]

-- | Compile a call to 'round'
--   A cast is added to the resulting expression to allow compatibility
--   between different integral and floating return types.
compRound :: (PrimType' a, Num a, RealFrac b, MonadC m) =>
    PrimTypeRep a -> ASTF PrimDomain b -> m C.Exp
compRound t a = do
    addInclude "<tgmath.h>"
    rounded <- case viewPrimTypeRep t of
        PrimTypeIntWord _  -> compFun "lround" (a :* Nil)
        PrimTypeFloating _ -> compFun "round"  (a :* Nil)
        PrimTypeComplex _  -> compFun "round"  (a :* Nil)
        _ -> error $ "compRound: type " ++ show t ++ " not supported"
    compCastExp t rounded

-- Note: There's no problem with including both `tgmath.h` and `math.h`. As long
-- as the former is included, including the latter (before or after) doesn't
-- make a difference.
--
-- See: <https://gist.github.com/emilaxelsson/51310b3353f96914cd9bdb18b10b3103>

div_def = [cedecl|
int feld_div(int x, int y) {
    int q = x/y;
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) --q;
    return q;
}
|]

ldiv_def = [cedecl|
long int feld_ldiv(long int x, long int y) {
    int q = x/y;
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) --q;
    return q;
}
|]

mod_def = [cedecl|
int feld_mod(int x, int y) {
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) { r += y; }
    return r;
}
|]

lmod_def = [cedecl|
long int feld_lmod(long int x, long int y) {
    int r = x%y;
    if ((r!=0) && ((r<0) != (y<0))) { r += y; }
    return r;
}
|]

-- The above C implementations are taken from
-- <http://www.microhowto.info/howto/round_towards_minus_infinity_when_dividing_integers_in_c_or_c++.html>

compDiv :: MonadC m =>
    PrimTypeRep a -> ASTF PrimDomain a -> ASTF PrimDomain b -> m C.Exp
compDiv t a b = case viewPrimTypeRep t of
    PrimTypeIntWord (WordType _) -> compBinOp C.Div a b
    PrimTypeIntWord (IntType Int64Type) -> do
        addGlobal ldiv_def
        compFun "feld_ldiv" (a :* b :* Nil)
    PrimTypeIntWord _ -> do
        addGlobal div_def
        compFun "feld_div" (a :* b :* Nil)
    _ -> error $ "compDiv: type " ++ show t ++ " not supported"

compMod :: MonadC m =>
    PrimTypeRep a -> ASTF PrimDomain a -> ASTF PrimDomain b -> m C.Exp
compMod t a b = case viewPrimTypeRep t of
    PrimTypeIntWord (WordType _) -> compBinOp C.Mod a b
    PrimTypeIntWord (IntType Int64Type) -> do
        addGlobal lmod_def
        compFun "feld_lmod" (a :* b :* Nil)
    PrimTypeIntWord _ -> do
        addGlobal mod_def
        compFun "feld_mod" (a :* b :* Nil)
    _ -> error $ "compMod: type " ++ show t ++ " not supported"

-- | Compile an expression
compPrim :: MonadC m => Prim a -> m C.Exp
compPrim = simpleMatch (\(s :&: t) -> go t s) . unPrim
  where
    go :: forall m sig . MonadC m
        => PrimTypeRep (DenResult sig)
        -> Primitive sig
        -> Args (AST PrimDomain) sig
        -> m C.Exp
    go _ (FreeVar v) Nil = touchVar v >> return [cexp| $id:v |]
    go t (Lit a) Nil
        | Dict <- witPrimType t
        = compLit (Proxy :: Proxy PrimType') a
    go _ Add  (a :* b :* Nil) = compBinOp C.Add a b
    go _ Sub  (a :* b :* Nil) = compBinOp C.Sub a b
    go _ Mul  (a :* b :* Nil) = compBinOp C.Mul a b
    go _ Neg  (a :* Nil)      = compUnOp C.Negate a
    go t Abs  (a :* Nil)      = compAbs t a
    go t Sign (a :* Nil)      = compSign t a

    go _ Quot (a :* b :* Nil) = compBinOp C.Div a b
    go _ Rem  (a :* b :* Nil) = compBinOp C.Mod a b
    go t Div  (a :* b :* Nil) = compDiv t a b
    go t Mod  (a :* b :* Nil) = compMod t a b
    go _ FDiv (a :* b :* Nil) = compBinOp C.Div a b

    go _ Pi Nil = addGlobal pi_def >> return [cexp| FELD_PI |]
      where pi_def = [cedecl|$esc:("#define FELD_PI 3.141592653589793")|]
              -- This is the value of `pi :: Double`.
              -- Apparently there is no standard C99 definition of pi.
    go _ Exp   args = addInclude "<tgmath.h>" >> compFun "exp" args
    go _ Log   args = addInclude "<tgmath.h>" >> compFun "log" args
    go _ Sqrt  args = addInclude "<tgmath.h>" >> compFun "sqrt" args
    go _ Pow   args = addInclude "<tgmath.h>" >> compFun "pow" args
    go _ Sin   args = addInclude "<tgmath.h>" >> compFun "sin" args
    go _ Cos   args = addInclude "<tgmath.h>" >> compFun "cos" args
    go _ Tan   args = addInclude "<tgmath.h>" >> compFun "tan" args
    go _ Asin  args = addInclude "<tgmath.h>" >> compFun "asin" args
    go _ Acos  args = addInclude "<tgmath.h>" >> compFun "acos" args
    go _ Atan  args = addInclude "<tgmath.h>" >> compFun "atan" args
    go _ Sinh  args = addInclude "<tgmath.h>" >> compFun "sinh" args
    go _ Cosh  args = addInclude "<tgmath.h>" >> compFun "cosh" args
    go _ Tanh  args = addInclude "<tgmath.h>" >> compFun "tanh" args
    go _ Asinh args = addInclude "<tgmath.h>" >> compFun "asinh" args
    go _ Acosh args = addInclude "<tgmath.h>" >> compFun "acosh" args
    go _ Atanh args = addInclude "<tgmath.h>" >> compFun "atanh" args

    go _ Complex (a :* b :* Nil) = do
        addInclude "<tgmath.h>"
        a' <- compPrim $ Prim a
        b' <- compPrim $ Prim b
        return $ case (viewLitPrim a, viewLitPrim b) of
            (Just 0, _) -> [cexp| I*$b'       |]
            (_, Just 0) -> [cexp| $a'         |]
            _           -> [cexp| $a' + I*$b' |]
      -- We assume that constant folding has been performed, so that not both
      -- `a` and `b` are constants
    go _ Polar (m :* p :* Nil)
        | Just 0 <- viewLitPrim m = return [cexp| 0 |]
        | Just 0 <- viewLitPrim p = do
            m' <- compPrim $ Prim m
            return [cexp| $m' |]
        | Just 1 <- viewLitPrim m = do
            p' <- compPrim $ Prim p
            return [cexp| exp(I*$p') |]
        | otherwise = do
            m' <- compPrim $ Prim m
            p' <- compPrim $ Prim p
            return [cexp| $m' * exp(I*$p') |]
      -- We assume that constant folding has been performed, so that not both
      -- `m` and `p` are constants
    go _ Real      args = addInclude "<tgmath.h>" >> compFun "creal" args
    go _ Imag      args = addInclude "<tgmath.h>" >> compFun "cimag" args
    go _ Magnitude args = addInclude "<tgmath.h>" >> compFun "cabs"  args
    go _ Phase     args = addInclude "<tgmath.h>" >> compFun "carg"  args
    go _ Conjugate args = addInclude "<tgmath.h>" >> compFun "conj"  args

    go t I2N   (a :* Nil) = compCast t a
    go t I2B   (a :* Nil) = compCast t a
    go t B2I   (a :* Nil) = compCast t a
    go t Round (a :* Nil) = compRound t a

    go _ Not  (a :* Nil)      = compUnOp C.Lnot a
    go _ And  (a :* b :* Nil) = compBinOp C.Land a b
    go _ Or   (a :* b :* Nil) = compBinOp C.Lor a b
    go _ Eq   (a :* b :* Nil) = compBinOp C.Eq a b
    go _ NEq  (a :* b :* Nil) = compBinOp C.Ne a b
    go _ Lt   (a :* b :* Nil) = compBinOp C.Lt a b
    go _ Gt   (a :* b :* Nil) = compBinOp C.Gt a b
    go _ Le   (a :* b :* Nil) = compBinOp C.Le a b
    go _ Ge   (a :* b :* Nil) = compBinOp C.Ge a b

    go _ BitAnd   (a :* b :* Nil) = compBinOp C.And a b
    go _ BitOr    (a :* b :* Nil) = compBinOp C.Or a b
    go _ BitXor   (a :* b :* Nil) = compBinOp C.Xor a b
    go _ BitCompl (a :* Nil)      = compUnOp C.Not a
    go _ ShiftL   (a :* b :* Nil) = compBinOp C.Lsh a b
    go _ ShiftR   (a :* b :* Nil) = compBinOp C.Rsh a b

    go _ (ArrIx arr) (i :* Nil) = do
        i' <- compPrim $ Prim i
        touchVar arr
        return [cexp| $id:arr[$i'] |]

    go _ Cond (c :* t :* f :* Nil) = do
        c' <- compPrim $ Prim c
        t' <- compPrim $ Prim t
        f' <- compPrim $ Prim f
        return $ C.Cond c' t' f' mempty

    go _ s _ = error $ "compPrim: no handling of symbol " ++ renderSym s
      -- Should not occur, but the completeness checker doesn't know that

instance CompExp Prim where compExp = compPrim

max_stars_repo_name: Abhiroop/mu-feldspar

max_issues_repo_name: Abhiroop/mu-feldspar

max_forks_repo_name: Abhiroop/mu-feldspar


text:
{-# LANGUAGE FlexibleContexts #-}
module Main where

import           System.Exit
import           System.Random
import           System.TimeIt
import           Control.Monad.State
import           Numeric.LinearAlgebra
import           Control.Monad.Trans.Maybe      ( runMaybeT )
import           Control.Monad.Writer           ( runWriter )
import           Criterion.Main
import           Criterion.Types

import           Neural.Network
import           Neural.Activation              ( sigmoid
                                                , sigmoid'
                                                )
import           Neural.Training
import           Data.MNIST                     ( loadData )

shallowShape :: [Int]
shallowShape = [784, 30, 10]

deepShape :: [Int]
deepShape = [784, 30, 30, 30, 10]

-- HYPER PARAMETERS
-- learing rate
μ :: Double
μ = 0.5

-- Mini batch size
mbs :: Int
mbs = 100

-- Amount of training epochs
epochs :: Int
epochs = 50

main :: IO ()
main = defaultMain
  [ env setupEnv $ \ ~(trainEx, testEx) ->
      let trainShallow = timeIt (randomIO >>= trainNetwork trainEx testEx shallowShape)
          trainDeep    = timeIt (randomIO >>= trainNetwork trainEx testEx deepShape)
      in  bgroup
            "main"
            [ bench "shallow" $ nfIO trainShallow
            , bench "deep"    $ nfIO trainDeep
            ]
  ]

trainNetwork
  :: [TrainingExample Double]
  -> [TestingExample Double]
  -> [Int]
  -> Seed
  -> IO (Network Double)
trainNetwork trainEx testEx shape seed = do
  let gen             = mkStdGen seed
  let trainFunc = train sigmoid sigmoid' μ trainEx (Just testEx) epochs mbs
  let (randNet, gen') = runState (randomNetwork shape) gen
  let (net, logs) = runWriter $ evalStateT (trainFunc randNet) gen'
  putStrLn "Training network..."

  mapM_ putStrLn logs
  return net

setupEnv
    :: IO ([TrainingExample Double], [TestingExample Double])
setupEnv = do
  mnistData <- runMaybeT loadData
  case mnistData of
    Nothing -> putStrLn "Failed to load MNIST data." >> exitFailure
    Just (trainEx, testEx) -> return (trainEx, testEx)

max_stars_repo_name: cornelius-sevald/nnhd

max_issues_repo_name: cornelius-sevald/nnhd

max_forks_repo_name: cornelius-sevald/nnhd


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

module Numeric.LinearAlgebra.Matrix (
    -- * Immutable matrices
    Matrix,
    dim,

    -- * Matrix construction
    fromList,
    fromCols,
    fromRows,
    constant,
    zero,

    -- * Accessing matrices
    at,
    indices,
    elems,
    assocs,

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

    -- * Derived matrices
    map,
    zipWith,

    -- * Matrix views
    slice,
    takeRows,
    dropRows,
    splitRowsAt,
    takeCols,
    dropCols,
    splitColsAt,

    -- * Matrix rows and columns
    col,
    cols,
    row,
    rows,

    -- * Matrix diagonals
    diag,

    -- * Conversions to vectors
    toVector,

    -- * Conversions from vectors
    fromVector,
    fromCol,
    fromRow,

    -- * Matrix math operations
    shiftDiag,
    shiftDiagWithScale,
    add,
    addWithScale,
    sub,
    scale,
    scaleRows,
    scaleCols,
    negate,
    conjugate,

    -- * Linear algebra
    trans,
    conjTrans,
    rank1Update,

    -- ** Matrix-Vector multiplication
    mulVector,
    mulVectorWithScale,
    addMulVectorWithScales,

    -- ** Matrix-Matrix multiplication
    mulMatrix,
    mulMatrixWithScale,
    addMulMatrixWithScales,

    -- * Conversions between foreign pointers
    unsafeFromForeignPtr,
    unsafeToForeignPtr,

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

    -- * Hermitian views
    module Numeric.LinearAlgebra.Matrix.Herm,

    -- * Triangular views
    module Numeric.LinearAlgebra.Matrix.Tri,

    -- * Cholesky factorizations
    module Numeric.LinearAlgebra.Matrix.Cholesky,

    -- * Eigenvalues and eigenvectors
    module Numeric.LinearAlgebra.Matrix.Eigen,

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

    ) where

import Prelude hiding ( map, zipWith, negate, )
import Control.Monad( zipWithM_ )
import Control.Monad.ST( ST )
import Foreign( Storable )
import Text.Printf( printf )

import Foreign.BLAS( BLAS1, BLAS2, BLAS3, Trans(..) )
import Foreign.VMath( VNum )

import Numeric.LinearAlgebra.Matrix.Base
import Numeric.LinearAlgebra.Matrix.Herm
import Numeric.LinearAlgebra.Matrix.Tri
import Numeric.LinearAlgebra.Matrix.ST
import Numeric.LinearAlgebra.Matrix.Cholesky
import Numeric.LinearAlgebra.Matrix.Eigen
import Numeric.LinearAlgebra.Matrix.Statistics
import Numeric.LinearAlgebra.Vector( Vector )
import qualified Numeric.LinearAlgebra.Vector as V

infixl 7 `scale`, `scaleRows`, `scaleCols`
infixl 6 `add`, `shiftDiag`, `sub`

-- | Create a matrix of the given dimension with the given vectors as
-- columns.
fromCols :: (Storable e) => (Int,Int) -> [Vector e] -> Matrix e
fromCols mn cs = create $ do
    a <- new_ mn
    withColsM a $ \cs' -> zipWithM_ V.copyTo cs' cs
    return a

-- | Create a matrix of the given dimension with the given vectors as
-- rows.
fromRows :: (Storable e) => (Int,Int) -> [Vector e] -> Matrix e
fromRows (m,n) rs = create $ do
    a <- new_ (m,n)
    sequence_ [ setRow a i r | (i,r) <- zip [ 0..m-1 ] rs ]
    return a

-- | Get the given row of the matrix.
row :: (BLAS1 e) => Matrix e -> Int -> Vector e
row a i
    | i < 0 || i >= m = error $
        printf ("row <matrix with dim (%d,%d)> %d:"
                ++ " index out of range") m n i
    | otherwise =
        unsafeRow a i
  where
    (m,n) = dim a
{-# INLINE row #-}

-- | Version of 'row' that doesn't range-check indices.
unsafeRow :: (BLAS1 e) => Matrix e -> Int -> Vector e
unsafeRow a@(Matrix v _ n lda) i
    | lda == 1 = v
    | otherwise = V.create $ do
        r <- V.new_ n
        unsafeRowTo r a i
        return r

-- | Get a list of the rows in a matrix
rows :: (BLAS1 e) => Matrix e -> [Vector e]
rows a = [ unsafeRow a i | i <- [ 0..m-1 ] ]
  where
    (m,_) = dim a

-- | Get the diagonal of the matrix.
diag :: (Storable e) => Matrix e -> Vector e
diag a = V.create $ do
    x <- V.new_ mn
    diagTo x a
    return x
  where
    (m,n) = dim a
    mn = min m n

-- | @shiftDiag d a@ returns @diag(d) + a@.
shiftDiag :: (BLAS1 e) => Vector e -> Matrix e -> Matrix e
shiftDiag s a = create $ do
    a' <- newCopy a
    shiftDiagM_ s a'
    return a'

-- | @shiftDiagWithScale alpha d a@ returns @alpha * diag(d) + a@.
shiftDiagWithScale :: (BLAS1 e) => e -> Vector e -> Matrix e -> Matrix e
shiftDiagWithScale e s a = create $ do
    a' <- newCopy a
    shiftDiagWithScaleM_ e s a'
    return a'

-- | @add a b@ returns @a + b@.
add :: (VNum e) => Matrix e -> Matrix e -> Matrix e
add = result2 addTo

-- | @sub a b@ returns @a - b@.
sub :: (VNum e) => Matrix e -> Matrix e -> Matrix e
sub = result2 subTo

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

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

-- | @scaleRows s a@ returns @diag(s) * a@.
scaleRows :: (BLAS1 e) => Vector e -> Matrix e -> Matrix e
scaleRows s a = create $ do
    a' <- newCopy a
    scaleRowsM_ s a'
    return a'

-- | @scaleCols s a@ returns @a * diag(s)@.
scaleCols :: (BLAS1 e) => Vector e -> Matrix e -> Matrix e
scaleCols s a = create $ do
    a' <- newCopy a
    scaleColsM_ s a'
    return a'

-- | @negate a@ returns @-a@.
negate :: (VNum e) => Matrix e -> Matrix e
negate = result negateTo

-- | @conjugate a@ returns @conjugate(a)@.
conjugate :: (VNum e) => Matrix e -> Matrix e
conjugate = result conjugateTo

-- | @trans a@ retunrs @trans(a)@.
trans :: (BLAS1 e)
      => Matrix e
      -> Matrix e
trans a = let
    (m,n) = dim a
    in create $ do
        a' <- new_ (n,m)
        transTo a' a
        return a'

-- | @conjTrans a@ retunrs @conj(trans(a))@.
conjTrans :: (BLAS1 e)
          => Matrix e
          -> Matrix e
conjTrans a = let
    (m,n) = dim a
    in create $ do
        a' <- new_ (n,m)
        conjTransTo a' a
        return a'

-- | @rank1Update alpha x y a@ returns @alpha * x * y^H + a@.
rank1Update :: (BLAS2 e)
            => e
            -> Vector e
            -> Vector e
            -> Matrix e
            -> Matrix e
rank1Update alpha x y a =
    create $ do
        a' <- newCopy a
        rank1UpdateM_ alpha x y a'
        return a'

-- | @mulVector transa a x@
-- returns @op(a) * x@, where @op(a)@ is determined by @transa@.
mulVector :: (BLAS2 e)
          => Trans -> Matrix e
          -> Vector e
          -> Vector e
mulVector transa a x = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    in V.create $ do
        y <- V.new_ m
        mulVectorTo y transa a x
        return y

-- | @mulVectorWithScale alpha transa a x@
-- retunrs @alpha * op(a) * x@, where @op(a)@ is determined by @transa@.
mulVectorWithScale :: (BLAS2 e)
                   => e
                   -> Trans -> Matrix e
                   -> Vector e
                   -> Vector e
mulVectorWithScale alpha transa a x = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    in V.create $ do
        y <- V.new_ m
        mulVectorWithScaleTo y alpha transa a x
        return y

-- | @addMulVectorWithScales alpha transa a x beta y@
-- returns @alpha * op(a) * x + beta * y@, where @op(a)@ is
-- determined by @transa@.
addMulVectorWithScales :: (BLAS2 e)
                       => e
                       -> Trans -> Matrix e
                       -> Vector e
                       -> e
                       -> Vector e
                       -> Vector e
addMulVectorWithScales alpha transa a x beta y =
    V.create $ do
        y' <- V.newCopy y
        addMulVectorWithScalesM_ alpha transa a x beta y'
        return y'

-- | @mulMatrix transa a transb b@
-- returns @op(a) * op(b)@, where @op(a)@ and @op(b)@ are determined
-- by @transa@ and @transb@.
mulMatrix :: (BLAS3 e)
          => Trans -> Matrix e
          -> Trans -> Matrix e
          -> Matrix e
mulMatrix transa a transb b = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    n = case transb of NoTrans -> (snd . dim) b
                       _       -> (fst . dim) b
    in create $ do
        c <- new_ (m,n)
        mulMatrixTo c transa a transb b
        return c

-- | @mulMatrixWithScale alpha transa a transb b@
-- returns @alpha * op(a) * op(b)@, where @op(a)@ and @op(b)@ are determined
-- by @transa@ and @transb@.
mulMatrixWithScale :: (BLAS3 e)
                   => e
                   -> Trans -> Matrix e
                   -> Trans -> Matrix e
                   -> Matrix e
mulMatrixWithScale alpha transa a transb b = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    n = case transb of NoTrans -> (snd . dim) b
                       _       -> (fst . dim) b
    in create $ do
        c <- new_ (m,n)
        mulMatrixWithScaleTo c alpha transa a transb b
        return c

-- | @addMulMatrixWithScales alpha transa a transb b beta c@
-- returns @alpha * op(a) * op(b) + beta * c@, where @op(a)@ and
-- @op(b)@ are determined by @transa@ and @transb@.
addMulMatrixWithScales :: (BLAS3 e)
                       => e
                       -> Trans -> Matrix e
                       -> Trans -> Matrix e
                       -> e
                       -> Matrix e
                       -> Matrix e
addMulMatrixWithScales alpha transa a transb b beta c =
    create $ do
        c' <- newCopy c
        addMulMatrixWithScalesM_ alpha transa a transb b beta c'
        return c'

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

result2 :: (Storable e, Storable f, Storable g)
        => (forall s . STMatrix s g -> Matrix e -> Matrix f -> ST s a)
        -> Matrix e
        -> Matrix f
        -> Matrix g
result2 f a1 a2 = create $ newResult2 f a1 a2
{-# INLINE result2 #-}

newResult :: (RMatrix m, Storable e, Storable f)
          => (STMatrix s f -> m e -> ST s a)
          -> m e
          -> ST s (STMatrix s f)
newResult f a = do
    mn <- getDim a
    c <- new_ mn
    _ <- f c a
    return c
{-# INLINE newResult #-}

newResult2 :: (RMatrix m1, RMatrix m2, Storable e, Storable f, Storable g)
           => (STMatrix s g -> m1 e -> m2 f -> ST s a)
           -> m1 e
           -> m2 f
           -> ST s (STMatrix s g)
newResult2 f a1 a2 = do
    mn <- getDim a1
    c <- new_ mn
    _ <- f c a1 a2
    return c
{-# INLINE newResult2 #-}

max_stars_repo_name: patperry/hs-linear-algebra

max_issues_repo_name: cartazio/hs-cblas

max_forks_repo_name: cartazio/hs-cblas


text:
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE RankNTypes            #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Recurrent.Layers.LSTM where

import           Hedgehog
import           Hedgehog.Internal.Source
import           Hedgehog.Internal.Show
import           Hedgehog.Internal.Property ( failWith, Diff (..) )

import           Data.Foldable ( toList )
import           Data.Singletons.TypeLits

import           Grenade
import           Grenade.Recurrent

import qualified Numeric.LinearAlgebra as H
import qualified Numeric.LinearAlgebra.Static as S

import qualified Test.Grenade.Recurrent.Layers.LSTM.Reference as Reference
import           Test.Hedgehog.Hmatrix

genLSTM :: forall i o. (KnownNat i, KnownNat o) => Gen (LSTM i o)
genLSTM = do
    let w = uniformSample
        u = uniformSample
        v = randomVector

        w0 = S.konst 0
        u0 = S.konst 0
        v0 = S.konst 0

    LSTM <$> (LSTMWeights <$> w <*> u <*> v <*> w <*> u <*> v <*> w <*> u <*> v <*> w <*> v)
         <*> pure (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)

prop_lstm_reference_forwards =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM

    let actual          = runRecurrentForwards net (S1D cell) (S1D input)
    case actual of
      (_, (S1D cellOut) :: S ('D1 2), (S1D output) :: S ('D1 2)) ->
        let cellOut'        = Reference.Vector . H.toList . S.extract $ cellOut
            output'         = Reference.Vector . H.toList . S.extract $ output
            refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            (refCO, refO)   = Reference.runLSTM refNet refCell refInput
        in do toList refCO ~~~ toList cellOut'
              toList refO ~~~ toList output'

prop_lstm_reference_backwards =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
    let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
                                          = runRecurrentForwards  net (S1D cell) (S1D input)
        actualBacks                       = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
    case actualBacks of
      (actualGradients, _, _ :: S ('D1 3)) ->
        let refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            refGradients    = Reference.runLSTMback refCell refInput refNet
        in toList refGradients ~~~ toList (Reference.lstmToReference actualGradients)

prop_lstm_reference_backwards_input =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
    let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
                                          = runRecurrentForwards  net (S1D cell) (S1D input)
        actualBacks                       = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
    case actualBacks of
      (_, _, S1D actualGradients :: S ('D1 3)) ->
        let refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            refGradients    = Reference.runLSTMbackOnInput refCell refNet refInput
        in toList refGradients ~~~ H.toList (S.extract actualGradients)

prop_lstm_reference_backwards_cell =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
    let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
                                          = runRecurrentForwards  net (S1D cell) (S1D input)
        actualBacks                       = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
    case actualBacks of
      (_, S1D actualGradients, _ :: S ('D1 3)) ->
        let refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            refGradients    = Reference.runLSTMbackOnCell refInput refNet refCell
        in toList refGradients ~~~ H.toList (S.extract actualGradients)

(~~~) :: (Monad m, Eq a, Ord a, Num a, Fractional a, Show a, HasCallStack) => [a] -> [a] -> PropertyT m ()
(~~~) x y =
  if all (< 1e-8) (zipWith (-) x y) then
    success
  else
    case valueDiff <$> mkValue x <*> mkValue y of
      Nothing ->
        withFrozenCallStack $
          failWith Nothing $ unlines [
              "━━━ Not Simliar ━━━"
            , showPretty x
            , showPretty y
            ]
      Just differ ->
        withFrozenCallStack $
          failWith (Just $ Diff "Failed (" "- lhs" "~/~" "+ rhs" ")" differ) ""
infix 4 ~~~

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

max_issues_repo_name: Alien-Inc/grenade

max_forks_repo_name: Alien-Inc/grenade