Separate "model" into distribution and priors/sampling

[swo]

I'm pretty sure your trouble are from trying to get samples for the simulation study from the numpyro.sample object, rather than from the dist object.

My sense is that you should parameterize the dist object, and then my_dist.sample() to get your random samples. Then you can wrap around the model function that gives you the parameterized dist with the priors/sampling function that returns the sample object. I call these functions one_division_dist and one_division_sample (or something).

This code doesn't work because I run into a problem when doing all the sampling in one big dist.Multinomial() call. I'm not sure if plating is supposed to make this easy when inference-sampling, but it makes it hard when you're trying to just draw samples from the dist object with dist.sample(), because in fact there are multiple multinomial samples, one per "observation" (lineage + time combination). I not sure I can make the z = beta_0 + beta_1 * time thing work, because the dimensions of the different things is so crazy.

[thanasibakis]

I've manually incorporated these ideas into the branch a bit, rather than done a proper merge. Thanks for thinking this through 🙂