Source file utils.ml

(*********************************************************************************)
(*                Statocaml                                                      *)
(*                                                                               *)
(*    Copyright (C) 2025 INRIA All rights reserved.                              *)
(*    Author: Maxence Guesdon (INRIA Saclay)                                     *)
(*      with Gabriel Scherer (INRIA Paris) and Florian Angeletti (INRIA Paris)   *)
(*                                                                               *)
(*    This program is free software; you can redistribute it and/or modify       *)
(*    it under the terms of the GNU General Public License as                    *)
(*    published by the Free Software Foundation, version 3 of the License.       *)
(*                                                                               *)
(*    This program is distributed in the hope that it will be useful,            *)
(*    but WITHOUT ANY WARRANTY; without even the implied warranty of             *)
(*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the               *)
(*    GNU General Public License for more details.                               *)
(*                                                                               *)
(*    You should have received a copy of the GNU General Public                  *)
(*    License along with this program; if not, write to the Free Software        *)
(*    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA                   *)
(*    02111-1307  USA                                                            *)
(*                                                                               *)
(*    As a special exception, you have permission to link this program           *)
(*    with the OCaml compiler and distribute executables, as long as you         *)
(*    follow the requirements of the GNU GPL in regard to all of the             *)
(*    software in the executable aside from the OCaml compiler.                  *)
(*                                                                               *)
(*    Contact: Maxence.Guesdon@inria.fr                                          *)
(*                                                                               *)
(*********************************************************************************)

(** *)

module S = Statocaml
module Log = Statocaml.Log
module Xr = Xtmpl.Rewrite


let gen_file ~file contents =
  Lwt_io.(with_file ~mode:Output file (fun oc -> write oc contents))

let gen_xml_file file xml =
  let str = Xr.to_string xml in
  gen_file ~file str

let cdata fmt = Printf.ksprintf Xr.cdata fmt

let lwt_temp_file ?prefix ?suffix () =
  let%lwt (f,oc) = Lwt_io.open_temp_file ?prefix ?suffix () in
  let%lwt () = Lwt_io.close oc in
  Lwt.return f

[@@@warning "-27"]
let set_gnuplot_svg_title xmls title =
  let f data _env ?loc _atts subs =
    match subs with
    | [Xr.D { text = "Gnuplot" }] ->
        data, [Xr.node ("","title") [cdata "%s" title]]
    | _ -> raise Xr.No_change
  in
  let env = Xr.(env_add_cb "title" f (Xr.env_empty())) in
  let (_,xmls) = Xr.apply_to_xmls () env xmls  in
  xmls
[@@@warning "+27"]

let svg_load_and_delete ?gnuplot_title svg_file =
  let%lwt xml = Lwt_io.(with_file ~mode:Input svg_file read) in
  try
    let doc = Xtmpl.Rewrite.doc_from_string xml in
    let%lwt () = try%lwt Lwt_unix.unlink svg_file with _ -> Lwt.return_unit in
    let%lwt xmls = Lwt.return doc.elements in
    let xmls = match gnuplot_title with
      | None -> xmls
      | Some t -> set_gnuplot_svg_title xmls t
    in
    Lwt.return xmls
  with e ->
      Log.err (fun m -> m "XMl code (in %s): %s" svg_file xml);
      raise e

let plot_to_svg ?title ?(w=300) ?(h=300) f =
  let%lwt svg_file = lwt_temp_file ~suffix:".svg" () in
  let gp = Plot.create ~terminal:"svg" ?title ~w ~h svg_file in
  let%lwt () = f gp in
  svg_load_and_delete svg_file

let bar_chart gp ?(with_titles=false) ?(legend_pos="right center outside") bars =
  let p fmt = Plot.p gp fmt in
  p "set style histogram clustered gap 2" ;
  if with_titles then p "set key %s" legend_pos else p "unset key";
  let (bars, max_v) =
    List.fold_right (fun (v,color,title) (acc,max_v) ->
       let id = List.length acc + 1 in
       Plot.define_float_data gp ("data" ^ string_of_int id) [ v ];
       (id,color,title) :: acc, max max_v v
    ) bars ([], 0.)
  in
  p "set xtics (' ' 1)";
  p "set yrange [0:%d]" (truncate max_v + 1);
  let hists = List.map
    (fun (id, color, title) ->
       Printf.sprintf "$data%d using 1 with histogram fs solid lc %S %s"
         id color (Printf.sprintf "title %S" title)
    )
      bars
  in
  p "plot %s" (String.concat ", " hists);
  Plot.run gp

let bar_chart_svg ?title ?w ?h ?with_titles ?legend_pos bars =
  let f gp = bar_chart gp ?with_titles ?legend_pos bars in
  plot_to_svg ?title ?w ?h f

let year_lines_chart gp ?(with_titles=false) ?(legend_pos="right center outside") lines =
  let p fmt = Plot.p gp fmt in
  if with_titles then p "set key %s" legend_pos else p "unset key";
  let (lines, max_v, min_x, max_x) =
    List.fold_right (fun (year_values,color,title) (acc,max_v,min_x,max_x) ->
       let id = List.length acc + 1 in
       let values = List.map snd year_values in
       let years = List.map fst year_values in
       let year_values = List.map (fun (y,v) -> float y, v) year_values in
       Plot.define_float_float_data gp ("data" ^ string_of_int id) year_values ;
       (id,color,title) :: acc,
       (List.fold_right max values max_v),
       (List.fold_right min years min_x),
       (List.fold_right max years max_x)
    ) lines ([], 0., max_int, 0)
  in
  p "set xrange [%d:%d]" min_x max_x ;
  p "set yrange [0:%d]" (truncate max_v + 1);
  let plines = List.map
    (fun (id, color, title) ->
       Printf.sprintf "$data%d with linespoints lc %S %s"
         id color (Printf.sprintf "title %S" title)
    )
      lines
  in
  p "plot %s" (String.concat ", " plines);
  Log.info (fun m -> m "gnuplot code =\n%s" (Plot.code gp));
  Plot.run gp

let year_lines_chart_svg ?title ?(w=600) ?(h=300) ?with_titles ?legend_pos lines =
  let f gp = year_lines_chart gp ?with_titles ?legend_pos lines in
  plot_to_svg ?title ~w ~h f


let add_events_by_month gp min_month max_month events =
  let module T = Statocaml.Types in
  let add m (ev:Statocaml.Types.event) map =
    let ev =
      match Plot.Mmap.find_opt m map with
    | None | Some None -> ev
    | Some (Some ev0) ->
        { ev0 with ev_label = ev0.Statocaml.Types.ev_label ^ " / " ^ ev.ev_label }
    in
    Plot.Mmap.add m (Some ev) map
  in
  let f ev acc =
    let (y,m,_d) = ev.T.ev_date in
    let month = (y,m) in
    if month >= min_month && month <= max_month then
      add month ev acc
    else
      acc
  in
  let map = List.fold_right f events Plot.Mmap.empty in
  let map = Plot.add_missing_months ~default:None min_month max_month map in
  let elts = Plot.Mmap.fold (fun _mon elt acc -> elt :: acc) map [] in
  let elts = List.rev elts in
  let draw_event i = function
  | None -> ()
  | Some ev ->
      let (_,m,d) = ev.T.ev_date in
      let color = Printf.sprintf "#%02x%02x%02x" (Random.int 256) (m * 12) (d * 8) in
      (*Log.info (fun p -> p "[%d] adding event %S (%d/%d/%d) (color=%S)" i T.ev_label d m y color);*)
      Plot.p gp
        "set arrow from %d, graph 0 to %d, graph 1 nohead lw 1 lc rgb %S" i i color ;
      Plot.p gp
        "set label %S at %d, graph 1 rotate by 90 right offset 1,0" ev.T.ev_label i
  in
  List.iteri draw_event elts

let add_events_for_periods gp ?font_size events periods =
  let module T = Statocaml.Types in
  let module P = Statocaml.Period in
  let concat_events = function
  | [] -> None
  | [ev] -> Some ev
  | (ev :: _) as l ->
      let ev_label = String.concat " / "
        (List.map (fun ev -> ev.T.ev_label) l)
      in
      Some { ev with ev_label }
  in
  let f period _ =
    concat_events (List.filter (fun ev -> P.date_in_period period ev.T.ev_date) events)
  in
  let map = P.Map.mapi f periods in
  let elts = P.Map.fold (fun _mon ev acc -> ev :: acc) map [] in
  let elts = List.rev elts in
  let font = match font_size with None -> "" | Some n -> Printf.sprintf " font \",%d\"" n in
  let draw_event i = function
  | None -> ()
  | Some ev ->
      let (_,m,d) = ev.T.ev_date in
      let color = Printf.sprintf "#%02x%02x%02x" (Random.int 256) (m * 12) (d * 8) in
      (*Log.info (fun p -> p "[%d] adding event %S (%d/%d/%d) (color=%S)" i T.ev_label d m y color);*)
      Plot.p gp
        "set arrow from %d, graph 0 to %d, graph 1 nohead lw 1 lc rgb %S" i i color ;
      Plot.p gp
        "set label %S%s at %d, graph 1 rotate by 90 right offset 1,0" ev.T.ev_label font i
  in
  List.iteri draw_event elts

let mk_year_steps ?(div12=false) by_month =
  let ymap = Plot.Mmap.fold
    (fun (y,_) n acc ->
       match S.Imap.find_opt y acc with
       | None -> S.Imap.add y n acc
       | Some yn -> S.Imap.add y (yn+n) acc)
      by_month S.Imap.empty
  in
  let map = Plot.Mmap.mapi (fun (y,_) _ -> float (S.Imap.find y ymap)) by_month in
  let f = if div12 then (fun (_,v) -> v /. 12.) else snd in
  List.map f (Plot.Mmap.bindings map)

let mk_sliding_sums ?(div=false) ?(window_size=12) by_month =
  let rec iter acc tmp sum = function
  | [] ->
      (match tmp with
       | [] -> acc
       | l -> (List.map (fun ym -> (ym,sum)) l) @ acc
      )
  | ((y,m), v) :: q ->
      let tmp = (y,m) :: tmp in
      let sum = sum + v in
      if List.length tmp < window_size then
        iter acc tmp sum q
      else
        let acc = (List.map (fun ym -> (ym,sum)) tmp) @ acc in
        iter acc [] 0 q
  in
  let l = iter [] [] 0 (List.rev (Plot.Mmap.bindings by_month)) in
  let f = if div then (fun (_,v) -> float v /. (float window_size)) else (fun (_,v) -> float v) in
  List.map f l

let float_mean = S.Misc.float_mean

let float_median ?(sorted=false) l =
  match List.length l with
  | 0 -> None
  | len ->
      let l = if sorted then l else List.sort Float.compare l in
      let half = len / 2 in
      let even = len mod 2 = 0 in
      let rec iter p = function
      | [] ->
          Log.err (fun m -> m "float_median: values=%s"
             (String.concat ", " (List.map string_of_float l)));
          assert false
      | v :: _ when not even && p = half + 1 -> Some v
      | v1 :: v2 :: _ when even && p = half -> Some ((v1 +. v2) /. 2.)
      | _ :: q -> iter (p+1) q
      in
      iter 1 l

let sd values avg =
  let sum = List.fold_left
    (fun acc v -> acc +. ((v -. avg) ** 2.))
      0. values
  in
  sqrt (sum /. (float (List.length values)))

let classes_of_data l =
  (* using formula from https://gilles-hunault.leria-info.univ-angers.fr/wstat/discr.php *)
  let l = List.sort Float.compare l in
  match float_mean l with
  | None -> ()
  | Some mean ->
      let a, b = match l, List.rev l with
        | [], _ | _, [] -> assert false
        | a::_, b::_ -> a, b
      in
      let sd = sd l mean in
      let n = float (List.length l) in
      let scott = (b -. a) /. (3.5 *. sd *. (Float.pow n (-1. /. 3.))) in
      let sturges = Float.log2 (n +. 1.) in
      prerr_endline (Printf.sprintf "Scott: mean=%.1f, sd=%.2f, nb_classes=%f" mean sd scott);
      prerr_endline (Printf.sprintf "Sturges: nb_classes=%f" sturges)

let to_quantiles ?(quantiles=10) values =
  let len = float (List.length values) in
  let quantiles = float quantiles in
  let values = List.sort Float.compare values in
  let rec iter limit i acc accp = function
  | [] -> List.rev ((List.rev accp) :: acc)
  | l when float i > limit ->
        iter (limit +. (len /. quantiles)) i ((List.rev accp)::acc) [] l
  | h :: q ->
      iter limit (i+1) acc (h::accp) q
  in
  let l = iter (len /. quantiles) 0 [] [] values in
  assert (List.length l = truncate quantiles);
  List.map (fun x -> match float_mean x with None -> assert false | Some v -> v) l