R/RMHMC_sim_proposal_sample.R

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

#This script does a Riemanniean Manifold Hamiltonian MCMC leapfrog proposal for the LOLOG model:
#simply updates the covariance matrix at each step compared to the regular HMC algorithm
#uses simulations to estimate the Fisher information matrix, rather than just using the observed information
#updates the momentum matrix at each step when the positions is updated
#note that this means the HMC dynamics are wrong - so should only be used for a small number of steps e.g. 1 
#Does not do fixed point iteration.


#' @export

RMHMC_sim_proposal_sample <- 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,net,change_on){(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, #sigma used in momentum function generator
                               change_off = NULL,
                               change_on = NULL,
                               proposal_sims = 40,#number of sims used for the calculation of Hessian
                               ...
){
  #Rename the theta_0 as q - in line with HMC literature
  q <- theta_0
  names(q) <- NULL
  #derivative of change statistics on top
  top_deriv <- lolog::calculateStatistics(as.formula(paste("net ~ ",formula_rhs))) #Used in q_grad
  e <- length(change_on)
  samples <- 100
  
  #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)
  }
  
  #From s get the permutation heads and tails:
  n <- network.size(net)
  edges <- combn(seq(1,n),2)
  tails <- edges[1,][s]-1
  heads <- edges[2,][s]-1
  
  #===== Sample the change stats probabilistically
  #Calculate the z-score for each change stat:
  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)
  change_on_sample <- sample(change_on,samples,prob = weights)
  
  #Calculate gradient of the log prior + log liklihood function for LOLOG
  q_grad <- function(q,change_on=NULL,network=NULL){
    if(!is.null(network)){
      network <- net}
    if(is.null(change_on)){
      formula <- as.formula(paste("net ~",formula_rhs,sep = ""))
      change_on <- lolog_change_stats(network,s,formula_rhs)
    }
    
    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)*(e/samples)
    
    # print("Prior_deriv")
    # print(prior_deriv)
    # print("Top Deriv")
    # print(top_deriv)
    # print("bottom deriv")
    # print(bottom_deriv)
    
    #return their sum with the correct signs
    return(-prior_deriv - top_deriv  + bottom_deriv)
  }
  
  #specify a local momentum matrix if no momentum is supplied:
  #it is the local covariance matrix at the starting point
  #"ideal" dynamics under the assumption that the proposal distribution is Gaussian are sampling the momentum from the inverse covariance matrix
  momentum_assign <- function(q){
    t<-proc.time()
    model <- lolog::createLatentOrderLikelihood(formula,theta = q)
    momentum_sigma <- lapply(1:proposal_sims,function(i){
      tmp = model$generateNetworkWithEdgeOrder(heads,tails)
      change_on <- tmp$changeStats
      #===== Sample the change stats probabilistically
      #Calculate the z-score for each change stat:
      
      #use weights from first sample to prevent having to calculate - not great!!!:
      change_on_sample <- sample(change_on,samples,prob = weights)
      
      grad = tryCatch({q_grad(q,change_on = change_on_sample, network = lolog::as.network(tmp$network))},error = function(e){return(rep(0,length(q)))})
      #grad = q_grad(q,change_on = tmp$changeStats, network = lolog::as.network(tmp$network))
      return(outer(grad,grad))
    })
    momentum_sigma <- Reduce("+",momentum_sigma)/length(momentum_sigma)
    #check if singular, if not add to singular values
    if(abs(det(momentum_sigma)) < 10**(-6)){
      for(i in 10**seq(-6,20,1)){
        if(abs(det(momentum_sigma))>10**(-6)){
          break
        }
        else{
          tmp  = svd(momentum_sigma)
          values = tmp$d
          values[which(values == min(values))] <- values[which(values == min(values))] + i
          momentum_sigma <- tmp$u%*%diag(values)%*%t(tmp$v)
          #momentum_sigma <- momentum_sigma + diag(min(diag(momentum_sigma))*i,dim(momentum_sigma)[1])
        }
      }
    }
    momentum_sigma_inv <- solve(momentum_sigma)
    assign("momentum_sigma",momentum_sigma, envir = parent.env(environment()))
    assign("momentum_sigma_inv",momentum_sigma_inv, envir = parent.env(environment()))
    # print("Momentum is :")
    # print(momentum_sigma)
    # print("Momentum inversed is:")
    # print(momentum_sigma_inv)
    # print("Momentum Assigned")
    
    print(paste("assigning the momentum took", (proc.time() - t)[3], "seconds",sep=""))
    return()
  }
  #debug(momentum_assign)
  momentum_assign(q)
  
  
  #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
  }
  

  
  #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)
  
  #Do L_steps full updates
  if(L_steps != 1){
    for(i in 1:(L_steps-1)){
      q <- q + epsilon*(momentum_sigma_inv%*%p)
      
      if(prior(q)==0){
        return(list(proposal = as.vector(q), prob_factor = 0))
      }
      p <- p - epsilon*q_grad(q,change_on = change_on_sample)
      
      
      #print(q)
      #print(p)
      #print(momentum_sigma)
      #print(det(momentum_sigma))
      #print(momentum_sigma_inv)
      #print(det(momentum_sigma_inv))
      
      
      #Update covariance matrix:
      momentum_assign(q)
    }
  }
  #Do final updates
  q <- q + epsilon*(momentum_sigma_inv%*%p)
  p <- p - (epsilon/2)*q_grad(q,change_on = change_on_sample)
  
  #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_init)%*%(momentum_sigma_inv%*%p_init) - 0.5*t(p)%*%(momentum_sigma_inv%*%p))
  
  return(list(proposal = as.vector(q), prob_factor = prob_factor))
}