Skip to content

Commit

Permalink
More relevant documentation examples for Model
Browse files Browse the repository at this point in the history
  • Loading branch information
ricardoV94 committed Jun 17, 2024
1 parent 5a26eaf commit b77747d
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 58 deletions.
143 changes: 85 additions & 58 deletions pymc/model/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -392,76 +392,103 @@ class Model(WithMemoization, metaclass=ContextMeta):
name : str
name that will be used as prefix for names of all random
variables defined within model
coords : dict
Xarray-like coordinate keys and values. These coordinates can be used
to specify the shape of random variables and to label (but not specify)
the shape of Determinsitic, Potential and Data objects.
Other than specifying the shape of random variables, coordinates have no
effect on the model. They can't be used for label-based broadcasting or indexing.
You must use numpy-like operations for those behaviors.
check_bounds : bool
Ensure that input parameters to distributions are in a valid
range. If your model is built in a way where you know your
parameters can only take on valid values you can set this to
False for increased speed. This should not be used if your model
contains discrete variables.
model : PyMC model, optional
A parent model that this model belongs to. If not specified and the current model
is created inside another model's context, the parent model will be set to that model.
If `None` the model will not have a parent.
Examples
--------
How to define a custom model
Use context manager to define model and respective variables
.. code-block:: python
class CustomModel(Model):
# 1) override init
def __init__(self, mean=0, sigma=1, name=''):
# 2) call super's init first, passing model and name
# to it name will be prefix for all variables here if
# no name specified for model there will be no prefix
super().__init__(name, model)
# now you are in the context of instance,
# `modelcontext` will return self you can define
# variables in several ways note, that all variables
# will get model's name prefix
# 3) you can create variables with the register_rv method
self.register_rv(Normal.dist(mu=mean, sigma=sigma), 'v1', initval=1)
# this will create variable named like '{name::}v1'
# and assign attribute 'v1' to instance created
# variable can be accessed with self.v1 or self['v1']
# 4) this syntax will also work as we are in the
# context of instance itself, names are given as usual
Normal('v2', mu=mean, sigma=sigma)
# something more complex is allowed, too
half_cauchy = HalfCauchy('sigma', beta=10, initval=1.)
Normal('v3', mu=mean, sigma=half_cauchy)
# Deterministic variables can be used in usual way
Deterministic('v3_sq', self.v3 ** 2)
# Potentials too
Potential('p1', pt.constant(1))
# After defining a class CustomModel you can use it in several
# ways
# I:
# state the model within a context
with Model() as model:
CustomModel()
# arbitrary actions
# II:
# use new class as entering point in context
with CustomModel() as model:
Normal('new_normal_var', mu=1, sigma=0)
# III:
# just get model instance with all that was defined in it
model = CustomModel()
# IV:
# use many custom models within one context
with Model() as model:
CustomModel(mean=1, name='first')
CustomModel(mean=2, name='second')
# variables inside both scopes will be named like `first::*`, `second::*`
import pymc as pm
with pm.Model() as model:
x = pm.Normal("x")
Use object API to define model and respective variables
.. code-block:: python
import pymc as pm
model = pm.Model()
x = pm.Normal("x", model=model)
Use coords for defining the shape of random variables and labeling other model variables
.. code-block:: python
import pymc as pm
import numpy as np
coords = {
"feature", ["A", "B", "C"],
"trial", [1, 2, 3, 4, 5],
}
with pm.Model(coords=coords) as model:
intercept = pm.Normal("intercept", shape=(3,)) # Variable will have default dim label `intercept__dim_0`
beta = pm.Normal("beta", dims=("feature",)) # Variable will have shape (3,) and dim label `feature`
# Dims below are only used for labeling, they have no effect on shape
idx = pm.Data("idx", np.array([0, 1, 1, 2, 2])) # Variable will have default dim label `idx__dim_0`
x = pm.Data("x", np.random.normal(size=(5, 3)), dims=("trial", "feature"))
mu = pm.Deterministic("mu", intercept[idx] + beta @ x, dims="trial") # single dim can be passed as string
# Dims controls the shape of the variable
# If not specified, it would be inferred from the shape of the observations
y = pm.Normal("y", mu=mu, observed=[-1, 0, 0, 1, 1], dims=("trial",))
Define nested models, and provide name for variable name prefixing
.. code-block:: python
import pymc as pm
with pm.Model(name="root") as root:
x = pm.Normal("x") # Variable wil be named "root::x"
with pm.Model(name='first') as first:
# Variable will belong to root and first
y = pm.Normal("y", mu=x) # Variable wil be named "root::first::y"
# Can pass parent model explicitly
with pm.Model(name='second', model=root) as second:
# Variable will belong to root and second
z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z"
Set `check_bounds` to False for models with only continuous variables and default transformers
PyMC will remove the bounds check from the model logp which can speed up sampling
.. code-block:: python
import pymc as pm
with pm.Model(check_bounds=False) as model:
sigma = pm.HalfNormal("sigma")
x = pm.Normal("x", sigma=sigma) # No bounds check will be performed on `sigma`
"""

if TYPE_CHECKING:
Expand Down
26 changes: 26 additions & 0 deletions tests/model/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,27 @@ def test_context_passes_vars_to_parent_model(self):
assert m["d"] is model["one_more::d"]
assert m["one_more::d"] is model["one_more::d"]

def test_docstring_example(self):
with pm.Model(name="root") as root:
x = pm.Normal("x") # Variable wil be named "root::x"

with pm.Model(name="first") as first:
# Variable will belong to root and first
y = pm.Normal("y", mu=x) # Variable wil be named "root::first::y"

# Can pass parent model explicitly
with pm.Model(name="second", model=root) as second:
# Variable will belong to root and second
z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z"

assert x.name == "root::x"
assert y.name == "root::first::y"
assert z.name == "root::second::z"

assert set(root.basic_RVs) == {x, y, z}
assert set(first.basic_RVs) == {y}
assert set(second.basic_RVs) == {z}


class TestNested:
def test_nest_context_works(self):
Expand Down Expand Up @@ -1084,7 +1105,12 @@ def test_model_parent_set_programmatically():
with pm.Model(model=model):
y = pm.Normal("y")

with model:
with pm.Model(model=None):
z = pm.Normal("z")

assert "y" in model.named_vars
assert "z" in model.named_vars


class TestModelContext:
Expand Down

0 comments on commit b77747d

Please sign in to comment.