ch5_multilinear.qmd

---
format: html
execute:
  fig-align: center
  echo: false
  warning: false
  message: false
  cache: true
  cache.lazy: false

editor: visual
---

```{r setup, include = FALSE}
knitr::opts_chunk$set(fig.align = "center")

library(rethinking)
data(WaffleDivorce)
data(milk)
data(Howell1)
kungHeight <- Howell1
detach(package:rethinking, unload = T)


library(tidyverse)
library(brms)
library(ggrepel)
library(tigris)
library(bayesplot)
library(patchwork)
library(dagitty)
library(ggdag)
library(tidybayes)
```

## TLDR

\[claude\] come up with 3-4 bullet points on what was accomplished in this chapter.

## Motivation

\[claude\] come up with a distrilled paragraph argument for why we should care about muli-variate regression. Here's a section of Richard McElreath's text statistical Rethinking

(1) Statistical “control” for confounds. A confound is something that misleads us about a causal influence—there will be a more precise definition in the next chapter. The spurious waffles and divorce correlation is one possible type of confound, where the confound (southernness) makes a variable with no real importance (Waffle House density) appear to be important. But confounds are diverse. They can hide real important variables just as easily as they can produce false ones.
(2) Multiple causation. A phenomenon may arise from multiple causes. Measurement of each cause is useful, so when we can use the same data to estimate more than one type of influence, we should. Furthermore, when causation is multiple, one cause can hide another.
(3) Interactions. The importance of one variable may depend upon another. For ex- ample, plants benefit from both light and water. But in the absence of either, the other is no benefit at all. Such interactions occur very often. Effective inference about one variable will often depend upon consideration of others.

Claude give a quick description of the waffle divorce data set : Data for the individual States of the United States, describing number of Waffle House diners and various marriage and demographic facts.

Format Location : State name Loc : State abbreviation Population : 2010 population in millions MedianAgeMarriage: 2005-2010 median age at marriage Marriage : 2009 marriage rate per 1000 adults Marriage.SE : Standard error of rate Divorce : 2009 divorce rate per 1000 adults Divorce.SE : Standard error of rate WaffleHouses : Number of diners South : 1 indicates Southern State Slaves1860 : Number of slaves in 1860 census Population1860 : Population from 1860 census PropSlaves1860 : Proportion of total population that were slaves in 1860

```{r plot waffle house vs divorce}
#| fig-width: 5
#| fig-height: 4

WaffleDivorce %>%
  ggplot(aes(x = WaffleHouses/Population, y = Divorce)) +
  stat_smooth(method = "lm", fullrange = T, linewidth = 1/2,
              color = "firebrick4", fill = "firebrick", alpha = 1/5) +
  geom_point(size = 1.5, color = "firebrick4", alpha = 1/2) +
  geom_text_repel(data = WaffleDivorce %>% filter(Loc %in% c("ME", "OK", "AR", "AL", "GA", "SC", "NJ")),  
                  aes(label = Loc), 
                  size = 3, seed = 1042) +  # this makes it reproducible
  scale_x_continuous("Waffle Houses per million", limits = c(0, 55)) +
  ylab("Divorce rate") +
  coord_cartesian(xlim = c(0, 50), ylim = c(5, 15)) +
  theme_bw() +
  theme(panel.grid = element_blank())  
```

Most likely a spurious correlation, waffles don't cause divorces nor vice versa

## Spurious Correlation

Excerpt from Richard McElreath: "Let’s leave waffles behind, at least for the moment. An example that is easier to understand is the correlation between divorce rate and marriage rate (Figure 5.2). The rate at which adults marry is a great predictor of divorce rate, as seen in the left-hand plot in the figure. But does marriage cause divorce? In a trivial sense it obviously does: One cannot get a divorce without first getting married. But there’s no reason high marriage rate must be correlated with divorce."

```{r plot marriage/divorce on US map}
#| fig-width: 8
#| fig-height: 2

WaffleDivorce_states <- 
  WaffleDivorce %>% 
  mutate(Divorce_z           = (Divorce - mean(Divorce))                     / sd(Divorce),
         MedianAgeMarriage_z = (MedianAgeMarriage - mean(MedianAgeMarriage)) / sd(MedianAgeMarriage),
         Marriage_z          = (Marriage - mean(Marriage))                   / sd(Marriage)) %>% 
  right_join(readRDS("externalData/us_states_geometry.rds"), 
             by = c("Location" = "NAME")) %>%
  pivot_longer(ends_with("_z"))

WaffleDivorce_states %>% 
  ggplot() +
  geom_sf(aes(fill = value, geometry = geometry),
          size = 0) +
  scale_fill_gradient(low = "#f8eaea", high = "firebrick4") +
  coord_sf(xlim = c(-29e5, 21e5),
           ylim = c(-2e6, 2e6)) +
  theme_void() +
  theme(legend.position = "none",
        strip.text = element_text(margin = margin(0, 0, 0.5, 0))) +
  facet_wrap(~name) 

```

We see that Divorce rate is high in the south and low in the north midwest and mid atlantic. Marriage rate is high in

```{r DAG representations}
#| fig-width: 8
#| fig-height: 3

dag_coords <-
  tibble(name = c("A", "M", "D"),
         x    = c(1, 3, 2),
         y    = c(2, 2, 1))

m1 <- dagify("M" ~ "A",
       "D" ~ "M" + "A",
       coords = dag_coords) %>%
  ggplot(aes(x = x, y = y, xend = xend, yend = yend)) +
  geom_dag_point(color = "firebrick", alpha = 1/4, size = 10) +
  geom_dag_text(color = "firebrick") +
  geom_dag_edges(edge_color = "firebrick")+
  scale_x_continuous(NULL, breaks = NULL, expand = c(0.1, 0.1)) +
  scale_y_continuous(NULL, breaks = NULL, expand = c(0.2, 0.2)) +
  theme_bw() +
  theme(panel.grid = element_blank())+
  labs(title = "Mediated Effect through Marriage Rate")


m2 <- dagify("M" ~ "A",
       "D" ~ "A",
       coords = dag_coords) %>%
  ggplot(aes(x = x, y = y, xend = xend, yend = yend)) +
  geom_dag_point(color = "firebrick", alpha = 1/4, size = 10) +
  geom_dag_text(color = "firebrick") +
  geom_dag_edges(edge_color = "firebrick")+
  scale_x_continuous(NULL, breaks = NULL, expand = c(0.1, 0.1)) +
  scale_y_continuous(NULL, breaks = NULL, expand = c(0.2, 0.2)) +
  theme_bw() +
  theme(panel.grid = element_blank())+
  labs(title = "Only Direct Effect")

m1 + m2 + 
  plot_annotation(caption = "A = Marriage Age, M = Marriage Rate, D = Divorce Rate",
                  theme = theme(plot.caption = element_text(hjust = 0.5, size = 12)))
```

```{r data wrangling, include = FALSE}
WaffleDivorce <- WaffleDivorce %>%
  mutate(MedianAgeMarriage_s = (MedianAgeMarriage - mean(MedianAgeMarriage)) / sd(MedianAgeMarriage),
         Marriage_s  = (Marriage - mean(Marriage)) / sd(Marriage))
```

### Median Marriage Age Model

$$ \textbf{b5.1} $$

$$ \text{Divorce}_i \sim \text{Intercept} + \beta_1 \text{Median Age Marriage}_i $$

```{r model - median marriage age, include = FALSE}
b5.1 <- 
  brm(data = WaffleDivorce, 
      family = gaussian,
      Divorce ~ 1 + MedianAgeMarriage_s,
      prior = c(prior(normal(10, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.01")

at <- c(-3, -2, -1, 0, 1, 2, 3)

simModel <- as_tibble(as_draws_df(b5.1)) %>% 
  mutate(simMarriageAge = seq(from = -3, to = 3.5, length.out = n()),
         simDivorceEst = Intercept + (b_MedianAgeMarriage_s * simMarriageAge),
         simDivorce = rnorm(n(), simDivorceEst, sd = sigma))

```

::: panel-tabset
#### 📈 μ vs full distribution

```{r mu vs full heat plot median marriage age}
#| fig-width: 8
#| fig-height: 4

# Find global limits across both datasets
ymin <- min(c(simModel$simDivorce, simModel$simDivorceEst, WaffleDivorce$Divorce))
ymax <- max(c(simModel$simDivorce, simModel$simDivorceEst, WaffleDivorce$Divorce))

modelEst_plot <- ggplot() +
  stat_density_2d(data = simModel, 
                  aes(x = simMarriageAge, y = simDivorceEst, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = MedianAgeMarriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(y = "Divorce rate per 1000 adults",
       title = "Divorce Rate ~ Median Marriage Age", subtitle = "Mu estimate")+
  theme_minimal()+
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$MedianAgeMarriage) + 
                                      mean(WaffleDivorce$MedianAgeMarriage), 1)) +
  scale_y_continuous(limits = c(ymin, ymax), expand = c(0, 0))+
  guides(fill = "none")

model_plot <- ggplot() +
  stat_density_2d(data = simModel, 
                  aes(x = simMarriageAge, y = simDivorce, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = MedianAgeMarriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(y = "Divorce rate per 1000 adults",  subtitle = "Full Distribution")+
  theme_minimal()+
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$MedianAgeMarriage) + 
                                      mean(WaffleDivorce$MedianAgeMarriage), 1)) +
  scale_y_continuous(limits = c(ymin, ymax), expand = c(0, 0))+
  guides(fill = "none")

modelEst_plot + model_plot + plot_layout(guides = "collect")

```

#### 🎛 Parameters

```{r summary median marriage age}
#| fig-width: 10
#| fig-height: 2

as_tibble(b5.1) %>% 
  rename("Median Marriage Age" = b_MedianAgeMarriage_s) %>% 
  dplyr::select(c(`Median Marriage Age`, `Intercept`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Covariate",
               values_to = "Effect") %>% 
  ggplot(aes(x = `Effect`, y = reorder(Covariate, `Effect`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "Effect on Divorce Rate",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```

#### ⚙️ Code

```{r show model, echo = TRUE}
b5.1 <- 
  brm(data = WaffleDivorce, 
      family = gaussian,
      Divorce ~ 1 + MedianAgeMarriage_s,
      prior = c(prior(normal(10, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.01")
```

#### 🌌 Prior Vs Posterior

```{r prior vs post heat plot - median marriage age}
#| fig-width: 8
#| fig-height: 4

sim_n <- 1e3
prior_sim <- tibble(simMarriageAge = seq(from = -3, to = 3.5, length.out = sim_n),
       simDivorceAge = rnorm(sim_n, 
                             mean = (rnorm(sim_n, 10, 10) + (rnorm(sim_n) * simMarriageAge)), 
                             sd = runif(sim_n, min = 0, max = 10)))

ymin <- min(simModel$Estimate, WaffleDivorce$Divorce, prior_sim$simDivorceAge)
ymax <- max(simModel$Estimate, WaffleDivorce$Divorce, prior_sim$simDivorceAge)


prior_plot <- ggplot(data = prior_sim, aes(x = simMarriageAge, y = simDivorceAge)) +
  stat_density_2d(aes(fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = MedianAgeMarriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+  
  labs(y = "Divorce rate per 1000 adults",
       title = "Divorce Rate ~ Median Marriage Age", subtitle = "Prior")+
  scale_y_continuous(limits = c(ymin, ymax), expand = c(0, 0)) +
  theme_minimal()+
  
  # here it is!
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$MedianAgeMarriage) + mean(WaffleDivorce$MedianAgeMarriage), 1))+
  guides(fill = "none")

post_plot  <- ggplot() +
  stat_density_2d(data = simModel, 
                  aes(x = simMarriageAge, y = simDivorce, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = MedianAgeMarriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(y = "Divorce rate per 1000 adults",
       subtitle = "Posterior")+
  scale_y_continuous(limits = c(ymin, ymax), expand = c(0, 0))+
  theme_minimal()+
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$MedianAgeMarriage) + mean(WaffleDivorce$MedianAgeMarriage), 1))+
  guides(fill = "none")

prior_plot + post_plot + plot_layout(guides = "collect")
```
:::

The regression of Divorce Rate on Age of Marriage, tells us only that the total influence of age at marriage is strongly negative with divorce rate.

### Marriage Rate

$$ \textbf{b5.2} $$

$$ \text{Divorce}_i \sim \text{Intercept} + \beta_1 \text{Marriage % Rate}_i $$

```{r marriage % rate model - marriage % rate, include = FALSE}
b5.2 <- 
  brm(data = WaffleDivorce, 
      family = gaussian,
      Divorce ~ 1 + Marriage_s,
      prior = c(prior(normal(10, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.02")


simModel.02 <- as_tibble(as_draws_df(b5.2)) %>% 
  mutate(simMarriage = seq(from = -3, to = 3.5, length.out = n()),
         simDivorceEst = Intercept + (b_Marriage_s * simMarriage),
         simDivorce = rnorm(n(), simDivorceEst, sd = sigma))
```

::: panel-tabset
#### 📈 μ vs full distribution

```{r mu -  marriage rate}


# Find global limits across both datasets
ymin_mRate <- min(c(simModel.02$simDivorce, simModel.02$simDivorceEst, WaffleDivorce$Divorce))
ymax_mRate <- max(c(simModel.02$simDivorce, simModel.02$simDivorceEst, WaffleDivorce$Divorce))

modelEst_plot_mRate <- ggplot() +
  stat_density_2d(data = simModel.02, 
                  aes(x = simMarriage, y = simDivorceEst, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = Marriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(y = "Divorce rate per 1000 adults",
       title = "Divorce Rate ~ Marriage Rate", subtitle = "Mu estimate")+
  theme_minimal()+
  scale_x_continuous("Marriage rate per 1000 adults ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$Marriage) + 
                                      mean(WaffleDivorce$Marriage), 1)) +
  scale_y_continuous(limits = c(ymin_mRate, ymax_mRate), expand = c(0, 0)) +
  guides(fill = "none")

```

```{r full heat plot marriage rate}
#| fig-width: 8
#| fig-height: 4

model_plot_mRate <- ggplot() +
  stat_density_2d(data = simModel.02, 
                  aes(x = simMarriage, y = simDivorce, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = Marriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(y = "Divorce rate per 1000 adults",  subtitle = "Full Distribution")+
  theme_minimal()+
  scale_x_continuous("Marriage rate per 1000 adults",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$Marriage) + 
                                      mean(WaffleDivorce$Marriage), 1)) +
  scale_y_continuous(limits = c(ymin_mRate, ymax_mRate), expand = c(0, 0))+
  guides(fill = "none")

combined_plot_mRate <- modelEst_plot_mRate + model_plot_mRate + plot_layout(guides = "collect")

combined_plot_mRate

```

#### 🎛 Parameters

```{r summary b5.2}
#| fig-width: 10
#| fig-height: 2

as_tibble(b5.2) %>% 
  rename("Marriage Rate" = b_Marriage_s) %>% 
  dplyr::select(c(`Marriage Rate`, `Intercept`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Covariate",
               values_to = "Effect") %>% 
  ggplot(aes(x = `Effect`, y = reorder(Covariate, `Effect`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "Effect on Divorce Rate",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```

#### ⚙️ Code

```{r show model 5.2, echo = TRUE}
b5.2  <- 
  brm(data = WaffleDivorce, 
      family = gaussian,
      Divorce ~ 1 + Marriage_s,
      prior = c(prior(normal(10, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.02")
```

#### 🌌 Prior Vs Posterior

```{r prior vs post heat plot - marriage rate}
#| fig-width: 8
#| fig-height: 4

sim_n <- 1e3
prior_sim <- tibble(simMarriageRate = seq(from = -3, to = 3.5, length.out = sim_n),
       simDivorceAge = rnorm(sim_n, 
                             mean = (rnorm(sim_n, 10, 10) + (rnorm(sim_n) * simMarriageRate)), 
                             sd = runif(sim_n, min = 0, max = 10)))

ymin <- min(simModel$Estimate, WaffleDivorce$Divorce, prior_sim$simDivorceAge)
ymax <- max(simModel$Estimate, WaffleDivorce$Divorce, prior_sim$simDivorceAge)


prior_plot <- ggplot(data = prior_sim, aes(x = simMarriageRate, y = simDivorceAge)) +
  stat_density_2d(aes(fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = Marriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+  
  labs(y = "Divorce rate per 1000 adults",
       title = "Divorce Rate ~ Median Marriage Age", subtitle = "Prior")+
  scale_y_continuous(limits = c(ymin, ymax), expand = c(0, 0)) +
  theme_minimal()+
  
  # here it is!
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$Marriage) + mean(WaffleDivorce$Marriage), 1))+
  guides(fill = "none")

post_plot  <- ggplot() +
  stat_density_2d(data = simModel.02, 
                  aes(x = simMarriage, y = simDivorce, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = WaffleDivorce,
             aes(x = Marriage_s, y = Divorce), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(y = "Divorce rate per 1000 adults",
       subtitle = "Posterior")+
  scale_y_continuous(limits = c(ymin, ymax), expand = c(0, 0))+
  theme_minimal()+
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$Marriage) + mean(WaffleDivorce$Marriage), 1))+
  guides(fill = "none")

prior_plot + post_plot + plot_layout(guides = "collect")
```
:::

### Combined Multi-linear model

$$ \textbf{b5.3} $$

$$ \text{Divorce}_i \sim \text{Intercept} + \beta_1 \text{Marriage % Rate}_i + \beta_2 \text{Median Marriage Age}_i $$

```{r model - marriage age + marriage rate, include = FALSE}
b5.3 <- brm(data = WaffleDivorce, 
      family = gaussian,
      Divorce ~ 1 + Marriage_s + MedianAgeMarriage_s,
      prior = c(prior(normal(10, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.031")

simModel.03 <- as_tibble(as_draws_df(b5.3)) %>% 
  mutate(simMarriageAge = rnorm(n()),
         simMarriageRate = rnorm(n()),
         simDivorceEst = Intercept + (b_MedianAgeMarriage_s * simMarriageAge) +  (b_Marriage_s * simMarriageRate),
         simDivorce = rnorm(n(), simDivorceEst, sd = sigma))
```

::: panel-tabset
#### 🎛️ Parameters

```{r b5.3 params}
#| fig-width: 10
#| fig-height: 2

as_tibble(b5.3) %>% 
  rename("Marriage Rate" = b_Marriage_s,
         "Median Marriage Age" = b_MedianAgeMarriage_s) %>% 
  dplyr::select(c(`Marriage Rate`, `Median Marriage Age`, `Intercept`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Covariate",
               values_to = "Effect") %>% 
  ggplot(aes(x = `Effect`, y = reorder(Covariate, `Effect`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "Effect on Divorce Rate",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```

#### ⚙️ Code

```{r show model 5.3, echo = TRUE}
b5.3 <- brm(data = WaffleDivorce, 
      family = gaussian,
      Divorce ~ 1 + Marriage_s + MedianAgeMarriage_s,
      prior = c(prior(normal(10, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.031")
```
:::

### Diagnosing Multi-linear models

:::: panel-tabset
#### Predictor Residual Plot

::: panel-tabset
##### Marriage Rate

```{r Marriage Rate Residuals}
#| fig-width: 10
#| fig-height: 4

b5.4 <- 
  brm(data = WaffleDivorce,
      family = gaussian,
      Marriage_s ~ 1 + MedianAgeMarriage_s,
      prior = c(prior(normal(0, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.04.2")

fitted_B5.4 <- 
  fitted(b5.4) %>%
  as_tibble() %>%
  bind_cols(WaffleDivorce)

rateFitted <- fitted_B5.4 %>% 
  ggplot(aes(x = MedianAgeMarriage_s, y = Marriage_s)) +
  geom_point(size = 2, shape = 1, color = "firebrick4") +
  geom_segment(aes(xend = MedianAgeMarriage_s, yend = Estimate), 
               linewidth = 1/4) +
  geom_line(aes(y = Estimate), 
            color = "firebrick4") +
  coord_cartesian(ylim = range(WaffleDivorce$Marriage_s)) +
  theme_bw() +
  theme(panel.grid = element_blank()) 

r_b5.4 <- 
  residuals(b5.4) %>%
  # to use this in ggplot2, we need to make it a tibble or data frame
  as_tibble() %>% 
  bind_cols(WaffleDivorce)


# for the annotation at the top
text <-
  tibble(Estimate = c(- 1, 1),
         Divorce  = 14,
         label    = c("More Marriages", "Less Marriages"))

# plot
rateResidual <- r_b5.4 %>% 
  ggplot(aes(x = Estimate, y = Divorce)) +
  stat_smooth(method = "lm", fullrange = T,
              color = "firebrick4", fill = "firebrick4", 
              alpha = 1/5, linewidth = 1/2) +
  geom_vline(xintercept = 0, linetype = 2, color = "grey50") +
  geom_point(size = 2, color = "firebrick4", alpha = 2/3) +
  geom_text(data = text,
            aes(label = label)) +
  scale_x_continuous("Marriage rate residuals", limits = c(-2, 2)) +
  coord_cartesian(xlim = range(r_b5.4$Estimate),
                  ylim = c(6, 14.1)) +
  theme_bw() +
  theme(panel.grid = element_blank())

rateFitted + rateResidual
```

Little relationship between divorce and marriage rates, once we have accounted for the effect of Median Marriage Age

##### Median Marriage Age

```{r Median Age Residuals, include = FALSE}

b5.4b <- 
  brm(data = WaffleDivorce, 
      family = gaussian,
      MedianAgeMarriage_s ~ 1 + Marriage_s,
      prior = c(prior(normal(0, 10), class = Intercept),
                prior(normal(0, 1), class = b),
                prior(uniform(0, 10), class = sigma, ub = 10)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.04b")

fitted_B5.4b <- 
  fitted(b5.4b) %>%
  as_tibble() %>%
  bind_cols(WaffleDivorce)

AgeFitted <- fitted_B5.4b %>% 
  ggplot(aes(x = Marriage_s, y = MedianAgeMarriage_s)) +
  geom_point(size = 2, shape = 1, color = "firebrick4") +
  geom_segment(aes(xend = Marriage_s, yend = Estimate), 
               linewidth = 1/4) +
  geom_line(aes(y = Estimate), 
            color = "firebrick4") +
  coord_cartesian(ylim = range(WaffleDivorce$MedianAgeMarriage_s)) +
  theme_bw() +
  theme(panel.grid = element_blank()) 

r_b5.4b <- 
  residuals(b5.4b) %>%
  # to use this in ggplot2, we need to make it a tibble or data frame
  as_tibble() %>% 
  bind_cols(WaffleDivorce)

```

```{r}
#| fig-width: 10
#| fig-height: 4
text <-
  tibble(Estimate = c(- 0.7, 0.5),
         Divorce  = 14.1,
         label    = c("younger", "older"))

# extract the residuals
AgeResidual <- residuals(b5.4b) %>%
  as_tibble() %>%
  bind_cols(WaffleDivorce) %>% 
  
  # plot
  ggplot(aes(x = Estimate, y = Divorce)) +
  stat_smooth(method = "lm", fullrange = T,
              color = "firebrick4", fill = "firebrick4", 
              alpha = 1/5, linewidth = 1/2) +
  geom_vline(xintercept = 0, linetype = 2, color = "grey50") +
  geom_point(size = 2, color = "firebrick4", alpha = 2/3) +
  geom_text(data = text,
            aes(label = label)) +
  scale_x_continuous("Age of marriage residuals", limits = c(-2, 3)) +
  coord_cartesian(xlim = range(r_b5.4b$Estimate),
                  ylim = c(6, 14.1)) +
  theme_bw() +
  theme(panel.grid = element_blank()) 

AgeFitted + AgeResidual
```

Negative relationship between divorce and marriage age, even after we have accounted for the effect of the Marriage % Rate
:::

#### Posterior Prediction Plots

```{r Observed vs Fitted plot 5.3}
#| fig-width: 6
#| fig-height: 4

observed_vs_fitted_plot <-
  fitted(b5.3) %>%
  as_tibble() %>%
  bind_cols(WaffleDivorce) %>%
  ggplot(aes(x = Divorce, y = Estimate)) +
  geom_abline(linetype = 2, color = "grey50", linewidth = 0.5) +
  geom_point(size = 1.5, color = "firebrick4", alpha = 3/4) +
  geom_linerange(aes(ymin = Q2.5, ymax = Q97.5),
                 linewidth = 1/4, color = "firebrick4") +
  geom_linerange(aes(ymin = Estimate - Est.Error, 
                     ymax = Estimate + Est.Error),
                 linewidth = 1/2, color = "firebrick4") +
  # Note our use of the dot placeholder, here: https://magrittr.tidyverse.org/reference/pipe.html
  geom_text(data = . %>% filter(Loc %in% c("ID", "UT", "NJ", "ME", "AR", "AL")),
            aes(label = Loc), 
            hjust = 0, nudge_x = - 0.65) +
  labs(x = "Observed divorce", 
       y = "Predicted divorce",) +
  theme_bw() +
  theme(panel.grid = element_blank())

observed_vs_fitted_plot

```

#### Causal Counterfactual Plots

```{r}
#| fig-width: 10
#| fig-height: 4

n_size = 300

simMarriages_age <- as_tibble(as_draws_df(b5.4)) %>% 
  mutate(MedianAgeMarriage_s = seq(from = -3, to = 3.5, length.out = n()),
         simMarriageRateEst = Intercept + (b_MedianAgeMarriage_s * MedianAgeMarriage_s),
         Marriage_s = rnorm(n(), simMarriageRateEst, sd = sigma)) %>% 
  dplyr::select(c(MedianAgeMarriage_s, Marriage_s)) %>% 
  slice_sample(n = n_size)

simMRatePred <- as_tibble(predict(b5.4, simMarriages_age)) %>% 
  cbind(simMarriages_age) 

simDRatePred <- as_tibble(predict(b5.3, simMarriages_age)) %>% 
  cbind(simMarriages_age) 


mcf_plot <- ggplot(data = simMRatePred, aes(x = MedianAgeMarriage_s, y = Marriage_s)) +
  geom_ribbon(aes(ymin = Q2.5, ymax = Q97.5),
              fill = "firebrick", alpha = .4)+
  geom_smooth(method = "lm", color = "red4")+
  scale_x_continuous("Set Median Age of Marriage",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$MedianAgeMarriage) + 
                                      mean(WaffleDivorce$MedianAgeMarriage))) +
  scale_y_continuous("Marriage Rate per 1000",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$Marriage) + 
                                      mean(WaffleDivorce$Marriage))) +
  labs(title = "Marriage Rate if you set Marriage Age")+
  theme_minimal()

dcf_plot <- ggplot(data = simDRatePred, aes(x = MedianAgeMarriage_s, y = Estimate)) +
  geom_ribbon(aes(ymin = Q2.5, ymax = Q97.5),
              fill = "firebrick", alpha = .4)+
  geom_smooth(method = "lm", color = "red4")+
  scale_x_continuous("Set Median Age of Marriage",
                     breaks = at,
                     labels = round(at * sd(WaffleDivorce$MedianAgeMarriage) + 
                                      mean(WaffleDivorce$MedianAgeMarriage))) +
  labs(y = "Divorce Rate per 1000", title = "Divorce Rate if you set Marriage Age")+
  theme_minimal()

mcf_plot + dcf_plot
  

```
::::

"This procedure also brings home the message that regression models measure the remaining association of each predictor with the out- come, after already knowing the other predictors. In computing the predictor residual plots, you had to perform those calculations yourself. In the unified multivariate model, it all hap- pens automatically. Nevertheless, it is useful to keep this fact in mind, because regressions can behave in surprising ways as a result."

### Solving the DAG

From McElreath, "Age of marriage influences divorce in two ways. First it can have a direct effect, perhaps because younger people change faster than older people and are therefore more likely to grow incompatible with a partner. Second, it can have an indirect effect by influencing the marriage rate. If people get married earlier, then the marriage rate may rise, because there are more young people. Consider for example if an evil dictator forced everyone to marry at age 65. Since a smaller fraction of the population lives to 65 than to 25, forcing delayed marriage will also reduce the marriage rate. If marriage rate itself has any direct effect on divorce, maybe by making marriage more or less normative, then some of that direct effect could be the indirect effect of age at marriage."

```{r parameter comparison}
posterior_comparison <- bind_cols(
  as_draws_df(b5.1) %>% 
    transmute(`b5.1_beta[" Median Marriage Age"]` = b_MedianAgeMarriage_s),
  as_draws_df(b5.2) %>% 
    transmute(`b5.2_beta[" Marriage % Rate"]` = b_Marriage_s),
  as_draws_df(b5.3) %>% 
    transmute(`b5.3_beta[" Median Marriage Age"]` = b_MedianAgeMarriage_s,
              `b5.3_beta[" Marriage % Rate"]` = b_Marriage_s)
) %>% 
  pivot_longer(everything()) %>% 
  separate(name, into = c("fit", "parameter"), sep = "_")


# Create the plot with distributions
posterior_comparison %>% 
  ggplot(aes(x = value, y = fit)) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4", alpha = 0.8) +
  labs(x = "Beta Coefficient Estimates", 
       y = "Model", 
       title = "Posterior Distributions for Different Models") +
  theme_bw() +
  theme(panel.grid = element_blank(),
        strip.background = element_rect(fill = "transparent", color = "transparent"),
        axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank()) +
  facet_wrap(~ parameter, ncol = 1, labeller = label_parsed, scales = "free_x")

```

Once we know median age at marriage for a State, there is little or no additional predictive power in also knowing the rate of marriage in that State.

::: panel-tabset
#### US Map Residual

```{r}
#| fig-width: 6
#| fig-height: 4

deviance_plot_map <-
  fitted(b5.3) %>%
  as_tibble() %>%
  bind_cols(WaffleDivorce) %>%
  right_join(readRDS("externalData/us_states_geometry.rds"), 
             by = c("Location" = "NAME")) %>%
  mutate(deviance =  Divorce -Estimate) %>% 
  pivot_longer("deviance") %>% 
  ggplot() +
  geom_sf(aes(fill = value, geometry = geometry),
          size = 0) +
  scale_fill_gradient(low = "#f8eaea", high = "firebrick4") +
  coord_sf(xlim = c(-29e5, 21e5),
           ylim = c(-2e6, 2e6)) +
  theme_void() +
  theme(legend.position = "none",
        strip.text = element_text(margin = margin(0, 0, 0.5, 0))) +
  labs(title = "Model Residuals By State", caption = "White = Over-estimateed Divorce,   Red = Under-estimated Divorce")


deviance_plot_map
```

#### States by Residual

```{r}
#| fig-width: 3
#| fig-height: 6

residuals_plot <-
  residuals(b5.3) %>% 
  as_tibble() %>% 
  rename(f_ll = Q2.5,
         f_ul = Q97.5) %>% 
  bind_cols(
    predict(b5.3) %>% 
      as_tibble() %>% 
      transmute(p_ll = Q2.5,
                p_ul = Q97.5),
    WaffleDivorce
  ) %>% 
  # here we put our `predict()` intervals into a deviance metric
  mutate(p_ll = Divorce - p_ll,
         p_ul = Divorce - p_ul) %>% 
  
  # now plot!
  ggplot(aes(x = reorder(Loc, Estimate), y = Estimate)) +
  geom_hline(yintercept = 0, linewidth = 1/2, 
             color = "firebrick4", alpha = 1/10) +
  geom_pointrange(aes(ymin = f_ll, ymax = f_ul),
                  linewidth = 2/5, shape = 20, color = "firebrick4") + 
  geom_segment(aes(y    = Estimate - Est.Error, 
                   yend = Estimate + Est.Error,
                   x    = Loc, 
                   xend = Loc),
               linewidth = 1, color = "firebrick4") +
  geom_segment(aes(y = p_ll, yend = p_ul,
                   x = Loc,  xend = Loc),
               linewidth = 3, color = "firebrick4", alpha = 1/10) +
  labs(x = NULL, y = NULL) +
  coord_flip(ylim = c(-6, 5)) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0, size = 7, margin = margin(r = 5)),
        axis.text.x = element_text(size = 7),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank(),
        plot.margin = margin(20, 20, 20, 20))


residuals_plot

```

#### Waffle House Effect

```{r}
#| fig-width: 6
#| fig-height: 4

waffleHouseEffect_plot <-
  residuals(b5.3) %>% 
  as_tibble() %>% 
  bind_cols(WaffleDivorce) %>% 
  mutate(wpc = WaffleHouses / Population) %>% 
  
  ggplot(aes(x = wpc, y = Estimate)) +
  geom_point(size = 1.5, color = "firebrick4", alpha = 1/2) +
  stat_smooth(method = "lm", fullrange = T,
              color = "firebrick4", linewidth = 1/2, 
              fill = "firebrick", alpha = 1/5) + 
  geom_text_repel(data = . %>% filter(Loc %in% c("ME", "AR", "MS", "AL", "GA", "SC", "ID")),
                  aes(label = Loc),
                  seed = 5.6) +
  scale_x_continuous("Waffles per capita", limits = c(0, 45)) +
  ylab("Divorce error") +
  coord_cartesian(xlim = range(0, 40)) +
  theme_bw() +
  theme(panel.grid = element_blank())

waffleHouseEffect_plot

```
:::

## Masked Relationships

New data set Comparative primate milk composition data, from Table 2 of Hinde and Milligan. 2011. Evolutionary Anthropology 20:9-23.

species: Species name

kcal.per.g: Kilocalories per gram of milk

mass: Body mass of mother, in kilograms

neocortex.perc: Percent of brain mass that is neocortex

```{r milk wrangling, include = FALSE}
milk %>% 
  select(kcal.per.g, mass, neocortex.perc) %>% 
  pairs(col = "firebrick4")

milk2 <- milk %>% 
  filter(!is.na(neocortex.perc)) %>% 
  mutate(neocortex.perc_s = (neocortex.perc - mean(neocortex.perc))/sd(neocortex.perc),
         kcal.per.g_s = (kcal.per.g - mean(kcal.per.g))/sd(kcal.per.g),
         mass_s = (mass - mean(mass))/sd(mass))


```

### Neocortex % Effect on Milk Calories

```{r neocortex percent model, include = FALSE}
milk.01 <- 
  brm(data = milk2, 
      family = gaussian,
      kcal.per.g_s ~ 1 + neocortex.perc_s,
      prior = c(prior(normal(0, .2), class = Intercept),
                prior(normal(0, .5), class = b),
                prior(exponential(1), class = sigma)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/milk.01.6")


at <- c(-3, -2, -1, 0, 1, 2, 3)

plot(milk.01)

simMilk <- as_tibble(as_draws_df(milk.01)) %>% 
  mutate(simNeocortexPerc = seq(from = -3, to = 3, length.out = n()),
         simKcalEst = Intercept + (b_neocortex.perc_s * simNeocortexPerc),
         simKcal = rnorm(n(), simKcalEst, sd = sigma))
```

::: panel-tabset
#### 📈 μ vs full distribution

```{r mu vs full heat plot milk neocortex}
#| fig-width: 10
#| fig-height: 4

# Find global limits across both datasets
ymin <- min(c(simMilk$simKcal, simMilk$simKcalEst, milk2$kcal.per.g_s))
ymax <- max(c(simMilk$simKcal, simMilk$simKcalEst, milk2$kcal.per.g_s))

modelEst_plot <- ggplot() +
  stat_density_2d(data = simMilk, 
                  aes(x = simNeocortexPerc, y = simKcalEst, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = milk2,
             aes(x = neocortex.perc_s, y = kcal.per.g_s), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  
  scale_y_continuous("Kilo-Calories Per Gram of Breast Milk",
                     breaks = at,
                     labels = round(at * sd(milk2$kcal.per.g) + 
                                      mean(milk2$kcal.per.g), 2), expand = c(0, 0)) +
  scale_x_continuous("% of Brain that is Neocortex",
                     breaks = at,
                     labels = round(at * sd(milk2$neocortex.perc) + 
                                      mean(milk2$neocortex.perc))) +
  labs(title = "Milk Calories ~ Neocortex", subtitle = "Mu estimate")+
  theme_minimal()+
  guides(fill = "none")

model_plot <- ggplot() +
  stat_density_2d(data = simMilk, 
                  aes(x = simNeocortexPerc, y = simKcal, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = milk2,
             aes(x = neocortex.perc_s, y = kcal.per.g_s), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(subtitle = "Full Distribution")+
  theme_minimal()+
  scale_x_continuous("% of Brain that is Neocortex",
                     breaks = at,
                     labels = round(at * sd(milk2$neocortex.perc) + 
                                      mean(milk2$neocortex.perc))) +
  scale_y_continuous("Kilo-Calories Per Gram of Breast Milk",
                     breaks = at,
                     labels = round(at * sd(milk2$kcal.per.g) + 
                                      mean(milk2$kcal.per.g), 2), expand = c(0, 0)) +
  guides(fill = "none")

modelEst_plot + model_plot + plot_layout(guides = "collect")

```

The posterior mean line is weakly positive, but it is highly imprecise. A lot of mildly positive and negative slopes are plausible, given this model and these data.

#### 🎛️ Parameters

```{r}
#| fig-width: 10
#| fig-height: 2

as_tibble(milk.01) %>% 
  rename("% of Brain Neocortex" = b_neocortex.perc_s) %>% 
  dplyr::select(c(`% of Brain Neocortex`, `Intercept`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Covariate",
               values_to = "Effect") %>% 
  ggplot(aes(x = `Effect`, y = reorder(Covariate, `Effect`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "Effect on Kilo-Calories per Gram of Milk",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```

#### ⚙️ Code

```{r show milk neocortex model, echo = TRUE}
milk.01 <- 
  brm(data = milk2, 
      family = gaussian,
      kcal.per.g_s ~ 1 + neocortex.perc_s,
      prior = c(prior(normal(0, .2), class = Intercept),
                prior(normal(0, .5), class = b),
                prior(exponential(1), class = sigma)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/milk.01.6")
```

#### 🌌 Prior Vs Posterior

```{r prior vs post heat plot - milk neocortex model}
#| fig-width: 10
#| fig-height: 4

sim_n <- 1e3
prior_sim <- tibble(simNeocortex = seq(from = -3, to = 3, length.out = sim_n),
                    simKCal = rnorm(sim_n,
                                    mean = (rnorm(sim_n, 0, 0.2) + (rnorm(sim_n, 0, 0.5) * simNeocortex)), 
                                    sd = rexp(sim_n, rate = 1)))

ymin <- min(c(simMilk$simKcal, prior_sim$simKCal, milk2$kcal.per.g_s))
ymax <- max(c(simMilk$simKcal, prior_sim$simKCal, milk2$kcal.per.g_s))


prior_plot <- ggplot(data = prior_sim, aes(x = simNeocortex, y = simKCal)) +
  stat_density_2d(aes(fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = milk2,
             aes(x = neocortex.perc_s , y = kcal.per.g_s), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+  
  labs(subtitle = "Prior")+
  scale_x_continuous("% of Brain that is Neocortex",
                     breaks = at,
                     labels = round(at * sd(milk2$neocortex.perc) + 
                                      mean(milk2$neocortex.perc))) +
  scale_y_continuous("Kilo-Calories Per Gram of Breast Milk",
                     breaks = seq(from = ymin, to = ymax, length.out = 8),
                     labels = round(seq(from = ymin, to = ymax, length.out = 8) * sd(milk2$kcal.per.g) + 
                                      mean(milk2$kcal.per.g), 1), expand = c(0, 0)) +
  guides(fill = "none")+
  theme_minimal()+
  
  # here it is!
  scale_x_continuous("Median Age ",
                     breaks = at,
                     labels = round(at * sd(milk2$neocortex.perc) + mean(milk2$neocortex.perc), 1))+
  guides(fill = "none")

post_plot <- ggplot() +
  stat_density_2d(data = simMilk, 
                  aes(x = simNeocortexPerc, y = simKcal, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = milk2,
             aes(x = neocortex.perc_s, y = kcal.per.g_s),
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(subtitle = "Posterior")+
  theme_minimal()+
  scale_x_continuous("% of Brain that is Neocortex",
                     breaks = at,
                     labels = round(at * sd(milk2$neocortex.perc) + 
                                      mean(milk2$neocortex.perc))) +
  scale_y_continuous("Kilo-Calories Per Gram of Breast Milk",
                     breaks = seq(from = ymin, to = ymax, length.out = 8),
                     labels = round(seq(from = ymin, to = ymax, length.out = 8) * sd(milk2$kcal.per.g) + 
                                      mean(milk2$kcal.per.g), 1), expand = c(0, 0)) +
  guides(fill = "none")



prior_plot + post_plot + plot_layout(guides = "collect")
```
:::

### Mass Effect on Milk Calories

Now consider another predictor variable, adult female body mass, mass in the data frame. Let’s use the logarithm of mass, log(mass), as a predictor as well. Why the logarithm of mass instead of the raw mass in kilograms? It is often true that scaling measurements like body mass are related by magnitudes to other variables. Taking the log of a measure trans- lates the measure into magnitudes. So by using the logarithm of body mass here, we’re saying that we suspect that the magnitude of a mother’s body mass is related to milk energy, in a linear fashion.

```{r Mass  model, include = FALSE}
milk2 <- milk2 %>% 
  mutate(logMass = log(mass),
         logMass_s = (logMass - mean(logMass))/sd(logMass))

milk.02 <- 
  brm(data = milk2, 
      family = gaussian,
      kcal.per.g_s ~ 1 + logMass_s,
      prior = c(prior(normal(0, .2), class = Intercept),
                prior(normal(0, .5), class = b),
                prior(exponential(1), class = sigma)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/milk.2.0")


at <- c(-1.5, -1, 0, 1, 1.5)

simMilkMass <- as_tibble(as_draws_df(milk.02)) %>% 
  mutate(simLogMass = seq(from = -3, to = 3, length.out = n()),
         simKcalEst = Intercept + (b_logMass_s * simLogMass),
         simKcal = rnorm(n(), simKcalEst, sd = sigma))



```

::: panel-tabset
#### 📈 μ vs full distribution

```{r mu vs full heat plot milk mass}
#| fig-width: 10
#| fig-height: 4
# Find global limits across both datasets
ymin <- min(c(simMilkMass$simKcal, simMilkMass$simKcalEst, milk2$kcal.per.g_s))
ymax <- max(c(simMilkMass$simKcal, simMilkMass$simKcalEst, milk2$kcal.per.g_s))

modelEst_plot <- ggplot() +
  stat_density_2d(data = simMilkMass, 
                  aes(x = simLogMass, y = simKcalEst, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = milk2,
             aes(x = logMass_s, y = kcal.per.g_s), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  
  scale_y_continuous("Kilo-Calories Per Gram of Breast Milk",
                     breaks = at,
                     labels = round(at * sd(milk2$kcal.per.g) + 
                                      mean(milk2$kcal.per.g), 2), expand = c(0, 0)) +
  scale_x_continuous("Body Mass of Mother Grams",
                     breaks = at,
                     labels = round(exp(at * sd(milk2$logMass) + 
                                      mean(milk2$logMass))), 1) +
  labs(title = "Milk Calories ~ Neocortex", subtitle = "Mu estimate")+
  theme_minimal()+
  guides(fill = "none")

model_plot <- ggplot() +
  stat_density_2d(data = simMilkMass, 
                  aes(x = simLogMass, y = simKcal, fill = after_stat(ndensity)),
                  geom = "raster", contour = FALSE) +
  scale_fill_viridis_c(option = "magma") +
  geom_point(data = milk2,
             aes(x = logMass_s, y = kcal.per.g_s), 
             shape = 21, color = "white", fill = "black", lwd = 3, alpha = .8)+
  labs(subtitle = "Full Distribution")+
  theme_minimal()+
  scale_x_continuous("Body Mass of Mother Grams",
                     breaks = at,
                     labels = round(exp(at * sd(milk2$logMass) + 
                                      mean(milk2$logMass))), 1) +
  scale_y_continuous("Kilo-Calories Per Gram of Breast Milk",
                     breaks = at,
                     labels = round(at * sd(milk2$kcal.per.g) + 
                                      mean(milk2$kcal.per.g), 2), expand = c(0, 0)) +
  guides(fill = "none")

modelEst_plot + model_plot + plot_layout(guides = "collect")

```

The posterior mean line is moderately negative, but it is highly imprecise.

#### 🎛️ Parameters

```{r}
#| fig-width: 10
#| fig-height: 2

as_tibble(milk.02) %>% 
  rename("Brain Mass" = b_logMass_s) %>% 
  dplyr::select(c(`Brain Mass`, `Intercept`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Covariate",
               values_to = "Effect") %>% 
  ggplot(aes(x = `Effect`, y = reorder(Covariate, `Effect`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "Effect on Kilo-Calories per Gram of Milk",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```

#### ⚙️ Code

```{r show milk mass model, echo = TRUE}
milk.02 <- 
  brm(data = milk2, 
      family = gaussian,
      kcal.per.g_s ~ 1 + logMass_s,
      prior = c(prior(normal(0, .2), class = Intercept),
                prior(normal(0, .5), class = b),
                prior(exponential(1), class = sigma)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/milk.2.0")

```
:::

### Multivariate Milk

```{r, include = FALSE}
milk.3 <- 
  brm(data = milk2, 
      family = gaussian,
      kcal.per.g ~ 1 + neocortex.perc_s + logMass_s,
      prior = c(prior(normal(0, 0.2), class = Intercept),
                prior(normal(0, .5), class = b),
                prior(exponential(1), class = sigma)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/milk.3")

```

```{r}
#| fig-width: 10
#| fig-height: 2

as_tibble(milk.3) %>% 
  rename("Brain Mass" = b_logMass_s,
         "% Brain Neocortex" = b_neocortex.perc_s) %>% 
  dplyr::select(c(`Brain Mass`, `% Brain Neocortex`, `Intercept`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Covariate",
               values_to = "Effect") %>% 
  ggplot(aes(x = `Effect`, y = reorder(Covariate, `Effect`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "Effect on Kilo-Calories per Gram of Milk",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```

```{r}
#| fig-width: 5
#| fig-height: 4

nd <- 
  tibble(neocortex.perc_s =  seq(from = -3, to = 3, length.out = 1e3),
         logMass_s = mean(milk2$logMass))

p1 <-
  milk.3 %>%
  fitted(newdata = nd, 
         probs = c(.025, .975, .25, .75)) %>%
  as_tibble() %>%
  bind_cols(nd) %>% 

  ggplot(aes(x = neocortex.perc_s, y = Estimate)) +
  geom_ribbon(aes(ymin = Q2.5, ymax = Q97.5),
              fill = "firebrick", alpha = 1/5) +
  geom_smooth(aes(ymin = Q25, ymax = Q75),
              stat = "identity",
              fill = "firebrick4", color = "firebrick4", alpha = 1/5, linewidth = 1/2) +
  geom_point(data = milk2, 
             aes(y = kcal.per.g_s),
             size = 2, color = "firebrick4") +
  ylab("kcal.per.g") +
  theme_minimal()
  # coord_cartesian(xlim = range(milk2$neocortex.perc), 
  #                 ylim = range(milk2$kcal.per.g))

p1
```

```{r}
dag5.7 <- dagitty( "dag{
    M -> K <- N
    M -> N }" )
coordinates(dag5.7) <- list( x=c(M=0,K=1,N=2) , y=c(M=0.5,K=1,N=0.5) )
MElist <- equivalentDAGs(dag5.7)

library(rethinking)
#drawdag(MElist)

dag5.8 <- dagitty( "dag{
    M -> K 
    M -> N }" )
coordinates(dag5.8) <- list( x=c(M=0,K=1,N=2) , y=c(M=0.5,K=1,N=0.5) )
MElist <- equivalentDAGs(dag5.8)

#drawdag(MElist)
```

## Categorical Variable

### Binary Categories

Two mutually exclusive categories make the interpretation of models quite simple because it's hard to forget that the lack of Category A implies Category B. An example, lets use male vs females as our categorical variable, with Males being represented as 1 and females as 0 in our data. The models interpretations is that the intercept is the effect of being female and the intercept + $\beta_\text{effect of male}$ is the males effect. The same math goes into the interpretations when our categorical variables have more than 2 options, but it's quite easy to forget that your comparing to one of the categories when you have a list of 20 categories. Always remember that you are comparing to some other category with factor effects.

Let's go back to our Kung! height dataset we used in the last chapter. Let's ignore the effect of weight and the other variables and focus only on sex. Our model is

$$ \text{Height}_i \sim \text{Normal(}\mu_i, \ \sigma_i\text{)}$$ $$ \mu_i = \alpha + \beta_\text{male} \mathbf{I}\text{(male = 1)}$$

$$\alpha \sim \text{Normal(178, 20)}$$

$$ \beta_\text{male} \sim \text{Normal(0, 10)}$$

$$ \sigma \sim \text{Uniform(0, 50)}$$

The parameter $\beta_\text{male}$ only influences prediction for the cases where the individual is a male. Using this approach means that $\beta_\text{male}$ represents the expected difference between males and females in height. This makes assigning our priors a bit harder. Our approach assumes there is more uncertainty about males than females, because a male includes two parameters and therefore has two parameters.

```{r prior height sim}
#| fig-width: 10
#| fig-height: 3

simHeight <- tibble(female = rnorm(1e5, 178, 20),
       male = female + rnorm(1e5, 0, 10)) %>% 
  pivot_longer(cols = everything(),
               names_to = "sex",
               values_to = "Height")

prior_plt_5.6 <- ggplot(data = simHeight)+
  geom_density(aes(x = Height, color = sex),  adjust = 2, lwd = 1)+
  theme_minimal() +
  theme(axis.text.y = element_blank(),
        axis.title.y = element_blank(),
        axis.ticks.y = element_blank())+
  scale_x_continuous(
    breaks = c(48, 60,  72, 84) * 2.54,
    labels = c("4'0\"",  "5'0\"", 
                "6'0\"",  "7'0\""),
    limits = c(40, 98) * 2.54
  )+
  labs(title = "Prior Mean")
  

b5.6 <- 
  brm(data = kungHeight, 
      family = gaussian,
      height ~ 1 + male,
      prior = c(prior(normal(178, 20), class = Intercept),
                prior(normal(0, 10), class = b),
                prior(uniform(0, 50), class = sigma, ub = 50)),
      iter = 2000, warmup = 500, chains = 4, cores = 4,
      seed = 5, backend = "cmdstanr", silent = 2,
      file = "fits/b05.6b")

post_5.6 <- as_tibble(b5.6) %>% 
  mutate(female_mu = b_Intercept,
         male_mu = b_Intercept + b_male,
         female_est = rnorm(n(), mean = female_mu, sd = sigma),
         male_est = rnorm(n(), mean = male_mu, sd = sigma)) %>% 
  dplyr::select(female_mu, male_mu, female_est, male_est) %>% 
  pivot_longer(cols = everything(),
               names_to = c("sex", ".value"),
               names_sep = "_",
               values_to = c("mu", "est"))

post_plt_5.6 <- ggplot(data = post_5.6)+
  geom_density(aes(x = mu, color = sex),  adjust = 2, lwd = 1)+
  theme_minimal() +
  theme(axis.text.y = element_blank(),
        axis.title.y = element_blank(),
        axis.ticks.y = element_blank())+
  scale_x_continuous(
    breaks = c(48, 60,  72, 84) * 2.54,
    labels = c("4'0\"",  "5'0\"", 
                "6'0\"",  "7'0\""),
    limits = c(40, 98) * 2.54,
    name = "Height"
  )+
  labs(title = "Posterior Mean")

prior_plt_5.6 + post_plt_5.6 + plot_layout(guides = "collect")

```

Our priors are a bit wide, and the extra variance on males doesn't make to much sense.

A better way is to use an index variable.

$$ \text{Height}_i \sim \text{Normal(}\mu_i, \ \sigma_i\text{)}$$ $$ \mu_i = \alpha_\text{Sex[i]}$$

$$\alpha_j \sim \text{Normal(178, 20),  for j  = 1..2}$$

$$ \sigma \sim \text{Uniform(0, 50)}$$

```{r}
#| fig-width: 10
#| fig-height: 3

kungHeight <- kungHeight %>% 
  mutate(sex = ifelse(male == 1, "male", "female"))

simHeight <- tibble(female = rnorm(1e5, 178, 20),
       male = rnorm(1e5, 178, 20)) %>% 
  pivot_longer(cols = everything(),
               names_to = "sex",
               values_to = "Height")

prior_plt_5.6b <- ggplot(data = simHeight)+
  geom_density(aes(x = Height, color = sex),  adjust = 2, lwd = 1, alpha = .5)+
  theme_minimal() +
  theme(axis.text.y = element_blank(),
        axis.title.y = element_blank(),
        axis.ticks.y = element_blank())+
  scale_x_continuous(
    breaks = c(48, 60,  72, 84) * 2.54,
    labels = c("4'0\"",  "5'0\"", 
                "6'0\"",  "7'0\""),
    limits = c(40, 98) * 2.54
  )+
  labs(title = "Prior Mean")
  
b5.6b <- brm(
  data = kungHeight,
  family = gaussian,
  height ~ 0 + sex,  # 0 removes intercept, sex creates separate parameters for each level
  prior = c(
    prior(normal(178, 20), class = b),     # prior for sex coefficients (a[sex])
    prior(uniform(0, 50), class = sigma, ub = 50)   # prior for sigma
  ),
  chains = 4, iter = 2000, warmup = 500, cores = 4,
  backend = "cmdstanr", silent = 2, file = "fits/b05.6b4")

post_5.6b <- as_tibble(b5.6b) %>% 
  mutate(female_mu = b_sexfemale,
         male_mu = b_sexmale,
         female_est = rnorm(n(), mean = female_mu, sd = sigma),
         male_est = rnorm(n(), mean = male_mu, sd = sigma)) %>% 
  dplyr::select(female_mu, male_mu, female_est, male_est) %>% 
  pivot_longer(cols = everything(),
               names_to = c("sex", ".value"),
               names_sep = "_",
               values_to = c("mu", "est"))

post_plt_5.6b <- ggplot(data = post_5.6b)+
  geom_density(aes(x = mu, color = sex),  adjust = 2, lwd = 1)+
  theme_minimal() +
  theme(axis.text.y = element_blank(),
        axis.title.y = element_blank(),
        axis.ticks.y = element_blank())+
  scale_x_continuous(
    breaks = c(48, 60,  72, 84) * 2.54,
    labels = c("4'0\"",  "5'0\"", 
                "6'0\"",  "7'0\""),
    limits = c(40, 98) * 2.54,
    name = "Height"
  )+
  labs(title = "Posterior Mean")

prior_plt_5.6b + post_plt_5.6b + plot_layout(guides = "collect")

```

We actually care about the differences in the genders in height in this case. So for this model we simply subtract the posterior intercepts from each other to show that difference.

```{r}
#| fig-width: 10
#| fig-height: 2 
as_tibble(b5.6b) %>% 
  mutate(female_mu = b_sexfemale,
         male_mu = b_sexmale,
         contrast = b_sexmale - b_sexfemale) %>% 
  ggplot()+
  geom_density(aes(x = contrast), color = "purple2", adjust = 2, lwd = 1)+
    theme_minimal() +
  theme(axis.text.y = element_blank(),
        axis.title.y = element_blank(),
        axis.ticks.y = element_blank())+
  scale_x_continuous(
    breaks = c(0, 2, 4, 6) * 2.54,
    labels = c("0'0\"",  "02\"", "0'4\"",  "0'6\""),
    name = "Height"
  )+
  labs(title = "Male Height - Female Height Mu")
```

Binary categories are easy, whether you use an indicator variable or instead an index variable. But when there are more than two categories, the indicator variable approach explodes. While the index approach does not change at all when you add more categories. You do get more parameters, of course, just as many as in the indicator variable approach. But the model specification looks just like it does in the binary case. And the priors continue to be easier, unless you really do have prior information about contrasts.

### Milk example

```{r}
b5.7 <- brm(
  data = milk,
  family = gaussian,
  kcal.per.g ~ 0 + clade,
  prior = c(
    prior(normal(0, 0.5), class = b),
    prior(exponential(1), class = sigma)
  ),
  seed = 5, iter = 2000, warmup = 500, cores = 4,
  backend = "cmdstanr", silent = 2,
  file = "fits/b05.7c")
```

```{r}
#| fig-width: 10
#| fig-height: 2

nd <- tibble(clade_nwm = c(1, 0, 0, 0),
             clade_owm = c(0, 1, 0, 0),
             clade_s   = c(0, 0, 1, 0),
             primate   = c("New World Monkey", "Old World Monkey", "Strepsirrhine", "Ape"))

as_tibble(b5.7) %>% 
  rename("Ape" = b_cladeApe,
         "New World Monkey" = b_cladeNewWorldMonkey,
         "Old World Monkey" = b_cladeOldWorldMonkey,
         "Strepsirrhine" = b_cladeStrepsirrhine) %>% 
  dplyr::select(c(`Ape`, `New World Monkey`, `Old World Monkey`, `Strepsirrhine`)) %>% 
  pivot_longer(cols = everything(),
               names_to = "Primate",
               values_to = "KCal Per Gram of Milk") %>% 
  ggplot(aes(x = `KCal Per Gram of Milk`, y = reorder(Primate, `KCal Per Gram of Milk`))) +
  stat_halfeye(point_interval = median_qi, .width = .95,
               fill = "firebrick4") +
  labs(x = "kcal.per.g",
       y = NULL) +
  theme_bw() +
  theme(axis.text.y = element_text(hjust = 0),
        axis.ticks.y = element_blank(),
        panel.grid = element_blank())
```