srtree-db-0.1.1.0: app/FitData.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module FitData
  ( FitDataOpts(..)
  , fitdataParser
  , runFitData
  ) where

import Control.Exception (bracket, SomeException, catch, displayException)
import Control.Monad.State.Strict (runStateT)
import Data.IORef
import Data.List (maximumBy)
import Data.Ord (comparing)
import qualified Data.Text as T
import qualified Data.Vector.Unboxed as VU
import Options.Applicative hiding (Const)
import System.Random (randomRIO)

import Data.SRTree (Fix(..), SRTree(..), Op(..), relabelParams, countParamsUniq, countNodes)
import Data.SRTree.Print (showExpr)
import Data.SRTree.Eval (Target)
import Data.SRTree.Datasets (loadDataset)
import Algorithm.SRTree.NonlinearOpt (minimizeNLL')
import Algorithm.SRTree.Likelihoods (Loss(..), Distribution(..), readLoss)
import Algorithm.SRTree.AD (ADBackEnd(..))
import Numeric.Optimization.NLOPT (LocalAlgorithm(..))
import Algorithm.EqSat.Egraph (EGraph(..), EClassId, getBestExpr, canonical)
import Algorithm.EqSat.Storage.Backend (SqlBackend(..), SqlValue(..), sqlToInt)
import Algorithm.EqSat.Storage.SQLite (loadGraphLazy, saveGraph, flushStore)
import Algorithm.EqSat.Storage.Query (getOrCreateDataset, writeDatasetFit)
import Algorithm.EqSat.Storage.Types (serializeTheta)
import Database.SQLite3 (Database, open, close)

-- | CLI options for the fitdata sub-command.
data FitDataOpts = FitDataOpts
  { fitdataDb        :: String
  , fitdataDataset   :: String
  , fitdataData      :: String
  , fitdataLoss      :: Loss
  , fitdataHasHeader :: Bool
  , fitdataNRep      :: Int
  , fitdataNIter     :: Int
  , fitdataBatchSize :: Int
  } deriving (Show)

fitdataParser :: Parser FitDataOpts
fitdataParser = FitDataOpts
  <$> strOption
      ( long "db"
      <> metavar "FILE"
      <> help "SQLite database file path" )
  <*> strOption
      ( long "dataset"
      <> metavar "NAME"
      <> help "Dataset name" )
  <*> strOption
      ( long "data"
      <> metavar "SPEC"
      <> help "Dataset CSV spec: file:start:end:target:features:yerr" )
  <*> option auto
      ( long "loss"
      <> value (NLL Gaussian)
      <> metavar "LOSS"
      <> help "Loss function (MSE, NLL Gaussian, etc.)" )
  <*> switch
      ( long "has-header"
      <> help "CSV has header row (default: True)" )
  <*> option auto
      ( long "n-rep"
      <> value 1
      <> metavar "N"
      <> help "Number of random restarts per expression" )
  <*> option auto
      ( long "n-iter"
      <> value 100
      <> metavar "N"
      <> help "Max NLopt iterations" )
  <*> option auto
      ( long "batch-size"
      <> value 100
      <> metavar "N"
      <> help "Fit N expressions per commit batch" )

-- | Run the fitdata sub-command.
runFitData :: FitDataOpts -> IO ()
runFitData opts = do
  let FitDataOpts{..} = opts
  -- 1. Load dataset
  putStrLn $ "Loading dataset: " ++ fitdataData
  ((xTrain, yTrain, _xVal, _yVal), (mYErr, _), _varnames, _target) <-
    loadDataset fitdataData fitdataHasHeader

  let nNoiseParams = case fitdataLoss of
        NLL Gaussian -> 1
        NLL ROXY     -> 3
        _            -> 0

  -- 2. Load paged graph and query for unfitted e-classes
  putStrLn $ "Loading paged graph from " ++ fitdataDb ++ "..."
  withSQLite fitdataDb $ \db -> do
    dsid <- getOrCreateDataset db fitdataDataset
    er <- loadGraphLazy db dsid
    case er of
      Left err -> putStrLn $ "loadGraphLazy failed: " ++ err
      Right eg -> do
        -- Query for unfitted e-classes
        unfitted <- queryUnfitted db dsid
        let total = length unfitted
        putStrLn $ "Found " ++ show total ++ " unfitted e-classes"

        if total == 0
          then putStrLn "Nothing to fit."
          else do
            -- Process in batches
            counter <- newIORef (0 :: Int)
            let processBatch [] = pure ()
                processBatch batch = do
                  mapM_ (fitOne db dsid eg xTrain yTrain mYErr fitdataLoss fitdataNIter fitdataNRep nNoiseParams counter) batch
                  -- Commit checkpoint
                  flushStore eg
                  putStrLn $ "  [checkpoint] committed batch"

            let batches = chunk fitdataBatchSize unfitted
            mapM_ processBatch batches

            fitted <- readIORef counter
            putStrLn $ "Fitted " ++ show fitted ++ "/" ++ show total
                     ++ " expressions"

-- | Fit a single e-class and write the result.
fitOne :: SqlBackend db
       => db -> Int -> EGraph
       -> [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double)
       -> Loss -> Int -> Int -> Int -> IORef Int -> EClassId -> IO ()
fitOne db dsid eg xTrain yTrain mYErr loss nIter nRep nNoiseParams counter eid = do
  n <- readIORef counter
  let label = "[" ++ show (n+1) ++ "]"
  -- Extract best expression from the e-class
  mTree <- runStateT (getBestExpr eid) eg >>= pure . Just . fst
  case mTree of
    Nothing -> do
      putStrLn $ label ++ " eclass " ++ show eid ++ ": could not extract expression (skipped)"
      modifyIORef' counter (+1)
    Just tree -> do
      let tree' = relabelParams tree
          np = countParamsUniq tree' + nNoiseParams
          sz = countNodes tree'
      -- Skip expressions with no parameters (pure constants/variables)
      if np == 0
        then do
          putStrLn $ label ++ " eclass " ++ show eid ++ ": no parameters (skipped)"
          modifyIORef' counter (+1)
        else do
          -- Try multiple random restarts, pick the best
          results <- mapM (\_ -> fitOneRandom xTrain yTrain mYErr loss nIter tree' np) [1..nRep]
          let (bestFitness, bestTheta) = maximumBy (comparing fst) results
          -- Write result to DB
          writeDatasetFit db dsid eid (Just bestFitness) Nothing
            (T.pack (serializeTheta [bestTheta])) sz
          putStrLn $ label ++ " eclass " ++ show eid ++ " (" ++ takeExpr tree' ++ "): fitness=" ++ showFit bestFitness
          modifyIORef' counter (+1)

-- | Fit an expression with a random initial theta.
fitOneRandom :: [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double)
             -> Loss -> Int -> Fix SRTree -> Int -> IO (Double, VU.Vector Double)
fitOneRandom xTrain yTrain mYErr loss nIter tree np = do
  theta0 <- VU.replicateM np (randomRIO (-1, 1))
  let (theta, lossVal, _) = minimizeNLL' VAR1 SingleThread loss mYErr nIter xTrain yTrain tree theta0
      fitness = negate lossVal
  pure (fitness, theta)

-- | Query for e-class IDs that are not yet fitted for a dataset.
queryUnfitted :: SqlBackend db => db -> Int -> IO [EClassId]
queryUnfitted db dsid = do
  rows <- queryDb db
    "SELECT e.eid FROM eclass e \
    \LEFT JOIN dataset_fit df ON df.eid = e.eid AND df.dataset_id = ? \
    \WHERE df.fitted IS NULL OR df.fitted = 0"
    [SqlInteger (fromIntegral dsid)]
  pure [ sqlToInt eid | [eid] <- rows ]

-- | Chunk a list into sub-lists of the given size.
chunk :: Int -> [a] -> [[a]]
chunk _ [] = []
chunk n xs = let (h, t) = splitAt n xs in h : chunk n t

-- | Take first few chars of an expression for display.
takeExpr :: Fix SRTree -> String
takeExpr t
  | length s > 40 = take 40 s ++ "..."
  | otherwise = s
  where s = showExpr t

-- | Format fitness for display.
showFit :: Double -> String
showFit f
  | f == (-1/0) = "-Infinity"
  | f == (1/0)  = "Infinity"
  | isNaN f     = "NaN"
  | otherwise   = show (fromIntegral (round (f * 1000) :: Int) / 1000 :: Double)

-- | Open a SQLite database, run an action, and close it.
withSQLite :: String -> (Database -> IO a) -> IO a
withSQLite path = bracket (open (T.pack path)) close