From 38ef6ccceffcb0b520d52bbf354313b3cd7d8cc0 Mon Sep 17 00:00:00 2001 From: Rob Zinkov Date: Sat, 3 Dec 2022 04:55:25 +0100 Subject: [PATCH 1/3] Add support for automatic naming of random variables --- pymc/distributions/distribution.py | 59 +++++++++++++++++-- pymc/tests/distributions/test_distribution.py | 20 +++++++ 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 8c7a841162..6a8f794269 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -22,6 +22,7 @@ from typing import Callable, Optional, Sequence, Tuple, Union import numpy as np +import opcode from aesara import tensor as at from aesara.compile.builders import OpFromGraph @@ -164,6 +165,45 @@ def fn(*args, **kwargs): return fn +# Helper function from pyprob +def _extract_target_of_assignment(depth): + frame = sys._getframe(depth) + code = frame.f_code + next_instruction = code.co_code[frame.f_lasti + 2] + instruction_arg = code.co_code[frame.f_lasti + 3] + instruction_name = opcode.opname[next_instruction] + if instruction_name == "STORE_FAST": + return code.co_varnames[instruction_arg] + elif instruction_name in ["STORE_NAME", "STORE_GLOBAL"]: + return code.co_names[instruction_arg] + elif ( + instruction_name in ["LOAD_FAST", "LOAD_NAME", "LOAD_GLOBAL"] + and opcode.opname[code.co_code[frame.f_lasti + 4]] in ["LOAD_CONST", "LOAD_FAST"] + and opcode.opname[code.co_code[frame.f_lasti + 6]] == "STORE_SUBSCR" + ): + if instruction_name == "LOAD_FAST": + base_name = code.co_varnames[instruction_arg] + else: + base_name = code.co_names[instruction_arg] + + second_instruction = opcode.opname[code.co_code[frame.f_lasti + 4]] + second_arg = code.co_code[frame.f_lasti + 5] + if second_instruction == "LOAD_CONST": + value = code.co_consts[second_arg] + elif second_instruction == "LOAD_FAST": + var_name = code.co_varnames[second_arg] + value = frame.f_locals[var_name] + else: + value = None + if value is not None: + index_name = repr(value) + return base_name + "[" + index_name + "]" + else: + return None + else: + return None + + class SymbolicRandomVariable(OpFromGraph): """Symbolic Random Variable @@ -216,7 +256,6 @@ class Distribution(metaclass=DistributionMeta): def __new__( cls, - name: str, *args, rng=None, dims: Optional[Dims] = None, @@ -234,8 +273,6 @@ def __new__( ---------- cls : type A PyMC distribution. - name : str - Name for the new model variable. rng : optional Random number generator to use with the RandomVariable. dims : tuple, optional @@ -277,6 +314,19 @@ def __new__( "for a standalone distribution." ) + if "name" in kwargs: + name = kwargs.pop("name") + elif len(args) > 0 and isinstance(args[0], string_types): + name = args[0] + args = args[1:] + else: + name = _extract_target_of_assignment(2) + if name is None: + raise TypeError("Name could not be inferred for variable") + + if not isinstance(name, string_types): + raise TypeError(f"Name needs to be a string but got: {name}") + if "testval" in kwargs: initval = kwargs.pop("testval") warnings.warn( @@ -285,9 +335,6 @@ def __new__( stacklevel=2, ) - if not isinstance(name, string_types): - raise TypeError(f"Name needs to be a string but got: {name}") - dims = convert_dims(dims) if observed is not None: observed = convert_observed_data(observed) diff --git a/pymc/tests/distributions/test_distribution.py b/pymc/tests/distributions/test_distribution.py index 80b91cb360..a1b351a990 100644 --- a/pymc/tests/distributions/test_distribution.py +++ b/pymc/tests/distributions/test_distribution.py @@ -416,3 +416,23 @@ def test_tag_future_warning_dist(): with pytest.warns(FutureWarning, match="Use model.rvs_to_values"): value_var = new_x.tag.value_var assert value_var == "1" + + +def test_autonaming(): + """Test that random variable ends up with same name as assignment variable""" + d = {} + with pm.Model() as m: + x = pm.Normal(0.0, 1.0) + d[2] = pm.Normal(0.0, 1.0) + + assert x.name == "x" + assert m["x"] == x + assert d[2].name == "d[2]" + assert m["d[2]"] == d[2] + + +def test_autonaming_noname(): + """Test that autonaming fails if no assignment can be found""" + with pytest.raises(TypeError, match="Name could not be inferred for variable"): + with pm.Model(): + pm.Normal() From 97b1f584007f3bfaf95c8e474f2bc46619a60caf Mon Sep 17 00:00:00 2001 From: Rob Zinkov Date: Sat, 3 Dec 2022 19:00:23 +0100 Subject: [PATCH 2/3] Provide more informative error message --- pymc/distributions/distribution.py | 6 +++++- pymc/tests/distributions/test_distribution.py | 9 ++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 6a8f794269..d06816e830 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -322,7 +322,11 @@ def __new__( else: name = _extract_target_of_assignment(2) if name is None: - raise TypeError("Name could not be inferred for variable") + raise TypeError( + "Name could not be inferred for variable from surrounding " + "context. Pass a name explicitly as the first argument to " + "the Distribution." + ) if not isinstance(name, string_types): raise TypeError(f"Name needs to be a string but got: {name}") diff --git a/pymc/tests/distributions/test_distribution.py b/pymc/tests/distributions/test_distribution.py index a1b351a990..4165bbe5a8 100644 --- a/pymc/tests/distributions/test_distribution.py +++ b/pymc/tests/distributions/test_distribution.py @@ -433,6 +433,13 @@ def test_autonaming(): def test_autonaming_noname(): """Test that autonaming fails if no assignment can be found""" - with pytest.raises(TypeError, match="Name could not be inferred for variable"): + with pytest.raises( + TypeError, + match=( + "Name could not be inferred for variable from surrounding " + "context. Pass a name explicitly as the first argument to " + "the Distribution." + ), + ): with pm.Model(): pm.Normal() From e835fa01deac82884810baea491b8494dd6862fc Mon Sep 17 00:00:00 2001 From: Rob Zinkov Date: Mon, 8 May 2023 00:49:38 +0200 Subject: [PATCH 3/3] Update ZeroInflated distribution --- pymc/distributions/discrete.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index a220b325d8..7718884af8 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -44,7 +44,7 @@ normal_lccdf, normal_lcdf, ) -from pymc.distributions.distribution import Discrete +from pymc.distributions.distribution import Distribution, Discrete from pymc.distributions.logprob import logp from pymc.distributions.mixture import Mixture from pymc.distributions.shape_utils import rv_size_is_none @@ -1254,7 +1254,7 @@ def _zero_inflated_mixture(*, name, nonzero_p, nonzero_dist, **kwargs): return Mixture.dist(weights, comp_dists, **kwargs) -class ZeroInflatedPoisson: +class ZeroInflatedPoisson(Distribution): R""" Zero-inflated Poisson log-likelihood. @@ -1306,11 +1306,6 @@ class ZeroInflatedPoisson: (mu >= 0). """ - def __new__(cls, name, psi, mu, **kwargs): - return _zero_inflated_mixture( - name=name, nonzero_p=psi, nonzero_dist=Poisson.dist(mu=mu), **kwargs - ) - @classmethod def dist(cls, psi, mu, **kwargs): return _zero_inflated_mixture(