1#   Copyright 2020 The PyMC Developers
2#
3#   Licensed under the Apache License, Version 2.0 (the "License");
4#   you may not use this file except in compliance with the License.
5#   You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#   Unless required by applicable law or agreed to in writing, software
10#   distributed under the License is distributed on an "AS IS" BASIS,
11#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#   See the License for the specific language governing permissions and
13#   limitations under the License.
14
15import numbers
16
17from copy import copy
18
19import numpy as np
20import theano.tensor as tt
21
22from pymc3 import distributions as pm_dists
23from pymc3.model import modelcontext
24
25__all__ = ["Normal", "StudentT", "Binomial", "Poisson", "NegativeBinomial"]
26
27# Define link functions
28
29# Hack as assigning a function in the class definition automatically binds
30# it as a method.
31
32
33class Identity:
34    def __call__(self, x):
35        return x
36
37
38identity = Identity()
39logit = tt.nnet.sigmoid
40inverse = tt.inv
41exp = tt.exp
42
43
44class Family:
45    """Base class for Family of likelihood distribution and link functions."""
46
47    priors = {}
48    link = None
49
50    def __init__(self, **kwargs):
51        # Overwrite defaults
52        for key, val in kwargs.items():
53            if key == "priors":
54                self.priors = copy(self.priors)
55                self.priors.update(val)
56            else:
57                setattr(self, key, val)
58
59    def _get_priors(self, model=None, name=""):
60        """Return prior distributions of the likelihood.
61
62        Returns
63        -------
64        dict: mapping name -> pymc3 distribution
65        """
66        if name:
67            name = f"{name}_"
68        model = modelcontext(model)
69        priors = {}
70        for key, val in self.priors.items():
71            if isinstance(val, (numbers.Number, np.ndarray, np.generic)):
72                priors[key] = val
73            else:
74                priors[key] = model.Var(f"{name}{key}", val)
75
76        return priors
77
78    def create_likelihood(self, name, y_est, y_data, model=None):
79        """Create likelihood distribution of observed data.
80
81        Parameters
82        ----------
83        y_est: theano.tensor
84            Estimate of dependent variable
85        y_data: array
86            Observed dependent variable
87        """
88        priors = self._get_priors(model=model, name=name)
89        # Wrap y_est in link function
90        priors[self.parent] = self.link(y_est)
91        if name:
92            name = f"{name}_"
93        return self.likelihood(f"{name}y", observed=y_data, **priors)
94
95    def __repr__(self):
96        return """Family {klass}:
97    Likelihood  : {likelihood}({parent})
98    Priors      : {priors}
99    Link function: {link}.""".format(
100            klass=self.__class__,
101            likelihood=self.likelihood.__name__,
102            parent=self.parent,
103            priors=self.priors,
104            link=self.link,
105        )
106
107
108class StudentT(Family):
109    link = identity
110    likelihood = pm_dists.StudentT
111    parent = "mu"
112    priors = {"lam": pm_dists.HalfCauchy.dist(beta=10, testval=1.0), "nu": 1}
113
114
115class Normal(Family):
116    link = identity
117    likelihood = pm_dists.Normal
118    parent = "mu"
119    priors = {"sd": pm_dists.HalfCauchy.dist(beta=10, testval=1.0)}
120
121
122class Binomial(Family):
123    link = logit
124    likelihood = pm_dists.Binomial
125    parent = "p"
126    priors = {"n": 1}
127
128
129class Poisson(Family):
130    link = exp
131    likelihood = pm_dists.Poisson
132    parent = "mu"
133    priors = {"mu": pm_dists.HalfCauchy.dist(beta=10, testval=1.0)}
134
135
136class NegativeBinomial(Family):
137    link = exp
138    likelihood = pm_dists.NegativeBinomial
139    parent = "mu"
140    priors = {
141        "mu": pm_dists.HalfCauchy.dist(beta=10, testval=1.0),
142        "alpha": pm_dists.HalfCauchy.dist(beta=10, testval=1.0),
143    }
144