Ensemble.hs

-- | M4 — ensemble search. Run an inner optimizer several times on different
-- /bootstrap resamples/ of the training set (sampling with replacement, the
-- classic bagging trick that makes the resulting candidates complementary),
-- collect the candidate programs, and combine them into one program with EP-5's
-- 'Shikumi.Combinator.ensemble' combinator under a majority-vote reducer. The
-- ensemble runs all members on the input and returns their modal answer, which is
-- correct more often than any single member when the members err on different
-- inputs.
--
-- Resampling is deterministic (a fixed linear-congruential stream seeded by the
-- member index), so the search is reproducible run to run.
--
-- The returned artifact is structure-changing: it is an @Ensemble@ combinator over
-- the optimized members, not just a parameter rewrite of the student. The member
-- parameters are persisted by 'Shikumi.Compile.Serialize.encodeCompiled', but the
-- reducer closure and exact member structure live in the program template held in
-- code. Load saved state onto the matching ensemble template, not onto the plain
-- student.
module Shikumi.Optimize.Ensemble
  ( ensembleSearch,
    ensembleSearchWith,
    majorityReducer,
  )
where

import Data.Vector qualified as V
import Shikumi.Combinator (ensemble)
import Shikumi.Compile.Types (compiledProgram)
import Shikumi.Eval (Dataset, dataset, datasetExamples)
import Shikumi.Optimize.Search (freezeProgram, withLmCallCount)
import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget)

-- | Build an @size@-member ensemble: run @inner@ on @size@ bootstrap resamples of
-- the training set and combine the resulting programs by majority vote. This
-- changes structure by returning an @Ensemble@; saved state must be decoded onto a
-- matching ensemble template.
ensembleSearch :: (Eq o) => Int -> Optimizer i o -> Optimizer i o
ensembleSearch = ensembleSearchWith defaultBudget

-- | Build a budgeted ensemble. The budget is enforced between members using exact
-- LLM-call counting; at least one member always runs, so the final spend may exceed
-- the bound by one member's cost.
ensembleSearchWith :: (Eq o) => Budget -> Int -> Optimizer i o -> Optimizer i o
ensembleSearchWith budget size inner = Optimizer $ \train metric student -> do
  let seeds = [1 .. max 1 size]
      go _total acc [] = pure (reverse acc)
      go total acc (seed : rest)
        | total >= maxLmCalls budget && not (null acc) = pure (reverse acc)
        | otherwise = do
            (cp, calls) <- withLmCallCount (runOptimizer inner (resample seed train) metric student)
            go (total + calls) (compiledProgram cp : acc) rest
  members <- go 0 [] seeds
  pure (freezeProgram (ensemble members majorityReducer))

-- | Sample a dataset with replacement, deterministically, seeded by @seed@. An
-- empty dataset is returned unchanged.
resample :: Int -> Dataset i o -> Dataset i o
resample seed ds =
  let exs = V.fromList (datasetExamples ds)
      m = V.length exs
   in if m == 0
        then ds
        else dataset [exs V.! (i `mod` m) | i <- take m (lcg seed)]

-- | A simple linear-congruential pseudo-random stream (glibc constants). Pure and
-- deterministic — there is no @Math.random@ in this environment, and reproducible
-- sampling is what the test relies on.
lcg :: Int -> [Int]
lcg s0 = drop 1 (iterate step s0)
  where
    step s = (1103515245 * s + 12345) `mod` 2147483648

-- | The modal value of a non-empty list under 'Eq': most frequent, ties broken by
-- first appearance. This is the ensemble's default vote. (EP-4's @modal@ is not
-- exported, so it is reproduced here.)
majorityReducer :: (Eq o) => [o] -> o
majorityReducer items = case counts of
  [] -> error "Shikumi.Optimize.Ensemble.majorityReducer: empty member list"
  (c : cs) -> fst (foldl' (\best cur -> if snd cur > snd best then cur else best) c cs)
  where
    counts = [(v, length (filter (== v) items)) | v <- distinct items]
    distinct = go []
      where
        go _ [] = []
        go seen (y : ys)
          | any (== y) seen = go seen ys
          | otherwise = y : go (seen ++ [y]) ys