Figure3_ExtraAnalysis.Rmd

---
title: "Untitled"
output: html_document
---

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



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

## Analysis of average speed vs location 

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

## plot average speed vs location 

1. Write function to add position to speed
```{r}
add_position_withoutcome <- function(hit, try, run) {
  df <- tibble(Speed = c(unlist(hit), unlist(try), unlist(run)), 
               Position = rep(1:200, times=3), 
               Indicator = c(rep("Hit", times=200), rep("Try", times=200), rep("Run", times=200)))
  return(df)
}
```

2. Run on dataframe : Average across trial types

input columns: 
Speed_mean_rewarded = speed on rewarded beaconed trials
Speed_mean_try =  speed ontry beaconed trials
Speed_mean_runthru = speed on run through trials
```{r}
spatial_firing <- spatial_firing %>%
  mutate(Speed_avg = pmap(list(Speed_mean_rewarded, Speed_mean_try, Speed_mean_run), add_position_withoutcome))
```

Then average over selected neurons 

Function to plot data mean and SD of the population data based on the classification, as per Figure 2, on beaconed and non-beaconed trials.

```{r}
  # The function to use at each step is `mean`.
  # The window size is 5
  rolling_mean <- rollify(mean, window = 2)

```

```{r}
mean_speed_plots <- function(df, x_start = 30, x_end = 90){
 rolling_mean <- rollify(mean, window = 2)
 
 df <- df %>%
   select(Speed_avg) %>%
   unnest((Speed_avg)) %>%
   mutate(Speed = rolling_mean(Speed)/2)

 df <- df %>%
   group_by(Position, Indicator) %>%
   summarise(mean_b = mean(Speed, na.rm = TRUE),
             se_b = std.error(Speed, na.rm = TRUE))
 ggplot(data=df) +
  annotate("rect", xmin=0, xmax=30, ymin=0,ymax=Inf, alpha=0.2, fill="Grey60") +
  annotate("rect", xmin=170, xmax=200, ymin=0,ymax=Inf, alpha=0.2, fill="Grey60") +
  annotate("rect", xmin=90, xmax=110, ymin=0,ymax=Inf, alpha=0.2, fill="Chartreuse4") +
  geom_ribbon(aes(x=Position, 
                  y=mean_b, ymin = mean_b - se_b, ymax = mean_b + se_b, 
                  fill=as.factor(Indicator)), alpha=0.1) +
   geom_line(aes(y=mean_b, 
                 x=Position, 
                 color=as.factor(Indicator)), alpha=0.5, n= 100, span=0.2) +
   scale_fill_manual(values=c("black", "blue", "red")) +
   scale_color_manual(values=c("black", "blue", "red")) +
   theme_classic() +
   labs(y = "Mean speed (cm/s)", x = "Position") +
   xlim(x_start, x_end) +
   #ylim(10,170) +
   theme(axis.text.x = element_text(size=12),
        axis.text.y = element_text(size=12),
        legend.title = element_blank(),
        text = element_text(size=12))
 
}
```


Subset data by group then average rates for plotting          
```{r}
(Pos_speed_plot <- mean_speed_plots(spatial_firing,  0, 200))
if (save_figures == 1) {
  ggsave(file = "plots/PI_speed.png", width = 4.5, height = 2.5)
}
```



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

## calculate a histogram of speed in the reward zone

Here we want to see the distribution of speed in the reward zone in the trial types (hit, run, try)

Function to make the plots
```{r}
speed_hist <- function(df, fill_colour = 'black') {
  ggplot(data=df, aes(x = speed)) +
  coord_cartesian(xlim=c(0,100)) +
  geom_histogram(aes(y=..density..),binwidth = 2, alpha=0.5, fill=fill_colour) +
  #scale_y_continuous(breaks = scales::pretty_breaks(n = 4)) +
  labs(y="Density", x="") +
  theme_classic() +
  theme(axis.text.x = element_text(size=16),
        axis.text.y = element_text(size=16),
        text = element_text(size=16), 
        axis.title.y = element_text(margin = margin(t = 0, r = 20, b = 0, l = 0)))
}
```



```{r}
(df <- spatial_firing %>%
  select(rewardzone_speed_hit) %>%
  unnest(rewardzone_speed_hit) %>%
  mutate(speed = as.numeric(rewardzone_speed_hit)) %>%
  speed_hist(fill_colour = 'black'))

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


4. Plot histogram for rewarded trials

```{r}
(df <- spatial_firing %>%
  select(rewardzone_speed_try) %>%
  unnest(rewardzone_speed_try)  %>%
  mutate(speed = as.numeric(rewardzone_speed_try)) %>%
  speed_hist(fill_colour = 'red'))
if (save_figures == 1) {
  ggsave(file = "plots/speedhist_try.png", width = 4, height = 2.5)
}
```



```{r}
(df <- spatial_firing %>%
  select(rewardzone_speed_run) %>%
  unnest(rewardzone_speed_run) %>%
  mutate(speed = as.numeric(rewardzone_speed_run)) %>%
  speed_hist(fill_colour = 'blue'))
if(save_figures==1){
  ggsave(file = "plots/speedhist_run.png", width = 4, height = 2.5)
}
```


Plot examples of speed or acceleration used to fit models in Figure 3 as a function of position. This is just as a check for obvious errors.
Relevant columns are:
spatial_firing$spikes_in_time
Column 1: Rates
Column 2: Position
Column 3: Speed 
Column 4: Acceleration
```{r}
trial_data_as_tibble <- function(df){
    df <-
    tibble(
      Rates = as.numeric(Re(df[[1]])),
      Position = as.numeric(Re(df[[2]])),
      Acceleration = as.numeric(Re(df[[4]])),
      Speed = as.numeric(Re(df[[3]])),
      Trials = as.factor(df[[5]]),
      Types = as.factor(df[[6]])
    )
}

a <- spatial_firing[!duplicated(spatial_firing$session_id),] %>%
  select(session_id, spikes_in_time) %>%
  mutate(spikes_in_time = map(spikes_in_time, trial_data_as_tibble))



ggplot(a$spikes_in_time[[1]], aes(Position, Speed)) +
  geom_point() +
  geom_smooth()


ggplot(a$spikes_in_time[[2]], aes(Position, Acceleration)) +
  geom_point() +
  geom_smooth()

# look at first 12 sessions to save time.
b <- unnest(a[1:12,], spikes_in_time)

ggplot(b, aes(Position, Speed)) +
  geom_point() +
  geom_smooth() +
  facet_wrap(vars(session_id), scales = "free_y") +
  theme_minimal() 
  
ggplot(b, aes(Position, Acceleration)) +
  geom_point() +
  geom_smooth() +
  facet_wrap(vars(session_id), scales = "free_y") +
  theme_minimal() 

  

```