Re-implementation of the pyrenews basicrenewlmodel using the proposed modularization

[ghost]

This PR corresponds to the materialization of the programming pattern proposal. The main changes are:

Structure

• It creates submodules to organize the code better, namely: processes, observations, latent, and models.
• It creates the meta-classes RandomProcess and Model, which are used by the new submodules. The Model class includes the following pre-defined methods:
  • _init_model(): Creates the NUTS kernel and MCMC sampler
  • run(): runs the MCMC.
  • print_summary(): A wrapper of MCMC.print_summary.
  • spread_draws(): A wrapper of mcmcutils.spread_draws() (see below).
• RandomProcess and Model have a sample() function with three parameters: self, random_variables (dict), and constants (dict).
• All instances of RandomProcess and Model return a collections.namedtuple object.

Re-factoring and organization

• It refactors the basic.py model into two different models: BasicRenewalModel and HospitalizationsModel, the latter composed by BasicRenewalModel.
• It moves the function spread_draws (from the pyrenew_demo.ipynp) to mcmcutils.py within the library.

New functionalities + code

• It creates two new RandomProcess' (latent): Infections and Hospitalizations (abstractions from the basic.py) model.
• It creates the process RtRandomWalkProcess (also from basic.py).
• It adds six new test files (which include 9 tests) for each one of the observation processes and model.
• Adds two notebooks: getting-started and pyrenew_demo (from the previous version of the project)

Rough edges

• The documentation needs work. Although all parameters and main methods are documented. We need to make a editorial review to make sure it is clear.
• Random tests are WIP. Since we still need to figure out how to properly test random functions (see Correctness unit tests for pyrenew.process #4,) the implemented tests make some simple checks, including: The function runs, the shape is as expected, etc.

[ghost]

OK, @dylanhmorris. This is ready for review! Please see the notes at the description of the PR.

[dylanhmorris]

Some comments and suggestions inline.

Additional general comment: Right now, we use the name <Quantity>Observation for classes that represent a key quantity and optionally allow us to observe that quantity with noise: e.g.InfectionsObservation for infections; HospitalizationsObservation for hospital admissions. I think this could be misleading, since we also use those classes to sample those quantities when they are latent. Some options:

Option 1: Base class Observable, with optional user-defined observation models

Have these quantities be of base class Observable or ObservableQuantity and name them things like

class Infections(Observable)
     ...
class HospitalAdmissions(Observable)
     ...

Observable classes are made actually observed giving them a non-null observation model at instantiation. This is close the to current pattern, as these lines of getting-started.ipynb:

infections_obs = InfectionsObservation(
    gen_int=jnp.array([0.25, 0.25, 0.25, 0.25]),
    inf_observation_model=PoissonObservation(
        rate_varname='infections_mean',
        counts_varname='infections_obs',
        )
    )

Option 2: Inheritance

You could also do this via inheritance, though it's not my preference

class Infections(Observable):
    ...

class ObservedInfections(Infections)
    def observe(self, data):
         ...

Option 3: Fully separate the latent observables from their observation processes

With this option, you'd use an object of class Infections to sample predicted infections and then (optionally) an object of class Observation (e.g. NegativeBinomialObservation, PoissonObservation) to observe it, but you'd do that at the level of the Model class model() function (or sample() function if we rename it per my suggestion above).

e.g. something like the following:

class MyModel(Model):
    def __init__(self):
        self.infections = Infections()
        self.obs_dist = NegativeBinomialObservation()
   def model(self, data):
        pred_infections = self.infections.sample()
        inf_obs = self.obs_dist.sample(
             "observed_infections",
             pred_infections)

Note that (especially if we go with option three) we go this route I'm not convinced that the base class for things like Infections, HospitalAdmissions etc should be named Observable. We can separately discuss the name.

[ghost]

OK @dylanhmorris, the latest version features the following changes:

1. The __init__.py script now uses __all__ to avoid the pre-commit comment.
2. HospitalizationsModel now uses BasicRenewalModel as a component, no inheritance.
3. All documentation now uses the numpy style.
4. Given your comments, @SamuelBrand1, and @seabbs regarding latent vs observed, we now follow approach 3 in your comment, i.e., latent and observed process are fully split.
5. Removed the plot method from Model (for now)
6. Removed the spread indices function from mcmcutils.py

We can discuss details tomorrow during the team meeting.

[ghost]

OK @dylanhmorris, the latest version features the following changes:

1. The __init__.py script now uses __all__ to avoid the pre-commit comment.
2. HospitalizationsModel now uses BasicRenewalModel as a component, no inheritance.
3. All documentation now uses the numpy style.
4. Given your comments, @SamuelBrand1, and @seabbs regarding latent vs observed, we now follow approach 3 in your comment, i.e., latent and observed process are fully split.
5. Removed the plot method from Model (for now)
6. Removed the spread indices function from mcmcutils.py

We can discuss details tomorrow during the team meeting.

Also:

• The Hospitalizations latent class is called HospitalAdmissions
• Previously, the HospitalizationsModel would receive a RandomProcess for observed infections. This is unlikely to happen and confusing so I've removed it.
• BasicRenewalModel is now called RtInfectionsRenewalModel.

[dylanhmorris]

We should think about module naming on the following fronts:

1. Plural versus singular
2. Noun-like versus adjective-like

Currently we have pyrenew.processes and pyrenew.observations (plural nouns) and pyrenew.latent (adjective)

[ghost]

Thank you, @dylanhmorris! That was 💨!