NegativeBinomialMixture.h

#ifndef theplu_yat_statistics_negative_binomial_mixture_
#define theplu_yat_statistics_negative_binomial_mixture_

// $Id: NegativeBinomialMixture.h 4192 2022-08-11 07:35:35Z peter $

/*
  Copyright (C) 2022 Peter Johansson

  This file is part of the yat library, https://dev.thep.lu.se/yat

  The yat library is free software; you can redistribute it and/or
  modify it under the terms of the GNU General Public License as
  published by the Free Software Foundation; either version 3 of the
  License, or (at your option) any later version.

  The yat library 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
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with yat. If not, see <https://www.gnu.org/licenses/>.
*/


#include <yat/utility/Matrix.h>
#include <yat/utility/Vector.h>

#include <map>

namespace theplu {
namespace yat {
  namespace utility {
    class Matrix;
    class MultiMinimizerDerivative;
  }
namespace statistics {


  class NegativeBinomialMixture
  {
  public:
    void add(unsigned long int k, unsigned long int n=1);

    const utility::VectorBase& alpha(void) const;

    void clear(void);

    void fit(size_t n);

    void fit(const utility::VectorBase& m,
             const utility::VectorBase& alpha,
             const utility::VectorBase& tau);

    double logL(void) const;

    double pdf(unsigned long int k) const;

    double pdf(unsigned long int k, size_t i) const;

    const utility::VectorBase& m(void) const;

    const utility::VectorBase& tau(void) const;
  private:
    std::map<unsigned long int, unsigned long int> count_;

    utility::Vector tau_;

    utility::Vector m_;
    utility::Vector alpha_;


    // p = (alpha-1.0) / alpha = 1.0 - 1.0/alpha
    double p(size_t i) const;
    // r = m*(1-p)/p = m / alpha^2
    double r(size_t i) const;

    void init_fit(size_t n);
    void optimize(void);

    bool expectation_step(utility::Matrix& H);
    bool optimize_tau(const utility::Matrix& H);
    bool optimize_param(const utility::Matrix& H,
                        utility::MultiMinimizerDerivative& solver);

    double pdf(unsigned long int k, double m, double alpha) const;

    class Q
    {
    public:
      Q(const std::map<unsigned long int, unsigned long int>& count,
        const utility::VectorBase& h);
      double operator()(const utility::VectorBase& x);
      void operator()(const utility::VectorBase& x,
                      utility::VectorMutable& gradient);

    private:
      std::map<unsigned long int, unsigned long int> count_;
      const utility::VectorBase& h_;

      double log_L(double m, double alpha, unsigned long int k) const;

      // return <dlogL/dm, dlogL/dalpha>
      utility::Vector
      calc_gradient(double m, double alpha, unsigned long int k) const;

      double calculate_dp(double p, double r, unsigned long int k) const;
      double calculate_dr(double p, double r, unsigned long int k) const;
    };
  };

}}}
#endif