dr.math.distributions.NegativeBinomialDistribution.java Source code

Introduction

Here is the source code for dr.math.distributions.NegativeBinomialDistribution.java

Source

/*
 * NegativeBinomialDistribution.java
 *
 * Copyright (c) 2002-2015 Alexei Drummond, Andrew Rambaut and Marc Suchard
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.math.distributions;

import dr.math.ErrorFunction;

import org.apache.commons.math.MathException;
import org.apache.commons.math.special.Beta;

import dr.math.*;

/**
 * @author Trevor Bedford
 * @version $Id$
 */
public class NegativeBinomialDistribution implements Distribution {

    double mean;
    double alpha;

    public NegativeBinomialDistribution(double mean, double alpha) {
        this.mean = mean;
        this.alpha = alpha;
    }

    public double pdf(double x) {
        return pdf(x, mean, alpha);
    }

    public double logPdf(double x) {
        return logPdf(x, mean, alpha);
    }

    public double cdf(double x) {
        return cdf(x, mean, alpha);
    }

    public double quantile(double y) {
        // TB - I'm having trouble implementing this
        // LM - A first stab using simple minimisation to invert the function (under absolute loss)
        // Implementation based on the qnbinom.c function used in R
        final double stdev = Math.sqrt(mean + (mean * mean * alpha));
        final double r = -1 * (mean * mean) / (mean - stdev * stdev);
        final double p = mean / (stdev * stdev);
        final double prob = y;

        final double Q = 1.0 / p;
        final double P = (1.0 - p) * Q;
        final double gamma = (Q + P) / stdev;
        final double z = Math.sqrt(2.0) * ErrorFunction.inverseErf(2.0 * y - 1.0);
        final double crudeY = mean + stdev * (z + gamma * (z * z - 1) / 6);

        UnivariateFunction f = new UnivariateFunction() {
            double tent = Double.NaN;

            public double evaluate(final double argument) {
                try {
                    tent = Beta.regularizedBeta(p, r, argument + 1);
                } catch (MathException e) {
                    return Double.NaN;
                }
                double score = Math.abs(tent - prob);
                return score;
            }

            public int getNumArguments() {
                return 1;
            }

            public double getLowerBound() { // 20% window should cut it. Probably too large even...
                return Math.min(crudeY - .2 * crudeY, 0);
            }

            public double getUpperBound() {
                return crudeY + .2 * crudeY;
            }
        };
        UnivariateMinimum minimum = new UnivariateMinimum();
        double q = minimum.findMinimum(f);
        return Math.ceil(q);
    }

    public double mean() {
        return mean;
    }

    public double variance() {
        return mean + (mean * mean * alpha);
    }

    public UnivariateFunction getProbabilityDensityFunction() {
        throw new RuntimeException();
    }

    public static double pdf(double x, double mean, double alpha) {
        if (x < 0)
            return 0;
        return Math.exp(logPdf(x, mean, alpha));
    }

    public static double logPdf(double x, double mean, double alpha) {
        if (x < 0)
            return Double.NEGATIVE_INFINITY;
        //        double r = -1 * (mean*mean) / (mean - stdev*stdev);
        //        double p = mean / (stdev*stdev);
        //        return Math.log(Math.pow(1-p,x)) + Math.log(Math.pow(p, r)) + GammaFunction.lnGamma(r+x) - GammaFunction.lnGamma(r) - GammaFunction.lnGamma(x+1);
        double theta = 1.0 / alpha;

        double p = theta / (theta + mean);
        return Math.log(1 - p) * x + Math.log(p) * theta + GammaFunction.lnGamma(theta + x)
                - GammaFunction.lnGamma(theta) - GammaFunction.lnGamma(x + 1);
    }

    public static double cdf(double x, double mean, double alpha) {
        double theta = 1.0 / alpha;
        double p = theta / (theta + mean);
        try {
            return Beta.regularizedBeta(p, theta, x + 1);
        } catch (MathException e) {
            // AR - throwing exceptions deep in numerical code causes trouble. Catching runtime
            // exceptions is bad. Better to return NaN and let the calling code deal with it.
            return Double.NaN;
            //                throw MathRuntimeException.createIllegalArgumentException(
            //                "Couldn't calculate beta cdf for alpha = " + alpha + ", beta = " + beta + ": " +e.getMessage());
        }
    }

    public static void main(String[] args) {
        System.out.println("Test negative binomial");
        double mean = 5;
        double stdev = 5;
        //         double r = -1 * (mean*mean) / (mean - stdev*stdev);
        double alpha = (stdev * stdev - mean) / (mean * mean);

        NegativeBinomialDistribution dist = new NegativeBinomialDistribution(5, alpha);
        System.out.println("Mean 5, sd 5, x 5, pdf 0.074487, logPdf -2.59713, median 4");
        System.out.println("pdf = " + dist.pdf(5));
        System.out.println("quantile(0.5) aka median = " + dist.quantile(0.5));
        System.out.println("logPdf = " + dist.logPdf(5));
    }

}