Feature: Hessian-vector product function

R/R/bridgestan.R

@@ -23,7 +23,7 @@ StanModel <- R6::R6Class("StanModel",
       private$seed <- seed
       private$lib <- tools::file_path_as_absolute(lib)
       private$lib_name <- tools::file_path_sans_ext(basename(lib))
-      if (is.loaded("construct_R", PACKAGE = private$lib_name)) {
+      if (is.loaded("bs_model_construct_R", PACKAGE = private$lib_name)) {
         warning(
           paste0("Loading a shared object '", lib, "' which is already loaded.\n",
                   "If the file has changed since the last time it was loaded, this load may not update the library!"
@@ -289,6 +289,32 @@ StanModel <- R6::R6Class("StanModel",
         stop(handle_error(private$lib_name, vars$err_msg, vars$err_ptr, "log_density_hessian"))
       }
       list(val = vars$val, gradient = vars$gradient, hessian = matrix(vars$hess, nrow = dims, byrow = TRUE))
+    },
+    #' @description
+    #' Return the log density and the product of the Hessian
+    #' with the specified vector.
+    #' @param theta_unc The vector of unconstrained parameters.
+    #' @param v The vector to multiply the Hessian by.
+    #' @param propto If `TRUE`, drop terms which do not depend on the parameters.
+    #' @param jacobian If `TRUE`, include change of variables terms for constrained parameters.
+    #' @return List containing entries `val` (the log density) and `Hvp` (the hessian-vector product).
+    log_density_hessian_vector_product = function(theta_unc, v, propto = TRUE, jacobian = TRUE){
+      dims <- self$param_unc_num()
+      vars <- .C("bs_log_density_hessian_vector_product_R",
+        as.raw(private$model), as.logical(propto), as.logical(jacobian),
+        as.double(theta_unc),
+        as.double(v),
+        val = double(1), Hvp = double(dims),
+        return_code = as.integer(0),
+        err_msg = as.character(""),
+        err_ptr = raw(8),
+        NAOK = TRUE,
+        PACKAGE = private$lib_name
+      )
+      if (vars$return_code) {
+        stop(handle_error(private$lib_name, vars$err_msg, vars$err_ptr, "log_density_hessian_vector_product"))
+      }
+      list(val = vars$val, Hvp = vars$Hvp)
     }
   ),
   private = list(

R/tests/testthat/test-bridgestan.R

@@ -55,6 +55,11 @@ test_that("simple_model Hessian is -I",{
     expect_equal(-diag(5), simple$log_density_hessian(x)$hessian)
 })
 
+test_that("simple_model Hvp returns -v",{
+    x <- runif(5)
+    v <- runif(5)
+    expect_equal(-v, simple$log_density_hessian_vector_product(x, v)$Hvp)
+})
 
 bernoulli <- load_model("bernoulli")
 

docs/languages/julia.md

@@ -190,6 +190,24 @@ This allocates new memory for the gradient and Hessian output each call. See `lo
 
 <a target='_blank' href='https://github.com/roualdes/bridgestan/blob/main/julia/src/model.jl#L662-L673' class='documenter-source'>source</a><br>
 
+<a id='BridgeStan.log_density_hessian_vector_product' href='#BridgeStan.log_density_hessian_vector_product'>#</a>
+**`BridgeStan.log_density_hessian_vector_product`** &mdash; *Function*.
+
+
+
+```julia
+log_density_hessian_vector_product(sm, q, v; propto=true, jacobian=true)
+```
+
+Returns log density and the product of the Hessian of the log density with the vector `v` at the specified unconstrained parameters.
+
+This calculation drops constant terms that do not depend on the parameters if `propto` is `true` and includes change of variables terms for constrained parameters if `jacobian` is `true`.
+
+This allocates new memory for the output each call. See `log_density_hessian_vector_product!` for a version which allows re-using existing memory.
+
+
+<a target='_blank' href='https://github.com/roualdes/bridgestan/blob/main/julia/src/model.jl#L731-L742' class='documenter-source'>source</a><br>
+
 <a id='BridgeStan.param_constrain' href='#BridgeStan.param_constrain'>#</a>
 **`BridgeStan.param_constrain`** &mdash; *Function*.
 
@@ -380,6 +398,24 @@ The gradient is stored in the vector `out_grad` and the Hessian is stored in `ou
 
 <a target='_blank' href='https://github.com/roualdes/bridgestan/blob/main/julia/src/model.jl#L598-L609' class='documenter-source'>source</a><br>
 
+<a id='BridgeStan.log_density_hessian_vector_product!' href='#BridgeStan.log_density_hessian_vector_product!'>#</a>
+**`BridgeStan.log_density_hessian_vector_product!`** &mdash; *Function*.
+
+
+
+```julia
+log_density_hessian_vector_product!(sm, q, v, out; propto=true, jacobian=true)
+```
+
+Returns log density and the product of the Hessian of the log density with the vector `v` at the specified unconstrained parameters.
+
+This calculation drops constant terms that do not depend on the parameters if `propto` is `true` and includes change of variables terms for constrained parameters if `jacobian` is `true`.
+
+The product is stored in the vector `out` and a reference is returned. See `log_density_hessian_vector_product` for a version which allocates fresh memory.
+
+
+<a target='_blank' href='https://github.com/roualdes/bridgestan/blob/main/julia/src/model.jl#L673-L684' class='documenter-source'>source</a><br>
+
 <a id='BridgeStan.param_constrain!' href='#BridgeStan.param_constrain!'>#</a>
 **`BridgeStan.param_constrain!`** &mdash; *Function*.
 

docs/languages/r.md

@@ -401,3 +401,33 @@ _Returns_
   List containing entries `val` (the log density), `gradient`
   (the gradient), and `hessian` (the Hessian).
 
+
+**Method** `log_density_hessian_vector_product()`:
+
+Return the log density and the product of the Hessian
+with the specified vector.
+
+_Usage_
+
+```R
+StanModel$log_density_hessian_vector_product(theta_unc, v, propto = TRUE, jacobian = TRUE)
+```
+
+
+_Arguments_
+
+  - `theta_unc` The vector of unconstrained parameters.
+
+  - `v` The vector to multiply the Hessian by.
+
+  - `propto` If `TRUE`, drop terms which do not depend on the
+      parameters.
+
+  - `jacobian` If `TRUE`, include change of variables terms for
+      constrained parameters.
+
+
+_Returns_
+
+  List containing entries `val` (the log density) and `Hvp`
+  (the hessian-vector product).

julia/docs/src/julia.md

@@ -75,6 +75,7 @@ StanModel
 log_density
 log_density_gradient
 log_density_hessian
+log_density_hessian_vector_product
 param_constrain
 param_unconstrain
 param_unconstrain_json
@@ -86,6 +87,7 @@ param_names
 param_unc_names
 log_density_gradient!
 log_density_hessian!
+log_density_hessian_vector_product!
 param_constrain!
 param_unconstrain!
 param_unconstrain_json!

julia/src/BridgeStan.jl

@@ -14,12 +14,14 @@ export StanModel,
     param_unconstrain_json!,
     log_density_gradient!,
     log_density_hessian!,
+    log_density_hessian_vector_product!,
     param_constrain,
     param_unconstrain,
     param_unconstrain_json,
     log_density,
     log_density_gradient,
     log_density_hessian,
+    log_density_hessian_vector_product,
     get_bridgestan_path,
     set_bridgestan_path!,
     compile_model,
@@ -49,10 +51,13 @@ function __init__()
     # On Windows, we may need to add TBB to %PATH%
     if Sys.iswindows()
         try
-            run(pipeline(`where.exe tbb.dll`, stdout=devnull, stderr=devnull))
+            run(pipeline(`where.exe tbb.dll`, stdout = devnull, stderr = devnull))
         catch
             # add TBB to %PATH%
-            ENV["PATH"] = joinpath(get_bridgestan_path(), "stan", "lib", "stan_math", "lib", "tbb") * ";" * ENV["PATH"]
+            ENV["PATH"] =
+                joinpath(get_bridgestan_path(), "stan", "lib", "stan_math", "lib", "tbb") *
+                ";" *
+                ENV["PATH"]
         end
     end
 end

julia/src/model.jl

@@ -683,6 +683,92 @@ function log_density_hessian(
     log_density_hessian!(sm, q, grad, hess; propto = propto, jacobian = jacobian)
 end
 
+"""
+    log_density_hessian_vector_product!(sm, q, v, out; propto=true, jacobian=true)
+
+Returns log density and the product of the Hessian of the log density with the vector `v`
+at the specified unconstrained parameters.
+
+This calculation drops constant terms that do not depend on the parameters if `propto` is `true`
+and includes change of variables terms for constrained parameters if `jacobian` is `true`.
+
+The product is stored in the vector `out` and a reference is returned. See
+`log_density_hessian_vector_product` for a version which allocates fresh memory.
+"""
+function log_density_hessian_vector_product!(
+    sm::StanModel,
+    q::Vector{Float64},
+    v::Vector{Float64},
+    out::Vector{Float64};
+    propto = true,
+    jacobian = true,
+)
+    dims = param_unc_num(sm)
+    if length(out) != dims
+        throw(
+            DimensionMismatch(
+                "out must be same size as number of unconstrained parameters",
+            ),
+        )
+    end
+    lp = Ref(0.0)
+    err = Ref{Cstring}()
+    rc = ccall(
+        Libc.Libdl.dlsym(sm.lib, "bs_log_density_hessian_vector_product"),
+        Cint,
+        (
+            Ptr{StanModelStruct},
+            Cint,
+            Cint,
+            Ref{Cdouble},
+            Ref{Cdouble},
+            Ref{Cdouble},
+            Ref{Cdouble},
+            Ref{Cstring},
+        ),
+        sm.stanmodel,
+        propto,
+        jacobian,
+        q,
+        v,
+        lp,
+        out,
+        err,
+    )
+    if rc != 0
+        error(handle_error(sm.lib, err, "log_density_hessian_vector_product"))
+    end
+    (lp[], out)
+end
+
+"""
+    log_density_hessian_vector_product(sm, q, v; propto=true, jacobian=true)
+
+Returns log density and the product of the Hessian of the log density with the vector `v`
+at the specified unconstrained parameters.
+
+This calculation drops constant terms that do not depend on the parameters if `propto` is `true`
+and includes change of variables terms for constrained parameters if `jacobian` is `true`.
+
+This allocates new memory for the output each call. See
+`log_density_hessian_vector_product!` for a version which allows re-using existing memory.
+"""
+function log_density_hessian_vector_product(
+    sm::StanModel,
+    q::Vector{Float64},
+    v::Vector{Float64};
+    propto = true,
+    jacobian = true,
+)
+    out = zeros(param_unc_num(sm))
+    log_density_hessian_vector_product!(sm, q, v, out; propto = propto, jacobian = jacobian)
+end
+
+"""
+    log_density_hessian_vector_product(sm, q, v; propto=true, jacobian=true)
+"""
+
+
 """
     handle_error(lib::Ptr{Nothing}, err::Ref{Cstring}, method::String)
 

julia/test/model_tests.jl

@@ -526,7 +526,7 @@ end
     ld = Vector{Bool}(undef, R)
     g = Vector{Bool}(undef, R)
 
-    @Threads.threads for it = 1:nt
+    Threads.@threads for it = 1:nt
         for r = it:nt:R
             x = randn(BridgeStan.param_num(model))
             (lp, grad) = BridgeStan.log_density_gradient(model, x)
@@ -550,24 +550,24 @@ end
     x = [0.5] # bernoulli parameter
 
     R = 1000
-    out_size = BridgeStan.param_num(model;include_tp=false, include_gq=true)
+    out_size = BridgeStan.param_num(model; include_tp = false, include_gq = true)
 
     # to test the thread safety of our RNGs, we do two runs
     # the first we do in parallel
     gq1 = zeros(Float64, out_size, R)
-    @Threads.threads for it = 1:nt
-        rng = StanRNG(model,seeds[it]) # RNG is created per-thread
+    Threads.@threads for it = 1:nt
+        rng = StanRNG(model, seeds[it]) # RNG is created per-thread
         for r = it:nt:R
-            gq1[:, r] = BridgeStan.param_constrain(model, x; include_gq=true, rng=rng)
+            gq1[:, r] = BridgeStan.param_constrain(model, x; include_gq = true, rng = rng)
         end
     end
 
     # the second we do sequentially
     gq2 = zeros(Float64, out_size, R)
     for it = 1:nt
-        rng = StanRNG(model,seeds[it])
+        rng = StanRNG(model, seeds[it])
         for r = it:nt:R
-            gq2[:, r] = BridgeStan.param_constrain(model, x; include_gq=true, rng=rng)
+            gq2[:, r] = BridgeStan.param_constrain(model, x; include_gq = true, rng = rng)
         end
     end
 
@@ -651,6 +651,9 @@ end
     using LinearAlgebra
     @test isapprox(-Matrix(1.0I, D, D), hess)
 
+    v = rand(D)
+    lp, Hvp = BridgeStan.log_density_hessian_vector_product(model, y, v)
+    @test isapprox(-v, Hvp)
 end
 
 @testset "printing" begin