text:
module STCR2Z1T0 where

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

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:initDistStr:histFilePath:alphaStr:pinwheelFlagStr: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]
      initDist = read initDistStr :: [R2S1RPPoint]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numThread = read numThreadStr :: Int
      sourceDist = L.take 1 initDist
      sinkDist = L.drop 1 initDist
      folderPath = "output/test/STCR2Z1T0"
  flag <- doesFileExist histFilePath
  -- arrR2Z1T0 <-
  --   if pinwheelFlag
  --     then computeR2Z1T0Array numPoint numPoint alpha thetaFreqs theta0Freqs
  --     else if flag
  --            then getNormalizedHistogramArr <$> decodeFile histFilePath
  --            else solveMonteCarloR2Z1T0
  --                   numThread
  --                   numTrail
  --                   maxTrail
  --                   numPoint
  --                   numPoint
  --                   sigma
  --                   tao
  --                   1
  --                   len
  --                   theta0Freqs
  --                   thetaFreqs
  --                   histFilePath
  --                   (emptyHistogram
  --                      [ numPoint
  --                      , numPoint
  --                      , L.length theta0Freqs
  --                      , L.length thetaFreqs
  --                      ]
  --                      0)
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z1T0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          sigma
          tao
          0
          theta0Freqs
          thetaFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length theta0Freqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z1T0 <-
    computeUnboxedP $
    computeR2Z1T0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      thetaFreqs
      theta0Freqs
  let arr3d =
        rotate3D . R.slice arrR2Z1T0 $
        (Z :. All :. (L.length theta0Freqs - 1) :. All :. All)
  createDirectoryIfMissing True folderPath
  MP.mapM_
    (\i ->
       plotImageRepaComplex
         (folderPath </> "GreensR2Z1T0_" L.++ show (i + 1) L.++ ".png") .
       ImageRepa 8 .
       computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr3d $
       (Z :. All :. All :. i))
    [0 .. (L.length thetaFreqs) - 1]
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  sourceDistArr <-
    computeInitialDistributionR2T0 plan numPoint numPoint theta0Freqs sourceDist
  sinkDistArr <-
    computeInitialDistributionR2T0 plan numPoint numPoint theta0Freqs sinkDist
  arrR2Z1T0F <- (computeP . makeFilter2D $ arrR2Z1T0) >>= dftR2Z1T0 plan
  arrR2Z1T0TRF <-
    (computeP . makeFilter2D . timeReverseR2Z1T0 thetaFreqs theta0Freqs $
     arrR2Z1T0) >>=
    dftR2Z1T0 plan
  -- Source field
  sourceArr <- convolveR2T0 plan arrR2Z1T0F sourceDistArr
  -- let sourceR2Z1 = R.sumS . rotateR2Z1T0Array $ sourceArr
  --     sourceField =
  --       computeS .
  --       R.extend (Z :. (1 :: Int) :. All :. All) .
  --       R.sumS . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
  --       sourceR2Z1
  sourceR2Z1 <- R.sumP . rotateR2Z1T0Array $ sourceArr
  sourceField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
    sourceR2Z1
  plotImageRepaComplex (folderPath </> "Source.png") . ImageRepa 8 $ sourceField
  -- Sink field
  sinkArr' <- convolveR2T0 plan arrR2Z1T0F sinkDistArr
  let   sinkArr  = computeSinkFromSourceR2Z1T0 thetaFreqs theta0Freqs sinkArr'
  -- let sinkR2Z1 = R.sumS . rotateR2Z1T0Array $ sinkArr
  --     sinkField =
  --       computeS .
  --       R.extend (Z :. (1 :: Int) :. All :. All) .
  --       R.sumS . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
  --       sinkR2Z1
  sinkR2Z1 <- R.sumP . rotateR2Z1T0Array $ sinkArr
  sinkField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
    sinkR2Z1
  plotImageRepaComplex (folderPath </> "Sink.png") . ImageRepa 8 $ sinkField
  -- Completion Filed
  completionFiled <-
    convolveR2Z1 plan thetaFreqs sourceR2Z1 sinkR2Z1
  let completionFiledR2 =
        R.sumS . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
        completionFiled
  -- completionFiledR2 <-
  --   R.sumP . R.map magnitude . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
  --   completionFiled
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.map magnitude $
    completionFiledR2

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:
module Mandelbruh
    ( C,
      mandelbruh,
      inMandelbrot,
      color,
      series
    ) where

import Data.Complex

type C = Complex Float

f :: C -> C -> C
f c z = z^2 + c

type Series = [C]

series :: C -> Series
series c = iterate (f c) c

inMandelbrot :: C -> Bool
inMandelbrot =
  let ceiling = 100 -- TODO: move to config file
      samples = 20
      in all (< ceiling) . map magnitude . take samples . series

mandelbruh = 42

readComplex :: String -> C
readComplex s = read s :: C

type Color = Char
type Palette = [Color]

color :: Palette -> Series -> Color
color palette = (palette !!) . length . take n . takeWhile fairlyClose
  where n = length palette - 1
        fairlyClose c = magnitude c < ceiling
        ceiling = 100

type Image color = C -> color

ext: hs

max_stars_repo_name: brunolange/mandelbruh

max_issues_repo_name: brunolange/mandelbruh

max_forks_repo_name: brunolange/mandelbruh


text:
{-# LANGUAGE CPP #-}

-- !!! Testing Typeable instances
module Main(main) where

import Data.Dynamic
#if MIN_VERSION_base(4,10,0)
import Data.Typeable (TyCon, TypeRep, typeOf)
#endif
import Data.Array
import Data.Array.MArray
import Data.Array.ST
import Data.Array.IO
import Data.Array.Unboxed
import Data.Complex
import Data.Int
import Data.Word
import Data.IORef
import System.IO
import Control.Monad.ST
import System.Mem.StableName
import System.Mem.Weak
import Foreign.StablePtr
import Control.Exception
import Foreign.C.Types

main :: IO ()
main = do
   print (typeOf (undefined :: [()]))
   print (typeOf (undefined :: ()))
   print (typeOf (undefined :: ((),())))
   print (typeOf (undefined :: ((),(),())))
   print (typeOf (undefined :: ((),(),(),())))
   print (typeOf (undefined :: ((),(),(),(),())))
   print (typeOf (undefined :: (() -> ())))
   print (typeOf (undefined :: (Array () ())))
   print (typeOf (undefined :: Bool))
   print (typeOf (undefined :: Char))
   print (typeOf (undefined :: (Complex ())))
   print (typeOf (undefined :: Double))
   print (typeOf (undefined :: (Either () ())))
   print (typeOf (undefined :: Float))
   print (typeOf (undefined :: Handle))
   print (typeOf (undefined :: Int))
   print (typeOf (undefined :: Integer))
   print (typeOf (undefined :: IO ()))
   print (typeOf (undefined :: (Maybe ())))
   print (typeOf (undefined :: Ordering))

   print (typeOf (undefined :: Dynamic))
   print (typeOf (undefined :: (IORef ())))
   print (typeOf (undefined :: Int8))
   print (typeOf (undefined :: Int16))
   print (typeOf (undefined :: Int32))
   print (typeOf (undefined :: Int64))
   print (typeOf (undefined :: (ST () ())))
   print (typeOf (undefined :: (StableName ())))
   print (typeOf (undefined :: (StablePtr ())))
   print (typeOf (undefined :: TyCon))
   print (typeOf (undefined :: TypeRep))
   print (typeOf (undefined :: Word8))
   print (typeOf (undefined :: Word16))
   print (typeOf (undefined :: Word32))
   print (typeOf (undefined :: Word64))

   print (typeOf (undefined :: ArithException))
   print (typeOf (undefined :: AsyncException))
   print (typeOf (undefined :: (IOArray () ())))
   print (typeOf (undefined :: (IOUArray () ())))
   print (typeOf (undefined :: (STArray () () ())))
   print (typeOf (undefined :: (STUArray () () ())))
   print (typeOf (undefined :: (StableName ())))
   print (typeOf (undefined :: (StablePtr ())))
   print (typeOf (undefined :: (UArray () ())))
   print (typeOf (undefined :: (Weak ())))

   print (typeOf (undefined :: CChar))
   print (typeOf (undefined :: CSChar))
   print (typeOf (undefined :: CUChar))
   print (typeOf (undefined :: CShort))
   print (typeOf (undefined :: CUShort))
   print (typeOf (undefined :: CInt))
   print (typeOf (undefined :: CUInt))
   print (typeOf (undefined :: CLong))
   print (typeOf (undefined :: CULong))
   print (typeOf (undefined :: CLLong))
   print (typeOf (undefined :: CULLong))
   print (typeOf (undefined :: CFloat))
   print (typeOf (undefined :: CDouble))

   print (typeOf (undefined :: CPtrdiff))
   print (typeOf (undefined :: CSize))
   print (typeOf (undefined :: CWchar))
   print (typeOf (undefined :: CSigAtomic))
   print (typeOf (undefined :: CClock))
   print (typeOf (undefined :: CTime))

max_stars_repo_name: JakobBruenker/ghc

max_issues_repo_name: JakobBruenker/ghc

max_forks_repo_name: JakobBruenker/ghc


text:
{-# LANGUAGE TypeOperators, GADTs, CPP #-}
#ifndef NO_SAFE_HASKELL
{-# LANGUAGE Safe #-}
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
{-# LANGUAGE PatternSynonyms, ViewPatterns #-}
#endif

#ifndef NO_GENERICS
{-# LANGUAGE DefaultSignatures, FlexibleContexts #-}
#endif

#ifndef NO_POLYKINDS
{-# LANGUAGE PolyKinds #-}
#endif

-- | Generation of random shrinkable, showable functions.
-- See the paper \"Shrinking and showing functions\" by Koen Claessen.
--
-- __Note__: most of the contents of this module are re-exported by
-- "Test.QuickCheck". You probably do not need to import it directly.
--
-- Example of use:
--
-- >>> :{
-- >>> let prop :: Fun String Integer -> Bool
-- >>>     prop (Fun _ f) = f "monkey" == f "banana" || f "banana" == f "elephant"
-- >>> :}
-- >>> quickCheck prop
-- *** Failed! Falsified (after 3 tests and 134 shrinks):
-- {"elephant"->1, "monkey"->1, _->0}
--
-- To generate random values of type @'Fun' a b@,
-- you must have an instance @'Function' a@.
-- If your type has a 'Show' instance, you can use 'functionShow' to write the instance; otherwise,
-- use 'functionMap' to give a bijection between your type and a type that is already an instance of 'Function'.
-- See the @'Function' [a]@ instance for an example of the latter.
module Test.QuickCheck.Function
  ( Fun(..)
  , applyFun
  , apply
  , applyFun2
  , applyFun3
  , (:->)
  , Function(..)
  , functionMap
  , functionShow
  , functionIntegral
  , functionRealFrac
  , functionBoundedEnum
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
  , pattern Fn
  , pattern Fn2
  , pattern Fn3
#endif
  )
 where

--------------------------------------------------------------------------
-- imports

import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Poly

import Control.Applicative
import Data.Char
import Data.Word
import Data.List( intersperse )
import Data.Ratio
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Sequence as Sequence
import Data.Int
import Data.Complex
import Data.Foldable(toList)
import Data.Functor.Identity
import qualified Data.Monoid as Monoid

#ifndef NO_FIXED
import Data.Fixed
#endif

#ifndef NO_GENERICS
import GHC.Generics hiding (C)
#endif

--------------------------------------------------------------------------
-- concrete functions

-- | The type of possibly partial concrete functions
data a :-> c where
  Pair  :: (a :-> (b :-> c)) -> ((a,b) :-> c)
  (:+:) :: (a :-> c) -> (b :-> c) -> (Either a b :-> c)
  Unit  :: c -> (() :-> c)
  Nil   :: a :-> c
  Table :: Eq a => [(a,c)] -> (a :-> c)
  Map   :: (a -> b) -> (b -> a) -> (b :-> c) -> (a :-> c)

instance Functor ((:->) a) where
  fmap f (Pair p)    = Pair (fmap (fmap f) p)
  fmap f (p:+:q)     = fmap f p :+: fmap f q
  fmap f (Unit c)    = Unit (f c)
  fmap f Nil         = Nil
  fmap f (Table xys) = Table [ (x,f y) | (x,y) <- xys ]
  fmap f (Map g h p) = Map g h (fmap f p)

instance (Show a, Show b) => Show (a:->b) where
  show p = showFunction p Nothing

-- only use this on finite functions
showFunction :: (Show a, Show b) => (a :-> b) -> Maybe b -> String
showFunction p md =
  "{" ++ concat (intersperse ", " ( [ show x ++ "->" ++ show c
                                    | (x,c) <- table p
                                    ]
                                 ++ [ "_->" ++ show d
                                    | Just d <- [md]
                                    ] )) ++ "}"

-- turning a concrete function into an abstract function (with a default result)
abstract :: (a :-> c) -> c -> (a -> c)
abstract (Pair p)    d (x,y) = abstract (fmap (\q -> abstract q d y) p) d x
abstract (p :+: q)   d exy   = either (abstract p d) (abstract q d) exy
abstract (Unit c)    _ _     = c
abstract Nil         d _     = d
abstract (Table xys) d x     = head ([y | (x',y) <- xys, x == x'] ++ [d])
abstract (Map g _ p) d x     = abstract p d (g x)

-- generating a table from a concrete function
table :: (a :-> c) -> [(a,c)]
table (Pair p)    = [ ((x,y),c) | (x,q) <- table p, (y,c) <- table q ]
table (p :+: q)   = [ (Left x, c) | (x,c) <- table p ]
                 ++ [ (Right y,c) | (y,c) <- table q ]
table (Unit c)    = [ ((), c) ]
table Nil         = []
table (Table xys) = xys
table (Map _ h p) = [ (h x, c) | (x,c) <- table p ]

--------------------------------------------------------------------------
-- Function

-- | The class @Function a@ is used for random generation of showable
-- functions of type @a -> b@.
--
-- There is a default implementation for 'function', which you can use
-- if your type has structural equality. Otherwise, you can normally
-- use 'functionMap' or 'functionShow'.
class Function a where
  function :: (a->b) -> (a:->b)
#ifndef NO_GENERICS
  default function :: (Generic a, GFunction (Rep a)) => (a->b) -> (a:->b)
  function = genericFunction
#endif

-- basic instances

-- | Provides a 'Function' instance for types with 'Bounded' and 'Enum'.
-- Use only for small types (i.e. not integers): creates
-- the list @['minBound'..'maxBound']@!
functionBoundedEnum :: (Eq a, Bounded a, Enum a) => (a->b) -> (a:->b)
functionBoundedEnum f = Table [(x,f x) | x <- [minBound..maxBound]]

-- | Provides a 'Function' instance for types with 'RealFrac'.
functionRealFrac :: RealFrac a => (a->b) -> (a:->b)
functionRealFrac = functionMap toRational fromRational

-- | Provides a 'Function' instance for types with 'Integral'.
functionIntegral :: Integral a => (a->b) -> (a:->b)
functionIntegral = functionMap fromIntegral fromInteger

-- | Provides a 'Function' instance for types with 'Show' and 'Read'.
functionShow :: (Show a, Read a) => (a->c) -> (a:->c)
functionShow f = functionMap show read f

-- | The basic building block for 'Function' instances.
-- Provides a 'Function' instance by mapping to and from a type that
-- already has a 'Function' instance.
functionMap :: Function b => (a->b) -> (b->a) -> (a->c) -> (a:->c)
functionMap = functionMapWith function

functionMapWith :: ((b->c) -> (b:->c)) -> (a->b) -> (b->a) -> (a->c) -> (a:->c)
functionMapWith function g h f = Map g h (function (\b -> f (h b)))

instance Function () where
  function f = Unit (f ())

instance Function a => Function (Const a b) where
  function = functionMap getConst Const

instance Function a => Function (Identity a) where
  function = functionMap runIdentity Identity

instance (Function a, Function b) => Function (a,b) where
  function = functionPairWith function function

functionPairWith :: ((a->b->c) -> (a:->(b->c))) -> ((b->c) -> (b:->c)) -> ((a,b)->c) -> ((a,b):->c)
functionPairWith func1 func2 f = Pair (func2 `fmap` func1 (curry f))

instance (Function a, Function b) => Function (Either a b) where
  function = functionEitherWith function function

functionEitherWith :: ((a->c) -> (a:->c)) -> ((b->c) -> (b:->c)) -> (Either a b->c) -> (Either a b:->c)
functionEitherWith func1 func2 f = func1 (f . Left) :+: func2 (f . Right)

-- tuple convenience instances

instance (Function a, Function b, Function c) => Function (a,b,c) where
  function = functionMap (\(a,b,c) -> (a,(b,c))) (\(a,(b,c)) -> (a,b,c))

instance (Function a, Function b, Function c, Function d) => Function (a,b,c,d) where
  function = functionMap (\(a,b,c,d) -> (a,(b,c,d))) (\(a,(b,c,d)) -> (a,b,c,d))

instance (Function a, Function b, Function c, Function d, Function e) => Function (a,b,c,d,e) where
  function = functionMap (\(a,b,c,d,e) -> (a,(b,c,d,e))) (\(a,(b,c,d,e)) -> (a,b,c,d,e))

instance (Function a, Function b, Function c, Function d, Function e, Function f) => Function (a,b,c,d,e,f) where
  function = functionMap (\(a,b,c,d,e,f) -> (a,(b,c,d,e,f))) (\(a,(b,c,d,e,f)) -> (a,b,c,d,e,f))

instance (Function a, Function b, Function c, Function d, Function e, Function f, Function g) => Function (a,b,c,d,e,f,g) where
  function = functionMap (\(a,b,c,d,e,f,g) -> (a,(b,c,d,e,f,g))) (\(a,(b,c,d,e,f,g)) -> (a,b,c,d,e,f,g))

-- other instances

instance Function a => Function [a] where
  function = functionMap g h
   where
    g []     = Left ()
    g (x:xs) = Right (x,xs)

    h (Left _)       = []
    h (Right (x,xs)) = x:xs

instance Function a => Function (Maybe a) where
  function = functionMap g h
   where
    g Nothing  = Left ()
    g (Just x) = Right x

    h (Left _)  = Nothing
    h (Right x) = Just x

instance Function Bool where
  function = functionMap g h
   where
    g False = Left ()
    g True  = Right ()

    h (Left _)  = False
    h (Right _) = True

instance Function Integer where
  function = functionMap gInteger hInteger
   where
    gInteger n | n < 0     = Left (gNatural (abs n - 1))
               | otherwise = Right (gNatural n)

    hInteger (Left ws)  = -(hNatural ws + 1)
    hInteger (Right ws) = hNatural ws

    gNatural 0 = []
    gNatural n = (fromIntegral (n `mod` 256) :: Word8) : gNatural (n `div` 256)

    hNatural []     = 0
    hNatural (w:ws) = fromIntegral w + 256 * hNatural ws

instance Function Int where
  function = functionIntegral

instance Function Word where
  function = functionIntegral

instance Function Char where
  function = functionMap ord chr

instance Function Float where
  function = functionRealFrac

instance Function Double where
  function = functionRealFrac

-- instances for assorted types in the base package

instance Function Ordering where
  function = functionMap g h
    where
      g LT = Left False
      g EQ = Left True
      g GT = Right ()

      h (Left False) = LT
      h (Left True)  = EQ
      h (Right _)    = GT

instance (Integral a, Function a) => Function (Ratio a) where
  function = functionMap g h
   where
     g r = (numerator r, denominator r)
     h (n, d) = n % d

#ifndef NO_FIXED
instance HasResolution a => Function (Fixed a) where
  function = functionRealFrac
#endif

instance (RealFloat a, Function a) => Function (Complex a) where
  function = functionMap g h
   where
     g (x :+ y) = (x,   y)
     h (x,   y) =  x :+ y

instance (Ord a, Function a) => Function (Set.Set a) where
  function = functionMap Set.toList Set.fromList

instance (Ord a, Function a, Function b) => Function (Map.Map a b) where
  function = functionMap Map.toList Map.fromList

instance Function IntSet.IntSet where
  function = functionMap IntSet.toList IntSet.fromList

instance Function a => Function (IntMap.IntMap a) where
  function = functionMap IntMap.toList IntMap.fromList

instance Function a => Function (Sequence.Seq a) where
  function = functionMap toList Sequence.fromList

instance Function Int8 where
  function = functionBoundedEnum

instance Function Int16 where
  function = functionIntegral

instance Function Int32 where
  function = functionIntegral

instance Function Int64 where
  function = functionIntegral

instance Function Word8 where
  function = functionBoundedEnum

instance Function Word16 where
  function = functionIntegral

instance Function Word32 where
  function = functionIntegral

instance Function Word64 where
  function = functionIntegral

-- instances for Data.Monoid newtypes

instance Function a => Function (Monoid.Dual a) where
  function = functionMap Monoid.getDual Monoid.Dual

instance Function Monoid.All where
  function = functionMap Monoid.getAll Monoid.All

instance Function Monoid.Any where
  function = functionMap Monoid.getAny Monoid.Any

instance Function a => Function (Monoid.Sum a) where
  function = functionMap Monoid.getSum Monoid.Sum

instance Function a => Function (Monoid.Product a) where
  function = functionMap Monoid.getProduct Monoid.Product

instance Function a => Function (Monoid.First a) where
  function = functionMap Monoid.getFirst Monoid.First

instance Function a => Function (Monoid.Last a) where
  function = functionMap Monoid.getLast Monoid.Last

#if MIN_VERSION_base(4,8,0)
instance Function (f a) => Function (Monoid.Alt f a) where
  function = functionMap Monoid.getAlt Monoid.Alt
#endif

-- poly instances

instance Function A where
  function = functionMap unA A

instance Function B where
  function = functionMap unB B

instance Function C where
  function = functionMap unC C

instance Function OrdA where
  function = functionMap unOrdA OrdA

instance Function OrdB where
  function = functionMap unOrdB OrdB

instance Function OrdC where
  function = functionMap unOrdC OrdC

-- instance Arbitrary

instance (Function a, CoArbitrary a, Arbitrary b) => Arbitrary (a:->b) where
  arbitrary = function `fmap` arbitrary
  shrink    = shrinkFun shrink

--------------------------------------------------------------------------
-- generic function instances

#ifndef NO_GENERICS
-- | Generic 'Function' implementation.
genericFunction :: (Generic a, GFunction (Rep a)) => (a->b) -> (a:->b)
genericFunction = functionMapWith gFunction from to

class GFunction f where
  gFunction :: (f a -> b) -> (f a :-> b)

instance GFunction U1 where
  gFunction = functionMap (\U1 -> ()) (\() -> U1)

instance (GFunction f, GFunction g) => GFunction (f :*: g) where
  gFunction = functionMapWith (functionPairWith gFunction gFunction) g h
   where
     g (x :*: y) = (x, y)
     h (x, y) = x :*: y

instance (GFunction f, GFunction g) => GFunction (f :+: g) where
  gFunction = functionMapWith (functionEitherWith gFunction gFunction) g h
   where
     g (L1 x) = Left x
     g (R1 x) = Right x
     h (Left x) = L1 x
     h (Right x) = R1 x

instance GFunction f => GFunction (M1 i c f) where
  gFunction = functionMapWith gFunction (\(M1 x) -> x) M1

instance Function a => GFunction (K1 i a) where
  gFunction = functionMap (\(K1 x) -> x) K1
#endif

--------------------------------------------------------------------------
-- shrinking

shrinkFun :: (c -> [c]) -> (a :-> c) -> [a :-> c]
shrinkFun shr (Pair p) =
  [ pair p' | p' <- shrinkFun (\q -> shrinkFun shr q) p ]
 where
  pair Nil = Nil
  pair p   = Pair p

shrinkFun shr (p :+: q) =
  [ p .+. Nil | not (isNil q) ] ++
  [ Nil .+. q | not (isNil p) ] ++
  [ p  .+. q' | q' <- shrinkFun shr q ] ++
  [ p' .+. q  | p' <- shrinkFun shr p ]
 where
  isNil :: (a :-> b) -> Bool
  isNil Nil = True
  isNil _   = False

  Nil .+. Nil = Nil
  p   .+. q   = p :+: q

shrinkFun shr (Unit c) =
  [ Nil ] ++
  [ Unit c' | c' <- shr c ]

shrinkFun shr (Table xys) =
  [ table xys' | xys' <- shrinkList shrXy xys ]
 where
  shrXy (x,y) = [(x,y') | y' <- shr y]

  table []  = Nil
  table xys = Table xys

shrinkFun shr Nil =
  []

shrinkFun shr (Map g h p) =
  [ mapp g h p' | p' <- shrinkFun shr p ]
 where
  mapp g h Nil = Nil
  mapp g h p   = Map g h p

--------------------------------------------------------------------------
-- the Fun modifier

-- | Generation of random shrinkable, showable functions.
--
-- To generate random values of type @'Fun' a b@,
-- you must have an instance @'Function' a@.
--
-- See also 'applyFun', and 'Fn' with GHC >= 7.8.
data Fun a b = Fun (a :-> b, b, Shrunk) (a -> b)
data Shrunk = Shrunk | NotShrunk deriving Eq

instance Functor (Fun a) where
  fmap f (Fun (p, d, s) g) = Fun (fmap f p, f d, s) (f . g)

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
-- | A modifier for testing functions.
--
-- > prop :: Fun String Integer -> Bool
-- > prop (Fn f) = f "banana" == f "monkey"
-- >            || f "banana" == f "elephant"
#if __GLASGOW_HASKELL__ >= 800
pattern Fn :: (a -> b) -> Fun a b
#endif
pattern Fn f <- (applyFun -> f)

-- | A modifier for testing binary functions.
--
-- > prop_zipWith :: Fun (Int, Bool) Char -> [Int] -> [Bool] -> Bool
-- > prop_zipWith (Fn2 f) xs ys = zipWith f xs ys == [ f x y | (x, y) <- zip xs ys]
#if __GLASGOW_HASKELL__ >= 800
pattern Fn2 :: (a -> b -> c) -> Fun (a, b) c
#endif
pattern Fn2 f <- (applyFun2 -> f)

-- | A modifier for testing ternary functions.
#if __GLASGOW_HASKELL__ >= 800
pattern Fn3 :: (a -> b -> c -> d) -> Fun (a, b, c) d
#endif
pattern Fn3 f <- (applyFun3 -> f)
#endif

mkFun :: (a :-> b) -> b -> Fun a b
mkFun p d = Fun (p, d, NotShrunk) (abstract p d)

-- | Alias to 'applyFun'.
apply :: Fun a b -> (a -> b)
apply = applyFun

-- | Extracts the value of a function.
--
-- 'Fn' is the pattern equivalent of this function.
--
-- > prop :: Fun String Integer -> Bool
-- > prop f = applyFun f "banana" == applyFun f "monkey"
-- >       || applyFun f "banana" == applyFun f "elephant"
applyFun :: Fun a b -> (a -> b)
applyFun (Fun _ f) = f

-- | Extracts the value of a binary function.
--
-- 'Fn2' is the pattern equivalent of this function.
--
--  > prop_zipWith :: Fun (Int, Bool) Char -> [Int] -> [Bool] -> Bool
--  > prop_zipWith f xs ys = zipWith (applyFun2 f) xs ys == [ applyFun2 f x y | (x, y) <- zip xs ys]
--
applyFun2 :: Fun (a, b) c -> (a -> b -> c)
applyFun2 (Fun _ f) a b = f (a, b)

-- | Extracts the value of a ternary function. 'Fn3' is the
-- pattern equivalent of this function.
applyFun3 :: Fun (a, b, c) d -> (a -> b -> c -> d)
applyFun3 (Fun _ f) a b c = f (a, b, c)

instance (Show a, Show b) => Show (Fun a b) where
  show (Fun (_, _, NotShrunk) _) = "<fun>"
  show (Fun (p, d, Shrunk) _)    = showFunction p (Just d)

instance (Function a, CoArbitrary a, Arbitrary b) => Arbitrary (Fun a b) where
  arbitrary =
    do p <- arbitrary
       d <- arbitrary
       return (mkFun p d)

  shrink (Fun (p, d, s) f) =
    [ mkFun p' d' | (p', d') <- shrink (p, d) ] ++
    [ Fun (p, d, Shrunk) f | s == NotShrunk ]

--------------------------------------------------------------------------
-- the end.

max_stars_repo_name: amigalemming/quickcheck

max_issues_repo_name: amigalemming/quickcheck

max_forks_repo_name: amigalemming/quickcheck


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds   #-}

module Synapse
  ( sigmoid
  , sigmoidm
  , sigmoidBackward
  , relu
  , relum
  , relumBackward
  , softmax
  , softmaxm
  , softmaxWithCross
  , softmaxWithCrossBackward
  , crossEntropy
  , crossEntropym
  , affinem
  , affinemBackward
  ) where

import           Debug.Trace
import           Numeric.LinearAlgebra
import           Prelude               hiding ((<>))

sigmoid :: (Floating a) => a -> a
sigmoid a = 1 / (1 + exp (-a))

sigmoidm :: (Floating a, Container Matrix a) => Matrix a -> Matrix a
sigmoidm = cmap sigmoid

sigmoidBackward ::
     (Floating a, Num (Vector a), Container Matrix a)
  => Matrix a
  -> Matrix a
  -> Matrix a
sigmoidBackward y d = d * (1 - y) * y

relu :: (Ord a, Num a) => a -> a
relu = max 0

relum :: (Ord a, Num a, Container Matrix a) => Matrix a -> Matrix a
relum = cmap relu

relumBackward ::
     (Ord a, Num a, Num (Matrix a), Container Vector a, Container Matrix a)
  => Matrix a
  -> Matrix a
  -> Matrix a
relumBackward x d = d * mask
  where
    mask = cmap (fromIntegral . onoff) x
    onoff a
      | 0 < a = 1
      | otherwise = 0

softmax :: (Floating a, Container Vector a) => Vector a -> Vector a
softmax v = cmap (/ s) v'
  where
    m = maxElement v
    v' = cmap (\a -> exp (a - m)) v
    s = sumElements v'

softmaxm :: (Floating a, Container Vector a) => Matrix a -> Matrix a
softmaxm m = fromRows $ map softmax $ toRows m

crossEntropy ::
     (Floating a, Num (Vector a), Container Vector a)
  => Vector a
  -> Vector a
  -> a
crossEntropy t y = -(sumElements $ cmap log (y + d) * t)
  where
    d = 1e-10

crossEntropym ::
     (Floating a, Num (Vector a), Container Vector a)
  => Matrix a
  -> Matrix a
  -> a
crossEntropym t m = sum vs / batchSize
  where
    vs = zipWith crossEntropy (toRows t) (toRows m)
    batchSize = fromIntegral $ rows t

softmaxWithCross ::
     (Floating a, Num (Vector a), Container Vector a)
  => Matrix a
  -> Matrix a
  -> (Matrix a, a)
softmaxWithCross t x = (y, crossEntropym t y)
  where
    y = softmaxm x

softmaxWithCrossBackward ::
     (Floating a, Num (Vector a), Container Vector a)
  => Matrix a
  -> Matrix a
  -> Matrix a
softmaxWithCrossBackward t y = cmap (/ batchSize) $ y - t
  where
    batchSize = fromIntegral $ rows t

affinem ::
     (Floating a, Numeric a, Num (Vector a))
  => Matrix a
  -> Vector a
  -> Matrix a
  -> Matrix a
affinem w b x = (x <> w) + b'
  where
    b' = fromRows $ replicate (rows x) b

affinemBackward ::
     (Floating a, Numeric a)
  => Matrix a
  -> Matrix a
  -> Matrix a
  -> (Matrix a, Matrix a, Vector a)
affinemBackward w x d = (dx, dw, db)
  where
    dx = d <> tr w
    dw = tr x <> d
    db = fromList $ sumElements `map` toColumns d


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE DeriveGeneric #-}

module Language.Grappa.Interp.PVIE where

import Data.List
import GHC.Generics hiding (R)
import Data.Functor.Const
import Data.Functor.Compose
-- import Data.Functor.Product
import Control.Lens
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Primitive
import Data.IORef
import qualified Numeric.Log as Log

import qualified Data.ByteString.Lazy as BS
import Data.Aeson
import System.IO
import System.IO.Unsafe
import System.Directory (doesFileExist)

import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra.Data as SV (Vector)
import qualified Data.Vector.Generic as SV hiding (Vector)
import qualified Data.Vector.Generic as SVGen
import qualified Data.Vector.Storable.Mutable as SMV

import qualified Numeric.AD.Mode.Reverse as ADR
import qualified Numeric.AD.Internal.Reverse as ADR
import qualified Data.Reflection as ADR (Reifies)

import Numeric.GSL.Minimization

import Language.Grappa.Distribution
import Language.Grappa.GrappaInternals

import Language.Grappa.Rand.MWCRandM
import qualified System.Random.MWC as MWC
import qualified System.Random.MWC.Distributions as MWC

import qualified Text.PrettyPrint.ANSI.Leijen as PP
import Text.PrettyPrint.ANSI.Leijen ((<+>), (<>))

import Options.Applicative

import Debug.Trace

-- | Read JSON file @filename@ and parse it to an @a@, where @filename@ can be
-- @"-"@ to denote 'stdin'
readJSONfile :: FromJSON a => String -> IO a
readJSONfile f =
  withFileContents f $ \contents ->
  case eitherDecode' contents of
    Left err -> error ("Could not parse JSON file " ++ f ++ ": " ++ err)
    Right a -> return a
  where
    withFileContents :: String -> (BS.ByteString -> IO a) -> IO a
    withFileContents "-" k = BS.hGetContents stdin >>= k
    withFileContents filename k =
      withFile filename ReadMode $ \h -> BS.hGetContents h >>= k

-- | Write JSON file @filename@, which can be @"-"@ to indicate 'stdout'
writeJSONfile :: ToJSON a => String -> a -> IO ()
writeJSONfile "-" a = BS.putStr $ encode a
writeJSONfile filename a = encodeFile filename a

----------------------------------------------------------------------
-- * Dimensionality Expressions
----------------------------------------------------------------------

-- | A "dimensionality variable"
-- (FIXME: explain what this means: an index into a 'VIDimAsgn'?)
newtype VIDimVar = VIDimVar { viDimVarIx :: Int } deriving Eq

-- | Build a list of the first @n@ dimensionality variables
viDimVars :: Int -> [VIDimVar]
viDimVars n = take n $ map VIDimVar $ [0 ..]

-- | The "first" dimensionality variable
zeroVar :: VIDimVar
zeroVar = VIDimVar 0

-- | Get the next dimensionality variable after the given one
nextVar :: VIDimVar -> VIDimVar
nextVar (VIDimVar i) = (VIDimVar (i+1))

maxVar :: VIDimVar -> VIDimVar -> VIDimVar
maxVar (VIDimVar v1) (VIDimVar v2) = VIDimVar $ max v1 v2

-- | An assignment to dimensionality variables, which should be non-negative
newtype VIDimAsgn = VIDimAsgn [Int] deriving Generic

instance FromJSON VIDimAsgn where

instance ToJSON VIDimAsgn where
  toEncoding = genericToEncoding defaultOptions

-- | Create an assignment where all variables below the given one are set to 0
zeroAsgn :: VIDimVar -> VIDimAsgn
zeroAsgn (VIDimVar i) = VIDimAsgn $ take i $ repeat 1

-- | Increment the assignment to a particular dimensionality variable
incrAsgn :: VIDimVar -> VIDimAsgn -> VIDimAsgn
incrAsgn (VIDimVar i) (VIDimAsgn asgn) = VIDimAsgn (over (ix i) (+ 1) asgn)

-- | Increment the assignment to a particular dimensionality variable
incrAsgnAll :: VIDimAsgn -> VIDimAsgn
incrAsgnAll (VIDimAsgn asgn) = VIDimAsgn (map (+1) asgn)

-- | Look up the value of a dimensionality variable in an assignment
viDimVarVal :: VIDimVar -> VIDimAsgn -> Int
viDimVarVal (VIDimVar i) (VIDimAsgn ns) = ns !! i

-- | An expression for the dimensionality of variational inference parameters in
-- terms of some dimensionality variables, given as a function on the values of
-- those variables along with an indicator of how many variables are
-- needed. Note that this function should always return a non-negative value.
data VIDim = VIDim {
  evalVIDim :: VIDimAsgn -> Int,
  viDimFirstUnusedVar :: VIDimVar
  }

-- | Build a constant dimensionality
constVIDim :: Int -> VIDim
constVIDim k = VIDim (const k) zeroVar

-- | Build a variable dimensionality
varVIDim :: VIDimVar -> VIDim
varVIDim v = VIDim (viDimVarVal v) (nextVar v)

-- | Apply a binary operation to a 'VIDim', ensuring the resulting dimension
-- evaluation is non-negative
viDimBinOp :: (Int -> Int -> Int) -> VIDim -> VIDim -> VIDim
viDimBinOp binop (VIDim eval1 v1) (VIDim eval2 v2) =
  VIDim (\asgn -> max (eval1 asgn `binop` eval2 asgn) 0) (maxVar v1 v2)

instance Num VIDim where
  (+) = viDimBinOp (+)
  (-) = viDimBinOp (-)
  (*) = viDimBinOp (*)
  abs dim = dim -- dimensions are always non-negative
  signum _ = 1 -- satisfies abs dim * signum dim = dim (up to equivalence)
  fromInteger = constVIDim . fromInteger

----------------------------------------------------------------------
-- * Differentiable Functions
----------------------------------------------------------------------

type Params = SV.Vector Double
type MutParams = SMV.MVector (PrimState IO) Double
type ParamsGrad = SMV.MVector (PrimState IO) Double

-- | Create a normal 'Vector' for a (small number of) 'Params'. This is needed
-- in order to use the reverse-mode AD 'ADR.grad' function.
vectorOfParams :: Int -> Params -> Vector Double
vectorOfParams n params = V.generate n (params SV.!)

-- | Helper class to write "differentiable" functions on our parameters
class FromDouble a where
  fromDouble :: Double -> a

instance FromDouble Double where
  fromDouble = id

instance ADR.Reifies s ADR.Tape => FromDouble (ADR.Reverse s Double) where
  fromDouble = ADR.auto

-- | All the typeclasses we want for types used in differentiable functions
type DiffConstr a = (RealFloat a, Ord a, Show a, FromDouble a,
                     HasGamma a, Log.Precise a)

-- | The type of functions that are differentiable using AD
type DiffFun = forall r. DiffConstr r => Vector r -> r

-- | Take a differentiable function on @n@ (the first argument) variables and
-- turn it into a function that computes its gradient, scaled by the given
-- factor, and accumulates the result into the supplied mutable vector
scaledGrad :: Int -> Double -> DiffFun -> Params -> ParamsGrad -> IO ()
scaledGrad n scale f params params_grad =
  let raw_grad = ADR.grad f (vectorOfParams n params) in
  forM_ [0 .. n-1] $ \i ->
  SMV.modify params_grad ((scale * raw_grad V.! i) +) i

-- | Take a differentiable function on @n@ (the first argument) variables and
-- turn it into a function that computes its gradient, accumulating the result
-- into the supplied mutable vector, and also returns the result of the function
gradAndRet :: Int -> DiffFun -> Params -> ParamsGrad -> IO Double
gradAndRet n f params params_grad =
  do scaledGrad n 1 f params params_grad
     return (f $ vectorOfParams n params)

----------------------------------------------------------------------
-- * Computations Over Parameters
----------------------------------------------------------------------

newtype ParamsT m a = ParamsT (ReaderT (IORef Params) m a)
                    deriving (Functor, Applicative, Monad, MonadTrans)

runParamsT :: MonadIO m => ParamsT m a -> Params -> m a
runParamsT (ParamsT m) params =
  liftIO (newIORef params) >>= runReaderT m

instance MonadIO m => MonadIO (ParamsT m) where
  liftIO m = ParamsT $ liftIO m

class Monad m => MonadParams m where
  getParams :: Int -> m Params

instance MonadParams m => MonadParams (ReaderT r m) where
  getParams n = lift $ getParams n

instance MonadIO m => MonadParams (ParamsT m) where
  getParams n =
    ParamsT $ do r <- ask
                 ps <- liftIO $ readIORef r
                 liftIO $ writeIORef r (SV.drop n ps)
                 return $ SV.take n ps

newtype MutParamsT m a = MutParamsT (ReaderT (IORef MutParams) m a)
                       deriving (Functor, Applicative, Monad, MonadTrans)

runMutParamsT :: MonadIO m => MutParamsT m a -> MutParams -> m a
runMutParamsT (MutParamsT m) mut_params =
  liftIO (newIORef mut_params) >>= runReaderT m

instance MonadIO m => MonadIO (MutParamsT m) where
  liftIO m = MutParamsT $ liftIO m

instance MonadParams m => MonadParams (MutParamsT m) where
  getParams n = MutParamsT $ getParams n

class Monad m => MonadMutParams m where
  getMutParams :: Int -> m MutParams

instance MonadIO m => MonadMutParams (MutParamsT m) where
  getMutParams n =
    MutParamsT $ do r <- ask
                    ps <- liftIO $ readIORef r
                    liftIO $ writeIORef r (SMV.drop n ps)
                    return (SMV.take n ps)

instance MonadMutParams m => MonadMutParams (ReaderT r m) where
  getMutParams n = lift $ getMutParams n

-- | A computation for generating samples of distribution families
type SamplingM = ReaderT VIDimAsgn (ParamsT MWCRandM)

-- | Run a 'SamplingM' computation
runSamplingM :: SamplingM a -> MWC.GenIO -> VIDimAsgn -> Params -> IO a
runSamplingM m gen asgn params =
  runRand gen (runParamsT (runReaderT m asgn) params)

embedSamplingFun :: VIDim -> (Params -> MWCRandM a) -> SamplingM a
embedSamplingFun dim f =
  do asgn <- ask
     params <- getParams (evalVIDim dim asgn)
     lift $ lift $ f params

-- | A computation that consumes a vector of parameters and a vector of mutable
-- parameters, the latter being used for accumulating gradients
type ParamsGradM = ReaderT VIDimAsgn (MutParamsT (ParamsT IO))

runParamsGradM :: ParamsGradM a -> VIDimAsgn -> Params -> MutParams -> IO a
runParamsGradM m asgn params mut_params =
  runParamsT (runMutParamsT (runReaderT m asgn) mut_params) params

-- | Take the scaled gradient of a differentiable function with @n@ parameters
-- and accumulate it into the gradient of the current computation
embedScaledGrad :: VIDim -> DiffFun -> Double -> ParamsGradM ()
embedScaledGrad dim f scale =
  do n <- evalVIDim dim <$> ask
     params <- getParams n
     params_grad <- getMutParams n
     liftIO $ scaledGrad n scale f params params_grad

-- | Evaluate a differentiable function with @n@ parameters and its gradient,
-- returning the value and accumulating the gradient into the gradient of the
-- current computation
embedGradAndRet :: VIDim -> DiffFun -> ParamsGradM Double
embedGradAndRet dim f =
  do n <- evalVIDim dim <$> ask
     params <- getParams n
     params_grad <- getMutParams n
     liftIO $ gradAndRet n f params params_grad

-- | A computation for initializing parameters
type InitM = ReaderT VIDimAsgn (MutParamsT MWCRandM)

runInitM :: InitM a -> MWC.GenIO -> VIDimAsgn -> MutParams -> IO a
runInitM m gen asgn mut_params =
  runRand gen $ runMutParamsT (runReaderT m asgn) mut_params

simpleInitFun :: (VIDimAsgn -> MWCRandM [Double]) -> InitM ()
simpleInitFun initFun =
  do asgn <- ask
     params <- lift $ lift $ initFun asgn
     let params_v = SV.fromList params
     mut_params <- getMutParams (length params)
     liftIO $ SV.copy mut_params params_v

-- | A computation for growing sequences of parameters
type GrowM = ReaderT (VIDimAsgn, VIDimAsgn) (MutParamsT (ParamsT MWCRandM))

runGrowM :: GrowM a -> MWC.GenIO -> VIDimAsgn -> VIDimAsgn ->
            Params -> MutParams -> IO a
runGrowM m gen asgn new_asgn params new_params =
  runRand gen $
  runParamsT (runMutParamsT (runReaderT m (asgn, new_asgn)) new_params) params

simpleGrowFun :: VIDim -> (VIDimAsgn -> MWCRandM [Double]) -> GrowM ()
simpleGrowFun dim initFun =
  do (old_asgn, new_asgn) <- ask
     params <- getParams (evalVIDim dim old_asgn)
     old_mut_params <- getMutParams (evalVIDim dim old_asgn)
     liftIO $ SV.copy old_mut_params params
     new_mut_params <-
       getMutParams (evalVIDim dim new_asgn - evalVIDim dim old_asgn)
     inits_l <- lift $ lift $ lift $ initFun new_asgn
     let inits_v = SV.fromList $ drop (SV.length params) inits_l
     liftIO $ SV.copy new_mut_params inits_v

type PPFun = ReaderT VIDimAsgn (ParamsT IO) PP.Doc

simplePPFun :: VIDim -> String -> PPFun
simplePPFun dim nm =
  do n <- evalVIDim dim <$> ask
     params <- getParams n
     return (PP.text nm <+> PP.tupled (map PP.pretty $ SV.toList params))

applyPPFun :: PPFun -> VIDimAsgn -> Params -> IO PP.Doc
applyPPFun f asgn params = runParamsT (runReaderT f asgn) params

----------------------------------------------------------------------
-- * Distribution Approximators for Variational Inference
----------------------------------------------------------------------

type EntropyFun = ParamsGradM Double
type ScaledGradPDFFun a = Double -> a -> ParamsGradM ()
type InitFun = InitM ()
type GrowFun = GrowM ()

{-
-- | The trivial entropy function for the only distribution over ()
unitEntropyFun :: EntropyFun
unitEntropyFun _ _ _ = return 0

-- | Add two independent entropy functions, where the first has the given
-- dimensionality expression
addEntropyFuns :: VIDim -> EntropyFun -> EntropyFun -> EntropyFun
addEntropyFuns dim1 f1 f2 =
  \asgn ps ps_grad ->
  (+) <$>
  f1 asgn (SV.take (evalVIDim dim1 asgn) ps)
  (SMV.take (evalVIDim dim1 asgn) ps_grad)
  <*>
  f2 asgn (SV.drop (evalVIDim dim1 asgn) ps)
  (SMV.drop (evalVIDim dim1 asgn) ps_grad)

-- | The trivial log PDF function with 0 gradient
unitGradPDFFun :: GradPDFFun a
unitGradPDFFun _ _ _ _ _ = return ()

-- | Add two log PDF functions, assuming they are independent, given a
-- volume-preserving map from the codomain type to the two domain types
-- (represented as a pair of functions instead of a function to a pair) along
-- with a dimensionality expression for the number of params of the first
addGradPDFFuns :: (a -> b) -> (a -> c) -> VIDim ->
                  GradPDFFun b -> GradPDFFun c -> GradPDFFun a
addGradPDFFuns split_l split_r dim1 f1 f2 =
  \scale a asgn ps ps_grad ->
  f1 scale (split_l a) asgn (SV.take (evalVIDim dim1 asgn) ps)
  (SMV.take (evalVIDim dim1 asgn) ps_grad)
  >>
  f2 scale (split_r a) asgn (SV.drop (evalVIDim dim1 asgn) ps)
  (SMV.drop (evalVIDim dim1 asgn) ps_grad)

-- | The trivial 'viDistGrowParams' function for no parameters
unitGrowFun :: GrowFun
unitGrowFun _ _ _ _ = return ()

-- | Combine two 'viDistGrowParams' fields, given the dimensionality expression
-- for the first one
combineGrowFuns :: VIDim -> GrowFun -> GrowFun -> GrowFun
combineGrowFuns dim1 f1 f2 =
  \old_asgn new_asgn params mut_params ->
  f1 old_asgn new_asgn (SV.take (evalVIDim dim1 old_asgn) params)
  (SMV.take (evalVIDim dim1 new_asgn) mut_params)
  >>
  f2 old_asgn new_asgn (SV.drop (evalVIDim dim1 old_asgn) params)
  (SMV.drop (evalVIDim dim1 new_asgn) mut_params)
-}

-- | A family of distributions for variational inference
data VIDistFam a =
  VIDistFam
  {
    -- | The dimensionality of the parameters
    viDistDim :: VIDim,

    -- | Sampling function to randomly generate an @a@ given 'viDistDim'
    -- 'Double'-valued parameters
    viDistSample :: SamplingM a,

    -- | Evaluate the entropy at the given parameters and also add the gradient
    -- of the entropy with respect to those parameters into the mutable vector
    viDistEntropy :: EntropyFun,

    -- | Evaluate the gradient with respect to the parameters of the log PDF of
    -- the distribution, scaled by the supplied factor, and add the result to
    -- the supplied mutable vector
    viDistScaledGradPDF :: ScaledGradPDFFun a,

    -- | Randomly initialize the parameters of this distribution family
    viDistInit :: InitFun,

    -- | Grow a vector of parameters from the dimensionality implied by one
    -- assignment to that of another, larger one, storing the result in the
    -- given mutable vector (assuming it has the correct size) and initializing
    -- any new parameters
    viDistGrowParams :: GrowFun,

    -- | Pretty-print a distribution given its parameters
    viDistPP :: PPFun
  }

-- | Build a distribution family for a "simple" distribution, meaning it is not
-- a composite of multiple distributions on sub-components.  Such a distribution
-- is defined by a dimensionality expression, a sampling function,
-- differentiable functions for the entropy and log PDF, and an initialization
-- function for the parameters.
simpleVIFam :: String -> VIDim -> (Params -> MWCRandM a) -> DiffFun ->
               (a -> DiffFun) -> (VIDimAsgn -> MWCRandM [Double]) ->
               VIDistFam a
simpleVIFam nm dim sampleFun entropyFun pdfFun initFun =
  VIDistFam
  { viDistDim = dim
  , viDistSample = embedSamplingFun dim sampleFun
  , viDistEntropy = embedGradAndRet dim entropyFun
  , viDistScaledGradPDF =
      (\scale a -> embedScaledGrad dim (pdfFun a) scale)
  , viDistInit = simpleInitFun initFun
  , viDistGrowParams = simpleGrowFun dim initFun
  , viDistPP = simplePPFun dim nm
  }

-- | Build a distribution family for tuples from a tuple of families
tupleVIFam :: TupleF ts (Compose VIDistFam f) (ADT (TupleF ts)) ->
              VIDistFam (TupleF ts f (ADT (TupleF ts)))
tupleVIFam ds =
  VIDistFam
  { viDistDim = foldrADT (viDistDim . getCompose) (+) 0 ds
  , viDistSample = traverseADT (viDistSample . getCompose) ds
  , viDistEntropy =
      foldrADT (viDistEntropy . getCompose)
      (\m1 m2 -> (+) <$> m1 <*> m2) (return 0) ds
  , viDistScaledGradPDF =
      (\scale tup ->
        foldrADT getConst (>>) (return ()) $
        mapTuple2 (\(Compose d) a ->
                    Const $ viDistScaledGradPDF d scale a) ds tup)
  , viDistInit = foldrADT (viDistInit . getCompose) (>>) (return ()) ds
  , viDistGrowParams =
    foldrADT (viDistGrowParams . getCompose) (>>) (return ()) ds
  , viDistPP =
    PP.tupled <$> sequence (ctorArgsADT (viDistPP . getCompose) ds)
  }

-- | Transform a distribution family over one type to another type using a
-- volume-preserving isomorphism. A volume-preserving map is one that does not
-- squeeze any objects closer together or move them farther apart; the technical
-- definition is that it preserves the measures of all sets with respect to the
-- reference measure on the two spaces.
xformVIDistFam :: (a -> b) -> (b -> a) -> VIDistFam a -> VIDistFam b
xformVIDistFam f_to f_from (VIDistFam {..}) =
  VIDistFam
  {
    viDistSample = fmap f_to $ viDistSample,
    viDistScaledGradPDF = (\scale a -> viDistScaledGradPDF scale (f_from a)),
    ..
  }

-- | Map a vector of distribution families to a distribution family of vectors
vecDistVIFam :: (SVGen.Vector f a, SVGen.Vector f (VIDistFam a)) =>
                f (VIDistFam a) -> VIDistFam (f a)
vecDistVIFam ds =
  VIDistFam
  { viDistDim = SVGen.foldl' (\dim d -> dim + viDistDim d) 0 ds
  , viDistSample = SVGen.mapM viDistSample ds
  , viDistEntropy = SVGen.foldM' (\e d -> (e +) <$> viDistEntropy d) 0 ds
  , viDistScaledGradPDF =
    (\scale v ->
      if SVGen.length v == SVGen.length ds then
        SVGen.zipWithM_ (\d x -> viDistScaledGradPDF d scale x) ds v
      else
        error "vecDist distribution: wrong size vector!")
  , viDistInit = SVGen.mapM_ viDistInit ds
  , viDistGrowParams = SVGen.mapM_ viDistGrowParams ds
  , viDistPP =
    do pps <- mapM viDistPP $ SVGen.toList ds
       return (PP.text "vecDist" <+> PP.list pps)
  }

-- | The distribution family over vectors with a given length whose elements are
-- drawn IID from the supplied distribution family
vecIIDVIFam :: SVGen.Vector f a => VIDim -> VIDistFam a -> VIDistFam (f a)
vecIIDVIFam len d =
  VIDistFam
  { viDistDim = len * viDistDim d
  , viDistSample =
    do asgn <- ask
       SVGen.replicateM (evalVIDim len asgn) (viDistSample d)
  , viDistEntropy =
    do asgn <- ask
       entropies <- replicateM (evalVIDim len asgn) (viDistEntropy d)
       return $ sum entropies
  , viDistScaledGradPDF =
    (\scale v ->
      do asgn <- ask
         if SVGen.length v == evalVIDim len asgn then
           SVGen.forM_ v (viDistScaledGradPDF d scale)
           else
           error "IID distribution: wrong size vector!")
  , viDistInit =
    do asgn <- ask
       replicateM_ (evalVIDim len asgn) (viDistInit d)
  , viDistGrowParams =
    do (old_asgn, new_asgn) <- ask
       replicateM_ (evalVIDim len old_asgn) (viDistGrowParams d)
       g <- lift $ lift $ lift $ nthGenIO 0
       let num_new_elems = evalVIDim len new_asgn - evalVIDim len old_asgn
       replicateM_ num_new_elems $
         do new_params <- getMutParams (evalVIDim (viDistDim d) new_asgn)
            liftIO $ runInitM (viDistInit d) g new_asgn new_params
  , viDistPP =
    do asgn <- ask
       let n = evalVIDim len asgn
       pps <- replicateM n (viDistPP d)
       return (PP.text ("IID(" ++ show n ++ ")") <+> PP.list pps)
  }

----------------------------------------------------------------------
-- * Distribution Approximator Expressions
----------------------------------------------------------------------

-- | An "expression" for building up a family of distributions, which keeps
-- track of how many dimensionality variables we have used so far and also a
-- list of the data file names to be used by 'readJSONVIDistFamExpr'
newtype VIDistFamExpr a =
  VIDistFamExpr { runVIDistFamExpr ::
                    StateT (VIDimVar, [String]) IO (VIDistFam a) }

-- | Evaluate a distribution family expression into a distribution family
evalVIDistFamExpr :: [String] -> VIDistFamExpr a -> IO (VIDistFam a)
evalVIDistFamExpr files (VIDistFamExpr m) = evalStateT m (VIDimVar 0, files)

entropyNaNProtect :: String -> DiffFun -> DiffFun
entropyNaNProtect s f ps =
  let ret = f ps in
  if isNaN ret then
    error ("NaN entropy for distribution " ++ s ++ ", params = " ++ show ps)
  else
    ret

-- | Build a distribution family expression for a "simple" distribution, meaning
-- it is not a composite of multiple distributions on sub-components.  Such a
-- distribution is defined by a dimensionality expression, a sampling function,
-- and differentiable functions for the entropy and log PDF.
simpleVIFamExpr :: String -> VIDim -> (Params -> MWCRandM a) -> DiffFun ->
                   (a -> DiffFun) -> (VIDimAsgn -> MWCRandM [Double]) ->
                   VIDistFamExpr a
simpleVIFamExpr nm dim sampleFun entropyFun pdfFun initFun =
  VIDistFamExpr $ return $
  simpleVIFam nm dim sampleFun (entropyNaNProtect nm entropyFun) pdfFun initFun

-- | The constant distribution (also known as the delta distribution), that
-- returns a single value with unit probability
deltaVIFamExpr :: (Eq a, GrappaShow a) => a -> VIDistFamExpr a
deltaVIFamExpr a =
  simpleVIFamExpr ("Delta(" ++ grappaShow a ++ ")")
  0 (\_ -> return a) (\_ -> 0)
  (\x _ -> if x == a then 0 else log 0)
  (const $ return [])

-- | Build a distribution family expression for the normal distribution, where
-- we use absolute value of @sigma@ to map it to the non-negative reals
normalVIFamExpr :: VIDistFamExpr R
normalVIFamExpr =
  simpleVIFamExpr "Normal" 2 (\ps -> mwcNormal (ps SV.! 0) (abs (ps SV.! 1)))
  (\ps ->  0.5 * (1 + log (2 * pi) + log ((ps V.! 1) * (ps V.! 1))))
  (\x ps -> Log.ln $ normalDensityUnchecked (ps V.! 0) (abs (ps V.! 1)) (fromDouble x))
  (\_ ->
    -- Init mu from the standard normal and sigma from the inverse gamma
    do mu <- random MWC.standard
       inv_sigma_sq <- mwcGamma 1 1
       let sigma = 1 / sqrt inv_sigma_sq
       return [mu,sigma])

-- | Build a distribution family expression for the uniform distribution over
-- the range @(min a b, max a b]@
uniformVIFamExpr :: VIDistFamExpr R
uniformVIFamExpr =
  simpleVIFamExpr "Uniform" 2 (\ps -> mwcUniform (ps SV.! 0) (ps SV.! 1))
  (\ps -> log $ abs (ps V.! 1 - ps V.! 0))
  (\x ps ->
    if min (ps V.! 0) (ps V.! 1) < fromDouble x
       && fromDouble x <= max (ps V.! 0) (ps V.! 1) then
      log $ abs (ps V.! 1 - ps V.! 0)
    else log 0)
  (\_ ->
    -- Initialize to a sub-interval of the unit interval by setting lower =
    -- uniform(0,1) and upper = lower + uniform(0,1) * (1 - lower)
    do lower <- mwcUniform 0 1
       upper_nonnorm <- mwcUniform 0 1
       let upper = lower + upper_nonnorm * (1 - lower)
       return [lower, upper])

-- | Build a distribution family expression for the categorical distribution
-- over @[0,..,n-1]@ with relative probabilities @[a1/a,..,an/a]@, where each
-- @ai@ is the absolute value of the @i@th parameter and @a@ is the sum of the
-- @ai@s. The special case that @n=0@ is treated like @n=1@, meaning it always
-- returns @0@.
categoricalVIFamExpr :: VIDim -> VIDistFamExpr Int
categoricalVIFamExpr dim =
  simpleVIFamExpr "Categorical" dim
  (\ps ->
    -- NOTE: MWC.categorical normalizes the weights for us
    random $ MWC.categorical $ SV.map abs ps)
  (\ps ->
    let ps' = V.map abs ps
        p_sum = V.sum ps' in
    V.sum $ V.map (\p -> (p / p_sum) * log (p / p_sum)) ps')
  (\x ps ->
    let ps' = V.map abs ps
        p_sum = V.sum ps'
        n = V.length ps in
    if n == 0 && x == 0 then 1 else
      if x < 0 || x >= n then 0 else
        log ((ps' V.! x) / p_sum))
  (\asgn ->
    -- Initialize the probabilities from a Dirichlet distribution
    let n = evalVIDim dim asgn in
    if n == 0 then return [] else
      mwcDirichlet (replicate n 1))

-- | Build a distribution family expression for the gamma distribution, where
-- we use absolute value of @k@ and @theta@ to them to the non-negative reals
gammaVIFamExpr :: VIDistFamExpr R
gammaVIFamExpr =
  simpleVIFamExpr "Gamma" 2 (\ps -> mwcGamma (abs (ps SV.! 0)) (abs (ps SV.! 1)))
  (\ps ->
    let k = abs (ps SV.! 0)
        theta = abs (ps SV.! 1) in
    k + log theta + logGamma k + (1-k) * digamma k)
  (\x ps ->
    Log.ln $
    gammaDensityUnchecked (abs (ps V.! 0)) (abs (ps V.! 1)) (fromDouble x))
  (\_ ->
    -- Initialize k and theta from exponentials (because their actual conjugate
    -- priors are messy and I'm too busy for that)
    do k <- mwcExponential 1
       theta <- mwcExponential 1
       return [k,theta])

-- | Build a distribution family expression for the gamma distribution over
-- probabilities, i.e., in log space, where the @k@ and @theta@ parameters are
-- also in log space
gammaProbVIFamExpr :: VIDistFamExpr Prob
gammaProbVIFamExpr =
  simpleVIFamExpr "GammaProb" 2
  (\ps -> Prob <$> mwcGammaLogLog (Log.Exp (ps SV.! 0)) (Log.Exp (ps SV.! 1)))
  (\ps ->
    let k = exp (ps SV.! 0)
        log_theta = ps SV.! 1 in
    k + log_theta + logGamma k + (1-k) * digamma k)
  (\x ps ->
    Log.ln $
    gammaDensityUnchecked (exp (ps V.! 0)) (exp (ps V.! 1))
    (fromDouble $ probToLogR x))
  (\_ ->
    -- Initialize (exp k) and (exp theta) from the exponential distribution, in
    -- keeping with the initialization for gammaVIFamExpr, above. NOTE: this
    -- kind of sucks, because we are just taking log (log uniform(0,1)), and
    -- taking two logarithms in a row probably loses a lot of precision.
    do exp_k <- mwcExponential 1
       exp_theta <- mwcExponential 1
       return [log exp_k, log exp_theta])

-- | Build a distribution family expression for the beta distribution, where
-- we use absolute value of @alpha@ and @beta@ to keep them non-negative
betaVIFamExpr :: VIDistFamExpr R
betaVIFamExpr =
  simpleVIFamExpr "Beta" 2 (\ps -> mwcBeta (abs (ps SV.! 0)) (abs (ps SV.! 1)))
  (\ps ->
    let alpha = abs (ps SV.! 0)
        beta = abs (ps SV.! 1) in
    logGamma alpha + logGamma beta - logGamma (alpha + beta)
    + (alpha + beta - 2) * digamma (alpha + beta)
    - (alpha - 1) * digamma alpha - (beta - 1) * digamma beta)
  (\x ps ->
    Log.ln $
    betaDensityUnchecked (abs (ps V.! 0)) (abs (ps V.! 1)) (fromDouble x))
  (\_ ->
    -- Initialize alpha and beta from the exponential distribution, because why
    -- not? And because, as Wikipedia says (for the Dirichlet distribution,
    -- which is similar), "In the published literature there is no practical
    -- algorithm to efficiently generate samples from" the conjugate prior.
    do alpha <- mwcExponential 1
       beta <- mwcExponential 1
       return [alpha,beta])

-- | Build a distribution family expression for the beta distribution over
-- probabilities, i.e., in log space, where the @alpha@ and @beta@ parameters
-- are also in log space
betaProbVIFamExpr :: VIDistFamExpr Prob
betaProbVIFamExpr =
  simpleVIFamExpr "BetaProb" 2
  (\ps -> Prob <$> mwcBetaLogLog (Log.Exp (ps SV.! 0)) (Log.Exp (ps SV.! 1)))
  (\ps ->
    let alpha = exp (ps SV.! 0)
        beta = exp (ps SV.! 1) in
    logGamma alpha + logGamma beta - logGamma (alpha + beta)
    + (alpha + beta - 2) * digamma (alpha + beta)
    - (alpha - 1) * digamma alpha - (beta - 1) * digamma beta)
  (\x ps ->
    Log.ln $
    betaDensityLog (exp (ps V.! 0)) (exp (ps V.! 1))
    (fmap fromDouble $ fromProb x))
  (\_ ->
    -- Initialize (exp alpha) and (exp beta) from the exponential distribution,
    -- in keeping with the initialization for betaVIFamExpr, above. NOTE: this
    -- kind of sucks, because we are just taking log (log uniform(0,1)), and
    -- taking two logarithms in a row probably loses a lot of precision.
    do exp_alpha <- mwcExponential 1
       exp_beta <- mwcExponential 1
       return [log exp_alpha, log exp_beta])

-- | Build a distribution family for the dirichlet distribution, where the
-- alphas are in log space, so they can never be negative and so that the
-- optimization can do better at picking very small alphas
dirichletVIFamExpr :: VIDim -> VIDistFamExpr [R]
dirichletVIFamExpr dim =
  simpleVIFamExpr "Dirichlet" dim
  (\ps -> mwcDirichlet (SV.toList $ SV.map exp ps))
  (\ps ->
    let alphas = SV.toList $ SV.map exp ps
        k = fromIntegral $ SV.length ps
        alpha0 = sum alphas in
    sum (map logGamma alphas) - logGamma alpha0
    + (alpha0 - k) * digamma alpha0
    - sum (flip map alphas $ \alpha -> (alpha - 1) * digamma alpha))
  (\x ps -> Log.ln $
            dirichletDensity (SV.toList $ SV.map exp ps) (map fromDouble x))
  (\asgn ->
    -- Initialize alphas from the exponential distribution, because why not? And
    -- because, as Wikipedia says, "In the published literature there is no
    -- practical algorithm to efficiently generate samples from" the conjugate
    -- prior.
    replicateM (evalVIDim dim asgn) (mwcExponential 1))

-- | Build a distribution family for the dirichlet distribution over
-- probabilities, i.e., over lists of reals in log space. The alphas are in log
-- space, so they can never be negative and so that the optimization can do
-- better at picking very small alphas
dirichletProbVIFamExpr :: VIDim -> VIDistFamExpr [Prob]
dirichletProbVIFamExpr dim =
  simpleVIFamExpr "DirichletProb" dim
  (\ps -> map Prob <$> mwcDirichletLog (SV.toList $ SV.map exp ps))
  (\ps ->
    let alphas = SV.toList $ SV.map exp ps
        k = fromIntegral $ SV.length ps
        alpha0 = sum alphas in
    sum (map logGamma alphas) - logGamma alpha0
    + (alpha0 - k) * digamma alpha0
    - sum (flip map alphas $ \alpha -> (alpha - 1) * digamma alpha))
  (\x ps -> Log.ln $
            dirichletDensityLog (SV.toList $ SV.map exp ps)
            (map (fmap fromDouble . fromProb) x))
  (\asgn ->
    -- Initialize (exp alphas) from the exponential distribution, because why
    -- not? And because, as Wikipedia says, "In the published literature there
    -- is no practical algorithm to efficiently generate samples from" the
    -- conjugate prior. NOTE: this kind of sucks, because we are just taking log
    -- (log uniform(0,1)), and taking two logarithms in a row probably loses a
    -- lot of precision.
    map log <$> replicateM (evalVIDim dim asgn) (mwcExponential 1))

-- | Bind a fresh dimensionality variable in a distribution family expression
bindVIDimFamExpr :: (VIDim -> VIDistFamExpr a) -> VIDistFamExpr a
bindVIDimFamExpr f =
  VIDistFamExpr $ do
  (v,files) <- get
  put $ (nextVar v, files)
  runVIDistFamExpr $ f (varVIDim v)

-- | Build a distribution family expression for a tuple from a tuple of
-- distribution family expressions
tupleVIFamExpr :: TupleF ts (Compose VIDistFamExpr f) (ADT (TupleF ts)) ->
                  VIDistFamExpr (TupleF ts f (ADT (TupleF ts)))
tupleVIFamExpr d_exprs =
  VIDistFamExpr $ fmap tupleVIFam $
  traverseADT ((Compose <$>) . runVIDistFamExpr . getCompose) d_exprs

-- | Transform a distribution family expression over one type to another type
-- using a volume-preserving isomorphism; see 'xformVIDistFam'.
xformVIDistFamExpr :: (a -> b) -> (b -> a) -> VIDistFamExpr a -> VIDistFamExpr b
xformVIDistFamExpr f_to f_from expr =
  VIDistFamExpr (xformVIDistFam f_to f_from <$> runVIDistFamExpr expr)

-- | The distribution family expression over vectors with a given length whose
-- elements are drawn IID from the supplied distribution family expression
vecIIDVIFamExpr :: VIDim -> VIDistFamExpr a -> VIDistFamExpr (Vector a)
vecIIDVIFamExpr len d_expr =
  VIDistFamExpr (vecIIDVIFam len <$> runVIDistFamExpr d_expr)

-- | Map a vector of distribution family expressions to a distribution family
-- expression over vectors
vecDistVIFamExpr :: (SVGen.Vector f a, SVGen.Vector f (VIDistFam a),
                     SVGen.Vector f (VIDistFamExpr a)) =>
                    f (VIDistFamExpr a) -> VIDistFamExpr (f a)
vecDistVIFamExpr d_exprs =
  VIDistFamExpr (vecDistVIFam <$> SVGen.mapM runVIDistFamExpr d_exprs)

-- | The distribution family expression over lists with a given length whose
-- elements are drawn IID from the supplied distribution family expression
iidVIFamExpr :: VIDim -> VIDistFamExpr a -> VIDistFamExpr [a]
iidVIFamExpr len d_expr =
  xformVIDistFamExpr V.toList V.fromList $ vecIIDVIFamExpr len d_expr

-- | The distribution family expression over unboxed 'ProbVector's with a given
-- length whose elements are drawn IID from the supplied distribution family
-- expression
iidPVVIFamExpr :: VIDim -> VIDistFamExpr Prob -> VIDistFamExpr ProbVector
iidPVVIFamExpr len d_expr =
  xformVIDistFamExpr ProbVector unProbVector $
  VIDistFamExpr (vecIIDVIFam len <$> runVIDistFamExpr
                 (xformVIDistFamExpr probToLogR logRToProb d_expr))

-- | The distribution family expression over unboxed 'RVector's with a given
-- length whose elements are drawn IID from the supplied distribution family
-- expression
iidVVIFamExpr :: VIDim -> VIDistFamExpr R -> VIDistFamExpr RVector
iidVVIFamExpr len d_expr =
  xformVIDistFamExpr RVector unRVector $
  VIDistFamExpr (vecIIDVIFam len <$> runVIDistFamExpr d_expr)

-- | This distribution family reads an input of type @a@ from @stdin@ as a JSON
-- file and then passes that input to the supplied function to build a
-- distribution family of type @b@
mapJSONVIDistFamExpr :: FromJSON a => (a -> VIDistFamExpr b) -> VIDistFamExpr b
mapJSONVIDistFamExpr f =
  VIDistFamExpr $
  do (v,files) <- get
     let (file, files') =
           if length files == 0 then
             error "mapJSONVIDistFamExpr: not enough data files specified"
           else (head files, tail files)
     put (v, files')
     a <- liftIO $ readJSONfile file
     runVIDistFamExpr $ f a

-- | This distribution family is a delta distribution (i.e., one that always
-- returns the same value, like 'deltaVIFamExpr') that reads its input from
-- stdin as a JSON file
readJSONVIDistFamExpr :: (Eq a, GrappaShow a, FromJSON a) => VIDistFamExpr a
readJSONVIDistFamExpr = mapJSONVIDistFamExpr deltaVIFamExpr

----------------------------------------------------------------------
-- * Variational Inference, Yay!
----------------------------------------------------------------------

-- | A PVIE model is a dimensionaltiy variable assignment + model parameters
data PVIEModel =
  PVIEModel { pvieModelAsgn :: VIDimAsgn,
              pvieModelParams :: Params }
  deriving Generic

instance FromJSON PVIEModel where

instance ToJSON PVIEModel where
  toEncoding = genericToEncoding defaultOptions

readPVIEModel :: MWC.GenIO -> VIDistFam a -> String -> IO PVIEModel
readPVIEModel g d filename =
  doesFileExist filename >>= \ex ->
  if ex then
    -- If filename exists, then read it
    readJSONfile filename
  else
    -- Otherwise, start with 0 for all dimensionality variables, and
    -- initialize all params to 1
    do let asgn = zeroAsgn $ viDimFirstUnusedVar $ viDistDim d
       mut_params <- SMV.new (evalVIDim (viDistDim d) asgn)
       runInitM (viDistInit d) g asgn mut_params
       params <- SV.freeze mut_params
       return $ PVIEModel asgn params

-- | The PVIE modes: training and evaluation
data PVIEMode = PVIETrainMode | PVIEEvalMode

-- | The PVIE command-line options
data PVIEOpts =
  PVIEOpts {
  pvieModelFile :: String,
  pvieDataFiles :: [String],
  pvieMode :: PVIEMode,
  pvieVerbosity :: Int,
  pvieDebugTruncate :: Maybe Int,
  pvieNumSamples :: Int,
  pvieGD :: Bool
  }

-- | Parser for PVIE command-line options
pvieOptsParser :: Parser PVIEOpts
pvieOptsParser =
  PVIEOpts
  <$> (strOption (long "model"
                  <> short 'm'
                  <> metavar "MODEL"
                  <> help "Model file"
                  <> value "-"
                  <> showDefault))
  <*> many (strOption (long "data"
                       <> short 'd'
                       <> metavar "DATAFILE"
                       <> help "Data file"))
  <*> (flag' PVIETrainMode (long "train" <> short 't'
                            <> help "Specifies training mode")
       <|>
       flag' PVIEEvalMode (long "eval" <> short 'e'
                            <> help "Specifies evaluation mode"))
  <*> (option auto (long "verbosity" <> short 'v'
                    <> help "Verbosity level for debugging"
                    <> showDefault <> value 0 <> metavar "VERBOSITY"))
  <*> (option (Just <$> auto)
       (long "truncate" <> short 'T'
        <> help "Truncate length for debugging strings"
        <> showDefault <> value Nothing <> metavar "TRUNCATE"))
  <*> (option auto (long "numsamples" <> short 'N'
                    <> help "Number of samples for approximating the variational family"
                    <> showDefault <> value 1000 <> metavar "NUMSAMPLES"))
  <*> (flag False True
       (long "gd" <>
        help "Use gradient descent instead of BFGS to optimize (only meaningful in training mode)"))

-- | Parse the command-line options
parsePVIEOpts :: IO PVIEOpts
parsePVIEOpts = execParser (info (pvieOptsParser <**> helper)
                            (fullDesc <>
                             progDesc "FIXME: description of PVIE"))

debugTruncate :: PVIEOpts -> String -> String
debugTruncate opts s
  | Just l <- pvieDebugTruncate opts
  , length s >= l = take l s ++ " ..."
debugTruncate _ s = s

-- | Print debugging info if the verbosity level is @>= level@
debugM :: PVIEOpts -> Int -> String -> IO ()
debugM opts level s
  | pvieVerbosity opts >= level = traceM (debugTruncate opts s)
debugM _ _ _ = return ()

-- | Find the parameters that /minimize/ the given differentiable function
optimize :: PVIEOpts -> (Params -> MutParams -> IO Double) ->
            MutParams -> IO Double
optimize opts f mut_params =
  do let len = SMV.length mut_params
     let meth = if pvieGD opts then SteepestDescent else VectorBFGS2
     params <- SV.freeze mut_params
     let eval_f ps =
           unsafePerformIO $
           do debugM opts 2 ("neg_elbo: params = " ++ show ps)
              grad <- SMV.replicate len 0
              ret <- f ps grad
              debugM opts 2 ("surprisal = " ++ show ret)
              return ret
     let grad_f ps =
           unsafePerformIO $
           do debugM opts 2 ("neg_elbo_grad: params = " ++ show ps)
              grad <- SMV.replicate len 0
              _ <- f ps grad
              ret <- SV.unsafeFreeze grad
              debugM opts 2 ("grad = " ++ show ret)
              return ret
     let (opt_params,_) =
           minimizeVD meth 0.0001 1000 100 0.1 eval_f grad_f params
     SV.copy mut_params opt_params
     return $ eval_f opt_params

-- FIXME HERE: make pvie_epsilon into command-line options

-- | The amount that PVIE has to improve in an interation to seem "better"
pvie_epsilon :: Double
pvie_epsilon = 1.0e-6

-- | The negative infinity value
negInfinity :: Double
negInfinity = log 0

-- | Test if a 'Double' is (positive or negative) infinity or NaN
isInfiniteOrNaN :: Double -> Bool
isInfiniteOrNaN x = isInfinite x || isNaN x

-- | Compute the negative Evidence Lower BOund (or ELBO) and its gradient
neg_elbo_with_grad :: GrappaShow a => PVIEOpts -> MWC.GenIO ->
                      VIDistFam a -> (a -> Double) -> VIDimAsgn ->
                      (Params -> MutParams -> IO Double)
neg_elbo_with_grad opts g d log_p asgn params grad =
  do ret <- elbo_with_grad opts g d log_p asgn params grad
     forM_ [0 .. (SMV.length grad)-1] (SMV.modify grad negate)
     return (-ret)

-- | Compute the Evidence Lower BOund (or ELBO) and its gradient
elbo_with_grad :: GrappaShow a => PVIEOpts -> MWC.GenIO ->
                  VIDistFam a -> (a -> Double) -> VIDimAsgn ->
                  (Params -> MutParams -> IO Double)
elbo_with_grad opts g d log_p asgn params grad =
  (replicateM (pvieNumSamples opts) $
   do s <- runSamplingM (viDistSample d) g asgn params
      return (s, log_p s)) >>= \samples_log_ps ->
  case find (isInfiniteOrNaN . snd) samples_log_ps of
    Just (bad_samp, p) ->
      -- If any of our samples have 0 probability in our model (i.e., from
      -- log_p), the elbo is -infinity, and the gradient is undefined. However,
      -- knowing our optimization algorithm, we cheat, and tell it that the
      -- gradient is just the sum of the negatives of the gradients of d at
      -- those 0 probability samples. That way, our optimization algorithm will
      -- keep trying to reduce the probabilities of generating bad samples until
      -- it finally succeeds. Note that we still return -infinity as the value.
      do debugM opts 3 ((if isNaN p then "NaN" else
                           if p < 0 then "Zero-probability"
                           else "Infinite (?) probability") ++
                        " sample: " ++ grappaShow bad_samp)
         forM_ samples_log_ps $ \(samp, samp_p) ->
           when (isInfiniteOrNaN samp_p) $
           runParamsGradM (viDistScaledGradPDF d (-1e9) samp) asgn params grad
         return negInfinity
    _ ->
      do let n = fromIntegral (pvieNumSamples opts)
         entr <- runParamsGradM (viDistEntropy d) asgn params grad
         debugM opts 2 ("Entropy: " ++ show entr)
         entr_grad <- SV.freeze grad
         debugM opts 3 ("Entropy grad: " ++ show entr_grad)
         forM_ samples_log_ps $ \(samp, p) ->
           runParamsGradM (viDistScaledGradPDF d (p / n) samp) asgn params grad
         -- grad_const <- SV.unsafeFreeze grad
         -- traceM ("grad = " ++ show grad_const)
         let ret = entr + (1/n) * sum (map snd samples_log_ps)
         -- traceM ("surprisal = " ++ show ret)
         debugM opts 4 ("Sample 0: " ++ grappaShow (fst $ head samples_log_ps))
         return ret

-- | The main entrypoint for the PVIE engine
pvie_main :: GrappaShow a => VIDistFamExpr a -> (a -> Double) -> IO ()
pvie_main dist_expr log_p =
  do opts <- parsePVIEOpts
     dist_fam <- evalVIDistFamExpr (pvieDataFiles opts) dist_expr
     case pvieMode opts of
       PVIETrainMode ->
         do (model@(PVIEModel asgn params), val) <-
              pvie_train opts dist_fam log_p
            pp <- applyPPFun (viDistPP dist_fam) asgn params
            debugM opts 1 ("Final model:\n" ++ show pp)
            debugM opts 1 ("Surprisal score: " ++ show val)
            writeJSONfile (pvieModelFile opts) model
       PVIEEvalMode ->
         do val <- pvie_eval opts dist_fam log_p
            putStrLn ("Surprisal score: " ++ show val)

-- | The main entrypoint for the PVIE engine with anomaly score
pvie_anom_main :: GrappaShow a => VIDistFamExpr a -> (a -> Double) ->
                  (a -> Prob) -> IO ()
pvie_anom_main dist_expr log_p anom_score =
  do opts <- parsePVIEOpts
     dist_fam <- evalVIDistFamExpr (pvieDataFiles opts) dist_expr
     case pvieMode opts of
       PVIETrainMode ->
         do (model@(PVIEModel asgn params), val) <-
              pvie_train opts dist_fam log_p
            pp <- applyPPFun (viDistPP dist_fam) asgn params
            debugM opts 1 ("Final model:\n" ++ show pp)
            debugM opts 1 ("Surprisal score: " ++ show val)
            writeJSONfile (pvieModelFile opts) model
       PVIEEvalMode ->
         do score <- pvie_eval_anom opts dist_fam anom_score
            putStrLn $ show (1 - min 1 (probToR score))

-- | The PVIE eval mode
pvie_eval :: GrappaShow a => PVIEOpts -> VIDistFam a -> (a -> Double) ->
             IO Double
pvie_eval opts d log_p =
  do PVIEModel asgn params <- readJSONfile $ pvieModelFile opts
     g <- MWC.createSystemRandom
     grad <- SMV.new (evalVIDim (viDistDim d) asgn)
     val <- neg_elbo_with_grad opts g d log_p asgn params grad
     return val

-- | Evaluate the anomaly score, which is the average over @N@ samples of the
-- supplied anomaly score function applied to each sample
pvie_eval_anom :: GrappaShow a => PVIEOpts -> VIDistFam a -> (a -> Prob) ->
                  IO Prob
pvie_eval_anom opts d anom_score =
  do PVIEModel asgn params <- readJSONfile $ pvieModelFile opts
     g <- MWC.createSystemRandom
     anom_scores <-
       replicateM (pvieNumSamples opts) $
       do samp <- runSamplingM (viDistSample d) g asgn params
          return (probToLogR $ anom_score samp)
     return $ logRToProb (sum anom_scores / fromIntegral (pvieNumSamples opts))

-- | The PVIE training mode
pvie_train :: GrappaShow a => PVIEOpts -> VIDistFam a -> (a -> Double) ->
              IO (PVIEModel, Double)
pvie_train opts d log_p = init_pvie where

  -- Initialize PVIE and start it running
  init_pvie :: IO (PVIEModel, Double)
  init_pvie = do
    g <- MWC.createSystemRandom
    PVIEModel asgn params <- readPVIEModel g d (pvieModelFile opts)
    mut_params <- SV.unsafeThaw params
    -- Generate the initial value to try to beat
    val <- optimize opts (neg_elbo_with_grad opts g d log_p asgn) mut_params
    params' <- SV.unsafeFreeze mut_params
    -- If there are no dimensionality variables in our distribution family, then
    -- there is nothing to increment, so we are done
    if viDimFirstUnusedVar (viDistDim d) == zeroVar then
      outer g asgn params' val False
      else
      -- Perform the first iteration of optimizing the dimensionality variables
      outer g asgn params' val True

  -- The main outer loop
  outer :: MWC.GenIO -> VIDimAsgn -> Params -> Double -> Bool ->
           IO (PVIEModel, Double)
  outer _ asgn params last_val False =
    -- If our last iteration of the outer loop did not improve, we are done, and
    -- return the current assignment, parameters, and value
    return (PVIEModel asgn params, last_val)

  outer g asgn params last_val True =
    do
      -- Otherwise, start by allocating two new arrays, one for the current
      -- parameters and one for the scratch space. The size should be big enough
      -- for if we increment all dimensionality variables by one, which is the
      -- most we will do in the upcoming iteration of the outer loop.
      new_mut_params <- SMV.new (evalVIDim (viDistDim d) $ incrAsgnAll asgn)
      scratch <- SMV.new (evalVIDim (viDistDim d) $ incrAsgnAll asgn)
      -- Copy the old parameters into the new mutable params
      let new_mut_params_slice = SMV.slice 0 (SV.length params) new_mut_params
      SV.copy new_mut_params_slice params
      new_params <- SV.unsafeFreeze new_mut_params
      -- Start the inner loop at the first dimensionality variable, telling it
      -- that we have not yet improved
      inner g asgn new_params last_val zeroVar scratch False

  -- The main inner loop: takes in the current assignment to dimensionality
  -- variables, the current "best" params for the associated size, the value for
  -- those params, the next dimensionality variable to try to increment,
  -- pre-allocated scratch space for parameters, and whether we have already
  -- seen improvement in this iteration of the outer loop
  inner :: MWC.GenIO -> VIDimAsgn -> Params -> Double -> VIDimVar ->
           MutParams -> Bool -> IO (PVIEModel, Double)
  inner g asgn last_params last_val next_var _ improved
    | next_var == viDimFirstUnusedVar (viDistDim d)
    = outer g asgn last_params last_val improved

  inner g asgn last_params last_val next_var scratch improved = do
    -- The next assignment we will try = increment next_var
    let new_asgn = incrAsgn next_var asgn
    -- Copy our current best parameters into our scratch area
    runGrowM (viDistGrowParams d) g asgn new_asgn last_params scratch
    -- Optimize those parameters
    new_val <- optimize opts (neg_elbo_with_grad opts g d log_p asgn) scratch
    -- Test how much we improved
    if last_val - new_val >= pvie_epsilon then
      -- If we did improve, swap last_params and scratch, then iterate
      do new_params <- SV.unsafeFreeze scratch
         new_scratch <- SV.unsafeThaw last_params
         inner g new_asgn new_params new_val (nextVar next_var) new_scratch True
      else
      -- Otherwise, keep the same assignment, params, and scratch
      inner g asgn last_params last_val (nextVar next_var) scratch improved

max_stars_repo_name: GaloisInc/grappa

max_issues_repo_name: GaloisInc/grappa

max_forks_repo_name: GaloisInc/grappa


text:
{-# LANGUAGE FlexibleContexts #-}
module Image.Transform where

import           Data.Array.Repa              as R
import           Data.Array.Repa.Stencil      as R
import           Data.Array.Repa.Stencil.Dim2 as R
import           Data.DList                   as DL
import           Data.List                    as L
import           Data.Vector.Unboxed          as VU
import           Numeric.LinearAlgebra        as NL

-- factor = 2^n, n = 0,1,..
-- the first factor in the list corresponds to the inner-most (right-most) dimension.
{-# INLINE downsample #-}
downsample ::
     (Source s e, Shape sh) => [Int] -> R.Array s sh e -> R.Array D sh e
downsample factorList arr
  | L.all (== 1) factorList = delay arr
  | L.any (< 1) newDList = error "Downsample factors are too large."
  | otherwise =
    R.backpermute
      (shapeOfList newDList)
      (shapeOfList . L.zipWith (*) factorList . listOfShape)
      arr
  where
    dList = listOfShape . extent $ arr
    newDList = L.zipWith div dList factorList

{-# INLINE downsampleUnsafe #-}
downsampleUnsafe ::
     (Source s e, Shape sh) => [Int] -> R.Array s sh e -> R.Array D sh e
downsampleUnsafe factorList arr =
  R.backpermute newSh (shapeOfList . L.zipWith (*) factorList . listOfShape) arr
  where
    dList = listOfShape $ extent arr
    newSh = shapeOfList $ L.zipWith div dList factorList

{-# INLINE upsample #-}
upsample ::
     (Source s e, Shape sh, Num e) => [Int] -> R.Array s sh e -> R.Array D sh e
upsample factorList arr
  | L.all (== 1) factorList = delay arr
  | L.any (< 1) factorList =
    error $ "upsample: factor must be >= 1.\n" L.++ show factorList
  | otherwise =
    R.backpermuteDft
      (fromFunction
         (shapeOfList $ L.zipWith (*) (listOfShape . extent $ arr) factorList)
         (const 0))
      (\shape ->
         if L.all (== 0) . L.zipWith (flip mod) factorList . listOfShape $ shape
           then Just .
                shapeOfList . L.zipWith (flip div) factorList . listOfShape $
                shape
           else Nothing)
      arr

{-# INLINE crop #-}
crop ::
     (Source s e, Shape sh)
  => [Int]
  -> [Int]
  -> R.Array s sh e
  -> R.Array D sh e
crop start len arr
  | L.any (< 0) start ||
      L.or (L.zipWith3 (\x y z -> x > (z - y)) start len dList) =
    error $
    "Crop out of boundary!\n" L.++ show start L.++ "\n" L.++ show len L.++ "\n" L.++
    show dList
  | L.length start /= L.length len || L.length start /= L.length dList =
    error $
    "crop: dimension error. \n start: " L.++ show (L.length start) L.++ " len:" L.++
    show (L.length len) L.++
    " arr:" L.++
    show (L.length dList)
  | otherwise =
    R.backpermute
      (shapeOfList len)
      (shapeOfList . L.zipWith (+) start . listOfShape)
      arr
  where
    dList = listOfShape $ extent arr

{-# INLINE cropUnsafe #-}
cropUnsafe ::
     (Source s e, Shape sh)
  => [Int]
  -> [Int]
  -> R.Array s sh e
  -> R.Array D sh e
cropUnsafe start len =
  R.backpermute
    (shapeOfList len)
    (shapeOfList . L.zipWith (+) start . listOfShape)

{-# INLINE pad #-}
pad :: (Source s e, Shape sh) => [Int] -> e -> R.Array s sh e -> R.Array D sh e
pad newDims padVal arr
  | L.all (== 0) diff = delay arr
  | otherwise =
    backpermuteDft
      (fromFunction (shapeOfList dimList) (const padVal))
      (\sh' ->
         let idx = L.zipWith (-) (listOfShape sh') diff
          in if L.or (L.zipWith (\i j -> i < 0 || (i >= j)) idx oldDimList)
               then Nothing
               else Just $ shapeOfList idx)
      arr
  where
    oldDimList = listOfShape . extent $ arr
    dimList = L.zipWith max newDims oldDimList
    diff =
      L.zipWith
        (\a b ->
           if a - b <= 0
             then 0
             else if a - b == 1
                    then 1
                    else div (a - b) 2)
        newDims
        oldDimList

{-# INLINE normalizeValueRange #-}
normalizeValueRange ::
     (Source r e, Shape sh, Num e, Fractional e, Ord e, Unbox e)
  => (e, e)
  -> R.Array r sh e
  -> R.Array D sh e
normalizeValueRange (minVal, maxVal) arr
  | maxV == minV = delay arr
  | otherwise =
    R.map (\x -> (x - minV) / (maxV - minV) * (maxVal - minVal) + minVal) $ arr
  where
    vec = computeUnboxedS . delay $ arr
    minV = VU.minimum . toUnboxed $ vec
    maxV = VU.maximum . toUnboxed $ vec

{-# INLINE computeDerivativeS #-}
computeDerivativeS ::
     (R.Source r e, Num e, Fractional e) => R.Array r DIM2 e -> R.Array D DIM3 e
computeDerivativeS arr =
  L.foldl' R.append (R.extend (Z :. R.All :. R.All :. (1 :: Int)) arr) ds
  where
    xStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. 0 :. (-1) -> Just (-1)
          Z :. 0 :. 1 -> Just 1
          _ -> Nothing
    yStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. -1 :. 0 -> Just (-1)
          Z :. 1 :. 0 -> Just 1
          _ -> Nothing
    xyStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. -1 :. -1 -> Just 1
          Z :. -1 :. 1 -> Just (-1)
          Z :. 1 :. -1 -> Just (-1)
          Z :. 1 :. 1 -> Just 1
          _ -> Nothing
    ds =
      L.map
        (\s ->
           extend (Z :. R.All :. R.All :. (1 :: Int)) . R.map (/ 2) $
           mapStencil2 BoundClamp s arr)
        [xStencil, yStencil, xyStencil]

{-# INLINE bicubicInterpolation #-}
bicubicInterpolation ::
     (R.Source r Double)
  => ((Int, Int) -> (Double, Double))
  -> (Int, Int)
  -> R.Array r DIM2 Double
  -> R.Array U DIM2 Double
bicubicInterpolation idxFunc (newCols, newRows) arr =
  let (Z :. cols :. rows) = extent arr
      derivative = computeDerivativeS arr
      derivative16 =
        R.traverse
          derivative
          (const (Z :. newCols :. newRows :. 4 :. (4 :: Int)))
          (\f (Z :. a :. b :. derivativeIdx :. pairIdx) ->
             let (x, y) = idxFunc (a, b)
                 (x', y') =
                   case pairIdx of
                     0 -> (floor x, floor y)
                     1 -> (floor x, ceiling y)
                     2 -> (ceiling x, floor y)
                     3 -> (ceiling x, ceiling y)
              in if (x' < 0) ||
                    (x' > (fromIntegral cols - 1)) ||
                    (y' < 0) || (y' > (fromIntegral rows - 1))
                   then 0
                   else f (Z :. x' :. y' :. derivativeIdx))
      mat = ((newRows * newCols) >< 16) . R.toList $ derivative16
      multiplicationResultArray =
        fromListUnboxed (Z :. newCols :. newRows :. 4 :. 4) .
        NL.toList . flatten $
        mat NL.<> maxtrixA
   in R.sumS . R.sumS $
      R.traverse
        multiplicationResultArray
        id
        (\f idx@(Z :. a :. b :. i :. j) ->
           let (x, y) = idxFunc (a, b)
               x' = x - fromIntegral (floor x :: Int)
               y' = y - fromIntegral (floor y :: Int)
            in f idx * (x' ^ i) * (y' ^ j))
  where
    maxtrixA =
      NL.fromColumns . L.map NL.fromList $
      [ [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [-3, 3, 0, 0, -2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [2, -2, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, -3, 3, 0, 0, -2, -1, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 2, -2, 0, 0, 1, 1, 0, 0]
      , [-3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, -3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0]
      , [9, -9, -9, 9, 6, 3, -6, -3, 6, -6, 3, -3, 4, 2, 2, 1]
      , [-6, 6, 6, -6, -3, -3, 3, 3, -4, 4, -2, 2, -2, -2, -1, -1]
      , [2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0]
      , [-6, 6, 6, -6, -4, -2, 4, 2, -3, 3, -3, 3, -2, -1, -2, -1]
      , [4, -4, -4, 4, 2, 2, -2, -2, 2, -2, 2, -2, 1, 1, 1, 1]
      ]

{-# INLINE bilinearInterpolation #-}
bilinearInterpolation ::
     (R.Source r Double)
  => ((Int, Int) -> (Double, Double))
  -> (Int, Int)
  -> R.Array r DIM2 Double
  -> R.Array U DIM2 Double
bilinearInterpolation idxFunc (newCols, newRows) arr =
  let (Z :. cols :. rows) = extent arr
  in computeS . R.traverse arr id $ \f (Z :. a :. b) ->
       let (x, y) = idxFunc (a, b)
           x1 = floor x
           x2 = ceiling x
           y1 = floor y
           y2 = ceiling y
           q11 = f (Z :. x1 :. y1)
           q12 = f (Z :. x1 :. y2)
           q21 = f (Z :. x2 :. y1)
           q22 = f (Z :. x2 :. y2)
           r1 =
             ((fromIntegral x2 - x) / fromIntegral (x2 - x1)) * q11 +
             ((x - fromIntegral x1) / fromIntegral (x2 - x1)) * q21
           r2 =
             ((fromIntegral x2 - x) / fromIntegral (x2 - x1)) * q12 +
             ((x - fromIntegral x1) / fromIntegral (x2 - x1)) * q22
           p =
             ((fromIntegral y2 - y) / fromIntegral (y2 - y1)) * r1 +
             ((y - fromIntegral y1) / fromIntegral (y2 - y1)) * r2
       in if (x < 0) ||
             (x > (fromIntegral cols - 1)) ||
             (y < 0) || (y > (fromIntegral rows - 1))
            then 0
            else if x1 == x2
                   then if y1 == y2
                          then q11
                          else ((fromIntegral y2 - y) / fromIntegral (y2 - y1)) *
                               q11 +
                               ((y - fromIntegral y1) / fromIntegral (y2 - y1)) *
                               q12
                   else if y1 == y2
                          then r1
                          else p

{-# INLINE resize2D #-}
resize2D ::
     (Source s Double)
  => (Int, Int)
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
resize2D newSize@(newNx, newNy) bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (ratioX * fromIntegral i, ratioY * fromIntegral j))
    newSize
    arr
  where
    (Z :. nx' :. ny') = extent arr
    ratioX = fromIntegral nx' / fromIntegral newNx
    ratioY = fromIntegral ny' / fromIntegral newNy

{-# INLINE resize25D #-}
resize25D ::
     (Source s Double)
  => (Int, Int)
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
resize25D newSize@(newNx, newNy) bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (ratioX * fromIntegral i, ratioY * fromIntegral j))
         newSize .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    ratioX = fromIntegral nx' / fromIntegral newNx
    ratioY = fromIntegral ny' / fromIntegral newNy

{-# INLINE resize2DFixedRatio #-}
resize2DFixedRatio ::
     (Source s Double)
  => Int
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
resize2DFixedRatio n bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (ratio * fromIntegral i, ratio * fromIntegral j))
    newSize
    arr
  where
    (Z :. nx' :. ny') = extent arr
    newSize =
      if nx' > ny'
        then (n, round $ fromIntegral ny' * ratio)
        else (round $ fromIntegral nx' * ratio, n)
    ratio = fromIntegral (max nx' ny') / fromIntegral n

{-# INLINE resize25DFixedRatio #-}
resize25DFixedRatio ::
     (Source s Double)
  => Int
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
resize25DFixedRatio n bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (ratio * fromIntegral i, ratio * fromIntegral j))
         newSize .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    newSize@(newNx, newNy) =
      if nx' > ny'
        then (n, round $ fromIntegral ny' * ratio)
        else (round $ fromIntegral nx' * ratio, n)
    ratio = fromIntegral (max nx' ny') / fromIntegral n

{-# INLINE rescale2D #-}
rescale2D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
rescale2D scale bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (fromIntegral i / scale, fromIntegral j / scale))
    (round $ scale * fromIntegral nx', round $ scale * fromIntegral ny')
    arr
  where
    (Z :. nx' :. ny') = extent arr

{-# INLINE rescale25D #-}
rescale25D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
rescale25D scale bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (fromIntegral i / scale, fromIntegral j / scale))
         (newNx, newNy) .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    newNx = round $ scale * fromIntegral nx'
    newNy = round $ scale * fromIntegral ny'

{-# INLINE rotate2D #-}
rotate2D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array U DIM2 Double
rotate2D theta (centerX, centerY) arr =
  bicubicInterpolation
    (\(i, j) ->
       let i' = fromIntegral i - centerX
           j' = fromIntegral j - centerY
        in ( i' * cos theta - j' * sin theta + centerX
           , i' * sin theta + j' * cos theta + centerY))
    (nx', ny')
    arr
  where
    (Z :. nx' :. ny') = extent arr

{-# INLINE rotate25D #-}
rotate25D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array U DIM3 Double
rotate25D theta (centerX, centerY) arr =
  fromUnboxed (Z :. nf' :. nx' :. ny') .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       -- bicubicInterpolation
       bilinearInterpolation
         (\(i, j) ->
            let i' = fromIntegral i - centerX
                j' = fromIntegral j - centerY
            in ( i' * cos theta - j' * sin theta + centerX
               , i' * sin theta + j' * cos theta + centerY))
         (nx', ny') .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr


text:
{-# LANGUAGE FlexibleContexts #-}

{-| FM demodulation pipes -}
module SDR.Demod (
    fmDemodStr,
    fmDemodVec,
    fmDemod
    ) where

import Data.Complex
import qualified Data.Vector.Generic               as VG
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM
import qualified Data.Vector.Fusion.Bundle.Monadic as VFBM
import qualified Data.Vector.Fusion.Bundle.Size    as VFBS

import SDR.VectorUtils
import Pipes

-- | FM demodulate a stream of complex samples
{-# INLINE fmDemodStr #-}
fmDemodStr :: (RealFloat a, Monad m)
           => Complex a          -- ^ The starting sample - i.e. the last sample in the last buffer
           -> VFSM.Stream m (Complex a) -- ^ The input stream
           -> VFSM.Stream m a           -- ^ The output stream
fmDemodStr = mapAccumMV func
    where
    {-# INLINE func #-}
    func last sample = return (sample, phase (sample * conjugate last))

-- | FM demodulate a vector of complex samples
{-# INLINE fmDemodVec #-}
fmDemodVec :: (RealFloat a, VG.Vector v (Complex a), VG.Vector v a)
           => Complex a     -- ^ The starting sample - i.e. the last sample in the last buffer
           -> v (Complex a) -- ^ The input Vector
           -> v a           -- ^ The output Vector
fmDemodVec init inp = VG.unstream $ flip VFBM.fromStream (VFBS.Exact $ VG.length inp) $ fmDemodStr init $ VFBM.elements $ VG.stream inp

-- | Pipe that performs FM demodulation
{-# INLINE fmDemod #-}
fmDemod :: (RealFloat a, VG.Vector v (Complex a), VG.Vector v a) => Pipe (v (Complex a)) (v a) IO ()
fmDemod = func 0
    where
    func lastSample = do
        dat <- await
        yield $ fmDemodVec lastSample dat
        func $ VG.unsafeIndex dat (VG.length dat - 1)

max_stars_repo_name: adamwalker/sdr

max_issues_repo_name: peixian/sdr

max_forks_repo_name: adamwalker/sdr