Refactor to store `stan.sampling` directly to enable diagnostic methods about orbit HOT 3 CLOSED

I agree with @steveyang90 that this is only applied to .sampling. Option 1 would trigger too many changes.

First, to have the samples with the drawing order and chain info preserved, we have to turn off permute=True, but it comes with a problem that

• .extract(permuted=True) returns an ordered dictionary
• .extract(permuted=False) returns a ndarray with 3 dims: iteration x chains x params

I found that .extract(pars = ['param_1', ..., ], permuted=False) will return an ordered dictionary, however, each keyed item has an extra dimension (corresponding to the chain number), compared to the return of .exract(permuted=True).

So, I'm thinking the following concrete plan for mcmc method, where fit is compiled_stan_file in our case

stan_extract = fit.extract(pars = fit._get_param_names(), permuted=False)
for idx, (key, val) in enumerate(stan_extract.items()):
    if len(fit._get_param_dims()[idx]) == 0:
        stan_extract[key] = val.flatten(order='F')## here `order` is important to make chains flattened one by one
    else:
        stan_extract[key] = val.reshape(-1, val.shape[-1], order='F') 

After this, we got stan_extract which has exactly the same structure as from .extract(permuted=True), but the sample order is preserved. Say, 4 chains, 500 samples each chain, we will have 2000 samples with order [500 in chain1, ..., 500 in chain4] and insider each chain the draw order is also preserved.

Then we can use this ordered samples for diagnostics viz (ofc, needs a bit processing to cut the samples into chains).

This seems to require minimal change in our code base.

from orbit.

I think this makes sense @wangzhishi

What is if len(fit._get_param_dims()[idx]) == 0 checking for?

Also, you might find np.transpose() useful here

from orbit.

for example, scalar parameter (with dim []) and vector parameter (size 8, with dim [8]) samples has shape (500, 4) and (500, 4, 8) for 4 chains. My proposal is to collapse the shape into (2000,) and (2000, 8), which are consistent with the return of .extract(permuted=True)

from orbit.