mirror of
https://gitlab.com/sheaf/metabrush.git
synced 2024-11-06 07:13:37 +00:00
320 lines
11 KiB
Haskell
320 lines
11 KiB
Haskell
{-# LANGUAGE AllowAmbiguousTypes #-}
|
||
{-# LANGUAGE NumericUnderscores #-}
|
||
{-# LANGUAGE ScopedTypeVariables #-}
|
||
{-# LANGUAGE TemplateHaskell #-}
|
||
|
||
module Main (main) where
|
||
|
||
-- base
|
||
import Prelude hiding
|
||
( Num(..), (^) )
|
||
import Data.Foldable
|
||
( toList )
|
||
import Data.List.NonEmpty
|
||
( NonEmpty(..) )
|
||
import Data.Maybe
|
||
( catMaybes )
|
||
import Data.Traversable
|
||
( for )
|
||
import Unsafe.Coerce
|
||
( unsafeCoerce )
|
||
|
||
-- brush-strokes
|
||
import Math.Algebra.Dual
|
||
import Math.Linear
|
||
import Math.Module
|
||
import Math.Monomial
|
||
( multiSubsetSum, multiSubsetsSum
|
||
, MonomialBasis ( monTabulate, monIndex )
|
||
)
|
||
import Math.Ring
|
||
|
||
-- hspray
|
||
import Math.Algebra.Hspray
|
||
( Spray )
|
||
import qualified Math.Algebra.Hspray as Spray
|
||
|
||
-- falsify
|
||
import Test.Tasty.Falsify
|
||
import qualified Test.Falsify.Generator as Falsify
|
||
( Gen )
|
||
import qualified Test.Falsify.Generator as Falsify.Gen
|
||
import Test.Falsify.Predicate
|
||
( (.$) )
|
||
import qualified Test.Falsify.Predicate as Falsify.Prop
|
||
import qualified Test.Falsify.Property as Falsify
|
||
( Property
|
||
, assert
|
||
, discard
|
||
, gen, genWith
|
||
)
|
||
import qualified Test.Falsify.Range as Falsify
|
||
|
||
-- tasty
|
||
import qualified Test.Tasty as Tasty
|
||
|
||
-- unordered-containers
|
||
import qualified Data.HashMap.Lazy as HashMap
|
||
|
||
--------------------------------------------------------------------------------
|
||
|
||
|
||
main :: IO ()
|
||
main =
|
||
Tasty.defaultMain $
|
||
Tasty.testGroup "brush-strokes property tests"
|
||
[ Tasty.testGroup "Automatic differentiation"
|
||
[ Tasty.testGroup "Monomial basis"
|
||
[ testProperty "Round trip D33" testMonomialBasisD33
|
||
]
|
||
, Tasty.testGroup "Monomials"
|
||
[ Tasty.testGroup "multiSubsetSum"
|
||
[ testProperty "multiSubsetSum valid" testMultiSubsetSumValid
|
||
, testProperty "multiSubsetSum exhaustive" testMultiSubsetSumExhaustive
|
||
]
|
||
-- , Tasty.testGroup "multiSubsetsSum"
|
||
-- [ testProperty "multiSubsetsSum exhaustive" testMultiSubsetsSumExhaustive
|
||
-- ]
|
||
]
|
||
, Tasty.testGroup "chainRule1NQ"
|
||
[ testProperty "chainRule1NQ_1" testChainRule1NQ_1
|
||
, testProperty "chainRule1NQ_2" testChainRule1NQ_2
|
||
, testProperty "chainRule1NQ_3" testChainRule1NQ_3
|
||
]
|
||
]
|
||
]
|
||
|
||
-- | Check that the 'multiSubsetSum' function returns valid answers, i.e.
|
||
-- all returned multisubsets have the desired size and sum.
|
||
testMultiSubsetSumValid :: Falsify.Property ()
|
||
testMultiSubsetSumValid = do
|
||
rg <- Falsify.genWith (\ rg -> Just $ "range = " ++ show rg ) $ Falsify.Gen.inRange $ Falsify.between ( 1, 6 )
|
||
sz <- Falsify.genWith (\ sz -> Just $ "size = " ++ show sz ) $ Falsify.Gen.inRange $ Falsify.between ( 0, 20 )
|
||
tot <- Falsify.genWith (\ tot -> Just $ "tot = " ++ show tot) $ Falsify.Gen.inRange $ Falsify.between ( sz, sz * rg )
|
||
let range = [ 1 .. rg ]
|
||
mss = multiSubsetSum sz tot range
|
||
case mss of
|
||
[] -> Falsify.discard
|
||
r:rs -> do
|
||
ms <- Falsify.gen $ Falsify.Gen.elem ( r :| rs )
|
||
Falsify.assert
|
||
$ Falsify.Prop.eq
|
||
.$ ("(sz, tot)", (sz, tot) )
|
||
.$ ("computed (sz, tot)", (size ms, total ms))
|
||
where
|
||
size, total :: [ ( Word, Word ) ] -> Word
|
||
size [] = 0
|
||
size ((_,n):ins) = n + size ins
|
||
total [] = 0
|
||
total ((i,n):ins) = i * n + total ins
|
||
|
||
-- | Check that the 'multiSubsetSum' function returns all multisubsets of
|
||
-- the given set, by generating a random multisubset, computing its size, and
|
||
-- checking it belongs to the output of the 'multiSubsetSum' function.
|
||
testMultiSubsetSumExhaustive :: Falsify.Property ()
|
||
testMultiSubsetSumExhaustive = do
|
||
rg <- Falsify.genWith (\ rg -> Just $ "range = " ++ show rg) $ Falsify.Gen.inRange $ Falsify.between ( 1, 6 )
|
||
sz <- Falsify.genWith (\ sz -> Just $ "size = " ++ show sz) $ Falsify.Gen.inRange $ Falsify.between ( 0, 10 )
|
||
let range = [ 1 .. rg ]
|
||
(multiSubset, tot) <- Falsify.genWith (\ ms -> Just $ "multisubset = " ++ show ms) $ genMultiSubset range sz
|
||
Falsify.assert
|
||
$ Falsify.Prop.elem
|
||
.$ ("all multisubsets", multiSubsetSum sz tot range )
|
||
.$ ("random multisubset", multiSubset)
|
||
|
||
genMultiSubset :: [ Word ] -> Word -> Falsify.Gen ( [ ( Word, Word ) ] , Word )
|
||
genMultiSubset [i] sz =
|
||
return $
|
||
if sz == 0
|
||
then ( [], 0 )
|
||
else ( [ ( i, sz ) ], i * sz )
|
||
genMultiSubset (i:is) sz = do
|
||
nb <- Falsify.Gen.inRange $ Falsify.between ( 0, sz )
|
||
(rest, tot) <- genMultiSubset is ( sz - nb )
|
||
return $ ( if nb == 0 then rest else ( i, nb ) : rest, tot + nb * i )
|
||
genMultiSubset [] _ = error "impossible"
|
||
|
||
coerceVec1 :: [ a ] -> Vec n a
|
||
coerceVec1 = unsafeCoerce
|
||
|
||
coerceVec2 :: Vec n a -> [ a ]
|
||
coerceVec2 = toList
|
||
|
||
-- | Check that the 'multiSubsetSums' function returns all collections of
|
||
-- multisubsets of the given set (see 'testMultiSubsetSumExhaustive').
|
||
testMultiSubsetsSumExhaustive :: Falsify.Property ()
|
||
testMultiSubsetsSumExhaustive = do
|
||
rg <- Falsify.genWith (\ rg -> Just $ "range = " ++ show rg) $ Falsify.Gen.inRange $ Falsify.between ( 1, 5 )
|
||
let range = [ 1 .. rg ]
|
||
n <- Falsify.genWith (\ n -> Just $ "n = " ++ show n ) $ Falsify.Gen.inRange $ Falsify.between ( 1, 10 )
|
||
multiSubsets <- for ( [ 0 .. n - 1 ] :: [ Word ] ) \ i -> do
|
||
sz <- Falsify.gen $ Falsify.Gen.inRange $ Falsify.between ( 0, 5 )
|
||
( ms, tot ) <- Falsify.genWith ( \ ms -> Just $ "ms_" ++ show i ++ " = " ++ show ms ) $ genMultiSubset range sz
|
||
return ( ms, sz, tot )
|
||
let mss = map ( \ (ms, _,_) -> ms ) multiSubsets
|
||
szs = map ( \ (_,sz,_) -> sz) multiSubsets
|
||
tot = sum $ map ( \(_,_,t) -> t) multiSubsets
|
||
Falsify.assert
|
||
$ Falsify.Prop.elem
|
||
.$ ("all multisubsets", map coerceVec2 $ multiSubsetsSum range tot $ coerceVec1 szs )
|
||
.$ ("random multisubset", mss)
|
||
|
||
|
||
testRoundTrip
|
||
:: ( Show a, Eq a )
|
||
=> Falsify.Gen a
|
||
-> ( a -> a )
|
||
-> Falsify.Property ()
|
||
testRoundTrip g roundTrip = do
|
||
d <- Falsify.gen g
|
||
Falsify.assert
|
||
$ Falsify.Prop.eq
|
||
.$ ("value", d )
|
||
.$ ("round tripped", roundTrip d )
|
||
|
||
testMonomialBasisD33 :: Falsify.Property ()
|
||
testMonomialBasisD33 =
|
||
testRoundTrip genD33 \ d -> $$( monTabulate \ mon -> monIndex [|| d ||] mon )
|
||
where
|
||
genD33 :: Falsify.Gen ( D3𝔸3 Double )
|
||
genD33 =
|
||
D33 <$> (unT <$> g)
|
||
<*> g <*> g <*> g
|
||
<*> g <*> g <*> g <*> g <*> g <*> g
|
||
<*> g <*> g <*> g <*> g <*> g <*> g <*> g <*> g <*> g <*> g
|
||
|
||
g :: Falsify.Gen ( T Double )
|
||
g = T . fromIntegral <$> Falsify.Gen.inRange ( Falsify.withOrigin ( -100, 100 ) ( 0 :: Int ) )
|
||
|
||
-- | Test the Faà di Bruno formula on polynomials, with a composition
|
||
-- \( g(f_1(x), f_2(x), .., f_n(x)) \).
|
||
testChainRule1NQ_1 :: Falsify.Property ()
|
||
testChainRule1NQ_1 = do
|
||
f <- genSpray "f" 1
|
||
g <- genSpray "g" 1
|
||
let gof_spray = Spray.composeSpray g [f]
|
||
gof_chain =
|
||
chain @_ @3 @( ℝ 1 ) ( ℝ1 <$> fromSpray @3 @( ℝ 1 ) f ) ( fromSpray @3 @( ℝ 1 ) g )
|
||
Falsify.assert
|
||
$ Falsify.Prop.eq
|
||
.$ ("direct", fromSpray @3 @( ℝ 1 ) gof_spray )
|
||
.$ ("chain rule", gof_chain )
|
||
|
||
-- | Test the Faà di Bruno formula on polynomials, with a composition
|
||
-- \( g(f_1(x), f_2(x), .., f_n(x)) \).
|
||
testChainRule1NQ_2 :: Falsify.Property ()
|
||
testChainRule1NQ_2 = do
|
||
f1 <- genSpray "f1" 1
|
||
f2 <- genSpray "f2" 1
|
||
g <- genSpray "g" 2
|
||
let gof_spray = Spray.composeSpray g [f1, f2]
|
||
f = ℝ2 <$> fromSpray @3 @( ℝ 1 ) f1
|
||
<*> fromSpray @3 @( ℝ 1 ) f2
|
||
gof_chain =
|
||
chain @_ @3 @( ℝ 2 ) f ( fromSpray @3 @( ℝ 2 ) g )
|
||
Falsify.assert
|
||
$ Falsify.Prop.eq
|
||
.$ ("direct", fromSpray @3 @( ℝ 1 ) gof_spray )
|
||
.$ ("chain rule", gof_chain )
|
||
|
||
-- | Test the Faà di Bruno formula on polynomials, with a composition
|
||
-- \( g(f_1(x), f_2(x), .., f_n(x)) \).
|
||
testChainRule1NQ_3 :: Falsify.Property ()
|
||
testChainRule1NQ_3 = do
|
||
f1 <- genSpray "f1" 1
|
||
f2 <- genSpray "f2" 1
|
||
f3 <- genSpray "f3" 1
|
||
g <- genSpray "g" 3
|
||
let gof_spray = Spray.composeSpray g [f1, f2, f3]
|
||
f = ℝ3 <$> fromSpray @3 @( ℝ 1 ) f1
|
||
<*> fromSpray @3 @( ℝ 1 ) f2
|
||
<*> fromSpray @3 @( ℝ 1 ) f3
|
||
gof_chain =
|
||
chain @_ @3 @( ℝ 3 ) f ( fromSpray @3 @( ℝ 3 ) g )
|
||
Falsify.assert
|
||
$ Falsify.Prop.eq
|
||
.$ ("direct", fromSpray @3 @( ℝ 1 ) gof_spray )
|
||
.$ ("chain rule", gof_chain )
|
||
|
||
class FromSpray v where
|
||
varFn :: Int -> v
|
||
linFn :: v -> Int -> Double
|
||
|
||
instance FromSpray ( ℝ 1 ) where
|
||
varFn = \case
|
||
0 -> ℝ1 1
|
||
i -> error $ "fromSpray in 1d but variable " ++ show i
|
||
linFn ( ℝ1 x ) = \case
|
||
0 -> x
|
||
i -> error $ "fromSpray in 1d but variable " ++ show i
|
||
instance FromSpray ( ℝ 2 ) where
|
||
varFn = \case
|
||
0 -> ℝ2 1 0
|
||
1 -> ℝ2 0 1
|
||
i -> error $ "fromSpray in 2d but variable " ++ show i
|
||
linFn ( ℝ2 x y ) = \case
|
||
0 -> x
|
||
1 -> y
|
||
i -> error $ "fromSpray in 2d but variable " ++ show i
|
||
instance FromSpray ( ℝ 3 ) where
|
||
varFn = \case
|
||
0 -> ℝ3 1 0 0
|
||
1 -> ℝ3 0 1 0
|
||
2 -> ℝ3 0 0 1
|
||
i -> error $ "fromSpray in 3d but variable " ++ show i
|
||
linFn ( ℝ3 x y z ) = \case
|
||
0 -> x
|
||
1 -> y
|
||
2 -> z
|
||
i -> error $ "fromSpray in 3d but variable " ++ show i
|
||
|
||
genSpray :: String -> Word -> Falsify.Property ( Spray Double )
|
||
genSpray lbl nbVars = Falsify.genWith (\ p -> Just $ lbl ++ " = " ++ Spray.prettySpray show "x" p) $ do
|
||
deg <- Falsify.Gen.inRange $ Falsify.between ( 0, 10 )
|
||
let mons = allMonomials deg nbVars
|
||
coeffs <-
|
||
for mons $ \ mon -> do
|
||
if all (== 0) mon
|
||
then return Nothing
|
||
else do
|
||
nonZero <- Falsify.Gen.bool False
|
||
if nonZero
|
||
then return Nothing
|
||
else do
|
||
-- Just use (small) integral values in tests for now,
|
||
-- to avoid errors arising from rounding.
|
||
c <- Falsify.Gen.inRange $ Falsify.withOrigin ( -100, 100 ) ( 0 :: Int )
|
||
return $ Just ( map fromIntegral mon, fromIntegral c )
|
||
return $ Spray.fromList $ catMaybes coeffs
|
||
|
||
allMonomials :: Word -> Word -> [ [ Word ] ]
|
||
allMonomials k _ | k < 0 = []
|
||
allMonomials _ 0 = [ [] ]
|
||
allMonomials 0 n = [ replicate ( fromIntegral n ) 0 ]
|
||
allMonomials k n = [ i : is | i <- reverse [ 0 .. k ], is <- allMonomials ( k - i ) ( n - 1 ) ]
|
||
|
||
-- | Convert a multivariate polynomial from the @hspray@ library to the dual algebra.
|
||
fromSpray
|
||
:: forall k v
|
||
. ( HasChainRule Double k v
|
||
, Module Double (T v)
|
||
, Applicative ( D k v )
|
||
, Ring ( D k v Double )
|
||
, FromSpray v
|
||
)
|
||
=> Spray Double
|
||
-> D k v Double
|
||
fromSpray coeffs = HashMap.foldlWithKey' addMonomial ( konst @Double @k @v $ HashMap.lookupDefault 0 (Spray.Powers mempty 0) coeffs ) coeffs
|
||
where
|
||
addMonomial :: D k v Double -> Spray.Powers -> Double -> D k v Double
|
||
addMonomial a xs c = a + monomial c ( toList $ Spray.exponents xs )
|
||
monomial :: Double -> [ Int ] -> D k v Double
|
||
monomial _ [] = konst @Double @k @v 0
|
||
monomial c is = fmap ( c * ) $ go 0 is
|
||
go :: Int -> [ Int ] -> D k v Double
|
||
go _ [] = konst @Double @k @v 1
|
||
go d (i : is) = pow d i * go ( d + 1 ) is
|
||
pow :: Int -> Int -> D k v Double
|
||
pow _ 0 = konst @Double @k @v 1
|
||
pow d i = linearD @Double @k @v ( \ x -> linFn @v x d ) ( unT origin :: v ) ^ ( fromIntegral i )
|