Skip to content

Commit 8d0c433

Browse files
WardBrianWardBrian
committed
Add transformation information to --info
1 parent 5692dc4 commit 8d0c433

File tree

3 files changed

+112
-33
lines changed

3 files changed

+112
-33
lines changed

‎src/frontend/Info.ml‎

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ open Middle
44
open Yojson.Basic
55

66
let rec unsized_basetype_json t =
7-
let to_json (type_, dim) : t =
8-
`Assoc [("type", `String type_); ("dimensions", `Int dim)] in
7+
let to_json (type_, dim) =
8+
[("type", `String type_); ("dimensions", `Int dim)] in
99
let internal, dims = UnsizedType.unwind_array_type t in
1010
match internal with
1111
| UnsizedType.UInt -> to_json ("int", dims)
@@ -16,23 +16,56 @@ let rec unsized_basetype_json t =
1616
| UMatrix -> to_json ("real", dims + 2)
1717
| UComplexMatrix -> to_json ("complex", dims + 2)
1818
| UTuple internals ->
19-
`Assoc
20-
[ ("type", `List (List.map ~f:unsized_basetype_json internals))
21-
; ("dimensions", `Int dims) ]
19+
[ ( "type"
20+
, `List
21+
(List.map ~f:(fun t -> `Assoc (unsized_basetype_json t)) internals)
22+
); ("dimensions", `Int dims) ]
2223
| UMathLibraryFunction | UFun _ | UArray _ -> assert false
2324

2425
let basetype_dims t = SizedType.to_unsized t |> unsized_basetype_json
2526

27+
let rec transformation t =
28+
let expr_string = Fmt.to_to_string Pretty_printing.pp_expression in
29+
let expr_string e =
30+
`String (expr_string (Ast.untyped_expression_of_typed_expression e)) in
31+
let transform details = [("constraint", details)] in
32+
match t with
33+
| Transformation.Identity -> transform (`String "none")
34+
| Lower e -> transform @@ `Assoc [("lower", expr_string e)]
35+
| Upper e -> transform @@ `Assoc [("upper", expr_string e)]
36+
| LowerUpper (e1, e2) ->
37+
transform @@ `Assoc [("lower", expr_string e1); ("upper", expr_string e2)]
38+
| Offset e -> transform @@ `Assoc [("offset", expr_string e)]
39+
| Multiplier e -> transform @@ `Assoc [("multiplier", expr_string e)]
40+
| OffsetMultiplier (e1, e2) ->
41+
transform
42+
@@ `Assoc [("offset", expr_string e1); ("multiplier", expr_string e2)]
43+
| Ordered -> transform (`String "ordered")
44+
| PositiveOrdered -> transform (`String "positive_ordered")
45+
| Simplex -> transform (`String "simplex")
46+
| UnitVector -> transform (`String "unit_vector")
47+
| SumToZero -> transform (`String "sum_to_zero")
48+
| CholeskyCorr -> transform (`String "cholesky_corr")
49+
| CholeskyCov -> transform (`String "cholesky_cov")
50+
| Correlation -> transform (`String "correlation")
51+
| Covariance -> transform (`String "covariance")
52+
| StochasticRow -> transform (`String "stochastic_row")
53+
| StochasticColumn -> transform (`String "stochastic_column")
54+
| TupleTransformation ts ->
55+
transform (`List (List.map ~f:(fun t -> `Assoc (transformation t)) ts))
56+
2657
let get_var_decl {stmts; _} : t =
2758
`Assoc
2859
(List.fold_right ~init:[]
2960
~f:(fun stmt acc ->
3061
match stmt.Ast.stmt with
3162
| Ast.VarDecl decl ->
3263
let type_info = basetype_dims decl.decl_type in
64+
let transform_info = transformation decl.transformation in
3365
let decl_info =
3466
List.map
35-
~f:(fun {identifier; _} -> (identifier.name, type_info))
67+
~f:(fun {identifier; _} ->
68+
(identifier.name, `Assoc (type_info @ transform_info)))
3669
decl.variables in
3770
decl_info @ acc
3871
| _ -> acc)

‎test/integration/cli-args/info/info.expected‎

Lines changed: 70 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,21 @@
11
$ ../../../../../install/default/bin/stanc --include-paths=.,includes --info info.stan
22
{
33
"inputs": {
4-
"a": { "type": "int", "dimensions": 0 },
5-
"b": { "type": "real", "dimensions": 0 },
6-
"c": { "type": "real", "dimensions": 1 },
7-
"d1": { "type": "real", "dimensions": 1 },
8-
"d2": { "type": "real", "dimensions": 1 },
9-
"e": { "type": "real", "dimensions": 2 },
10-
"f": { "type": "int", "dimensions": 1 },
11-
"g": { "type": "real", "dimensions": 1 },
12-
"h": { "type": "real", "dimensions": 2 },
13-
"i": { "type": "real", "dimensions": 3 },
14-
"j": { "type": "int", "dimensions": 3 },
15-
"cplx": { "type": "complex", "dimensions": 1 },
16-
"cplx_vec": { "type": "complex", "dimensions": 1 },
17-
"cplx_row": { "type": "complex", "dimensions": 1 },
18-
"cplx_mat": { "type": "complex", "dimensions": 2 },
4+
"a": { "type": "int", "dimensions": 0, "constraint": "none" },
5+
"b": { "type": "real", "dimensions": 0, "constraint": "none" },
6+
"c": { "type": "real", "dimensions": 1, "constraint": "none" },
7+
"d1": { "type": "real", "dimensions": 1, "constraint": "none" },
8+
"d2": { "type": "real", "dimensions": 1, "constraint": "none" },
9+
"e": { "type": "real", "dimensions": 2, "constraint": "none" },
10+
"f": { "type": "int", "dimensions": 1, "constraint": "none" },
11+
"g": { "type": "real", "dimensions": 1, "constraint": "none" },
12+
"h": { "type": "real", "dimensions": 2, "constraint": "none" },
13+
"i": { "type": "real", "dimensions": 3, "constraint": "none" },
14+
"j": { "type": "int", "dimensions": 3, "constraint": "none" },
15+
"cplx": { "type": "complex", "dimensions": 1, "constraint": "none" },
16+
"cplx_vec": { "type": "complex", "dimensions": 1, "constraint": "none" },
17+
"cplx_row": { "type": "complex", "dimensions": 1, "constraint": "none" },
18+
"cplx_mat": { "type": "complex", "dimensions": 2, "constraint": "none" },
1919
"tuples": {
2020
"type": [
2121
{ "type": "int", "dimensions": 0 },
@@ -28,22 +28,65 @@
2828
"dimensions": 0
2929
}
3030
],
31-
"dimensions": 1
31+
"dimensions": 1,
32+
"constraint": [
33+
{ "constraint": "none" },
34+
{ "constraint": "none" },
35+
{
36+
"constraint": [
37+
{ "constraint": "none" }, { "constraint": "none" }
38+
]
39+
}
40+
]
3241
}
3342
},
3443
"parameters": {
35-
"l": { "type": "real", "dimensions": 1 },
36-
"m": { "type": "real", "dimensions": 1 },
37-
"n": { "type": "real", "dimensions": 1 },
38-
"o": { "type": "real", "dimensions": 1 },
39-
"p": { "type": "real", "dimensions": 2 },
40-
"q": { "type": "real", "dimensions": 2 },
41-
"r": { "type": "real", "dimensions": 2 },
42-
"s": { "type": "real", "dimensions": 2 },
43-
"y": { "type": "real", "dimensions": 0 }
44+
"low": {
45+
"type": "real",
46+
"dimensions": 0,
47+
"constraint": { "lower": "0" }
48+
},
49+
"l": { "type": "real", "dimensions": 1, "constraint": "simplex" },
50+
"m": { "type": "real", "dimensions": 1, "constraint": "unit_vector" },
51+
"n": { "type": "real", "dimensions": 1, "constraint": "ordered" },
52+
"o": {
53+
"type": "real",
54+
"dimensions": 1,
55+
"constraint": "positive_ordered"
56+
},
57+
"p": { "type": "real", "dimensions": 2, "constraint": "covariance" },
58+
"q": { "type": "real", "dimensions": 2, "constraint": "correlation" },
59+
"r": { "type": "real", "dimensions": 2, "constraint": "cholesky_cov" },
60+
"s": { "type": "real", "dimensions": 2, "constraint": "cholesky_corr" },
61+
"y": { "type": "real", "dimensions": 0, "constraint": "none" },
62+
"parameter_transforms": {
63+
"type": [
64+
{ "type": "real", "dimensions": 0 },
65+
{
66+
"type": [
67+
{ "type": "real", "dimensions": 1 },
68+
{ "type": "real", "dimensions": 2 }
69+
],
70+
"dimensions": 0
71+
}
72+
],
73+
"dimensions": 0,
74+
"constraint": [
75+
{ "constraint": { "lower": "y" } },
76+
{
77+
"constraint": [
78+
{ "constraint": "simplex" }, { "constraint": "covariance" }
79+
]
80+
}
81+
]
82+
}
83+
},
84+
"transformed parameters": {
85+
"t": { "type": "real", "dimensions": 2, "constraint": "none" }
86+
},
87+
"generated quantities": {
88+
"u": { "type": "real", "dimensions": 0, "constraint": "none" }
4489
},
45-
"transformed parameters": { "t": { "type": "real", "dimensions": 2 } },
46-
"generated quantities": { "u": { "type": "real", "dimensions": 0 } },
4790
"functions": [
4891
"fatal_error", "log", "print", "reduce_sum", "reject", "sin", "square"
4992
],

‎test/integration/cli-args/info/info.stan‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ transformed data {
3737
}
3838

3939
parameters {
40+
real <lower=0> low;
4041
simplex[10] l;
4142
unit_vector[11] m;
4243
ordered[12] n;
@@ -46,6 +47,8 @@ parameters {
4647
cholesky_factor_cov[16] r;
4748
cholesky_factor_corr[17] s;
4849
real y;
50+
51+
tuple(real<lower=y>, tuple(simplex[18], cov_matrix[19])) parameter_transforms;
4952
}
5053

5154
transformed parameters {

0 commit comments

Comments
 (0)