Clever way of writing a loop to calculate Jensen–Shannon divergence

Question

I have tried to use apply or aggregate functions to calculate the Jensen–Shannon divergence (JS.dist) between some distributions.

I am simulating some data under four different models, and for each data I calculate a series of statistics.

So imagine I have the following data.frame:

dataframe1:
Model Factor1 Factor2 stats1 stats2
M1    0.0001  0.2     -1.0   0.9
M1    0.0001  0.2     -1.3   0.5
M1    0.0002  0.3     -1.9   0.2
M2    0.0001  0.2     -2.0   0.2
M2    0.0001  0.2     -2.0   0.2
M2    0.0002  0.3     -2.1   0.4
M3    0.0001  0.2      9.9   0.4
M3    0.0001  0.2      8.3   0.4
M3    0.0002  0.3      8.0   0.4
M4    0.0001  0.2      3.0   0.1
M4    0.0001  0.2      3.5   0.3
M4    0.0002  0.3      3.2   0.3

The funtion to calculate the JS.dist is as follow:

Function to change Inf or -Inf to zero in a log. It takes as argument the log of a number

 test.logs <- function(log.num){

  log.num[log.num == -Inf | log.num == Inf] <- 0
  return (log.num)

}

# function to calculate the kl.dist (Kullback–Leibler divergence). It takes as arguments the probability vectors (see below) for two distributions (x.p and y.p)

kl.dist <- function(x.p, y.p) {
  # x.p, y.p: probability vectors for x and y distributions

  log.x <- test.logs(log(x.p))
  log.y <- test.logs(log(y.p))

  sum(x.p * (log.x - log.y))

}

# function to calculate the js.dist. It takes as arguments the probability vectors for x, y and M. M is the middle distribution

js.dist <- function(x.p, y.p, M.p){
  0.5 * kl.dist(x.p, M.p) + .5 * kl.dist(y.p, M.p)
}

To use the above functions I have to calculate the density curves of my distributions (for stats 1 and stats2 by Model and factors).

To calculate this I have to set a minmum and maximum value that the density curves will be calculate for each stats.

For example:

x.d <- density(x, n=512, from=min, to=max)
y.d <- density(y, n=512, from=min, to=max)
M.d <- (x.d$y + y.d$y)/2

# width of the histogram
w <- x.d$x[2] - x.d$x[1]

# probability of x value in n-th bin
x.p <- x.d$y * w # (hist hight) * (bin width)
y.p <- y.d$y * w
M.p <- M.d * w

I tried to write a code in which I have two for loops (for each of the factors) and I subset my data by model and I calculate the min and max for each of my stats. And finally I calculate the densities and probabilities and only after I can calculate the JS.dist.

As for example of R code:

density_js.dist <- function(data.df){
# gets the unique values for mutation rate
factor1 <- unique(data.df$Factor1)
# gets the unique values for rate of new copies
factor2 <- unique(data.df$factor2)

# calculates the minimum and maximum value for each of the statistics
# showing only for stats1
stats1.min <- min(data.df$stats1)
stats1.max <- max(data.df$stats1)



# for loop to calculate the densities and probabilities and JS distance for each combination of factor1 and factor2

for (f1 in factor1){
  for (f2 in factor2){

  new.df <- subset(data.df, factor1 == f1 & factor2 == f2)

  # subsetting data. One data frame for each of the four models
  MM.df <- subset(new.df, Model == "M1")
  TM.df <- subset(new.df, Model == "M2")

  MI.df <- subset(new.df, Model == "M3")
  TI.df <- subset(new.df, Model == "M4")

  # densitiy and probability for each stats

  #1.stats1
  # calculating densities for model M1 and M2
  MM1.d <- density(MM.df$stats1, n=512, from=stats1.min, to=stats1.max)
  TM1.d <- density(TM.df$stats1, n=512, from=stats1.min, to=stats1.max)

  # Density for the middle distribution between models M1 and M2 
  Middle12.d <- (MM1.d$y + TM1.d$y)/2

  # width for models
  w12 <- MM1.d$x[2] - MM1.d$x[1]

  # calculating probabilities for each models
  MM1.p <- MM1.d$y * w12 # (hist hight) * (bin width)
  TM1.p <- TM1.d$y * w12
  Middle12.p <- Middle12.d * w12 

  # calculating densities for models M3 and M4
  MI1.d <- density(MI.df$stats1, n=512, from=stats1.min, to=stats1.max)
  TI1.d <- density(TI.df$stats1, n=512, from=stats1.min, to=stats1.max)
  Middle34.d <- (MI1.d$y + TI1.d$y)/2

  w34 <- MI1.d$x[2] - MI1.d$x[1]

  # calculating probabilities for M3 and M4 models
  MI1.p <- MM1.d$y * w34 
  TI1.p <- TM1.d$y * w34
  Middle34.p <- Middle34.d * w34 


 js.dist(MM1.p, TM1.p, Middle12.p)
 js.dist(MI1.p, TI1.p, Middle34.p)
  }
 }
}

My question is:

I have tried to use apply or aggregate, however I can't figure out how to pass as an argument the min and max for each statistics to be able to create the density curves? Note that this min and max is calculate over all combinations of factors and models rather than for each subset. For example, in order to allow comparison, I cannot calculate the min and max for the subset by factors and models.

My data is actulally much more complex. I have 10 different statistics that I want to calculate the JS.dist between two distributions by factors. My two distributions would be M1 with M2, and M3 with M4. The code above works, but it requires me to write more the 700 lines which I really don't think is very clever.

I would really appreciate if someone could please help me with this.

Answer 1

Here is a hack-ish way of doing the calculation for all 10 series at once using lists. Due to the length and verboseness of the code, it would require a complete rewrite if you want a one function solution. Was only able to test the code on the first two series (not even fully since more than one factor1:factor2 combination had only 1 observation, making density calculation impossible). Also removed the function since it does absolutely nothing for you.

library(dplyr)
L = list()

  # gets the unique values for mutation rate
  factor1 <- unique(data.df$Factor1)
  # gets the unique values for rate of new copies
  factor2 <- unique(data.df$Factor2)

  # calculates the minimum and maximum value for each of the statistics
  # Store all 10 min and max in a vector
  vector.min <- lapply(data.df %>% select(stats1:stats10), min)
  vector.max <- lapply(data.df %>% select(stats1:stats10), max)

  # for loop to calculate the densities and probabilities and JS distance for each combination of factor1 and factor2

  for (f1 in factor1){
    for (f2 in factor2){
      new.df <- subset(data.df, factor1 == f1 & factor2 == f2)    
      # subsetting data. One data frame for each of the four models
      MM.df <- subset(new.df, Model == "M1")
      TM.df <- subset(new.df, Model == "M2")
      MI.df <- subset(new.df, Model == "M3")
      TI.df <- subset(new.df, Model == "M4")

      # densitiy and probability for each stats

      # calculating densities for model M1 and M2
      MM.d = lapply(1:10, function(i) density(MM.df %>% select(i+3) %>% unlist, n = 512, from = vector.min[[i]], to = vector.min[[i]]))
      TM.d = lapply(1:10, function(i) density(TM.df %>% select(i+3) %>% unlist, n = 512, from = vector.min[[i]], to = vector.min[[i]]))

      # Density for the middle distribution between models M1 and M2 
      Middle12.d <- mapply(function(d1, d2) (d1$y+d2$y)/2, MM.d, TM.d, SIMPLIFY = F)

      # width for models
      w12 = lapply(MM.d, function(y) {y$x[2] - y$x[1]})

      # calculating probabilities for each models
      MM1.p = mapply(function(arg1, arg2) {arg1$y * arg2}, MM.d, w12)  # (hist hight) * (bin width)
      TM1.p = mapply(function(arg1, arg2) {arg1$y * arg2}, TM.d, w12)
      Middle12.p = mapply("*", Middle12.d, w12)

      # calculating densities for models M3 and M4
      MI.d = lapply(1:10, function(i) density(MI.df %>% select(i+3) %>% unlist, n = 512, from = vector.min[[i]], to = vector.min[[2]]))
      TI.d = lapply(1:10, function(i) density(TI.df %>% select(i+3) %>% unlist, n = 512, from = vector.min[[i]], to = vector.min[[2]]))
      Middle34.d <- mapply(function(d1, d2) (d1$y+d2$y)/2, MI.d, TI.d)

      w34 = lapply(MI.d, function(y) {y$x[2] - y$x[1]})      

      # calculating probabilities for M3 and M4 models
      MI1.p = mapply(function(arg1, arg2) {arg1$y * arg2}, MI.d, w34)  # (hist hight) * (bin width)
      TI1.p = mapply(function(arg1, arg2) {arg1$y * arg2}, TI.d, w34)
      Middle34.p = mapply("*", Middle34.d, w34)

      L = c(L, list(mapply(js.dist, MM1.p, TM1.p, Middle12.p), mapply(js.dist, MI1.p, TI1.p, Middle34.p)))
    }
  }