text:
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}

module Vector where

import Tensor
import Data.Monoid
import Data.Maybe
import Data.Complex
import Helper
import Data.Either
import System.Random

data Vector a = Scalar a | Row [Vector a] | Column [Vector a]

-- This Eq implementation ignores structural differences
instance (Eq a) => Eq (Vector a) where
    x == y = wx == wy && hx == hy && go 0 0
        where
            wx = width x
            wy = width y
            hx = height x
            hy = height y
            go i j
                | i == hx = True
                | j == wx = go (i + 1) 0
                | otherwise = atM x (i, j) == atM y (i, j) && go i (j + 1)

instance Functor Vector where
    fmap f (Scalar s) = Scalar $ f s
    fmap f (Row t) = Row $ map (fmap f) t
    fmap f (Column t) = Column $ map (fmap f) t

instance Applicative Vector where
    pure = Scalar
    Scalar s <*> t = fmap s t
    t <*> (Scalar s) = fmap ($s) t
    Row r <*> t = Row $ map (<*> t) r
    Column c <*> t = Column $ map (<*> t) c

instance Monad Vector where
    Scalar s >>= f = f s
    Row r >>= f = Row $ map (>>= f) r
    Column c >>= f = Column $ map (>>= f) c

instance Num a => Num (Vector a) where
    Row x + (Row y) = Row $ zipWith (+) x y
    Column x + (Column y) = Column $ zipWith (+) x y
    Scalar x + y = fmap (+ x) y
    x + (Scalar y) = fmap (+ y) x
    _+_ = error "Vector Addition: Dimension mismatch"
    Scalar x * y = fmap (* x) y
    x * Scalar y = fmap (* y) x
    Row x * (Column y) = getSum $ foldMap Sum $ zipWith (*) x y
    Column x * (Column y) = Column $ map (* Column y) x
    x * (Row y) = Row $ map (x *) y
    abs = fmap abs
    signum = fmap signum
    fromInteger = Scalar . fromInteger
    negate = fmap negate

instance Foldable Vector where
  foldr f acc (Scalar s) = f s acc
  foldr f acc (Row r) = foldr (flip (foldr f)) acc r
  foldr f acc (Column c) = foldr (flip (foldr f)) acc c

instance Traversable Vector where
    traverse f (Scalar s) = Scalar <$> f s
    traverse f (Row r) = Row <$> traverse (traverse f) r
    traverse f (Column c) = Column <$> traverse (traverse f) c

instance (Num a) => Semigroup (Vector a) where
    x <> y = (*) <$> x <*> y

instance (Num a) => Monoid (Vector a) where
    mempty = Scalar 1

instance (Cmplx a) => Cmplx (Vector a) where
    conj = fmap conj . transposeStruct

instance (Num a, Cmplx a) => Tensor Vector a where
    inner = dot
    -- TODO: norm for matrices
    norm t = (/ (sqrt.getSum $ foldMap (Sum . (**2) . abs) t)) <$> t
    fromList xs dim = go (foldr (con . Scalar) (Column []) xs) dim
        where
            go (Column c) n
                | n == 1 = Column [go (Row c) dim]
                | l <= fromIntegral dim = Column c
                | otherwise = go (Row $ take (l `div` fromIntegral n) c) dim
                    `con` go (Column $ drop (l `div` fromIntegral n) c) (n - 1)
                where l = length c

            go (Row r) n
                | n == 1 = Row [go (Column r) dim]
                | l <= fromIntegral dim = Row r
                | otherwise = go (Column $ take (l `div` fromIntegral n) r) dim
                    `con` go (Row $ drop (l `div` fromIntegral n) r) (n - 1)
                where l = length r

            go s n = error "error when adding tensor rank"
    rank (Scalar s) = 0
    rank (Row []) = 1
    rank (Column []) = 1
    rank (Row (r:rs)) = rank r + 1
    rank (Column (c:cs)) = rank c + 1

    isSquare t
        | (Scalar _) <- t = True
        | (Column c) <- t = getAll $ go (Column c) (length c)
        | (Row r) <- t = getAll $ go (Row r) (length r)
        where
            go (Column c) l
                | length c == l = foldMap (`go` l) c
                | otherwise = All False
            go (Row r) l
                | length r == l = foldMap (`go` l) r
                | otherwise = All False
            go (Scalar _) _ = All True

    conjTr = conj

instance (Floating a, Eq a, Show a, Ord a, UniformRange a, Cmplx a) => QuantumRegister Vector a where

    densityOp x = transposeStruct x * x

    initQubit zero one
        | abs zero ** 2 + abs one ** 2 == 1 = Right $ fromList [zero, one] 2
        | otherwise = Left "Sum of probability of |0> and |1> state does not equal to 1"

    initQubit0 = fromList [1, 0] 2

    initNumQubit0 n
        | n < 1 = Left "Invalid qubit number"
        | otherwise = Right $ foldr (|*|) initQubit0 $ replicate (fromIntegral n - 1) initQubit0

    toQuantumRegister q
        | isLeft num = Left "Input list length is not a power of 2"
        | otherwise = Right . norm $ fromList q 2
        where num = integralLog 2 (length q)

    illegalPeek (Scalar s) = show s
    illegalPeek q = go 0 0
        where
            w = width q
            h = height q
            go i j
                | i == h = mempty
                | j == w = "\t|\n" ++ go (i + 1) 0
                | j == 0 = "|\t" <> current <> "\t" <> go i (j + 1)
                | otherwise = current <> "\t" <> go i (j + 1)
                where
                    current = maybe "null" show (atM q (i, j))

    subSystem lst q
        | getAny $ foldMap (Any . (\ x -> x >= rank q || x < 0)) lst = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid Quantum Register"
        | otherwise = Right . fmap sqrt $ go lst (fmap (\x -> conj x * x) q)
        where
            go [] (Scalar s) = Scalar s
            go is (Row r)
                | null is || Prelude.head is /= 0 = foldr (+) (Scalar 0) preSum
                | otherwise = Row $ map (go $ fmap (flip (-) 1) (Prelude.tail is)) r
                where
                    preSum = map (go $ fmap (flip (-) 1) is) r
            go is (Column c)
                | null is || Prelude.head is /= 0 = foldr (+) (Scalar 0) preSum
                | otherwise = Column $ map (go $ fmap (flip (-) 1) (Prelude.tail is)) c
                where
                    preSum = map (go $ fmap (flip (-) 1) is) c
            go _ _ = error "Exceptions not handled by subSystem function"

    isEntangled lst q = do
            reduced <- reducedDensity lst q
            return $ abs (trace (reduced * reduced) - 1) > 1e-13

    mask0 = Row [
                Column[1, 0],
                Column[0, 0]
            ]

    mask1 = Row [
                Column[0, 0],
                Column[0, 1]
            ]

    pauliId = Row [
                Column [1, 0],
                Column [0, 1]
            ]

instance (RealFloat a, Show a, UniformRange a) => Gates Vector (Complex a) where
    pauliX = Row [
            Column [Scalar (0 :+ 0), Scalar (1 :+ 0)],
            Column [Scalar (1 :+ 0), Scalar (0 :+ 0)]
        ]

    pauliY = Row [
                Column [Scalar (0 :+ 0), Scalar (0 :+ 1)],
                Column [Scalar (0 :+ (-1)), Scalar (0 :+ 0)]
            ]

    pauliZ = Row [
                Column [Scalar (1 :+ 0), Scalar (0 :+ 0)],
                Column [Scalar (0 :+ 0), Scalar ((-1) :+ 0)]
            ]

    hadamard = Row [
                Column [Scalar (1 / sqrt 2 :+ 0), Scalar (1 / sqrt 2 :+ 0)],
                Column [Scalar (1 / sqrt 2 :+ 0), Scalar ((-1) / sqrt 2 :+ 0)]
            ]

at :: Integral b => a -> Vector a -> (b, b) -> a
at x y z = fromMaybe x (atM y z)

atM :: Integral b => Vector a -> (b, b) -> Maybe a
atM (Scalar s) (0, 0) = Just s
atM (Column (c:cs)) (i, j)
    | height (Column (c:cs)) > i = let offset = height c in
                                    if offset > i
                                    then atM c (i, j)
                                    else atM (Column cs) (i - offset, j)
    | otherwise = Nothing
atM (Row (r:rs)) (i, j)
    | width (Row (r:rs)) > j = let offset = width r in
                                    if offset > j
                                    then atM r (i, j)
                                    else atM (Row rs) (i, j - offset)
    | otherwise = Nothing
atM _ _ = Nothing

con :: Vector a -> Vector a -> Vector a
con x (Row y) = Row $ x : y
con x (Column y) = Column $ x : y
con _ _ = error "Appending to Scalar is prohibited"

append :: Vector a -> Vector a -> Vector a
append (Row x) (Row y) = Row $ x ++ y
append (Column x) (Column y) = Column $ x ++ y
append _ _ = error "Vector direction mismatch"

width :: Integral b => Vector a -> b
width (Scalar s) = 1
width (Row r) = getSum $ foldMap (Sum . width) r
width (Column c) = width $ Prelude.head c

height :: Integral b => Vector a -> b
height (Scalar s) = 1
height (Column c) = getSum $ foldMap (Sum . height) c
height (Row r) = height $ Prelude.head r

-- Transform structure retaining the original rank and shape
transformStruct :: Vector a -> Vector a
transformStruct (Row (Column c:cs)) = Column . go $ Column c:cs
    where
        go (Row [] : _) = []
        go (Column [] : _) = []
        go cs = Row (map (transformStruct . Vector.head) cs)
            : go (map Vector.tail cs)

transformStruct (Column (Row r:rs)) = Row . go $ Row r:rs
    where
        go (Row [] : _) = []
        go (Column [] : _) = []
        go rs = Column (map (transformStruct . Vector.head) rs)
            : go (map Vector.tail rs)

transformStruct (Column cs) = Column $ map transformStruct cs
transformStruct (Row rs) = Row $ map transformStruct rs
transformStruct t = t

-- transpose via structure
transposeStruct :: Vector a -> Vector a
transposeStruct (Scalar s) = Scalar s
transposeStruct (Row r) = Column $ map transposeStruct r
transposeStruct (Column c) = Row $ map transposeStruct c

-- transpose retaining structure
transpose :: Vector a -> Vector a
transpose = transformStruct . transposeStruct

head :: Vector a -> Vector a
head (Scalar s) = Scalar s
head (Row r) = Prelude.head r
head (Column c) = Prelude.head c

last :: Vector a -> Vector a
last (Scalar s) = Scalar s
last (Row r) = Prelude.last r
last (Column c) = Prelude.last c

tail :: Vector a -> Vector a
tail (Scalar s) = Scalar s
tail (Row r) = Row $ Prelude.tail r
tail (Column c) = Column $ Prelude.tail c

init :: Vector a -> Vector a
init (Scalar s) = Scalar s
init (Row r) = Row $ Prelude.init r
init (Column c) = Column $ Prelude.init c

dot :: Num a => Vector a -> Vector a -> Vector a
dot x y = transposeStruct x * y

reducedDensity :: (Num a, Cmplx a, Floating a) => [Int] -> Vector a -> Either String (Vector a)
reducedDensity lst q
    | getAny $ foldMap (Any . (\ x -> x >= rank q || x < 0)) lst = Left "Index out of bounds"
    | not (isSquare q) = Left "Invalid Quantum Register"
    | otherwise = Right $ go lst q (transposeStruct q)
    where
        go [] (Scalar s) (Scalar t) = Scalar (s * t)
        go is (Column c) (Row r)
            | null is || Prelude.head is /= 0 = (getSum . foldMap Sum) $ zipWith (go (flip (-) 1 <$> is)) c r
            | otherwise = Row $ map (Column . ($ c) . (\ b a -> map (\ x -> go (flip (-) 1 <$> Prelude.tail is) x b) a)) r
        go _ _ _ = error "Input is not a valid statevector"

isColumn :: Vector a -> Bool
isColumn (Column _) = True
isColumn _ = False

isRow :: Vector a -> Bool
isRow (Row _) = True
isRow _ = False

isScalar :: Vector a -> Bool
isScalar (Scalar _) = True
isScalar _ = False

trace :: (Num a) => Vector a -> a
trace (Scalar s) = s
trace (Row r) = (getSum . foldMap Sum) $ zipWith go [0, 1 ..] r
    where
        go i (Row r) = trace (r !! i)
        go i (Column c) = trace (c !! i)
        go i (Scalar s) = s
{-
    Tensor type:
    All one-dimensional tensors are column vectors. This means matrices,
    which are rank 2 tensors, compose of multiple column tensors, which
    may not follow conventional representation in other traversable types.
-}

-- data Tensor a = Scalar a | Tensor [Tensor a]
--     deriving(Eq)

-- instance (Show a) => Show (Tensor a) where
--     show (Scalar s) = show s
--     show (Tensor ((Tensor t):ts)) = "[\n" ++ unlines (go (Tensor t:ts)) ++ "]"
--         where go = fmap ("  " ++).concatMap(lines . show)
--     show (Tensor s) = show s ++ "\n"

-- instance Functor Tensor where
--     fmap f (Scalar s) = Scalar $ f s
--     fmap f (Tensor t) = Tensor $ map (fmap f) t

-- instance Applicative Tensor where
--     pure = Scalar
--     (Scalar s) <*> t = fmap s t
--     t <*> (Scalar s) = fmap ($ s) t
--     (Tensor []) <*> _ = Tensor []
--     (Tensor (Scalar s : ss)) <*> t = Tensor $ fmap s t : tensorList
--         where Tensor tensorList = Tensor ss <*> t
--     (Tensor (Tensor t : ts)) <*> t2 = Tensor $ (Tensor t <*> t2) : tensorList
--         where Tensor tensorList = Tensor ts <*> t2

-- instance Monad Tensor where
--     (Scalar s) >>= f = f s
--     (Tensor t) >>= f = Tensor $ map (>>= f) t

-- instance (Num a) => Num (Tensor a) where
--     (Scalar s) + t = fmap (+ s) t
--     t + (Scalar s) = fmap (+ s) t
--     (Tensor xs) + (Tensor ys) = Tensor $ zipWith (+) xs ys
--     (Scalar s) * t = fmap (* s) t
--     t * (Scalar s) = fmap (* s) t
--     _ * Tensor[] = Tensor []
--     (Tensor x) * (Tensor y) = getSum $ foldMap Sum (zipWith go x y)
--         where
--             go (Tensor a) b = Tensor $ map (* b) a
--             go a b = a * b
--     abs = fmap abs
--     signum = fmap signum
--     negate = fmap negate
--     fromInteger i = Scalar $ fromInteger i

-- instance Foldable Tensor where
--     foldr f acc (Tensor t) = foldr (flip $ foldr f) acc t
--     foldr f acc (Scalar s) = f s acc

-- instance Traversable Tensor where
--     traverse f (Scalar s) = Scalar <$> f s
--     traverse f (Tensor t) = Tensor <$> traverse (traverse f) t

-- -- TODO: revise function for higher-dimension tensors
-- fromList :: [a] -> Int -> Tensor a
-- fromList xs dim = go (foldr (con . Scalar) (Tensor[]) xs) dim
--     where
--         go (Tensor t) n
--             | n == 1 = Tensor [go (Tensor t) dim]
--             | length t <= dim = Tensor t
--             | otherwise = go (Tensor $ take (length t `div` n) t) dim
--                 `con` go (Tensor $ drop (length t `div` n) t) (n - 1)
--         go s n = error "error when adding tensor rank"

-- product :: (Num a) => Tensor a -> Tensor a -> Tensor a
-- product (Scalar s) t = Scalar s * t
-- product t (Scalar s) = t * Scalar s
-- product (Tensor []) _ = Tensor []
-- product _ (Tensor []) = Tensor []
-- product (Tensor (x:xs)) y = Tensor $ current : rest
--     where
--         current = Tensor.product x y
--         Tensor rest = Tensor.product (Tensor xs) y

-- norm :: (Floating a) => Tensor a -> Tensor a
-- norm t= (/ (sqrt.getSum $ foldMap (Sum . (**2) . abs) t)) <$> t

-- append :: Tensor a -> Tensor a -> Tensor a
-- append (Tensor x) (Tensor y) = Tensor $ x ++ y
-- append _ _ = error "Concatenating scalars is prohibited"

-- con :: Tensor a -> Tensor a -> Tensor a
-- con x (Tensor y) = Tensor $ x : y
-- con _ _ = error "Appending to Scalar is prohibited"

-- dot :: Num a => Tensor a -> Tensor a -> a
-- dot (Tensor x) (Tensor y) = getSum $ foldMap Sum (zipWith dot x y)
-- dot (Scalar x) (Scalar y) = x * y
-- dot _ _ = error "Length mismatch"

-- flatten2d :: Tensor a -> Tensor a
-- flatten2d (Tensor ((Tensor ((Tensor (t:ts)):ss)):xs)) = Tensor . go $ Tensor (Tensor (t:ts):ss):xs
--     where
--         go :: [Tensor a] -> [Tensor a]
--         go (Tensor (Tensor (Scalar s:ss):us):ys) = foldr (\a acc ->
--             let (Tensor result) = Tensor.transpose a
--                 in result ++ acc
--             ) [] (Tensor (Tensor (Scalar s:ss):us):ys)

--         go y = go flattened
--             where flattened = foldr (\a acc ->
--                     let (Tensor next) = a
--                         in Tensor (go next) : acc
--                     ) [] y
-- flatten2d t = t

-- head :: Tensor a -> Tensor a
-- head (Tensor ts) = Prelude.head ts
-- head t = t

-- last :: Tensor a -> Tensor a
-- last (Tensor ts) = Prelude.last ts
-- last t = t

-- tail :: Tensor a -> Tensor a
-- tail (Tensor ts) = Tensor $ Prelude.tail ts
-- tail t = t

-- init :: Tensor a -> Tensor a
-- init (Tensor ts) = Tensor $ Prelude.init ts
-- init t = t

-- transpose :: Tensor a -> Tensor a
-- transpose (Tensor (Tensor t:ts)) = Tensor $ go $ Tensor t:ts
--     where
--         go :: [Tensor a] -> [Tensor a]
--         go (Tensor [] : _) = []
--         go xs = Tensor (map Tensor.head xs) : go (map Tensor.tail xs)
-- transpose (Tensor (Scalar s:ss)) = Tensor $ map (Tensor.(:[])) (Scalar s:ss)
-- transpose t = t

-- isSquare :: Tensor a -> Bool
-- isSquare (Scalar _) = True
-- isSquare (Tensor []) = True
-- isSquare (Tensor t) = getAll $ go (Tensor t) (length t)
--     where
--         go (Tensor x) l
--             | length x == l = foldMap (`go` l) x
--             | otherwise = All False
--         go (Scalar _) _ = All True

-- rank :: Integral b => Tensor a -> b
-- rank (Scalar s) = 0
-- rank (Tensor []) = 1
-- rank (Tensor (t:ts)) = rank t + 1


text:
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict #-}
module FokkerPlanck.GreensFunction where

import           Control.DeepSeq
import           Control.Monad                  as M
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           FokkerPlanck.Analytic
import           FokkerPlanck.FourierSeriesGPU
import           FokkerPlanck.Histogram
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.FourierSeries2D
import           Pinwheel.Transform
import           Sparse.Vector
import           STC.Utils
import           System.FilePath
import           Utils.Array
import           Utils.List
import           Utils.Parallel
import           Utils.SimpsonRule
import           Utils.Time
import Text.Printf
import           Utils.Distribution

sampleCartesian ::
     FilePath
  -> FilePath
  -> [PTX]
  -> Int
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int -> Double
  -> IO (Histogram (Complex Double))
sampleCartesian filePath folderPath ptxs numPoints period delta oris gamma thetaSigma tau threshold r2Sigma maxPhiFreqs maxRhoFreqs maxThetaFreqs maxRFreqs stdR2 = do
  let center = div numPoints 2
      deltaTheta = 2 * pi / fromIntegral oris
      origin = R2S1RP 0 0 0 gamma
      idx =
        L.filter
          (\(R2S1RP c r _ g) -> (sqrt (c ^ 2 + r ^ 2) > 0) && g > 0)
          [ R2S1RP
            ((fromIntegral $ c - center) * delta)
            ((fromIntegral $ r - center) * delta)
            0
            gamma
          | c <- [0 .. numPoints - 1]
          , r <- [0 .. numPoints - 1]
          ]
      logGamma = 0 -- log gamma
      -- simpsonWeights = weightsSimpsonRule oris
      xs =
        L.concat -- .
        -- L.zipWith
        --   (\w -> L.map (\(a, b, c, d, v) -> (a, b, c, d, v * w)))
        --   simpsonWeights
         $
        parMap
          rdeepseq
          (\o ->
             let ori = fromIntegral o * deltaTheta
              in L.filter (\(_, _, _, _, v) -> v > threshold) .
                 L.map
                   (\point@(R2S1RP x y _ g) ->
                      let v =
                            computePji thetaSigma tau origin (R2S1RP x y ori g)
                          phi = atan2 y x
                          rho = sqrt $ (x ^ 2 + y ^ 2)
                       in ( phi
                          , log rho
                          , ori
                          , logGamma
                          , v * (1 - gaussian2DPolar rho stdR2)
                          )) $
                 idx)
          [0 .. oris - 1]
  -- arr <-
  --   computeUnboxedP . fromFunction (Z :. oris :. numPoints :. numPoints) $ \(Z :. o :. i :. j) ->
  --     let x = (fromIntegral $ i - center) * delta
  --         y = (fromIntegral $ j - center) * delta
  --         phi = atan2 y x
  --         rho = sqrt $ (x ^ 2 + y ^ 2)
  --         theta = fromIntegral o * deltaTheta
  --         v = computePji thetaSigma tau origin (R2S1RP x y theta gamma)
  --     in if rho <= 1.5
  --          then (1, 1, 1, 1, 0)
  --          else (phi, log rho, theta, logGamma, v)
  -- plotImageRepa (folderPath </> "GreensFunction.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   -- reduceContrast 20 .
  --   extend (Z :. (1 :: Int) :. All :. All) .
  --   sumS . rotate3D . R.map (\(_, _, _, _, v) -> v) $
  --   arr
      -- simpsonWeights =
      --   computeWeightArrFromListOfShape [numPoints, numPoints, oris]
      -- xs =
      --   L.filter (\(_, _, _, _, v) -> v > threshold) .
      --   R.toList -- .
      --   -- R.zipWith (\w (a, b, c, d, v) -> (a, b, c, d, w * v)) simpsonWeights
      --   $
      --   arr
  let
  printCurrentTime $
    printf
      "Sparsity: %f%%\n"
      (100 * (fromIntegral . L.length $ xs) /
       (fromIntegral $ oris * numPoints ^ 2) :: Double)
  let phiFreqs = L.map fromIntegral [-maxPhiFreqs .. maxPhiFreqs]
      rhoFreqs = L.map fromIntegral [-maxRhoFreqs .. maxRhoFreqs]
      thetaFreqs = L.map fromIntegral [-maxThetaFreqs .. maxThetaFreqs]
      rFreqs = L.map fromIntegral [-maxRFreqs .. maxRFreqs]
      hist =
        L.foldl1' (addHistogram) .
        parZipWith
          rdeepseq
          (\ptx ys ->
             computeFourierCoefficientsGPU'
               r2Sigma
               period
               phiFreqs
               rhoFreqs
               thetaFreqs
               rFreqs
               ptx
               ys)
          ptxs .
        divideListN (L.length ptxs) $
        xs
  encodeFile filePath hist
  return hist

sampleCartesianCorner ::
     FilePath
  -> FilePath
  -> [PTX]
  -> Int
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> IO (Histogram (Complex Double))
sampleCartesianCorner filePath folderPath ptxs numPoints period delta oris gamma thetaSigma tau threshold r2Sigma maxPhiFreqs maxRhoFreqs maxThetaFreqs maxRFreqs stdR2 weight = do
  let center = div numPoints 2
      deltaTheta = 2 * pi / fromIntegral oris
  --     origin = R2S1RP 0 0 0 gamma
      logGamma = 0
      idxFunc c r o =
        ( fromIntegral o * deltaTheta
        , fromIntegral (c - center) * delta
        , fromIntegral (r - center) * delta)
  --     arr1 =
  --       fromFunction (Z :. oris :. numPoints :. numPoints) $ \(Z :. o :. c :. r) ->
  --         if c == center && r == center
  --           then 0
  --           else let (ori, x, y) = idxFunc c r o
  --                 in computePji thetaSigma tau origin (R2S1RP x y ori gamma)
  --     orientations = [fromIntegral i * deltaTheta | i <- [0 .. oris - 1]]
  --     arr2 =
  --       fromFunction (Z :. oris :. numPoints :. numPoints) $ \(Z :. o :. c :. r) ->
  --         if c == center && r == center
  --           then 0
  --           else let (ori, x, y) = idxFunc c r o
  --                 in computePjiCorner'
  --                      thetaSigma
  --                      tau
  --                      threshold
  --                      orientations
  --                      (R2S1RP x y ori gamma)
  -- arr <- computeUnboxedP $ arr1 +^ (R.map (* weight) arr2)
  arr <- sampleR2S1Corner numPoints numPoints delta deltaTheta gamma thetaSigma tau threshold [0 .. oris - 1] weight
  let idx =
        [ idxFunc c r o
        | o <- [0 .. oris - 1]
        , c <- [0 .. numPoints - 1]
        , r <- [0 .. numPoints - 1]
        ]
      xs =
        L.map
          (\((ori, x, y), v) ->
             let phi = atan2 y x
                 rho = sqrt (x ^ 2 + y ^ 2)
              in ( phi
                 , log rho
                 , ori
                 , logGamma
                 , v * (1 - gaussian2DPolar rho stdR2))) .
        L.filter (\((_, x, y), v) -> v > threshold && (x /= 0 || y /= 0)) .
        L.zip idx . R.toList $
        arr
  printCurrentTime $
    printf
      "Sparsity: %f%%\n"
      (100 * (fromIntegral . L.length $ xs) /
       (fromIntegral $ oris * numPoints ^ 2) :: Double)
  let phiFreqs = L.map fromIntegral [-maxPhiFreqs .. maxPhiFreqs]
      rhoFreqs = L.map fromIntegral [-maxRhoFreqs .. maxRhoFreqs]
      thetaFreqs = L.map fromIntegral [-maxThetaFreqs .. maxThetaFreqs]
      rFreqs = L.map fromIntegral [-maxRFreqs .. maxRFreqs]
      hist =
        L.foldl1' (addHistogram) .
        parZipWith
          rdeepseq
          (\ptx ys ->
             computeFourierCoefficientsGPU'
               r2Sigma
               period
               phiFreqs
               rhoFreqs
               thetaFreqs
               rFreqs
               ptx
               ys)
          ptxs .
        divideListN (L.length ptxs) $
        xs
  encodeFile filePath hist
  return hist

-- sampleLogpolar ::
--      FilePath
--   -> [PTX]
--   -> Int
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Int
--   -> Int
--   -> Int
--   -> Int
--   -> IO (Histogram (Complex Double))
-- sampleLogpolar filePath ptxs phis rhos oris maxScale gamma sigma tau threshold r2Sigma maxPhiFreqs maxRhoFreqs maxThetaFreqs maxRFreqs = do
--   let deltaPhi = 2 * pi / fromIntegral phis
--       deltaTheta = 2 * pi / fromIntegral oris
--       logMaxScale = log maxScale
--       deltaLogRho = logMaxScale / fromIntegral rhos
--       logGamma = 0 -- log gamma
--       origin = R2S1RP 0 0 0 gamma
--       -- xs =
--       --   parMap
--       --     rdeepseq
--       --     (\(phi, logRho, theta) ->
--       --        let x = (exp logRho) * cos phi
--       --            y = (exp logRho) * sin phi
--       --            v = computePji sigma tau origin (R2S1RP x y theta gamma)
--       --        in (phi, logRho, theta, 0, v))
--       --     [ ( deltaTheta * fromIntegral iPhi
--       --       , deltaLogRho * fromIntegral iLogRho - logMaxScale
--       --       , deltaTheta * fromIntegral iTheta)
--       --     | iPhi <- [0 .. oris - 1]
--       --     , iLogRho <- [0 .. rhos - 1]
--       --     , iTheta <- [0 .. oris - 1]
--       --     ]
--   arr <-
--     computeUnboxedP . R.fromFunction (Z :. oris :. rhos :. phis) $ \(Z :. iTheta :. iLogRho :. iPhi) ->
--       let theta = deltaTheta * fromIntegral iTheta
--           phi = deltaPhi * fromIntegral iPhi
--           logRho = deltaLogRho * fromIntegral iLogRho -- logMaxScale
--           x = (exp logRho) * cos phi
--           y = (exp logRho) * sin phi
--           v = computePji sigma tau origin (R2S1RP x y theta gamma)
--       in (phi, logRho, theta, logGamma, v)
--   let simpsonWeights = computeWeightArrFromListOfShape [phis, rhos, oris]
--       xs =
--         L.filter (\(_, _, _, _, v) -> v > threshold) .
--         R.toList .
--         R.zipWith (\w (a, b, c, d, v) -> (a, b, c, d, w * v)) simpsonWeights $
--         arr
--   printCurrentTime $
--     printf
--       "Sparsity: %f%%\n"
--       (100 * (fromIntegral . L.length $ xs) / (fromIntegral $ oris * phis * rhos) :: Double)
--   let phiFreqs = L.map fromIntegral [-maxPhiFreqs .. maxPhiFreqs]
--       rhoFreqs = L.map fromIntegral [-maxRhoFreqs .. maxRhoFreqs]
--       thetaFreqs = L.map fromIntegral [-maxThetaFreqs .. maxThetaFreqs]
--       rFreqs = L.map fromIntegral [-maxRFreqs .. maxRFreqs]
--       hist =
--         L.foldl1' (addHistogram) .
--         parZipWith
--           rdeepseq
--           (\ptx ys ->
--              computeFourierCoefficientsGPU'
--                r2Sigma
--                phiFreqs
--                rhoFreqs
--                thetaFreqs
--                rFreqs
--                ptx
--                ys)
--           ptxs .
--         divideListN (L.length ptxs) $
--         xs
--   encodeFile filePath hist
--   return hist

{-# INLINE sampleR2S1 #-}
sampleR2S1 ::
     Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Int]
  -> IO (R.Array U DIM3 Double)
sampleR2S1 rows cols delta deltaTheta gamma sigma tau oris = do
  let centerRow = div rows 2
      centerCol = div cols 2
      origin = R2S1RP 0 0 0 gamma
  computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length oris)) oris)
      (\(Z :. o) -> (Z :. o :. cols :. rows)) $ \f (Z :. o :. c :. r) ->
    if (c == centerCol) && (r == centerRow)
      then 0
      else computePji
             sigma
             tau
             origin
             (R2S1RP
                (delta * fromIntegral (c - centerCol))
                (delta * fromIntegral (r - centerRow))
                (fromIntegral (f (Z :. o)) * deltaTheta)
                gamma)

{-# INLINE sampleR2S1' #-}
sampleR2S1' ::
     Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array U DIM2 (Int, Double)
sampleR2S1' rows cols delta gamma sigma tau ori =
  let centerRow = div rows 2
      centerCol = div cols 2
      origin = R2S1RP 0 0 0 gamma
  in computeUnboxedS . R.fromFunction (Z :. cols :. rows) $ \(Z :. c :. r) ->
       let x = c - centerCol
           y = r - centerRow
       in if x == 0 && y == 0
            then (0, 0)
            else ( getVectorIndex2D rows c r
                 , computePji
                     sigma
                     tau
                     origin
                     (R2S1RP
                        (delta * fromIntegral (c - centerCol))
                        (delta * fromIntegral (r - centerRow))
                        ori
                        gamma))

{-# INLINE sampleR2S1Corner #-}
sampleR2S1Corner ::
     Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Int]
  -> Double
  -> IO (R.Array U DIM3 Double)
sampleR2S1Corner rows cols delta deltaTheta gamma sigma tau threshold oris weight = do
  let centerRow = div rows 2
      centerCol = div cols 2
      origin = R2S1RP 0 0 0 gamma
      orientations = L.map (\o -> deltaTheta * fromIntegral o) oris
  arr1 <-
    computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length oris)) oris)
      (\(Z :. o) -> (Z :. o :. cols :. rows)) $ \f (Z :. o :. c :. r) ->
      if (c == centerCol) && (r == centerRow)
        then 0
        else computePji
               sigma
               tau
               origin
               (R2S1RP
                  (delta * fromIntegral (c - centerCol))
                  (delta * fromIntegral (r - centerRow))
                  (fromIntegral (f (Z :. o)) * deltaTheta)
                  gamma)
  arr2 <-
    computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length oris)) oris)
      (\(Z :. o) -> (Z :. o :. cols :. rows)) $ \f (Z :. o :. c :. r) ->
      if (c == centerCol) && (r == centerRow)
        then 0
        else computePjiCorner'
               delta
               sigma
               tau
               threshold
               orientations
               (R2S1RP
                  (delta * fromIntegral (c - centerCol))
                  (delta * fromIntegral (r - centerRow))
                  (fromIntegral (f (Z :. o)) * deltaTheta)
                  gamma)
  computeUnboxedP $ (arr1 +^ (R.map (* weight) arr2))

cornerDistribution ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Double]
  -> Double
  -> (Double, Double)
  -> [Double]
cornerDistribution delta gamma sigma tau threshold oris weight (x, y) =
  let origin = R2S1RP 0 0 0 gamma
   in L.map
        (\ori ->
           computePji sigma tau origin (R2S1RP x y ori gamma) +
           weight *
           computePjiCorner'
             delta
             sigma
             tau
             threshold
             oris
             (R2S1RP x y ori gamma))
        oris

{-# INLINE cornerDistribution' #-}
cornerDistribution' ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Double]
  -> Double
  -> (Double, Double)
  -> Double
  -> IO (Double, Double, Double)
cornerDistribution' delta gamma sigma tau threshold oris weight (x, y) ori = do
  let origin = R2S1RP 0 0 0 gamma
      a = computePji sigma tau origin (R2S1RP x y ori gamma)
      b = computePjiCorner' delta sigma tau threshold oris (R2S1RP x y ori gamma)
  printf "sigma = %f\ntau = %f\nthreshold = %f\noris = %s\n(R2S1RP %f %f %f %f)\n" sigma tau threshold (show oris) x y ori gamma
  return (a, b, a + weight * b)

computeFourierCoefficients ::
     FilePath
  -> [Int]
  -> [PTX]
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> IO (Histogram (Complex Double))
computeFourierCoefficients histFilePath deviceIDs ptxs numPoints oris delta deltaFreq gamma sigma tau numR2Freqs period std numBatchR2Freqs numBatchOri maxPhiFreq maxRhoFreq maxThetaFreq maxRFreq = do
  when
    (even oris)
    (error "computeFourierCoefficients: the number of orientations is not odd.")
  let deltaTheta = 2 * pi / fromIntegral oris
      simpsonWeights =
        computeUnboxedS $
        (computeWeightArrFromListOfShape [numPoints, numPoints] :: R.Array D DIM2 (Complex Double))
      oriIdxs = divideListN numBatchOri [0 .. oris - 1]
  xs <-
    M.mapM
      (sampleR2S1 numPoints numPoints delta deltaTheta gamma sigma tau)
      oriIdxs
  let weightedVecs =
        parMap
          rdeepseq
          (\arr ->
             let (Z :. n :. m :. c) = extent arr
             in CuMat (n * m) c .
                CuVecHost .
                VU.convert .
                toUnboxed . computeS . R.traverse2 arr simpsonWeights const $ \f fW idx@(Z :. i :. j :. _) ->
                  (f idx :+ 0) * fW (Z :. i :. j))
          xs
  coefficients <-
    computeFourierCoefficientsR2Stream
      deviceIDs
      ptxs
      numR2Freqs
      numPoints
      period
      delta
      deltaFreq
      numBatchR2Freqs
      weightedVecs
  print "coefficients done"
  let centerR2Freq = div numR2Freqs 2
      centerPoints = div numPoints 2
      logGamma = log gamma
      gaussianWeightedCoef =
        computeUnboxedS $ applyGaussian deltaFreq std coefficients
      -- gaussianWeightedCoef = coefficients
      simpsonWeightsOri =
        computeUnboxedS $
        (computeWeightArrFromListOfShape [oris] :: R.Array D DIM1 (Complex Double))
      weightedCoef =
        R.traverse2 gaussianWeightedCoef simpsonWeightsOri const $ \fG fS idx@(Z :. o :. _ :. _) ->
          fG idx * fS (Z :. o)
      folderPath = "output/test/STCPinwheel"
  -- print . extent $ gaussianWeightedCoef
  -- inverseR2Harmonics <-
  --   createInverseHarmonicMatriesGPU ptxs 1 numPoints numR2Freqs period delta
  -- sourceR2 <-
  --   computeFourierSeriesR2
  --     deviceIDs
  --     numR2Freqs
  --     numPoints
  --     period
  --     inverseR2Harmonics
  --     [ CuMat oris (numR2Freqs ^ 2) . CuVecHost . VU.convert . toUnboxed $
  --       gaussianWeightedCoef
  --     ]
  plotImageRepa (folderPath </> "Green.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt . toUnboxed . sumS . R.map (\x -> (x) ** 2) . L.head $
    xs
  plotImageRepa (folderPath </> "Coefficients.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freqs :. numR2Freqs) .
    VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
    gaussianWeightedCoef
  -- plotImageRepa (folderPath </> "GreenRecon.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   sourceR2
  ys <-
    (computeUnboxedP . R.traverse weightedCoef id $ \f idx@(Z :. o :. i :. j) ->
       if i == centerR2Freq && j == centerR2Freq
         then (0, 0, 0, 0, 0)
         else let x = deltaFreq * (fromIntegral $ i - centerR2Freq)
                  y = deltaFreq * (fromIntegral $ j - centerR2Freq)
                  phi = atan2 y x
                  logRho = log . sqrt $ (x ^ 2 + y ^ 2)
              in (phi, logRho, fromIntegral o * deltaTheta, logGamma, f idx))
  -- ys <-
  --   (computeUnboxedP .
  --    R.traverse
  --      (R.traverse2 (L.head xs) simpsonWeightsOri const $ \fG fS idx@(Z :. _ :. _ :. o) ->
  --         (fG idx :+ 0) * fS (Z :. o))
  --      id $ \f idx@(Z :. i :. j :. o) ->
  --      if i == centerPoints && j == centerPoints
  --        then (0, 0, 0, 0, 0)
  --        else let x = delta * (fromIntegral $ i - centerPoints)
  --                 y = delta * (fromIntegral $ j - centerPoints)
  --                 phi = atan2 x y
  --                 logRho = log . sqrt $ (x ^ 2 + y ^ 2)
  --             in ( phi
  --                , logRho
  --                , fromIntegral o * deltaTheta
  --                , logGamma
  --                , f idx))
  let hist =
        L.foldl1' (addHistogram) .
        parZipWith
          rdeepseq
          (\ptx ->
             computePinwheelCoefficients
               maxPhiFreq
               maxRhoFreq
               maxThetaFreq
               maxRFreq
               ptx)
          ptxs .
        divideListN (L.length ptxs) . R.toList $
        ys
  encodeFile histFilePath hist
  return hist

computePinwheelTransformCoefficients ::
     Int
  -> FilePath
  -> [Int]
  -> [PTX]
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> IO (Histogram (Complex Double))
computePinwheelTransformCoefficients numThread histFilePath deviceIDs ptxs numPoints oris delta threshold gamma sigma tau numR2Freqs periodR2 std numBatchR2 numBatchPinwheelFreqs maxPhiFreq maxRhoFreq maxThetaFreq maxRFreq batchSize s = do
  printCurrentTime "compute pinwheelArr"
  pinwheelArr <-
    pinwheelFourierSeries
      deviceIDs
      ptxs
      numPoints -- numR2Freqs
      numPoints
      delta
      periodR2
      maxPhiFreq
      maxRhoFreq
      maxThetaFreq
      maxRFreq
      s
      numBatchR2
      numBatchPinwheelFreqs :: IO (IA.Array (Int, Int) (VU.Vector (Complex Double)))
  printCurrentTime "pinwheelArr done"
  let folderPath = "output/test/STCPinwheel"
  let deltaTheta = 2 * pi / (fromIntegral oris)
      greensFunc =
        parMap
          rdeepseq
          (\o ->
             let theta = fromIntegral o * deltaTheta
                 vec =
                   toUnboxed $
                   sampleR2S1' numPoints numPoints delta gamma sigma tau theta
             in PinwheelTransformData (log gamma) theta .
                uncurry SparseVector .
                VU.unzip .
                VU.map (\(i, v) -> (i, v :+ 0)) .
                VU.filter (\(_, v) -> v > threshold) $
                vec)
          [0 .. oris - 1]
      len =
        L.foldl'
          (\s (PinwheelTransformData _ _ (SparseVector vec _)) ->
             s + VU.length vec)
          0
          greensFunc
  printCurrentTime $
    printf
      "Sparsity: %f%%\n"
      (100 * (fromIntegral len) / (fromIntegral $ oris * numPoints ^ 2) :: Double)
  let hist =
        pinwheelTransform
          numThread
          maxPhiFreq
          maxRhoFreq
          maxThetaFreq
          maxRFreq
          s
          pinwheelArr
          greensFunc
  encodeFile histFilePath hist
  return hist

ext: hs

max_stars_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_issues_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_forks_repo_name: XinhuaZhang/Stochastic-Completion-Field


text:
{- |
   Description :  Parse arguments needed for the bmm simulator
   Copyright   :  (c) Dominik Schrempf 2017
   License     :  GPLv3

   Maintainer  :  dominik.schrempf@gmail.com
   Stability   :  unstable
   Portability :  non-portable (not tested)

Provides a data type with all options and a parser.

* Changelog

-}

module ArgParse
  ( BMMArgs(..)
  , parseBMMArgs
  ) where

import qualified Data.Attoparsec.Text         as A
import           Data.Semigroup               ((<>))
import           Data.Text                    (pack, unpack)
import qualified Defaults                     as Def
import           DNAModel                     (DNAModelSpec(..))
import           RateMatrix                   (StationaryDist)
import           Numeric.LinearAlgebra        (norm_1, size, vector)
import           Options.Applicative
import qualified Text.PrettyPrint.ANSI.Leijen as Doc
import           Tools                        (nearlyEq)

-- Convenience function to read in more complicated command line options with
-- attoparsec and optparse
-- (https://github.com/pcapriotti/optparse-applicative#option-readers).
attoReadM :: A.Parser a -> ReadM a
attoReadM p = eitherReader (A.parseOnly p . pack)

-- | Combined option and argument data structure.
data BMMArgs = BMMArgs
  { -- | Name of output counts file.
    outFileName    :: String
    -- | Mutation model specification.
  , dnaModelSpec   :: DNAModelSpec
    -- | Shape parameter of gamma rate heterogeneity.
  , gammaShape     :: Maybe Double
    -- | Number of gamma rate heterogeneity categories.
  , gammaNCat      :: Maybe Int
    -- | Virtual population size of boundary mutation model.
  , popSize        :: Int
    -- | Heterozygosity value (2Nu).
  , heterozygosity :: Double
    -- | Tree height in boundary mutation model units.
  , treeHeight     :: Double
    -- | Tree type or scenario (ILS, Yule).
  , treeType       :: String
    -- | Yule speciation rate.
  , treeYuleRate   :: Maybe Double
    -- | Number of sites to simulate.
  , nSites         :: Int
    -- | Seed of the random number generator.
  , seed           :: String }

-- Composition of all options.
bmSimOptions :: Parser BMMArgs
bmSimOptions = BMMArgs
  <$> outFileNameOpt
  <*> dnaModelOpt
  <*> gammaShapeOpt
  <*> gammaNCatOpt
  <*> popSizeOpt
  <*> heterozygosityOpt
  <*> treeHeightOpt
  <*> treeTypeOpt
  <*> treeYuleRateOpt
  <*> nSitesOpt
  <*> seedOpt

-- | The impure IO action that reads the arguments and prints out help if
-- needed. Maybe put this into Main.hs?
parseBMMArgs :: IO BMMArgs
parseBMMArgs = execParser $
  info (helper <*> bmSimOptions)
  (fullDesc
    <> progDesc "Simulate count files using the boundary mutation model."
    <> header "Boundary mutation model simulator"
    <> footerDoc models )
  where
    models = Just $ foldl1 (Doc.<$>) (map Doc.text strs)
    strs   = [ "Available mutation models:"

             , "  - HKY model with transition to transversion ratio kappa and a state frequency vector."
             , "    Specified with \"-m HKY[DOUBLE][DOUBLE,DOUBLE,DOUBLE,DOUBLE]\"."
             , "  - GTR model with five rate parameters and state frequency vector."
             , "    Specified with \"-m HKY[DOUBLE,DOUBLE,DOUBLE,DOUBLE,DOUBLE][DOUBLE,DOUBLE,DOUBLE,DOUBLE]\"."
             , ""
             , "Note: The state frequency vector has to sum up to 1.0 and only has three free parameters." ]

-- General things and options.
outFileNameOpt :: Parser String
outFileNameOpt = strOption
  ( long "output"
    <> short 'o'
    <> metavar "NAME"
    <> value Def.outFileName
    <> showDefault
    <> help "Write output files to NAME.[cf.gz|log|tree]" )

-- Read a stationary frequency of the form `pi_A,pi_C,pi_G,...`.
parseStateFreq :: A.Parser StationaryDist
parseStateFreq = do
  _ <- A.char '['
  f <- vector <$> A.sepBy A.double (A.char ',')
  _ <- A.char ']'
  if size f /= nAlleles
    then error "Length of stationary frequency vector is faulty, only DNA models are supported."
  else if nearlyEq 1e-6 (norm_1 f) 1.0 then return f
    else error $ "Stationary frequencies sum to " ++ show (norm_1 f) ++ " but should sum to 1.0."
    -- Nucleotide count hard coded. See `BndState.nAlleles`. However, I cannot
    -- take different numbers depending on the mutation model because optparse has
    -- an applicative interface and not a monadic one. Also, allowing vectors of
    -- arbitrary size is undesirable because of meaningless error messages.
    where nAlleles = 4

dnaModelOpt :: Parser DNAModelSpec
dnaModelOpt = option (attoReadM parseDNAModelSpec)
  ( long "mutation-model"
    <> short 'm'
    <> metavar "MODEL"
    <> value Def.dnaModelSpec
    <> showDefault
    <> help "Set the mutation model; available models are shown below" )

parseParams :: A.Parser [Double]
parseParams = do
  _ <- A.char '['
  params <- A.sepBy1 A.double (A.char ',')
  _ <- A.char ']'
  return params

parseDNAModelSpec :: A.Parser DNAModelSpec
parseDNAModelSpec = do
  m  <- A.takeWhile (/= '[')
  case unpack m of
       "JC" -> return JC
       "HKY" -> do
         ps <- parseParams
         f  <- parseStateFreq
         if length ps /= 1
           then error "HKY model only has one parameter, kappa."
           else return $ HKY (head ps) f
       "GTR" -> do
         ps <- parseParams
         f  <- parseStateFreq
         if length ps /= 5
           then error "GTR model has five parameters."
           else return $ GTR (head ps) (ps !! 1) (ps !! 2) (ps !! 3) (ps !! 4) f
       _ -> error "Model string could not be parsed."

gammaShapeOpt :: Parser (Maybe Double)
gammaShapeOpt = optional $ option auto
  ( long "gamma-shape"
    <> metavar "DOUBLE"
    <> help "Activate gamma rate heterogeneity and set gamma shape parameter (default: off)" )

gammaNCatOpt :: Parser (Maybe Int)
gammaNCatOpt = optional $ option auto
  ( long "gamma-ncat"
    <> metavar "INT"
    <> help "Set the number of gamma rate categories (no default value)" )

popSizeOpt :: Parser Int
popSizeOpt = option auto
  ( short 'N'
    <> metavar "DOUBLE"
    <> value Def.popSize
    <> showDefault
    <> help "Set the virtual population size" )

heterozygosityOpt :: Parser Double
heterozygosityOpt = option auto
  ( long "heterozygosity"
    <> short 't'
    <> metavar "DOUBLE"
    <> value Def.heterozygosity
    <> showDefault
    <> help "Set heterozygosity" )

treeHeightOpt :: Parser Double
treeHeightOpt = option auto
  ( long "tree-height"
    <> short 'H'
    <> metavar "DOUBLE"
    <> value Def.treeHeight
    <> showDefault
    <> help "Set tree height [average number of substitutions]" )

treeTypeOpt :: Parser String
treeTypeOpt = strOption
  ( long "tree-type"
    <> metavar "TYPE"
    <> value Def.treeType
    <> showDefault
    <> help "Set tree type; ILS or Yule")

treeYuleRateOpt :: Parser (Maybe Double)
treeYuleRateOpt = optional $ option auto
  ( long "tree-yule-rate"
  <> metavar "DOUBLE"
  <> help "Set the speciation rate of Yule tree (no default value)")

nSitesOpt :: Parser Int
nSitesOpt = option auto
  ( long "nsites"
    <> short 'n'
    <> metavar "INT"
    <> value Def.nSites
    <> showDefault
    <> help "Set number of sites to simulate" )

seedOpt :: Parser String
seedOpt = strOption
  ( long "seed"
    <> short 's'
    <> metavar "INT"
    <> value "random"
    <> showDefault
    <> help "Set seed for the random number generator" )

max_stars_repo_name: dschrempf/bmm-simulate

max_issues_repo_name: pomo-dev/bmm-simulate

max_forks_repo_name: pomo-dev/bmm-simulate


text:
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

import CPVO.Numeric
import CPVO.IO
import CPVO.IO.Reader.Ecalj.Common
import CPVO.IO.Reader.Ecalj.DOS

import System.Environment (getArgs)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.List
-------------------------

import Numeric.LinearAlgebra
-------------------------

rydberg :: Double
rydberg=1/13.605

testArgs :: [String]
testArgs = concat [ [ "GWtable01.tex"
                   , "Total and partial DOS at *E*~F~ in states/eV/unit-cell, QSGW method."
                   , "|*N*$_\\uparrow$(*E*~F~)|*N*$_\\downarrow$(*E*~F~)|*N*(*E*~F~)"
                     ]
                   , words "@{}lSSS@{} -9:6 25 T o flipSpin nico2o4 extendedNiCo2O4.normal/nico2o4.6G10 Ni:Ni#3d:6:7:9:8:10 Co:Co#3d:6:7:8:9:10 O:O#2p:3:4:5"
                   ]

main :: IO ()
main = do
    allArgs <- getArgs
    --Right (invStat, ymax, xmin, xmax, ctrlAtoms, uniqAtoms, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer, colAlign, texFile) <- readHeaderData allArgs
    Right (invStat,_,_,_,_,_, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer,_, texFile) <- readHeaderData allArgs
    totalDOS <- readTotalDOSText tailer foldernya
    let resTotY0 = map (\i -> getY0 $ totalDOS ¿ [0,i]) $ flipBy invStat [1,2]
    putStrLn "===========NF total DOS========="
    putStrLn $ show resTotY0
    pdosAtAll <- readPDOS invStat tailer foldernya ctrlAtomicAOs
    let rs' = rendertable
              $ (:) jdHeads
              $ (:)  ("Total" : (map (showDouble 3) $ processCols $ map ((* rydberg) . last . toList) $ resTotY0))
              $ map (\a -> (((fst $ head a):(map (showDouble 3) $ processCols $ sumIt [0,0] $ map snd a))))
              $ groupBy (\a b -> fst a == fst b)
              $ (\[us,ds] -> zipWith (\(yu,(_,(j,_))) (yd,_) -> (j, [yu,yd])) us ds )
              $ groupBy (\(_,(s,_)) (_,(s',_)) ->  s == s')
              $ map (\(mp,b) -> ((* rydberg) $ last $ toList $ getY0 mp, b)) pdosAtAll
    let rs = unlines [ rs' , jdTable ]
    T.writeFile texFile $ T.pack rs
{-
    res' <- markdownToTex rs
    let res = T.replace "\\}" "}"
            $ T.replace "\\{" "{"
            $ T.pack $ unlines  [ "\\begin{longtable}[]{" ++ colAlign ++ "}"
                                , unlines $ tail $ lines $ T.unpack res'
                                ]
    T.writeFile texFile res
    -}
    putStrLn $ texFile ++ "=============done:t1dosEf.hs=========="
      where
        processCols [u,d] = [u,d,u+d]
        processCols _ = []
        sumIt x [] = x
        sumIt [u0,d0] ([up,dn]:xs) = sumIt [u0+up,d0+dn] xs
        sumIt _ _ = []

max_stars_repo_name: hasanalrasyid/cpvoh

max_issues_repo_name: hasanalrasyid/cpvoh

max_forks_repo_name: hasanalrasyid/cpvoh


text:
{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random
import           Control.Monad.Trans.Except

import qualified Data.Attoparsec.Text as A
import           Data.List ( foldl' )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif

import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Vector.Storable as V

import           Numeric.LinearAlgebra ( maxIndex )
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative

import           Grenade
import           Grenade.Utils.OneHot

-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger MNIST
-- type.
type FL i o =
  Network
    '[ FullyConnected i o, Logit ]
    '[ 'D1 i, 'D1 o, 'D1 o ]

-- The definition of our convolutional neural network.
-- In the type signature, we have a type level list of shapes which are passed between the layers.
-- One can see that the images we are inputing are two dimensional with 28 * 28 pixels.

-- It's important to keep the type signatures, as there's many layers which can "squeeze" into the gaps
-- between the shapes, so inference can't do it all for us.

-- With the mnist data from Kaggle normalised to doubles between 0 and 1, learning rate of 0.01 and 15 iterations,
-- this network should get down to about a 1.3% error rate.
--
-- /NOTE:/ This model is actually too complex for MNIST, and one should use the type given in the readme instead.
--         This one is just here to demonstrate Inception layers in use.
--
type MNIST =
  Network
    '[ Reshape,
       Concat ('D3 28 28 1) Trivial ('D3 28 28 14) (InceptionMini 28 28 1 5 9),
       Pooling 2 2 2 2, Relu,
       Concat ('D3 14 14 3) (Convolution 15 3 1 1 1 1) ('D3 14 14 15) (InceptionMini 14 14 15 5 10), Crop 1 1 1 1, Pooling 3 3 3 3, Relu,
       Reshape, FL 288 80, FL 80 10 ]
    '[ 'D2 28 28, 'D3 28 28 1,
       'D3 28 28 15, 'D3 14 14 15, 'D3 14 14 15, 'D3 14 14 18,
       'D3 12 12 18, 'D3 4 4 18, 'D3 4 4 18,
       'D1 288, 'D1 80, 'D1 10 ]

randomMnist :: MonadRandom m => m MNIST
randomMnist = randomNetwork

convTest :: Int -> FilePath -> FilePath -> LearningParameters -> ExceptT String IO ()
convTest iterations trainFile validateFile rate = do
  net0         <- lift randomMnist
  trainData    <- readMNIST trainFile
  validateData <- readMNIST validateFile
  lift $ foldM_ (runIteration trainData validateData) net0 [1..iterations]

    where
  trainEach rate' !network (i, o) = train rate' network i o

  runIteration trainRows validateRows net i = do
    let trained' = foldl' (trainEach ( rate { learningRate = learningRate rate * 0.9 ^ i} )) net trainRows
    let res      = fmap (\(rowP,rowL) -> (rowL,) $ runNet trained' rowP) validateRows
    let res'     = fmap (\(S1D label, S1D prediction) -> (maxIndex (SA.extract label), maxIndex (SA.extract prediction))) res
    print trained'
    putStrLn $ "Iteration " ++ show i ++ ": " ++ show (length (filter ((==) <$> fst <*> snd) res')) ++ " of " ++ show (length res')
    return trained'

data MnistOpts = MnistOpts FilePath FilePath Int LearningParameters

mnist' :: Parser MnistOpts
mnist' = MnistOpts <$> argument str (metavar "TRAIN")
                   <*> argument str (metavar "VALIDATE")
                   <*> option auto (long "iterations" <> short 'i' <> value 15)
                   <*> (LearningParameters
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )

main :: IO ()
main = do
    MnistOpts mnist vali iter rate <- execParser (info (mnist' <**> helper) idm)
    putStrLn "Training convolutional neural network..."

    res <- runExceptT $ convTest iter mnist vali rate
    case res of
      Right () -> pure ()
      Left err -> putStrLn err

readMNIST :: FilePath -> ExceptT String IO [(S ('D2 28 28), S ('D1 10))]
readMNIST mnist = ExceptT $ do
  mnistdata <- T.readFile mnist
  return $ traverse (A.parseOnly parseMNIST) (T.lines mnistdata)

parseMNIST :: A.Parser (S ('D2 28 28), S ('D1 10))
parseMNIST = do
  Just lab <- oneHot <$> A.decimal
  pixels   <- many (A.char ',' >> A.double)
  image    <- maybe (fail "Parsed row was of an incorrect size") pure (fromStorable . V.fromList $ pixels)
  return (image, lab)

max_stars_repo_name: claudeha/grenade

max_issues_repo_name: Alien-Inc/grenade

max_forks_repo_name: Alien-Inc/grenade


text:
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
module Q.Options.ImpliedVol.Surface
  (
    Surface(..)
  , totalVarKT
  , fwdTotalVarKT
  , volKT)

where

import qualified Data.Map                                 as M
import qualified Data.Map.Strict                          as M
import           Data.Maybe                               (fromJust)
import           Numeric.LinearAlgebra                    hiding (maxElement,
                                                           minElement)
import qualified Q.Options.Bachelier                      as Bacherlier
import qualified Q.Options.Black76                        as B76
import           Q.Options.ImpliedVol
import           Q.Options.ImpliedVol.InterpolatingSmile
import           Q.Options.ImpliedVol.StrikeInterpolation
import           Q.Options.ImpliedVol.TimeInterpolation
import           Q.Options.ImpliedVol.TimeSlice
import           Q.SortedVector
import           Q.Types

-- | Implied volatility surface where the strikes are in the space of 'k' and
-- implied volatility time slice is 'v'.
data Surface v k = Surface
  {
    surfaceSpot              :: Spot                   -- ^ Spot.
  , surfaceTenors            :: SortedVector YearFrac  -- ^ Ordered list of tenors.
  , surfaceForwardCurve      :: YearFrac -> Forward    -- ^ The forward curve.
  , surfaceDiscountCurve     :: YearFrac -> DF         -- ^ The discount curve.
  , surfaceAtmTotalVar       :: YearFrac -> TotalVar   -- ^ A spline of the at the money total variance.
  , surfaceVols              :: M.Map YearFrac v       -- ^ Map from tenor to 'TimeSlice'
  , surfaceTimeInterpolation :: TimeInterpolation      -- ^ Method of interpolation between tenors.
  , surfaceType              :: VolType                -- ^ The type of surface.
  }

totalVarKT :: (StrikeSpace k, TimeSlice v k) => Surface v k -> Strike -> YearFrac -> TotalVar
totalVarKT surface@Surface{..} k t | t <= minElement surfaceTenors =
                                       extrapolateTotalVarFrom (minElement surfaceTenors) surface k t
                                   | t >= maxElement surfaceTenors =
                                       extrapolateTotalVarFrom (maxElement surfaceTenors) surface k t
                                   | otherwise =
                                       timeInterpolate surfaceTimeInterpolation surface k t

volKT :: (StrikeSpace k, TimeSlice v k) => Surface v k -> Strike -> YearFrac ->  Vol
volKT surface k t = totalVarToVol (totalVarKT surface k t) t

fwdTotalVarKT :: ( StrikeSpace k, TimeSlice v k) => Surface v k -> Strike -> YearFrac -> Strike -> YearFrac -> TotalVar
fwdTotalVarKT surface@Surface{..} k1 t1 k2 t2 = TotalVar $ (totalVarKT2 - totalVarKT1) / (unYearFrac t2 - unYearFrac t1)
  where (TotalVar totalVarKT1) = totalVarKT surface k1 t1
        (TotalVar totalVarKT2) = totalVarKT surface k2 t2

class StrikeSpace k where
  strikeSpaceToCash :: k -> YearFrac -> Spot -> Forward -> Vol -> VolShift -> Strike
  cashToStrikeSpace :: Strike -> YearFrac -> Spot -> Forward -> Vol -> VolShift -> k

instance StrikeSpace Strike where
  strikeSpaceToCash x _ _ _ _ _ = x
  cashToStrikeSpace k _ _ _ _ _ = k

instance StrikeSpace AbsRelStrike where
  strikeSpaceToCash (AbsRel x) _ _ (Forward f) _ _ = Strike $ x + f
  cashToStrikeSpace (Strike k) _ _ (Forward f) _ _ = AbsRel $ k - f

instance StrikeSpace LogRelStrike where
  strikeSpaceToCash (LogRel x) _ _ (Forward f) _ _ = Strike $ f * exp x
  cashToStrikeSpace (Strike k) _ _ (Forward f) _ _ = LogRel $ log $ k - f

instance StrikeSpace MoneynessForwardStrike  where
  strikeSpaceToCash (MoneynessForward x) (YearFrac t) _ (Forward f) (Vol atmVol) _ =
    Strike $ x * sqrt t * atmVol + f

  cashToStrikeSpace (Strike k) (YearFrac t) _ (Forward f) (Vol atmVol) _ =
    MoneynessForward $ (k - f) / (atmVol * sqrt t)

instance StrikeSpace LogMoneynessForwardStrike  where
  strikeSpaceToCash (LogMoneynessForward x) (YearFrac t) _ (Forward f) (Vol atmVol) (VolShift slnShift) =
    Strike $ (f + slnShift) * exp (x * (sqrt t) * atmVol) - slnShift

  cashToStrikeSpace (Strike k) (YearFrac t)  _ (Forward f) (Vol atmVol) (VolShift slnShift) =
    LogMoneynessForward $ (log ((k - slnShift) / (f + slnShift))) / (atmVol * sqrt t)

instance StrikeSpace LogMoneynessSpotStrike  where
  strikeSpaceToCash (LogMoneynessSpot x) (YearFrac t) (Spot s) _ (Vol atmVol) (VolShift slnShift) =
    Strike $ (s + slnShift) * exp (x * (sqrt t) * atmVol) - slnShift

  cashToStrikeSpace (Strike k) (YearFrac t)  (Spot s) _ (Vol atmVol) (VolShift slnShift) =
    LogMoneynessSpot $ (log ((k - slnShift) / (s + slnShift))) / (atmVol * sqrt t)

instance StrikeSpace MoneynessSpotStrike  where
  strikeSpaceToCash (MoneynessSpot x) (YearFrac t) (Spot s) _ (Vol atmVol) _ =
    Strike $ x * (sqrt t) * atmVol + s

  cashToStrikeSpace (Strike k) (YearFrac t) (Spot s) _ (Vol atmVol) _ =
    MoneynessSpot $ (k - s) / (atmVol * sqrt t)

slnShift Surface{..} = case surfaceType of
  (ShiftedLogNormal shift) -> shift
  _                        -> VolShift 0
extrapolateTotalVarFrom :: forall v k. (StrikeSpace k, TimeSlice v k) => YearFrac -> Surface v k -> Strike -> YearFrac -> TotalVar
extrapolateTotalVarFrom t0 surface@Surface{..} k t = let
  f0          = surfaceForwardCurve t0
  atmVol0     = totalVarToVol (surfaceAtmTotalVar t0) t0
  f           = surfaceForwardCurve t
  spot        = surfaceSpot
  atmTotalVar = surfaceAtmTotalVar t
  atmVol      = totalVarToVol atmTotalVar t
  x           = cashToStrikeSpace k t spot f atmVol (slnShift surface)::k
  k'          = strikeSpaceToCash x t0 spot f0 atmVol0 (slnShift surface)
  x'          = cashToStrikeSpace k' t0 spot f0 atmVol (slnShift surface)::k
  in totalVar (surfaceVols M.! t0) x'

timeInterpolate :: forall v k. (StrikeSpace k, TimeSlice v k) => TimeInterpolation -> Surface v k -> Strike -> YearFrac -> TotalVar
timeInterpolate Gatheral surface@Surface{..} k11 t =
  let (t1, smile1) = fromJust $ M.lookupLE t surfaceVols
      (t2, smile2) = fromJust $ M.lookupGE t surfaceVols
      (TotalVar thetaT)  = surfaceAtmTotalVar t
      (TotalVar thetaT1) = surfaceAtmTotalVar t1
      (TotalVar thetaT2) = surfaceAtmTotalVar t2
      alphaT  = (sqrt thetaT2 - sqrt thetaT1 ) / (sqrt thetaT1 - sqrt thetaT)
      atmVol  = totalVarToVol (TotalVar thetaT) t
      (Forward f)  = surfaceForwardCurve t
      (Forward f1) = surfaceForwardCurve t1
      (Forward f2) = surfaceForwardCurve t2
      df  = surfaceDiscountCurve t
      df1 = surfaceDiscountCurve t1
      df2 = surfaceDiscountCurve t2
      k1      = k11 $*$ (f1 / f)
      k2      = k11 $*$ (f2 / f)
      x1      = cashToStrikeSpace k1 t1 surfaceSpot (Forward f1) atmVol (slnShift surface)::k
      x2      = cashToStrikeSpace k2 t2 surfaceSpot (Forward f2) atmVol (slnShift surface)::k
      vol1 = totalVarToVol (totalVar smile1 x1) t1
      vol2 = totalVarToVol (totalVar smile2  x2) t2
      (Premium premium1) = vPremium $ euOption surfaceType k1 f1 df1 t1 vol1
      (Premium premium2) = vPremium $ euOption surfaceType k1 f2 df2 t1 vol1
      premiumT  = Premium $ (alphaT * premium1 $/$ k1 + (1- alphaT)* premium2 $/$ k2) $*$ k11
      x1 :: k
      x2 :: k
  in if surfaceType == Normal then
       volToTotalVar (euImpliedVol Normal Call (Forward f) k11 t df premiumT) t
     else
       volToTotalVar (euImpliedVol LogNormal Call (Forward f) k11 t df premiumT) t

timeInterpolate LinearInVol surface@Surface{..} k t =
  let f = surfaceForwardCurve t
      atmVol = totalVarToVol (surfaceAtmTotalVar t) t
      x      = cashToStrikeSpace k t surfaceSpot f atmVol (slnShift surface)::k
      (t1, smile1) = fromJust $ M.lookupLE t surfaceVols
      (t2, smile2) = fromJust $ M.lookupGE t surfaceVols
      (Vol vol1)         = totalVarToVol (totalVar smile1 x) t1
      (Vol vol2)         = totalVarToVol (totalVar smile2 x) t2
  in volToTotalVar (Vol $ linearInterpolate (t1, vol1) (t2, vol2) t) t

timeInterpolate LinearInTotalVar surface@Surface{..} k t =
  let f = surfaceForwardCurve t
      atmVol = totalVarToVol (surfaceAtmTotalVar t) t
      x      = cashToStrikeSpace k t surfaceSpot f atmVol (slnShift surface)::k
      (t1, smile1) = fromJust $ M.lookupLE t surfaceVols
      (t2, smile2) = fromJust $ M.lookupGE t surfaceVols
      (TotalVar tv1)         = totalVar smile1 x
      (TotalVar tv2)         = totalVar smile2 x
  in TotalVar $ linearInterpolate (t1, tv1) (t2, tv2) t

euOption Normal k f (DF df) (YearFrac t) vol =
  let r = Rate $ -(log df) / t
      bacherlier = Bacherlier.Bachelier (Forward f) r vol
  in Bacherlier.eucall bacherlier (YearFrac t) k

euOption _ k f df t vol =
  let b76 = B76.Black76 (Forward f) df t vol
  in B76.eucall b76 k

linearInterpolate (YearFrac t1, v1) (YearFrac t2, v2) (YearFrac t) =
  v1 + (v2 - v1)*(t - t1) / (t2 - t1)

max_stars_repo_name: ghais/lowq

max_issues_repo_name: ghais/lowq

max_forks_repo_name: ghais/lowq


text:
module OrientationAnalysis where

import           Control.Monad                as M
import           Data.Array.Repa              as R
import           Data.Binary                  (decodeFile)
import           Data.Complex
import           Data.List                    as L
import           DFT.Plan
import           Filter.Pinwheel
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC.CompletionField          (makeImagePlan, makeR2Z1T0Plan)
import           STC.CompletionField
import           STC.OrientationScaleAnalysis
import           System.Directory
import           System.Environment
import           System.FilePath

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:alphaStr:pinwheelFlagStr:imagePath:idxStr:numOrientationSampleStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      init = read initStr :: (Double, Double, Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      idx = read idxStr :: (Int, Int)
      numOrientationSample = read numOrientationSampleStr :: Int
      numThread = read numThreadStr :: Int
      pinwheelParams =
        PinwheelParams numPoint numPoint alpha (exp 1) theta0Freqs [0]
      folderPath = "output/test/OrientationAnalysis"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  arrR2Z1T0 <-
    if pinwheelFlag
      then error
             "Using pinwheels to construct the Green's function has not been implemented yet."
           --computeR2Z1T0Array numPoint numPoint alpha thetaFreqs theta0Freqs
      else if flag
             then do
               getNormalizedHistogramArr <$> decodeFile histFilePath
             else solveMonteCarloR2Z1T0
                    numThread
                    numTrail
                    maxTrail
                    numPoint
                    numPoint
                    sigma
                    tao
                    len
                    theta0Freqs
                    thetaFreqs
                    ""
                    (emptyHistogram
                       [ numPoint
                       , numPoint
                       , L.length theta0Freqs
                       , L.length thetaFreqs
                       ]
                       0)
  (ImageRepa _ img) <- readImageRepa imagePath False
  plotImageRepa (folderPath </> "input.png") (ImageRepa 8 img)
  plan0 <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  (plan1, imgF) <- makeImagePlan plan0 img
  (plan, filterF, filterPIF) <- pinwheelFilter plan1 pinwheelParams
  arrR2Z1T0F <- dftR2Z1T0 plan . makeFilterR2Z1T0 $ arrR2Z1T0
  convolvedImg <-
    fmap (\x -> R.slice x (Z :. All :. (0 :: Int) :. All :. All)) $
    convolvePinwheel plan filterF (R.slice imgF (Z :. (0 :: Int) :. All :. All))
  plotMagnitudeOrientationSource
    plan
    folderPath
    numOrientationSample
    numOrientation
    thetaFreqs
    arrR2Z1T0F
    convolvedImg
    idx
  plotMagnitudeOrientation
    folderPath
    numOrientationSample
    numOrientation
    thetaFreqs
    convolvedImg
    idx


text:
-- |
-- Module:       Math.NumberTheory.DirichletCharactersTests
-- Copyright:    (c) 2018 Bhavik Mehta
-- License:      MIT
-- Maintainer:   Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Tests for Math.NumberTheory.DirichletCharacters
--

{-# LANGUAGE GADTs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}

module Math.NumberTheory.DirichletCharactersTests where

import Test.Tasty

import Data.Complex
import Data.List (genericLength)
import Data.Maybe (isJust, mapMaybe)
import Data.Proxy
import Data.Semigroup
import qualified Data.Vector as V
import Numeric.Natural

import GHC.TypeNats (SomeNat(..), someNatVal, KnownNat, natVal, sameNat)
import Data.Type.Equality

import Math.NumberTheory.ArithmeticFunctions (totient, divisorsList)
import Math.NumberTheory.DirichletCharacters
import qualified Math.NumberTheory.Moduli.Sqrt as J
import Math.NumberTheory.Moduli.Class (SomeMod(..), modulo)
import Math.NumberTheory.TestUtils (testSmallAndQuick, Positive(..))

-- | This tests property 6 from https://en.wikipedia.org/wiki/Dirichlet_character#Axiomatic_definition
dirCharOrder :: forall n. KnownNat n => DirichletCharacter n -> Bool
dirCharOrder chi = isPrincipal (totient n `stimes` chi)
  where n = natVal @n Proxy

-- | Tests wikipedia's property 3 (note 1,2,5 are essentially enforced by the type system).
testMultiplicative :: KnownNat n => DirichletCharacter n -> Natural -> Natural -> Bool
testMultiplicative chi (fromIntegral -> a) (fromIntegral -> b) = chiAB == chiAchiB
  where chi' = evalGeneral chi
        chiAB = chi' (a*b)
        chiAchiB = (<>) <$> chi' a <*> chi' b

-- | Test property 4 from wikipedia
testAtOne :: KnownNat n => DirichletCharacter n -> Bool
testAtOne chi = eval chi mempty == mempty

dirCharProperty :: (forall n. KnownNat n => DirichletCharacter n -> a) -> Positive Natural -> Natural -> a
dirCharProperty test (Positive n) i =
  case someNatVal n of
    SomeNat (Proxy :: Proxy n) -> test chi
      where chi = indexToChar @n (i `mod` totient n)

realCharProperty :: (forall n. KnownNat n => RealCharacter n -> a) -> Positive Natural -> Int -> a
realCharProperty test (Positive n) i =
  case someNatVal n of
    SomeNat (Proxy :: Proxy n) -> test chi
      where chi = chars !! (i `mod` length chars)
            chars = mapMaybe isRealCharacter [principalChar @n .. maxBound]

-- | There should be totient(n) characters
countCharacters :: Positive Natural -> Bool
countCharacters (Positive n) =
  case someNatVal n of
    SomeNat (Proxy :: Proxy n) ->
      genericLength (allChars @n) == totient n

-- | The principal character should be 1 if gcd k n is 1 and 0 otherwise
principalCase :: Positive Natural -> Positive Integer -> Bool
principalCase (Positive n) (Positive k) =
  case k `modulo` n of
    SomeMod a -> evalGeneral chi a == if gcd k (fromIntegral n) > 1
                                         then Zero
                                         else mempty
      where chi = principalChar
    InfMod{} -> False

-- | Test the orthogonality relations https://en.wikipedia.org/wiki/Dirichlet_character#Character_orthogonality
orthogonality1 :: forall n. KnownNat n => DirichletCharacter n -> Bool
orthogonality1 chi = magnitude (total - correct) < (1e-13 :: Double)
  where n = natVal @n Proxy
        total = sum [orZeroToNum toComplex (evalGeneral chi a) | a <- [0 .. maxBound]]
        correct = if isPrincipal chi
                     then fromIntegral $ totient n
                     else 0

orthogonality2 :: Positive Natural -> Integer -> Bool
orthogonality2 (Positive n) a =
  case a `modulo` n of
    SomeMod a' -> magnitude (total - correct) < (1e-13 :: Double)
      where total = sum [orZeroToNum toComplex (evalGeneral chi a') | chi <- allChars]
            correct = if a' == 1
                         then fromIntegral $ totient n
                         else 0
    InfMod {} -> False

-- | Manually confirm isRealCharacter is correct (in both directions)
realityCheck :: KnownNat n => DirichletCharacter n -> Bool
realityCheck chi = isJust (isRealCharacter chi) == isReal'
  where isReal' = and [real (evalGeneral chi t) | t <- [minBound..maxBound]]
        real Zero = True
        real (NonZero t) = t <> t == mempty

-- | Check real character evaluation matches normal evaluation
realEvalCheck :: KnownNat n => RealCharacter n -> Int -> Bool
realEvalCheck chi i' = fromIntegral (toRealFunction chi i) == (orZeroToNum toComplex (evalGeneral (getRealChar chi) i) :: Complex Double)
  where i = fromIntegral i'

-- | The jacobi character agrees with the jacobi symbol
jacobiCheck :: Positive Natural -> Bool
jacobiCheck (Positive n) =
  case someNatVal (2*n+1) of
    SomeNat (Proxy :: Proxy n) ->
      case jacobiCharacter @n of
        Just chi -> and [toRealFunction chi (fromIntegral j) == J.symbolToNum (J.jacobi j (2*n+1)) | j <- [0..2*n]]
        _ -> False

-- | Bulk evaluation agrees with pointwise evaluation
evalAllCheck :: forall n. KnownNat n => DirichletCharacter n -> Bool
evalAllCheck chi = V.generate (fromIntegral $ natVal @n Proxy) (evalGeneral chi . fromIntegral) == evalAll chi

-- | Induced characters agree with the original character.
-- (Except for when d=1, where chi(0) = 1, which is true for no other d)
inducedCheck :: forall d. KnownNat d => DirichletCharacter d -> Positive Natural -> Bool
inducedCheck chi (Positive k) =
  case someNatVal (d*k) of
    SomeNat (Proxy :: Proxy n) ->
      case induced @n chi of
        Just chi2 -> and (V.izipWith matchedValue (V.concat (replicate (fromIntegral k) (evalAll chi))) (evalAll chi2))
        Nothing -> False
  where d = natVal @d Proxy
        matchedValue i x1 x2 = if gcd (fromIntegral i) (d*k) > 1
                                  then x2 == Zero
                                  else x2 == x1

-- | Primitive checker is correct (in both directions)
primitiveCheck :: forall n. KnownNat n => DirichletCharacter n -> Bool
primitiveCheck chi = isJust (isPrimitive chi) == isPrimitive'
  where isPrimitive' = all testModulus possibleModuli
        n = fromIntegral (natVal @n Proxy) :: Int
        possibleModuli = init (divisorsList n)
        table = evalAll chi
        testModulus d = not $ null [a | a <- [1..n-1], gcd a n == 1, a `mod` d == 1 `mod` d, table V.! a /= mempty]

-- | Ensure that makePrimitive gives primitive characters
makePrimitiveCheck :: DirichletCharacter n -> Bool
makePrimitiveCheck chi = case makePrimitive chi of
                            WithNat chi' -> isJust (isPrimitive (getPrimitiveChar chi'))

-- | sameNat also ensures the two new moduli are the same
makePrimitiveIdem :: DirichletCharacter n -> Bool
makePrimitiveIdem chi = case makePrimitive chi of
                          WithNat (chi' :: PrimitiveCharacter n') ->
                            case makePrimitive (getPrimitiveChar chi') of
                              WithNat (chi'' :: PrimitiveCharacter n'') ->
                                case sameNat (Proxy :: Proxy n') (Proxy :: Proxy n'') of
                                  Just Refl -> chi' == chi''
                                  Nothing -> False

orderCheck :: DirichletCharacter n -> Bool
orderCheck chi = isPrincipal (n `stimes` chi) && and [not (isPrincipal (i `stimes` chi)) | i <- [1..n-1]]
  where n = orderChar chi

fromTableCheck :: forall n. KnownNat n => DirichletCharacter n -> Bool
fromTableCheck chi = isJust (fromTable @n (evalAll chi))

-- A bunch of functions making sure that every function which can produce a character (in
-- particular by fiddling internal representation) produces a valid character
indexToCharValid :: KnownNat n => DirichletCharacter n -> Bool
indexToCharValid = validChar

principalCharValid :: Positive Natural -> Bool
principalCharValid (Positive n) =
  case someNatVal n of
    SomeNat (Proxy :: Proxy n) -> validChar (principalChar @n)

mulCharsValid :: KnownNat n => DirichletCharacter n -> DirichletCharacter n -> Bool
mulCharsValid chi1 chi2 = validChar (chi1 <> chi2)

mulCharsValid' :: Positive Natural -> Natural -> Natural -> Bool
mulCharsValid' (Positive n) i j =
  case someNatVal n of
    SomeNat (Proxy :: Proxy n) ->
      mulCharsValid (indexToChar @n (i `mod` totient n)) (indexToChar @n (j `mod` totient n))

stimesCharValid :: KnownNat n => DirichletCharacter n -> Int -> Bool
stimesCharValid chi n = validChar (n `stimes` chi)

succValid :: KnownNat n => DirichletCharacter n -> Bool
succValid = validChar . succ

inducedValid :: forall d. KnownNat d => DirichletCharacter d -> Positive Natural -> Bool
inducedValid chi (Positive k) = case someNatVal (natVal @d Proxy * k) of
  SomeNat (Proxy :: Proxy n) -> maybe False validChar (induced @n chi)

jacobiValid :: Positive Natural -> Bool
jacobiValid (Positive n) =
  case someNatVal (2*n+1) of
    SomeNat (Proxy :: Proxy n) ->
      case jacobiCharacter @n of
        Just chi -> validChar (getRealChar chi)
        _ -> False

makePrimitiveValid :: DirichletCharacter n -> Bool
makePrimitiveValid chi = case makePrimitive chi of
                            WithNat chi' -> validChar (getPrimitiveChar chi')

testSuite :: TestTree
testSuite = testGroup "DirichletCharacters"
  [ testSmallAndQuick "Dirichlet characters divide the right order" (dirCharProperty dirCharOrder)
  , testSmallAndQuick "Dirichlet characters are multiplicative" (dirCharProperty testMultiplicative)
  , testSmallAndQuick "Dirichlet characters are 1 at 1" (dirCharProperty testAtOne)
  , testSmallAndQuick "Right number of Dirichlet characters" countCharacters
  , testSmallAndQuick "Principal character behaves as expected" principalCase
  , testSmallAndQuick "Orthogonality relation 1" (dirCharProperty orthogonality1)
  , testSmallAndQuick "Orthogonality relation 2" orthogonality2
  , testSmallAndQuick "Real character checking is correct" (dirCharProperty realityCheck)
  , testSmallAndQuick "Real character evaluation is accurate" (realCharProperty realEvalCheck)
  , testSmallAndQuick "Jacobi character matches symbol" jacobiCheck
  , testSmallAndQuick "Bulk evaluation matches pointwise" (dirCharProperty evalAllCheck)
  , testSmallAndQuick "Induced character is correct" (dirCharProperty inducedCheck)
  , testSmallAndQuick "Primitive character checking is correct" (dirCharProperty primitiveCheck)
  , testSmallAndQuick "makePrimitive produces primitive character" (dirCharProperty makePrimitiveCheck)
  , testSmallAndQuick "makePrimitive is idempotent" (dirCharProperty makePrimitiveIdem)
  , testSmallAndQuick "Calculates correct order" (dirCharProperty orderCheck)
  , testSmallAndQuick "Can construct from table" (dirCharProperty fromTableCheck)
  , testGroup "Creates valid characters"
    [ testSmallAndQuick "indexToChar" (dirCharProperty indexToCharValid)
    , testSmallAndQuick "principalChar" principalCharValid
    , testSmallAndQuick "mulChars" mulCharsValid'
    , testSmallAndQuick "stimesChar" (dirCharProperty stimesCharValid)
    , testSmallAndQuick "succ" (dirCharProperty succValid)
    , testSmallAndQuick "induced" (dirCharProperty inducedValid)
    , testSmallAndQuick "jacobi" jacobiValid
    , testSmallAndQuick "makePrimitive" (dirCharProperty makePrimitiveValid)
    ]
  ]

max_stars_repo_name: cjlarose/arithmoi

max_issues_repo_name: cjlarose/arithmoi

max_forks_repo_name: cjlarose/arithmoi


text:
module Numeric.LinearAlgebra.Matrix.Mat33 where

import Numeric.LinearAlgebra.Matrix.Class

import Data.List ( foldl1' )

data Mat33 a = Mat33
  { m00 :: !a, m01 :: !a, m02 :: !a
  , m10 :: !a, m11 :: !a, m12 :: !a
  , m20 :: !a, m21 :: !a, m22 :: !a
  }
  deriving (Read, Show, Eq, Ord)

instance Matrix Mat33 where
  {-# INLINE mDim #-}
  mDim _ = 3
  {-# INLINE mElement #-}
  mElement m 0 0 = m00 m
  mElement m 0 1 = m01 m
  mElement m 0 2 = m02 m
  mElement m 1 0 = m10 m
  mElement m 1 1 = m11 m
  mElement m 1 2 = m12 m
  mElement m 2 0 = m20 m
  mElement m 2 1 = m21 m
  mElement m 2 2 = m22 m
  mElement _ i j = error ("Index " ++ show i ++ ", " ++ show j ++ ": out of range, must be 0,0 to 2,2")
  {-# INLINE mZip #-}
  mZip f m n = Mat33 (f (m00 m) (m00 n)) (f (m01 m) (m01 n)) (f (m02 m) (m02 n))
                     (f (m10 m) (m10 n)) (f (m11 m) (m11 n)) (f (m12 m) (m12 n))
                     (f (m20 m) (m20 n)) (f (m21 m) (m21 n)) (f (m22 m) (m22 n))
  {-# INLINE mFold #-}
  mFold f m = foldl1' f [ mElement m i j | i <- [ 0 .. 2 ], j <- [ 0 .. 2 ] ]
  {-# INLINE mIndexOf #-}
  mIndexOf p m = fst (foldl1' p' [ ((i,j), mElement m i j) | j <- [ 2, 1 .. 0 ], i <- [ 2, 1 .. 0 ] ])
    where
    p' x@(_, a) y@(_, a') | a `p` a'  = x
                          | otherwise = y
  {-# INLINE det #-}
  det (Mat33 a b c d e f g h i) = det33 a b c d e f g h i

instance Functor Mat33 where
  {-# INLINE fmap #-}
  fmap f m = Mat33 (f (m00 m)) (f (m01 m)) (f (m02 m))
                   (f (m10 m)) (f (m11 m)) (f (m12 m))
                   (f (m20 m)) (f (m21 m)) (f (m22 m))

{-# SPECIALIZE det33 :: Double -> Double -> Double
                     -> Double -> Double -> Double
                     -> Double -> Double -> Double
                     -> Double #-}
{-# SPECIALIZE det33 :: Float -> Float -> Float
                     -> Float -> Float -> Float
                     -> Float -> Float -> Float
                     -> Float #-}
det33 :: Num a
      => a -> a -> a
      -> a -> a -> a
      -> a -> a -> a
      -> a
det33 a b c d e f g h i = a*e*i + d*h*c + g*b*f - g*e*c - d*b*i - a*h*f

{- SPECIALIZE inv33 :: Mat33 Double -> Mat33 Double -}
{- SPECIALIZE inv33 :: Mat33 Float  -> Mat33 Float -}
{-
inv33 :: Fractional a => Mat33 a -> Mat33 a
inv33 m = m'
  where
  d  = det m
  m' = Mat33 (m11 * m22 - m21 * m21)
  m' = Mat44 (det3 (m11 m) (m12 m) (m13 m)
                   (m21 m) (m22 m) (m23 m)
                   (m31 m) (m32 m) (m33 m) / d)
           (-(det3 (m01 m) (m02 m) (m03 m)
                   (m21 m) (m22 m) (m23 m)
                   (m31 m) (m32 m) (m33 m) / d))
             (det3 (m01 m) (m02 m) (m03 m)
                   (m11 m) (m12 m) (m13 m)
                   (m31 m) (m32 m) (m33 m) / d)
           (-(det3 (m01 m) (m02 m) (m03 m)
                   (m11 m) (m12 m) (m13 m)
                   (m21 m) (m22 m) (m23 m) / d))

           (-(det3 (m10 m) (m12 m) (m13 m)
                   (m20 m) (m22 m) (m23 m)
                   (m30 m) (m32 m) (m33 m) / d))
             (det3 (m00 m) (m02 m) (m03 m)
                   (m20 m) (m22 m) (m23 m)
                   (m30 m) (m32 m) (m33 m) / d)
           (-(det3 (m00 m) (m02 m) (m03 m)
                   (m10 m) (m12 m) (m13 m)
                   (m30 m) (m32 m) (m33 m) / d))
             (det3 (m00 m) (m02 m) (m03 m)
                   (m10 m) (m12 m) (m13 m)
                   (m20 m) (m22 m) (m23 m) / d)

             (det3 (m10 m) (m11 m) (m13 m)
                   (m20 m) (m21 m) (m23 m)
                   (m30 m) (m31 m) (m33 m) / d)
           (-(det3 (m00 m) (m01 m) (m03 m)
                   (m20 m) (m21 m) (m23 m)
                   (m30 m) (m31 m) (m33 m) / d))
             (det3 (m00 m) (m01 m) (m03 m)
                   (m10 m) (m11 m) (m13 m)
                   (m30 m) (m31 m) (m33 m) / d)
           (-(det3 (m00 m) (m01 m) (m03 m)
                   (m10 m) (m11 m) (m13 m)
                   (m20 m) (m21 m) (m23 m) / d))

           (-(det3 (m10 m) (m11 m) (m12 m)
                   (m20 m) (m21 m) (m22 m)
                   (m30 m) (m31 m) (m32 m) / d))
             (det3 (m00 m) (m01 m) (m02 m)
                   (m20 m) (m21 m) (m22 m)
                   (m30 m) (m31 m) (m32 m) / d)
           (-(det3 (m00 m) (m01 m) (m02 m)
                   (m10 m) (m11 m) (m12 m)
                   (m30 m) (m31 m) (m32 m) / d))
             (det3 (m00 m) (m01 m) (m02 m)
                   (m10 m) (m11 m) (m12 m)
                   (m20 m) (m21 m) (m22 m) / d)
-}

max_stars_repo_name: dagit/lin-alg

max_issues_repo_name: dagit/lin-alg

max_forks_repo_name: dagit/lin-alg


text:
{-# OPTIONS_GHC -Wall #-}
module Diagrams.Spirograph
  ( Spiro(Spiro)
  , drawSpiros
  , drawSpirosIntermediate
  ) where

import           Data.Complex                 (Complex(..), mkPolar)
import           Diagrams.Prelude             (Diagram, P2)
import qualified Diagrams.Backend.SVG.CmdLine as SVG
import qualified Diagrams.Prelude             as D
import           Data.Monoid                  (Sum(..))

data Spiro = Spiro
  { offset ∷ Double
  , speed  ∷ Double
  }

newtype Path w = Path (Double → w)

deriving newtype instance Semigroup w ⇒ Semigroup (Path w)
deriving newtype instance Monoid w    ⇒ Monoid    (Path w)

newtype P = P (Sum (P2 Double))

runP ∷ P → P2 Double
runP (P (Sum p)) = p

deriving newtype instance Semigroup P
deriving newtype instance Monoid P

-- | @'disk' d@ is the path that @d@ will travel.
disk ∷ Spiro → Path P
disk Spiro{..} = Path $ \t → P $ Sum $ p2 $ mkPolar offset (speed * t)

disks ∷ [Spiro] → Path P
disks = foldMap disk

allSpiros ∷ [Spiro] → [Path P]
allSpiros = \case
  []     → mempty
  (x:xs) → px : map f pxs
    where
    px  = disk x
    pxs = allSpiros xs
    f   ∷ Path P → Path P
    f p = px <> p

-- | Form a point in diagrams in terms of a complex number.
p2 ∷ Complex a → P2 a
p2 (x :+ y) = D.p2 (x, y)

-- | Convert a path into a path.
sample ∷ Path P → [P2 Double]
sample (Path f) = runP . f <$> enumFromThenTo 0 0.1 100

drawSpiros ∷ [Spiro] → Diagram SVG.B
drawSpiros xs = D.fromVertices $ sample (disks xs)

drawSpirosIntermediate ∷ [Spiro] → Diagram SVG.B
drawSpirosIntermediate xs = foldMap step (allSpiros xs)
  where
  step p = D.fromVertices $ sample p

max_stars_repo_name: fredefox/spirograph

max_issues_repo_name: fredefox/spirograph

max_forks_repo_name: fredefox/spirograph


text:
module STCR2Z2T0S0PointSetBinary where

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

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:hollowRadiusStr:cutoffRadiusStr:histFilePath:filterFileFolder:numIterationStr:writeSourceFlagStr:saveEdgeDataFlagStr:loadEdgeDataFlagStr:shape2DStr:useFFTWWisdomFlagStr:fftwWisdomFileName:batchSizeStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      hollowRadius = read hollowRadiusStr :: Double
      cutoffRadius = read cutoffRadiusStr :: Double
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      saveEdgeDataFlag = read saveEdgeDataFlagStr :: Bool
      loadEdgeDataFlag = read loadEdgeDataFlagStr :: Bool
      shape2D@(Points _ minDist _) = read shape2DStr :: Points Shape2D
      useFFTWWisdomFlag = read useFFTWWisdomFlagStr :: Bool
      batchSize = read batchSizeStr :: Int
      numThread = read numThreadStr :: Int
      parallelParams = ParallelParams numThread batchSize
      folderPath = "output/test/STCR2Z2T0S0PointSetBinary"
      fftwWisdomFilePath = folderPath </> fftwWisdomFileName
      filterFileName =
        printf
          "Filter_%.0f_%.0f_%s"
          hollowRadius
          cutoffRadius
          (takeFileName histFilePath)
      filterFilePath = filterFileFolder </> filterFileName
      eigenVecFilePath = printf "%s/EigenVec.dat" filterFileFolder
  createDirectoryIfMissing True folderPath
  createDirectoryIfMissing True filterFileFolder
  plan <-
    makePlanBinary
      emptyPlan
      useFFTWWisdomFlag
      fftwWisdomFilePath
      (L.length thetaFreqs)
      (L.length scaleFreqs)
      numPoint
      numPoint
  filterFlag <- doesFileExist filterFilePath
  flag <-
    if filterFlag
      then do
        size <- getFileSize filterFilePath
        return $
          if size == 0
            then False
            else True
      else return False
  unless
    flag
    (do histFlag <- doesFileExist histFilePath
        radialArr <-
          if histFlag
            then R.map magnitude . getNormalizedHistogramArr <$>
                 decodeFile histFilePath
            else do
              putStrLn
                "Couldn't find a Green's function data. Start simulation..."
              solveMonteCarloR2Z2T0S0Radial
                numThread
                numTrail
                maxTrail
                numPoint
                numPoint
                thetaSigma
                scaleSigma
                maxScale
                tao
                thetaFreqs
                thetaFreqs
                scaleFreqs
                scaleFreqs
                histFilePath
                (emptyHistogram
                   [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
                   , L.length scaleFreqs
                   , L.length thetaFreqs
                   , L.length scaleFreqs
                   , L.length thetaFreqs
                   ]
                   0)
        putStrLn "Write filter to disk..."
        writeDFTPinwheel
          parallelParams
          plan
          radialArr
          hollowRadius
          cutoffRadius
          thetaFreqs
          scaleFreqs
          maxScale
          (numPoint, numPoint)
          filterFilePath)
  let pointSet = makeShape2D shape2D
      points =
        L.map (\(x, y) -> R2S1RPPoint (x, y, 0, 1)) . getShape2DIndexList $
        pointSet
      bias = computeBiasR2T0S0 numPoint numPoint thetaFreqs scaleFreqs points
      eigenVec =
        computeInitialEigenVectorBinary
          numPoint
          numPoint
          thetaFreqs
          scaleFreqs
          points
  eigenVec <-
    if loadEdgeDataFlag
      then readRepaArray eigenVecFilePath
      else return $
           computeInitialEigenVectorBinary
             numPoint
             numPoint
             thetaFreqs
             scaleFreqs
             points
  powerMethodBinary
    parallelParams
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    numScale
    scaleFreqs
    maxScale
    filterFilePath
    numIteration
    writeSourceFlag
    (printf
       "_%.2f_%.2f_%d_%d_%d_%d"
       thetaSigma
       scaleSigma
       (round maxScale :: Int)
       (round tao :: Int)
       (round thetaFreq :: Int)
       (round scaleFreq :: Int))
    saveEdgeDataFlag
    eigenVecFilePath
    bias
    eigenVec


text:
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
module Myelin.Model where

import Control.Lens hiding ((.=), (*~))
import Data.Aeson
import Numeric.LinearAlgebra

import Myelin.Neuron

{-- TODO: Implement units
import Numeric.Units.Dimensional
import Numeric.Units.Dimensional.UnitNames
import Numeric.Units.Dimensional.Quantities
import Numeric.Units.Dimensional.SIUnits

-- TODO(Christian): Move Units somewhere else
type Siemens = Unit Metric DElectricConductance Float
type Volt = Unit Metric DElectricPotential Float
type Farrad = Unit Metric DCapacitance Float
type Seconds = Unit Metric DTime Float
type Ampere = Unit Metric DElectricCurrent Float
-}

type Label = String

-- | A Node is a vertex in the connectivity graph.
data Node = Population {
        _numNeurons :: Integer, -- ^ number of neurons in the population
        _neuronType :: NeuronType,
        _label :: Label, -- ^ human readable label for the population
        _id :: Int -- ^ internal identifier to keep track of the population
    }
    | SpikeSourceArray {
        _spikeTimes :: [Float],
        _id :: Int
    }
    | SpikeSourcePoisson {
        _rate :: Float,
        _start :: Integer,
        _id :: Int
    }
    deriving (Eq, Show)

makeLenses ''Node
makePrisms ''Node

type Delay = Float -- TODO: This is platform specific
type Probability = Float -- TODO: Very unsophisticated representation

data ParameterGenerator
  = Constant Float
  | GaussianRandom
    { mean :: Float
    , scale :: Float
    }
  deriving (Eq, Show)

-- | Biases that biases the output of a single projection
data Biases
  = BiasGenerator ParameterGenerator
  | Biases [Float]
  deriving (Eq, Show)

-- | Weight types that describes weights for a projection
data Weights
  = WeightGenerator ParameterGenerator
  | Weights (Matrix Float)
  deriving (Eq, Show)

makeLenses ''Weights
makePrisms ''Weights

-- | Types of connections in a projection from one neuron population to another
data ConnectionType
  = AllToAll
    { _biases :: Biases
    , _weights :: Weights
    }
  | OneToOne
    { _biases :: Biases
    , _weights :: Weights
    }
  deriving (Eq, Show)

makeLenses ''ConnectionType
makePrisms ''ConnectionType

-- | Determines the effect of a synapse to be excitatory of inhibitory
data SynapseEffect
  = Inhibitory
  | Excitatory
  deriving (Eq, Show)

makeLenses '' SynapseEffect
makePrisms '' SynapseEffect

-- | Effect of projections that determines the method with which 'Node's connect
data ProjectionEffect
  = Static SynapseEffect ConnectionType
-- | Dynamic .. TODO (Jens)
  deriving (Eq, Show)

makeLenses ''ProjectionEffect
makePrisms ''ProjectionEffect

-- | An edge between two units in the network with a given 'ProjectionEffect'
data Edge
  = DenseProjection
    { _effect :: ProjectionEffect
    , _input :: Node
    , _output :: Node
    }
  | MergeProjection
    { _effect :: ProjectionEffect
    , _inputs :: (Node, Node)
    , _output :: Node
    }
  | ReplicateProjection
    { _effect :: ProjectionEffect
    , _input :: Node
    , _outputs :: (Node, Node)
    }
  deriving (Eq, Show)

max_stars_repo_name: volr/myelin

max_issues_repo_name: volr/myelin

max_forks_repo_name: volr/myelin


text:
{-# LANGUAGE FlexibleContexts #-}

module School.Utils.LinearAlgebra
( compareDoubleMatrix
, compareDoubleVector
, constVector
, oneMatrix
, oneVector
, mapCols
, mapRows
, sumCols
, sumRows
, zeroMatrix
, zeroVector
) where

import Numeric.LinearAlgebra
import School.Types.FloatEq (compareDouble)

mapCols :: (Element a)
        => (Vector a -> Vector a)
        -> Matrix a
        -> Matrix a
mapCols f = fromColumns . map f . toColumns

mapRows :: (Element a)
        => (Vector a -> Vector a)
        -> Matrix a
        -> Matrix a
mapRows f = fromRows . map f . toRows

sumCols :: Matrix R -> Vector R
sumCols m = n |> summedCols where
  n = rows m
  summedCols = map sumElements $ toRows m

sumRows :: Matrix R -> Vector R
sumRows m = n |> summedRows where
  n = cols m
  summedRows = map sumElements $ toColumns m

oneVector :: Int -> Vector R
oneVector = konst 1

constVector :: Int -> Double -> Vector R
constVector d v = konst v d

zeroVector :: Int -> Vector R
zeroVector = konst 0

zeroMatrix :: (Container Vector a, Num a)
           => Int
           -> Int
           -> Matrix a
zeroMatrix r c = konst 0 (r, c)

oneMatrix :: Int -> Int -> Matrix R
oneMatrix r c = konst 1 (r, c)

compareDoubleMatrix :: Double
                    -> Matrix R
                    -> Matrix R
                    -> Bool
compareDoubleMatrix prec m1 m2 =
  and $ zipWith (compareDouble prec) l1 l2
  where l1 = concat . toLists $ m1
        l2 = concat . toLists $ m2

compareDoubleVector :: Double
                    -> Vector R
                    -> Vector R
                    -> Bool
compareDoubleVector prec v1 v2 =
  and $ zipWith (compareDouble prec) l1 l2
  where l1 = toList v1
        l2 = toList v2

max_stars_repo_name: jfulseca/School

max_issues_repo_name: jfulseca/School

max_forks_repo_name: jfulseca/School