Skip to content

PyMC3 Jupyter Notebook Style Guide

Oriol Abril-Pla edited this page Apr 5, 2021 · 27 revisions

PyMC3 has a very rich notebook (NB) gallery. With the goal of standardizing and giving an identity to this gallery, here are a few steps to check when you create or update a NB:

  1. In a cell just below the cell where you imported matplotlib (usually the first one), set the ArviZ style and display format (this has to be in another cell than the MPL import because of the way MPL sets its defaults):

    RANDOM_SEED = 8927
    rng = np.random.default_rng(RANDOM_SEED)
    az.style.use("arviz-darkgrid")

    A good practice when generating synthetic data is also to set a random seed as above, to improve reproducibility. Also, please check convergence (e.g. assert all(r_hat < 1.03)) because we sometime re-run notebooks automatically without carefully checking each one.

  2. We run some code-quality checks on our notebooks during Continuous Integration. The easiest way to make sure your notebook(s) pass the CI checks is using pre-commit. You can install it with

    pip install -U pre-commit

    and then enable it with

    pre-commit install

    Then, the code-quality checks will run automatically whenever you commit any changes. To run the code-quality checks manually, you can do, e.g.:

    pre-commit run --files notebook1.ipynb notebook2.ipynb

    replacing notebook1.ipynb and notebook2.ipynb with any notebook you've modified.

    NB: sometimes, Black will be frustrating (well, who isn't?). In these cases, you can disable its magic for specific lines of code: just write #fmt: on/off to disable/re-enable it, like this:

    # fmt: off
    np.array(
        [
            [1, 0, 0, 0],
            [0, -1, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, -1],
        ]
    )
    # fmt: on
  3. Once you're finished with your NB, add a very last cell with the watermark package. This will automatically print the versions of Python and the packages you used to run the NB -- reproducibility rocks!

    %load_ext watermark
    %watermark -n -u -v -iv -w (-p theano xarray)

    watermark should be in your virtual environment if you installed our requirements-dev.txt. Otherwise, just run pip install watermark. The p flag is optional but should be added if Theano (or Aesara if in v4) or xarray are not imported explicitly. This will also be checked by pre-commit (because we all forget to do things sometimes 😳).

You're all set now 🎉 You can push your changes, open a pull request, and, once it's merged, rest with the feeling of a job well done 👏 Thanks a lot for your contribution to open-source, we really appreciate it!