dataframe-1.0.0.0: app/Synthesis.hs

{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}

import qualified Data.Text as T
import qualified DataFrame as D
import qualified DataFrame.Functions as F

import Data.Char
import DataFrame.DecisionTree
import DataFrame.Operators hiding (name)

import System.Random

$(F.declareColumnsFromCsvFile "./data/titanic/train.csv")

main :: IO ()
main = do
    train <- D.readCsv "./data/titanic/train.csv"
    test <- D.readCsv "./data/titanic/test.csv"

    -- Apply the same transformations to training and test.
    let combined =
            (train <> test)
                |> D.deriveMany
                    [ "Ticket" .= F.lift (T.filter isAlpha) ticket
                    , "Name" .= F.match "\\s*([A-Za-z]+)\\." name
                    , "Cabin" .= F.whenPresent (T.take 1) cabin
                    ]
                |> D.renameMany
                    [ (F.name name, "title")
                    , (F.name cabin, "cabin_prefix")
                    , (F.name pclass, "passenger_class")
                    , (F.name sibsp, "number_of_siblings_and_spouses")
                    , (F.name parch, "number_of_parents_and_children")
                    ]
    print combined

    let (train', validation) =
            D.take
                (D.nRows train)
                combined
                |> D.filterJust (F.name survived)
                |> D.randomSplit (mkStdGen 4232) 0.7
        -- Split the test out again.
        test' =
            D.drop
                (D.nRows train)
                combined

        model =
            fitDecisionTree
                ( defaultTreeConfig
                    { maxTreeDepth = 5
                    , minSamplesSplit = 5
                    , minLeafSize = 3
                    , taoIterations = 100
                    , synthConfig =
                        defaultSynthConfig
                            { complexityPenalty = 0.1
                            , maxExprDepth = 3
                            , disallowedCombinations =
                                [ (F.name age, F.name fare)
                                , ("passenger_class", "number_of_siblings_and_spouses")
                                , ("passenger_class", "number_of_parents_and_children")
                                ]
                            }
                    }
                )
                survived -- Label to predict
                ( train'
                    |> D.exclude [F.name passengerid]
                )

    print model

    putStrLn "Training accuracy: "
    print $
        computeAccuracy
            (train' |> D.derive (F.name prediction) model)

    putStrLn "Validation accuracy: "
    print $
        computeAccuracy
            ( validation
                |> D.derive (F.name prediction) model
            )

    let predictions = D.derive (F.name survived) model test'
    D.writeCsv
        "./predictions.csv"
        (predictions |> D.select [F.name passengerid, F.name survived])

prediction :: D.Expr Int
prediction = F.col @Int "prediction"

computeAccuracy :: D.DataFrame -> Double
computeAccuracy df =
    let
        tp =
            fromIntegral $
                D.nRows
                    ( D.filterWhere
                        (survived .== F.lit (1 :: Int) .&& prediction .== F.lit (1 :: Int))
                        df
                    )
        tn =
            fromIntegral $
                D.nRows
                    ( D.filterWhere
                        (survived .== F.lit (0 :: Int) .&& prediction .== F.lit (0 :: Int))
                        df
                    )
        fp =
            fromIntegral $
                D.nRows
                    ( D.filterWhere
                        (survived .== F.lit (0 :: Int) .&& prediction .== F.lit (1 :: Int))
                        df
                    )
        fn =
            fromIntegral $
                D.nRows
                    ( D.filterWhere
                        (survived .== F.lit (1 :: Int) .&& prediction .== F.lit (0 :: Int))
                        df
                    )
     in
        (tp + tn) / (tp + tn + fp + fn)