Update jupyter style guide

docs/source/conf.py

@@ -153,6 +153,8 @@
     "aesara": ("https://aesara.readthedocs.io/en/latest/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "nb": ("https://pymc-examples.readthedocs.io/en/latest/", None),
+    "myst": ("https://myst-parser.readthedocs.io/en/latest", None),
+    "myst-nb": ("https://myst-nb.readthedocs.io/en/latest/", None),
 }
 
 

docs/source/contributing/jupyter_style.md

@@ -1,21 +1,49 @@
 (jupyter_style)=
 # Jupyter Style Guide
 
-These guidelines should be followed by all notebooks in the documentation, independently of
-the repository where the notebook is in (pymc or pymc-examples).
+These guidelines should be followed by notebooks in the documentation.
+All notebooks in pymc-examples must follow this to the letter, the style
+is more permissive for the ones on pymc where not everything is available.
+
+The documentation websites are generated by Sphinx, which uses
+{doc}`myst:index` and {doc}`myst-nb:index`
+to parse the notebooks.
 
 ## General guidelines
 
 * Don't use abbreviations or acronyms whenever you can use complete words. For example, write "random variables" instead of "RVs".
 
 * Explain the reasoning behind each step.
 
-* Use the glossary whenever possible. If you use a term that is defined in the Glossary, link to it the first time that term appears in a significant manner. Use [this syntax](https://jupyterbook.org/content/content-blocks.html?highlight=glossary#glossaries) to add a term reference. [Link to glossary source](https://github.com/pymc-devs/pymc/blob/main/docs/source/glossary.md) where new terms should be added.
-
 * Attribute quoted text or code, and link to relevant references.
 
 * Keep notebooks short: 20/30 cells for content aimed at beginners or intermediate users, longer notebooks are fine at the advanced level.
 
+### MyST guidelines
+Using MyST allows taking advantage of all sphinx features from markdown cells in the notebooks.
+All markdown should be valid MyST (note that MyST is a superset of recommonmark).
+This guide does not teach nor cover MyST extensively, only gives some opinionated guidelines.
+
+* **Never** use url links to refer to other notebooks, PyMC documentation or other python
+  libraries documentations. Use [sphinx cross-references](https://docs.readthedocs.io/en/stable/guides/cross-referencing-with-sphinx.html)
+  instead.
+
+  :::{caution}
+  Using urls links breaks self referencing in versioned docs! And at the same time they are
+  less robust than sphinx cross-references.
+  :::
+
+  * When linking to other notebooks, always use a `ref` type cross-reference pointing
+    to the target in the {ref}`jupyter_style_first_cell`.
+
+* If the output (or even code and output) of a cell is not necessary to follow the
+  notebook or it is very long and can break the flow of reading, consider hiding
+  it with a {doc}`toggle button <myst-nb:use/hiding>`
+
+* Consider using {ref}`myst:syntax/figures` to add captions to images used in the notebook.
+
+* Use the glossary whenever possible. If you use a term that is defined in the Glossary, link to it the first time that term appears in a significant manner. Use [this syntax](https://jupyterbook.org/content/content-blocks.html?highlight=glossary#glossaries) to add a term reference. [Link to glossary source](https://github.com/pymc-devs/pymc/blob/main/docs/source/glossary.md) where new terms should be added.
+
 ### Variable names
 
 * Above all, stay consistent with variable names within the notebook. Notebooks using multiple names for the same variable will not be merged.
@@ -28,7 +56,119 @@ the repository where the notebook is in (pymc or pymc-examples).
 
 * When using non meaningful names such as single letters, add bullet points with a 1-2 sentence description of each variable below the equation where they are first introduced.
 
+Choosing variable names can sometimes be difficult, tedious or annoying.
+In case it helps, the dropdown below has some suggestions so you can focus on writing the actual content
+
+:::::::{dropdown} Variable name suggestions
+:icon: light-bulb
+
+**Models and sampling results**
+* Use `idata` for sampling results, always containing a variable of type InferenceData.
+* Store inferecedata groups as variables to ease writing and reading of code operating on sampling results.
+  Use underscore separated 3-5 word abbreviations or the group name. Some examples of `abbrebiation`/`group_name`:
+  `post`/`posterior`, `const`/`constant_data`, `post_pred`/`posterior_predictive` or `obs_data`/`observed_data`
+* For stats and diagnostics, use the ArviZ function name as variable name: `ess = az.ess(...)`, `loo = az.loo(...)`
+* If there are multiple models in a notebook, assign a prefix to each model,
+  and use it throughout to identify which variables map to each model.
+  Taking the famous eight school as example, with a `centered` and `non_centered` model
+  to compare parametrizations, use `centered_model` (pm.Model object), `centered_idata`, `centered_post`, `centered_ess`... and `non_centered_model`, `non_centered_idata`...
+
+**Dimension and random variable names**
+* Use singular dimension names, following ArviZ `chain` and `draw`.
+  For example `cluster`, `axis`, `component`, `forest`, `time`...
+* If you can't think of a meaningful name for the dimension representing the number of observations such as time, fall back to `obs_id`.
+* For matrix dimensions, as xarray doesn't allow repeated dimension names, add a `_bis` suffix. i.e. `param, param_bis`.
+* For the dimension resulting from stacking `chain` and `draw` use `sample`, that is `.stack(sample=("chain", "draw"))`.
+* We often need to encode a categorical variable as integers. add `_idx` to the name of the variable it's encoding.
+  i.e. from `floor` and `county` to `floor_idx` and `county_idx`.
+* To avoid clashes and overwriting variables when using `pm.Data`, use the following pattern:
+
+  ```
+  x = np.array(...)
+  with pm.Model():
+      x_ = pm.Data("x", x)
+      ...
+  ```
+
+  This avoids overwriting the original `x` while having `idata.constant_data["x"]`,
+  and within the model `x_` is still available to play the role of `x`.
+  Otherwise, always try to use the same variable name as the string name given to the PyMC random variable.
+
+**Plotting**
+* Matplotlib figures and axes. Use:
+  * `fig` for matplotlib figures
+  * `ax` for a single matplotib axes object
+  * `axs` for arrays of matplotlib axes objects
+
+  When manually working with multiple matplotlib axes, use local `ax` variables:
+
+  ::::{tab-set}
+  :::{tab-item} Local `ax` variables
+  ```{code-block} python
+  :emphasize-lines: 3, 7
+  fig, axs = pyplot.subplots()
+
+  ax = axs[0, 1]
+  ax.plot(...)
+  ax.set(...)
+
+  ax = axs[1, 2]
+  ax.scatter(...)
+  ```
+  :::
+  :::{tab-item} Instead of subsetting every time
+  ```
+  fig, axs = pyplot.subplots()
 
+  axs[0, 1].plot(...)
+  axs[0, 1].set(...)
+
+  axs[1. 2].scatter(...)
+  ```
+  :::
+  ::::
+
+  This makes editing the code if restructuring the subplots easier, only one change per subplot
+  is needed instead of one change per matplotlib function call.
+
+* It is often useful to make a numpy linspace into an {class}`~xarray.DataArray`
+  for xarray to handle aligning and broadcasing automatically and ease computation.
+  * If a dimension name is needed, use `x_plot`
+  * If a variable name is needed for the original array and DataArray to coexist, add `_da` suffix
+
+  Thus, ending up with code like:
+
+  ```
+  x = xr.DataArray(np.linspace(0, 10, 100), dims=["x_plot"])
+  # or
+  x = np.linspace(0, 10, 100)
+  x_da = xr.DataArray(x)
+  ```
+
+**Looping**
+* When using enumerate, take the first letter of the variable as the count:
+
+  ```
+  for p, person in enumerate(persons)
+  ```
+
+* When looping, if you need to store a variable after subsetting with the loop index,
+  append the index variable used for looping to the original variable name:
+
+  ```{code-block} python
+  :emphasize-lines: 4, 6
+  variable = np.array(...)
+  x = np.array(...)
+  for i in range(N):
+      variable_i = variable[i]
+      for j in range(K):
+          x_j = x[j]
+          ...
+  ```
+
+:::::::
+
+(jupyter_style_first_cell)=
 ## First cell
 The first cell of all example notebooks should have a MyST target, a level 1 markdown title (that is a title with a single `#`) followed by the post directive.
 The syntax is as follows:
@@ -46,7 +186,7 @@ The syntax is as follows:
 
 The date should correspond to the latest update/execution date, at least roughly (it's not a problem if the date is a few days off due to the review process before merging the PR). This will allow users to see which notebooks have been updated lately and will help the PyMC team make sure no notebook is left outdated for too long.
 
-The [MyST target](https://myst-parser.readthedocs.io/en/latest/syntax/syntax.html#targets-and-cross-referencing)
+The {ref}`MyST target <myst:syntax/targets>`
 is important to ease referencing and linking notebooks between each other.
 
 Tags can be anything, but we ask you to try to use [existing tags](https://github.com/pymc-devs/pymc/wiki/Categories-and-Tags-for-PyMC-Examples)
@@ -58,6 +198,124 @@ Choose a category from [existing categories](https://github.com/pymc-devs/pymc/w
 Authors should list people who authored, adapted or updated the notebook. See {ref}`jupyter_authors`
 for more details.
 
+## Extra dependencies
+If the notebook uses libraries that are not PyMC dependencies, these extra dependencies should
+be indicated together with some advise on how to install them.
+This ensures readers know what they'll need to install beforehand and can for example
+decide between running it locally or on binder.
+
+To make things easier for notebook writers and maintainers, pymc-examples contains
+a template for this that warns about the extra dependencies and provides specific
+installation instructions inside a dropdown.
+
+Thus, notebooks with extra dependencies should:
+
+1.  list the extra dependencies as notebook metadata using the `myst_substitutions` category
+    and then either the `extra_dependencies` or the `pip_dependencies` and `conda_dependencies`.
+    In addition, there is also an `extra_install_notes` to include custom text inside the dropdown.
+
+    * notebook metadata can be edited from the menubar `Edit` -> `Edit notebook metadata`
+      in the dropdown
+
+      This will open a window with json formatted text that might look a bit like:
+
+      ::::{tab-set}
+      :::{tab-item} No myst_substitutions
+
+      ```json
+      {
+        "kernelspec": {
+          "name": "python3",
+          "display_name": "Python 3 (ipykernel)",
+          "language": "python"
+        },
+        "language_info": {
+          "name": "python",
+          "version": "3.9.7",
+          "mimetype": "text/x-python",
+          "codemirror_mode": {
+            "name": "ipython",
+            "version": 3
+          },
+          "pygments_lexer": "ipython3",
+          "nbconvert_exporter": "python",
+          "file_extension": ".py"
+        }
+      }
+      ```
+      :::
+
+      :::{tab-item} extra_dependencies key
+
+      ```{code-block} json
+      :emphasize-lines: 19-21
+      {
+        "kernelspec": {
+          "name": "python3",
+          "display_name": "Python 3 (ipykernel)",
+          "language": "python"
+        },
+        "language_info": {
+          "name": "python",
+          "version": "3.9.7",
+          "mimetype": "text/x-python",
+          "codemirror_mode": {
+            "name": "ipython",
+            "version": 3
+          },
+          "pygments_lexer": "ipython3",
+          "nbconvert_exporter": "python",
+          "file_extension": ".py"
+        },
+        "myst_substitutions": {
+          "extra_dependencies": "bambi seaborn"
+        }
+      }
+      ```
+      :::
+
+      :::{tab-item} pip and conda specific keys
+      ```{code-block} json
+      :emphasize-lines: 19-22
+      {
+        "kernelspec": {
+          "name": "python3",
+          "display_name": "Python 3 (ipykernel)",
+          "language": "python"
+        },
+        "language_info": {
+          "name": "python",
+          "version": "3.9.7",
+          "mimetype": "text/x-python",
+          "codemirror_mode": {
+            "name": "ipython",
+            "version": 3
+          },
+          "pygments_lexer": "ipython3",
+          "nbconvert_exporter": "python",
+          "file_extension": ".py"
+        },
+        "myst_substitutions": {
+          "pip_dependencies": "graphviz",
+          "conda_dependencies": "python-graphviz",
+        }
+      }
+      ```
+
+      The pip and conda spcific keys overwrite the `extra_installs` one, so it doesn't make
+      sense to use `extra_installs` if using them. Either both pip and conda substitutions
+      are defined or none of them is.
+      :::
+      ::::
+
+1.  include the warning and installation advise template with the following markdown right before
+    the extra dependencies are imported:
+
+    ```markdown
+    :::{include} ../extra_installs.md
+    :::
+    ```
+
 ## Code preamble
 
 In a cell just below the cell where you imported matplotlib and/or ArviZ (usually the first one),
@@ -185,7 +443,7 @@ References can be cited twice within a single notebook. Two common reference for
 
 which can be added inline, within the text itself. At the end of the notebook, add the bibliography with the following markdown
 
-```
+```markdown
 ## References
 
 :::{bibliography}
@@ -195,7 +453,7 @@ which can be added inline, within the text itself. At the end of the notebook, a
 
 or alternatively, if you wanted to add extra references that have not been cited within the text, use:
 
-```
+```markdown
 ## References
 
 :::{bibliography}
@@ -211,12 +469,14 @@ Once you're finished with your NB, add a very last cell with [the watermark pack
 
 ```python
 %load_ext watermark
-%watermark -n -u -v -iv -w -p theano,xarray
+%watermark -n -u -v -iv -w -p aesara,aeppl,xarray
 ```
 
 This second to last code cell should be preceded by a markdown cell with the `## Watermark` title only so it appears in the table of contents.
 
-`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.
+`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 Aesara or xarray are not imported explicitly.
 This will also be checked by `pre-commit` (because we all forget to do things sometimes 😳).
 
 ## Epilogue