MeasurementTree.hs 2.63 KB
Newer Older
Ben Gamari's avatar
Ben Gamari committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}

module MeasurementTree where

import Control.Applicative
import Control.Monad
import Data.Aeson
import qualified Data.Map.Strict as M
import GHC.Generics
import qualified Data.ByteString.Lazy as BSL

type MeasurementTree = MeasurementTree' Double

data MeasurementTree' a = Sample a
                          -- ^ A measurement result
                        | Group (M.Map String (MeasurementTree' a))
                          -- ^ A group of related measurements.
  deriving (Functor, Generic)

instance ToJSON a => ToJSON (MeasurementTree' a) where
    toJSON (Sample x) = toJSON x
    toJSON (Group x) = object [ "type" .= ("group" :: String)
                              , "children" .= x
                              ]
instance FromJSON a => FromJSON (MeasurementTree' a) where
    parseJSON v = parseGroup v <|> fmap Sample (parseJSON v)
      where
        parseGroup = withObject "group" $ \o -> do
          typ <- o .: "type"
          guard $ typ == ("group" :: String)
          Group <$> o .: "children"

--instance ToJSON a => ToJSON (MeasurementTree' a)
--instance FromJSON a => FromJSON (MeasurementTree' a)

singleton :: String -> MeasurementTree' a -> MeasurementTree' a
singleton name = Group . M.singleton name

singletonSample :: String -> a -> MeasurementTree' a
singletonSample name = singleton name . Sample

writeFile' :: ToJSON a => FilePath -> MeasurementTree' a -> IO ()
writeFile' fname = BSL.writeFile fname . encode

writeFile :: FilePath -> MeasurementTree -> IO ()
writeFile = writeFile'

readFile' :: FromJSON a => FilePath -> IO (Maybe (MeasurementTree' a))
readFile' fname = decode <$> BSL.readFile fname

readFile :: FilePath -> IO (Maybe MeasurementTree)
readFile = readFile'

group :: [(String, MeasurementTree' a)] -> MeasurementTree' a
group = Group . M.fromList

collect :: RealFrac a => [MeasurementTree' a] -> MeasurementTree' [a]
collect = go [""]
  where
    isSample (Sample _) = True
    isSample _ = False
    isGroup = not . isSample

    go path trees
      | all isSample trees = Sample [ x | Sample x <- trees ]
      | all isGroup  trees = 
        Group $ M.mapWithKey (\k ts -> go (k:path) ts) $ M.fromListWith (<>)
        [ (k,[t])
        | Group xs <- trees 
        , (k,t) <- M.toList xs
        ]
      | otherwise = error $ "Measurement tree type mismatch at " ++ show path

mean :: RealFrac a => [a] -> a
mean xs = sum xs / realToFrac (length xs)

geomMean :: RealFloat a => [a] -> a
geomMean = exp . mean . map log

stdDev :: RealFloat a => [a] -> a
stdDev xs = sqrt $ mean $ map (\x -> (x-m)^(2::Int)) xs
  where m = mean xs