Synthesis.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

{- | Feature-synthesis enumerator: it should recover small exact features
(@x²@, @a/b@) from the example rows, score them near-perfectly, and be
deterministic.
-}
module Learn.Synthesis (tests) where

import qualified DataFrame as D
import DataFrame.Model (fit)
import DataFrame.Synthesis

import Test.HUnit

quad :: D.DataFrame
quad =
    D.fromNamedColumns
        [ ("x", D.fromList xs)
        , ("y", D.fromList (map (\x -> x * x) xs))
        ]
  where
    xs = map fromIntegral [1 .. 12 :: Int] :: [Double]

ratio :: D.DataFrame
ratio =
    D.fromNamedColumns
        [ ("a", D.fromList ([2, 6, 12, 20, 30, 42] :: [Double]))
        , ("b", D.fromList ([1, 2, 3, 4, 5, 6] :: [Double]))
        , ("y", D.fromList ([2, 3, 4, 5, 6, 7] :: [Double]))
        ]

-- | Pearson r²: enumeration should find x·x and score it ~1.
recoversQuadratic :: Test
recoversQuadratic = TestCase $ do
    let sf = fit defaultSynthesisConfig (D.col @Double "y") quad
    assertBool
        ("quadratic r2 = " ++ show (sfScore sf))
        (sfScore sf > 0.999 && sfScore sf <= 1.0001)

-- | MSE: the best feature reproduces the target exactly, so -MSE ~ 0.
exactRecoveryMSE :: Test
exactRecoveryMSE = TestCase $ do
    let sf =
            fit defaultSynthesisConfig{synLoss = MeanSquaredError} (D.col @Double "y") quad
    assertBool ("exact -mse = " ++ show (sfScore sf)) (sfScore sf > -1.0e-6)

-- | Division is enumerated (with the denominator guard): a/b is recovered.
recoversRatio :: Test
recoversRatio = TestCase $ do
    let sf = fit defaultSynthesisConfig (D.col @Double "y") ratio
    assertBool
        ("ratio r2 = " ++ show (sfScore sf))
        (sfScore sf > 0.999 && sfScore sf <= 1.0001)

-- | The bank is non-trivial and its ranked features are distinct expressions.
distinctFeatures :: Test
distinctFeatures = TestCase $ do
    let sf = fit defaultSynthesisConfig (D.col @Double "y") quad
        names = [D.prettyPrint e | (e, _) <- sfFeatures sf]
    assertBool "synthesizes more than one feature" (length names > 1)
    assertBool "ranked features are distinct" (length names == length (dedup names))
  where
    dedup = foldr (\x acc -> if x `elem` acc then acc else x : acc) []

-- | Same config and data give the same best expression.
deterministic :: Test
deterministic = TestCase $ do
    let a = fit defaultSynthesisConfig (D.col @Double "y") quad
        b = fit defaultSynthesisConfig (D.col @Double "y") quad
    assertEqual
        "same best expression"
        (D.prettyPrint (sfExpr a))
        (D.prettyPrint (sfExpr b))

tests :: [Test]
tests =
    [ recoversQuadratic
    , exactRecoveryMSE
    , recoversRatio
    , distinctFeatures
    , deterministic
    ]