text:
{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample
-- Copyright : (c) 2008 Don Stewart, 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Commonly used sample statistics, also known as descriptive
-- statistics.

module Statistics.Sample
    (
    -- * Types
      Sample
    , WeightedSample
    -- * Descriptive functions
    , range

    -- * Statistics of location
    , mean
    , welfordMean
    , meanWeighted
    , harmonicMean
    , geometricMean

    -- * Statistics of dispersion
    -- $variance

    -- ** Functions over central moments
    , centralMoment
    , centralMoments
    , skewness
    , kurtosis

    -- ** Two-pass functions (numerically robust)
    -- $robust
    , variance
    , varianceUnbiased
    , meanVariance
    , meanVarianceUnb
    , stdDev
    , varianceWeighted
    , stdErrMean

    -- ** Single-pass functions (faster, less safe)
    -- $cancellation
    , fastVariance
    , fastVarianceUnbiased
    , fastStdDev

    -- * Joint distributions
    , covariance
    , correlation
    , pair
    -- * References
    -- $references
    ) where

import Statistics.Function (minMax)
import Statistics.Sample.Internal (robustSumVar, sum)
import Statistics.Types.Internal  (Sample,WeightedSample)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- Operator ^ will be overridden
import Prelude hiding ((^), sum)

-- | /O(n)/ Range. The difference between the largest and smallest
-- elements of a sample.
range :: (G.Vector v Double) => v Double -> Double
range s = hi - lo
    where (lo , hi) = minMax s
{-# INLINE range #-}

-- | /O(n)/ Arithmetic mean.  This uses Kahan-Babuška-Neumaier
-- summation, so is more accurate than 'welfordMean' unless the input
-- values are very large.
mean :: (G.Vector v Double) => v Double -> Double
mean xs = sum xs / fromIntegral (G.length xs)
{-# SPECIALIZE mean :: U.Vector Double -> Double #-}
{-# SPECIALIZE mean :: V.Vector Double -> Double #-}

-- | /O(n)/ Arithmetic mean.  This uses Welford's algorithm to provide
-- numerical stability, using a single pass over the sample data.
--
-- Compared to 'mean', this loses a surprising amount of precision
-- unless the inputs are very large.
welfordMean :: (G.Vector v Double) => v Double -> Double
welfordMean = fini . G.foldl' go (T 0 0)
  where
    fini (T a _) = a
    go (T m n) x = T m' n'
        where m' = m + (x - m) / fromIntegral n'
              n' = n + 1
{-# SPECIALIZE welfordMean :: U.Vector Double -> Double #-}
{-# SPECIALIZE welfordMean :: V.Vector Double -> Double #-}

-- | /O(n)/ Arithmetic mean for weighted sample. It uses a single-pass
-- algorithm analogous to the one used by 'welfordMean'.
meanWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> Double
meanWeighted = fini . G.foldl' go (V 0 0)
    where
      fini (V a _) = a
      go (V m w) (x,xw) = V m' w'
          where m' | w' == 0   = 0
                   | otherwise = m + xw * (x - m) / w'
                w' = w + xw
{-# INLINE meanWeighted #-}

-- | /O(n)/ Harmonic mean.  This algorithm performs a single pass over
-- the sample.
harmonicMean :: (G.Vector v Double) => v Double -> Double
harmonicMean = fini . G.foldl' go (T 0 0)
  where
    fini (T b a) = fromIntegral a / b
    go (T x y) n = T (x + (1/n)) (y+1)
{-# INLINE harmonicMean #-}

-- | /O(n)/ Geometric mean of a sample containing no negative values.
geometricMean :: (G.Vector v Double) => v Double -> Double
geometricMean = exp . mean . G.map log
{-# INLINE geometricMean #-}

-- | Compute the /k/th central moment of a sample.  The central moment
-- is also known as the moment about the mean.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
centralMoment :: (G.Vector v Double) => Int -> v Double -> Double
centralMoment a xs
    | a < 0  = error "Statistics.Sample.centralMoment: negative input"
    | a == 0 = 1
    | a == 1 = 0
    | otherwise = sum (G.map go xs) / fromIntegral (G.length xs)
  where
    go x = (x-m) ^ a
    m    = mean xs
{-# SPECIALIZE centralMoment :: Int -> U.Vector Double -> Double #-}
{-# SPECIALIZE centralMoment :: Int -> V.Vector Double -> Double #-}

-- | Compute the /k/th and /j/th central moments of a sample.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
centralMoments :: (G.Vector v Double) => Int -> Int -> v Double -> (Double, Double)
centralMoments a b xs
    | a < 2 || b < 2 = (centralMoment a xs , centralMoment b xs)
    | otherwise      = fini . G.foldl' go (V 0 0) $ xs
  where go (V i j) x = V (i + d^a) (j + d^b)
            where d  = x - m
        fini (V i j) = (i / n , j / n)
        m            = mean xs
        n            = fromIntegral (G.length xs)
{-# SPECIALIZE
    centralMoments :: Int -> Int -> U.Vector Double -> (Double, Double) #-}
{-# SPECIALIZE
    centralMoments :: Int -> Int -> V.Vector Double -> (Double, Double) #-}

-- | Compute the skewness of a sample. This is a measure of the
-- asymmetry of its distribution.
--
-- A sample with negative skew is said to be /left-skewed/.  Most of
-- its mass is on the right of the distribution, with the tail on the
-- left.
--
-- > skewness $ U.to [1,100,101,102,103]
-- > ==> -1.497681449918257
--
-- A sample with positive skew is said to be /right-skewed/.
--
-- > skewness $ U.to [1,2,3,4,100]
-- > ==> 1.4975367033335198
--
-- A sample's skewness is not defined if its 'variance' is zero.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
skewness :: (G.Vector v Double) => v Double -> Double
skewness xs = c3 * c2 ** (-1.5)
    where (c3 , c2) = centralMoments 3 2 xs
{-# SPECIALIZE skewness :: U.Vector Double -> Double #-}
{-# SPECIALIZE skewness :: V.Vector Double -> Double #-}

-- | Compute the excess kurtosis of a sample.  This is a measure of
-- the \"peakedness\" of its distribution.  A high kurtosis indicates
-- that more of the sample's variance is due to infrequent severe
-- deviations, rather than more frequent modest deviations.
--
-- A sample's excess kurtosis is not defined if its 'variance' is
-- zero.
--
-- This function performs two passes over the sample, so is not subject
-- to stream fusion.
--
-- For samples containing many values very close to the mean, this
-- function is subject to inaccuracy due to catastrophic cancellation.
kurtosis :: (G.Vector v Double) => v Double -> Double
kurtosis xs = c4 / (c2 * c2) - 3
    where (c4 , c2) = centralMoments 4 2 xs
{-# SPECIALIZE kurtosis :: U.Vector Double -> Double #-}
{-# SPECIALIZE kurtosis :: V.Vector Double -> Double #-}

-- $variance
--
-- The variance — and hence the standard deviation — of a
-- sample of fewer than two elements are both defined to be zero.

-- $robust
--
-- These functions use the compensated summation algorithm of Chan et
-- al. for numerical robustness, but require two passes over the
-- sample data as a result.
--
-- Because of the need for two passes, these functions are /not/
-- subject to stream fusion.

data V = V {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Maximum likelihood estimate of a sample's variance.  Also known
-- as the population variance, where the denominator is /n/.
variance :: (G.Vector v Double) => v Double -> Double
variance samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral n
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE variance :: U.Vector Double -> Double #-}
{-# SPECIALIZE variance :: V.Vector Double -> Double #-}

-- | Unbiased estimate of a sample's variance.  Also known as the
-- sample variance, where the denominator is /n/-1.
varianceUnbiased :: (G.Vector v Double) => v Double -> Double
varianceUnbiased samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral (n-1)
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE varianceUnbiased :: U.Vector Double -> Double #-}
{-# SPECIALIZE varianceUnbiased :: V.Vector Double -> Double #-}

-- | Calculate mean and maximum likelihood estimate of variance. This
-- function should be used if both mean and variance are required
-- since it will calculate mean only once.
meanVariance ::  (G.Vector v Double) => v Double -> (Double,Double)
meanVariance samp
  | n > 1     = (m, robustSumVar m samp / fromIntegral n)
  | otherwise = (m, 0)
    where
      n = G.length samp
      m = mean samp
{-# SPECIALIZE meanVariance :: U.Vector Double -> (Double,Double) #-}
{-# SPECIALIZE meanVariance :: V.Vector Double -> (Double,Double) #-}

-- | Calculate mean and unbiased estimate of variance. This
-- function should be used if both mean and variance are required
-- since it will calculate mean only once.
meanVarianceUnb :: (G.Vector v Double) => v Double -> (Double,Double)
meanVarianceUnb samp
  | n > 1     = (m, robustSumVar m samp / fromIntegral (n-1))
  | otherwise = (m, 0)
    where
      n = G.length samp
      m = mean samp
{-# SPECIALIZE meanVarianceUnb :: U.Vector Double -> (Double,Double) #-}
{-# SPECIALIZE meanVarianceUnb :: V.Vector Double -> (Double,Double) #-}

-- | Standard deviation.  This is simply the square root of the
-- unbiased estimate of the variance.
stdDev :: (G.Vector v Double) => v Double -> Double
stdDev = sqrt . varianceUnbiased
{-# SPECIALIZE stdDev :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdDev :: V.Vector Double -> Double #-}

-- | Standard error of the mean. This is the standard deviation
-- divided by the square root of the sample size.
stdErrMean :: (G.Vector v Double) => v Double -> Double
stdErrMean samp = stdDev samp / (sqrt . fromIntegral . G.length) samp
{-# SPECIALIZE stdErrMean :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdErrMean :: V.Vector Double -> Double #-}

robustSumVarWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> V
robustSumVarWeighted samp = G.foldl' go (V 0 0) samp
    where
      go (V s w) (x,xw) = V (s + xw*d*d) (w + xw)
          where d = x - m
      m = meanWeighted samp
{-# INLINE robustSumVarWeighted #-}

-- | Weighted variance. This is biased estimation.
varianceWeighted :: (G.Vector v (Double,Double)) => v (Double,Double) -> Double
varianceWeighted samp
    | G.length samp > 1 = fini $ robustSumVarWeighted samp
    | otherwise         = 0
    where
      fini (V s w) = s / w
{-# SPECIALIZE varianceWeighted :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE varianceWeighted :: V.Vector (Double,Double) -> Double #-}

-- $cancellation
--
-- The functions prefixed with the name @fast@ below perform a single
-- pass over the sample data using Knuth's algorithm. They usually
-- work well, but see below for caveats. These functions are subject
-- to array fusion.
--
-- /Note/: in cases where most sample data is close to the sample's
-- mean, Knuth's algorithm gives inaccurate results due to
-- catastrophic cancellation.

fastVar :: (G.Vector v Double) => v Double -> T1
fastVar = G.foldl' go (T1 0 0 0)
  where
    go (T1 n m s) x = T1 n' m' s'
      where n' = n + 1
            m' = m + d / fromIntegral n'
            s' = s + d * (x - m')
            d  = x - m

-- | Maximum likelihood estimate of a sample's variance.
fastVariance :: (G.Vector v Double) => v Double -> Double
fastVariance = fini . fastVar
  where fini (T1 n _m s)
          | n > 1     = s / fromIntegral n
          | otherwise = 0
{-# INLINE fastVariance #-}

-- | Unbiased estimate of a sample's variance.
fastVarianceUnbiased :: (G.Vector v Double) => v Double -> Double
fastVarianceUnbiased = fini . fastVar
  where fini (T1 n _m s)
          | n > 1     = s / fromIntegral (n - 1)
          | otherwise = 0
{-# INLINE fastVarianceUnbiased #-}

-- | Standard deviation.  This is simply the square root of the
-- maximum likelihood estimate of the variance.
fastStdDev :: (G.Vector v Double) => v Double -> Double
fastStdDev = sqrt . fastVariance
{-# INLINE fastStdDev #-}

-- | Covariance of sample of pairs. For empty sample it's set to
--   zero
covariance :: (G.Vector v (Double,Double), G.Vector v Double)
           => v (Double,Double)
           -> Double
covariance xy
  | n == 0    = 0
  | otherwise = mean $ G.zipWith (*)
                         (G.map (\x -> x - muX) xs)
                         (G.map (\y -> y - muY) ys)
  where
    n       = G.length xy
    (xs,ys) = G.unzip xy
    muX     = mean xs
    muY     = mean ys
{-# SPECIALIZE covariance :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE covariance :: V.Vector (Double,Double) -> Double #-}

-- | Correlation coefficient for sample of pairs. Also known as
--   Pearson's correlation. For empty sample it's set to zero.
correlation :: (G.Vector v (Double,Double), G.Vector v Double)
           => v (Double,Double)
           -> Double
correlation xy
  | n == 0    = 0
  | otherwise = cov / sqrt (varX * varY)
  where
    n       = G.length xy
    (xs,ys) = G.unzip xy
    (muX,varX) = meanVariance xs
    (muY,varY) = meanVariance ys
    cov = mean $ G.zipWith (*)
            (G.map (\x -> x - muX) xs)
            (G.map (\y -> y - muY) ys)
{-# SPECIALIZE correlation :: U.Vector (Double,Double) -> Double #-}
{-# SPECIALIZE correlation :: V.Vector (Double,Double) -> Double #-}

-- | Pair two samples. It's like 'G.zip' but requires that both
--   samples have equal size.
pair :: (G.Vector v a, G.Vector v b, G.Vector v (a,b)) => v a -> v b -> v (a,b)
pair va vb
  | G.length va == G.length vb = G.zip va vb
  | otherwise = error "Statistics.Sample.pair: vector must have same length"
{-# INLINE pair #-}

------------------------------------------------------------------------
-- Helper code. Monomorphic unpacked accumulators.

-- (^) operator from Prelude is just slow.
(^) :: Double -> Int -> Double
x ^ 1 = x
x ^ n = x * (x ^ (n-1))
{-# INLINE (^) #-}

-- don't support polymorphism, as we can't get unboxed returns if we use it.
data T = T {-# UNPACK #-}!Double {-# UNPACK #-}!Int

data T1 = T1 {-# UNPACK #-}!Int {-# UNPACK #-}!Double {-# UNPACK #-}!Double

{-

Consider this core:

with data T a = T !a !Int

$wfold :: Double#
               -> Int#
               -> Int#
               -> (# Double, Int# #)

and without,

$wfold :: Double#
               -> Int#
               -> Int#
               -> (# Double#, Int# #)

yielding to boxed returns and heap checks.

-}

-- $references
--
-- * Chan, T. F.; Golub, G.H.; LeVeque, R.J. (1979) Updating formulae
--   and a pairwise algorithm for computing sample
--   variances. Technical Report STAN-CS-79-773, Department of
--   Computer Science, Stanford
--   University. <ftp://reports.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf>
--
-- * Knuth, D.E. (1998) The art of computer programming, volume 2:
--   seminumerical algorithms, 3rd ed., p. 232.
--
-- * Welford, B.P. (1962) Note on a method for calculating corrected
--   sums of squares and products. /Technometrics/
--   4(3):419–420. <http://www.jstor.org/stable/1266577>
--
-- * West, D.H.D. (1979) Updating mean and variance estimates: an
--   improved method. /Communications of the ACM/
--   22(9):532–535. <http://doi.acm.org/10.1145/359146.359153>

max_stars_repo_name: vaerksted/statistics

max_issues_repo_name: vaerksted/statistics

max_forks_repo_name: vaerksted/statistics


text:
-- https://github.com/minoki/my-atcoder-solutions
{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
import           Control.Exception           (assert)
import           Control.Monad
import           Data.Bits
import qualified Data.ByteString.Builder     as BSB
import qualified Data.ByteString.Char8       as BS
import           Data.Char                   (isSpace)
import           Data.Coerce
import           Data.Complex
import           Data.Int                    (Int64)
import           Data.List                   (unfoldr)
import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Unboxed         as U
import qualified Data.Vector.Unboxed.Mutable as UM
import           System.IO                   (stdout)

main = do
  n <- readLn @Int -- n <= 10^5
  (as,bs) <- fmap U.unzip $ U.replicateM n $ do
    [a,b] <- unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine
    return (a,b)
  let p, q :: Poly U.Vector Int
      p = Poly $ normalizePoly (0 `U.cons` as)
      q = Poly $ normalizePoly (0 `U.cons` bs)
      -- v = coeffAsc (p * q)
      !v = coeffAsc p `mulFFT` coeffAsc q
      !l = U.length v
  BSB.hPutBuilder stdout $ mconcat
    [ if k < l then
        BSB.intDec (v U.! k) <> BSB.char8 '\n' -- <= 10^9
      else
        BSB.string8 "0\n"
    | k <- [1..2*n]
    ]

--
-- Fast Fourier Transform (FFT)
--

halve :: G.Vector vec a => vec a -> vec a
halve v = let n = G.length v
          in G.generate (n `quot` 2) $ \j -> v G.! (j * 2)

fft :: forall vec a. (Num a, G.Vector vec a)
    => [vec a] -- ^ For a primitive n-th root of unity @u@, @iterate halve [1,u,u^2 .. u^(n-1)]@
    -> vec a -- ^ a polynomial of length n (= 2^k for some k)
    -> vec a
fft (u:u2) f | n == 1 = f
             | otherwise = let !n2 = n `quot` 2
                               r0, r1', t0, t1' :: vec a
                               r0 = G.generate n2 $ \j -> (f G.! j) + (f G.! (j + n2))
                               r1' = G.generate n2 $ \j -> ((f G.! j) - (f G.! (j + n2))) * u G.! j
                               !t0 = fft u2 r0
                               !t1' = fft u2 r1'
                           in G.generate n $ \j -> if even j then t0 G.! (j `quot` 2) else t1' G.! (j `quot` 2)
  where n = G.length f

mulFFT :: U.Vector Int -> U.Vector Int -> U.Vector Int
mulFFT !f !g = let n' = U.length f + U.length g - 2
                   k = finiteBitSize n' - countLeadingZeros n'
                   !_ = assert (n' < 2^k) ()
                   n = bit k
                   u :: U.Vector (Complex Double)
                   u = U.generate n $ \j -> cis (fromIntegral j * (2 * pi / fromIntegral n))
                   us = iterate halve u
                   f' = U.generate n $ \j -> if j < U.length f then
                                               fromIntegral (f U.! j)
                                             else
                                               0
                   g' = U.generate n $ \j -> if j < U.length g then
                                               fromIntegral (g U.! j)
                                             else
                                               0
                   f'' = fft us f'
                   g'' = fft us g'
                   fg = U.generate n $ \j -> (f'' U.! j) * (g'' U.! j)
                   fg' = fft (map (U.map conjugate) us) fg
               in U.generate n $ \j -> round (realPart (fg' U.! j) / fromIntegral n)

--
-- Univariate polynomial
--

newtype Poly vec a = Poly { coeffAsc :: vec a } deriving Eq

normalizePoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a
normalizePoly v | G.null v || G.last v /= 0 = v
                | otherwise = normalizePoly (G.init v)

addPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
addPoly v w = case compare n m of
                LT -> G.generate m $ \i -> if i < n
                                           then v G.! i + w G.! i
                                           else w G.! i
                GT -> G.generate n $ \i -> if i < m
                                           then v G.! i + w G.! i
                                           else v G.! i
                EQ -> normalizePoly $ G.zipWith (+) v w
  where n = G.length v
        m = G.length w

subPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
subPoly v w = case compare n m of
                LT -> G.generate m $ \i -> if i < n
                                           then v G.! i - w G.! i
                                           else negate (w G.! i)
                GT -> G.generate n $ \i -> if i < m
                                           then v G.! i - w G.! i
                                           else v G.! i
                EQ -> normalizePoly $ G.zipWith (-) v w
  where n = G.length v
        m = G.length w

naiveMulPoly :: (Num a, G.Vector vec a) => vec a -> vec a -> vec a
naiveMulPoly v w = G.generate (n + m - 1) $
                   \i -> sum [(v G.! (i-j)) * (w G.! j) | j <- [max (i-n+1) 0..min i (m-1)]]
  where n = G.length v
        m = G.length w

doMulP :: (Eq a, Num a, G.Vector vec a) => Int -> vec a -> vec a -> vec a
doMulP n !v !w | n <= 16 = naiveMulPoly v w
doMulP n !v !w
  | G.null v = v
  | G.null w = w
  | G.length v < n2 = let (w0, w1) = G.splitAt n2 w
                          u0 = doMulP n2 v w0
                          u1 = doMulP n2 v w1
                      in G.generate (G.length v + G.length w - 1)
                         $ \i -> case () of
                                   _ | i < n2     ->  u0 `at` i
                                     | i < n      -> (u0 `at` i) + (u1 `at` (i - n2))
                                     | i < n + n2 -> (u1 `at` (i - n2))
  | G.length w < n2 = let (v0, v1) = G.splitAt n2 v
                          u0 = doMulP n2 v0 w
                          u1 = doMulP n2 v1 w
                      in G.generate (G.length v + G.length w - 1)
                         $ \i -> case () of
                                   _ | i < n2     ->  u0 `at` i
                                     | i < n      -> (u0 `at` i) + (u1 `at` (i - n2))
                                     | i < n + n2 -> (u1 `at` (i - n2))
  | otherwise = let (v0, v1) = G.splitAt n2 v
                    (w0, w1) = G.splitAt n2 w
                    v0_1 = v0 `addPoly` v1
                    w0_1 = w0 `addPoly` w1
                    p = doMulP n2 v0_1 w0_1
                    q = doMulP n2 v0 w0
                    r = doMulP n2 v1 w1
                    -- s = (p `subPoly` q) `subPoly` r -- p - q - r
                    -- q + s*X^n2 + r*X^n
                in G.generate (G.length v + G.length w - 1)
                   $ \i -> case () of
                             _ | i < n2     ->   q `at` i
                               | i < n      -> ((q `at` i) + (p `at` (i - n2))) - ((q `at` (i - n2)) + (r `at` (i - n2)))
                               | i < n + n2 -> ((r `at` (i - n)) + (p `at` (i - n2))) - ((q `at` (i - n2)) + (r `at` (i - n2)))
                               | otherwise  ->   r `at` (i - n)
  where n2 = n `quot` 2
        at :: (Num a, G.Vector vec a) => vec a -> Int -> a
        at v i = if i < G.length v then v G.! i else 0
{-# INLINE doMulP #-}

mulPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
mulPoly !v !w = let k = ceiling ((log (fromIntegral (max n m)) :: Double) / log 2) :: Int
                in doMulP (2^k) v w
  where n = G.length v
        m = G.length w
{-# INLINE mulPoly #-}

zeroPoly :: (G.Vector vec a) => Poly vec a
zeroPoly = Poly G.empty

constPoly :: (Eq a, Num a, G.Vector vec a) => a -> Poly vec a
constPoly 0 = Poly G.empty
constPoly x = Poly (G.singleton x)

scalePoly :: (Eq a, Num a, G.Vector vec a) => a -> Poly vec a -> Poly vec a
scalePoly a (Poly xs)
  | a == 0 = zeroPoly
  | otherwise = Poly $ G.map (* a) xs

valueAtPoly :: (Num a, G.Vector vec a) => Poly vec a -> a -> a
valueAtPoly (Poly xs) t = G.foldr' (\a b -> a + t * b) 0 xs

instance (Eq a, Num a, G.Vector vec a) => Num (Poly vec a) where
  (+) = coerce (addPoly :: vec a -> vec a -> vec a)
  (-) = coerce (subPoly :: vec a -> vec a -> vec a)
  negate (Poly v) = Poly (G.map negate v)
  (*) = coerce (mulPoly :: vec a -> vec a -> vec a)
  fromInteger = constPoly . fromInteger
  abs = undefined; signum = undefined

divModPoly :: (Eq a, Fractional a, G.Vector vec a) => Poly vec a -> Poly vec a -> (Poly vec a, Poly vec a)
divModPoly f g@(Poly w)
  | G.null w = error "divModPoly: divide by zero"
  | degree f < degree g = (zeroPoly, f)
  | otherwise = loop zeroPoly (scalePoly (recip b) f)
  where
    g' = toMonic g
    b = leadingCoefficient g
    -- invariant: f == q * g + scalePoly b p
    loop q p | degree p < degree g = (q, scalePoly b p)
             | otherwise = let q' = Poly (G.drop (degree' g) (coeffAsc p))
                           in loop (q + q') (p - q' * g')

    toMonic :: (Fractional a, G.Vector vec a) => Poly vec a -> Poly vec a
    toMonic f@(Poly xs)
      | G.null xs = zeroPoly
      | otherwise = Poly $ G.map (* recip (leadingCoefficient f)) xs

    leadingCoefficient :: (Num a, G.Vector vec a) => Poly vec a -> a
    leadingCoefficient (Poly xs)
      | G.null xs = 0
      | otherwise = G.last xs

    degree :: G.Vector vec a => Poly vec a -> Maybe Int
    degree (Poly xs) = case G.length xs - 1 of
      -1 -> Nothing
      n  -> Just n

    degree' :: G.Vector vec a => Poly vec a -> Int
    degree' (Poly xs) = case G.length xs of
      0 -> error "degree': zero polynomial"
      n -> n - 1

-- 組立除法
-- second constPoly (divModByDeg1 f t) = divMod f (Poly (G.fromList [-t, 1]))
divModByDeg1 :: (Eq a, Num a, G.Vector vec a) => Poly vec a -> a -> (Poly vec a, a)
divModByDeg1 f t = let w = G.postscanr (\a b -> a + b * t) 0 $ coeffAsc f
                   in (Poly (G.tail w), G.head w)

ext: hs

max_stars_repo_name: minoki/my-atcoder-solutions

max_issues_repo_name: minoki/my-atcoder-solutions

max_forks_repo_name: minoki/my-atcoder-solutions


text:
module FFT where

import Control.Applicative

import Data.Matrix as M
import Data.Complex as C

import Utils

-- 1-D Discrete Fourier Transform (DFT), the naive way.
dft :: (RealFloat a) => [Complex a] -> [Complex a]
dft xs =
  let n = length xs
      theta i j = 2 * pi * (fromIntegral i) * (fromIntegral j) / (fromIntegral n)
      w = M.matrix n n (\(i,j) -> C.mkPolar 1 $ theta (i-1) (j-1))
      xs' = M.fromList n 1 xs
  in M.toList $ vdot w xs'

-- Inverse transform for `dft`.
idft :: (RealFloat a) => [Complex a] -> [Complex a]
idft fs =
  let n = length fs
      theta i j = -2 * pi * (fromIntegral i) * (fromIntegral j) / (fromIntegral n)
      norm = fromIntegral n
      winv = M.matrix n n (\(i,j) -> C.mkPolar (1/norm) $ theta (i-1) (j-1))
      fs' = M.fromList n 1 fs
  in M.toList $ vdot winv fs'

-- 1-D DFT for real-valued vector. Unlike the numpy equivalent, returns all
-- frequency components including the negative frequency ones (which are simply
-- complex conjugates of the positive frequency components).
rdft :: (RealFloat a) => [a] -> [Complex a]
rdft xs = dft $ toComplex <$> xs
  where
    toComplex magnitude = C.mkPolar magnitude 0

-- Inverse transform for `rdft`.
irdft :: (RealFloat a) => [Complex a] -> [a]
irdft = map C.realPart . idft

-- numpy.fft.fft
-- 1-D Fast Fourier Transform (FFT).
fft :: (RealFloat a) => [Complex a] -> [Complex a]
fft xs
  | n == 0 = []
  | n == 1 = xs
  | even n = positiveFreqs ++ negativeFreqs
  | otherwise = dft xs
  where
    n = length xs
    positiveFreqs = zipWith3 (\evenX oddX w -> evenX + w * oddX) evenFreqs oddFreqs weights
    negativeFreqs = zipWith3 (\evenX oddX w -> evenX - w * oddX) evenFreqs oddFreqs weights
    evenFreqs = fft $ [xi | (i,xi) <- zip [0..] xs, even i]
    oddFreqs = fft $ [xi | (i,xi) <- zip [0..] xs, odd i]
    weights = [C.conjugate . C.mkPolar 1 $ theta k | k <- [0 .. (n `div` 2)-1]]
    theta k = 2 * pi * (fromIntegral k) / (fromIntegral n)

-- numpy.fft.ifft
-- Inverse transform for `fft`.
ifft :: (RealFloat a) => [Complex a] -> [Complex a]
ifft fs = normalize . ifft' $ fs
  where
    normalize xs = map (/ (fromIntegral (length xs))) xs
    ifft' fs
      | n == 0 = []
      | n == 1 = fs
      | even n = positiveFreqs ++ negativeFreqs
      | otherwise =  unnormalize . idft $ fs -- undo normalization done by idft do avoid double normalization
      where
        n = length fs
        positiveFreqs = zipWith3 (\evenX oddX w -> evenX + w * oddX) evenFreqs oddFreqs weights
        negativeFreqs = zipWith3 (\evenX oddX w -> evenX - w * oddX) evenFreqs oddFreqs weights
        evenFreqs = ifft' $ [xi | (i,xi) <- zip [0..] fs, even i]
        oddFreqs = ifft' $ [xi | (i,xi) <- zip [0..] fs, odd i]
        weights = [C.mkPolar 1 $ theta k | k <- [0 .. (n `div` 2)-1]]
        theta k = 2 * pi * (fromIntegral k) / (fromIntegral n)
        unnormalize = map (* (fromIntegral n))

-- numpy.fft.rfft
-- 1-D FFT for real-valued vector. Unlike the numpy equivalent, returns all
-- frequency components including the negative frequency ones (which are simply
-- complex conjugates of the positive frequency components).
rfft :: (RealFloat a) => [a] -> [Complex a]
rfft xs = fft $ toComplex <$> xs
  where
    toComplex magnitude = C.mkPolar magnitude 0

-- numpy.fft.irfft
-- Inverse transform for `rfft`.
irfft :: (RealFloat a) => [Complex a] -> [a]
irfft = map C.realPart . ifft

-- numpy.fft.fftfreq
-- Computes frequency bins.
-- n: Window length
-- d: Sample spacing (inverse of sampling rate)
fftfreq :: (Fractional a) => Int -> a -> [a]
fftfreq n d = map ((/ norm) . fromIntegral) $ freqs n
  where
    freqs n
      | n == 0 = []
      | otherwise = [0..n_2-1] ++ [i - n | i <- [n_2..n-1]]
    n_2 = if even n then n `div` 2 else (n + 1) `div` 2
    norm = (fromIntegral n) * d

max_stars_repo_name: viswanathgs/haskell-fft

max_issues_repo_name: viswanathgs/haskell-fft

max_forks_repo_name: viswanathgs/haskell-fft


text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TemplateHaskell #-}
-- | a variation on <http://hackage.haskell.org/package/hmatrix-syntax>,
-- which constructs a matrix which has dimensions stored.
module Numeric.LinearAlgebra.Dimensional.DK.QuasiQuotes
(
  vec,
  vecShape,
  mat,
) where

import Language.Haskell.TH

import Language.Haskell.TH as TH
import Language.Haskell.TH.Quote as TH
import Language.Haskell.TH.Syntax as TH
import Language.Haskell.TH.Quote
import Language.Haskell.Meta.Parse as Meta

import Numeric.Units.Dimensional.DK (Quantity, Dimension, DLength, DMass)
import Numeric.LinearAlgebra.Dimensional.DK.Internal (DimMat(..), vecSingleton, vecCons, fromRowVector, vconcat, vconcat', vconcat'')
import Numeric.LinearAlgebra.Dimensional.DK.Shapes

import Data.List.Split

vec = QuasiQuoter {
  quoteExp = parseVectorExp,
  quotePat = error "vec",
  quoteDec = error "vec",
  quoteType = parseVectorType
}

vecShape = QuasiQuoter {
  quoteExp = error "vecShape",
  quotePat = error "vecShape",
  quoteDec = error "vecShape",
  quoteType = parseVectorShapeType
}

mat = QuasiQuoter {
  quoteExp = parseMatrixExp,
  quotePat = error "mat",
  quoteDec = error "mat",
  quoteType = error "mat"
}

parseMatrixExp :: String -> Q Exp
parseMatrixExp s = let rs = fmap parseVectorExp $ splitSemicolonList s
                    in makeMatrixExp rs

parseVectorExp :: String -> Q Exp
parseVectorExp s = let qs = fmap parseQuantityExp $ splitCommaList s
                    in makeVectorExp qs

parseQuantityExp :: String -> Q Exp
parseQuantityExp s = do
                       case (Meta.parseExp s) of
                         Left err -> fail $ show err
                         Right e -> return e

parseVectorType :: String -> Q Type
parseVectorType s = [t| DimMat $(parseVectorShapeType s) |]

parseVectorShapeType :: String -> Q Type
parseVectorShapeType s = let dimTypes = fmap parseDimensionType $ splitCommaList s
                          in makeVectorShapeType dimTypes

parseDimensionType :: String -> Q Type
parseDimensionType s = do
                          case (Meta.parseType s) of
                            Left err -> fail $ show err
                            Right t -> return t

makeMatrixExp :: [Q Exp] -> Q Exp
makeMatrixExp [] = fail "Empty matrices not permitted."
makeMatrixExp [e] = [| fromRowVector $(e) |]
--makeMatrixExp (e:[e2]) = [| vconcat'' $(e) $(e2) |]
makeMatrixExp (e:es) = [| vconcat' $(e) $(makeMatrixExp es) |]

makeVectorExp :: [Q Exp] -> Q Exp
makeVectorExp [] = fail "Empty vectors not permitted."
makeVectorExp [e] = [| vecSingleton $(e) |]
makeVectorExp (e:es) = [| vecCons $(e) $(makeVectorExp es) |]

makeVectorShapeType :: [Q Type] -> Q Type
makeVectorShapeType (d : ds) = [t| 'VectorShape $(d) $(makeTypeLevelList ds) |]
makeVectorShapeType _ = fail "Empty vectors not permitted."

makeTypeLevelList :: [Q Type] -> Q Type
makeTypeLevelList (t : ts) = [t| $(t) ': $(makeTypeLevelList ts) |]
makeTypeLevelList [] = [t| '[] |]

splitCommaList :: String -> [String]
splitCommaList = splitOn ","

splitSemicolonList :: String -> [String]
splitSemicolonList = splitOn ";"

max_stars_repo_name: dmcclean/dimensional-dk-linalg

max_issues_repo_name: dmcclean/dimensional-dk-linalg

max_forks_repo_name: dmcclean/dimensional-dk-linalg


text:
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
module Numeric.Matrix.SVD
  ( MatrixSVD (..), SVD (..)
  , svd1, svd2, svd3, svd3q
  ) where

import Control.Monad
import Control.Monad.ST
import Data.Kind
import Numeric.Basics
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.ST
import Numeric.DataFrame.SubSpace
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Matrix.Bidiagonal
import Numeric.Matrix.Internal
import Numeric.Quaternion.Internal
import Numeric.Scalar.Internal
import Numeric.Subroutine.Sort
import Numeric.Tuple
import Numeric.Vector.Internal

-- | Result of SVD factorization
--   @ M = svdU %* asDiag svdS %* transpose svdV @.
--
--   Invariants:
--
--   * Singular values `svdS` are in non-increasing order and are non-negative.
--   * svdU and svdV are orthogonal matrices
--   * det svdU == 1
--
--   NB: <https://en.wikipedia.org/wiki/Singular_value_decomposition SVD on wiki>
data SVD (t :: Type) (n :: Nat) (m :: Nat)
  = SVD
  { svdU :: Matrix t n n
    -- ^ Left-singular basis matrix
  , svdS :: Vector t (Min n m)
    -- ^ Vector of singular values
  , svdV :: Matrix t m m
    -- ^ Right-singular basis matrix
  }

deriving instance ( Show t, PrimBytes t
                  , KnownDim n, KnownDim m, KnownDim (Min n m))
                  => Show (SVD t n m)
deriving instance ( Eq (Matrix t n n)
                  , Eq (Matrix t m m)
                  , Eq (Vector t (Min n m)))
                  => Eq (SVD t n m)

class RealFloatExtras t
    => MatrixSVD (t :: Type) (n :: Nat) (m :: Nat) where
    -- | Compute SVD factorization of a matrix
    svd :: IterativeMethod => Matrix t n m -> SVD t n m

-- | Obvious dummy implementation of SVD for 1x1 matrices
svd1 :: (PrimBytes t, Num t, Eq t) => Matrix t 1 1 -> SVD t 1 1
svd1 m = SVD
    { svdU = 1
    , svdS = broadcast $ abs x
    , svdV = broadcast $ if x == 0 then 1 else signum x
    }
  where
    x = ixOff 0 m

-- | SVD of a 2x2 matrix can be computed analytically
--
--   Related discussion:
--
--   https://scicomp.stackexchange.com/questions/8899/robust-algorithm-for-2-times-2-svd/
--
--   https://ieeexplore.ieee.org/document/486688
svd2 :: forall t . RealFloatExtras t => Matrix t 2 2 -> SVD t 2 2
svd2 (DF2 (DF2 m00 m01) (DF2 m10 m11)) =
    SVD
    { svdU = DF2 (DF2         uc  us)
                 (DF2 (negate us) uc)
    , svdS = DF2 sigma1 (abs sigma2)
    , svdV = DF2 (DF2 vc  (sg2 (negate vs)))
                 (DF2 vs  (sg2 vc         ))
    }
  where
    x1 = m00 - m11 -- 2F
    x2 = m00 + m11 -- 2E
    y1 = m01 + m10 -- 2G
    y2 = m01 - m10 -- 2H
    yy = y1*y2
    h1 = hypot x1 y1 -- h1 >= abs x1
    h2 = hypot x2 y2 -- h2 >= abs x2
    sigma1 = 0.5 * (h2 + h1)  -- sigma1 >= abs sigma2
    sigma2 = 0.5 * (h2 - h1)  -- can be negative, which is accounted by sg2
    sg2 = negateUnless (sigma2 >= 0)
    hx1 = h1 + x1
    hx2 = h2 + x2
    hxhx = hx1*hx2
    hxy  = hx1*y2
    yhx  = y1*hx2
    (uc', us', vc', vs') = case (x1 > 0 || y1 /= 0, x2 > 0 || y2 /= 0) of
      (True , True ) -> (hxhx + yy, hxy - yhx, hxhx - yy, hxy + yhx)
      (True , False) -> (y1,  hx1, -y1, hx1)
      (False, True ) -> (y2, -hx2, -y2, hx2)
      (False, False) -> (1, 0, -1, 0)
    ru = recip $ hypot uc' us'
    rv = recip $ hypot vc' vs'
    uc = uc' * ru
    us = us' * ru
    vc = vc' * rv
    vs = vs' * rv

-- | Get SVD decomposition of a 3x3 matrix using `svd3q` function.
--
--   This function reorders the singular components under the hood to make sure
--   @s1 >= s2 >= s3 >= 0@.
--   Thus, it has some overhead on top of `svd3q`.
svd3 :: forall t . (Quaternion t, RealFloatExtras t) => Matrix t 3 3 -> SVD t 3 3
svd3 m = SVD
        { svdU = toMatrix33 u
        , svdS = DF3 s1 s2 s3'
        , svdV = neg3If (s3 < 0) (toMatrix33 v)
        }
  where
    (u, (DF3 s1 s2 s3), v) = svd3q m
    s3' = abs s3
    neg3If :: Bool -> Matrix t 3 3 -> Matrix t 3 3
    neg3If False = id
    neg3If True  = ewmap @t @'[3] neg3
    neg3 :: Vector t 3 -> Vector t 3
    neg3 (DF3 a b c) = DF3 a b (negate c)

-- | Get SVD decomposition of a 3x3 matrix, with orthogonal matrices U and V
--   represented as quaternions.
--   Important: U and V are bound to be rotations at the expense of the last
--              singular value being possibly negative.
--
--   This is an adoptation of a specialized 3x3 SVD algorithm described in
--     "Computing the Singular Value Decomposition of 3x3 matrices
--      with minimal branching and elementary floating point operations",
--   by  A. McAdams, A. Selle, R. Tamstorf, J. Teran, E. Sifakis.
--
--   http://pages.cs.wisc.edu/~sifakis/papers/SVD_TR1690.pdf
svd3q :: forall t . (Quaternion t, RealFloatExtras t)
      => Matrix t 3 3 -> (Quater t, Vector t 3, Quater t)
svd3q m = (u, s, v)
  where
    v = jacobiEigenQ (transpose m %* m)
    (s, u) = uncurry fixSigns $ qrDecomposition3 (m %* toMatrix33 v)
    -- last bit: make sure    s1 >= s2 >= 0
    fixSigns :: Vector t 3 -> Quater t -> (Vector t 3, Quater t)
    fixSigns (DF3 s1 s2 s3) q@(Quater a b c d) = case (s1 >= 0, s2 >= 0) of
      (True , True ) -> (mk3 s1 s2 s3, q)
      (False, True ) -> (mk3 (negate s1) s2 (negate s3), Quater (-c)  d   a  (-b))
      (True , False) -> (mk3 s1 (negate s2) (negate s3), Quater   d   c (-b) (-a))
      (False, False) -> (mk3 (negate s1) (negate s2) s3, Quater   b (-a)  d  (-c))
    -- one more thing:
    --   the singular values are ordered, but may have small errors;
    --   as a result adjacent values may seem to be out of order by a very small number
    mk3 :: Scalar t -> Scalar t -> Scalar t -> Vector t 3
    mk3 s1 s2 s3' = case (s1 >= s2, s1 >= abs s3, s2 >= abs s3) of
        (True , True , True ) -> DF3 s1 s2     s3' -- s1 >= s2 >= s3
        (True , True , False) -> DF3 s1 s3 (cs s2) -- s1 >= s3 >  s2
        (True , False, _    ) -> DF3 s3 s1 (cs s2) -- s3 >  s1 >= s2
        (False, True , True ) -> DF3 s2 s1     s3' -- s2 >  s1 >= s3
        (False, _    , False) -> DF3 s3 s2 (cs s1) -- s3 >  s2 >  s1
        (False, False, True ) -> DF3 s2 s3 (cs s1) -- s2 >= s3 >  s1
      where
        s3 = abs s3'
        cs = negateUnless (s3' >= 0)

-- | Approximate values for cos (a/2) and sin (a/2) of a Givens rotation for
--    a 2x2 symmetric matrix. (Algorithm 2)
jacobiGivensQ :: forall t . RealFloatExtras t => t -> t -> t -> (t, t)
jacobiGivensQ aii aij ajj
    | g*sh*sh < ch*ch = (w * ch, w * sh)
    | otherwise       = (c', s')
  where
    ch = 2 * (aii-ajj)
    sh = aij
    w = recip $ hypot ch sh
    g  = 5.82842712474619 :: t  -- 3 + sqrt 8
    c' = 0.9238795325112867 :: t -- cos (pi/8)
    s' = 0.3826834323650898 :: t -- sin (pi/8)

-- | A quaternion for a QR Givens iteration
qrGivensQ :: forall t . RealFloatExtras t => t -> t -> (t, t)
qrGivensQ a1 a2
    | a1 < 0    = (sh * w, ch * w)
    | otherwise = (ch * w, sh * w)
  where
    rho2 = a1*a1 + a2*a2
    sh = if rho2 > M_EPS then a2 else 0
    ch = abs a1 + sqrt (max rho2 M_EPS)
    w = recip $ hypot ch sh -- TODO: consider something like a hypot

-- | One iteration of the Jacobi algorithm on a symmetric 3x3 matrix
--
--   The three words arguments are indices:
--     0 <= i /= j /= k <= 2
jacobiEigen3Iteration :: (Quaternion t, RealFloatExtras t)
                     => Int -> Int -> Int
                     -> STDataFrame s t '[3,3]
                     -> ST s (Quater t)
jacobiEigen3Iteration i j k sPtr = do
    sii <- readDataFrameOff sPtr ii
    sij <- readDataFrameOff sPtr ij
    sjj <- readDataFrameOff sPtr jj
    sik <- readDataFrameOff sPtr ik
    sjk <- readDataFrameOff sPtr jk
    -- Coefficients for a quaternion corresponding to a Givens rotation
    let (ch, sh) = jacobiGivensQ sii sij sjj
        a = ch*ch - sh*sh
        b = 2 * sh*ch
        aa = a * a
        ab = a * b
        bb = b * b
    -- update the matrix
    writeDataFrameOff sPtr ii $
      aa * sii + 2 * ab * sij + bb * sjj
    writeDataFrameOff sPtr ij $
      ab * (sjj - sii) + (aa - bb) * sij
    writeDataFrameOff sPtr jj $
      bb * sii - 2 * ab * sij + aa * sjj
    writeDataFrameOff sPtr ik $ a * sik + b * sjk
    writeDataFrameOff sPtr jk $ a * sjk - b * sik

    -- write the quaternion
    qPtr <- unsafeThawDataFrame 0
    writeDataFrameOff qPtr k (negate sh)
    writeDataFrameOff qPtr 3 ch
    fromVec4 <$> unsafeFreezeDataFrame qPtr
  where
    ii = i*3 + i
    ij = if i < j then i*3 + j else j*3 + i
    jj = j*3 + j
    ik = if i < k then i*3 + k else k*3 + i
    jk = if j < k then j*3 + k else k*3 + j

-- | Total number of the Givens rotations during the Jacobi eigendecomposition
--   part of the 3x3 SVD equals eigenItersX3*3.
--   Value `eigenItersX3 = 6` corresponds to 18 iterations and gives a good precision.
eigenItersX3 :: Int
eigenItersX3 = 12

-- | Run a few iterations of the Jacobi algorithm on a real-valued 3x3 symmetric matrix.
--   The eigenvectors basis of such matrix is orthogonal, and can be represented as
--   a quaternion.
jacobiEigenQ :: forall t
              . (Quaternion t, RealFloatExtras t)
             => Matrix t 3 3 -> Quater t
jacobiEigenQ m = runST $ do
    mPtr <- thawDataFrame m
    q  <- go eigenItersX3 mPtr 1
    s1 <- readDataFrameOff mPtr 0
    s2 <- readDataFrameOff mPtr 4
    s3 <- readDataFrameOff mPtr 8
    return $ sortQ s1 s2 s3 * q
  where
    go :: Int -> STDataFrame s t '[3,3] -> Quater t -> ST s (Quater t)
    go 0 _ q = pure q

    -- -- primitive cyclic iteration;
    -- --   fast, but the convergence is not perfect
    -- --
    -- -- set eigenItersX3 = 6 for good precision
    -- go n p q = do
    --   q1 <- jacobiEigen3Iteration 0 1 2 p
    --   q2 <- jacobiEigen3Iteration 1 2 0 p
    --   q3 <- jacobiEigen3Iteration 2 0 1 p
    --   go (n - 1) p (q3 * q2 * q1 * q)

    -- Pick the largest element on lower triangle;
    --   slow because of branching, but has a better convergence
    --
    --  set eigenItersX3 = 12 for good precision
    --    (slightly faster than the cyclic version with -O0)
    go n p q = do
      a10 <- abs <$> readDataFrameOff p 1
      a20 <- abs <$> readDataFrameOff p 2
      a21 <- abs <$> readDataFrameOff p 5
      q' <- jiter n p a10 a20 a21
      go (n - 1) p (q' * q)
    jiter :: Int -> STDataFrame s t '[3,3]
          -> Scalar t -> Scalar t -> Scalar t -> ST s (Quater t)
    jiter n p a10 a20 a21
      | gt2 a10 a20 a21
        = jacobiEigen3Iteration 0 1 2 p
      | gt2 a20 a10 a21
        = jacobiEigen3Iteration 2 0 1 p
      | gt2 a21 a10 a20
        = jacobiEigen3Iteration 1 2 0 p
      | otherwise
        = case mod n 3 of
            0 -> jacobiEigen3Iteration 0 1 2 p
            1 -> jacobiEigen3Iteration 2 0 1 p
            _ -> jacobiEigen3Iteration 1 2 0 p
    gt2 :: Scalar t -> Scalar t -> Scalar t -> Bool
    gt2 a b c = case compare a b of
                  GT -> a >= c
                  EQ -> a >  c
                  LT -> False

    -- Make such a quaternion that rotates the matrix so that:
    -- abs s1 >= abs s2 >= abs s3
    -- Note, the corresponding singular values may be negative, which must be
    -- taken into account later.
    sortQ :: Scalar t -> Scalar t -> Scalar t -> Quater t
    sortQ s1 s2 s3 = sortQ' (s1 >= s2) (s1 >= s3) (s2 >= s3)
    sortQ' :: Bool -> Bool -> Bool -> Quater t
    sortQ' True  True  True  = Quater 0 0 0 1                    -- s1 >= s2 >= s3
    sortQ' True  True  False = Quater M_SQRT1_2 0 0 (-M_SQRT1_2) -- s1 >= s3 >  s2
    sortQ' True  False _     = Quater 0.5 0.5 0.5 0.5            -- s3 >  s1 >= s2
    sortQ' False True  True  = Quater 0 0 M_SQRT1_2 (-M_SQRT1_2) -- s2 >  s1 >= s3
    sortQ' False _     False = Quater 0 M_SQRT1_2 0 (-M_SQRT1_2) -- s3 >  s2 >  s1
    sortQ' False False True  = Quater 0.5 0.5 0.5 (-0.5)         -- s2 >= s3 >  s1

-- | One Givens rotation for a QR algorithm on a 3x3 matrix
--
--   The three words arguments are indices:
--     0 <= i /= j /= k <= 2
--
--     if i < j then the eigen values are already sorted!
qrDecomp3Iteration :: (Quaternion t, RealFloatExtras t)
                   => Int -> Int -> Int
                   -> STDataFrame s t '[3,3]
                   -> ST s (Quater t)
qrDecomp3Iteration i j k sPtr = do
    sii <- readDataFrameOff sPtr ii
    sij <- readDataFrameOff sPtr ij
    sji <- readDataFrameOff sPtr ji
    sjj <- readDataFrameOff sPtr jj
    sik <- readDataFrameOff sPtr ik
    sjk <- readDataFrameOff sPtr jk
    -- Coefficients for a quaternion corresponding to a Givens rotation
    let (ch, sh) = qrGivensQ sii sji
        a = ch*ch - sh*sh
        b = 2 * sh*ch
    -- update the matrix
    writeDataFrameOff sPtr ii $ a * sii + b * sji
    writeDataFrameOff sPtr ij $ a * sij + b * sjj
    writeDataFrameOff sPtr ik $ a * sik + b * sjk
    writeDataFrameOff sPtr ji 0 --  a * sji - b * sii
    writeDataFrameOff sPtr jj $ a * sjj - b * sij
    writeDataFrameOff sPtr jk $ a * sjk - b * sik

    -- write the quaternion
    qPtr <- unsafeThawDataFrame 0
    writeDataFrameOff qPtr k (negateUnless leftTriple sh)
    writeDataFrameOff qPtr 3 ch
    fromVec4 <$> unsafeFreezeDataFrame qPtr
  where
    leftTriple = (j - i) /= 1 && (k - j) /= 1
    i3 = i*3
    j3 = j*3
    ii = i3 + i
    ij = i3 + j
    ik = i3 + k
    ji = j3 + i
    jj = j3 + j
    jk = j3 + k

-- | Run QR decomposition in context of 3x3 svd: AV = US = QR
--   The input here is matrix AV
--   The R upper-triangular matrix here is in fact a diagonal matrix Sigma;
--   The Q orthogonal matrix is matrix U in the svd decomposition,
--     represented here as a quaternion.
qrDecomposition3 :: (Quaternion t, RealFloatExtras t)
                 => Matrix t 3 3 -> (Vector t 3, Quater t)
qrDecomposition3 m = runST $ do
    mPtr <- thawDataFrame m
    q1 <- qrDecomp3Iteration 0 1 2 mPtr
    q2 <- qrDecomp3Iteration 0 2 1 mPtr
    q3 <- qrDecomp3Iteration 1 2 0 mPtr
    sig0 <- readDataFrameOff mPtr 0
    sig1 <- readDataFrameOff mPtr 4
    sig2 <- readDataFrameOff mPtr 8
    return (DF3 sig0 sig1 sig2, q3 * q2 * q1)

instance RealFloatExtras t => MatrixSVD t 1 1 where
    svd = svd1

instance RealFloatExtras t => MatrixSVD t 2 2 where
    svd = svd2

instance (RealFloatExtras t, Quaternion t) => MatrixSVD t 3 3 where
    svd = svd3

instance {-# INCOHERENT #-}
         ( RealFloatExtras t, KnownDim n, KnownDim m)
         => MatrixSVD t n m where
    svd a = runST $ do
      D <- pure dnm
      Dict <- pure $ minIsSmaller dn dm -- GHC is not convinced :(
      alphas <- unsafeThawDataFrame bdAlpha
      betas <- unsafeThawDataFrame bdBeta
      uPtr <- unsafeThawDataFrame bdU
      vPtr <- unsafeThawDataFrame bdV

      -- remove last beta if m > n
      bLast <- readDataFrameOff betas nm1
      when (abs bLast > M_EPS) $
        svdGolubKahanZeroCol alphas betas vPtr nm1

      -- main routine for a bidiagonal matrix
      let maxIter = 3*nm -- number of tries
      withinIters <- svdBidiagonalInplace alphas betas uPtr vPtr nm maxIter
      unless withinIters . tooManyIterations
        $ "SVD - Givens rotation sweeps for a bidiagonal matrix ("
           ++ show maxIter ++ " sweeps max)."

      -- sort singular values
      sUnsorted <- unsafeFreezeDataFrame alphas
      let sSorted :: Vector (Tuple '[t, Word]) (Min n m)
          sSorted = sortBy (\(S (x :! _)) (S (y :! _)) -> compare y x)
                  $ iwmap @_ @_ @'[] (\(Idx i :* U) (S x) -> S (abs x :! i :! U) ) sUnsorted
          svdS = ewmap @t @_ @'[] (\(S (x :! _)) -> S x) sSorted
          perm = ewmap @Word @_ @'[] (\(S (_ :! i :! U)) -> S i) sSorted
          pCount =
             if nm < 2
             then 0 :: Word
             else foldl (\s (i, j) -> if perm!i > perm!j then succ s else s)
                        0 [(i, j) | i <- [0..nm2w], j <- [i+1..nm2w+1]]
          pPositive = even pCount

      -- alphas and svdS are now out of sync, but that is not a problem

      -- make sure det U == 1
      when ((bdUDet < 0) == pPositive) $ do
        readDataFrameOff alphas 0 >>= writeDataFrameOff alphas 0 . negate
        forM_ [0..n - 1] $ \i ->
          readDataFrameOff uPtr (i*n) >>= writeDataFrameOff uPtr (i*n) . negate

      -- negate negative singular values
      forM_ [0..nm1] $ \i -> do
        s <- readDataFrameOff alphas i
        when (s < 0) $ do
          writeDataFrameOff alphas i $ negate s
          forM_ [0..m - 1] $ \j ->
            readDataFrameOff vPtr (j*m + i)
              >>= writeDataFrameOff vPtr (j*m + i) . negate

      -- apply permutations if necessary
      if pCount == 0
      then do
        svdU <- unsafeFreezeDataFrame uPtr
        svdV <- unsafeFreezeDataFrame vPtr
        return SVD {..}
      else do
        svdU' <- unsafeFreezeDataFrame uPtr
        svdV' <- unsafeFreezeDataFrame vPtr
        let svdU = iwgen @_ @_ @'[] $ \(i :* Idx j :* U) ->
              if j >= dimVal dnm
              then index (i :* Idx j :* U) svdU'
              else index (i :* Idx (unScalar $ perm!j) :* U) svdU'
            svdV = iwgen @_ @_ @'[] $ \(i :* Idx j :* U) ->
              if j >= dimVal dnm
              then index (i :* Idx j :* U) svdV'
              else index (i :* Idx (unScalar $ perm!j) :* U) svdV'
        return SVD {..}
      where
        n = fromIntegral $ dimVal dn :: Int
        m = fromIntegral $ dimVal dm :: Int
        dn = dim @n
        dm = dim @m
        dnm = minDim dn dm
        nm1 = nm - 1
        nm = fromIntegral (dimVal dnm) :: Int
        nm2w = fromIntegral (max (nm - 2) 0) :: Word
        -- compute the bidiagonal form b first, solve svd for b.
        BiDiag {..} = bidiagonalHouseholder a

{- Compute svd for a square bidiagonal matrix inplace
   \( B = U S V^\intercal \)

@
  B = | a1 b1 0     ... 0 |
      | 0  a2 b2 0  ... 0 |
      | 0  0 a3 b3  ... 0 |
      | ................. |
      | 0  0  ... an1 bn1 |
      | 0  0  ...  0  an  | bn? (in case if n > m)
@
 -}
svdBidiagonalInplace ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (nm :: Nat)
     . ( IterativeMethod, RealFloatExtras t
       , KnownDim n, KnownDim m, KnownDim nm, nm ~ Min n m)
    => STDataFrame s t '[nm] -- ^ the main diagonal of B and then the singular values.
    -> STDataFrame s t '[nm] -- ^ first upper diagonal of B
    -> STDataFrame s t '[n,n] -- ^ U
    -> STDataFrame s t '[m,m] -- ^ V
    -> Int -- ^ 0 < q <= nm -- size of a reduced matrix, such that leftover is diagonal
    -> Int -- iters
    -> ST s Bool -- whether the algorithm succeeds within the maxIter
svdBidiagonalInplace _ _ _ _ 0 _ = pure True
svdBidiagonalInplace _ _ _ _ 1 _ = pure True
svdBidiagonalInplace _ _ _ _ _ 0 = pure False
svdBidiagonalInplace aPtr bPtr uPtr vPtr q' iter = do
    Dict <- pure $ minIsSmaller (dim @n) (dim @m)
    (p, q) <- findCounters q'
    if (q /= 0)
    then do
      findZeroDiagonal p q >>= \case
        Just k
          | k == q-1  -> svdGolubKahanZeroCol aPtr bPtr vPtr (k-1)
          | otherwise -> svdGolubKahanZeroRow aPtr bPtr uPtr k
        Nothing -> svdGolubKahanStep aPtr bPtr uPtr vPtr p q
      svdBidiagonalInplace aPtr bPtr uPtr vPtr q (iter - 1)
    else return True
  where
    -- nm = fromIntegral $ dimVal' @nm :: Int

    -- Check if off-diagonal elements are close to zero and nullify them
    -- if they are small along the way.
    -- And find such indices p and q that satisfy condition in alg. 8.6.2
    -- on p. 492. of "Matrix Computations " (4-th edition).
    -- Except these are inverted:
    --   p -- is the starting index of B22 (last submatrix with non-zero superdiagonal)
    --   q -- is the starting index of B33 (diagonal submatrix)
    --
    -- that is, p and q determine the index and the size of next work piece.
    findCounters :: Int -> ST s (Int, Int)
    findCounters = goQ
      where
        checkEps :: Int -> ST s Bool
        checkEps k = do
          b <- abs <$> readDataFrameOff bPtr (k-1)
          if b == 0
          then return True
          else do
            a1 <- abs <$> readDataFrameOff aPtr (k-1)
            a2 <- abs <$> readDataFrameOff aPtr  k
            if b <= M_EPS * (max (a1 + a2) 1)
            then True <$ writeDataFrameOff bPtr (k-1) 0
            else return False
        goQ :: Int -> ST s (Int, Int)
        goQ 0 = pure (0, 0) -- guard against calling with q == 0
        goQ 1 = pure (0, 0) -- 1x1 matrix is always diagonal
        goQ k = checkEps (k-1) >>= \case
          True  -> goQ (k-1)
          False -> flip (,) k <$> goP (k-2)
        goP :: Int -> ST s Int
        goP 0 = pure 0
        goP k = checkEps k >>= \case
          True  -> return k
          False -> goP (k-1)

    -- For indices p and q (p < q), find the biggest index (< q) such that
    --  a[p] == 0
    findZeroDiagonal :: Int -> Int -> ST s (Maybe Int)
    findZeroDiagonal p q
      | k < p     = pure Nothing
      | otherwise = do
        ak <- readDataFrameOff aPtr k
        if ak == 0
        then pure $ Just k
        else if abs ak <= M_EPS
             then Just k <$ writeDataFrameOff aPtr k 0
             else findZeroDiagonal p k
      where
        k = q - 1

-- | Apply a series of column transformations to make b[k] (and whole column k+1) zero
--    (page 491, 1st paragraph) when a[k+1] == 0.
--   To make this element zero, I apply a series of Givens transforms on columns
--   (multiply on the right).
--
--   Prerequisites:
--     * a[k+1] == 0
--     * 0 <= k < min n (m-1)
--   Invariants:
--     * matrix \(B :: n \times m \) is bidiagonal, represented by two diagonals;
--     * matrix V is orthogonal
--     * \( B = A V^\intercal \), where \(A :: n \times m\) is an implicit original matrix
--   Results:
--     * Same bidiagonal matrix with b[k] == 0; i.e. (k+1)-th column is zero.
--     * matrix V is updated (multiplied on the right)
--
--   NB: All changes are made inplace.
--
svdGolubKahanZeroCol ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat)
     . (RealFloatExtras t, KnownDim n, KnownDim m, n <= m)
    => STDataFrame s t '[n] -- ^ the main diagonal of \(B\)
    -> STDataFrame s t '[n] -- ^ first upper diagonal of \(B\)
    -> STDataFrame s t '[m,m] -- ^ \(V\)
    -> Int -- ^ 0 <= k < min (n+1) m
    -> ST s ()
svdGolubKahanZeroCol aPtr bPtr vPtr k
  | k < 0 || k >= lim = error $ unwords
      [ "svdGolubKahanZeroCol: k =", show k
      , "is outside of a valid range 0 <= k <", show lim]
    -- this trick is to convince GHC that constraint (n <= m) is not redundant
  | Dict <- Dict @(n <= m) = do
    b <- readDataFrameOff bPtr k
    writeDataFrameOff bPtr k 0
    foldM_ goGivens b [k, k-1 .. 0]
  where
    n = fromIntegral $ dimVal' @n :: Int
    m = fromIntegral $ dimVal' @m :: Int
    lim = min n (m-1)
    goGivens :: Scalar t -> Int -> ST s (Scalar t)
    goGivens 0 _ = return 0 -- non-diagonal element is nullified prematurely
    goGivens b i = do
      ai <- readDataFrameOff aPtr i
      let rab = recip $ hypot b ai
          c = ai*rab
          s = b *rab
      updateGivensMat vPtr i (k+1) c s
      writeDataFrameOff aPtr i $ ai*c + b*s -- B[i,i]
      if i == 0
      then return 0
      else do
        bi1 <- readDataFrameOff bPtr (i - 1)  -- B[i,i-1]
        writeDataFrameOff bPtr (i - 1) $ bi1 * c
        return $ negate (bi1 * s)

-- | Apply a series of row transformations to make b[k] (and whole column k) zero
--    (page 490, last paragraph) when a[k] == 0.
--   To make this element zero, I apply a series of Givens transforms on rows
--   (multiply on the left).
--
--   Prerequisites:
--     * a[k] == 0
--     * 0 <= k < n - 1
--   Invariants:
--     * matrix \(B :: m \times n \) is bidiagonal, represented by two diagonals;
--     * matrix U is orthogonal
--     * \( B = U A \), where \(A :: m \times n \) is an implicit original matrix
--   Results:
--     * Same bidiagonal matrix with b[k] == 0; i.e. k-th column is zero.
--     * matrix U is updated (multiplied on the right)
--
--   NB: All changes are made inplace.
--
svdGolubKahanZeroRow ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat)
     . (RealFloatExtras t, KnownDim n, KnownDim m, n <= m)
    => STDataFrame s t '[n] -- ^ the main diagonal of B
    -> STDataFrame s t '[n] -- ^ first upper diagonal of B
    -> STDataFrame s t '[m,m] -- ^ U
    -> Int -- ^ 0 <= k < n - 1
    -> ST s ()
svdGolubKahanZeroRow aPtr bPtr uPtr k
  | k < 0 || k >= n1 = error $ unwords
      [ "svdGolubKahanZeroRow: k =", show k
      , "is outside of a valid range 0 <= k <", show n1]
    -- this trick is to convince GHC that constraint (n <= m) is not redundant
  | Dict <- Dict @(n <= m) = do
    b <- readDataFrameOff bPtr k
    writeDataFrameOff bPtr k 0
    foldM_ goGivens b [k+1..n1]
  where
    n = fromIntegral $ dimVal' @n :: Int
    n1 = n - 1
    goGivens :: Scalar t -> Int -> ST s (Scalar t)
    goGivens 0 _ = return 0 -- non-diagonal element is nullified prematurely
    goGivens b j = do
      aj <- readDataFrameOff aPtr j
      bj <- readDataFrameOff bPtr j
      let rab = recip $ hypot b aj
          c = aj*rab
          s =  b*rab
      updateGivensMat uPtr k j c (negate s)
      writeDataFrameOff aPtr j $ b*s + aj*c
      writeDataFrameOff bPtr j $ bj*c
      return $ negate (bj * s)

-- | A Golub-Kahan bidiagonal SVD step on an unreduced matrix
svdGolubKahanStep ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (nm :: Nat)
     . ( RealFloatExtras t
       , KnownDim n, KnownDim m, KnownDim nm, nm ~ Min n m)
    => STDataFrame s t '[nm] -- ^ the main diagonal of B and then the singular values.
    -> STDataFrame s t '[nm] -- ^ first upper diagonal of B
    -> STDataFrame s t '[n,n] -- ^ U
    -> STDataFrame s t '[m,m] -- ^ V
    -> Int -- ^ p : 0 <= p < q <= nm; p <= q - 2
    -> Int -- ^ q : 0 <= p < q <= nm; p <= q - 2
    -> ST s ()
svdGolubKahanStep aPtr bPtr uPtr vPtr p q
  | p > q - 2 || p < 0 || q > nm
    = error $ unwords
        [ "svdGolubKahanStep: p =", show p, "and q =", show q
        , "do not satisfy p <= q - 2 or 0 <= p < q <=", show nm]
  | Dict <- Dict @(nm ~ Min n m) = do
    (y,z) <- getWilkinsonShiftYZ
    goGivens2 y z p
  where
    nm = fromIntegral $ dimVal' @nm :: Int

    -- get initial values for one recursion sweep.
    -- Note, input must satisfy: q >= p+2
    getWilkinsonShiftYZ :: ST s (Scalar t, Scalar t)
    getWilkinsonShiftYZ  = do
      a1 <- readDataFrameOff aPtr p
      b1 <- readDataFrameOff bPtr p
      am <- readDataFrameOff aPtr (q-2)
      an <- readDataFrameOff aPtr (q-1)
      bm <- if q >= p + 3
            then readDataFrameOff bPtr (q-3)
            else pure 0
      bn <- readDataFrameOff bPtr (q-2)
      let t11 = a1*a1
          t12 = a1*b1
          tmm = am*am + bm*bm
          tnn = an*an + bn*bn
          tnm = am*bn
          d   = 0.5*(tmm - tnn)
          mu  = tnn + d - negateUnless (d >= 0) (hypot d tnm)
      return (t11 - mu, t12)

    -- yv = b[k-1]; zv = B[k-1,k+1] -- to be eliminated by 1st Givens r
    -- yu = a[k];   zu = B[k+1,k-1] -- to be eliminated by 2nd Givens r
    goGivens2 :: Scalar t -> Scalar t -> Int -> ST s ()
    goGivens2 yv zv k = do
          a1 <- readDataFrameOff aPtr k     -- B[k,k]
          a2 <- readDataFrameOff aPtr (k+1) -- B[k+1,k+1]
          b1 <- readDataFrameOff bPtr k     -- B[k,k+1]
          let a1' = a1*cv + b1*sv  -- B[k,k] == yu
              a2' = a2*cv          -- B[k+1,k+1]
              b0' = yv*cv + zv*sv  -- B[k-1,k]
              b1' = b1*cv - a1*sv  -- B[k,k+1]
              yu  = a1'            -- B[k,k]
              zu  = a2*sv          -- B[k+1,k]
              ryzu = recip $ hypot yu zu
              cu = yu * ryzu
              su = zu * ryzu
              a1'' = yu *cu + zu *su
              a2'' = a2'*cu - b1'*su
              b1'' = b1'*cu + a2'*su
          updateGivensMat vPtr k (k+1) cv sv
          updateGivensMat uPtr k (k+1) cu su

          when (k > p) $ writeDataFrameOff bPtr (k-1) b0'
          writeDataFrameOff bPtr k b1''
          writeDataFrameOff aPtr k a1''
          writeDataFrameOff aPtr (k+1) a2''
          when (k < q - 2) $ do
            b2 <- readDataFrameOff bPtr (k+1) -- B[k+1,k+2]
            let b2'' = b2*cu
                zvn  = b2*su
            writeDataFrameOff bPtr (k+1) b2''
            goGivens2 b1'' zvn (k+1)
        where
          ryzv = recip $ hypot yv zv
          cv = yv * ryzv
          sv = zv * ryzv

-- | Update a transformation matrix with a Givens transform (on the right)
updateGivensMat ::
       forall (s :: Type) (t :: Type) (n :: Nat)
     . (PrimBytes t, Num t, KnownDim n)
    => STDataFrame s t '[n,n]
    -> Int -> Int
    -> Scalar t -> Scalar t -> ST s ()
updateGivensMat p i j c s = forM_ [0..n-1] $ \k -> do
    let nk = n*k
        ioff = nk + i
        joff = nk + j
    uki <- readDataFrameOff p ioff
    ukj <- readDataFrameOff p joff
    writeDataFrameOff p ioff $ uki*c + ukj*s
    writeDataFrameOff p joff $ ukj*c - uki*s
  where
    n = fromIntegral $ dimVal' @n :: Int

minIsSmaller :: forall (n :: Nat) (m :: Nat)
              . Dim n -> Dim m -> Dict (Min n m <= n, Min n m <= m)
minIsSmaller dn dm
  | Just Dict <- lessOrEqDim dnm dn
  , Just Dict <- lessOrEqDim dnm dm
    = Dict
  | otherwise
    = error "minIsSmaller: impossible type-level comparison"
  where
    dnm = minDim dn dm

max_stars_repo_name: achirkin/fasttensor

max_issues_repo_name: achirkin/fasttensor

max_forks_repo_name: achirkin/fasttensor


text:
{-# LANGUAGE TemplateHaskell, OverloadedLists #-}
module Myelin.PyNN.Types.Node where

import Data.Aeson.TH
import Control.Lens
import Numeric.LinearAlgebra

-- | pyNN population node
data Node =
    IFCurrentAlpha {
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_refrac :: Vector Float, -- ^ refractory time
        _v_thresh :: Vector Float, -- ^ threshhold potential
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _v_rest :: Vector Float, -- ^ resting potential
        _cm :: Vector Float, -- ^ membrane capactitance
        _v_reset :: Vector Float, -- ^ reset potential
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _i_offset :: Vector Float -- ^ offset current
    }
    | IFCondAlpha {
        _v_rest :: Vector Float, -- ^ resting membrane potential
        _cm :: Vector Float, -- ^ capacity of the membrane
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_refrac :: Vector Float, -- ^ duration of refractory period
        _tau_syn_E :: Vector Float, -- ^ rise time of the excitatory synaptic alpha function
        _tau_syn_I :: Vector Float, -- ^ rise time of the inhibitory synaptic alpha function
        _e_rev_E :: Vector Float, -- ^ reversal potential for excitatory input
        _e_rev_I :: Vector Float, -- ^ reversal potential for inhibitory input
        _v_thresh :: Vector Float, -- ^ spike threshold
        _v_reset :: Vector Float, -- ^ reset potential after a spike
        _i_offset :: Vector Float -- ^ offset current
    }
    | IFSpikey {
        _e_rev_I :: Vector Float, -- ^ excitatory reversal current
        _g_leak :: Vector Float, -- ^ leak conductance
        _tau_refrac :: Vector Float, -- ^ refractory time
        _v_reset :: Vector Float, -- ^ reset potential
        _v_rest :: Vector Float, -- ^ resting potential
        _v_thresh :: Vector Float -- ^ threshhold potential
    }
    | IFCurrExp {
        _cm :: Vector Float, -- ^ membrane capacitance
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _tau_refrac :: Vector Float, -- ^ refractory time
        _v_thresh :: Vector Float, -- ^ threshhold voltage
        _v_rest :: Vector Float, -- ^ resting potential
        _v_reset :: Vector Float, -- ^ reset potential
        _i_offset :: Vector Float -- ^ offset current
    }
    | IFCondExp {
        _v_rest :: Vector Float, -- ^ resting potential
        _cm :: Vector Float, -- ^ membrane capacitance
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_refrac :: Vector Float, -- ^ refractory time
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _e_rev_E :: Vector Float, -- ^ excitatory reversal potential
        _e_rev_I :: Vector Float, -- ^ inhibitory reversal potential
        _v_thresh :: Vector Float, -- ^ spike initiation threshold
        _v_reset :: Vector Float, -- ^ reset value for membrane potential after a spike
        _i_offset :: Vector Float -- ^ offset current
    }
    | Izhikevich {
        _a :: Vector Float,
        _i_offset :: Vector Float, -- ^ offset current
        _c :: Vector Float,
        _d :: Vector Float,
        _b :: Vector Float
    }
    | EIFCondExp {
        _v_reset :: Vector Float, -- ^ reset potential
        _i_offset :: Vector Float, -- ^ offset current
        _tau_w :: Vector Float,
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _e_rev_E :: Vector Float,
        _v_rest :: Vector Float, -- ^ resting potential
        _cm :: Vector Float, -- ^ capactivie memory
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _tau_m :: Vector Float, -- ^ membrane timeconstant
        _a :: Vector Float,
        _delta_T :: Vector Float,
        _v_thresh :: Vector Float, -- ^ threshhold potential
        _b :: Vector Float,
        _v_spike :: Vector Float,
        _e_rev_I :: Vector Float,
        _tau_refrac :: Vector Float -- ^ refractory time constant
    }
    | EIFCondAlpha {
        _cm :: Vector Float, -- ^ capacity of the membrane
        _tau_refrac :: Vector Float, -- ^ duration of the refractory period
        _v_spike :: Vector Float, -- ^ spike detection threshold
        _v_reset :: Vector Float, -- ^ reset value for membrane potential after a spike
        _v_rest :: Vector Float, -- ^ resting membrane potential (Leak reversal potential)
        _tau_m :: Vector Float, -- ^ membrane time constant
        _i_offset :: Vector Float, -- ^ offset current
        _a :: Vector Float, -- ^ subthreshold adaptation conductance
        _b :: Vector Float, -- ^ spike-triggered adaptation
        _delta_T :: Vector Float, -- ^ slope factor
        _tau_w :: Vector Float, -- ^ adaptation time constant
        _v_thresh :: Vector Float, -- ^ spike initiation threshold
        _e_rev_E :: Vector Float, -- ^ excitatory reversal potential
        _tau_syn_E :: Vector Float, -- ^ rise time of excitatory synaptic conductance (alpha function)
        _e_rev_I :: Vector Float, -- ^ inhibitory reversal potential
        _tau_syn_I :: Vector Float -- ^ rise time of the inhibitory synaptic conductance (alpha function)
    }
    | HHCondExp {
        _cm :: Vector Float, -- ^ capacity of the membrane
        _e_rev_E :: Vector Float, -- ^ excitatory reversal potential
        _e_rev_I :: Vector Float, -- ^ inhibitory reversal potential
        _e_rev_K :: Vector Float,
        _e_rev_Na :: Vector Float,
        _e_rev_leak :: Vector Float,
        _g_leak :: Vector Float,
        _gbar_K :: Vector Float,
        _gbar_Na :: Vector Float,
        _i_offset :: Vector Float,
        _tau_syn_E :: Vector Float,
        _tau_syn_I :: Vector Float,
        _v_offset :: Vector Float
    }
    | SpikeSourcePoisson {
        _start :: Float,
        _duration :: Float,
        _rate :: Float
    }
    | SpikeSourceArray {
        _spike_times :: Vector Float
    }
    | SpikeSourceInhGamma {
        _a :: Vector Float,
        _tbins :: Vector Float,
        _duration :: Float,
        _b :: Vector Float,
        _start :: Float
    }
    deriving (Eq, Show, Ord, Read)

-- |  Defaults taken from http://neuralensemble.org/docs/PyNN/standardmodels.html

-- | default parameters for the integrate and fire conductance based neuron with exponential term
if_cond_exp :: Node
if_cond_exp = IFCondExp {
    _v_rest = [-65.0],
    _cm = [1.0],
    _tau_m = [20.0],
    _tau_refrac = [0.0],
    _tau_syn_E = [5.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _e_rev_I = [-70.0],
    _v_thresh = [-50.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

if_current_alpha :: Node
if_current_alpha = IFCurrentAlpha {
    _tau_m = [20.0],
    _tau_refrac = [0.1],
    _v_thresh = [-50.0],
    _tau_syn_E = [ 0.5],
    _tau_syn_I = [0.5],
    _v_rest = [-65.0],
    _cm = [1.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

if_cond_alpha :: Node
if_cond_alpha = IFCondAlpha {
    _v_rest = [-65.0],
    _cm = [1.0],
    _tau_m = [20.0],
    _tau_refrac = [0.0],
    _tau_syn_E = [5.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _e_rev_I = [-70.0],
    _v_thresh = [-50.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

if_spikey :: Node
if_spikey = IFSpikey {
    _e_rev_I = [-80.0],
    _g_leak = [20.0],
    _tau_refrac = [1.0],
    _v_reset = [-80.0],
    _v_rest = [-75.0],
    _v_thresh = [-55.0]
}

if_current_exponential :: Node
if_current_exponential = IFCurrExp {
    _cm = [1.0],
    _tau_m = [20.0],
    _tau_syn_E = [5.0],
    _tau_syn_I = [5.0],
    _tau_refrac = [0.1],
    _v_thresh = [-50.0],
    _v_rest = [-65.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

izhikevich :: Node
izhikevich = Izhikevich {
    _a = [0.02],
    _i_offset = [0.0],
    _c = [-65.0],
    _d = [2.0],
    _b = [0.2]
}

eif_cond_exp :: Node
eif_cond_exp = EIFCondExp {
    _v_reset = [-70.6],
    _i_offset = [0.0],
    _tau_w = [144.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _v_rest = [-70.6],
    _cm = [0.281],
    _tau_syn_E = [5.0],
    _tau_m = [9.3667],
    _a = [4.0],
    _delta_T = [2.0],
    _v_thresh = [-50.4],
    _b = [0.0805],
    _v_spike = [-40.0],
    _e_rev_I = [-80.0],
    _tau_refrac = [0.1]
}

eif_cond_alpha :: Node
eif_cond_alpha = EIFCondAlpha {
    _v_reset = [-70.6],
    _i_offset = [0.0],
    _tau_w = [144.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _v_rest = [-70.6],
    _cm = [0.281],
    _tau_syn_E = [5.0],
    _tau_m = [9.3667],
    _a = [4.0],
    _delta_T = [2.0],
    _v_thresh = [-50.4],
    _b = [0.0805],
    _v_spike = [-40.0],
    _e_rev_I = [-80.0],
    _tau_refrac = [0.1]
}

spike_source_poisson = SpikeSourcePoisson {
    _start = 0.0,
    _duration = 100.0,
    _rate = 0.1
}

spike_source_array = SpikeSourceArray {
    _spike_times = []
}

deriveJSON defaultOptions ''Node
makeLenses ''Node
makePrisms ''Node

max_stars_repo_name: volr/myelin

max_issues_repo_name: volr/myelin

max_forks_repo_name: volr/myelin


text:
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Lib

import Numeric.AD
import Numeric.AD.Internal.Forward
import Numeric.AD.Rank1.Tower(Tower)
import Numeric.SpecFunctions
import Data.Csv
import qualified Data.ByteString.Lazy.Char8 as C

--c :: Num a => (a -> a) -> a -> a -> a
difference :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> (forall s. AD s (Tower a) -> AD s (Tower a))
difference v base = \x -> v x - v (auto base)

difference' :: Num a => (t -> a) -> t -> t -> a
difference' v base x = v x - v base

optValLin :: Num a => a -> a
optValLin r = 1000 * r

optValQuad :: Num a => a -> a
--optValQuad r = 1000 * (r + r*r)
optValQuad r = 10000 * (r*r) + 10* r

-- Truncate a taylor exapansion
-- If nothing don't take a taylor expansion just evaluate the 0th derivative a.k.a zero
trunTay :: Fractional a => Maybe Int -> (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> a
trunTay Nothing  f centre x =  (diffs f x)!!0
trunTay (Just n) f centre x = foldr (+) 0 $ take (n+1) $ (taylor0 f centre (x-centre))

-- create convexity correction function
-- takes a function and two points and calculates the adjustment needed for the difference
convexityAdj :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> a
convexityAdj f centre b = difference' (\x -> trunTay (Just x) (difference f centre) centre b ) 1 2

main :: IO ()
main = do
    --tr (Just 1) (c d 0.02) 0.02 0.04
    print $ trunTay (Just 0) (difference optValLin 0.02) 0.02 0.04
    --print $ trunTay (Just 1) (difference optionVal 0.02) 0.02 0.04
    --print $ trunTay (Just 2) (difference optionVal 0.02) 0.02 0.04
    --print $ trunTay Nothing (difference optionVal 0.02) 0.02 0.04

    print $ "convexity"
    print $ convexityAdj optValLin 0.02 0.04
    print $ convexityAdj optValQuad 0.02 0.04
    -- just get the second derivative of the adjustment

    -- plot a function trunTay for a range from like 0 -> 0.1 for examples with Nothing
    let xFunValues = [ x*0.001 | x <- [1..100] ]
        linValues = map (\x -> trunTay Nothing optValLin 0.02 x) xFunValues
        quadValues = map (\x -> trunTay Nothing optValQuad 0.02 x) xFunValues

--    print $ (linValues :: [(Double,Double)])
    C.writeFile "./plotting/funVals.csv" $ encode (zip3 xFunValues linValues quadValues :: [(Double,Double,Double)])

    -- Then plot first order derivaitve trunTay (Just 1) centred at 0.02 up to 0.04 to get straight line
    let xApproxValues = [ x*0.001 + 0.02 | x <- [0..60] ]
        quadApproxVals1 = map (\x -> trunTay (Just 1) optValQuad  0.02 x) xApproxValues
        quadApproxVals2 = map (\x -> trunTay (Just 2) optValQuad  0.02 x) xApproxValues

    C.writeFile "./plotting/approxVals.csv" $ encode (zip3 xApproxValues quadApproxVals1 quadApproxVals2 :: [(Double,Double,Double)])

    -- Now plot the convexity adustment as a constant values starting at 0.8
    let sPoint = trunTay (Just 1) optValQuad  0.02 0.08
        convexAdj = (convexityAdj optValQuad 0.02 0.08)
    writeFile "./plotting/convexAdj.csv" $ "0.08," ++ show sPoint ++ ",0," ++ show convexAdj

    print "done"

max_stars_repo_name: collinstmatthew/taylorConvexity

max_issues_repo_name: collinstmatthew/taylorConvexity

max_forks_repo_name: collinstmatthew/taylorConvexity


text:
{-# LANGUAGE TemplateHaskell #-}
module AI.NeuralNetwork
    ( Network
    , layerSizes
    , biases
    , weights
    , sigmoid
    , newNetwork
    , feedforward
    , toFile
    , fromFile ) where

import           Control.Lens
import           Control.Monad
import           Control.Monad.State
import           Data.Random.Normal
import           Numeric.LinearAlgebra
import           System.IO
import           System.Random

data Network = Network { _layerSizes :: [Int]
                       , _biases     :: [Vector Double]
                       , _weights    :: [Matrix Double] }
                       deriving (Eq, Show)

makeLenses ''Network

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

newNetwork :: Monad m => [Int] -> StateT StdGen m Network
newNetwork sizes = do
    let matrixSizes = zip (init sizes) (tail sizes)

    biasesL  <- mapM sequence
        [[state normal | _ <- [1..y]] | y <- tail sizes]
    weightsL <- mapM sequence $
        [[state normal | _ <- [1..x], x <- [1..y]] | (x, y) <- matrixSizes]

    let randBiases  = map fromList biasesL
    let randWeights = map (\((x, y), l) -> (y><x) l) (zip matrixSizes weightsL)

    return Network { _layerSizes = sizes
                   , _biases = randBiases
                   , _weights = randWeights }

feedforward :: Vector Double -> Network -> Vector Double
feedforward input nn = foldl (\a (b, w) -> cmap sigmoid ((w #> a) + b)) input bw
    where bw = zip (nn^.biases) (nn^.weights)

toFile :: FilePath -> Network -> IO ()
toFile filepath nn =
    let headerStr  = show (nn^.layerSizes) ++ "\n"
        biasesStr  = unlines $ show . toList  <$> (nn^.biases)
        weightsStr = unlines $ show . toLists <$> (nn^.weights)
        nnStr      = headerStr ++ biasesStr ++ weightsStr
     in writeFile filepath nnStr

fromFile :: FilePath -> IO Network
fromFile filepath = withFile filepath ReadMode (\handle -> do
    layerSizes' <- (read :: String -> [Int]) <$> hGetLine handle
    biases'     <- forM (tail layerSizes') $ \_ -> do
        line <- hGetLine handle
        return $ fromList $ (read :: String -> [Double]) line
    weights'    <- forM (tail layerSizes') $ \_ -> do
        line <- hGetLine handle
        return $ fromLists $ (read :: String -> [[Double]]) line
    return $ Network { _layerSizes = layerSizes'
                     , _biases     = biases'
                     , _weights    = weights' })

max_stars_repo_name: cornelius-sevald/neurocar

max_issues_repo_name: cornelius-sevald/neurocar

max_forks_repo_name: cornelius-sevald/neurocar


text:
module Statistics.Information.Continuous.TotalCorrelation where

import Data.Matrix
import Statistics.Information.Continuous.Entropy
import Statistics.Information.Continuous.MutualInfo
import Statistics.Information.Utils.Matrix
import Statistics.Information.Utils.Random
import System.Random

tc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
tc g k base xs = sum hxs - hx where
  hxs = [entropy gcol k base col | (col, gcol) <- cols]
  hx = entropy g2 k base xs
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

ctc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
       Double
ctc g k base xs ys = sum hxs - hx where
  hxs = [centropy gcol k base col ys | (col, gcol) <- cols]
  hx = centropy g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

corex_tcs :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
             Double
corex_tcs g k base xs ys = tc g1 k base xs - ctc g2 k base xs ys where
  (g1, g2) = split g

corex_mis :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
             Double
corex_mis g k base xs ys = sum mixs - mi_all where
  mixs = [mi gcol k base col ys | (col, gcol) <- cols]
  mi_all = mi g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

residual :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
residual g k base xs = sum $ [hlessxi i gi | (i, gi) <- igs] where
  hlessxi i gi = let xi = (colVector $ getCol i xs) in
    centropy gi k base xi $ residualMatrix xs i
  igs = [1..(ncols xs)] `zip` (splitN (ncols xs) g)

dtc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
dtc g k base xs = entropy g1 k base xs - residual g2 k base xs where
  (g1, g2) = split g

max_stars_repo_name: eligottlieb/Shannon

max_issues_repo_name: eligottlieb/Shannon

max_forks_repo_name: eligottlieb/Shannon


text:
module Wigner where

import Operators
import BraKet
import LinAlg
import QUtil

import Data.Complex
import Data.List
import Data.List.Split

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic.Mutable as M
import Numeric.FFT.Vector.Invertible

-- Expected value of displacement operator
-- for a given density matrix
-- Params:
-- - Number of terms to use for displacement operator expansion
-- - Displacement parameter
-- - Probabilities per state
-- - Basis in which to evaluate trace
-- - States for which to compute density matrix
expectDisp :: Int -> Complex Double
                -> [Complex Double] -> [Ket (Complex Double)] -> Complex Double
expectDisp n alpha = expectedValue (displacement n alpha)

-- 2D Wigner transform
-- fourier transform all rows, then inversely FT all columns
-- Then normalize
-- unfortunately this doesn't work that well, so manual transform is
-- implemented below with crude riemann integration
wt2d:: [[Complex Double]] -> [[Complex Double]]
wt2d = map ((/pi^2)<$>) . transpose . map (U.toList . run dft . U.fromList) .
    transpose . map (U.toList . run idft . U.fromList)

-- manual W transform
wignerTransform' :: Complex Double -> Double -> [Complex Double] -> [Complex Double] -> Complex Double
wignerTransform' alpha dbeta betas ds = sum $ map (*(dbeta :+ 0.0)) integrand
    where
        integrand = zipWith (*) ds $ (\beta -> exp $ alpha * conjugate beta - conjugate alpha * beta) <$> betas

wignerTransform :: [Complex Double] -> Double -> [Complex Double] -> [Complex Double] -> [Complex Double]
wignerTransform alphas dbeta betas ds = map (/pi^2) $ (\a -> wignerTransform' a dbeta betas ds) <$> alphas

max_stars_repo_name: Jvanrhijn/optiqs

max_issues_repo_name: Jvanrhijn/optiqs

max_forks_repo_name: Jvanrhijn/optiqs


text:
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Recurrent.Layers.BasicRecurrent (
    BasicRecurrent (..)
  , randomBasicRecurrent
  ) where

import           Control.Monad.Random ( MonadRandom, getRandom )

import           Data.Singletons.TypeLits

#if MIN_VERSION_base(4,9,0)
import           Data.Kind (Type)
#endif

import           Numeric.LinearAlgebra.Static

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Recurrent.Core

data BasicRecurrent :: Nat -- Input layer size
                    -> Nat -- Output layer size
                    -> Type where
  BasicRecurrent :: ( KnownNat input
                    , KnownNat output
                    , KnownNat matrixCols
                    , matrixCols ~ (input + output))
                 => !(R output)   -- Bias neuron weights
                 -> !(R output)   -- Bias neuron momentum
                 -> !(L output matrixCols) -- Activation
                 -> !(L output matrixCols) -- Momentum
                 -> BasicRecurrent input output

data BasicRecurrent' :: Nat -- Input layer size
                     -> Nat -- Output layer size
                     -> Type where
  BasicRecurrent' :: ( KnownNat input
                     , KnownNat output
                     , KnownNat matrixCols
                     , matrixCols ~ (input + output))
                  => !(R output)   -- Bias neuron gradients
                  -> !(L output matrixCols)
                  -> BasicRecurrent' input output

instance Show (BasicRecurrent i o) where
  show BasicRecurrent {} = "BasicRecurrent"

instance (KnownNat i, KnownNat o, KnownNat (i + o)) => UpdateLayer (BasicRecurrent i o) where
  type Gradient (BasicRecurrent i o) = (BasicRecurrent' i o)

  runUpdate lp (BasicRecurrent oldBias oldBiasMomentum oldActivations oldMomentum) (BasicRecurrent' biasGradient activationGradient) =
    let newBiasMomentum = konst (learningMomentum lp) * oldBiasMomentum - konst (learningRate lp) * biasGradient
        newBias         = oldBias + newBiasMomentum
        newMomentum     = konst (learningMomentum lp) * oldMomentum - konst (learningRate lp) * activationGradient
        regulariser     = konst (learningRegulariser lp * learningRate lp) * oldActivations
        newActivations  = oldActivations + newMomentum - regulariser
    in BasicRecurrent newBias newBiasMomentum newActivations newMomentum

  createRandom = randomBasicRecurrent

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentUpdateLayer (BasicRecurrent i o) where
  type RecurrentShape (BasicRecurrent i o) = S ('D1 o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentLayer (BasicRecurrent i o) ('D1 i) ('D1 o) where

  type RecTape (BasicRecurrent i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
  -- Do a matrix vector multiplication and return the result.
  runRecurrentForwards (BasicRecurrent wB _ wN _) (S1D lastOutput) (S1D thisInput) =
    let thisOutput = S1D $ wB + wN #> (thisInput # lastOutput)
    in ((S1D lastOutput, S1D thisInput), thisOutput, thisOutput)

  -- Run a backpropogation step for a full connected layer.
  runRecurrentBackwards (BasicRecurrent _ _ wN _) (S1D lastOutput, S1D thisInput) (S1D dRec) (S1D dEdy) =
    let biasGradient        = (dRec + dEdy)
        layerGrad           = (dRec + dEdy) `outer` (thisInput # lastOutput)
        -- calcluate derivatives for next step
        (backGrad, recGrad) = split $ tr wN #> (dRec + dEdy)
    in  (BasicRecurrent' biasGradient layerGrad, S1D recGrad, S1D backGrad)

randomBasicRecurrent :: (MonadRandom m, KnownNat i, KnownNat o, KnownNat x, x ~ (i + o))
                     => m (BasicRecurrent i o)
randomBasicRecurrent = do
    seed1 <- getRandom
    seed2 <- getRandom
    let wB = randomVector seed1 Uniform * 2 - 1
        wN = uniformSample seed2 (-1) 1
        bm = konst 0
        mm = konst 0
    return $ BasicRecurrent wB bm wN mm

max_stars_repo_name: LuisChDev/grenade

max_issues_repo_name: LuisChDev/grenade

max_forks_repo_name: LuisChDev/grenade