From b77747d7a41a5df767f7bdfe40622f26ccdb0901 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 10 Jun 2024 17:33:44 +0200 Subject: [PATCH] More relevant documentation examples for Model --- pymc/model/core.py | 143 +++++++++++++++++++++++---------------- tests/model/test_core.py | 26 +++++++ 2 files changed, 111 insertions(+), 58 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 5ad35498fa..7aec544d79 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -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: diff --git a/tests/model/test_core.py b/tests/model/test_core.py index fe4932b572..484ac76d9c 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -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): @@ -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: