module ELynx.Data.Tree.BranchSupportTree
( BranchSupport
, BranchSupportLabel (..)
, normalize
, collapse
) where
import Data.List
import Data.Maybe
import Data.Tree
type BranchSupport = Maybe Double
class BranchSupportLabel a where
getBranchSupport :: a -> BranchSupport
setBranchSupport :: BranchSupport -> a -> a
apply :: BranchSupportLabel a => (Double -> Double) -> a -> a
apply f l = setBranchSupport (f <$> s) l
where s = getBranchSupport l
normalize :: BranchSupportLabel a => Tree a -> Tree a
normalize t = if isNothing m then t else fmap (apply (/ fromJust m)) t
where m = maximum $ fmap getBranchSupport t
accept :: Double -> Maybe Double -> Bool
accept _ Nothing = True
accept thresh (Just s) = s > thresh
collapse :: BranchSupportLabel a => Double -> Tree a -> Tree a
collapse _ n@(Node _ []) = n
collapse thresh (Node l xs) = Node l $ map (collapse thresh) (highS ++ lowSubForest)
where (highS, lowS) = partition (accept thresh . getBranchSupport . rootLabel) xs
lowSubForest = concatMap subForest lowS