text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE TupleSections #-}

module Main where

import Data.Random
import Data.Random.Distribution.MultiNormal
import Numeric.Kalman

import Control.Monad (replicateM)
--import GHC.TypeLits

import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo

import Numeric.LinearAlgebra.Static hiding (mean)

--import Control.Monad (replicateM)
--import Data.Maybe (fromJust)

import qualified Data.List as L
import qualified Numeric.LinearAlgebra as LA

main :: IO ()
--main = putStrLn "Test suite not yet implemented"

--main = runUKFTest1
--main = runNoisySine1
--main = runGSTest1
main = runGSTest2

---------------------------
--- Gaussian sum filter ---
---------------------------
runGSTest2 :: IO ()
runGSTest2 = do
  putStrLn "Running Gaussian sum test 2"
  let numSamples = 1000
      dt = 0.01
      tVals = take numSamples [0, dt ..]
      truey = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truem = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      trueb = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truew = [(head tVals, 1), (last tVals, 1)] :: [(Double, Double)]

  let intErr = 0.04
      slpErr = 0.01
      yErr   = 0.03

  let initEst = [(1/2, MultiNormal (vector [2,-0.5,0.1])  (sym . diag . vector $ [intErr, slpErr, yErr])),
                 (1/2, MultiNormal (vector [-2,0.5,0.1]) (sym . diag . vector $ [intErr, slpErr, yErr]))]
                :: [(Double, MultiNormal (R 3))]

      measMat = row . vector $ [0,0,1]
      measModel = const (const measMat)
      measErr = diag . vector $ [0.0001] :: Sq 1

      sysModel :: Double -> R 3 -> R 3
      sysModel t cur =
        let [b,m,y] = LA.toList (extract cur)
        in vector [b, m, m*t + (2*b^2+b-1)]
      sysErr = sym . diag . vector $ [0,0,0]

  measurements <- replicateM numSamples (sample $ MultiNormal (vector [0]) (sym measErr))

  let samples = [(t, head . LA.toList . extract $ s) | t <- tVals | s <- measurements]
      gsfilter = scanl
                 (\[(_w1, est1), (_w2, est2)] (t, meas) ->
                    let pred1 = runUKFPrediction sysModel (const sysErr) t est1
                        pmu1  = mu pred1
                        pcov1 = unSym . cov $ pred1
                        like1 = sym $ measMat <> pcov1 <> tr measMat + measErr
                        neww1 = pdf (MultiNormal (measModel t pmu1 #> pmu1) like1) meas
                        updt1 = runKFUpdate measModel (const (sym measErr)) t pred1 meas

                        pred2 = runUKFPrediction sysModel (const sysErr) t est2
                        pmu2  = mu pred2
                        pcov2 = unSym . cov $ pred2
                        like2 = sym $ measMat <> pcov2 <> tr measMat + measErr
                        neww2 = pdf (MultiNormal (measModel t pmu2 #> pmu2) like2) meas
                        updt2 = runKFUpdate measModel (const (sym measErr)) t pred2 meas

                        scale = neww1 + neww2

                    in [(neww1 / scale, updt1), (neww2 / scale, updt2)]
                 ) initEst $ zip tVals measurements

  let gsy1 = [(t, last . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsy2 = [(t, last . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsy  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = last . LA.toList . extract . mu $ est1
                            p2 = last . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsm1 = [(t, (!!1) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsm2 = [(t, (!!1) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsm  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!1) . LA.toList . extract . mu $ est1
                            p2 = (!!1) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsb1 = [(t, (!!0) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsb2 = [(t, (!!0) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsb  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!0) . LA.toList . extract . mu $ est1
                            p2 = (!!0) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsw1 = zipWith (\t [(w1,_),_] -> (t,w1)) tVals gsfilter
      gsw2 = zipWith (\t [_,(w2,_)] -> (t,w2)) tVals gsfilter
      gsw  = zipWith (\t [(w1, _), (w2, _)] -> (t,w1 + w2))
             tVals gsfilter

  let [(w1,cur1), (w2, cur2)] = last gsfilter
      final1 = (w1, LA.toList . extract . mu $ cur1)
      final2 = (w2, LA.toList . extract . mu $ cur2)

  print final1
  print final2
  print $ w1 + w2

  toFile def "gs2-y.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truey])
    plot (points "Measurements" samples)
    plot (points "G1 estimates" gsy1)
    plot (points "G2 estimates" gsy2)
    plot (points "GS estimates" gsy)

  toFile def "gs2-m.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truem])
    plot (points "G1 estimates" gsm1)
    plot (points "G2 estimates" gsm2)
    plot (points "GS estimates" gsm)

  toFile def "gs2-b.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [trueb])
    plot (points "G1 estimates" gsb1)
    plot (points "G2 estimates" gsb2)
    plot (points "GS estimates" gsb)

  toFile def "gs2-w.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truew])
    plot (points "G1 estimates" gsw1)
    plot (points "G2 estimates" gsw2)
    plot (points "GS estimates" gsw)

runGSTest1 :: IO ()
runGSTest1 = do
  putStrLn "Running Gaussian sum test 1"
  let numSamples = 1000
      dt = 0.01
      tVals = take numSamples [0, dt ..]
      truey = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truem = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      trueb = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truew = [(head tVals, 1), (last tVals, 1)] :: [(Double, Double)]

  let intErr = 0.04
      slpErr = 0.01
      yErr   = 0.03

  let initEst = [(1/2, MultiNormal (vector [2,-0.5,0.1])  (sym . diag . vector $ [intErr, slpErr, yErr])),
                 (1/2, MultiNormal (vector [-2,0.5,0.1]) (sym . diag . vector $ [intErr, slpErr, yErr]))]
                :: [(Double, MultiNormal (R 3))]

      measMat = row . vector $ [0,0,1]
      measModel = const (const measMat)
      measErr = diag . vector $ [0.0001] :: Sq 1

      sysModel :: Double -> R 3 -> Sq 3
      sysModel t _cur = matrix [1,0,0, 0,1,0, 1,t,0]
      sysErr = sym . diag . vector $ [0,0,0]

  measurements <- replicateM numSamples (sample $ MultiNormal (vector [0]) (sym measErr))

  let samples = [(t, head . LA.toList . extract $ s) | t <- tVals | s <- measurements]
      gsfilter = scanl
                 (\[(_w1, est1), (_w2, est2)] (t, meas) ->
                    let pred1 = runKFPrediction sysModel (const sysErr) t est1
                        pmu1  = mu pred1
                        pcov1 = unSym . cov $ pred1
                        like1 = sym $ measMat <> pcov1 <> tr measMat + measErr
                        neww1 = pdf (MultiNormal (measModel t pmu1 #> pmu1) (like1)) meas
                        updt1 = runKFUpdate measModel (const (sym measErr)) t pred1 meas

                        pred2 = runKFPrediction sysModel (const sysErr) t est2
                        pmu2  = mu pred2
                        pcov2 = unSym . cov $ pred2
                        like2 = sym $ measMat <> pcov2 <> tr measMat + measErr
                        neww2 = pdf (MultiNormal (measModel t pmu2 #> pmu2) (like2)) meas
                        updt2 = runKFUpdate measModel (const (sym measErr)) t pred2 meas

                        scale = neww1 + neww2

                    in [(neww1 / scale, updt1), (neww2 / scale, updt2)]
                 ) initEst $ zip tVals measurements

  let gsy1 = [(t, last . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsy2 = [(t, last . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsy  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = last . LA.toList . extract . mu $ est1
                            p2 = last . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsm1 = [(t, (!!1) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsm2 = [(t, (!!1) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsm  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!1) . LA.toList . extract . mu $ est1
                            p2 = (!!1) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsb1 = [(t, (!!0) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsb2 = [(t, (!!0) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsb  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!0) . LA.toList . extract . mu $ est1
                            p2 = (!!0) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsw1 = zipWith (\t [(w1,_),_] -> (t,w1)) tVals gsfilter
      gsw2 = zipWith (\t [_,(w2,_)] -> (t,w2)) tVals gsfilter
      gsw  = zipWith (\t [(w1, _), (w2, _)] -> (t,w1 + w2))
             tVals gsfilter

  let [(w1,cur1), (w2, cur2)] = last gsfilter
      final1 = (w1, LA.toList . extract . mu $ cur1)
      final2 = (w2, LA.toList . extract . mu $ cur2)

  print final1
  print final2
  print $ w1 + w2

  toFile def "gs1-y.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truey])
    plot (points "Measurements" samples)
    plot (points "G1 estimates" gsy1)
    plot (points "G2 estimates" gsy2)
    plot (points "GS estimates" gsy)

  toFile def "gs1-m.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truem])
    plot (points "G1 estimates" gsm1)
    plot (points "G2 estimates" gsm2)
    plot (points "GS estimates" gsm)

  toFile def "gs1-b.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [trueb])
    plot (points "G1 estimates" gsb1)
    plot (points "G2 estimates" gsb2)
    plot (points "GS estimates" gsb)

  toFile def "gs1-w.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truew])
    plot (points "G1 estimates" gsw1)
    plot (points "G2 estimates" gsw2)
    plot (points "GS estimates" gsw)

---------------------------------------
--- Unscented Kalman filter testing ---
---------------------------------------
runUKFTest1 :: IO ()
runUKFTest1 = do
  -- polar coordinate transform test
  let numSamples = 500
      dt = 0.01
      tVals = take (numSamples + 1) [0,dt..] :: [Double] -- +1 because t = 0 is our initial estimate

      initSys = MultiNormal (vector [0, 1]) (sym . diag . vector $ [4e-4, 2.7e-5] :: Sym 2)
      initEst = MultiNormal (vector [1, pi/2]) (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2)

  let coordTrans :: R 2 -> R 2
      coordTrans polar =
        let [r,t] = LA.toList $ extract polar
        in vector [r * cos t, r * sin t]

  measurements <- map coordTrans <$> replicateM numSamples (sample initEst)

  let samples = map (\meas ->
                       let [x,y] = LA.toList $ extract meas
                       in (x,y)
                    ) measurements

      ukf   = scanl (runUKF (const coordTrans)
                            (const (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2))
                            (const id)
                            (const (sym . diag . vector $ [0,0])) 1)
                    initSys measurements

      ukpts = map (\kf ->
                       let [x,y] = LA.toList . extract $ mu kf
                       in (x,y)
                  ) ukf

  let linMeas :: R 2 -> L 2 2
      linMeas est =
        let [x, _] = LA.toList $ extract est
        in matrix [1,0, 0,x]

  let ekf = scanl (runEKF (const coordTrans)
                          (const linMeas)
                          (const (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2))
                          (const id)
                          (const (const eye))
                          (const (sym . diag . vector $ [0,0])) 1)
                  initSys measurements

      ekpts = map (\kf ->
                       let [x,y] = LA.toList . extract $ mu kf
                       in (x,y)
                    ) ekf

      -- predTest = runUKFPrediction id (sym . diag . vector $ [0,0]) initSys
      -- updtTest = runUKFUpdate coordTrans (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2) initSys (head measurements)

  -- print $ extract . mu $ predTest
  -- print $ extract . unSym . cov $ predTest
  -- putStrLn "\n"
  -- print $ extract . mu $ updtTest
  -- print $ extract . unSym . cov $ updtTest
  -- putStrLn "\n"
  -- print $ last ukpts
  -- print $ last ekpts

  toFile def "noisy-coords.png" $ do
    layout_title .= "Noisy coord transform"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (points "True signal"  [(0,1)])
    plot (points "Measurements" samples)
    plot (points "UKF estimates" ukpts)
    plot (points "EKF estimates" ekpts)

  toFile def "noisy-params.png" $ do
    layout_title .= "Noisy coord transform"
    setColors [opaque black, opaque black
              ,opaque green, opaque green
              ,opaque red, opaque red
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line "True X signal"  [[(head tVals, 0), (last tVals, 0)]])
    plot (line "True Y signal"  [[(head tVals, 1), (last tVals, 1)]])
    plot (points "UKF X estimates" $ zipWith (\t (x,_) -> (t, x)) tVals ukpts)
    plot (points "UKF Y estimates" $ zipWith (\t (_,y) -> (t, y)) tVals ukpts)
    plot (points "EKF X estimates" $ zipWith (\t (x,_) -> (t, x)) tVals ekpts)
    plot (points "EKF Y estimates" $ zipWith (\t (_,y) -> (t, y)) tVals ekpts)

---------------------
--- Noisy sine # 1---
---------------------
-- In this example, we estimate the parameters (amplitude, period, vertical shift, horizontal shift)
-- defining sinusoidal evolution.
-- Figure 3 in my notes

runNoisySine1 :: IO ()
runNoisySine1 = do

  let (a, p, v) = (0.7, 2.2*pi, -0.5) -- (hidden) internal system parameters
      (a0, p0, v0) = (0.8, 2*pi, -0.6)    -- initial estimates
      --(aSig, pSig, vSig) = (0.2, 0.2*pi, 0.001) -- initial uncertainties
      measCov = 0.01

      --- Figure 1 in my notes
      --numSamples = 499
      --dt = 0.002
      --- Figure 2 in my notes
      numSamples = 199
      dt = 0.02
      tVals = take (numSamples + 1) [0,dt..] -- +1 because t = 0 is our initial estimate

      --initSys = vector [a, p, v] -- initial (true) system: (parameter1, parameter2, state)
      --initEst = MultiNormal (vector [a0, p0, v0, a0*p0, v0]) (sym . diag . vector $ [0.04, 0.4, 0.04, 0.001, 0.001] :: Sym 5)
      initEst = MultiNormal (vector [a0, p0, v0, a0*p0, v0])
                (sym . matrix $ [0.04,  0,    0,    1e-5, -2e-3,
                                 0,     0.4,  0,    1e-2, 12e-3,
                                 0,     0,    0.04, 1e-4, 1e-3,
                                 1e-5,  1e-2, 1e-4, 1e-3, 1e-4,
                                 -2e-3, 12e-3, 1e-3, 1e-4, 1e-3] :: Sym 5)

      evolMat :: R 5 -> Sq 5
      evolMat cur =
        let pCur  = (extract cur) LA.! 1
            p2dt  = pCur * pCur * dt
        in matrix [1,0,0,0,0, 0,1,0,0,0, 0,0,1,0,0, 0,0,p2dt,1,-p2dt, 0,0,0,dt,1]

      evolSys :: Double -> R 5 -> R 5
      evolSys t cur =
        let [a,p,v,ap,b] = LA.toList $ extract cur
        in vector [a, p, v, a * p * cos (p * t), a * sin (p * t) + v]

      sysCovariance = sym . diag . vector $ [0,0,0,5e-4,7e-5] :: Sym 5

      measMat = row (vector [0,0,0,0,1]) :: L 1 5
      measTrans = (measMat #>)
      measCovariance = sym . diag . vector $ [measCov] :: Sym 1
      measurementNoise = MultiNormal (vector [0]) measCovariance

      trueSystem = [vector [a, p, v, a * p * cos(p*t), a * sin(p*t) + v] | t <- tVals] :: [R 5]

  --trueSystem   <- simulateSystem numSamples sysModel initSys
  measurements <- mapM (\sys -> do
                            noise <- sample measurementNoise
                            return $ measTrans sys + noise
                       ) $ tail trueSystem

  --mapM_ print $ map (LA.toList . extract) trueSystem
  -- print $ map ((LA.! 0) . extract) measurements

  let truesine = zip tVals $ map ((LA.! 4) . extract) trueSystem
      samples  = zip (tail tVals) $ map ((LA.! 0) . extract) measurements

      kfiltered = scanl (runKF (const (const measMat)) (const measCovariance) (const evolMat) (const sysCovariance) 1) initEst measurements
      ekf = scanl (\sys (t,meas) ->
                     runEKF (\_t sys -> measMat #> sys)
                     (\_t _sys -> measMat)
                     (const measCovariance)
                     evolSys
                     (const evolMat)
                     (const sysCovariance)
                     t sys meas
                  ) initEst $ zip (tail tVals) measurements

      ukf = scanl (\sys (t,meas) ->
                     runUKF (\_t sys -> measMat #> sys)
                     (const measCovariance)
                     evolSys
                     (const sysCovariance)
                     t sys meas
                  ) initEst $ zip (tail tVals) measurements

      kfstates  = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals kfiltered
      ekfstates = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals ekf
      ukfstates = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals ukf

      [kfa,kfp,kfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals kfiltered
      [ekfa,ekfp,ekfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals ekf
      [ukfa,ukfp,ukfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals ukf

      [kfaCovs, kfpCovs, kfvCovs, kfsCovs] = L.transpose $
                                             map (\kf -> let kfcov = extract . unSym $ cov kf
                                                         in [kfcov LA.! 0 LA.! 0, kfcov LA.! 1 LA.! 1,
                                                             kfcov LA.! 2 LA.! 2, kfcov LA.! 4 LA.! 4])
                                             kfiltered

      kfsUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfstates kfsCovs
      kfsLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfstates kfsCovs

      kfsUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfstates kfsCovs
      kfsLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfstates kfsCovs

      kfaUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfa kfaCovs
      kfaLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfa kfaCovs

      kfaUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfa kfaCovs
      kfaLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfa kfaCovs

      kfpUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfp kfpCovs
      kfpLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfp kfpCovs

      kfpUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfp kfpCovs
      kfpLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfp kfpCovs

      kfvUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfv kfvCovs
      kfvLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfv kfvCovs

      kfvUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfv kfvCovs
      kfvLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfv kfvCovs

      ksmoothed = scanr1 (\pres fut -> runKS (const evolMat) (const sysCovariance) 1 fut pres) kfiltered
      ksstates  = zipWith (\t ks -> (t, (extract . mu $ ks) LA.! 4)) tVals ksmoothed

  --mapM_ print $ map (LA.toList . extract . mu) kfiltered
  --mapM_ print kfa

  toFile def "noisy-sine-points.png" $ do
    layout_title .= "Noisy sine (state estimate)"
    setColors [opaque black, opaque blue
              ,opaque red, opaque purple, opaque hotpink, opaque green
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey]
    plot (line   "True signal"  [truesine])
    plot (points "Measurements" samples)
    plot (line   "KF estimated signal" [kfstates])
    plot (line   "EKF estimated signal" [ekfstates])
    plot (line   "UKF estimated signal" [ukfstates])
    plot (line   "KS smoothed signal"  [ksstates])
    plot (line   "Filtered 1 standard deviation" [kfsUpper1])
    plot (line   "Filtered 1 standard deviation" [kfsLower1])
    plot (line   "Filtered 2 standard deviations" [kfsUpper2])
    plot (line   "Filtered 2 standard deviations" [kfsLower2])

  let truea = [(head tVals, a), (last tVals, a)]
      truep = [(head tVals, p), (last tVals, p)]
      truev = [(head tVals, v), (last tVals, v)]

  toFile def "noisy-sine-amplitude.png" $ do
    layout_title .= "Noisy sine (amplitude estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truea])
    plot (line   "KF estimated signal" [kfa])
    plot (line   "EKF estimated signal" [ekfa])
    plot (line   "UKF estimated signal" [ukfa])
    plot (line   "Filtered 1 standard deviation" [kfaUpper1])
    plot (line   "Filtered 1 standard deviation" [kfaLower1])
    plot (line   "Filtered 2 standard deviations" [kfaUpper2])
    plot (line   "Filtered 2 standard deviations" [kfaLower2])

  toFile def "noisy-sine-period.png" $ do
    layout_title .= "Noisy sine (period estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truep])
    plot (line   "KF estimated signal" [kfp])
    plot (line   "EKF estimated signal" [ekfp])
    plot (line   "UKF estimated signal" [ukfp])
    plot (line   "Filtered 1 standard deviation" [kfpUpper1])
    plot (line   "Filtered 1 standard deviation" [kfpLower1])
    plot (line   "Filtered 2 standard deviations" [kfpUpper2])
    plot (line   "Filtered 2 standard deviations" [kfpLower2])

  toFile def "noisy-sine-vertical.png" $ do
    layout_title .= "Noisy sine (vertical estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truev])
    plot (line   "KF estimated signal" [kfv])
    plot (line   "EKF estimated signal" [ekfv])
    plot (line   "UKF estimated signal" [ukfv])
    plot (line   "Filtered 1 standard deviation" [kfvUpper1])
    plot (line   "Filtered 1 standard deviation" [kfvLower1])
    plot (line   "Filtered 2 standard deviations" [kfvUpper2])
    plot (line   "Filtered 2 standard deviations" [kfvLower2])

  putStrLn $ "Kalman processing complete."

max_stars_repo_name: idontgetoutmuch/Kalman

max_issues_repo_name: f-o-a-m/Kalman

max_forks_repo_name: f-o-a-m/Kalman


text:
module Data.MyPolynomial.Print
	where

import Data.MyPolynomial.Type
import Data.Complex (realPart, imagPart, Complex)
import Data.List (sort)
import Data.IntMap.Lazy (lookupLT, lookupGE, IntMap)

 -- `procelain' functions are intended for canonical printing :
 -- i.e. Show instances, i.e. serialization, i.e. Read reciprocal

 -- `pretty' functions are intended for pretty printing only :
 -- i.e. absolutly no consistent Read <=> Show reciprocity.

 -- `prints' functions are general purpose printing functions;
 -- they either use `pretty' or `porcelain' to generate output.
 -- The 3rd person-like terminal 's' is a general Haskell convention
 -- for functions using a provided interface to do a specific job.

 -- S suffixed functions return a String difference list (ShowS)
 -- for efficient arbitrary concatenation.

 -- M suffixed functions take a polynomial represented by
 -- an IntMap Float instance.

---------------------------------------------------------------------------
----  Abstract printers
---------------------------------------------------------------------------

printsPolynomialS :: (Monomial -> ShowS) -> Polynomial -> ShowS
printsPolynomialS _ []  = ('0':)
printsPolynomialS pfS (mn:[]) = pfS mn
printsPolynomialS pfS (m1:(m2@(c2 :*^: p2)):ms) = pfS m1 . signBridge . next
 where
  signBridge
   | c2 < 0 = (" - " ++)
   | otherwise = (" + " ++)
  next
   | c2 < 0 = printsPolynomialS pfS (((-c2) :*^: p2):ms)
   | otherwise = printsPolynomialS pfS (m2:ms)

printsPolynomialSM :: (Polynomial -> ShowS) -> IntMap Float -> ShowS
printsPolynomialSM pfS mmap = ( pfS . sort ) $ consUp (smallest 0 mmap) mmap ([]++) []
  where
    smallest k map = case lookupLT k map of { Nothing -> k;	Just (k, _) -> smallest k map; }
    consUp k map dl = case lookupGE k map of	Nothing -> dl
						Just (kn, v) -> consUp (kn + 1) map (((v :*^: kn):) . dl)

printsEquationS :: (Polynomial -> ShowS) -> Equation -> ShowS
printsEquationS pfS (Eq (l, r)) = memb l . showString " = " . memb r
 where
  memb [] = ('0':)
  memb p  = pfS p

---------------------------------------------------------------------------
----  Canonical printing ( Read <=> Show )
---------------------------------------------------------------------------

porcelainMonomialS :: Monomial -> ShowS
porcelainMonomialS (c :*^: p) = shows c . showString " * X^" . shows p

porcelainPolynomialS :: Polynomial -> ShowS
porcelainPolynomialS = printsPolynomialS porcelainMonomialS

porcelainEquationS :: Equation -> ShowS
porcelainEquationS = printsEquationS porcelainPolynomialS

porcelainPolynomialSM :: IntMap Float -> ShowS
porcelainPolynomialSM = printsPolynomialSM porcelainPolynomialS

---------------------------------------------------------------------------
----  Pretty printing
---------------------------------------------------------------------------

prettyMonomialS :: Monomial -> ShowS
prettyMonomialS (0 :*^: _) = showString ""
prettyMonomialS (c :*^: 0) = shows c
prettyMonomialS (c :*^: 1) = shows c . showString " X"
prettyMonomialS (c :*^: p) = prettyMonomialS (c :*^: 1) . ('^':) . shows p

prettyMonomial :: Monomial -> String
prettyMonomial mn =  prettyMonomialS mn ""

prettyPolynomialS :: Polynomial -> ShowS
prettyPolynomialS = ( printsPolynomialS prettyMonomialS ) . purgeHollow
  where purgeHollow = filter ( \(c :*^: _) -> c /= 0 )

prettyPolynomial :: Polynomial -> String
prettyPolynomial p = prettyPolynomialS p ""

prettyPolynomialSM :: IntMap Float -> ShowS
prettyPolynomialSM = printsPolynomialSM prettyPolynomialS

prettyPolynomialM :: IntMap Float -> String
prettyPolynomialM m = prettyPolynomialSM m []

prettyEquationS :: Equation -> ShowS
prettyEquationS = printsEquationS (printsPolynomialS porcelainMonomialS)

prettyEquation :: Equation -> String
prettyEquation eq = prettyEquationS eq []

prettyComplexS :: (Ord c, Num c, Show c, RealFloat c) => Complex c -> ShowS
prettyComplexS c = shows a . img
	where
	 a = realPart c
	 img
	  | b < 0 = (" - " ++) . shows (-b) . ('i':)
	  | b > 0 = (" + " ++) . shows b . ('i':)
	  | b == 0 = showString ""
	  where
	   b = imagPart c

prettyComplex :: (Ord c, Num c, Show c, RealFloat c) => Complex c -> String
prettyComplex c = prettyComplexS c ""

---------------------------------------------------------------------------
----  Show instances
---------------------------------------------------------------------------

instance Show Monomial where
  showsPrec _ = porcelainMonomialS
  showList = porcelainPolynomialS

instance Show Equation where
  showsPrec _ = porcelainEquationS

ext: hs


text:
-----------------------------------------------------------------------------
-- |
-- Module      :  Finance.Hqfl.Pricer.Black
-- Copyright   :  (C) 2016 Mika'il Khan
-- License     :  (see the file LICENSE)
-- Maintainer  :  Mika'il Khan <co.kleisli@gmail.com>
-- Stability   :  stable
-- Portability :  portable
--
----------------------------------------------------------------------------
{-# LANGUAGE FlexibleInstances #-}

module Finance.Hqfl.Pricer.Black where

import Finance.Hqfl.Instrument
import Statistics.Distribution.Normal
import Data.Random

class Black a where
  price :: a -> Double -> Double -> Double

instance Black (Option Future) where
  price (Option (Future f) m European k t) r v =
    case m of
       Call -> exp (-r * t) * (f * cdf normal d1 - k * cdf normal d2)
       Put  -> exp (-r * t) * (k * cdf normal (-d2) - f * cdf normal (-d1))
    where d1 = (log (f / k) + ((v * v) / 2) * t) / (v * sqrt t)
          d2 = d1 - v * sqrt t
          normal = Normal (0 :: Double) 1

max_stars_repo_name: co-category/hqfl

max_issues_repo_name: cokleisli/hqfl

max_forks_repo_name: cokleisli/hqfl


text:
-- MandelBrot implementation in Haskell
import Data.Complex
import System.Environment
import Control.Parallel.Strategies
import Criterion.Measurement
--import Graphics.UI.SDL

--initSDL :: IO Surface
--initSDL = do
--    SDL.init [SDL.InitEverything]
--    SDL.setVideoMode (floor imageWidth) (floor imageHeight) 32 []
--    SDL.setCaption "Mandelbrot" "Mandelbrot"-   screen <- getVideoSurface
--    white <- SDL.mapRGB (SDL.surfaceGetPixelFormat screen) 255 255 255
--    fillRect screen Nothing white >> SDL.flip screen
--    return screen

mandelbrot :: Num a => a -> Int -> a
mandelbrot a n = iterate (\z -> z^2 + a) 0 !! n

mbrot_seq :: Int -> Int -> Int -> [[Bool]]
mbrot_seq iter w h = [[magnitude (mandelbrot (x :+ y) iter) < 2
                       | x <- [-2, (-2 + incW) ..   2]]
                       | y <- [ 2,  (2 - incH) ..  -2]]
 where incW = 4.0 / (fromIntegral w)
       incH = 4.0 / (fromIntegral h)

mbrot_par :: Int -> Int -> Int  -> [[Bool]]
mbrot_par iter w h = (mbrot_seq iter w h) `using` parList rdeepseq

draw_image :: Int -> Int -> [[Bool]] -> IO()
draw_image x y xs = print "to be implemented"

main = do
    initializeTime
    start_time <- getTime
    args <- getArgs
    if (length args < 4)
        then print "Expected \"par/seq\"iterations, image width, and image height"
        else do
            let nIter = read (args !! 1) :: Int
            let nWidth = read (args !! 2) :: Int
            let nHeight = read (args !! 3) :: Int
            case (args !! 0) of
                "seq" -> print $ mbrot_seq nIter nWidth nHeight
                "par" -> print $ mbrot_par nIter nWidth nHeight
                otherwise -> print "expected first argument to be `par` or `seq`"
            end_time <- getTime
            print $ "Time elapsed : " ++ show (floor $ 10000 * (end_time - start_time)) ++ "ms"

max_stars_repo_name: RossMeikleham/Parallel-Mandelbrot

max_issues_repo_name: RossMeikleham/Parallel-Mandelbrot

max_forks_repo_name: RossMeikleham/Parallel-Mandelbrot


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
module BrainWall.Random where

import BrainWall.Edge
import qualified BrainWall.Main.GenerateHole as GenHole
import qualified BrainWall.Main.GenerateFigure as GenFig
import BrainWall.Json
import BrainWall.Polygon
import BrainWall.Polygon.ContainsEdge
import BrainWall.Problem
import BrainWall.V2
import Control.Monad
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Types as Aeson
import Data.List
import Data.Ord
import Data.Ratio
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import Debug.Trace
import Numeric.GSL.Minimization
import Numeric.LinearAlgebra
import System.Directory
import System.Random.MWC as MWC
import System.Random.Stateful

chooseFrom :: StatefulGen g m => V.Vector x -> g -> m x
chooseFrom v gen = do
  idx <- uniformRM (0, V.length v - 1) gen
  return $ v V.! idx

genFigureFromPolygon :: StatefulGen g m => Polygon Integer -> g -> m Figure
genFigureFromPolygon polygon gen = do
  edges <- replicateM 50 genEdge
  return $ figureFromEdges $ filter isGood edges --
  where
    vertices = unPolygon polygon
    genEdge = do
      v1 <- chooseFrom vertices gen
      v2 <- chooseFrom vertices gen
      return $ Edge v1 v2
    isGood edge@(Edge v1 v2) = v1 /= v2 && containsEdge edge polygon

data EpsilonType
  = ForwardEpsilon -- The mutated figure will be the target figure
  | ReverseEpsilon -- The mutated figure will be the initial figure
  deriving (Enum, Bounded, Eq, Show)

data InitialPositionType
  = StartFigurePosition
  | RandomPosition
  deriving (Enum, Bounded, Eq, Show)

data ObjectiveType
  = Crumple
  | Stretch
  deriving (Enum, Bounded, Eq, Show)

data FigureSource
  = HandDrawn
  | PolyTree
  | FromHoleGenerator
  deriving (Enum, Bounded, Eq, Show)

randomGenerateFigure :: StatefulGen g IO => FigureSource -> g -> IO Figure
randomGenerateFigure PolyTree gen =
  figureFromEdges . GenFig.bfEdges <$> GenFig.generateFigure GenFig.defaultOptions gen
randomGenerateFigure FromHoleGenerator gen = do
  poly <- GenHole.generateHole GenHole.defaultOptions gen
  let problemHole = case GenHole.polyToHole poly of
        Just p  -> p
        Nothing -> error "error: no polygon"
  GenHole.genFigureInHole (GenHole.optionsFigure GenHole.defaultOptions) problemHole gen
randomGenerateFigure HandDrawn gen = do
  examples <- listDirectory "examples"
  let problems = V.fromList $ map ("examples/" ++) $ filter isProblem examples
  problemPath <- chooseFrom problems gen
  Just value <- Aeson.decodeFileStrict problemPath
  let figure = case Aeson.parse (\(Aeson.Object obj) -> obj Aeson..: T.pack "figure" >>= decodeFigure ) value of
        Aeson.Success p -> p
        Aeson.Error str -> error str
  pure figure
  where
    isProblem "" = False
    isProblem ".problem" = True
    isProblem (_:xs) = isProblem xs

randomMutateFigure ::
  forall m g.
  StatefulGen g m =>
  InitialPositionType ->
  EpsilonType ->
  ObjectiveType ->
  Figure ->
  (Integer, Integer) ->
  g ->
  m (Figure, Epsilon)
randomMutateFigure initialPositionType epsilonType objectiveType startFigure (canvasWidth, canvasHeight) gen = do
  initialVec <- case initialPositionType of
    StartFigurePosition -> pure $ V.map (fmap fromIntegral) $ figureVertices startFigure
    RandomPosition -> V.replicateM numVertices genPosition
  objective <- case objectiveType of
    Crumple -> do
      targetPositions <- V.replicateM numVertices genPosition
      pure $ negateObjective $ targetObjective targetPositions
    Stretch -> do
      vertexPairs <- genVertexPairs (V.length $ figureVertices startFigure) numPairs gen
      pure $ negateObjective $ distanceSumObjective vertexPairs
  return $ randomMutateFigureWithObjective initialVec epsilonType startFigure objective
  where
    numPairs = 50
    numVertices = V.length $ figureVertices startFigure
    genPosition :: m (V2 Double)
    genPosition = do
      x <- uniformRM (0, fromIntegral canvasWidth) gen
      y <- uniformRM (0, fromIntegral canvasHeight) gen
      return $ V2 x y

genVertexPairs :: forall m g. StatefulGen g m => Int -> Int -> g -> m (V.Vector (Int, Int))
genVertexPairs numVertices numPairs gen = do
  pairs <- replicateM numPairs genPair
  return $ V.fromList $ Set.toAscList $ Set.fromList $ filter (\(x, y) -> x /= y) pairs
  where
    genPair :: m (Int, Int)
    genPair = do
      x <- uniformRM (0, numVertices - 1) gen
      y <- uniformRM (0, numVertices - 1) gen
      return $ if x < y then (x, y) else (y, x)

data PositionObjectiveWithGradientHessian = PositionObjectiveWithGradientHessian
  { posObjective :: V.Vector (V2 Double) -> Double,
    posGradient :: V.Vector (V2 Double) -> VS.Vector Double,
    posHessian :: V.Vector (V2 Double) -> Matrix Double
  }

targetObjective :: V.Vector (V2 Double) -> PositionObjectiveWithGradientHessian
targetObjective targetPositions =
  PositionObjectiveWithGradientHessian
    { posObjective,
      posGradient,
      posHessian
    }
  where
    posObjective positions = V.sum $ V.zipWith dist targetPositions positions
    numVertices = V.length targetPositions
    posGradient positions =
      VS.convert $
        V.concat
          [ V.zipWith (\(V2 tx _) (V2 x _) -> 2 * (x - tx)) targetPositions positions,
            V.zipWith (\(V2 _ ty) (V2 _ y) -> 2 * (y - ty)) targetPositions positions
          ]
    posHessian _ = build (2 * numVertices, 2 * numVertices) $ \i j -> if i == j then 2 else 0
    dist v1 v2 = sq (v2X v1 - v2X v2) + sq (v2Y v1 - v2Y v2)
    sq x = x * x

negateObjective :: PositionObjectiveWithGradientHessian -> PositionObjectiveWithGradientHessian
negateObjective = scaleObjective (-1)

scaleObjective :: Double -> PositionObjectiveWithGradientHessian -> PositionObjectiveWithGradientHessian
scaleObjective mult oldObj =
  PositionObjectiveWithGradientHessian
    { posObjective = (*) mult . posObjective oldObj,
      posGradient = VS.map ((*) mult) . posGradient oldObj,
      posHessian = scale mult . posHessian oldObj
    }

distanceSumObjective :: V.Vector (Int, Int) -> PositionObjectiveWithGradientHessian
distanceSumObjective vertexPairs =
  PositionObjectiveWithGradientHessian
    { posObjective,
      posGradient,
      posHessian
    }
  where
    posObjective positions = V.sum $ V.map go vertexPairs
      where
        go (i, j) = dist posI posJ
          where
            posI = positions V.! i
            posJ = positions V.! j
    posGradient positions = VS.fromList $ map (go v2X) [0 .. numVertices - 1] ++ map (go v2Y) [0 .. numVertices - 1]
      where
        numVertices = V.length positions
        go get k = V.sum $ V.map go2 vertexPairs
          where
            go2 (i, j)
              | posI == posJ = 0
              | k == i = (get posI - get posJ) / dist posI posJ
              | k == j = (get posJ - get posI) / dist posI posJ
              | otherwise = 0
              where
                posI = positions V.! i
                posJ = positions V.! j
    posHessian positions = fromBlocks [[makeHessian goXX, makeHessian goXY], [makeHessian goXY, makeHessian goYY]]
      where
        numVertices = V.length positions
        get = (positions V.!)
        makeHessian aux =
          fromLists [[aux rowIdx colIdx | colIdx <- [0 .. numVertices - 1]] | rowIdx <- [0 .. numVertices - 1]]

        goXX rowIdx colIdx
          | rowIdx == colIdx =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else 1 / d - sq (v2X $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == rowIdx || j == rowIdx) vertexPairs
          | otherwise =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else - sq (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == min colIdx rowIdx && j == max colIdx rowIdx) vertexPairs
        goYY rowIdx colIdx
          | rowIdx == colIdx =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else 1 / d - sq (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == rowIdx || j == rowIdx) vertexPairs
          | otherwise =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else - sq (v2X $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == min colIdx rowIdx && j == max colIdx rowIdx) vertexPairs
        goXY rowIdx colIdx
          | rowIdx == colIdx =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else - (v2X $ get i .-. get j) * (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == rowIdx || j == rowIdx) vertexPairs
          | otherwise =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else (v2X $ get i .-. get j) * (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == min colIdx rowIdx && j == max colIdx rowIdx) vertexPairs

    dist v1 v2 = sqrt $ sq (v2X v1 - v2X v2) + sq (v2Y v1 - v2Y v2)
    sq x = x * x
    cube x = x * x * x

randomMutateFigureWithObjective :: V.Vector (V2 Double) -> EpsilonType -> Figure -> PositionObjectiveWithGradientHessian -> (Figure, Epsilon)
randomMutateFigureWithObjective initialPos epsilonType startFigure PositionObjectiveWithGradientHessian {posObjective, posGradient, posHessian} =
  traceShow ("rows", rows path) $
    traceShow ("objective", objective solution, epsilon) $
      traceShow ("distances", V.zip4 distances finalDistances roundedDistances epsilons) $
        traceShow (newPositions, newMultipliers) $
          traceShow ("gradients", lagrangian (newPositions, newMultipliers), lagrangianGradient (newPositions, newMultipliers), lagrangianHessian (newPositions, newMultipliers) #> lagrangianGradient (newPositions, newMultipliers)) $
            (Figure roundedPositions (figureEdges startFigure), epsilon)
  where
    initialVec = VS.convert $ V.map v2X initialPos V.++ V.map v2Y initialPos V.++ V.replicate numEdges 0
    lagrangian :: (V.Vector (V2 Double), V.Vector Double) -> Double
    lagrangian (positions, multipliers) =
      posObjective positions
        + V.sum (V.zipWith3 (\mult edge realDist -> mult * (distIdx positions edge - fromIntegral realDist) / sqrt (fromIntegral realDist)) multipliers (figureEdges startFigure) distances)
    lagrangianGradient :: (V.Vector (V2 Double), V.Vector Double) -> VS.Vector Double
    lagrangianGradient (positions, multipliers) =
      ( posGradient positions
          + VS.concat
            [ VS.map (edgeGradient v2X) (VS.enumFromN 0 numVertices),
              VS.map (edgeGradient v2Y) (VS.enumFromN 0 numVertices)
            ]
      )
        VS.++ V.convert (V.zipWith (\edge realDist -> (distIdx positions edge - fromIntegral realDist) / sqrt (fromIntegral realDist)) (figureEdges startFigure) distances)
      where
        edgeGradient get vIdx =
          V.sum $
            V.zipWith3
              ( \(v1Idx, v2Idx) mult realDist ->
                  case () of
                    _
                      | vIdx == v1Idx -> mult * 2 * (get (positions V.! vIdx) - get (positions V.! v2Idx)) / sqrt (fromIntegral realDist)
                      | vIdx == v2Idx -> mult * 2 * (get (positions V.! vIdx) - get (positions V.! v1Idx)) / sqrt (fromIntegral realDist)
                      | otherwise -> 0
              )
              (figureEdges startFigure)
              multipliers
              distances
    lagrangianHessian (positions, multipliers) =
      fromBlocks
        [ [topLeftHessian, tr' bottomLeftHessian],
          [bottomLeftHessian, multHessian]
        ]
      where
        multHessian :: Matrix Double
        multHessian = build (numEdges, numEdges) $ \_ _ -> 0
        xMultHessian get = build (numEdges, numVertices) $ \eIdx vIdx ->
          let realDist = distances V.! round eIdx
              (v1Idx, v2Idx) = figureEdges startFigure V.! round eIdx
           in case () of
                _
                  | round vIdx == v1Idx -> 2 / sqrt (fromIntegral realDist) * (get (positions V.! v1Idx) - get (positions V.! v2Idx))
                  | round vIdx == v2Idx -> 2 / sqrt (fromIntegral realDist) * (get (positions V.! v2Idx) - get (positions V.! v1Idx))
                  | otherwise -> 0
        topLeftHessian =
          fromBlocks
            [ [xxHessian, 0],
              [0, xxHessian]
            ]
            + posHessian positions
        bottomLeftHessian = fromBlocks [[xMultHessian v2X, xMultHessian v2Y]]
        xxHessian :: Matrix Double
        xxHessian = build (numVertices, numVertices) $ \i j ->
          V.sum $
            V.zipWith3
              ( \(v1Idx, v2Idx) mult realDist ->
                  case () of
                    _
                      | round i == v1Idx && round j == v1Idx -> mult * 2 / sqrt (fromIntegral realDist)
                      | round i == v2Idx && round j == v2Idx -> mult * 2 / sqrt (fromIntegral realDist)
                      | round i == v1Idx && round j == v2Idx -> - mult * 2 / sqrt (fromIntegral realDist)
                      | round i == v2Idx && round j == v1Idx -> - mult * 2 / sqrt (fromIntegral realDist)
                      | otherwise -> 0
              )
              (figureEdges startFigure)
              multipliers
              distances

    objective = VS.sum . VS.map sq . lagrangianGradient . splitParams . VS.convert
    -- (solution, path) =
    --   minimizeV
    --     NMSimplex2
    --     1e-9
    --     1000000
    --     (VS.replicate numParams 1)
    --     objective
    --     initialVec
    (solution, path) =
      minimizeVD
        VectorBFGS2
        1e-9
        1000
        0.01
        0.1
        objective
        ((\(positions, multipliers) -> scale 2 $ lagrangianHessian (positions, multipliers) #> lagrangianGradient (positions, multipliers)) . splitParams . VS.convert)
        initialVec
    -- initialVec = posInitial (V.map (fmap fromIntegral) $ figureVertices startFigure) VS.++ VS.replicate numEdges 0
    (newPositions, newMultipliers) = splitParams $ VS.convert solution
    numRotations = 200 :: Int
    numShifts = 20 :: Int
    (roundedPositions, roundedDistances, epsilons, epsilon) =
      minimumBy (comparing $ \(_, _, _, x4) -> x4) $
        [tryShiftRotation (V2 shiftX shiftY) angle | angle <- angles, shiftX <- shifts, shiftY <- shifts]
      where
        angles = [fromIntegral x / fromIntegral numRotations * pi / 2 | x <- [0 .. (numRotations - 1)]]
        shifts = [fromIntegral x / fromIntegral numShifts | x <- [0 .. (numShifts - 1)]]
    tryShiftRotation shift angle = (pos, roundedDists, epss, eps)
      where
        pos = V.map (fmap round) shifted
        shifted = V.map (.+. shift) rotated
        rotated = V.map (\(V2 x y) -> (V2 (cos angle * x - sin angle * y) (sin angle * x + cos angle * y))) newPositions
        roundedDists = V.map (distIdx pos) (figureEdges startFigure)
        eps :: Epsilon
        eps = ceiling $ V.maximum epss * fromIntegral epsilonDenominator
        epss = case epsilonType of
          ForwardEpsilon -> V.zipWith getEpsilon distances roundedDists
          ReverseEpsilon -> V.zipWith (flip getEpsilon) distances roundedDists
        reallyBig = fromIntegral (1000000000000 :: Int)
        getEpsilon refDist finalDistance = if refDist == 0 then reallyBig else
          (case compare finalDistance refDist of
            EQ -> 0
            GT -> finalDistance % refDist - 1
            LT -> negate $ finalDistance % refDist - 1)
    sq x = x * x
    numVertices = V.length $ figureVertices startFigure
    numEdges = V.length $ figureEdges startFigure
    splitParams params = (positions, multipliers)
      where
        multipliers = V.drop (2 * numVertices) params
        xs = V.take numVertices params
        ys = V.take numVertices $ V.drop numVertices params
        positions = V.zipWith V2 xs ys
    distances = V.map (distIdx (figureVertices startFigure)) (figureEdges startFigure)
    finalDistances = V.map (distIdx newPositions) (figureEdges startFigure)
    dist v1 v2 = sq (v2X v1 - v2X v2) + sq (v2Y v1 - v2Y v2)
    distIdx positions (v1Idx, v2Idx) = dist v1 v2
      where
        v1 = positions V.! v1Idx
        v2 = positions V.! v2Idx


text:
{-# OPTIONS_GHC -fwarn-unused-imports #-}

module Main(main) where

import Control.Monad.Par
import Control.Monad.Par.Internal (runParAsync)
import Control.Monad.Par.Stream as S
import Control.Monad.Par.Logging
import Control.Exception
import Data.Complex
import GHC.Conc as Conc
import Debug.Trace
import Math.FFT (dft)

type Elt = Complex Double

fft_kern :: Window Elt -> Window Elt
fft_kern arr = dft arr

--   -- TEMP, evaluate one element to make sure the fft really gets called:
-- --  trace ("One elt sample: "++ show (arr!10)) $
--   case arr2 ! 10 of _ -> arr2
--  where arr2 = dft arr

--------------------------------------------------------------------------------
-- Main script

-- target  = maxBound
target  = 10 * 1000 * 1000
bufsize = 1024

main = do
  putStrLn$ "numCapabilities: "++ show numCapabilities
  putStrLn$ "  Frequency in measurable ticks:  "++ commaint oneSecond ++ "\n"

  putStrLn$ "Performing FFT of "++ commaint target ++" numbers windowed into buffers of size "++ show bufsize ++"\n"

  results <- evaluate $ runParAsync$ do

       strm1 <- countupWin bufsize target :: Par (WStream Elt)

       print_$ "\n Next, applying FFT filter... "
       strm2 <- streamMap fft_kern strm1

-- Make a pipeline of 10 stages:
--       strm2 <- foldl (\ s _ -> streamMap kern0) strm1 [1..10]

       print_$ "\n Stream graph constructed, returning from Par computation... "
       return strm2

  measureRate results
  putStrLn$ "End of stream reached.  All done."

print_ msg = trace msg $ return ()

-- work pop 1 peek N push 1
-- float->float filter
-- firFilter n coefs =
-- {

--     float sum = 0;
--     for (int i = 0; i < N; i++)
--       sum += peek(i) * COEFF[N-1-i];
--     pop();
--     push(sum);
--   }
-- }

{-
 Notes:

 -}


text:
{-# LANGUAGE ForeignFunctionInterface #-}

module Haskstat where

import Data.List
import Data.Maybe
import Data.Complex

foreign import ccall "erf" c_erf :: Double -> Double

floatLength :: Fractional a => [b] -> a
floatLength xs = fromIntegral $ length xs

mean :: Fractional a => [a] -> a
mean xs = (foldl' (+) 0 xs) / floatLength xs

trimmedMean :: (Ord a, Fractional a) => [a] -> Int -> a
trimmedMean xs 0 = mean xs
trimmedMean xs m = mean $ take (n-2*m) $ drop m (sort xs)
    where
        n = length xs

var :: Fractional a => [a] -> a
var xs =  (x2Sum/n - (xSum/n)^2)
    where
        (xSum, x2Sum) = foldl' (\(f,s) x -> (f+x,s+x^2)) (0,0) xs
        n = floatLength xs

varSample :: Fractional a => [a] -> a
varSample xs = n/(n-1) * var xs
    where
        n = floatLength xs

std :: Floating a => [a] -> a
std x = sqrt $ var x

stdSample :: Floating a => [a] -> a
stdSample x = sqrt $ varSample x

skewness :: Floating a => [a] -> a
skewness xs = (mu3/n) / (mu2/n)**(1.5)
    where
        mu = mean xs
        (mu2, mu3) = foldl' (\(s,t) x -> (s + (x-mu)^2, t + (x-mu)^3)) (0, 0) xs
        n = floatLength xs

quartileSkewness :: (RealFrac a, Floating a) => [a] -> a
quartileSkewness xs = (q3 + q1 - 2.0*q2)/(q3 - q1)
    where
        q1 = percentile xs 25
        q2 = median xs
        q3 = percentile xs 75

octileSkewness :: (RealFrac a, Floating a) => [a] -> a
octileSkewness xs = (q875 + q125 - 2.0*q2)/(q875 - q125)
    where
        q125 = percentile xs 12.5
        q2 = median xs
        q875 = percentile xs 87.5

substractValue :: Floating a => [a] -> a -> [a]
substractValue (x:xs) v = abs(x-v):(substractValue xs v)
substractValue [] v = []

mad :: (RealFrac a, Floating a) => [a] -> a
mad xs = median dev
    where
        dev = substractValue xs $ median xs

--medcouple :: (Ord a, Floating a, RealFrac a) => [a] -> a
--medcouple xs = med
--	where
--		sorted_xs = reverse $ sort xs
--		med = median sorted_xs
--		zs = sorted_xs - med
--		h_zi_zj = [
--		x_plus = filter (med<) xs
--		x_minus = filter (med>) xs

rank :: (Enum a, Num a, Ord a, Fractional a) => [a] -> [a]
rank xs = fromJust <$> map (\k -> lookup k pairs) xs
    where
            g = groupBy (\a b -> fst a == fst b) $ zip (sort xs) [1..]
            r = map (\k -> mean $ map snd k) g
            pairs = zip (nub (sort xs)) r

kurtosis :: Floating a => [a] -> a
kurtosis xs = n * mu4 / mu2^2
    where
        mu = mean xs
        (mu2, mu4) = foldl' (\(s, f) x -> (s + (x-mu)^2, f + (x-mu)^4)) (0, 0) xs
        n = floatLength xs

robustKurtosis :: (RealFrac a, Floating a) => [a] -> a
robustKurtosis xs = ((q875 - q625) + (q375 - q125))/(q75 - q25)
    where
        q125 = percentile xs 12.5
        q25  = percentile xs 25
        q375 = percentile xs 37.5
        q625 = percentile xs 62.5
        q75  = percentile xs 75
        q875 = percentile xs 87.5

percentile :: (Ord a, Fractional a, RealFrac a) => [a] -> a -> a
percentile xs 0 = minimum xs
percentile xs 100 = maximum xs
percentile xs perc = (xSorted!!idx) * (1-fraction) + (xSorted!!(idx+1)) * (fraction)
    where
        xSorted = sort xs
        perc' = ( perc)/100.0
        n = floatLength xs
        idx = floor (perc'*(n-1)) :: Int
        fraction = perc'*(n-1) - fromIntegral idx

median :: (Ord a, Fractional a, RealFrac a) => [a] -> a
median xs = percentile xs 50

data Summary a = Summary {
        min_ :: a,
        p25 :: a,
        med :: a,
        p75 :: a,
        max_ :: a
        } deriving (Show)

summary :: (Ord a, Fractional a, RealFrac a) => [a] -> Summary a
summary xs = Summary {
        min_ = minimum xs,
        p25 = percentile xs 25,
        med = median xs,
        p75 = percentile xs 75,
        max_ = maximum xs
        }

count :: (Ord a, Num a) => [a] -> [(a,Int)]
count x = zip (nub x_sorted) (map length (group x_sorted))
    where
        x_sorted = sort x

binning :: (Ord a, Fractional a) => [a] -> Int -> [((a,a), Int)]
binning [] _ = []
binning x nbrBin = zip step $ map length filters
    where
        step = steps ((maxX - minX) / fromIntegral nbrBin) minX maxX
        maxX = maximum x
        minX   = minimum x
        bounds  = map (\(a,b) -> (&&) <$> (>=a) <*> (<b)) (init step) ++
                     [(\(a,b) -> (&&) <$> (>=a) <*> (<=b)) (last step)]
        filters = map (flip filter x) bounds

steps :: (Ord a, Num a) => a -> a -> a -> [(a,a)]
steps step start stop
    | start < (stop-step) = (start, start+step) : steps step (start+step) stop
    | otherwise = [(start,stop)]

jarqueBera :: [Double] -> (Double, Double)
jarqueBera xs = (value, exp (-0.5*value))
    where
        n = fromIntegral $ length xs
        mu = mean xs
        (mu2, mu3, mu4) = foldl' (\(s,t,f) x -> (s+(x-mu)^2, t+(x-mu)^3, f+(x-mu)^4)) (0,0,0) xs
        s = (mu3/n)/(mu2/n)**1.5
        k = n * mu4 /mu2^2
        value = n / 6.0 * (s^2 + 0.25*(k - 3.0)^2)

fft :: (Enum a, Floating a, RealFloat a) => [Complex a] -> [Complex a]
fft xs@(y:ys)
    | n > 1 = zipWith (+) even_part factor1 ++ zipWith (-) even_part factor1
    | otherwise = [y]
    where
        n = floatLength xs
        factor = expBasis n
        factor1 = zipWith (*) factor odd_part
        even_part = fft $ evens xs
        odd_part = fft $ odds xs

ifftUnnormalised :: (Enum a, Floating a, RealFloat a) => [Complex a] -> [Complex a]
ifftUnnormalised xs@(y:ys)
    | n > 1 = zipWith (+) even_part factor1 ++ zipWith (-) even_part factor1
    | otherwise = [y]
    where
        n = floatLength xs
        factor = conjugate <$> expBasis n
        factor1 = zipWith (*) factor odd_part
        even_part = ifftUnnormalised $ evens xs
        odd_part = ifftUnnormalised $ odds xs

ifft :: (Enum a, Floating a, RealFloat a) => [Complex a] -> [Complex a]
ifft xs = (/n) <$> (ifftUnnormalised xs)
    where
        n = floatLength xs

autocorr :: (Enum a, Floating a, RealFloat a) => Int -> [Complex a] -> [a]
autocorr nlags xs = (/(head acovf)) <$> acovf
    where
        n = floatLength xs
        fr = fft $ zeroPadding (nextRegular (2*n+1)) (normalise xs)
        acovf = take nlags (realPart <$> (ifft $ zipWith (*) fr (map conjugate fr)))

zeroPadding :: (Num a) => Double -> [a] -> [a]
zeroPadding m xs
    | n < m = zeroPadding m (xs ++ [0])
    | otherwise = xs
    where
        n = floatLength xs

nextRegular :: (RealFrac a, Floating a) => a -> a
nextRegular m = 2 ^ exponent
    where
        exponent = ceiling $ logBase 2 m

-- zeroPadding :: [Double] -> [Double]
-- zeroPadding xs
--      | n < m = zeroPadding (xs ++ [0])
--      | otherwise = xs
--      where
--         n = fromIntegral $ length xs
--         exponent =  ceiling $ logBase 2 n
--         m = 2.0 ^ exponent

evens :: [a] -> [a]
evens (x:z:xs) = x : evens xs
evens (x:xs) = [x]
evens _ = []

odds :: [a] -> [a]
odds (x:z:xs) = z : odds xs
odds _ = []

expBasis :: (Enum a, RealFloat a, Num a) => a -> [Complex a]
expBasis n = map g [0..(n-1)]
    where
        exponent = (0 :+ ) <$> ((-2.0)*pi / n *)
        g = exp . exponent

normalise :: (Floating a) => [a] -> [a]
normalise xs = [(x - mu) / sigma | x <- xs]
    where
        mu = mean xs
        sigma = std xs

max_stars_repo_name: kevintyloo/haskstat

max_issues_repo_name: vekrt/haskstat

max_forks_repo_name: vekrt/haskstat


text:
import ImageLoader
import Brain

--------------------------------------------------------------------------------

import Control.Exception (SomeException (..))
import System.Random     (mkStdGen, randomR)
import Numeric.LinearAlgebra
import Data.List
import qualified Numeric.LinearAlgebra.Data as V
import Debug.Trace

--------------------------------------------------------------------------------

type Samples = [(V.Vector Double, V.Vector Double)]

main :: IO ()
main = loadFiles

loadFiles :: IO ()
loadFiles = do
    ex    <- loadExamples
    notEx <- loadNonExamples
    case (ex, notEx)
      of (Left (SomeException e), _) -> printError e
         (_, Left (SomeException e)) -> printError e
         (Right v1, Right v2)        -> runTraining v1 v2
  where
    printError e = putStrLn $ "Exception: " ++ (show e)

runTraining :: Samples -> Samples -> IO ()
runTraining s1 s2 = do
    net <- constructNeuralNetwork 256 [512, 256, 128, 64] 1 all_
    putStrLn "----"
    printSamples all_ net
    putStrLn "----"
    let net' = trainN net 800
    printSamples all_ net'
  where all_   = s1 ++ s2

printSamples :: Samples -> NetworkWithSamples -> IO ()
printSamples [] _       = return ()
printSamples (x:xs) net = do
    putStrLn $ vid ++ " --> " ++ output ++ " ? " ++ expectedOut
    printSamples xs net
  where in_    = fst x
        expectedOut = show $ snd x
        vid    = show $ (vecId in_)
        output = show $ (activate net in_)

-- Given a list of examples, separates into a list of
-- test data (fst) and training data (snd)
sepExTest :: [a] -> ([a], [a])
sepExTest xs = go xs ([], []) (mkStdGen 10)
  where
        go [] a _            = a
        go (v:vs) (a, b) gen = case randomR (0, 99) gen
                               of (n, gen') -> if n < (20 :: Int)
                                                  then go vs (v:a, b) gen'
                                                  else go vs (a, v:b) gen'

vecId :: V.Vector Double -> Double
vecId v = v <.> V.fromList (replicate 256 (1 :: Double))


text:
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DataKinds #-}

module SimpleR.Interpreter.Commons.Vector
  ( Type (..)
  -- , SInt (..)
  -- , SDouble (..)
  -- , SComplex (..)
  -- , SString (..)
  -- , SBool (..)
  , Atom (..)
  , Vector (..)
  , idFromAtomString
  , vecFromConst
  , vecInts
  , vecDoubles
  , vecComplexes
  , vecStrings
  , vecBools
  , vecLength
  , vecResize
  , vecType
  , vecToType
  , joinType
  , vecPairJoinType
  , vecPairLengthen
  , vecPairLengthenJoinType
  ) where

import Data.Maybe
import qualified Data.Complex as C

import SimpleR.Language
import SimpleR.Smt

data Type = BoolTy | IntTy | DoubleTy | ComplexTy | StringTy | NullTy
  deriving (Eq, Show, Read)

data Atom a where
  Atom :: a -> Atom a
  NAtom :: Atom a

instance (Eq a) => Eq (Atom a) where
  (Atom a) == (Atom b) = a == b
  (NAtom) == (NAtom) = True
  _ == _ = False

instance (Show a) => Show (Atom a) where
  show (Atom a) = "Atom (" ++ show a ++ ")"
  show (NAtom) = "NAtom"

instance (Read a) => Read (Atom a) where
  readsPrec _ = error "readsPrec: Atom read unimplemented"

instance Functor Atom where
  fmap _ NAtom = NAtom
  fmap f (Atom a) = Atom (f a)

instance Applicative Atom where
  pure a = Atom a

  (NAtom) <*> _ = NAtom
  (Atom f) <*> atom = fmap f atom

instance Monad Atom where
  return a = Atom a

  (NAtom) >>= _ = NAtom
  (Atom a) >>= f = f a

data Vector =
    IntVec [Atom Int]
  | DoubleVec [Atom Double]
  | ComplexVec [Atom Complex]
  | BoolVec [Atom Bool]
  | StringVec [Atom String]
  | SymVec SmtIdent Type Int
  | NilVec
  deriving (Eq, Show, Read)

idFromAtomString :: Atom String -> Ident
idFromAtomString (Atom str) = idFromString str

-- Vector conversion
vecFromConst :: Const -> Vector
vecFromConst (IntConst int) = IntVec [Atom int]
vecFromConst (DoubleConst double) = DoubleVec [Atom double]
vecFromConst (ComplexConst complex) = ComplexVec [Atom complex]
vecFromConst (BoolConst bool) = BoolVec [Atom bool]
vecFromConst (StringConst str) = StringVec [Atom str]
vecFroMConst (NaConst) = NilVec

vecLength :: Vector -> Int
vecLength (IntVec xs) = length xs
vecLength (DoubleVec xs) = length xs
vecLength (ComplexVec xs) = length xs
vecLength (BoolVec xs) = length xs
vecLength (StringVec xs) = length xs
vecLength (SymVec _ _ len) = len -- error $ "vecLength: called with SymVec"

vecType :: Vector -> Type
vecType (IntVec _) = IntTy
vecType (DoubleVec _) = DoubleTy
vecType (ComplexVec _) = ComplexTy
vecType (BoolVec _) = BoolTy
vecType (StringVec _) = StringTy
vecType (SymVec _ ty _) = ty
vecType (NilVec) = NullTy

vecInts :: Vector -> Maybe [Atom Int]
vecInts (IntVec xs) = Just xs
vecInts _ = Nothing

vecDoubles :: Vector -> Maybe [Atom Double]
vecDoubles (DoubleVec xs) = Just xs
vecDoubles _ = Nothing

vecComplexes :: Vector -> Maybe [Atom Complex]
vecComplexes (ComplexVec xs) = Just xs
vecComplexes _ = Nothing

vecStrings :: Vector -> Maybe [Atom String]
vecStrings (StringVec xs) = Just xs
vecStrings _ = Nothing

vecBools :: Vector -> Maybe [Atom Bool]
vecBools (BoolVec xs) = Just xs
vecBools _ = Nothing

------
-- Resizing

snip :: Int -> [a] -> [a]
snip n xs = take n $ concat $ repeat xs

vecResize :: Int -> Vector -> Vector
vecResize n (IntVec (x : xs)) = IntVec $ snip n (x : xs)
vecResize n (DoubleVec (x : xs)) = DoubleVec $ snip n (x : xs)
vecResize n (ComplexVec (x : xs)) = ComplexVec $ snip n (x : xs)
vecResize n (StringVec (x : xs)) = StringVec $ snip n (x : xs)
vecResize n (BoolVec (x : xs)) = BoolVec $ snip n (x : xs)
vecResize _ (IntVec []) = error $ "vecLength: IntVec has length 0"
vecResize _ (DoubleVec []) = error $ "vecLength: DoubleVec has length 0"
vecResize _ (ComplexVec []) = error $ "vecLength: ComplexVec has length 0"
vecResize _ (StringVec []) = error $ "vecLength: StringVec has length 0"
vecResize _ (BoolVec []) = error $ "vecLength: BoolVec has length 0"
vecResize _ (NilVec) = error $ "vecLength: called with NilVec"
vecResize n (SymVec id ty len) =
  SymVec id ty n
  -- error $ "vecLength: called with SymVec"

-----
-- Conversion

maybeRead :: (Read a) => String -> Maybe a
maybeRead str = fmap fst $ listToMaybe $ reads str

-- Int
intFromInt :: Int -> Int
intFromInt int = int

intFromDouble :: Double -> Int
intFromDouble double = round double

intFromComplex :: Complex -> Int
intFromComplex complex = round $ C.realPart complex

intFromString :: String -> int
intFromString = error $ "intFromString: TODO"

intFromBool :: Bool -> Int
intFromBool bool = if bool then 1 else 0

-- Double
doubleFromInt :: Int -> Double
doubleFromInt int = fromIntegral int

doubleFromDouble :: Double -> Double
doubleFromDouble double = double

doubleFromComplex :: Complex -> Double
doubleFromComplex complex = C.realPart complex

doubleFromString :: String -> Double
doubleFromString = error "doubleFromString: TODO"

doubleFromBool :: Bool -> Double
doubleFromBool bool = if bool then 1.0 else 0.0

-- Complex
complexFromInt :: Int -> Complex
complexFromInt int = (fromIntegral int) C.:+ 0.0

complexFromDouble :: Double -> Complex
complexFromDouble double = double C.:+ 0.0

complexFromComplex :: Complex -> Complex
complexFromComplex complex = complex

complexFromString :: String -> Complex
complexFromString = error "complexFromString: TODO"

complexFromBool :: Bool -> Complex
complexFromBool bool = if bool then 1.0 C.:+ 0.0 else 0.0 C.:+ 0.0

-- String
stringFromInt :: Int -> String
stringFromInt int = show int

stringFromDouble :: Double -> String
stringFromDouble double = show double

stringFromComplex :: Complex -> String
stringFromComplex complex = show complex

stringFromString :: String -> String
stringFromString string = string

stringFromBool :: Bool -> String
stringFromBool bool = show bool

-- Bool
boolFromInt :: Int -> Bool
boolFromInt int = int /= 0

boolFromDouble :: Double -> Bool
boolFromDouble double = double /= 0.0

boolFromComplex :: Complex -> Bool
boolFromComplex complex = complex /= (0.0 C.:+ 0.0)

boolFromString :: String -> Bool
boolFromString = error $ "boolFromString: TODO"

boolFromBool :: Bool -> Bool
boolFromBool bool = bool

vecToType :: Vector -> Type -> Vector
vecToType (IntVec xs) IntTy = IntVec $ (map . fmap) intFromInt xs
vecToType (IntVec xs) DoubleTy = DoubleVec $ (map . fmap) doubleFromInt xs
vecToType (IntVec xs) ComplexTy = ComplexVec $ (map . fmap) complexFromInt xs
vecToType (IntVec xs) StringTy = StringVec $ (map . fmap) stringFromInt xs
vecToType (IntVec xs) BoolTy = BoolVec $ (map . fmap) boolFromInt xs

vecToType (DoubleVec xs) IntTy = IntVec $ (map . fmap) intFromDouble xs
vecToType (DoubleVec xs) DoubleTy =
  DoubleVec $ (map . fmap) doubleFromDouble xs
vecToType (DoubleVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromDouble xs
vecToType (DoubleVec xs) StringTy =
  StringVec $ (map . fmap) stringFromDouble xs
vecToType (DoubleVec xs) BoolTy = BoolVec $ (map . fmap) boolFromDouble xs

vecToType (ComplexVec xs) IntTy = IntVec $ (map . fmap) intFromComplex xs
vecToType (ComplexVec xs) DoubleTy =
  DoubleVec $ (map . fmap) doubleFromComplex xs
vecToType (ComplexVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromComplex xs
vecToType (ComplexVec xs) StringTy =
  StringVec $ (map . fmap) stringFromComplex xs
vecToType (ComplexVec xs) BoolTy = BoolVec $ (map . fmap) boolFromComplex xs

vecToType (StringVec xs) IntTy = IntVec $ (map . fmap) intFromString xs
vecToType (StringVec xs) DoubleTy =
  DoubleVec $ (map . fmap) doubleFromString xs
vecToType (StringVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromString xs
vecToType (StringVec xs) StringTy =
  StringVec $ (map . fmap) stringFromString xs
vecToType (StringVec xs) BoolTy = BoolVec $ (map . fmap) boolFromString xs

vecToType (BoolVec xs) IntTy = IntVec $ (map . fmap) intFromBool xs
vecToType (BoolVec xs) DoubleTy = DoubleVec $ (map . fmap) doubleFromBool xs
vecToType (BoolVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromBool xs
vecToType (BoolVec xs) StringTy = StringVec $ (map . fmap) stringFromBool xs
vecToType (BoolVec xs) BoolTy = BoolVec $ (map . fmap) boolFromBool xs

vecToType (SymVec sym _ len) ty = SymVec sym ty len
vecToType (NilVec) _ = NilVec

-- Type lattice
--  BoolTy <= IntTy <= Double <= Complex <= String <= Null
instance Ord Type where
  BoolTy `compare` BoolTy = EQ
  BoolTy `compare` _ = LT

  IntTy `compare` IntTy = EQ
  IntTy `compare` _ = LT

  DoubleTy `compare` DoubleTy = EQ
  DoubleTy `compare` _ = LT

  ComplexTy `compare` ComplexTy = EQ
  ComplexTy `compare` _ = LT

  StringTy `compare` StringTy = EQ
  StringTy `compare` _ = LT

  NullTy `compare` NullTy = EQ
  NullTy `compare` _ = GT

joinType :: Type -> Type -> Type
joinType ty1 ty2 = if ty1 < ty2 then ty2 else ty1

vecPairJoinType :: (Vector, Vector) -> (Vector, Vector)
vecPairJoinType (vec1, vec2) =
  let joinTy = joinType (vecType vec1) (vecType vec2) in
    (vecToType vec1 joinTy, vecToType vec2 joinTy)

vecPairLengthen :: (Vector, Vector) -> (Vector, Vector)
vecPairLengthen (vec1, vec2) =
  let len1 = vecLength vec1 in
  let len2 = vecLength vec2 in
    if len1 > len2 then
      (vec1, vecResize len1 vec2)
    else
      (vecResize len2 vec1, vec2)

vecPairLengthenJoinType :: (Vector, Vector) -> (Vector, Vector)
vecPairLengthenJoinType (vec1, vec2) =
  vecPairLengthen $ vecPairJoinType (vec1, vec2)

max_stars_repo_name: AntonXue/simple-r

max_issues_repo_name: AntonXue/simple-r

max_forks_repo_name: AntonXue/simple-r


text:
import Data.Complex

#include "../Prelude.hs"

f :: Int -> Complex Double
f n = mkPolar 1 ((2 Prelude.* Prelude.pi) / Prelude.fromIntegral n) Prelude.^ n

realPart'Double = realPart :: Complex Double -> Double
add'ComplexDouble = (Prelude.+) :: Complex Double -> Complex Double -> Complex Double
lit0'ComplexDouble = 0 :: Complex Double

max_stars_repo_name: batterseapower/chsc

max_issues_repo_name: batterseapower/chsc

max_forks_repo_name: batterseapower/chsc


text:
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2014
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- Root finding using Halley's rational method (the second in
-- the class of Householder methods). Assumes the function is three
-- times continuously differentiable and converges cubically when
-- progress can be made.
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Halley
  (
  -- * Halley's Method (Tower AD)
    findZero
  , inverse
  , fixedPoint
  , extremum
  ) where

import Prelude hiding (all)
import Numeric.AD.Internal.Forward (Forward)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Tower (Tower)
import Numeric.AD.Mode
import Numeric.AD.Rank1.Tower (diffs0)
import Numeric.AD.Rank1.Forward (diff)

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,1.8571428571428572,1.9997967892704736,1.9999999999994755,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
findZero f = go where
  go x = x : if x == xn then [] else go xn where
    (y:y':y'':_) = diffs0 f x
    xn = x - 2*y*y'/(2*y'*y'-y*y'')
{-# INLINE findZero #-}

-- | The 'inverse' function inverts a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Note: the @take 10 $ inverse sqrt 1 (sqrt 10)@ example that works for Newton's method
-- fails with Halley's method because the preconditions do not hold!
inverse :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> a -> [a]
inverse f x0 y = findZero (\x -> f x - auto y) x0
{-# INLINE inverse  #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Halley's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
fixedPoint f = findZero (\x -> f x - x)
{-# INLINE fixedPoint #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Halley's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- >>> take 10 $ extremum cos 1
-- [1.0,0.29616942658570555,4.59979519460002e-3,1.6220740159042513e-8,0.0]
extremum :: (Fractional a, Eq a) => (On (Forward (Tower a)) -> On (Forward (Tower a))) -> a -> [a]
extremum f = findZero (diff (off . f . On))
{-# INLINE extremum #-}

max_stars_repo_name: silky/ad

max_issues_repo_name: silky/ad

max_forks_repo_name: silky/ad


text:
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeSynonymInstances #-}

#if __GLASGOW_HASKELL__ >= 701
{-# LANGUAGE DefaultSignatures #-}
#endif

#if __GLASGOW_HASKELL__ >= 705
{-# LANGUAGE PolyKinds #-}
#endif

#if __GLASGOW_HASKELL__ >= 708
{-# LANGUAGE EmptyCase #-}
#endif

#if __GLASGOW_HASKELL__ >= 710
{-# LANGUAGE Safe #-}
#elif __GLASGOW_HASKELL__ >= 701
{-# LANGUAGE Trustworthy #-}
#endif

module Generics.Deriving.Functor (
  -- * Generic Functor class
    GFunctor(..)

  -- * Default method
  , gmapdefault

  -- * Internal Functor class
  , GFunctor'(..)

  ) where

import           Control.Applicative (Const, ZipList)

import qualified Data.Monoid as Monoid (First, Last, Product, Sum)
import           Data.Monoid (Dual)

import           Generics.Deriving.Base

#if MIN_VERSION_base(4,4,0)
import           Data.Complex (Complex)
#endif

#if MIN_VERSION_base(4,6,0)
import           Data.Ord (Down)
#else
import           GHC.Exts (Down)
#endif

#if MIN_VERSION_base(4,7,0)
import           Data.Proxy (Proxy)
#endif

#if MIN_VERSION_base(4,8,0)
import           Data.Functor.Identity (Identity)
import           Data.Monoid (Alt)
#endif

#if MIN_VERSION_base(4,9,0)
import qualified Data.Functor.Product as Functor (Product)
import qualified Data.Functor.Sum as Functor (Sum)
import           Data.List.NonEmpty (NonEmpty)
import qualified Data.Semigroup as Semigroup (First, Last)
import           Data.Semigroup (Arg, Max, Min, Option, WrappedMonoid)
#endif

--------------------------------------------------------------------------------
-- Generic fmap
--------------------------------------------------------------------------------

class GFunctor' f where
  gmap' :: (a -> b) -> f a -> f b

instance GFunctor' V1 where
  gmap' _ x = case x of
#if __GLASGOW_HASKELL__ >= 708
                {}
#else
                !_ -> error "Void gmap"
#endif

instance GFunctor' U1 where
  gmap' _ U1 = U1

instance GFunctor' Par1 where
  gmap' f (Par1 a) = Par1 (f a)

instance GFunctor' (K1 i c) where
  gmap' _ (K1 a) = K1 a

instance (GFunctor f) => GFunctor' (Rec1 f) where
  gmap' f (Rec1 a) = Rec1 (gmap f a)

instance (GFunctor' f) => GFunctor' (M1 i c f) where
  gmap' f (M1 a) = M1 (gmap' f a)

instance (GFunctor' f, GFunctor' g) => GFunctor' (f :+: g) where
  gmap' f (L1 a) = L1 (gmap' f a)
  gmap' f (R1 a) = R1 (gmap' f a)

instance (GFunctor' f, GFunctor' g) => GFunctor' (f :*: g) where
  gmap' f (a :*: b) = gmap' f a :*: gmap' f b

instance (GFunctor f, GFunctor' g) => GFunctor' (f :.: g) where
  gmap' f (Comp1 x) = Comp1 (gmap (gmap' f) x)

-- instance GFunctor' UAddr where
--   gmap' _ (UAddr a) = UAddr a

-- instance GFunctor' UChar where
--   gmap' _ (UChar c) = UChar c

-- instance GFunctor' UDouble where
--   gmap' _ (UDouble d) = UDouble d

-- instance GFunctor' UFloat where
--   gmap' _ (UFloat f) = UFloat f

-- instance GFunctor' UInt where
--   gmap' _ (UInt i) = UInt i

-- instance GFunctor' UWord where
--   gmap' _ (UWord w) = UWord w

class GFunctor f where
  gmap :: (a -> b) -> f a -> f b
#if __GLASGOW_HASKELL__ >= 701
  default gmap :: (Generic1 f, GFunctor' (Rep1 f))
               => (a -> b) -> f a -> f b
  gmap = gmapdefault
#endif

gmapdefault :: (Generic1 f, GFunctor' (Rep1 f))
            => (a -> b) -> f a -> f b
gmapdefault f = to1 . gmap' f . from1

-- Base types instances
instance GFunctor ((->) r) where
  gmap = fmap

instance GFunctor ((,) a) where
  gmap = gmapdefault

instance GFunctor [] where
  gmap = gmapdefault

#if MIN_VERSION_base(4,8,0)
instance GFunctor f => GFunctor (Alt f) where
  gmap = gmapdefault
#endif

#if MIN_VERSION_base(4,9,0)
instance GFunctor (Arg a) where
  gmap = gmapdefault
#endif

#if MIN_VERSION_base(4,4,0)
instance GFunctor Complex where
  gmap = gmapdefault
#endif

instance GFunctor (Const m) where
  gmap = gmapdefault

instance GFunctor Down where
  gmap = gmapdefault

instance GFunctor Dual where
  gmap = gmapdefault

instance GFunctor (Either a) where
  gmap = gmapdefault

instance GFunctor Monoid.First where
  gmap = gmapdefault

#if MIN_VERSION_base(4,9,0)
instance GFunctor (Semigroup.First) where
  gmap = gmapdefault
#endif

#if MIN_VERSION_base(4,8,0)
instance GFunctor Identity where
  gmap = gmapdefault
#endif

instance GFunctor IO where
  gmap = fmap

instance GFunctor Monoid.Last where
  gmap = gmapdefault

#if MIN_VERSION_base(4,9,0)
instance GFunctor Semigroup.Last where
  gmap = gmapdefault

instance GFunctor Max where
  gmap = gmapdefault
#endif

instance GFunctor Maybe where
  gmap = gmapdefault

#if MIN_VERSION_base(4,9,0)
instance GFunctor Min where
  gmap = gmapdefault

instance GFunctor NonEmpty where
  gmap = gmapdefault

instance GFunctor Option where
  gmap = gmapdefault
#endif

instance GFunctor Monoid.Product where
  gmap = gmapdefault

#if MIN_VERSION_base(4,9,0)
instance (GFunctor f, GFunctor g) => GFunctor (Functor.Product f g) where
  gmap = gmapdefault
#endif

#if MIN_VERSION_base(4,7,0)
instance GFunctor Proxy where
  gmap = gmapdefault
#endif

instance GFunctor Monoid.Sum where
  gmap = gmapdefault

#if MIN_VERSION_base(4,9,0)
instance (GFunctor f, GFunctor g) => GFunctor (Functor.Sum f g) where
  gmap = gmapdefault

instance GFunctor WrappedMonoid where
  gmap = gmapdefault
#endif

instance GFunctor ZipList where
  gmap = gmapdefault

max_stars_repo_name: prozacchiwawa/generic-deriving

max_issues_repo_name: prozacchiwawa/generic-deriving

max_forks_repo_name: prozacchiwawa/generic-deriving


text:
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

module SSSFN
  ( -- * Domain
    np,

    -- * Grid
    Grid (..),
    gunit,
    (!),
    gcoerce,

    -- * Operators
    Op (..),
    ounit,
    ocoerce,

    -- * Grids are functors etc.
    gmap,
    gpure,
    gzipWith,
    gfoldMap,
    gcount,
    gmaxabs,
    gsum,

    -- * Operators are functors etc.
    omap,
    opure,
    ozipWith,
    ofoldMap,
    ocount,
    omaxabs,
    osum,

    -- * Operators
    (#>),
    transpose,
    inv,

    -- * Coordinates
    isOrigin,
    coordU,
    coordV,
    coordT,
    coordR,
    coordR1, -- remove?

    -- * Sampling functions
    gsample,

    -- * Evaluating functions
    geval,
    gevalCoeffs,

    -- * Derivatives
    vandermonde,
    expfilter,
    delta,
    symmetrizeU,
    symmetrizeV,
    symmetrizeR,
    derivU,
    derivV,
    derivUU,
    derivUV,
    derivVV,
    derivT,
    derivR,
    derivTT,
    derivRR,
    r1derivR,
    r1derivU,
    r1derivV,
    r2project,
    zproject,
    bndU,
    bndV,
    metric,

    -- * Systems of equations
    solve,
    solveLS,

    -- * MaxAbs
    maxabs,
    getMaxAbs,
  )
where

import Control.Exception (assert)
import Data.Maybe
import Data.Poly hiding (scale)
import Data.Semigroup hiding ((<>))
import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import Foreign
import qualified GHC.Exts as GHC
import Numeric.LinearAlgebra hiding ((!), (#>), inv)
import qualified Numeric.LinearAlgebra as L
import Numeric.LinearAlgebra.Devel
import qualified Test.QuickCheck as QC
import Prelude hiding ((<>))
import qualified Prelude as P

default (Int)

--------------------------------------------------------------------------------

-- | Grid size
np :: Int
np = 33

--------------------------------------------------------------------------------

-- | Grid datatype
newtype Grid a = Grid {getGrid :: Vector a}
  deriving (Eq, Ord, Read, Show)

instance (QC.Arbitrary a, Storable a) => QC.Arbitrary (Grid a) where

  arbitrary = Grid . fromList <$> QC.vector (np ^ 2)

  shrink = map (Grid . fromList) . traverse QC.shrink . toList . getGrid

instance (Storable a) => GHC.IsList (Grid a) where

  type Item (Grid a) = a

  fromList = Grid . fromList

  toList = toList . getGrid

(!) :: (Container Vector a) => Grid a -> (Int, Int) -> a
(!) (Grid x) (i, j) = x `atIndex` (i * np + j)

gunit :: (Num a, Storable a) => (Int, Int) -> Grid a
gunit (i, j) =
  let ij = i * np + j
   in Grid $ fromList [if k == ij then 1 else 0 | k <- [0 .. np ^ 2 -1]]

gcoerce :: forall a b. (Storable a, Storable b) => Grid a -> Grid b
gcoerce (Grid xss) =
  assert (na * sa `mod` sb == 0) $
    (Grid $ fromByteString $ toByteString xss)
  where
    na = cols (asRow xss)
    sa = sizeOf @a undefined
    sb = sizeOf @b undefined

--------------------------------------------------------------------------------

-- | Operator datatype
newtype Op a = Op {getOp :: Matrix a}
  deriving (Read, Show)

instance (QC.Arbitrary a, Element a) => QC.Arbitrary (Op a) where

  arbitrary = Op . reshape (np ^ 2) . fromList <$> QC.vector (np ^ 4)

  shrink =
    map (Op . reshape (np ^ 2) . fromList)
      . traverse QC.shrink
      . (toList . flatten . getOp)

instance Element a => GHC.IsList (Op a) where

  type Item (Op a) = a

  fromList = Op . reshape (np ^ 2) . fromList

  toList = toList . flatten . getOp

ounit :: (Element a, Num a) => Op a
ounit = Op $ ident (np ^ 2)

ocoerce :: forall a b. (Container Vector a, Num a, Storable b) => Op a -> Op b
ocoerce (Op x) =
  assert (ma * sa `mod` sb == 0 && na * sa `mod` sb == 0) $
    (Op $ reshape nb $ getGrid $ gcoerce $ Grid $ flatten x)
  where
    (ma, na) = size x
    sa = sizeOf @a undefined
    sb = sizeOf @b undefined
    nb = na * sa `div` sb

--------------------------------------------------------------------------------

-- | Grids are functors etc.
-- gmap :: (Container Vector a, Element b) => (a -> b) -> Grid a -> Grid b
-- gmap f (Grid x) = Grid (cmap f x)
gmap :: (Storable a, Storable b) => (a -> b) -> Grid a -> Grid b
gmap f (Grid x) = Grid (mapVectorWithIndex (\i a -> f a) x)

gpure :: Container Vector a => a -> Grid a
gpure a = Grid $ konst a (np ^ 2)

gzipWith ::
  (Storable a, Storable b, Storable c) =>
  (a -> b -> c) ->
  Grid a ->
  Grid b ->
  Grid c
gzipWith f (Grid x) (Grid y) = Grid (zipVectorWith f x y)

gfoldMap :: (Storable a, Monoid b) => (a -> b) -> Grid a -> b
gfoldMap f (Grid x) = foldVector (\a b -> f a P.<> b) mempty x

gcount :: Container Vector a => Grid a -> Int
gcount (Grid x) = size x

gmaxabs :: (Container Vector a, Num a, Ord a) => Grid a -> a
gmaxabs = getMaxAbs . gfoldMap maxabs

gsum :: (Container Vector a) => Grid a -> a
gsum (Grid x) = sumElements x

--------------------------------------------------------------------------------

-- | Operators are functors etc.
-- omap :: (Container Vector a, Num a, Element b) => (a -> b) -> Op a -> Op b
-- omap f (Op x) = Op (cmap f x)
omap :: (L.Element a, Storable b) => (a -> b) -> Op a -> Op b
omap f (Op x) = Op $ reshape n $ getGrid $ gmap f $ Grid $ flatten x
  where
    n = cols x

opure :: (Container Vector a, Num a) => a -> Op a
opure a = Op (konst a (np ^ 2, np ^ 2))

ozipWith ::
  (Element a, Element b, Storable c) =>
  (a -> b -> c) ->
  Op a ->
  Op b ->
  Op c
ozipWith f (Op x) (Op y) =
  Op $ reshape (np ^ 2) $ zipVectorWith f (flatten x) (flatten y)

ofoldMap :: (Element a, Monoid b) => (a -> b) -> Op a -> b
ofoldMap f (Op x) = foldVector (\a b -> f a P.<> b) mempty (flatten x)

ocount :: (Container Vector a, Num a) => Op a -> Int
ocount (Op x) = let (ni, nj) = size x in ni * nj

omaxabs :: (Container Vector a, Num a, Ord a) => Op a -> a
omaxabs = getMaxAbs . ofoldMap maxabs

osum :: (Container Vector a, Num a) => Op a -> a
osum (Op x) = sumElements x

--------------------------------------------------------------------------------

-- | Grids are an additive group
instance
  (Container Vector a, Num a) =>
  AdditiveGroup (Grid a)
  where

  zeroV = gpure 0

  (^+^) = gzipWith (+)

  negateV = gmap negate

-- | Grids are a vector space
instance
  (Container Vector a, Num a) =>
  VectorSpace (Grid a)
  where

  type Scalar (Grid a) = a

  (*^) a = gmap (a *)

--------------------------------------------------------------------------------

-- | Operators form an additive group
instance (Container Vector a, Num a) => AdditiveGroup (Op a) where

  zeroV = opure 0

  (^+^) = ozipWith (+)

  negateV = omap negate

-- | Operators form a vector space
instance (Container Vector a, Num a) => VectorSpace (Op a) where

  type Scalar (Op a) = a

  (*^) a = omap (a *)

-- | Operators form a ring
instance (Container Vector a, Num a, Numeric a) => Num (Op a) where

  (+) = (^+^)

  negate = negateV

  (-) = (^-^)

  Op x * Op y = Op (x <> y)

  abs = omap abs

  signum = omap signum

  fromInteger i = fromInteger i *^ ounit

-- | Matrix-vector multiplication
infixr 8 #>

(#>) :: Numeric a => Op a -> Grid a -> Grid a
Op x #> Grid y = Grid (x L.#> y)

-- | Matrix transpose
transpose :: Transposable (Matrix a) (Matrix a) => Op a -> Op a
transpose (Op x) = Op (tr x)

-- | Matrix inverse
inv :: Field a => Op a -> Op a
inv (Op x) = Op (L.inv x)

--------------------------------------------------------------------------------

-- | Coordinate of a grid point
coord :: Floating a => Int -> a
coord i =
  let imin = 0
      imax = np - 1
      umin = -1
      umax = 1
      u =
        fromIntegral (i - imax) / fromIntegral (imin - imax) * umin
          + fromIntegral (i - imin) / fromIntegral (imax - imin) * umax
   in sin (pi / 2 * u)

-- | Coordinates of all grid points
coords :: (Floating a, Storable a) => Vector a
coords = fromList [coord i | i <- [0 .. np -1]]

poly2vector :: (Num a, Storable a, U.Unbox a) => UPoly a -> Vector a
poly2vector p =
  let v = unPoly p
      n = U.length v
   in fromList [if i < n then v U.! i else 0 | i <- [0 .. np -1]]

-- | Evalute a polynomial at a grid point
evalcoeffs :: (Eq a, Num a, Storable a, U.Unbox a) => a -> Vector a
evalcoeffs u = fromList [eval (mon i) u | i <- [0 .. np -1]]
  where
    mon :: (Eq a, Num a, U.Unbox a) => Int -> UPoly a
    mon i = monomial (fromIntegral i) 1

-- | Evaluate a polynomial at all grid points
gvandermonde :: (Element a, Eq a, Floating a, U.Unbox a) => Matrix a
gvandermonde = fromRows [evalcoeffs u | u <- toList coords]

--------------------------------------------------------------------------------

-- Note: hmatrix is row-major by default
mkGrid :: (Num a, Storable a) => (Int -> a) -> (Int -> a) -> Grid a
mkGrid fx fy =
  Grid $ fromList [fx i * fy j | i <- [0 .. np -1], j <- [0 .. np -1]]

isOrigin :: Grid Bool
isOrigin =
  Grid $
    fromList
      [ i == j
        | i <- [0 .. np -1],
          j <- [0 .. np -1]
      ]

-- | Coordinates of all grid points
coordU :: (Container Vector a, Floating a) => Grid a
coordU = mkGrid (coords `atIndex`) (const 1)

-- | Coordinates of all grid points
coordV :: (Container Vector a, Floating a) => Grid a
coordV = mkGrid (const 1) (coords `atIndex`)

-- | Combined coordinates
coordT :: (Container Vector a, Floating a) => Grid a
coordT = coordV ^+^ coordU

coordR :: (Container Vector a, Floating a) => Grid a
coordR = coordV ^-^ coordU

coordR1 :: (Container Vector a, RealFloat a) => Grid a
coordR1 = gzipWith (\isOrig r -> if isOrig then 0 else recip r) isOrigin coordR

--------------------------------------------------------------------------------

-- | Sample a function at grid points
gsample ::
  (Container Vector a, Floating a, Storable b) =>
  ((a, a) -> b) ->
  Grid b
gsample f = gzipWith (curry f) (coordU) (coordV)

--------------------------------------------------------------------------------

-- | Evaluate a grid at a point
-- f / Identity
geval ::
  ( Container Vector a,
    Eq a,
    Field a,
    U.Unbox a,
    Container Vector b,
    VectorSpace b,
    Scalar b ~ a
  ) =>
  Grid b ->
  (a, a) ->
  b
geval g (u, v) =
  let ucoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs u
      vcoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs v
   in sumV
        [ ((ucoeffs `atIndex` i) * (vcoeffs `atIndex` j)) *^ (g ! (i, j))
          | i <- [0 .. np - 1],
            j <- [0 .. np - 1]
        ]

gevalCoeffs ::
  ( Container Vector a,
    Eq a,
    Field a,
    U.Unbox a
  ) =>
  (a, a) ->
  Grid a
gevalCoeffs (u, v) =
  let ucoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs u
      vcoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs v
   in Grid $
        fromList
          [ (ucoeffs `atIndex` i) * (vcoeffs `atIndex` j)
            | i <- [0 .. np - 1],
              j <- [0 .. np - 1]
          ]

--------------------------------------------------------------------------------

-- | Derivative of a basis function
deriv1 :: (Eq a, Num a, Storable a, U.Unbox a) => Int -> Vector a
deriv1 i = poly2vector $ deriv $ monomial (fromIntegral i) 1

-- | Derivatives of all basis functions
derivs1 :: (Element a, Eq a, Num a, U.Unbox a) => Matrix a
derivs1 = fromColumns [deriv1 i | i <- [0 .. np -1]]

derivs :: (Eq a, Field a, U.Unbox a) => Matrix a
derivs = gvandermonde <> derivs1 <> L.inv gvandermonde

-- derivs2 :: (Eq a, Field a, U.Unbox a) => Matrix a
-- derivs2 = derivs <> derivs

expfilter1 :: (Container Vector a, RealFloat a) => Matrix a
expfilter1 =
  diag $
    fromList
      [ let n0 = np `div` 2
         in if i <= n0
              then 1
              else peps ** (fromIntegral (i - n0) / fromIntegral (np - n0))
        | i <- [0 .. np -1]
      ]

expfilters :: (Container Vector a, Field a, RealFloat a, U.Unbox a) => Matrix a
expfilters = gvandermonde <> expfilter1 <> L.inv gvandermonde

bnd :: (Element a, Fractional a) => Matrix a
bnd =
  diag $
    fromList
      [ if i == 0
          then -1 / fromIntegral (np - 1)
          else if i == np - 1 then 1 / fromIntegral (np - 1) else 0
        | i <- [0 .. np - 1]
      ]

weights :: (Element a, Fractional a) => Matrix a
weights =
  diag $
    fromList
      [ if i == 0 || i == np - 1
          then 1 / fromIntegral (np - 1)
          else 2 / fromIntegral (np - 1)
        | i <- [0 .. np - 1]
      ]

--------------------------------------------------------------------------------

mkOp :: (Element a, Num a) => ((Int, Int) -> a) -> ((Int, Int) -> a) -> Op a
mkOp fx fy =
  Op $
    fromLists
      [ [ fx (ir, ic) * fy (jr, jc)
          | ic <- [0 .. np - 1],
            jc <- [0 .. np - 1]
        ]
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

vandermonde :: (Container Vector a, Eq a, Floating a, U.Unbox a) => Op a
vandermonde = mkOp (gvandermonde `atIndex`) (gvandermonde `atIndex`)

expfilter :: (Container Vector a, Field a, RealFloat a, U.Unbox a) => Op a
expfilter = mkOp (expfilters `atIndex`) (expfilters `atIndex`)

delta :: Field a => Op a
delta = mkOp delta1 delta1
  where
    delta1 (i, j) = if i == j then 1 else 0

diagonal :: (Element a, Num a) => Grid a -> Op a
diagonal (Grid xs) = Op $ diag xs

originR :: Field a => Op a
originR =
  Op $
    fromLists
      [ [ if ic == ir && jc == jr && ic == jc then 1 else 0
          | ic <- [0 .. np - 1],
            jc <- [0 .. np - 1]
        ]
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

valueAtOriginR :: (Eq a, Field a, U.Unbox a) => Op a
valueAtOriginR =
  Op $
    fromRows
      [ let t = coordT ! (ir, jr)
            r = 0
            v = (t + r) / 2
            u = (t - r) / 2
         in getGrid (gevalCoeffs (u, v))
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

reverseU :: (Element a, Num a) => Op a
reverseU = mkOp reverse1 delta1
  where
    reverse1 (i, j) = if i == np -1 - j then 1 else 0
    delta1 (i, j) = if i == j then 1 else 0

reverseV :: (Element a, Num a) => Op a
reverseV = mkOp delta1 reverse1
  where
    reverse1 (i, j) = if i == np -1 - j then 1 else 0
    delta1 (i, j) = if i == j then 1 else 0

-- change (t,r) -> (t,-r); (u,v) -> (v,u)
reverseR :: (Element a, Num a) => Op a
reverseR =
  Op $
    fromLists
      [ [ if ic == jr && jc == ir then 1 else 0
          | ic <- [0 .. np - 1],
            jc <- [0 .. np - 1]
        ]
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

symmetrizeU :: Field a => Op a
symmetrizeU = (1 / 2) *^ (delta ^+^ reverseU)

symmetrizeV :: Field a => Op a
symmetrizeV = (1 / 2) *^ (delta ^+^ reverseV)

symmetrizeR :: Field a => Op a
symmetrizeR = (1 / 2) *^ (delta ^+^ reverseR)

-- | Derivatives of a grid
derivU :: (Eq a, Field a, U.Unbox a) => Op a
derivU = mkOp (derivs `atIndex`) delta1
  where
    delta1 (i, j) = if i == j then 1 else 0

derivV :: (Eq a, Field a, U.Unbox a) => Op a
derivV = mkOp delta1 (derivs `atIndex`)
  where
    delta1 (i, j) = if i == j then 1 else 0

derivUU :: (Eq a, Field a, U.Unbox a) => Op a
derivUU = derivU * derivU

derivUV :: (Eq a, Field a, U.Unbox a) => Op a
derivUV = derivU * derivV

derivVV :: (Eq a, Field a, U.Unbox a) => Op a
derivVV = derivV * derivV

-- | Combined derivatives
-- > t = v + u
-- > r = v - u
-- > t^2 = v^2 + 2 v u + u^2
-- > r^2 = v^2 - 2 v u + u^2
derivT :: (Eq a, Field a, U.Unbox a) => Op a
derivT = (derivV ^+^ derivU) ^/ 2

derivR :: (Eq a, Field a, U.Unbox a) => Op a
derivR = (derivV ^-^ derivU) ^/ 2

derivTT :: (Eq a, Field a, U.Unbox a) => Op a
derivTT = derivT * derivT

derivRR :: (Eq a, Field a, U.Unbox a) => Op a
derivRR = derivR * derivR

-- | 1/r dr
r1derivR :: (Field a, RealFloat a, U.Unbox a) => Op a
r1derivR =
  ((delta - originR) * diagonal coordR1 * derivR + originR * derivRR)
    * symmetrizeR

-- t = v + u   v = (t + r) / 2
-- r = v - u   u = (t - r) / 2

-- | 1/r du
r1derivU :: (Field a, RealFloat a, U.Unbox a) => Op a
r1derivU =
  ((delta - originR) * diagonal coordR1 * derivU - originR * derivUU)
    * symmetrizeR

-- | 1/r dv
r1derivV :: (Field a, RealFloat a, U.Unbox a) => Op a
r1derivV =
  ((delta - originR) * diagonal coordR1 * derivV + originR * derivVV)
    * symmetrizeR

-- | 1/r^2
r2project :: (Field a, RealFloat a, U.Unbox a) => Op a
r2project =
  ( (delta - originR) * diagonal (gmap (^ 2) coordR1)
      + originR * ((1 / 2) *^ derivRR)
  )
    * (delta - valueAtOriginR)
    * symmetrizeR

-- | zero at origin
zproject :: (Eq a, Field a, U.Unbox a) => Op a
zproject = (delta - valueAtOriginR) * symmetrizeR

-- | Boundary locations
bndU :: (Container Vector a, Fractional a) => Op a
bndU = mkOp (bnd `atIndex`) delta1
  where
    delta1 (i, j) = if i == j then 1 else 0

bndV :: (Container Vector a, Fractional a) => Op a
bndV = mkOp delta1 (bnd `atIndex`)
  where
    delta1 (i, j) = if i == j then 1 else 0

metric :: (Container Vector a, Fractional a) => Op a
metric = mkOp (weights `atIndex`) (weights `atIndex`)

--------------------------------------------------------------------------------

solve :: Field a => Op a -> Grid a -> Grid a
solve (Op a) (Grid x) =
  Grid $ head $ toColumns $ fromJust $ linearSolve a (asColumn x)

solveLS :: Field a => Op a -> Grid a -> Grid a
solveLS (Op a) (Grid x) =
  Grid $ head $ toColumns $ linearSolveLS a (asColumn x)

--------------------------------------------------------------------------------

newtype MaxAbs a = MaxAbsInternal a
  deriving (Eq, Ord, Read, Show)

maxabs :: Num a => a -> MaxAbs a
maxabs x = MaxAbsInternal (abs x)

getMaxAbs :: MaxAbs a -> a
getMaxAbs (MaxAbsInternal x) = x

instance Ord a => Semigroup (MaxAbs a) where
  MaxAbsInternal x <> MaxAbsInternal y = MaxAbsInternal $ max x y

instance (Num a, Ord a) => Monoid (MaxAbs a) where
  mempty = MaxAbsInternal 0

max_stars_repo_name: eschnett/sssfn

max_issues_repo_name: eschnett/sssfn

max_forks_repo_name: eschnett/sssfn