cpt.ml 1.92 KB
Newer Older
Philippe Veber's avatar
Philippe Veber committed
1 2 3 4
open Rresult
open Core_kernel
open Let_syntax.Result

5 6
(* FIXME: rewrite based on Dataframe? *)

Philippe Veber's avatar
Philippe Veber committed
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27
type t = {
  nrows : int ;
  cols : (string * float option array) list ;
}

let of_file fn =
  let* df = Dataframe.from_file fn in
  let+ cols =
    Dataframe.columns df
    |> List.map ~f:(fun (label, col) ->
        let+ data = match (col : Dataframe.column) with
          | Floats xs -> Ok (Array.map xs ~f:Option.some)
          | Float_opts xs -> Ok xs
          | Ints xs -> Ok (Array.map xs ~f:(fun i -> Some (Float.of_int i)))
          | Int_opts xs -> Ok (Array.map xs ~f:(Option.map ~f:Float.of_int))
          | _ -> R.error_msgf "column %s is not numeric" label
        in
        label, data
      )
    |> Result.all
  in
Philippe Veber's avatar
Philippe Veber committed
28
  let cols = List.filter cols ~f:(fun (label, _) -> not (String.equal label "Sites")) in
Philippe Veber's avatar
Philippe Veber committed
29 30 31 32 33 34
  let nrows = match cols with
    | [] -> assert false
    | (_, h) :: _ -> Array.length h
  in
  { nrows ; cols }

35 36
let of_file_exn fn = of_file fn |> Rresult.R.failwith_error_msg

Philippe Veber's avatar
Philippe Veber committed
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
let columns x = x.cols

let make = function
  | [] -> Error (`Msg "empty column list")
  | (_, h) :: t as cols ->
    let nrows = Array.length h in
    if List.for_all t ~f:(fun (_, col) -> Array.length col = nrows) then
      Ok { nrows ; cols }
    else R.error_msg "Columns with different lengths"

let write_row oc fields =
  String.concat ~sep:"\t" fields
  |> Out_channel.fprintf oc "%s\n"

let write_columns oc res =
52
  Range.iter 0 res.nrows ~f:(fun i ->
Philippe Veber's avatar
Philippe Veber committed
53 54 55 56 57 58 59 60 61
      Int.to_string i
      :: List.map res.cols ~f:(fun (_, col) ->
          Option.value_map ~default:"NA" ~f:Float.to_string col.(i)
        )
      |> write_row oc
    )

let to_file table ~output =
  Out_channel.with_file output ~f:(fun oc ->
Philippe Veber's avatar
Philippe Veber committed
62
      write_row oc ("Sites" :: List.map table.cols ~f:fst) ;
Philippe Veber's avatar
Philippe Veber committed
63 64
      write_columns oc table
    )
65 66 67 68 69 70

let get_col cpt colname = 
  List.Assoc.find ~equal:String.equal cpt.cols colname

let get_col_exn cpt colname = 
  get_col cpt colname |> Option.value_exn