Repository: pyro-ppl/numpyro

Issue:
Higher level interface for predictions and arviz integration
### ArviZ integration

In https://github.com/arviz-devs/arviz/pull/811, I already made an attempt for this integration. As demonstrated in [InferenceData introduction](https://arviz-devs.github.io/arviz/notebooks/XarrayforArviZ.html), there are several groups of an inference data (prior/posterior/predictive/log_lik/inferece_stats/observed_data/...). Currently, users can do something like [NumPyro section in ArviZ cookbook](https://arviz-devs.github.io/arviz/notebooks/InferenceDataCookbook.html#From-NumPyro). The `az.from_numpyro` requires 3 args `posterior` (which is `mcmc`), `prior` (i.e. prior_predictive), `posterior_predictive`. Currently, we lack a handy utility for prior predictive.

Internally, we lack `observed_data` and `log_lik` groups. The reason is given a `mcmc` instance, we do not store `args`, `kwargs`, so even that we can access to model, we still unable to run it from `arviz`. `observed_data` is required for [ppcplot](https://arviz-devs.github.io/arviz/examples/plot_ppc.html) or [plot_loo_pit](https://arviz-devs.github.io/arviz/examples/plot_loo_pit_overlay.html). `log_lik` is useful for many stuffs, including some bayesian criterions.

**Solutions:**

+ store args, kwargs of `run` method to MCMC state.
+ make prior_predictive, log_like_predictive utilities.

### SVI

I think that current `predictive` utility is enough to get samples from `guide` (similar to the sample_posterior method in autoguide).

### What predictive likes in other frameworks

**Stan** has [generated quantities](https://mc-stan.org/docs/2_18/reference-manual/program-block-generated-quantities.html) block, which is used to
+ forward sampling to generate simulated data for model testing,
+ generating predictions for new data,
+ calculating posterior event probabilities, including multiple comparisons, sign tests, etc.,
+ calculating posterior expectations,
+ transforming parameters for reporting,
+ applying full Bayesian decision theory,
+ calculating log likelihoods, deviances, etc. for model comparison.

**PyMC3** has [deterministic nodes](https://docs.pymc.io/notebooks/api_quickstart.html#Deterministic-transforms) to keep track of transformed variables and [posterior_predictive](https://docs.pymc.io/notebooks/api_quickstart.html#4.-Posterior-Predictive-Sampling) to get values at observed nodes.

Typically, for prediction/forecasting, I imagine we will do something like
```
samples = mcmc.get_samples()
def forecasting_model(*args, **kwargs):
    vars_from_model = model(*args, **kwargs)
    forecast_values = do_something_else(vars_from_model)
    return forecast_values
forecast_values = predictive(rng, forecasting_model, samples)
```

I think we can do all the above tasks by either
+ having _RETURN node in trace as in Pyro and using `predictive` util to get values of this node
+ having a flag `include_returned_values` in predictive (or having another predictive utility for this purpose) which will run `condition(model, preditive_sites + posterior_sites)(*args, **kwargs)` to get those returned values. I prefer this. The minor drawback of this way is we will run the model 2 times: one to get predictive samples, another one to get returned values.


Discussion:
> The reason is given a mcmc instance, we do not store args, kwargs, so even that we can access to model, we still unable to run it from arviz.

Could you point to the arviz code that needs access to args / kwargs to run it? I wonder if there is a way to include it in the wrapper for arviz rather than storing it in MCMC. 

> having _RETURN node in trace as in Pyro and using predictive util to get values of this node

Could you clarify why we need this? This would have been observed by some observed site like `pyro.sample('obs', .., obs=y)`. Is it not possible to extract this from the trace instead? For instance, your forecasting model above is doing something on top of the model over which we run inference. A concrete example will help.

> make prior_predictive, log_like_predictive utilities.

These should be straightforward to add, once we clarify the issues above. I think predictive can simply be extended because when `posterior_samples = {}`, we will not be conditioning any latent sites and be effectively sampling from the prior predictive.
We need arg, kwargs to compute log_likelihood  at observe site. The PR
above does not compute log_lik so currently, arviz does not need to access
args, kwargs. Otherwise, numpyro users will need to compute obs_data,
log_likehood themselves and  provide them to arviz through from_numpyro
method. That’s fine but redundant to me.

About using sample sites to store transformed values, it is doable if we
relax sample_shape, random_state from  fn call (at sample site). Currently
I guess we can do it with: sample(‘return’, lambda sample_shape,
random_state: forecasted_values). Are you suggesting to follow this line? A
concrete example is to get log(y) where y is the returned value from
model(). Another example is to assume that model return loc, scale, we want
to get a sample from Normal(loc, scale) distribution. The generated
quantities from Stan are best to illustrate what I need: get some inferred
values (either stochastic or not) from posterior samples.


On Tue, Sep 24, 2019 at 7:55 PM Neeraj Pradhan <notifications@github.com>
wrote:

> The reason is given a mcmc instance, we do not store args, kwargs, so even
> that we can access to model, we still unable to run it from arviz.
>
> Could you point to the arviz code that needs access to args / kwargs to
> run it? I wonder if there is a way to include it in the wrapper for arviz
> rather than storing it in MCMC.
>
> having _RETURN node in trace as in Pyro and using predictive util to get
> values of this node
>
> Could you clarify why we need this? This would have been observed by some
> observed site like pyro.sample('obs', .., obs=y). Is it not possible to
> extract this from the trace instead? For instance, your forecasting model
> above is doing something on top of the model over which we run inference. A
> concrete example will help.
>
> make prior_predictive, log_like_predictive utilities.
>
> These should be straightforward to add, once we clarify the issues above.
> I think predictive can simply be extended because when posterior_samples
> = {}, we will not be conditioning any latent sites and be effectively
> sampling from the prior predictive.
>
> —
> You are receiving this because you were assigned.
>
>
> Reply to this email directly, view it on GitHub
> <https://github.com/pyro-ppl/numpyro/issues/349?email_source=notifications&email_token=ABEEKVTB3CQGG2466OJJLALQLKSIXA5CNFSM4I2C5EQ2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD7QER4A#issuecomment-534792432>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/ABEEKVW2MUTQEGSHFNHROYLQLKSIXANCNFSM4I2C5EQQ>
> .
>

> Otherwise, numpyro users will need to compute obs_data,
log_likehood themselves and  provide them to arviz through from_numpyro
method. That’s fine but redundant to me.

I think we should provide those utility functions, but I'm trying to understand the constraints here. Based on what you say, we could provide a `log_likelihood` function that depends on `args, kwargs`, but I'm guessing that arviz requires this function to only depend on the model? 

> The generated quantities from Stan are best to illustrate what I need: get some inferred
values (either stochastic or not) from posterior samples.

Can we use something like `numpyro.sample('log_y', dist.Delta(v=trans_x), obs=trans_x)` to record these values? We can probably have a helper function that adds the obs kwarg directly.

> arviz requires this function to only depend on the model

I guess I have given a bad explanation. `arxiv` only needs to access log_likelihood value at the observed node. So there are two ways to do it:
+ compute the value in numpyro: this requires users to provide `log_likelihood` value to `az.from_numpyro` method.
+ compute the value inside the `az.from_numpyro`: this requires access to `args`, `kwargs` so that we can trace the model to get log_likelihood.

For both ways, having `log_lik_predictive` utility is helpful:
+ for the first case, it is helpful for users to compute `log_likelihood` value
+ for the second case, it is helpful to not populate handler codes to arviz (so inside `az.from_numpyro`, we just need to call `log_lik_predictive(rng, mcmc.sampler.model, *args,**kwargs)`.

My proposal is to store `args`, `kwargs` in `MCMC` for the second way.

> Can we use something like numpyro.sample('log_y', dist.Delta(v=trans_x), obs=trans_x) to record these values?

Thanks, I didn't know that we can use Delta distribution for this purpose! :D That would be more than enough! (`obs=trans_x` might not be necessary unless we want to run some inference with `forecast_model`).
> for the second case, it is helpful to not populate handler codes to arviz (so inside az.from_numpyro, we just need to call log_lik_predictive(rng, mcmc.sampler.model, *args,**kwargs).

Sure, that sounds reasonable. I suppose we can store the args, kwargs in MCMC and provide a `log_likelihood(model, posterior_samples, *args, **kwargs)` function. 

Hmm..`obs=` will be necessary I think, otherwise we'll run HMC on those delta distributions and that will forever be rejected.
> obs= will be necessary I think

Sorry, my bad. I should have distinguished `model` and `forecast_model`. The former is the one which we run MCMC, while the latter is just useful for predicting/forecasting. But as you said, having `obs=...` will be helpful if we want to use `model` for both inference and forecasting (but doing so will slow down MCMC I guess).

In addition, now I think that having `return` node which stores a pytree is a bad idea because we can't run `summary` for this node. Your Delta solution can record return values in separated keys, which is much better!
> But as you said, having obs=... will be helpful if we want to use model for both inference and forecasting (but doing so will slow down MCMC I guess).

I don't think there should be any significant slowdown since the delta won't add anything to the PE term, but adding the same obs keyword isn't a very elegant solution. Do you expect the users to write this forecast function? 
I don't mean the slowdown is caused by using Delta distribution. I mean the slowdown is caused by calculating the extra terms (e.g. log(y)) which does not contribute to joint density computation. But if those extra calculations go inside `if forecast` block then there won't be any slowdown at all. For complicated models, I still prefer to create a separate model for prediction, even though there will be duplicated code for `model` and `forecast_model`. :)

> Do you expect the users to write this forecast function?

I am not sure for other users, but it is enough for my purpose. For example, in the [ts forecasting tutorial](http://pyro.ai/numpyro/time_series_forecasting.html#Forecasting), I did use the low-level handlers `substitute`, `seed`, `trace` for forecasting. I think that with your Delta solution, all I need is to create a forecast_model and use `predictive` utility. :D
Are there any remaining todos on this one?
Yeah, I need to verify that we fully support arviz. I will make a notebook for verification tomorrow.

Resolving patch:
diff --git a/numpyro/hmc_util.py b/numpyro/hmc_util.py
--- a/numpyro/hmc_util.py
+++ b/numpyro/hmc_util.py
@@ -22,7 +22,7 @@
 
 TreeInfo = namedtuple('TreeInfo', ['z_left', 'r_left', 'z_left_grad',
                                    'z_right', 'r_right', 'z_right_grad',
-                                   'z_proposal', 'z_proposal_pe', 'z_proposal_grad',
+                                   'z_proposal', 'z_proposal_pe', 'z_proposal_grad', 'z_proposal_energy',
                                    'depth', 'weight', 'r_sum', 'turning', 'diverging',
                                    'sum_accept_probs', 'num_proposals'])
 
@@ -506,10 +506,10 @@ def _combine_tree(current_tree, new_tree, inverse_mass_matrix, going_right, rng,
         turning = current_tree.turning
 
     transition = random.bernoulli(rng, transition_prob)
-    z_proposal, z_proposal_pe, z_proposal_grad = cond(
+    z_proposal, z_proposal_pe, z_proposal_grad, z_proposal_energy = cond(
         transition,
-        new_tree, lambda tree: (tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad),
-        current_tree, lambda tree: (tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad)
+        new_tree, lambda tree: (tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad, tree.z_proposal_energy),
+        current_tree, lambda tree: (tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad, tree.z_proposal_energy)
     )
 
     tree_depth = current_tree.depth + 1
@@ -520,7 +520,7 @@ def _combine_tree(current_tree, new_tree, inverse_mass_matrix, going_right, rng,
     num_proposals = current_tree.num_proposals + new_tree.num_proposals
 
     return TreeInfo(z_left, r_left, z_left_grad, z_right, r_right, r_right_grad,
-                    z_proposal, z_proposal_pe, z_proposal_grad,
+                    z_proposal, z_proposal_pe, z_proposal_grad, z_proposal_energy,
                     tree_depth, tree_weight, r_sum, turning, diverging,
                     sum_accept_probs, num_proposals)
 
@@ -543,7 +543,7 @@ def _build_basetree(vv_update, kinetic_fn, z, r, z_grad, inverse_mass_matrix, st
     diverging = delta_energy > max_delta_energy
     accept_prob = np.clip(np.exp(-delta_energy), a_max=1.0)
     return TreeInfo(z_new, r_new, z_new_grad, z_new, r_new, z_new_grad,
-                    z_new, potential_energy_new, z_new_grad,
+                    z_new, potential_energy_new, z_new_grad, energy_new,
                     depth=0, weight=tree_weight, r_sum=r_new, turning=False,
                     diverging=diverging, sum_accept_probs=accept_prob, num_proposals=1)
 
@@ -654,7 +654,7 @@ def _body_fn(state):
     # update depth and turning condition
     return TreeInfo(tree.z_left, tree.r_left, tree.z_left_grad,
                     tree.z_right, tree.r_right, tree.z_right_grad,
-                    tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad,
+                    tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad, tree.z_proposal_energy,
                     depth, tree.weight, tree.r_sum, turning, tree.diverging,
                     tree.sum_accept_probs, tree.num_proposals)
 
@@ -689,7 +689,7 @@ def build_tree(verlet_update, kinetic_fn, verlet_state, inverse_mass_matrix, ste
     r_ckpts = np.zeros((max_tree_depth, inverse_mass_matrix.shape[-1]))
     r_sum_ckpts = np.zeros((max_tree_depth, inverse_mass_matrix.shape[-1]))
 
-    tree = TreeInfo(z, r, z_grad, z, r, z_grad, z, potential_energy, z_grad,
+    tree = TreeInfo(z, r, z_grad, z, r, z_grad, z, potential_energy, z_grad, energy_current,
                     depth=0, weight=0., r_sum=r, turning=False, diverging=False,
                     sum_accept_probs=0., num_proposals=0)
 
diff --git a/numpyro/mcmc.py b/numpyro/mcmc.py
--- a/numpyro/mcmc.py
+++ b/numpyro/mcmc.py
@@ -27,7 +27,7 @@
 )
 from numpyro.util import cond, copy_docs_from, fori_collect, fori_loop, identity
 
-HMCState = namedtuple('HMCState', ['i', 'z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',
+HMCState = namedtuple('HMCState', ['i', 'z', 'z_grad', 'potential_energy', 'energy', 'num_steps', 'accept_prob',
                                    'mean_accept_prob', 'diverging', 'adapt_state', 'rng'])
 """
 A :func:`~collections.namedtuple` consisting of the following fields:
@@ -37,6 +37,7 @@
    the posterior) at latent sites.
  - **z_grad** - Gradient of potential energy w.r.t. latent sample sites.
  - **potential_energy** - Potential energy computed at the given value of ``z``.
+ - **energy** - Sum of potential energy and kinetic energy of the current state.
  - **num_steps** - Number of steps in the Hamiltonian trajectory (for diagnostics).
  - **accept_prob** - Acceptance probability of the proposal. Note that ``z``
    does not correspond to the proposal if it is rejected.
@@ -236,8 +237,9 @@ def init_kernel(init_params,
         wa_state = wa_init(z, rng_wa, step_size, mass_matrix_size=np.size(z_flat))
         r = momentum_generator(wa_state.mass_matrix_sqrt, rng)
         vv_state = vv_init(z, r)
-        hmc_state = HMCState(0, vv_state.z, vv_state.z_grad, vv_state.potential_energy, 0, 0., 0.,
-                             False, wa_state, rng_hmc)
+        energy = kinetic_fn(wa_state.inverse_mass_matrix, vv_state.r)
+        hmc_state = HMCState(0, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
+                             0, 0., 0., False, wa_state, rng_hmc)
 
         # TODO: Remove; this should be the responsibility of the MCMC class.
         if run_warmup and num_warmup > 0:
@@ -263,10 +265,10 @@ def _hmc_next(step_size, inverse_mass_matrix, vv_state, rng):
         accept_prob = np.clip(np.exp(-delta_energy), a_max=1.0)
         diverging = delta_energy > max_delta_energy
         transition = random.bernoulli(rng, accept_prob)
-        vv_state = cond(transition,
-                        vv_state_new, lambda state: state,
-                        vv_state, lambda state: state)
-        return vv_state, num_steps, accept_prob, diverging
+        vv_state, energy = cond(transition,
+                                (vv_state_new, energy_new), lambda args: args,
+                                (vv_state, energy_old), lambda args: args)
+        return vv_state, energy, num_steps, accept_prob, diverging
 
     def _nuts_next(step_size, inverse_mass_matrix, vv_state, rng):
         binary_tree = build_tree(vv_update, kinetic_fn, vv_state,
@@ -279,7 +281,7 @@ def _nuts_next(step_size, inverse_mass_matrix, vv_state, rng):
                                    r=vv_state.r,
                                    potential_energy=binary_tree.z_proposal_pe,
                                    z_grad=binary_tree.z_proposal_grad)
-        return vv_state, num_steps, accept_prob, binary_tree.diverging
+        return vv_state, binary_tree.z_proposal_energy, num_steps, accept_prob, binary_tree.diverging
 
     _next = _nuts_next if algo == 'NUTS' else _hmc_next
 
@@ -295,9 +297,9 @@ def sample_kernel(hmc_state):
         rng, rng_momentum, rng_transition = random.split(hmc_state.rng, 3)
         r = momentum_generator(hmc_state.adapt_state.mass_matrix_sqrt, rng_momentum)
         vv_state = IntegratorState(hmc_state.z, r, hmc_state.potential_energy, hmc_state.z_grad)
-        vv_state, num_steps, accept_prob, diverging = _next(hmc_state.adapt_state.step_size,
-                                                            hmc_state.adapt_state.inverse_mass_matrix,
-                                                            vv_state, rng_transition)
+        vv_state, energy, num_steps, accept_prob, diverging = _next(hmc_state.adapt_state.step_size,
+                                                                    hmc_state.adapt_state.inverse_mass_matrix,
+                                                                    vv_state, rng_transition)
         # not update adapt_state after warmup phase
         adapt_state = cond(hmc_state.i < wa_steps,
                            (hmc_state.i, accept_prob, vv_state.z, hmc_state.adapt_state),
@@ -309,7 +311,7 @@ def sample_kernel(hmc_state):
         n = np.where(hmc_state.i < wa_steps, itr, itr - wa_steps)
         mean_accept_prob = hmc_state.mean_accept_prob + (accept_prob - hmc_state.mean_accept_prob) / n
 
-        return HMCState(itr, vv_state.z, vv_state.z_grad, vv_state.potential_energy, num_steps,
+        return HMCState(itr, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy, num_steps,
                         accept_prob, mean_accept_prob, diverging, adapt_state, rng)
 
     # Make `init_kernel` and `sample_kernel` visible from the global scope once


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -31,7 +31,7 @@ def model(labels):
     init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(2), model, labels)
     samples = mcmc(warmup_steps, num_samples, init_params, sampler='hmc', algo=algo,
                    potential_fn=potential_fn, trajectory_length=10, constrain_fn=constrain_fn)
-    assert_allclose(np.mean(samples['coefs'], 0), true_coefs, atol=0.21)
+    assert_allclose(np.mean(samples['coefs'], 0), true_coefs, atol=0.22)
 
     if 'JAX_ENABLE_x64' in os.environ:
         assert samples['coefs'].dtype == np.float64