73 lines
1.9 KiB
Python
Executable File
73 lines
1.9 KiB
Python
Executable File
import numpy as np
|
|
from prml.nn.math.gamma import gamma
|
|
from prml.nn.math.log import log
|
|
from prml.nn.math.product import prod
|
|
from prml.nn.math.sum import sum
|
|
from prml.nn.random.random import RandomVariable
|
|
from prml.nn.tensor.tensor import Tensor
|
|
|
|
|
|
class Dirichlet(RandomVariable):
|
|
"""
|
|
Dirichlet distribution
|
|
Parameters
|
|
----------
|
|
alpha : (..., K) tensor_like
|
|
pseudo-count of each outcome
|
|
axis : int
|
|
axis along which represents each outcome
|
|
data : tensor_like
|
|
realization
|
|
p : RandomVariable
|
|
original distribution of a model
|
|
Attributes
|
|
----------
|
|
n_category : int
|
|
number of categories
|
|
"""
|
|
|
|
def __init__(self, alpha, axis=-1, data=None, p=None):
|
|
super().__init__(data, p)
|
|
assert axis == -1
|
|
self.axis = axis
|
|
self.alpha = self._convert2tensor(alpha)
|
|
|
|
@property
|
|
def alpha(self):
|
|
return self.parameter["alpha"]
|
|
|
|
@alpha.setter
|
|
def alpha(self, alpha):
|
|
self._atleast_ndim(alpha, 1)
|
|
if (alpha.value <= 0).any():
|
|
raise ValueError("alpha must all be positive")
|
|
self.parameter["alpha"] = alpha
|
|
|
|
def forward(self):
|
|
if self.alpha.ndim == 1:
|
|
return Tensor(np.random.dirichlet(self.alpha.value), function=self)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def backward(self):
|
|
raise NotImplementedError
|
|
|
|
def _pdf(self, x):
|
|
return (
|
|
gamma(self.alpha.sum(axis=self.axis))
|
|
* prod(
|
|
x ** (self.alpha - 1)
|
|
/ gamma(self.alpha),
|
|
axis=self.axis
|
|
)
|
|
)
|
|
|
|
def _log_pdf(self, x):
|
|
return (
|
|
log(gamma(self.alpha.sum(axis=self.axis)))
|
|
+ sum(
|
|
(self.alpha - 1) * log(x) - log(gamma(self.alpha)),
|
|
axis=self.axis
|
|
)
|
|
)
|