NegativeBinomialExtendedMixture.h

#ifndef theplu_yat_statistics_negative_binomial_extended_mixture_
#define theplu_yat_statistics_negative_binomial_extended_mixture_

// $Id: NegativeBinomialExtendedMixture.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/MultiMinimizerDerivative.h>
#include <yat/utility/Vector.h>

#include <vector>

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

  class NegativeBinomialExtendedMixture
  {
  public:
    void add(unsigned long int k, double y);

    const utility::Vector& 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;

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

    const utility::Vector& tau(void) const;
  private:
    std::vector<std::pair<unsigned long int, double>> data_;

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

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

    // 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);

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

    private:
      const std::vector<std::pair<unsigned long int, double>>& data_;
      const yat::utility::VectorBase& h_;

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

      // return <dlogL/dm, dlogL/dalpha>
      yat::utility::Vector
      calc_gradient(double m, double alpha, unsigned long int k,
                    double y) 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