text:
{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_HADDOCK prune #-}

module Penrose.Serializer where

import           Penrose.Env
import           Penrose.GenOptProblem
import           Penrose.Optimizer
-- import           Penrose.Plugins
import           Penrose.Style
import           Penrose.SubstanceTokenizer
import           Penrose.Util

import           Data.Aeson
import           Data.Aeson.TH
import           Data.Char                  (toLower)
import           GHC.Generics
import qualified Numeric.LinearAlgebra      as L
import           System.Random              (StdGen)
import           Text.Megaparsec

--------------------------------------------------------------------------------
-- Packet serialization for server
data Packet a = Packet
  { packetType     :: String
  , packetSession  :: Maybe String
  , packetContents :: a
  }

$(deriveJSON
    defaultOptions
    {fieldLabelModifier = map toLower . drop 6, omitNothingFields = True}
    ''Packet)

data Request a = Request
  { requestSession :: Maybe String
  , requestCall    :: a
  }

$(deriveJSON
    defaultOptions
    {fieldLabelModifier = map toLower . drop 7, omitNothingFields = True}
    ''Request)

--------------------------------------------------------------------------------
instance ToJSONKey Name

instance FromJSONKey Name

instance ToJSONKey TypeVar

instance FromJSONKey TypeVar

instance ToJSONKey T

instance FromJSONKey T

instance ToJSONKey Var

instance FromJSONKey Var

--------------------------------------------------------------------------------
-- TODO: slowly move all the JSON decls other than the ones in Server to
--       this module
deriveJSON defaultOptions ''SourcePosition

deriveJSON defaultOptions ''Pos

deriveJSON defaultOptions ''SourcePos

deriveJSON defaultOptions ''Translation

deriveJSON defaultOptions ''Name

deriveJSON defaultOptions ''FieldExpr

deriveJSON defaultOptions ''TagExpr

deriveJSON defaultOptions ''Expr

deriveJSON defaultOptions ''AnnoFloat

deriveJSON defaultOptions ''BinaryOp

deriveJSON defaultOptions ''UnaryOp

deriveJSON defaultOptions ''PropertyDecl

deriveJSON defaultOptions ''Path

deriveJSON defaultOptions ''Var

deriveJSON defaultOptions ''BindingForm

deriveJSON defaultOptions ''StyVar

deriveJSON defaultOptions ''VarEnv

deriveJSON defaultOptions ''TypeVar

deriveJSON defaultOptions ''T

deriveJSON defaultOptions ''TypeCtorApp

deriveJSON defaultOptions ''TypeConstructor

deriveJSON defaultOptions ''Arg

deriveJSON defaultOptions ''K

deriveJSON defaultOptions ''ValConstructor

deriveJSON defaultOptions ''Y

deriveJSON defaultOptions ''Type

deriveJSON defaultOptions ''Operator

deriveJSON defaultOptions ''PredicateEnv

deriveJSON defaultOptions ''StmtNotationRule

deriveJSON defaultOptions ''Predicate1

deriveJSON defaultOptions ''Prop

deriveJSON defaultOptions ''Predicate2

deriveJSON defaultOptions ''Penrose.SubstanceTokenizer.Token

-- TODO: de-lambdaize this to make it serializable
-- deriveJSON defaultOptions ''Policy
deriveJSON defaultOptions ''Params

deriveJSON defaultOptions ''Fn

deriveJSON defaultOptions ''StdGen

deriveJSON defaultOptions ''PolicyParams

deriveJSON defaultOptions ''OptType

deriveJSON defaultOptions ''OptStatus

deriveJSON defaultOptions ''BfgsParams

deriveJSON defaultOptions ''Penrose.GenOptProblem.State

--------------------------------------------------------------------------------
-- Interface
deriveJSON defaultOptions ''CompilerError

deriveJSON defaultOptions ''RuntimeError

--------------------------------------------------------------------------------
-- Plugins

-- deriveJSON defaultOptions ''PluginInput

ext: hs

max_stars_repo_name: apatriarca/penrose

max_issues_repo_name: apatriarca/penrose

max_forks_repo_name: apatriarca/penrose


text:
{-# LANGUAGE CPP               #-}
{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveTraversable #-}
module Main (main) where

import Data.List        (zip4)
import Test.Tasty       (TestTree, defaultMain, testGroup, withResource)
import Test.Tasty.HUnit (assertEqual, testCase)

import qualified Data.Foldable                      as F
import qualified Data.List.NonEmpty                 as NE
import qualified Data.Traversable                   as T

#if __GLASGOW_HASKELL__ >= 704
import qualified Numeric.AD.Mode.Reverse.Double     as AD
import qualified Statistics.Distribution            as S
import qualified Statistics.Distribution.ChiSquared as S
#endif

import Math.Regression.Simple
import Numeric.KBN            (sumKBN)

-------------------------------------------------------------------------------
-- Main
-------------------------------------------------------------------------------

main :: IO ()
main = defaultMain $ testGroup "regression-simple"
    [ linearTests
    , quadraticTests
    , lm1Tests
    , lm2Tests
    ]

-------------------------------------------------------------------------------
-- data
-------------------------------------------------------------------------------

withData :: FilePath -> (IO [(Double, Double, Double, Double)] -> TestTree) -> TestTree
withData fp = withResource acquire release where
    acquire = do
        contents <- readFile ("gnuplot/" ++ fp)
        return $ map (quad . map read . words) $ lines contents

    release _ = return ()

    quad :: [Double] -> (Double, Double, Double, Double)
    quad (x:y:dx:dy:_) = (x,y,dx,dy)
    quad _             = error "invalid data"

-------------------------------------------------------------------------------
-- Linear
-------------------------------------------------------------------------------

linearTests :: TestTree
linearTests = withData "linear.dat" $ \load -> testGroup "linear"
    [ testCase "no-errors" $ do
        linearData <- load
        let fit = linearFit (\(x,y,_,_) -> (x,y)) linearData
        assertEqual "params" (V2 2.95689 6.04617)   (round' (fitParams fit))
        assertEqual "errors" (V2 7.9788e-2 0.95195) (round' (fitErrors fit))
        assertEqual "ndf"    18                     (round' (fitNDF fit))
        assertEqual "wssr"   75.6356                (round' (fitWSSR fit))

    , testCase "y-errors" $ do
        linearData <- load
        let fit = linearWithYerrors (\(x,y,_,dy) -> (x,y,dy)) linearData
        assertEqual "params" (V2 2.97271 5.91878)  (round' (fitParams fit))
        assertEqual "errors" (V2 7.722e-2 0.91882) (round' (fitErrors fit))
        assertEqual "ndf"    18                    (round' (fitNDF fit))
        assertEqual "wssr"   38.8345               (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      2.999e-3              (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif

    , testCase "xy-errors" $ do
        linearData <- load
        let fit = nth 5 $ linearWithXYerrors (\(x,y,dx,dy) -> (x,y,dx,dy)) linearData
        assertEqual "params" (V2 2.97021 5.99061)   (round' (fitParams fit))
        assertEqual "errors" (V2 7.6542e-2 0.90917) (round' (fitErrors fit))
        assertEqual "ndf"    18                     (round' (fitNDF fit))
        assertEqual "wssr"   29.141                 (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      4.6683e-2              (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif

    , testCase "yx-errors" $ do
        linearData <- load
        let fit = nth 5 $ linearWithXYerrors (\(x,y,dx,dy) -> (y,x,dy,dx)) linearData
        assertEqual "params" (V2 0.33271 (-1.87107)) (round' (fitParams fit))
        assertEqual "errors" (V2 8.5724e-3 0.34855)  (round' (fitErrors fit))
        assertEqual "ndf"    18                      (round' (fitNDF fit))
        assertEqual "wssr"   29.3171                 (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      4.4639e-2               (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif
    ]

nth :: Int -> NE.NonEmpty a -> a
nth n (x NE.:| xs) = go n x xs where
    go _ z []     = z
    go m z (y:ys) = if m <= 0 then z else go (m - 1) y ys

-------------------------------------------------------------------------------
-- Quad
-------------------------------------------------------------------------------

quadraticTests :: TestTree
quadraticTests = withData "quad.dat" $ \load -> testGroup "quad"
    [ testCase "no-errors" $ do
        quadraticData <- load
        let fit = quadraticFit (\(x,y,_,_) -> (x,y)) quadraticData
        assertEqual "params" (V3 0.11487 (-3.34246) 6.63601) (round' (fitParams fit))
        assertEqual "errors" (V3 1.0297e-2 0.22674 1.07032)  (round' (fitErrors fit))
        assertEqual "ndf"    17                              (round' (fitNDF fit))
        assertEqual "wssr"   33.9104                         (round' (fitWSSR fit))

    , testCase "y-errors" $ do
        quadraticData <- load
        let fit = quadraticWithYerrors (\(x,y,_,dy) -> (x,y,dy)) quadraticData
        assertEqual "params" (V3 0.11156 (-3.27481) 6.25286) (round' (fitParams fit))
        assertEqual "errors" (V3 9.7603e-3 0.21331 0.99362)  (round' (fitErrors fit))
        assertEqual "ndf"    17                              (round' (fitNDF fit))
        assertEqual "wssr"   16.793                          (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        let q = S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)
        assertEqual "P"      0.46847                         (round' (1 - q))
#endif

    , testCase "xy-errors" $ do
        quadraticData <- load
        let fit = nth 5 $ quadraticWithXYerrors (\(x,y,dx,dy) -> (x,y,dx,dy)) quadraticData
        assertEqual "params" (V3 0.11222 (-3.29575) 6.39876) (round' (fitParams fit))
        assertEqual "errors" (V3 9.9372e-3 0.22027 1.06196)  (round' (fitErrors fit))
        assertEqual "ndf"    17                              (round' (fitNDF fit))
        assertEqual "wssr"   15.6318                         (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        let q = S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)
        assertEqual "P"      0.55007                         (round' (1 - q))
#endif
    ]

-------------------------------------------------------------------------------
-- LM
-------------------------------------------------------------------------------

lm1Tests :: TestTree
lm1Tests = withData "linear.dat" $ \load -> testGroup "lm1"
    [ testCase "no-errors" $ do
        linearData <- load

        let scale a (x, y, _, _) = case scaleGrad' (H2 a x) of
                (f, H2 da _) -> (y, f, da)

        let fit = NE.last $ levenbergMarquardt1 scale 1 linearData
        assertEqual "params" 3.03374   (round' (fitParams fit))
        assertEqual "errors" 3.8628e-2 (round' (fitErrors fit))
        assertEqual "ndf"    19        (round' (fitNDF fit))
        assertEqual "wssr"   80.7105   (round' (fitWSSR fit))

    , testCase "y-errors" $ do
        linearData <- load
        let scaleY a (x, y, _, dy) = case scaleGrad' (H2 a x) of
                (f, H2 da _) -> (y, f, da, dy)
        let fit = NE.last $ levenbergMarquardt1WithYerrors scaleY 1 linearData
        assertEqual "params" 3.04015   (round' (fitParams fit))
        assertEqual "errors" 3.7604e-2 (round' (fitErrors fit))
        assertEqual "ndf"    19        (round' (fitNDF fit))
        assertEqual "wssr"   40.9918   (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      2.4195e-3 (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif

    , testCase "xy-errors" $ do
        linearData <- load
        let scaleXY a (x, y, dx, dy) = case scaleGrad' (H2 a x) of
                (f, H2 da f') -> (y, f, da, f', dx, dy)
        let fit = NE.last $ levenbergMarquardt1WithXYerrors scaleXY 1 linearData
        assertEqual "params" 3.04315   (round' (fitParams fit))
        assertEqual "errors" 3.7477e-2 (round' (fitErrors fit))
        assertEqual "ndf"    19        (round' (fitNDF fit))
        assertEqual "wssr"   30.7021   (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      4.3516e-2 (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif
    ]

lm2Tests :: TestTree
lm2Tests = withData "linear.dat" $ \load -> testGroup "lm2"
    [ testCase "no-errors" $ do
        linearData <- load
        let lin (V2 a b) (x, y, _, _) = case linearGrad' (H3 a b x) of
                (f, H3 da db _) -> (y, f, V2 da db)

        let fit = NE.last $ levenbergMarquardt2 lin (V2 1 1) linearData
        assertEqual "params" (V2 2.95689 6.04617)   (round' (fitParams fit))
        assertEqual "errors" (V2 7.9788e-2 0.95195) (round' (fitErrors fit))
        assertEqual "ndf"    18                     (round' (fitNDF fit))
        assertEqual "wssr"   75.6356                (round' (fitWSSR fit))

    , testCase "y-errors" $ do
        linearData <- load
        let linY (V2 a b) (x, y, _, dy) = case linearGrad' (H3 a b x) of
                (f, H3 da db _) -> (y, f, V2 da db, dy)
        let fit = NE.last $ levenbergMarquardt2WithYerrors linY (V2 1 1) linearData
        assertEqual "params" (V2 2.97271 5.91882)  (round' (fitParams fit))
        assertEqual "errors" (V2 7.722e-2 0.91882) (round' (fitErrors fit))
        assertEqual "ndf"    18                    (round' (fitNDF fit))
        assertEqual "wssr"   38.8345               (round' (fitWSSR fit))

    , testCase "xy-errors" $ do
        linearData <- load
        let linXY (V2 a b) (x, y, dx, dy) = case linearGrad' (H3 a b x) of
                (f, H3 da db f') -> (y, f, V2 da db, f', dx, dy)
        let fit = NE.last $ levenbergMarquardt2WithXYerrors linXY (V2 1 1) linearData
        assertEqual "params" (V2 2.97021 5.99061)   (round' (fitParams fit))
        assertEqual "errors" (V2 7.6542e-2 0.90917) (round' (fitErrors fit))
        assertEqual "ndf"    18                     (round' (fitNDF fit))
        assertEqual "wssr"   29.141                 (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      4.6683e-2              (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif

    , testCase "yx-errors" $ do
        -- with x and y flipped:
        linearData <- load
        let linYX (V2 a b) (y, x, dy, dx) = case linearGrad' (H3 a b x) of
                (f, H3 da db f') -> (y, f, V2 da db, f', dx, dy)
        let fit = NE.last $ levenbergMarquardt2WithXYerrors linYX (V2 1 1) linearData
        assertEqual "params" (V2 0.33402 (-1.92156)) (round' (fitParams fit))
        assertEqual "errors" (V2 8.5785e-3 0.3488)   (round' (fitErrors fit))
        assertEqual "ndf"    18                      (round' (fitNDF fit))
        assertEqual "wssr"   29.1822                 (round' (fitWSSR fit))

#if __GLASGOW_HASKELL__ >= 704
        assertEqual "P"      4.6197e-2               (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif

#if __GLASGOW_HASKELL__ >= 704
    , testCase "orear-example" $ do
        let orearData :: [(Double, Double, Double, Double)]
            orearData = zip4
                [22000, 22930,23880,25130,26390]
                [-4.017,-2.742,-1.1478,1.491,6.873]
                [440,470,500,530,540]
                [0.50,0.25,0.08,0.09,1.90]

        let orearXY (V2 a b) (x, y, dx, dy) = case AD.grad' orearF (H3 a b x) of
                (f, H3 da db f') -> (y, f, V2 da db, f', dx, dy)

        let wssr0 = sumKBN
                [ sq (y - f) * w
                | d <- orearData
                , let a1 = 1e-3
                , let a2 = 6e5
                , let (y, f, _, f', dx, dy) = orearXY (V2 a1 a2) d
                , let w = recip $ sq (f' * dx) + sq dy
                ]

        assertEqual "wssr0"  3.82243                 (round' wssr0)

        let fit = NE.last $ levenbergMarquardt2WithXYerrors orearXY (V2 1e-3 6e5) orearData

        assertEqual "params" (V2 1.0163e-3 593725.0) (round' (fitParams fit))
        assertEqual "errors" (V2 1.7025e-4 95284.8)  (round' (fitErrors fit))
        assertEqual "ndf"    3                       (round' (fitNDF fit))
        assertEqual "wssr"   2.18668                 (round' (fitWSSR fit))

        assertEqual "P"      0.53458                 (round' (1 - S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)))
#endif

    ]

data H2 a = H2 a a   deriving (Functor, F.Foldable, T.Traversable)
data H3 a = H3 a a a deriving (Functor, F.Foldable, T.Traversable)

scaleF :: Num a => H2 a -> a
scaleF (H2 a x) = a * x + 5

scaleGrad' :: H2 Double -> (Double, H2 Double)
#if __GLASGOW_HASKELL__ >= 704
scaleGrad' = AD.grad' scaleF
#else
scaleGrad' (H2 a x) = (a * x + 5, H2 x a)
#endif

linearF :: Num a => H3 a -> a
linearF (H3 a b x) = a * x + b

linearGrad' :: H3 Double -> (Double, H3 Double)
#if __GLASGOW_HASKELL__ >= 704
linearGrad' = AD.grad' linearF
#else
linearGrad' (H3 a b x) = (a * x + b, H3 x 1 a)
#endif

orearF :: Fractional a => H3 a -> a
orearF (H3 a b x) = a * x - b / x

sq :: Num a => a -> a
sq x = x * x

-------------------------------------------------------------------------------
-- Round
-------------------------------------------------------------------------------

class Round a where
    round' :: a -> a

instance Round Double where
    round' 0 = 0
    round' x = fromInteger (round (x * rat)) / rat
      where
        mag = truncate (logBase 10 (abs x)) :: Int
        rat = 10 ^ (5 - mag)

instance Round Int where
    round' = id

instance Round V2 where
    round' (V2 x y) = V2 (round' x) (round' y)

instance Round V3 where
    round' (V3 x y z) = V3 (round' x) (round' y) (round' z)

max_stars_repo_name: phadej/regression-simple

max_issues_repo_name: phadej/regression-simple

max_forks_repo_name: phadej/regression-simple


text:
module Data.MNIST
  ( loadData
  )
where

import           System.IO
import qualified Data.ByteString.Lazy          as B
import           Codec.Compression.GZip
import           Data.IDX
import           Control.Error.Util
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.Class
import qualified Data.Vector.Unboxed           as V
import           Numeric.LinearAlgebra

import           Neural.Training                ( TrainingExample(..)
                                                , TestingExample(..)
                                                )

trainingImagesPath = "./data/train-images-idx3-ubyte.gz" :: FilePath
testingImagesPath = "./data/t10k-images-idx3-ubyte.gz" :: FilePath
trainingLabelsPath = "./data/train-labels-idx1-ubyte.gz" :: FilePath
testingLabelsPath = "./data/t10k-labels-idx1-ubyte.gz" :: FilePath

loadData :: MaybeT IO ([TrainingExample Double], [TestingExample Double])
loadData = do
  -- Uncompress the data.
  lift $ putStrLn "Uncompressing files..."
  trainingImagesData <- lift $ readCompressedFile trainingImagesPath
  testingImagesData  <- lift $ readCompressedFile testingImagesPath
  trainingLabelsData <- lift $ readCompressedFile trainingLabelsPath
  testingLabelsData  <- lift $ readCompressedFile testingLabelsPath
  -- Decode IDX format.
  lift $ putStrLn "Decoding data..."
  trainingImages <- hoistMaybe $ decodeIDX trainingImagesData
  testingImages  <- hoistMaybe $ decodeIDX testingImagesData
  trainingLabels <- hoistMaybe $ decodeIDXLabels trainingLabelsData
  testingLabels  <- hoistMaybe $ decodeIDXLabels testingLabelsData
  -- Pair images and label.
  lift $ putStrLn "Pairing images and labels..."
  labeledTrainingImages <- hoistMaybe
    $ labeledDoubleData trainingLabels trainingImages
  labeledTestingImages <- hoistMaybe
    $ labeledDoubleData testingLabels testingImages
  -- Convert to a proper type.
  lift $ putStrLn "Formatting data..."
  let trainingExamples =
        map labeledImageToTrainingExample labeledTrainingImages
  let testingExamples = map labeledImageToTestingExample labeledTestingImages

  lift $ putStrLn "Done converting MNIST data."
  return (trainingExamples, testingExamples)

labeledImageToTrainingExample
  :: (Int, V.Vector Double) -> TrainingExample Double
labeledImageToTrainingExample (label, image) = TrainingExample iv lv
 where
  lv = fromList $ replicate label 0 ++ [1] ++ replicate (10 - label - 1) 0
  iv = fromList $ V.toList image

labeledImageToTestingExample :: (Int, V.Vector Double) -> TestingExample Double
labeledImageToTestingExample (label, image) = TestingExample iv label
  where iv = fromList $ V.toList image

readCompressedFile :: FilePath -> IO B.ByteString
readCompressedFile fp = decompress <$> B.readFile fp

max_stars_repo_name: cornelius-sevald/nnhd

max_issues_repo_name: cornelius-sevald/nnhd

max_forks_repo_name: cornelius-sevald/nnhd


text:
{-# LANGUAGE TupleSections, DeriveGeneric, DeriveAnyClass #-}

module Main where

-- base
import Prelude hiding (pi)
import Control.Monad (replicateM, foldM)
import Data.List (sort, foldl', transpose, maximumBy)
import Data.Tuple (swap)
import Text.Printf (printf)
import System.CPUTime (getCPUTime)
import System.IO (hFlush, stdout)
import Control.Exception (evaluate)
import GHC.Generics (Generic)
-- mwc-random
import System.Random.MWC (initialize, GenIO)
-- vector
import Data.Vector.Unboxed (singleton)
-- deepseq
import Control.DeepSeq (NFData, rnf, ($!!))
-- proto-lens
import Data.ProtoLens (defMessage, encodeMessage)
-- microlens
import Lens.Micro ((&), (.~), traversed, (^..), _2, to)
-- bytestring
import qualified Data.ByteString as B
-- containers
import Data.Map.Strict (Map, (!), fromList, toList, adjust, fromAscList, toAscList)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IM
-- statistics
import Statistics.Distribution (cumulative, complCumulative, logDensity)
import Statistics.Distribution.Normal (normalDistr, NormalDistribution)
-- optparse-applicative
import Options.Applicative ((<**>), Parser, argument, str, metavar, helper, briefDesc, info, execParser)

import qualified Proto.Results as P
import qualified Proto.Results_Fields as P
import Lib

args :: Parser String
args = argument str (metavar "ENV")

main :: IO ()
main = do
  env <- execParser $ info (args <**> helper) briefDesc
  let (otype, ptype) = case env of "os-pf" -> (OStable, PFixed)
                                   "os-ps" -> (OStable, PStable)
                                   "os-pv" -> (OStable, PVar)
                                   "ov-pf" -> (OVar, PFixed)
                                   "ov-ps" -> (OVar, PStable)
                                   "ov-pv" -> (OVar, PVar)
                                   _ -> error $ "Invalid environment: " ++ env
  start <- getCPUTime
  gen <- initialize (singleton 3509)
  result <- runTournament gen otype ptype
  putChar '\n'
  putStrLn . matrixToStr 3 $ meanMatrix result
  putChar '\n'
  putStrLn . matrixToStr 3 $ semMatrix result
  B.writeFile ("output/"++env++".bin") . encodeMessage $ genMessage result
  end <- getCPUTime
  let diff = fromIntegral (end - start) / (10^(12 :: Int))
  printf "\nComputation time: %0.3f sec\n" (diff :: Double)

type Choice = Int
type Choices = (Choice, Choice)
type Payoffs = (Double, Double)
type Role = Int

type Posterior = [PostEntry]
type PostEntry = (PostPoint, Double)
type PostPoint = (Double, Double)

type QTables = [(Int, QTable)]
type QTable = Map QKey Double
type QKey = (QState, Choice)
type QState = [Choices] -- (choice_self, choice_opp), latest first

data Agent = Random
           | TfT { _lastOppChoice :: Choice }
           | NUM { _lambda :: Double }
           | ToM { _lambda :: Double
                 , _posterior :: Posterior
                 }
           | RL { _lambda :: Double
                , _history :: QState
                , _tables :: QTables
                }
           | FBM { _lambda :: Double }
           deriving (Generic, NFData, Eq, Ord, Show)

type Id = Int
type Sample = Bool
type Pair = ((Id, Agent), (Id, Agent))
type Game = [Double]
data OppStability = OStable | OVar
data PayStability = PFixed | PStable | PVar
data GameResult = GameResult { _game :: Game
                             , _choices :: Choices
                             , _agents :: Pair
                             , _payoffs :: Payoffs
                             }
                | GameResult0 { _payoffs :: Payoffs }
                deriving (Generic, NFData, Show)
data PairGamesResult = PairGamesResult [GameResult] Pair
                     | PairGamesResult0 Payoffs
                     deriving (Generic, NFData)
type PairResult = (Pair, [PairGamesResult])
data OVarGamesResult = OVarGamesResult [GameResult] [Agent] [Agent]
                     | OVarGamesResult0 [Payoffs]
                     deriving (Generic, NFData, Show)
data AllResult = OStableResult [PairResult]
               | OVarResult [OVarGamesResult]

mus, sigmaExps :: [Double]
mus = [-2,-1.8..2]
sigmaExps = map (+ logBase 2 0.75) [-3..3]

muPrior, sigmaExpPrior :: NormalDistribution
muPrior = normalDistr 0 1
sigmaExpPrior = normalDistr (logBase 2 0.75) 1.5

initialPosterior :: Posterior
initialPosterior = [((m, 2**s), logDensity muPrior m + logDensity sigmaExpPrior s) | m <- mus, s <- sigmaExps]

initialTables :: [(Int, QTable)]
initialTables = [(l, initialTable l) | l <- [0..3]]

initialTable :: Int -> QTable
initialTable l = fromList [((s, a), 0) | s <- possibleStates l, a <- [0, 1]]

possibleStates :: Int -> [QState]
possibleStates l = replicateM l [(x, y) | x <- [0, 1], y <- [0, 1]]

agents :: [Agent]
agents = [ Random, TfT 0 ]
  ++ map (\l -> RL l [] initialTables) [-1..1]
  ++ map (\l -> NUM l) [-1..1]
  ++ map (\l -> FBM l) [-1..1]
  ++ map (\l -> ToM l initialPosterior) [-1..1]
agentNums :: [Int]
agentNums = [4,4,1,2,1,1,2,1,1,2,1,1,2,1]
agents' :: [Agent]
agents' = concat $ zipWith replicate agentNums agents

pairs :: [Pair]
pairs = comb $ zip [0..] agents

runTournament :: GenIO -> OppStability -> PayStability -> IO AllResult
runTournament gen OStable gs = OStableResult <$> mapM (runPair gen gs) pairs
runTournament gen OVar gs = do
  let (n, n0) = case gs of PFixed -> (0, 5)
                           PStable -> (0, 500)
                           PVar -> (0, 500)
  r <- replicateM n $ runGamesOV gen gs True
  r0 <- fmap concat . replicateM 100 $ do
    r0' <- replicateM n0 $ runGamesOV gen gs False
    logProgress r0'
    return r0'
  return $ OVarResult (r ++ r0)

runPair :: GenIO -> PayStability -> Pair -> IO (Pair, [PairGamesResult])
runPair gen gs p = do
  let (n, n0) = case gs of PFixed -> (0, 100)
                           PStable -> (0, 10000)
                           PVar -> (0, 100)
  r <- replicateM n $ runGamesOS gen gs True p
  r0 <- replicateM n0 $ runGamesOS gen gs False p
  logProgress (r, r0)
  return (p, r ++ r0)

runGamesOS :: GenIO -> PayStability -> Sample -> Pair -> IO PairGamesResult
runGamesOS gen gs sp p = do
  game <- case gs of PStable -> genGame gen
                     _ -> return []
  let gameGen = case gs of PFixed -> return pd
                           PStable -> return game
                           PVar -> genGame gen
  (rs, p') <- foldM (folderOS gen gameGen sp) ([], p) [1..100]
  return $!! if sp
             then PairGamesResult (reverse rs) p'
             else PairGamesResult0 . meanPairs . map _payoffs $ rs

folderOS :: GenIO -> IO Game -> Sample -> ([GameResult], Pair) -> Int -> IO ([GameResult], Pair)
folderOS gen gameGen sp (rs, p) _ = do
  game <- gameGen
  (r, p') <- runGame gen updateOS game sp p
  return (r:rs, p')

runGamesOV :: GenIO -> PayStability -> Sample -> IO OVarGamesResult
runGamesOV gen gs sp = do
  game <- case gs of PStable -> genGame gen
                     _ -> return []
  let gameGen = case gs of PFixed -> return pd
                           PStable -> return game
                           PVar -> genGame gen
      t = [1..100*length agents']  -- folded structure
      am = IM.fromAscList $ zip [0..] agents'  -- agent map
  result <- if sp
            then do (rs, (as0, as1)) <- foldM (folderOVSp gen gameGen) ([], (am, am)) t
                    let ea = map snd . IM.toAscList  -- extract agents
                    return $ OVarGamesResult (reverse rs) (ea as0) (ea as1)
            else do let pm = fromAscList . map (,[]) $ comb [0..length agents-1]  -- payoffs map
                    (rs, _) <- foldM (folderOV gen gameGen) (pm, (am, am)) t
                    return $ OVarGamesResult0 . map (meanPairs . snd) $ toAscList rs
  return $!! result

type AccOVSp = ([GameResult], (IntMap Agent, IntMap Agent))
folderOVSp :: GenIO -> IO Game -> AccOVSp -> Int -> IO AccOVSp
folderOVSp gen gameGen (rs, (as0, as1)) _ = do
  game <- gameGen
  (i, j) <- samplePairWRUnif gen (length agents')
  (r, ((_, ai), (_, aj))) <- runGame gen updateOV game True ((i, as0 IM.! i), (j, as1 IM.! j))
  return (r:rs, (IM.insert i ai as0, IM.insert j aj as1))

type AccOV = (Map (Id, Id) [Payoffs], (IntMap Agent, IntMap Agent))
folderOV :: GenIO -> IO Game -> AccOV -> Int -> IO AccOV
folderOV gen gameGen (rs, (as0, as1)) _ = do
  game <- gameGen
  (i, j) <- samplePairWRUnif gen (length agents')
  (GameResult0 (pi, pj), ((_, ai), (_, aj))) <- runGame gen updateOV game False ((i, as0 IM.! i), (j, as1 IM.! j))
  let i' = reverseId i
      j' = reverseId j
      rs' = if i' <= j'
            then adjust ((pi, pj):) (i', j') rs
            else adjust ((pj, pi):) (j', i') rs
  return (rs', (IM.insert i ai as0, IM.insert j aj as1))

reverseId :: Int -> Int
reverseId i = length . takeWhile (<= i) $ scanl1 (+) agentNums

logProgress :: NFData a => a -> IO ()
logProgress a = do
  evaluate (rnf a)
  putChar '.'
  hFlush stdout

type Updater = Agent -> Game -> Role -> Choices -> Agent
runGame :: GenIO -> Updater -> Game -> Sample -> Pair -> IO (GameResult, Pair)
runGame g u game sp p@((i0, a0), (i1, a1)) = do
  c0 <- choose g a0 game 0
  c1 <- choose g a1 game 1
  let cs = (c0, c1)
      a0' = u a0 game 0 cs
      a1' = u a1 game 1 cs
      r = if sp
          then GameResult game cs p (payoffs game cs)
          else GameResult0 (payoffs game cs)
  return (r, ((i0, a0'), (i1, a1')))

choose :: GenIO -> Agent -> Game -> Role -> IO Choice
choose g Random _ _ = unifB g
choose _ (TfT c) _ _ = return c
choose _ (NUM l) m r = do
  let [p0,p1,p2,p3,p4,p5,p6,p7] = if r == 0
                                  then m
                                  else transposeGame m
  return $ if (p1+p3)*l+p0+p2 > (p5+p7)*l+p4+p6 then 0 else 1
choose _ a@ToM{} m r = do
  let (u0, u1) = tomUtil m r a
  return $ if u0 > u1 then 0 else 1
choose g a@RL{} _ _ = do
  let (u0, u1) = qUtil a
  if u0 == u1
  then unifB g
  else return $ if u0 > u1 then 0 else 1
choose _ (FBM l) m r = do
  let (u0, u1) = fbmUtil m r l
  return $ if u0 > u1 then 0 else 1

-- likelihood of choosing 0
tomLambdaLlhd :: Game -> Role -> PostPoint -> (Double, Double)
tomLambdaLlhd g r (mu, sigma) =
  let (cl, o) = cLambda g (1-r)
      [p0, p1] = (\f -> f (normalDistr mu sigma) cl) <$> [cumulative, complCumulative]
  in if o == LT then (p0, p1) else (p1, p0)

cLambda :: Game -> Role -> (Double, Ordering) -- LT: when lambda < cLambda, choose 0
cLambda g r = let [p0,p1,p2,p3,p4,p5,p6,p7] = if r == 0
                                              then g
                                              else transposeGame g
                  cl = (p0+p2-p4-p6)/(p5+p7-p1-p3)
                  o = compare (p1+p3) (p5+p7)
              in (cl, o)

transposeGame :: Game -> Game
transposeGame = map snd . sort . zip ([1,0,5,4,3,2,7,6] :: [Int])

tomUtil :: Game -> Role -> Agent -> (Double, Double)
tomUtil g r (ToM l p) = ((p0+p1*l)*c0+(p2+p3*l)*c1, (p4+p5*l)*c0+(p6+p7*l)*c1)
  where [p0,p1,p2,p3,p4,p5,p6,p7] = if r == 0
                                    then g
                                    else transposeGame g
        (c0, c1) = tomLambdaLlhd g r pp
        (pp, _) = maximumBy compareSnd p
tomUtil _ _ _ = error "tomUtil must be applied to a ToM agent"

fbmUtil :: Game -> Role -> Double -> (Double, Double)
fbmUtil g r l = ((p0+p1*l)*c0+(p2+p3*l)*c1, (p4+p5*l)*c0+(p6+p7*l)*c1)
  where [p0,p1,p2,p3,p4,p5,p6,p7] = if r == 0
                                    then g
                                    else transposeGame g
        (c0, c1) = tomLambdaLlhd g r (0,0.75)

qLRBase, qLRStep, qDiscount :: Double
qLRBase = 0.05
qLRStep = 0.25
qDiscount = 0.8

qUtil :: Agent -> (Double, Double)
qUtil (RL _ h ts) = let ts' = filter (\(len, _) -> len <= length h) ts
                        [q0, q1] = map mean . transpose $ getQValues h <$> ts'
                    in (q0, q1)
qUtil _ = error "qUtil must be applied to an RL agent"

getQValues :: QState -> (Int, QTable) -> [Double]
getQValues h (l, t) = [t ! (take l h, c) | c <- [0, 1]]

updateOS :: Updater
updateOS (TfT _) _ r (c0, c1) = TfT $ if r == 0 then c1 else c0
updateOS (ToM l p) g r cs = ToM l p'
  where p' = updatePostEntry g r cs <$> p
updateOS (RL l h ts) g r cs@(c0, c1) = RL l h' ts'
  where h' = cs':h
        cs' = if r == 0 then cs else swap cs
        css = if r == 0 then [(0, c1), (1, c1)] else [(c0, 0), (c0, 1)]
        pss = map (payoffs g) css
        rewards = map (util . if r == 0 then id else swap) pss
        util (p0, p1) = p0 + l*p1
        ts' = (\(len, t) -> (len, updateTable h rewards cs' len t)) <$> ts
updateOS a _ _ _ = a

updateOV :: Updater
updateOV a@(ToM _ _) _ _ _ = a
updateOV a b c d = updateOS a b c d

updatePostEntry :: Game -> Role -> Choices -> PostEntry -> PostEntry
updatePostEntry g r (c0, c1) (pp, v) = (pp, v + if log p' < -745 then -745 else log p')
  where c = if r == 0 then c1 else c0
        (p0, p1) = tomLambdaLlhd g r pp
        p' = if c == 0 then p0 else p1

updateTable :: QState -> [Double] -> Choices -> Int -> QTable -> QTable
updateTable h rewards (_, c1) l t =
  if l > length h
  then t
  else let futures = map (\c -> maximum $ getQValues ((c, c1):h) (l, t)) [0, 1]
           updateEntry t' c = adjust f (take l h, c) t'
             where f q = q + lr * (rewards!!c + qDiscount * futures!!c - q)
                   lr = qLRBase + qLRStep * fi l
       in foldl' updateEntry t [0, 1]

payoffs :: Game -> Choices -> Payoffs
payoffs g (c0, c1) = (g!!(cell * 2), g!!(cell * 2 + 1))
  where cell = c0 * 2 + c1

pd :: Game
pd = normalize4 [3,0,5,1]

gameSpace :: String
gameSpace = "all"

genGame :: GenIO -> IO Game
genGame = case gameSpace of "vpd" -> genVpd
                            "sym" -> genSym
                            "all" -> genAll
                            _ -> error $ "Unknow game space: " ++ gameSpace

genVpd :: GenIO -> IO Game
genVpd g = do
  v <- replicateM 4 $ unifD g
  let [b, d, a, c] = sort v
  if b + c < a * 2
    then return $ normalize4 [a,b,c,d]
    else genVpd g

genSym :: GenIO -> IO Game
genSym g = do
  v <- replicateM 4 $ unifD g
  return $ normalize4 v

genAll :: GenIO -> IO Game
genAll g = do
  v <- replicateM 8 $ unifD g
  return $ normalize v

normalize4 :: Game -> Game
normalize4 v = [a,a,b,c,c,b,d,d]
  where [a,b,c,d] = normalize v

-- print results
foldPairPayoffs :: ([Double] -> Double) -> [PairResult] -> [Payoffs]
foldPairPayoffs f rs = rs ^.. traversed . _2 . to (foldPairs f . map fromPairGamesResult)

fromPairGamesResult :: PairGamesResult -> Payoffs
fromPairGamesResult (PairGamesResult rs _) = meanPairs $ map _payoffs rs
fromPairGamesResult (PairGamesResult0 p) = p

foldOVarPayoffs :: ([Double] -> Double) -> [OVarGamesResult] -> [Payoffs]
foldOVarPayoffs f rs = map (foldPairs f) . transpose $ map fromOVarGamesResult rs

fromOVarGamesResult :: OVarGamesResult -> [Payoffs]
fromOVarGamesResult (OVarGamesResult rs _ _) = map (meanPairs . snd) $ toAscList tallied
  where tallied = foldl' f (fromAscList $ map (,[]) (comb [0..length agents-1])) rs
        f as (GameResult _ _ ((i, _), (j, _)) (pi, pj)) =
          let i' = reverseId i
              j' = reverseId j
          in if i' <= j'
             then adjust ((pi, pj):) (i', j') as
             else adjust ((pj, pi):) (j', i') as
        f _ (GameResult0 _) = error "OVarGamesResult can only contain GameResult, not GameResult0"
fromOVarGamesResult (OVarGamesResult0 rs) = rs

resultMatrix :: ([Double] -> Double) -> AllResult -> [[Double]]
resultMatrix f (OStableResult rs) = toMatrix . foldPairPayoffs f $ rs
resultMatrix f (OVarResult rs) = toMatrix . foldOVarPayoffs f $ rs

meanMatrix, semMatrix :: AllResult -> [[Double]]
meanMatrix = resultMatrix mean
semMatrix = resultMatrix sem

-- protobuf
genMessage :: AllResult -> P.AllResult
genMessage rs = defMessage
  & P.allResult .~ f1 rs
  & P.meanMatrix .~ map makeRow (meanMatrix rs)
  & P.semMatrix .~ map makeRow (semMatrix rs)
  where f1 :: AllResult -> P.AllResultOneof
        f1 (OStableResult r) = defMessage & P.oStableResult .~ f2 r
        f1 (OVarResult r) = defMessage & P.oVarResult .~ f3 r
        f2 :: [PairResult] -> P.OStableResult
        f2 r = defMessage & P.pairResults .~ map f4 r
        f3 :: [OVarGamesResult] -> P.OVarResult
        f3 r = defMessage
          & P.oVarAgents .~ map makeAgent agents
          & P.oVarGamesResults .~ map f5 r
        f4 :: PairResult -> P.PairResult
        f4 (p, r) = defMessage
          & P.pair .~ makePair p
          & P.pairGamesResults .~ map f6 r
        f5 :: OVarGamesResult -> P.OVarGamesResultOneof
        f5 (OVarGamesResult r ai aj) = defMessage & P.oVarGamesResult .~ f7 r ai aj
        f5 (OVarGamesResult0 r) = defMessage & P.oVarGamesResult0 .~ f8 r
        f6 :: PairGamesResult -> P.PairGamesResultOneof
        f6 (PairGamesResult r p) = defMessage & P.pairGamesResult .~ f9 r p
        f6 (PairGamesResult0 p) = defMessage & P.pairMeanPayoffs .~ makePayoffs p
        f7 :: [GameResult] -> [Agent] -> [Agent] -> P.OVarGamesResult
        f7 r ai aj = defMessage
          & P.oVarGameResults .~ map makeGameResult r
          & P.finalAgentsi .~ map makeAgent ai
          & P.finalAgentsj .~ map makeAgent aj
        f8 :: [Payoffs] -> P.OVarGamesResult0
        f8 r = defMessage & P.oVarPayoffs .~ map makePayoffs r
        f9 :: [GameResult] -> Pair -> P.PairGamesResult
        f9 r p = defMessage
          & P.pairGameResults .~ map makeGameResult r
          & P.finalPair .~ makePair p

makeRow :: [Double] -> P.Row
makeRow r = defMessage & P.row .~ r

makePair :: Pair -> P.Pair
makePair (a0, a1) = defMessage
  & P.agent0 .~ makeIndexedAgent a0
  & P.agent1 .~ makeIndexedAgent a1

makePairS :: Pair -> P.Pair
makePairS (a0, a1) = defMessage
  & P.agent0 .~ makeIndexedAgentS a0
  & P.agent1 .~ makeIndexedAgentS a1

makePayoffs :: Payoffs -> P.Payoffs
makePayoffs (p0, p1) = defMessage
  & P.payoff0 .~ p0
  & P.payoff1 .~ p1

makeIndexedAgent :: (Id, Agent) -> P.IndexedAgent
makeIndexedAgent (i, a) = defMessage
  & P.id .~ fi i
  & P.iagent .~ makeAgent a

makeIndexedAgentS :: (Id, Agent) -> P.IndexedAgent
makeIndexedAgentS (i, a) = defMessage
  & P.id .~ fi i
  & P.iagent .~ makeAgentS a

makeAgent :: Agent -> P.Agent
makeAgent Random = defMessage & P.random .~ defMessage
makeAgent (TfT c) = defMessage & P.tft .~ m
  where m = defMessage & P.lastOppChoice .~ fi c
makeAgent (NUM l) = defMessage & P.num .~ m
  where m = defMessage & P.lambda .~ l
makeAgent (ToM l p) = defMessage & P.tom .~ m
  where m = defMessage
          & P.lambda .~ l
          & P.posterior .~ map makePostEntry p
makeAgent (RL l h ts) = defMessage & P.rl .~ m
  where m = defMessage
          & P.lambda .~ l
          & P.history .~ map makeChoices h
          & P.tables .~ map makeQTable ts
makeAgent (FBM l) = defMessage & P.fbm .~ m
  where m = defMessage & P.lambda .~ l

makePostEntry :: PostEntry -> P.PostEntry
makePostEntry (pp, v) = defMessage
  & P.postPoint .~ makePostPoint pp
  & P.postValue .~ v

makePostPoint :: PostPoint -> P.PostPoint
makePostPoint (mu, sigma) = defMessage
  & P.mu .~ mu
  & P.sigma .~ sigma

makeAgentS :: Agent -> P.Agent
makeAgentS (ToM _ p) = defMessage & P.tomS .~ m
  where m = defMessage & P.postEntryS .~ makePostEntry (maximumBy compareSnd p)
makeAgentS (RL _ _ ts) = defMessage & P.rlS .~ m
  where m = defMessage & P.qStd .~ map qDiff ts
makeAgentS a = makeAgent a

qStd :: (Int, QTable) -> Double
qStd (_, t) = std (toList t ^.. traversed . _2)

qDiff :: (Int, QTable) -> Double
qDiff (l, t) = rms . map (\s -> t ! (s, 0) - t ! (s, 1)) $ possibleStates l

makeGameResult :: GameResult -> P.GameResult
makeGameResult (GameResult g cs a ps) = defMessage
  & P.game .~ g
  & P.choices .~ makeChoices cs
  & P.agents .~ makePairS a
  & P.payoffs .~ makePayoffs ps
makeGameResult _ = error "makeGameResult can only accept GameResult, not GameResult0"

makeChoices :: Choices -> P.Choices
makeChoices (c0, c1) = defMessage
  & P.choice0 .~ fi c0
  & P.choice1 .~ fi c1

makeQTable :: (Int, QTable) -> P.QTable
makeQTable (l, t) = defMessage
  & P.stateLength .~ fi l
  & P.table .~ map makeQEntry (toList t)

makeQEntry :: (QKey, Double) -> P.QEntry
makeQEntry (k, q) = defMessage
  & P.qKey .~ makeQKey k
  & P.qValue .~ q

makeQKey :: QKey -> P.QKey
makeQKey (s, c) = defMessage
  & P.qState .~ map makeChoices s
  & P.choice .~ fi c

max_stars_repo_name: jameswhqi/evo-tom

max_issues_repo_name: jameswhqi/evo-tom

max_forks_repo_name: jameswhqi/evo-tom


text:
{-|
Module      : Control.Monad.Bayes.Class
Description : Types for probabilistic modelling
Copyright   : (c) Adam Scibior, 2016
License     : MIT
Maintainer  : ams240@cam.ac.uk
Stability   : experimental
Portability : GHC

-}

module Control.Monad.Bayes.Class (
  MonadSample,
  random,
  uniform,
  normal,
  gamma,
  beta,
  bernoulli,
  categorical,
  logCategorical,
  uniformD,
  geometric,
  poisson,
  dirichlet,
  MonadCond,
  score,
  factor,
  condition,
  MonadInfer,
  normalPdf
) where

import Control.Monad.Trans.Class
import Control.Monad.Trans.Identity
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State
import Control.Monad.Trans.Writer
import Control.Monad.Trans.Reader
import Control.Monad.Trans.RWS hiding (tell)
import Control.Monad.Trans.List
import Control.Monad.Trans.Cont

import Statistics.Distribution
import Statistics.Distribution.Uniform (uniformDistr)
import Statistics.Distribution.Normal (normalDistr)
import Statistics.Distribution.Gamma (gammaDistr)
import Statistics.Distribution.Beta (betaDistr)
import Statistics.Distribution.Geometric (geometric0)
import qualified Statistics.Distribution.Poisson as Poisson

import Numeric.Log

import Data.Vector.Generic as VG
import qualified Data.Vector as V
import Control.Monad (when)

-- | Class of monads that can draw random variables.
class Monad m => MonadSample m where
  -- | A random variable distributed uniformly on [0,1].
  random :: m Double

  uniform :: Double -> Double -> m Double
  uniform a b = draw (uniformDistr a b)
  normal :: Double -> Double -> m Double
  normal m s = draw (normalDistr m s)
  gamma :: Double -> Double -> m Double
  gamma shape scale = draw (gammaDistr shape scale)
  beta :: Double -> Double -> m Double
  beta a b = draw (betaDistr a b)

  bernoulli :: Double -> m Bool
  bernoulli p = fmap (< p) random
  categorical :: Vector v Double => v Double -> m Int
  categorical ps = fromPMF (ps !)
  uniformD :: [a] -> m a
  uniformD xs = do
    let n = Prelude.length xs
    i <- categorical $ V.replicate n (1 / fromIntegral n)
    return (xs !! i)
  logCategorical :: (Vector v (Log Double), Vector v Double) => v (Log Double) -> m Int
  logCategorical = categorical . VG.map (exp . ln)
  geometric :: Double -> m Int
  geometric p = discrete (geometric0 p)
  poisson :: Double -> m Int
  poisson lambda = discrete (Poisson.poisson lambda)

  dirichlet :: Vector v Double => v Double -> m (v Double)
  dirichlet as = do
    xs <- VG.mapM (`gamma` 1) as
    let s = VG.sum xs
    let ys = VG.map (/ s) xs
    return ys

-- | Draw a value from a continuous distribution using iverse CDF.
draw :: (ContDistr d, MonadSample m) => d -> m Double
draw d = fmap (quantile d) random

-- | Draw a value from a discrete distribution using a sequence of draws from Bernoulli.
fromPMF :: MonadSample m => (Int -> Double) -> m Int
fromPMF p = f 0 1 where
  f i r = do
    when (r < 0) $ error "fromPMF: total PMF above 1"
    let q = p i
    when (q < 0 || q > 1) $ error "fromPMF: invalid probability value"
    b <- bernoulli (q / r)
    if b then pure i else f (i+1) (r-q)

-- | Draw a value from a discrete distributions using the PMF.
discrete :: (DiscreteDistr d, MonadSample m) => d -> m Int
discrete d = fromPMF (probability d)

-- | Monads that can score different execution paths.
class Monad m => MonadCond m where
  score :: Log Double -> m ()

-- | Synonym for 'score'.
factor :: MonadCond m => Log Double -> m ()
factor = score

-- | Hard conditioning.
condition :: MonadCond m => Bool -> m ()
condition b = score $ if b then 1 else 0

-- | Monads that support both sampling and scoring.
class (MonadSample m, MonadCond m) => MonadInfer m

normalPdf :: Double -> Double -> Double -> Log Double
normalPdf mu sigma x = Exp $ logDensity (normalDistr mu sigma) x

----------------------------------------------------------------------------
-- Instances that lift probabilistic effects to standard tranformers.

instance MonadSample m => MonadSample (IdentityT m) where
  random = lift random
  bernoulli = lift . bernoulli

instance MonadCond m => MonadCond (IdentityT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (IdentityT m)

instance MonadSample m => MonadSample (MaybeT m) where
  random = lift random

instance MonadCond m => MonadCond (MaybeT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (MaybeT m)

instance MonadSample m => MonadSample (ReaderT r m) where
  random = lift random
  bernoulli = lift . bernoulli

instance MonadCond m => MonadCond (ReaderT r m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ReaderT r m)

instance (Monoid w, MonadSample m) => MonadSample (WriterT w m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance (Monoid w, MonadCond m) => MonadCond (WriterT w m) where
  score = lift . score

instance (Monoid w, MonadInfer m) => MonadInfer (WriterT w m)

instance MonadSample m => MonadSample (StateT s m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance MonadCond m => MonadCond (StateT s m) where
  score = lift . score

instance MonadInfer m => MonadInfer (StateT s m)

instance (MonadSample m, Monoid w) => MonadSample (RWST r w s m) where
  random = lift random

instance (MonadCond m, Monoid w) => MonadCond (RWST r w s m) where
  score = lift . score

instance (MonadInfer m, Monoid w) => MonadInfer (RWST r w s m)

instance MonadSample m => MonadSample (ListT m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance MonadCond m => MonadCond (ListT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ListT m)

instance MonadSample m => MonadSample (ContT r m) where
  random = lift random

instance MonadCond m => MonadCond (ContT r m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ContT r m)

max_stars_repo_name: eligottlieb/monad-bayes

max_issues_repo_name: eligottlieb/monad-bayes

max_forks_repo_name: eligottlieb/monad-bayes


text:
{-# LANGUAGE Rank2Types, DeriveDataTypeable #-}
{-# OPTIONS_HADDOCK hide #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Vector.STBase
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : asinerimental
--

module Numeric.LinearAlgebra.Vector.STBase (
    STVector,
    IOVector,
    RVector(..),

    create,
    freeze,

    new_,
    new,
    newCopy,

    clear,
    copyTo,
    unsafeCopyTo,
    swap,
    unsafeSwap,

    getIndices,
    getElems,
    getElems',
    getAssocs,
    getAssocs',
    setElems,
    setAssocs,
    unsafeSetAssocs,

    read,
    unsafeRead,
    write,
    unsafeWrite,
    modify,
    unsafeModify,
    unsafeSwapElems,

    mapTo,
    unsafeMapTo,
    zipWithTo,
    unsafeZipWithTo,

    withSlice,
    withDrop,
    withTake,
    withSplitAt,

    withSliceM,
    withDropM,
    withTakeM,
    withSplitAtM,

    getSumAbs,
    getNorm2,
    getWhichMaxAbs,
    getDot,
    unsafeGetDot,
    scaleM_,
    addWithScaleM_,
    unsafeAddWithScaleM_,
    kroneckerTo,

    addTo,
    subTo,
    mulTo,
    negateTo,
    conjugateTo,
    absTo,
    signumTo,
    divTo,
    recipTo,
    sqrtTo,
    expTo,
    logTo,
    powTo,
    sinTo,
    cosTo,
    tanTo,
    asinTo,
    acosTo,
    atanTo,
    sinhTo,
    coshTo,
    tanhTo,
    asinhTo,
    acoshTo,
    atanhTo,
    ) where

import Prelude hiding ( drop, read, splitAt, take )

import Control.Monad( when, liftM2 )
import Control.Monad.ST( RealWorld, ST, runST, stToIO, unsafeIOToST )
import Data.Complex( Complex )
import Data.Typeable( Typeable )
import Foreign( Ptr, Storable, advancePtr, peek, poke, peekElemOff,
    pokeElemOff, copyArray, mallocForeignPtrArray )
import System.IO.Unsafe( unsafeInterleaveIO )
import Text.Printf( printf )
import Unsafe.Coerce( unsafeCoerce )

import Foreign.BLAS( BLAS1, BLAS2 )
import qualified Foreign.BLAS as BLAS
import Foreign.VMath( VNum, VFractional, VFloating )
import qualified Foreign.VMath as VMath

import Numeric.LinearAlgebra.Internal( clearArray )
import Numeric.LinearAlgebra.Vector.Base hiding ( unsafeWith )
import qualified Numeric.LinearAlgebra.Vector.Base as V

-- | Mutable vectors in the 'ST' monad.
newtype STVector s e = STVector { unSTVector :: Vector e }
    deriving (Typeable)

-- | Mutable vectors in the 'IO' monad.  Note that 'IO' operations
-- aren't directly supported; to perform an operation in the 'IO'
-- monad, perform the action in 'ST' 'RealWorld' and then convert
-- it via 'stToIO'.
type IOVector s = STVector RealWorld

-- | A safe way to create and work with a mutable vector before returning
-- an immutable vector for later perusal. This function avoids copying
-- the vector before returning it - it uses 'unsafeFreeze' internally,
-- but this wrapper is a safe interface to that function.
create :: (Storable e) => (forall s . ST s (STVector s e)) -> Vector e
create stmv = runST $ do
    mv <- stmv
    unsafeFreeze mv
{-# INLINE create #-}

-- | Converts a mutable vector to an immutable one by taking a complete
-- copy of it.
freeze :: (RVector v, Storable e) => v e -> ST s (Vector e)
freeze mv = do
    mv' <- newCopy mv
    unsafeFreeze mv'
{-# INLINE freeze #-}

-- | Read-only vectors
class RVector v where
    -- | Get the dimension of the vector.  This is equal to the number of
    -- elements in the vector.
    getDim :: (Storable e) => v e -> ST s Int

    -- | Same as 'withSlice' but does not range-check indices.
    unsafeWithSlice :: (Storable e)
                    => Int -> Int -> v e
                    -> (forall v'. RVector v' => v' e -> ST s a)
                    -> ST s a

    -- | Execute an 'IO' action with a pointer to the first element in the
    -- vector.
    unsafeWith :: (Storable e) => v e -> (Ptr e -> IO a) -> IO a

    -- | Converts a read-only vector into an immutable vector. This simply
    -- casts the vector from one type to the other without copying the vector.
    -- Note that because the vector is possibly not copied, any subsequent
    -- modifications made to the mutable version of the vector may be shared
    -- with the immutable version. It is safe to use, therefore, if the
    -- mutable version is never modified after the freeze operation.
    unsafeFreeze :: (Storable e) => v e -> ST s (Vector e)

    -- | Unsafe cast from a read-only vector to a mutable vector.
    unsafeThaw :: (Storable e)
               => v e -> ST s (STVector s e)

instance RVector Vector where
    getDim = return . dim
    {-# INLINE getDim #-}

    unsafeWith = V.unsafeWith
    {-# INLINE unsafeWith #-}

    unsafeWithSlice i n' v f =
        f (V.unsafeSlice i n' v)
    {-# INLINE unsafeWithSlice #-}

    unsafeFreeze = return . id
    {-# INLINE unsafeFreeze #-}

    unsafeThaw = return . STVector
    {-# INLINE unsafeThaw #-}

instance RVector (STVector s) where
    getDim = return . dim . unSTVector
    {-# INLINE getDim #-}

    unsafeWith v f = V.unsafeWith (unSTVector v) f
    {-# INLINE unsafeWith #-}

    unsafeWithSlice i n' v f =
        f $ unsafeSlice i n' (unSTVector v)
    {-# INLINE unsafeWithSlice #-}

    unsafeFreeze = return . unSTVector
    {-# INLINE unsafeFreeze #-}

    unsafeThaw v = return $ cast v
      where
        cast :: STVector s e -> STVector s' e
        cast = unsafeCoerce
    {-# INLINE unsafeThaw #-}

-- | @withSlice i n v@ performs an action with a view of the
-- @n@-dimensional subvector of @v@ starting at index @i@.
withSlice :: (RVector v, Storable e)
          => Int
          -> Int
          -> v e
          -> (forall v'. RVector v' => v' e -> ST s a)
          -> ST s a
withSlice i n' v f = do
    n <- getDim v
    when (i < 0 || n' < 0 || i + n' > n) $ error $
        printf "withSlice %d %d <vector with dim %d>: index out of range"
               i n' n
    unsafeWithSlice i n' v f

-- | Same as 'withSliceM' but does not range-check indices.
unsafeWithSliceM :: (Storable e)
                 => Int
                 -> Int
                 -> STVector s e
                 -> (STVector s e -> ST s a)
                 -> ST s a
unsafeWithSliceM i n' v f =
    f $ STVector $ unsafeSlice i n' (unSTVector v)

-- | Like 'withSlice', but perform the action with a mutable view
-- of the vector.
withSliceM :: (Storable e)
           => Int
           -> Int
           -> STVector s e
           -> (STVector s e -> ST s a)
           -> ST s a
withSliceM i n' v f = do
    n <- getDim v
    when (i < 0 || n' < 0 || i + n' > n) $ error $
        printf "withSlice %d %d <vector with dim %d>: index out of range"
               i n' n
    unsafeWithSliceM i n' v f

-- | Like 'withDrop', but perform the action with a mutable view
-- of the vector.
withDropM :: (Storable e)
          => Int
          -> STVector s e
          -> (STVector s e -> ST s a)
          -> ST s a
withDropM i v f = do
    n <- getDim v
    withSliceM i (n-i) v f

-- | Like 'withTake', but perform the action with a mutable view
-- of the vector.
withTakeM :: (Storable e)
          => Int
          -> STVector s e
          -> (STVector s e -> ST s a)
          -> ST s a
withTakeM = withSliceM 0

-- | Like 'withSplitAt' but perform the action with mutable views
-- of the vector.
withSplitAtM :: (Storable e)
             => Int
             -> STVector s e
             -> (STVector s e -> STVector s e -> ST s a)
             -> ST s a
withSplitAtM i v f = do
    n <- getDim v
    withSliceM 0 i v $ \v1 ->
        withSliceM i (n-i) v $ \v2 ->
            f v1 v2

-- | Perform an action the a view gotten from dropping the given
-- number of elements from the start of the vector.
withDrop :: (RVector v, Storable e)
         => Int
         -> v e
         -> (forall v'. RVector v' => v' e -> ST s a)
         -> ST s a
withDrop i v f = do
    mv <- unsafeThaw v
    withDropM i mv f

-- | Perform an action with a view gotten from taking the given
-- number of elements from the start of the vector.
withTake :: (RVector v, Storable e)
         => Int
         -> v e
         -> (forall v'. RVector v' => v' e -> ST s a)
         -> ST s a
withTake n v f = do
    mv <- unsafeThaw v
    withTakeM n mv f

-- | Perform an action with views from splitting the vector at the
-- given index.
withSplitAt :: (RVector v, Storable e)
            => Int
            -> v e
            -> (forall v1' v2'. (RVector v1', RVector v2') => v1' e -> v2' e -> ST s a)
            -> ST s a
withSplitAt i v f = do
    mv <- unsafeThaw v
    withSplitAtM i mv f

-- | Creates a new vector of the given length.  The elements will be
-- uninitialized.
new_ :: (Storable e) => Int -> ST s (STVector s e)
new_ n
    | n < 0 =  error $
        printf "new_ %d: invalid dimension" n
    | otherwise = unsafeIOToST $ do
        f <- mallocForeignPtrArray n
        return $ STVector $ V.unsafeFromForeignPtr f 0 n

-- | Create a vector with every element initialized to the same value.
new :: (Storable e) => Int -> e -> ST s (STVector s e)
new n e = do
    x <- new_ n
    setElems x $ replicate n e
    return x

-- | Creates a new vector by copying another one.
newCopy :: (RVector v, Storable e) => v e -> ST s (STVector s e)
newCopy x = do
    n <- getDim x
    y <- new_ n
    unsafeCopyTo y x
    return y

-- | @copyTo dst src@ replaces the values in @dst@ with those in
-- source.  The operands must be the same shape.
copyTo :: (RVector v, Storable e) => STVector s e -> v e -> ST s ()
copyTo = checkOp2 "copyTo" unsafeCopyTo
{-# INLINE copyTo #-}

-- | Same as 'copyTo' but does not check the dimensions.
unsafeCopyTo :: (RVector v, Storable e) => STVector s e -> v e -> ST s ()
unsafeCopyTo dst src = do
    n <- getDim dst
    unsafeIOToST $
        unsafeWith dst $ \pdst ->
        unsafeWith src $ \psrc ->
            copyArray pdst psrc n
{-# INLINE unsafeCopyTo #-}

-- | Swap the values stored in two vectors.
swap :: (BLAS1 e) => STVector s e -> STVector s e -> ST s ()
swap = checkOp2 "swap" unsafeSwap
{-# INLINE swap #-}

-- | Same as 'swap' but does not check the dimensions.
unsafeSwap :: (BLAS1 e) => STVector s e -> STVector s e -> ST s ()
unsafeSwap = strideCall2 BLAS.swap
{-# INLINE unsafeSwap #-}

-- | Get the indices of the elements in the vector, @[ 0..n-1 ]@, where
-- @n@ is the dimension of the vector.
getIndices :: (RVector v, Storable e) => v e -> ST s [Int]
getIndices v = do
    n <- getDim v
    return $ [ 0..n-1 ]
{-# INLINE getIndices #-}

-- | Lazily get the elements of the vector.
getElems :: (RVector v, Storable e) => v e -> ST s [e]
getElems v = let
    go end p' | p' == end = do
                  touch v
                  return []
              | otherwise = unsafeInterleaveIO $ do
                  e   <- peek p'
                  es  <- go end (p' `advancePtr` 1)
                  return $ e `seq` (e:es)
    in do
        n <- getDim v
        unsafeIOToST $
            unsafeWith v $ \p ->
                go (p `advancePtr` n) p
  where
    touch v' = unsafeWith v' $ const (return ())
{-# SPECIALIZE INLINE getElems :: STVector s Double -> ST s [Double] #-}
{-# SPECIALIZE INLINE getElems :: STVector s (Complex Double) -> ST s [Complex Double] #-}

-- | Get the elements of the vector.
getElems' :: (RVector v, Storable e) => v e -> ST s [e]
getElems' v = let
    go end p' es | p' == end =
                     return es
                 | otherwise = do
                     e <- peek p'
                     go end (p' `advancePtr` (-1)) (e:es)
    in do
        n <- getDim v
        unsafeIOToST $
            unsafeWith v $ \p ->
                go (p `advancePtr` (-1)) (p `advancePtr` (n-1)) []
{-# SPECIALIZE INLINE getElems' :: STVector s Double -> ST s [Double] #-}
{-# SPECIALIZE INLINE getElems' :: STVector s (Complex Double) -> ST s [Complex Double] #-}

-- | Lazily get the association list of the vector.
getAssocs :: (RVector v, Storable e) => v e -> ST s [(Int,e)]
getAssocs x = liftM2 zip (getIndices x) (getElems x)
{-# INLINE getAssocs #-}

-- | Get the association list of the vector.
getAssocs' :: (RVector v, Storable e) => v e -> ST s [(Int,e)]
getAssocs' x = liftM2 zip (getIndices x) (getElems' x)
{-# INLINE getAssocs' #-}

-- | Set all of the values of the vector from the elements in the list.
setElems  :: (Storable e) => STVector s e -> [e] -> ST s ()
setElems x es = let
    go n [] i _ | i < n = error $
                    printf ("setElems <vector with dim %d>"
                            ++ " <list with length %d>:"
                            ++ " not enough elements") n i
    go n (_:_) i _ | i == n = error $
                    printf ("setElems <vector with dim %d>"
                            ++ " <list with length at least %d>:"
                            ++ " too many elements") n (i+1)
    go _ []     _ _ = return ()
    go n (f:fs) i p = do
        poke p f
        go n fs (i+1) (p `advancePtr` 1)

    in do
        n <- getDim x
        unsafeIOToST $ unsafeWith x $ go n es 0

-- | Set the given values in the vector.  If an index is repeated twice,
-- the value is implementation-defined.
setAssocs :: (Storable e) => STVector s e -> [(Int,e)] -> ST s ()
setAssocs x ies =
    let go n p ((i,e):ies') = do
            when (i < 0 || i >= n) $ error $
                printf ("setAssocs <vector with dim %d>"
                        ++ " [ ..., (%d,_), ... ]: invalid index") n i
            pokeElemOff p i e
            go n p ies'
        go _ _ [] = return ()
    in do
        n <- getDim x
        unsafeIOToST $ unsafeWith x $ \p -> go n p ies

unsafeSetAssocs :: (Storable e) => STVector s e -> [(Int,e)] -> ST s ()
unsafeSetAssocs x ies =
    let go p ((i,e):ies') = do
            pokeElemOff p i e
            go p ies'
        go _ [] = return ()
    in unsafeIOToST $ unsafeWith x $ \p -> go p ies

-- | Get the element stored at the given index.
read :: (RVector v, Storable e) => v e -> Int -> ST s e
read x i = do
    n <- getDim x
    when (i < 0 || i >= n) $ error $
        printf ("read <vector with dim %d> %d:"
                ++ " invalid index") n i
    unsafeRead x i
{-# SPECIALIZE INLINE read :: STVector s Double -> Int -> ST s (Double) #-}
{-# SPECIALIZE INLINE read :: STVector s (Complex Double) -> Int -> ST s (Complex Double) #-}

-- | Same as 'read' but does not range check the index.
unsafeRead :: (RVector v, Storable e) => v e -> Int -> ST s e
unsafeRead x i =
    unsafeIOToST $ unsafeWith x $ \p -> peekElemOff p i
{-# SPECIALIZE INLINE unsafeRead :: STVector s Double -> Int -> ST s (Double) #-}
{-# SPECIALIZE INLINE unsafeRead :: STVector s (Complex Double) -> Int -> ST s (Complex Double) #-}

-- | Set the element stored at the given index.
write :: (Storable e) => STVector s e -> Int -> e -> ST s ()
write x i e = do
    n <- getDim x
    when (i < 0 || i >= n) $ error $
        printf ("write <vector with dim %d> %d:"
                ++ " invalid index") n i
    unsafeWrite x i e
{-# SPECIALIZE INLINE write :: STVector s Double -> Int -> Double -> ST s () #-}
{-# SPECIALIZE INLINE write :: STVector s (Complex Double) -> Int -> Complex Double -> ST s () #-}

-- | Same as 'write' but does not range check the index.
unsafeWrite :: (Storable e) => STVector s e -> Int -> e -> ST s ()
unsafeWrite x i e =
    unsafeIOToST $ unsafeWith x $ \p -> pokeElemOff p i e
{-# SPECIALIZE INLINE unsafeWrite :: STVector s Double -> Int -> Double -> ST s () #-}
{-# SPECIALIZE INLINE unsafeWrite :: STVector s (Complex Double) -> Int -> Complex Double -> ST s () #-}

-- | Modify the element stored at the given index.
modify :: (Storable e) => STVector s e -> Int -> (e -> e) -> ST s ()
modify x i f = do
    n <- getDim x
    when (i < 0 || i >= n) $ error $
        printf ("modify <vector with dim %d> %d:"
                ++ " invalid index") n i
    unsafeModify x i f
{-# SPECIALIZE INLINE modify :: STVector s Double -> Int -> (Double -> Double) -> ST s () #-}
{-# SPECIALIZE INLINE modify :: STVector s (Complex Double) -> Int -> (Complex Double -> Complex Double) -> ST s () #-}

-- | Same as 'modify' but does not range check the index.
unsafeModify :: (Storable e) => STVector s e -> Int -> (e -> e) -> ST s ()
unsafeModify x i f =
    unsafeIOToST $ unsafeWith x $ \p -> do
        e <- peekElemOff p i
        pokeElemOff p i $ f e
{-# SPECIALIZE INLINE unsafeModify :: STVector s Double -> Int -> (Double -> Double) -> ST s () #-}
{-# SPECIALIZE INLINE unsafeModify :: STVector s (Complex Double) -> Int -> (Complex Double -> Complex Double) -> ST s () #-}

unsafeSwapElems :: (Storable e) => STVector s e -> Int -> Int -> ST s ()
unsafeSwapElems x i1 i2 = unsafeIOToST $ unsafeWith x $ \p ->
    let p1 = p `advancePtr` i1
        p2 = p `advancePtr` i2
    in do
        e1  <- peek p1
        e2  <- peek p2
        poke p2 e1
        poke p1 e2
{-# SPECIALIZE INLINE unsafeSwapElems :: STVector s Double -> Int -> Int -> ST s () #-}
{-# SPECIALIZE INLINE unsafeSwapElems :: STVector s (Complex Double) -> Int -> Int -> ST s () #-}

-- | @mapTo dst f src@ replaces @dst@ elementwise with @f(src)@.
mapTo :: (RVector v, Storable e, Storable f)
      => STVector s f
      -> (e -> f)
      -> v e
      -> ST s ()
mapTo dst f src = (checkOp2 "mapTo _" $ \z x -> unsafeMapTo z f x) dst src
{-# INLINE mapTo #-}

-- | Same as 'mapTo' but does not check dimensions.
unsafeMapTo :: (RVector v, Storable e, Storable f)
            => STVector s f
            -> (e -> f)
            -> v e
            -> ST s ()
unsafeMapTo dst f src =
    let go end pdst psrc
            | pdst == end =
                return ()
            | otherwise = do
                e <- peek psrc
                poke pdst (f e)
                go end (pdst `advancePtr` 1) (psrc `advancePtr` 1)
    in do
        ndst <- getDim dst
        unsafeIOToST $
           unsafeWith dst $ \pdst ->
           unsafeWith src $ \psrc ->
               go (pdst `advancePtr` ndst) pdst psrc
  where

{-# INLINE unsafeMapTo #-}

-- | @zipWithTo dst f x y@ replaces @dst@ elementwise with @f(x,y)@.
zipWithTo :: (RVector v1, RVector v2, Storable e1, Storable e2, Storable f)
          => STVector s f
          -> (e1 -> e2 -> f)
          -> v1 e1
          -> v2 e2
          -> ST s ()
zipWithTo dst f x y =
    (checkOp3 "zipWithTo _" $ \dst1 x1 y1 -> unsafeZipWithTo dst1 f x1 y1)
        dst x y
{-# INLINE zipWithTo #-}

-- | Same as 'zipWithTo' but does not range-check dimensions.
unsafeZipWithTo :: (RVector v1, RVector v2, Storable e1, Storable e2, Storable f)
                => STVector s f
                -> (e1 -> e2 -> f)
                -> v1 e1
                -> v2 e2
                -> ST s ()
unsafeZipWithTo dst f src1 src2 =
    let go end pdst psrc1 psrc2
            | pdst == end =
                return ()
            | otherwise = do
                e1 <- peek psrc1
                e2 <- peek psrc2
                poke pdst (f e1 e2)
                go end (pdst `advancePtr` 1) (psrc1 `advancePtr` 1)
                   (psrc2 `advancePtr` 1)
    in do
        ndst <- getDim dst
        unsafeIOToST $
           unsafeWith dst $ \pdst ->
           unsafeWith src1 $ \psrc1 ->
           unsafeWith src2 $ \psrc2 ->
               go (pdst `advancePtr` ndst) pdst psrc1 psrc2
{-# INLINE unsafeZipWithTo #-}

-- | Set every element in the vector to a default value.  For
-- standard numeric types (including 'Double', 'Complex Double', and 'Int'),
-- the default value is '0'.
clear :: (Storable e) => STVector s e -> ST s ()
clear x = do
    n <- getDim x
    unsafeIOToST $ unsafeWith x $ \p -> clearArray p n

-- | @negateTo dst x@ replaces @dst@ with @negate(x)@.
negateTo :: (RVector v, VNum e) => STVector s e -> v e -> ST s ()
negateTo = checkOp2 "negateTo" $ \dst x ->
    call2 VMath.vNeg x dst
{-# INLINE negateTo #-}

-- | @absTo dst x@ replaces @dst@ with @abs(x)@.
absTo :: (RVector v, VNum e) => STVector s e -> v e -> ST s ()
absTo = checkOp2 "absTo" $ \dst x ->
    call2 VMath.vAbs x dst
{-# INLINE absTo #-}

-- | @signumTo dst x@ replaces @dst@ with @signum(x)@.
signumTo :: (RVector v, VNum e) => STVector s e -> v e -> ST s ()
signumTo = checkOp2 "signumTo" $ \dst x ->
    call2 VMath.vSgn x dst
{-# INLINE signumTo #-}

-- | @conjugateTo dst x@ replaces @dst@ with @conjugate(x)@.
conjugateTo :: (RVector v, VNum e) => STVector s e -> v e -> ST s ()
conjugateTo = checkOp2 "conjugateTo" $ \dst x ->
    call2 VMath.vConj x dst
{-# INLINE conjugateTo #-}

-- | @addTo dst x y@ replaces @dst@ with @x+y@.
addTo :: (RVector v1, RVector v2, VNum e)
      =>  STVector s e -> v1 e -> v2 e -> ST s ()
addTo = checkOp3 "addTo" $ \dst x y -> call3 VMath.vAdd x y dst
{-# INLINE addTo #-}

-- | @subTo dst x y@ replaces @dst@ with @x-y@.
subTo :: (RVector v1, RVector v2, VNum e)
      => STVector s e -> v1 e -> v2 e -> ST s ()
subTo = checkOp3 "subTo" $ \dst x y -> call3 VMath.vSub x y dst
{-# INLINE subTo #-}

-- | @mulTo dst x y@ replaces @dst@ with @x*y@.
mulTo :: (RVector v1, RVector v2, VNum e)
      => STVector s e -> v1 e -> v2 e -> ST s ()
mulTo = checkOp3 "mulTo" $ \dst x y -> call3 VMath.vMul x y dst
{-# INLINE mulTo #-}

-- | @divTo dst x y@ replaces @dst@ with @x/y@.
divTo :: (RVector v1, RVector v2, VFractional e)
      => STVector s e -> v1 e -> v2 e -> ST s ()
divTo = checkOp3 "divTo" $ \dst x y -> call3 VMath.vDiv x y dst
{-# INLINE divTo #-}

-- | @recipTo dst x@ replaces @dst@ with @1/x@.
recipTo :: (RVector v, VFractional e)
        => STVector s e -> v e -> ST s ()
recipTo = checkOp2 "recipTo" $ \dst x -> call2 VMath.vInv x dst
{-# INLINE recipTo #-}

-- | @sqrtTo dst x@ replaces @dst@ with @sqrt(x)@.
sqrtTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
sqrtTo = checkOp2 "sqrtTo" $ \dst x -> call2 VMath.vSqrt x dst
{-# INLINE sqrtTo #-}

-- | @expTo dst x@ replaces @dst@ with @exp(x)@.
expTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
expTo = checkOp2 "expTo" $ \dst x -> call2 VMath.vExp x dst
{-# INLINE expTo #-}

-- | @logTo dst x@ replaces @dst@ with @log(x)@.
logTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
logTo = checkOp2 "logTo" $ \dst x -> call2 VMath.vLog x dst
{-# INLINE logTo #-}

-- | @powTo dst x y@ replaces @dst@ with @x ** y@.
powTo :: (RVector v1, RVector v2, VFloating e)
      => STVector s e -> v1 e -> v2 e -> ST s ()
powTo = checkOp3 "powTo" $ \dst x y -> call3 VMath.vPow x y dst
{-# INLINE powTo #-}

-- | @sinTo dst x@ replaces @dst@ with @sin(x)@.
sinTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
sinTo = checkOp2 "sinTo" $ \dst x -> call2 VMath.vSin x dst
{-# INLINE sinTo #-}

-- | @cosTo dst x@ replaces @dst@ with @cos(x)@.
cosTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
cosTo = checkOp2 "cosTo" $ \dst x -> call2 VMath.vCos x dst
{-# INLINE cosTo #-}

-- | @tanTo dst x@ replaces @dst@ with @tan(x)@.
tanTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
tanTo = checkOp2 "tanTo" $ \dst x -> call2 VMath.vTan x dst
{-# INLINE tanTo #-}

-- | @asinTo dst x@ replaces @dst@ with @asin(x)@.
asinTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
asinTo = checkOp2 "asinTo" $ \dst x -> call2 VMath.vASin x dst
{-# INLINE asinTo #-}

-- | @acosTo dst x@ replaces @dst@ with @acos(x)@.
acosTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
acosTo = checkOp2 "acosTo" $ \dst x -> call2 VMath.vACos x dst
{-# INLINE acosTo #-}

-- | @atanTo dst x@ replaces @dst@ with @atan(x)@.
atanTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
atanTo = checkOp2 "atanTo" $ \dst x -> call2 VMath.vATan x dst
{-# INLINE atanTo #-}

-- | @sinhTo dst x@ replaces @dst@ with @sinh(x)@.
sinhTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
sinhTo = checkOp2 "sinhTo" $ \dst x -> call2 VMath.vSinh x dst
{-# INLINE sinhTo #-}

-- | @coshTo dst x@ replaces @dst@ with @cosh(x)@.
coshTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
coshTo = checkOp2 "coshTo" $ \dst x -> call2 VMath.vCosh x dst
{-# INLINE coshTo #-}

-- | @tanhTo dst x@ replaces @dst@ with @tanh(x)@.
tanhTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
tanhTo = checkOp2 "tanhTo" $ \dst x -> call2 VMath.vTanh x dst
{-# INLINE tanhTo #-}

-- | @asinhTo dst x@ replaces @dst@ with @asinh(x)@.
asinhTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
asinhTo = checkOp2 "asinhTo" $ \dst x -> call2 VMath.vASinh x dst
{-# INLINE asinhTo #-}

-- | @acoshTo dst x@ replaces @dst@ with @acosh(x)@.
acoshTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
acoshTo = checkOp2 "acoshTo" $ \dst x -> call2 VMath.vACosh x dst
{-# INLINE acoshTo #-}

-- | @atanhTo dst x@ replaces @dst@ with @atanh(x)@.
atanhTo :: (RVector v, VFloating e) => STVector s e -> v e -> ST s ()
atanhTo = checkOp2 "atanhTo" $ \dst x -> call2 VMath.vATanh x dst
{-# INLINE atanhTo #-}

-- | Gets the sum of the absolute values of the vector entries.
getSumAbs :: (RVector v, BLAS1 e) => v e -> ST s Double
getSumAbs = strideCall BLAS.asum
{-# INLINE getSumAbs #-}

-- | Gets the 2-norm of a vector.
getNorm2 :: (RVector v, BLAS1 e) => v e -> ST s Double
getNorm2 = strideCall BLAS.nrm2
{-# INLINE getNorm2 #-}

-- | Gets the index and norm of the element with maximum magnitude.  This is
-- undefined if any of the elements are @NaN@.  It will throw an exception if
-- the dimension of the vector is 0.
getWhichMaxAbs :: (RVector v, BLAS1 e) => v e -> ST s (Int, e)
getWhichMaxAbs x = do
    n <- getDim x
    when (n == 0) $ error $
        "getWhichMaxAbs <vector with dim 0>: empty vector"

    i <- strideCall BLAS.iamax x
    e <- unsafeRead x i
    return (i,e)
{-# INLINE getWhichMaxAbs #-}

-- | Computes the dot product of two vectors.
getDot :: (RVector v, RVector v', BLAS1 e)
       => v e -> v' e -> ST s e
getDot = checkOp2 "getDot" unsafeGetDot
{-# INLINE getDot #-}

-- | Same as 'getDot' but does not check dimensions.
unsafeGetDot :: (RVector x, RVector y, BLAS1 e)
             => x e -> y e -> ST s e
unsafeGetDot x y = (strideCall2 BLAS.dotc) y x
{-# INLINE unsafeGetDot #-}

-- | @scaleM k x@ sets @x := k * x@.
scaleM_ :: (Storable e, BLAS1 e) => e -> STVector s e -> ST s ()
scaleM_ k x = do
    n <- getDim x
    unsafeIOToST $
        unsafeWith x $ \px ->
            BLAS.scal n k px 1
{-# INLINE scaleM_ #-}

-- | @addWithScaleM_ alpha x y@ sets @y := alpha * x + y@.
addWithScaleM_ :: (RVector v, BLAS1 e) => e -> v e -> STVector s e -> ST s ()
addWithScaleM_ alpha x y =
    (checkOp2 "addWithScaleM_" $ \x1 y1 -> unsafeAddWithScaleM_ alpha x1 y1)
        x y
{-# INLINE addWithScaleM_ #-}

-- | Same as 'addWithScaleM_' but does not check dimensions.
unsafeAddWithScaleM_ :: (RVector v, BLAS1 e)
                     => e -> v e -> STVector s e -> ST s ()
unsafeAddWithScaleM_ alpha x y =
    (strideCall2 $ flip BLAS.axpy alpha) x y
{-# INLINE unsafeAddWithScaleM_ #-}

-- | @kroneckerTo dst x y@ sets @dst := x \otimes y@.
kroneckerTo :: (RVector v1, RVector v2, BLAS2 e)
            => STVector s e -> v1 e -> v2 e -> ST s ()
kroneckerTo dst x y = do
    m <- getDim x
    n <- getDim y
    dimdst <- getDim dst

    when (dimdst /= m * n) $ error $
        printf ("kroneckerTo"
                ++ " <vector with dim %d>"
                ++ " <vector with dim %d>"
                ++ " <vector with dim %d>:"
                ++ " dimension mismatch") dimdst m n

    clear dst
    unsafeIOToST $
        unsafeWith dst $ \pdst ->
        unsafeWith x $ \px ->
        unsafeWith y $ \py ->
            BLAS.geru n m 1 py 1 px 1 pdst (max n 1)

call2 :: (RVector x, RVector y, Storable e, Storable f)
      => (Int -> Ptr e -> Ptr f -> IO a)
      -> x e -> y f -> ST s a
call2 f x y = do
    n <- getDim x
    unsafeIOToST $
        unsafeWith x $ \pX ->
        unsafeWith y $ \pY ->
            f n pX pY
{-# INLINE call2 #-}

call3 :: (RVector x, RVector y, RVector z, Storable e, Storable f, Storable g)
      => (Int -> Ptr e -> Ptr f -> Ptr g -> IO a)
      -> x e -> y f -> z g -> ST s a
call3 f x y z = do
    n <- getDim x
    unsafeIOToST $
        unsafeWith x $ \pX ->
        unsafeWith y $ \pY ->
        unsafeWith z $ \pZ ->
            f n pX pY pZ
{-# INLINE call3 #-}

strideCall :: (RVector x, Storable e)
           => (Int -> Ptr e -> Int -> IO a)
           ->  x e -> ST s a
strideCall f x = do
    n <- getDim x
    unsafeIOToST $
        unsafeWith x $ \pX ->
            f n pX incX
  where
    incX = 1
{-# INLINE strideCall #-}

strideCall2 :: (RVector x, RVector y, Storable e, Storable f)
            => (Int -> Ptr e -> Int -> Ptr f -> Int -> IO a)
            -> x e -> y f -> ST s a
strideCall2 f x y = do
    n <- getDim x
    unsafeIOToST $
       unsafeWith x $ \pX ->
       unsafeWith y $ \pY ->
           f n pX incX pY incY
  where
    incX = 1
    incY = 1
{-# INLINE strideCall2 #-}

checkOp2 :: (RVector x, RVector y, Storable e, Storable f)
         => String
         -> (x e -> y f -> ST s a)
         -> x e
         -> y f
         -> ST s a
checkOp2 str f x y = do
    n1 <- getDim x
    n2 <- getDim y

    when (n1 /= n2) $ error $
        printf ("%s <vector with dim %d> <vector with dim %d>:"
                ++ " dimension mismatch") str n1 n2

    f x y
{-# INLINE checkOp2 #-}

checkOp3 :: (RVector x, RVector y, RVector z, Storable e, Storable f, Storable g)
         => String
         -> (x e -> y f -> z g -> ST s a)
         -> x e
         -> y f
         -> z g
         -> ST s a
checkOp3 str f x y z = do
    n1 <- getDim x
    n2 <- getDim y
    n3 <- getDim z

    when (n1 /= n2 || n1 /= n3) $ error $
        printf ("%s <vector with dim %d> <vector with dim %d>"
                ++ " <vector with dim %d>:"
                ++ " dimension mismatch") str n1 n2 n3

    f x y z
{-# INLINE checkOp3 #-}

max_stars_repo_name: patperry/hs-linear-algebra

max_issues_repo_name: cartazio/hs-cblas

max_forks_repo_name: cartazio/hs-cblas


text:
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS -Wall -fwarn-tabs -fno-warn-type-defaults -fno-warn-unused-do-bind #-}
-- The package juicypixels needs to have been installed.
import Data.Complex
import Codec.Picture

type C = Complex Double
type Pnt = (Double, Double)
type Grid a = [[a]]

w :: Int
w = 500

h :: Int
h = 500

-- Complex constants for the Julia set rendering.
fixC1, fixC2, fixC3, fixC4 :: C
fixC1 = (-1.037) :+ 0.17
fixC2 = (-0.52) :+ 0.57
fixC3 = 0.295 :+ 0.55
fixC4 = (-0.624) :+ 0.435

juliaConstants :: [(Int, C)]
juliaConstants =[ (1, fixC1), (2, fixC2), (3, fixC3), (4, fixC4)]

--iterate f x == [x, f x, f (f x), ...]
-- in mandel the seed is fixed at 0
mandel :: C -> [C]
mandel c = iterate (\z -> z * z + c)  (0.0 :+ 0.0)

--iterate f x == [x, f x, f (f x), ...]
-- in julia, for a given fixed constant,  the seed varies
julia :: C ->  C -> [C]
julia fixc  = iterate (\z -> z * z + fixc)

closeToZero :: C -> Bool
closeToZero z = magnitude z < 5

-- get the pallete value corresponding to the 'depth' at which the value fails the test
chooseColor  :: [color] -> ( C -> [C]) -> C -> color
chooseColor colors fz z = (colors !!) . length . take n . takeWhile closeToZero $ fz z
  where n = length colors - 1

-- A grid of points that is defined by its  bottom left, top right and the number of subdivision per row and per column
grid :: Int -> Int -> Pnt -> Pnt -> Grid Pnt
grid c r (xmin, ymin) (xmax, ymax ) = [[(x, y) | x <- for c xmin xmax ] | y <- for r ymin ymax ] where
    for :: Int ->Double ->Double ->[Double]
    for n mn mx = take n [mn, mn + delta ..]
     where delta = (mx - mn) / fromIntegral (n - 1)

-- a grid
grid1 :: Grid Pnt
grid1 = grid w h (-1.5, -1.5) (1.5, 1.5)

-- calculate sequences over the grid and for each sequence get an 'end' colour
mapOverGrid :: Grid Pnt -> (C -> [C]) -> [color] -> [[color]]
mapOverGrid grd fz colors = [ [chooseColor colors  fz  (x :+ y) | (x,y) <- row] | row <- grd ]

-- used by juicy pixels to get a pixel value at x, y
gridFunc ::Grid Pnt ->  (C -> [C]) -> Int -> Int -> PixelRGB8
gridFunc grd fz  x y = ((mapOverGrid grd fz rgb8) !! y) !! x

generateImg :: (Int -> Int -> PixelRGB8) -> DynamicImage
generateImg gf = ImageRGB8 (generateImage gf w h)

-- Juicy pixels type
rgb8 :: [PixelRGB8]
rgb8 = foldr (\x ac  -> PixelRGB8 x x x  : ac) [] [255, 254..0]

allJulia :: IO()
allJulia =
    mapM_ makeSaveImage juliaConstants where
        makeSaveImage :: (Int, C) -> IO()
        makeSaveImage (n, c) =
         savePngImage name $ generateImg . gridFunc grid1 . julia $ c where
            name = "imageJ" ++ show n ++ ".png"

main :: IO ()
main = do
    allJulia
    savePngImage "imageM.png" $ generateImg . gridFunc grid1 $ mandel
    print "Done."

max_stars_repo_name: mike-k-houghton/fractals

max_issues_repo_name: mike-k-houghton/fractals

max_forks_repo_name: mike-k-houghton/fractals


text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}

{-
This allows us to simplify the constraints we need to specify up front for
KnownNat constraints. For example, the following function

f :: forall n . (KnownNat n, KnownNat (n+2)) => Proxy n -> Integer
f _ = natVal (Proxy :: Proxy n) +
      natVal (Proxy :: Proxy (n+2))

can be simplified to

f :: forall n . KnownNat n => Proxy n -> Integer
f _ = natVal (Proxy :: Proxy n) +
      natVal (Proxy :: Proxy (n+2))
-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

-- Allows us to equate terms like (a + b) + c and a + (b + c) as the same thing.
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

module Main where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as LA
import           Data.Proxy
import qualified Numeric.LinearProgramming.L1 as L1
import qualified Numeric.LinearProgramming.L1.Static as L1S
import Data.Maybe (fromJust)
import System.Random
import qualified Data.Vector.Storable as V

f :: forall n . KnownNat n => Proxy n -> Integer
f _ = natVal (Proxy :: Proxy n) +
      natVal (Proxy :: Proxy (n+2))

l1_example :: (R 1000, LA.Vector ℝ, R 1000, LA.Vector ℝ)
l1_example = (signal_static, x, x_static, lsq_sol)
  where
    m = 200
    n = 1000
    s = 30 -- Sparsity, number of non-zero elements in the signal

    -- Random state transition matrix
    a = LA.reshape n $ LA.randomVector 0 LA.Gaussian (n*m)

    -- Create a signal with `s` elements being either -1 or 1.
    g = mkStdGen 0
    indx = take s $ randomRs (0, n - 1) g :: [Int]
    r = take s $ fmap ([-1, 1] !!) $ randomRs (0, 1) g
    signal = V.replicate n 0 V.// zip indx r
    b = a LA.#> signal

    -- Make dependently typed versions of the signal.
    a_static = fromJust $ create a :: L 200 1000
    signal_static = fromJust $ create signal :: R 1000
    b_static = a_static #> signal_static

    λ = 1
    x = L1.l1Solve λ a b
    x_static = L1S.l1Solve λ a_static b_static

    lsq_sol = a LA.<\> b

originalSignal :: R 1000
l1Sol :: LA.Vector ℝ
l1SolStatic :: R 1000
lsqComparisonSol :: V.Vector ℝ
(originalSignal, l1Sol, l1SolStatic, lsqComparisonSol) = l1_example

main :: IO ()
main = do
  putStrLn "Hello, World!"
  print $ f (Proxy @3)

max_stars_repo_name: ryanorendorff/convex

max_issues_repo_name: ryanorendorff/convex

max_forks_repo_name: ryanorendorff/convex


text:
module Language.Type
  ( Type(..)
  , type HaskellType
  , Scalar(..)
  , TypeProxy(..)
  , SomeType(..)
  , sameHaskellType
  , Int64
  , Symbol
  , Complex(..)
  , pattern Boolean_
  , pattern Integer_
  , pattern Real_
  , pattern Complex_
  , pattern Rational_
  , pattern Color_
  , showType
  , showValue
  , KnownType(..)
  , withKnownType
  ) where

import Data.Int
import GHC.TypeLits
import Data.Type.Equality ((:~:)(..))
import Data.Color (Color, colorToRGB)
import Data.Complex
import Data.List (intercalate)

data Type
  = VoidT
  | BooleanT
  | IntegerT
  | RealT
  | ComplexT
  | RationalT
  | Pair Type Type
  | ColorT
  | ImageT
  | ListT Type
  deriving (Eq, Ord, Show)

-- | Constant values for scalar types
type family HaskellType (t :: Type) :: * where
  HaskellType 'BooleanT   = Bool
  HaskellType 'IntegerT   = Int64
  HaskellType 'RealT      = Double
  HaskellType 'ComplexT   = Complex Double
  HaskellType 'RationalT  = (Int64, Int64)
  HaskellType 'ColorT     = Color
  HaskellType ('Pair x y) = (HaskellType x, HaskellType y)
  HaskellType 'VoidT      = ()
  HaskellType 'ImageT     = Int
  HaskellType ('ListT x)  = [HaskellType x]

-- | Constant values for scalar types. Match on the
-- first argument to make the type of the second argument
-- known.
data Scalar (t :: Type) where
  Scalar :: forall t. TypeProxy t -> HaskellType t -> Scalar t

instance Eq (Scalar t) where
  Scalar t x == Scalar _ y = case t of
    BooleanType  -> x == y
    IntegerType  -> x == y
    RealType     -> x == y
    ComplexType  -> x == y
    RationalType -> x == y
    ColorType    -> x == y
    PairType t1 t2 ->
      let (x1, x2) = x
          (y1, y2) = y
      in (Scalar t1 x1, Scalar t2 x2) == (Scalar t1 y1, Scalar t2 y2)
    VoidType     -> x == y
    ImageType    -> x == y
    ListType it -> map (Scalar it) x == map (Scalar it) y

instance Ord (Scalar t) where
  compare (Scalar t x) (Scalar _ y) = case t of
    BooleanType  -> compare x y
    IntegerType  -> compare x y
    RealType     -> compare x y
    ComplexType  -> compare (realPart x, imagPart x) (realPart y, imagPart y)
    RationalType -> compare x y
    ColorType    -> compare x y
    PairType t1 t2 ->
      let (x1, x2) = x
          (y1, y2) = y
      in compare (Scalar t1 x1, Scalar t2 x2) (Scalar t1 y1, Scalar t2 y2)
    VoidType     -> compare x y
    ImageType    -> compare x y
    ListType it -> compare (map (Scalar it) x) (map (Scalar it) y)

sameHaskellType :: TypeProxy t1 -> TypeProxy t2 -> Maybe (t1 :~: t2)
sameHaskellType v1 v2 = case v1 of
  BooleanType  -> case v2 of { BooleanType  -> Just Refl; _ -> Nothing }
  IntegerType  -> case v2 of { IntegerType  -> Just Refl; _ -> Nothing }
  RealType     -> case v2 of { RealType     -> Just Refl; _ -> Nothing }
  ComplexType  -> case v2 of { ComplexType  -> Just Refl; _ -> Nothing }
  RationalType -> case v2 of { RationalType -> Just Refl; _ -> Nothing }
  ColorType    -> case v2 of { ColorType    -> Just Refl; _ -> Nothing }
  VoidType     -> case v2 of { VoidType     -> Just Refl; _ -> Nothing }
  ImageType    -> case v2 of { ImageType    -> Just Refl; _ -> Nothing }
  PairType x y -> case v2 of
    PairType x' y' -> case (,) <$> sameHaskellType x x' <*> sameHaskellType y y' of
      Just (Refl, Refl) -> Just Refl
      Nothing           -> Nothing
    _ -> Nothing
  ListType x -> case v2 of
    ListType x' -> case sameHaskellType x x' of
      Just Refl -> Just Refl
      Nothing   -> Nothing
    _ -> Nothing

-- | Singleton values reflecting the type-level parameter @t@.
data TypeProxy (t :: Type) where
  BooleanType  :: TypeProxy 'BooleanT
  IntegerType  :: TypeProxy 'IntegerT
  RealType     :: TypeProxy 'RealT
  ComplexType  :: TypeProxy 'ComplexT
  RationalType :: TypeProxy 'RationalT
  ColorType    :: TypeProxy 'ColorT
  PairType     :: (KnownType x, KnownType y) => TypeProxy x -> TypeProxy y -> TypeProxy ('Pair x y)
  VoidType     :: TypeProxy 'VoidT
  ImageType    :: TypeProxy 'ImageT
  ListType     :: KnownType x => TypeProxy x -> TypeProxy ('ListT x)

data SomeType where
  SomeType :: forall t. TypeProxy t -> SomeType

instance Show SomeType where
  show (SomeType t) = showType t

instance Eq SomeType where
  SomeType t1 == SomeType t2 = maybe False (const True) (sameHaskellType t1 t2)

class KnownType (t :: Type)   where typeProxy :: TypeProxy t
instance KnownType 'BooleanT  where typeProxy = BooleanType
instance KnownType 'IntegerT  where typeProxy = IntegerType
instance KnownType 'RealT     where typeProxy = RealType
instance KnownType 'ComplexT  where typeProxy = ComplexType
instance KnownType 'RationalT where typeProxy = RationalType
instance KnownType 'ColorT    where typeProxy = ColorType
instance KnownType 'VoidT     where typeProxy = VoidType
instance KnownType 'ImageT    where typeProxy = ImageType
instance KnownType x => KnownType ('ListT x) where
  typeProxy = ListType (typeProxy @x)
instance (KnownType x, KnownType y) => KnownType ('Pair x y) where
  typeProxy = PairType (typeProxy @x) (typeProxy @y)

withKnownType :: TypeProxy t -> (KnownType t => a) -> a
withKnownType ty k = case ty of
  BooleanType  -> k
  IntegerType  -> k
  RealType     -> k
  ComplexType  -> k
  RationalType -> k
  ColorType    -> k
  VoidType     -> k
  ImageType    -> k
  PairType {}  -> k
  ListType {}  -> k

pattern Boolean_ :: forall (t :: Type). () => (t ~ 'BooleanT) => HaskellType t -> Scalar t
pattern Boolean_ x = Scalar BooleanType x

pattern Integer_ :: forall (t :: Type). () => (t ~ 'IntegerT) => HaskellType t -> Scalar t
pattern Integer_ x = Scalar IntegerType x

pattern Real_ :: forall (t :: Type). () => (t ~ 'RealT) => HaskellType t -> Scalar t
pattern Real_ x    = Scalar RealType x

pattern Complex_ :: forall (t :: Type). () => (t ~ 'ComplexT) => HaskellType t -> Scalar t
pattern Complex_  pair = Scalar ComplexType pair

pattern Rational_ :: forall (t :: Type). () => (t ~ 'RationalT) => HaskellType t -> Scalar t
pattern Rational_ pair = Scalar RationalType pair

pattern Color_ :: forall (t :: Type). () => (t ~ 'ColorT) => HaskellType t -> Scalar t
pattern Color_ c = Scalar ColorType c

showType :: TypeProxy t -> String
showType = \case
  BooleanType  -> "Boolean"
  IntegerType  -> "Integer"
  RealType     -> "Real"
  ComplexType  -> "Complex"
  RationalType -> "Rational"
  ColorType    -> "Color"
  VoidType     -> "Unit"
  ImageType    -> "Image"
  PairType x y -> "(" <> showType x <> " x " <> showType y <> ")"
  ListType x   -> "List " <> showType x

showValue :: TypeProxy t -> HaskellType t -> String
showValue ty v = case ty of
  BooleanType  -> if v then "true" else "false"
  IntegerType  -> show v
  RealType     -> show v
  ComplexType  -> let x :+ y = v
                   in show x <> " + " <> show y <> "i"
  RationalType -> let (x, y) = v
                   in show x <> " / " <> show y
  ColorType    -> show (colorToRGB v)
  VoidType     -> "n/a"
  ImageType    -> "(image)"
  PairType xt yt -> let (x, y) = v
                     in showValue xt x <> " , " <> showValue yt y
  ListType xt -> "[" ++ intercalate ", " (map (showValue xt) v) ++ "]"

max_stars_repo_name: matt-noonan/FractalStream

max_issues_repo_name: matt-noonan/FractalStream

max_forks_repo_name: matt-noonan/FractalStream


text:
-- % Haskell for Numerics?
-- % Dominic Steinitz
-- % 2nd June 2017

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE NumDecimals         #-}

{-# OPTIONS_GHC -Wall #-}

module Main (main) where

import           Numeric.Sundials
import           Numeric.LinearAlgebra
import           Data.Csv
import           Data.Char
import           Data.ByteString.Lazy (putStr, writeFile)
import           Prelude hiding (putStr, writeFile)
import           Control.Exception
import           Data.Coerce
import           Katip.Monadic
import           GHC.Int

bigN :: Int
bigN = 201

k :: Double
k = 0.5

deltaX :: Double
deltaX = 1.0 / (fromIntegral bigN - 1)
c1, c2 :: Double
c1 = k / deltaX / deltaX
c2 = (-2.0) * k / deltaX / deltaX

t0 :: Double
t0 = 0.0

tf :: Double
tf =1.0

bigNt :: Int
bigNt = 10

rtol :: Double
rtol = 1.0e-6

atol :: Double
atol = 1.0e-10

bigA :: Matrix Double
bigA = assoc (bigN, bigN) 0.0 [ ((i, j), f (i, j)) | i <- [0 .. bigN - 1]
                                                   , j <- [0 .. bigN - 1]
                        ]
 where
   f (i, j) | i     == 0        = 0.0    -- left boundary condition
            | i     == bigN - 1 = 0.0    -- right boundary condition
            | i     == j        = c2
            | i - 1 == j        = c1
            | i + 1 == j        = c1
            | otherwise         = 0.0

b :: Vector Double
b = assoc bigN 0.0 [ (iSource, 0.01 / deltaX) ]
  where
    iSource = bigN `div` 2

bigU0 :: Vector Double
bigU0 = assoc bigN 0.0 []

deltaT :: Double
deltaT = (tf - t0) / (fromIntegral bigNt)

sol :: IO (Matrix Double)
sol = do
  x <- runNoLoggingT $ solve (defaultOpts SDIRK_5_3_4) heat1D
  case x of
    Left e  -> error $ show e
    Right y -> return (solutionMatrix y)

heat1D :: OdeProblem
heat1D = emptyOdeProblem
  { odeRhs = odeRhsPure $ \_t x -> coerce (bigA #> (coerce x) + b)
  , odeJacobian = Just (\_t _x -> bigA)
  , odeEvents = mempty
  , odeEventHandler = nilEventHandler
  , odeMaxEvents = 0
  , odeInitCond = bigU0
  , odeSolTimes =vector $ map (deltaT *) [0 .. 10]
  , odeTolerances = defaultTolerances
  }

myOptions :: EncodeOptions
myOptions = defaultEncodeOptions {
      encDelimiter = fromIntegral (ord ' ')
    }

main :: IO ()
main = do
  x <- sol
  writeFile "heat1G.txt" $ encodeWith myOptions $ map toList $ toRows x

defaultOpts :: method -> ODEOpts method
defaultOpts method = ODEOpts
  { maxNumSteps = 1e5
  , minStep     = 1.0e-14
  , fixedStep   = 0
  , maxFail     = 10
  , odeMethod   = method
  , initStep    = Nothing
  , jacobianRepr = SparseJacobian
                 $ SparsePattern
                 $ cmap (fromIntegral :: I -> Int8)
                 $ cmap (\x -> case x of 0 -> 0; _ -> 1)
                 $ flatten
                 $ toInt bigA
  }

instance Element Int8

emptyOdeProblem :: OdeProblem
emptyOdeProblem = OdeProblem
      { odeRhs = error "emptyOdeProblem: no odeRhs provided"
      , odeJacobian = Nothing
      , odeInitCond = error "emptyOdeProblem: no odeInitCond provided"
      , odeEvents = mempty
      , odeTimeBasedEvents = TimeEventSpec $ return $ 1.0 / 0.0
      , odeEventHandler = nilEventHandler
      , odeMaxEvents = 100
      , odeSolTimes = error "emptyOdeProblem: no odeSolTimes provided"
      , odeTolerances = defaultTolerances
      }

nilEventHandler :: EventHandler
nilEventHandler _ _ _ = throwIO $ ErrorCall "nilEventHandler"

defaultTolerances :: Tolerances
defaultTolerances = Tolerances
  { absTolerances = Left 1.0e-6
  , relTolerance = 1.0e-10
  }

max_stars_repo_name: idontgetoutmuch/Diffusions

max_issues_repo_name: idontgetoutmuch/Diffusions

max_forks_repo_name: idontgetoutmuch/Diffusions


text:
{-# LANGUAGE ForeignFunctionInterface #-}
module Q.Options.ImpliedVol.LetsBeRational (
  euImpliedVol
) where

import           Data.Coerce                    (coerce)
import           Data.Number.Erf
import           Foreign.C.Types
import           Numeric.IEEE                   (epsilon, maxFinite, minNormal)
import           Q.Options.BlackScholes
import           Q.Options
import           Q.Types
import           Statistics.Distribution        (cumulative, density, quantile)
import           Statistics.Distribution.Normal (standard)

foreign import ccall
   "lets_be_rational.h implied_volatility_from_a_transformed_rational_guess" c_lbr ::
     CDouble -> CDouble  -> CDouble -> CDouble  -> CDouble  -> CDouble

euImpliedVol :: OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVol cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p) =
  coerce $ c_lbr (CDouble p) (CDouble f) (CDouble k) (CDouble t) (CDouble (cpi cp))

max_stars_repo_name: ghais/lowq

max_issues_repo_name: ghais/lowq

max_forks_repo_name: ghais/lowq


text:
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE OverloadedStrings   #-}
{-|
Module      : Grenade.Demos.MNIST
Description : Example NN recognizing handwritten digits

The MNIST database, is a database of handwritten digits
with 700000 examples. A more comprehensive introduction
to the database together with the full dataset can be
found here: http://yann.lecun.com/exdb/mnist/
-}
module Grenade.Demos.MNIST where

import qualified Data.ByteString              as B
import           Data.List                    (maximumBy)
import           Data.Maybe                   (fromMaybe)
import           Data.Ord                     (comparing)
import           Data.Serialize               (Get, runGet, get)
import qualified Data.Vector.Storable         as V
import qualified Data.Vector.Unboxed          as UB
import qualified Data.Vector                  as DV
import           Data.Word

import           Grenade
import           Grenade.Layers.Internal.Shrink (shrink_2d_rgba)

import qualified Numeric.LinearAlgebra.Static as SA

netLoad :: FilePath -> IO MNIST
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $ runGet (get :: Get MNIST) modelData

runNet'' :: MNIST -> UB.Vector Word8 -> Int -> String
runNet'' net vec = runNet' net (DV.convert vec)

runNet' :: MNIST -> V.Vector Word8 -> Int -> String
runNet' net m d = (\(S1D ps) -> let (p, i) = (getProb . V.toList) (SA.extract ps)
                              in "This number is " ++ show i ++ " with probability " ++ (show (round (p * 100) :: Int)) ++ "%") $ runNet net (conv m)
  where
    getProb :: (Ord a) => [a] -> (a, Int)
    getProb xs = maximumBy (comparing fst) (Prelude.zip xs [0..])

    conv :: V.Vector Word8 -> S ('D2 28 28)
    conv m = S2D $ fromMaybe (error "") $ SA.create $ shrink_2d_rgba d d 28 28 m

type MNIST
  = Network
    '[ Convolution 'WithoutBias 'NoPadding 1 10 5 5 1 1
     , Pooling 2 2 2 2
     , Relu
     , Convolution 'WithoutBias 'NoPadding 10 16 5 5 1 1
     , Pooling 2 2 2 2
     , Reshape
     , Relu
     , FullyConnected 256 80
     , Logit
     , FullyConnected 80 10
     , Logit]
    '[ 'D2 28 28
     , 'D3 24 24 10
     , 'D3 12 12 10
     , 'D3 12 12 10
     , 'D3 8 8 16
     , 'D3 4 4 16
     , 'D1 256
     , 'D1 256
     , 'D1 80
     , 'D1 80
     , 'D1 10
     , 'D1 10]

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade


text:
import Numeric.GSL
import Numeric.GSL.Special
import Numeric.LinearAlgebra
import Prelude hiding (catch)
import Control.Exception

test x = catch
       (print x)
       (\e -> putStrLn $ "captured ["++ show (e :: SomeException) ++"]")

main = do
    setErrorHandlerOff

    test $ log_e (-1)
    test $ 5 + (fst.exp_e) 1000
    test $ bessel_zero_Jnu_e (-0.3) 2

    test $ (linearSolve 0 4 :: Matrix Double)
    test $ (linearSolve 5 (sqrt (-1)) :: Matrix Double)

    putStrLn "Bye"

max_stars_repo_name: curiousleo/liquidhaskell

max_issues_repo_name: curiousleo/liquidhaskell

max_forks_repo_name: curiousleo/liquidhaskell