|
| 1 | +import numpy as np |
| 2 | +import theano.tensor as tt |
| 3 | + |
| 4 | +from pymc3.distributions.distribution import ( |
| 5 | + Distribution, Discrete, draw_values, generate_samples) |
| 6 | +from pymc3.distributions import transforms |
| 7 | +from pymc3.distributions.dist_math import bound |
| 8 | + |
| 9 | +__all__ = ['Bound'] |
| 10 | + |
| 11 | + |
| 12 | +class _Bounded(Distribution): |
| 13 | + R""" |
| 14 | + An upper, lower or upper+lower bounded distribution |
| 15 | +
|
| 16 | + Parameters |
| 17 | + ---------- |
| 18 | + distribution : pymc3 distribution |
| 19 | + Distribution to be transformed into a bounded distribution |
| 20 | + lower : float (optional) |
| 21 | + Lower bound of the distribution, set to -inf to disable. |
| 22 | + upper : float (optional) |
| 23 | + Upper bound of the distribibution, set to inf to disable. |
| 24 | + tranform : 'infer' or object |
| 25 | + If 'infer', infers the right transform to apply from the supplied bounds. |
| 26 | + If transform object, has to supply .forward() and .backward() methods. |
| 27 | + See pymc3.distributions.transforms for more information. |
| 28 | + """ |
| 29 | + |
| 30 | + def __init__(self, distribution, lower, upper, |
| 31 | + transform='infer', *args, **kwargs): |
| 32 | + if lower == -np.inf: |
| 33 | + lower = None |
| 34 | + if upper == np.inf: |
| 35 | + upper = None |
| 36 | + |
| 37 | + if lower is not None: |
| 38 | + lower = tt.as_tensor_variable(lower) |
| 39 | + if upper is not None: |
| 40 | + upper = tt.as_tensor_variable(upper) |
| 41 | + |
| 42 | + self.lower = lower |
| 43 | + self.upper = upper |
| 44 | + |
| 45 | + if transform == 'infer': |
| 46 | + if lower is None and upper is None: |
| 47 | + transform = None |
| 48 | + default = None |
| 49 | + elif lower is not None and upper is not None: |
| 50 | + transform = transforms.interval(lower, upper) |
| 51 | + default = 0.5 * (lower + upper) |
| 52 | + elif upper is not None: |
| 53 | + transform = transforms.upperbound(upper) |
| 54 | + default = upper - 1 |
| 55 | + else: |
| 56 | + transform = transforms.lowerbound(lower) |
| 57 | + default = lower + 1 |
| 58 | + else: |
| 59 | + default = None |
| 60 | + |
| 61 | + # We don't use transformations for dicrete variables |
| 62 | + if issubclass(distribution, Discrete): |
| 63 | + transform = None |
| 64 | + |
| 65 | + kwargs['transform'] = transform |
| 66 | + self._wrapped = distribution.dist(*args, **kwargs) |
| 67 | + self._default = default |
| 68 | + |
| 69 | + if issubclass(distribution, Discrete) and default is not None: |
| 70 | + default = default.astype(str(self._wrapped.default().dtype)) |
| 71 | + |
| 72 | + if default is None: |
| 73 | + defaults = self._wrapped.defaults |
| 74 | + for name in defaults: |
| 75 | + setattr(self, name, getattr(self._wrapped, name)) |
| 76 | + else: |
| 77 | + defaults = ('_default',) |
| 78 | + |
| 79 | + super(_Bounded, self).__init__( |
| 80 | + shape=self._wrapped.shape, |
| 81 | + dtype=self._wrapped.dtype, |
| 82 | + testval=self._wrapped.testval, |
| 83 | + defaults=defaults, |
| 84 | + transform=self._wrapped.transform) |
| 85 | + |
| 86 | + def _random(self, lower, upper, point=None, size=None): |
| 87 | + lower = np.asarray(lower) |
| 88 | + upper = np.asarray(upper) |
| 89 | + if lower.size > 1 or upper.size > 1: |
| 90 | + raise ValueError('Drawing samples from distributions with ' |
| 91 | + 'array-valued bounds is not supported.') |
| 92 | + samples = np.zeros(size, dtype=self.dtype).flatten() |
| 93 | + i, n = 0, len(samples) |
| 94 | + while i < len(samples): |
| 95 | + sample = self._wrapped.random(point=point, size=n) |
| 96 | + select = sample[np.logical_and(sample >= lower, sample <= upper)] |
| 97 | + samples[i:(i + len(select))] = select[:] |
| 98 | + i += len(select) |
| 99 | + n -= len(select) |
| 100 | + if size is not None: |
| 101 | + return np.reshape(samples, size) |
| 102 | + else: |
| 103 | + return samples |
| 104 | + |
| 105 | + def random(self, point=None, size=None, repeat=None): |
| 106 | + if self.lower is None and self.upper is None: |
| 107 | + return self._wrapped.random(point=point, size=size) |
| 108 | + elif self.lower is not None and self.upper is not None: |
| 109 | + lower, upper = draw_values([self.lower, self.upper], point=point) |
| 110 | + return generate_samples(self._random, lower, upper, point, |
| 111 | + dist_shape=self.shape, |
| 112 | + size=size) |
| 113 | + elif self.lower is not None: |
| 114 | + lower = draw_values([self.lower], point=point) |
| 115 | + return generate_samples(self._random, lower, np.inf, point, |
| 116 | + dist_shape=self.shape, |
| 117 | + size=size) |
| 118 | + else: |
| 119 | + upper = draw_values([self.upper], point=point) |
| 120 | + return generate_samples(self._random, -np.inf, upper, point, |
| 121 | + dist_shape=self.shape, |
| 122 | + size=size) |
| 123 | + |
| 124 | + def logp(self, value): |
| 125 | + logp = self._wrapped.logp(value) |
| 126 | + bounds = [] |
| 127 | + if self.lower is not None: |
| 128 | + bounds.append(value >= self.lower) |
| 129 | + if self.upper is not None: |
| 130 | + bounds.append(value <= self.upper) |
| 131 | + if len(bounds) > 0: |
| 132 | + return bound(logp, *bounds) |
| 133 | + else: |
| 134 | + return logp |
| 135 | + |
| 136 | + |
| 137 | +class Bound(object): |
| 138 | + R""" |
| 139 | + Create a new upper, lower or upper+lower bounded distribution. |
| 140 | +
|
| 141 | + The resulting distribution is not normalized anymore. This |
| 142 | + is usually fine if the bounds are constants. If you need |
| 143 | + truncated distributions, use `Bound` in combination with |
| 144 | + a `pm.Potential` with the cumulative probability function. |
| 145 | +
|
| 146 | + The bounds are inclusive for discrete distributions. |
| 147 | +
|
| 148 | + Parameters |
| 149 | + ---------- |
| 150 | + distribution : pymc3 distribution |
| 151 | + Distribution to be transformed into a bounded distribution. |
| 152 | + lower : float or array like, optional |
| 153 | + Lower bound of the distribution. |
| 154 | + upper : float or array like, optional |
| 155 | + Upper bound of the distribution. |
| 156 | +
|
| 157 | + Example |
| 158 | + ------- |
| 159 | + # Bounded distribution can be defined before the model context |
| 160 | + PositiveNormal = pm.Bound(pm.Normal, lower=0.0) |
| 161 | + with pm.Model(): |
| 162 | + par1 = PositiveNormal('par1', mu=0.0, sd=1.0, testval=1.0) |
| 163 | + # or within the model context |
| 164 | + NegativeNormal = pm.Bound(pm.Normal, upper=0.0) |
| 165 | + par2 = NegativeNormal('par2', mu=0.0, sd=1.0, testval=1.0) |
| 166 | +
|
| 167 | + # or you can define it implicitly within the model context |
| 168 | + par3 = pm.Bound(pm.Normal, lower=-1.0, upper=1.0)( |
| 169 | + 'par3', mu=0.0, sd=1.0, testval=1.0) |
| 170 | + """ |
| 171 | + |
| 172 | + def __init__(self, distribution, lower=None, upper=None): |
| 173 | + self.distribution = distribution |
| 174 | + self.lower = lower |
| 175 | + self.upper = upper |
| 176 | + |
| 177 | + def __call__(self, *args, **kwargs): |
| 178 | + if 'observed' in kwargs: |
| 179 | + raise ValueError('Observed Bound distributions are not allowed. ' |
| 180 | + 'If you want to model truncated data ' |
| 181 | + 'you can use a pm.Potential in combination ' |
| 182 | + 'with the cumulative probability function. See ' |
| 183 | + 'pymc3/examples/censored_data.py for an example.') |
| 184 | + first, args = args[0], args[1:] |
| 185 | + |
| 186 | + return _Bounded(first, self.distribution, self.lower, self.upper, |
| 187 | + *args, **kwargs) |
| 188 | + |
| 189 | + def dist(self, *args, **kwargs): |
| 190 | + return _Bounded.dist(self.distribution, self.lower, self.upper, |
| 191 | + *args, **kwargs) |
0 commit comments