diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 1eeea5a9ea..d231a8ea3c 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -49,7 +49,6 @@ jobs: --ignore=pymc/tests/test_updates.py --ignore=pymc/tests/test_gp.py --ignore=pymc/tests/test_model.py - --ignore=pymc/tests/test_model_func.py --ignore=pymc/tests/test_ode.py --ignore=pymc/tests/test_posdef_sym.py --ignore=pymc/tests/test_quadpotential.py @@ -82,7 +81,6 @@ jobs: pymc/tests/test_distributions_timeseries.py pymc/tests/test_gp.py pymc/tests/test_model.py - pymc/tests/test_model_func.py pymc/tests/test_model_graph.py pymc/tests/test_ode.py pymc/tests/test_posdef_sym.py @@ -166,7 +164,6 @@ jobs: pymc/tests/test_ode.py - | pymc/tests/test_model.py - pymc/tests/test_model_func.py pymc/tests/test_modelcontext.py pymc/tests/test_model_graph.py pymc/tests/test_pickling.py diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 386a1fcb0d..cbe530721a 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -39,9 +39,13 @@ All of the above apply to: - `pm.sample_prior_predictive`, `pm.sample_posterior_predictive` and `pm.sample_posterior_predictive_w` now return an `InferenceData` object by default, instead of a dictionary (see [#5073](https://github.com/pymc-devs/pymc/pull/5073)). - `pm.sample_prior_predictive` no longer returns transformed variable values by default. Pass them by name in `var_names` if you want to obtain these draws (see [4769](https://github.com/pymc-devs/pymc/pull/4769)). - `pm.sample(trace=...)` no longer accepts `MultiTrace` or `len(.) > 0` traces ([see 5019#](https://github.com/pymc-devs/pymc/pull/5019)). +- `logpt`, `logpt_sum`, `logp_elemwiset` and `nojac` variations were removed. Use `Model.logpt(jacobian=True/False, sum=True/False)` instead. +- `dlogp_nojact` and `d2logp_nojact` were removed. Use `Model.dlogpt` and `d2logpt` with `jacobian=False` instead. +- `logp`, `dlogp`, and `d2logp` and `nojac` variations were removed. Use `Model.compile_logp`, `compile_dlgop` and `compile_d2logp` with `jacobian` keyword instead. +- `model.makefn` is now called `Model.compile_fn`, and `model.fn` was removed. +- Methods starting with `fast_*`, such as `Model.fast_logp`, were removed. Same applies to `PointFunc` classes - The GLM submodule was removed, please use [Bambi](https://bambinos.github.io/bambi/) instead. - `pm.Bound` interface no longer accepts a callable class as argument, instead it requires an instantiated distribution (created via the `.dist()` API) to be passed as an argument. In addition, Bound no longer returns a class instance but works as a normal PyMC distribution. Finally, it is no longer possible to do predictive random sampling from Bounded variables. Please, consult the new documentation for details on how to use Bounded variables (see [4815](https://github.com/pymc-devs/pymc/pull/4815)). -- `pm.logpt(transformed=...)` kwarg was removed (816b5f). - `Model(model=...)` kwarg was removed - `Model(theano_config=...)` kwarg was removed - `Model.size` property was removed (use `Model.ndim` instead). diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 2663344bcb..9d872e0724 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -246,12 +246,26 @@ def _extract_log_likelihood(self, trace): # TODO: We no longer need one function per observed variable if self.log_likelihood is True: cached = [ - (var, self.model.fn(self.model.logp_elemwiset(var)[0])) + ( + var, + self.model.compile_fn( + self.model.logpt(var, sum=False)[0], + inputs=self.model.value_vars, + on_unused_input="ignore", + ), + ) for var in self.model.observed_RVs ] else: cached = [ - (var, self.model.fn(self.model.logp_elemwiset(var)[0])) + ( + var, + self.model.compile_fn( + self.model.logpt(var, sum=False)[0], + inputs=self.model.value_vars, + on_unused_input="ignore", + ), + ) for var in self.model.observed_RVs if var.name in self.log_likelihood ] diff --git a/pymc/backends/base.py b/pymc/backends/base.py index 6d8197d890..55b859760c 100644 --- a/pymc/backends/base.py +++ b/pymc/backends/base.py @@ -65,7 +65,7 @@ def __init__(self, name, model=None, vars=None, test_point=None): self.vars = vars self.varnames = [var.name for var in vars] - self.fn = model.fastfn(vars) + self.fn = model.compile_fn(vars, inputs=model.value_vars, on_unused_input="ignore") # Get variable shapes. Most backends will need this # information. diff --git a/pymc/distributions/__init__.py b/pymc/distributions/__init__.py index 2aff0c18c0..e1e899a3e0 100644 --- a/pymc/distributions/__init__.py +++ b/pymc/distributions/__init__.py @@ -16,8 +16,7 @@ logcdf, logp, logp_transform, - logpt, - logpt_sum, + joint_logpt, ) from pymc.distributions.bound import Bound @@ -191,9 +190,8 @@ "Censored", "CAR", "PolyaGamma", - "logpt", + "joint_logpt", "logp", "logp_transform", "logcdf", - "logpt_sum", ] diff --git a/pymc/distributions/logprob.py b/pymc/distributions/logprob.py index 956e555d1f..9fe2b94b99 100644 --- a/pymc/distributions/logprob.py +++ b/pymc/distributions/logprob.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import warnings from collections.abc import Mapping from functools import singledispatch @@ -118,7 +117,7 @@ def _get_scaling(total_size, shape, ndim): ) -def logpt( +def joint_logpt( var: Union[TensorVariable, List[TensorVariable]], rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None, *, @@ -264,17 +263,3 @@ def logcdf(rv, value): value = at.as_tensor_variable(value, dtype=rv.dtype) return logcdf_aeppl(rv, value) - - -def logpt_sum(*args, **kwargs): - """Return the sum of the logp values for the given observations. - - Subclasses can use this to improve the speed of logp evaluations - if only the sum of the logp values is needed. - """ - warnings.warn( - "logpt_sum has been deprecated, you can use logpt instead, which now defaults" - "to the same behavior of logpt_sum", - DeprecationWarning, - ) - return logpt(*args, sum=True, **kwargs) diff --git a/pymc/model.py b/pymc/model.py index dc43ecc338..6c90f1a42a 100644 --- a/pymc/model.py +++ b/pymc/model.py @@ -22,6 +22,7 @@ from typing import ( TYPE_CHECKING, Any, + Callable, Dict, List, Optional, @@ -57,7 +58,7 @@ ) from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.data import GenTensorVariable, Minibatch -from pymc.distributions import logp_transform, logpt +from pymc.distributions import joint_logpt, logp_transform from pymc.exceptions import ImputationWarning, SamplingError, ShapeError from pymc.initial_point import make_initial_point_fn from pymc.math import flatten_list @@ -73,15 +74,12 @@ __all__ = [ "Model", - "Factor", - "compilef", - "fn", - "fastfn", "modelcontext", - "Point", "Deterministic", "Potential", "set_data", + "Point", + "compile_fn", ] FlatView = collections.namedtuple("FlatView", "input, replacements") @@ -271,90 +269,6 @@ def modelcontext(model: Optional["Model"]) -> "Model": return model -class Factor: - """Common functionality for objects with a log probability density - associated with them. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - @property - def logp(self): - """Compiled log probability density function""" - return self.model.fn(self.logpt) - - def logp_elemwise(self, vars=None, jacobian=True): - return self.model.fn(self.logp_elemwiset(vars=vars, jacobian=jacobian)) - - def dlogp(self, vars=None): - """Compiled log probability density gradient function""" - return self.model.fn(gradient(self.logpt, vars)) - - def d2logp(self, vars=None): - """Compiled log probability density hessian function""" - return self.model.fn(hessian(self.logpt, vars)) - - @property - def logp_nojac(self): - return self.model.fn(self.logp_nojact) - - def dlogp_nojac(self, vars=None): - """Compiled log density gradient function, without jacobian terms.""" - return self.model.fn(gradient(self.logp_nojact, vars)) - - def d2logp_nojac(self, vars=None): - """Compiled log density hessian function, without jacobian terms.""" - return self.model.fn(hessian(self.logp_nojact, vars)) - - @property - def fastlogp(self): - """Compiled log probability density function""" - return self.model.fastfn(self.logpt) - - def fastdlogp(self, vars=None): - """Compiled log probability density gradient function""" - return self.model.fastfn(gradient(self.logpt, vars)) - - def fastd2logp(self, vars=None): - """Compiled log probability density hessian function""" - return self.model.fastfn(hessian(self.logpt, vars)) - - @property - def fastlogp_nojac(self): - return self.model.fastfn(self.logp_nojact) - - def fastdlogp_nojac(self, vars=None): - """Compiled log density gradient function, without jacobian terms.""" - return self.model.fastfn(gradient(self.logp_nojact, vars)) - - def fastd2logp_nojac(self, vars=None): - """Compiled log density hessian function, without jacobian terms.""" - return self.model.fastfn(hessian(self.logp_nojact, vars)) - - @property - def logpt(self): - """Aesara scalar of log-probability of the model""" - if getattr(self, "total_size", None) is not None: - logp = self.logp_sum_unscaledt * self.scaling - else: - logp = self.logp_sum_unscaledt - if self.name is not None: - logp.name = f"__logp_{self.name}" - return logp - - @property - def logp_nojact(self): - """Aesara scalar of log-probability, excluding jacobian terms.""" - if getattr(self, "total_size", None) is not None: - logp = at.sum(self.logp_nojac_unscaledt) * self.scaling - else: - logp = at.sum(self.logp_nojac_unscaledt) - if self.name is not None: - logp.name = f"__logp_{self.name}" - return logp - - class ValueGradFunction: """Create an Aesara function that computes a value and its gradient. @@ -510,7 +424,7 @@ def profile(self): return self._aesara_function.profile -class Model(Factor, WithMemoization, metaclass=ContextMeta): +class Model(WithMemoization, metaclass=ContextMeta): """Encapsulates the variables and likelihood factors of a model. Model class can be used for creating class based models. To create @@ -721,7 +635,7 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs): # where the potential terms are added to the observations? costs = [self.varlogpt + self.potentiallogpt, self.observedlogpt] else: - costs = [self.logpt] + costs = [self.logpt()] input_vars = {i for i in graph_inputs(costs) if not isinstance(i, Constant)} extra_vars = [self.rvs_to_values.get(var, var) for var in self.free_RVs] @@ -731,11 +645,67 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs): } return ValueGradFunction(costs, grad_vars, extra_vars_and_values, **kwargs) - def logp_elemwiset( + def compile_logp( self, - vars: Optional[Union[Variable, List[Variable]]] = None, + vars: Optional[Union[Variable, Sequence[Variable]]] = None, jacobian: bool = True, - ) -> List[Variable]: + sum: bool = True, + ): + """Compiled log probability density function. + + Parameters + ---------- + vars: list of random variables or potential terms, optional + Compute the gradient with respect to those variables. If None, use all + free and observed random variables, as well as potential terms in model. + jacobian: + Whether to include jacobian terms in logprob graph. Defaults to True. + sum: + Whether to sum all logp terms or return elemwise logp for each variable. + Defaults to True. + """ + return self.model.compile_fn(self.logpt(vars=vars, jacobian=jacobian, sum=sum)) + + def compile_dlogp( + self, + vars: Optional[Union[Variable, Sequence[Variable]]] = None, + jacobian: bool = True, + ): + """Compiled log probability density gradient function. + + Parameters + ---------- + vars: list of random variables or potential terms, optional + Compute the gradient with respect to those variables. If None, use all + free and observed random variables, as well as potential terms in model. + jacobian: + Whether to include jacobian terms in logprob graph. Defaults to True. + """ + return self.model.compile_fn(self.dlogpt(vars=vars, jacobian=jacobian)) + + def compile_d2logp( + self, + vars: Optional[Union[Variable, Sequence[Variable]]] = None, + jacobian: bool = True, + ): + """Compiled log probability density hessian function. + + Parameters + ---------- + vars: list of random variables or potential terms, optional + Compute the gradient with respect to those variables. If None, use all + free and observed random variables, as well as potential terms in model. + jacobian: + Whether to include jacobian terms in logprob graph. Defaults to True. + """ + return self.model.compile_fn(self.d2logpt(vars=vars, jacobian=jacobian)) + + def logpt( + self, + vars: Optional[Union[Variable, Sequence[Variable]]] = None, + jacobian: bool = True, + sum: bool = True, + ) -> Union[Variable, List[Variable]]: """Elemwise log-probability of the model. Parameters @@ -743,12 +713,15 @@ def logp_elemwiset( vars: list of random variables or potential terms, optional Compute the gradient with respect to those variables. If None, use all free and observed random variables, as well as potential terms in model. - jacobian + jacobian: Whether to include jacobian terms in logprob graph. Defaults to True. + sum: + Whether to sum all logp terms or return elemwise logp for each variable. + Defaults to True. Returns ------- - Elemwise logp terms for ecah requested variable, in the same order of input. + Logp graph(s) """ if vars is None: vars = self.free_RVs + self.observed_RVs + self.potentials @@ -776,7 +749,7 @@ def logp_elemwiset( rv_logps = [] if rv_values: - rv_logps = logpt(list(rv_values.keys()), rv_values, sum=False, jacobian=jacobian) + rv_logps = joint_logpt(list(rv_values.keys()), rv_values, sum=False, jacobian=jacobian) if not isinstance(rv_logps, list): rv_logps = [rv_logps] @@ -785,82 +758,121 @@ def logp_elemwiset( if potentials: potential_logps, _ = rvs_to_value_vars(potentials, apply_transforms=True) - logp_elemwise = [None] * len(vars) + logp_factors = [None] * len(vars) for logp_order, logp in zip((rv_order + potential_order), (rv_logps + potential_logps)): - logp_elemwise[logp_order] = logp + logp_factors[logp_order] = logp - return logp_elemwise - - @property - def logpt(self): - """Aesara scalar of log-probability of the model""" - logp_var = self.varlogpt + self.datalogpt + if not sum: + return logp_factors + logp_scalar = at.sum([at.sum(factor) for factor in logp_factors]) + logp_scalar_name = "__logp" if jacobian else "__logp_nojac" if self.name: - logp_var.name = f"__logp_{self.name}" - else: - logp_var.name = "__logp" - return logp_var + logp_scalar_name = f"{logp_scalar_name}_{self.name}" + logp_scalar.name = logp_scalar_name + return logp_scalar - @property - def logp_nojact(self): - """Aesara scalar of log-probability of the model but without the jacobian - if transformed Random Variable is presented. + def dlogpt( + self, + vars: Optional[Union[Variable, Sequence[Variable]]] = None, + jacobian: bool = True, + ) -> Variable: + """Gradient of the models log-probability w.r.t. ``vars``. - Note that if there is no transformed variable in the model, logp_nojact - will be the same as logpt as there is no need for Jacobian correction. + Parameters + ---------- + vars: list of random variables or potential terms, optional + Compute the gradient with respect to those variables. If None, use all + free and observed random variables, as well as potential terms in model. + jacobian: + Whether to include jacobian terms in logprob graph. Defaults to True. + + Returns + ------- + dlogp graph """ - logp_var = self.varlogp_nojact + self.datalogpt + if vars is None: + value_vars = None + else: + if not isinstance(vars, (list, tuple)): + vars = [vars] + + value_vars = [] + for i, var in enumerate(vars): + value_var = self.rvs_to_values.get(var) + if value_var is not None: + value_vars.append(value_var) + else: + raise ValueError( + f"Requested variable {var} not found among the model variables" + ) - if self.name: - logp_var.name = f"__logp_nojac_{self.name}" + cost = self.logpt(jacobian=jacobian) + return gradient(cost, value_vars) + + def d2logpt( + self, + vars: Optional[Union[Variable, Sequence[Variable]]] = None, + jacobian: bool = True, + ) -> Variable: + """Hessian of the models log-probability w.r.t. ``vars``. + + Parameters + ---------- + vars: list of random variables or potential terms, optional + Compute the gradient with respect to those variables. If None, use all + free and observed random variables, as well as potential terms in model. + jacobian: + Whether to include jacobian terms in logprob graph. Defaults to True. + + Returns + ------- + d²logp graph + """ + if vars is None: + value_vars = None else: - logp_var.name = "__logp_nojac" - return logp_var + if not isinstance(vars, (list, tuple)): + vars = [vars] + + value_vars = [] + for i, var in enumerate(vars): + value_var = self.rvs_to_values.get(var) + if value_var is not None: + value_vars.append(value_var) + else: + raise ValueError( + f"Requested variable {var} not found among the model variables" + ) + + cost = self.logpt(jacobian=jacobian) + return hessian(cost, value_vars) @property - def datalogpt(self): + def datalogpt(self) -> Variable: """Aesara scalar of log-probability of the observed variables and potential terms""" return self.observedlogpt + self.potentiallogpt @property - def varlogpt(self): + def varlogpt(self) -> Variable: """Aesara scalar of log-probability of the unobserved random variables (excluding deterministic).""" - rv_values = {} - for var in self.free_RVs: - rv_values[var] = self.rvs_to_values[var] - if rv_values: - return logpt(self.free_RVs, rv_values) - else: - return 0 + return self.logpt(vars=self.free_RVs) @property - def varlogp_nojact(self): + def varlogp_nojact(self) -> Variable: """Aesara scalar of log-probability of the unobserved random variables (excluding deterministic) without jacobian term.""" - rv_values = {} - for var in self.free_RVs: - rv_values[var] = self.rvs_to_values[var] - if rv_values: - return logpt(self.free_RVs, rv_values, jacobian=False) - else: - return 0 + return self.logpt(vars=self.free_RVs, jacobian=False) @property - def observedlogpt(self): + def observedlogpt(self) -> Variable: """Aesara scalar of log-probability of the observed variables""" - obs_values = {} - for obs in self.observed_RVs: - obs_values[obs] = obs.tag.observations - if obs_values: - return logpt(self.observed_RVs, obs_values) - else: - return 0 + return self.logpt(vars=self.observed_RVs) @property - def potentiallogpt(self): + def potentiallogpt(self) -> Variable: """Aesara scalar of log-probability of the Potential terms""" # Convert random variables in Potential expression into their log-likelihood # inputs and apply their transforms, if any @@ -868,7 +880,7 @@ def potentiallogpt(self): if potentials: return at.sum([at.sum(factor) for factor in potentials]) else: - return 0 + return at.constant(0.0) @property def vars(self): @@ -1464,63 +1476,48 @@ def __getitem__(self, key): except KeyError: raise e - def makefn(self, outs, mode=None, *args, **kwargs): - """Compiles an Aesara function which returns ``outs`` and takes the variable - ancestors of ``outs`` as inputs. + def compile_fn( + self, + outs: Sequence[Variable], + *, + inputs: Optional[Sequence[Variable]] = None, + mode=None, + point_fn: bool = True, + **kwargs, + ) -> Union["PointFunc", Callable[[Sequence[np.ndarray]], Sequence[np.ndarray]]]: + """Compiles an Aesara function Parameters ---------- outs: Aesara variable or iterable of Aesara variables + inputs: Aesara input variables, defaults to aesaraf.inputvars(outs). mode: Aesara compilation mode, default=None + point_fn: + Whether to wrap the compiled function in a PointFunc, which takes a Point + dictionary with model variable names and values as input. Returns ------- Compiled Aesara function """ + if inputs is None: + inputs = inputvars(outs) + with self: - return compile_pymc( - self.value_vars, + fn = compile_pymc( + inputs, outs, allow_input_downcast=True, - on_unused_input="ignore", accept_inplace=True, mode=mode, - *args, **kwargs, ) - def fn(self, outs, mode=None, *args, **kwargs): - """Compiles an Aesara function which returns the values of ``outs`` - and takes values of model vars as arguments. + if point_fn: + return PointFunc(fn) + return fn - Parameters - ---------- - outs: Aesara variable or iterable of Aesara variables - mode: Aesara compilation mode - - Returns - ------- - Compiled Aesara function - """ - return LoosePointFunc(self.makefn(outs, mode, *args, **kwargs), self) - - def fastfn(self, outs, mode=None, *args, **kwargs): - """Compiles an Aesara function which returns ``outs`` and takes values - of model vars as a dict as an argument. - - Parameters - ---------- - outs: Aesara variable or iterable of Aesara variables - mode: Aesara compilation mode - - Returns - ------- - Compiled Aesara function as point function. - """ - f = self.makefn(outs, mode, *args, **kwargs) - return FastPointFunc(f) - - def profile(self, outs, n=1000, point=None, profile=True, *args, **kwargs): + def profile(self, outs, *, n=1000, point=None, profile=True, **kwargs): """Compiles and profiles an Aesara function which returns ``outs`` and takes values of model vars as a dict as an argument. @@ -1540,7 +1537,8 @@ def profile(self, outs, n=1000, point=None, profile=True, *args, **kwargs): ProfileStats Use .summary() to print stats. """ - f = self.makefn(outs, profile=profile, *args, **kwargs) + kwargs.setdefault("on_unused_input", "ignore") + f = self.compile_fn(outs, inputs=self.value_vars, point_fn=False, profile=profile, **kwargs) if point is None: point = self.recompute_initial_point() @@ -1709,7 +1707,9 @@ def point_logps(self, point=None, round_vals=2): factor.name: np.round(np.asarray(factor_logp), round_vals) for factor, factor_logp in zip( factors, - self.fn([at.sum(factor) for factor in self.logp_elemwiset(factors)])(point), + self.compile_fn([at.sum(factor) for factor in self.logpt(factors, sum=False)])( + point + ), ) }, name="Point log-probability", @@ -1761,39 +1761,22 @@ def set_data(new_data, model=None): model.set_data(variable_name, new_value) -def fn(outs, mode=None, model=None, *args, **kwargs): - """Compiles an Aesara function which returns the values of ``outs`` and - takes values of model vars as arguments. - - Parameters - ---------- - outs: Aesara variable or iterable of Aesara variables - mode: Aesara compilation mode, default=None - model: Model, default=None - - Returns - ------- - Compiled Aesara function - """ - model = modelcontext(model) - return model.fn(outs, mode, *args, **kwargs) - - -def fastfn(outs, mode=None, model=None): +def compile_fn(outs, mode=None, point_fn=True, model=None, **kwargs): """Compiles an Aesara function which returns ``outs`` and takes values of model vars as a dict as an argument. - Parameters ---------- outs: Aesara variable or iterable of Aesara variables mode: Aesara compilation mode - + point_fn: + Whether to wrap the compiled function in a PointFunc, which takes a Point + dictionary with model variable names and values as input. Returns ------- Compiled Aesara function as point function. """ model = modelcontext(model) - return model.fastfn(outs, mode) + return model.compile_fn(outs, mode, point_fn=point_fn, **kwargs) def Point(*args, filter_model_vars=False, **kwargs) -> Dict[str, np.ndarray]: @@ -1820,7 +1803,7 @@ def Point(*args, filter_model_vars=False, **kwargs) -> Dict[str, np.ndarray]: } -class FastPointFunc: +class PointFunc: """Wraps so a function so it takes a dict of arguments instead of arguments.""" def __init__(self, f): @@ -1830,22 +1813,6 @@ def __call__(self, state): return self.f(**state) -class LoosePointFunc: - """Wraps so a function so it takes a dict of arguments instead of arguments - but can still take arguments.""" - - def __init__(self, f, model): - self.f = f - self.model = model - - def __call__(self, *args, **kwargs): - point = Point(model=self.model, *args, filter_model_vars=True, **kwargs) - return self.f(**point) - - -compilef = fastfn - - def Deterministic(name, var, model=None, dims=None, auto=False): """Create a named deterministic variable diff --git a/pymc/model_graph.py b/pymc/model_graph.py index bfb9134c11..0dc3bb6efb 100644 --- a/pymc/model_graph.py +++ b/pymc/model_graph.py @@ -91,6 +91,7 @@ def _filter_parents(self, var, parents) -> Set[VarName]: def get_parents(self, var: TensorVariable) -> Set[VarName]: """Get the named nodes that are direct inputs to the var""" + # TODO: Update these lines, variables no longer have a `logpt` attribute if hasattr(var, "transformed"): func = var.transformed.logpt elif hasattr(var, "logpt"): diff --git a/pymc/sampling.py b/pymc/sampling.py index ec917801bd..5f6ae0c051 100644 --- a/pymc/sampling.py +++ b/pymc/sampling.py @@ -210,13 +210,14 @@ def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None # Use competence classmethods to select step methods for remaining # variables selected_steps = defaultdict(list) + model_logpt = model.logpt() for var in model.value_vars: if var not in assigned_vars: # determine if a gradient can be computed has_gradient = var.dtype not in discrete_types if has_gradient: try: - tg.grad(model.logpt, var) + tg.grad(model_logpt, var) except (NotImplementedError, tg.NullTypeGradError): has_gradient = False # select the best method diff --git a/pymc/sampling_jax.py b/pymc/sampling_jax.py index 13380c83c8..a2bc2cd42f 100644 --- a/pymc/sampling_jax.py +++ b/pymc/sampling_jax.py @@ -28,7 +28,6 @@ from pymc import Model, modelcontext from pymc.backends.arviz import find_observations -from pymc.distributions import logpt from pymc.util import get_default_varnames warnings.warn("This module is experimental.") @@ -73,7 +72,7 @@ def replace_shared_variables(graph: List[TensorVariable]) -> List[TensorVariable def get_jaxified_logp(model: Model) -> Callable: """Compile model.logpt into an optimized jax function""" - logpt = replace_shared_variables([model.logpt])[0] + logpt = replace_shared_variables([model.logpt()])[0] logpt_fgraph = FunctionGraph(outputs=[logpt], clone=False) optimize_graph(logpt_fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"]) @@ -123,7 +122,7 @@ def _get_log_likelihood(model, samples): "Compute log-likelihood for all observations" data = {} for v in model.observed_RVs: - logp_v = replace_shared_variables([model.logp_elemwiset(v)[0]]) + logp_v = replace_shared_variables([model.logpt(v, sum=False)[0]]) fgraph = FunctionGraph(model.value_vars, logp_v, clone=False) optimize_graph(fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"]) jax_fn = jax_funcify(fgraph) diff --git a/pymc/step_methods/elliptical_slice.py b/pymc/step_methods/elliptical_slice.py index 6a55c86a7b..ad95f831c1 100644 --- a/pymc/step_methods/elliptical_slice.py +++ b/pymc/step_methods/elliptical_slice.py @@ -93,7 +93,7 @@ def __init__(self, vars=None, prior_cov=None, prior_chol=None, model=None, **kwa vars = [self.model.rvs_to_values.get(var, var) for var in vars] vars = inputvars(vars) - super().__init__(vars, [self.model.fastlogp], **kwargs) + super().__init__(vars, [self.model.compile_logp()], **kwargs) def astep(self, q0, logp): """q0: current state diff --git a/pymc/step_methods/metropolis.py b/pymc/step_methods/metropolis.py index 79ad7b826e..e4156f3782 100644 --- a/pymc/step_methods/metropolis.py +++ b/pymc/step_methods/metropolis.py @@ -202,7 +202,7 @@ def __init__( self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) - self.delta_logp = delta_logp(initial_values, model.logpt, vars, shared) + self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) super().__init__(vars, shared) def reset_tuning(self): @@ -340,7 +340,7 @@ def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None): if not all([v.dtype in pm.discrete_types for v in vars]): raise ValueError("All variables must be Bernoulli for BinaryMetropolis") - super().__init__(vars, [model.fastlogp]) + super().__init__(vars, [model.compile_logp()]) def astep(self, q0: RaveledVars, logp) -> Tuple[RaveledVars, List[Dict[str, Any]]]: @@ -441,7 +441,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None): if not all([v.dtype in pm.discrete_types for v in vars]): raise ValueError("All variables must be binary for BinaryGibbsMetropolis") - super().__init__(vars, [model.fastlogp]) + super().__init__(vars, [model.compile_logp()]) def astep(self, q0: RaveledVars, logp: Callable[[RaveledVars], np.ndarray]) -> RaveledVars: @@ -527,7 +527,9 @@ def __init__(self, vars, proposal="uniform", order="random", model=None): if isinstance(distr, CategoricalRV): k_graph = rv_var.owner.inputs[3].shape[-1] (k_graph,), _ = rvs_to_value_vars((k_graph,), apply_transforms=True) - k = model.fn(k_graph)(initial_point) + k = model.compile_fn(k_graph, inputs=model.value_vars, on_unused_input="ignore")( + initial_point + ) elif isinstance(distr, BernoulliRV): k = 2 else: @@ -554,7 +556,7 @@ def __init__(self, vars, proposal="uniform", order="random", model=None): else: raise ValueError("Argument 'proposal' should either be 'uniform' or 'proportional'") - super().__init__(vars, [model.fastlogp]) + super().__init__(vars, [model.compile_logp()]) def astep_unif(self, q0: RaveledVars, logp) -> RaveledVars: @@ -742,7 +744,7 @@ def __init__( self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) - self.delta_logp = delta_logp(initial_values, model.logpt, vars, shared) + self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) super().__init__(vars, shared) def astep(self, q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]: @@ -905,7 +907,7 @@ def __init__( self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) - self.delta_logp = delta_logp(initial_values, model.logpt, vars, shared) + self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) super().__init__(vars, shared) def reset_tuning(self): diff --git a/pymc/step_methods/mlda.py b/pymc/step_methods/mlda.py index af618e71e3..fa52bac44d 100644 --- a/pymc/step_methods/mlda.py +++ b/pymc/step_methods/mlda.py @@ -538,7 +538,7 @@ def __init__( # Construct Aesara function for current-level model likelihood # (for use in acceptance) shared = pm.make_shared_replacements(initial_values, vars, model) - self.delta_logp = delta_logp(initial_values, model.logpt, vars, shared) + self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) # Construct Aesara function for below-level model likelihood # (for use in acceptance) @@ -547,7 +547,7 @@ def __init__( vars_below = pm.inputvars(vars_below) shared_below = pm.make_shared_replacements(initial_values, vars_below, model_below) self.delta_logp_below = delta_logp( - initial_values, model_below.logpt, vars_below, shared_below + initial_values, model_below.logpt(), vars_below, shared_below ) super().__init__(vars, shared) diff --git a/pymc/step_methods/sgmcmc.py b/pymc/step_methods/sgmcmc.py index b839edf2d3..88e238fd5a 100644 --- a/pymc/step_methods/sgmcmc.py +++ b/pymc/step_methods/sgmcmc.py @@ -45,7 +45,7 @@ def _check_minibatches(minibatch_tensors, minibatches): def prior_dlogp(vars, model, flat_view): """Returns the gradient of the prior on the parameters as a vector of size D x 1""" - terms = at.concatenate([aesara.grad(var.logpt, var).flatten() for var in vars], axis=0) + terms = at.concatenate([aesara.grad(var.logpt(), var).flatten() for var in vars], axis=0) dlogp = aesara.clone_replace(terms, flat_view.replacements, strict=False) return dlogp diff --git a/pymc/step_methods/slicer.py b/pymc/step_methods/slicer.py index a75b07324d..d8bbc9d2c4 100644 --- a/pymc/step_methods/slicer.py +++ b/pymc/step_methods/slicer.py @@ -61,7 +61,7 @@ def __init__(self, vars=None, w=1.0, tune=True, model=None, iter_limit=np.inf, * vars = [self.model.rvs_to_values.get(var, var) for var in vars] vars = inputvars(vars) - super().__init__(vars, [self.model.fastlogp], **kwargs) + super().__init__(vars, [self.model.compile_logp()], **kwargs) def astep(self, q0, logp): q0_val = q0.data diff --git a/pymc/tests/test_data_container.py b/pymc/tests/test_data_container.py index e8776a04d7..d1f42ebe1c 100644 --- a/pymc/tests/test_data_container.py +++ b/pymc/tests/test_data_container.py @@ -24,7 +24,6 @@ import pymc as pm from pymc.aesaraf import floatX -from pymc.distributions import logpt from pymc.exceptions import ShapeError from pymc.tests.helpers import SeededTest @@ -35,7 +34,7 @@ def test_deterministic(self): with pm.Model() as model: X = pm.MutableData("X", data_values) pm.Normal("y", 0, 1, observed=X) - model.logp(model.recompute_initial_point()) + model.compile_logp()(model.recompute_initial_point()) def test_sample(self): x = np.random.normal(size=100) @@ -206,8 +205,10 @@ def test_shared_scalar_as_rv_input(self): shared_var = shared(5.0) v = pm.Normal("v", mu=shared_var, size=1) + m_logp_fn = m.compile_logp() + np.testing.assert_allclose( - logpt(v, np.r_[5.0]).eval(), + m_logp_fn({"v": np.r_[5.0]}), -0.91893853, rtol=1e-5, ) @@ -215,7 +216,7 @@ def test_shared_scalar_as_rv_input(self): shared_var.set_value(10.0) np.testing.assert_allclose( - logpt(v, np.r_[10.0]).eval(), + m_logp_fn({"v": np.r_[10.0]}), -0.91893853, rtol=1e-5, ) diff --git a/pymc/tests/test_dist_math.py b/pymc/tests/test_dist_math.py index c2f4147fc1..528d4b3ea6 100644 --- a/pymc/tests/test_dist_math.py +++ b/pymc/tests/test_dist_math.py @@ -119,7 +119,9 @@ def test_multinomial_check_parameters(): p_b = pm.Dirichlet("p", floatX(np.ones(2))) MultinomialB("x", n, p_b, observed=x) - assert np.isclose(modelA.logp({"p_simplex__": [0]}), modelB.logp({"p_simplex__": [0]})) + assert np.isclose( + modelA.compile_logp()({"p_simplex__": [0]}), modelB.compile_logp()({"p_simplex__": [0]}) + ) class TestMvNormalLogp: diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index bff0b86dbf..ad2e3831b0 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -50,7 +50,7 @@ def polyagamma_cdf(*args, **kwargs): from aesara.graph.basic import ancestors from aesara.tensor.random.op import RandomVariable from aesara.tensor.var import TensorVariable -from numpy import array, inf, log +from numpy import array, inf from numpy.testing import assert_almost_equal, assert_equal from scipy import integrate from scipy.special import erf, gammaln, logit @@ -71,7 +71,6 @@ def polyagamma_cdf(*args, **kwargs): Cauchy, ChiSquared, Constant, - DensityDist, Dirichlet, DirichletMultinomial, DiscreteUniform, @@ -121,10 +120,9 @@ def polyagamma_cdf(*args, **kwargs): ZeroInflatedBinomial, ZeroInflatedNegativeBinomial, ZeroInflatedPoisson, + joint_logpt, logcdf, logp, - logpt, - logpt_sum, ) from pymc.distributions.shape_utils import to_tuple from pymc.math import kronecker @@ -602,7 +600,7 @@ def test_hierarchical_logpt(): x = pm.Uniform("x", lower=0, upper=1) y = pm.Uniform("y", lower=0, upper=x) - logpt_ancestors = list(ancestors([m.logpt])) + logpt_ancestors = list(ancestors([m.logpt()])) ops = {a.owner.op for a in logpt_ancestors if a.owner} assert len(ops) > 0 assert not any(isinstance(o, RandomVariable) for o in ops) @@ -617,7 +615,7 @@ def test_hierarchical_obs_logpt(): x = pm.Uniform("x", 0, 1, observed=obs) pm.Uniform("y", x, 2, observed=obs) - logpt_ancestors = list(ancestors([model.logpt])) + logpt_ancestors = list(ancestors([model.logpt()])) ops = {a.owner.op for a in logpt_ancestors if a.owner} assert len(ops) > 0 assert not any(isinstance(o, RandomVariable) for o in ops) @@ -705,7 +703,7 @@ def _model_input_dict(model, param_vars, pt): return pt_d model, param_vars = build_model(pymc_dist, domain, paramdomains, extra_args) - logp_pymc = model.fastlogp_nojac + logp_pymc = model.compile_logp(jacobian=False) # Test supported value and parameters domain matches scipy domains = paramdomains.copy() @@ -839,7 +837,7 @@ def check_logcdf( model, param_vars = build_model(pymc_dist, domain, paramdomains) rv = model["value"] value = model.rvs_to_values[rv] - pymc_logcdf = model.fastfn(logcdf(rv, value)) + pymc_logcdf = model.compile_fn(logcdf(rv, value)) if decimal is None: decimal = select_by_precision(float64=6, float32=3) @@ -944,8 +942,8 @@ def check_selfconsistency_discrete_logcdf( model, param_vars = build_model(distribution, domain, paramdomains) rv = model["value"] value = model.rvs_to_values[rv] - dist_logcdf = model.fastfn(logcdf(rv, value)) - dist_logp = model.fastfn(logp(rv, value)) + dist_logcdf = model.compile_fn(logcdf(rv, value)) + dist_logp = model.compile_fn(logp(rv, value)) for pt in product(domains, n_samples=n_samples): params = dict(pt) @@ -964,25 +962,6 @@ def check_selfconsistency_discrete_logcdf( err_msg=str(pt), ) - def check_int_to_1(self, model, value, domain, paramdomains, n_samples=10): - pdf = model.fastfn(exp(model.logpt)) - for pt in product(paramdomains, n_samples=n_samples): - pt = Point(pt, value=value.tag.test_value, model=model) - bij = DictToVarBijection(value, (), pt) - pdfx = bij.mapf(pdf) - area = integrate_nd(pdfx, domain, value.dshape, value.dtype) - assert_almost_equal(area, 1, err_msg=str(pt)) - - def checkd(self, distfam, valuedomain, vardomains, checks=None, extra_args=None): - if checks is None: - checks = (self.check_int_to_1,) - - if extra_args is None: - extra_args = {} - m = build_model(distfam, valuedomain, vardomains, extra_args=extra_args) - for check in checks: - check(m, m.named_vars["value"], valuedomain, vardomains) - def test_uniform(self): self.check_logp( Uniform, @@ -1232,7 +1211,7 @@ def test_wald_logp_custom_points(self, value, mu, lam, phi, alpha, logp): Wald("wald", mu=mu, lam=lam, phi=phi, alpha=alpha, transform=None) pt = {"wald": value} decimals = select_by_precision(float64=6, float32=1) - assert_almost_equal(model.fastlogp(pt), logp, decimal=decimals, err_msg=str(pt)) + assert_almost_equal(model.compile_logp()(pt), logp, decimal=decimals, err_msg=str(pt)) def test_beta_logp(self): self.check_logp( @@ -1665,15 +1644,6 @@ def test_binomial(self): {"n": NatSmall, "p": Unit}, ) - @pytest.mark.xfail(reason="checkd tests have not been refactored") - @pytest.mark.skipif(condition=(aesara.config.floatX == "float32"), reason="Fails on float32") - def test_beta_binomial_distribution(self): - self.checkd( - BetaBinomial, - Nat, - {"alpha": Rplus, "beta": Rplus, "n": NatSmall}, - ) - def test_beta_binomial(self): self.check_logp( BetaBinomial, @@ -1772,18 +1742,6 @@ def test_poisson(self): def test_constantdist(self): self.check_logp(Constant, I, {"c": I}, lambda value, c: np.log(c == value)) - @pytest.mark.xfail(reason="Test has not been refactored") - @pytest.mark.skipif( - condition=(aesara.config.floatX == "float32"), - reason="Fails on float32 due to inf issues", - ) - def test_zeroinflatedpoisson_distribution(self): - self.checkd( - ZeroInflatedPoisson, - Nat, - {"theta": Rplus, "psi": Unit}, - ) - def test_zeroinflatedpoisson(self): def logp_fn(value, psi, theta): if value == 0: @@ -1814,18 +1772,6 @@ def logcdf_fn(value, psi, theta): {"theta": Rplus, "psi": Unit}, ) - @pytest.mark.xfail(reason="Test not refactored yet") - @pytest.mark.skipif( - condition=(aesara.config.floatX == "float32"), - reason="Fails on float32 due to inf issues", - ) - def test_zeroinflatednegativebinomial_distribution(self): - self.checkd( - ZeroInflatedNegativeBinomial, - Nat, - {"mu": Rplusbig, "alpha": Rplusbig, "psi": Unit}, - ) - def test_zeroinflatednegativebinomial(self): def logp_fn(value, psi, mu, alpha): n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha) @@ -1874,14 +1820,6 @@ def logcdf_fn(value, psi, mu, alpha): {"psi": Unit, "mu": Rplusbig, "alpha": Rplusbig}, ) - @pytest.mark.xfail(reason="Test not refactored yet") - def test_zeroinflatedbinomial_distribution(self): - self.checkd( - ZeroInflatedBinomial, - Nat, - {"n": NatSmall, "p": Unit, "psi": Unit}, - ) - def test_zeroinflatedbinomial(self): def logp_fn(value, psi, n, p): if value == 0: @@ -2162,7 +2100,7 @@ def test_lkj(self, x, eta, n, lp): pt = {"lkj": x} decimals = select_by_precision(float64=6, float32=4) - assert_almost_equal(model.fastlogp(pt), lp, decimal=decimals, err_msg=str(pt)) + assert_almost_equal(model.compile_logp()(pt), lp, decimal=decimals, err_msg=str(pt)) @pytest.mark.parametrize("n", [1, 2, 3]) def test_dirichlet(self, n): @@ -2378,13 +2316,6 @@ def test_orderedprobit(self, n): lambda value, eta, cutpoints: orderedprobit_logpdf(value, eta, cutpoints), ) - @pytest.mark.xfail(reason="checkd tests have not been refactored") - def test_densitydist(self): - def logp(x): - return -log(2 * 0.5) - abs(x - 0.5) / 0.5 - - self.checkd(DensityDist, R, {}, extra_args={"logp": logp}) - def test_get_tau_sigma(self): sigma = np.array(2) assert_almost_equal(get_tau_sigma(sigma=sigma), [1.0 / sigma ** 2, sigma]) @@ -2423,7 +2354,7 @@ def test_ex_gaussian(self, value, mu, sigma, nu, logp): ExGaussian("eg", mu=mu, sigma=sigma, nu=nu) pt = {"eg": value} assert_almost_equal( - model.fastlogp(pt), + model.compile_logp()(pt), logp, decimal=select_by_precision(float64=6, float32=2), err_msg=str(pt), @@ -2596,11 +2527,11 @@ def test_interpolated_transform(self, transform): x = pm.Interpolated("x", x_points, pdf_points, transform=transform) if transform is UNSET: - assert np.isfinite(m.logp({"x_interval__": -1.0})) - assert np.isfinite(m.logp({"x_interval__": 11.0})) + assert np.isfinite(m.compile_logp()({"x_interval__": -1.0})) + assert np.isfinite(m.compile_logp()({"x_interval__": 11.0})) else: - assert not np.isfinite(m.logp({"x": -1.0})) - assert not np.isfinite(m.logp({"x": 11.0})) + assert not np.isfinite(m.compile_logp()({"x": -1.0})) + assert not np.isfinite(m.compile_logp()({"x": 11.0})) class TestBound: @@ -2618,29 +2549,29 @@ def test_continuous(self): UpperNormalTransform = Bound("uppertrans", dist, upper=10) BoundedNormalTransform = Bound("boundedtrans", dist, lower=1, upper=10) - assert logpt(LowerNormal, -1).eval() == -np.inf - assert logpt(UpperNormal, 1).eval() == -np.inf - assert logpt(BoundedNormal, 0).eval() == -np.inf - assert logpt(BoundedNormal, 11).eval() == -np.inf + assert joint_logpt(LowerNormal, -1).eval() == -np.inf + assert joint_logpt(UpperNormal, 1).eval() == -np.inf + assert joint_logpt(BoundedNormal, 0).eval() == -np.inf + assert joint_logpt(BoundedNormal, 11).eval() == -np.inf - assert logpt(UnboundedNormal, 0).eval() != -np.inf - assert logpt(UnboundedNormal, 11).eval() != -np.inf - assert logpt(InfBoundedNormal, 0).eval() != -np.inf - assert logpt(InfBoundedNormal, 11).eval() != -np.inf + assert joint_logpt(UnboundedNormal, 0).eval() != -np.inf + assert joint_logpt(UnboundedNormal, 11).eval() != -np.inf + assert joint_logpt(InfBoundedNormal, 0).eval() != -np.inf + assert joint_logpt(InfBoundedNormal, 11).eval() != -np.inf value = model.rvs_to_values[LowerNormalTransform] - assert logpt(LowerNormalTransform, value).eval({value: -1}) != -np.inf + assert joint_logpt(LowerNormalTransform, value).eval({value: -1}) != -np.inf value = model.rvs_to_values[UpperNormalTransform] - assert logpt(UpperNormalTransform, value).eval({value: 1}) != -np.inf + assert joint_logpt(UpperNormalTransform, value).eval({value: 1}) != -np.inf value = model.rvs_to_values[BoundedNormalTransform] - assert logpt(BoundedNormalTransform, value).eval({value: 0}) != -np.inf - assert logpt(BoundedNormalTransform, value).eval({value: 11}) != -np.inf + assert joint_logpt(BoundedNormalTransform, value).eval({value: 0}) != -np.inf + assert joint_logpt(BoundedNormalTransform, value).eval({value: 11}) != -np.inf ref_dist = Normal.dist(mu=0, sigma=1) - assert np.allclose(logpt(UnboundedNormal, 5).eval(), logpt(ref_dist, 5).eval()) - assert np.allclose(logpt(LowerNormal, 5).eval(), logpt(ref_dist, 5).eval()) - assert np.allclose(logpt(UpperNormal, -5).eval(), logpt(ref_dist, 5).eval()) - assert np.allclose(logpt(BoundedNormal, 5).eval(), logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(UnboundedNormal, 5).eval(), joint_logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(LowerNormal, 5).eval(), joint_logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(UpperNormal, -5).eval(), joint_logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(BoundedNormal, 5).eval(), joint_logpt(ref_dist, 5).eval()) def test_discrete(self): with Model() as model: @@ -2650,19 +2581,19 @@ def test_discrete(self): UpperPoisson = Bound("upper", dist, upper=10) BoundedPoisson = Bound("bounded", dist, lower=1, upper=10) - assert logpt(LowerPoisson, 0).eval() == -np.inf - assert logpt(UpperPoisson, 11).eval() == -np.inf - assert logpt(BoundedPoisson, 0).eval() == -np.inf - assert logpt(BoundedPoisson, 11).eval() == -np.inf + assert joint_logpt(LowerPoisson, 0).eval() == -np.inf + assert joint_logpt(UpperPoisson, 11).eval() == -np.inf + assert joint_logpt(BoundedPoisson, 0).eval() == -np.inf + assert joint_logpt(BoundedPoisson, 11).eval() == -np.inf - assert logpt(UnboundedPoisson, 0).eval() != -np.inf - assert logpt(UnboundedPoisson, 11).eval() != -np.inf + assert joint_logpt(UnboundedPoisson, 0).eval() != -np.inf + assert joint_logpt(UnboundedPoisson, 11).eval() != -np.inf ref_dist = Poisson.dist(mu=4) - assert np.allclose(logpt(UnboundedPoisson, 5).eval(), logpt(ref_dist, 5).eval()) - assert np.allclose(logpt(LowerPoisson, 5).eval(), logpt(ref_dist, 5).eval()) - assert np.allclose(logpt(UpperPoisson, 5).eval(), logpt(ref_dist, 5).eval()) - assert np.allclose(logpt(BoundedPoisson, 5).eval(), logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(UnboundedPoisson, 5).eval(), joint_logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(LowerPoisson, 5).eval(), joint_logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(UpperPoisson, 5).eval(), joint_logpt(ref_dist, 5).eval()) + assert np.allclose(joint_logpt(BoundedPoisson, 5).eval(), joint_logpt(ref_dist, 5).eval()) def create_invalid_distribution(self): class MyNormal(RandomVariable): @@ -2766,19 +2697,19 @@ def test_array_bound(self): UpperPoisson = Bound("upper", dist, upper=[np.inf, 10], transform=None) BoundedPoisson = Bound("bounded", dist, lower=[1, 2], upper=[9, 10], transform=None) - first, second = logpt(LowerPoisson, [0, 0], sum=False).eval() + first, second = joint_logpt(LowerPoisson, [0, 0], sum=False).eval() assert first == -np.inf assert second != -np.inf - first, second = logpt(UpperPoisson, [11, 11], sum=False).eval() + first, second = joint_logpt(UpperPoisson, [11, 11], sum=False).eval() assert first != -np.inf assert second == -np.inf - first, second = logpt(BoundedPoisson, [1, 1], sum=False).eval() + first, second = joint_logpt(BoundedPoisson, [1, 1], sum=False).eval() assert first != -np.inf assert second == -np.inf - first, second = logpt(BoundedPoisson, [10, 10], sum=False).eval() + first, second = joint_logpt(BoundedPoisson, [10, 10], sum=False).eval() assert first == -np.inf assert second != -np.inf @@ -3027,8 +2958,8 @@ def test_orderedlogistic_dimensions(shape): p=p, observed=obs, ) - ologp = logpt_sum(ol, np.ones_like(obs)).eval() * loge - clogp = logpt_sum(c, np.ones_like(obs)).eval() * loge + ologp = joint_logpt(ol, np.ones_like(obs), sum=True).eval() * loge + clogp = joint_logpt(c, np.ones_like(obs), sum=True).eval() * loge expected = -np.prod((size,) + shape) assert c.owner.inputs[3].ndim == (len(shape) + 1) @@ -3171,15 +3102,15 @@ def test_issue_4499(self): # https://github.com/pymc-devs/pymc/issues/4499 with pm.Model(check_bounds=False) as m: x = pm.Uniform("x", 0, 2, size=10, transform=None) - assert_almost_equal(m.logp({"x": np.ones(10)}), -np.log(2) * 10) + assert_almost_equal(m.compile_logp()({"x": np.ones(10)}), -np.log(2) * 10) with pm.Model(check_bounds=False) as m: x = pm.DiscreteUniform("x", 0, 1, size=10) - assert_almost_equal(m.logp({"x": np.ones(10)}), -np.log(2) * 10) + assert_almost_equal(m.compile_logp()({"x": np.ones(10)}), -np.log(2) * 10) with pm.Model(check_bounds=False) as m: x = pm.Constant("x", 1, size=10) - assert_almost_equal(m.logp({"x": np.ones(10)}), 0 * 10) + assert_almost_equal(m.compile_logp()({"x": np.ones(10)}), 0 * 10) def test_serialize_density_dist(): @@ -3268,7 +3199,7 @@ def logp(value, mu): a_val = np.random.normal(loc=mu_val, scale=1, size=to_tuple(size) + (supp_shape,)).astype( aesara.config.floatX ) - log_densityt = pm.logpt(a, a.tag.value_var, sum=False) + log_densityt = joint_logpt(a, a.tag.value_var, sum=False) assert ( log_densityt.eval( {a.tag.value_var: a_val, mu.tag.value_var: mu_val}, diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 4a6361504e..295f1e53da 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -64,7 +64,7 @@ ZeroInflatedPoisson, ) from pymc.distributions.distribution import _get_moment, get_moment -from pymc.distributions.logprob import logpt +from pymc.distributions.logprob import joint_logpt from pymc.distributions.multivariate import MvNormal from pymc.distributions.shape_utils import rv_size_is_none, to_tuple from pymc.initial_point import make_initial_point_fn @@ -163,7 +163,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True): assert np.allclose(moment, expected) if check_finite_logp: - logp_moment = logpt(model["x"], at.constant(moment), transformed=False).eval() + logp_moment = joint_logpt(model["x"], at.constant(moment), transformed=False).eval() assert np.isfinite(logp_moment) diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index ddba55eec5..07ed32e957 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -811,7 +811,7 @@ def setup_method(self): gp = pm.gp.Marginal(mean_func=mean_func, cov_func=cov_func) f = gp.marginal_likelihood("f", X, y, noise=0.0) p = gp.conditional("p", Xnew) - self.logp = model.logp({"p": pnew}) + self.logp = model.compile_logp()({"p": pnew}) self.X = X self.Xnew = Xnew self.y = y @@ -824,7 +824,7 @@ def testLatent1(self): gp = pm.gp.Latent(mean_func=mean_func, cov_func=cov_func) f = gp.prior("f", self.X, reparameterize=False) p = gp.conditional("p", self.Xnew) - latent_logp = model.logp({"f": self.y, "p": self.pnew}) + latent_logp = model.compile_logp()({"f": self.y, "p": self.pnew}) npt.assert_allclose(latent_logp, self.logp, atol=0, rtol=1e-2) def testLatent2(self): @@ -836,7 +836,7 @@ def testLatent2(self): p = gp.conditional("p", self.Xnew) chol = np.linalg.cholesky(cov_func(self.X).eval()) y_rotated = np.linalg.solve(chol, self.y - 0.5) - latent_logp = model.logp({"f_rotated_": y_rotated, "p": self.pnew}) + latent_logp = model.compile_logp()({"f_rotated_": y_rotated, "p": self.pnew}) npt.assert_allclose(latent_logp, self.logp, atol=5) @@ -858,7 +858,7 @@ def setup_method(self): sigma = 0.1 f = gp.marginal_likelihood("f", X, y, noise=sigma) p = gp.conditional("p", Xnew) - self.logp = model.logp({"p": pnew}) + self.logp = model.compile_logp()({"p": pnew}) self.X = X self.Xnew = Xnew self.y = y @@ -874,7 +874,7 @@ def testApproximations(self, approx): gp = pm.gp.MarginalApprox(mean_func=mean_func, cov_func=cov_func, approx=approx) f = gp.marginal_likelihood("f", self.X, self.X, self.y, self.sigma) p = gp.conditional("p", self.Xnew) - approx_logp = model.logp({"f": self.y, "p": self.pnew}) + approx_logp = model.compile_logp()({"p": self.pnew}) npt.assert_allclose(approx_logp, self.logp, atol=0, rtol=1e-2) @pytest.mark.parametrize("approx", ["FITC", "VFE", "DTC"]) @@ -922,14 +922,14 @@ def testAdditiveMarginal(self): gpsum = gp1 + gp2 + gp3 fsum = gpsum.marginal_likelihood("f", self.X, self.y, noise=self.noise) - model1_logp = model1.logp({"fsum": self.y}) + model1_logp = model1.compile_logp()({}) with pm.Model() as model2: gptot = pm.gp.Marginal( mean_func=reduce(add, self.means), cov_func=reduce(add, self.covs) ) fsum = gptot.marginal_likelihood("f", self.X, self.y, noise=self.noise) - model2_logp = model2.logp({"fsum": self.y}) + model2_logp = model2.compile_logp()({}) npt.assert_allclose(model1_logp, model2_logp, atol=0, rtol=1e-2) with model1: @@ -940,8 +940,8 @@ def testAdditiveMarginal(self): fp2 = gptot.conditional("fp2", self.Xnew) fp = np.random.randn(self.Xnew.shape[0]) - logp1 = model1.logp({"fp1": fp}) - logp2 = model2.logp({"fp2": fp}) + logp1 = model1.compile_logp()({"fp1": fp}) + logp2 = model2.compile_logp()({"fp2": fp}) npt.assert_allclose(logp1, logp2, atol=0, rtol=1e-2) @pytest.mark.parametrize("approx", ["FITC", "VFE", "DTC"]) @@ -961,14 +961,14 @@ def testAdditiveMarginalApprox(self, approx): gpsum = gp1 + gp2 + gp3 fsum = gpsum.marginal_likelihood("f", self.X, Xu, self.y, noise=sigma) - model1_logp = model1.logp({"fsum": self.y}) + model1_logp = model1.compile_logp()({}) with pm.Model() as model2: gptot = pm.gp.MarginalApprox( mean_func=reduce(add, self.means), cov_func=reduce(add, self.covs), approx=approx ) fsum = gptot.marginal_likelihood("f", self.X, Xu, self.y, noise=sigma) - model2_logp = model2.logp({"fsum": self.y}) + model2_logp = model2.compile_logp()({}) npt.assert_allclose(model1_logp, model2_logp, atol=0, rtol=1e-2) with model1: @@ -982,8 +982,8 @@ def testAdditiveMarginalApprox(self, approx): fp = np.random.randn(self.Xnew.shape[0]) - model1_logp = model1.logp({"fp1": fp, "fsum": self.y}) - model2_logp = model2.logp({"fp2": fp, "fsum": self.y}) + model1_logp = model1.compile_logp()({"fp1": fp}) + model2_logp = model2.compile_logp()({"fp2": fp}) npt.assert_allclose(model1_logp, model2_logp, atol=0, rtol=1e-2) @@ -995,12 +995,12 @@ def testAdditiveLatent(self): gpsum = gp1 + gp2 + gp3 fsum = gpsum.prior("fsum", self.X, reparameterize=False) - model1_logp = model1.logp({"fsum": self.y}) + model1_logp = model1.compile_logp()({"fsum": self.y}) with pm.Model() as model2: gptot = pm.gp.Latent(mean_func=reduce(add, self.means), cov_func=reduce(add, self.covs)) fsum = gptot.prior("fsum", self.X, reparameterize=False) - model2_logp = model2.logp({"fsum": self.y}) + model2_logp = model2.compile_logp()({"fsum": self.y}) npt.assert_allclose(model1_logp, model2_logp, atol=0, rtol=1e-2) with model1: @@ -1009,8 +1009,8 @@ def testAdditiveLatent(self): fp2 = gptot.conditional("fp2", self.Xnew) fp = np.random.randn(self.Xnew.shape[0]) - logp1 = model1.logp({"fsum": self.y, "fp1": fp}) - logp2 = model2.logp({"fsum": self.y, "fp2": fp}) + logp1 = model1.compile_logp()({"fsum": self.y, "fp1": fp}) + logp2 = model2.compile_logp()({"fsum": self.y, "fp2": fp}) npt.assert_allclose(logp1, logp2, atol=0, rtol=1e-2) def testAdditiveSparseRaises(self): @@ -1055,7 +1055,7 @@ def setup_method(self): gp = pm.gp.Latent(cov_func=cov_func) f = gp.prior("f", X, reparameterize=False) p = gp.conditional("p", Xnew) - self.gp_latent_logp = model1.logp({"f": y, "p": pnew}) + self.gp_latent_logp = model1.compile_logp()({"f": y, "p": pnew}) self.X = X self.y = y self.Xnew = Xnew @@ -1068,7 +1068,7 @@ def testTPvsLatent(self): tp = pm.gp.TP(cov_func=cov_func, nu=self.nu) f = tp.prior("f", self.X, reparameterize=False) p = tp.conditional("p", self.Xnew) - tp_logp = model.logp({"f": self.y, "p": self.pnew}) + tp_logp = model.compile_logp()({"f": self.y, "p": self.pnew}) npt.assert_allclose(self.gp_latent_logp, tp_logp, atol=0, rtol=1e-2) def testTPvsLatentReparameterized(self): @@ -1079,7 +1079,7 @@ def testTPvsLatentReparameterized(self): p = tp.conditional("p", self.Xnew) chol = np.linalg.cholesky(cov_func(self.X).eval()) f_rotated = np.linalg.solve(chol, self.y) - tp_logp = model.logp({"f_rotated_": f_rotated, "p": self.pnew}) + tp_logp = model.compile_logp()({"f_rotated_": f_rotated, "p": self.pnew}) npt.assert_allclose(self.gp_latent_logp, tp_logp, atol=0, rtol=1e-2) def testAdditiveTPRaises(self): @@ -1122,14 +1122,14 @@ def setup_method(self): p = gp.conditional("p", self.Xnew) chol = np.linalg.cholesky(cov_func(self.X).eval()) self.y_rotated = np.linalg.solve(chol, self.y - 0.5) - self.logp = latent_model.logp({"f_rotated_": self.y_rotated, "p": self.pnew}) + self.logp = latent_model.compile_logp()({"f_rotated_": self.y_rotated, "p": self.pnew}) def testLatentKronvsLatent(self): with pm.Model() as kron_model: kron_gp = pm.gp.LatentKron(mean_func=self.mean, cov_funcs=self.cov_funcs) f = kron_gp.prior("f", self.Xs) p = kron_gp.conditional("p", self.Xnew) - kronlatent_logp = kron_model.logp({"f_rotated_": self.y_rotated, "p": self.pnew}) + kronlatent_logp = kron_model.compile_logp()({"f_rotated_": self.y_rotated, "p": self.pnew}) npt.assert_allclose(kronlatent_logp, self.logp, atol=0, rtol=1e-3) def testLatentKronRaisesAdditive(self): @@ -1178,7 +1178,7 @@ def setup_method(self): f = gp.marginal_likelihood("f", self.X, self.y, noise=self.sigma) p = gp.conditional("p", self.Xnew) self.mu, self.cov = gp.predict(self.Xnew) - self.logp = model.logp({"p": self.pnew}) + self.logp = model.compile_logp()({"p": self.pnew}) def testMarginalKronvsMarginalpredict(self): with pm.Model() as kron_model: @@ -1197,7 +1197,7 @@ def testMarginalKronvsMarginal(self): kron_gp = pm.gp.MarginalKron(mean_func=self.mean, cov_funcs=self.cov_funcs) f = kron_gp.marginal_likelihood("f", self.Xs, self.y, sigma=self.sigma) p = kron_gp.conditional("p", self.Xnew) - kron_logp = kron_model.logp({"p": self.pnew}) + kron_logp = kron_model.compile_logp()({"p": self.pnew}) npt.assert_allclose(kron_logp, self.logp, atol=0, rtol=1e-2) def testMarginalKronRaises(self): diff --git a/pymc/tests/test_idata_conversion.py b/pymc/tests/test_idata_conversion.py index 240976856d..b846831f52 100644 --- a/pymc/tests/test_idata_conversion.py +++ b/pymc/tests/test_idata_conversion.py @@ -388,63 +388,6 @@ def test_multiple_observed_rv(self, log_likelihood): fails = check_multiple_attrs(test_dict, inference_data) assert not fails - @pytest.mark.xfail(reason="MultiObservedRV is no longer used in v4") - def test_multiple_observed_rv_without_observations(self): - with pm.Model(): - mu = pm.Normal("mu") - x = pm.DensityDist( # pylint: disable=unused-variable - "x", mu, logp=lambda value, mu: pm.Normal.logp(value, mu, 1), observed=0.1 - ) - inference_data = pm.sample(100, chains=2, return_inferencedata=True) - test_dict = { - "posterior": ["mu"], - "sample_stats": ["lp"], - "log_likelihood": ["x"], - "observed_data": ["value", "~x"], - } - fails = check_multiple_attrs(test_dict, inference_data) - assert not fails - assert inference_data.observed_data.value.dtype.kind == "f" - - @pytest.mark.xfail(reason="MultiObservedRV is no longer used in v4") - @pytest.mark.parametrize("multiobs", (True, False)) - def test_multiobservedrv_to_observed_data(self, multiobs): - # fake regression data, with weights (W) - np.random.seed(2019) - N = 100 - X = np.random.uniform(size=N) - W = 1 + np.random.poisson(size=N) - a, b = 5, 17 - Y = a + np.random.normal(b * X) - - with pm.Model(): - a = pm.Normal("a", 0, 10) - b = pm.Normal("b", 0, 10) - mu = a + b * X - sigma = pm.HalfNormal("sigma", 1) - w = W - - def weighted_normal(value, mu, sigma, w): - return w * pm.Normal.logp(value, mu, sigma) - - y_logp = pm.DensityDist( # pylint: disable=unused-variable - "y_logp", mu, sigma, w, logp=weighted_normal, observed=Y, size=N - ) - idata = pm.sample( - 20, tune=20, return_inferencedata=True, idata_kwargs={"density_dist_obs": multiobs} - ) - multiobs_str = "" if multiobs else "~" - test_dict = { - "posterior": ["a", "b", "sigma"], - "sample_stats": ["lp"], - "log_likelihood": ["y_logp"], - f"{multiobs_str}observed_data": ["y", "w"], - } - fails = check_multiple_attrs(test_dict, idata) - assert not fails - if multiobs: - assert idata.observed_data.y.dtype.kind == "f" - def test_single_observation(self): with pm.Model(): p = pm.Uniform("p", 0, 1) diff --git a/pymc/tests/test_logprob.py b/pymc/tests/test_logprob.py index 53a1061cb0..9b336b0023 100644 --- a/pymc/tests/test_logprob.py +++ b/pymc/tests/test_logprob.py @@ -31,7 +31,7 @@ from pymc.aesaraf import floatX, walk_model from pymc.distributions.continuous import HalfFlat, Normal, TruncatedNormal, Uniform from pymc.distributions.discrete import Bernoulli -from pymc.distributions.logprob import logcdf, logp, logpt +from pymc.distributions.logprob import joint_logpt, logcdf, logp from pymc.model import Model, Potential from pymc.tests.helpers import select_by_precision @@ -41,7 +41,7 @@ def assert_no_rvs(var): return var -def test_logpt_basic(): +def test_joint_logpt_basic(): """Make sure we can compute a log-likelihood for a hierarchical model with transforms.""" with Model() as m: @@ -58,7 +58,7 @@ def test_logpt_basic(): c_value_var = m.rvs_to_values[c] - b_logp = logpt(b, b_value_var, sum=False) + b_logp = joint_logpt(b, b_value_var, sum=False) res_ancestors = list(walk_model((b_logp,), walk_past_rvs=True)) res_rv_ancestors = [ @@ -81,7 +81,7 @@ def test_logpt_basic(): ((np.array([0, 1, 4]), np.array([0, 1, 4])), (5, 5)), ], ) -def test_logpt_incsubtensor(indices, size): +def test_joint_logpt_incsubtensor(indices, size): """Make sure we can compute a log-likelihood for ``Y[idx] = data`` where ``Y`` is univariate.""" mu = floatX(np.power(10, np.arange(np.prod(size)))).reshape(size) @@ -102,7 +102,7 @@ def test_logpt_incsubtensor(indices, size): a_idx_value_var = a_idx.type() a_idx_value_var.name = "a_idx_value" - a_idx_logp = logpt(a_idx, {a_idx: a_value_var}, sum=False) + a_idx_logp = joint_logpt(a_idx, {a_idx: a_value_var}, sum=False) logp_vals = a_idx_logp.eval({a_value_var: a_val}) @@ -116,7 +116,7 @@ def test_logpt_incsubtensor(indices, size): np.testing.assert_almost_equal(logp_vals, exp_obs_logps) -def test_logpt_subtensor(): +def test_joint_logpt_subtensor(): """Make sure we can compute a log-likelihood for ``Y[I]`` where ``Y`` and ``I`` are random variables.""" size = 5 @@ -144,7 +144,7 @@ def test_logpt_subtensor(): I_value_var = I_rv.type() I_value_var.name = "I_value" - A_idx_logps = logpt(A_idx, {A_idx: A_idx_value_var, I_rv: I_value_var}, sum=False) + A_idx_logps = joint_logpt(A_idx, {A_idx: A_idx_value_var, I_rv: I_value_var}, sum=False) A_idx_logp = at.add(*A_idx_logps) logp_vals_fn = aesara.function([A_idx_value_var, I_value_var], A_idx_logp) @@ -201,7 +201,7 @@ def test_logcdf_transformed_argument(): sigma_value = np.exp(sigma_value_log) x_value = 0.5 - observed = m.logp_nojac({"sigma_log__": sigma_value_log, "x": x_value}) + observed = m.compile_logp(jacobian=False)({"sigma_log__": sigma_value_log, "x": x_value}) expected = logp(TruncatedNormal.dist(0, sigma_value, lower=None, upper=1.0), x_value).eval() assert np.isclose(observed, expected) @@ -214,6 +214,6 @@ def test_model_unchanged_logprob_access(): original_inputs = set(aesara.graph.graph_inputs([c])) # Extract model.logpt - model.logpt + model.logpt() new_inputs = set(aesara.graph.graph_inputs([c])) assert original_inputs == new_inputs diff --git a/pymc/tests/test_minibatches.py b/pymc/tests/test_minibatches.py index a9c9874798..0fe142881c 100644 --- a/pymc/tests/test_minibatches.py +++ b/pymc/tests/test_minibatches.py @@ -170,11 +170,11 @@ class TestScaling: def test_density_scaling(self): with pm.Model() as model1: Normal("n", observed=[[1]], total_size=1) - p1 = aesara.function([], model1.logpt) + p1 = aesara.function([], model1.logpt()) with pm.Model() as model2: Normal("n", observed=[[1]], total_size=2) - p2 = aesara.function([], model2.logpt) + p2 = aesara.function([], model2.logpt()) assert p1() * 2 == p2() def test_density_scaling_with_generator(self): @@ -189,12 +189,12 @@ def true_dens(): # We have same size models with pm.Model() as model1: Normal("n", observed=gen1(), total_size=100) - p1 = aesara.function([], model1.logpt) + p1 = aesara.function([], model1.logpt()) with pm.Model() as model2: gen_var = generator(gen2()) Normal("n", observed=gen_var, total_size=100) - p2 = aesara.function([], model2.logpt) + p2 = aesara.function([], model2.logpt()) for i in range(10): _1, _2, _t = p1(), p2(), next(t) @@ -208,12 +208,12 @@ def test_gradient_with_scaling(self): genvar = generator(gen1()) m = Normal("m") Normal("n", observed=genvar, total_size=1000) - grad1 = aesara.function([m.tag.value_var], at.grad(model1.logpt, m.tag.value_var)) + grad1 = aesara.function([m.tag.value_var], at.grad(model1.logpt(), m.tag.value_var)) with pm.Model() as model2: m = Normal("m") shavar = aesara.shared(np.ones((1000, 100))) Normal("n", observed=shavar) - grad2 = aesara.function([m.tag.value_var], at.grad(model2.logpt, m.tag.value_var)) + grad2 = aesara.function([m.tag.value_var], at.grad(model2.logpt(), m.tag.value_var)) for i in range(10): shavar.set_value(np.ones((100, 100)) * i) @@ -224,27 +224,27 @@ def test_gradient_with_scaling(self): def test_multidim_scaling(self): with pm.Model() as model0: Normal("n", observed=[[1, 1], [1, 1]], total_size=[]) - p0 = aesara.function([], model0.logpt) + p0 = aesara.function([], model0.logpt()) with pm.Model() as model1: Normal("n", observed=[[1, 1], [1, 1]], total_size=[2, 2]) - p1 = aesara.function([], model1.logpt) + p1 = aesara.function([], model1.logpt()) with pm.Model() as model2: Normal("n", observed=[[1], [1]], total_size=[2, 2]) - p2 = aesara.function([], model2.logpt) + p2 = aesara.function([], model2.logpt()) with pm.Model() as model3: Normal("n", observed=[[1, 1]], total_size=[2, 2]) - p3 = aesara.function([], model3.logpt) + p3 = aesara.function([], model3.logpt()) with pm.Model() as model4: Normal("n", observed=[[1]], total_size=[2, 2]) - p4 = aesara.function([], model4.logpt) + p4 = aesara.function([], model4.logpt()) with pm.Model() as model5: Normal("n", observed=[[1]], total_size=[2, Ellipsis, 2]) - p5 = aesara.function([], model5.logpt) + p5 = aesara.function([], model5.logpt()) _p0 = p0() assert ( np.allclose(_p0, p1()) @@ -258,27 +258,27 @@ def test_common_errors(self): with pytest.raises(ValueError) as e: with pm.Model() as m: Normal("n", observed=[[1]], total_size=[2, Ellipsis, 2, 2]) - m.logpt + m.logpt() assert "Length of" in str(e.value) with pytest.raises(ValueError) as e: with pm.Model() as m: Normal("n", observed=[[1]], total_size=[2, 2, 2]) - m.logpt + m.logpt() assert "Length of" in str(e.value) with pytest.raises(TypeError) as e: with pm.Model() as m: Normal("n", observed=[[1]], total_size="foo") - m.logpt + m.logpt() assert "Unrecognized" in str(e.value) with pytest.raises(TypeError) as e: with pm.Model() as m: Normal("n", observed=[[1]], total_size=["foo"]) - m.logpt + m.logpt() assert "Unrecognized" in str(e.value) with pytest.raises(ValueError) as e: with pm.Model() as m: Normal("n", observed=[[1]], total_size=[Ellipsis, Ellipsis]) - m.logpt + m.logpt() assert "Double Ellipsis" in str(e.value) def test_mixed1(self): @@ -296,11 +296,11 @@ def test_mixed2(self): def test_free_rv(self): with pm.Model() as model4: Normal("n", observed=[[1, 1], [1, 1]], total_size=[2, 2]) - p4 = aesara.function([], model4.logpt) + p4 = aesara.function([], model4.logpt()) with pm.Model() as model5: n = Normal("n", total_size=[2, Ellipsis, 2], size=(2, 2)) - p5 = aesara.function([n.tag.value_var], model5.logpt) + p5 = aesara.function([n.tag.value_var], model5.logpt()) assert p4() == p5(pm.floatX([[1]])) assert p4() == p5(pm.floatX([[1, 1], [1, 1]])) diff --git a/pymc/tests/test_missing.py b/pymc/tests/test_missing.py index 2160ebf6cb..b64d105118 100644 --- a/pymc/tests/test_missing.py +++ b/pymc/tests/test_missing.py @@ -21,7 +21,7 @@ from aesara.graph import graph_inputs from numpy import array, ma -from pymc import logpt +from pymc import joint_logpt from pymc.distributions import Dirichlet, Gamma, Normal, Uniform from pymc.exceptions import ImputationWarning from pymc.model import Model @@ -41,7 +41,7 @@ def test_missing(data): assert "y_missing" in model.named_vars test_point = model.recompute_initial_point() - assert not np.isnan(model.logp(test_point)) + assert not np.isnan(model.compile_logp()(test_point)) with model: prior_trace = sample_prior_predictive(return_inferencedata=False) @@ -59,7 +59,7 @@ def test_missing_with_predictors(): assert "y_missing" in model.named_vars test_point = model.recompute_initial_point() - assert not np.isnan(model.logp(test_point)) + assert not np.isnan(model.compile_logp()(test_point)) with model: prior_trace = sample_prior_predictive(return_inferencedata=False) @@ -144,13 +144,13 @@ def test_double_counting(): with Model(check_bounds=False) as m1: x = Gamma("x", 1, 1, size=4) - logp_val = m1.logp({"x_log__": np.array([0, 0, 0, 0])}) + logp_val = m1.compile_logp()({"x_log__": np.array([0, 0, 0, 0])}) assert logp_val == -4.0 with Model(check_bounds=False) as m2: x = Gamma("x", 1, 1, observed=[1, 1, 1, np.nan]) - logp_val = m2.logp({"x_missing_log__": np.array([0])}) + logp_val = m2.compile_logp()({"x_missing_log__": np.array([0])}) assert logp_val == -4.0 @@ -158,12 +158,14 @@ def test_missing_logp(): with Model() as m: theta1 = Normal("theta1", 0, 5, observed=[0, 1, 2, 3, 4]) theta2 = Normal("theta2", mu=theta1, observed=[0, 1, 2, 3, 4]) - m_logp = m.logp() + m_logp = m.compile_logp()({}) with Model() as m_missing: theta1 = Normal("theta1", 0, 5, observed=np.array([0, 1, np.nan, 3, np.nan])) theta2 = Normal("theta2", mu=theta1, observed=np.array([np.nan, np.nan, 2, np.nan, 4])) - m_missing_logp = m_missing.logp({"theta1_missing": [2, 4], "theta2_missing": [0, 1, 3]}) + m_missing_logp = m_missing.compile_logp()( + {"theta1_missing": [2, 4], "theta2_missing": [0, 1, 3]} + ) assert m_logp == m_missing_logp @@ -188,8 +190,8 @@ def test_missing_multivariate(): # # inp_vals = simplex.forward(np.array([0.3, 0.3, 0.4])).eval() # assert np.isclose( - # m_miss.logp({"x_missing_simplex__": inp_vals}), - # m_unobs.logp_nojac({"x_simplex__": inp_vals}) * 2, + # m_miss.compile_logp()({"x_missing_simplex__": inp_vals}), + # m_unobs.compile_logp(jacobian=False)({"x_simplex__": inp_vals}) * 2, # ) @@ -206,7 +208,7 @@ def test_missing_vector_parameter(): assert np.all(x_draws[:, 0] < 0) assert np.all(x_draws[:, 1] > 0) assert np.isclose( - m.logp({"x_missing": np.array([-10, 10, -10, 10])}), + m.compile_logp()({"x_missing": np.array([-10, 10, -10, 10])}), scipy.stats.norm(scale=0.1).logpdf(0) * 6, ) @@ -228,7 +230,7 @@ def test_missing_symmetric(): x_unobs_rv = m["x_missing"] x_unobs_vv = m.rvs_to_values[x_unobs_rv] - logp = logpt([x_obs_rv, x_unobs_rv], {x_obs_rv: x_obs_vv, x_unobs_rv: x_unobs_vv}) + logp = joint_logpt([x_obs_rv, x_unobs_rv], {x_obs_rv: x_obs_vv, x_unobs_rv: x_unobs_vv}) logp_inputs = list(graph_inputs([logp])) assert x_obs_vv in logp_inputs assert x_unobs_vv in logp_inputs diff --git a/pymc/tests/test_model.py b/pymc/tests/test_model.py index adabb05abd..65f3d00132 100644 --- a/pymc/tests/test_model.py +++ b/pymc/tests/test_model.py @@ -25,6 +25,7 @@ import pandas as pd import pytest import scipy.sparse as sps +import scipy.stats as st from aesara.tensor.random.op import RandomVariable from aesara.tensor.var import TensorConstant @@ -33,7 +34,7 @@ from pymc import Deterministic, Potential from pymc.blocking import DictToArrayBijection, RaveledVars -from pymc.distributions import Normal, logpt_sum, transforms +from pymc.distributions import Normal, transforms from pymc.exceptions import ShapeError from pymc.model import Point, ValueGradFunction from pymc.tests.helpers import SeededTest @@ -327,7 +328,7 @@ def test_aesara_switch_broadcast_edge_cases_1(self): obs = pm.Bernoulli("obs", p=p, observed=data) npt.assert_allclose( - logpt_sum(obs).eval({p.tag.value_var: pm.floatX(np.array(0.0))}), + m.compile_logp(obs)({"p_logodds__": pm.floatX(np.array(0.0))}), np.log(0.5) * 10, ) @@ -343,7 +344,7 @@ def test_aesara_switch_broadcast_edge_cases_2(self): mu = pm.Normal("mu", 0, 5) obs = pm.TruncatedNormal("obs", mu=mu, sigma=1, lower=-1, upper=2, observed=data) - npt.assert_allclose(m.dlogp([m.rvs_to_values[mu]])({"mu": 0}), 2.499424682024436, rtol=1e-5) + npt.assert_allclose(m.compile_dlogp(mu)({"mu": 0}), 2.499424682024436, rtol=1e-5) def test_multiple_observed_rv(): @@ -665,3 +666,87 @@ def test_nested_model_coords(): assert m1.coords is m2.coords assert m1.dim_lengths is m2.dim_lengths assert set(m2.RV_dims) < set(m1.RV_dims) + + +@pytest.mark.parametrize("jacobian", [True, False]) +def test_model_logp(jacobian): + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + y = pm.LogNormal("y", 0, 1, size=2) + + test_vals = np.array([0.0, 1.0]) + + expected_x_logp = st.norm().logpdf(test_vals) + expected_y_logp = expected_x_logp.copy() + if not jacobian: + expected_y_logp -= np.array([0.0, 1.0]) + + x_logp, y_logp = m.compile_logp(sum=False, jacobian=jacobian)( + {"x": test_vals, "y_log__": test_vals} + ) + assert np.all(np.isclose(x_logp, expected_x_logp)) + assert np.all(np.isclose(y_logp, expected_y_logp)) + + x_logp2 = m.compile_logp(vars=[x], sum=False, jacobian=jacobian)({"x": test_vals}) + assert np.all(np.isclose(x_logp2, expected_x_logp)) + + y_logp2 = m.compile_logp(vars=[y], sum=False, jacobian=jacobian)({"y_log__": test_vals}) + assert np.all(np.isclose(y_logp2, expected_y_logp)) + + logp_sum = m.compile_logp(sum=True, jacobian=jacobian)({"x": test_vals, "y_log__": test_vals}) + assert np.isclose(logp_sum, expected_x_logp.sum() + expected_y_logp.sum()) + + +@pytest.mark.parametrize("jacobian", [True, False]) +def test_model_dlogp(jacobian): + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + y = pm.LogNormal("y", 0, 1, size=2) + + test_vals = np.array([0.0, -1.0]) + state = {"x": test_vals, "y_log__": test_vals} + + expected_x_dlogp = expected_y_dlogp = np.array([0.0, 1.0]) + if not jacobian: + expected_y_dlogp = np.array([-1.0, 0.0]) + + dlogps = m.compile_dlogp(jacobian=jacobian)(state) + assert np.all(np.isclose(dlogps[:2], expected_x_dlogp)) + assert np.all(np.isclose(dlogps[2:], expected_y_dlogp)) + + x_dlogp2 = m.compile_dlogp(vars=[x], jacobian=jacobian)(state) + assert np.all(np.isclose(x_dlogp2, expected_x_dlogp)) + + y_dlogp2 = m.compile_dlogp(vars=[y], jacobian=jacobian)(state) + assert np.all(np.isclose(y_dlogp2, expected_y_dlogp)) + + +@pytest.mark.parametrize("jacobian", [True, False]) +def test_model_d2logp(jacobian): + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + y = pm.LogNormal("y", 0, 1, size=2) + + test_vals = np.array([0.0, -1.0]) + state = {"x": test_vals, "y_log__": test_vals} + + expected_x_d2logp = expected_y_d2logp = np.eye(2) + + dlogps = m.compile_d2logp(jacobian=jacobian)(state) + assert np.all(np.isclose(dlogps[:2, :2], expected_x_d2logp)) + assert np.all(np.isclose(dlogps[2:, 2:], expected_y_d2logp)) + + x_dlogp2 = m.compile_d2logp(vars=[x], jacobian=jacobian)(state) + assert np.all(np.isclose(x_dlogp2, expected_x_d2logp)) + + y_dlogp2 = m.compile_d2logp(vars=[y], jacobian=jacobian)(state) + assert np.all(np.isclose(y_dlogp2, expected_y_d2logp)) + + +def test_deterministic(): + with pm.Model() as model: + x = pm.Normal("x", 0, 1) + y = pm.Deterministic("y", x ** 2) + + assert model.y == y + assert model["y"] == y diff --git a/pymc/tests/test_model_func.py b/pymc/tests/test_model_func.py deleted file mode 100644 index 38c2329081..0000000000 --- a/pymc/tests/test_model_func.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright 2020 The PyMC Developers -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import numpy as np -import scipy.stats as sp - -import pymc as pm - -from pymc.tests.checks import close_to -from pymc.tests.models import mv_simple, simple_model - -tol = 2.0 ** -11 - - -def test_logp(): - start, model, (mu, sig) = simple_model() - lp = model.fastlogp - lp(start) - close_to(lp(start), sp.norm.logpdf(start["x"], mu, sig).sum(), tol) - - -def test_dlogp(): - start, model, (mu, sig) = simple_model() - dlogp = model.fastdlogp() - close_to(dlogp(start), -(start["x"] - mu) / sig ** 2, 1.0 / sig ** 2 / 100.0) - - -def test_dlogp2(): - start, model, (_, sig) = mv_simple() - H = np.linalg.inv(sig) - d2logp = model.fastd2logp() - close_to(d2logp(start), H, np.abs(H / 100.0)) - - -def test_deterministic(): - with pm.Model() as model: - x = pm.Normal("x", 0, 1) - y = pm.Deterministic("y", x ** 2) - - assert model.y == y - assert model["y"] == y diff --git a/pymc/tests/test_ode.py b/pymc/tests/test_ode.py index a4afb59783..a57638af7d 100644 --- a/pymc/tests/test_ode.py +++ b/pymc/tests/test_ode.py @@ -200,7 +200,7 @@ def system_1(y, t, p): with pm.Model() as model_1: forward = ode_model(theta=[alpha], y0=[y0]) y = pm.Normal("y", mu=forward, sd=1, observed=yobs) - pymc_logp = model_1.logp() + pymc_logp = model_1.compile_logp()({}) np.testing.assert_allclose(manual_logp, pymc_logp) diff --git a/pymc/tests/test_profile.py b/pymc/tests/test_profile.py index b47f3d9fe4..60c9260122 100644 --- a/pymc/tests/test_profile.py +++ b/pymc/tests/test_profile.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pymc as pm - from pymc.tests.models import simple_model @@ -22,12 +20,12 @@ def setup_method(self): _, self.model, _ = simple_model() def test_profile_model(self): - assert self.model.profile(self.model.logpt).fct_call_time > 0 + assert self.model.profile(self.model.logpt()).fct_call_time > 0 def test_profile_variable(self): rv = self.model.basic_RVs[0] - assert self.model.profile(pm.logpt(rv, self.model.rvs_to_values[rv])).fct_call_time + assert self.model.profile(self.model.logpt(vars=[rv], sum=False)).fct_call_time def test_profile_count(self): count = 1005 - assert self.model.profile(self.model.logpt, n=count).fct_callcount == count + assert self.model.profile(self.model.logpt(), n=count).fct_callcount == count diff --git a/pymc/tests/test_shared.py b/pymc/tests/test_shared.py index b4f0101685..b5c5ec0220 100644 --- a/pymc/tests/test_shared.py +++ b/pymc/tests/test_shared.py @@ -14,6 +14,7 @@ import aesara import numpy as np +import scipy.stats as st import pymc as pm @@ -26,7 +27,9 @@ def test_deterministic(self): data_values = np.array([0.5, 0.4, 5, 2]) X = aesara.shared(np.asarray(data_values, dtype=aesara.config.floatX), borrow=True) pm.Normal("y", 0, 1, observed=X) - model.logp(model.recompute_initial_point()) + assert np.all( + np.isclose(model.compile_logp(sum=False)({}), st.norm().logpdf(data_values)) + ) def test_sample(self): x = np.random.normal(size=100) diff --git a/pymc/tests/test_smc.py b/pymc/tests/test_smc.py index 2b700afe64..21cec4acd6 100644 --- a/pymc/tests/test_smc.py +++ b/pymc/tests/test_smc.py @@ -144,7 +144,16 @@ def test_marginal_likelihood(self): marginals.append(trace.report.log_marginal_likelihood) # compare to the analytical result - assert abs(np.exp(np.nanmean(marginals[1]) - np.nanmean(marginals[0])) - 4.0) <= 1 + assert ( + np.abs( + np.exp( + np.nanmean(np.array(marginals[1], dtype=float)) + - np.nanmean(np.array(marginals[0], dtype=float)) + - 4.0 + ) + ) + <= 1 + ) def test_start(self): with pm.Model() as model: @@ -299,7 +308,7 @@ def setup_class(self): s = pm.Simulator("s", self.normal_sim, a, b, observed=self.data) def test_one_gaussian(self): - assert self.count_rvs(self.SMABC_test.logpt) == 1 + assert self.count_rvs(self.SMABC_test.logpt()) == 1 with self.SMABC_test: trace = pm.sample_smc(draws=1000, chains=1, return_inferencedata=False) @@ -333,7 +342,7 @@ def test_custom_dist_sum_stat(self): observed=self.data, ) - assert self.count_rvs(m.logpt) == 1 + assert self.count_rvs(m.logpt()) == 1 with m: pm.sample_smc(draws=100) @@ -354,7 +363,7 @@ def test_custom_dist_sum_stat_scalar(self): sum_stat=self.quantiles, observed=scalar_data, ) - assert self.count_rvs(m.logpt) == 1 + assert self.count_rvs(m.logpt()) == 1 with pm.Model() as m: s = pm.Simulator( @@ -366,10 +375,10 @@ def test_custom_dist_sum_stat_scalar(self): sum_stat="mean", observed=scalar_data, ) - assert self.count_rvs(m.logpt) == 1 + assert self.count_rvs(m.logpt()) == 1 def test_model_with_potential(self): - assert self.count_rvs(self.SMABC_potential.logpt) == 1 + assert self.count_rvs(self.SMABC_potential.logpt()) == 1 with self.SMABC_potential: trace = pm.sample_smc(draws=100, chains=1, return_inferencedata=False) @@ -413,17 +422,17 @@ def test_multiple_simulators(self): observed=data2, ) - assert self.count_rvs(m.logpt) == 2 + assert self.count_rvs(m.logpt()) == 2 # Check that the logps use the correct methods a_val = m.rvs_to_values[a] sim1_val = m.rvs_to_values[sim1] - logp_sim1 = pm.logpt(sim1, sim1_val) + logp_sim1 = pm.joint_logpt(sim1, sim1_val) logp_sim1_fn = aesara.function([a_val], logp_sim1) b_val = m.rvs_to_values[b] sim2_val = m.rvs_to_values[sim2] - logp_sim2 = pm.logpt(sim2, sim2_val) + logp_sim2 = pm.joint_logpt(sim2, sim2_val) logp_sim2_fn = aesara.function([b_val], logp_sim2) assert any( @@ -463,7 +472,7 @@ def test_nested_simulators(self): observed=data, ) - assert self.count_rvs(m.logpt) == 2 + assert self.count_rvs(m.logpt()) == 2 with m: trace = pm.sample_smc(return_inferencedata=False) diff --git a/pymc/tests/test_starting.py b/pymc/tests/test_starting.py index e4b88bec4b..a989ed2466 100644 --- a/pymc/tests/test_starting.py +++ b/pymc/tests/test_starting.py @@ -46,7 +46,7 @@ def test_accuracy_non_normal(): close_to(newstart["x"], mu, select_by_precision(float64=1e-5, float32=1e-4)) -@pytest.mark.xfail(reason="find_MAP fails with derivatives") +@pytest.mark.xfail(reason="first call to find_MAP is failing") def test_find_MAP_discrete(): tol = 2.0 ** -11 alpha = 4 @@ -68,15 +68,12 @@ def test_find_MAP_discrete(): assert map_est2["ss"] == 14 -@pytest.mark.xfail(reason="find_MAP fails with derivatives") def test_find_MAP_no_gradient(): _, model = simple_arbitrary_det() with model: find_MAP() -@pytest.mark.skip(reason="test is slow because it's failing") -@pytest.mark.xfail(reason="find_MAP fails with derivatives") def test_find_MAP(): tol = 2.0 ** -11 # 16 bit machine epsilon, a low bar data = np.random.randn(100) diff --git a/pymc/tests/test_step.py b/pymc/tests/test_step.py index 246fff59db..8f23f72497 100644 --- a/pymc/tests/test_step.py +++ b/pymc/tests/test_step.py @@ -1030,8 +1030,13 @@ def test_sampler_stats(self): # Assert model logp is computed correctly: computing post-sampling # and tracking while sampling should give same results. + model_logp_fn = model.compile_logp() model_logp_ = np.array( - [model.logp(trace.point(i, chain=c)) for c in trace.chains for i in range(len(trace))] + [ + model_logp_fn(trace.point(i, chain=c)) + for c in trace.chains + for i in range(len(trace)) + ] ) assert (trace.model_logp == model_logp_).all() diff --git a/pymc/tests/test_transforms.py b/pymc/tests/test_transforms.py index d38300cdc9..df1fb57096 100644 --- a/pymc/tests/test_transforms.py +++ b/pymc/tests/test_transforms.py @@ -24,7 +24,7 @@ import pymc.distributions.transforms as tr from pymc.aesaraf import floatX, jacobian -from pymc.distributions import logpt +from pymc.distributions import joint_logpt from pymc.tests.checks import close_to, close_to_logical from pymc.tests.helpers import SeededTest from pymc.tests.test_distributions import ( @@ -296,10 +296,12 @@ def check_transform_elementwise_logp(self, model): x_val_untransf = at.constant(test_array_untransf).type() jacob_det = transform.log_jac_det(test_array_transf, *x.owner.inputs) - assert logpt(x, sum=False).ndim == x.ndim == jacob_det.ndim + assert joint_logpt(x, sum=False).ndim == x.ndim == jacob_det.ndim - v1 = logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf}) - v2 = logpt(x, x_val_untransf, transformed=False).eval({x_val_untransf: test_array_untransf}) + v1 = joint_logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf}) + v2 = joint_logpt(x, x_val_untransf, transformed=False).eval( + {x_val_untransf: test_array_untransf} + ) close_to(v1, v2, tol) def check_vectortransform_elementwise_logp(self, model): @@ -317,13 +319,15 @@ def check_vectortransform_elementwise_logp(self, model): jacob_det = transform.log_jac_det(test_array_transf, *x.owner.inputs) # Original distribution is univariate if x.owner.op.ndim_supp == 0: - assert logpt(x, sum=False).ndim == x.ndim == (jacob_det.ndim + 1) + assert joint_logpt(x, sum=False).ndim == x.ndim == (jacob_det.ndim + 1) # Original distribution is multivariate else: - assert logpt(x, sum=False).ndim == (x.ndim - 1) == jacob_det.ndim + assert joint_logpt(x, sum=False).ndim == (x.ndim - 1) == jacob_det.ndim - a = logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf}) - b = logpt(x, x_val_untransf, transformed=False).eval({x_val_untransf: test_array_untransf}) + a = joint_logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf}) + b = joint_logpt(x, x_val_untransf, transformed=False).eval( + {x_val_untransf: test_array_untransf} + ) # Hack to get relative tolerance close_to(a, b, np.abs(0.5 * (a + b) * tol)) diff --git a/pymc/tuning/scaling.py b/pymc/tuning/scaling.py index 432c72ea76..28471a0dda 100644 --- a/pymc/tuning/scaling.py +++ b/pymc/tuning/scaling.py @@ -59,7 +59,7 @@ def find_hessian(point, vars=None, model=None): Variables for which Hessian is to be calculated. """ model = modelcontext(model) - H = model.fastd2logp(vars) + H = model.compile_d2logp(vars) return H(Point(point, filter_model_vars=True, model=model)) @@ -75,7 +75,7 @@ def find_hessian_diag(point, vars=None, model=None): Variables for which Hessian is to be calculated. """ model = modelcontext(model) - H = model.fastfn(hessian_diag(model.logpt, vars)) + H = model.compile_fn(hessian_diag(model.logpt(), vars)) return H(Point(point, model=model)) diff --git a/pymc/tuning/starting.py b/pymc/tuning/starting.py index 3cd3613a35..aedeb8fbc0 100644 --- a/pymc/tuning/starting.py +++ b/pymc/tuning/starting.py @@ -95,6 +95,9 @@ def find_MAP( vars = model.cont_vars if not vars: raise ValueError("Model has no unobserved continuous variables.") + else: + vars = [model.rvs_to_values.get(var, var) for var in vars] + vars = inputvars(vars) disc_vars = list(typefilter(vars, discrete_types)) allinmodel(vars, model) @@ -113,18 +116,15 @@ def find_MAP( # TODO: If the mapping is fixed, we can simply create graphs for the # mapping and avoid all this bijection overhead - def logp_func(x): - return DictToArrayBijection.mapf(model.fastlogp_nojac)(RaveledVars(x, x0.point_map_info)) + compiled_logp_func = DictToArrayBijection.mapf(model.compile_logp(jacobian=False)) + logp_func = lambda x: compiled_logp_func(RaveledVars(x, x0.point_map_info)) + rvs = [model.values_to_rvs[value] for value in vars] try: # This might be needed for calls to `dlogp_func` # start_map_info = tuple((v.name, v.shape, v.dtype) for v in vars) - - def dlogp_func(x): - return DictToArrayBijection.mapf(model.fastdlogp_nojac(vars))( - RaveledVars(x, x0.point_map_info) - ) - + compiled_dlogp_func = DictToArrayBijection.mapf(model.compile_dlogp(rvs, jacobian=False)) + dlogp_func = lambda x: compiled_dlogp_func(RaveledVars(x, x0.point_map_info)) compute_gradient = True except (AttributeError, NotImplementedError, tg.NullTypeGradError): compute_gradient = False @@ -166,7 +166,7 @@ def dlogp_func(x): vars = get_default_varnames(model.unobserved_value_vars, include_transformed) mx = { var.name: value - for var, value in zip(vars, model.fastfn(vars)(DictToArrayBijection.rmap(mx0))) + for var, value in zip(vars, model.compile_fn(vars)(DictToArrayBijection.rmap(mx0))) } if return_raw: