Debug data clean up

Author: WardBrian

src/analysis_and_optimization/Debug_data_generation.ml

@@ -85,71 +85,70 @@ let gen_num_real m t =
               (t : Expr.Typed.t Transformation.t)] in
   Random.float_range low up
 
-let rec repeat n e =
-  match n with n when n <= 0 -> [] | m -> e :: repeat (m - 1) e
+let repeat n e = List.init n ~f:(Fn.const e)
+let repeat_th n f = List.init n ~f:(fun _ -> f ())
+let random_floats n = repeat_th n (fun () -> Random.float 2.)
 
-let rec repeat_th n f =
-  match n with n when n <= 0 -> [] | m -> f () :: repeat_th (m - 1) f
+let unpack_or_repeat n (e : Expr.Typed.t) : Expr.Typed.t list =
+  match e.pattern with
+  | FunApp (CompilerInternal (FnMakeRowVec | FnMakeArray), l) -> l
+  | FunApp
+      ( StanLib ("Transpose__", FnPlain, _)
+      , [{pattern= FunApp (CompilerInternal FnMakeRowVec, l); _}] ) ->
+      l
+  | _ -> repeat n e
 
 let gen_bounded m n gen e =
-  let evaled = eval_expr m e in
-  let exprs =
-    match Expr.Helpers.try_unpack evaled with
-    | Some unpacked -> unpacked
-    | None -> List.init n ~f:(Fn.const evaled) in
+  let exprs = unpack_or_repeat n (eval_expr m e) in
   List.map ~f:gen exprs
 
 let gen_ul_bounded m n gen e1 e2 =
   let create_bounds l u =
     List.map2_exn ~f:(fun x y -> gen (Transformation.LowerUpper (x, y))) l u
   in
-  let e1 = eval_expr m e1 in
-  let es1 =
-    match Expr.Helpers.try_unpack e1 with
-    | Some unpacked -> unpacked
-    | None -> List.init n ~f:(Fn.const e1) in
-  let e2 = eval_expr m e2 in
-  let es2 =
-    match Expr.Helpers.try_unpack e2 with
-    | Some unpacked -> unpacked
-    | None -> List.init n ~f:(Fn.const e2) in
+  let es1 = unpack_or_repeat n (eval_expr m e1) in
+  let es2 = unpack_or_repeat n (eval_expr m e2) in
   create_bounds es1 es2
 
 let gen_row_vector m n t =
   match (t : Expr.Typed.t Transformation.t) with
-  | Transformation.Lower
-      ({meta= {type_= UVector | URowVector | UReal | UInt; _}; _} as e) ->
+  | Transformation.Lower e ->
       gen_bounded m n (fun x -> gen_num_real m (Lower x)) e
       |> Expr.Helpers.row_vector
-  | Upper ({meta= {type_= UVector | URowVector | UReal | UInt; _}; _} as e) ->
+  | Upper e ->
       gen_bounded m n (fun x -> gen_num_real m (Upper x)) e
       |> Expr.Helpers.row_vector
-  | LowerUpper
-      ( ({meta= {type_= UVector | URowVector | UReal | UInt; _}; _} as e1)
-      , ({meta= {type_= UVector | URowVector | UReal | UInt; _}; _} as e2) ) ->
+  | LowerUpper (e1, e2) ->
       gen_ul_bounded m n (gen_num_real m) e1 e2 |> Expr.Helpers.row_vector
   | Identity | Offset _ | Multiplier _ | OffsetMultiplier _ ->
-      Expr.Helpers.row_vector (repeat_th n (fun _ -> gen_num_real m t))
+      Expr.Helpers.row_vector (repeat_th n (fun () -> gen_num_real m t))
   | Ordered | PositiveOrdered | Simplex | UnitVector | SumToZero
    |CholeskyCorr | CholeskyCov | Correlation | Covariance | StochasticRow
-   |StochasticColumn | TupleTransformation _ | Lower _ | LowerUpper _ | Upper _
-    ->
+   |StochasticColumn | TupleTransformation _ ->
       Common.ICE.internal_compiler_error
         [%message
-          "Unknown transformation for (row) vector"
+          "Unknown transformation for row_vector"
             (t : Expr.Typed.t Transformation.t)]
 
+let sum_to_zero_floats n =
+  let l = random_floats n in
+  let sum = List.fold l ~init:0. ~f:( +. ) in
+  List.map l ~f:(fun x -> x -. (sum /. float_of_int n))
+
+let simplex_floats n =
+  let l = random_floats n in
+  let sum = List.fold l ~init:0. ~f:( +. ) in
+  List.map l ~f:(fun x -> x /. sum)
+
 let gen_vector m n t =
   let gen_ordered n =
-    let l = repeat_th n (fun _ -> Random.float 1.) in
+    let l = random_floats n in
     List.fold_map l ~init:0. ~f:(fun accum elt ->
         let elt = accum +. elt in
         (elt, elt)) in
   match t with
   | Transformation.Simplex ->
-      let l = repeat_th n (fun _ -> Random.float 1.) in
-      let sum = List.fold l ~init:0. ~f:( +. ) in
-      let l = List.map l ~f:(fun x -> x /. sum) in
+      let l = simplex_floats n in
       Expr.Helpers.vector l
   | Ordered ->
       let max, l = gen_ordered n in
@@ -159,25 +158,28 @@ let gen_vector m n t =
       let _, l = gen_ordered n in
       Expr.Helpers.vector l
   | UnitVector ->
-      let l = repeat_th n (fun _ -> Random.float 1.) in
+      let l = random_floats n in
       let sum =
         Float.sqrt
           (List.fold l ~init:0. ~f:(fun accum elt -> accum +. (elt ** 2.)))
       in
       let l = List.map l ~f:(fun x -> x /. sum) in
       Expr.Helpers.vector l
   | SumToZero ->
-      let l = repeat_th n (fun _ -> Random.float 1.) in
-      let sum = List.fold l ~init:0. ~f:( +. ) in
-      let l = List.map l ~f:(fun x -> x -. (sum /. float_of_int n)) in
+      let l = sum_to_zero_floats n in
       Expr.Helpers.vector l
-  | Identity | Offset _ | Multiplier _ | OffsetMultiplier _ | CholeskyCorr
-   |CholeskyCov | Correlation | Covariance | StochasticRow | StochasticColumn
-   |TupleTransformation _ | Lower _ | LowerUpper _ | Upper _ ->
+  | Identity | Offset _ | Multiplier _ | OffsetMultiplier _ | Lower _
+   |LowerUpper _ | Upper _ ->
       Expr.Helpers.transpose (gen_row_vector m n t)
+  | CholeskyCorr | CholeskyCov | Correlation | Covariance | StochasticRow
+   |StochasticColumn | TupleTransformation _ ->
+      Common.ICE.internal_compiler_error
+        [%message
+          "Unknown transformation for vector"
+            (t : Expr.Typed.t Transformation.t)]
 
 let gen_cov_unwrapped n =
-  let l = repeat_th (n * n) (fun _ -> Random.float 2.) in
+  let l = random_floats (n * n) in
   let l_mat = vect_to_mat l n in
   matprod l_mat (transpose l_mat)
 
@@ -195,22 +197,20 @@ let gen_diag_mat l =
 let fill_lower_triangular m =
   let fill_row i l =
     let _, tl = List.split_n l i in
-    List.init ~f:(fun _ -> Random.float 2.) i @ tl in
+    random_floats i @ tl in
   List.mapi ~f:fill_row m
 
 let pad_mat mm m n =
-  let padding_mat =
-    List.init (m - n) ~f:(fun _ -> List.init n ~f:(fun _ -> Random.float 2.))
-  in
+  let padding_mat = repeat_th (m - n) (fun () -> random_floats n) in
   wrap_real_mat (mm @ padding_mat)
 
 let gen_cov_cholesky m n =
-  let diag_mat = gen_diag_mat (List.init ~f:(fun _ -> Random.float 2.) n) in
+  let diag_mat = gen_diag_mat (random_floats n) in
   let filled_mat = fill_lower_triangular diag_mat in
   if m <= n then wrap_real_mat filled_mat else pad_mat filled_mat m n
 
 let gen_corr_cholesky_unwrapped n =
-  let diag_mat = gen_diag_mat (List.init ~f:(fun _ -> Random.float 2.) n) in
+  let diag_mat = gen_diag_mat (random_floats n) in
   let filled_mat = fill_lower_triangular diag_mat in
   let row_normalizer l =
     let row_norm =
@@ -232,11 +232,7 @@ let gen_corr_matrix n =
 let gen_sum_to_zero_matrix m n =
   (* to make each row and column sum to zero: - make each row sum to zero - add
      a new column which is - sum(rest of row) *)
-  let rows =
-    repeat_th (m - 1) (fun _ ->
-        let row = repeat_th (n - 1) (fun _ -> Random.float 2.) in
-        let row_sum = List.fold row ~init:0. ~f:( +. ) in
-        row @ [-.row_sum]) in
+  let rows = repeat_th (m - 1) (fun () -> sum_to_zero_floats n) in
   let col_sums =
     List.fold rows ~init:(repeat n 0.) ~f:(fun accum row ->
         List.map2_exn accum row ~f:( +. )) in
@@ -260,14 +256,9 @@ let gen_matrix mm m n t =
       Expr.Helpers.matrix_from_rows
         (gen_ul_bounded mm m (gen_row_vector mm n) e1 e2)
   | StochasticRow ->
-      Expr.Helpers.matrix_from_rows
-        (repeat_th m (fun () ->
-             Expr.Helpers.transpose (gen_vector mm n Simplex)))
+      Expr.Helpers.matrix (repeat_th m (fun () -> simplex_floats n))
   | StochasticColumn ->
-      Expr.Helpers.transpose
-        (Expr.Helpers.matrix_from_rows
-           (repeat_th m (fun () ->
-                Expr.Helpers.transpose (gen_vector mm n Simplex))))
+      Expr.Helpers.matrix (transpose (repeat_th n (fun () -> simplex_floats m)))
   | SumToZero -> gen_sum_to_zero_matrix m n
   | Identity | Lower _ | Upper _ | LowerUpper _ | Offset _ | Multiplier _
    |OffsetMultiplier _ ->

src/middle/Expr.ml

@@ -177,15 +177,6 @@ module Helpers = struct
         ; adlevel= TupleAD (List.map ~f:Typed.adlevel_of l) }
     ; pattern= FunApp (CompilerInternal FnMakeTuple, l) }
 
-  let try_unpack e =
-    match e.pattern with
-    | FunApp (CompilerInternal (FnMakeRowVec | FnMakeArray), l) -> Some l
-    | FunApp
-        ( StanLib ("Transpose__", FnPlain, _)
-        , [{pattern= FunApp (CompilerInternal FnMakeRowVec, l); _}] ) ->
-        Some l
-    | _ -> None
-
   let transpose e =
     let new_type =
       match Typed.type_of e with

src/middle/Expr.mli

@@ -75,7 +75,6 @@ module Helpers : sig
   val transpose : Typed.t -> Typed.t
   val array_expr : Typed.t list -> Typed.t
   val tuple_expr : Typed.t list -> Typed.t
-  val try_unpack : Typed.t -> Typed.t list option
   val loop_bottom : Typed.t
   val internal_funapp : 'a t Internal_fun.t -> 'a t list -> 'a -> 'a t
   val contains_fn_kind : ('a t Fun_kind.t -> bool) -> ?init:bool -> 'a t -> bool

test/unit/Debug_data_generation_tests.ml

@@ -213,10 +213,10 @@ let%expect_test "whole program data generation check" =
           0.090541396473787575, 0.043192465951095979, 0.12555826151802352,
           0.10919243895118017, 0.0040762414708886826, 0.015162211321058825
         ],
-        "j": [ -0.11142279798014265, 0.35444916424955808 ],
+        "j": [ -0.2228455959602853, 0.70889832849911616 ],
         "k": [
-          0.896716387037453, 1.7463321550725821, 1.8890769581764726,
-          2.6141237355067388
+          1.7934327740749061, 3.4926643101451642, 3.7781539163529452,
+          5.2282474710134776
         ],
         "l": [
           [ 0.38021657248665253, 0.0, 0.0, 0.0 ],
@@ -232,59 +232,59 @@ let%expect_test "whole program data generation check" =
           ]
         ],
         "zv": [
-          0.131724853063783, 0.056365587188466293, -0.45727981825495007,
-          0.030082007965824054, -0.081295883564592941, 0.23855136841195879,
-          0.12102997636010093, 0.31252035187713811, -0.084656401249884272,
-          -0.26704204179784391
+          0.263449706127566, 0.11273117437693259, -0.91455963650990013,
+          0.060164015931648107, -0.16259176712918588, 0.47710273682391757,
+          0.24205995272020187, 0.62504070375427623, -0.16931280249976854,
+          -0.53408408359568782
         ],
         "zm": [
           [
-            1.7903779233006683, 1.5847867503908164, 0.68134888421621276,
-            -4.0565135579076976
+            -0.09855574381348281, 0.51592882639360615, -0.087489119393748871,
+            -0.32988396318637464
           ],
           [
-            1.1194692100409906, 0.88890763701559339, 1.0216031524558131,
-            -3.0299799995123973
+            0.49638223131349624, 0.29079105840364439, -0.61264680777095926,
+            -0.17452648194618137
           ],
           [
-            0.32715908021625711, 0.35152264435373926, 1.5732645147953737,
-            -2.25194623936537
+            0.24160950850524265, 0.37430502394546239, -0.32013904829409362,
+            -0.29577548415661148
           ],
           [
-            0.59535627722869866, 0.0662964423288532, 0.4481524391087563,
-            -1.1098051586663082
+            0.90249709642995324, -0.07541114113672176, -0.60447097603656719,
+            -0.22261497925666413
           ],
           [
-            -3.8323624907866147, -2.8915134740890021, -3.7243689905761559,
-            10.448244955451774
+            -1.5419330924352093, -1.1056137676059912, 1.6247459514953688,
+            1.0228009085458316
           ]
         ],
         "srm": [
+          [
+            0.14524735713946751, 0.2676928246643322, 0.2298449764268167,
+            0.35721484176938356
+          ],
           [
             0.013902633661123317, 0.32771602782761389, 0.39502784967666338,
             0.26335348883459936
           ],
           [
             0.26002298003803503, 0.18673815713314645, 0.33907506675093968,
             0.21416379607787894
-          ],
-          [
-            0.18154796808376553, 0.6083390995215554, 0.18923430780048708,
-            0.020878624594191989
           ]
         ],
         "scm": [
           [
-            0.34768699481768295, 0.0513471326409491, 0.32947167889815893,
-            0.27149419364320893
+            0.49112547258566897, 0.43806687475397765, 0.36384673991338035,
+            0.25606541787153714
           ],
           [
-            0.23768237689080784, 0.19520318071229228, 0.27412016057501803,
-            0.29299428182188186
+            0.28461848662832567, 0.48962438588802881, 0.29982023824581161,
+            0.35958785716771219
           ],
           [
-            0.14524735713946751, 0.2676928246643322, 0.2298449764268167,
-            0.35721484176938356
+            0.22425604078600533, 0.072308739357993479, 0.3363330218408081,
+            0.38434672496075062
           ]
         ]
       } |}]