NumPyro for Sampling

[willGraham01]

Is Your Feature Request Related to a Problem? Please Describe

We should investigate NumPyro as a possible library for handling the (Python object that describes) the distributions, and samples from them.

Describe the Solution You'd Like

Node (or DistributionNode) should have a method of constructing the "random variable" (or more precisely, the distribution describing that RV) from the information stored in the Node.

For now, we will assume we're only using NumPyro distributions, and we may re-introduce the backend-agnosticity later.

Describe Alternatives You've Considered

No response

Additional Context

[willGraham01]

#59 now lets us create "parameterised" models from our Graph. That is to say, one can call

constructor = graph.model_constructor()

The constructor obtained is itself another function, which is called using

realisation = constructor(**parameter_node_values)

and will produce realisation, which is the graph model with the parameter nodes set to the param_node_values, coded up in a way that numpyro understands. realisation can thus be used can be used as a sampler for MCMC algorithms, for example:

constructor = graph.model_constructor()
realisation = constructor(**parameter_node_values)

kernel = numpyro.infer.NUTS(realisation)
mcmc = numpyro.mcmc(kernel)
mcmc.run()

# Samples can then be fetched using
mcmc.get_samples()[node_label]

This means that we do now have a way to draw samples from a model, if we have been told the parameter node values beforehand. In particular, we can now empirically estimate the moments of RVs using our newfound ability to generate appropriate samples.

This is still not ideal for our needs when looking at the optimisation problem. We have to run an MCMC sampler every time we want to evaluate a particular moment of a random variable, since we have to re-build the realisation each time the parameter nodes change value.

[willGraham01]

@mscroggs In light of the above, we should now be able to make functions like

moment(graph, node, order, sample_generation_options, **param_node_values)
# graph & node are self explanatory. node likely only needs to be node_label too
# order; 1 = mean, 2 = variance etc
# sample_generation_options = things we need to pass to the underlying sampler to get it to generate samples
# param_node_values = the values the parameter nodes are taking, which we need to then generate samples and thus provide an estimate

that can estimate the order-moment of a distribution attached to a node (likely multiple nodes at the same in fact), if I'm not mistaken?

If so, we might want to re-visit the current algorithms submodule and update it to use this sampling method?

[willGraham01]

All of these issues: #47 #49 #50 #25 #26 #29, will be closed by virtue of solving this issue.

We decided that Node.sample should just run an MCMC on the graph model to generate a batch of samples, and then Node.moment (and other derived methods) should just use these samples to give an empirical estimation.

(In practice, we should probably have a .sample method for the Graph class that just generates MCMC samples for all RVs in the graph, then Node.sample can just pluck the samples it cares about. We could even have Graph.sample take a node_labels argument so we don't have to run one MCMC per node we want to sample from.)

This will be different to what we implement when solving the causal problem itself, which will need wrapped functions and likely some effect handlers, which will not be intended to be evaluated manually by the user.