Revert "Revert "adds row and column stochastic matrices""

This reverts commit 126d326.
stan-dev · Jul 16, 2024 · 03c4492 · 03c4492
1 parent 126d326
commit 03c4492
Show file tree

Hide file tree

Showing 13 changed files with 690 additions and 334 deletions.
diff --git a/scripts/hooks/install_hooks.sh b/scripts/hooks/install_hooks.sh
diff --git a/src/analysis_and_optimization/Mir_utils.ml b/src/analysis_and_optimization/Mir_utils.ml
@@ -55,8 +55,8 @@ let trans_bounds_values (trans : Expr.Typed.t Transformation.t) : bound_values =
   | PositiveOrdered -> {lower= `Lit 0.; upper= `None}
   | UnitVector -> {lower= `Lit (-1.); upper= `Lit 1.}
   | CholeskyCorr | CholeskyCov | Correlation | Covariance | Ordered | Offset _
-   |Multiplier _ | OffsetMultiplier _
-   |Identity
+   |Multiplier _ | OffsetMultiplier _ | Identity | RowStochastic
+   |ColumnStochastic
     (* This is a stub, but,
        until we define a distribution which accepts a tuple,
        this doesn't matter.

diff --git a/src/analysis_and_optimization/Pedantic_dist_warnings.ml b/src/analysis_and_optimization/Pedantic_dist_warnings.ml
@@ -36,6 +36,8 @@ type var_constraint =
   | CholeskyCov
   | Correlation
   | Covariance
+  | RowStochastic
+  | ColumnStochastic
 
 (** Constraint paired with a name for user messages *)
 type var_constraint_named = {name: string; constr: var_constraint}
@@ -90,6 +92,8 @@ let transform_mismatch_constraint (constr : var_constraint)
   | CholeskyCov -> trans <> CholeskyCov && trans <> CholeskyCorr
   | Correlation -> trans <> Correlation
   | Covariance -> trans <> Covariance && trans <> Correlation
+  | RowStochastic -> trans <> RowStochastic
+  | ColumnStochastic -> trans <> ColumnStochastic
 
 (** Check for inconsistency between a distribution argument's range and
    a literal value *)

diff --git a/src/frontend/Ast_to_Mir.ml b/src/frontend/Ast_to_Mir.ml
@@ -271,7 +271,8 @@ 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 | TupleTransformation _ ->
+   |PositiveOrdered | Simplex | UnitVector | Identity | TupleTransformation _
+   |RowStochastic | ColumnStochastic ->
       []
 
 let copy_indices indexed (var : Expr.Typed.t) =
@@ -294,7 +295,8 @@ 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 | TupleTransformation _ ->
+   |PositiveOrdered | Simplex | UnitVector | Identity | TupleTransformation _
+   |RowStochastic | ColumnStochastic ->
       []
 
 let rec param_size transform sizedtype =
@@ -326,7 +328,7 @@ let rec param_size transform sizedtype =
    |LowerUpper (_, _)
    |Offset _ | Multiplier _
    |OffsetMultiplier (_, _)
-   |Ordered | PositiveOrdered | UnitVector ->
+   |Ordered | PositiveOrdered | UnitVector | RowStochastic | ColumnStochastic ->
       sizedtype
   | TupleTransformation tms ->
       let _, dims = SizedType.get_array_dims sizedtype in

diff --git a/src/frontend/Pretty_printing.ml b/src/frontend/Pretty_printing.ml
@@ -322,7 +322,9 @@ let pp_bracketed_transform ppf = function
       pf ppf "<@[offset=%a,@ multiplier=%a@]>" pp_expression e1 pp_expression e2
   | Identity | Ordered | PositiveOrdered | Simplex | UnitVector | CholeskyCorr
    |CholeskyCov | Correlation | Covariance | TupleTransformation _
-  (* tuple transformations are handled in pp_transformed_type *) ->
+   |RowStochastic
+   |ColumnStochastic
+    (* tuple transformations are handled in pp_transformed_type *) ->
       ()
 
 let rec pp_transformed_type ppf (st, trans) =
@@ -362,6 +364,8 @@ let rec pp_transformed_type ppf (st, trans) =
     | CholeskyCov -> pf ppf "cholesky_factor_cov%a" cov_sizes_fmt ()
     | Correlation -> pf ppf "corr_matrix%a" cov_sizes_fmt ()
     | Covariance -> pf ppf "cov_matrix%a" cov_sizes_fmt ()
+    | RowStochastic -> pf ppf "row_stochastic%a" sizes_fmt ()
+    | ColumnStochastic -> pf ppf "column_stochastic%a" sizes_fmt ()
     | TupleTransformation transforms ->
         (* NB this calls the top-level function to handle internal arrays etc *)
         let transTypes = Middle.Utils.zip_stuple_trans_exn st transforms in

diff --git a/src/frontend/Typechecker.ml b/src/frontend/Typechecker.ml
@@ -1657,6 +1657,8 @@ and check_transformation cf tenv ut trans =
         List.map typesTrans ~f:(fun (ut, tm) ->
             check_transformation cf tenv ut tm) in
       TupleTransformation tes
+  | RowStochastic -> RowStochastic
+  | ColumnStochastic -> ColumnStochastic
 
 and check_var_decl loc cf tenv sized_ty trans
     (variables : untyped_expression Ast.variable list) is_global =

diff --git a/src/frontend/lexer.mll b/src/frontend/lexer.mll
@@ -139,6 +139,8 @@ rule token = parse
                                 Parser.CHOLESKYFACTORCOV }
   | "corr_matrix"             { lexer_logger "corr_matrix" ; Parser.CORRMATRIX }
   | "cov_matrix"              { lexer_logger "cov_matrix" ; Parser.COVMATRIX }
+  | "row_stochastic"          { lexer_logger "row_stochastic" ; Parser.ROWSTOCHASTIC }
+  | "column_stochastic"       { lexer_logger "column_stochastic" ; Parser.COLUMNSTOCHASTIC }
 (* Transformation keywords *)
   | "lower"                   { lexer_logger "lower" ; Parser.LOWER }
   | "upper"                   { lexer_logger "upper" ; Parser.UPPER }