[WIP] Composable transforms

src/analysis_and_optimization/Mir_utils.ml

@@ -41,22 +41,34 @@ type bound_values =
   { lower: [`None | `Nonlit | `Lit of float]
   ; upper: [`None | `Nonlit | `Lit of float] }
 
-let trans_bounds_values (trans : Expr.Typed.t Transformation.t) : bound_values
-    =
-  let bound_value e =
-    match num_expr_value e with None -> `Nonlit | Some (f, _) -> `Lit f
+let trans_bounds_values (trans : Expr.Typed.t Transformation.t) :
+    bound_values list =
+  let single_bound (t : 'e Transformation.primitive) =
+    let bound_value e =
+      match num_expr_value e with None -> `Nonlit | Some (f, _) -> `Lit f
+    in
+    match t with
+    | Lower lower -> {lower= bound_value lower; upper= `None}
+    | Upper upper -> {lower= `None; upper= bound_value upper}
+    | LowerUpper (lower, upper) ->
+        {lower= bound_value lower; upper= bound_value upper}
+    | Simplex -> {lower= `Lit 0.; upper= `Lit 1.}
+    | PositiveOrdered -> {lower= `Lit 0.; upper= `None}
+    | UnitVector -> {lower= `Lit (-1.); upper= `Lit 1.}
+    | CholeskyCorr | CholeskyCov | Correlation | Covariance | Ordered
+     |Offset _ | Multiplier _ | OffsetMultiplier _ ->
+        {lower= `None; upper= `None}
   in
   match trans with
-  | Lower lower -> {lower= bound_value lower; upper= `None}
-  | Upper upper -> {lower= `None; upper= bound_value upper}
-  | LowerUpper (lower, upper) ->
-      {lower= bound_value lower; upper= bound_value upper}
-  | Simplex -> {lower= `Lit 0.; upper= `Lit 1.}
-  | PositiveOrdered -> {lower= `Lit 0.; upper= `None}
-  | UnitVector -> {lower= `Lit (-1.); upper= `Lit 1.}
-  | CholeskyCorr | CholeskyCov | Correlation | Covariance | Ordered
-   |Offset _ | Multiplier _ | OffsetMultiplier _ | Identity ->
-      {lower= `None; upper= `None}
+  | Identity -> [{lower= `None; upper= `None}]
+  | Single t -> [single_bound t]
+  | Chain ts -> List.map ~f:single_bound ts
+
+let trans_domain_bounds trans : bound_values =
+  (* in distribution constraints, we only care that the final constraint
+     applied satisfies
+  *)
+  List.last_exn (trans_bounds_values trans)
 
 let chop_dist_name (fname : string) : string Option.t =
   (* Slightly inefficient, would be better to short-circuit *)

src/analysis_and_optimization/Mir_utils.mli

@@ -9,7 +9,8 @@ type bound_values =
   { lower: [`None | `Nonlit | `Lit of float]
   ; upper: [`None | `Nonlit | `Lit of float] }
 
-val trans_bounds_values : Expr.Typed.t Transformation.t -> bound_values
+val trans_bounds_values : Expr.Typed.t Transformation.t -> bound_values list
+val trans_domain_bounds : Expr.Typed.t Transformation.t -> bound_values
 val chop_dist_name : string -> string Option.t
 val top_var_declarations : Stmt.Located.t -> string Set.Poly.t
 

src/analysis_and_optimization/Pedantic_analysis.ml

@@ -37,12 +37,12 @@ let list_hard_constrained (mir : Program.Typed.t) :
     | {lower= `Lit _; upper= `Lit _} -> Some `HardConstraint
     | _ -> None
   in
-  Set.Poly.filter_map
-    ~f:(fun (name, trans) ->
-      Option.map
-        ~f:(fun c -> (name, c))
-        (constrained (trans_bounds_values trans)) )
-    (parameter_set mir)
+  parameter_set mir |> Set.Poly.to_list
+  |> List.concat_map ~f:(fun (name, trans) ->
+         trans_bounds_values trans
+         |> List.filter_map ~f:constrained
+         |> List.map ~f:(fun el -> (name, el)) )
+  |> Set.Poly.of_list
 
 let list_multi_twiddles (mir : Program.Typed.t) :
     (string * Location_span.t Set.Poly.t) Set.Poly.t =

src/analysis_and_optimization/Pedantic_dist_warnings.ml

@@ -80,16 +80,21 @@ let bounds_out_of_range (range : range) (bounds : bound_values) : bool =
    constraint transformation of a variable *)
 let transform_mismatch_constraint (constr : var_constraint)
     (trans : Expr.Typed.t Transformation.t) : bool =
+  let open Transformation in
   match constr with
-  | Range range -> bounds_out_of_range range (trans_bounds_values trans)
-  | Ordered -> trans <> Transformation.Ordered
-  | PositiveOrdered -> trans <> PositiveOrdered
-  | Simplex -> trans <> Simplex
-  | UnitVector -> trans <> UnitVector
-  | CholeskyCorr -> trans <> CholeskyCorr
-  | CholeskyCov -> trans <> CholeskyCov && trans <> CholeskyCorr
-  | Correlation -> trans <> Correlation
-  | Covariance -> trans <> Covariance && trans <> Correlation
+  | Range range -> bounds_out_of_range range (trans_domain_bounds trans)
+  | Ordered -> not (domains_match trans Ordered)
+  | PositiveOrdered -> not (domains_match trans PositiveOrdered)
+  | Simplex -> not (domains_match trans Simplex)
+  | UnitVector -> not (domains_match trans UnitVector)
+  | CholeskyCorr -> not (domains_match trans CholeskyCorr)
+  | CholeskyCov ->
+      (not (domains_match trans CholeskyCov))
+      && not (domains_match trans CholeskyCorr)
+  | Correlation -> not (domains_match trans Correlation)
+  | Covariance ->
+      (not (domains_match trans Covariance))
+      && not (domains_match trans Correlation)
 
 (* Check for inconsistency between a distribution argument's range and
    a literal value *)
@@ -204,7 +209,7 @@ let uniform_dist_warning (dist_info : dist_info) :
   match dist_info with
   | {args= (Param (pname, trans), _) :: (arg1, _) :: (arg2, _) :: _; _} -> (
       let warning = Some (dist_info.loc, uniform_dist_message pname) in
-      match (arg1, arg2, trans_bounds_values trans) with
+      match (arg1, arg2, trans_domain_bounds trans) with
       | _, _, {upper= `None; _} | _, _, {lower= `None; _} ->
           (* the variate is unbounded *)
           warning

src/frontend/Ast.ml

@@ -313,15 +313,24 @@ let get_loc_dt (t : untyped_expression Type.t) =
       Some e.emeta.loc.begin_loc
 
 let get_loc_tf (t : untyped_expression Transformation.t) =
+  let get_loc_prim (t : untyped_expression Transformation.primitive) =
+    match t with
+    | Lower e
+     |Upper e
+     |LowerUpper (e, _)
+     |Offset e
+     |Multiplier e
+     |OffsetMultiplier (e, _) ->
+        Some e.emeta.loc.begin_loc
+    | _ -> None
+  in
   match t with
-  | Lower e
-   |Upper e
-   |LowerUpper (e, _)
-   |Offset e
-   |Multiplier e
-   |OffsetMultiplier (e, _) ->
-      Some e.emeta.loc.begin_loc
-  | _ -> None
+  | Identity -> None
+  | Single t -> get_loc_prim t
+  | Chain ts ->
+      let locs = List.map ~f:get_loc_prim ts in
+      (* return first location encountered *)
+      List.fold ~f:(Option.merge ~f:(fun x _ -> x)) ~init:None locs
 
 let get_first_loc (s : untyped_statement) =
   match s.stmt with

src/frontend/Ast_to_Mir.ml

@@ -198,8 +198,8 @@ type decl_context =
   {transform_action: transform_action; dadlevel: UnsizedType.autodifftype}
 
 let constraint_forl = function
-  | Transformation.Identity | Offset _ | Multiplier _ | OffsetMultiplier _
-   |Lower _ | Upper _ | LowerUpper _ ->
+  | Transformation.Offset _ | Multiplier _ | OffsetMultiplier _ | Lower _
+   |Upper _ | LowerUpper _ ->
       Stmt.Helpers.for_scalar
   | Ordered | PositiveOrdered | Simplex | UnitVector | CholeskyCorr
    |CholeskyCov | Correlation | Covariance ->
@@ -216,7 +216,7 @@ let same_shape decl_id decl_var id var meta =
                   [str "constraint"; str decl_id; decl_var; str id; var] )
         ; meta } ]
 
-let check_transform_shape decl_id decl_var meta = function
+let check_transform_shape_prim decl_id decl_var meta = function
   | Transformation.Offset e -> same_shape decl_id decl_var "offset" e meta
   | Multiplier e -> same_shape decl_id decl_var "multiplier" e meta
   | Lower e -> same_shape decl_id decl_var "lower" e meta
@@ -228,9 +228,15 @@ let check_transform_shape decl_id decl_var meta = function
       same_shape decl_id decl_var "lower" e1 meta
       @ same_shape decl_id decl_var "upper" e2 meta
   | Covariance | Correlation | CholeskyCov | CholeskyCorr | Ordered
-   |PositiveOrdered | Simplex | UnitVector | Identity ->
+   |PositiveOrdered | Simplex | UnitVector ->
       []
 
+let check_transform_shape decl_id decl_var meta = function
+  | Transformation.Identity -> []
+  | Single t -> check_transform_shape_prim decl_id decl_var meta t
+  | Chain ts ->
+      List.concat_map ~f:(check_transform_shape_prim decl_id decl_var meta) ts
+
 let copy_indices indexed (var : Expr.Typed.t) =
   if UnsizedType.is_scalar_type var.meta.type_ then var
   else
@@ -253,7 +259,7 @@ let extract_transform_args var = function
   | LowerUpper (a1, a2) | OffsetMultiplier (a1, a2) ->
       [copy_indices var a1; copy_indices var a2]
   | Covariance | Correlation | CholeskyCov | CholeskyCorr | Ordered
-   |PositiveOrdered | Simplex | UnitVector | Identity ->
+   |PositiveOrdered | Simplex | UnitVector ->
       []
 
 let param_size transform sizedtype =
@@ -278,30 +284,48 @@ let param_size transform sizedtype =
   let k_choose_2 k =
     Expr.Helpers.(binop (binop k Times (binop k Minus (int 1))) Divide (int 2))
   in
-  match transform with
-  | Transformation.Identity | Lower _ | Upper _
-   |LowerUpper (_, _)
-   |Offset _ | Multiplier _
-   |OffsetMultiplier (_, _)
-   |Ordered | PositiveOrdered | UnitVector ->
-      sizedtype
-  | Simplex ->
-      shrink_eigen (fun d -> Expr.Helpers.(binop d Minus (int 1))) sizedtype
-  | CholeskyCorr | Correlation -> shrink_eigen k_choose_2 sizedtype
-  | CholeskyCov ->
-      (* (N * (N + 1)) / 2 + (M - N) * N *)
-      shrink_eigen_mat
-        (fun m n ->
-          Expr.Helpers.(
-            binop
-              (binop (k_choose_2 n) Plus n)
-              Plus
-              (binop (binop m Minus n) Times n)) )
-        sizedtype
-  | Covariance ->
-      shrink_eigen
-        (fun k -> Expr.Helpers.(binop k Plus (k_choose_2 k)))
+  let resize sizedtype trans =
+    match trans with
+    | Transformation.Lower _ | Upper _
+     |LowerUpper (_, _)
+     |Offset _ | Multiplier _
+     |OffsetMultiplier (_, _)
+     |Ordered | PositiveOrdered | UnitVector ->
         sizedtype
+    | Simplex ->
+        shrink_eigen (fun d -> Expr.Helpers.(binop d Minus (int 1))) sizedtype
+    | CholeskyCorr | Correlation -> shrink_eigen k_choose_2 sizedtype
+    | CholeskyCov ->
+        (* (N * (N + 1)) / 2 + (M - N) * N *)
+        shrink_eigen_mat
+          (fun m n ->
+            Expr.Helpers.(
+              binop
+                (binop (k_choose_2 n) Plus n)
+                Plus
+                (binop (binop m Minus n) Times n)) )
+          sizedtype
+    | Covariance ->
+        shrink_eigen
+          (fun k -> Expr.Helpers.(binop k Plus (k_choose_2 k)))
+          sizedtype
+  in
+  match transform with
+  | Transformation.Identity -> sizedtype
+  | Single t -> resize sizedtype t
+  | Chain ts ->
+      (* Finds the first constraint which changes sizes and goes with that
+       * Similar to Transform_Mir.read_constrain_dims, this won't work in all cases,
+       * like if you ever did a simplex corr_matrix. In our standard library
+       * This is fine, may need more complicated logic for user defined transforms.
+       *)
+      let rec loop st = function
+        | [] -> st
+        | t :: ts ->
+            let st' = resize st t in
+            if st' = st then loop st ts else st'
+      in
+      loop sizedtype ts
 
 let remove_possibly_exn pst action loc =
   match pst with
@@ -313,20 +337,36 @@ let remove_possibly_exn pst action loc =
 
 let rec check_decl var decl_type' decl_id decl_trans smeta adlevel =
   let decl_type = remove_possibly_exn decl_type' "check" smeta in
+  let check_single trans =
+    match trans with
+    | Transformation.LowerUpper (lb, ub) ->
+        check_decl var decl_type' decl_id (Transformation.Single (Lower lb))
+          smeta adlevel
+        @ check_decl var decl_type' decl_id (Single (Upper ub)) smeta adlevel
+    | _ when Transformation.primitive_has_check trans ->
+        let check_id id =
+          let var_name = Fmt.strf "%a" Expr.Typed.pp id in
+          let args = extract_transform_args id trans in
+          Stmt.Helpers.internal_nrfunapp
+            (FnCheck {trans; var_name; var= id})
+            args smeta
+        in
+        [(constraint_forl trans) decl_type check_id var smeta]
+    | _ -> []
+  in
   match decl_trans with
-  | Transformation.LowerUpper (lb, ub) ->
-      check_decl var decl_type' decl_id (Lower lb) smeta adlevel
-      @ check_decl var decl_type' decl_id (Upper ub) smeta adlevel
-  | _ when Transformation.has_check decl_trans ->
-      let check_id id =
-        let var_name = Fmt.strf "%a" Expr.Typed.pp id in
-        let args = extract_transform_args id decl_trans in
-        Stmt.Helpers.internal_nrfunapp
-          (FnCheck {trans= decl_trans; var_name; var= id})
-          args smeta
-      in
-      [(constraint_forl decl_trans) decl_type check_id var smeta]
-  | _ -> []
+  | Transformation.Identity -> []
+  | Single t -> check_single t
+  (* NB: We only allow chain transforms in parameters, which are never checked. 
+   * REM: A naive attempt at doing this would be `List.concat_map ~f:check_single ts`
+   * This currently concatinates all checks in a way that is invalid, 
+   * e.g. lower(0), upper(10), will not yield values in [0,10],
+   * one must perform unconstraining transform between sucessive checks in general
+   *)
+  | Chain _ ->
+      raise_s
+        [%message
+          "Attempting to check a chained transform. This should never happen"]
 
 let check_sizedtype name =
   let check x = function
@@ -362,8 +402,8 @@ let check_sizedtype name =
       (ll, Type.Sized st)
   | Unsized ut -> ([], Unsized ut)
 
-let trans_decl {transform_action; dadlevel} smeta decl_type transform
-    identifier initial_value =
+let trans_decl {transform_action; dadlevel} smeta decl_type
+    (transform : Expr.Typed.t Transformation.t) identifier initial_value =
   let decl_id = identifier.Ast.name in
   let rhs = Option.map ~f:trans_expr initial_value in
   let size_checks, dt = check_sizedtype identifier.name decl_type in
@@ -641,9 +681,9 @@ let trans_sizedtype_decl declc tr name =
     | SVector (mem_pattern, s) ->
         let fn =
           match (declc.transform_action, tr) with
-          | Constrain, Transformation.Simplex ->
+          | Constrain, Transformation.Single Simplex ->
               Internal_fun.FnValidateSizeSimplex
-          | Constrain, UnitVector -> FnValidateSizeUnitVector
+          | Constrain, Single UnitVector -> FnValidateSizeUnitVector
           | _ -> FnValidateSize
         in
         let l, s = grab_size fn n s in
@@ -656,7 +696,7 @@ let trans_sizedtype_decl declc tr name =
         let l2, c = grab_size FnValidateSize (n + 1) c in
         let cf_cov =
           match (declc.transform_action, tr) with
-          | Constrain, CholeskyCov ->
+          | Constrain, Single CholeskyCov ->
               [ { Stmt.Fixed.pattern=
                     NRFunApp
                       ( StanLib ("check_greater_or_equal", FnPlain, AoS)