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

module Graphical where

import Control.Monad
import Graphics.X11.Xlib
import System.Exit (exitWith, ExitCode(..))
import Control.Concurrent
import Data.Bits
import Data.Word
import Data.Scientific as Scientific
import Data.Complex
import GHC.Conc (numCapabilities)
import qualified Control.Concurrent.MVar.Strict as SM
import GHC.Generics (Generic)
import Control.DeepSeq
import Mand

data Ret a = Ret [Word64] Int
  deriving (Eq, Generic, NFData)

data Ven a = Ven a Int
  deriving (Eq, Generic, NFData)

startGraphical :: MVar Bool -> IO ()
startGraphical mv = do
  print "graphical started"
  dpy <- openDisplay ""
  let dflt = defaultScreen dpy
  let scr = defaultScreenOfDisplay dpy
  let background = whitePixel dpy dflt
  let border = blackPixel dpy dflt
  let max = 500000
  rootw <- rootWindow dpy dflt
  win  <- createSimpleWindow dpy rootw 0 0 600 600 1 border background
  --win <- mkUnmanagedWindow dpy scr rootw 0 0 500 500
  setTextProperty dpy win "this will not be seen" wM_NAME
  mapWindow dpy win
  --drawKey dpy win max
  --drawManyInWin 10 10 50 ["green","yellow", "blue","black"] dpy win 0 0
  --callNTime 50 10 10 ["green","yellow", "blue","black"] (drawManyInWin 10 10 10 dpy win)
  --drawInWin dpy win "blue"
  col <- initColor dpy "red"
  --drawAPoint 400 400 dpy win col
  drawSetParalell dpy win (0.250001 :: Double) (-0.002 :: Double) 500 500 0.000008 mand_iteration max
  --drawSetParalell dpy win (-2 :: Double) (-2 :: Double) 500 500 0.008 mand_iteration max

  --drawSet dpy win (-1.4) (-0.005) 500 500 0.00002 mand_iteration
  --mapM_ (printAllCol dpy win) [(x,y) | x <- [0..255], y <- [0..255]]
  --drawRGBRec dpy win $ pixelFromRGB 150 200 25
  --drawRGBRec dpy win $ pixelFromRGB 0 0 255
  --allocaXEvent (eventThing dpy win)
  sync dpy False
  return ()
  --threadDelay (10 * 1000000)
  --exitWith ExitSuccess

printAllCol :: Display -> Window -> (Int, Int) -> IO ()
printAllCol dpy win (x, y) = drawAPoint x y dpy win $ scaleUpToPixel (256*255 + 255) $ 256 * x + y

scaleUpToPixel :: Int -> Int -> Word64
scaleUpToPixel max n = floor $ (256*256*256 -1) * ((fromIntegral n) / (fromIntegral max))

mkUnmanagedWindow :: Display -> Screen -> Window -> Position -> Position -> Dimension -> Dimension -> IO Window
mkUnmanagedWindow dpy scr rw x y w h = do
 let visual = defaultVisualOfScreen scr
     attrmask = cWOverrideRedirect .|. cWBorderPixel .|. cWBackPixel
 win <- allocaSetWindowAttributes $ \attributes -> do
     set_override_redirect attributes True
     set_background_pixel attributes $ whitePixel dpy (defaultScreen dpy)
     set_border_pixel attributes $ blackPixel dpy (defaultScreen dpy)
     createWindow dpy rw x y w h 1 (defaultDepthOfScreen scr) inputOutput visual attrmask attributes
 return win

-- without mask:
-- took 1:51 to draw set
-- with mask:
-- took 0:31
drawSet :: RealFloat a => Display -> Window -> a -> a -> Int -> Int -> a -> (Complex a -> Complex a -> Complex a) -> Int ->IO ()
drawSet dpy win startRe startIm numberRe numberIm step nextIteration max = do
  gc <- createGC dpy win
  mapM_ (stuff gc) [(x,y) | x <- [0..(numberRe-1)] , y <- [0..(numberIm -1)] ]
  freeGC dpy gc
    where
      --stuff ::  -> (Int, Int) -> IO ()
      stuff gc (x, y) = do
        let re = startRe + (fromIntegral x) * step
        let im = startIm + (fromIntegral y) * step
        --let pixel = scaleIntToRGB (general nextIteration (re :+ im) 1000) 10
        let pixel = scaleUpToPixel (max -1) $ general nextIteration (Just zzcosMask) (re :+ im) max
        setForeground dpy gc pixel
        drawPoint dpy win gc (fromIntegral x) (fromIntegral y)

--without mask:
--took 56 seconds with 2 processing threads
--took 55 with 3 processing threads
--took 57 with 4 processing threads
--with mask
--took 17 with 2
--took 18 with 3
--took 30 with 4
drawSetParalell ::(NFData a, RealFloat a) => Display -> Window -> a -> a -> Int -> Int -> a -> (Complex a -> Complex a -> Complex a) -> Int ->IO ()
drawSetParalell dpy win startRe startIm numberRe numberIm step nextIteration max = do
  gc <- createGC dpy win
  case numCapabilities of
    1 -> drawSet dpy win startRe startIm numberRe numberIm step nextIteration max
    otherwise -> do

      vendor <- SM.newEmptyMVar
      ret <- SM.newEmptyMVar
      v1 <- forkIO $ vendorThread vendor startRe step max numberRe
      gc <- createGC dpy win
      t1 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "1"
      t2 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "2"
      t3 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "3"
      t4 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t5 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t6 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t7 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t8 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t9 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t10 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t11 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t12 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      mapM_ (\bob -> do
        Ret iteration_arr x <- SM.takeMVar ret
        drawLineOfPoints dpy win gc x 0 iteration_arr
        print bob ) [1..max]
      killThread v1
      killThread t1
      killThread t2
      killThread t3
      killThread t4
      killThread t5
      killThread t6
      killThread t7
      killThread t8
      killThread t9
      killThread t10
      killThread t11
      killThread t12
      freeGC dpy gc

processingThread :: (NFData a, RealFloat a) => SM.MVar (Ven a) -> SM.MVar (Ret a) -> a -> a -> (Complex a -> Complex a -> Complex a) -> Int -> Int -> String -> IO ()
processingThread vendor ret im_start step nextIteration max numberIm name = do
  forever (do
    Ven re x <- SM.takeMVar vendor
    --print name
    SM.putMVar ret $ Ret (calcuateLine re im_start step nextIteration max numberIm) x)

vendorThread :: (NFData a, RealFloat a)=> SM.MVar (Ven a) -> a -> a -> Int -> Int -> IO ()
vendorThread vend reStart step max maxn = do mapM_ (\n -> SM.putMVar vend $ Ven (reStart + step * fromIntegral n) n) [0.. maxn]

calcuateLine :: RealFloat a => a -> a -> a -> (Complex a -> Complex a -> Complex a) -> Int -> Int -> [Word64]
calcuateLine re im_start step nextIteration max numberIm =
  map (\ x -> scaleUpToPixel (max-1)  $ general nextIteration (Nothing ) (re :+ (im_start + (fromIntegral x) * step)) max) [0..(numberIm-1)]

drawLineOfPoints :: Display -> Window -> GC -> Int -> Int -> [Word64] -> IO ()
drawLineOfPoints dpy win gc x y [] = return ()
drawLineOfPoints dpy win gc x y (a:as) = do
  setForeground dpy gc a
  drawPoint dpy win gc (fromIntegral x) (fromIntegral y)
  drawLineOfPoints dpy win gc x (y + 1) as

drawKey :: Display -> Window -> Int -> IO ()
drawKey dpy win max = do
  gc <- createGC dpy win
  mapM_ (f gc) [0..500]
  freeGC dpy gc
  where
    f gc p = do
      drawRGBRec' dpy win  (scaleUpToPixel max $ round (fromInteger p * ((fromIntegral max) / 500))) 200 200 200 200

      setForeground dpy gc $ scaleUpToPixel max $ round (fromInteger p * ((fromIntegral max) / 500))
      fillRectangle dpy win gc 510 (fromIntegral p) 20 1
      print $ scaleUpToPixel max $ round (fromInteger p * ((fromIntegral max) / 500))

--the peramiter max should be the third root of the maximum value that n can take
--scaleIntToRGB n max = if n == max * max * max then 0 else pixelFromRGB c1 c2 c3
scaleIntToRGB :: Int -> Int -> Word64
scaleIntToRGB n max = pixelFromRGB c1 c2 c3
  where
    scale x = round $ 255 * ((fromIntegral x) / fromIntegral (max))
    c1 = scale $ div n (max* max)
    c2 = scale $ mod (div n max) (max * max)
    c3 = scale $ mod n max

scaleIntToWord :: Int -> Int -> Word64
scaleIntToWord max i = floor $ 16777215 * fromIntegral i / fromIntegral max

scaleInt :: Int -> Int -> Word64
scaleInt i max = fromIntegral . round $ ((fromIntegral i) / (fromIntegral max)) * (16777215 / fromIntegral max)
{-
scaleInt2 :: Int -> Int -> Word64
scaleInt2 i max
  | i < mrt = fromIntegral i
  | i < mrt * 2 = fromIntegral 256 * i
  | otherwise = fromIntegral 256 * 256 * i
  where
    mrt = max ** (1/3)
-}
drawInWin :: Display -> Window -> String -> IO ()
drawInWin dpy win col = do
  fgcolour <- initColor dpy col
  gc <- createGC  dpy win
  setForeground dpy gc fgcolour
  fillRectangle dpy win gc 0 0 50 50
  freeGC dpy gc

drawAPoint :: Int -> Int -> Display -> Window -> Graphics.X11.Xlib.Pixel -> IO ()
drawAPoint x y dpy win pixel = do
  gc <- createGC dpy win
  setForeground dpy gc pixel
  drawPoint dpy win gc (fromIntegral x) (fromIntegral y)
  freeGC dpy gc

drawRGBRec :: Display -> Window -> Graphics.X11.Xlib.Pixel -> IO ()
drawRGBRec dpy win pixel = do
  gc <- createGC dpy win
  setForeground dpy gc pixel
  fillRectangle dpy win gc 100 100 200 300
  freeGC dpy gc

drawRGBRec' :: Display -> Window -> Graphics.X11.Xlib.Pixel -> Int -> Int -> Int -> Int -> IO ()
drawRGBRec' dpy win pixel x y width height= do
  gc <- createGC dpy win
  setForeground dpy gc pixel
  fillRectangle dpy win gc (fromIntegral x) (fromIntegral y)  (fromIntegral width) (fromIntegral height)
  freeGC dpy gc

drawManyInWin :: Int -> Int -> Int -> Display -> Window -> Int -> Int -> [String] -> IO ()
drawManyInWin width height number dpy win x y (c:cols)= do
  col <- initColor dpy c
  gc <- createGC dpy win
  setForeground dpy gc col
  fillRectangle dpy win gc (fromIntegral x) (fromIntegral y) (fromIntegral width) (fromIntegral height)
  freeGC dpy gc
  if number == 0 then return () else drawManyInWin width height (number -1) dpy win (x + width) y (cols ++ [c])

callNTime :: Int -> Int -> Int -> [String] -> (Int -> Int -> [String] -> IO ()) -> IO ()
callNTime 0 _ _ _ _ = return ()
callNTime number startx starty colar@(c:cols) mon = (mon startx (starty * number) colar) >> callNTime (number -1) startx starty (cols ++ [c]) mon

initColor :: Display -> String -> IO Graphics.X11.Xlib.Pixel
initColor dpy colour = do
  let colourmap = defaultColormap dpy (defaultScreen dpy)
  (apros, real) <- allocNamedColor dpy colourmap colour
  return $ color_pixel apros

printBits :: Word64 -> IO()
printBits word = print "adsadfalkdfdas dadad af" -- $ wordToBinary word --map (\ a -> if a then '1' else '0') (reverse $ wordToBinary word)

wordToBinary :: Word64 -> [Bool]
wordToBinary = work 63
  where
    work word n
      | n >= 0 = if(word `div` (2 ^ n) > 0) then True:(work (word - (2^n)) (n-1)) else False:(work word (n-1))
      | otherwise = []

pixelFromRGB :: Word8 -> Word8 -> Word8 -> Word64
pixelFromRGB r g b = ((shiftL ((shiftL (zeroBits .|. toWord64 r ) 8 ).|. toWord64 g) 8) .|. toWord64 b)
  where
    toWord64 :: Word8 -> Word64
    toWord64 = fromInteger . toInteger
--pixelFromRGB r g b = (fromIntegral (65025 * r) + fromIntegral (256 * g) + fromIntegral b)

max_stars_repo_name: AlwinHughes/fractal

max_issues_repo_name: AlwinHughes/fractal

max_forks_repo_name: AlwinHughes/fractal


text:
import Test.HUnit
import FFT.DFTControl
import FFT.Orig
import FFT.Samples
import System.Environment
import Data.Complex

testDcDft = TestCase (assertEqual "DFTDc" expected actual)
    where
        expected = (sum $ dcdft samples512)
        actual = (sum $ dft samples512)
        samples512 = samples 1 512

tests = TestList [TestLabel "testDcDFT" testDcDft
                    ]

main = runTestTT tests

ext: hs

max_stars_repo_name: BNJHope/parallel-fourier-transforms

max_issues_repo_name: BNJHope/parallel-fourier-transforms

max_forks_repo_name: BNJHope/parallel-fourier-transforms


text:
{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE BangPatterns         #-}

module TestChapter4DP (
       testChapter4
      )where

import           Control.Concurrent
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.State

import           Data.Configurator
import           Data.Configurator.Types
import           Data.List.Split (chunksOf)
import           Data.Foldable (toList)

import           Data.Random
import           Data.Random.Distribution
import           Data.Random.Distribution.Poisson

import           Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import           Data.Text (Text)
import qualified Data.Vector.Unboxed as VU

import           Graphics.Matplotlib hiding (def)

import           Numeric.LinearAlgebra (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA
import           System.Console.AsciiProgress(Options(..), Stats(..),
                                              displayConsoleRegions, complete,
                                              getProgressStats, def, newProgressBar, tickN)
import           Text.Printf
-- project
import           Utils
import           Chapter4DP

testChapter4 :: FilePath -> IO ()
testChapter4 configPath = do
  print "Chapter 4 Experiment Starting "
  (config, _) <- autoReload autoConfig [Required configPath]
  (bCarRental :: Bool) <- require config "enable.bCarRental"
  when bCarRental (doCarRentalTest config)
  (bGambler :: Bool) <- require config "enable.bGamblerProblem"
  when bGambler (doGamblerTest config)

----------------------------------------------------------------------------------------------------
-- Car Rental Test
doCarRentalTest :: Config -> IO ()
doCarRentalTest config = do
  (experimentName::String) <- require config "carRental.experiment"
  (theTheta::Double) <- require config "carRental.theta"
  (discountGamma::Double) <- require config "carRental.discount"
  (maxCars::[Int]) <- require config "carRental.maxCars"
  (rentalCredit::[Double]) <- require config "carRental.rentalCredit"
  (transCost::[Double]) <- require config "carRental.transferCost"
  (maxTransferCars::[Int]) <- require config "carRental.maxTransferCars"
  (freeParkingLimit::[Int]) <- require config "carRental.freeParkingLimit"
  (additionalParkingCost::[Double]) <- require config "carRental.additionalParkingCost"
  (additionalTransferSaving::[Double]) <- require config "carRental.additionalTransferSaving"
  (rentalCars::[Double]) <- require config "carRental.rentalCars"
  (returnCars::[Double]) <- require config "carRental.returnCars"

  let carRental = mkCarRental (length maxCars) theTheta discountGamma maxCars rentalCredit transCost
                              maxTransferCars freeParkingLimit additionalParkingCost
                              additionalTransferSaving rentalCars returnCars
  -- do experiments
  carRental' <- goLoop carRental
  when (_locationNum carRental == 2) (drawTwoLocationCarRental carRental' experimentName)
  threadDelay 100000
  where
  goLoop carRental = displayConsoleRegions $ do
    putStrLn "Will do car rental experiment " >> putStrLn (show carRental)
    pg <- newProgressBar def { pgWidth = 80
                             , pgTotal = 100
                             , pgOnCompletion = Just "Done :percent after :elapsed seconds"
                             }
    carRental' <- loop pg carRental
    putStrLn "car rental experiment finish. Final Results: "
    putStrLn $ showCarRentalResult carRental'
    pure carRental'
    where
    loop pg carRental = do
      let (!(bFinish, percent), !carRental') = runState carRentalStep carRental
      case bFinish of
         False -> do
           stat <- getProgressStats pg
           let ticked = fromInteger $ stCompleted stat
               willTick = percent - ticked
           if (willTick > 0) then(tickN pg willTick)
              else print ("Finish One Improvement " ++ show percent)
           loop pg carRental'
         True -> complete pg >> pure carRental'

drawTwoLocationCarRental :: CarRental -> String -> IO ()
drawTwoLocationCarRental carRental experimentName = do
  putStrLn "Draw Two Location Car Rental Result Graph"
  -- dataX is the second location, whereas dataY is the first location
  let dataX = toList $ fmap (\ (f:s:[]) -> s) (_states carRental)
      dataY = toList $ fmap (\ (f:s:[]) -> f) (_states carRental)
      (dataZ::[Int]) = toList (Seq.zipWith (\ i saPair -> calcZ . foldl (zipWith (+)) [0, 0] $ Seq.index saPair i)
                                           (_actions carRental) (_possibleActions carRental))
      dataValue = toList $ _stateValues carRental
      maxZ = maximum dataZ
      figure = readData ([dataX], [dataY], [dataZ])
                 % mp # "ax = plot.gca(projection='3d')"
                 % mp # "ax.scatter(np.array(data[" # 0 # "]), np.array(data[" # 1 # "]), np.array(data[" # 2 # "]))"
                 % xlabel "#Cars at second location"
                 % ylabel "#Cars at first location"
                 % zlabel "#Cars to move (action)"
                 % xticks [0, 1..(_maxCars carRental !! 1)]
                 % yticks [0, 1..(_maxCars carRental !! 0)]
                 % zticks [negate maxZ, (negate $ maxZ - 1)..maxZ]
                 % title experimentName
      valueFigure = readData ([dataX], [dataY], [dataValue])
                      % mp # "ax = plot.gca(projection='3d')"
                      % mp # "ax.scatter(np.array(data[" # 0 # "]), np.array(data[" # 1 # "]), np.array(data[" # 2 # "]))"
                      % xlabel "#Cars at second location"
                      % ylabel "#Cars at first location"
                      % zlabel "State Values"
                      % xticks [0, 1..(_maxCars carRental !! 1)]
                      % yticks [0, 1..(_maxCars carRental !! 0)]
                      % title (experimentName ++ " - State Values")
  -- avoid Matplotlib's bug
  code figure >> code valueFigure >> onscreen figure >> onscreen valueFigure
  threadDelay 100000
  -- also output table format
  let moveResult = showAsTable [0..(_maxCars carRental)!!0]
                               [0..(_maxCars carRental)!!1] (map fromIntegral dataZ) False
      stateResult = showAsTable [0..(_maxCars carRental)!!0]
                                [0..(_maxCars carRental)!!1] dataValue True
  putStrLn moveResult
  putStrLn stateResult
  where
  calcZ :: [Int] -> Int
  calcZ (firstLocation:secondLocation:[]) =
            (firstLocation > 0) ? (negate firstLocation, secondLocation)
  calcZ _ = error "Not correct action move"

-- output in github markdown format
showAsTable :: [Int] -> [Int] -> [Double] -> Bool -> String
showAsTable col row vals bFloat =
  let colLen = length col
      header = "|First/second| " ++ (concat $ map ((++ "|") . show) col) ++ "\n"
      alignHeader = (concat . take (colLen + 1) $ repeat "|:-----:") ++ "|\n"
      showRows = concat $ map (go colLen vals) row
  in  header ++ alignHeader ++ showRows
  where
  go len vals rowIdx =
    let first = show rowIdx ++ "|"
        rows = take len $ drop (len * rowIdx) vals
        format = bFloat ? ("%7.2f", "%7.0f")
    in  first ++ (concat $ map (\ x -> (printf format x :: String) ++ "|") rows) ++ "|\n"

----------------------------------------------------------------------------------------------------
-- Gambler Problem Test

doGamblerTest :: Config -> IO ()
doGamblerTest config = do
  (prob::Double) <- require config "gambler.headProb"
  (gamblerGoal::Int) <- require config "gambler.goal"
  (theTheta::Double) <- require config "gambler.theta"
  let gambler = mkGambler prob theTheta gamblerGoal
      (_, !gambler') = runState gamblerStep gambler
  drawGamblerGraph gambler'
  threadDelay 100000
  pure ()

drawGamblerGraph :: Gambler -> IO ()
drawGamblerGraph gambler = do
  putStrLn "Draw Gambler Problem Result Graph"
  -- dataX is the second location, whereas dataY is the first location
  let policyFigure = plot [0..(goal gambler)] (stateVals gambler)
                       % xlabel ("Capital (head probability = " ++ (show $ headProb gambler) ++ ")")
                       % ylabel "Value Estimates"
                       % xticks [0, 10..(goal gambler)]
                       % title ("Figure 4.3 - State Values")
      valueFigure = scatter [0..(goal gambler - 1)] (stateActs gambler)
                      % xlabel ("Capital (head probability = " ++ (show $ headProb gambler) ++ ")")
                      % ylabel "Final Policy (Stack)"
                      % xticks [0, 10..(goal gambler)]
                      % title ("Figure 4.3 - Optimal Actions")
  -- avoid Matplotlib's bug
  code policyFigure >> code valueFigure >> onscreen policyFigure >> onscreen valueFigure
  threadDelay 100000

max_stars_repo_name: Xingtao/ReinforceLearningIntro

max_issues_repo_name: Xingtao/ReinforceLearningIntro

max_forks_repo_name: Xingtao/ReinforceLearningIntro


text:
-- |
-- Module    : Statistics.Test.MannWhitneyU
-- Copyright : (c) 2010 Neil Brown
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Mann-Whitney U test (also know as Mann-Whitney-Wilcoxon and
-- Wilcoxon rank sum test) is a non-parametric test for assessing
-- whether two samples of independent observations have different
-- mean.
module Statistics.Test.MannWhitneyU (
    -- * Mann-Whitney U test
    mannWhitneyUtest
  , mannWhitneyU
  , mannWhitneyUCriticalValue
  , mannWhitneyUSignificant
    -- ** Wilcoxon rank sum test
  , wilcoxonRankSums
  , module Statistics.Test.Types
    -- * References
    -- $references
  ) where

import Control.Applicative ((<$>))
import Data.List (findIndex)
import Data.Ord (comparing)
import Numeric.SpecFunctions (choose)
import Prelude hiding (sum)
import Statistics.Distribution (quantile)
import Statistics.Distribution.Normal (standard)
import Statistics.Function (sortBy)
import Statistics.Sample.Internal (sum)
import Statistics.Test.Internal (rank, splitByTags)
import Statistics.Test.Types (TestResult(..), PositionTest(..), significant)
import Statistics.Types (PValue,pValue)
import qualified Data.Vector.Unboxed as U

-- | The Wilcoxon Rank Sums Test.
--
-- This test calculates the sum of ranks for the given two samples.
-- The samples are ordered, and assigned ranks (ties are given their
-- average rank), then these ranks are summed for each sample.
--
-- The return value is (W₁, W₂) where W₁ is the sum of ranks of the first sample
-- and W₂ is the sum of ranks of the second sample.  This test is trivially transformed
-- into the Mann-Whitney U test.  You will probably want to use 'mannWhitneyU'
-- and the related functions for testing significance, but this function is exposed
-- for completeness.
wilcoxonRankSums :: (Ord a, U.Unbox a) => U.Vector a -> U.Vector a -> (Double, Double)
wilcoxonRankSums xs1 xs2 = (sum ranks1, sum ranks2)
  where
    -- Ranks for each sample
    (ranks1,ranks2) = splitByTags $ U.zip tags (rank (==) joinSample)
    -- Sorted and tagged sample
    (tags,joinSample) = U.unzip
                      $ sortBy (comparing snd)
                      $ tagSample True xs1 U.++ tagSample False xs2
    -- Add tag to a sample
    tagSample t = U.map (\x -> (t,x))

-- | The Mann-Whitney U Test.
--
-- This is sometimes known as the Mann-Whitney-Wilcoxon U test, and
-- confusingly many sources state that the Mann-Whitney U test is the same as
-- the Wilcoxon's rank sum test (which is provided as 'wilcoxonRankSums').
-- The Mann-Whitney U is a simple transform of Wilcoxon's rank sum test.
--
-- Again confusingly, different sources state reversed definitions for U₁
-- and U₂, so it is worth being explicit about what this function returns.
-- Given two samples, the first, xs₁, of size n₁ and the second, xs₂,
-- of size n₂, this function returns (U₁, U₂)
-- where U₁ = W₁ - (n₁(n₁+1))\/2
-- and U₂ = W₂ - (n₂(n₂+1))\/2,
-- where (W₁, W₂) is the return value of @wilcoxonRankSums xs1 xs2@.
--
-- Some sources instead state that U₁ and U₂ should be the other way round, often
-- expressing this using U₁' = n₁n₂ - U₁ (since U₁ + U₂ = n₁n₂).
--
-- All of which you probably don't care about if you just feed this into 'mannWhitneyUSignificant'.
mannWhitneyU :: (Ord a, U.Unbox a) => U.Vector a -> U.Vector a -> (Double, Double)
mannWhitneyU xs1 xs2
  = (fst summedRanks - (n1*(n1 + 1))/2
    ,snd summedRanks - (n2*(n2 + 1))/2)
  where
    n1 = fromIntegral $ U.length xs1
    n2 = fromIntegral $ U.length xs2

    summedRanks = wilcoxonRankSums xs1 xs2

-- | Calculates the critical value of Mann-Whitney U for the given sample
-- sizes and significance level.
--
-- This function returns the exact calculated value of U for all sample sizes;
-- it does not use the normal approximation at all.  Above sample size 20 it is
-- generally recommended to use the normal approximation instead, but this function
-- will calculate the higher critical values if you need them.
--
-- The algorithm to generate these values is a faster, memoised version of the
-- simple unoptimised generating function given in section 2 of \"The Mann Whitney
-- Wilcoxon Distribution Using Linked Lists\"
mannWhitneyUCriticalValue
  :: (Int, Int)     -- ^ The sample size
  -> PValue Double  -- ^ The p-value (e.g. 0.05) for which you want the critical value.
  -> Maybe Int      -- ^ The critical value (of U).
mannWhitneyUCriticalValue (m, n) p
  | m < 1 || n < 1 = Nothing    -- Sample must be nonempty
  | p' <= 1        = Nothing    -- p-value is too small. Null hypothesis couldn't be disproved
  | otherwise      = findIndex (>= p')
                   $ take (m*n)
                   $ tail
                   $ alookup !! (m+n-2) !! (min m n - 1)
  where
    mnCn = (m+n) `choose` n
    p'   = mnCn * pValue p

{-
-- Original function, without memoisation, from Cheung and Klotz:
-- Double is needed to avoid integer overflows.
a :: Int -> Int -> Int -> Double
a u bigN m
  | u < 0            = 0
  | u >= m * n       = bigN `choose` m
  | m == 1 || n == 1 = fromIntegral (u + 1)
  | otherwise        = a  u      (bigN - 1)  m
                     + a (u - n) (bigN - 1) (m-1)
  where
    n = bigN - m
-}

-- Memoised version of the original a function, above.
--
-- Doubles are stored to avoid integer overflow. 32-bit Ints begin to
-- overflow for bigN as small as 33 (64-bit one at 66) while Double to
-- go to infinity till bigN=1029
--
--
-- outer list is indexed by big N - 2
-- inner list by (m-1) (we know m < bigN)
-- innermost list by u
--
-- So: (alookup !! (bigN - 2) !! (m - 1) ! u) == a u bigN m
alookup :: [[[Double]]]
alookup = gen 2 [1 : repeat 2]
  where
    gen bigN predBigNList
       = let bigNlist = [ [ amemoed u m
                          | u <- [0 .. m*(bigN-m)]
                          ] ++ repeat (bigN `choose` m)
                        | m <- [1 .. (bigN-1)]] -- has bigN-1 elements
         in bigNlist : gen (bigN+1) bigNlist
      where
        amemoed :: Int -> Int -> Double
        amemoed u m
          | m == 1 || n == 1 = fromIntegral (u + 1)
          | otherwise        = mList !! u
                             + if u < n then 0 else predmList !! (u-n)
          where
            n = bigN - m
            (predmList : mList : _) = drop (m-2) predBigNList
            -- Lists for m-1 and m respectively. i-th list correspond to m=i+1
            --
            -- We know that predBigNList has bigN - 2 elements
            -- (and we know that n > 1 therefore bigN > m + 1)
            -- So bigN - 2 >= m, i.e. predBigNList must have at least m elements
            -- elements, so dropping (m-2) must leave at least 2

-- | Calculates whether the Mann Whitney U test is significant.
--
-- If both sample sizes are less than or equal to 20, the exact U critical value
-- (as calculated by 'mannWhitneyUCriticalValue') is used.  If either sample is
-- larger than 20, the normal approximation is used instead.
--
-- If you use a one-tailed test, the test indicates whether the first sample is
-- significantly larger than the second.  If you want the opposite, simply reverse
-- the order in both the sample size and the (U₁, U₂) pairs.
mannWhitneyUSignificant
  :: PositionTest     -- ^ Perform one-tailed test (see description above).
  -> (Int, Int)       -- ^ The samples' size from which the (U₁,U₂) values were derived.
  -> PValue Double    -- ^ The p-value at which to test (e.g. 0.05)
  -> (Double, Double) -- ^ The (U₁, U₂) values from 'mannWhitneyU'.
  -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too
                      --   small to make a decision.
mannWhitneyUSignificant test (in1, in2) pVal (u1, u2)
  -- Use normal approximation
  | in1 > 20 || in2 > 20 =
    let mean  = n1 * n2 / 2     -- (u1+u2) / 2
        sigma = sqrt $ n1*n2*(n1 + n2 + 1) / 12
        z     = (mean - u1) / sigma
    in Just $ case test of
                AGreater      -> significant $ z     < quantile standard p
                BGreater      -> significant $ (-z)  < quantile standard p
                SamplesDiffer -> significant $ abs z > abs (quantile standard (p/2))
  -- Use exact critical value
  | otherwise = do crit <- fromIntegral <$> mannWhitneyUCriticalValue (in1, in2) pVal
                   return $ case test of
                              AGreater      -> significant $ u2        <= crit
                              BGreater      -> significant $ u1        <= crit
                              SamplesDiffer -> significant $ min u1 u2 <= crit
  where
    n1 = fromIntegral in1
    n2 = fromIntegral in2
    p  = pValue pVal

-- | Perform Mann-Whitney U Test for two samples and required
-- significance. For additional information check documentation of
-- 'mannWhitneyU' and 'mannWhitneyUSignificant'. This is just a helper
-- function.
--
-- One-tailed test checks whether first sample is significantly larger
-- than second. Two-tailed whether they are significantly different.
mannWhitneyUtest
  :: (Ord a, U.Unbox a)
  => PositionTest     -- ^ Perform one-tailed test (see description above).
  -> PValue Double    -- ^ The p-value at which to test (e.g. 0.05)
  -> U.Vector a       -- ^ First sample
  -> U.Vector a       -- ^ Second sample
  -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too small to
                      --   make a decision.
mannWhitneyUtest ontTail p smp1 smp2 =
  mannWhitneyUSignificant ontTail (n1,n2) p $ mannWhitneyU smp1 smp2
    where
      n1 = U.length smp1
      n2 = U.length smp2

-- $references
--
-- * Cheung, Y.K.; Klotz, J.H. (1997) The Mann Whitney Wilcoxon
--   distribution using linked lists. /Statistica Sinica/
--   7:805–813.
-- <http://www3.stat.sinica.edu.tw/statistica/oldpdf/A7n316.pdf>.


text:
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Data as D (tr,  (><))

type X_Train = LA.Matrix Double
type Y_Train = LA.Matrix Double
type W = Double
type B = Double
type Output = LA.Matrix Double
type LearningRate = Double
-- | y = w * x + b for a single training example

ones :: Int -> Int -> LA.Matrix Double
ones x y = (x D.>< y) [1,1 ..]

output' :: W -> X_Train -> B -> Output
output' w xTrain b = w `LA.scale` xTrain + b `LA.scale` ones (fst (LA.size xTrain)) 1

cost :: Output -> Y_Train -> Double
cost output yTrain = LA.sumElements $ (1/2) * (output - yTrain) ^ 2

deltaCW' :: W -> X_Train -> B -> Double
deltaCW' w xTrain b = LA.sumElements $ D.tr (output' w xTrain b) <> xTrain

deltaCB' :: W -> X_Train -> Y_Train -> B -> Double
deltaCB' w xTrain yTrain b = LA.sumElements $ (output' w xTrain b - yTrain)

updateW' :: W -> X_Train -> B -> LearningRate -> Double
updateW' w xTrain b alpha = w - alpha * deltaCW' w xTrain b

updateB' :: W -> X_Train -> Y_Train -> B -> LearningRate -> Double
updateB' w xTrain yTrain b alpha = b - alpha * deltaCB' w xTrain yTrain b

test = do
  xTrain <- LA.rand 3 1
  yTrain <- LA.rand 3 1
  let w = 0
  let b = 0
  let alpha = 0.00001
  let loop i w' b' | i > 600 = return ()
      loop i w' b' | i <= 600 = do
         let c = cost (output' w' xTrain b') yTrain
         putStrLn $ concat ["Cost: ", show c]
         let w'' = updateW' w' xTrain b' alpha
         let b'' = updateB' w' xTrain yTrain b' alpha
         loop (i + 1) w'' b''
  loop 0 w b

{- ==== BAD CODE ====
Only for my own learning

type W = LA.Matrix Double
type B = LA.Matrix Double

output :: W -> X_Train -> B -> Output
output w xTrain b = (D.tr (w) <> xTrain) + b

deltaCW :: W -> X_Train -> Y_Train -> B -> LA.Matrix Double
deltaCW w xTrain yTrain b = ((D.tr (w) <> xTrain) + b) <> xTrain

deltaCB :: W -> X_Train -> Y_Train -> B -> LA.Matrix Double
deltaCB w xTrain yTrain b = (D.tr (w) <> xTrain) + b

updateW :: W -> X_Train -> Y_Train -> B -> LearningRate -> W
updateW w xTrain yTrain b alpha = w - alpha `LA.scale` deltaCW w xTrain yTrain b

updateB :: W -> X_Train -> Y_Train -> B -> LearningRate -> W
updateB w xTrain yTrain b alpha = b - alpha `LA.scale` deltaCB w xTrain yTrain b

main = do
  xTrain <- LA.rand 3 1
  yTrain <- LA.rand 3 1
  w <- LA.rand 3 1
  b <- LA.rand 1 1
  let alpha = 0.0001
  let loop i w' b' | i > 1000 = return ()
      loop i w' b' | i <= 1000 = do
         let c = cost (output w' xTrain b') yTrain
         putStrLn $ concat ["Cost: ", show c]
  --       putStrLn "\n"
         let w'' = updateW w' xTrain yTrain b' alpha
         let b'' = updateB w' xTrain yTrain b' alpha
--         putStrLn $ concat ["Weights: ", show w'']
--         putStrLn "\n"
--         putStrLn $ concat ["Bias: ", show b'']
         loop (i + 1) w'' b''
  loop 0 w b
-}

max_stars_repo_name: Lockne/neural-network-from-scratch

max_issues_repo_name: Lockne/neural-network-from-scratch

max_forks_repo_name: Lockne/neural-network-from-scratch


text:
{-# LANGUAGE RecordWildCards #-}

import Control.Monad
import Data.Complex
import Control.Monad.Trans.Either
import Foreign.C.Types
import Data.Word

import qualified Data.Vector.Storable as VS hiding ((++))
import qualified Data.Vector.Generic  as VG hiding ((++))
import           Pipes
import qualified Pipes.Prelude        as P
import           Options.Applicative

import SDR.Filter
import SDR.RTLSDRStream
import SDR.Util
import SDR.FFT
import SDR.Plot
import SDR.Demod
import SDR.Pulse
import SDR.PipeUtils
import SDR.ArgUtils
import SDR.CPUID

import Graphics.DynamicGraph.Waterfall
import Graphics.DynamicGraph.Util

--The filter coefficients are stored in another module
import Coeffs

data Options = Options {
    frequency :: Word32
}

optParser :: Parser Options
optParser = Options
          <$> option (fmap fromIntegral parseSize) (
                 long "frequency"
              <> short 'f'
              <> metavar "FREQUENCY"
              <> help "Frequency to tune to"
              )

opt :: ParserInfo Options
opt = info (helper <*> optParser) (fullDesc <> progDesc "SDR library demo. Receive FM broadcast band radio." <> header "SDR demo")

samples      = 8192
audioSamples = 1024

{-
    sampling frequency of the input is 1280 khz
    this is decimated by a factor of 8 and then demodulated
    sampling frequency of demodulated signal is 160 khz
    need audio output at 48 khz
    resampling factor is 48/160 == 3/10
    FM pilot tone at 19khz (0.3958 * 48)
    start audio filter cutoff at 15khz (0.3125 * 48)
-}

doIt Options{..} = do

    info <- lift getCPUInfo

    --Initialize GLFW
    res            <- lift setupGLFW
    unless res (left "Unable to initilize GLFW")

    --Initialize the RTLSDR device
    str            <- sdrStream frequency 1280000 1 (fromIntegral samples * 2)

    --Initialize the graphs and FFTs
    rfFFT          <- lift $ fftw samples
    rfSpectrum     <- plotWaterfall 1024 480 samples 1000 jet_mod
    audioFFT       <- lift $ fftwReal audioSamples
    audioSpectrum  <- plotFillAxes 1024 480 ((audioSamples `quot` 2) + 1) jet (zeroAxes 1024 480 48 10)

    --Initialize the sound sink
    pulseSink      <- lift pulseAudioSink

    --Initialize the filters
    deci <- lift $ fastDecimatorC info 8 coeffsRFDecim
    resp <- lift $ fastResamplerR info 3 10 coeffsAudioResampler
    filt <- lift $ fastFilterSymR info coeffsAudioFilter

    let window0 = hanning samples :: VS.Vector Float
        window  = hanning samples

    --Build the pipeline
    lift $ runEffect $   str
                     >-> P.map convertCAVX
                     >-> foldl1 combine [

                            P.map (VG.zipWith (flip mult) window0 . VG.zipWith mult (fftFixup samples))
                        >-> P.map (VG.map (cplxMap (realToFrac :: Float -> CDouble)))
                        >-> rfFFT
                        >-> P.map (VG.map ((* (32 / fromIntegral samples)) . realToFrac . magnitude))
                        >-> rfSpectrum,

                            firDecimator deci samples
                        >-> fmDemod
                        >-> firResampler resp samples
                        >-> firFilter filt audioSamples
                        >-> foldl1 combine [

                                P.map (VG.zipWith (*) window)
                            >-> P.map (VG.map (realToFrac :: Float -> CDouble))
                            >-> audioFFT
                            >-> P.map (VG.map ((/ 10) . realToFrac . magnitude))
                            >-> audioSpectrum,

                                P.map (VG.map ((* 0.2) . realToFrac))
                            >-> pulseSink
                        ]
                    ]

main = execParser opt >>= eitherT putStrLn return . doIt

max_stars_repo_name: adamwalker/sdr-demo

max_issues_repo_name: adamwalker/sdr-demo

max_forks_repo_name: adamwalker/sdr-demo


text:
{-# OPTIONS_HADDOCK prune #-}
module ForSyDe.Atom.Skel.FastVector.DSP where

import Data.Complex
import qualified Data.Number.FixedFunctions as F
import ForSyDe.Atom.MoC as MoC
import ForSyDe.Atom.Skel.FastVector.Lib as V hiding (duals, unduals)
import ForSyDe.Atom.Skel.FastVector.Matrix as M
import ForSyDe.Atom ((><))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.taylor'.
taylor :: (Eq a, Floating a)
       => Int      -- ^ number of points in the output window.
       -> Int      -- ^ Number of nearly constant level sidelobes adjacent to the mainlobe
       -> a        -- ^ Desired peak sidelobe level in decibels (db) relative to the mainlobe
       -> Vector a -- ^ The window, with the center value normalized to one (the value
                   -- one appears only if the number of samples is odd).
taylor n nbar level = V.farm11 (*scale) w
  where
    -- explicit conversions to floating
    bN   = fromIntegral n
    nBar = fromIntegral nbar
    ma   = map fromIntegral [1 .. nbar-1]
    -- calculate intermediate values
    b  = 10**((-level) / 20)
    a  = log(b + sqrt(b**2 - 1)) / pi
    s2 = nBar ** 2 / (a**2 + (nBar - 0.5)**2)
    -- functions for calculating coefficients
    fmcalc m  = let numer = (-1)**(m+1) * product[1.0 - m**2/s2/(a**2 + (j - 0.5)**2) | j <- ma]
                    denom = 2 * product[1 - m**2/j**2 | j <- ma, j /= m]
                in numer / denom
    ccalc m x = cos(2 * pi * x * (m - bN/2 + 1/2) / bN)
    wcalc m   = 2 * dotvv (vector $ map fmcalc ma) (vector $ map (ccalc m) ma) + 1
    -- calculate window coefficients
    w  = vector $ map (wcalc . fromIntegral) [0..n-1]
    -- normalize (Note that this is not described in the original text [1])
    scale = 1 / wcalc (bN - 1) / 2

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.taylor''.
taylor' n = taylor 4 (-30)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvm'.
dotvm :: Num a => Vector a -> Vector (Vector a) -> Vector a
dotvm = dotvm' (+) (*)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvv'.
dotvv :: Num a => Vector a -> Vector a -> a
dotvv a b
  | V.length a == V.length b = V.reduce (+) (V.farm21 (*) a b)
  | otherwise                = error "Vector sizes must match"

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvv''.
dotvv' :: Num a
     => ((a -> a -> a) -> f a -> f a -> f a)
     -- ^ higher-order function that can wrap the (*) operation.
     -> Vector (f a) -> Vector (f a) -> f a
dotvv' wrap a b
  | V.length a == V.length b = V.reduce (wrap (+)) $ V.farm21 (wrap (*)) a b
  | otherwise                = error "Vector sizes must match"

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotmv''.
dotmv' :: (a -> a -> a) -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (b -> a -> a) -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector a      -- ^ /length/ = @xa@
       -> Vector a      -- ^ /length/ = @xa@
dotmv' f g mA y = V.farm11 (\x -> V.reduce f $ V.farm21 g x y) mA

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvm''.
dotvm' :: (b -> b -> b)     -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (a -> b -> b)     -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector a          -- ^ /length/ = @xa@
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector b          -- ^ /length/ = @ya@
dotvm' f g vs = V.reduce (V.farm21 f) . V.farm21 (\x -> V.farm11 (g x)) vs

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.hanning'.
hanning :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hanning size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.5 * (1 - cos((2 * pi * i) / (n - 1)))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.hamming'.
hamming :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hamming size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.54 - 0.46 * cos((2 * pi * i) / (n - 1))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.blackman'.
blackman :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
blackman size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.42 - 0.5 * cos((2 * pi * i) / (n - 1)) + 0.08 * cos((4 * pi * i) / (n - 1))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fir'.
fir :: Num a
    => Vector a  -- ^ vector of coefficients
    -> Vector a  -- ^ input vector of numbers; /size/ = @n@
    -> Vector a  -- ^ output vector of numbers; /size/ = @n@
fir coefs = V.reverse . V.farm11 applyFilter . tails . V.reverse
  where
    applyFilter = V.reduce (+) . V.farm21 (*) coefs

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fir''.
fir' :: (a -> a -> a)  -- ^ process/operation replacing '+'
     -> (c -> a -> a)  -- ^ process/operation replacing '*'
     -> (a -> a)       -- ^ delay process
     -> Vector c       -- ^ vector of coefficients
     -> a              -- ^ input signal/structure
     -> a              -- ^ output signal/structure
fir' plus times delay coefs =
  -- V.reduce plus . V.farm21 (\c -> times c) coefs . V.recur (V.fanoutn n delay <: id)
  V.reduce plus . V.farm21 times coefs . V.recuri (V.fanoutn n delay)
  where n = V.length coefs - 1

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.twiddles'.
twiddles :: Floating a => Int -> V.Vector (Complex a)
twiddles bN = (bitrev . V.take (bN `div` 2)) (V.farm11 bW $ vector [0..])
  where bW x = (cis . negate) (-2 * pi * fromIntegral x / fromIntegral bN)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fft'.
fft :: RealFloat a
     => Int -> V.Vector (Complex a) -> V.Vector (Complex a)
fft k vs | n == 2^k = V.reverse $ bitrev $ (stage `V.pipe1` V.iterate k (*2) 2) vs
         | otherwise = error $ "n=" ++ show n ++ "; k=" ++ show k
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input
    -------------------------------------------------
    butterfly w x0 x1 = let t = w * x1 in (x0 + t, x0 - t) -- kernel function

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fft''.
fft' :: Floating a
     => (Complex a -> a -> a -> (a, a))
     -> Int -> V.Vector a -> V.Vector a
fft' butterfly k vs | n == 2^k = bitrev $ (stage `V.pipe1` (V.iterate k (*2) 2)) vs
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.duals'.
duals    :: Vector a -> (Vector a, Vector a)
duals v  = (V.take k v, V.drop k v)
  where
    k = V.length v `div` 2

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.unduals'.
unduals  :: Vector a -> Vector a -> Vector a
unduals x y =  x <++> y

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.bitrev'.
bitrev = V.unsafeLift bitrevF
  where
    bitrevF [x] = [x]
    bitrevF xs  = bitrevF (V.evensF xs) ++ bitrevF (V.oddsF xs)

max_stars_repo_name: forsyde/forsyde-atom

max_issues_repo_name: forsyde/forsyde-atom

max_forks_repo_name: forsyde/forsyde-atom


text:
{-# LANGUAGE BangPatterns #-}
-- file: Astro/Misc.hs
-- Miscellaneous utilities

module Astro.Misc where

import Data.List
import Debug.Trace
import Numeric.LinearAlgebra
import Numeric (showEFloat)
import Text.Printf

{-
 - General math
 -}

type Polynomial = [Double]

-- Evaluate real polynomial given as list of coefficients [a0,a1,...]
evalPoly :: Polynomial->Double->Double
evalPoly p x = sum $ zipWith (*) p $ map (x^) [0..]

-- Derivative of a (real) polynomial
diffPoly :: Polynomial->Polynomial
diffPoly [] = []
diffPoly (a0:[]) = []
diffPoly (a0:as) = zipWith (*) as [1..]

-- Degree of a polynomial
degPoly :: Polynomial->Int
degPoly [] = 0
degPoly (a:[]) = if a /= 0 then 1 else 0
degPoly p = if l == [] then 0 else last l
	where
	l = Data.List.findIndices (/= 0) p

-- Laguerre polynomial root finder
-- Returns Nothing if fails to converge to a real root within maxIt
rootLaguerre (tol,maxIt) p x0  = result
	where
	result = helper x0 0
	p' = diffPoly p
	p'' = diffPoly p'
	deg = fromIntegral $ degPoly p
	helper x n
		| isNaN x = Nothing
		| abs (evalPoly p x) <= tol = Just x
		| n >= maxIt = Nothing -- `debug` ("maxit x,nx:"++show(x,nx))
		| otherwise = helper nx (n+1)
                {-
                        `debug` (printf ("g:n %g h: %g sqa: %g sq: %g x:"
                                ++"%g y: %g") g h sqarg sq nx y)
                                -}
		where
		g = evalPoly p' x / evalPoly p x
		h = g^2 - evalPoly p'' x / evalPoly p x
		sqarg = (deg-1)*(deg*h-g^2)
		sq = sqrt sqarg
		denom = if g>0 then g+sq else g-sq
		nx = x - deg/denom -- `debug` ("n,denom:"++show(n,denom))
                y = evalPoly p nx

{-
 - Vector Math
 -}

-- Standard vector norms
norm :: Vector Double -> Double
--norm v = pnorm PNorm2 v
-- XXX: For 3-vectors this is quite a bit faster
norm !v = sqrt $ v@>0*v@>0 + v@>1*v@>1 + v@>2*v@>2

norm3D :: Vector Double -> Double
norm3D = sqrt . sqrNorm3D

sqrNorm3D :: Vector Double -> Double
sqrNorm3D v = v@>0*v@>0 + v@>1*v@>1 + v@>2*v@>2

-- Normalize vector
unitV v = scale (recip $ norm v) v

-- 3-vector cross product
cross :: Vector Double->Vector Double->Vector Double
cross !a !b =  3|>[a@>1 * b@>2 - a@>2 * b@>1
		,-a@>0 * b@>2 + a@>2 * b@>0
		,a@>0 * b@>1 - a@>1 * b@>0]

-- Rodrigues' rotation formula
rodr :: Vector Double->Double->Vector Double->Vector Double
rodr z th vec = scale (cos th) vec + scale (sin th) (cross z vec)
		+ scale ((z<.>vec)*(1-cos th)) z

{-
 - List and tuple manipulations
 -}

thrd (_,_,a) = a
frth (_,_,_,a) = a

-- Give all pairings, excluding (p,p) and when (p1,p2) excluding (p2,p1)
-- Courtesy of Jussi Knuuttila
chooseTwo :: [a] -> [(a,a)]
chooseTwo as = concatMap pairs (tails as)
    where pairs [] = []
          pairs (a:as) = zip (repeat a) as

-- Take every nth of list
takenth 0 _ = []
takenth _ [] = []
takenth n list =
	iter list []
	where
	iter [] acc = acc
	iter list acc =
		iter (tail newList) (acc ++ [head newList])
		where newList = drop (n-1) list

-- Debugging and printing functions
debug = flip trace
showFloat f = showEFloat (Just 4) f ""

max_stars_repo_name: sageh/AstroHaskell

max_issues_repo_name: sageh/AstroHaskell

max_forks_repo_name: sageh/AstroHaskell


text:
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-|

Module      : Grenade.Layers.Dropout
Description : Defines the Dropout Layer

Dropout is a regularization method used to prevent overfitting
of the network to the training dataset.
It works by randomly "dropping out", i.e. ignoring, outputs of
a percentage of neurons.
This is strategy is effective in preventing overfitting as it
forces a NN not to rely on a single feature for computing the
output of a neural network
-}
module Grenade.Layers.Dropout (
    Dropout (..)
  , randomDropout
  ) where

import           Control.DeepSeq
import           Control.Monad.Primitive        (PrimBase, PrimState)
import           Data.Proxy
import           Data.Serialize
import           GHC.Generics                   hiding (R)
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static   hiding (Seed)
import           System.Random.MWC

import           Grenade.Core

-- Dropout layer help to reduce overfitting.
-- Idea here is that the vector is a shape of 1s and 0s, which we multiply the input by.
-- After backpropogation, we return a new matrix/vector, with different bits dropped out.
-- The provided argument is the proportion to drop in each training iteration (like 1% or
-- 5% would be reasonable).
data Dropout (pct :: Nat) =
  Dropout
    { dropoutActive :: Bool     -- ^ Add possibility to deactivate dropout
    , dropoutSeed   :: !Int     -- ^ Seed
    }
  deriving (Generic)

instance NFData (Dropout pct) where rnf (Dropout a s) = rnf a `seq` rnf s
instance Show (Dropout pct) where show (Dropout _ _) = "Dropout"
instance Serialize (Dropout pct) where
  put (Dropout act seed) = put act >> put seed
  get = Dropout <$> get <*> get

instance UpdateLayer (Dropout pct) where
  type Gradient (Dropout pct) = ()
  runUpdate _ (Dropout act seed) _ = Dropout act (seed+1)
  runSettingsUpdate set (Dropout _ seed) = Dropout (setDropoutActive set) seed
  reduceGradient _ = ()

instance RandomLayer (Dropout pct) where
  createRandomWith _ = randomDropout

randomDropout :: (PrimBase m) => Gen (PrimState m) -> m (Dropout pct)
randomDropout gen = Dropout True <$> uniform gen

instance (KnownNat pct, KnownNat i) => Layer (Dropout pct) ('D1 i) ('D1 i) where
  type Tape (Dropout pct) ('D1 i) ('D1 i) = R i
  runForwards (Dropout act seed) (S1D x)
    | not act = (v, S1D $ dvmap (rate *) x) -- multily with rate to normalise throughput
    | otherwise = (v, S1D $ v * x)
    where
      rate = (/100) $ fromIntegral $ max 0 $ min 100 $ natVal (Proxy :: Proxy pct)
      v = dvmap mask $ randomVector seed Uniform
      mask r
        | not act || r < rate = 1
        | otherwise = 0
  runBackwards (Dropout _ _) v (S1D x) = ((), S1D $ x * v)

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade


text:
module Legacy.Base {-# DEPRECATED "Use Vector instead" #-} where

import Data.Complex
import Data.Either
import Data.List
import Control.Applicative

newtype StateVector = StateVector [Complex Double]
    deriving (Eq, Show)

newtype TransformTensor = TransformTensor [[Complex Double]]
    deriving (Eq, Show)

{-|
    integerLog2 determines if an integer is a power of 2
    If it is, the log base 2 of the integer will be wrapped
    by Right constructor; if not, Left constructor with
    corresponding error will be returned.
    The function takes 1 argument of type 'Int'
-}
integerLog2 :: Int -> Either String Int
integerLog2 x
    | x < 1 = Left "Input less than 0"
    | x == 1 = Right 0
    | y * 2 == x = (+1) <$> integerLog2 y
    | otherwise = Left "Input is not a power of 2"
    where
        y = quot x 2

initQubit :: Complex Double -> Complex Double -> Either String StateVector
initQubit a b
    | magnitude a ** 2 + magnitude b ** 2 == 1 = Right $ StateVector [a, b]
    | otherwise = Left "Magnitude Error"

initQubit0 :: StateVector
initQubit0 = StateVector [1:+0, 0:+0]

initNumQubits :: Int -> Either String StateVector
initNumQubits n
    | n < 1 = Left "Number of Qubits Must be Greater than 1"
    | otherwise = Right $ foldr mergeState (StateVector [1 :+ 0]) $ replicate n initQubit0

initState :: [Complex Double] -> Either String StateVector
initState q
    | isLeft numQubits = Left "Input State Dimension Invalid"
    | otherwise = Right $ StateVector q
    where
        numQubits = integerLog2 $ length q

mergeState :: StateVector -> StateVector -> StateVector
mergeState (StateVector x) (StateVector y) = StateVector $ liftA2 (*) x y

mergeTensor :: TransformTensor -> TransformTensor -> TransformTensor
mergeTensor (TransformTensor x) (TransformTensor y) = TransformTensor $ fmap ((<*>) . fmap (*)) x <*> y

tensorAdd :: TransformTensor -> TransformTensor -> Either String TransformTensor
tensorAdd (TransformTensor i) (TransformTensor j)
    | m == x && n == y = Right $ TransformTensor $ zipWith (zipWith (+)) i j
    | otherwise = Left "Adding Two Tensors with Different Dimensions"
    where
        m = length i
        n = length $ head i
        x = length j
        y = length $ head j

doubleBindEither :: Either String TransformTensor -> Either String TransformTensor -> (TransformTensor -> TransformTensor -> Either String TransformTensor) -> Either String TransformTensor
doubleBindEither x y f
    | isLeft x && isRight y = Left $ fromLeft [] x ++ "\n" ++ fromLeft [] y
    | isLeft x = x
    | isLeft y = y
    | otherwise = f (fromRight (TransformTensor [[]]) x) (fromRight (TransformTensor [[]]) y)

applyTensorToState :: TransformTensor -> StateVector -> Either String StateVector
applyTensorToState (TransformTensor g) (StateVector q)
    | length f == length q && length q == length g =
        Right $ StateVector $ foldr (zipWith (+)) acc $ zipWith fmap (fmap (*) q) f
    | otherwise = Left "Transformation Matrix Dimension Mismatch"
    where
        acc = replicate (length q) $0:+0
        f = transpose g

-- The qubits are 1-indexed
applyGateToQubit :: TransformTensor -> Int -> StateVector -> Either String StateVector
applyGateToQubit g n (StateVector q) = t >>= (`applyTensorToState` StateVector q)
    where
        numQubits = integerLog2 $ length q
        t = numQubits >>= flip (genTensor g) n

-- Gate to insert -> Number of Qubits -> Index of Qubit
genTensor :: TransformTensor -> Int -> Int -> Either String TransformTensor
genTensor g n i
    | n < i || i < 1 = Left "Index Out of Range"
    | otherwise = Right $ foldr mergeTensor (TransformTensor [[1:+0]]) expGate
    where
        expGate = fstId ++ [g] ++ lstId
        fstId = replicate (i - 1) pauliId
        lstId = replicate (n - i) pauliId

getQubitState :: Int -> StateVector -> Either String (Double, Double)
getQubitState x (StateVector q)
    | x < 1 = Left "Index Out of Bounds"
    | otherwise = Right (zeroP, oneP)
    where
        numRow = 2 ^ (x - 1)
        q2 = fmap (** 2) q
        zeroP = magnitude $ sum $ takeHalf q2 numRow
        oneP = magnitude $ sum $ takeHalf (drop numRow q2) numRow
        takeHalf :: [Complex Double] -> Int -> [Complex Double]
        takeHalf xs n
            | length xs < n = []
            | otherwise = take n xs ++ takeHalf (drop (2 * n) xs) n

-- Gate to Apply -> Control Bit Index -> Application Bit Index -> Circuit To Apply
applyControlToQubits :: TransformTensor -> Int -> Int -> StateVector -> Either String StateVector
applyControlToQubits g m n (StateVector q)
    | m < n = t m n mask1 g >>= (`applyTensorToState` StateVector q)
    | m > n = t n m g mask1 >>= (`applyTensorToState` StateVector q)
    | otherwise = Left "Control and Application Qubit Cannot be Identical"
    where
        zeroMask = numQubits >>= flip (genTensor mask0) m
        numQubits = integerLog2 $ length q
        t :: Int -> Int -> TransformTensor -> TransformTensor -> Either String TransformTensor
        t i j fstGate sndGate = doubleBindEither (fmap mergeTensor fstTensor <*> sndTensor) zeroMask tensorAdd
            where
                fstTensor = genTensor fstGate i i
                sndTensor = numQubits >>= flip (genTensor sndGate) (j - i) . flip (-) i

pauliX :: TransformTensor
pauliX = TransformTensor [[0 :+ 0, 1 :+ 0],
                      [1 :+ 0, 0 :+ 0]]

pauliY :: TransformTensor
pauliY = TransformTensor [[0 :+ 0, 0 :+ (-1)],
                      [0 :+ 1, 0 :+ 0]]

pauliZ :: TransformTensor
pauliZ = TransformTensor [[1 :+ 0, 0 :+ 0],
                      [0 :+ 0, (-1) :+ 0]]

pauliId :: TransformTensor
pauliId = TransformTensor [[1 :+ 0, 0 :+ 0],
                      [0 :+ 0, 1 :+ 0]]

hadamard :: TransformTensor
hadamard = TransformTensor [[1 / sqrt 2 :+ 0, 1 / sqrt 2 :+ 0],
                        [1 / sqrt 2 :+ 0, (-1) / sqrt 2 :+ 0]]

mask0 :: TransformTensor
mask0 = TransformTensor [[1 :+ 0, 0 :+ 0],
                    [0 :+ 0, 0 :+ 0]]

mask1 :: TransformTensor
mask1 = TransformTensor [[0 :+ 0, 0 :+ 0],
                    [0 :+ 0, 1 :+ 0]]

-- main :: IO()
-- main = print result
--     where q0 = initQubit0
--           crt = mergeState q0 q0
--           flip1 = applyGateToQubit hadamard 1 crt
--           flip2 = flip1 >>= applyControlToQubits pauliY 1 2
--           --result = flip2 >>= (getQubitState 2)
--           result = TransformTensor [[0 :+ 0, 0 :+ 0],
--                                     [0 :+ 0, 1 :+ 0]]

-- Linear algebra versions

-- safeMultiply :: Num a => T.Tensor a -> T.Tensor a -> Either String (T.Tensor a)
-- safeMultiply (T.Tensor x) (T.Tensor y)
--     | T.isSquare (T.Tensor x)
--         && T.isSquare (T.Tensor y)
--         && length x == length y = Right $ T.Tensor x * T.Tensor y
--     | otherwise = Left "Dimension mismatch"
-- safeMultiply x y = Right $ x * y

-- -- combine two qubit circuits
-- addQubit :: Num a => T.Tensor a -- ^ first half statevector
--   -> T.Tensor a -- ^ second half statevector
--   -> T.Tensor a
-- addQubit (T.Tensor x) (T.Tensor y)
--     | T.rank (T.Tensor x) /= 1 || T.rank (T.Tensor y) /= 1 = error "One of the elements is not a valid statevector"
--     | otherwise = T.Tensor $ liftA2 (*) x y
-- addQubit _ _ = error "One of the elements is not a valid statevector"

-- Applies input gate to qubit of index

-- -- Builds a transformation tensor to apply input gate to selected qubit
-- build :: RealFloat a => Int -- ^ Number of qubits this tensor applies to
--   -> [(Int, T.Tensor (Complex a))]-- ^ Index / Gate pair (Pre-sorted)
--   -> T.Tensor (Complex a)-- ^ Default Gate
--   -> T.Tensor (Complex a)
-- build len lst dft = T.flatten2d $ foldr T.product (T.fromList [1]) (go len lst 0)
--     where
--         go l (x:xs) offset
--             | fst x >= l = go l [] offset
--             | otherwise = replicate (fst x - offset) dft ++ [snd x] ++ go l xs (fst x + 1)
--         go l [] offset = replicate (l - offset) dft

-- applyGate :: RealFloat a => T.Tensor (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Index
--   -> T.Tensor (Complex a) -- ^ State vector
--   -> Either String (T.Tensor (Complex a))
-- applyGate (T.Tensor xs) i (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | i < 0 || i >= numV = Left "Index out of bounds"
--     | otherwise = Right $ build numV [(i, x)] pauliId * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)
-- applyGate _ _ _ = Left "Input is not a Tensor"

-- applyGateAll :: RealFloat a => T.Tensor (Complex a) -> T.Tensor (Complex a) -> Either String (T.Tensor (Complex a))
-- applyGateAll (T.Tensor xs) (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | otherwise = Right $ build numV [] x * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)
-- applyGateAll _ _ = Left "Input is not a Tensor"

-- applyControl :: RealFloat a => T.Tensor (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Control qubit index
--   -> Int -- ^ Apply qubit index
--   -> T.Tensor (Complex a) -- ^ State vector
--   -> Either String (T.Tensor (Complex a))
-- applyControl (T.Tensor xs) ctl i (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | i < 0 || i >= numV = Left "Application qubit index out of bounds"
--     | ctl < 0 || ctl >= numV = Left "Control qubit index out of bounds"
--     | ctl == i = Left "Control and application qubit cannot be the same"
--     | otherwise = Right $ (build numV [(ctl, mask0)] pauliId
--         + build numV (sortBy (compare `on` fst) [(ctl, mask1), (i, x)]) pauliId) * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)

-- applyControl _ _ _ _ = Left "Input is not a Tensor"


text:
{-# LANGUAGE GADTs, DataKinds, KindSignatures, ScopedTypeVariables, TypeApplications, TypeOperators, RankNTypes, FlexibleContexts #-}

import Numeric.LinearAlgebra.Static
import GHC.TypeLits
import Data.Proxy
import Data.Maybe (mapMaybe)
import Data.Monoid (getSum, Sum(..))
import Data.Coerce
import MLP.Types (Pattern(..))
import MLP.Network (Net(..), AllCon, ArrList(..), Weights(..), MLP(..), Learn(..), fromScalarR, netUpd)
import Data.Singletons

data EncPat where
    EncPat :: forall i o. (KnownNat i, KnownNat o) => Pattern i o -> EncPat

data EncNet where
    EncNet :: forall i hs o. (SingI (Topo (Net i hs o)), Learn (Net i hs o), AllCon KnownNat (i ': o ': hs), MLP (Net i hs o)) => Net i hs o -> EncNet

readMat :: IO (EncPat, EncNet)
readMat = do
    i <- read <$> (putStr "inputs: " >> getLine)
    o <- read <$> (putStr "outputs: " >> getLine)

    let Just someNatI = someNatVal (ceiling i)
        Just someNatO = someNatVal (ceiling o)

    case someNatI of
       SomeNat (_ :: Proxy i) ->
         case someNatO of
            SomeNat (_ :: Proxy o) ->
               let n = EncNet . SLayer $ Weights (matrix @o @i [1..(i * o)] ) (vector @o [1..o])
                   p = EncPat $ Pattern (vector @i [1..i]) (vector @o [1..o])
                 in pure (p,n)

io :: forall i o hs. (KnownNat i, KnownNat o) => Net i hs o -> [Integer]
io _ = map fromIntegral . fromSing $ sing @'[i, o]

main :: IO ()
main = do
    (EncPat (Pattern i o), EncNet n) <- readMat
    let arrl = mkDeltas i o n

    case netUpd 1.0 i arrl n of
        SLayer l -> print l
        _ -> print ""

max_stars_repo_name: ovanr/Typed-MLP

max_issues_repo_name: ovanr/Typed-MLP

max_forks_repo_name: ovanr/Typed-MLP


text:
{-# LANGUAGE RecordWildCards, ScopedTypeVariables #-}
module Data.Eigen.Matrix.Mutable (
    MMatrix(..),
    MMatrixXf,
    MMatrixXd,
    MMatrixXcf,
    MMatrixXcd,
    IOMatrix,
    STMatrix,
    -- * Construction
    new,
    replicate,
    -- * Consistency check
    valid,
    -- * Accessing individual elements
    read,
    write,
    unsafeRead,
    unsafeWrite,
    -- * Modifying matrices
    set,
    copy,
    unsafeCopy,
    -- * Raw pointers
    unsafeWith
) where

import Prelude hiding (read, replicate)
import Control.Monad.Primitive
import Foreign.Ptr
import Foreign.C.Types
import Data.Complex
import Text.Printf
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.Eigen.Internal as I

-- | Mutable matrix. You can modify elements
data MMatrix a b s = MMatrix {
    mm_rows :: Int,
    mm_cols :: Int,
    mm_vals :: VSM.MVector s b
}

-- | Alias for single precision mutable matrix
type MMatrixXf = MMatrix Float CFloat
-- | Alias for double precision mutable matrix
type MMatrixXd = MMatrix Double CDouble
-- | Alias for single previsiom mutable matrix of complex numbers
type MMatrixXcf = MMatrix (Complex Float) (I.CComplex CFloat)
-- | Alias for double prevision mutable matrix of complex numbers
type MMatrixXcd = MMatrix (Complex Double) (I.CComplex CDouble)

type IOMatrix a b = MMatrix a b RealWorld
type STMatrix a b s = MMatrix a b s

-- | Verify matrix dimensions and memory layout
valid :: I.Elem a b => MMatrix a b s -> Bool
valid MMatrix{..} = mm_rows >= 0 && mm_cols >= 0 && VSM.length mm_vals == mm_rows * mm_cols

-- | Create a mutable matrix of the given size and fill it with 0 as an initial value.
new :: (PrimMonad m, I.Elem a b) => Int -> Int -> m (MMatrix a b (PrimState m))
new rows cols = replicate rows cols 0

-- | Create a mutable matrix of the given size and fill it with as an initial value.
replicate :: (PrimMonad m, I.Elem a b) => Int -> Int -> a -> m (MMatrix a b (PrimState m))
replicate rows cols val = do
    vals <- VSM.replicate (rows * cols) (I.cast val)
    return $ MMatrix rows cols vals

-- | Set all elements of the matrix to the given value
set :: (PrimMonad m, I.Elem a b) => (MMatrix a b (PrimState m)) -> a -> m ()
set MMatrix{..} val = VSM.set mm_vals (I.cast val)

-- | Copy a matrix. The two matrices must have the same size and may not overlap.
copy :: (PrimMonad m, I.Elem a b) => (MMatrix a b (PrimState m)) -> (MMatrix a b (PrimState m)) -> m ()
copy m1 m2
    | not (valid m1) = fail "MMatrix.copy: lhs matrix layout is invalid"
    | not (valid m2) = fail "MMatrix.copy: rhs matrix layout is invalid"
    | mm_rows m1 /= mm_rows m2 = fail "MMatrix.copy: matrices have different number of cols"
    | mm_cols m1 /= mm_cols m2 = fail "MMatrix.copy: matrices have different number of rows"
    | otherwise = VSM.copy (mm_vals m1) (mm_vals m2)

-- | Yield the element at the given position.
read :: (PrimMonad m, I.Elem a b) => MMatrix a b (PrimState m) -> Int -> Int -> m a
read mm@MMatrix{..} row col
    | not (valid mm) = fail "MMatrix.read: matrix layout is invalid"
    | row < 0 || row >= mm_rows = fail $ printf "MMatrix.read: row %d is out of bounds [0..%d)" row mm_rows
    | col < 0 || col >= mm_cols = fail $ printf "MMatrix.read: col %d is out of bounds [0..%d)" col mm_cols
    | otherwise = unsafeRead mm row col

-- | Replace the element at the given position.
write :: (PrimMonad m, I.Elem a b) => MMatrix a b (PrimState m) -> Int -> Int -> a -> m ()
write mm@MMatrix{..} row col val
    | not (valid mm) = fail "MMatrix.write: matrix layout is invalid"
    | row < 0 || row >= mm_rows = fail $ printf "MMatrix.write: row %d is out of bounds [0..%d)" row mm_rows
    | col < 0 || col >= mm_cols = fail $ printf "MMatrix.write: col %d is out of bounds [0..%d)" col mm_cols
    | otherwise = unsafeWrite mm row col val

-- | Copy a matrix. The two matrices must have the same size and may not overlap however no bounds check performaned to it may SEGFAULT for incorrect input.
unsafeCopy :: (PrimMonad m, I.Elem a b) => (MMatrix a b (PrimState m)) -> (MMatrix a b (PrimState m)) -> m ()
unsafeCopy m1 m2 = VSM.unsafeCopy (mm_vals m1) (mm_vals m2)

-- | Yield the element at the given position. No bounds checks are performed.
unsafeRead :: (PrimMonad m, I.Elem a b) => MMatrix a b (PrimState m) -> Int -> Int -> m a
unsafeRead MMatrix{..} row col = VSM.unsafeRead mm_vals (col * mm_rows + row) >>= \val -> return (I.cast val)

-- | Replace the element at the given position. No bounds checks are performed.
unsafeWrite :: (PrimMonad m, I.Elem a b) => MMatrix a b (PrimState m) -> Int -> Int -> a -> m ()
unsafeWrite MMatrix{..} row col val = VSM.unsafeWrite mm_vals (col * mm_rows + row) (I.cast val)

-- | Pass a pointer to the matrix's data to the IO action. Modifying data through the pointer is unsafe if the matrix could have been frozen before the modification.
unsafeWith :: I.Elem a b => IOMatrix a b -> (Ptr b -> CInt -> CInt -> IO c) -> IO c
unsafeWith mm@MMatrix{..} f
    | not (valid mm) = fail "mutable matrix layout is invalid"
    | otherwise = VSM.unsafeWith mm_vals $ \p -> f p (I.cast mm_rows) (I.cast mm_cols)

max_stars_repo_name: osidorkin/haskell-eigen

max_issues_repo_name: osidorkin/haskell-eigen

max_forks_repo_name: osidorkin/haskell-eigen


text:
{-# LANGUAGE ApplicativeDo         #-}
{-# LANGUAGE EmptyCase             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module TensorOps.Backend.BTensor
  ( BTensor
  , BTensorL, BTensorV
  , HMat
  , HMatD
  ) where

import           Control.Applicative
import           Control.DeepSeq
import           Data.Distributive
import           Data.Kind
import           Data.List.Util
import           Data.Monoid
import           Data.Nested hiding             (unScalar, unVector, gmul')
import           Data.Singletons
import           Data.Singletons.Prelude hiding (Reverse, Head, sReverse, (:-))
import           Data.Type.Combinator
import           Data.Type.Combinator.Util
import           Data.Type.Length               as TCL
import           Data.Type.Length.Util          as TCL
import           Data.Type.Nat
import           Data.Type.Product              as TCP
import           Data.Type.Product.Util         as TCP
import           Data.Type.Sing
import           Data.Type.Uniform
import           Statistics.Distribution
import           TensorOps.BLAS
import           TensorOps.BLAS.HMat
import           TensorOps.NatKind
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Witness
import           Type.Family.List
import           Type.Family.List.Util
import           Type.Family.Nat
import qualified Data.Type.Vector               as TCV
import qualified Data.Type.Vector.Util          as TCV
import qualified Data.Vector.Sized              as VS

data BTensor :: (k -> Type -> Type) -> (BShape k -> Type) -> [k] -> Type where
    BTS :: { unScalar :: !(ElemB b)     } -> BTensor v b '[]
    BTV :: { unVector :: !(b ('BV n))   } -> BTensor v b '[n]
    BTM :: { unMatrix :: !(b ('BM n m)) } -> BTensor v b '[n,m]
    BTN :: { unNested :: !(v n (BTensor v b (o ': m ': ns))) }
        -> BTensor v b (n ': o ': m ': ns)

type BTensorL = BTensor (Flip2 TCV.VecT   I)
type BTensorV = BTensor (Flip2 VS.VectorT I)

instance (Nesting Proxy Show v, Show1 b, Show (ElemB b)) => Show (BTensor v b s) where
    showsPrec p = \case
      BTS x  -> showParen (p > 10) $ showString "BTS "
                                   . showsPrec 11 x
      BTV xs -> showParen (p > 10) $ showString "BTV "
                                   . showsPrec1 11 xs
      BTM xs -> showParen (p > 10) $ showString "BTM "
                                   . showsPrec1 11 xs
      BTN xs -> showParen (p > 10) $ showString "BTN "
                                   . showsPrec' 11 xs
      where
        showsPrec' :: forall n s'. Int -> v n (BTensor v b s') -> ShowS
        showsPrec' p' xs = showsPrec p' xs
            \\ (nesting Proxy :: Show (BTensor v b s')
                              :- Show (v n (BTensor v b s'))
               )

instance (Nesting Proxy Show v, Show1 b, Show (ElemB b)) => Show1 (BTensor v b)

instance (NFData (ElemB b), NFData1 b, Nesting Proxy NFData v) => Nesting1 Proxy NFData (BTensor v b) where
    nesting1 _ = Wit

instance (NFData (ElemB b), NFData1 b, Nesting Proxy NFData v) => NFData (BTensor v b js) where
    rnf = \case
      BTS x  -> rnf  x
      BTV xs -> rnf1 xs
      BTM xs -> rnf1 xs
      BTN (xs :: v n (BTensor v b (o ': m ': ns))) ->
        rnf xs \\ (nesting Proxy :: NFData (BTensor v b (o ': m ': ns))
                                 :- NFData (v n (BTensor v b (o ': m ': ns)))
                  )

instance (NFData (ElemB b), NFData1 b, Nesting Proxy NFData v) => NFData1 (BTensor v b)

instance ( BLAS b
         , Vec v
         , Nesting1 Proxy Functor v
         , Nesting1 Sing Applicative v
         , SingI ns
         , Num (ElemB b)
         )
        => Num (BTensor v b ns) where
    (+) = zipBase sing
                  (+)
                  (\_          xs ys -> axpy 1 xs (Just ys))
                  (\(SBM _ sM) xs ys -> gemm 1 xs (eye sM) (Just (1, ys)))
    {-# INLINE (+) #-}
    (-) = zipBase sing
                  (-)
                  (\_          xs ys -> axpy (-1) ys (Just xs))
                  (\(SBM _ sM) xs ys -> gemm 1 xs (eye sM) (Just (-1, ys)))
    {-# INLINE (-) #-}
    (*) = zipBTensorElems sing (*)
    {-# INLINE (*) #-}
    negate = mapBase sing
                     negate
                     (\_          xs -> axpy (-1) xs Nothing)
                     (\(SBM _ sM) xs -> gemm 1 xs (eye sM) Nothing)
    {-# INLINE negate #-}
    abs    = mapBTensorElems abs
    {-# INLINE abs #-}
    signum = mapBTensorElems signum
    {-# INLINE signum #-}
    fromInteger i = genBTensor sing $ \_ -> fromInteger i
    {-# INLINE fromInteger #-}

-- | TODO: add RULES pragmas so that this can be done without checking
-- lengths at runtime in the common case that the lengths are known at
-- compile-time.
--
-- Also, totally forgot about matrix-scalar multiplication here, but there
-- isn't really any way of making it work without a lot of empty cases.
-- should probably handle one level up.
dispatchBLAS
    :: forall b ms os ns v. (RealFloat (ElemB b), BLAS b)
    => MaxLength N1 ms
    -> MaxLength N1 os
    -> MaxLength N1 ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
dispatchBLAS lM lO lN v r = case (lM, lO, lN) of
    (MLZ    , MLZ    , MLZ    ) -> case (v, r) of
      -- scalar-scalar
      (BTS x, BTS y) -> BTS $ x * y
    (MLZ    , MLZ    , MLS MLZ) -> case (v, r) of
      -- scalar-vector
      (BTS x, BTV y) -> BTV $ axpy x y Nothing
    (MLZ    , MLS MLZ, MLZ    ) -> case (v, r) of
      -- dot
      (BTV x, BTV y) -> BTS $ x `dot` y
    (MLZ    , MLS MLZ, MLS MLZ) -> case (v, r) of
      -- vector-matrix
      -- TODO: transpose?
      (BTV x, BTM y) -> BTV $ gemv 1 (transpB y) x Nothing
    (MLS MLZ, MLZ    , MLZ    ) -> case (v, r) of
      -- vector-scalar
      (BTV x, BTS y) -> BTV $ axpy y x Nothing
    (MLS MLZ, MLZ    , MLS MLZ) -> case (v, r) of
      -- vector-scalar
      (BTV x, BTV y) -> BTM $ ger x y
    (MLS MLZ, MLS MLZ, MLZ    ) -> case (v, r) of
      -- matrx-vector
      (BTM x, BTV y) -> BTV $ gemv 1 x y Nothing
    (MLS MLZ, MLS MLZ, MLS MLZ) -> case (v, r) of
      -- matrix-matrix
      (BTM x, BTM y) -> BTM $ gemm 1 x y Nothing
{-# INLINE dispatchBLAS #-}

mapRowsBTensor
    :: forall k (v :: k -> Type -> Type) ns ms os b. (Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (BTensor v b ms -> BTensor v b os)
    -> BTensor v b (ns ++ ms)
    -> BTensor v b (ns ++ os)
mapRowsBTensor sN lO f = getI . bRows sN lO (I . f)
{-# INLINE mapRowsBTensor #-}

bRows
    :: forall k (v :: k -> Type -> Type) ns ms os b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (BTensor v b ms -> f (BTensor v b os))
    -> BTensor v b (ns ++ ms)
    -> f (BTensor v b (ns ++ os))
bRows sN lO f = bIxRows sN lO (\_ -> f)
{-# INLINE bRows #-}

mapIxRows
    :: forall k (v :: k -> Type -> Type) ns ms os b. (Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (Prod (IndexN k) ns -> BTensor v b ms -> BTensor v b os)
    -> BTensor v b (ns ++ ms)
    -> BTensor v b (ns ++ os)
mapIxRows sN lO f = getI . bIxRows sN lO (\i -> I . f i)
{-# INLINE mapIxRows #-}

foldMapIxRows
    :: forall k (v :: k -> Type -> Type) ns ms m b. (Vec v, Monoid m, BLAS b)
    => Sing ns
    -> (Prod (IndexN k) ns -> BTensor v b ms -> m)
    -> BTensor v b (ns ++ ms)
    -> m
foldMapIxRows s f = getConst . bIxRows s LZ (\i -> Const . f i)
{-# INLINE foldMapIxRows #-}

bIxRows
    :: forall k (v :: k -> Type -> Type) ns ms os b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (Prod (IndexN k) ns -> BTensor v b ms -> f (BTensor v b os))
    -> BTensor v b (ns ++ ms)
    -> f (BTensor v b (ns ++ os))
bIxRows = \case
    SNil   -> \_  f -> f Ø
    s `SCons` ss -> \lO f -> \case
      BTV xs -> case ss of
        -- ns ~ '[n]
        -- ms ~ '[]
        SNil -> case lO of
          -- ns ++ os ~ '[n]
          LZ        -> BTV <$> iElemsB (\i -> fmap unScalar . f (pbvProd i) . BTS) xs
          -- ns ++ os ~ '[n,m]
          LS LZ     -> BTM <$> bgenRowsA (\i -> unVector <$> f (i :< Ø) (BTS $ indexB (PBV i) xs))
                         \\ s
          LS (LS _) -> BTN <$> vGenA s (\i -> f (i :< Ø) (BTS $ indexB (PBV i) xs))
      BTM xs -> case ss of
        -- ns ~ '[n]
        -- ms ~ '[m]
        SNil -> case lO of
          -- ns ++ os ~ '[n]
          LZ        -> BTV <$> bgenA (SBV s) (\(PBV i) -> unScalar <$> f (i :< Ø) (BTV (indexRowB i xs)))
          -- ns ++ os ~ '[n,o]
          LS LZ     -> BTM <$> iRowsB (\i -> fmap unVector . f (i :< Ø) . BTV) xs
          LS (LS _) -> BTN <$> vGenA s (\i -> f (i :< Ø) (BTV (indexRowB i xs)))
        -- ns ~ '[n,m]
        -- ms ~ '[]
        s' `SCons` ss' -> (\\ s') $ case ss' of
          SNil -> case lO of
            LZ   -> BTM <$> iElemsB (\i -> fmap unScalar . f (pbmProd i) . BTS) xs
            LS _ -> BTN <$>
                      vGenA s (\i ->
                          btn lO <$>
                            vGenA s' (\j ->
                                f (i :< j :< Ø) (BTS (indexB (PBM i j) xs))
                              )
                        )
      BTN xs -> (\\ s) $
          fmap (btn (singLength ss `TCL.append'` lO))
        . vITraverse (\i -> bIxRows ss lO (\is -> f (i :< is)))
        $ xs

indexRowBTensor
    :: forall k (b :: BShape k -> Type) v ns ms.
     ( BLAS b
     , Vec v
     )
    => Prod (IndexN k) ns
    -> BTensor v b (ns ++ ms)
    -> BTensor v b ms
indexRowBTensor = \case
    Ø       -> id
    i :< is -> \case
      BTV xs -> case is of
        Ø      -> BTS $ indexB    (PBV i)   xs
      BTM xs -> case is of
        Ø      -> BTV $ indexRowB i         xs
        j :< Ø -> BTS $ indexB    (PBM i j) xs
      BTN xs -> indexRowBTensor is (vIndex i xs)
{-# INLINE indexRowBTensor #-}

mapBTensorElems
    :: (Vec v, BLAS b)
    => (ElemB b -> ElemB b)
    -> BTensor v b ns
    -> BTensor v b ns
mapBTensorElems f = getI . bTensorElems (I . f)
{-# INLINE mapBTensorElems #-}

bTensorElems
    :: forall k (v :: k -> Type -> Type) ns b f. (Applicative f, Vec v, BLAS b)
    => (ElemB b -> f (ElemB b))
    -> BTensor v b ns
    -> f (BTensor v b ns)
bTensorElems f = \case
    BTS x  -> BTS <$> f x
    BTV xs -> BTV <$> elemsB f xs
    BTM xs -> BTM <$> elemsB f xs
    BTN xs -> BTN <$> vITraverse (\_ x -> bTensorElems f x) xs
{-# INLINE bTensorElems #-}

ifoldMapBTensor
    :: forall k (v :: k -> Type -> Type) ns m b. (Monoid m, Vec v, BLAS b)
    => (Prod (IndexN k) ns -> ElemB b -> m)
    -> BTensor v b ns
    -> m
ifoldMapBTensor f = getConst . bTensorIxElems (\i -> Const . f i)
{-# INLINE ifoldMapBTensor #-}

bTensorIxElems
    :: forall k (v :: k -> Type -> Type) ns b f. (Applicative f, Vec v, BLAS b)
    => (Prod (IndexN k) ns -> ElemB b -> f (ElemB b))
    -> BTensor v b ns
    -> f (BTensor v b ns)
bTensorIxElems f = \case
    BTS x  -> BTS <$> f Ø x
    BTV xs -> BTV <$> iElemsB (f . pbvProd) xs
    BTM xs -> BTM <$> iElemsB (f . pbmProd) xs
    BTN xs -> BTN <$> vITraverse (\i -> bTensorIxElems (\is -> f (i :< is))) xs
{-# INLINE bTensorIxElems #-}

zipBTensorElems
    :: forall v b ns. (BLAS b, Nesting1 Sing Applicative v)
    => Sing ns
    -> (ElemB b -> ElemB b -> ElemB b)
    -> BTensor v b ns
    -> BTensor v b ns
    -> BTensor v b ns
zipBTensorElems = \case
    SNil -> \f -> \case
      BTS x -> \case
        BTS y -> BTS (f x y)
    sN `SCons` SNil -> \f -> \case
      BTV xs -> \case
        BTV ys -> BTV (zipB (SBV sN) f xs ys)
    sN `SCons` (sM `SCons` SNil) -> \f -> \case
      BTM xs -> \case
        BTM ys -> BTM (zipB (SBM sN sM) f xs ys)
    (s :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f -> \case
      BTN xs -> \case
        BTN ys -> BTN (zipBTensorElems ss f <$> xs <*> ys)
                    \\ (nesting1 s :: Wit (Applicative (v k)))
{-# INLINE zipBTensorElems #-}

liftBTensor
    :: forall v b ns n.
     ( BLAS b
     , Nesting1 Proxy Functor      v
     , Nesting1 Sing  Distributive v
     )
    => Sing ns
    -> (TCV.Vec n (ElemB b) -> ElemB b)
    -> TCV.Vec n (BTensor v b ns)
    -> BTensor v b ns
liftBTensor = \case
    SNil                         -> \f xs ->
        let xs' = unScalar <$> xs
        in  BTS $ f xs'
    sN `SCons` SNil              -> \f xs ->
        let xs' = unVector <$> xs
        in  BTV $ liftB (SBV sN) f xs'
    sN `SCons` (sM `SCons` SNil) -> \f xs ->
        let xs' = unMatrix <$> xs
        in  BTM $ liftB (SBM sN sM) f xs'
    (s :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f xs ->
        let xs' = unNested <$> xs
        in  BTN $ TCV.liftVecD (liftBTensor ss f) xs'
              \\ (nesting1 s     :: Wit (Distributive (v k)))
{-# INLINE liftBTensor #-}

mapBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b. (Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (b ('BM n m) -> BTensor v b ms)
    -> BTensor v b (ns ++ [n,m])
    -> BTensor v b (ns ++ ms)
mapBTM sN lM f = getI . traverseBTM sN lM (I . f)
{-# INLINE mapBTM #-}

foldMapBTM
    :: (Monoid a, Vec v, BLAS b)
    => Length ns
    -> (b ('BM n m) -> a)
    -> BTensor v b (ns ++ [n,m])
    -> a
foldMapBTM l f = ifoldMapBTM l (\_ -> f)
{-# INLINE foldMapBTM #-}

traverseBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (b ('BM n m) -> f (BTensor v b ms))
    -> BTensor v b (ns ++ [n,m])
    -> f (BTensor v b (ns ++ ms))
traverseBTM = \case
    SNil         -> \_ f -> \case
      BTM x  -> f x
    s `SCons` ss -> \lM f -> \case
      BTV _  -> case ss of
      BTM _  -> case ss of
      BTN xs -> (\\ s) $
          fmap (btn (singLength ss `TCL.append'` lM))
        . vITraverse (\_ -> traverseBTM ss lM f)
        $ xs
{-# INLINE traverseBTM #-}

imapBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b. (Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (Prod (IndexN k) ns -> b ('BM n m) -> BTensor v b ms)
    -> BTensor v b (ns ++ [n,m])
    -> BTensor v b (ns ++ ms)
imapBTM sN lM f = getI . itraverseBTM sN lM (\i -> I . f i)
{-# INLINE imapBTM #-}

ifoldMapBTM
    :: (Vec v, Monoid a, BLAS b)
    => Length ns
    -> (Prod (IndexN k) ns -> b ('BM n m) -> a)
    -> BTensor v b (ns ++ [n,m])
    -> a
ifoldMapBTM = \case
    LZ -> \f -> \case
      BTM xs -> f Ø xs
    LS l -> \f -> \case
      BTV _  -> case l of
      BTM _  -> case l of
      BTN xs -> vIFoldMap (\i -> ifoldMapBTM l (\is -> f (i :< is))) xs
{-# INLINE ifoldMapBTM #-}

itraverseBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (Prod (IndexN k) ns -> b ('BM n m) -> f (BTensor v b ms))
    -> BTensor v b (ns ++ [n,m])
    -> f (BTensor v b (ns ++ ms))
itraverseBTM = \case
    SNil         -> \_ f -> \case
      BTM x  -> f Ø x
    s `SCons` ss -> \lM f -> \case
      BTV _  -> case ss of
      BTM _  -> case ss of
      BTN xs -> (\\ s) $
          fmap (btn (singLength ss `TCL.append'` lM))
        . vITraverse (\i -> itraverseBTM ss lM (\is ys -> f (i :< is) ys))
        $ xs
{-# INLINE itraverseBTM #-}

mapBase
    :: forall v b ns. (Nesting1 Proxy Functor v)
    => Sing ns
    -> (ElemB b -> ElemB b)
    -> (forall n. Sing n -> b ('BV n) -> b ('BV n))
    -> (forall n m. Sing ('BM n m) -> b ('BM n m) -> b ('BM n m))
    -> BTensor v b ns
    -> BTensor v b ns
mapBase = \case
    SNil -> \f _ _ -> \case
      BTS x  -> BTS (f x)
    sN `SCons` SNil -> \_ g _ -> \case
      BTV xs -> BTV (g sN          xs)
    sN `SCons` (sM `SCons` SNil) -> \_ _ h -> \case
      BTM xs -> BTM (h (SBM sN sM) xs)
    (_ :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f g h -> \case
      BTN xs -> BTN (mapBase ss f g h <$> xs)
                  \\ (nesting1 Proxy :: Wit (Functor (v k)))
{-# INLINE mapBase #-}

zipBase
    :: forall v b ns. (Nesting1 Sing Applicative v)
    => Sing ns
    -> (ElemB b -> ElemB b -> ElemB b)
    -> (forall n. Sing n -> b ('BV n) -> b ('BV n) -> b ('BV n))
    -> (forall n m. Sing ('BM n m) -> b ('BM n m) -> b ('BM n m) -> b ('BM n m))
    -> BTensor v b ns
    -> BTensor v b ns
    -> BTensor v b ns
zipBase = \case
    SNil -> \f _ _ -> \case
      BTS x -> \case
        BTS y -> BTS (f x y)
    sN `SCons` SNil -> \_ g _ -> \case
      BTV xs -> \case
        BTV ys -> BTV (g sN          xs ys)
    sN `SCons` (sM `SCons` SNil) -> \_ _ h -> \case
      BTM xs -> \case
        BTM ys -> BTM (h (SBM sN sM) xs ys)
    (s :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f g h -> \case
      BTN xs -> \case
        BTN ys -> BTN $ zipBase ss f g h <$> xs <*> ys
                    \\ (nesting1 s :: Wit (Applicative (v k)))
{-# INLINE zipBase #-}

genBTensorA
    :: forall k (b :: BShape k -> Type) v (ns :: [k]) f. (Applicative f, BLAS b, Vec v)
    => Sing ns
    -> (Prod (IndexN k) ns -> f (ElemB b))
    -> f (BTensor v b ns)
genBTensorA = \case
    SNil                                   -> \f ->
        BTS <$> f Ø
    sN `SCons` SNil                        -> \f ->
        BTV <$> bgenA (SBV sN)    (f . pbvProd)
    sN `SCons` (sM `SCons` SNil)           -> \f ->
        BTM <$> bgenA (SBM sN sM) (f . pbmProd)
    s `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f ->
        BTN <$> vGenA s (\i -> genBTensorA ss (\is -> f (i :< is)))
{-# INLINE genBTensorA #-}

genBTensor
    :: forall k (b :: BShape k -> Type) v (ns :: [k]). (BLAS b, Vec v)
    => Sing ns
    -> (Prod (IndexN k) ns -> ElemB b)
    -> BTensor v b ns
genBTensor s f = getI $ genBTensorA s (I . f)
{-# INLINE genBTensor #-}

indexBTensor
    :: forall k (b :: BShape k -> Type) v ns. (BLAS b, Vec v)
    => Prod (IndexN k) ns
    -> BTensor v b ns
    -> ElemB b
indexBTensor = \case
    Ø      -> \case
      BTS x  -> x
    i :< Ø -> \case
      BTV xs -> indexB (PBV i) xs
    i :< j :< Ø -> \case
      BTM xs -> indexB (PBM i j) xs
    i :< js@(_ :< _ :< _) -> \case
      BTN xs -> indexBTensor js (vIndex i xs)
{-# INLINE indexBTensor #-}

btn :: (BLAS b, Vec v, SingI n)
    => Length ns
    -> v n (BTensor v b ns)
    -> BTensor v b (n ': ns)
btn = \case
    LZ        -> \xs ->
      BTV $ bgen sing (unScalar . (`vIndex` xs) . unPBV)
    LS LZ     -> \xs ->
      BTM $ bgenRows  (unVector . (`vIndex` xs))
    LS (LS _) -> BTN
{-# INLINE btn #-}

gmul'
    :: forall v b ms os ns.
     ( SingI (ms ++ ns)
     , RealFloat (ElemB b)
     , Vec v
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     , BLAS b
     )
    => Length ms
    -> Length os
    -> Length ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
gmul' lM lO lN = gmulB sM lO lN \\ sN
  where
    sM :: Sing ms
    sN :: Sing ns
    (sM, sN) = splitSing lM sing
{-# INLINE[0] gmul' #-}

{-# RULES
"gmul'/SS"  gmul' = dispatchSS
"gmul'/SV"  gmul' = dispatchSV
"gmul'/dot" gmul' = dispatchDot
"gmul'/VM"  gmul' = dispatchVM
"gmul'/VS"  gmul' = dispatchVS
"gmul'/out" gmul' = dispatchOut
"gmul'/MV"  gmul' = dispatchMV
"gmul'/MM"  gmul' = dispatchMM
  #-}

-- | General strategy:
--
-- *   We can only outsource to BLAS (using 'dispatchBLAS') in the case
--     that @os@ and @ns@ have length 0 or 1.  Anything else, fall back to
--     the basic reverse-indexing method in "Data.Nested".
-- *   If @ms@ is length 2 or higher, "traverse down" to the length 0 or
--     1 tail...and then sum them up.
gmulB
    :: forall k (b :: BShape k -> Type) v ms os ns.
     ( RealFloat (ElemB b)
     , SingI ns
     , BLAS b
     , Vec v
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     )
    => Sing ms
    -> Length os
    -> Length ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
gmulB sM lO lN v r = case splitting (S_ Z_) (lengthProd lN) of
    Fewer mlN _ -> case splittingEnd (S_ (S_ Z_)) (lengthProd lO) of
      FewerEnd MLZ             _ -> gmulBLAS sM MLZ       mlN v r
      FewerEnd (MLS MLZ)       _ -> gmulBLAS sM (MLS MLZ) mlN v r
      FewerEnd (MLS (MLS MLZ)) _ -> case mlN of
        MLZ -> case r of
          BTM ys -> mapBTM sM LZ (\xs -> BTS $ traceB (gemm 1 xs ys Nothing)) v
        MLS MLZ -> naiveGMul sM lO lN v r
      SplitEnd _ _ _ -> naiveGMul sM lO lN v r
    Split _ _ _ -> naiveGMul sM lO lN v r
{-# INLINE[0] gmulB #-}

-- | Naive implementation of 'gmul' (based on the implementation for
-- 'NTensor') that does not utilize any BLAS capabilities.
naiveGMul
    :: forall k (b :: BShape k -> Type) v ms os ns.
     ( BLAS b
     , Vec v
     , Num (ElemB b)
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     , SingI ns
     )
    => Sing ms
    -> Length os
    -> Length ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
naiveGMul sM _ lN v r =
    mapRowsBTensor sM lN (getSum . ifoldMapBTensor (\i -> Sum . f i)) v
  where
    f  :: Prod (IndexN k) os
       -> ElemB b
       -> BTensor v b ns
    f is x = mapBase sing
                     (x *)
                     (\_ ys -> scaleB x ys)
                     (\_ ys -> scaleB x ys)
                     (indexRowBTensor (TCP.reverse' is) r)

-- | A 'gmul' that runs my dispatching BLAS commands when it can.
-- Contains the type-level constraint that @os@ and @ns@ have to have
-- either length 0 or 1.
--
-- TODO: no longer needs Sing ms
gmulBLAS
    :: forall b ms os ns v.
     ( RealFloat (ElemB b)
     , BLAS b
     , Vec v
     , SingI ns
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     )
    => Sing ms
    -> MaxLength N1 os
    -> MaxLength N1 ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
gmulBLAS sM mlO mlN v r = case mlO of
    MLZ -> case splittingEnd (S_ (S_ Z_)) spM of
      FewerEnd MLZ             _ -> dispatchBLAS MLZ       mlO  mlN v r
      FewerEnd (MLS MLZ)       _ -> dispatchBLAS (MLS MLZ) mlO mlN v r
      FewerEnd (MLS (MLS MLZ)) _ -> case v of
        BTM xs -> case mlN of
          MLZ     -> case r of
            BTS y  -> BTM $ scaleB y xs
          -- TODO: can this be made non-naive?
          -- ms ~ '[m1,m2]
          -- os ~ '[]
          -- ns ~ '[n]
          MLS MLZ -> naiveGMul sM LZ (fromMaxLength mlN) v r
      SplitEnd (ELS (ELS ELZ)) spM0 spM1 -> case mlN of
        MLZ -> case r of
          BTS y -> mapBTM (prodSing   spM0)
                          (prodLength spM1)
                          (\xs -> BTM $ scaleB y xs)
                          v
                     \\ appendNil lM
        -- TODO: can this be made non-naive?
        -- ms ~ (ms0 ++ '[m1,m2])
        -- os ~ '[]
        -- ns ~ '[n]
        MLS MLZ -> naiveGMul sM LZ (fromMaxLength mlN) v r
    MLS MLZ -> case splittingEnd (S_ Z_) spM of
      FewerEnd mlM       _         -> dispatchBLAS mlM mlO mlN v r
      SplitEnd (ELS ELZ) spM0 spM1 ->
        let sM0 = prodSing spM0
            lM0 = prodLength spM0
            lM1 = prodLength spM1
        in  (\\ appendAssoc (TCL.tail' lM0)
                            lM1
                            (LS LZ :: Length os)
            ) $ case mlN of
          MLZ -> case r of
            BTV ys -> mapBTM sM0 lM1 (\xs -> BTV $ gemv 1 xs ys Nothing) v
                        \\ appendNil lM
          MLS MLZ -> case r of
            BTM ys -> mapBTM sM0
                            (lM1 `TCL.append'` (LS LZ :: Length ns))
                            (\xs -> BTM $ gemm 1 xs ys Nothing)
                            v
                        \\ appendAssoc (TCL.tail' lM0)
                                       lM1
                                       (LS LZ :: Length ns)
  where
    spM = singProd sM
    lM  = singLength sM

diagBTensor
    :: forall k (b :: BShape k -> Type) v n ns.
     ( SingI (n ': ns)
     , BLAS b
     , Vec v
     , Num (ElemB b)
     , Eq (IndexN k n)
     )
    => Uniform n ns
    -> BTensor v b '[n]
    -> BTensor v b (n ': ns)
diagBTensor = \case
    UØ    -> id
    US UØ -> \case
      BTV xs -> BTM $ diagB xs
    u@(US (US _)) -> \(BTV xs) ->
      genBTensor sing (\i -> case TCV.uniformVec (prodToVec I (US u) i) of
                               Nothing -> 0
                               Just (I i') -> indexB (PBV i') xs
                      )
{-# INLINE diagBTensor #-}

transpBTensor
    :: (BLAS b, Vec v)
    => Sing ns
    -> BTensor v b ns
    -> BTensor v b (Reverse ns)
transpBTensor s = \case
    BTS x      -> BTS x
    BTV xs     -> BTV xs
    BTM xs     -> BTM $ transpB xs
    xs@(BTN _) -> (\\ reverseReverse (singLength s)) $
                    genBTensor (sReverse s) $ \i ->
                      indexBTensor (TCP.reverse' i) xs
{-# INLINE transpBTensor #-}

sumBTensor
    :: forall v b n ns.
     ( BLAS b
     , Vec v
     , Num (ElemB b)
     , Foldable (v n)
     , SingI ns
     , SingI n
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     )
    => BTensor v b (n ': ns)
    -> BTensor v b ns
sumBTensor = \case
    BTV xs  -> BTS $ sumB xs
    BTM (xs :: b ('BM n m))
            -> BTV $ gemv 1 (transpB xs)
                          (bgen (SBV (sing :: Sing n)) (\_ -> 1))
                          Nothing
    BTN xs  -> sum xs

instance
      ( Vec (v :: k -> Type -> Type)
      , BLAS b
      , RealFloat (ElemB b)
      , Nesting1 Proxy Functor      v
      , Nesting1 Proxy Foldable     v
      , Nesting1 Sing  Applicative  v
      , Nesting1 Sing  Distributive v
      , Eq1 (IndexN k)
      )
  => Tensor (BTensor v b) where
    type ElemT (BTensor v b) = ElemB b

    liftT
        :: SingI ns
        => (TCV.Vec n (ElemB b) -> ElemB b)
        -> TCV.Vec n (BTensor v b ns)
        -> BTensor v b ns
    liftT = liftBTensor sing
    {-# INLINE liftT #-}

    sumT = sum'
    {-# INLINE sumT #-}

    scaleT α = mapBase sing (α*) (\_ -> scaleB α) (\_ -> scaleB α)
    {-# INLINE scaleT #-}

    gmul
        :: forall ms os ns. SingI (ms ++ ns)
        => Length ms
        -> Length os
        -> Length ns
        -> BTensor v b (ms         ++ os)
        -> BTensor v b (Reverse os ++ ns)
        -> BTensor v b (ms         ++ ns)
    gmul = gmul'
    {-# INLINE gmul #-}

    diag
        :: forall n ns. SingI (n ': ns)
        => Uniform n ns
        -> BTensor v b '[n]
        -> BTensor v b (n ': ns)
    diag = diagBTensor
             \\ (produceEq1 :: Eq1 (IndexN k) :- Eq (IndexN k n))
    {-# INLINE diag #-}

    getDiag
        :: SingI n
        => Uniform n ns
        -> BTensor v b (n ': n ': ns)
        -> BTensor v b '[n]
    getDiag = \case
      UØ   -> \case
        BTM xs -> BTV $ getDiagB xs
      u@(US _) -> \xs ->
        genBTensor sing $ \(i :< Ø) ->
          indexBTensor (TCP.replicate i (US (US u))) xs
    {-# INLINE getDiag #-}

    transp = transpBTensor sing
    {-# INLINE transp #-}

    generateA = genBTensorA sing
    {-# INLINE generateA #-}

    genRand d g = generateA (\_ -> realToFrac <$> genContVar d g)
    {-# INLINE genRand #-}

    ixRows
        :: forall f ms os ns. (Applicative f, SingI (ms ++ os))
        => Length ms
        -> Length os
        -> (Prod (IndexN k) ms -> BTensor v b ns -> f (BTensor v b os))
        -> BTensor v b (ms ++ ns)
        -> f (BTensor v b (ms ++ os))
    ixRows lM lO = bIxRows sM lO
      where
        sM :: Sing ms
        sM = takeSing lM lO (sing :: Sing (ms ++ os))
    {-# INLINE ixRows #-}

    (!) = flip indexBTensor
    {-# INLINE (!) #-}

    sumRows
        :: forall n ns. (SingI (n ': ns), SingI ns)
        => BTensor v b (n ': ns)
        -> BTensor v b ns
    sumRows = sumBTensor
                \\ (nesting1 Proxy :: Wit (Foldable (v n)))
                \\ sHead (sing :: Sing (n ': ns))
    {-# INLINE sumRows #-}

    mapRows :: forall ns ms. SingI (ns ++ ms)
            => Length ns
            -> (BTensor v b ms -> BTensor v b ms)
            -> BTensor v b (ns ++ ms)
            -> BTensor v b (ns ++ ms)
    mapRows l f = mapRowsBTensor sN (singLength sM) f
      where
        sN :: Sing ns
        sM :: Sing ms
        (sN, sM) = splitSing l (sing :: Sing (ns ++ ms))
    {-# INLINE mapRows #-}

-- * Boring dispatches

dispatchSS
    :: Num (ElemB b)
    => Length '[]
    -> Length '[]
    -> Length '[]
    -> BTensor v b '[]
    -> BTensor v b '[]
    -> BTensor v b '[]
dispatchSS _ _ _ (BTS x) (BTS y) = BTS (x * y)
{-# INLINE dispatchSS #-}

dispatchSV
    :: BLAS b
    => Length '[]
    -> Length '[]
    -> Length '[n]
    -> BTensor v b '[]
    -> BTensor v b '[n]
    -> BTensor v b '[n]
dispatchSV _ _ _ (BTS x) (BTV y) = BTV $ axpy x y Nothing
{-# INLINE dispatchSV #-}

dispatchDot
    :: BLAS b
    => Length '[]
    -> Length '[n]
    -> Length '[]
    -> BTensor v b '[n]
    -> BTensor v b '[n]
    -> BTensor v b '[]
dispatchDot _ _ _ (BTV x) (BTV y) = BTS $ x `dot` y
{-# INLINE dispatchDot #-}

dispatchVM
    :: (Num (ElemB b), BLAS b)
    => Length '[]
    -> Length '[n]
    -> Length '[m]
    -> BTensor v b '[n]
    -> BTensor v b '[n,m]
    -> BTensor v b '[m]
dispatchVM _ _ _ (BTV x) (BTM y) = BTV $ gemv 1 (transpB y) x Nothing
{-# INLINE dispatchVM #-}

dispatchVS
    :: BLAS b
    => Length '[n]
    -> Length '[]
    -> Length '[]
    -> BTensor v b '[n]
    -> BTensor v b '[]
    -> BTensor v b '[n]
dispatchVS _ _ _ (BTV x) (BTS y) = BTV $ axpy y x Nothing
{-# INLINE dispatchVS #-}

dispatchOut
    :: BLAS b
    => Length '[n]
    -> Length '[]
    -> Length '[m]
    -> BTensor v b '[n]
    -> BTensor v b '[m]
    -> BTensor v b '[n,m]
dispatchOut _ _ _ (BTV x) (BTV y) = BTM $ ger x y
{-# INLINE dispatchOut #-}

dispatchMV
    :: (Num (ElemB b), BLAS b)
    => Length '[n]
    -> Length '[m]
    -> Length '[]
    -> BTensor v b '[n,m]
    -> BTensor v b '[m]
    -> BTensor v b '[n]
dispatchMV _ _ _ (BTM x) (BTV y) = BTV $ gemv 1 x y Nothing
{-# INLINE dispatchMV #-}

dispatchMM
    :: (Num (ElemB b), BLAS b)
    => Length '[m]
    -> Length '[o]
    -> Length '[n]
    -> BTensor v b '[m,o]
    -> BTensor v b '[o,n]
    -> BTensor v b '[m,n]
dispatchMM _ _ _ (BTM x) (BTM y) = BTM $ gemm 1 x y Nothing
{-# INLINE dispatchMM #-}

max_stars_repo_name: mstksg/tensor-ops

max_issues_repo_name: mstksg/tensor-ops

max_forks_repo_name: mstksg/tensor-ops