d_tree.ml

open Proto_t
open Model_t

type m = {
  max_depth : int;
  feature_map : Feat_map.t;
  splitter : Loss.splitter
}

let string_of_split { s_gamma ; s_n ; s_loss } =
  Printf.sprintf "gamma=%f n=%d loss=%f" s_gamma s_n s_loss

let directions_of_split s_to_k s_pivot =
  let num_categories = Array.length s_to_k in
  let directions = Array.create num_categories `Right in
  assert (0 <= s_pivot && s_pivot < num_categories );
  for s = 0 to s_pivot do
    let k = s_to_k.(s) in
    directions.(k) <- `Left
  done;
    (*
      for s = s_pivot to num_directions-1 do
       let k = s_to_k.(s) in
        directions.(k) <- `Right
      done;
    *)
  directions

let timeout = 1000.

let best_split_of_splits best_splits =
  List.fold_left (
    fun best_opt (_,_, response) ->
      let s_opt =
        match response with
          | `AckBestSplit s_opt -> s_opt
          | _ -> assert false
      in
      match best_opt, s_opt with
        | Some (best_loss, best_split), Some (loss, split) ->

          if best_loss < loss then
            (* still superior *)
            best_opt
          else
            (* new champ *)
            Some (loss, split)

        | None, Some (loss, split) ->
          (* first guy's always champ *)
          Some (loss, split)

        | Some _, None -> best_opt
        | None, None -> None

  ) None best_splits

let rec terminate (best_split : Proto_t.split) =
  let difference_in_gamma =
    let os =
      match best_split with
        | `OrdinalSplit os -> os
        | `CategoricalSplit (os, _) -> os
    in
    os.os_left.s_gamma -. os.os_right.s_gamma
  in
  if abs_float difference_in_gamma < -1e8 then
    None (* split not found *)
  else
    let node =
      match best_split with
        | `CategoricalSplit (os, s_to_k) ->
          let directions = directions_of_split s_to_k os.os_split in
          `CategoricalNode {
            cn_feature_id = os.os_feature_id;
            cn_category_directions = directions;
            cn_left_tree = `Leaf os.os_left.s_gamma;
            cn_right_tree = `Leaf os.os_right.s_gamma;
          }

        | `OrdinalSplit os -> `OrdinalNode {
            on_feature_id = os.os_feature_id;
            on_split = os.os_split;
            on_left_tree = `Leaf os.os_left.s_gamma;
            on_right_tree = `Leaf os.os_right.s_gamma
          }
    in
    Some node

and make task_id max_depth workers depth =
  let open Worker_client in
  let request = `Learning (task_id, `Sample) in
  let is_response_valid = function
    | `AckSample -> true
    | _ -> false
  in
  lwt result = broad_send_recv workers timeout request is_response_valid in

  let request = `Learning (task_id, `BestSplit) in
  let is_response_valid = function
    | `AckBestSplit _ -> true
    | _ -> false
  in
  lwt result = broad_send_recv workers timeout request is_response_valid in

  match best_split_of_splits result with
    | None -> Lwt.return None

    | Some (loss, split) ->

      if depth + 1 >= max_depth then
        Lwt.return (terminate split)

      else
        let in_subset_left, in_subset_right =
          partition_observations in_subset splitting_feature split in

        let os =
          match split with
            | `CategoricalSplit (os,_) -> os
            | `OrdinalSplit os -> os
        in

        let side_left =
          match make m (depth+1) in_subset_left with
            | None -> `Leaf os.os_left.s_gamma
            | Some tree -> tree
        in

        let side_right =
          match make m (depth+1) in_subset_right with
            | None -> `Leaf os.os_right.s_gamma
            | Some tree -> tree
        in

        let node =
          match split with
            | `CategoricalSplit (os, s_to_k) ->
              let directions = directions_of_split s_to_k os.os_split in
              `CategoricalNode {
                cn_feature_id = os.os_feature_id;
                cn_category_directions = directions;
                cn_left_tree = side_left;
                cn_right_tree = side_right
              }

            | `OrdinalSplit os -> `OrdinalNode {
                on_feature_id = os.os_feature_id;
                on_split = os.os_split;
                on_left_tree = side_left;
                on_right_tree = side_right;
              }
        in
        Some node


(* partially evaluate a tree when the value of feature with id
   [feature_id] is [value] *)
let rec eval_partially_1 feature_id value = function
  | `OrdinalNode {
      on_split;
      on_left_tree;
      on_right_tree;
      on_feature_id
    } ->
    if on_feature_id = feature_id then
      let sub_tree =
        if value <= on_split then
          on_left_tree
        else
          on_right_tree
      in
      eval_partially_1 feature_id value sub_tree
    else
      (* create a new tree, the result of simplifying each subtree *)
      let on_left_tree  = eval_partially_1 feature_id value on_left_tree in
      let on_right_tree = eval_partially_1 feature_id value on_right_tree in
      `OrdinalNode { on_split; on_left_tree; on_right_tree; on_feature_id }

  | `CategoricalNode {
      cn_category_directions;
      cn_left_tree;
      cn_right_tree;
      cn_feature_id
    } ->
    if cn_feature_id = feature_id then
      let sub_tree =
        match cn_category_directions.(value) with
          | `Left -> cn_left_tree
          | `Right -> cn_right_tree
      in
      eval_partially_1 feature_id value sub_tree
    else
      let cn_left_tree  = eval_partially_1 feature_id value cn_left_tree in
      let cn_right_tree = eval_partially_1 feature_id value cn_right_tree in
      `CategoricalNode {
        cn_category_directions;
        cn_left_tree;
        cn_right_tree;
        cn_feature_id
      }

  | (`Leaf _) as leaf -> leaf


let is_leaf = function
  | `Leaf _ -> true
  | _ -> false

let eval_partially trees cardinality feature_id = function
  | `Dense v ->
    let width_num_bytes = Utils.num_bytes cardinality in
    let num_leaves = ref 0 in
    Dense.iter ~width:width_num_bytes v (
      fun ~index ~value ->
        let tree = eval_partially_1 feature_id value trees.(index) in
        trees.(index) <- tree;

        match tree with
          | `Leaf _ -> incr num_leaves
          | _ -> ()
    );
    !num_leaves

  | `RLE v ->
    let num_leaves = ref 0 in
    Rlevec.iter v (
      fun ~index ~length ~value ->
        for i = index to index + length - 1 do
          trees.(i) <- eval_partially_1 feature_id value trees.(i);
          match trees.(i) with
            | `Leaf _ -> incr num_leaves
            | _ -> ()
        done
    );
    !num_leaves


module IntSet = Set.Make( Utils.Int )

(* accumulate all the feature_id's referenced by a tree *)
let rec add_feature_id_to_set set = function
  | `Leaf _ -> set

  | `OrdinalNode { on_feature_id; on_left_tree; on_right_tree } ->
    let set = add_feature_id_to_set set on_left_tree in
    let set = add_feature_id_to_set set on_right_tree in
    IntSet.add on_feature_id set

  | `CategoricalNode { cn_feature_id; cn_left_tree; cn_right_tree } ->
    let set = add_feature_id_to_set set cn_left_tree in
    let set = add_feature_id_to_set set cn_right_tree in
    IntSet.add cn_feature_id set

let feature_id_set_of_tree tree =
  add_feature_id_to_set IntSet.empty tree


(*
  let nl tree_array =
    Array.fold_left (
      fun count tree ->
        match tree with
          | `Leaf _ -> count + 1
          | _ -> count
    ) 0 tree_array
*)

let mk_eval num_observations =
  let gamma = Array.create num_observations nan in
  let gamma_leaf = Array.create num_observations (`Leaf nan) in
  fun find_by_id tree ->
    let feature_id_set = feature_id_set_of_tree tree in
    Array.fill gamma_leaf 0 num_observations tree;

    let num_leaves = ref 0 in
    IntSet.iter (
      fun feature_id ->
        let a_feature = find_by_id feature_id in
        (*
          let i_feature =
          Feat_map.i_to_a feature_map i_feature
        in *)
        let open Dog_t in
        let cardinality = Feat_utils.cardinality_of_feature a_feature in
        let vector = Feat_utils.vector_of_feature a_feature in
        let num_leaves_new = eval_partially gamma_leaf cardinality
            feature_id vector in
        assert ( num_leaves_new >= !num_leaves );
        num_leaves := num_leaves_new

    ) feature_id_set;
    assert ( if !num_leaves <> 0 then !num_leaves = num_observations else true );
    for i = 0 to num_observations-1 do
      match gamma_leaf.(i) with
        | `Leaf g ->  gamma.(i) <- g
        | `OrdinalNode _
        | `CategoricalNode _ -> assert false
    done;
    gamma


let rec shrink alpha = function
  | `Leaf gamma ->  `Leaf (alpha *. gamma)

  | `CategoricalNode cn ->
    let cn_left_tree = shrink alpha cn.cn_left_tree in
    let cn_right_tree = shrink alpha cn.cn_right_tree in
    `CategoricalNode { cn with cn_left_tree; cn_right_tree }

  | `OrdinalNode on ->
    let on_left_tree = shrink alpha on.on_left_tree in
    let on_right_tree = shrink alpha on.on_right_tree in
    `OrdinalNode { on with on_left_tree; on_right_tree }