R/HMC_proposal.R

In duncan-clark/Blolog: Bayesian Inference for LOLOG Network Models

#This script does a Hamiltonian MCMC leapfrog proposal with sampling for gradient for the LOLOG model:

#' @export
#' 
HMC_proposal <- function(theta_0, #the current paramter value - starting position for HMC
                         net, #the observed network,
                         s, #the fixed edge permutation
                         formula_rhs,    #rhs formula of model
                         prior = function(theta){(sum(abs(theta)<10) == length(theta))*((1/20)**length(theta))}, #prior function for theta
                         prior_grad = NULL, #specify prior derivative to get speed up
                         L_steps, #the number of leapfrog steps
                         epsilon = NULL, #size of the leapfrog step
                         epsilon_factor = 1, #used when epsilon is null to amend the generate "ideal epsilon" 
                         momentum_sigma = NULL, #sigma used in momentum function generator
                         change_off = NULL,
                         change_on = NULL,
                         HMC_grad_samples = NULL,
                         ...
){
  #Rename the theta_0 as q - in line with HMC literature
  q <- theta_0
  names(q) <- NULL
  top_deriv <- lolog::calculateStatistics(as.formula(paste("net ~ ",formula_rhs))) #used inside q_grad function
  e <- length(change_on)
  
  #Make the lolog formula
  formula <- as.formula(paste("net ~",formula_rhs,sep = ""))
  
  #calculate the change stats for the given permutaion and graph
  if(is.null(change_on)){
    change_on <- lolog_change_stats(net,s,formula_rhs)
  }
  
  #Calculate gradient of the log liklihood function for LOLOG
  q_grad <- function(q,change_on=NULL,network=NULL,sample_weights){
    if(is.null(prior_grad)){
      prior_deriv <- (numDeriv::grad(prior,q,net=net,change_on = change_on)/prior(q,net=net,change_on = change_on))
    }
    else{prior_deriv <- prior_grad(q,net=net,change_on = change_on)/prior(q,net=net,change_on = change_on)}
    
    if(sum(is.na(prior_deriv)) != 0){
      prior_deriv <- rep(0,length(q))
    }
    
    bottom_deriv <- bottom_deriv_helper_cpp(change_on,q,weights = sample_weights)
    
    #return their sum with the correct signs
    return(-prior_deriv - top_deriv  + bottom_deriv)
  }
  
  #If epsilon is null put it as the mimimum eigen value of the momentum matrix:
  if(is.null(epsilon)){
    epsilon <- (min(eigen(momentum_sigma)$values)**0.5)*epsilon_factor
  }
  
  #If momentum is not invertible, make it invertible and define the inverse matrix:r
  if(abs(det(momentum_sigma)) < 10**(-6)){
    for(i in 10**seq(-6,20,1)){
      if(abs(det(momentum_sigma))>10**(-6)){
        break
      }
      else{
        momentum_sigma <- momentum_sigma + diag(min(diag(momentum_sigma))*i,dim(momentum_sigma)[1])
      }
    }
  }
  momentum_sigma_inv <- solve(momentum_sigma)
  
  #===== Sample the change stats probabilistically
  #Calculate the z-score for each change stat:
  if(!is.null(HMC_grad_samples)){
    tmp <- do.call(rbind,change_on)
    tmp <- scale(tmp)
    weights <- abs(rowSums(tmp[,2:dim(tmp)[2]]))
    rm(list=c("tmp"))
    weights <- weights/sum(weights)
    if(HMC_grad_samples < 1){
      HMC_grad_samples <- HMC_grad_samples*e
    }
    samp <- sample(1:e,ceiling(HMC_grad_samples),prob = weights)
    sample_weights <- weights[samp]
    change_on_sample <- change_on[samp]
  }else{
    change_on_sample <- change_on
    sample_weights <- NULL
  }

  
  #Generate initial momentum
  p_init <- MASS::mvrnorm(1,rep(0,length(q)),Sigma = momentum_sigma)
  p <- p_init
  
  #Do half momentum update fist
  p <- p - (epsilon/2)*q_grad(q,change_on = change_on_sample,sample_weights = sample_weights)
  
  #Do L_steps full updates
  if(L_steps != 1){
    for(i in 1:(L_steps-1)){
      q <- q + epsilon*(momentum_sigma_inv%*%p)
      p <- p - epsilon*q_grad(q,change_on = change_on_sample,sample_weights = sample_weights)
    }
  }
  #Do final updates
  p_old <- p #used for prob factor
  q <- q + epsilon*(momentum_sigma_inv%*%p)
  p <- p - (epsilon/2)*q_grad(q,change_on = change_on_sample,sample_weights = sample_weights)
  
  #negate the momentum variable - does not actaully effect anything
  p <- -p
  
  #prob factor takes account of the joint distribution in the metropolis step
  prob_factor = exp(0.5*t(p_old)%*%momentum_sigma_inv%*%p_old - 0.5*t(p)%*%momentum_sigma_inv%*%p)
  if(is.nan(prob_factor)){prob_factor = 0}
  
  #print(paste("prob factor was : ",prob_factor))
  
  return(list(proposal = as.vector(q), prob_factor = prob_factor))
}