module Bio.Phylogeny.PhyBin.PreProcessor
( collapseBranches,
collapseBranchLenThresh, collapseBranchBootStrapThresh,
pruneTreeLeaves
)
where
import qualified Data.Set as S
import Data.Maybe (catMaybes)
import Bio.Phylogeny.PhyBin.CoreTypes
pruneTreeLeaves :: S.Set Label -> NewickTree DefDecor -> Maybe (NewickTree DefDecor)
pruneTreeLeaves set tr = loop tr
where
loop orig@(NTLeaf _ lab)
| S.member lab set = Just orig
| otherwise = Nothing
loop (NTInterior dec@(_,blen) ls) =
case catMaybes $ map loop ls of
[] -> Nothing
[one] -> Just$ fmap (\(bs,bl) -> (bs,blen + bl)) one
ls' -> Just (NTInterior dec ls')
collapseBranches :: forall a . (a -> Bool) -> (a -> a -> a) -> NewickTree a -> NewickTree a
collapseBranches isCollapsable collapse origtr = final
where
(_,_, final) = loop origtr
loop :: NewickTree a -> ([(a,Label)], [NewickTree a], NewickTree a)
loop lf@(NTLeaf dec lb) | isCollapsable dec = ([(dec,lb)], [], lf)
| otherwise = ([], [lf], lf)
loop (NTInterior dec children) =
let (floats, anchors, _) = unzip3 $ map loop children
thenode = NTInterior dec $ concat anchors ++
map (uncurry NTLeaf) (concat floats)
in
if isCollapsable dec then
(map (\ (a,l) -> (collapse dec a, l))
(concat floats),
concat anchors,
thenode)
else
([], [thenode], thenode)
collapseBranchLenThresh :: Double -> NewickTree DefDecor -> NewickTree DefDecor
collapseBranchLenThresh thr tr =
collapseBranches ((< thr) . getBranchLen) collapser tr
where
collapser _intermediate@(_bt1, len1) _floatee@(_bt2, len2) =
(Nothing, len1 + len2)
collapseBranchBootStrapThresh :: Int -> NewickTree DefDecor -> NewickTree DefDecor
collapseBranchBootStrapThresh thr tr =
collapseBranches ((< thr) . getBoot) collapser tr
where
getBoot (Nothing,_) = error$"collapseBranchBootStrapThresh: bootstrap value missing on tree node!\n"++
"They must be present if --minbootstrap is used."
getBoot (Just boot,_) = boot
collapser (_,len1) (_,len2) = (Nothing, len1+len2)