Figure3_Analysis.Rmd

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

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



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

## Differences in the positional dependance of firing between rewarded trials, run through trials and trials where the animal "tries" 

trial types :
1. rewarded/hit = animal recieves a reward
2. try = speed in the reward zone is within the 95% confidence interval of speed in the reward zone in
rewarded trials
3. run through = speed in the reward zone is outside the 95% confidence interval of speed in the reward zone in rewarded trials



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


Now we want to examine if there are differences in the reset activity between trial types (hit, try, run). To do this, we can plot population rate for rewarded and failed trials across whole track. 

Function to join firing rates from different trial types and add indicator of the type of trial

First combine hit and miss trials together with an indicator of trial type (hit, try, run)



1. Generate firing tibble for each cell containing firing rates for all three trial outcomes. The rates for each outcome are combined in to a single data frame in tidy format and then normalized.
```{r}
spatial_firing <- spatial_firing %>%
  mutate(avg_both_asr_b = pmap(list(Avg_FiringRate_HitTrials_smoothed,
                                    Avg_FiringRate_RunTrials_smoothed,
                                    Avg_FiringRate_TryTrials_smoothed,
                                    session_id,
                                    cluster_id),
                               join_average_rates))
```





_plots for conjunctive position encoding neurons (PA, PS, PSA)_

Look at all position encoding neurons (P,PS, PA, PSA)
```{r}
position_all_neurons <- 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")
# Total number of cells (not necessarily all used for plots as not all neurons have all outcomes)
# Next to do is to remove these neurons
dim(position_all_neurons)[1]

# Cells with data for all outcomes
sum(sapply(position_all_neurons$avg_both_asr_b, anyNA))

not_anyNA <- function(df) {
  !anyNA(df)
}
position_all_neurons <- position_all_neurons[sapply(position_all_neurons$avg_both_asr_b, not_anyNA),]

# Cells with data for all outcomes
sum(sapply(position_all_neurons$avg_both_asr_b, anyNA))
```
  
```{r}
position_all_neurons_plots <- all_plots_by_outome(position_all_neurons)
plot_grid(position_all_neurons_plots[[1]][[1]],
          position_all_neurons_plots[[1]][[2]],
          position_all_neurons_plots[[1]][[3]],
          position_all_neurons_plots[[1]][[4]])

as_vector(position_all_neurons_plots[[2]])

if (save_figures==1) {
 ggsave(file = "plots/TrialOutcome_NegNeg_position_all.png",
        plot = position_all_neurons_plots[[1]][[1]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_NegPos_position_all.png",
         plot = position_all_neurons_plots[[1]][[2]], width = 3.6, height = 2.9)
  ggsave(file = "plots/TrialOutcome_PosPos_position_all.png",
         plot = position_all_neurons_plots[[1]][[3]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_PosNeg_position_all.png",
         plot = position_all_neurons_plots[[1]][[4]], width = 3.6, height = 2.9) 
}
```

Split data based on conjunctive position encoding (PA, PS, PSA) (Figure 2):
```{r}
conjunctive_neurons <- filter(spatial_firing, lm_group_b == "Positive" | lm_group_b == "Negative", final_model_o_b == "PS" | final_model_o_b == "PA" | final_model_o_b == "PSA")
# Total number of cells
dim(conjunctive_neurons)[1]
# Cells with data for all outcomes
sum(sapply(conjunctive_neurons$avg_both_asr_b, anyNA))
```

```{r}
conjunctive_neurons_plots <- all_plots_by_outome(conjunctive_neurons)
plot_grid(conjunctive_neurons_plots[[1]][[1]],
          conjunctive_neurons_plots[[1]][[2]],
          conjunctive_neurons_plots[[1]][[3]],
          conjunctive_neurons_plots[[1]][[4]])

as_vector(conjunctive_neurons_plots[[2]])

if (save_figures==1) {
 ggsave(file = "plots/TrialOutcome_NegNeg_conjunctive.png",
        plot = conjunctive_neurons_plots[[1]][[1]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_NegPos_conjunctive.png",
         plot = conjunctive_neurons_plots[[1]][[2]], width = 3.6, height = 2.9)
  ggsave(file = "plots/TrialOutcome_PosPos_conjunctive.png",
         plot = conjunctive_neurons_plots[[1]][[3]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_PosNeg_conjunctive.png",
         plot = conjunctive_neurons_plots[[1]][[4]], width = 3.6, height = 2.9) 
}
```

_plots for pure position encoding neurons (P)_

Repeat above for just position neurons

1. Split data based position encoding group (P) (Figure 2) as we are only interested in position neurons here

```{r}
position_neurons <- filter(spatial_firing, lm_group_b == "Positive" | lm_group_b == "Negative", final_model_o_b == "P")
# Total number of cells
dim(position_neurons)[1]
# Cells with data for all outcomes
sum(sapply(position_neurons$avg_both_asr_b, anyNA))
```

```{r}
position_neurons_plots <- all_plots_by_outome(position_neurons)
plot_grid(position_neurons_plots[[1]][[1]],
          position_neurons_plots[[1]][[2]],
          position_neurons_plots[[1]][[3]],
          position_neurons_plots[[1]][[4]])

as_vector(position_neurons_plots[[2]])

if (save_figures==1) {
 ggsave(file = "plots/TrialOutcome_NegNeg_justposition.png",
        plot = position_neurons_plots[[1]][[1]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_NegPos_justposition.png",
         plot = position_neurons_plots[[1]][[2]], width = 3.6, height = 2.9)
  ggsave(file = "plots/TrialOutcome_PosPos_justposition.png",
         plot = position_neurons_plots[[1]][[3]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_PosNeg_justposition.png",
         plot = position_neurons_plots[[1]][[4]], width = 3.6, height = 2.9) 
}
```




_plots for speed encoding neurons (S) with POSITIVE or NEGATIVE slopes_
Repeat above plots but for speed encoding neurons

Split data based on speed encoding group (Figure 2) as we are only interested in these neurons here
Need to make this selective for neurons with positive or negative slopes?
```{r}
speed_neurons <- filter(spatial_firing, lm_group_b == "Negative", final_model_o_b == "S" | final_model_o_b == "SA")
# Total number of cells
dim(speed_neurons)[1]
# Cells with data for all outcomes
sum(sapply(speed_neurons$avg_both_asr_b, anyNA))
```

```{r}
speed_neurons_plots <- all_plots_by_outome(speed_neurons)
plot_grid(speed_neurons_plots[[1]][[1]],
          speed_neurons_plots[[1]][[2]],
          speed_neurons_plots[[1]][[3]],
          speed_neurons_plots[[1]][[4]])

as_vector(speed_neurons_plots[[2]])

if (save_figures==1) {
 ggsave(file = "plots/TrialOutcome_NegNeg_speed.png", plot = speed_neurons_plots[[1]][[1]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_NegPos_speed.png", plot = speed_neurons_plots[[1]][[2]], width = 3.6, height = 2.9)
  ggsave(file = "plots/TrialOutcome_PosPos_speed.png", plot = speed_neurons_plots[[1]][[3]], width = 3.6, height = 2.9) 
  ggsave(file = "plots/TrialOutcome_PosNeg_speed.png", plot = speed_neurons_plots[[1]][[4]], width = 3.6, height = 2.9) 
}
```







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

Examine slopes and firing rate offsets 

These analyses start with data in the following column of spatial firing:
Avg_FiringRate_TryTrials
Avg_FiringRate_RunTrials
Avg_FiringRate_HitTrials


We also want to fit a simple linear model for each trial outcome separately so we can assess the slope and reward zone offset for each neuron.

1. Add position to data for try and run through trials (we already have coefficients for rewarded data from Figure 1). Function is the same as used in Figure 1.
```{r}
spatial_firing <- spatial_firing %>%
  mutate(asr_b_hit = pmap(list(Avg_FiringRate_HitTrials), add_position),
         asr_b_try = pmap(list(Avg_FiringRate_TryTrials), add_position),
         asr_b_run = pmap(list(Avg_FiringRate_RunTrials), add_position))
```

2. Fit lm to data from all cells.
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_hit_fit_')) %>%
  select(-contains('asr_b_try_fit_')) %>%
  select(-contains('asr_b_run_fit_')) %>%
  mutate(asr_b_hit_fit = pmap(list(asr_b_hit, 30, 90), lm_tidy_helper)) %>%
  mutate(asr_b_try_fit = pmap(list(asr_b_try, 30, 90), lm_tidy_helper)) %>%
  mutate(asr_b_run_fit = pmap(list(asr_b_run, 30, 90), lm_tidy_helper)) %>%
  unnest_wider(asr_b_hit_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_b_try_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_b_run_fit, names_sep = "_", names_repair = "universal")

spatial_firing <- spatial_firing %>%
  select(-contains('asr_b_h_hit_fit_')) %>%
  select(-contains('asr_b_h_try_fit_')) %>%
  select(-contains('asr_b_h_run_fit_')) %>%
  mutate(asr_b_h_hit_fit = pmap(list(asr_b_hit, 110, 170), lm_tidy_helper)) %>%
  mutate(asr_b_h_try_fit = pmap(list(asr_b_try, 110, 170), lm_tidy_helper)) %>%
  mutate(asr_b_h_run_fit = pmap(list(asr_b_run, 110, 170), lm_tidy_helper)) %>%
  unnest_wider(asr_b_h_hit_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_b_h_try_fit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(asr_b_h_run_fit, names_sep = "_", names_repair = "universal")

```

Linear model results are stored in:
spatial_firing$asr_b_try_fit_pval
spatial_firing$asr_b_try_fit_slope
spatial_firing$asr_b_try_fit_r.squared


Trials with all try outcomes:
spatial_firing[!sapply(spatial_firing$Avg_FiringRate_TryTrials, anyNA),]
Trials with all run outcomes:
spatial_firing[!sapply(spatial_firing$Avg_FiringRate_RunTrials, anyNA),]

Trials with all run and try outcomes:
spatial_firing[!(sapply(spatial_firing$Avg_FiringRate_TryTrials, anyNA) | sapply(spatial_firing$Avg_FiringRate_RunTrials, anyNA)),]

To check consistency with smoothed data:
spatial_firing[!sapply(spatial_firing$Avg_FiringRate_TryTrials, anyNA)
== sapply(spatial_firing$Avg_FiringRate_TryTrials_smoothed, anyNA),]


## Compare slopes before the reward zone as a function of trial outcome.
⎄ 
```{r}
spatial_firing %>%
  filter(lm_group_b == "Positive", final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PSA" | final_model_o_b == "PA") %>%
  slopes_by_outcome(-0.4, 0.8)

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

spatial_firing %>%
  filter(lm_group_b == "Positive", final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PSA" | final_model_o_b == "PA") %>%
  h_slopes_by_outcome(-0.8, 0.8)

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

spatial_firing %>%
  filter(lm_group_b == "Negative", final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PSA" | final_model_o_b == "PA") %>%
  slopes_by_outcome(-2.5, 0.5)


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

spatial_firing %>%
  filter(lm_group_b == "Negative", final_model_o_b == "P" | final_model_o_b == "PS" | final_model_o_b == "PSA" | final_model_o_b == "PA") %>%
  h_slopes_by_outcome(-0.8, 0.8)


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

```

```{r}
one.way <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Positive") %>%
  slopes_by_outcome_aov
summary(one.way)

spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Positive") %>%
  slopes_by_outcome_t %>%
  mutate(ttest_adj = p.adjust(ttest, "bonferroni"))
  
one.way <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Negative") %>%
  slopes_by_outcome_aov
summary(one.way)

spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Negative") %>%
  slopes_by_outcome_t %>%
  mutate(ttest_adj = p.adjust(ttest, "bonferroni"))


# homebound comparison
one.way <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Positive") %>%
  h_slopes_by_outcome_aov
summary(one.way)


spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Positive") %>%
  h_slopes_by_outcome_t %>%
  mutate(ttest_adj = p.adjust(ttest, "bonferroni"))
  
one.way <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Negative") %>%
  h_slopes_by_outcome_aov
summary(one.way)

spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Negative") %>%
  h_slopes_by_outcome_t %>%
  mutate(ttest_adj = p.adjust(ttest, "bonferroni"))


```



## Calculate offset for run and try trials

We want to find out if firing rate reset depends on trial outcome.

These analyses use the following columns od sptial_firing as a starting point:
Avg_FiringRate_TryTrials
Avg_FiringRate_RunTrials
Avg_FiringRate_HitTrials

As in Figure 1 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. 

First, normalise firing rates. 
```{r}
spatial_firing <- spatial_firing %>%
  mutate(normalised_rates_try = map(Avg_FiringRate_TryTrials, normalise_rates), 
         normalised_rates_run = map(Avg_FiringRate_RunTrials, normalise_rates),
         normalised_rates_hit = map(Avg_FiringRate_HitTrials, normalise_rates))
```

Then, predict firing rate in homebound region based on fit from real data in outbound region Predict mean and confidence intervals for firing rate at the start of the homebound zone (track positions 110 to 115 cm) based on firing in the outbound zone (30 to 90 cm). The function predict_homebound was used previously for analyses for Figure 1.
```{r}
spatial_firing <- spatial_firing %>%
  select(-contains('predict_params_hit_')) %>%
  select(-contains('predict_params_try_')) %>%
  select(-contains('predict_params_run_')) %>%
  mutate(predict_params_hit = map(normalised_rates_hit, predict_homebound),
         predict_params_try = map(normalised_rates_try, predict_homebound),
         predict_params_run = map(normalised_rates_run, predict_homebound)) %>%
  unnest_wider(predict_params_hit, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(predict_params_try, names_sep = "_", names_repair = "universal") %>%
  unnest_wider(predict_params_run, names_sep = "_", names_repair = "universal") 

spatial_firing <- spatial_firing %>%
  mutate(
    predict_diff_hit = map2_dbl(normalised_rates_hit, predict_params_hit_fit, calc_predict_diff),
         predict_diff_try = map2_dbl(normalised_rates_try, predict_params_try_fit, calc_predict_diff),
         predict_diff_run = map2_dbl(normalised_rates_run, predict_params_run_fit, calc_predict_diff))

```




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

## plot violin plots of mean apsolute difference between predicted and real in the first 5 cm of the homebound region for hit, run and try trials. 

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



Make and save plots
```{r}
spatial_firing %>% filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                          lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  offsets_by_outcome(-4, 2)

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


spatial_firing %>% filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                          lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  offsets_by_outcome(-1.5, 4.5)

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

_run statistical tests_
Run anova to compare effect of trial type (hit, try, run) on offset values for ++ and -- neurons 
```{r}
one.way <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  offsets_by_outcome_aov
summary(one.way)

spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Positive" & lm_group_b_h == "Positive") %>%
  offsets_by_outcome_t %>%
  mutate(ttest_adj = p.adjust(ttest, "bonferroni"))
  
one.way <- spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  offsets_by_outcome_aov
summary(one.way)

spatial_firing %>%
  filter(final_model_o_b == "P" | final_model_o_b == "PA"  | final_model_o_b == "PS"  | final_model_o_b == "PSA",
                lm_group_b == "Negative" & lm_group_b_h == "Negative") %>%
  offsets_by_outcome_t %>%
  mutate(ttest_adj = p.adjust(ttest, "bonferroni"))
```