I have a model which is based on integrating a d-dimensional ordinary differential equation (ODE), where the dimensionality is given by the parameter
sigma = numpyro.sample("sigma", dist.LogNormal(-1, 1).expand([n_components])) y = odeint(dydt, ...) numpyro.sample("y", dist.Normal(y, sigma), obs=y)
The result of the ODE integration (
y) is an array of shape (n_timesteps x n_components). And in the code written above inference is performed over the whole output array
y. In my case I want to further filter this, meaning that depending on the component, I want to consider a different subset of timesteps in the inference.
An example would be a system with
n_components = 3 and 10 time steps in the ODE integration leading to 10x3 array for
y. Now I want to only consider for component 0 the timesteps [0,1,2,5], for component 1 the timesteps [1,2,8,9] and for component 3 [0,2,3,4,5,7,8] in the inference.
What is the best way to do this?