module ELynx.Data.Tree.MeasurableTree
( Measurable (..)
, distancesRootLeaves
, averageDistanceRootLeaves
, height
, lengthenRoot
, shortenRoot
, summarize
, totalBranchLength
, normalize
, prune
) where
import qualified Data.ByteString.Lazy.Char8 as L
import Data.Foldable
import Data.Tree
import ELynx.Data.Tree.Tree
class Measurable a where
getLen :: a -> Double
setLen :: Double -> a -> a
lengthen :: Double -> a -> a
lengthen dl l = setLen (dl + getLen l) l
shorten :: Double -> a -> a
shorten dl = lengthen (-dl)
distancesRootLeaves :: (Measurable a) => Tree a -> [Double]
distancesRootLeaves (Node l []) = [getLen l]
distancesRootLeaves (Node l f ) = concatMap (map (+ getLen l) . distancesRootLeaves) f
averageDistanceRootLeaves :: (Measurable a) => Tree a -> Double
averageDistanceRootLeaves tr = sum ds / fromIntegral n
where ds = distancesRootLeaves tr
n = length ds
height :: (Measurable a) => Tree a -> Double
height = maximum . distancesRootLeaves
lengthenRoot :: (Measurable a) => Double -> Tree a -> Tree a
lengthenRoot dl (Node lbl chs) = Node (lengthen dl lbl) chs
shortenRoot :: (Measurable a) => Double -> Tree a -> Tree a
shortenRoot dl = lengthenRoot (-dl)
summarize :: (Measurable a) => Tree a -> L.ByteString
summarize t = L.unlines $ map L.pack
[ "Leaves: " ++ show n ++ "."
, "Height: " ++ show h ++ "."
, "Average distance root to leaves: " ++ show h' ++ "."
, "Total branch length: " ++ show b ++ "." ]
where n = length . leaves $ t
h = height t
b = totalBranchLength t
h' = sum (distancesRootLeaves t) / fromIntegral n
totalBranchLength :: (Measurable a) => Tree a -> Double
totalBranchLength = foldl' (\acc n -> acc + getLen n) 0
normalize :: (Measurable a) => Tree a -> Tree a
normalize t = fmap (\n -> setLen (getLen n / s) n) t
where s = totalBranchLength t
prune :: (Measurable a) => Tree a -> Tree a
prune = pruneWith f
where f da pa = lengthen (getLen pa) da