{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Nudge (
nudgeProp
, nudgeProp2
) where
import Control.Monad
import Data.Bifunctor
import Data.Finite
import Data.Kind
import Data.Maybe
import Data.Proxy
import GHC.TypeLits
import Hedgehog
import Lens.Micro
import Lens.Micro.Platform ()
import Numeric.Backprop
import qualified Data.Ix as Ix
import qualified Data.Vector.Sized as SV
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import qualified Numeric.LinearAlgebra as HU
import qualified Numeric.LinearAlgebra.Static as H
import qualified Numeric.LinearAlgebra.Static.Backprop as B
nudge :: Double
nudge = 1e-6
eps :: Double
eps = 1e-11
class (Backprop c, Show c, Show (TIx c)) => Testing c where
type TIx c :: Type
allIx :: c -> [TIx c]
ixLens :: TIx c -> Lens' c Double
scalarize :: Reifies s W => BVar s c -> BVar s Double
genTest :: Gen c
sized
:: forall s t d. H.Sized t s d
=> Lens' s (d t)
sized f = fmap (fromJust . H.create) . f . H.extract
ixContainer
:: forall t d. HU.Container d t
=> HU.IndexOf d
-> Lens' (d t) t
ixContainer i = lens (`HU.atIndex` i)
(\xs x -> HU.accum xs (\_ _ -> x) [(i, x)])
instance Testing Double where
type TIx Double = ()
allIx _ = [()]
ixLens _ = id
scalarize = abs
genTest = Gen.filter ((> eps) . (**2)) $
Gen.double (Range.linearFracFrom 0 (-5) 5)
instance KnownNat n => Testing (H.R n) where
type TIx (H.R n) = Int
allIx v = [0 .. H.size v - 1]
ixLens i = sized . ixContainer i
scalarize = B.norm_2V
genTest = H.vector <$> replicateM n genTest
where
n = fromInteger $ natVal (Proxy @n)
instance (KnownNat n, KnownNat m) => Testing (H.L n m) where
type TIx (H.L n m) = (Int, Int)
allIx m = Ix.range ((0,0), bimap pred pred (H.size m))
ixLens i = sized . ixContainer i
scalarize = sqrt . B.sumElements . (**2)
genTest = H.matrix <$> replicateM nm genTest
where
nm = fromInteger $ natVal (Proxy @n) * natVal (Proxy @m)
instance Testing (HU.Vector Double) where
type TIx (HU.Vector Double) = Int
allIx v = [0 .. HU.size v - 1]
ixLens = ixContainer
scalarize = liftOp1 . op1 $ \xs -> (HU.sumElements xs, (`HU.konst` HU.size xs))
genTest = HU.fromList <$> replicateM 3 genTest
instance Testing (HU.Matrix Double) where
type TIx (HU.Matrix Double) = (Int, Int)
allIx m = Ix.range ((0,0), bimap pred pred (HU.size m))
ixLens = ixContainer
scalarize = liftOp1 . op1 $ \xs -> (HU.sumElements xs, (`HU.konst` HU.size xs))
genTest = HU.fromLists <$> (replicateM 3 . replicateM 2) genTest
instance (KnownNat n, Testing a, Num a) => Testing (SV.Vector n a) where
type TIx (SV.Vector n a) = (Finite n, TIx a)
allIx = fst . SV.imapM (\i x -> ((fromIntegral i,) <$> allIx x , x))
ixLens (i,j) = SV.ix i . ixLens j
scalarize = scalarize . liftOp1 o . (^ (2 :: Int))
where
o :: Op '[SV.Vector n a] a
o = op1 $ \xs -> (SV.sum xs, SV.replicate)
genTest = SV.replicateM genTest
instance (Testing a, Testing b) => Testing (a, b) where
type TIx (a, b) = Either (TIx a) (TIx b)
allIx (x, y) = (Left <$> allIx x)
++ (Right <$> allIx y)
ixLens (Left i) = _1 . ixLens i
ixLens (Right j) = _2 . ixLens j
scalarize t = B.norm_2V (B.vec2 (scalarize (t ^^. _1))
(scalarize (t ^^. _2))
)
genTest = (,) <$> genTest <*> genTest
instance (Testing a, Testing b, Testing c) => Testing (a, b, c) where
type TIx (a, b, c) = Either (TIx a) (Either (TIx b) (TIx c))
allIx (x, y, z) = (Left <$> allIx x)
++ (Right . Left <$> allIx y)
++ (Right . Right <$> allIx z)
ixLens (Left i ) = _1 . ixLens i
ixLens (Right (Left j)) = _2 . ixLens j
ixLens (Right (Right k)) = _3 . ixLens k
scalarize t = B.norm_2V (B.vec3 (scalarize (t ^^. _1))
(scalarize (t ^^. _2))
(scalarize (t ^^. _3))
)
genTest = (,,) <$> genTest <*> genTest <*> genTest
validGrad
:: Monad m
=> Lens' c Double
-> c
-> c
-> (c -> Double)
-> PropertyT m (Double, Double)
validGrad l x0 g f = forAll $ Gen.double (Range.constantFrom 0 (-nudge) nudge) <&> \d ->
let x = x0 & l %~ (+d)
old = f x0 + (g ^. l) * d
new = f x
in (old, new)
nudgeProp
:: forall c d. (Testing c, Testing d)
=> (forall s. Reifies s W => BVar s c -> BVar s d)
-> Property
nudgeProp f = property $ do
(inp, i) <- forAll $ do
inp <- genTest
i <- Gen.element (allIx inp)
return (inp, i)
let (r,gr) = backprop (scalarize . f) inp
when (r**2 < eps) discard
(old, new) <- validGrad (ixLens i) inp gr (evalBP (scalarize . f))
footnoteShow (r, gr, old, new, (old - new)**2, ((old - new)/old)**2)
assert $ ((old - new)/old)**2 < eps
nudgeProp2
:: forall c d e. (Testing c, Testing d, Testing e)
=> (forall s. Reifies s W => BVar s c -> BVar s d -> BVar s e)
-> Property
nudgeProp2 f = property $ do
(inpC, inpD, i) <- forAll $ do
inpC <- genTest
inpD <- genTest
i <- Gen.element (allIx (inpC, inpD))
return (inpC, inpD, i)
let (r, gr) = backprop2 (\x -> scalarize . f x) inpC inpD
when (r**2 < eps) discard
(old, new) <- validGrad (ixLens i) (inpC, inpD) gr
(evalBP (\t -> scalarize $ f (t ^^. _1) (t ^^. _2)))
footnoteShow (r, gr, old, new, (old - new)**2, ((old - new)/old)**2)
assert $ ((old - new)/old)**2 < eps
instance (HU.Container HU.Vector a, Num a) => Backprop (HU.Matrix a) where
-- TODO: make more efficient?
zero = HU.cmap (const 0)
add = HU.add
one = HU.cmap (const 1)
instance (KnownNat n, Num a) => Backprop (SV.Vector n a) where
zero = (0 <$)
add = (+)
one = (1 <$)