Skip to content

Commit 8fb1ed4

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

File tree

3 files changed

+100
-33
lines changed

3 files changed

+100
-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 f -> `Assoc (unsized_basetype_json f)) 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 = [("transform", 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: 58 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, "transform": "none" },
5+
"b": { "type": "real", "dimensions": 0, "transform": "none" },
6+
"c": { "type": "real", "dimensions": 1, "transform": "none" },
7+
"d1": { "type": "real", "dimensions": 1, "transform": "none" },
8+
"d2": { "type": "real", "dimensions": 1, "transform": "none" },
9+
"e": { "type": "real", "dimensions": 2, "transform": "none" },
10+
"f": { "type": "int", "dimensions": 1, "transform": "none" },
11+
"g": { "type": "real", "dimensions": 1, "transform": "none" },
12+
"h": { "type": "real", "dimensions": 2, "transform": "none" },
13+
"i": { "type": "real", "dimensions": 3, "transform": "none" },
14+
"j": { "type": "int", "dimensions": 3, "transform": "none" },
15+
"cplx": { "type": "complex", "dimensions": 1, "transform": "none" },
16+
"cplx_vec": { "type": "complex", "dimensions": 1, "transform": "none" },
17+
"cplx_row": { "type": "complex", "dimensions": 1, "transform": "none" },
18+
"cplx_mat": { "type": "complex", "dimensions": 2, "transform": "none" },
1919
"tuples": {
2020
"type": [
2121
{ "type": "int", "dimensions": 0 },
@@ -28,22 +28,53 @@
2828
"dimensions": 0
2929
}
3030
],
31-
"dimensions": 1
31+
"dimensions": 1,
32+
"transform": [
33+
{ "transform": "none" },
34+
{ "transform": "none" },
35+
{ "transform": [ { "transform": "none" }, { "transform": "none" } ] }
36+
]
3237
}
3338
},
3439
"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 }
40+
"low": { "type": "real", "dimensions": 0, "transform": { "lower": "0" } },
41+
"l": { "type": "real", "dimensions": 1, "transform": "simplex" },
42+
"m": { "type": "real", "dimensions": 1, "transform": "unit_vector" },
43+
"n": { "type": "real", "dimensions": 1, "transform": "ordered" },
44+
"o": { "type": "real", "dimensions": 1, "transform": "positive_ordered" },
45+
"p": { "type": "real", "dimensions": 2, "transform": "covariance" },
46+
"q": { "type": "real", "dimensions": 2, "transform": "correlation" },
47+
"r": { "type": "real", "dimensions": 2, "transform": "cholesky_cov" },
48+
"s": { "type": "real", "dimensions": 2, "transform": "cholesky_corr" },
49+
"y": { "type": "real", "dimensions": 0, "transform": "none" },
50+
"parameter_transforms": {
51+
"type": [
52+
{ "type": "real", "dimensions": 0 },
53+
{
54+
"type": [
55+
{ "type": "real", "dimensions": 1 },
56+
{ "type": "real", "dimensions": 2 }
57+
],
58+
"dimensions": 0
59+
}
60+
],
61+
"dimensions": 0,
62+
"transform": [
63+
{ "transform": { "lower": "y" } },
64+
{
65+
"transform": [
66+
{ "transform": "simplex" }, { "transform": "covariance" }
67+
]
68+
}
69+
]
70+
}
71+
},
72+
"transformed parameters": {
73+
"t": { "type": "real", "dimensions": 2, "transform": "none" }
74+
},
75+
"generated quantities": {
76+
"u": { "type": "real", "dimensions": 0, "transform": "none" }
4477
},
45-
"transformed parameters": { "t": { "type": "real", "dimensions": 2 } },
46-
"generated quantities": { "u": { "type": "real", "dimensions": 0 } },
4778
"functions": [
4879
"fatal_error", "log", "print", "reduce_sum", "reject", "sin", "square"
4980
],

‎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)