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:

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

[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.