Figure1_Analysis.Rmd

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

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


## Identification of ramping activity of neurons recorded from the parahippocampal areas during a location memory task

The aim of this analysis is:
1. Identify cells that show ramping changes in their firing rate.
2. Classify cells according to the profile of their activity.
3. Examine the population average activity of different cell groups.


To set up, including loading packages and data, first run SetUp.Rmd.

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

## Prepare data for fitting linear models

Make a copy of the average firing rates for each cell that has position added.

input columns: 
Rates_averaged_rewarded_b = beaconed trials
Rates_averaged_rewarded_nb = non-beaconed and probe trials
Rates_averaged_rewarded_p = probe trials only
```{r}
spatial_firing <- spatial_firing %>%
  mutate(asr_b_rewarded = future_pmap(list(Rates_averaged_rewarded_b, session_id, cluster_id), add_position),
         asr_nb_rewarded = future_pmap(list(Rates_averaged_rewarded_nb, session_id, cluster_id), add_position),
         asr_p_rewarded = future_pmap(list(Rates_averaged_rewarded_p, session_id, cluster_id), add_position)
         )
```


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

# Fit linear models to examine relationship between firing rate and position


# Focus on the region of the track before the reward zone (labelled as 30-90 cm here)

Removes any previously generated results (select), fits the data (mutate) and then adds model outputs as columns to spatial firing (unnest_wider). The function lm_tidy_helper fits the linear model and then returns data extracted from the model fit using glance and tidy tools.
```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('asr_b_o_rewarded_fit')) %>%
  select(-contains('asr_nb_o_rewarded_fit')) %>%
  select(-contains('asr_p_o_rewarded_fit')) %>%
  mutate(asr_b_o_rewarded_fit = pmap(list(asr_b_rewarded, 30, 90), lm_tidy_helper),
         asr_nb_o_rewarded_fit = pmap(list(asr_nb_rewarded, 30, 90), lm_tidy_helper),
         asr_p_o_rewarded_fit = pmap(list(asr_p_rewarded, 30, 90), lm_tidy_helper))

spatial_firing <- spatial_firing %>%
  unnest_wider(asr_b_o_rewarded_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_nb_o_rewarded_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_p_o_rewarded_fit, names_sep = "_", names_repair = "universal")
```


  
Linear model results are stored in:
spatial_firing$asr_b_o_rewarded_fit_slope
spatial_firing$asr_b_o_rewarded_fit_intercept
spatial_firing$asr_b_o_rewarded_fit_p.value
spatial_firing$asr_b_o_rewarded_fit_r.squared


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

Generate or load shuffled data.
```{r}
# Make this smaller for testing.
shuffles <- 1000

# To load circular shuffle results
if(shuffle_type=="circular") {
  cs_path = "data_in/all_mice_concatenated_shuffle_data_rewarded_unsmoothened.feather"
  spatial_firing_circ <- local_circ_shuffles(spatial_firing, cs_path)
  
  spatial_firing <- spatial_firing %>%
    select(-contains('shuffle_results_b_o')) %>%
    select(-contains('shuffle_results_nb_o')) %>%
    select(-contains('shuffle_results_pb_o')) %>%
    select(-contains('shuffle_results_b_h')) %>%
    select(-contains('shuffle_results_nb_h')) %>%
    select(-contains('shuffle_results_pb_h')) %>%
    mutate(shuffle_results_b_o = spatial_firing_circ$shuffle_results_b_o,
           shuffle_results_nb_o = spatial_firing_circ$shuffled_results_nb_o,
           shuffle_results_p_o = spatial_firing_circ$shuffled_results_p_o,
           shuffle_results_b_h = spatial_firing_circ$shuffle_results_b_h,
           shuffle_results_nb_h = spatial_firing_circ$shuffled_results_nb_h,
           shuffle_results_p_h = spatial_firing_circ$shuffled_results_p_h)
  # Remove unused frame
  rm(spatial_firing_circ)
}

# Check to see if the column shuffle_results_b_o exists.
# It will exist if shuffles have been pre-loaded or if circular shuffle results have been loaded.
# If it doesn't exist then the mean firing rate as a function of position will be shuffled.
if(!"shuffle_results_b_o" %in% colnames(spatial_firing)) {
  spatial_firing <- spatial_firing %>%
    mutate(shuffle_results_b_o = future_pmap(list(Rates_averaged_rewarded_b, 30, 90, shuffles), shuffle_rates)) %>%
    mutate(shuffle_results_nb_o = future_pmap(list(Rates_averaged_rewarded_nb, 30, 90, shuffles), shuffle_rates)) %>%
    mutate(shuffle_results_p_o = future_pmap(list(Rates_averaged_rewarded_p, 30, 90, shuffles), shuffle_rates))
}


if (save_results == 1) {
  saveRDS(spatial_firing, "data_out/SpatialFiring_with_1000_shuffles_of.Rda")
  # And use this to save a truncated version. Useful for testing code.
  saveRDS(slice_head(spatial_firing, n = 5), "SpatialFiring_with_1000_shuffles_trunc.Rda")
}

```



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

### Identify ramping activity by whether the coefficients of the linear model lie outside the 5th-95th percentiles of the same result from 1000 shuffled data sets.


First, extract the 5 % and 95 % quantiles of 1000 shuffles for each neuron.
```{r}
spatial_firing <- spatial_firing %>%

  mutate(shuffle_min_slope_b_o = future_map2_dbl(shuffle_results_b_o, 0.05, extract_quantile_shuffle_slopes),
         shuffle_max_slope_b_o = future_map2_dbl(shuffle_results_b_o, 0.95, extract_quantile_shuffle_slopes),
         shuffle_min_slope_nb_o = future_map2_dbl(shuffle_results_nb_o, 0.05, extract_quantile_shuffle_slopes),
         shuffle_max_slope_nb_o = future_map2_dbl(shuffle_results_nb_o, 0.95, extract_quantile_shuffle_slopes),
         shuffle_min_slope_p_o = future_map2_dbl(shuffle_results_p_o, 0.05, extract_quantile_shuffle_slopes),
         shuffle_max_slope_p_o = future_map2_dbl(shuffle_results_p_o, 0.95, extract_quantile_shuffle_slopes),
         shuffle_min_slope_b_h = future_map2_dbl(shuffle_results_b_h, 0.05, extract_quantile_shuffle_slopes),
         shuffle_max_slope_b_h = future_map2_dbl(shuffle_results_b_h, 0.95, extract_quantile_shuffle_slopes),
         shuffle_min_slope_nb_h = future_map2_dbl(shuffle_results_nb_h, 0.05, extract_quantile_shuffle_slopes),
         shuffle_max_slope_nb_h = future_map2_dbl(shuffle_results_nb_h, 0.95, extract_quantile_shuffle_slopes),
         shuffle_min_slope_p_h = future_map2_dbl(shuffle_results_p_h, 0.05, extract_quantile_shuffle_slopes),
         shuffle_max_slope_p_h = future_map2_dbl(shuffle_results_p_h, 0.95, extract_quantile_shuffle_slopes))
```


  
  
We also want to extract slopes, r2 and pvalues of the 1000 shuffles for each neuron

1. Extract shuffle results (slopes and r2 for each shuffle)

2. run on all neurons
```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('shuffle_results_b_o_')) %>%
  select(-contains('shuffle_results_nb_o_')) %>%
  select(-contains('shuffle_results_p_o_')) %>%
  select(-contains('adjust_pval_')) %>%
  unnest_wider(shuffle_results_b_o, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(shuffle_results_nb_o, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(shuffle_results_p_o, names_sep = "_", names_repair = "universal")

spatial_firing <- spatial_firing %>%
  unnest_wider(shuffle_results_b_h, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(shuffle_results_nb_h, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(shuffle_results_p_h, names_sep = "_", names_repair = "universal")
```

We also want to correct the p values of the lm fits to the experimental data to account for multiple comparisons.


```{r}
spatial_firing <- spatial_firing %>%
  mutate(adjust_pval_b_o = p.adjust(asr_b_o_rewarded_fit_p.value, "BH"),
         adjust_pval_nb_o = p.adjust(asr_nb_o_rewarded_fit_p.value, "BH"),
         adjust_pval_p_o = p.adjust(asr_p_o_rewarded_fit_p.value, "BH"))
```


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

Now we want to classify neurons based on:
1. Whether their slopes are outside the 5-95% range of the shuffled data.
2. Whether the adjusted p-value of the linear model fit is <= 0.01.

```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('lm_group_')) %>%
  mutate(
    lm_group_b = pmap(
      list(
        shuffle_min_slope_b_o,
        shuffle_max_slope_b_o,
        asr_b_o_rewarded_fit_slope,
        adjust_pval_b_o
      ),
      compare_slopes
    ),
    lm_group_nb = pmap(
      list(
        shuffle_min_slope_nb_o,
        shuffle_max_slope_nb_o,
        asr_nb_o_rewarded_fit_slope,
        adjust_pval_nb_o
      ),
      compare_slopes
    ),
    lm_group_p = pmap(
      list(
        shuffle_min_slope_p_o,
        shuffle_max_slope_p_o,
        asr_p_o_rewarded_fit_slope,
        adjust_pval_p_o
      ),
      compare_slopes
    )
  )
```


Linear model classification is stored in:
spatial_firing$lm_group_b
spatial_firing$lm_group_nb
spatial_firing$lm_group_p


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

### Classify neurons based on their ramp activity in the homebound region

We use a similar strategy as for the outbound track region.

Because the homebound zone is the same length as the outbound zone, and as shuffling is across the full track length, we can use the shuffle results previously generated.

# Fit linear model to examine relationship between firing rate and position
_note:for now we are only interested in the homebound region of the track (110 - 170 cm)_

Removes any previously generated results (select), fits the data (mutate) and then adds model outputs as columns to spatial firing (unnest_wider).
```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('asr_b_h_rewarded_fit_')) %>%
  select(-contains('asr_nb_h_rewarded_fit_')) %>%
  select(-contains('asr_p_h_rewarded_fit_')) %>%
  mutate(asr_b_h_rewarded_fit = pmap(list(asr_b_rewarded, 110, 170), lm_tidy_helper),
         asr_nb_h_rewarded_fit = pmap(list(asr_nb_rewarded, 110, 170), lm_tidy_helper),
         asr_p_h_rewarded_fit = pmap(list(asr_p_rewarded, 110, 170), lm_tidy_helper)) %>%
  unnest_wider(asr_b_h_rewarded_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_nb_h_rewarded_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_p_h_rewarded_fit, names_sep = "_", names_repair = "universal")
```

Linear model results are stored in:
spatial_firing$asr_b_h_rewarded_fit_pval
spatial_firing$asr_b_h_rewarded_fit_slope
spatial_firing$asr_b_h_rewarded_fit_r.squared


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



# Correct p-values from fits of the experimental data

```{r}
spatial_firing <- spatial_firing %>%
  mutate(adjust_pval_b_h = p.adjust(asr_b_h_rewarded_fit_p.value, "BH"),
         adjust_pval_nb_h = p.adjust(asr_nb_h_rewarded_fit_p.value, "BH"),
         adjust_pval_p_h = p.adjust(asr_p_h_rewarded_fit_p.value, "BH"))
```


# Assign classifications to neurons based on homebound firing

Now we want to classify neurons, taking the adjusted significance into account

Uses the function compare_slopes from above.

```{r}
spatial_firing <- spatial_firing %>%
  mutate(
    lm_group_b_h = pmap(
      list(
        shuffle_min_slope_b_o,
        shuffle_max_slope_b_o,
        asr_b_h_rewarded_fit_slope,
        adjust_pval_b_h
      ),
      compare_slopes
    ),
    lm_group_nb_h = pmap(
      list(
        shuffle_min_slope_nb_o,
        shuffle_max_slope_nb_o,
        asr_nb_h_rewarded_fit_slope,
        adjust_pval_nb_h
      ),
      compare_slopes
    ),
    lm_group_p_h = pmap(
      list(
        shuffle_min_slope_p_o,
        shuffle_max_slope_p_o,
        asr_p_h_rewarded_fit_slope,
        adjust_pval_p_h
      ),
      compare_slopes
    )
  )

```

Linear model classification is stored in:
spatial_firing$lm_group_b_h
spatial_firing$lm_group_nb_h
spatial_firing$lm_group_p_h



### ----------------------------------------------------------------------------------------- ###
### How many neurons passed criteria in the linear model ? 

1. Numbers of cells for which the classification above is positive or negative for segments of the track before (outbound) or after (homebound) the reward zone.
```{r}
(ramps_beaconed_o <- nrow(subset(spatial_firing, lm_group_b == "Negative" | lm_group_b == "Positive")))

(ramps_beaconed_h <- nrow(subset(spatial_firing, lm_group_b_h == "Negative" | lm_group_b_h == "Positive")))

(ramps_beaconed_o_h <- nrow(subset(spatial_firing, lm_group_b == "Negative" | lm_group_b == "Positive" | lm_group_b_h == "Negative" | lm_group_b_h == "Positive")))

ramps_beaconed_o_h/nrow(spatial_firing)

# Calculate by mouse (:a quick hack, there are better ways to do this)
ramps_mice_n <- spatial_firing %>%
  subset((lm_group_b == "Negative" | lm_group_b == "Positive" | lm_group_b_h == "Negative" | lm_group_b_h == "Positive")) %>%
  count(Mouse, cohort) %>%
  select(n)
mice_n <- spatial_firing %>% count(Mouse, cohort) %>%
  select(n)
ramps_by_mouse = tibble(ramps = ramps_mice_n$n, mice = mice_n$n) %>%
  mutate(percent = 100 * ramps / mice)

# Mean
mean(ramps_by_mouse$percent)
# SEM
sd(ramps_by_mouse$percent)/sqrt(length(ramps_by_mouse))

# min and max
min(ramps_by_mouse$percent)
max(ramps_by_mouse$percent)
```

### How much of the shuffled dataset is past criteria? 

1. Extract shuffled slopes and rsquared values. 
```{r}
shuffles = 1000
shuff_slopes <- tibble(slopes_o = unlist(spatial_firing$shuffle_results_b_o_slope), 
                       r2_o = unlist(spatial_firing$shuffle_results_b_o_r.squared), 
                       pval_o = unlist(spatial_firing$shuffle_results_b_o_p.value), 
                       min_slope_o = rep(spatial_firing$shuffle_min_slope_b_o, each = shuffles), 
                       max_slope_o = rep(spatial_firing$shuffle_max_slope_b_o, each = shuffles),
                       slopes_h = unlist(spatial_firing$shuffle_results_b_h_slope), 
                       r2_h = unlist(spatial_firing$shuffle_results_b_h_r.squared), 
                       pval_h = unlist(spatial_firing$shuffle_results_b_h_p.value), 
                       min_slope_h = rep(spatial_firing$shuffle_min_slope_b_h, each = shuffles), 
                       max_slope_h = rep(spatial_firing$shuffle_max_slope_b_h, each = shuffles))
```


2. Classify activity of shuffled cells on the track segment before and after the reward zone.
```{r}
shuff_slopes <- shuff_slopes %>%
  mutate(
    shuff_lm_group_b = pmap(
      list(min_slope_o,
           max_slope_o,
           slopes_o,
           pval_o),
      compare_slopes
    ),
    shuff_lm_group_b_h = pmap(
      list(min_slope_h,
           max_slope_h,
           slopes_h,
           pval_h),
      compare_slopes
    )
  )
```



3. Calculate proportion of cells in the shuffled datasets that pass criteria
```{r}
(total_shuffles <- nrow(shuff_slopes))
(shuff_ramps_o <- nrow(subset(shuff_slopes, shuff_lm_group_b == "Positive" | shuff_lm_group_b == "Negative" )))
(shuff_ramps_h <- nrow(subset(shuff_slopes, shuff_lm_group_b_h == "Positive" | shuff_lm_group_b_h == "Negative" )))
(shuff_ramps_o_h <- nrow(subset(shuff_slopes, shuff_lm_group_b == "Positive" | shuff_lm_group_b == "Negative" | shuff_lm_group_b_h == "Positive" | shuff_lm_group_b_h == "Negative" )))
(shuff_ramps_percentage <- (shuff_ramps_o_h / total_shuffles)*100)
```

Calculate how many shuffles are unclassified
```{r}
(non_shuff_ramps <- nrow(subset(shuff_slopes,shuff_lm_group_b == "Unclassified" | shuff_lm_group_b == "NA" )))
(non_shuff_ramps_percentage <- (non_shuff_ramps / total_shuffles)*100)
```


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


Now, classify cells based on their activity in the outbound and homebound region

Homebound   Outbound    Text label    Numeric label
Positive    Positive    pospos        1
Positive    Negative    posneg        2
Positive    Unlcassified    posnon    3
Negative    Negative    negneg        4
Negative    Positive    negpos        5
Negative    Unclassified    negnon    6
Other                                 0

```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('track_category')) %>%
  mutate(track_category = future_map2(lm_group_b, lm_group_b_h, mark_track_category),
         track_category_numeric = map2(lm_group_b, lm_group_b_h, mark_numeric_track_category))

```



1. How many non-other neurons in the dataset?
```{r}
(ramps <- nrow(subset(spatial_firing, track_category_numeric != 0)))
```


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

## Does firing rate reset or continue across the reward zone region? 

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

Now, we want to find out if within pospos and negneg groups - does their firing rate reset across the reward region?

To do this, we will predict the firing rate on the homebound zone from the activity in the outbound. Then find the difference between the predicted and real data to determine if cells have reset or continued. 


1. Normalize rates for all cells. This is to facilitate comparison later.
```{r}
spatial_firing <- spatial_firing %>%
  mutate(normalised_rates = map(Rates_averaged_rewarded_b, normalise_rates))

spatial_firing <- spatial_firing %>%
  mutate(normalised_rates_smoothed = map(Rates_averaged_rewarded_smoothed_b, normalise_smooth_rates))

```

2. Then, predict firing rate in homebound region based on fit from real data in outbound region.
Run on all neurons
```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('predict_params_')) %>%
  mutate(predict_params = map(normalised_rates, predict_homebound)) %>%
  unnest_wider(predict_params, names_sep = "_", names_repair = "universal") 

spatial_firing <- spatial_firing %>%
  mutate(offset = pmap_chr(list(normalised_rates, predict_params_lwr, predict_params_upr), offset_test),
         predict_diff = map2_dbl(normalised_rates, predict_params_fit, calc_predict_diff))

spatial_firing <- spatial_firing %>%
  mutate(reset_group = map(offset, mark_reset_group_predict))
```

predict_params_fit is the prediction for the firing rate after fitting.
predict_params_upr and predict_params_lwr are the 99% confidence intervals of the prediction.

offset is either 'pos', 'neg' or 'none' depending on whether the difference between the predicted and actual firing rate is above, below or inside the intervals specified by predict_params_lwr and predict_params_upr.

predict_diff is the difference between the predicted and actual firing rates at the start of the track segment that follows the reward zone.



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



## plot bar chart of mean apsolute difference between predicted and real

- do only for neurons that have slopes in outbound zone
- do this for both negative and positive slopes in the outbound zone

Plot for neurons classified as ++
```{r}
spatial_firing %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  offset_ggplot()

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

Plot neurons classified as --
```{r}
spatial_firing %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  offset_ggplot(colour_2 = "violetred2")


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



Do the means of the distributions differ from zero?

Treat each observation as independent and use a t-test
```{r}
spatial_firing %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  select(predict_diff) %>%
  t.test(mu = 0)

spatial_firing %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  select(predict_diff) %>%
  t.test(mu = 0)
```

Or consider after averaging data from the same mouse. This seems overly conservative as neurons are sampled independently but nevertheless useful to look at.
```{r}
spatial_firing %>%
  filter(lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  group_by(Mouse) %>%
  summarise(m_p_d = mean(predict_diff)) %>%
  select(m_p_d) %>%
  t.test(mu = 0)

spatial_firing %>%
  filter(lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  group_by(Mouse) %>%
  summarise(m_p_d = mean(predict_diff)) %>%
  select(m_p_d) %>%
  t.test(mu = 0)
```
as.integer(factor(spatial_firing$session_id, levels = unique(spatial_firing$session_id)))





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


## now plot scatter of slopes on track segements before and after the reward zone.

1. Here we subset neurons based on positive, negative or unclassified in the linear model and either reset or continuous firing. 
```{r}
position_neurons_all <- spatial_firing %>% 
  filter(lm_group_b == "Positive" | lm_group_b == "Negative") %>%
  select(asr_b_o_rewarded_fit_slope,
         asr_b_h_rewarded_fit_slope,
         track_category,
         lm_group_b)
```

2. Plot scatter plot
```{r}

ggplot() + 
    geom_point(data=subset(position_neurons_all, track_category == "pospos" | track_category == "negneg"),
               aes(x = as.numeric(unlist(asr_b_o_rewarded_fit_slope)), 
                   y = as.numeric(unlist(asr_b_h_rewarded_fit_slope)), 
                   color=factor(unlist(lm_group_b))), alpha=0.8) +
    geom_point(data=subset(position_neurons_all, track_category == "posneg" | track_category == "negpos"),
               aes(x = as.numeric(unlist(asr_b_o_rewarded_fit_slope)), 
                   y = as.numeric(unlist(asr_b_h_rewarded_fit_slope)), 
                   color=factor(unlist(lm_group_b))), shape=2, alpha=0.8) +
    geom_point(data=subset(position_neurons_all, track_category == "posnon" | track_category == "negnon"),
               aes(x = as.numeric(unlist(asr_b_o_rewarded_fit_slope)), 
                   y = as.numeric(unlist(asr_b_h_rewarded_fit_slope)), 
                   color=factor(unlist(lm_group_b))), shape=3, alpha=0.8) +  
    geom_point(data=spatial_firing %>% filter(lm_group_b == "Unclassified"),
               aes(x = as.numeric(unlist(asr_b_o_rewarded_fit_slope)), 
                   y = as.numeric(unlist(asr_b_h_rewarded_fit_slope)), 
                   color=factor(unlist(lm_group_b))), shape=4, alpha=0.8) +     
    coord_cartesian(ylim = c(-.45,.61), xlim = c(-.45,.45)) +
    geom_abline(intercept = 0, slope = 1, colour = "grey", linetype = "dashed") +
    geom_abline(intercept = 0, slope = -1, colour = "grey", linetype = "dashed") +
    xlab("Outbound slope") +
    ylab("Homebound slope") +
    theme_classic() +
    scale_color_manual(values=c("violetred2", "chartreuse3", "grey81")) +
    theme(axis.text.x = element_text(size=18),
          axis.text.y = element_text(size=18),
          legend.position="bottom", 
          legend.title = element_blank(),
          text = element_text(size=17), 
          legend.text=element_text(size=16), 
          axis.title.y = element_text(margin = margin(t = 0, r = 20, b = 0, l = 0))) 

if (save_figures == 1) {
 ggsave(file = "plots/slope_comparison_reset.png", width = 4, height = 4) 
}
```
Corresponding plot for shuffled data from one neuron (plotting all of the shuffled data this way gives too many points to be meaningful).
```{r}
plot(spatial_firing$shuffle_results_b_o_slope[[1]],
     spatial_firing$shuffle_results_b_h_slope[[1]])
```




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

## Plot population rates 

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


Now we want to plot population rate across whole track for diff groups so we can visualise the average firing rate

groups are as follows :

outbound homebound  reset
    +       +         n
    +       +         y
    +       -         -
    +      non        -
    -       +         -
    -       -         n
    -       -         y
    -      non        -
    
    

1. make tibble with average firing rates and classifications : 
```{r}
bin <- 200
df <- tibble(session_id = rep(spatial_firing$session_id, each=bin),
             cluster = rep(as.character(spatial_firing$cluster_id), each=bin),
             Position = rep(-30:169, times=nrow(spatial_firing)), 
             Rates = unlist(spatial_firing$normalised_rates_smoothed), 
             Outbound_beaconed = rep(spatial_firing$lm_group_b, each=bin), 
             Homebound_beaconed = rep(spatial_firing$lm_group_b_h, each=bin))

```



Subset data by group then average rates for plotting

**Negative Negative**
```{r}
(plot_neg_neg <- df %>%
  subset(Outbound_beaconed == "Negative" & Homebound_beaconed == "Negative") %>%
  group_by(Position) %>% 
  mean_SEM_plots_prep %>%
  mean_SEM_plots(colour1 = "black"))

if (save_figures == 1) {
 ggsave(file = "plots/negneg_mean_Hz.png",  width = 3.6, height = 2.9) 
}
```




**Positive Positive **
```{r}
(plot_pos_pos <- df %>%
  subset(Outbound_beaconed == "Positive" & Homebound_beaconed == "Positive") %>%
  group_by(Position) %>%
  mean_SEM_plots_prep %>%
  mean_SEM_plots(colour1 = "black"))


if (save_figures == 1) {
  ggsave(file = "plots/pospos_mean_Hz.png", width = 3.6, height = 2.9) 
}
```


**Positive Negative**
```{r}
(plot_pos_neg <- df %>%
  subset(Outbound_beaconed == "Positive" & Homebound_beaconed == "Negative") %>%
  group_by(Position) %>%
  mean_SEM_plots_prep %>%
  mean_SEM_plots(colour1 = "black"))


if (save_figures == 1) {
  ggsave(file = "plots/posneg_mean.png", width = 3.6, height = 2.9)
}
```



**Negative Positive**
```{r}
(plot_neg_pos <- df %>%
  subset(Outbound_beaconed == "Negative" & Homebound_beaconed == "Positive") %>%
  group_by(Position) %>%
  mean_SEM_plots_prep %>%
  mean_SEM_plots(colour1 = "black"))


if (save_figures == 1) {
  ggsave(file = "plots/negpos_mean.png", width = 3.6, height = 2.9)
}
```


**Positive Unclassified**
```{r}
(plot_pos_uc <- df %>%
  subset(Outbound_beaconed == "Positive" & Homebound_beaconed == "Unclassified") %>%
  group_by(Position) %>%
  mean_SEM_plots_prep %>%
  mean_SEM_plots(colour1 = "black"))

if (save_figures == 1) {
  ggsave(file = "plots/posnon_mean.png", width = 3.6, height = 2.9)
}
```

**Negative Unclassified**
```{r}
(plot_neg_uc <- df %>%
  subset(Outbound_beaconed == "Negative" & Homebound_beaconed == "Unclassified") %>%
  group_by(Position)  %>%
  mean_SEM_plots_prep %>%
  mean_SEM_plots(colour1 = "black"))

if (save_figures == 1) {
  ggsave(file = "plots/negnon_mean.png", width = 3.6, height = 2.9)
}
```
# get numbers for the plots
```{r}
table(unlist(spatial_firing$track_category_numeric))
```


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


# Plot heat map of firing rate across location for all neurons

First, reorder the dataframe with ramps according to slope.
_For start ramps, steepest slope should be negative - thus will have the highest cluster id_
```{r}
pospos <-spatial_firing %>% filter(track_category_numeric == 1)
posneg <-spatial_firing %>% filter(track_category_numeric == 2)
posnon <-spatial_firing %>% filter(track_category_numeric == 3)
negneg <-spatial_firing %>% filter(track_category_numeric == 4)
negpos <-spatial_firing %>% filter(track_category_numeric == 5)
negnon <-spatial_firing %>% filter(track_category_numeric == 6)

pospos<-pospos[order(-rank(pospos$asr_b_o_rewarded_fit_slope)),decreasing = TRUE]
posneg<-posneg[order(-rank(posneg$asr_b_o_rewarded_fit_slope)),decreasing = TRUE]
posnon<-posnon[order(-rank(posnon$asr_b_o_rewarded_fit_slope)),decreasing = TRUE]
negneg<-negneg[order(negneg$asr_b_o_rewarded_fit_slope),decreasing = TRUE]
negpos<-negpos[order(negpos$asr_b_o_rewarded_fit_slope),decreasing = TRUE]
negnon<-negnon[order(negnon$asr_b_o_rewarded_fit_slope),decreasing = TRUE]

#start_ramps <- rbind(pospos,posneg,posnon,negneg,negpos,negnon)
start_ramps<-spatial_firing[order(unlist(spatial_firing$track_category_numeric),spatial_firing$asr_b_o_rewarded_fit_slope),]
start_ramps<-start_ramps %>% filter(track_category_numeric != 0)
start_ramp_number = nrow(start_ramps)
new_cluster_id = seq(from = 1, to = start_ramp_number, by = 1)
start_ramps <- cbind(start_ramps, new_cluster_id)
```

Then, scale firing rate for all neurons

1. make function to load rates and normalise
2. Run on dataframe 
```{r}
normalise_rates_outbound <- function(df){
  df <- tibble(Rates = unlist(df), Position = rep(1:200))
  df <- df %>%
    filter(Position >=30, Position <= 170)
  x <- normalit(df$Rates)
  return(x)
}

start_ramps <- start_ramps %>%
  mutate(normalised_rates_o = map(Rates_averaged_rewarded_smoothed_b, normalise_rates_outbound))
```

Add position to normalised rates for plotting
```{r}
# Might need to modify this so position covers range 31:171.
start_ramps <- start_ramps %>%
  mutate(normalised_rates_o = map(normalised_rates_o, add_position)) 

```

Extract columns (normalised rates) for plotting into a tibble
```{r}
concat_firing_start <- unnest(select(start_ramps, new_cluster_id, normalised_rates_o),
                              cols = c(normalised_rates_o))
```

Now the firing rates have been normalised and the data in the right format we want to make annotations for the heatmap

First, extract ramp score from the dataframe for annotating heatmap
_since its a list of three (outbound/homebound/all) we extract the first one_
```{r}
start_ramps <- start_ramps %>%
  mutate(start_ramp_score = map(ramp_score, ~.x[1]))
  
```

Put ramp scores alongisde lm results and brain region classifier in tibble for annotating heatmap
```{r, warning=FALSE}
ramp_result <- tibble(ramp_score = as.numeric(start_ramps$start_ramp_score))
brain_region <- tibble(region = as.character(start_ramps$brain_region))
lm_result <- tibble(result = as.character(start_ramps$lm_group_b))
lm_result_homebound <- tibble(result_homebound = as.character(start_ramps$lm_group_b_h))
track_result <- tibble(track_cat = as.numeric(start_ramps$track_category_numeric))
cluster_result <- tibble(result = as.character(start_ramps$new_cluster_id))
```

Now we can plot the heatmap with annotations using pheatmap
```{r}
library(viridis)
#convert data to wide format
wide_DF <- concat_firing_start %>% spread(Position, Rates)

#remove unused column
wide_DF <- subset(wide_DF, select=-c(new_cluster_id))

#rownames(wide_DF) <- seq(length=nrow(wide_DF))
# Generte data (modified the mydf slightly)
rownames(wide_DF) <- paste("neuron", 1:max(start_ramps$new_cluster_id), sep="_")

# data for annotation rows in seperate dataframe
mydf <- data.frame(row.names = paste("neuron", 1:max(start_ramps$new_cluster_id), sep="_"), region = brain_region, ramp_score=ramp_result, track_catagory=track_result)
#mydf <- data.frame(row.names = paste("neuron", 1:max(start_ramps$new_cluster_id), sep="_"), region = brain_region, track_catagory=track_result)

# change the color of annotation to what you want: (eg: "navy", "darkgreen")
Var1        <- c("violetred2", "black", "chartreuse3")
names(Var1) <- c("Negative", "Unclassified", "Positive")

Var2        <- c("coral2", "deepskyblue2", "blueviolet" , "grey29")
names(Var2) <- c("PS", "RH", "MEC", "V1")

#Var3        <- c("violetred2", "black", "chartreuse3")
#names(Var3) <- c("Negative", "Unclassified", "Positive")


anno_col <- list(region = Var2, ramp_score = brewer.pal(15,"RdBu"), track_cat = viridis(7))
#anno_col <- list(track_cat = viridis(7))
#anno_col <- list(region = Var2, track_cat = viridis(7))

myheatmap<-pheatmap(wide_DF,cluster_cols = F, cluster_rows = F, annotation_row = mydf, annotation_colors = anno_col, show_rownames = F, show_colnames = F )

```

Save the heatmap
```{r}
save_pheatmap_png <- function(x, filename, width=1300, height=2500, res = 250) {
  png(filename, width = width, height = height, res = res)
  grid::grid.newpage()
  grid::grid.draw(x$gtable)
  dev.off()
}

if (save_figures == 1) {
  save_pheatmap_png(myheatmap, "plots/my_heatmap_all_update.png")
}
```




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


# Plot ramp scores for high and low firing rate neurons
First, extract ramp score from the dataframe for annotating heatmap
_since its a list of three (outbound/homebound/all) we extract the first one_
```{r}
spatial_firing <- spatial_firing %>%
  mutate(start_ramp_score = map(ramp_score, ~.x[1]))
  
```


```{r}
position_neurons_all <- spatial_firing  %>% 
  filter(lm_group_b == "Positive" | lm_group_b == "Negative") 

```

2. Plot scatter plot
```{r}

ggplot() + 
    geom_point(data=spatial_firing,
               aes(x = as.numeric(unlist(mean_firing_rate)), 
                   y = as.numeric(unlist(start_ramp_score)),
                   color=as.factor(unlist(lm_group_b))), alpha=0.8) +
    ylab("Ramp score (0 - 60 cm)") +
    xlab("Mean firing rate (Hz)") +
    theme_classic() +
    scale_color_manual(values=c( "violetred2", "chartreuse3", "grey32")) +
    theme(axis.text.x = element_text(size=18),
          axis.text.y = element_text(size=18),
          legend.position="bottom", 
          legend.title = element_blank(),
          text = element_text(size=17), 
          legend.text=element_text(size=16), 
          axis.title.y = element_text(margin = margin(t = 0, r = 20, b = 0, l = 0))) 

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

2. Plot scatter plot
```{r}

ggplot() + 
    geom_point(data=spatial_firing  %>% filter(lm_group_b == "Unclassified"),
               aes(x = as.numeric(unlist(mean_firing_rate)), 
                   y = as.numeric(unlist(asr_b_o_rewarded_fit_slope)),
                   color="grey32"), alpha=0.8) +
    geom_point(data=subset(position_neurons_all, start_ramp_score > 0),
               aes(x = as.numeric(unlist(mean_firing_rate)), 
                   y = as.numeric(unlist(start_ramp_score)),
                   color=as.factor(unlist(lm_group_b))), alpha=0.8) +
    #coord_cartesian(ylim = c(-.45,.61), xlim = c(-.45,.45)) +
    ylab("Ramp score (0 - 60 cm)") +
    xlab("Mean firing rate (Hz)") +
    theme_classic() +
    scale_color_manual(values=c("grey32", "chartreuse3","chartreuse3")) +
    theme(axis.text.x = element_text(size=18),
          axis.text.y = element_text(size=18),
          legend.position="bottom", 
          legend.title = element_blank(),
          text = element_text(size=17), 
          legend.text=element_text(size=16), 
          axis.title.y = element_text(margin = margin(t = 0, r = 20, b = 0, l = 0))) 

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


```{r}
position_neurons_all <- spatial_firing  %>% 
  filter(lm_group_b == "Positive" | lm_group_b == "Negative") 

```

2. Plot scatter plot
```{r}

ggplot() + 
    geom_point(data=spatial_firing,
               aes(x = as.numeric(unlist(mean_firing_rate)), 
                   y = as.numeric(unlist(asr_b_o_rewarded_fit_slope)),
                   color=factor(unlist(lm_group_b))), alpha=0.8) +
    #coord_cartesian(ylim = c(-.45,.61), xlim = c(-.45,.45)) +
    #geom_abline(intercept = 0, slope = 1, colour = "grey", linetype = "dashed") +
    #geom_abline(intercept = 0, slope = -1, colour = "grey", linetype = "dashed") +
    ylab("Slope (0 - 60 cm)") +
    xlab("Mean firing rate (Hz)") +
    theme_classic() +
    scale_color_manual(values=c("violetred2", "chartreuse3", "grey32")) +
    theme(axis.text.x = element_text(size=18),
          axis.text.y = element_text(size=18),
          legend.position="bottom", 
          legend.title = element_blank(),
          text = element_text(size=17), 
          legend.text=element_text(size=16), 
          axis.title.y = element_text(margin = margin(t = 0, r = 20, b = 0, l = 0))) 

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

2. Plot scatter plot
```{r}

ggplot() + 
    geom_point(data=spatial_firing  %>% filter(lm_group_b == "Unclassified"),
               aes(x = as.numeric(unlist(mean_firing_rate)), 
                   y = as.numeric(unlist(asr_b_o_rewarded_fit_slope)),
                   color=factor(unlist(lm_group_b))), alpha=0.8) +
    geom_point(data=subset(spatial_firing  %>% filter(lm_group_b == "Positive" | lm_group_b == "Negative", lm_group_b == "Positive")),
               aes(x = as.numeric(unlist(mean_firing_rate)), 
                   y = as.numeric(unlist(asr_b_o_rewarded_fit_slope)),
                   color=factor(unlist(lm_group_b))), alpha=0.8) +
    ylab("Slope (0 - 60 cm)") +
    xlab("Mean firing rate (Hz)") +
    theme_classic() +
    scale_color_manual(values=c( "chartreuse3", "grey32")) +
    theme(axis.text.x = element_text(size=18),
          axis.text.y = element_text(size=18),
          legend.position="bottom", 
          legend.title = element_blank(),
          text = element_text(size=17), 
          legend.text=element_text(size=16), 
          axis.title.y = element_text(margin = margin(t = 0, r = 20, b = 0, l = 0))) 

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



## Find average and STD number of rewarded sessions in the data

1. select columns we want
2. find distinct columns (this is because multiple days have many cells, an we only want unique days so the mean is accurate)
3. find mean number of rewarded trials and std of rewarded trials

```{r}
spatial_firing %>%
  select(Day, Mouse, cohort, max_trial_number, number_of_rewards) %>%
  distinct() %>%
  dplyr::summarise(mean_r = mean(as.numeric(number_of_rewards), na.rm =TRUE), sem_r = std.error(as.numeric(number_of_rewards))) 

```