reproduction.Rmd

---
title: "R based Reproduction of Saxe et al. 2013"
output: html_notebook
---

```{r}
library(reshape2)
library(ggplot2)
library(tidyverse)
library(grid)
library(gridExtra)
library(png)
library(pbapply)

FONT_FAMILY <- "times new roman"
TITLE_FONT_SIZE <- 30
TICK_FONT_SIZE <- 22
```

# Reproduce figure 2

## figure 2 IO

```{r}
reproduce_figure_2_io <- function(d, gradient_colours) {
  
  # Set column and row names
  colnames(d) <- c("C", "S", "O", "R")
  df <- melt(as.matrix(d))
  row_names <- c("M", "F", "S", "B", "P")
  df$Var1 <- factor(df$Var1, levels = 1:length(row_names), labels = row_names)
  df$Var1 <- factor(df$Var1, levels = rev(levels(df$Var1)))
  
  ggplot(df, aes(Var2, Var1, fill = value)) + 
    geom_tile(colour = "black", size = 0.7) +
    gradient_colours +
    scale_x_discrete(position = "top") +
    theme_minimal() +
    theme(
      axis.title.x.top = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.title.x.bottom = element_blank(),
      axis.title.y.left = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.ticks.x = element_blank(),
      axis.ticks.y = element_blank(),
      axis.text.y = element_text(size = TICK_FONT_SIZE, family = FONT_FAMILY, face = "italic"),
      axis.text.x = element_text(size = TICK_FONT_SIZE, family = FONT_FAMILY, face = "italic"),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      legend.position = "none",
    ) +
    labs(x = "Items", y = "Properties", fill = "") +
    coord_fixed() +
    annotate("rect", xmin = 0.5, xmax = ncol(as.matrix(d)) + 0.5, ymin = 0.5, ymax = nrow(as.matrix(d)) + 0.5, 
             colour = "white", fill = NA, size = 1)
  
  ggsave("figures/fig2_io.png", height = 7, width = 5)
  
}
```

## figure 2 U

```{r}
reproduce_figure_2_u <- function(u, gradient_colours) {
  
  df <- melt(u)
  df <- df[order(df$Var1, decreasing=FALSE),]
  
  row_names <- c("M", "F", "S", "B", "P")
  df$Var1 <- factor(df$Var1, levels = 1:length(row_names), labels = row_names)
  df$Var1 <- factor(df$Var1, levels = rev(levels(df$Var1)))
  
  ggplot(df, aes(Var2, Var1, fill = value)) + 
    geom_tile(colour = "black", size = 0.7) +
    gradient_colours +
    scale_x_continuous(position = "top") +
    theme_minimal() +
    theme(
      axis.title.x.top = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.title.x.bottom = element_blank(),
      axis.title.y.left = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.ticks.x = element_blank(),
      axis.ticks.y = element_blank(),
      axis.text.y = element_text(size = TICK_FONT_SIZE, family = FONT_FAMILY, face = "italic"),
      axis.text.x = element_text(size = TICK_FONT_SIZE, family = FONT_FAMILY, face = "italic"),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      legend.position = "none",
    ) +
    labs(x = "Modes", y = "Properties", fill = "") +
    coord_fixed() +
    annotate("rect", xmin = 0.5, xmax = ncol(u) + 0.5, ymin = 0.5, ymax = nrow(u) + 0.5, 
             colour = "white", fill = NA, size = 1)
  
  ggsave("figures/fig2_u.png", height = 7, width = 6)
  
}

```

## figure 2 SV

```{r}
reproduce_figure_2_sv <- function(sv, gradient_colours) {
  
  df <- melt(sv)
  df <- df[order(df$Var1, decreasing=FALSE),]
  
  ggplot(df, aes(Var2, Var1, fill = value)) + 
    geom_tile(colour = "black", size = 0.7) +
    gradient_colours +
    scale_y_reverse() + 
    scale_x_continuous(position = "top") +
    theme_minimal() +
    theme(
      axis.title.x.top = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.title.x.bottom = element_blank(),
      axis.title.y.left = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.ticks.x = element_blank(),
      axis.ticks.y = element_blank(),
      axis.text.y = element_blank(),
      axis.text.x = element_text(size = TICK_FONT_SIZE, family = FONT_FAMILY, face = "italic"),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      legend.position = "none",
    ) +
    labs(x = "Modes", y = "Modes", fill = "") +
    coord_fixed() +
    annotate("rect", xmin = 0.5, xmax = ncol(sv) + 0.5, ymin = 0.5, ymax = nrow(sv) + 0.5, 
             colour = "white", fill = NA, size = 1)
  
  ggsave("figures/fig2_sv.png", height = 7, width = 4)
}

reproduce_figure_2_sv(sv, gradient_colours)

```

## figure 2 Vt

```{r}
reproduce_figure_2_vt <- function(vt, gradient_colours) {
  
  df <- melt(vt)
  df <- df[order(df$Var1, decreasing=FALSE),]
  
  # Set column names
  column_names <- c("C", "S", "O", "R")  
  df$Var2 <- factor(df$Var2, levels = 1:length(column_names), labels = column_names)
  
  ggplot(df, aes(Var2, Var1, fill = value)) + 
    geom_tile(colour = "black", size = 0.7) +
    gradient_colours +
    scale_y_reverse() +
    scale_x_discrete(position = "top") +
    theme_minimal() +
    theme(
      axis.title.x.top = element_text(vjust = 2, size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.title.x.bottom = element_blank(),
      axis.title.y.left = element_text(size = TITLE_FONT_SIZE, family = FONT_FAMILY),
      axis.ticks.x = element_blank(),
      axis.ticks.y = element_blank(),
      axis.text.y = element_blank(),
      axis.text.x = element_text(size = TICK_FONT_SIZE, family = FONT_FAMILY, face = "italic", vjust = -2),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      legend.position = "right",
      legend.justification = "top",
      legend.box.just = "top",
      legend.text = element_text(size = 50, family = FONT_FAMILY),
      legend.key.size = unit(2.675, "cm")
    ) +
    labs(x = "Items", y = "Modes", fill = "") +
    coord_fixed() +
    annotate("rect", xmin = 0.5, xmax = ncol(vt) + 0.5, ymin = 0.5, ymax = nrow(vt) + 0.5, 
             colour = "white", fill = NA, size = 1)
  
  ggsave("figures/fig2_vt.png", width = 10, height = 8)
  
}
```

## Put figure 2 together 

```{r}
d <- read_csv("data/natural_dataset.csv", col_names = FALSE)
s <- svd(d)

norm_matrix <- function(matrix) {
  # max norm
  norm <- max(abs(matrix))
  return(matrix/norm)
}

# Unit norm and restrict to first three modes, as it is in Saxe et al. figure 2
# The paper multiplies both vt and u by -1 to get more intuitive visualisations
vt <- norm_matrix(t(s$v)[1:3,]) * -1
u <- norm_matrix(s$u[,1:3]) * -1
sv <- norm_matrix(diag(s$d)[1:3,1:3])

gradient_colours <- scale_fill_gradient2(low = "blue", high = "red", mid = "white", 
                                        midpoint = 0, limit = c(-1, 1), 
                                        breaks = c(-1, 0, 1))

# Reproduce individual components of figure 2
reproduce_figure_2_io(d, gradient_colours)
reproduce_figure_2_u(u, gradient_colours)
reproduce_figure_2_sv(sv, gradient_colours)
reproduce_figure_2_vt(vt, gradient_colours)

# Fuck it, I'm just going to put the rest together in Canva. So annoying to do this programmatically with grids.
# Link: https://www.canva.com/design/DAF0GHxXVdU/XqJsZz267k3bf0sJekw6GQ/edit?utm_content=DAF0GHxXVdU&utm_campaign=designshare&utm_medium=link2&utm_source=sharebutton

```

# Sample hierarchical features from a branching diffusion process
Helper functions for figures 3, 4 and 5

```{r}

flip_feature <- function(flip_probability) {
  # Return -1 with probability flip_probability, otherwise return 1
  if (runif(1) < flip_probability) {
    return(-1)
  } else {
    return(1)
  }
}

sample_feature_vector <- function(node_value, depth, max_depth, e) {
  # Sample feature vector from a branching diffusion process with depth 'max_depth' and flip probability 'e'
  
  left_node_value <- flip_feature(e) * node_value
  right_node_value <- flip_feature(e) * node_value
  
  if (depth == max_depth) {
    return(c(node_value))
  } else if (depth + 1== max_depth) {
    return(c(left_node_value, right_node_value))
  } else {
    return(c(sample_feature_vector(left_node_value, depth + 1, max_depth, e), sample_feature_vector(right_node_value, depth + 1, max_depth, e)))
  }
}

sample_feature_matrix <- function(d, e, n) {
  # Sample 'n' features from a branching diffusion process of depth 'd' with flip probability 'e'
  
  features <- c()
  pblapply(1:n, function(x) {
    root_node_value <- flip_feature(0.5)
    feature <- list(sample_feature_vector(root_node_value, 0, d-1, e))
    features <<- c(features, feature)
  })
  #for (i in 1:n) {
  #  root_node_value <- flip_feature(0.5)
  #  feature <- list(sample_feature_vector(root_node_value, 0, d-1, e))
  #  features <- c(features, feature)
  #}
  
  feature_matrix <- do.call(rbind, features)
  
  return(feature_matrix)
}

```

# Deep linear network training
Helper function for figure 3

```{r}

sample_train_data <- function(D, e, N) {
  # Sample training dataset from a hierarchical diffusion process
  
  train_Y <- t(sample_feature_matrix(D, e, N))
  train_X <- diag(4)
  
  train_data <- c()
  for (i in 1:4) {
    train_data <- c(train_data, list(list(unlist(train_X[i,]), unlist(train_Y[i,]))))
  }
  
  return(train_data)
  
}

initialise_network <- function(N1, N2, N3) {
  # Initialise a three layer linear network with layer sizes N1, N2, N3
  
  network <- list(
    W21 = matrix(rnorm(N2 * N1, mean = 0, sd = sqrt(0.001)), nrow = N2, ncol = N1),
    W32 = matrix(rnorm(N3 * N2, mean = 0, sd = sqrt(0.001)), nrow = N3, ncol = N2)
  )
  
  return(network)
}

forward <- function(network, x) {
  # Forward propagation through our network
  
  z2 <- network$W21 %*% x
  z3 <- network$W32 %*% z2
  
  return(z3)
}

backward <- function(network, x, y_pred, y_true, learning_rate) {
  # Back propagation through our network
  
  # Calculate gradients
  dW21 <- t(network$W32) %*% (y_true - y_pred) %*% t(x)
  dW32 <- (y_true - y_pred) %*% t(network$W21 %*% x)
  
  # Update weights
  network$W21 <- network$W21 + learning_rate * dW21
  network$W32 <- network$W32 + learning_rate * dW32
  
  return(network)
  
}

squared_error_loss <- function(y, y_true) {
  error <- y - y_true
  loss <- sum(error^2)
  return(loss)
}

train_network <- function(network, data, learning_rate, num_epochs) {
  # Train our three layer deep linear network, return singular values of weights across training
  
  singular_values <- c()
    
  for (epoch in 1:num_epochs) {
    
    loss_sum <- 0
    
    for (d in data) {
      
      x <- d[1][[1]]
      y_true <- d[2][[1]]
      
      # Forward pass
      y_pred <- forward(network, x)
      loss <- squared_error_loss(y_pred, y_true)
      loss_sum <- loss_sum + loss
      
      # Backprop
      network <- backward(network, x, y_pred, y_true, learning_rate)
      
    }
    
    if (epoch %% 100 == 0) {
      cat(sprintf("Epoch: %d, Average loss: %.3f\n", epoch, loss_sum / length(data)))
    }
    
    # Calculate singular values
    sv <- svd(network$W32 %*% network$W21)$d[1:3]
    singular_values <- c(singular_values, list(sv))

  }
  
  return(singular_values)
}

simulate_weight_sv <- function(train_data, N1, N2, N3, learning_rate, num_epochs) {
  # Single simulation for fig 3 (the red lines)
  
  network <- initialise_network(N1, N2, N3)
  singular_values <- train_network(network, train_data, learning_rate, num_epochs)
  
  return(singular_values)
}

```

```{r}
# GENERATE TRAINING DATASET
D <- 3 # Depth of tree
e <- 0.1 # Flip probability
N <- 10000 # Number of features to sample

train_data <- sample_train_data(D, e, N)

# DEEP LINEAR NETWORK ARCHITECTURE
N1 <- 4  # Input layer size
N2 <- 10  # Hidden layer size
N3 <- 10000  # Output layer size
LEARNING_RATE <- 0.00025
NUM_EPOCHS <- 600

# RUN TRAINING SIMS
NUM_SIMULATIONS <- 10
singular_values <- c()
for (i in 1:NUM_SIMULATIONS) {
  cat(sprintf("Simulation %d\n", i))
  svs <- simulate_weight_sv(train_data, N1, N2, N3, LEARNING_RATE, NUM_EPOCHS)
  svs <- do.call(rbind, svs)
  singular_values <- c(singular_values, list(svs))
  cat("--------------------------------------------------\n")
}

```

```{r}

# CONVERT SIM SV's INTO PLOTTABLE FORMAT
long_data <- lapply(seq_along(singular_values), function(i) {
  mat <- singular_values[[i]]
  df <- melt(mat)
  df$matrix_id <- i
  df$line_type <- "Simulation"
  names(df)[1:2] <- c("X", "Line")
  df
})
simulation_df <- do.call(rbind, long_data)

# THEORETICAL CURVES
SCALE <- 3.25 # This is so suss lmao
tau <- 4 / LEARNING_RATE
a0 <- 0.1
ss <- svd(training_data)$d[1:3]
theoretical_curves <- c()
for (s in ss) {
  tc <- c()
  for (t in 1:NUM_EPOCHS) {
    tc <- c(tc, s*exp(2*s*t*SCALE/tau)/(exp(2*s*t*SCALE/tau)-1+(s*SCALE/a0)))
  }
  theoretical_curves <- c(theoretical_curves, list(tc))
}
theoretical_matrix <- t(do.call(rbind, theoretical_curves))
theoretical_df <- melt(theoretical_matrix)
names(theoretical_df)[1:2] <- c("X", "Line")
theoretical_df$matrix_id <- NUM_SIMULATIONS + 1
theoretical_df$line_type <- "Theory"

# COMBINE
combined_df <- rbind(theoretical_df, simulation_df)

# PLOT!
ggplot(combined_df, aes(x = X, y = value, group = interaction(factor(line_type, levels = c("Simulation", "Theory")), matrix_id, Line), color = line_type, size = line_type)) +
  geom_line() +
  theme_minimal() +
  scale_size_manual(values = c("Theory" = 2, "Simulation" = 0.5)) +
  scale_color_manual(values = c("Theory" = "blue", "Simulation" = "red")) +
  scale_x_continuous(breaks = seq(0, 601, by = 100), expand = c(0, 0), limits = c(0, 601)) +
  scale_y_continuous(breaks = seq(0, 160, by = 50)) +
  theme(
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    axis.text.x = element_text(size = TICK_FONT_SIZE*0.7), 
    axis.text.y = element_text(size = TICK_FONT_SIZE*0.7), 
    axis.title.x = element_text(size = TITLE_FONT_SIZE*0.7), 
    axis.title.y = element_text(size = TITLE_FONT_SIZE*0.7),
    axis.line = element_line(color = "black", size = 0.7),
    axis.ticks = element_line(color = "black", size = 0.7),
    legend.position = "none"
  ) +
  labs(color = "Line", x = "t (Epochs)", y = "Input-output mode strength")

ggsave("figures/fig3.png", width = 10, height = 7)
```

# Reproduce figure 4b

```{r}

# GENERATE SAMPLE DATASET
D <- 4 # Depth of tree
e <- 0.1 # Flip probability
N <- 5000 # Number of features to sample

for (n in c(10, 100, 1000, 10000)) {
  for (k in 1:10) {

    sample_data <- sample_feature_matrix(D, e, n)
    object_analyser <- svd(sample_data)$v
    
    oa_df <- reshape2::melt(round(object_analyser, 2))
    custom_colors <- scale_fill_gradientn(colors = c("#0a2ca0", "#9af482", "#94251c"),  
                            values = scales::rescale(c(-1, 0, 1)),
                            limits = c(-1, 1))
    
    ggplot(oa_df, aes(Var1, Var2, fill = value)) +
      geom_tile() +
      geom_text(aes(label=value), color="black", size=4, family="Times New Roman") +
      geom_hline(yintercept = seq(0.5, 8.5, by = 1), color = "black") +  # Adding horizontal black lines
      custom_colors +
      theme_minimal() +
      scale_x_continuous(breaks = 1:8) +  # Adjust x-axis ticks  # Flip and adjust y-axis ticks
      scale_y_reverse(breaks = 1:8, labels = 1:8) + 
      labs(
            fill=NULL,
            x="Labels",
            y="Modes"
        ) +
      theme(
            panel.grid.major = element_blank(),
            panel.grid.minor = element_blank(),
        ) +
      coord_fixed(ratio = 1)
    
    ggsave(sprintf("figures/fig4/fig4_%d_%d.png", n, k), width = 10, height = 7)
    
  }
}

```

# Reproduce figure 5

```{r}
D <- 6 # Depth of tree
e <- 0.1 # Flip probability
N <- 200 # Number of features to sample
  
simulate_singular_values <- function() {
  # Generate simulated singular values for figure 5, ie. the red points
  
  feature_matrix <- sample_feature_matrix(D, e, N)
  
  svd_result <- svd(feature_matrix)
  sigma <- svd_result$d
  
  return(sigma[1:6])
}

# Simulated singular values
num_simulations <- 100 # Number of simulations to run
singular_values <- c()
for (i in 1:num_simulations) {
  ssv <- simulate_singular_values()
  singular_values <- c(singular_values, list(ssv))
}
simulated_sv_matrix <- do.call(rbind, singular_values)
melted_ssv_matrix <- melt(simulated_sv_matrix)

# Theoretically derived singular values
P <- 2**(D-1) # Number of leaves in our tree
qs <- c()
for (k in 0:(D-1)) {
  qk <- (1-4*e*(1-e))**(D-1-k)
  qs <- c(qs, qk)
}
theoretical_sv <- c()
for (k in 0:(D-1)) {
  s <- 0
  for (l in k:(D-1)) {
    if (l == 0) {
      q1 <- 0
    } else {
      q1 <- qs[l]
    }
    delta_l <- qs[l+1] - q1
    M_l <- 2**l
    s <- s + delta_l / M_l
  }
  tsv <- sqrt(N*P*s)
  theoretical_sv <- c(theoretical_sv, tsv)
}

png("figures/fig5.png", width = 1100, height = 600)

# Create figure 5
plot(melted_ssv_matrix$Var2-1, melted_ssv_matrix$value, type = "p", pch = 16, col = "red", 
     xlab = "Hierarchy Level", ylab = "Singular Value", ylim = c(0, 50))
points(0:5, theoretical_sv, col = "black", pch = 19, cex = 1.5)
axis(2, at = seq(0, 50, by = 10))

dev.off()

```

```{r}
t(svd_result$v[1:8,1:8])
```

```{r}

features <- c()
for (i in 1:N) {
  root_node_value <- flip_feature(0.5)
  features <- c(features, list(sample_features(root_node_value,1,D)))
}
feature_matrix <- do.call(rbind, features)

svd_result <- svd(feature_matrix)
sigma <- svd_result$d

sigma[1:6]
```