{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

{- | Feature synthesis by bottom-up enumerative search with observational
equivalence — the canonical enumerative method from Solar-Lezama's
/Introduction to Program Synthesis/.

Given a frame and a numeric target column, it searches for a small, interpretable
arithmetic expression over the other columns whose values track the target.

The engine:

  * enumerates programs by increasing AST size (so the first representative of any
    behaviour is the smallest — interpretability for free);
  * evaluates each candidate /incrementally/ by combining the cached result
    vectors of its subprograms (one vector op), never re-interpreting the whole
    tree;
  * keeps exactly one program per /observational-equivalence/ class — candidates
    producing the same column (up to a float tolerance) are interchangeable, so
    duplicates are dropped rather than re-explored;
  * breaks commutative symmetry (never both @a+b@ and @b+a@) and uses protected
    operators (@sqrt|x|@, @log(|x|+1)@) plus a denominator guard so domain errors
    never arise;
  * caps each size layer by fit score when it grows large (a cost-guided
    tractability bound over /distinct/ behaviours, not a lossy beam over raw
    syntax).

'fit' returns the best 'SynthesizedFeature'; 'predict' is its expression.
'synthesizeFeatures' returns the whole ranked, deduplicated feature bank — useful
as automated feature engineering feeding a downstream model.

Deferred (documented next steps, not yet implemented): skeleton enumeration with
closed-form least-squares coefficient fitting, hard-row counterexample sampling
for very large frames, and piecewise (condition-abduction) features.
-}
module DataFrame.Synthesis (
    module DataFrame.Model,
    LossFunction (..),
    SynthesisConfig (..),
    defaultSynthesisConfig,
    SynthesizedFeature (..),
    synthesizeFeatures,
) where

import Control.Exception (throw)
import Data.Bits (xor)
import Data.Either (fromRight)
import Data.List (sortBy)
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import Data.Ord (Down (..), comparing)
import qualified Data.Text as T
import qualified Data.Vector.Unboxed as VU
import Data.Word (Word64)
import GHC.Float (castDoubleToWord64)

import DataFrame.Errors (DataFrameException (..))
import DataFrame.Featurize.Internal (featureNames)
import qualified DataFrame.Functions as F
import DataFrame.Internal.DataFrame (DataFrame)
import DataFrame.Internal.Expression (Expr (..))
import DataFrame.Internal.Statistics (
    meanSquaredError,
    mutualInformationBinned,
    percentile',
    variance',
 )
import DataFrame.Model
import DataFrame.Operations.Core (columnAsDoubleVector)

-- | How a candidate's output column is scored against the target (higher is better).
data LossFunction
    = -- | Pearson @r²@: scale-invariant, the default for derived features.
      PearsonCorrelation
    | -- | Binned mutual information: captures nonlinear association.
      MutualInformation
    | -- | Negative mean squared error: for reproducing a target exactly.
      MeanSquaredError
    deriving (LossFunction -> LossFunction -> Bool
(LossFunction -> LossFunction -> Bool)
-> (LossFunction -> LossFunction -> Bool) -> Eq LossFunction
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LossFunction -> LossFunction -> Bool
== :: LossFunction -> LossFunction -> Bool
$c/= :: LossFunction -> LossFunction -> Bool
/= :: LossFunction -> LossFunction -> Bool
Eq, Int -> LossFunction -> ShowS
[LossFunction] -> ShowS
LossFunction -> String
(Int -> LossFunction -> ShowS)
-> (LossFunction -> String)
-> ([LossFunction] -> ShowS)
-> Show LossFunction
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LossFunction -> ShowS
showsPrec :: Int -> LossFunction -> ShowS
$cshow :: LossFunction -> String
show :: LossFunction -> String
$cshowList :: [LossFunction] -> ShowS
showList :: [LossFunction] -> ShowS
Show)

-- | Search hyperparameters.
data SynthesisConfig = SynthesisConfig
    { SynthesisConfig -> Int
synMaxSize :: !Int
    -- ^ Largest AST (node count) to enumerate.
    , SynthesisConfig -> Int
synBankCap :: !Int
    -- ^ Max observationally-distinct programs kept per size layer.
    , SynthesisConfig -> LossFunction
synLoss :: !LossFunction
    , SynthesisConfig -> Int
synTopK :: !Int
    -- ^ How many ranked features to return in the bank.
    , SynthesisConfig -> Int
synMaxAllocBytes :: !Int
    -- ^ Refuse a search whose largest layer would allocate beyond this.
    }
    deriving (SynthesisConfig -> SynthesisConfig -> Bool
(SynthesisConfig -> SynthesisConfig -> Bool)
-> (SynthesisConfig -> SynthesisConfig -> Bool)
-> Eq SynthesisConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SynthesisConfig -> SynthesisConfig -> Bool
== :: SynthesisConfig -> SynthesisConfig -> Bool
$c/= :: SynthesisConfig -> SynthesisConfig -> Bool
/= :: SynthesisConfig -> SynthesisConfig -> Bool
Eq, Int -> SynthesisConfig -> ShowS
[SynthesisConfig] -> ShowS
SynthesisConfig -> String
(Int -> SynthesisConfig -> ShowS)
-> (SynthesisConfig -> String)
-> ([SynthesisConfig] -> ShowS)
-> Show SynthesisConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SynthesisConfig -> ShowS
showsPrec :: Int -> SynthesisConfig -> ShowS
$cshow :: SynthesisConfig -> String
show :: SynthesisConfig -> String
$cshowList :: [SynthesisConfig] -> ShowS
showList :: [SynthesisConfig] -> ShowS
Show)

defaultSynthesisConfig :: SynthesisConfig
defaultSynthesisConfig :: SynthesisConfig
defaultSynthesisConfig =
    SynthesisConfig
        { synMaxSize :: Int
synMaxSize = Int
6
        , synBankCap :: Int
synBankCap = Int
500
        , synLoss :: LossFunction
synLoss = LossFunction
PearsonCorrelation
        , synTopK :: Int
synTopK = Int
16
        , synMaxAllocBytes :: Int
synMaxAllocBytes = Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024
        }

{- | A synthesized feature. 'sfExpr' is the best-scoring expression and 'sfFeatures'
is the ranked, observationally-distinct bank (expression and its score).
-}
data SynthesizedFeature = SynthesizedFeature
    { SynthesizedFeature -> Expr Double
sfExpr :: !(Expr Double)
    , SynthesizedFeature -> Double
sfScore :: !Double
    , SynthesizedFeature -> [(Expr Double, Double)]
sfFeatures :: ![(Expr Double, Double)]
    }

instance Fit SynthesisConfig (Expr Double) where
    type ModelOf SynthesisConfig (Expr Double) = SynthesizedFeature
    fit :: CheckFrame
  (FrameReq SynthesisConfig (Expr Double))
  (FrameFor (Expr Double)) =>
SynthesisConfig
-> Expr Double
-> FrameFor (Expr Double)
-> FitResult
     (FrameFor (Expr Double)) (ModelOf SynthesisConfig (Expr Double))
fit = SynthesisConfig -> Expr Double -> DataFrame -> SynthesizedFeature
SynthesisConfig
-> Expr Double
-> FrameFor (Expr Double)
-> FitResult
     (FrameFor (Expr Double)) (ModelOf SynthesisConfig (Expr Double))
synthesizeFeatures

instance Predict SynthesizedFeature where
    type Prediction SynthesizedFeature = Expr Double
    predict :: SynthesizedFeature -> Prediction SynthesizedFeature
predict = SynthesizedFeature -> Expr Double
SynthesizedFeature -> Prediction SynthesizedFeature
sfExpr

-- | A candidate's evaluated column over the example rows.
type Output = VU.Vector Double

data Prog = Prog
    { Prog -> Expr Double
progExpr :: !(Expr Double)
    , Prog -> Int
progSize :: !Int
    , Prog -> Output
progOut :: !Output
    }

-- | Search for expressions over the non-target columns that track @target@.
synthesizeFeatures ::
    SynthesisConfig -> Expr Double -> DataFrame -> SynthesizedFeature
synthesizeFeatures :: SynthesisConfig -> Expr Double -> DataFrame -> SynthesizedFeature
synthesizeFeatures SynthesisConfig
cfg Expr Double
target DataFrame
df
    | [Prog] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Prog]
leaves Bool -> Bool -> Bool
|| Output -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Output
tgt = Expr Double
-> Double -> [(Expr Double, Double)] -> SynthesizedFeature
SynthesizedFeature (Double -> Expr Double
forall a. Columnable a => a -> Expr a
Lit Double
0) (Double -> Double
forall a. Num a => a -> a
negate (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)) []
    | Just DataFrameException
err <- SynthesisConfig -> Int -> Int -> Maybe DataFrameException
oversizedSearch SynthesisConfig
cfg ([Prog] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Prog]
leaves) Int
n = DataFrameException -> SynthesizedFeature
forall a e. Exception e => e -> a
throw DataFrameException
err
    | Bool
otherwise = Expr Double
-> Double -> [(Expr Double, Double)] -> SynthesizedFeature
SynthesizedFeature Expr Double
best Double
bestScore [(Expr Double, Double)]
ranked
  where
    feats :: [Text]
feats = Expr Double -> DataFrame -> [Text]
forall a. Expr a -> DataFrame -> [Text]
featureNames Expr Double
target DataFrame
df
    tgt :: Output
tgt = Output -> Either DataFrameException Output -> Output
forall b a. b -> Either a b -> b
fromRight Output
forall a. Unbox a => Vector a
VU.empty (Expr Double -> DataFrame -> Either DataFrameException Output
forall a.
(Columnable a, Num a) =>
Expr a -> DataFrame -> Either DataFrameException Output
columnAsDoubleVector Expr Double
target DataFrame
df)
    n :: Int
n = Output -> Int
forall a. Unbox a => Vector a -> Int
VU.length Output
tgt
    leaves :: [Prog]
leaves = DataFrame -> [Text] -> Int -> [Prog]
mkLeaves DataFrame
df [Text]
feats Int
n
    bank :: [Prog]
bank = SynthesisConfig -> Output -> [Prog] -> [Prog]
grow SynthesisConfig
cfg Output
tgt [Prog]
leaves
    scored :: [(Expr Double, Int, Double)]
scored =
        [ (Prog -> Expr Double
progExpr Prog
p, Prog -> Int
progSize Prog
p, Double
s)
        | Prog
p <- [Prog]
bank
        , Just Double
s <- [LossFunction -> Output -> Output -> Maybe Double
scoreOf (SynthesisConfig -> LossFunction
synLoss SynthesisConfig
cfg) Output
tgt (Prog -> Output
progOut Prog
p)]
        ]
    sorted :: [(Expr Double, Int, Double)]
sorted = ((Expr Double, Int, Double)
 -> (Expr Double, Int, Double) -> Ordering)
-> [(Expr Double, Int, Double)] -> [(Expr Double, Int, Double)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Expr Double, Int, Double) -> (Down Double, Int))
-> (Expr Double, Int, Double)
-> (Expr Double, Int, Double)
-> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (\(Expr Double
_, Int
sz, Double
s) -> (Double -> Down Double
forall a. a -> Down a
Down Double
s, Int
sz))) [(Expr Double, Int, Double)]
scored
    ranked :: [(Expr Double, Double)]
ranked = [(Expr Double
e, Double
s) | (Expr Double
e, Int
_, Double
s) <- Int -> [(Expr Double, Int, Double)] -> [(Expr Double, Int, Double)]
forall a. Int -> [a] -> [a]
take (SynthesisConfig -> Int
synTopK SynthesisConfig
cfg) [(Expr Double, Int, Double)]
sorted]
    (Expr Double
best, Double
bestScore) = case [(Expr Double, Double)]
ranked of
        ((Expr Double
e, Double
s) : [(Expr Double, Double)]
_) -> (Expr Double
e, Double
s)
        [] -> (Double -> Expr Double
forall a. Columnable a => a -> Expr a
Lit Double
0, Double -> Double
forall a. Num a => a -> a
negate (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0))

-- | Size-1 programs: numeric feature columns and a pool of constants, OE-deduped.
mkLeaves :: DataFrame -> [T.Text] -> Int -> [Prog]
mkLeaves :: DataFrame -> [Text] -> Int -> [Prog]
mkLeaves DataFrame
df [Text]
feats Int
n = ([Prog], Seen) -> [Prog]
forall a b. (a, b) -> a
fst (Seen -> [Prog] -> ([Prog], Seen)
dedupProgs Seen
forall k a. Map k a
M.empty [Prog]
candidates)
  where
    candidates :: [Prog]
candidates =
        [ Expr Double -> Int -> Output -> Prog
Prog (Text -> Expr Double
forall a. Columnable a => Text -> Expr a
Col Text
name) Int
1 Output
o
        | Text
name <- [Text]
feats
        , Right Output
o <- [Expr Double -> DataFrame -> Either DataFrameException Output
forall a.
(Columnable a, Num a) =>
Expr a -> DataFrame -> Either DataFrameException Output
columnAsDoubleVector (Text -> Expr Double
forall a. Columnable a => Text -> Expr a
Col Text
name :: Expr Double) DataFrame
df]
        ]
            [Prog] -> [Prog] -> [Prog]
forall a. [a] -> [a] -> [a]
++ [Expr Double -> Int -> Output -> Prog
Prog (Double -> Expr Double
forall a. Columnable a => a -> Expr a
Lit Double
v) Int
1 (Int -> Double -> Output
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
n Double
v) | Double
v <- DataFrame -> [Text] -> [Double]
constantPool DataFrame
df [Text]
feats]

{- | Domain-informed constants: per-column quartiles, variance, and std, plus a few
small integers. (Duplicates collapse under observational equivalence.)
-}
constantPool :: DataFrame -> [T.Text] -> [Double]
constantPool :: DataFrame -> [Text] -> [Double]
constantPool DataFrame
df [Text]
feats =
    [Double
0, Double
1, Double
2, -Double
1]
        [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [ Int -> Double -> Double
roundSig Int
2 Double
v
           | Text
name <- [Text]
feats
           , Right Output
c <- [Expr Double -> DataFrame -> Either DataFrameException Output
forall a.
(Columnable a, Num a) =>
Expr a -> DataFrame -> Either DataFrameException Output
columnAsDoubleVector (Text -> Expr Double
forall a. Columnable a => Text -> Expr a
Col Text
name :: Expr Double) DataFrame
df]
           , Double
v <-
                [Int -> Output -> Double
forall a. (Unbox a, Num a, Real a) => Int -> Vector a -> Double
percentile' Int
p Output
c | Int
p <- [Int
1, Int
25, Int
75, Int
99]] [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Output -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
variance' Output
c, Double -> Double
forall a. Floating a => a -> a
sqrt (Output -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
variance' Output
c)]
           ]

-- | Grow the bank one size layer at a time, keeping one program per OE class.
grow :: SynthesisConfig -> Output -> [Prog] -> [Prog]
grow :: SynthesisConfig -> Output -> [Prog] -> [Prog]
grow SynthesisConfig
cfg Output
tgt [Prog]
leaves = Int -> [Prog] -> Seen -> [Prog]
go Int
2 [Prog]
leaves ((Prog -> Seen -> Seen) -> Seen -> [Prog] -> Seen
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Output -> Seen -> Seen
seenInsert (Output -> Seen -> Seen)
-> (Prog -> Output) -> Prog -> Seen -> Seen
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog -> Output
progOut) Seen
forall k a. Map k a
M.empty [Prog]
leaves)
  where
    go :: Int -> [Prog] -> Seen -> [Prog]
go Int
size [Prog]
bank Seen
seen
        | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> SynthesisConfig -> Int
synMaxSize SynthesisConfig
cfg = [Prog]
bank
        | Bool
otherwise =
            let ([Prog]
kept, Seen
seen') = SynthesisConfig -> Output -> Seen -> [Prog] -> ([Prog], Seen)
absorb SynthesisConfig
cfg Output
tgt Seen
seen (Int -> [Prog] -> [Prog]
layer Int
size [Prog]
bank)
             in Int -> [Prog] -> Seen -> [Prog]
go (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([Prog]
bank [Prog] -> [Prog] -> [Prog]
forall a. [a] -> [a] -> [a]
++ [Prog]
kept) Seen
seen'

-- | All candidate programs of exactly @size@ nodes, built from smaller ones.
layer :: Int -> [Prog] -> [Prog]
layer :: Int -> [Prog] -> [Prog]
layer Int
size [Prog]
bank = [Prog]
unaries [Prog] -> [Prog] -> [Prog]
forall a. [a] -> [a] -> [a]
++ [Prog]
pows [Prog] -> [Prog] -> [Prog]
forall a. [a] -> [a] -> [a]
++ [Prog]
comms [Prog] -> [Prog] -> [Prog]
forall a. [a] -> [a] -> [a]
++ [Prog]
subs [Prog] -> [Prog] -> [Prog]
forall a. [a] -> [a] -> [a]
++ [Prog]
divs
  where
    atSize :: Int -> [Prog]
atSize Int
s = (Prog -> Bool) -> [Prog] -> [Prog]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s) (Int -> Bool) -> (Prog -> Int) -> Prog -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog -> Int
progSize) [Prog]
bank
    args1 :: [Prog]
args1 = Int -> [Prog]
atSize (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    unaries :: [Prog]
unaries =
        [ Expr Double -> Int -> Output -> Prog
Prog (Expr Double -> Expr Double
mk Expr Double
e) Int
size ((Double -> Double) -> Output -> Output
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map Double -> Double
f Output
o)
        | (Expr Double -> Expr Double
mk, Double -> Double
f) <- [(Expr Double -> Expr Double, Double -> Double)]
unaryProds
        , Prog Expr Double
e Int
_ Output
o <- [Prog]
args1
        ]
    pows :: [Prog]
pows =
        [ Expr Double -> Int -> Output -> Prog
Prog (Expr Double -> Int -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Int -> Expr a
F.pow Expr Double
e Int
k) Int
size ((Double -> Double) -> Output -> Output
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
k) Output
o)
        | Prog Expr Double
e Int
_ Output
o <- [Prog]
args1
        , Int
k <- [Int
2 .. Int
6 :: Int]
        ]
    comms :: [Prog]
comms =
        [ Expr Double -> Int -> Output -> Prog
Prog (Expr Double -> Expr Double -> Expr Double
mk Expr Double
ea Expr Double
eb) Int
size ((Double -> Double -> Double) -> Output -> Output -> Output
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith Double -> Double -> Double
f Output
oa Output
ob)
        | (Expr Double -> Expr Double -> Expr Double
mk, Double -> Double -> Double
f) <- [(Expr Double -> Expr Double -> Expr Double,
  Double -> Double -> Double)]
commutativeProds
        , (Prog Expr Double
ea Int
_ Output
oa, Prog Expr Double
eb Int
_ Output
ob) <- Int -> [Prog] -> [(Prog, Prog)]
unorderedPairs Int
size [Prog]
bank
        ]
    subs :: [Prog]
subs =
        [ Expr Double -> Int -> Output -> Prog
Prog (Expr Double
ea Expr Double -> Expr Double -> Expr Double
forall a. Num a => a -> a -> a
- Expr Double
eb) Int
size ((Double -> Double -> Double) -> Output -> Output -> Output
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (-) Output
oa Output
ob)
        | (Prog Expr Double
ea Int
_ Output
oa, Prog Expr Double
eb Int
_ Output
ob) <- Int -> [Prog] -> [(Prog, Prog)]
orderedPairs Int
size [Prog]
bank
        ]
    divs :: [Prog]
divs =
        [ Expr Double -> Int -> Output -> Prog
Prog (Expr Double
ea Expr Double -> Expr Double -> Expr Double
forall a. Fractional a => a -> a -> a
/ Expr Double
eb) Int
size ((Double -> Double -> Double) -> Output -> Output -> Output
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith Double -> Double -> Double
forall a. Fractional a => a -> a -> a
(/) Output
oa Output
ob)
        | (Prog Expr Double
ea Int
_ Output
oa, Prog Expr Double
eb Int
_ Output
ob) <- Int -> [Prog] -> [(Prog, Prog)]
orderedPairs Int
size [Prog]
bank
        , (Double -> Bool) -> Output -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1e-9) (Double -> Bool) -> (Double -> Double) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Num a => a -> a
abs) Output
ob
        ]

-- | Protected unary operators: total on all reals (no NaN/domain errors).
unaryProds :: [(Expr Double -> Expr Double, Double -> Double)]
unaryProds :: [(Expr Double -> Expr Double, Double -> Double)]
unaryProds =
    [ (Expr Double -> Expr Double
forall a. Floating a => a -> a
sqrt (Expr Double -> Expr Double)
-> (Expr Double -> Expr Double) -> Expr Double -> Expr Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr Double -> Expr Double
forall a. Num a => a -> a
abs, Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Num a => a -> a
abs)
    , (Expr Double -> Expr Double
forall a. Num a => a -> a
abs, Double -> Double
forall a. Num a => a -> a
abs)
    , (\Expr Double
e -> Expr Double -> Expr Double
forall a. Floating a => a -> a
log (Expr Double -> Expr Double
forall a. Num a => a -> a
abs Expr Double
e Expr Double -> Expr Double -> Expr Double
forall a. Num a => a -> a -> a
+ Expr Double
1), \Double
x -> Double -> Double
forall a. Floating a => a -> a
log (Double -> Double
forall a. Num a => a -> a
abs Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1))
    , (Expr Double -> Expr Double
forall a. Floating a => a -> a
exp, Double -> Double
forall a. Floating a => a -> a
exp)
    , (Expr Double -> Expr Double
forall a. Floating a => a -> a
sin, Double -> Double
forall a. Floating a => a -> a
sin)
    , (Expr Double -> Expr Double
forall a. Floating a => a -> a
cos, Double -> Double
forall a. Floating a => a -> a
cos)
    , (Expr Double -> Expr Double
forall a. (Columnable a, Num a, Ord a) => Expr a -> Expr a
F.relu, Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0)
    , (Expr Double -> Expr Double
forall a. Num a => a -> a
signum, Double -> Double
forall a. Num a => a -> a
signum)
    ]

-- | Commutative binary operators (enumerated over unordered operand pairs).
commutativeProds ::
    [(Expr Double -> Expr Double -> Expr Double, Double -> Double -> Double)]
commutativeProds :: [(Expr Double -> Expr Double -> Expr Double,
  Double -> Double -> Double)]
commutativeProds =
    [ (Expr Double -> Expr Double -> Expr Double
forall a. Num a => a -> a -> a
(+), Double -> Double -> Double
forall a. Num a => a -> a -> a
(+))
    , (Expr Double -> Expr Double -> Expr Double
forall a. Num a => a -> a -> a
(*), Double -> Double -> Double
forall a. Num a => a -> a -> a
(*))
    , (Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr a
F.min, Double -> Double -> Double
forall a. Ord a => a -> a -> a
min)
    , (Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr a
F.max, Double -> Double -> Double
forall a. Ord a => a -> a -> a
max)
    ]

-- | Ordered operand pairs whose sizes sum to @size-1@ (for non-commutative ops).
orderedPairs :: Int -> [Prog] -> [(Prog, Prog)]
orderedPairs :: Int -> [Prog] -> [(Prog, Prog)]
orderedPairs Int
size [Prog]
bank =
    [ (Prog
a, Prog
b)
    | Int
sa <- [Int
1 .. Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
    , let sb :: Int
sb = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sa
    , Int
sb Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
    , Prog
a <- Int -> [Prog]
atSize Int
sa
    , Prog
b <- Int -> [Prog]
atSize Int
sb
    ]
  where
    atSize :: Int -> [Prog]
atSize Int
s = (Prog -> Bool) -> [Prog] -> [Prog]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s) (Int -> Bool) -> (Prog -> Int) -> Prog -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog -> Int
progSize) [Prog]
bank

-- | Unordered operand pairs (for commutative ops): each pair once.
unorderedPairs :: Int -> [Prog] -> [(Prog, Prog)]
unorderedPairs :: Int -> [Prog] -> [(Prog, Prog)]
unorderedPairs Int
size [Prog]
bank =
    [ (Prog
a, Prog
b)
    | Int
sa <- [Int
1 .. Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
    , let sb :: Int
sb = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sa
    , Int
sb Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
    , Int
sa Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
sb
    , (Int
i, Prog
a) <- [Int] -> [Prog] -> [(Int, Prog)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] (Int -> [Prog]
atSize Int
sa)
    , (Int
j, Prog
b) <- [Int] -> [Prog] -> [(Int, Prog)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] (Int -> [Prog]
atSize Int
sb)
    , Int
sa Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sb Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
j
    ]
  where
    atSize :: Int -> [Prog]
atSize Int
s = (Prog -> Bool) -> [Prog] -> [Prog]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s) (Int -> Bool) -> (Prog -> Int) -> Prog -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog -> Int
progSize) [Prog]
bank

{- | Keep the valid, observationally-novel candidates of a layer, then cap by fit
score (cost-guided). Returns the kept programs and the updated OE-class set.
-}
absorb ::
    SynthesisConfig -> Output -> Seen -> [Prog] -> ([Prog], Seen)
absorb :: SynthesisConfig -> Output -> Seen -> [Prog] -> ([Prog], Seen)
absorb SynthesisConfig
cfg Output
tgt Seen
seen0 [Prog]
cands = (SynthesisConfig -> Output -> [Prog] -> [Prog]
capLayer SynthesisConfig
cfg Output
tgt [Prog]
fresh, Seen
seen')
  where
    ([Prog]
fresh, Seen
seen') = Seen -> [Prog] -> ([Prog], Seen)
dedupProgs Seen
seen0 [Prog]
cands

oversizedSearch :: SynthesisConfig -> Int -> Int -> Maybe DataFrameException
oversizedSearch :: SynthesisConfig -> Int -> Int -> Maybe DataFrameException
oversizedSearch SynthesisConfig
cfg Int
nLeaves Int
n
    | SynthesisConfig -> Int
synMaxSize SynthesisConfig
cfg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
3 = Maybe DataFrameException
forall a. Maybe a
Nothing
    | Int
estimate Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
budget = Maybe DataFrameException
forall a. Maybe a
Nothing
    | Bool
otherwise =
        DataFrameException -> Maybe DataFrameException
forall a. a -> Maybe a
Just
            ( Text -> DataFrameException
InternalException
                ( Text
"synthesizeFeatures: a search to synMaxSize="
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (SynthesisConfig -> Int
synMaxSize SynthesisConfig
cfg))
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" over "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
nLeaves)
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" leaves and "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
n)
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" rows would allocate about "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (Int
estimate Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024)))
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" GiB, past the "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (Int
budget Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024)))
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" GiB synMaxAllocBytes budget. Lower synMaxSize to "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
largestFittingSize)
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", narrow the column set, or raise synMaxAllocBytes."
                )
            )
  where
    budget :: Int
budget = SynthesisConfig -> Int
synMaxAllocBytes SynthesisConfig
cfg
    bytesPerProg :: Int
bytesPerProg = Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n
    -- 4 commutative ops over unordered pairs plus sub and div over ordered
    -- pairs, all quadratic in the bank; unaries and powers are lower order.
    binaryOpCount :: Int
binaryOpCount = Int
4 :: Int
    -- Layer k pairs the capped bank with itself over the binary operators.
    candidatesAt :: Int -> Int
candidatesAt Int
k
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2 = Int
nLeaves
        | Bool
otherwise =
            Int
binaryOpCount Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (SynthesisConfig -> Int
synBankCap SynthesisConfig
cfg) (Int -> Int
candidatesAt (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)
    estimate :: Int
estimate = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Int
candidatesAt Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bytesPerProg | Int
k <- [Int
2 .. SynthesisConfig -> Int
synMaxSize SynthesisConfig
cfg]]
    largestFittingSize :: Int
largestFittingSize =
        [Int] -> Int
forall a. HasCallStack => [a] -> a
last
            ( Int
3
                Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [ Int
k
                  | Int
k <- [Int
3 .. SynthesisConfig -> Int
synMaxSize SynthesisConfig
cfg]
                  , [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Int
candidatesAt Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bytesPerProg | Int
j <- [Int
2 .. Int
k]] Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
budget
                  ]
            )

-- | When a layer has more distinct programs than the cap, keep the best-scoring.
capLayer :: SynthesisConfig -> Output -> [Prog] -> [Prog]
capLayer :: SynthesisConfig -> Output -> [Prog] -> [Prog]
capLayer SynthesisConfig
cfg Output
tgt [Prog]
progs
    | [Prog] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Prog]
progs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= SynthesisConfig -> Int
synBankCap SynthesisConfig
cfg = [Prog]
progs
    | Bool
otherwise = Int -> [Prog] -> [Prog]
forall a. Int -> [a] -> [a]
take (SynthesisConfig -> Int
synBankCap SynthesisConfig
cfg) ((Prog -> Prog -> Ordering) -> [Prog] -> [Prog]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((Prog -> Down Double) -> Prog -> Prog -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Double -> Down Double
forall a. a -> Down a
Down (Double -> Down Double) -> (Prog -> Double) -> Prog -> Down Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog -> Double
rank)) [Prog]
progs)
  where
    rank :: Prog -> Double
rank Prog
p = Double -> Maybe Double -> Double
forall a. a -> Maybe a -> a
fromMaybe (Double -> Double
forall a. Num a => a -> a
negate (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)) (LossFunction -> Output -> Output -> Maybe Double
scoreOf (SynthesisConfig -> LossFunction
synLoss SynthesisConfig
cfg) Output
tgt (Prog -> Output
progOut Prog
p))

-- | Fit score of an output against the target (higher is better), or @Nothing@.
scoreOf :: LossFunction -> Output -> Output -> Maybe Double
scoreOf :: LossFunction -> Output -> Output -> Maybe Double
scoreOf LossFunction
lf Output
tgt Output
out
    | Output -> Int
forall a. Unbox a => Vector a -> Int
VU.length Output
out Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Output -> Int
forall a. Unbox a => Vector a -> Int
VU.length Output
tgt = Maybe Double
forall a. Maybe a
Nothing
    | Bool
otherwise = Maybe Double -> Maybe Double
forall {a}. RealFloat a => Maybe a -> Maybe a
finite (Maybe Double -> Maybe Double) -> Maybe Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ case LossFunction
lf of
        LossFunction
PearsonCorrelation -> Output -> Output -> Maybe Double
pearsonR2 Output
tgt Output
out
        LossFunction
MutualInformation -> Int -> Output -> Output -> Maybe Double
mutualInformationBinned Int
bins Output
tgt Output
out
        LossFunction
MeanSquaredError -> Double -> Double
forall a. Num a => a -> a
negate (Double -> Double) -> Maybe Double -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Output -> Output -> Maybe Double
meanSquaredError Output
tgt Output
out
  where
    bins :: Int
bins = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
10 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Output -> Int
forall a. Unbox a => Vector a -> Int
VU.length Output
tgt) :: Double)))
    -- Belt-and-suspenders: drop any non-finite score so it cannot win the ranking.
    finite :: Maybe a -> Maybe a
finite (Just a
s) | a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
s Bool -> Bool -> Bool
|| a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
s = Maybe a
forall a. Maybe a
Nothing
    finite Maybe a
ms = Maybe a
ms

{- | Pearson @r²@ via the numerically stable centered two-pass formula. Returns
'Nothing' when the feature (or target) is constant, and is bounded by
Cauchy–Schwarz to @[0,1]@ — unlike the one-pass @n·Σxy − Σx·Σy@ form, which
cancels catastrophically for low-variance features and can report @r² > 1@.
-}
pearsonR2 :: Output -> Output -> Maybe Double
pearsonR2 :: Output -> Output -> Maybe Double
pearsonR2 Output
ys Output
xs
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = Maybe Double
forall a. Maybe a
Nothing
    | Double
sxx Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Double
syy Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = Maybe Double
forall a. Maybe a
Nothing
    | Bool
otherwise = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 (Double
sxy Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sxy Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
sxx Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
syy)))
  where
    n :: Int
n = Output -> Int
forall a. Unbox a => Vector a -> Int
VU.length Output
xs
    nf :: Double
nf = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    mx :: Double
mx = Output -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Output
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nf
    my :: Double
my = Output -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Output
ys Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nf
    sxy :: Double
sxy = Output -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum ((Double -> Double -> Double) -> Output -> Output -> Output
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (\Double
x Double
y -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
my)) Output
xs Output
ys)
    sxx :: Double
sxx = Output -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum ((Double -> Double) -> Output -> Output
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx)) Output
xs)
    syy :: Double
syy = Output -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum ((Double -> Double) -> Output -> Output
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Double
y -> (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
my) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
my)) Output
ys)

-- | An output is usable iff it is non-empty and free of NaN/±Inf.
valid :: Output -> Bool
valid :: Output -> Bool
valid Output
o = Bool -> Bool
not (Output -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Output
o) Bool -> Bool -> Bool
&& (Double -> Bool) -> Output -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all (\Double
x -> Bool -> Bool
not (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x)) Output
o

{- | Quantize an output to nine significant digits, so float noise (@x*2@ vs
@x+x@) collapses while genuinely distinct features stay apart.
-}
quantize :: Output -> Output
quantize :: Output -> Output
quantize = (Double -> Double) -> Output -> Output
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Double
x -> if Double
x Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double -> Double
forall a. Num a => a -> a
signum Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double -> Double
roundSig Int
9 (Double -> Double
forall a. Num a => a -> a
abs Double
x))

{- | The observational-equivalence class set: a map from a 64-bit FNV-1a
fingerprint of the quantized output to the (usually one) quantized outputs with
that fingerprint. Bucketing on the fingerprint keeps membership cheap, and
verifying exact equality within the bucket makes a hash collision harmless — two
genuinely different columns that happen to collide are kept apart, not merged.
-}
type Seen = M.Map Int [Output]

-- | FNV-1a fingerprint of an already-quantized output's bit patterns.
fpOf :: Output -> Int
fpOf :: Output -> Int
fpOf = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int) -> (Output -> Word64) -> Output -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64 -> Double -> Word64) -> Word64 -> Output -> Word64
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' Word64 -> Double -> Word64
step (Word64
1469598103934665603 :: Word64)
  where
    step :: Word64 -> Double -> Word64
step !Word64
h Double
x = (Word64
h Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Double -> Word64
castDoubleToWord64 Double
x) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
1099511628211

-- | Record an output's observational-equivalence class.
seenInsert :: Output -> Seen -> Seen
seenInsert :: Output -> Seen -> Seen
seenInsert Output
o = ([Output] -> [Output] -> [Output])
-> Int -> [Output] -> Seen -> Seen
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [Output] -> [Output] -> [Output]
forall a. [a] -> [a] -> [a]
(++) (Output -> Int
fpOf Output
q) [Output
q]
  where
    q :: Output
q = Output -> Output
quantize Output
o

-- | Round a positive double to @n@ significant digits.
roundSig :: Int -> Double -> Double
roundSig :: Int -> Double -> Double
roundSig Int
n Double
x
    | Double
x Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = Double
0
    | Bool
otherwise =
        let magnitude :: Int
magnitude = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
10 (Double -> Double
forall a. Num a => a -> a
abs Double
x)) :: Int
            scale :: Double
scale = Double
10 Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
magnitude)
         in Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scale) :: Integer) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
scale

{- | Keep the first valid program of each observational-equivalence class,
preserving order; returns the kept programs and the grown class set.
-}
dedupProgs :: Seen -> [Prog] -> ([Prog], Seen)
dedupProgs :: Seen -> [Prog] -> ([Prog], Seen)
dedupProgs = [Prog] -> Seen -> [Prog] -> ([Prog], Seen)
go []
  where
    go :: [Prog] -> Seen -> [Prog] -> ([Prog], Seen)
go [Prog]
acc Seen
s [] = ([Prog] -> [Prog]
forall a. [a] -> [a]
reverse [Prog]
acc, Seen
s)
    go [Prog]
acc Seen
s (Prog
p : [Prog]
ps)
        | Bool -> Bool
not (Output -> Bool
valid Output
o) = [Prog] -> Seen -> [Prog] -> ([Prog], Seen)
go [Prog]
acc Seen
s [Prog]
ps
        | Bool
member = [Prog] -> Seen -> [Prog] -> ([Prog], Seen)
go [Prog]
acc Seen
s [Prog]
ps
        | Bool
otherwise = [Prog] -> Seen -> [Prog] -> ([Prog], Seen)
go (Prog
p Prog -> [Prog] -> [Prog]
forall a. a -> [a] -> [a]
: [Prog]
acc) (([Output] -> [Output] -> [Output])
-> Int -> [Output] -> Seen -> Seen
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [Output] -> [Output] -> [Output]
forall a. [a] -> [a] -> [a]
(++) Int
fp [Output
q] Seen
s) [Prog]
ps
      where
        o :: Output
o = Prog -> Output
progOut Prog
p
        q :: Output
q = Output -> Output
quantize Output
o
        fp :: Int
fp = Output -> Int
fpOf Output
q
        member :: Bool
member = Bool -> ([Output] -> Bool) -> Maybe [Output] -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Output
q Output -> [Output] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem`) (Int -> Seen -> Maybe [Output]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
fp Seen
s)