Update run_inference_algorithm to split initial_position and initial_state

[reubenharry]

Improvement to run_inference_algorithm

Description:

run_inference_algorithm currently uses a try-except clause to allow the user to either provide an initial position or an initial state. This has led to some problems, when the except clause fails to trigger.

run_inference_algorithm also produces an array of [num_steps, n_dims], which for long chains on high-dim problems can be prohibitively large.

Solution:

Make run_inference_algorithm take initial_state or initial_position explicitly, so that the caller of run_inference_algorithm is responsible for providing initial_state or initial_position, rather than deferring this task.

Add a streaming mode for run_inference_algorithm.

A few important guidelines and requirements before we can merge your PR:

• The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• There are tests covering the changes;
• The doc is up-to-date;

[junpenglao]

Actually, a few more nit

[junpenglao]

Should we combine the kwargs of transform and expectation for now? Otherwise we need a better name for expectation

[reubenharry]

expectation seems to me like an appropriate name, since indeed the value calculated is the expectation. For example, if expectation=lambda x: x**2, then you get back $E[x^2]$.

By contrast, transform operates on the full history of samples. In the future, I think it would make sense for it to also take Info, so that the user can choose to dispense with (part of) the diagnostic info. That is why I thought it was better to keep them separate.

I think it would in theory be possible to make them the same, if necessary.

[junpenglao]

That makes sense, but looking at the implementation, the return state is always transform(state), which means for a expectation transformation, we should probably do expectation(transform(x))?

[junpenglao]

That makes sense, but looking at the implementation, the return state is always transform(state), which means for a expectation transformation, we should probably do expectation(transform(x))?

[reubenharry]

The other argument returned, average, is the expectation. I'm not sure expectation(transform(x)) would work, because they are both functions of x (the state).

[junpenglao]

The other argument returned, average, is the expectation. I'm not sure expectation(transform(x)) would work, because they are both functions of x (the state).

What i meant here is that:

1. The function intially returns transform(x), which is a function of x (the state). For examples with transform is lambda x: x.position. If user wants to compute expectation eg $E[x^2]$, they need to do np.mean(mcmc_samples ** 2)
2. Now with the new functionarity that computes the expectation and not storing the full trace, my arguement here is that we should have trasnform being a function of the state, but function expectation being a function of the output of transform.

Does this makes sense?

[reubenharry]

Yeah, that makes sense. Updated.

[junpenglao]

Thank you!