From 496cdcc52f6fcb95cae33b4be45d6f13a81a6e78 Mon Sep 17 00:00:00 2001 From: Manuel Vera Date: Mon, 3 Oct 2022 20:53:34 -0400 Subject: [PATCH] poisson binomial rough draft --- .../python/distributions/poisson_binomial.py | 101 ++++++++++++++++++ .../distributions/poisson_binomial_test.py | 42 ++++++++ 2 files changed, 143 insertions(+) create mode 100644 tensorflow_probability/python/distributions/poisson_binomial.py create mode 100644 tensorflow_probability/python/distributions/poisson_binomial_test.py diff --git a/tensorflow_probability/python/distributions/poisson_binomial.py b/tensorflow_probability/python/distributions/poisson_binomial.py new file mode 100644 index 0000000000..1eae5c63ed --- /dev/null +++ b/tensorflow_probability/python/distributions/poisson_binomial.py @@ -0,0 +1,101 @@ +# Copyright 2022 The TensorFlow Probability Authors. +# +# 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. +# ============================================================================ +"""The Poisson Binomial distribution class.""" +import tensorflow.compat.v2 as tf +import numpy as np + +from tensorflow_probability.python.internal import dtype_util +from tensorflow_probability.python.distributions import distribution +from tensorflow_probability.python.internal import parameter_properties +from tensorflow_probability.python.internal import tensor_util +from scipy import fft +from math import pi + + +class PoissonBinomial(distribution.Distribution): + + def __init__(self, + probs=None, + validate_args=False, + allow_nan_stats=True, + name='PoissonBinomial'): + """ + Initialize a Poisson Binomial distribution. + + The Poisson Binomial distribution is used to calculate the probability + of the number of successes in independent Bernoulli trials that are not necessarily + identically distributed. Each of the trials has its own probability of success, + so the parameters of this distribution are p1, p2, ..., pn for n Bernoulli trials. + + This naive version uses the fast fourier transform method to obtain the PMF + of the Poisson Binomial distribution [1]. The CDF is obtained through the cumulative + sum of this pdf at a given value of successes. + + ### References + + [1]: Yili Hong. On computing the distribution function for the Poisson binomial distribution. + _Computational Statistics & Data Analysis_, 59:41-51, 2013. + https://doi.org/10.1016/j.csda.2012.10.006 + + """ + self._dtype = dtype_util.common_dtype([probs], tf.float32) + self._probs = tensor_util.convert_nonref_to_tensor( + probs, dtype=self._dtype) + self._total_count = tf.shape(self._probs) + self._pmf_list = self.get_pmf() + self._cdf_list = tf.math.cumsum(self._pmf_list) + + + + """Cumulative density function""" + def _cdf(self, count): + return self._cdf_list[count] + + """Probability mass function""" + def _prob(self, count): + return self._pmf_list[count] + + """PMF helper""" + def get_pmf(self): + + def xi(period, l): + if l == 0: + return tf.dtypes.complex(1, 0) + else: + wl = period*l + real = 1+self._probs*(tf.math.cos(wl)-1) + imag = self._probs*tf.math.sin(wl) + mod = tf.math.atan2(imag**2+real**2) + arg = tf.math.atan2(imag,real) + d = tf.math.exp((tf.math.log(mod)).sum()) + arg_sum = arg.sum() + a = d*tf.math.cos(arg_sum) + b = d*tf.math.sin(arg_sum) + return tf.dtypes.complex(a,b) + + period = 2*pi/(1+n) + n = self._total_count + + x = [xi(period,i) for i in range((n+1)//2+1)] + for i in range((n+1)//2+1,n+1): + c = x[n+1-i] + x.append(tf.math.conj(c)) + + return tf.math.real(tf.signal.fft(x))/(n+1) + + + + + diff --git a/tensorflow_probability/python/distributions/poisson_binomial_test.py b/tensorflow_probability/python/distributions/poisson_binomial_test.py new file mode 100644 index 0000000000..cc40aa2525 --- /dev/null +++ b/tensorflow_probability/python/distributions/poisson_binomial_test.py @@ -0,0 +1,42 @@ +# Copyright 2022 The TensorFlow Probability Authors. +# +# 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. +# ============================================================================ +"""Tests for poisson binomial distribution.""" + +# Dependency imports + +from importlib.metadata import distribution +import numpy as np +from scipy import stats +import tensorflow.compat.v2 as tf +from tensorflow_probability.python.distributions import discrete_rejection_sampling +from tensorflow_probability.python.distributions import geometric +from tensorflow_probability.python.distributions.poisson_binomial import PoissonBinomial +from tensorflow_probability.python.internal import test_util + + +@test_util.test_all_tf_execution_regimes +class PoissonBinomialTest(test_util.TestCase): + + def testFunctionalityPoissonBinomial(self): + + probs = np.float32([0.2, 0.5, 0.3, 0.4]) + dist = PoissonBinomial(probs=probs) + + self.evaluate(dist.cdf(3)) + self.evaluate(dist.pdf(3)) + + +if __name__ == '__main__': + test_util.main() \ No newline at end of file