Source file owl_countmin_table.ml

# 1 "src/base/misc/owl_countmin_table.ml"
(* A module type for the tables used to implement the CountMin sketch. *)
module type Sig = sig
  (** {5 Type definition} *)

  type t
  (** The type of count-min tables *)

  (** {5 Core functions} *)

  val init : int -> int -> t
  (** 
  [init l w] generates a table with length [l] and width [w],
  all counters initialized to 0.
  *)

  val incr : int -> int -> t -> unit
  (**
  [incr i j t] increments the counter at length index [i] 
  and width index [j] in table [t].
  *)

  val get : int -> int -> t -> int
  (**
  [get i j t] gets the value of the counter at length index [i] 
  and width index [j] in table [t].
  *)

  val clone : t -> t
  (**
  [clone t] returns a new table with the same contents as [t].
  *)

  val merge : t -> t -> t
  (**
  [merge t1 t2] merges tables [t1] and [t2] element-wise.  If [t1] and [t2]
  have the same dimensions, returns a new table whose elements are the sums of corresponding
  elements from [t1] and [t2].  If dimensions do not match, raises [INVALID_ARGUMENT].
  *)
end

(* Implementation of the CountMin sketch table using OCaml native arrays *)
module Native : Sig = struct
  type t = int array array

  let init l w = Array.make_matrix l w 0

  let incr i j t = t.(i).(j) <- t.(i).(j) + 1

  let get i j t = t.(i).(j)

  let clone = Array.map (Array.map (fun x -> x))

  let merge t1 t2 =
    try Array.map2 (Array.map2 ( + )) t1 t2 with
    | Invalid_argument _ ->
      raise
        (Owl_exception.INVALID_ARGUMENT
           "Cannot merge countmin tables of different dimensions")
end

(* Implementation of the CountMin sketch table using Owl ndarrays *)
module Owl : Sig = struct
  type t = Owl_base_dense_ndarray_s.arr

  let init l w = Owl_base_dense_ndarray_s.zeros [| l; w |]

  let incr i j t = Owl_base_dense_ndarray_s.(set t [| i; j |] (get t [| i; j |] +. 1.))

  let get i j t = Owl_base_dense_ndarray_s.get t [| i; j |] |> int_of_float

  let clone t = Owl_base_dense_ndarray_s.copy t

  let merge t1 t2 =
    let open Owl_base_dense_ndarray_s in
    if shape t1 = shape t2
    then add t1 t2
    else
      raise
        (Owl_exception.INVALID_ARGUMENT
           "Cannot merge countmin tables of different dimensions")
end