Bayesian Modeling and Probabilistic Programming in Python

RAW Doc

Source/ Templates/Autosummary/Class

{{ fullname | escape | underline}}

.. currentmodule:: {{ module }}

.. autoclass:: {{ objname }}

{% block methods %}
{% if methods %}

.. rubric:: Methods

.. autosummary::
:toctree: classmethods

{% for item in methods %}
{{ objname }}.{{ item }}
{%- endfor %}
{% endif %}
{% endblock %}

{% block attributes %}
{% if attributes %}
.. rubric:: Attributes

.. autosummary::
{% for item in attributes %}
~{{ name }}.{{ item }}
{%- endfor %}
{% endif %}
{% endblock %}


Source/ Templates/Distribution

{{ fullname | escape | underline}}

.. currentmodule:: {{ module }}

{% if objtype == "class" %}
.. autoclass:: {{ objname }}

.. rubric:: {{ _('Methods') }}

.. autosummary::
:toctree: classmethods

{{ objname }}.dist
{% else %}
.. autofunction:: {{ objname }}
{% endif %}


Source/Api/Dims/Distributions

*
Distributions
*

Scalar distributions
====================

.. currentmodule:: pymc.dims
.. autosummary::
:toctree: generated/
:template: distribution.rst

Flat
HalfFlat
Uniform,
Normal
HalfNormal
TruncatedNormal
LogNormal
StudentT
HalfStudentT
Cauchy
HalfCauchy
Beta
Laplace
Gamma
InverseGamma
Weibull
Poisson
NegativeBinomial
DiracDelta

Vector distributions
====================

.. currentmodule:: pymc.dims
.. autosummary::
:toctree: generated/
:template: distribution.rst

Categorical
MvNormal
ZeroSumNormal

Higher-Order distributions
==========================

.. currentmodule:: pymc.dims
.. autosummary::
:toctree: generated/
:template: distribution.rst

Censored


Source/Api/Dims/Math

*
Mathematical operations with dimensions
*

This module wraps all the mathematical operations defined in :ref:pytensor.xtensor.math .

It includes a `linalg submodule that wraps all the operations defined in :ref:pytensor.xtensor.linalg `.

Operations defined at the module level in :ref:pytensor.xtensor are available at the `pymc.dims` module level.


Source/Api/Dims/Model


Model constructors

.. currentmodule:: pymc.dims
.. autosummary::
:toctree: generated/

Data
Deterministic
Potential


Source/Api/Dims/Transforms

*
Distribution Transforms
*

.. currentmodule:: pymc.dims.transforms
.. autosummary::
:toctree: generated/

LogTransform
LogOddsTransform
ZeroSumTransform


Source/Api/Distributions/Censored


Censored

..
Manually follow the template in _templates/distribution.rst.
If at any point, multiple objects are listed here,
the pattern should instead be modified to that of the
other API files such as api/distributions/continuous.rst

.. currentmodule:: pymc

.. autoclass:: Censored

.. rubric:: Methods

.. autosummary::
:toctree: classmethods

Censored.dist


Source/Api/Distributions/Continuous


Continuous

.. currentmodule:: pymc
.. autosummary::
:toctree: generated/
:template: distribution.rst

AsymmetricLaplace
Beta
Cauchy
ChiSquared
ExGaussian
Exponential
Flat
Gamma
Gumbel
HalfCauchy
HalfFlat
HalfNormal
HalfStudentT
Interpolated
InverseGamma
Kumaraswamy
Laplace
Logistic
LogitNormal
LogNormal
Moyal
Normal
Pareto
PolyaGamma
Rice
SkewNormal
SkewStudentT
StudentT
Triangular
TruncatedNormal
Uniform
VonMises
Wald
Weibull


Source/Api/Distributions/Custom


CustomDist

..
Manually follow the template in _templates/distribution.rst.
If at any point, multiple objects are listed here,
the pattern should instead be modified to that of the
other API files such as api/distributions/continuous.rst

.. currentmodule:: pymc

.. autoclass:: CustomDist

.. rubric:: Methods

.. autosummary::
:toctree: classmethods

CustomDist.dist


Source/Api/Distributions/Discrete


Discrete

.. currentmodule:: pymc
.. autosummary::
:toctree: generated
:template: distribution.rst

Bernoulli
BetaBinomial
Binomial
Categorical
DiscreteUniform
DiscreteWeibull
Geometric
HyperGeometric
NegativeBinomial
OrderedLogistic
OrderedProbit
Poisson

.. note::

OrderedLogistic and OrderedProbit:
The `OrderedLogistic and OrderedProbit distributions expect the observed values to be 0-based, i.e., they should range from 0 to K-1. Using 1-based indexing (like 1, 2, 3,...K`) can result in errors.


Source/Api/Distributions/Mixture

*
Mixture
*

.. currentmodule:: pymc
.. autosummary::
:toctree: generated
:template: distribution.rst

Mixture
NormalMixture
ZeroInflatedBinomial
ZeroInflatedNegativeBinomial
ZeroInflatedPoisson
HurdlePoisson
HurdleNegativeBinomial
HurdleGamma
HurdleLogNormal


Source/Api/Distributions/Multivariate


Multivariate

.. currentmodule:: pymc
.. autosummary::
:toctree: generated
:template: distribution.rst

CAR
Dirichlet
DirichletMultinomial
ICAR
KroneckerNormal
LKJCholeskyCov
LKJCorr
MatrixNormal
Multinomial
MvNormal
MvStudentT
OrderedMultinomial
StickBreakingWeights
Wishart
WishartBartlett
ZeroSumNormal


Source/Api/Distributions/Simulator

*
Simulator
*

..
Manually follow the template in _templates/distribution.rst.
If at any point, multiple objects are listed here,
the pattern should instead be modified to that of the
other API files such as api/distributions/continuous.rst

.. currentmodule:: pymc

.. autoclass:: Simulator

.. rubric:: Methods

.. autosummary::
:toctree: classmethods

Simulator.dist


Source/Api/Distributions/Timeseries


Timeseries

.. currentmodule:: pymc
.. autosummary::
:toctree: generated
:template: distribution.rst

AR
EulerMaruyama
GARCH11
GaussianRandomWalk
MvGaussianRandomWalk
MvStudentTRandomWalk


Source/Api/Distributions/Transforms

*
Transformations
*

.. currentmodule:: pymc.distributions.transforms

While many distributions are defined on constrained spaces (e.g. intervals), MCMC samplers typically perform best when sampling on the unconstrained real line; this is especially true of HMC samplers. PyMC balances this through the use of transforms. A transform instance can be passed to the constructor of a random variable to tell the sampler how to move between the underlying unconstrained space where the samples are actually drawn and the transformed space constituting the support of the random variable. Transforms are not currently implemented for discrete random variables.

All transforms have three core methods:

  • `forward`: The map from a constrained space to the unconstrained space.
  • `backward`: The inverse map from the unconstrained space to a constrained space.
  • `log_jac_det: The log of the determinant of the Jacobian of the backward` map. This is used to account for the transformed random variable correctly in the posterior log-probability.

.. note::
Transforms are principally intended for internal use and in most cases users do not need to change them. In particular, all continuous distributions on a constrained domain that are implemented in PyMC have a `default_transform` that will automatically transform the random variables as required without needing any extra work from the user.

The main use-cases for setting custom transforms include the following:

#. The `default_transform may need to be replaced with an alternative transform on the same constained space. For example, the default_transform for positive-valued random variables is the :class:log transform but in some cases it may be advantageous to use the :class:log_exp_m1 transform instead.
#. The
default_transform may be removed entirely in some cases when using non-HMC samplers.
#. Exceptionally, transforms can be used to add constraints to the model specification without modifying the
default_transform. This can be done by specifying the additional transform via the transform parameter. However this should not be viewed as a default use-case and, in practice, this is mostly limited to using :class:ordered` in mixture models.

* NB: :class:ordered is not guaranteed to work correctly when used in combination with other transforms, such as :class:simplex and :class:ZeroSumTransform.

.. warning::
Transforms are only applied when sampling unobserved random variables with :func:pymc.sample. In particular:

* Transforms are not applied during forward sampling, i.e. :func:pymc.draw, :func:pymc.sample_prior_predictive and :func:pymc.sample_posterior_predictive
Transforms are not applied when sampling observed* random variables with :func:pymc.sample

Since transforms are not applied during :func:pymc.sample_prior_predictive, a workaround to carry out prior predictive checks is to remove observations from the likelihood and use :func:pymc.sample instead.

Transforms are not usually the correct tool to represent transformations that are part of the generative specification of the model. Such transformations should be included explicitly in the model, typically via :class:pymc.Deterministic. Doing so allows such transformed random variables to be sampled by forward samplers.

Transform Instances
~~~~~~~~~~~~~~~~~~~

Transform instances are the entities that should be used in the
`default_transform or transform` parameters to a random variable
constructor.

.. autosummary::
:toctree: generated

circular
log
log_exp_m1
logodds
ordered
simplex

Specific Transform Classes
~~~~~~~~~~~~~~~~~~~~~~~~~~

An instance of these classes needs to be created before being used
in the `default_transform or transform` parameters to a random variable
constructor.

.. autosummary::
:toctree: generated

CholeskyCovPacked
CircularTransform
Interval
LogExpM1
LogOddsTransform
LogTransform
Ordered
SimplexTransform
ZeroSumTransform

Transform Composition Classes
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

An instance of this class needs to be created from a list of transforms before
being used in the `transform` parameter to a random variable constructor.

If a random variable has a `default_transform and an additional transform
is provided through the
transform parameter, PyMC will automatically
create an instance of the :class:
Chain` transform that applies the
user-provided transform on top of the default one.

.. autosummary::
:toctree: generated

Chain


Source/Api/Distributions/Truncated

*
Truncated
*

..
Manually follow the template in _templates/distribution.rst.
If at any point, multiple objects are listed here,
the pattern should instead be modified to that of the
other API files such as api/distributions/continuous.rst

.. currentmodule:: pymc

.. autoclass:: Truncated

.. rubric:: Methods

.. autosummary::
:toctree: classmethods

Truncated.dist


Source/Api/Distributions/Utilities


Distribution utilities

.. currentmodule:: pymc
.. autosummary::
:toctree: generated/

Continuous
Discrete
Distribution
SymbolicRandomVariable
DiracDelta


Source/Api/Gp/Cov


Covariance Functions

.. automodule:: pymc.gp.cov
.. autosummary::
:toctree: generated

Constant
WhiteNoise
ExpQuad
RatQuad
Exponential
Matern52
Matern32
Linear
Polynomial
Cosine
Periodic
WarpedInput
Gibbs
Coregion
ScaledCov
Kron


Source/Api/Gp/Implementations

*
Implementations
*

.. currentmodule:: pymc.gp
.. autosummary::
:toctree: generated

HSGP
HSGPPeriodic
Latent
LatentKron
Marginal
MarginalKron
MarginalApprox
TP


Source/Api/Gp/Mean


Mean Functions

.. automodule:: pymc.gp.mean
.. autosummary::
:toctree: generated

Zero
Constant
Linear


Source/Api/Gp/Util


GP Utilities

.. automodule:: pymc.gp.util
.. autosummary::
:toctree: generated

plot_gp_dist


Source/Api/Model/Conditioning

Model Conditioning
------------------

.. currentmodule:: pymc.model.transform
.. autosummary::
:toctree: generated/

do
observe
change_value_transforms
remove_value_transforms


Source/Api/Model/Core

Model creation and inspection
-----------------------------

.. currentmodule:: pymc.model.core
.. autosummary::
:toctree: generated/

BaseModel
FrozenModel
Model
modelcontext

Others
------

.. currentmodule:: pymc.model.core
.. autosummary::
:toctree: generated/

Deterministic
Potential
set_data
Point
compile_fn

Graph visualization
-------------------

.. currentmodule:: pymc.model_graph
.. autosummary::
:toctree: generated/

model_to_graphviz
model_to_mermaid
model_to_networkx


Source/Api/Model/Deterministics

Deterministic Surgery
---------------------

.. currentmodule:: pymc.model.transform
.. autosummary::
:toctree: generated/

extract_deterministics
insert_deterministics


Source/Api/Model/Fgraph

FunctionGraph
-------------

.. currentmodule:: pymc.model.fgraph
.. autosummary::
:toctree: generated/

clone_model
fgraph_from_model
model_from_fgraph


Source/Api/Model/Optimization

Model Optimization
------------------
.. currentmodule:: pymc.model.transform
.. autosummary::
:toctree: generated/

freeze_dims_and_data
freeze_model


Source/Api/Model/Transform Values

Value Transforms
----------------

.. currentmodule:: pymc.model.transform_values
.. autosummary::
:toctree: generated/

constrain_values
unconstrain_values


Source/Api/Backends

Storage backends

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

to_inference_data
predictions_to_inference_data

Internal structures
-------------------

.. automodule:: pymc.backends

.. autosummary::
:toctree: generated/

NDArray
base.BaseTrace
base.MultiTrace
zarr.ZarrTrace
zarr.ZarrChain


Source/Api/Data

Data

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

Data
get_data
Minibatch


Source/Api/Dims

.. _api_dims:

Dims
====

.. warning:: This module is experimental and may contain critical breaks. API changes are expected in future releases.

This submodule contains functions for defining distributions and operations that use explicit dimensions.

The module is presented in :ref:dims_module.

.. toctree::

dims/model
dims/math
dims/distributions
dims/transforms


Source/Api/Distributions

.. _api_distributions:

*
Distributions
*

.. toctree::
:maxdepth: 2

distributions/continuous
distributions/discrete
distributions/multivariate
distributions/mixture
distributions/timeseries
distributions/truncated
distributions/censored
distributions/custom
distributions/simulator
distributions/transforms
distributions/utilities


Source/Api/Gp

Gaussian Processes
------------------

.. automodule:: pymc.gp

.. toctree::
:maxdepth: 2

gp/implementations
gp/mean
gp/cov
gp/util


Source/Api/Logprob

*
Probability
*

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

logp
logcdf
icdf

Conditional probability
-----------------------

.. currentmodule:: pymc.logprob

.. autosummary::
:toctree: generated/

conditional_logp
transformed_conditional_logp


Source/Api/Math

====
Math
====

This submodule contains various mathematical functions. Most are re-exported directly from
:mod:pytensor.tensor and :mod:pytensor.tensor.linalg (see there for full signatures and
details). Doing any kind of math with PyMC random variables, or defining custom likelihoods
or priors, requires you to use these PyTensor expressions rather than NumPy or Python code.

.. automodule:: pymc.math
:members:


Source/Api/Misc

Other utils
*

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

find_constrained_prior

.. currentmodule:: pymc.blocking

.. autosummary::
:toctree: generated/

DictToArrayBijection
model_table


Source/Api/Model

Model
------

.. automodule:: pymc.model

.. toctree::
:maxdepth: 2

model/core
model/conditioning
model/transform_values
model/deterministics
model/optimization
model/fgraph


Source/Api/Ode


Ordinary differential equations (ODEs)

.. automodule:: pymc.ode

.. autosummary::
:toctree: generated/

DifferentialEquation


Source/Api/Pytensorf

PyTensor utils

.. currentmodule:: pymc.pytensorf

.. autosummary::
:toctree: generated/

compile
gradient
hessian
hessian_diag
jacobian
inputvars
cont_inputs
floatX
intX
constant_fold
CallableTensor
join_nonshared_inputs
make_shared_replacements
convert_data


Source/Api/Samplers

Samplers
========

This submodule contains functions for MCMC and forward sampling.

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

sample
sample_prior_predictive
sample_posterior_predictive
draw
compute_deterministics
vectorize_over_posterior
init_nuts
sampling.jax.sample_blackjax_nuts
sampling.jax.sample_numpyro_nuts

Step methods

HMC family
----------
.. currentmodule:: pymc.step_methods.hmc

.. autosummary::
:toctree: generated/

NUTS
HamiltonianMC

Metropolis family
-----------------
.. currentmodule:: pymc.step_methods

.. autosummary::
:toctree: generated/

BinaryGibbsMetropolis
BinaryMetropolis
CategoricalGibbsMetropolis
CauchyProposal
DEMetropolis
DEMetropolisZ
LaplaceProposal
Metropolis
MultivariateNormalProposal
NormalProposal
PoissonProposal
UniformProposal

Other step methods
------------------
.. currentmodule:: pymc.step_methods

.. autosummary::
:toctree: generated/

CompoundStep
Slice


Source/Api/Shape Utils

*
shape_utils
*

This submodule contains various functions that apply numpy's broadcasting rules to shape tuples, and also to samples drawn from probability distributions.

The main challenge when broadcasting samples drawn from a generative model, is that each random variate has a core shape. When we draw many i.i.d samples from a given RV, for example if we ask for `size_tuple i.i.d draws, the result usually is a size_tuple + RV_core_shape`. In the generative model's hierarchy, the downstream RVs that are conditionally dependent on our above sampled values, will get an array with a shape that is inconsistent with the core shape they expect to see for their parameters. This is a problem sometimes because it prevents regular broadcasting in complex hierarchical models, and thus make prior and posterior predictive sampling difficult.

This module introduces functions that are made aware of the requested `size_tuple of i.i.d samples, and does the broadcasting on the core shapes, transparently ignoring or moving the i.i.d size_tuple` prepended axes around.

.. currentmodule:: pymc.distributions.shape_utils

.. autosummary::
:toctree: generated/

to_tuple
rv_size_is_none
change_dist_size


Source/Api/Smc

Sequential Monte Carlo

.. automodule:: pymc.smc

.. autosummary::
:toctree: generated/

sample_smc

.. _smc_kernels:

SMC kernels
-----------

.. currentmodule:: pymc.smc.kernels
.. autosummary::
:toctree: generated/

SMC_KERNEL
IMH
MH


Source/Api/Stats

Stats
*

.. currentmodule:: pymc.stats

.. autosummary::
:toctree: generated/

compute_log_prior
compute_log_likelihood

PyMC re-exports functions from the `arviz_stats library under the pymc.stats
namespace, allowing functions like
summary, ess, rhat, loo etc. to be
accessed as
pymc.stats.<function>. For the API documentation of those functions,
see the :doc:
arviz_stats documentation <arviz_stats:index>`.


Source/Api/Testing

=======
Testing
=======

This submodule contains tools to help with testing PyMC code.

.. currentmodule:: pymc.testing

.. autosummary::
:toctree: generated/

mock_sample
mock_sample_setup_and_teardown


Source/Api/Tuning

Tuning
------

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

find_hessian
find_MAP


Source/Api/Vi

*
Variational Inference
*

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

ADVI
ASVGD
SVGD
FullRankADVI
fit

.. currentmodule:: pymc.variational

.. autosummary::
:toctree: generated/

ImplicitGradient
Inference
KLqp

Approximations
--------------

.. currentmodule:: pymc

.. autosummary::
:toctree: generated/

Empirical
FullRank
MeanField
sample_approx

OPVI
----

.. autosummary::
:toctree: generated/

Group

.. currentmodule:: pymc.variational

.. autosummary::
:toctree: generated/

Approximation

Operators
---------

.. automodule:: pymc.variational.operators
.. autosummary::
:toctree: generated/

KL
KSD

Special
-------
.. currentmodule:: pymc.variational
.. autosummary::
:toctree: generated/

Stein

.. currentmodule:: pymc
.. autosummary::
:toctree: generated/

adadelta
adagrad
adagrad_window
adam
adamax
apply_momentum
apply_nesterov_momentum
momentum
nesterov_momentum
norm_constraint
rmsprop
sgd
total_norm_constraint


Source/Contributing/Build Docs

Build documentation locally

:::{warning}
Docs build is not supported on Windows.
To build docs on Windows we recommend running inside a Docker container.
:::

To build the docs, first install dependencies by running these commands at the PyMC repo root:

shell
# create the pymc-docs conda env, or equivalently make
# sure all dependencies listed in this file are installed
conda env create -f conda-envs/environment-docs.yml

# Install local pymc version in editable mode
pip install -e .

Building the documentation

There is a Makefile in the pymc repo to help with the doc building process.

shell
make clean
make html

make html is the command that builds the documentation with sphinx-build.
make clean deletes caches and intermediate files.

The make clean step is not always necessary. If you are working on a specific page,
for example, then you can rebuild the docs without the clean step, and everything should
work fine. If you are restructuring the content or editing toctrees, then you'll need
to execute make clean.

A good approach is generally to skip make clean, which makes
the make html faster, and see how everything looks. If something
looks strange, run make clean and make html one after the other
to see if it fixes the issue before checking anything else.

Emulate building on readthedocs

The target rtd is also available to chain make clean with sphinx-build
setting also some extra options and environment variables to instruct
sphinx to simulate a readthedocs build as much as possible.

shell
make rtd

:::{important}
This won't reinstall or update any dependencies, unlike on readthedocs where
all dependencies are installed in a clean env before each build.

But it will execute all notebooks inside the core_notebooks folder,
which by default are not executed. Executing the notebooks will add several minutes
to the doc build, as there are 6 notebooks which take between 20s to 5 minutes
to run.
:::

View the generated docs

shell
make view

This will use Python's webbrowser module to open the generated website on your browser.
The generated website is static, so there is no need to set a server to preview it.


Source/Contributing/Developer Guide


orphan: true


PyMC Developer Guide

{doc}PyMC is a Python package for Bayesian statistical modeling built on top of {doc}PyTensor .
This document aims to explain the design and implementation of probabilistic programming in PyMC, with comparisons to other PPLs like TensorFlow Probability (TFP) and Pyro.
A user-facing API introduction can be found in the {ref}API quickstart .
A more accessible, user facing deep introduction can be found in Peadar Coyle's probabilistic programming primer.

Distribution

Probability distributions in PyMC are implemented as classes that inherit from {class}~pymc.Continuous or {class}~pymc.Discrete.
Either of these inherit {class}~pymc.Distribution which defines the high level API.

For a detailed introduction on how a new distribution should be implemented check out the {ref}guide on implementing distributions .

Reflection

How tensor/value semantics for probability distributions are enabled in PyMC:

In PyMC, model variables are defined by calling probability distribution classes with parameters:

python
z = Normal("z", 0, 5)

This is done inside the context of a `pm.Model`, which intercepts some information, for example to capture known dimensions.
The notation aligns with the typically used math notation:

$$
z \sim \text{Normal}(0, 5)
$$

A call to a {class}~pymc.Distribution constructor as shown above returns a PyTensor {class}~pytensor.tensor.TensorVariable, which is a symbolic representation of the model variable and the graph of inputs it depends on.
Under the hood, the variables are created through the {meth}~pymc.Distribution.dist API, which calls the {class}~pytensor.tensor.random.basic.RandomVariable {class}~pytensor.graph.op.Op corresponding to the distribution.

At a high level of abstraction, the idea behind `RandomVariable Ops is to create symbolic variables (TensorVariables) that can be associated with the properties of a probability distribution.
For example, the
RandomVariable Op` which becomes part of the symbolic computation graph is associated with the random number generators or probability mass/density functions of the distribution.

In the example above, where the `TensorVariable z is created to be {math}\text{Normal}(0, 5) random variable, we can get a handle on the corresponding RandomVariable Op` instance:

python
with pm.Model():
    z = pm.Normal("z", 0, 5)
print(type(z.owner.op))
# ==> pytensor.tensor.random.basic.NormalRV
isinstance(z.owner.op, pytensor.tensor.random.basic.RandomVariable)
# ==> True

Now, because the `NormalRV can be associated with the probability density function of the Normal distribution, we can now evaluate it through the special pm.logp` function:

python
with pm.Model():
    z = pm.Normal("z", 0, 5)
symbolic = pm.logp(z, 2.5)
numeric = symbolic.eval()
# array(-2.65337645)

We can, of course, also work out the math by hand:

$$
\begin{aligned}
pdf_{\mathcal{N}}(\mu, \sigma, x) &= \frac{1}{\sigma \sqrt{2 \pi}} \exp^{- 0.5 (\frac{x - \mu}{\sigma})^2} \\
pdf_{\mathcal{N}}(0, 5, 2.5) &= 0.070413 \\
ln(0.070413) &= -2.6533
\end{aligned}
$$

In the probabilistic programming context, this enables PyMC and its backend PyTensor to create and evaluate computation graphs to compute, for example log-prior or log-likelihood values.

PyMC in Comparison

Within the PyMC model context, random variables are essentially PyTensor tensors that can be used in all kinds of operations as if they were NumPy arrays.
This is different compared to TFP and pyro, where one needs to be more explicit about the conversion from random variables to tensors.

Consider the following examples, which implement the below model.

$$
\begin{aligned}
z &\sim \mathcal{N}(0, 5) \\
x &\sim \mathcal{N}(z, 1) \\
\end{aligned}
$$

PyMC

python
with pm.Model() as model:
    z = pm.Normal('z', mu=0., sigma=5.)             # ==> pytensor.tensor.var.TensorVariable
    x = pm.Normal('x', mu=z, sigma=1., observed=5.) # ==> pytensor.tensor.var.TensorVariable
# The log-prior of z=2.5
pm.logp(z, 2.5).eval()                              # ==> -2.65337645
# ???????
x.logp({'z': 2.5})                                  # ==> -4.0439386
# ???????
model.logp({'z': 2.5})                              # ==> -6.6973152

Tensorflow Probability

python
import tensorflow.compat.v1 as tf
from tensorflow_probability import distributions as tfd

with tf.Session() as sess:
    z_dist = tfd.Normal(loc=0., scale=5.)            # ==> <class 'tfp.python.distributions.normal.Normal'>
    z = z_dist.sample()                              # ==> <class 'tensorflow.python.framework.ops.Tensor'>
    x = tfd.Normal(loc=z, scale=1.).log_prob(5.)     # ==> <class 'tensorflow.python.framework.ops.Tensor'>
    model_logp = z_dist.log_prob(z) + x
    print(sess.run(x, feed_dict={z: 2.5}))           # ==> -4.0439386
    print(sess.run(model_logp, feed_dict={z: 2.5}))  # ==> -6.6973152

Pyro

python
z_dist = dist.Normal(loc=0., scale=5.)           # ==> <class 'pyro.distributions.torch.Normal'>
z = pyro.sample("z", z_dist)                     # ==> <class 'torch.Tensor'>
# reset/specify value of z
z.data = torch.tensor(2.5)
x = dist.Normal(loc=z, scale=1.).log_prob(5.)    # ==> <class 'torch.Tensor'>
model_logp = z_dist.log_prob(z) + x
x                                                # ==> -4.0439386
model_logp                                       # ==> -6.6973152

Behind the scenes of the ``logp`` function

The `logp` function is straightforward - it is a PyTensor function within each distribution.
It has the following signature:

:::{warning}
The code block is outdated.
:::

python
def logp(self, value):
    # GET PARAMETERS
    param1, param2, ... = self.params1, self.params2, ...
    # EVALUATE LOG-LIKELIHOOD FUNCTION, all inputs are (or array that could be convert to) PyTensor tensor
    total_log_prob = f(param1, param2, ..., value)
    return total_log_prob

In the `logp method, parameters and values are either PyTensor tensors, or could be converted to tensors.
It is rather convenient as the evaluation of logp is represented as a tensor (
RV.logpt), and when we linked different logp together (e.g., summing all RVs.logpt to get the model total logp) the dependence is taken care of by PyTensor when the graph is built and compiled.
Again, since the compiled function depends on the nodes that already in the graph, whenever you want to generate a new function that takes new input tensors you either need to regenerate the graph with the appropriate dependencies, or replace the node by editing the existing graph.
In PyMC we use the second approach by using
pytensor.clone_replace()` when it is needed.

As explained above, distribution in a `pm.Model() context automatically turn into a tensor with distribution property (PyMC random variable).
To get the logp of a free\_RV is just evaluating the
logp()` on itself:

python
# self is a pytensor.tensor with a distribution attached
self.logp_sum_unscaledt = distribution.logp_sum(self)
self.logp_nojac_unscaledt = distribution.logp_nojac(self)

Or for an observed RV. it evaluate the logp on the data:

python
self.logp_sum_unscaledt = distribution.logp_sum(data)
self.logp_nojac_unscaledt = distribution.logp_nojac(data)

Model context and Random Variable

I like to think that the `with pm.Model() ...` is a key syntax feature and the signature of PyMC model language, and in general a great out-of-the-box thinking/usage of the context manager in Python (with some critics, of course).

Essentially what a context manager does is:

python
with EXPR as VAR:
    USERCODE

which roughly translates into this:

python
VAR = EXPR
VAR.__enter__()
try:
    USERCODE
finally:
    VAR.__exit__()

or conceptually:

python
with EXPR as VAR:
    # DO SOMETHING
    USERCODE
    # DO SOME ADDITIONAL THINGS

So what happened within the `with pm.Model() as model: ... block, besides the initial set up model = pm.Model()`?
Starting from the most elementary:

Random Variable

From the above session, we know that when we call e.g. `pm.Normal('x', ...)` within a Model context, it returns a random variable.
Thus, we have two equivalent ways of adding random variable to a model:

python
with pm.Model() as m:
    x = pm.Normal('x', mu=0., sigma=1.)

print(type(x))                              # ==> <class 'pytensor.tensor.var.TensorVariable'>
print(m.free_RVs)                           # ==> [x]
print(logpt(x, 5.0))                        # ==> Elemwise{switch,no_inplace}.0
print(logpt(x, 5.).eval({}))                # ==> -13.418938533204672
print(m.logp({'x': 5.}))                    # ==> -13.418938533204672

In general, if a variable has observations (`observed parameter), the RV is an observed RV, otherwise if it has a transformed (transform` parameter) attribute, it is a transformed RV otherwise, it will be the most elementary form: a free RV.
Note that this means that random variables with observations cannot be transformed.

Additional things that ``pm.Model`` does

In a way, `pm.Model` is a tape machine that records what is being added to the model, it keeps track the random variables (observed or unobserved) and potential term (additional tensor that to be added to the model logp), and also deterministic transformation (as bookkeeping):

  • named\_vars
  • free\_RVs
  • observed\_RVs
  • deterministics
  • potentials
  • missing\_values

The model context then computes some simple model properties, builds a bijection mapping that transforms between dictionary and numpy/PyTensor ndarray, thus allowing the `logp/dlogp functions to have two equivalent versions:
One takes a
dict as input and the other takes an ndarray as input.
More importantly, a
pm.Model()` contains methods to compile PyTensor functions that take Random Variables (that are also initialised within the same model) as input, for example:

python
with pm.Model() as m:
    z = pm.Normal('z', 0., 10., shape=10)
    x = pm.Normal('x', z, 1., shape=10)

print(m.initial_point)
print(m.dict_to_array(m.initial_point))  # ==> m.bijection.map(m.initial_point)
print(m.bijection.rmap(np.arange(20)))
# {'z': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), 'x': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])}
# [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
# {'z': array([10., 11., 12., 13., 14., 15., 16., 17., 18., 19.]), 'x': array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])}
python
list(filter(lambda x: "logp" in x, dir(pm.Model)))
#['d2logp',
# 'd2logp_nojac',
# 'datalogpt',
# 'dlogp',
# 'dlogp_array',
# 'dlogp_nojac',
# 'fastd2logp',
# 'fastd2logp_nojac',
# 'fastdlogp',
# 'fastdlogp_nojac',
# 'fastlogp',
# 'fastlogp_nojac',
# 'logp',
# 'logp_array',
# 'logp_dlogp_function',
# 'logp_elemwise',
# 'logp_nojac',
# 'logp_nojact',
# 'logpt',
# 'varlogpt']

Logp and dlogp

The model collects all the random variables (everything in `model.free_RVs and model.observed_RVs`) and potential term, and sum them together to get the model logp:

python
@property
def logpt(self):
    """PyTensor scalar of log-probability of the model"""
    with self:
        factors = [var.logpt for var in self.basic_RVs] + self.potentials
        logp = at.sum([at.sum(factor) for factor in factors])
        ...
        return logp

which returns a PyTensor tensor that its value depends on the free parameters in the model (i.e., its parent nodes from the PyTensor graph).
You can evaluate or compile into a python callable (that you can pass numpy as input args).
Note that the logp tensor depends on its input in the PyTensor graph, thus you cannot pass new tensor to generate a logp function.
For similar reason, in PyMC we do graph copying a lot using pytensor.clone_replace to replace the inputs to a tensor.

python
with pm.Model() as m:
    z = pm.Normal('z', 0., 10., shape=10)
    x = pm.Normal('x', z, 1., shape=10)
    y = pm.Normal('y', x.sum(), 1., observed=2.5)

print(m.basic_RVs)    # ==> [z, x, y]
print(m.free_RVs)     # ==> [z, x]

type(m.logp)
# pytensor.tensor.var.TensorVariable

m.logpt.eval(m.initial_point())
# array(-51.25369126)

PyMC then compiles a logp function with gradient that takes `model.free_RVs as input and model.logpt as output.
It could be a subset of tensors in
model.free_RVs` if we want a conditional logp/dlogp function:

python
def logp_dlogp_function(self, grad_vars=None, **kwargs):
    if grad_vars is None:
        grad_vars = list(typefilter(self.free_RVs, continuous_types))
    else:
        ...
    varnames = [var.name for var in grad_vars]  # In a simple case with only continuous RVs,
                                                # this is all the free_RVs
    extra_vars = [var for var in self.free_RVs if var.name not in varnames]
    return ValueGradFunction(self.logpt, grad_vars, extra_vars, **kwargs)

`ValueGradFunction is a callable class which isolates part of the PyTensor graph to compile additional PyTensor functions.
PyMC relies on
pytensor.clone_replace to copy the model.logpt` and replace its input.
It does not edit or rewrite the graph directly.

The important parts of the above function is highlighted and commented.
On a high level, it allows us to build conditional logp function and its gradient easily.
Here is a taste of how it works in action:

text
/* Detailed source-code truncated for AI context efficiency. */

So why is this necessary?
One can imagine that we just compile one logp function, and do bookkeeping ourselves.
For example, we can build the logp function in PyTensor directly:

python
import pytensor
func = pytensor.function(m.free_RVs, m.logpt)
func(*inputlist)
# array(-51.0769075)

logpt_grad = pytensor.grad(m.logpt, m.free_RVs)
func_d = pytensor.function(m.free_RVs, logpt_grad)
func_d(*inputlist)
# [array([ 0.74230226,  0.01658948,  1.38606194,  0.11253699, -1.07003284,
#          2.64302891,  1.12497754, -0.35967542, -1.18117557, -1.11489642]),
#  array([ 0.98281586,  1.69545542,  0.34626619,  1.61069443,  2.79155183,
#         -0.91020295,  0.60094326,  2.08022672,  2.8799075 ,  2.81681213])]

Similarly, build a conditional logp:

python
shared = pytensor.shared(inputlist[1])
func2 = pytensor.function([m.free_RVs[0]], m.logpt, givens=[(m.free_RVs[1], shared)])
print(func2(inputlist[0]))
# -51.07690750130328

logpt_grad2 = pytensor.grad(m.logpt, m.free_RVs[0])
func_d2 = pytensor.function([m.free_RVs[0]], logpt_grad2, givens=[(m.free_RVs[1], shared)])
print(func_d2(inputlist[0]))
# [ 0.74230226  0.01658948  1.38606194  0.11253699 -1.07003284  2.64302891
#   1.12497754 -0.35967542 -1.18117557 -1.11489642]

The above also gives the same logp and gradient as the output from `model.logp_dlogp_function`.
But the difficulty is to compile everything into a single function:

python
func_logp_and_grad = pytensor.function(m.free_RVs, [m.logpt, logpt_grad])
# ==> ERROR

We want to have a function that return the evaluation and its gradient re each input:
`value, grad = f(x), but the naive implementation does not work.
We can of course wrap 2 functions - one for logp one for dlogp - and output a list.
But that would mean we need to call 2 functions.
In addition, when we write code using python logic to do bookkeeping when we build our conditional logp.
Using
pytensor.clone_replace`, we always have the input to the PyTensor function being a 1d vector (instead of a list of RV that each can have very different shape), thus it is very easy to do matrix operation like rotation etc.

Notes

The current setup is quite powerful, as the PyTensor compiled function is fairly fast to compile and to call.
Also, when we are repeatedly calling a conditional logp function, external RV only need to reset once.
However, there are still significant overheads when we are passing values between PyTensor graph and NumPy.
That is the reason we often see no advantage in using GPU, because the data is copying between GPU and CPU at each function call - and for a small model, the result is a slower inference under GPU than CPU.

Also, `pytensor.clone_replace` is too convenient (PyMC internal joke is that it is like a drug - very addictive).
If all the operation happens in the graph (including the conditioning and setting value), I see no need to isolate part of the graph (via graph copying or graph rewriting) for building model and running inference.

Moreover, if we are limiting to the problem that we can solved most confidently - model with all continuous unknown parameters that could be sampled with dynamic HMC, there is even less need to think about graph cloning/rewriting.

Inference

MCMC

The ability for model instance to generate conditional logp and dlogp function enable one of the unique feature of PyMC - {class}~pymc.step_methods.CompoundStep method.
On a conceptual level it is a Metropolis-within-Gibbs sampler.
Users can specify different sampler for different RVs.
Alternatively, it is implemented as yet another interceptor:
The `pm.sample(...)` call will try to assign the best step methods to different free\_RVs (e.g., NUTS if all free\_RVs are continuous).
Then, (conditional) logp function(s) are compiled, and the sampler called each sampler within the list of CompoundStep in a for-loop for one sample circle.

For each sampler, it implements a `step.step method to perform MH updates.
Each time a dictionary (
point in PyMC land, same structure as model.initial_point) is passed as input and output a new dictionary with the free\_RVs being sampled now has a new value (if accepted, see here and here).
There are some example in the
CompoundStep` doc:

Transition kernel

The base class for most MCMC sampler (except SMC) is in ArrayStep.
You can see that the `step.step() is mapping the point into an array, and call self.astep(), which is an array in, array out function.
A PyMC model compiles a conditional logp/dlogp function that replace the input RVs with a shared 1D tensor (flatten and stack view of the original RVs).
And the transition kernel (i.e.,
.astep()`) takes an array as input and outputs an array.
For example, see the MH sampler.

This is of course very different compared to the transition kernel in e.g. TFP, which is a tenor in tensor out function.
Moreover, transition kernels in TFP do not flatten the tensors, see eg docstring of tensorflow\_probability/python/mcmc/random\_walk\_metropolis.py:

python
new_state_fn: Python callable which takes a list of state parts and a
        seed; returns a same-type `list` of `Tensor`s, each being a perturbation
        of the input state parts. The perturbation distribution is assumed to be
        a symmetric distribution centered at the input state part.
        Default value: `None` which is mapped to
            `tfp.mcmc.random_walk_normal_fn()`.

Dynamic HMC

We love NUTS, or to be more precise Dynamic HMC with complex stopping rules.
This part is actually all done outside of PyTensor, for NUTS, it includes:
The leapfrog, dual averaging, tuning of mass matrix and step size, the tree building, sampler related statistics like divergence and energy checking.
We actually have a PyTensor version of HMC, but it has never been used, and has been removed from the main repository.
It can still be found in the git history, though.

Variational Inference (VI)

The design of the VI module takes a different approach than MCMC - it has a functional design, and everything is done within PyTensor (i.e., Optimization and building the variational objective).
The base class of variational inference is pymc.variational.Inference, where it builds the objective function by calling:

python
...
self.objective = op(approx, **kwargs)(tf)
...

Where:

text
op     : Operator class
approx : Approximation class or instance
tf     : TestFunction instance
kwargs { kwargs passed to :class}`Operator`

The design is inspired by the great work Operator Variational Inference.
`Inference` object is a very high level of VI implementation.
It uses primitives: Operator, Approximation, and Test functions to combine them into single objective function.
Currently we do not care too much about the test function, it is usually not required (and not implemented).
The other primitives are defined as base classes in this file.
We use inheritance to easily implement a broad class of VI methods leaving a lot of flexibility for further extensions.

For example, consider ADVI.
We know that in the high-level, we are approximating the posterior in the latent space with a diagonal Multivariate Gaussian.
In another word, we are approximating each elements in `model.free_RVs` with a Gaussian.
Below is what happen in the set up:

python
def __init__(self, *args, **kwargs):
    super(ADVI, self).__init__(MeanField(*args, **kwargs))
# ==> In the super class KLqp
    super(KLqp, self).__init__(KL, MeanField(*args, **kwargs), None, beta=beta)
# ==> In the super class Inference
    ...
    self.objective = KL(MeanField(*args, **kwargs))(None)
    ...

where `KL` is Operator based on Kullback Leibler Divergence (it does not need any test function).

python
...
def apply(self, f):
    return -self.datalogp_norm + self.beta * (self.logq_norm - self.varlogp_norm)

Since the logp and logq are from the approximation, let's dive in further on it (there is another abstraction here - `Group - that allows you to combine approximation into new approximation, but we will skip this for now and only consider SingleGroupApproximation like MeanField):
The definition of
datalogp_norm, logq_norm, varlogp_norm are in variational/opvi, strip away the normalizing term, datalogp and varlogp are expectation of the variational free\_RVs and data logp - we clone the datalogp and varlogp from the model, replace its input with PyTensor tensor that samples from the variational posterior.
For ADVI, these samples are from a Gaussian.
Note that the samples from the posterior approximations are usually 1 dimension more, so that we can compute the expectation and get the gradient of the expectation (by computing the expectation of the gradient!).
As for the
logq` since it is a Gaussian it is pretty straightforward to evaluate.

##### Some challenges and insights from implementing VI.

  • Graph based approach was helpful, but PyTensor had no direct access to previously created nodes in the computational graph.

You can find a lot of `@node_property` usages in implementation.
This is done to cache nodes.
TensorFlow has graph utils for that that could potentially help in doing this.
On the other hand graph management in Tensorflow seemed to more tricky than expected.
The high level reason is that graph is an add only container.

  • There were few fixed bugs not obvious in the first place.

PyTensor has a tool to manipulate the graph (`pytensor.clone_replace`) and this tool requires extremely careful treatment when doing a lot of graph replacements at different level.

  • We coined a term `pytensor.clone_replace` curse.

We got extremely dependent on this feature.
Internal usages are uncountable:
- We use this to vectorize the model for both MCMC and VI to speed up computations
- We use this to create sampling graph for VI. This is the case you want posterior predictive as a part of computational graph.

As this is the core of the VI process, we were trying to replicate this pattern in TF.
However, when `pytensor.clone_replace` is called, PyTensor creates a new part of the graph that can be collected by garbage collector, but TF's graph is add only.
So we should solve the problem of replacing input in a different way.

Forward sampling

As explained above, in distribution we have method to walk the model dependence graph and generate forward random sample in scipy/numpy.
This allows us to do prior predictive samples using pymc.sampling.sample_prior_predictive see code.
It is a fairly fast batch operation, but we have quite a lot of bugs and edge case especially in high dimensions.
The biggest pain point is the automatic broadcasting.
As in the batch random generation, we want to generate (n\_sample, ) + RV.shape random samples.
In some cases, where we broadcast RV1 and RV2 to create a RV3 that has one more batch shape, we get error (even worse, wrong answer with silent error).

The good news is, we are fixing these errors with the amazing works from lucianopaz and others.
The challenge and some summary of the solution could be found in Luciano's blog post

python
with pm.Model() as m:
    mu = pm.Normal('mu', 0., 1., shape=(5, 1))
    sigma = pm.HalfNormal('sigma', 5., shape=(1, 10))
    pm.Normal('x', mu=mu, sigma=sigma, observed=np.random.randn(2, 5, 10))
    trace = pm.sample_prior_predictive(100)

trace['x'].shape # ==> should be (100, 2, 5, 10)

pm.Normal.dist(mu=np.zeros(2), sigma=1).random(size=(10, 4))

There are also other error related random sample generation (e.g., Mixture is currently broken).

Extending PyMC

  • Custom Inference method

- Inferencing Linear Mixed Model with EM.ipynb
- Laplace approximation in pymc.ipynb

  • Connecting it to other library within a model

- Using "black box" likelihood function by creating a custom PyTensor Op.
- Using emcee

  • Using other library for inference

- Connecting to Julia for solving ODE (with gradient for solution that can be used in NUTS)

What we got wrong

Shape

One of the pain point we often face is the issue of shape.
The approach in TFP and pyro is currently much more rigorous.
Adrian’s PR (https://github.com/pymc-devs/pymc/pull/2833) might fix this problem, but likely it is a huge effort of refactoring.
I implemented quite a lot of patches for mixture distribution, but still they are not done very naturally.

Random methods in numpy

There is a lot of complex logic for sampling from random variables, and because it is all in Python, we can't transform a sampling graph further.
Unfortunately, PyTensor does not have code to sample from various distributions and we didn't want to write that our own.

Samplers are in Python

While having the samplers be written in Python allows for a lot of flexibility and intuitive for experiment (writing e.g. NUTS in PyTensor is also very difficult), it comes at a performance penalty and makes sampling on the GPU very inefficient because memory needs to be copied for every logp evaluation.


Source/Contributing/Docker Container

(docker_container)=

Running PyMC in Docker

We have provided a Dockerfile which helps for isolating build problems, and local development.
Install Docker for your operating system, clone this repo, then
run the following commands to build a pymc docker image.

bash
cd pymc
bash scripts/docker_container.sh build

After successfully building the docker image, you can start a local docker container called pymc either from bash or from jupyter notebook server running on port 8888.

bash
bash scripts/docker_container.sh bash # running the container with bash
bash scripts/docker_container.sh jupyter # running the container with jupyter notebook

Source/Contributing/Implementing Distribution

(implementing-a-distribution)=

Implementing a RandomVariable Distribution

This guide provides an overview on how to implement a distribution for PyMC.
It is designed for developers who wish to add a new distribution to the library.
Users will not be aware of all this complexity and should instead make use of helper methods such as ~pymc.CustomDist.

PyMC {class}~pymc.Distribution builds on top of PyTensor's {class}~pytensor.tensor.random.op.RandomVariable, and implements logp, logcdf, icdf and support_point methods as well as other initialization and validation helpers.
Most notably shape/dims/observed kwargs, alternative parametrizations, and default transform.

Here is a summary check-list of the steps needed to implement a new distribution.
Each section will be expanded below:

  1. Creating a new RandomVariable Op
  2. Implementing the corresponding Distribution class
  3. Adding tests for the new RandomVariable
  4. Adding tests for logp / logcdf / icdf and support_point methods
  5. Documenting the new Distribution.

This guide does not attempt to explain the rationale behind the Distributions current implementation, and details are provided only insofar as they help to implement new "standard" distributions.

1. Creating a new `RandomVariable` `Op`

{class}~pytensor.tensor.random.op.RandomVariable are responsible for implementing the random sampling methods.
The RandomVariable is also responsible for parameter broadcasting and shape inference.

Before creating a new RandomVariable make sure that it is not already offered in the {mod}NumPy library <numpy.random>.
If it is, it should be added to the {doc}PyTensor library first and then imported into the PyMC library.

In addition, it might not always be necessary to implement a new RandomVariable.
For example if the new Distribution is just a special parametrization of an existing Distribution.
This is the case of the OrderedLogistic and OrderedProbit, which are just special parametrizations of the Categorical distribution.

The following snippet illustrates how to create a new RandomVariable:

text
/* Detailed source-code truncated for AI context efficiency. */

Some important things to keep in mind:

  1. Everything inside the rng_fn method is pure Python code (as are the inputs) and should __not__ make use of other PyTensor symbolic ops. The random method should make use of the rng which is a NumPy {class}~numpy.random.RandomGenerator, so that samples are reproducible.
  2. Non-default RandomVariable dimensions will end up in the rng_fn via the size kwarg. The rng_fn will have to take this into consideration for correct output. size is the specification used by NumPy and SciPy and works like PyMC shape for univariate distributions, but is different for multivariate distributions. For multivariate distributions the __size excludes the support dimensions__, whereas the __shape of the resulting TensorVariable or ndarray includes the support dimensions__. For more context check {ref}The dimensionality notebook <dimensionality>.
  3. PyTensor can automatically infer the output shape of univariate RandomVariables. For multivariate distributions, the method _supp_shape_from_params must be implemented in the new RandomVariable class. This method returns the support dimensionality of an RV given its parameters. In some cases this can be derived from the shape of one of its parameters, in which case the helper {func}pytensor.tensor.random.utils.supp_shape_from_ref_param_shape cand be used as is in {class}~pymc.DirichletMultinomialRV. In other cases the argument values (and not their shapes) may determine the support shape of the distribution, as happens in the ~pymc.distributions.multivarite._LKJCholeskyCovRV. In simpler cases they may be constant.
  4. It's okay to use the rng_fn classmethods of other PyTensor and PyMC RandomVariables inside the new rng_fn. For example if you are implementing a negative HalfNormal RandomVariable, your rng_fn can simply return - halfnormal.rng_fn(rng, scale, size).

Note: In addition to size, the PyMC API also provides shape, dims and observed as alternatives to define a distribution dimensionality, but this is taken care of by {class}~pymc.Distribution, and should not require any extra changes.

For a quick test that your new RandomVariable Op is working, you can call the Op with the necessary parameters and then call {class}~pymc.draw on the returned object:

python
# blah = pytensor.tensor.random.uniform in this example
# multiple calls with the same seed should return the same values
pm.draw(blah([0, 0], [1, 2], size=(10, 2)), random_seed=1)

# array([[0.83674527, 0.76593773],
#    [0.00958496, 1.85742402],
#    [0.74001876, 0.6515534 ],
#    [0.95134629, 1.23564938],
#    [0.41460156, 0.33241175],
#    [0.66707807, 1.62134924],
#    [0.20748312, 0.45307477],
#    [0.65506507, 0.47713784],
#    [0.61284429, 0.49720329],
#    [0.69325978, 0.96272673]])

2. Inheriting from a PyMC base `Distribution` class

After implementing the new RandomVariable Op, it's time to make use of it in a new PyMC {class}~pymc.Distribution.
PyMC works in a very {term}functional <Functional Programming> way, and the distribution classes are there mostly to add PyMC API features and keep related methods organized together.
In practice, they take care of:

  1. Linking ({term}Dispatching) an rv_op class with the corresponding support_point, logp, logcdf and icdf methods.
  2. Defining a standard transformation (for continuous distributions) that converts a bounded variable domain (e.g., positive line) to an unbounded domain (i.e., the real line), which many samplers prefer.
  3. Validating the parametrization of a distribution and converting non-symbolic inputs (i.e., numeric literals or NumPy arrays) to symbolic variables.
  4. Converting multiple alternative parametrizations to the standard parametrization that the RandomVariable is defined in terms of.

Here is how the example continues:

text
/* Detailed source-code truncated for AI context efficiency. */

Some notes:

  1. A distribution should at the very least inherit from {class}~pymc.Discrete or {class}~pymc.Continuous. For the latter, more specific subclasses exist: PositiveContinuous, UnitContinuous, BoundedContinuous, CircularContinuous, SimplexContinuous, which specify default transformations for the variables. If you need to specify a one-time custom transform you can also create a _default_transform dispatch function as is done for the {class}~pymc.distributions.multivariate.LKJCholeskyCov.
  2. If a distribution does not have a corresponding rng_fn implementation, a RandomVariable should still be created to raise a NotImplementedError. This is, for example, the case in {class}~pymc.distributions.continuous.Flat. In this case it will be necessary to provide a support_point method, because without a rng_fn, PyMC can't fall back to a random draw to use as an initial point for MCMC.
  3. As mentioned above, PyMC works in a very {term}functional <Functional Programming> way, and all the information that is needed in the logp, logcdf, icdf and support_point methods is expected to be "carried" via the RandomVariable inputs. You may pass numerical arguments that are not strictly needed for the rng_fn method but are used in the those methods. Just keep in mind whether this affects the correct shape inference behavior of the RandomVariable.
  4. The logcdf, and icdf methods is not a requirement, but it's a nice plus!
  5. Currently, only one moment is supported in the support_point method, and probably the "higher-order" one is the most useful (that is mean > median > mode)... You might need to truncate the moment if you are dealing with a discrete distribution. support_point should return a valid point for the random variable (i.e., it always has non-zero probability when evaluated at that point)
  6. When creating the support_point method, be careful with size != None and broadcast properly also based on parameters that are not necessarily used to calculate the moment. For example, the sigma in pm.Normal.dist(mu=0, sigma=np.arange(1, 6)) is irrelevant for the moment, but may nevertheless inform about the shape. In this case, the support_point should return [mu, mu, mu, mu, mu].

For a quick check that things are working you can try the following:

python
import pymc as pm
from pymc.distributions.distribution import support_point

# pm.blah = pm.Normal in this example
blah = pm.blah.dist(mu=0, sigma=1)

# Test that the returned blah_op is still working fine
pm.draw(blah, random_seed=1)
# array(-1.01397228)

# Test the support_point method
support_point(blah).eval()
# array(0.)

# Test the logp method
pm.logp(blah, [-0.5, 1.5]).eval()
# array([-1.04393853, -2.04393853])

# Test the logcdf method
pm.logcdf(blah, [-0.5, 1.5]).eval()
# array([-1.17591177, -0.06914345])

3. Adding tests for the new `RandomVariable`

Tests for new RandomVariables are mostly located in tests/distributions/test_*.py.
Most tests can be accommodated by the default BaseTestDistributionRandom class, which provides default tests for checking:

  1. Expected inputs are passed to the rv_op by the dist classmethod, via check_pymc_params_match_rv_op
  2. Expected (exact) draws are being returned, via check_pymc_draws_match_reference
  3. Shape variable inference is correct, via check_rv_size
python
from pymc.testing import BaseTestDistributionRandom, seeded_scipy_distribution_builder


class TestBlah(BaseTestDistributionRandom):
    pymc_dist = pm.Blah
    # Parameters with which to test the blah pymc Distribution
    pymc_dist_params = {"param1": 0.25, "param2": 2.0}
    # Parameters that are expected to have passed as inputs to the RandomVariable op
    expected_rv_op_params = {"param1": 0.25, "param2": 2.0}
    # If the new `RandomVariable` is simply calling a `numpy`/`scipy` method,
    # we can make use of `seeded_[scipy|numpy]_distribution_builder` which
    # will prepare a seeded reference distribution for us.
    reference_dist_params = {"mu": 0.25, "loc": 2.0}
    reference_dist = seeded_scipy_distribution_builder("blah")
    tests_to_run = [
        "check_pymc_params_match_rv_op",
        "check_pymc_draws_match_reference",
        "check_rv_size",
    ]

Additional tests should be added for each optional parametrization of the distribution.
In this case it's enough to include the test check_pymc_params_match_rv_op since only this differs.

Make sure the tested alternative parameter value would lead to a different value for the associated default parameter.
For instance, if it's just the inverse, testing with 1.0 is not very informative, since the conversion would return 1.0 as well, and we can't be (as) sure that is working correctly.

python
class TestBlahAltParam2(BaseTestDistributionRandom):

    pymc_dist = pm.Blah
    # param2 is equivalent to 1 / alt_param2
    pymc_dist_params = {"param1": 0.25, "alt_param2": 4.0}
    expected_rv_op_params = {"param1": 0.25, "param2": 2.0}
    tests_to_run = ["check_pymc_params_match_rv_op"]

Custom tests can also be added to the class as is done for the {class}~tests.distributions.test_continuous.TestFlat.

Note on `check_rv_size` test:

Custom input sizes (and expected output shapes) can be defined for the check_rv_size test, by adding the optional class attributes sizes_to_check and sizes_expected:

python
sizes_to_check = [None, (1), (2, 3)]
sizes_expected = [(3,), (1, 3), (2, 3, 3)]
tests_to_run = ["check_rv_size"]

This is usually needed for Multivariate distributions.
You can see an example in {class}~tests.distributions.test_multivariate.TestDirichlet.

Notes on `check_pymcs_draws_match_reference` test

The check_pymcs_draws_match_reference is a very simple test for the equality of draws from the RandomVariable and the exact same python function, given the same inputs and random seed.
A small number (size=15) is checked. This is not supposed to be a test for the correctness of the random number generator.
The latter kind of test (if warranted) can be performed with the aid of pymc_random and pymc_random_discrete methods, which will perform an expensive statistical comparison between the RandomVariable.rng_fn and a reference Python function.
This kind of test only makes sense if there is a good independent generator reference (i.e., not just the same composition of NumPy / SciPy calls that is done inside rng_fn).

Finally, when your rng_fn is doing something more than just calling a NumPy or SciPy method, you will need to set up an equivalent seeded function with which to compare for the exact draws (instead of relying on seeded_[scipy|numpy]_distribution_builder).
You can find an example in {class}~tests.distributions.test_continuous.TestWeibull, whose rng_fn returns beta * np.random.weibull(alpha, size=size).

4. Adding tests for the `logp` / `logcdf` / `icdf` methods

Tests for the logp, logcdf and icdf mostly make use of the helpers check_logp, check_logcdf, check_icdf and
check_selfconsistency_discrete_logcdf implemented in ~testing

python
from pymc.testing import Domain, check_logp, check_logcdf, select_by_precision

R = Domain([-np.inf, -2.1, -1, -0.01, 0.0, 0.01, 1, 2.1, np.inf])
Rplus = Domain([0, 0.01, 0.1, 0.9, 0.99, 1, 1.5, 2, 100, np.inf])


def test_blah():
    check_logp(
        pymc_dist=pm.Blah,
        # Domain of the distribution values
        domain=R,
        # Domains of the distribution parameters
        paramdomains={"mu": R, "sigma": Rplus},
        # Reference scipy (or other) logp function
        scipy_logp=lambda value, mu, sigma: sp.norm.logpdf(value, mu, sigma),
        # Number of decimal points expected to match between the pymc and reference functions
        decimal=select_by_precision(float64=6, float32=3),
        # Maximum number of combinations of domain * paramdomains to test
        n_samples=100,
    )

    check_logcdf(
        pymc_dist=pm.Blah,
        domain=R,
        paramdomains={"mu": R, "sigma": Rplus},
        scipy_logcdf=lambda value, mu, sigma: sp.norm.logcdf(value, mu, sigma),
        decimal=select_by_precision(float64=6, float32=1),
        n_samples=-1,
    )

These methods will perform a grid evaluation on the combinations of domain and paramdomains values, and check that the PyMC methods and the reference functions match.
There are a couple of details worth keeping in mind:

  1. By default, the first and last values (edges) of the Domain are not compared (they are used for other things). If it is important to test the edge of the Domain, the edge values can be repeated. This is done by the Bool: Bool = Domain([0, 0, 1, 1], "int64")
  2. There are some default domains (such as R and Rplus) that you can use for testing your new distribution, but it's also perfectly fine to create your own domains inside the test function if there is a good reason for it (e.g., when the default values lead too many extreme unlikely combinations that are not very informative about the correctness of the implementation).
  3. By default, a random subset of 100 param x paramdomain combinations is tested, to keep the test runtime under control. When testing your shiny new distribution, you can temporarily set n_samples=-1 to force all combinations to be tested. This is important to avoid your PR leading to surprising failures in future runs whenever some bad combinations of parameters are randomly tested.
  4. On GitHub some tests run twice, under the pytensor.config.floatX flags of "float64" and "float32". However, the reference Python functions will run in a pure "float64" environment, which means the reference and the PyMC results can diverge quite a lot (e.g., underflowing to -np.inf for extreme parameters). You should therefore make sure you test locally in both regimes. A quick and dirty way of doing this is to temporarily add pytensor.config.floatX = "float32" at the very top of file, immediately after import pytensor. Remember to set n_samples=-1 as well to test all combinations. The test output will show what exact parameter values lead to a failure. If you are confident that your implementation is correct, you may opt to tweak the decimal precision with select_by_precision, or adjust the tested Domain values. In extreme cases, you can mark the test with a conditional xfail (if only one of the sub-methods is failing, they should be separated, so that the xfail is as narrow as possible):
python
def test_blah_logp(self):
    ...


@pytest.mark.xfail(
   condition=(pytensor.config.floatX == "float32"),
   reason="Fails on float32 due to numerical issues",
)
def test_blah_logcdf(self):
    ...

5. Adding tests for the `support_point` method

Tests for the support_point make use of the function assert_support_point_is_expected
which checks if:

  1. Moments return the expected values
  2. Moments have the expected size and shape
  3. Moments have a finite logp
python
import pytest
from pymc.distributions import Blah
from pymc.testing import assert_support_point_is_expected


@pytest.mark.parametrize(
    "param1, param2, size, expected",
    [
        (0, 1, None, 0),
        (0, np.ones(5), None, np.zeros(5)),
        (np.arange(5), 1, None, np.arange(5)),
        (np.arange(5), np.arange(1, 6), (2, 5), np.full((2, 5), np.arange(5))),
    ],
)
def test_blah_support_point(param1, param2, size, expected):
    with Model() as model:
        Blah("x", param1=param1, param2=param2, size=size)
    assert_support_point_is_expected(model, expected)

Here are some details worth keeping in mind:

  1. In the case where you have to manually broadcast the parameters with each other it's important to add test conditions that would fail if you were not to do that. A straightforward way to do this is to make the used parameter a scalar, the unused one(s) a vector (one at a time) and size None.
  2. In other words, make sure to test different combinations of size and broadcasting to cover these cases.

6. Documenting the new `Distribution`

New distributions should have a rich docstring, following the same format as that of previously implemented distributions.
It generally looks something like this:

python
r"""Univariate blah distribution.

 The pdf of this distribution is

 .. math::

    f(x \mid \param1, \param2) = \exp{x * (param1 + \log{param2})}

 .. plot::

     import matplotlib.pyplot as plt
     import numpy as np
     import scipy.stats as st
     import arviz as az
     x = np.linspace(-5, 5, 1000)
     params1 = [0., 0., 0., -2.]
     params2 = [0.4, 1., 2., 0.4]
     for param1, param2 in zip(params1, params2):
         pdf = st.blah.pdf(x, param1, param2)
         plt.plot(x, pdf, label=r'$\param1$ = {}, $\param2$ = {}'.format(param1, param2))
     plt.xlabel('x', fontsize=12)
     plt.ylabel('f(x)', fontsize=12)
     plt.legend(loc=1)
     plt.show()

 ========  ==========================================
 Support   :math:`x \in [0, \infty)`
 ========  ==========================================

 Blah distribution can be parameterized either in terms of param2 or
 alt_param2. The link between the two parametrizations is
 given by

 .. math::

    \param2 = \dfrac{1}{\alt_param2}


 Parameters
 ----------
 param1: float
     Interpretation of param1.
 param2: float
     Interpretation of param2 (param2 > 0).
 alt_param2: float
     Interpretation of alt_param2 (alt_param2 > 0) (alternative to param2).

 Examples
 --------
 .. code-block:: python

     with pm.Model():
         x = pm.Blah('x', param1=0, param2=10)
 """

The new distribution should be referenced in the respective API page in the docs module (e.g., pymc/docs/api/distributions.continuous.rst).
If appropriate, a new notebook example should be added to pymc-examples illustrating how this distribution can be used and how it relates (and/or differs) from other distributions that users are more likely to be familiar with.


Source/Contributing/Index

Contributing

PyMC is an open source, collective effort.
There are many ways in which you can help make it better.
And all of them are welcome!

Contribute as an individual

PyMC is a joint effort of many people, each contributing to the areas they like
and have some expertise in, coordinating to try and cover all tasks.

Coding and documentation are the most common types of contributions, but
there are many more things that you can do to help PyMC which are just as
important. Moreover, both code and docs require submitting PRs via GitHub
to some of the repositories under the pymc-devs organization, and
while we have a {ref}pr_tutorial guide available, GitHub might not be
everyone's cup of tea. If that is your case, don't worry, you will be
more than welcome if you want to help.

:::{tip}
Contact us on Discourse if you want to contribute to the project but are not sure where you can contribute or how to start.

We also host office hours regularly to provide more support, especially to contributors.
If you are interested in participating subscribe to the office-hours tag on Discourse.
:::

Below there are some examples of non code nor doc contributions that could serve as an inspiration.
If you have other ideas let us know on Discourse to see if we can make it happen too.

  • Report a bug or make a suggestion for improvement by opening an issue in Github
  • Answer questions on Discourse
  • Teach about PyMC and advertise best practices by writing blogs or giving talks
  • Help plan PyMCon
  • Help with outreach and marketing. This could include for example reaching out to potential sponsor

companies, to people who could use PyMC in their work or making sure that academics who use PyMC
cite it correctly in their work

Contribute via Pull Requests on GitHub

We have a {ref}pr_tutorial and a {ref}pr_checklist page to help in all the steps of the contributing
process, from before your first ever contribution to regular contributions as a core contributor.

(pr_etiquette)=

Etiquette for code contributions

  • When you start working working on an issue, open a Draft pull request as soon as you make your first commit (see {ref}pr_tutorial).
  • Before opening a PR with a new feature, please make a proposal by opening an issue or Discussion with the maintainers. Depending on the proposal we might direct you to other places such as pymc-experimental or pymc-examples.
  • Any issue without an open pull request is available for work.

* If a pull request has no recent activity it may be closed, or taken over by someone else.
* The specific timeframe for "recent" is hard to define as it depends on the contributor the specific code change, and other contextual factors. As a rule of thumb in a normal pull request with no other blockers there is typically activity every couple of days.
* The core devs will make their best judgement when opting to close PRs or reassign them to others.

  • If unsure if an issue ticket is available feel free to ask in the issue ticket. Note however, that per the previous point an open pull request is way to claim an issue ticket. Please do not make unrealistic pledges in the issue tickets.
  • It's okay if you are delayed or need to take a break, but please leave a comment in the pull request if you cannot get it to a state where it can be merged. Depending on the change (urgent bugfix vs. new feature) the core devs can determine if the PR needs to be reassigned to get the work done.

Join the discussion or submit a solution for an open issue. See open issues

See all open issues in documentation here

:::{admonition} New to the open source space?
:class: tip

If you are not sure where or how to start, take a look at the sprint materials
(even if you plan on contributing on your own outside of sprint events).
They are the most detailed guide available on contributing to PyMC also
with advice on which contributions are good starting points.
:::

Contribute as an institution

Institutions can contribute in the following ways:

Contact PyMC at [email protected] for more information.

:::{toctree}
:hidden:
:maxdepth: 1
:caption: Tutorials

pr_tutorial
:::

:::{toctree}
:hidden:
:maxdepth: 1
:caption: How-to guides

build_docs
docker_container
running_the_test_suite
review_pr_pymc_examples
using_gitpod
implementing_distribution
:::

:::{toctree}
:hidden:
:maxdepth: 1
:caption: Reference content

python_style
jupyter_style
pr_checklist
release_checklist
:::

:::{toctree}
:hidden:
:maxdepth: 1
:caption: In depth explanations

versioning_schemes_explanation
:::


Source/Contributing/Jupyter Style

(jupyter_style)=

Jupyter Style Guide

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 <myst:index> and {doc}MYST-NB <myst-nb:index>
to parse the notebooks.

:::{tip}
There is a webinar available on contributing to the PyMC example gallery
:::

Template Notebook

There is a template Jupyter notebook to be used for new 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.
  • 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.

When linking to other notebooks, always use a ref type cross-reference pointing to the target in the {ref}jupyter_style_first_cell.

:::{caution}
Using urls links breaks self referencing in versioned docs! And at the same time they are
less robust than sphinx cross-references.
:::

::::{dropdown} Examples of cross-references

References to targets within the current project

That is, notebooks in pymc-examples referring to other notebooks in pymc-examples.

Pattern:

text
{ref}`explicit text <anchor_id>`

Example source:

text
{ref}`Kronecker product <GP-Kron>`

Rendered example: {ref}Kronecker product <GP-Kron>

References to targets of other projects

Here "other projects" means any sphinx documentation site that was build independently
of the current one. Therefore, this includes linking to pymc-examples notebooks
from the pymc documentation or vice versa, or linking to other libraries like
arviz, numpy, matplotlib...

Pattern:

text
{ref}`explicit text <key:anchor_id>`

Example source:

text
{ref}`how to use InferenceData <arviz:working_with_InferenceData>`

Rendered example: {ref}how to use InferenceData <arviz:working_with_InferenceData>

where key in the pattern (arviz in the example) is one of the keys defined in
the intersphinx_mapping variable of conf.py such as arviz, numpy, mpl...
For the main pymc repo it is located in docs/source/conf.py, for pymc-examples it is
in examples/conf.py.

To identify which anchor_id to use, you need to either look at the source of the document, or use sphobjinv.

References to python objects

Pattern

text
{type}`import.path`  # to show full import path
  {type}`~import.path`  # to show only object name

where type is func for functions, meth for methods, class for classes, prop for property, etc.

Example source:

text
{class}`~pymc.gp.HSGP`

Rendered example: {class}~pymc.gp.HSGP

:::{seealso}
* ReadTheDocs page on sphinx cross-references instead.
* {ref}MyST docs on cross-references <myst:syntax/referencing>.
:::
::::

  • 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 {ref}toggle button <myst-nb:use/hiding/code>

  • Consider using {ref}myst:syntax/md-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 to add a term reference. Link to glossary source 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.
  • Use meaningful variable names wherever possible. Our users come from different backgrounds and not everyone is familiar with the same naming conventions.

- Annotate dimensions too. Notebooks are published to be read, so even if the shape is derived
from the inputs or you don't like to use named dims and don't use them in your personal
code, notebooks must use dims, even if annotating and not setting the shape.
It makes the code easier to follow, especially for newcomers.

  • Sometimes it makes sense to use Greek letters to refer to variables, for example when writing equations, as this makes them easier to read. In that case, use LaTeX to insert the Greek letter like this $\theta$ instead of using Unicode like θ.
  • If you need to use Greek letter variable names inside the code, please spell them out instead of using unicode. For example, theta, not θ.
  • 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:
text
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(...)

text
:::
  :::{tab-item} Instead of subsetting every time

fig, axs = pyplot.subplots()

axs[0, 1].plot(...)
axs[0, 1].set(...)

axs[1. 2].scatter(...)

text
:::
  ::::

  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 broadcasting 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)

text
**Looping**
* When using enumerate, take the first letter of the variable as the count:

for p, person in enumerate(persons)

text
* 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:

markdown
(notebook_name)=
# Notebook Title

:::{post} Aug 31, 2021
:tags: tag1, tag2, tags can have spaces, tag4
:category: level
:author: Alice Abat, Bob Barceló
:::

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.

:::{important}
The {ref}MyST target <myst:syntax/targets> (the (notebook_name)= bit)
is used to link notebooks between each other.
It must be notebook specific, for example its file name.
Do not copy paste this and leave notebook_name unmodified
:::

Tags can be anything, but we ask you to try to use {doc}existing tags <nb:gallery>
to avoid the tag list from getting too long.

Each notebook should have a one or two categories indicating:

  • the level of the notebook (required):

- beginner (standing crow icon)
- intermediate (flying dove icon)
- advanced (dragon icon)

  • the diataxis type (optional for old notebooks):

- tutorial
- how-to
- explanation
- reference

Authors should list people who authored, adapted or updated the notebook, excluding those
who only re-executed a notebook with little to no code or wording changes.
Only author names should be added here as this is only metadata of the notebook,
self-promotion links and details on the changes should be added at the "Authors" section,
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 menu with {menuselection}Edit --> Edit notebook metadata

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-23
{
"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"
}
}
}

text
:::

      :::{tab-item} pip and conda specific keys
      ```{code-block} json
      :emphasize-lines: 19-24
      {
        "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 specific 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),
set the ArviZ style to darkgrid (this has to be in another cell than the matplotlib import because of the way matplotlib sets its defaults):

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

Reading from file

Use a try... except clause to load the data and use pm.get_data in the except path. This will ensure that users who have cloned pymc-examples repo will read their local copy of the data while also downloading the data from github for those who don't have a local copy. Here is one example:

python
try:
    df_all = pd.read_csv(os.path.join("..", "data", "file.csv"), ...)
except FileNotFoundError:
    df_all = pd.read_csv(pm.get_data("file.csv"), ...)

pre-commit and code formatting

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

bash
pip install -U pre-commit

and then enable it with

bash
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.:

bash
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:

python
# fmt: off
np.array(
    [
        [1, 0, 0, 0],
        [0, -1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, -1],
    ]
)
# fmt: on

(jupyter_authors)=

Authorship and attribution

After the notebook content finishes, there should be an ## Authors section with bullet points
to provide attribution to the people who contributed to the notebook. The general pattern should be:

markdown
## Authors

* <verb> by <author> in <date> ([repo#PR](https://link-to.pr))

where <author> should be the name (multiple people allowed) which can be formatted as
a hyperlink to the personal site or GitHub profile of the person,
and <date> should preferably be month and year.

The <verb> part should aim to be descriptive of the changes done, for example
"updated", "re-executed", "authored" or "adapted", but it is not restricted to anything.

Authors with significant contributions should also be included in the post metadata as indicated in
{ref}jupyter_style_first_cell. There are no general and strict guidelines on that, if in doubt,
add yourself and ask reviewers for a second opinion. The main reason for that is the authors
section here at the bottom aims to be a log of all changes that happen to the notebook,
whereas the metadata at the top is used for rendering citation recommendation,
so for example, re-executing a notebook that requires no changes is a valuable contribution
which will be logged both here in this section and on GitHub but does not match authorship criteria
for citing. On the other hand, updating the wording and rendering of a notebook to make
it clearer and more friendly to the reader is something that should be added in both places,
even if the notebook is not re-executed.

some examples:

markdown
## Authors

* Authored by Chris Fonnesbeck in May, 2017 ([pymc#2124](https://github.com/pymc-devs/pymc/pull/2124))
* Updated by Colin Carroll in June, 2018 ([pymc#3049](https://github.com/pymc-devs/pymc/pull/3049))
* Updated by Alex Andorra in January, 2020 ([pymc#3765](https://github.com/pymc-devs/pymc/pull/3765))
* Updated by Oriol Abril in June, 2020 ([pymc#3963](https://github.com/pymc-devs/pymc/pull/3963))
* Updated by Farhan Reynaldo in November 2021 ([pymc-examples#246](https://github.com/pymc-devs/pymc-examples/pull/246))

and

markdown
## Authors

* Adapted from chapter 5 of Bayesian Data Analysis 3rd Edition {cite:p}`gelman2013bayesian`
  by Demetri Pananos and Junpeng Lao in June, 2018 ([pymc#3054](https://github.com/pymc-devs/pymc/pull/3054))
* Reexecuted by Ravin Kumar with PyMC 3.6 in March, 2019 ([pymc#3397](https://github.com/pymc-devs/pymc/pull/3397))
* Reexecuted by Alex Andorra and Michael Osthege with PyMC 3.9 in June, 2020 ([pymc#3955](https://github.com/pymc-devs/pymc/pull/3955))
* Updated by Raúl Maldonado in 2021 ([pymc-examples#24](https://github.com/pymc-devs/pymc-examples/pull/24), [pymc-examples#45](https://github.com/pymc-devs/pymc-examples/pull/45) and [pymc-examples#147](https://github.com/pymc-devs/pymc-examples/pull/147))

References

References should be added to the references.bib file in bibtex format, and cited with sphinxcontrib-bibtex within the notebook text wherever they are relevant.

The references in the .bib file should have as id something along the lines authorlastnameYEARkeyword or libraryYEARkeyword for documentation pages, and they should be alphabetically sorted by this id in order to ease finding references within the file and preventing adding duplicate ones.

References can be cited twice within a single notebook. Two common reference formats are:

text
{cite:p}`bibtex_id`  # shows the reference author and year between parenthesis
{cite:t}`bibtex_id`  # textual cite, shows author and year without parenthesis

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}
:filter: docname in docnames
:::

or alternatively, if you wanted to add extra references that have not been cited within the text, use:

markdown
## References

:::{bibliography}
:filter: docname in docnames

extra_bibtex_id_1
extra_bibtex_id_2
:::

Watermark

watermark is a library which automatically prints the versions of Python and the packages you used to run the NB -- reproducibility rocks!

This library should be in your virtual environment if you installed our requirements-dev.txt. Otherwise, run pip install watermark.

First, add a Markdown cell with the ## Watermark title only so it appears in the table of contents. This is the second to last section, above the epilogue/footer. Then, add a code cell to print the versions of Python and packages used in the notebook. This is the last code cell in the notebook.

The p flag is optional (or it may need to have different libraries as input), but should be added if PyTensor or xarray are not imported explicitly. This will also be checked by pre-commit (because we all forget to do things sometimes 😳).

markdown
## Watermark
python
%load_ext watermark
%watermark -n -u -v -iv -w -p pytensor,xarray

Epilogue

The last cell in the notebooks should be a markdown cell with exactly the following content:

text
:::{include} ../page_footer.md
:::

The only exception being notebooks that are not on the usual place and therefore need to
update the path to page footer for the include to work.


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!


Source/Contributing/Pr Checklist

(pr_checklist)=

Pull request checklist

We recommended that your contribution complies with the following guidelines before you submit a pull request:

  • If your pull request addresses an issue, use the pull request title to describe the issue and mention the issue number in the pull request _description_.

This will make sure a link back to the original issue is created.

:::{caution}
Adding the related issue in the PR title generates no link and is therefore
not useful as nobody knows issue numbers. Please mention all related
issues in the PR but do so only in the PR description.
:::

  • All public methods must have informative docstrings with sample usage when appropriate.

Docstrings should follow the numpydoc style

  • Please select "Create draft pull request" in the dropdown menu when opening your pull request to indicate a work in progress. This is to avoid duplicated work, to get early input on implementation details or API/functionality, or to seek collaborators.
  • Documentation and high-coverage tests are necessary for enhancements to be accepted.
  • When adding additional functionality, consider adding also one example notebook at pymc-examples.

Open a proposal issue in the example repo to discuss the specific scope of the notebook.

  • Run any of the pre-existing examples in pymc-examples that contain analyses that would be affected by your changes to ensure that nothing breaks. This is a useful opportunity to not only check your work for bugs that might not be revealed by unit test, but also to show how your contribution improves PyMC for end users.
  • No pre-commit errors: see the {ref}python_style and {ref}jupyter_style page on how to install and run it.
  • All other tests pass when everything is rebuilt from scratch. See {ref}running_the_test_suite