log_mix.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_FUN_LOG_MIX_HPP
#define STAN_MATH_REV_FUN_LOG_MIX_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/fun/value_of.hpp>
#include <stan/math/prim/fun/log_mix.hpp>
#include <stan/math/prim/fun/value_of.hpp>
#include <stan/math/rev/functor/partials_propagator.hpp>
#include <cmath>
 
namespace stan {
namespace math {
 
/* Computes shared terms in log_mix partial derivative calculations
 *
 * @param[in] theta_val value of mixing proportion theta.
 * @param[in] lambda1_val value of log density multiplied by theta.
 * @param[in] lambda2_val value of log density multiplied by 1 - theta.
 * @param[out] one_m_exp_lam2_m_lam1 shared term in deriv calculation.
 * @param[out] one_m_t_prod_exp_lam2_m_lam1 shared term in deriv calculation.
 * @param[out] one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1 shared term in deriv
 * calculation.
 */
inline void log_mix_partial_helper(
    double theta_val, double lambda1_val, double lambda2_val,
    double& one_m_exp_lam2_m_lam1, double& one_m_t_prod_exp_lam2_m_lam1,
    double& one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1) {
  using std::exp;
  double lam2_m_lam1 = lambda2_val - lambda1_val;
  double exp_lam2_m_lam1 = exp(lam2_m_lam1);
  one_m_exp_lam2_m_lam1 = 1 - exp_lam2_m_lam1;
  double one_m_t = 1 - theta_val;
  one_m_t_prod_exp_lam2_m_lam1 = one_m_t * exp_lam2_m_lam1;
  one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1
      = 1 / (theta_val + one_m_t_prod_exp_lam2_m_lam1);
}
 
template <typename T_theta, typename T_lambda1, typename T_lambda2,
          require_any_var_t<T_theta, T_lambda1, T_lambda2>* = nullptr>
inline return_type_t<T_theta, T_lambda1, T_lambda2> log_mix(
    const T_theta& theta, const T_lambda1& lambda1, const T_lambda2& lambda2) {
  using std::log;
 
  auto ops_partials = make_partials_propagator(theta, lambda1, lambda2);
 
  double theta_double = value_of(theta);
  const double lambda1_double = value_of(lambda1);
  const double lambda2_double = value_of(lambda2);
 
  double log_mix_function_value
      = log_mix(theta_double, lambda1_double, lambda2_double);
 
  double one_m_exp_lam2_m_lam1(0.0);
  double one_m_t_prod_exp_lam2_m_lam1(0.0);
  double one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1(0.0);
 
  if (lambda1 > lambda2) {
    log_mix_partial_helper(theta_double, lambda1_double, lambda2_double,
                           one_m_exp_lam2_m_lam1, one_m_t_prod_exp_lam2_m_lam1,
                           one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1);
  } else {
    log_mix_partial_helper(1.0 - theta_double, lambda2_double, lambda1_double,
                           one_m_exp_lam2_m_lam1, one_m_t_prod_exp_lam2_m_lam1,
                           one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1);
    one_m_exp_lam2_m_lam1 = -one_m_exp_lam2_m_lam1;
    theta_double = one_m_t_prod_exp_lam2_m_lam1;
    one_m_t_prod_exp_lam2_m_lam1 = 1.0 - value_of(theta);
  }
 
  if (!is_constant_all<T_theta>::value) {
    partials<0>(ops_partials)[0]
        = one_m_exp_lam2_m_lam1 * one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1;
  }
  if (!is_constant_all<T_lambda1>::value) {
    partials<1>(ops_partials)[0]
        = theta_double * one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1;
  }
  if (!is_constant_all<T_lambda2>::value) {
    partials<2>(ops_partials)[0] = one_m_t_prod_exp_lam2_m_lam1
                                   * one_d_t_plus_one_m_t_prod_exp_lam2_m_lam1;
  }
 
  return ops_partials.build(log_mix_function_value);
}
 
}  // namespace math
}  // namespace stan
#endif