cladogram.ml

   1 (* virt-similarity
   2  * Copyright (C) 2013 Red Hat Inc.
   3  *
   4  * This program is free software; you can redistribute it and/or modify
   5  * it under the terms of the GNU General Public License as published by
   6  * the Free Software Foundation; either version 2 of the License, or
   7  * (at your option) any later version.
   8  *
   9  * This program is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12  * GNU General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU General Public License along
  15  * with this program; if not, write to the Free Software Foundation, Inc.,
  16  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  17  *)
  18
  19 open Printf
  20
  21 open Utils
  22
  23 type t =
  24   | Leaf of int                  (* A single disk image (index of). *)
  25   | Node of t list * int         (* An interior node in the tree. *)
  26
  27 let rec images_in_subtree = function
  28   | Leaf i -> [i]
  29   | Node (xs, _) -> List.concat (List.map images_in_subtree xs)
  30
  31 let max_list = List.fold_left max min_int
  32
  33 let mapi f xs =
  34   let rec loop i = function
  35     | [] -> []
  36     | x :: xs -> let r = f i x in r :: loop (i+1) xs
  37   in
  38   loop 0 xs
  39
  40 (* Compute the minimum distance between subtrees t1 and t2.  'matrix'
  41  * is the distance matrix between leaves.
  42  *)
  43 let min_distance_between_subtrees matrix t1 t2 =
  44   let min_distance = ref max_int in
  45
  46   let xs = images_in_subtree t1 in
  47   let ys = images_in_subtree t2 in
  48
  49   List.iter (
  50     fun (i, j) ->
  51       let d = matrix.(i).(j) in
  52       if d < !min_distance then min_distance := d
  53   ) (pairs xs ys);
  54
  55   !min_distance
  56
  57 (* Find the closest subtrees and combine them. *)
  58 let combine_closest_subtrees matrix trees =
  59   let trees = Array.of_list trees in
  60   let n = Array.length trees in
  61
  62   (* Find the minimum distance between any two subtrees. *)
  63   let min_distance = ref max_int in
  64   List.iter (
  65     fun (i, j) ->
  66       let d = min_distance_between_subtrees matrix trees.(i) trees.(j) in
  67       if d < !min_distance then min_distance := d
  68   ) (pairs_of_ints n);
  69
  70   let min_distance = !min_distance in
  71
  72   (* For each subtree that differs from another by exactly the
  73    * minimum distance, group them together into a single subtree.
  74    *)
  75   let in_group = Array.make n false in
  76   List.iter (
  77     fun (i, j) ->
  78       let d = min_distance_between_subtrees matrix trees.(i) trees.(j) in
  79       if d = min_distance then (
  80         in_group.(i) <- true;
  81         in_group.(j) <- true
  82       )
  83   ) (pairs_of_ints n);
  84
  85   let group = ref [] and rest = ref [] in
  86   Array.iteri (
  87     fun i in_group ->
  88       if in_group then
  89         group := trees.(i) :: !group
  90       else
  91         rest := trees.(i) :: !rest
  92   ) in_group;
  93
  94   !rest @ [Node (!group, min_distance)]
  95
  96 let construct_cladogram matrix n =
  97   (* At the bottom level, every disk image is in its own leaf. *)
  98   let leaves =
  99     let rec loop i = if i < n then Leaf i :: loop (i+1) else [] in
 100     loop 0 in
 101
 102   (* Work up the tree, combining subtrees together, until we end
 103    * up with one tree (ie. the top node of the final tree).
 104    *)
 105   let rec loop trees =
 106     match trees with
 107     | [] -> assert false
 108     | [x] -> x (* finished *)
 109     | xs -> loop (combine_closest_subtrees matrix xs)
 110   in
 111   loop leaves
 112
 113 let format_cladogram ?format_leaf t =
 114   let format_leaf = match format_leaf with
 115     | None -> string_of_int
 116     | Some f -> f
 117   in
 118   let rec format = function
 119     | Leaf i ->
 120       let s = "--- " ^ format_leaf i in
 121       [s; ""], String.length s
 122     | Node (xs, _) ->
 123       let xs = List.map format xs in
 124       let n = List.length xs in
 125       let w = 7 + max_list (List.map snd xs) in
 126       let xs = mapi (
 127         fun row (ss, _) ->
 128           let s, ss = match ss with
 129             | s :: ss -> s, ss
 130             | [] -> assert false in
 131           if row = 0 then (
 132             ("---+---" ^ s) ::
 133               List.map (fun s -> "   |   " ^ s) ss
 134           ) else if row < n-1 then (
 135             ("   +---" ^ s) ::
 136               List.map (fun s -> "   |   " ^ s) ss
 137           ) else (
 138             ("   +---" ^ s) ::
 139               List.map (fun s -> "       " ^ s) ss
 140           )
 141       ) xs in
 142       List.concat xs, w
 143   in
 144   let strs, _ = format t in
 145   strs