Attempt at coding the logbessel K according to https://github.com/tk2lab/logbesselk

logbesselk.R

library(pracma) # for hyperbolic functions
library(Rcpp)
library(rethinking) # log_sum_exp
cppFunction('NumericVector clip( NumericVector x, double a, double b){
    return clamp( a, x, b ) ;
}')


log_cosh <- function(x) {
    return( x + log1p(expm1(-2 * x) / 2) )
}

g <- function(v, x, t) {
  return( log_cosh(v * t) - x * cosh(t) )
}

gp <- function(v, x, t) {
  return( v * tanh(v * t) - x * sinh(t) )
}

gpp <- function(v, x, t) {
  return( v^2 * sech(v * t)^2 - x * cosh(t) )
}

gppp <- function(v, x, t) {
  return( -v^3 * sech(v * t)^2  * tanh(v * t) - x * sinh(t))
}

find_range <- function(gp_func, v, x, t) {
  if (gp_func(v, x, t) < 0) stop("gp_func must be greater than or equal to 0")

  m <- 0
  t0 <- t
  while (gp_func(v, x, t0 + 2^(m + 1)) >= 0 ) {
    m <- m + 1
  }
  return(c(t0 + 2^m, t0 + 2^(m + 1)))
}

find_zero <- function(gp_func, gpp_func, v, x, te, ts, tol = 1, max_iter = 1000) {
  # if(gp_func(v, x, te) <= 0 || gp_func(v, x, ts) >= 0)
  #   stop("gp_func(v, te) must be greater than 0 and 
  #         gp_func(v, ts) must be less than 0")
  # 
  t0p1 <- te
  t0 <- te
  t1 <- ts
  i <- 0
  for (i in 1:max_iter) {
    t_shrink <- t0 + 0.5 * (t1 - t0)
    # dx <- (-te / gpp_func(v, x, t0)) * (t1 - t0)
    t_newton <- clip(t0 - gp_func(v, x, t0) / gpp_func(v, x, t0), te, t_shrink)
    
    if(gp_func(v, x, t_shrink) < 0) {
      t0p1 <- t_shrink
    } else if (gp_func(v, x, t_newton) < 0){
      t0p1 <- t_newton
      t1 <- t_shrink
    } else {
      t0p1 <- t0
      t1 <- t_newton
    }
    if ( abs(t0 - t0p1) < tol ) return(t0p1)
      t0 <- t0p1
  }
  return(t0)
}

#
# I have to increase n here a lot to get something 
# resembling the answer compared to what's in TF 
# default in TF is n = 128
#
integ_logbesselk <- function(v, x, tp, t0, t1, n = 5000) {
  h <- (t1 - t0) / n
  c0 <- cn <- 0.5
  
  y <- 0
  
  for (i in 0:n) {
    tm <- t0 + i * h
    if (i == 0 || i == n) y <- y +  h * exp(0.5 * g(v, x, tm) - g(v, x, tp))
     else  y <- y + h * exp(g(v, x, tm) - g(v, x, tp)) 
  }
  
   return(g(v, x, tp) + log(y))
}

log_besselk <- function(v, x, dt0 = 0.1, log_e = .Machine$double.min.exp) {
  tp <- 0
  gpp_out <- gpp(v, x, tp)
  
  if(x == 0) return(Inf)
  if(x < 0) return(NaN)
  if(v < 0) v <- -v
  
  shape <- v * x

  
  if (gpp_out > 0) {
    tste <- find_range(gp, v, x, 0)
    tp <- find_zero(gp, gpp, v, x, tste[1], tste[2])
  }
  t0 <- 0
  
  if(g(v, x, 0) - g(v, x, tp) < 0) {
    t0 <- find_zero(function(v, x, t, p = tp, le = log_e) g(v, x, t) -  g(v, x, p) - le
      , gp, v = v, x, 0, tp)
  }
  tste <- find_range(
    function(v, x, t, p = tp, le = log_e) {
    g(v, x, t) - g(v, x, p) - le 
      }, 
    v,
    x, tp)
  
  t1 <- find_zero(function(v, x, t, p = tp, le = log_e) g(v, x, t) -  g(v, x, p) - le,
                  gp,
                  v, x,
                  tste[1], tste[2])
  return(integ_logbesselk(v, x, tp, t0, t1))
 #  return (c(t0, t1))
}
v <- 2.5
x <- 5

log_besselk(v, x)
log(besselK(x, v))
# log(x) - log(2 * v) + log( besselK(x, v + 1) - besselK(x, abs(v - 1)) )