File: stanmath.Rmd

package info (click to toggle)
r-cran-stanheaders 2.32.10-1
  • links: PTS, VCS
  • area: main
  • in suites: sid, trixie
  • size: 19,652 kB
  • sloc: cpp: 125,776; ansic: 73,706; f90: 30,959; ml: 243; sh: 15; makefile: 5
file content (500 lines) | stat: -rw-r--r-- 19,581 bytes parent folder | download | duplicates (2)
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
---
title: "Using the Stan Math C++ Library"
author: "Stan Development Team"
date: "`r Sys.Date()`"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Using the Stan Math C++ Library}
  %\VignetteEncoding{UTF-8}
  %\VignetteEngine{knitr::rmarkdown}
editor_options: 
  chunk_output_type: inline
---

```{r setup, include = FALSE}
options(width = 100)
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
local({
  hook_output <- knitr::knit_hooks$get('output')
  knitr::knit_hooks$set(output = function(x, options) {
    if (!is.null(options$max.height)) options$attr.output <- c(
      options$attr.output,
      sprintf('style="max-height: %s;"', options$max.height)
    )
    hook_output(x, options)
  })
})
Sys.setenv(USE_CXX17 = "1")
set.seed(12345)
```

# Using the **StanHeaders** Package from Other R Packages

The **StanHeaders** package contains no R functions. To use the Stan Math Library
in other packages, it is often sufficient to specify
```
LinkingTo: StanHeaders (>= 2.26.0), RcppParallel (>= 5.0.1)
```
in the DESCRIPTION file of another package and put something like
```
CXX_STD = CXX17
PKG_CXXFLAGS = $(shell "$(R_HOME)/bin$(R_ARCH_BIN)/Rscript" -e "RcppParallel::CxxFlags()") \
               $(shell "$(R_HOME)/bin$(R_ARCH_BIN)/Rscript" -e "StanHeaders:::CxxFlags()")
PKG_LIBS = $(shell "$(R_HOME)/bin$(R_ARCH_BIN)/Rscript" -e "RcppParallel::RcppParallelLibs()") \
           $(shell "$(R_HOME)/bin$(R_ARCH_BIN)/Rscript" -e "StanHeaders:::LdFlags()")
```
in the src/Makevars and src/Makevars.win files and put `GNU make` in the `SystemRequirements:` 
field of the package's DESCRIPTION file. If, in addition, the other package needs to utilize 
the MCMC, optimization, variational inference, or parsing facilities of the Stan Library, then it is 
also necessary to include the `src` directory of **StanHeaders** in the other package's `PKG_CXXFLAGS`
in the src/Makevars and src/Makevars.win files with something like
```
STANHEADERS_SRC = $(shell "$(R_HOME)/bin$(R_ARCH_BIN)/Rscript" -e "message()" \
  -e "cat(system.file('include', 'src', package = 'StanHeaders', mustWork = TRUE))" \
  -e "message()" | grep "StanHeaders")
PKG_CXXFLAGS += -I"$(STANHEADERS_SRC)"
```

# Calling functions in the **StanHeaders** Package from R

The only exposed R function in the in the **StanHeaders** package is `stanFunction`, which
can be used to call most functions in the Stan Math Library.

```{r}
example(stanFunction, package = "StanHeaders", run.dontrun = TRUE)
```
```{css, echo=FALSE}
.scroll-100 {
  max-height: 100px;
  overflow-y: auto;
  background-color: inherit;
}
```
The `functions` object defined in this example lists the many Stan functions that could be called
(if all their arguments are numeric, see `help(stanFunction, package = "StanHeaders")` for details)
```{r, echo = FALSE, warning = FALSE, class.output="scroll-100"}
if (length(functions) %% 2 == 1) {
  functions <- c(functions, "")
}
functions <- matrix(functions, ncol = 2, byrow = TRUE)
print(functions)
```


# Using Higher-Order Functions in the **StanHeaders** Package

This section will demonstrate how to use some of the C++ functions in the **StanHeaders** package
whose first argument is another C++ function, in which case the `stanFunction` in the previous
section will not work and you have to write your own C++.

## Derivatives and Minimization

The following is a toy example of using the Stan Math library via `Rcpp::sourceCpp`:
to minimize the function 
$$\left(\mathbf{x} - \mathbf{a}\right)^\top \left(\mathbf{x} - \mathbf{a}\right)$$
which has a global minimum when $\mathbf{x} = \mathbf{a}$. To find this minimum with autodifferentiation, 
we need to define the objective function. Then, its gradient with respect to $\mathbf{x}$, which we know is 
$2\left(\mathbf{x} - \mathbf{a}\right)$ in this case, can be calculated by autodifferentiation. At the 
optimum (or on the way to the optimum), we might want to evaluate the Hessian matrix, which we know is 
$2\mathbf{I}$, but would need an additional function to evaluate it via autodifferentiation. Finally, one 
could reconceptualize the problem as solving a homogeneous system of equations where the gradient is set 
equal to a vector of zeros. The `stan::math::algebra_solver` function can solve such a system using 
autodifferentiation to obtain the Jacobian, which we know to be the identity matrix in this case.

```{r}
Sys.setenv(PKG_CXXFLAGS = StanHeaders:::CxxFlags(as_character = TRUE))
SH <- system.file(ifelse(.Platform$OS.type == "windows", "libs", "lib"), .Platform$r_arch,
                  package = "StanHeaders", mustWork = TRUE)
Sys.setenv(PKG_LIBS = paste0(StanHeaders:::LdFlags(as_character = TRUE),
                             " -L", shQuote(SH), " -lStanHeaders"))
                            
```

Here is C++ code that does all of the above, except for the part of finding the optimum, which is done
using the R function `optim` below.
```{Rcpp}
// [[Rcpp::depends(BH)]]
// [[Rcpp::depends(RcppEigen)]]
// [[Rcpp::depends(RcppParallel)]]
// [[Rcpp::depends(StanHeaders)]]
#include <stan/math/mix.hpp> // stuff from mix/ must come first
#include <stan/math.hpp>     // finally pull in everything from rev/ and prim/
#include <Rcpp.h>
#include <RcppEigen.h>       // do this AFTER including stuff from stan/math

// [[Rcpp::plugins(cpp17)]]

/* Objective function */

// [[Rcpp::export]]
auto f(Eigen::VectorXd x, Eigen::VectorXd a) { // objective function in doubles
  using stan::math::dot_self;                  // dot_self() is a dot product with self
  return dot_self( (x - a) );
}

/* Gradient */

// [[Rcpp::export]]
auto g(Eigen::VectorXd x, Eigen::VectorXd a) {  // gradient by AD using Stan
  double fx;
  Eigen::VectorXd grad_fx;
  using stan::math::dot_self;
  stan::math::gradient([&a](auto x) { return dot_self( (x - a) ); },
                       x, fx, grad_fx);
  return grad_fx;
}

/* Hessian */

// [[Rcpp::export]]
auto H(Eigen::VectorXd x, Eigen::VectorXd a) { // Hessian by AD using Stan
  double fx;
  Eigen::VectorXd grad_fx;
  Eigen::MatrixXd H;
  using stan::math::dot_self;
  stan::math::hessian([&a](auto x) { return dot_self(x - a); },
                      x, fx, grad_fx, H);
  return H;
}

/* Jacobian */

// [[Rcpp::export]]
auto J(Eigen::VectorXd x, Eigen::VectorXd a) { // not actually used
  Eigen::VectorXd fx;
  Eigen::MatrixXd J;
  using stan::math::dot_self;
  stan::math::jacobian([&a](auto x) {
    return (2 * (x - a));
  }, x, fx, J);
  return J;
}

struct equations_functor {
  template <typename T0, typename T1>
  inline Eigen::Matrix<T0, Eigen::Dynamic, 1>
  operator()(const Eigen::Matrix<T0, Eigen::Dynamic, 1>& x,
             const Eigen::Matrix<T1, Eigen::Dynamic, 1>& theta,
             const std::vector<double>& x_r, const std::vector<int>& x_i,
             std::ostream* pstream__) const {
    return 2 * (x - stan::math::to_vector(x_r));
  }
};

// [[Rcpp::export]]
auto solution(Eigen::VectorXd a, Eigen::VectorXd guess) {
  Eigen::VectorXd theta;
  auto x_r = stan::math::to_array_1d(a);
  equations_functor f;
  auto x = stan::math::algebra_solver(f, guess, theta, x_r, {});
  return x;
}
```

In this compiled RMarkdown document, the **knitr** package has exported functions `f`, `g`, `H`, `J` 
and `solution` (but not `equations_functor`) to R's global environment using the `sourceCpp` function 
in the **Rcpp** package, so that they can now be called from R. Here we find the optimum starting from 
a random point in three dimensions:
```{r}
x <- optim(rnorm(3), fn = f, gr = g, a = 1:3, method = "BFGS", hessian = TRUE)
x$par
x$hessian
H(x$par, a = 1:3)
J(x$par, a = 1:3)
solution(a = 1:3, guess = rnorm(3))
```

# Integrals and Ordinary Differential Equations

The Stan Math library can do one-dimensional numerical integration and can solve stiff
and non-stiff systems of differential equations, such as the harmonic oscillator example
below. Solving stiff systems utilizes the CVODES library, which is included in **StanHeaders**.
```{Rcpp}
// [[Rcpp::depends(BH)]]
// [[Rcpp::depends(RcppEigen)]]
// [[Rcpp::depends(RcppParallel)]]
// [[Rcpp::depends(StanHeaders)]]
#include <stan/math.hpp>                         // pulls in everything from rev/ and prim/
#include <Rcpp.h>
#include <RcppEigen.h>                           // do this AFTER including stan/math

// [[Rcpp::plugins(cpp17)]]

/* Definite integrals */

// [[Rcpp::export]]
double Cauchy(double scale) {
  std::vector<double> theta;
  auto half = stan::math::integrate_1d([](auto x, auto xc, auto theta,
                                          auto x_r, auto x_i, auto msgs) {
    return exp(stan::math::cauchy_lpdf(x, 0, x_r[0]));
  }, -scale, scale, theta, {scale}, {}, nullptr, 1e-7);
  return half * 2; // should equal 1 for any positive scale
}

/* Ordinary Differential Equations */

// [[Rcpp::export]]
auto nonstiff(Eigen::MatrixXd A, Eigen::VectorXd y0) {
  using stan::math::integrate_ode_rk45;
  using stan::math::to_vector;
  using stan::math::to_array_1d;
  std::vector<double> theta;
  std::vector<double> times = {1, 2};
  auto y = integrate_ode_rk45([&A](auto t, auto y, 
                                   auto theta, auto x_r, auto x_i, std::ostream *msgs) {
    return to_array_1d( (A * to_vector(y)).eval() );
  }, to_array_1d(y0), 0, times, theta, {}, {});
  Eigen::VectorXd truth = stan::math::matrix_exp(A) * y0;
  return (to_vector(y[0]) - truth).eval(); // should be "zero"
}

// [[Rcpp::export]]
auto stiff(Eigen::MatrixXd A, Eigen::VectorXd y0) { // not actually stiff
  using stan::math::integrate_ode_bdf;              // but use the stiff solver anyways
  using stan::math::to_vector;
  using stan::math::to_array_1d;
  std::vector<double> theta;
  std::vector<double> times = {1, 2};
  auto y = integrate_ode_bdf([&A](auto t, auto y, 
                                  auto theta, auto x_r, auto x_i, std::ostream *msgs) {
    return to_array_1d( (A * to_vector(y)).eval() );
  }, to_array_1d(y0), 0, times, theta, {}, {});
  Eigen::VectorXd truth = stan::math::matrix_exp(A) * y0;
  return (to_vector(y[0]) - truth).eval(); // should be "zero"
}
```

Again, in this compiled RMarkdown document, the **knitr** package has exported the
`Cauchy`, `nonstiff` and `stiff` functions to R's global environment using the
`sourceCpp` function in the **Rcpp** package so that they can be called from R.

First, we numerically integrate the Cauchy PDF over its interquartile range --- which has
an area of $\frac{1}{2}$ --- that we then double to verify that it is almost within machine
precision of $1$.

```{r}
all.equal(1, Cauchy(rexp(1)), tol = 1e-15)
```

Next, we consider the system of differential equations
$$\frac{d}{dt}\mathbf{y} = \mathbf{A}\mathbf{y}$$
where $\mathbf{A}$ is a square matrix such as that for a simple harmonic oscillator

$$\mathbf{A} = \begin{bmatrix}0 & 1 \\ -1 & -\theta\end{bmatrix}$$
for $\theta \in \left(0,1\right)$. The solution for $\mathbf{y}_t = e^{t\mathbf{A}}\mathbf{y}_0$ can 
be obtained via the matrix exponential function, which is available in the Stan Math Library, but it
can also be obtained numerically using a fourth-order Runge-Kutta solver, which is 
appropriate for non-stiff systems of ODEs, such as this one. However, it is possible,
albeit less efficient in this case, to use the backward-differentiation formula solver
for stiff systems of ODEs. In both cases, we calculate the difference between the
analytical solution and the numerical one, and the stiff version does produce somewhat
better accuracy in this case.
```{r}
A <- matrix(c(0, -1, 1, -runif(1)), nrow = 2, ncol = 2)
y0 <- rexp(2)
all.equal(nonstiff(A, y0), c(0, 0), tol = 1e-5)
all.equal(   stiff(A, y0), c(0, 0), tol = 1e-8)
```

# Parellelization

## `map_rect` Function

The Stan Math Library includes the `map_rect` function, which applies a function to each element
of rectangular arrays and returns a vector, making it a bit like a restricted version of R's `sapply` 
function. However, `map_rect` can also be executed in parallel by defining the pre-processor
directive `STAN_THREADS` and then setting the `STAN_NUM_THREADS` environmental variable to be the 
number of threads to use, as in
```{r}
Sys.setenv(STAN_NUM_THREADS = 2) # specify -1 to use all available cores
```

Below is C++ code to test whether an integer is prime, using a rather brute-force algorithm
and running it in parallel via `map_rect`.
```{Rcpp}
// [[Rcpp::depends(BH)]]
// [[Rcpp::depends(RcppEigen)]]
// [[Rcpp::depends(RcppParallel)]]
// [[Rcpp::depends(StanHeaders)]]
#include <stan/math.hpp>                         // pulls in everything from rev/ and prim/
#include <Rcpp.h>
#include <RcppEigen.h>                           // do this AFTER including stan/math

// [[Rcpp::plugins(cpp17)]]

// see https://en.wikipedia.org/wiki/Primality_test#Pseudocode
struct is_prime {
  is_prime() {}
  template <typename T1, typename T2>
  auto
  operator()(const Eigen::Matrix<T1, Eigen::Dynamic, 1>& eta,
             const Eigen::Matrix<T2, Eigen::Dynamic, 1>& theta,
             const std::vector<double>& x_r, const std::vector<int>& x_i,
             std::ostream* msgs = nullptr) const {
    Eigen::VectorXd res(1); // can only return double or var vectors
    int n = x_i[0];
    if (n <= 3) {
      res.coeffRef(0) = n > 1;
      return res;
    } else if ( (n % 2 == 0) || (n % 3 == 0) ) {
      res.coeffRef(0) = false;
      return res;
    }
    int i = 5;
    while (i * i <= n) {
      if ( (n % i == 0) || (n % (i + 2) == 0) ) {
        res.coeffRef(0) = false;
        return res;
      }
      i += 6;
    }
    res.coeffRef(0) = true;
    return res;
  }
};

/* parallelization */
// [[Rcpp::export]]
auto psapply(std::vector<std::vector<int> > n) {
  std::vector<Eigen::VectorXd> eta(n.size()); // these all have to be the same size
  Eigen::VectorXd theta;
  std::vector<std::vector<double> > x_d(n.size());
  return stan::math::map_rect<0, is_prime>(theta, eta, x_d, n, &Rcpp::Rcout);
}
```

Since the signature for `n` is a `std::vector<std::vector<int> >`, we have to pass
it from R as a list (which is converted to the outer `std::vector<>`) of integer
vectors (which is converted to the inner `std::vector<int>`) that happen to be of
size one in this case.
```{r}
odd <- seq.int(from = 2^25 - 1, to = 2^26 - 1, by = 2)
tail(psapply(n = as.list(odd))) == 1 # check your process manager while this is running
```
Thus, $2^{26} - 5 = 67,108,859$ is a prime number.

## `reduce_sum` Function

The `reduce_sum` function can be used to sum a function of recursively-partitioned
data in parallel. The Probability Mass Function (PMF) of the logarithmic distribution is

$$\Pr\left(x = k \mid p\right) = \frac{p^k}{-k\ln\left(1 - p\right)} $$
for a positive integer $k$ and $0 < p < 1$. To verify that this PMF sums to $1$ over
the natural numbers, we could accumulate it for a large finite number of terms, but
would like to do so in parallel. To do so, we need to define a `partial_sum` function
that adds the PMF for some subset of natural numbers and pass it to the `reduce_sum`
function like

```{Rcpp}
// [[Rcpp::depends(BH)]]
// [[Rcpp::depends(RcppEigen)]]
// [[Rcpp::depends(RcppParallel)]]
// [[Rcpp::depends(StanHeaders)]]
#include <stan/math.hpp>                         // pulls in everything from rev/ and prim/
#include <Rcpp.h>
#include <RcppEigen.h>                           // do this AFTER including stan/math

// [[Rcpp::plugins(cpp17)]]

template <typename T>
struct partial_sum {
  partial_sum() {}
  T operator()(const std::vector<int>& k_slice, 
               int start, int end, // these are ignored in this example
               std::ostream* msgs, double p) {
    double S = 0;
    for (int n = 0; n < k_slice.size(); n++) {
      int k = k_slice[n];
      S += stan::math::pow(p, k) / k;
    }
    return S;
  }
};

// [[Rcpp::export]]
double check_logarithmic_PMF(const double p, const int N = 1000) {
  using stan::math::reduce_sum;
  std::vector<int> k(N);
  std::iota(k.begin(), k.end(), 1);
  return reduce_sum<partial_sum<double> >(k, 1, nullptr, p) / -std::log1p(-p);
}
```

The second argument passed to `reduce_sum` is the `grainsize`, which governs
the size (and number) of the recursive subsets of `k` that are passed to the
`partial_sum` function. A `grainsize` of $1$ indicates that the software will
try to automatically tune the subset size for good performance, but that may
be worse than choosing a `grainsize` by hand in a problem-specific fashion.

In any event, we can call the wrapper function `check_logarithmic_PMF` in R
to verify that it returns $1$ to numerical tolerance:
```{r} 
stopifnot(all.equal(1, check_logarithmic_PMF(p = 1 / sqrt(2))))
```


# Defining a Stan Model in C++

The Stan _language_ does not have much support for sparse matrices for a variety of
reasons. Essentially the only applicable function is `csr_matrix_times_vector`, which 
pre-multiplies a vector by a sparse matrix in compressed row storage by taking as arguments 
its number of rows, columns, non-zero values, column indices of non-zero values, and
locations where the non-zero values start in each row. While the 
`csr_matrix_times_vector` function could be used to implement the example below,
we illustrate how to use the sparse data structures in the **Matrix** and **RcppEigen**
packages in a Stan model written in C++, which could easily be extended to more complicated
models with sparse data structures.

Our C++ file for the log-likelihood of a linear model with a sparse design matrix reads as
```{r, echo = FALSE, comment = ""}
cat(readLines("sparselm_stan.hpp"), sep = "\n")
```

To use it from R, we call the `exposeClass` function in the **Rcpp** package with
the necessary arguments and then call `sourceCpp` on the file it wrote in the
temporary directory:
```{r, message = FALSE}
library(Rcpp)
tf <- tempfile(fileext = "Module.cpp")
exposeClass("sparselm_stan",
      constructors = list(c("Eigen::Map<Eigen::SparseMatrix<double> >", 
                            "Eigen::VectorXd")),
      fields = c("X", "y"),
      methods = c("log_prob<>", "gradient<>"),
      rename = c(log_prob = "log_prob<>", gradient = "gradient<>"),
      header = c("// [[Rcpp::depends(BH)]]",
                 "// [[Rcpp::depends(RcppEigen)]]",
                 "// [[Rcpp::depends(RcppParallel)]",
                 "// [[Rcpp::depends(StanHeaders)]]",
                 "// [[Rcpp::plugins(cpp17)]]",
                 paste0("#include <", file.path(getwd(), "sparselm_stan.hpp"), ">")),
      file = tf,
      Rfile = FALSE)
Sys.setenv(PKG_CXXFLAGS = paste0(Sys.getenv("PKG_CXXFLAGS"), " -I",
                                 system.file("include", "src", 
                                             package = "StanHeaders", mustWork = TRUE)))
sourceCpp(tf)
sparselm_stan
```
At this point, we need a sparse design matrix and (dense) outcome vector to pass to
the constructor. The former can be created with a variety of functions in the **Matrix**
package, such as
```{r}
dd <- data.frame(a = gl(3, 4), b = gl(4, 1, 12))
X <- Matrix::sparse.model.matrix(~ a + b, data = dd)
X
```
Finally, we call the `new` function in the **methods** package, which essentially
calls our C++ constructor and provides an R interface to the instantiated object,
which contains the `log_prob` and `gradient` methods we defined and can be called
with arbitrary inputs.
```{r}
sm <- new(sparselm_stan, X = X, y = rnorm(nrow(X)))
sm$log_prob(c(beta = rnorm(ncol(X)), log_sigma = log(pi)))
round(sm$gradient(c(beta = rnorm(ncol(X)), log_sigma = log(pi))), digits = 4)
```