likelihood.py: update Poisson likelihood

 - allow likelihood to use a rate function with a
   dependent variable
 - refs Monash/tupak!132

tupak/core/likelihood.py

@@ -82,7 +82,7 @@ class GaussianLikelihood(Likelihood):
     @staticmethod
     def _infer_parameters_from_function(func):
         """ Infers the arguments of function (except the first arg which is
-            assumed to be the dep. variable
+            assumed to be the dep. variable)
         """
         parameters = inspect.getargspec(func).args
         parameters.pop(0)
@@ -112,43 +112,47 @@ class GaussianLikelihood(Likelihood):
 
 
 class PoissonLikelihood(Likelihood):
-    def __init__(self, x, func):
+    def __init__(self, x, counts, func):
         """
         A general Poisson likelihood for a rate - the model parameters are
         inferred from the arguments of function, which provides a rate.
 
-        Note: This can only be used for fitting a single rate and not a
-        changing rate through a generalised linear model (GLM).
-
         Parameters
         ----------
+
         x: array_like
+            A dependent variable at which the Poisson rates will be calculated
+        counts: array_like
             The data to analyse - this must be a set of non-negative integers,
             each being the number of events within some interval.
         func:
             The python function providing the rate of events per interval to
-            fit to the data. The arguments will require priors and will be
-            sampled over (unless a fixed value is given).
+            fit to the data. The function must be defined with the first
+            argument being a dependent parameter (although this does not have
+            to be used by the function if not required). The subsequent
+            arguments will require priors and will be sampled over (unless a
+            fixed value is given).
         """
 
         parameters = self._infer_parameters_from_function(func)
         Likelihood.__init__(self, dict.fromkeys(parameters))
 
-        self.x = x
+        self.x = x           # the dependent variable
+        self.counts = counts # the counts
 
         # check values are non-negative integers
-        if isinstance(self.x, int):
+        if isinstance(self.counts, int):
             # convert to numpy array if passing a single integer
-            self.x = np.array(self.x)
+            self.counts = np.array(self.counts)
 
         try:
             # use bincount to check all values are non-negative integers
-            _ = np.bincount(self.x)
+            _ = np.bincount(self.counts)
         except ValueError:
             raise ValueError("Data must be non-negative integers")
 
         # save sum of log factorial of counts
-        self.sumlogfactorial = np.sum(gammaln(self.x + 1))
+        self.sumlogfactorial = np.sum(gammaln(self.counts + 1))
 
         self.function = func
 
@@ -158,7 +162,7 @@ class PoissonLikelihood(Likelihood):
     @staticmethod
     def _infer_parameters_from_function(func):
         """ Infers the arguments of function (except the first arg which is
-            assumed to be the dep. variable
+            assumed to be the dep. variable)
         """
         parameters = inspect.getargspec(func).args
         parameters.pop(0)
@@ -167,18 +171,32 @@ class PoissonLikelihood(Likelihood):
     @property
     def N(self):
         """ The number of data points """
-        return len(self.x)
+        return len(self.counts)
 
     def log_likelihood(self):
         # This sets up the function only parameters (i.e. not sigma)
         model_parameters = {k: self.parameters[k] for k in self.function_keys}
 
         # Calculate the rate
-        rate = self.function(**model_parameters)
-
-        # check rate is positive
-        if rate < 0.:
-            raise ValueError("Poisson rate function returns a negative value!")
-
-        # Return the summed log likelihood
-        return -self.N*rate + np.sum(self.x*np.log(rate)) - self.sumlogfactorial
+        rate = self.function(self.x, **model_parameters)
+
+        # check if rate is a single value
+        if isinstance(rate, float):
+            # check rate is positive
+            if rate < 0.:
+                raise ValueError(("Poisson rate function returns a negative ",
+                                  "value!"))
+
+            # Return the summed log likelihood
+            return -self.N*rate + np.sum(self.counts*np.log(rate)) -
+                   self.sumlogfactorial
+        elif isinstance(rate, np.ndarray):
+            # check rates are positive
+            if np.any(rate < 0.):
+                raise ValueError(("Poisson rate function returns a negative")),
+                                  " value!"))
+
+            return np.sum(-rate + self.counts*np.log(rate)) - 
+                   self.sumlogfactorial
+        else:
+            raise ValueError("Poisson rate function returns wrong value type!")
\ No newline at end of file