Proposal: Restructure the repository

[rlouf]

The problem

As discussed in #441 there are a few problems with the current organisation of the repository:

• kernels.py is getting quite big; contributions to unrelated algorithms have to make substantial modifications to the same big file. The construction could be generalized per algorithm type.
• Abstractions (e.g. MCMCAlgorithm) are defined at the top level of the library when they don't generalize past their related module (here /mcmc)
• It is not clear where meta-algorithms should live. SMC has its own module, but what about replica exchange, MCMC couplings, etc?

Proposal

We should have one type.py per class of algorithm. The /mcmc/types.py module would contain, for instance:

class MCMCInitFn(Protocol):
    def __call__(self, position: Position, logdensity_fn: Callable) -> State:
        """Initialize the algorithm's state."""

class MCMCStepFn(Protocol):
    def __call__(
        self,
        rng_key: PRNGKey,
        state: State,
        logdensity_fn: Callable,
        *parameters
    ) -> Tuple[State, Info]:
        """Update the current state using the sampling algorithm."""

We should have on alias.py per module. Working with algorithms of the same type means we can generalize the alias construction, and it thus makes sens to regroup them in the same file.

hmc = build_mcmc_algorithm(hmc.init, hmc.step)  # Returns `MCMCAlgorithm`
nuts = build_mcmc_algorithm(nuts.init, nuts.step)

The structure would thus look like

/mcmc
  /adaptation
    types.py
    alias.py
    meads.py
    pathfinder.py
    window.py
    ...
  /diagnostics
    ess.py
    rhat.py
  types.py
  alias.py
  hmc.py
  ...
/sgmcmc
  types.py
  alias.py
  slgd.py
  ...
/vi
  types.py
  alias.py
  meanfield_vi.py
  pathfinder.py
  ..
/meta
  types.py
  alias.py
  smc.py

We should aim to remove the /optimizer folder, either by upstreaming what is not imported from JaxOpt or Optax, or leaving the algorithms where they are needed.

[johannahaffner]

Did this go anywhere?

It sounds very sensible to me. I'm new to Blackjax and I looked for a structure like this when trying to figure out where the relevant bits of code live.

[junpenglao]

I think the structure now mostly follows (with some modification/changes)

[johannahaffner]

I agree with respect to the folder structure, and maybe I just don't have a good enough overview. Please excuse me if that is the case!

For instance in the MCMC folder, almost all algorithms define State and Info, inheriting from NamedTuple, in some cases in very similar or identical ways.

I don't know enough about sampling to judge if there is potential for abstractions according to families of sampling algorithms here. But if there is, then this would be a feature that I would find helpful as a user.