PGMax
PGMax copied to clipboard
Incomplete documentation for evidence
I would like to do the following, which is one step of the EM algorithm:
- set up a factor graph corresponding to a simple Bayesian network (I've got this far)
- specify evidence in the form of specified outcomes or marginal distributions for certain nodes
- perform belief propagation, get the resulting factors and use them to construct a new Bayesian network
I'm currently stuck at the second bullet point. If I search the docs for "evidence" I find pgmax.fg.graph.Evidence
and pgmax.fg.graph.update_evidence
, but the documentation for these is very incomplete and doesn't specify the format in which the inputs should be provided or the procedure for using them.
The blog post doesn't use these but something else - it uses an argument evidence_updates
to run_bp
, but this is not mentioned in the documentation. I can't really understand what this argument is doing in the example, but I suspect it's not what I'm looking for.
So I'm a bit stuck on understanding the basic usage, and I'm wondering if I could either get a hint about it, and/or whether the documentation could be updated to make it a bit clearer.
This is definitely another part of the documentation we need to improve.
You can update evidence using intuitive assignment syntax. https://github.com/vicariousinc/PGMax/blob/f7f3e3fe7fbb3a3c7582b7093c6ecd56eb04963a/examples/ising_model.py#L96 Here is an example for an Ising model. You index the evidence with variable names (tuples of the type (group_name, variable_name_in_group)
) and do assignment. If you are working with variable groups (e.g. an NDVariableArray
), you can also use intuitive assignment syntax with arrays (e.g. https://github.com/vicariousinc/PGMax/blob/f7f3e3fe7fbb3a3c7582b7093c6ecd56eb04963a/examples/ising_model.py#L103).
The evidence_updates
argument to run_bp
updates evidence on the fly for the given belief propagation run. You can use the elements in this dictionary in the above assignment syntax to achieve the same kind of evidence updates.
Some additional comments for this two different ways of doing evidence updates:
- Using the assignment syntax the updates are persistent. When you do
run_bp
later you will always work with the updated evidence. - Using
evidence_updates
the updates are not persistent. It only updates the evidence for that particularrun_bp
run. But by structuringevidence_updates
as a dictionary of jnp arrays and havingrun_bp
as pure python functions, we obtain the ability to vectorize over/get gradients w.r.t. the evidence updates usingjax.vmap
/jax.grad
. This is the main motivation for organizing things this way.
Also you are more than welcome to contribute an example with directed models (since you obviously are working on one). Please refer to the developer installation instructions and contribution guide, and feel free to open a PR when it's ready! I can make more comments on the PR and make it part of our examples.