From 185bb3cd0504ee94fcce6de811f289764951cb36 Mon Sep 17 00:00:00 2001 From: Till Hoffmann Date: Fri, 15 Dec 2023 16:22:01 +0100 Subject: [PATCH] Add R package. --- .github/workflows/R-package.yml | 20 + .github/workflows/main.yml | 2 +- R-package/.Rbuildignore | 3 + R-package/.gitignore | 5 + R-package/DESCRIPTION | 17 + R-package/Makefile | 23 ++ R-package/NAMESPACE | 3 + R-package/R/model.R | 6 + R-package/inst/extdata/gptools/fft.stan | 2 + R-package/inst/extdata/gptools/fft1.stan | 173 +++++++++ R-package/inst/extdata/gptools/fft2.stan | 244 ++++++++++++ R-package/inst/extdata/gptools/graph.stan | 437 ++++++++++++++++++++++ R-package/inst/extdata/gptools/util.stan | 350 +++++++++++++++++ R-package/man/gptools_include_path.Rd | 11 + R-package/vignettes/getting_started.Rmd | 65 ++++ R-package/vignettes/getting_started.stan | 23 ++ 16 files changed, 1383 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/R-package.yml create mode 100644 R-package/.Rbuildignore create mode 100644 R-package/.gitignore create mode 100644 R-package/DESCRIPTION create mode 100644 R-package/Makefile create mode 100644 R-package/NAMESPACE create mode 100644 R-package/R/model.R create mode 100644 R-package/inst/extdata/gptools/fft.stan create mode 100644 R-package/inst/extdata/gptools/fft1.stan create mode 100644 R-package/inst/extdata/gptools/fft2.stan create mode 100644 R-package/inst/extdata/gptools/graph.stan create mode 100644 R-package/inst/extdata/gptools/util.stan create mode 100644 R-package/man/gptools_include_path.Rd create mode 100644 R-package/vignettes/getting_started.Rmd create mode 100644 R-package/vignettes/getting_started.stan diff --git a/.github/workflows/R-package.yml b/.github/workflows/R-package.yml new file mode 100644 index 0000000..483e5e0 --- /dev/null +++ b/.github/workflows/R-package.yml @@ -0,0 +1,20 @@ +name: CI - R + +on: + push: + branches: ["main"] + pull_request: + branches: ["main"] + workflow_dispatch: + +env: + cmdstanVersion: "2.31.0" + +jobs: + build: + name: R-package + runs-on: "ubuntu-latest" + steps: + - uses: "actions/checkout@v2" + - uses: "r-lib/actions/setup-r@v2" + \ No newline at end of file diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 66dfce3..d48e208 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -1,4 +1,4 @@ -name: CI +name: CI - Python on: push: diff --git a/R-package/.Rbuildignore b/R-package/.Rbuildignore new file mode 100644 index 0000000..61a8251 --- /dev/null +++ b/R-package/.Rbuildignore @@ -0,0 +1,3 @@ +^gptoolsStan\.Rproj$ +^R-package\.Rproj$ +^\.Rproj\.user$ diff --git a/R-package/.gitignore b/R-package/.gitignore new file mode 100644 index 0000000..0a0ebd5 --- /dev/null +++ b/R-package/.gitignore @@ -0,0 +1,5 @@ +.Rproj.user +*.Rproj +inst/doc +vignettes/getting_started +*.html diff --git a/R-package/DESCRIPTION b/R-package/DESCRIPTION new file mode 100644 index 0000000..83f8ab5 --- /dev/null +++ b/R-package/DESCRIPTION @@ -0,0 +1,17 @@ +Package: gptoolsStan +Title: Gaussian Processes on Graphs and Lattices in Stan. +Version: 0.1 +Authors@R: + person("Till", "Hoffmann", , "t.hoffmann@hsph.harvard.edu", role = c("aut", "cre"), + comment = c(ORCID = "0000-0003-4403-0722")) + person("Jukka-Pekka", "Onnela", , "onnela@hsph.harvard.edu", role = c("aut", "cre"), + comment = c(ORCID = "0000-0001-6613-8668")) +Description: Gaussian processes are flexible distributions to model functional data. Whilst theoretically appealing, they are computationally cumbersome except for small datasets. This package implements two methods for scaling GP inference in Stan: (a) a sparse approximation of the likelihood that is generally applicable and (b) an exact method for regularly spaced data modeled by stationary kernels using fast Fourier methods. +License: BSD-3-Clause +Encoding: UTF-8 +Roxygen: list(markdown = TRUE) +RoxygenNote: 7.2.3 +Suggests: + knitr, + rmarkdown +VignetteBuilder: knitr diff --git a/R-package/Makefile b/R-package/Makefile new file mode 100644 index 0000000..d30be61 --- /dev/null +++ b/R-package/Makefile @@ -0,0 +1,23 @@ +.PHONY : stan_files man vignettes +SOURCE_DIR = ../stan/gptools/stan/gptools +SOURCE_FILES = $(wildcard $(SOURCE_DIR)/*.stan) +TARGET_DIR = inst/extdata +TARGET_FILES = $(foreach f,$(SOURCE_FILES),$(TARGET_DIR)/gptools/$(notdir $f)) + +# Targets to copy over the files whenever they need updating. +stan_files : $(TARGET_FILES) vignettes/getting_started.stan + +$(TARGET_FILES) : $(TARGET_DIR)/gptools/%.stan : $(SOURCE_DIR)/%.stan + mkdir -p $(dir $@) + cp $< $@ + +vignettes/getting_started.stan : ../stan/docs/getting_started/getting_started.stan + mkdir -p $(dir $@) + cp $< $@ + +# Build the vignette. +vignettes/getting_started.html : vignettes/getting_started.Rmd vignettes/getting_started.stan stan_files + Rscript -e "devtools::build_rmd('$<')" + +man : + Rscript -e 'devtools::document()' diff --git a/R-package/NAMESPACE b/R-package/NAMESPACE new file mode 100644 index 0000000..41ed28a --- /dev/null +++ b/R-package/NAMESPACE @@ -0,0 +1,3 @@ +# Generated by roxygen2: do not edit by hand + +export(gptools_include_path) diff --git a/R-package/R/model.R b/R-package/R/model.R new file mode 100644 index 0000000..20bd390 --- /dev/null +++ b/R-package/R/model.R @@ -0,0 +1,6 @@ +#' Get the gptools include path for compiling Stan programs with `cmdstanr` or `RStan`. +#' +#' @export +gptools_include_path <- function() { + system.file("extdata", package="gptoolsStan") +} diff --git a/R-package/inst/extdata/gptools/fft.stan b/R-package/inst/extdata/gptools/fft.stan new file mode 100644 index 0000000..27d99b3 --- /dev/null +++ b/R-package/inst/extdata/gptools/fft.stan @@ -0,0 +1,2 @@ +#include gptools/fft1.stan +#include gptools/fft2.stan diff --git a/R-package/inst/extdata/gptools/fft1.stan b/R-package/inst/extdata/gptools/fft1.stan new file mode 100644 index 0000000..66710ab --- /dev/null +++ b/R-package/inst/extdata/gptools/fft1.stan @@ -0,0 +1,173 @@ +// IMPORTANT: stan uses the questionable R indexing which is one-based and inclusive on both ends. +// I.e., x[1:3] includes x[1] to x[3]. More generally, x[i:j] comprises j - i + 1 elements. It could +// at least have been exclusive on the right... + +/** +Evaluate the scale of Fourier coefficients. + +:param cov_rfft: Precomputed real fast Fourier transform of the kernel with shape + :code:`(..., n %/% 2 + 1)`. +:param n: Size of the real signal. Necessary because the size cannot be inferred from + :code:`cov_rfft`. +:returns: Scale of Fourier coefficients with shape :code:`(..., n %/% 2 + 1)`. +*/ +vector gp_evaluate_rfft_scale(vector cov_rfft, int n) { + int nrfft = n %/% 2 + 1; + vector[nrfft] result = n * cov_rfft / 2; + // Check positive-definiteness. + real minval = min(result); + if (minval < 0) { + reject("covariance matrix is not positive-definite (minimum eigenvalue is ", minval, ")"); + } + // The first element has larger scale because it only has a real part but must still have the + // right variance. The same applies to the last element if the number of elements is even + // (Nyqvist frequency). + result[1] *= 2; + if (n % 2 == 0) { + result[nrfft] *= 2; + } + return sqrt(result); +} + + +/** +Unpack the Fourier coefficients of a real Fourier transform with :code:`n %/% 2 + 1` elements to a +vector of :code:`n` elements. + +:param z: Real Fourier transform coefficients. +:param n: Size of the real signal. Necessary because the size cannot be inferred from :code:`rfft`. + +:returns: Unpacked vector of :code:`n` elements comprising the :code:`n %/% 2 + 1` real parts of the + zero frequency term, complex terms, and Nyqvist frequency term (for even :code:`n`). The + subsequent :code:`(n - 1) %/% 2` elements are the imaginary parts of complex coefficients. +*/ +vector gp_unpack_rfft(complex_vector x, int n) { + vector[n] z; + int ncomplex = (n - 1) %/% 2; + int nrfft = n %/% 2 + 1; + z[1:nrfft] = get_real(x); + z[1 + nrfft:n] = get_imag(x[2:1 + ncomplex]); + return z; +} + + +/** +Transform a Gaussian process realization to white noise in the Fourier domain. + +:param y: Realization of the Gaussian process with shape :code:`(..., n)`. +:param loc: Mean of the Gaussian process with shape :code:`(..., n)`. +:param cov_rfft: Precomputed real fast Fourier transform of the kernel with shape + :code:`(..., n // 2 + 1)`. +:returns: Fourier-domain white noise with shape :code:`(..., n)`. See :stan:func:`gp_unpack_rfft` + for details on the data structure. +*/ +vector gp_rfft(vector y, vector loc, vector cov_rfft) { + return gp_unpack_rfft(rfft(y - loc) ./ gp_evaluate_rfft_scale(cov_rfft, size(y)), size(y)); +} + + +/** +Evaluate the log absolute determinant of the Jacobian associated with +:stan:func:`gp_rfft`. + +:param cov_rfft: Precomputed real fast Fourier transform of the kernel with shape + :code:`(..., n %/% 2 + 1)`. +:param n: Size of the real signal. Necessary because the size cannot be inferred from :code:`rfft`. +:returns: Log absolute determinant of the Jacobian. +*/ +real gp_rfft_log_abs_det_jacobian(vector cov_rfft, int n) { + vector[n %/% 2 + 1] rfft_scale = gp_evaluate_rfft_scale(cov_rfft, n); + return - sum(log(rfft_scale[1:n %/% 2 + 1])) -sum(log(rfft_scale[2:(n + 1) %/% 2])) + - log(2) * ((n - 1) %/% 2) + n * log(n) / 2; +} + + +/** +Evaluate the log probability of a one-dimensional Gaussian process realization in Fourier space. + +:param y: Realization of a Gaussian process with :code:`n` grid points. +:param loc: Mean of the Gaussian process of size :code:`n`. +:param cov_rfft: Precomputed real fast Fourier transform of the kernel of size :code:`n %/% 2 + 1`. +:returns: Log probability of the Gaussian process realization. +*/ +real gp_rfft_lpdf(vector y, vector loc, vector cov_rfft) { + int n = size(y); + int nrfft = n %/% 2 + 1; + vector[n] z = gp_rfft(y, loc, cov_rfft); + return std_normal_lpdf(z) + gp_rfft_log_abs_det_jacobian(cov_rfft, n); +} + + +/** +Transform a real vector with :code:`n` elements to a vector of complex Fourier coefficients with +:code:`n` elements ready for inverse real fast Fourier transformation. +*/ +complex_vector gp_pack_rfft(vector z) { + int n = size(z); // Number of observations. + int ncomplex = (n - 1) %/% 2; // Number of complex Fourier coefficients. + int nrfft = n %/% 2 + 1; // Number of elements in the real FFT. + int neg_offset = (n + 1) %/% 2; // Offset at which the negative frequencies start. + // Zero frequency, real part of positive frequency coefficients, and Nyqvist frequency. + complex_vector[nrfft] fft = z[1:nrfft]; + // Imaginary part of positive frequency coefficients. + fft[2:ncomplex + 1] += 1.0i * z[nrfft + 1:n]; + return fft; +} + + +/** +Transform white noise in the Fourier domain to a Gaussian process realization, i.e., a +*non-centered* parameterization in the Fourier domain. + +The :code:`n` real white noise variables must be assembled into a length-:code:`n %/% 2 + 1` complex +vector with structure expected by the fast Fourier transform. The input vector :code:`z` comprises + +- the real zero frequency term, +- :code:`(n - 1) %/% 2` real parts of positive frequency terms, +- the real Nyqvist frequency term if :code:`n` is even, +- and :code:`(n - 1) %/% 2` imaginary parts of positive frequency terms. + +:param z: Fourier-domain white noise comprising :code:`n` elements. +:param loc: Mean of the Gaussian process. +:param cov_rfft: Real fast Fourier transform of the covariance kernel. + +:returns: Realization of the Gaussian process with :code:`n` elements. +*/ +vector gp_inv_rfft(vector z, vector loc, vector cov_rfft) { + int n = size(z); + vector[n %/% 2 + 1] rfft_scale = gp_evaluate_rfft_scale(cov_rfft, n); + return get_real(inv_rfft(rfft_scale .* gp_pack_rfft(z), n)) + loc; +} + +/** +Evaluate the real fast Fourier transform of the periodic squared exponential kernel. + +:param n: Number of grid points. +:param sigma: Scale of the covariance. +:param length_scale: Correlation length. +:param period: Period for circular boundary conditions. +:returns: Fourier transform of the squared exponential kernel of size :code:`n %/% 2 + 1`. +*/ +vector gp_periodic_exp_quad_cov_rfft(int n, real sigma, real length_scale, real period) { + int nrfft = n %/% 2 + 1; + return n * sigma ^ 2 * length_scale / period * sqrt(2 * pi()) + * exp(-2 * (pi() * linspaced_vector(nrfft, 0, nrfft - 1) * length_scale / period) ^ 2); +} + +/** +Evaluate the real fast Fourier transform of the periodic Matern kernel. + +:param nu: Smoothness parameter (1 / 2, 3 / 2, and 5 / 2 are typical values). +:param n: Number of grid points. +:param sigma: Scale of the covariance. +:param length_scale: Correlation length. +:param period: Period for circular boundary conditions. +:returns: Fourier transform of the squared exponential kernel of size :code:`n %/% 2 + 1`. +*/ +vector gp_periodic_matern_cov_rfft(real nu, int n, real sigma, real length_scale, real period) { + int nrfft = n %/% 2 + 1; + vector[nrfft] k = linspaced_vector(nrfft, 0, nrfft - 1); + return sigma ^ 2 * n * sqrt(2 * pi() / nu) * tgamma(nu + 0.5) / tgamma(nu) + * (1 + 2 / nu * (pi() * length_scale / period * k) ^ 2) ^ -(nu + 0.5) * length_scale + / period; +} diff --git a/R-package/inst/extdata/gptools/fft2.stan b/R-package/inst/extdata/gptools/fft2.stan new file mode 100644 index 0000000..9b810fc --- /dev/null +++ b/R-package/inst/extdata/gptools/fft2.stan @@ -0,0 +1,244 @@ +// IMPORTANT: stan uses the questionable R indexing which is one-based and inclusive on both ends. +// I.e., x[1:3] includes x[1] to x[3]. More generally, x[i:j] comprises j - i + 1 elements. It could +// at least have been exclusive on the right... + +/** +Evaluate the scale of Fourier coefficients. + +:param cov_rfft2: Precomputed real fast Fourier transform of the kernel with shape + :code:`(height, width %/% 2 + 1)`. +:param width: Number of columns of the signal (cannot be inferred from the Fourier coefficients). +:returns: Scale of Fourier coefficients with shape :code:`(height, width %/% 2 + 1)`. +*/ +matrix gp_evaluate_rfft2_scale(matrix cov_rfft2, int width) { + int height = rows(cov_rfft2); + int n = width * height; + int fftwidth = width %/% 2 + 1; + int fftheight = height %/% 2 + 1; + matrix[height, fftwidth] fftscale = n * cov_rfft2 / 2; + // Check positive-definiteness. + real minval = min(fftscale); + if (minval < 0) { + reject("covariance matrix is not positive-definite (minimum eigenvalue is ", minval, ")"); + } + + // Adjust the scale for the zero-frequency (and Nyqvist) terms in the first column. + fftscale[1, 1] *= 2; + if (height % 2 == 0) { + fftscale[fftheight, 1] *= 2; + } + // If the width is even, the last column has the same structure as the first column. + if (width % 2 == 0) { + fftscale[1, fftwidth] *= 2; + } + // If the number of rows and columns is even, we also have a Nyqvist frequency term in the last + // column. + if (width % 2 == 0 && height % 2 == 0) { + fftscale[fftheight, fftwidth] *= 2; + } + return sqrt(fftscale); +} + + +/** +Unpack the complex Fourier coefficients of a two-dimensional real Fourier transform with shape to a +real matrix with shape :code:`(height, width)`. + +TODO: add details on packing structure. + +:param z: Two-dimensional real Fourier transform coefficients with shape + :code:`(height, width %/% 2 + 1)`. +:param width: Number of columns of the signal (cannot be inferred from the Fourier coefficients + :code:`z`). +:returns: Unpacked matrix with shape :code:`(height, width)`. +*/ +matrix gp_unpack_rfft2(complex_matrix z, int m) { + int height = rows(z); + int n = m * height; + int fftwidth = m %/% 2 + 1; + int fftheight = height %/% 2 + 1; + int wcomplex = (m - 1) %/% 2; + + matrix[height, m] result; + + // First column is always real. + result[:, 1] = gp_unpack_rfft(z[:fftheight, 1], height); + // Real and imaginary parts of complex coefficients. + result[:, 2:wcomplex + 1] = get_real(z[:, 2:wcomplex + 1]); + result[:, 2 + wcomplex:2 * wcomplex + 1] = get_imag(z[:, 2:wcomplex + 1]); + // Nyqvist frequency if the number of columns is even. + if (m % 2 == 0) { + result[:, m] = gp_unpack_rfft(z[:fftheight, fftwidth], height); + } + return result; +} + + +/** +Transform a real matrix with shape :code:`(height, width)` to a matrix of complex Fourier +coefficients with shape :code:`(height, width %/% 2 + 1)` ready for inverse real fast Fourier +transformation in two dimensions. + +:param z: Unpacked matrices with shape :code:`(height, width)`. +:returns: Two-dimensional real Fourier transform coefficients. +*/ +complex_matrix gp_pack_rfft2(matrix z) { + int height = rows(z); + int width = cols(z); + int ncomplex = (width - 1) %/% 2; + complex_matrix[height, width %/% 2 + 1] result; + // Real FFT in the first column due to zero-frequency terms for the row-wise Fourier transform. + result[:, 1] = expand_rfft(gp_pack_rfft(z[:, 1]), height); + // Complex Fourier coefficients. + result[:, 2:ncomplex + 1] = z[:, 2:ncomplex + 1] + 1.0i * z[:, ncomplex + 2:2 * ncomplex + 1]; + // Real FFT in the last column due to the Nyqvist frequency terms for the row-wise Fourier + // transform if the number of columns is even. + if (width % 2 == 0) { + result[:, width %/% 2 + 1] = expand_rfft(gp_pack_rfft(z[:, width]), height); + } + return result; +} + + +/** +Transform a Gaussian process realization to white noise in the Fourier domain. + +:param y: Realization of the Gaussian process with shape :code:`(height, width)`. +:param loc: Mean of the Gaussian process with shape :code:`(height, width)`. +:param cov_rfft2: Precomputed real fast Fourier transform of the kernel with shape + :code:`(height, width %/% 2 + 1)`. +:returns: Unpacked matrix with shape :code:`(height, width)`. +*/ +matrix gp_rfft2(matrix y, matrix loc, matrix cov_rfft2) { + return gp_unpack_rfft2(rfft2(y - loc) ./ gp_evaluate_rfft2_scale(cov_rfft2, cols(y)), cols(y)); +} + + +/** +Transform white noise in the Fourier domain to a Gaussian process realization. + +:param z: Unpacked matrix with shape :code:`(height, width)`. +:param loc: Mean of the Gaussian process with shape :code:`(height, width)`. +:param cov_rfft2: Precomputed real fast Fourier transform of the kernel with shape + :code:`(height, width %/% 2 + 1)`. +:returns: Realization of the Gaussian process. +*/ +matrix gp_inv_rfft2(matrix z, matrix loc, matrix cov_rfft2) { + complex_matrix[rows(z), cols(z) %/% 2 + 1] y = gp_pack_rfft2(z) + .* gp_evaluate_rfft2_scale(cov_rfft2, cols(z)); + return inv_rfft2(y, cols(z)) + loc; +} + + +/** +Evaluate the log absolute determinant of the Jacobian associated with +:stan:func:`gp_rfft2`. + +:param cov_rfft2: Precomputed real fast Fourier transform of the kernel with shape + :code:`(height, width %/% 2 + 1)`. +:param width: Number of columns of the signal (cannot be inferred from the precomputed kernel + Fourier transform :code:`cov_rfft2`). +:returns: Log absolute determinant of the Jacobian. +*/ +real gp_rfft2_log_abs_det_jacobian(matrix cov_rfft2, int width) { + int height = rows(cov_rfft2); + int n = width * height; + int fftwidth = width %/% 2 + 1; + int fftheight = height %/% 2 + 1; + matrix[height, fftwidth] log_rfft2_scale = log(gp_evaluate_rfft2_scale(cov_rfft2, width)); + real ladj = 0; + + // For the real part, we always use the full height of the non-redundant part. For the imaginary + // part, we discard the last element if the number of rows is even because it's the real Nyqvist + // frequency. + int idx = (height % 2) ? fftheight : fftheight - 1; + ladj -= sum(log_rfft2_scale[:fftheight, 1]) + sum(log_rfft2_scale[2:idx, 1]); + + // Evaluate the "bulk" likelihood that needs no adjustment. + ladj -= 2 * sum(to_vector(log_rfft2_scale[:, 2:fftwidth - 1])); + + if (width % 2) { + // If the width is odd, the last column comprises all-independent terms. + ladj -= 2 * sum(log_rfft2_scale[:, fftwidth]); + } else { + ladj -= sum(log_rfft2_scale[:fftheight, fftwidth]) + sum(log_rfft2_scale[2:idx, fftwidth]); + } + // Correction terms from the transform that only depend on the shape. + int nterms = (n - 1) %/% 2; + if (height % 2 == 0 && width % 2 == 0) { + nterms -=1; + } + ladj += - log2() * nterms + n * log(n) / 2; + return ladj; +} + + +/** +Evaluate the log probability of a two-dimensional Gaussian process realization in Fourier space. + +:param y: Realization of a Gaussian process with shape :code:`(height, width)`, where :code:`height` + is the number of rows, and :code:`width` is the number of columns. +:param loc: Mean of the Gaussian process with shape :code:`(height, width)`. +:param cov_rfft2: Precomputed real fast Fourier transform of the kernel with shape + :code:`(height, width %/% 2 + 1)`. +:returns: Log probability of the Gaussian process realization. +*/ +real gp_rfft2_lpdf(matrix y, matrix loc, matrix cov_rfft2) { + return std_normal_lpdf(to_vector(gp_rfft2(y, loc, cov_rfft2))) + + gp_rfft2_log_abs_det_jacobian(cov_rfft2, cols(y)); +} + + +/** +Evaluate the two-dimensional real fast Fourier transform of the periodic squared exponential kernel. + +:param height: Number of grid rows. +:param width: Number of grid columns. +:param sigma: Scale of the covariance. +:param length_scale: Correlation lengths in each dimension. +:param period: Periods for circular boundary conditions in each dimension. +:returns: Fourier transform of the periodic squared exponential kernel with shape + :code:`(height, width %/% 2 + 1`). +*/ +matrix gp_periodic_exp_quad_cov_rfft2(int height, int width, real sigma, vector length_scale, + vector period) { + vector[height %/% 2 + 1] rfftm = gp_periodic_exp_quad_cov_rfft(height, sigma, length_scale[1], + period[1]); + vector[width %/% 2 + 1] rfftn = gp_periodic_exp_quad_cov_rfft(width, 1, length_scale[2], + period[2]); + return get_real(expand_rfft(rfftm, height)) * rfftn'; +} + + +/** +Evaluate the two-dimensional real fast Fourier transform of the periodic Matern kernel. + +:param nu: Smoothness parameter (1 / 2, 3 / 2, and 5 / 2 are typical values). +:param height: Number of grid rows. +:param width: Number of grid columns. +:param sigma: Scale of the covariance. +:param length_scale: Correlation lengths in each dimension. +:param period: Periods for circular boundary conditions in each dimension. +:returns: Fourier transform of the periodic Matern kernel with shape + :code:`(height, width %/% 2 + 1`). +*/ +matrix gp_periodic_matern_cov_rfft2(real nu, int height, int width, real sigma, vector length_scale, + vector period) { + int nrfft = width %/% 2 + 1; + matrix[height, nrfft] result; + real ndim = 2; + row_vector[nrfft] col_part = (linspaced_row_vector(nrfft, 0, nrfft - 1) * length_scale[2] + / period[2]) ^ 2; + // We only iterate up to height %/% 2 + 1 because the kernel is symmetric in positive and + // negative frequencies. + for (i in 1:height %/% 2 + 1) { + int krow = i - 1; + result[i] = 1 + 2 / nu * pi() ^ 2 * ((krow * length_scale[1] / period[1]) ^ 2 + col_part); + if (i > 1) { + result[height - i + 2] = result[i]; + } + } + return sigma ^ 2 * height * width * 2 ^ ndim * (pi() / (2 * nu)) ^ (ndim / 2) + * tgamma(nu + ndim / 2) / tgamma(nu) + * result .^ -(nu + ndim / 2) * prod(to_array_1d(length_scale ./ period)); +} diff --git a/R-package/inst/extdata/gptools/graph.stan b/R-package/inst/extdata/gptools/graph.stan new file mode 100644 index 0000000..ae98e58 --- /dev/null +++ b/R-package/inst/extdata/gptools/graph.stan @@ -0,0 +1,437 @@ +// General functions for graph Gaussian processes -------------------------------------------------- + +/** +Evaluate the out-degree of nodes in the directed acyclic graph induced by :code:`edge_index`. The +node labels of successors must be ordered and predecessors must have index less than successors. + +:param n: Number of nodes. +:param edge_index: Directed edges as a tuple of predecessor and successor indices, i.e., the node + labelled :code:`edge_index[1, i]` is the predecessor of the node labelled + :code:`edge_index[2, i]`. +:returns: Out-degree of each node. +*/ +array [] int out_degrees(int n, array [,] int edge_index) { + array [n] int count = rep_array(0, n); + int previous = 0; + for (i in 1:size(edge_index[2])) { + int current = edge_index[2, i]; + if (previous > current) { + reject("nodes are not ordered: ", previous, " > ", current, " at ", i); + } + if (edge_index[1, i] == current) { + reject("self-loops are not allowed: ", current, " at ", i); + } + if (edge_index[1, i] > current) { + reject("predecessor is greater than successor: ", edge_index[1, i], " > ", current, + " at ", i); + } + count[current] += 1; + previous = current; + } + return count; +} + + +/** +Evaluate the location and scale for a node given its predecessors, assuming zero mean. + +:param y: State of each node. +:param x: Array of :code:`k + 1` locations in :code:`p` dimensions. +:param kernel: Kernel to use (0 for squared exponential, 1 for Matern 3/2, 2 for Matern 5/2). +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param predecessors: Labels of :code:`k` predecessors of the target node. +:param epsilon: Nugget variance for numerical stability. +:returns: Location and scale parameters for the distribution of the node given its predecessors. +*/ +vector gp_graph_conditional_loc_scale(vector y, array[] vector x, int kernel, real sigma, + real length_scale, int node, array [] int predecessors, + real epsilon) { + int k = size(predecessors); + matrix[1, k] cov12; + matrix[k, k] cov22; + if (kernel == 0) { + cov12 = gp_exp_quad_cov({x[node]}, x[predecessors], sigma, length_scale); + cov22 = gp_exp_quad_cov(x[predecessors], sigma, length_scale); + } else if (kernel == 1) { + cov12 = gp_matern32_cov({x[node]}, x[predecessors], sigma, length_scale); + cov22 = gp_matern32_cov(x[predecessors], sigma, length_scale); + } else if (kernel == 2) { + cov12 = gp_matern52_cov({x[node]}, x[predecessors], sigma, length_scale); + cov22 = gp_matern52_cov(x[predecessors], sigma, length_scale); + } else { + reject("invalid kernel indicator: ", kernel); + } + return gp_conditional_loc_scale(y[predecessors], sigma ^ 2 + epsilon, to_vector(cov12), + add_diag(cov22, epsilon)); +} + + +/** +Evaluate the log probability of a graph Gaussian process. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param kernel: Kernel to use (0 for squared exponential, 1 for Matern 3/2, 2 for Matern 5/2). +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_lpdf(vector y, vector loc, array [] vector x, int kernel, real sigma, + real length_scale, array [,] int edges, array[] int degrees, real epsilon) { + real lpdf = 0; + int offset_ = 1; + vector[size(y)] z = y - loc; + for (i in 1:size(x)) { + vector[2] loc_scale = gp_graph_conditional_loc_scale( + z, x, kernel, sigma, length_scale, i, segment(edges[1], offset_, degrees[i]), epsilon); + lpdf += normal_lpdf(z[i] | loc_scale[1], loc_scale[2]); + offset_ += degrees[i]; + } + return lpdf; +} + + +/** +Evaluate the log probability of a graph Gaussian process. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param kernel: Kernel to use (0 for squared exponential, 1 for Matern 3/2, 2 for Matern 5/2). +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_lpdf(vector y, vector loc, array [] vector x, int kernel, real sigma, + real length_scale, array [,] int edges) { + return gp_graph_lpdf(y | loc, x, kernel, sigma, length_scale, edges, + out_degrees(size(y), edges), 1e-12); +} + + +/** +Transform white noise to a sample from a graph Gaussian process + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param kernel: Kernel to use (0 for squared exponential, 1 for Matern 3/2, 2 for Matern 5/2). +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph(vector z, vector loc, array [] vector x, int kernel, real sigma, + real length_scale, array [,] int edges, array [] int degrees, real epsilon) { + vector[size(z)] y; + int offset_ = 1; + for (i in 1:size(x)) { + vector[2] loc_scale = gp_graph_conditional_loc_scale( + y, x, kernel, sigma, length_scale, i, segment(edges[1], offset_, degrees[i]), epsilon); + y[i] = loc_scale[1] + loc_scale[2] * z[i]; + offset_ += degrees[i]; + } + return y + loc; +} + + +/** +Transform white noise to a sample from a graph Gaussian process + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param kernel: Kernel to use (0 for squared exponential, 1 for Matern 3/2, 2 for Matern 5/2). +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph(vector z, vector loc, array [] vector x, int kernel, real sigma, + real length_scale, array [,] int edges) { + return gp_inv_graph(z, loc, x, kernel, sigma, length_scale, edges, + out_degrees(size(z), edges), 1e-12); +} + +// Functions for graph Gaussian processes with squared exponential kernel -------------------------- + +/** +Evaluate the log probability of a graph Gaussian with squared exponential kernel. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_exp_quad_cov_lpdf(vector y, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges, array[] int degrees, + real epsilon) { + return gp_graph_lpdf(y | loc, x, 0, sigma, length_scale, edges, degrees, epsilon); +} + + +/** +Evaluate the log probability of a graph Gaussian with squared exponential kernel. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_exp_quad_cov_lpdf(vector y, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges) { + return gp_graph_lpdf(y | loc, x, 0, sigma, length_scale, edges); +} + + +/** +Transform white noise to a sample from a graph Gaussian process with squared exponential kernel. + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph_exp_quad_cov(vector z, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges, array [] int degrees, + real epsilon) { + return gp_inv_graph(z, loc, x, 0, sigma, length_scale, edges, degrees, epsilon); +} + + +/** +Transform white noise to a sample from a graph Gaussian process with squared exponential kernel. + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph_exp_quad_cov(vector z, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges) { + return gp_inv_graph(z, loc, x, 0, sigma, length_scale, edges); +} + + +// Functions for graph Gaussian processes with Matern 3 / 2 kernel --------------------------------- + +/** +Evaluate the log probability of a graph Gaussian with Matern 3 / 2 kernel. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_matern32_cov_lpdf(vector y, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges, array[] int degrees, + real epsilon) { + return gp_graph_lpdf(y | loc, x, 1, sigma, length_scale, edges, degrees, epsilon); +} + + +/** +Evaluate the log probability of a graph Gaussian with Matern 3 / 2 kernel. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_matern32_cov_lpdf(vector y, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges) { + return gp_graph_lpdf(y | loc, x, 1, sigma, length_scale, edges); +} + + +/** +Transform white noise to a sample from a graph Gaussian process with Matern 3 / 2 kernel. + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph_matern32_cov(vector z, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges, array [] int degrees, + real epsilon) { + return gp_inv_graph(z, loc, x, 1, sigma, length_scale, edges, degrees, epsilon); +} + + +/** +Transform white noise to a sample from a graph Gaussian process with Matern 3 / 2 kernel. + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph_matern32_cov(vector z, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges) { + return gp_inv_graph(z, loc, x, 1, sigma, length_scale, edges); +} + + +// Functions for graph Gaussian processes with Matern 5 / 2 kernel --------------------------------- + +/** +Evaluate the log probability of a graph Gaussian with Matern 5 / 2 kernel. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_matern52_cov_lpdf(vector y, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges, array[] int degrees, + real epsilon) { + return gp_graph_lpdf(y | loc, x, 2, sigma, length_scale, edges, degrees, epsilon); +} + + +/** +Evaluate the log probability of a graph Gaussian with Matern 5 / 2 kernel. + +:param y: State of each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Log probability of the graph Gaussian process. +*/ +real gp_graph_matern52_cov_lpdf(vector y, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges) { + return gp_graph_lpdf(y | loc, x, 2, sigma, length_scale, edges); +} + + +/** +Transform white noise to a sample from a graph Gaussian process with Matern 5 / 2 kernel. + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:param degrees: Out-degree of each node. +:param epsilon: Nugget variance for numerical stability. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph_matern52_cov(vector z, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges, array [] int degrees, + real epsilon) { + return gp_inv_graph(z, loc, x, 2, sigma, length_scale, edges, degrees, epsilon); +} + + +/** +Transform white noise to a sample from a graph Gaussian process with Matern 5 / 2 kernel. + +:param z: White noise for each node. +:param loc: Mean of each node. +:param x: Position of each node. +:param sigma: Marginal scale of the kernel. +:param length_scale: Correlation length of the kernel. +:param edges: Directed edges between nodes constituting a directed acyclic graph. Edges are stored + as a matrix with shape :code:`(2, m)`, where :code:`m` is the number of edges. The first row + comprises parents of children in the second row. The first row can have arbitrary order, but the + second row must be sorted. +:returns: Sample from the Graph gaussian process. +*/ +vector gp_inv_graph_matern52_cov(vector z, vector loc, array [] vector x, real sigma, + real length_scale, array [,] int edges) { + return gp_inv_graph(z, loc, x, 2, sigma, length_scale, edges); +} diff --git a/R-package/inst/extdata/gptools/util.stan b/R-package/inst/extdata/gptools/util.stan new file mode 100644 index 0000000..ad5d6f1 --- /dev/null +++ b/R-package/inst/extdata/gptools/util.stan @@ -0,0 +1,350 @@ +// Scalars ----------------------------------------------------------------------------------------- + +/** +Assert that two integers are equal. + +:param actual: Actual value. +:param desired: Desired value. +*/ +void assert_equal(int actual, int desired) { + if (actual != desired) { + reject(actual, " is not equal to ", desired); + } +} + + +/** +Check whether two values are close. + +The actual value :code:`x` and desired value :code:`y` may differ by at most +:code:`tol = rtol * y + atol`, where :code:`rtol` is the relative tolerance, and :code:`atol` is the +absolute tolerance. The tolerance :code:`tol` is clipped below at :math:`10^{-15}` to avoid +rejection due to rounding errors. + +:param actual: Actual value :code:`x`. +:param desired: Desired value :code:`y`. +:param rtol: Relative tolerance :code:`r`. +:param atol: Absolute tolerance :code:`a`. +:returns: :code:`1` if the values are close, :code:`0` otherwise. +*/ +int is_close(real actual, real desired, real rtol, real atol) { + // We always allow a tolerance of at least 1e-15 in case there are rounding errors. + real tol = fmax(atol + rtol * abs(desired), 1e-15); + if (abs(actual - desired) <= tol) { + return 1; + } + return 0; +} + +/** +Assert that two values are close. See :stan:func:`is_close(real, real, real, real)` for a +description of the parameters. +*/ +void assert_close(real actual, real desired, real rtol, real atol) { + if (!is_close(actual, desired, rtol, atol)) { + reject(actual, " is not close to ", desired); + } +} + +/** +Assert that two values are close. See :stan:func:`is_close(real, real, real, real)` for a +description of the parameters. +*/ +void assert_close(real actual, real desired) { + assert_close(actual, desired, 1e-6, 0); +} + +/** +Check whether a possibly complex value is finite. + +:param x: Value to check. +:returns: :code:`1` if the value is finite, :code:`0` otherwise. +*/ +int is_finite(complex x) { + real rx = get_real(x); + real ix = get_imag(x); + if (is_nan(rx) || is_nan(ix) || is_inf(rx) || is_inf(ix)) { + return 0; + } + return 1; +} + + +// Vectors ----------------------------------------------------------------------------------------- + +/** +Assert that two vectors are close. See :stan:func:`is_close(real, real, real, real)` for a +description of the parameters. +*/ +void assert_close(vector actual, vector desired, real rtol, real atol) { + int n = size(desired); + int m = size(actual); + if (m != n) { + reject("number of elements are not equal: size(desired)=", n, "; size(actual)=", m); + } + for (i in 1:size(actual)) { + if (!is_close(actual[i], desired[i], rtol, atol)) { + reject(actual[i], " is not close to ", desired[i], " at position ", i); + } + } +} + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(vector actual, vector desired) { + assert_close(actual, desired, 1e-6, 0); +} + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(vector actual, real desired, real rtol, real atol) { + assert_close(actual, rep_vector(desired, size(actual)), rtol, atol); +} + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(vector actual, real desired) { + assert_close(actual, desired, 1e-6, 0); +} + +/** +Check whether all elements of a vector are finite. + +:param x: Vector to check. +:returns: :code:`1` if all elements of the vector are finite, :code:`0` otherwise. +*/ +int is_finite(vector x) { + for (i in 1:size(x)) { + if(!is_finite(x[i])) { + return 0; + } + } + return 1; +} + +/** +Assert that all elements of a vector are finite. + +:param: Vector to check. +*/ +void assert_finite(vector x) { + int n = size(x); + for (i in 1:n) { + if (!is_finite(x[i])) { + reject(x[i], " at index ", i, " is not finite"); + } + } +} + +// Matrices ---------------------------------------------------------------------------------------- + +/** +Check whether all elements of a matrix are finite. + +:param: Vector to check. +*/ +int is_finite(matrix x) { + for (i in 1:rows(x)) { + for (j in 1:cols(x)) { + if (!is_finite(x[i, j])) { + return 0; + } + } + } + return 1; +} + + +/** +Pretty-print a matrix. +*/ +void print_matrix(complex_matrix x) { + print("matrix with ", rows(x), " rows and ", cols(x), " columns"); + for (i in 1:rows(x)) { + print(x[i]); + } +} + +/** +Assert that two matrices are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(matrix actual, matrix desired, real rtol, real atol) { + array [2] int nshape = dims(desired); + array [2] int mshape = dims(actual); + if (mshape[1] != nshape[1]) { + reject("number of rows are not equal: dims(desired)[1]=", nshape[1], "; size(actual)[1]=", + mshape[1]); + } + if (mshape[2] != nshape[2]) { + reject("number of columns are not equal: dims(desired)[2]=", nshape[2], + "; size(actual)[2]=", mshape[2]); + } + for (i in 1:nshape[1]) { + for (j in 1:nshape[2]) { + if (!is_close(actual[i, j], desired[i, j], rtol, atol)) { + reject(actual[i, j], " is not close to ", desired[i, j], " at row ", i, ", column ", + j); + } + } + } +} + +/** +Assert that two matrices are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(matrix actual, matrix desired) { + assert_close(actual, desired, 1e-6, 0); +} + +/** +Assert that two matrices are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(matrix actual, real desired, real rtol, real atol) { + array [2] int shape = dims(actual); + assert_close(actual, rep_matrix(desired, shape[1], shape[2]), rtol, atol); +} + +/** +Assert that two matrices are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(matrix actual, real desired) { + assert_close(actual, desired, 1e-6, 0); +} + +// Complex vectors --------------------------------------------------------------------------------- + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(complex_vector actual, complex_vector desired, real rtol, real atol) { + int n = size(desired); + int m = size(actual); + if (m != n) { + reject("number of elements are not equal: size(desired)=", n, "; size(actual)=", m); + } + assert_close(get_real(actual), get_real(desired), rtol, atol); + assert_close(get_imag(actual), get_imag(desired), rtol, atol); +} + + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(complex_vector actual, complex_vector desired) { + assert_close(actual, desired, 1e-6, 0); +} + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(complex_vector actual, complex desired, real rtol, real atol) { + assert_close(actual, rep_vector(desired, size(actual)), rtol, atol); +} + +/** +Assert that two vectors are close. See :stan:func:`is_close` for description of parameters. +*/ +void assert_close(complex_vector actual, complex desired) { + assert_close(actual, desired, 1e-6, 0); +} + + +// Real Fourier transforms ------------------------------------------------------------------------- + +/** +Compute the one-dimensional discrete Fourier transform for real input. + +:param y: Real signal with :code:`n` elements to transform. +:returns: Truncated vector of Fourier coefficients with :code:`n %/% 2 + 1` elements. +*/ +complex_vector rfft(vector y) { + return fft(y)[:size(y) %/% 2 + 1]; +} + +/** +Expand truncated one-dimensional discrete Fourier transform coefficients for real input to full +Fourier coefficients. +*/ +complex_vector expand_rfft(complex_vector y, int n) { + complex_vector[n] result; + int nrfft = n %/% 2 + 1; + if (size(y) != nrfft) { + reject("expected complex vector with ", nrfft, " elements but got ", size(y)); + } + int ncomplex = (n - 1) %/% 2; + result[:nrfft] = y; + result[nrfft + 1:n] = conj(reverse(y[2:1 + ncomplex])); + return result; +} + +/** +Compute the one-dimensional inverse discrete Fourier transform for real output. + +:param z: Truncated vector of Fourier coefficents with :code:`n %/% 2 + 1` elements. +:param n: Length of the signal (required because the length of the signal cannot be determined from + :code:`z` alone). +:returns: Real signal with :code:`n` elements. +*/ +vector inv_rfft(complex_vector z, int n) { + return get_real(inv_fft(expand_rfft(z, n))); +} + +/** +Compute the two-dimensional discrete Fourier transform for real input. + +:param y: Real signal with :code:`n` rows and :code:`m` columns to transform. +:returns: Truncated vector of Fourier coefficients with :code:`n` rows and :code:`m %/% 2 + 1` + elements. +*/ +complex_matrix rfft2(matrix y) { + return fft2(y)[:, :cols(y) %/% 2 + 1]; +} + +/** +Compute the two-dimensional inverse discrete Fourier transform for real output. + +:param z: Truncated vector of Fourier coefficients with :code:`n` rows and :code:`m %/% 2 + 1` + elements. +:param m: Number of columns of the signal (required because the number of columns cannot be + determined from :code:`z` alone). +:returns: Real signal with :code:`n` rows and :code:`m` columns. +*/ +matrix inv_rfft2(complex_matrix z, int m) { + int n = rows(z); + complex_matrix[n, m] x; + int mrfft = m %/% 2 + 1; + int mcomplex = (m - 1) %/% 2; + x[:, 1:mrfft] = z[:, 1:mrfft]; + // Fill redundant values. + for (i in 1:n) { + x[i, mrfft + 1:m] = conj(reverse(z[i, 2:1 + mcomplex])); + } + // Reverse the order to account for negative frequencies. + for (i in mrfft + 1:mrfft + mcomplex) { + x[2:, i] = reverse(x[2:, i]); + } + return get_real(inv_fft2(x)); +} + +// Conditional location and scale parameters for multivariate normal distributions ----------------- + +/** +Evaluate the conditional location and scale parameter of a univariate normal random variable given +correlated observations from a multivariate normal distribution. + +:param y: Observation to condition on. +:param cov11: Marginal variance of the target random variable. +:param cov21: Covariance between :code:`y` and the target random variable. +:param cov22: Covariance amongst the elements of :code:`y`. +:returns: Location and scale as a vector. +*/ +vector gp_conditional_loc_scale(vector y, real cov11, vector cov21, matrix cov22) { + if (size(y) == 0) { + return [0, sqrt(cov11)]'; + } + vector[size(y)] v = mdivide_left_spd(cov22, cov21); + return [dot_product(v, y), sqrt(cov11 - dot_product(v, cov21))]'; +} diff --git a/R-package/man/gptools_include_path.Rd b/R-package/man/gptools_include_path.Rd new file mode 100644 index 0000000..4e86853 --- /dev/null +++ b/R-package/man/gptools_include_path.Rd @@ -0,0 +1,11 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/model.R +\name{gptools_include_path} +\alias{gptools_include_path} +\title{Get the gptools include path for compiling Stan programs with \code{cmdstanr} or \code{RStan}.} +\usage{ +gptools_include_path() +} +\description{ +Get the gptools include path for compiling Stan programs with \code{cmdstanr} or \code{RStan}. +} diff --git a/R-package/vignettes/getting_started.Rmd b/R-package/vignettes/getting_started.Rmd new file mode 100644 index 0000000..c8c4481 --- /dev/null +++ b/R-package/vignettes/getting_started.Rmd @@ -0,0 +1,65 @@ +--- +title: Getting Started with gptools in R +output: rmarkdown::html_vignette +vignette: > + %\VignetteIndexEntry{Getting Started with gptools in R} + %\VignetteEngine{knitr::rmarkdown} + %\VignetteEncoding{UTF-8} +--- + +```{r, include = FALSE} +knitr::opts_chunk$set( + collapse = TRUE, + comment = "#>" +) +``` + +`gptoolsStan` is a minimal package to publish Stan code for efficient Gaussian process inference. The package can be used with any R interface for Stan, such as [`cmdstanr`](https://mc-stan.org/cmdstanr/) or [`Rstan`](https://mc-stan.org/rstan/). If you're already familiar with Stan and know which interface you want to use, dive in below. If not, have a look at the getting started guides for the corresponding interface (see [here](https://mc-stan.org/cmdstanr/articles/cmdstanr.html) for `cmdstanr` and [here](https://github.com/stan-dev/rstan/wiki/RStan-Getting-Started) for `RStan`). + +This vignette demonstrates the package by sampling from a simple Gaussian process using Fourier methods (see the accompanying publication ["Scalable Gaussian Process Inference with Stan"](https://arxiv.org/abs/2301.08836) for background on the approach). Here is the model definition in Stan. + +```{r, results='markup', comment='', echo=FALSE} +cat(readLines("getting_started.stan"), sep = "\n") +``` + +In this vignette, we compile and later fit the model using `cmdstanr`. + +```{r} +model <- cmdstanr::cmdstan_model( + stan_file="getting_started.stan", + include_paths=gptoolsStan::gptools_include_path(), +) +``` + +As indicated in the `data` section of the Stan program, we need to define the number of elements `n` of the Gaussian process and the [real fast Fourier transform](https://en.wikipedia.org/wiki/Fast_Fourier_transform#FFT_algorithms_specialized_for_real_or_symmetric_data) (RFFT) of the covariance kernel `cov_rfft`. We'll use 100 elements and a [squared exponential covariance kernel](https://en.wikipedia.org/wiki/Gaussian_process#Usual_covariance_functions) which allows us to evaluate the RFFT directly. + +```{r} +n <- 100 +length_scale <- n / 10 +freq <- 1:(n %/% 2 + 1) - 1 +# See appendix B of https://arxiv.org/abs/2301.08836 for details on the expression. +cov_rfft <- exp(- 2 * (pi * freq * length_scale / n) ^ 2) + 1e-9 +print(freq) +cov_rfft +``` + +Let's obtain prior realization by sampling from the model. + +```{r} +fit <- model$sample( + data=list(n=n, cov_rfft=cov_rfft), + seed=123, + chains=1, + iter_warmup=1000, + iter_sampling=5, +) +``` + +We extract the draws from the `fit` object and plot a realization of the process. + +```{r, fig.width=6, fig.height=5} +f <- fit$draws("f", format="draws_matrix") +plot(1:n, f[1,], type="l", xlab="covariate x", ylab="Gaussian process f(x)") +``` + +This vignette illustrates the use of gptools with an elementary example. Further details can be found in the [more comprehensive documentation](http://gptools-stan.readthedocs.io/) although using the [`cmdstanpy`](https://mc-stan.org/cmdstanpy/) interface. diff --git a/R-package/vignettes/getting_started.stan b/R-package/vignettes/getting_started.stan new file mode 100644 index 0000000..71d0aa3 --- /dev/null +++ b/R-package/vignettes/getting_started.stan @@ -0,0 +1,23 @@ +functions { + // Include utility functions, such as real fast Fourier transforms. + #include gptools/util.stan + // Include functions to evaluate GP likelihoods with Fourier methods. + #include gptools/fft.stan +} + +data { + // The number of sample points. + int n; + // Real fast Fourier transform of the covariance kernel. + vector[n %/% 2 + 1] cov_rfft; +} + +parameters { + // GP value at the `n` sampling points. + vector[n] f; +} + +model { + // Sampling statement to indicate that `f` is a GP. + f ~ gp_rfft(zeros_vector(n), cov_rfft); +}