bssm/src/model_bsm_ng.cpp at master · learner-python-R/bssm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#include "model_bsm_ng.h"

// from Rcpp::List
bsm_ng::bsm_ng(const Rcpp::List model, const unsigned int seed) :
  ssm_ung(model, seed),
  prior_distributions(Rcpp::as<arma::uvec>(model["prior_distributions"])),
  prior_parameters(Rcpp::as<arma::mat>(model["prior_parameters"])),
  slope(Rcpp::as<bool>(model["slope"])),
  seasonal(Rcpp::as<bool>(model["seasonal"])),
  noise(Rcpp::as<bool>(model["noise"])),
  fixed(Rcpp::as<arma::uvec>(model["fixed"])), level_est(fixed(0) == 0),
  slope_est(slope && fixed(1) == 0), seasonal_est(seasonal && fixed(2) == 0),
  phi_est(Rcpp::as<bool>(model["phi_est"])) {
}

// used in parallel regions, does not depend on R
void bsm_ng::update_model(const arma::vec& new_theta) {

  if (arma::accu(fixed) < 3 || noise) {

    // sd_level
    if (level_est) {
      R(0, 0, 0) = std::exp(new_theta(0));
    }
    // sd_slope
    if (slope_est) {
      R(1, 1, 0) = std::exp(new_theta(level_est));
    }
    // sd_seasonal
    if (seasonal_est) {
      R(1 + slope, 1 + slope, 0) =
        std::exp(new_theta(level_est + slope_est));
    }
    if(noise) {
      R(m - 1, 1 + slope + seasonal, 0) =
        std::exp(new_theta(level_est + slope_est + seasonal_est));
      P1(m - 1, m - 1) = std::pow(R(m - 1, 1 + slope + seasonal, 0), 2.0);
    }
    compute_RR();
  }
  if(phi_est) {
    phi = std::exp(new_theta(level_est + slope_est + seasonal_est + noise));
  }
  if(xreg.n_cols > 0) {
    beta = new_theta.subvec(new_theta.n_elem - xreg.n_cols, new_theta.n_elem - 1);
    compute_xbeta();
  }
  theta = new_theta;
  // approximation does not match theta anymore (keep as -1 if so)
  if (approx_state > 0) approx_state = 0;
}

// used in mcmc, latter argument is not actually used
void bsm_ng::update_model(const arma::vec& new_theta, const Rcpp::Function update_fn) {

  if (arma::accu(fixed) < 3 || noise) {

    // sd_level
    if (level_est) {
      R(0, 0, 0) = std::exp(new_theta(0));
    }
    // sd_slope
    if (slope_est) {
      R(1, 1, 0) = std::exp(new_theta(level_est));
    }
    // sd_seasonal
    if (seasonal_est) {
      R(1 + slope, 1 + slope, 0) =
        std::exp(new_theta(level_est + slope_est));
    }
    if(noise) {
      R(m - 1, 1 + slope + seasonal, 0) =
        std::exp(new_theta(level_est + slope_est + seasonal_est));
      P1(m - 1, m - 1) = std::pow(R(m - 1, 1 + slope + seasonal, 0), 2.0);
    }
    compute_RR();
  }
  if(phi_est) {
    phi = std::exp(new_theta(level_est + slope_est + seasonal_est + noise));
  }
  if(xreg.n_cols > 0) {
    beta = new_theta.subvec(new_theta.n_elem - xreg.n_cols, new_theta.n_elem - 1);
    compute_xbeta();
  }
  theta = new_theta;
  // approximation does not match theta anymore (keep as -1 if so)
  if (approx_state > 0) approx_state = 0;
}

double bsm_ng::log_prior_pdf(const arma::vec& x, const Rcpp::Function prior_fn) const {

  double log_prior = 0.0;
  arma::vec pars = x;
  if (arma::accu(fixed) < 3 || noise) {
    pars.subvec(0, pars.n_elem - xreg.n_cols - 1) =
      arma::exp(pars.subvec(0, pars.n_elem - xreg.n_cols - 1));
    // add jacobian
    log_prior += arma::accu(x.subvec(0, x.n_elem - xreg.n_cols - 1));
  }

  for(unsigned int i = 0; i < pars.n_elem; i++) {
    switch(prior_distributions(i)) {
    case 0  :
      if (pars(i) < prior_parameters(0, i) || pars(i) > prior_parameters(1, i)) {
        return -std::numeric_limits<double>::infinity();
      }
      break;
    case 1  :
      if (pars(i) < 0) {
        return -std::numeric_limits<double>::infinity();
      } else {
        log_prior -= 0.5 * std::pow(pars(i) / prior_parameters(0, i), 2);
      }
      break;
    case 2  :
      log_prior -= 0.5 * std::pow((pars(i) - prior_parameters(0, i)) / prior_parameters(1, i), 2);
      break;
    case 3 : // truncated normal
      if (pars(i) < prior_parameters(2, i) || pars(i) > prior_parameters(3, i)) {
        return -std::numeric_limits<double>::infinity();
      } else {
        log_prior -= 0.5 * std::pow((pars(i) - prior_parameters(0, i)) / prior_parameters(1, i), 2);
      }
      break;
    case 4 : // gamma
      if (pars(i) < 0) {
        return -std::numeric_limits<double>::infinity();
      } else {
        log_prior += (prior_parameters(0, i) - 1) * log(pars(i)) - prior_parameters(1, i) * pars(i);
      }
      break;
    }

  }

  return log_prior;
}