Stan: arrays with nested tuples break stan interface

[aseyboldt]

This currently throws an error when parsing the stan variable names:

import nutpie

code = """
parameters {
    array[2, 3] tuple(matrix[5, 7], real) yi;
}
model {
    for (i in 1:2) {
        for (j in 1:3) {
            yi[i, j].2 ~ normal(-j - 3 * i, 0.1);
            for (k in 1:5) {
                for (m in 1:7) {
                    yi[i, j].1[k, m] ~ normal(m + 7 * k + 7 * 5 * j + 7 * 5 * 3 * i, 0.1);
                }
            }
        }
    }
    
}
"""

compiled = nutpie.compile_stan_model(code=code)

tr = nutpie.sample(compiled)

means = tr.posterior.mean(["draw", "chain"])

[WardBrian]

Oh, I was about to respond to your comment in #205 with a similar example.

Tuples are broadly fine on their own, but arrays of tuples (... of arrays.. of tuples...) is where the real complexity starts to be a bear to support. Of course, in exchange for all that effort, you get to cover maybe 1% of additional use cases 😉.
I won't blame you for a half-way solution or just throwing a better 'not yet supported error'.

StanIO handles this up in Python, with logic that might be adaptable to Rust because it primarily works by computing indexes and shapes. It also can produce a numpy custom dtype for each tuple. Not sure how xarray likes those or not, so you can always use object arrays.

[aseyboldt]

I think it shouldn't be too difficult to support this, as long as I just reproduce the nested structure in terms of variable names. So any compound variable would just stored as a different variable with a name that includes an index like var1:1, until the remaining variable is a relular array of reals, a matrix or a vector.
It is possible to use custom dtypes in xarray, but I think that would probably make life harder for users, they aren't particular well supported in the whole ecosystem. (maybe polars would be fine with them, but that's a whole different storage idea).

I guess the rule is that indices or arrays, matrices and vectors are separated by . in the variable name, and tuple by :?

Edit: Found the logic in stanio: https://github.com/stan-dev/stanio/blob/main/stanio/reshape.py#L155 That will be helpful :-)

[WardBrian]

Yes. Ideally we’d have , for arrays and . for tuples, but the fact that we knew we’d be using these names in a CSV file mandated other choices…

[WardBrian]

I hope the stanio logic is decently clear, I tried to design it in a principled fashion (given the constraints of a very unprincipled output format). The tests for accessing the output are also I think the most rigorous out there -- to my knowledge, none of our R packages would even pass them.

So if do you get comparable handling into nutpie, it would be better than many of the "first party" Stan packages on both speed and usability!