---
title: "Supplemental_Figure_4_and_5"
author: "Meredith Hanlon"
date: "2023-04-13"
output: html_document
---

```{r global_options, echo = FALSE}
knitr::opts_chunk$set(warning=FALSE, message=FALSE, echo=FALSE, error = TRUE)
```

```{r}
library(broom)
library(tidyverse)
library(rstatix)
library(ggpubr)
library(readxl)
library(stats)
library(viridis)
library(patchwork)

stderror <- function(x) sd(x, na.rm = TRUE) / sqrt(length(na.omit(x)))

theme_mh <- function(){ 
    font <- "sans"   #assign font family up front
    theme(
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      panel.background = element_blank(),
      panel.border = element_rect(color = "black", fill = NA),
      axis.ticks = element_line(color = "black"),   
      axis.line = element_line(color = "black"),
      strip.background = element_blank(),
      legend.key = element_rect(fill = NA),
      text = element_text(family = font, size = 10, color = "black"),
      axis.text = element_text(family = font, size = 8, color = "black")
    )
}
```

```{r}
x4 <- read_excel("Hanlon_2023_data.xlsx", sheet = "SFigure_4_1")
lx <- read_excel("Hanlon_2023_data.xlsx", sheet = "SFigure_4_2")

is <- read_excel("Hanlon_2023_data.xlsx", sheet = "SFigure_5_1")
irl <- read_excel("Hanlon_2023_data.xlsx", sheet = "SFigure_5_2")
ta  <- read_excel("Hanlon_2023_data.xlsx", sheet = "SFigure_5_3")
ss4 <-  read_excel("Hanlon_2023_data.xlsx", sheet = "SFigure_5_4")
```


```{r}
x5 <- x4 %>% pivot_wider(names_from = Element, values_from = value) %>% 
  select(week, Genotype, Depth, Mg = MgO, Al = Al2O3, Si = SiO2, P, S, Cl, K = K2O, Ca = Ca, Ti = Ti, Cr = Cr, Mn = Mn, Fe = Fe, Co = Co, Ni = Ni, Cu = Cu, Zn = Zn, Ga = Ga, As = As, Rb = Rb, Sr = Sr, Y = Y, Zr = Zr, Ba = Ba, Pb = Pb, Th = Th) %>% 
  pivot_longer(Mg:Th, names_to = "Element", values_to = "soil_XRF") %>% 
  dplyr::select(Week = week, Genotype, Depth, Element, soil_XRF) %>% 
  mutate(Genotype = case_when(Genotype == "85" ~ "85d",
                              Genotype == "99" ~ "99d",
                              TRUE ~ Genotype)) 
  


xrl <- lx %>%
  dplyr::select(Element, XRF, Genotype, Week, leaf, seg_length_seg_1:seg_length_seg_6) %>% 
  pivot_wider(names_from = Element, values_from = XRF) %>% 
  group_by(Genotype, Week, leaf) %>% 
  summarise(across(where(is.numeric), mean)) %>% 
  pivot_longer(cols = Al:Sn, names_to = "Element", values_to = "XRF") %>% 
  pivot_longer(cols = seg_length_seg_1:seg_length_seg_6, names_to = "RL_depth", values_to = "root_length") %>% 
  separate("RL_depth", into = c("a", "b", "c", "depth")) %>% 
  mutate(depth = as.numeric(depth)) %>% 
  dplyr::select(-a, -b, -c) %>% 
  mutate(Depth = depth * 10) %>% 
  dplyr::select(-depth) %>% 
  full_join(., x5 %>% dplyr::select(Week, Genotype, Element, soil_xrf = soil_XRF, Depth)) %>% 
  ungroup()
```


Build XRF correlation matrix

```{r}
rd <- xrl %>% 
  filter(Week %in% c("6w", "8w", "10w")) %>% 
  dplyr::select(Genotype, Week, root_length, Depth) %>% 
  mutate(depth_group = case_when(Depth %in% c(10, 20) ~ "shallow",
                                 Depth %in% c(30, 40) ~ "mid",
                                 TRUE ~ "deep")) %>% 
  dplyr::select(-Depth) %>% 
  group_by(Genotype, Week, depth_group) %>% 
  summarize(roots = sum(root_length, na.rm = TRUE)) %>% 
  ungroup() 
rd2 <- rd %>% 
  group_by(Genotype, Week) %>% 
  summarize(roots = sum(roots)) %>% 
  mutate(depth_group = "total") %>% 
  bind_rows(rd)


leaves <- xrl %>% 
  filter(Week %in% c("6w", "8w", "10w")) %>% 
  dplyr::select(Genotype, Week, leaf, Element, XRF) %>% 
  group_by(Genotype, Week, Element) %>% 
  summarize(XRF = sum(XRF)) %>% 
  mutate(leaf = "all") %>% 
  bind_rows(xrl %>% dplyr::select(Genotype, Week, leaf, Element, XRF))


soilx <- xrl %>% 
  filter(Week %in% c("6w", "8w", "10w")) %>% 
  dplyr::select(Genotype, Week, soil_xrf, Depth) %>% 
  mutate(depth_group = case_when(Depth %in% c(10, 20) ~ "shallow",
                                 Depth %in% c(30, 40) ~ "mid",
                                 TRUE ~ "deep")) %>% 
  dplyr::select(-Depth) %>% 
  group_by(Genotype, Week, depth_group) %>% 
  summarize(soil_xrf = mean(soil_xrf, na.rm = TRUE)) %>% 
  ungroup() 
soilx2 <- soilx %>% 
  group_by(Genotype, Week) %>% 
  summarize(soil_xrf = mean(soil_xrf, na.rm = TRUE)) %>% 
  mutate(depth_group = "total") %>% 
  bind_rows(soilx)

xrf_cors <- full_join(rd2, soilx2) %>% 
  full_join(., leaves) %>% 
  drop_na() %>% 
  group_by(depth_group, Element, leaf) %>% 
  dplyr::summarize(cos = stats::cor.test(XRF, soil_xrf)$estimate,
                   p = stats::cor.test(XRF, soil_xrf)$p.value) %>% 
  mutate(r2 = round((cos)^2, 2)) %>% 
  dplyr::select(-cos) 
```


```{r}
xrf_cors$leaf <- factor(xrf_cors$leaf,
                        levels = c("all", "old", "ear", "young"))
xrf_cors$depth_group <- factor(xrf_cors$depth_group,
                               levels = c("total", "shallow", "mid", "deep"))

xcp <- xrf_cors %>% 
  mutate(sig = case_when(p < 0.05 ~ "sig1",
                         p >= 0.05 & p < 0.1 ~ "sig2")) %>% 
  mutate(r3 = case_when(sig %in% c("sig1", "sig2") ~ r2)) %>%
  unite(leaf_by_depth, leaf, depth_group, remove = F) %>% 
  ggplot(aes(x = depth_group, y = Element, fill = r2)) +
  geom_tile(color = "white", linewidth = 0.1) +
  scale_fill_viridis(name = "r2", limits = c(0,1), expand = c(0,0), na.value = "white") +
  geom_text(aes(label = round(r3, 2), color = sig),  size = 3) +
  scale_color_manual(values = c("black", "white", NA), guide = "none") +
  theme_mh() +
  facet_wrap(leaf~., nrow =1) +
  scale_x_discrete(position = "top") +
  xlab(NULL) +
  theme(legend.position = "top",
        panel.spacing.y = unit(0, "lines"),
        legend.box.margin = margin(b = -10),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(angle = 90),
        axis.title.y = element_blank(),
        axis.title.x = element_blank(),
        axis.ticks.y = element_blank(),
        plot.margin = margin(l = 10, r = 10),
        legend.key.height = unit(2, "mm"),
        legend.title = element_text(margin = margin(b = 0, t = -2)),
        strip.placement = "outside") 
xcp

ggsave("Supplemental_Figure_4A.png", width = 7, height = 4)
```

Add root length in to the model

```{r}
xls <- full_join(rd2, soilx2) %>% 
  full_join(., leaves) %>% 
  drop_na() %>% 
  group_by(depth_group, Element, leaf) %>% 
  group_modify(~ glance(lm(XRF ~ soil_xrf + roots, data = .))) %>% 
  dplyr::select(Element, leaf, depth_group, r2 = adj.r.squared, p = p.value)  %>% 
  mutate(r2 = round(r2, 2))
```
```{r}
xls$leaf <- factor(xls$leaf,
                        levels = c("all", "old", "ear", "young"))
xls$depth_group <- factor(xls$depth_group,
                               levels = c("total", "shallow", "mid", "deep"))

xcp <- xls %>% 
  mutate(sig = case_when(p < 0.05 ~ "sig1",
                         p >= 0.05 & p < 0.1 ~ "sig2")) %>% 
  mutate(r3 = case_when(sig %in% c("sig1", "sig2") ~ r2)) %>%
  unite(leaf_by_depth, leaf, depth_group, remove = F) %>% 
  ggplot(aes(x = depth_group, y = Element, fill = r2)) +
  geom_tile(color = "white", linewidth = 0.1) +
  scale_fill_viridis(name = "r2", limits = c(0,1), expand = c(0,0), na.value = "white") +
  geom_text(aes(label = round(r3, 2), color = sig),  size = 3) +
  scale_color_manual(values = c("black", "white", NA), guide = "none") +
  theme_mh() +
  facet_wrap(leaf~., nrow =1) +
  scale_x_discrete(position = "top") +
  xlab(NULL) +
  theme(legend.position = "top",
        panel.spacing.y = unit(0, "lines"),
        legend.box.margin = margin(b = -10),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(angle = 90),
        axis.title.y = element_blank(),
        axis.title.x = element_blank(),
        axis.ticks.y = element_blank(),
        plot.margin = margin(l = 10, r = 10),
        legend.key.height = unit(2, "mm"),
        legend.title = element_text(margin = margin(b = 0, t = -2)),
        strip.placement = "outside") 
xcp

ggsave("Supplemental_Figure_4B.png", width = 7, height = 4)
```




Build ICP correlations

```{r}
is2 <- is %>% 
  group_by(Week, Genotype, Rep) %>% 
  mutate(leaf = case_when(leaf_no == min(leaf_no) ~ "old",
                          leaf_no == max(leaf_no) ~ "young",
                          TRUE ~ "ear")) %>% 
  left_join(., ta )%>% 
  # remove Cr and Pb which aren't measured well on ICP/below LOD
  filter(!Element %in% c("Cr", "Pb"))

irl <- is2 %>% 
  ungroup() %>% 
  dplyr::select(Element, ICP, Genotype, Week, repid, leaf, seg_length_seg_1:seg_length_seg_6) %>% 
  pivot_wider(names_from = Element, values_from = ICP) %>% 
  group_by(Genotype, Week, leaf) %>% 
  summarise(across(where(is.numeric), mean)) %>% 
  pivot_longer(cols = Al:Zn, names_to = "Element", values_to = "ICP") %>% 
  pivot_longer(cols = seg_length_seg_1:seg_length_seg_6, names_to = "RL_depth", values_to = "root_length") %>% 
  separate("RL_depth", into = c("a", "b", "c", "depth")) %>% 
  mutate(depth = as.numeric(depth)) %>% 
  dplyr::select(-a, -b, -c) %>% 
  mutate(Depth = depth * 10) %>% 
  dplyr::select(-depth) %>% 
  full_join(., ss4 %>% dplyr::select(Week, Genotype, Element = Elem, soil_icp = ppm, Depth)) %>% 
  ungroup() %>% 
  # remove Cr and Pb which aren't measured well on ICP/below LOD
  filter(!Element %in% c("Cr", "Pb"))
```


```{r}
rd <- irl %>% 
  filter(Week %in% c("6w", "8w", "10w")) %>% 
  dplyr::select(Genotype, Week, root_length, Depth) %>% 
  mutate(depth_group = case_when(Depth %in% c(10, 20) ~ "shallow",
                                 Depth %in% c(30, 40) ~ "mid",
                                 TRUE ~ "deep")) %>% 
  dplyr::select(-Depth) %>% 
  group_by(Genotype, Week, depth_group) %>% 
  summarize(roots = sum(root_length, na.rm = TRUE)) %>% 
  ungroup() 
rd2 <- rd %>% 
  group_by(Genotype, Week) %>% 
  summarize(roots = sum(roots)) %>% 
  mutate(depth_group = "total") %>% 
  bind_rows(rd)


leaves <- irl %>% 
  filter(Week %in% c("6w", "8w", "10w")) %>% 
  dplyr::select(Genotype, Week, leaf, Element, ICP) %>% 
  group_by(Genotype, Week, Element) %>% 
  summarize(ICP = sum(ICP)) %>% 
  mutate(leaf = "all") %>% 
  bind_rows(irl %>% dplyr::select(Genotype, Week, leaf, Element, ICP))


soilx <- irl %>% 
  filter(Week %in% c("6w", "8w", "10w")) %>% 
  dplyr::select(Genotype, Week, soil_icp, Depth) %>% 
  mutate(depth_group = case_when(Depth %in% c(10, 20) ~ "shallow",
                                 Depth %in% c(30, 40) ~ "mid",
                                 TRUE ~ "deep")) %>% 
  dplyr::select(-Depth) %>% 
  group_by(Genotype, Week, depth_group) %>% 
  summarize(soil_icp = mean(soil_icp, na.rm = TRUE)) %>% 
  ungroup() 
soilx2 <- soilx %>% 
  group_by(Genotype, Week) %>% 
  summarize(soil_icp = mean(soil_icp, na.rm = TRUE)) %>% 
  mutate(depth_group = "total") %>% 
  bind_rows(soilx)

icp_cors <- full_join(rd2, soilx2) %>% 
  full_join(., leaves) %>% 
  drop_na() %>% 
  group_by(depth_group, Element, leaf) %>% 
  dplyr::summarize(cos = stats::cor.test(ICP, soil_icp)$estimate,
                   p = stats::cor.test(ICP, soil_icp)$p.value) %>% 
  mutate(r2 = round((cos)^2, 2)) %>% 
  dplyr::select(-cos) 
```


```{r}
icp_cors$leaf <- factor(icp_cors$leaf,
                        levels = c("all", "old", "ear", "young"))
icp_cors$depth_group <- factor(icp_cors$depth_group,
                               levels = c("total", "shallow", "mid", "deep"))

xcp <- icp_cors %>% 
  mutate(sig = case_when(p < 0.05 ~ "sig1",
                         p >= 0.05 & p < 0.1 ~ "sig2")) %>% 
  mutate(r3 = case_when(sig %in% c("sig1", "sig2") ~ r2)) %>%
  unite(leaf_by_depth, leaf, depth_group, remove = F) %>% 
  ggplot(aes(x = depth_group, y = Element, fill = r2)) +
  geom_tile(color = "white", linewidth = 0.1) +
  scale_fill_viridis(name = "r2", limits = c(0,1), expand = c(0,0), na.value = "white") +
  geom_text(aes(label = round(r3, 2), color = sig),  size = 3) +
  scale_color_manual(values = c("black", "white", NA), guide = "none") +
  theme_mh() +
  facet_wrap(leaf~., nrow =1) +
  scale_x_discrete(position = "top") +
  xlab(NULL) +
  theme(legend.position = "top",
        panel.spacing.y = unit(0, "lines"),
        legend.box.margin = margin(b = -10),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(angle = 90),
        axis.title.y = element_blank(),
        axis.title.x = element_blank(),
        axis.ticks.y = element_blank(),
        plot.margin = margin(l = 10, r = 10),
        legend.key.height = unit(2, "mm"),
        legend.title = element_text(margin = margin(b = 0, t = -2)),
        strip.placement = "outside") 
xcp

ggsave("Supplemental_Figure_5A.png", width = 7, height = 4)
```

Add root length in to the model

```{r}
xls <- full_join(rd2, soilx2) %>% 
  full_join(., leaves) %>% 
  drop_na() %>% 
  group_by(depth_group, Element, leaf) %>% 
  group_modify(~ glance(lm(ICP ~ soil_icp + roots, data = .))) %>% 
  dplyr::select(Element, leaf, depth_group, r2 = adj.r.squared, p = p.value)  %>% 
  mutate(r2 = round(r2, 2))
```
```{r}
xls$leaf <- factor(xls$leaf,
                        levels = c("all", "old", "ear", "young"))
xls$depth_group <- factor(xls$depth_group,
                               levels = c("total", "shallow", "mid", "deep"))

xcp <- xls %>% 
  mutate(sig = case_when(p < 0.05 ~ "sig1",
                         p >= 0.05 & p < 0.1 ~ "sig2")) %>% 
  mutate(r3 = case_when(sig %in% c("sig1", "sig2") ~ r2)) %>%
  unite(leaf_by_depth, leaf, depth_group, remove = F) %>% 
  ggplot(aes(x = depth_group, y = Element, fill = r2)) +
  geom_tile(color = "white", linewidth = 0.1) +
  scale_fill_viridis(name = "r2", limits = c(0,1), expand = c(0,0), na.value = "white") +
  geom_text(aes(label = round(r3, 2), color = sig),  size = 3) +
  scale_color_manual(values = c("black", "white", NA), guide = "none") +
  theme_mh() +
  facet_wrap(leaf~., nrow =1) +
  scale_x_discrete(position = "top") +
  xlab(NULL) +
  theme(legend.position = "top",
        panel.spacing.y = unit(0, "lines"),
        legend.box.margin = margin(b = -10),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(angle = 90),
        axis.title.y = element_blank(),
        axis.title.x = element_blank(),
        axis.ticks.y = element_blank(),
        plot.margin = margin(l = 10, r = 10),
        legend.key.height = unit(2, "mm"),
        legend.title = element_text(margin = margin(b = 0, t = -2)),
        strip.placement = "outside") 
xcp

ggsave("Supplemental_Figure_5B.png", width = 7, height = 4)
```