Skip to content

Instantly share code, notes, and snippets.

@maxbiostat
Last active February 6, 2019 13:02
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save maxbiostat/d380b6734b5776673ee40c481cbeef2f to your computer and use it in GitHub Desktop.
Save maxbiostat/d380b6734b5776673ee40c481cbeef2f to your computer and use it in GitHub Desktop.
Implementing zeta function in Stan

Stan model

functions{
    real zeta(real s);
    /*TODO: implement rng for prior and posterior predictive checks*/
}
data{
  int<lower=0> K; // number of unique values
  int values[K];
  int<lower=0> frequencies[K];
}
parameters{
  real <lower=1> alpha;
}
model{
  real constant = log(zeta(alpha));
  for (k in 1:K) {
    target += frequencies[k] * (-alpha * log(values[k]) - constant);
  }
} 

And zeta.hpp reads

#include <boost/math/special_functions/zeta.hpp>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/quadrature/exp_sinh.hpp>
using boost::math::isfinite;

inline double zeta(const double& s, std::ostream* pstream__){
  return boost::math::zeta(s);
}

inline var zeta(const var& s_var, std::ostream* pstream__) {
  double s = s_var.val();
  double f = zeta(s, pstream__);
  boost::math::quadrature::exp_sinh<double> integrator;
  auto int_rep_zeta_prime = [&](double x) {
    double ans = std::pow(x, s-1) *( std::log(x) - boost::math::digamma(s))/(std::exp(x) -1);
    if(!isfinite(ans)) return 0.0; // weird patching to get quadrature working
    return ans;
  };
  double tolerance = std::sqrt(std::numeric_limits<double>::epsilon());
  double error = 0.0;
  double L1 = 0.0;
  size_t levels;
  double deriv = integrator.integrate(int_rep_zeta_prime, tolerance, &error, &L1, &levels)/boost::math::tgamma(s);
  return var(new precomp_v_vari(f, s_var.vi_, deriv));
}

Compiler call is (extracted from Rstan):

/usr/local/lib/R/bin/R CMD SHLIB file7cba24006b1c.cpp 2> file7cba24006b1c.cpp.err.txt 
g++ -std=gnu++14 -I"/usr/local/lib/R/include" -DNDEBUG   -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/Rcpp/include/"  -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/RcppEigen/include/"  -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/RcppEigen/include/unsupported"  -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/BH/include" -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/StanHeaders/include/src/"  -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/StanHeaders/include/"  -I"/home/max/R/x86_64-pc-linux-gnu-library/3.5/rstan/include" -DEIGEN_NO_DEBUG  -DBOOST_DISABLE_ASSERTS  -I/usr/local/include   -fpic  -g -O2 -c file7cba24006b1c.cpp -o file7cba24006b1c.o

Doing

writeLines(plaw.model@model_cpp$model_cppcode)

gives

// Code generated by Stan version 2.18.0

#include <stan/model/model_header.hpp>

namespace model7cba1e0291c9_discrete_power_law_namespace {

using std::istream;
using std::string;
using std::stringstream;
using std::vector;
using stan::io::dump;
using stan::math::lgamma;
using stan::model::prob_grad;
using namespace stan::math;

static int current_statement_begin__;

stan::io::program_reader prog_reader__() {
    stan::io::program_reader reader;
    reader.add_event(0, 0, "start", "model7cba1e0291c9_discrete_power_law");
    reader.add_event(20, 18, "end", "model7cba1e0291c9_discrete_power_law");
    return reader;
}

template <typename T0__>
typename boost::math::tools::promote_args<T0__>::type
zeta(const T0__& s, std::ostream* pstream__);

class model7cba1e0291c9_discrete_power_law : public prob_grad {
private:
    int K;
    vector<int> values;
    vector<int> frequencies;
public:
    model7cba1e0291c9_discrete_power_law(stan::io::var_context& context__,
        std::ostream* pstream__ = 0)
        : prob_grad(0) {
        ctor_body(context__, 0, pstream__);
    }

    model7cba1e0291c9_discrete_power_law(stan::io::var_context& context__,
        unsigned int random_seed__,
        std::ostream* pstream__ = 0)
        : prob_grad(0) {
        ctor_body(context__, random_seed__, pstream__);
    }

    void ctor_body(stan::io::var_context& context__,
                   unsigned int random_seed__,
                   std::ostream* pstream__) {
        typedef double local_scalar_t__;

        boost::ecuyer1988 base_rng__ =
          stan::services::util::create_rng(random_seed__, 0);
        (void) base_rng__;  // suppress unused var warning

        current_statement_begin__ = -1;

        static const char* function__ = "model7cba1e0291c9_discrete_power_law_namespace::model7cba1e0291c9_discrete_power_law";
        (void) function__;  // dummy to suppress unused var warning
        size_t pos__;
        (void) pos__;  // dummy to suppress unused var warning
        std::vector<int> vals_i__;
        std::vector<double> vals_r__;
        local_scalar_t__ DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());
        (void) DUMMY_VAR__;  // suppress unused var warning

        // initialize member variables
        try {
            current_statement_begin__ = 6;
            context__.validate_dims("data initialization", "K", "int", context__.to_vec());
            K = int(0);
            vals_i__ = context__.vals_i("K");
            pos__ = 0;
            K = vals_i__[pos__++];
            current_statement_begin__ = 7;
            validate_non_negative_index("values", "K", K);
            context__.validate_dims("data initialization", "values", "int", context__.to_vec(K));
            validate_non_negative_index("values", "K", K);
            values = std::vector<int>(K,int(0));
            vals_i__ = context__.vals_i("values");
            pos__ = 0;
            size_t values_limit_0__ = K;
            for (size_t i_0__ = 0; i_0__ < values_limit_0__; ++i_0__) {
                values[i_0__] = vals_i__[pos__++];
            }
            current_statement_begin__ = 8;
            validate_non_negative_index("frequencies", "K", K);
            context__.validate_dims("data initialization", "frequencies", "int", context__.to_vec(K));
            validate_non_negative_index("frequencies", "K", K);
            frequencies = std::vector<int>(K,int(0));
            vals_i__ = context__.vals_i("frequencies");
            pos__ = 0;
            size_t frequencies_limit_0__ = K;
            for (size_t i_0__ = 0; i_0__ < frequencies_limit_0__; ++i_0__) {
                frequencies[i_0__] = vals_i__[pos__++];
            }

            // validate, data variables
            current_statement_begin__ = 6;
            check_greater_or_equal(function__,"K",K,0);
            current_statement_begin__ = 7;
            current_statement_begin__ = 8;
            for (int k0__ = 0; k0__ < K; ++k0__) {
                check_greater_or_equal(function__,"frequencies[k0__]",frequencies[k0__],0);
            }
            // initialize data variables


            // validate transformed data

            // validate, set parameter ranges
            num_params_r__ = 0U;
            param_ranges_i__.clear();
            current_statement_begin__ = 11;
            ++num_params_r__;
        } catch (const std::exception& e) {
            stan::lang::rethrow_located(e, current_statement_begin__, prog_reader__());
            // Next line prevents compiler griping about no return
            throw std::runtime_error("*** IF YOU SEE THIS, PLEASE REPORT A BUG ***");
        }
    }

    ~model7cba1e0291c9_discrete_power_law() { }


    void transform_inits(const stan::io::var_context& context__,
                         std::vector<int>& params_i__,
                         std::vector<double>& params_r__,
                         std::ostream* pstream__) const {
        stan::io::writer<double> writer__(params_r__,params_i__);
        size_t pos__;
        (void) pos__; // dummy call to supress warning
        std::vector<double> vals_r__;
        std::vector<int> vals_i__;

        if (!(context__.contains_r("alpha")))
            throw std::runtime_error("variable alpha missing");
        vals_r__ = context__.vals_r("alpha");
        pos__ = 0U;
        context__.validate_dims("initialization", "alpha", "double", context__.to_vec());
        double alpha(0);
        alpha = vals_r__[pos__++];
        try {
            writer__.scalar_lb_unconstrain(1,alpha);
        } catch (const std::exception& e) { 
            throw std::runtime_error(std::string("Error transforming variable alpha: ") + e.what());
        }

        params_r__ = writer__.data_r();
        params_i__ = writer__.data_i();
    }

    void transform_inits(const stan::io::var_context& context,
                         Eigen::Matrix<double,Eigen::Dynamic,1>& params_r,
                         std::ostream* pstream__) const {
      std::vector<double> params_r_vec;
      std::vector<int> params_i_vec;
      transform_inits(context, params_i_vec, params_r_vec, pstream__);
      params_r.resize(params_r_vec.size());
      for (int i = 0; i < params_r.size(); ++i)
        params_r(i) = params_r_vec[i];
    }


    template <bool propto__, bool jacobian__, typename T__>
    T__ log_prob(vector<T__>& params_r__,
                 vector<int>& params_i__,
                 std::ostream* pstream__ = 0) const {

        typedef T__ local_scalar_t__;

        local_scalar_t__ DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());
        (void) DUMMY_VAR__;  // suppress unused var warning

        T__ lp__(0.0);
        stan::math::accumulator<T__> lp_accum__;

        try {
            // model parameters
            stan::io::reader<local_scalar_t__> in__(params_r__,params_i__);

            local_scalar_t__ alpha;
            (void) alpha;  // dummy to suppress unused var warning
            if (jacobian__)
                alpha = in__.scalar_lb_constrain(1,lp__);
            else
                alpha = in__.scalar_lb_constrain(1);


            // transformed parameters



            // validate transformed parameters

            const char* function__ = "validate transformed params";
            (void) function__;  // dummy to suppress unused var warning

            // model body
            {
            current_statement_begin__ = 14;
            local_scalar_t__ constant;
            (void) constant;  // dummy to suppress unused var warning

            stan::math::initialize(constant, DUMMY_VAR__);
            stan::math::fill(constant,DUMMY_VAR__);
            stan::math::assign(constant,stan::math::log(zeta(alpha, pstream__)));


            current_statement_begin__ = 15;
            for (int k = 1; k <= K; ++k) {

                current_statement_begin__ = 16;
                lp_accum__.add((get_base1(frequencies,k,"frequencies",1) * ((-(alpha) * stan::math::log(get_base1(values,k,"values",1))) - constant)));
            }
            }

        } catch (const std::exception& e) {
            stan::lang::rethrow_located(e, current_statement_begin__, prog_reader__());
            // Next line prevents compiler griping about no return
            throw std::runtime_error("*** IF YOU SEE THIS, PLEASE REPORT A BUG ***");
        }

        lp_accum__.add(lp__);
        return lp_accum__.sum();

    } // log_prob()

    template <bool propto, bool jacobian, typename T_>
    T_ log_prob(Eigen::Matrix<T_,Eigen::Dynamic,1>& params_r,
               std::ostream* pstream = 0) const {
      std::vector<T_> vec_params_r;
      vec_params_r.reserve(params_r.size());
      for (int i = 0; i < params_r.size(); ++i)
        vec_params_r.push_back(params_r(i));
      std::vector<int> vec_params_i;
      return log_prob<propto,jacobian,T_>(vec_params_r, vec_params_i, pstream);
    }


    void get_param_names(std::vector<std::string>& names__) const {
        names__.resize(0);
        names__.push_back("alpha");
    }


    void get_dims(std::vector<std::vector<size_t> >& dimss__) const {
        dimss__.resize(0);
        std::vector<size_t> dims__;
        dims__.resize(0);
        dimss__.push_back(dims__);
    }

    template <typename RNG>
    void write_array(RNG& base_rng__,
                     std::vector<double>& params_r__,
                     std::vector<int>& params_i__,
                     std::vector<double>& vars__,
                     bool include_tparams__ = true,
                     bool include_gqs__ = true,
                     std::ostream* pstream__ = 0) const {
        typedef double local_scalar_t__;

        vars__.resize(0);
        stan::io::reader<local_scalar_t__> in__(params_r__,params_i__);
        static const char* function__ = "model7cba1e0291c9_discrete_power_law_namespace::write_array";
        (void) function__;  // dummy to suppress unused var warning
        // read-transform, write parameters
        double alpha = in__.scalar_lb_constrain(1);
        vars__.push_back(alpha);

        // declare and define transformed parameters
        double lp__ = 0.0;
        (void) lp__;  // dummy to suppress unused var warning
        stan::math::accumulator<double> lp_accum__;

        local_scalar_t__ DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());
        (void) DUMMY_VAR__;  // suppress unused var warning

        try {



            // validate transformed parameters

            // write transformed parameters
            if (include_tparams__) {
            }
            if (!include_gqs__) return;
            // declare and define generated quantities



            // validate generated quantities

            // write generated quantities
        } catch (const std::exception& e) {
            stan::lang::rethrow_located(e, current_statement_begin__, prog_reader__());
            // Next line prevents compiler griping about no return
            throw std::runtime_error("*** IF YOU SEE THIS, PLEASE REPORT A BUG ***");
        }
    }

    template <typename RNG>
    void write_array(RNG& base_rng,
                     Eigen::Matrix<double,Eigen::Dynamic,1>& params_r,
                     Eigen::Matrix<double,Eigen::Dynamic,1>& vars,
                     bool include_tparams = true,
                     bool include_gqs = true,
                     std::ostream* pstream = 0) const {
      std::vector<double> params_r_vec(params_r.size());
      for (int i = 0; i < params_r.size(); ++i)
        params_r_vec[i] = params_r(i);
      std::vector<double> vars_vec;
      std::vector<int> params_i_vec;
      write_array(base_rng,params_r_vec,params_i_vec,vars_vec,include_tparams,include_gqs,pstream);
      vars.resize(vars_vec.size());
      for (int i = 0; i < vars.size(); ++i)
        vars(i) = vars_vec[i];
    }

    static std::string model_name() {
        return "model7cba1e0291c9_discrete_power_law";
    }


    void constrained_param_names(std::vector<std::string>& param_names__,
                                 bool include_tparams__ = true,
                                 bool include_gqs__ = true) const {
        std::stringstream param_name_stream__;
        param_name_stream__.str(std::string());
        param_name_stream__ << "alpha";
        param_names__.push_back(param_name_stream__.str());

        if (!include_gqs__ && !include_tparams__) return;

        if (include_tparams__) {
        }


        if (!include_gqs__) return;
    }


    void unconstrained_param_names(std::vector<std::string>& param_names__,
                                   bool include_tparams__ = true,
                                   bool include_gqs__ = true) const {
        std::stringstream param_name_stream__;
        param_name_stream__.str(std::string());
        param_name_stream__ << "alpha";
        param_names__.push_back(param_name_stream__.str());

        if (!include_gqs__ && !include_tparams__) return;

        if (include_tparams__) {
        }


        if (!include_gqs__) return;
    }

}; // model

}

typedef model7cba1e0291c9_discrete_power_law_namespace::model7cba1e0291c9_discrete_power_law stan_model;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment