BoxPlot.hs

{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ImportQualifiedPost #-}

module Main (main) where

import Chart
import Data.Text (Text)
import Prelude as P
import Prettychart
import DataFrame qualified as D
import DataFrame.Internal.Statistics qualified as D
import qualified Data.Vector.Algorithms.Intro as VA
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import GHC.Exts
import GHC.Generics
import Optics.Core
import Data.Text qualified as T
import qualified Data.Vector.Unboxed as VU

data BoxPlotOptions =
  BoxPlotOptions
    { q2Style :: Style
    , q3Style :: Style
    , q1Style :: Style
    , q4Style :: Style
    }
  deriving (Generic, Show, Eq)

defaultBoxPlotOptions :: BoxPlotOptions
defaultBoxPlotOptions =
  BoxPlotOptions
    { q2Style = defaultRectStyle
    , q3Style = defaultRectStyle
    , q1Style = (defaultRectStyle & #color .~ Colour 1 0 0 1)
    , q4Style = (defaultRectStyle & #color .~ Colour 1 0 0 1)
    }

boxAt :: BoxPlotOptions -> Double -> VU.Vector Double -> [Chart]
boxAt o x v =
  [ LineChart (view #q1Style o) [[Point x q0, Point x q1]]
  , RectChart (view #q2Style o) [Rect (x - w) (x + w) q1 q2]
  , RectChart (view #q3Style o) [Rect (x - w) (x + w) q2 q3]
  , LineChart (view #q4Style o) [[Point x q3, Point x q4]]
  , LineChart defaultLineStyle [[Point (x - w) q2, Point (x + w) q2]]
  ]
 where
  qs@[q0,q1,q2,q3,q4] =
    VU.toList $ D.quantiles' (VU.fromList [0,1,2,3,4]) 4 v
  w = 0.25


boxHud :: HudOptions
boxHud =
  defaultHudOptions
    & set #titles
        [ Priority 10 (defaultTitleOptions "Boxplot of Score"
                         & #place .~ PlaceTop
                         & #anchoring .~ 0.5
                         & #style % #size .~ 0.07)
        , Priority 8  (defaultTitleOptions "Score"
                         & #place .~ PlaceLeft
                         & #anchoring .~ 0.5
                         & #style % #size .~ 0.045)
        , Priority 8  (defaultTitleOptions "Teaching Method"
                         & #place .~ PlaceBottom
                         & #anchoring .~ 0.5
                         & #style % #size .~ 0.045)
        ]
    & set #frames
        [ Priority 1  (defaultFrameOptions
                        & #frame ?~ (defaultRectStyle
                                      & #color .~ paletteO 0 0.0
                                      & #borderColor .~ paletteO 0 0.5)
                        & #buffer .~ 0.02)
        , Priority 10 (defaultFrameOptions
                        & #frame ?~ (defaultRectStyle
                                      & #color .~ paletteO 4 0.10
                                      & #borderColor .~ paletteO 4 0.4)
                        & #buffer .~ 0.10)
        ]
    & set #legends []

boxPlot
  :: BoxPlotOptions
  -> [(Text, VU.Vector Double)]
  -> ChartOptions
boxPlot o xs =
  mempty
    & set #hudOptions boxHud
    & set (#markupOptions % #chartAspect) (FixedAspect 0.5)
    & set #chartTree boxPlotTree
 where
  xsWithX = zip [1 :: Double ..] xs

  allCharts =
    [ Chart (defaultRectStyle & #color .~ paletteO 1 0.4)
        (TextData [(lbl, (Point x ymin))])
    | (x,(lbl,_)) <- xsWithX
    ] ++
    concat
      [ boxAt o x v
      | (x,(_,v)) <- xsWithX
      ]

  boxPlotTree = named "boxplot" allCharts

  allVals = mconcat [ v | (_,v) <- xs ]
  (ymin, _)
    | VU.null allVals = (0,1)
    | otherwise       =
        let mn = VU.minimum allVals
            mx = VU.maximum allVals
            pad = 0.05 * (mx - mn)
        in (mn - pad, mx + pad)


main :: IO ()
main = do
  let m1 = VU.fromList [21, 25, 27, 18, 22, 30, 32]
      m2 = VU.fromList [23, 24, 26, 21, 25, 27, 30, 35]
      m3 = VU.fromList [10, 18, 20, 22, 28, 31, 32, 39]
      m4 = VU.fromList [24, 29, 30, 31, 33, 36, 37, 40]

      chart =
        boxPlot defaultBoxPlotOptions
          [ ("Method 1", m1)
          , ("Method 2", m2)
          , ("Method 3", m3)
          , ("Method 4", m4)
          ]

  writeChartOptions "./charts/boxplot.svg" chart

  (display, quit) <- startChartServer (Just "boxplot-demo")
  _ <- display $
      chart
        & set (#markupOptions % #markupHeight) (Just 600)
        & set (#markupOptions % #chartAspect) (FixedAspect 3)
        & set (#hudOptions % #frames % ix 1 % #item % #buffer) 0.10
  _ <- getLine
  quit