Figure4_Analysis.Rmd

---
title: "RampCodes_Figure4"
author: "Sarah Tennant & Matt Nolan"
date: "20/10/2021"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```


### ------------------------------------------------------------------------------------------ ### 


## Script aims to evaluate the dependence of firing rates on the reward zone cues and to evaluate properties of neurons with similar firing rate trajectories in the presence and absence of these cues.


### ------------------------------------------------------------------------------------------ ### 

First, classify neurons based on their slope activity in the outbound and homebound region in probe trials (similar to beaconed trials in Figure1_Analysis.Rmd)

This analysis isn't so useful as there are relatively few probe trials and therefore insufficient power for classification. This may be useful for future experiments that optimise probe trial number.

```{r}
# Use function 'mark_track_category' previously used in Figure 1
spatial_firing <- spatial_firing %>%
  mutate(track_category_nb = map2(lm_group_nb, lm_group_nb_h, mark_track_category),
         track_category_p = map2(lm_group_p, lm_group_p_h, mark_track_category))

# Find predicted rates after the reward zone based on rates before the reward zone
spatial_firing <- spatial_firing %>%
  mutate(normalised_rates_nb = map(Rates_averaged_rewarded_nb, normalise_rates),
         normalised_rates_p = map(Rates_averaged_rewarded_p, normalise_rates))

spatial_firing <- spatial_firing %>%
  mutate(normalised_rates_nb_smoothed = map(Rates_averaged_rewarded_smoothed_nb, normalise_rates),
         normalised_rates_p_smoothed = map(Rates_averaged_rewarded_smoothed_p, normalise_rates))

spatial_firing <- spatial_firing %>%
  select(-contains('predict_params_nb_')) %>%
  select(-contains('predict_params_p_')) %>%
  mutate(predict_params_nb = map(normalised_rates_nb, predict_homebound),
         predict_params_p = map(normalised_rates_p, predict_homebound)) %>%
  unnest_wider(predict_params_nb,
               names_sep = "_",
               names_repair = "universal") %>%
  unnest_wider(predict_params_p,
               names_sep = "_",
               names_repair = "universal")

spatial_firing <- spatial_firing %>%
  mutate(offset_nb = pmap_chr(list(normalised_rates_nb, predict_params_nb_lwr, predict_params_nb_upr), offset_test),
         predict_diff_nb = map2_dbl(normalised_rates_nb, predict_params_nb_fit, calc_predict_diff),
         offset_p = pmap_chr(list(normalised_rates_p, predict_params_p_lwr, predict_params_p_upr), offset_test),
         predict_diff_p = map2_dbl(normalised_rates_p, predict_params_p_fit, calc_predict_diff))

# Classify cells based on their predicted activity using function 'mark_reset_group_predict' defined in Figure 1.

spatial_firing <- spatial_firing %>%
  mutate(reset_group_nb = map(offset_nb, mark_reset_group_predict),
         reset_group_p = map(offset_p, mark_reset_group_predict))

table(spatial_firing$offset_nb, useNA = "always")
table(spatial_firing$offset_p, useNA = "always")
```



### ------------------------------------------------------------------------------------------ ### 

Does the firing rate slope before or after the reward zone differ between beaconed and probe trials?

Subset spatial firing to contain only cells that have probe trial data and where the number of trials in the session  is >= 120, and therefore probe trials >= 12.

Focus on positional neurons (P, PA, PS, PSA) that on beaconed trials are classified with positive or negative slopes on the track segment before the reward zone
```{r}
sum_unlist <- function(df) {sum(unlist(df), na.rm=TRUE)}
spatial_firing_probe_only <-
  subset(spatial_firing, max_trial_number >= 120) %>%
  mutate(asr_p_sum = map_dbl(Rates_averaged_p, sum_unlist)) %>%
  subset(asr_p_sum != 0) %>%
  filter(lm_group_b == "Positive" | lm_group_b == "Negative",
         final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" |  final_model_o_b == "PSA") 

# Number of eligible probe neurons
dim(spatial_firing_probe_only)

# Number of positional neurons
dim(filter(spatial_firing, lm_group_b == "Positive" | lm_group_b == "Negative",
         final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" |  final_model_o_b == "PSA"))

# Number of sessions
length(unique(spatial_firing_probe_only$session_id))

# Number of mice
unique(spatial_firing_probe_only$unique_mouse)
```



Plot absolute slope values for each trial type
```{r}
# Plots
probe_out_slope_plot(spatial_firing_probe_only, "Positive", -0.1, 0.45)

if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_beaconed_vs_probe_outbound_pos.png", width = 3.2, height = 3.5) 
}

probe_out_slope_plot(spatial_firing_probe_only, "Negative", -0.65, 0.2)


if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_beaconed_vs_probe_outbound_neg.png", width = 3.2, height = 3.5) 
}

probe_home_slope_plot(spatial_firing_probe_only, "Positive", -0.35, 0.8)


if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_beaconed_vs_probe_homebound_pos.png", width = 3.2, height = 3.5) 
}

probe_home_slope_plot(spatial_firing_probe_only, "Negative", -0.4, 0.75)


if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_beaconed_vs_probe_homebound_neg.png", width = 3.2, height = 3.5) 
}


# Plots with non-beaconed trials
nb_probe_out_slope_plot(spatial_firing_probe_only, "Positive", -0.1, 0.45)
nb_probe_out_slope_plot(spatial_firing_probe_only, "Negative", -0.65, 0.2)

nb_probe_home_slope_plot(spatial_firing_probe_only, "Positive", -0.45, 0.8)
nb_probe_home_slope_plot(spatial_firing_probe_only, "Negative", -0.6, 0.8)


# Compare probe vs beaconed with a paired t-test
probe_slope_tt_p <- spatial_firing_probe_only %>%
  filter(lm_group_b == "Positive") %>%
  summarise(ttest = list(t.test(asr_b_o_rewarded_fit_slope, asr_p_o_rewarded_fit_slope, paired = TRUE),
                         t.test(asr_b_h_rewarded_fit_slope, asr_p_h_rewarded_fit_slope, paired = TRUE))) 

probe_slope_tt_n <- spatial_firing_probe_only %>%
  filter(lm_group_b == "Negative") %>%
  summarise(ttest = list(t.test(asr_b_o_rewarded_fit_slope, asr_p_o_rewarded_fit_slope, paired = TRUE),
                         t.test(asr_b_h_rewarded_fit_slope, asr_p_h_rewarded_fit_slope, paired = TRUE))) 

probe_slope_tt_p[[1]]
probe_slope_tt_n[[1]]

# Test for probe slope different from zero
probe_slope_tt_p_0 <- spatial_firing_probe_only %>%
  filter(lm_group_b == "Positive") %>%
  summarise(ttest = list(t.test(asr_p_o_rewarded_fit_slope, mu = 0),
                         t.test(asr_p_h_rewarded_fit_slope, mu = 0))) 

probe_slope_tt_n_0 <- spatial_firing_probe_only %>%
  filter(lm_group_b == "Negative") %>%
  summarise(ttest = list(t.test(asr_p_o_rewarded_fit_slope, mu = 0),
                         t.test(asr_p_h_rewarded_fit_slope, mu = 0))) 

probe_slope_tt_p_0[[1]]
probe_slope_tt_n_0[[1]]
```



Plot scatter of slopes after the reward zone on beaconed & probe trials.
```{r}
spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") %>%
  b_vs_p_o_slope_plot()

if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_outbound.png", width = 4, height = 4) 
}


spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") %>%
  b_vs_p_h_slope_plot()

if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_homebound.png", width = 4, height = 4) 
}
```



Calculate r2 and p value for liner correlation of the above data
```{r}
model <- lm(asr_b_o_rewarded_fit_slope ~ asr_p_o_rewarded_fit_slope, data = spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") )
summary(model)
# get the model parameters
params <- select(glance(model), r.squared, p.value)

model <- lm(asr_b_h_rewarded_fit_slope ~ asr_p_h_rewarded_fit_slope, data = spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") )
summary(model)
# get the model parameters
params <- select(glance(model), r.squared, p.value)
```


Now compare offsets on beaconed vs probe trials

Plot absolute offset values for each trial type
```{r}
spatial_firing_probe_only %>% filter(lm_group_b == "Positive" & lm_group_b_h == "Positive" | lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  probe_offset_plot("Positive", -3, 1)

if (save_figures == 1) {
 ggsave(file = "plots/offset_comparison_beaconed_vs_probe_pospos.png", width =  3.2, height = 3.5) 
}

spatial_firing_probe_only %>% filter(lm_group_b == "Positive" & lm_group_b_h == "Positive" | lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  probe_offset_plot("Negative", -1, 5)

if (save_figures == 1) {
 ggsave(file = "plots/offset_comparison_beaconed_vs_probe_negneg.png", width = 3.2, height = 3.5) 
}


spatial_firing_probe_only %>% filter(lm_group_b == "Positive" & lm_group_b_h == "Positive" | lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  nb_probe_offset_plot("Positive", -3, 1)

spatial_firing_probe_only %>% filter(lm_group_b == "Positive" & lm_group_b_h == "Positive" | lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  nb_probe_offset_plot("Negative", -1, 5)
```


Look at possible correlations
```{r}
spatial_firing_probe_only %>% filter(lm_group_b == "Positive" & lm_group_b_h == "Positive" | lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  b_vs_p_offset_plot()


if (save_figures == 1) {
 ggsave(file = "plots/offset_comparison_beaconed_vs_probe.png", width = 4, height = 4) 
}

```


Plot distribution of offsets on probe trials.
```{r}
# All position types
spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA" | final_model_o_b == "PS" | final_model_o_b == "PSA") %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  offset_ggplot(diff_colname = "predict_diff_p", group_colname = "reset_group_p", colour_2 = "violetred2")

spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA" | final_model_o_b == "PS" | final_model_o_b == "PSA") %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  offset_ggplot(diff_colname = "predict_diff_p", group_colname = "reset_group_p", colour_2 = "violetred2")
```

Do the offsets differ between beaconed and probe trials?
```{r}
offset_b_vs_p_pos <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA" | final_model_o_b == "PS" | final_model_o_b == "PSA") %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  summarise(ttest = list(t.test(predict_diff, predict_diff_p, paired = TRUE)))

offset_b_vs_p_neg <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA" | final_model_o_b == "PS" | final_model_o_b == "PSA") %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  summarise(ttest = list(t.test(predict_diff, predict_diff_p, paired = TRUE)))

offset_b_vs_p_pos[[1]]
offset_b_vs_p_neg[[1]]
```


Does the distribution mean differ from zero?

Treat each observation as independent and use a t-test
```{r}
spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA" | final_model_o_b == "PS" | final_model_o_b == "PSA") %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  select(predict_diff_p) %>%
  t.test(mu = 0)

spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA" | final_model_o_b == "PS" | final_model_o_b == "PSA") %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  select(predict_diff_p) %>%
  t.test(mu = 0)
```
Do offsets on beaconed trials predict offsets on probe trials?
```{r}
model <- lm(predict_diff_p ~ predict_diff, data = spatial_firing_probe_only %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") )
summary(model)
# get the model parameters
params <- select(glance(model), r.squared, p.value)
```


Does the change in slope between beaconed and probe trials predict the change in offset? E.g. If a neuron's slope is stable then is it's offset? We might expect this for recall of a positional memory.
```{r}
# Look only at P group
spatial_firing %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  ggplot(aes(x = asr_p_o_rewarded_fit_slope, y = predict_diff_p)) +
  geom_point()


spatial_firing %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  ggplot(aes(x = asr_p_o_rewarded_fit_slope, y = predict_diff_p)) +
  geom_point()
```



### -------------------------------------------------------------------------------------------------------------------------------------------------- ###


## Plot average firing rates for neurons with different firing rate profiles (e.g. ++, +-, etc).

```{r}
# all P cells (used for the manuscript)
spatial_firing_probe_only %>% filter(final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") %>% comp_beacon_probe_rate_plots
spatial_firing_probe_only_P <- filter(spatial_firing_probe_only, final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA")

P_plots <- comp_beacon_probe_rate_plots(spatial_firing_probe_only_P)

if (save_figures==1) {
  ggsave(file = "plots/beaconed_vs_probe_negneg_position_all.png", plot = P_plots[[1]], width = 3.6, height = 2.9)
  ggsave(file = "plots/beaconed_vs_probe_negpos_position_all.png", plot = P_plots[[2]], width = 3.6, height = 2.9)
  ggsave(file = "plots/beaconed_vs_probe_pospos_position_all.png", plot = P_plots[[3]], width = 3.6, height = 2.9)
  ggsave(file = "plots/beaconed_vs_probe_posneg_position_all.png", plot = P_plots[[4]], width = 3.6, height = 2.9)
}


nrow(filter(spatial_firing_probe_only, lm_group_b == "Negative" & lm_group_b_h == "Negative"))
nrow(filter(spatial_firing_probe_only, lm_group_b == "Negative" & lm_group_b_h == "Positive"))
nrow(filter(spatial_firing_probe_only, lm_group_b == "Positive" & lm_group_b_h == "Positive"))
nrow(filter(spatial_firing_probe_only, lm_group_b == "Positive" & lm_group_b_h == "Negative"))


# only P cells
spatial_firing_probe_only %>% filter(final_model_o_b == "P") %>% comp_beacon_probe_rate_plots

# conjunctive P cells
spatial_firing_probe_only %>% filter(final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA") %>% comp_beacon_probe_rate_plots

```


What is the distribution of firing rate profiles of previously positional ramping neurons?
To look at distributions in the data superimpose mean firing rate as a function of positon for each neuron in each group and colour code according to the offset, slope or change in slope / offset.
```{r}
comp_plot <- function(df, plot_column_1 = "Rates_averaged_rewarded_smoothed_b", plot_column_2 = "Rates_averaged_rewarded_smoothed_p", group1 = "Positive", group2 = "Positive") {
  df$pc_1 <- df[[plot_column_1]]
  df$pc_2 <- df[[plot_column_2]]
  df %>% filter(lm_group_b == group1 & lm_group_b_h == group2) %>%
    select(unique_id, pc_1, pc_2) %>%
    mutate(Rates_1 = map(pc_1, unlist),
           Rates_2 = map(pc_2, unlist)) %>%
    unnest(c(unique_id, Rates_1, Rates_2)) %>%
    mutate(Pos = rep(-30:169, length(unique(unique_id)))) %>%
    ggplot() +
    geom_line(aes(x = Pos, y = Rates_1), colour = "blue")+
    geom_line(aes(x = Pos, y = Rates_2), colour = "red") +
    facet_wrap(~unique_id, scales = "free")
}

arp_P <- spatial_firing_probe_only %>%
    filter(final_model_o_b == "P") %>%
  comp_plot("Rates_averaged_rewarded_smoothed_b", "Rates_averaged_rewarded_smoothed_p", "Positive", "Positive")

arp_N <- spatial_firing_probe_only %>%
    filter(final_model_o_b == "P") %>%
  comp_plot("Rates_averaged_rewarded_smoothed_b","Rates_averaged_rewarded_smoothed_p", "Negative", "Negative")

add_track(arp_P)
add_track(arp_N)

```



### ------------------------------------------------------------------------------------------------------------------- ###

We want to know about whether the slope classification of neurons on beaconed trials is maintained on probe trials.

## Calculate numbers of neurons with given outbound and homebound classifications on beaconed and probe trials.
```{r}
# How many cells are in each group?
(b_vs_p_nos <- table(
  mapply(
    paste,
    spatial_firing_probe_only$lm_group_b,
    spatial_firing_probe_only$lm_group_b_h
  ),
  mapply(
    paste,
    spatial_firing_probe_only$lm_group_p,
    spatial_firing_probe_only$lm_group_p_h
  )
))

```



### ------------------------------------------------------------------------------------------ ### 




### Calculate average stop histogram for all mice, sessions and plot

        
```{r}
df <- spatial_firing_probe_only %>%
  select(average_stops, average_stops_p, session_id) %>%
  distinct %>%
  mutate(average_stops = map(average_stops, unlist), 
         average_stops_p = map(average_stops_p, unlist)) %>%
  unnest(c(session_id, average_stops, average_stops_p)) %>%
  mutate(Position = rep(-30:169, length(unique(session_id)))) %>%
  group_by(Position) %>%
  dplyr::summarise(mean_r = mean(as.numeric(average_stops), na.rm =TRUE), sem_r = std.error(as.numeric(average_stops)), 
                   mean_p = mean(as.numeric(average_stops_p), na.rm =TRUE), sem_p = std.error(as.numeric(average_stops_p))) 

bp <- ggplot(data=df) +
  geom_line(aes(y=mean_r, x=Position), color = "Black") +
  geom_line(aes(y=mean_p, x=Position), color = "Blue")

(add_track(bp))


```

```{r}
stops_plot <-
  function(df, cn1 = "average_stops", cn2 = "average_stops_p") {
    df$cn1 <- df[[cn1]]
    df$cn2 <- df[[cn2]]
    df %>% select(session_id, cn1, cn2) %>%
      distinct() %>%
      mutate(Stops1 = map(cn1, unlist),
             Stops2 = map(cn2, unlist)) %>%
      unnest(c(session_id, Stops1, Stops2)) %>%
      mutate(Pos = rep(1:200, length(unique(session_id)))) %>%
      ggplot() +
      geom_line(aes(x = Pos, y = Stops1), colour = "Red") +
      geom_line(aes(x = Pos, y = Stops2), colour = "Blue") +
      facet_wrap(vars(session_id), ncol = 5, scales = "free_y")
}

(pl_p <- spatial_firing_probe_only %>% filter(final_model_o_b == "P") %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  stops_plot(cn1 = "average_stops", cn2 = "average_stops_p"))
  
(pl_n <- spatial_firing_probe_only %>% filter(final_model_o_b == "P") %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  stops_plot(cn1 = "average_stops", cn2 = "average_stops_p"))
  
```