mirror of
https://gitlab.com/sheaf/metabrush.git
synced 2024-11-05 23:03:38 +00:00
0160081e80
The code is now generic over the dimension. There is a slight performance loss that I need to investigate; perhaps some things are not getting specialised? Maybe it is better to be more explicit about staging and splice in the functions with fixed dimensions.
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
|
||
, MonomialBasisQ ( monTabulateQ, monIndexQ )
|
||
)
|
||
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" testMonomialBasisQD33
|
||
]
|
||
, 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 )
|
||
|
||
testMonomialBasisQD33 :: Falsify.Property ()
|
||
testMonomialBasisQD33 =
|
||
testRoundTrip genD33 \ d -> $$( monTabulateQ \ mon -> monIndexQ [|| 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 )
|