Merge pull request #1443 from stan-dev/feature/sum-to-zero

Add sum_to_zero transform
stan-dev · Aug 13, 2024 · 7865be8 · 7865be8
2 parents 0ea83bb + 13938a8
commit 7865be8
Show file tree

Hide file tree

Showing 16 changed files with 1,124 additions and 353 deletions.
diff --git a/src/analysis_and_optimization/Mir_utils.ml b/src/analysis_and_optimization/Mir_utils.ml
@@ -56,8 +56,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
+   |SumToZero
     (* This is a stub, but,
        until we define a distribution which accepts a tuple,
        this doesn't matter.

diff --git a/src/frontend/Ast_to_Mir.ml b/src/frontend/Ast_to_Mir.ml
@@ -271,8 +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 _
-   |StochasticRow | StochasticColumn ->
+   |PositiveOrdered | Simplex | UnitVector | SumToZero | Identity
+   |TupleTransformation _ | StochasticRow | StochasticColumn ->
       []
 
 let copy_indices indexed (var : Expr.Typed.t) =
@@ -295,8 +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 _
-   |StochasticRow | StochasticColumn ->
+   |PositiveOrdered | Simplex | UnitVector | SumToZero | Identity
+   |TupleTransformation _ | StochasticRow | StochasticColumn ->
       []
 
 let rec param_size transform sizedtype =
@@ -349,7 +349,7 @@ let rec param_size transform sizedtype =
         (SizedType.STuple
            (List.map subtypes_transforms ~f:(fun (st, trans) ->
                 param_size trans st)))
-  | Simplex ->
+  | Simplex | SumToZero ->
       shrink_eigen (fun d -> Expr.Helpers.(binop d Minus (int 1))) sizedtype
   | CholeskyCorr | Correlation -> shrink_eigen k_choose_2 sizedtype
   | StochasticRow -> stoch_size Fn.id min_one sizedtype

diff --git a/src/frontend/Pretty_printing.ml b/src/frontend/Pretty_printing.ml
@@ -320,9 +320,9 @@ let pp_bracketed_transform ppf = function
   | Multiplier e -> pf ppf "<@[multiplier=%a@]>" pp_expression e
   | OffsetMultiplier (e1, e2) ->
       pf ppf "<@[offset=%a,@ multiplier=%a@]>" pp_expression e1 pp_expression e2
-  | Identity | Ordered | PositiveOrdered | Simplex | UnitVector | CholeskyCorr
-   |CholeskyCov | Correlation | Covariance | TupleTransformation _
-   |StochasticColumn
+  | Identity | Ordered | PositiveOrdered | Simplex | UnitVector | SumToZero
+   |CholeskyCorr | CholeskyCov | Correlation | Covariance
+   |TupleTransformation _ | StochasticColumn
    |StochasticRow (* tuple transformations are handled in pp_transformed_type *)
     ->
       ()
@@ -360,6 +360,7 @@ let rec pp_transformed_type ppf (st, trans) =
     | PositiveOrdered -> pf ppf "positive_ordered%a" sizes_fmt ()
     | Simplex -> pf ppf "simplex%a" sizes_fmt ()
     | UnitVector -> pf ppf "unit_vector%a" sizes_fmt ()
+    | SumToZero -> pf ppf "sum_to_zero_vector%a" sizes_fmt ()
     | CholeskyCorr -> pf ppf "cholesky_factor_corr%a" cov_sizes_fmt ()
     | CholeskyCov -> pf ppf "cholesky_factor_cov%a" cov_sizes_fmt ()
     | Correlation -> pf ppf "corr_matrix%a" cov_sizes_fmt ()

diff --git a/src/frontend/Typechecker.ml b/src/frontend/Typechecker.ml
@@ -1647,6 +1647,7 @@ and check_transformation cf tenv ut trans =
   | PositiveOrdered -> PositiveOrdered
   | Simplex -> Simplex
   | UnitVector -> UnitVector
+  | SumToZero -> SumToZero
   | CholeskyCorr -> CholeskyCorr
   | CholeskyCov -> CholeskyCov
   | Correlation -> Correlation

diff --git a/src/frontend/lexer.mll b/src/frontend/lexer.mll
@@ -133,6 +133,7 @@ rule token = parse
                                 Parser.POSITIVEORDERED }
   | "simplex"                 { lexer_logger "simplex" ; Parser.SIMPLEX }
   | "unit_vector"             { lexer_logger "unit_vector" ; Parser.UNITVECTOR }
+  | "sum_to_zero_vector"      { lexer_logger "sum_to_zero_vector" ; Parser.SUMTOZERO }
   | "cholesky_factor_corr"    { lexer_logger "cholesky_factor_corr" ;
                                 Parser.CHOLESKYFACTORCORR }
   | "cholesky_factor_cov"     { lexer_logger "cholesky_factor_cov" ;

diff --git a/src/frontend/parser.messages b/src/frontend/parser.messages
diff --git a/src/frontend/parser.mly b/src/frontend/parser.mly
@@ -83,9 +83,10 @@ let nest_unsized_array basic_type n =
        ROWVECTOR "row_vector" ARRAY "array" TUPLE "tuple" MATRIX "matrix" ORDERED "ordered"
        COMPLEXVECTOR "complex_vector" COMPLEXROWVECTOR "complex_row_vector"
        POSITIVEORDERED "positive_ordered" SIMPLEX "simplex" UNITVECTOR "unit_vector"
-       CHOLESKYFACTORCORR "cholesky_factor_corr" CHOLESKYFACTORCOV "cholesky_factor_cov"
-       CORRMATRIX "corr_matrix" COVMATRIX "cov_matrix" COMPLEXMATRIX "complex_matrix"
-       STOCHASTICCOLUMNMATRIX "column_stochastic_matrix" STOCHASTICROWMATRIX "row_stochastic_matrix"
+       SUMTOZERO "sum_to_zero_vector" CHOLESKYFACTORCORR "cholesky_factor_corr"
+       CHOLESKYFACTORCOV "cholesky_factor_cov" CORRMATRIX "corr_matrix" COVMATRIX "cov_matrix"
+       COMPLEXMATRIX "complex_matrix" STOCHASTICCOLUMNMATRIX "column_stochastic_matrix"
+       STOCHASTICROWMATRIX "row_stochastic_matrix"
 %token LOWER "lower" UPPER "upper" OFFSET "offset" MULTIPLIER "multiplier"
 %token JACOBIAN "jacobian"
 %token <string> INTNUMERAL "24"
@@ -254,6 +255,7 @@ reserved_word:
   | POSITIVEORDERED { "positive_ordered", $loc, true }
   | SIMPLEX { "simplex", $loc, true }
   | UNITVECTOR { "unit_vector", $loc, true }
+  | SUMTOZERO  { "sum_to_zero_vector", $loc, true }
   | CHOLESKYFACTORCORR { "cholesky_factor_corr", $loc, true }
   | CHOLESKYFACTORCOV { "cholesky_factor_cov", $loc, true }
   | CORRMATRIX { "corr_matrix", $loc, true }
@@ -520,6 +522,8 @@ top_var_type:
     { grammar_logger "SIMPLEX_top_var_type" ; (SVector (AoS, e), Simplex) }
   | UNITVECTOR LBRACK e=expression RBRACK
     { grammar_logger "UNITVECTOR_top_var_type" ; (SVector (AoS, e), UnitVector) }
+  | SUMTOZERO LBRACK e=expression RBRACK
+    { grammar_logger "SUMTOZERO_top_var_type" ; (SVector (AoS, e), SumToZero) }
   | CHOLESKYFACTORCORR LBRACK e=expression RBRACK
     {
       grammar_logger "CHOLESKYFACTORCORR_top_var_type" ;

diff --git a/src/middle/Transformation.ml b/src/middle/Transformation.ml
@@ -16,6 +16,7 @@ type 'e t =
   | PositiveOrdered
   | Simplex
   | UnitVector
+  | SumToZero
   | CholeskyCorr
   | CholeskyCov
   | Correlation

diff --git a/src/stan_math_backend/Lower_expr.ml b/src/stan_math_backend/Lower_expr.ml
@@ -27,6 +27,7 @@ let constraint_to_string = function
   | PositiveOrdered -> Some "positive_ordered"
   | Simplex -> Some "simplex"
   | UnitVector -> Some "unit_vector"
+  | SumToZero -> Some "sum_to_zero"
   | CholeskyCorr -> Some "cholesky_factor_corr"
   | CholeskyCov -> Some "cholesky_factor_cov"
   | Correlation -> Some "corr_matrix"

diff --git a/test/integration/bad/reserved/stanc.expected b/test/integration/bad/reserved/stanc.expected
@@ -348,6 +348,16 @@ Semantic error in 'struct.stan', line 2, column 7 to column 13:
    -------------------------------------------------
 
 Identifier 'struct' clashes with reserved keyword.
+  $ ../../../../../install/default/bin/stanc sum_to_zero.stan
+Syntax error in 'sum_to_zero.stan', line 2, column 7 to column 25, parsing error:
+   -------------------------------------------------
+     1:  data {
+     2:    real sum_to_zero_vector;
+                ^
+     3:  }
+   -------------------------------------------------
+
+Expected a new identifier but found reserved keyword 'sum_to_zero_vector'.
   $ ../../../../../install/default/bin/stanc then.stan
 Semantic error in 'then.stan', line 2, column 7 to column 11:
    -------------------------------------------------

diff --git a/test/integration/bad/reserved/sum_to_zero.stan b/test/integration/bad/reserved/sum_to_zero.stan
@@ -0,0 +1,3 @@
+data {
+  real sum_to_zero_vector;
+}
diff --git a/test/integration/bad/stanc.expected b/test/integration/bad/stanc.expected
@@ -20,7 +20,7 @@ Syntax error in 'array-expr-decl-bad2.stan', line 2, column 2 to column 10, pars
 
 Invalid type in declaration. Valid types:
   int, real, vector, row_vector, matrix,
-  unit_vector, simplex, ordered, positive_ordered,
+  unit_vector, simplex, sum_to_zero_vector, ordered, positive_ordered,
   corr_matrix, cov_matrix, cholesky_factor_corr, cholesky_factor_cov,
   row_stochastic_matrix, column_stochastic_matrix, tuple(...)
 optionally preceded by a single array[...]
@@ -35,7 +35,7 @@ Syntax error in 'array-expr-decl-bad3.stan', line 2, column 2 to column 6, parsi
 
 Invalid type in declaration. Valid types:
   int, real, vector, row_vector, matrix,
-  unit_vector, simplex, ordered, positive_ordered,
+  unit_vector, simplex, sum_to_zero_vector, ordered, positive_ordered,
   corr_matrix, cov_matrix, cholesky_factor_corr, cholesky_factor_cov,
   row_stochastic_matrix, column_stochastic_matrix, tuple(...)
 optionally preceded by a single array[...]
@@ -782,7 +782,7 @@ Syntax error in 'err-decl-double.stan', line 2, column 2 to column 8, parsing er
 
 Invalid type in declaration. Valid types:
   int, real, vector, row_vector, matrix,
-  unit_vector, simplex, ordered, positive_ordered,
+  unit_vector, simplex, sum_to_zero_vector, ordered, positive_ordered,
   corr_matrix, cov_matrix, cholesky_factor_corr, cholesky_factor_cov,
   row_stochastic_matrix, column_stochastic_matrix, tuple(...)
 optionally preceded by a single array[...]