3_plot_model_predictions.R

#!/usr/bin/env Rscript
library(tidyverse)
library(brms)
library(pals)
library(patchwork)
library(matrixStats)
source("util.R")

set.seed(6789)

dir.create("plots", showWarnings=FALSE)

#load child data
d_child <- read_csv("data/combined_long_data.csv") %>%
	filter(AgeGroup == "children") %>%
	mutate(response_has = factor(response_has, levels = c("No", "Partially", "Yes"), labels=c("No", "Partially", "Yes"), exclude=NA, ordered = TRUE),
	       response_similar = factor(response_similar, levels = c("Different", "Partially", "Same"), labels=c("Different", "Partially", "Same"), exclude=NA, ordered = TRUE))
AGE_MEAN <- mean(d_child$Age)
AGE_SD <- sd(d_child$Age)

# Load models
m_children_has <- readRDS("m_children_has.rds")
m_children_sim <- readRDS("m_children_sim.rds")

# Some utility functions
recode_for_models <- function(d) {
	d %>%
		mutate(
		       Age = (Age - AGE_MEAN) / AGE_SD,
		       urb_rural_thoughts=case_when(type == "feelings" ~ 0,
                                                    Urb_rural == "urban" ~ -0.5,
                                                    Urb_rural == "rural" ~ +0.5,
                                                    Urb_rural == 0 ~ 0),
		       urb_rural_feelings=case_when(type == "thoughts" ~ 0,
                                                    Urb_rural == "urban" ~ -0.5,
                                                    Urb_rural == "rural" ~ +0.5,
                                                    Urb_rural == 0 ~ 0),
		       type=if_else(type == "thoughts", -0.5, 0.5),
		       ) %>%
		return()
}

tidy_up <- function(d) {
	d %>%
		mutate(question = if_else(question == "has",
				  "Do animals have ...?",
				  "Are animals' ... the same as humans'?"),
	               question = factor(question, c("Do animals have ...?", "Are animals' ... the same as humans'?"), ordered=T),
 	               type = str_to_title(type),
	               type = factor(type, c("Thoughts", "Feelings"), ordered=T)) %>%
		return()
}

## This is black magic from https://stackoverflow.com/questions/70977700/creating-completely-customized-legends-in-ggplot2
dummy_guide <- function(
    labels = NULL,
    ...,
    title = NULL,
    key   = draw_key_point,
    guide_args = list()
) {
  # Capture arguments
  aesthetics <- list(...)
  n <- max(lengths(aesthetics), 0)
  labels <- labels %||%  seq_len(n)

  # Overrule the alpha = 0 that we use to hide the points
  aesthetics$alpha <- aesthetics$alpha %||% rep(1, n)

  # Construct guide
  guide_args$override.aes <- guide_args$override.aes %||% aesthetics
  guide <- do.call(guide_legend, guide_args)

  # Allow dummy aesthetic
  update_geom_defaults("point", list(dummy = "x"))

  dummy_geom <- geom_point(
    data = data.frame(x = rep(Inf, n), y = rep(Inf, n),
                      dummy = factor(labels)),
    aes(x, y, dummy = dummy), alpha = 0, key_glyph = key
  )
  dummy_scale <- discrete_scale(
    "dummy", palette = scales::identity_pal(), name = title,
    guide = guide
  )
  list(dummy_geom, dummy_scale)
}

# Build raw data summaries
raw_has <- d_child %>% 
	filter(!is.na(response_has)) %>%
	group_by(type, Country_full, Urb_rural, Age) %>%
	summarise(n = sum(response_has == "Yes"),
		  N = length(response_has)) %>%
	mutate(p = n / N) %>%
	rename(Country = Country_full)
raw_has$question <- "has"

raw_sim <- d_child %>% 
	filter(!is.na(response_similar)) %>%
	group_by(type, Country_full, Urb_rural, Age) %>%
	summarise(n = sum(response_similar == "Same"),
		  N = length(response_similar)) %>%
	mutate(p = n / N) %>%
	rename(Country = Country_full)
raw_sim$question <- "sim"

raw <- rbind(raw_has, raw_sim) %>%
	tidy_up()

raw_unsplit <- raw %>%
	group_by(type, Country, Age, question) %>%
	summarise(N = sum(N), n=sum(n)) %>%
	mutate(p = n / N)

# Build posterior mean prediction tables
age_range <- seq(4, 17, 0.25)
if(!file.exists("plot_tables.csv")) {
	# First run of this script, crunch the numbers to get the predictions
	# but save the results to speed up subsequent re-plots
	thoughts_has_plot <- get_country_mean_preds(d_child, m_children_has, type="thoughts", ages=age_range)
	thoughts_has_plot$question <- "has"
	thoughts_has_plot$type <- "thoughts"
	thoughts_has_plot$Urb_rural <- "combined"


	feelings_has_plot <- get_country_mean_preds(d_child, m_children_has, type="feelings", ages=age_range)
	feelings_has_plot$question <- "has"
	feelings_has_plot$type <- "feelings"
	feelings_has_plot$Urb_rural <- "combined"

	thoughts_sim_plot <- get_country_mean_preds(filter(d_child, !is.na(response_similar)), m_children_sim, type="thoughts", ages=age_range)
	thoughts_sim_plot$question <- "sim"
	thoughts_sim_plot$type <- "thoughts"
	thoughts_sim_plot$Urb_rural <- "combined"

	feelings_sim_plot <- get_country_mean_preds(filter(d_child, !is.na(response_similar)), m_children_sim, type="feelings", ages=age_range)
	feelings_sim_plot$question <- "sim"
	feelings_sim_plot$type <- "feelings"
	feelings_sim_plot$Urb_rural <- "combined"

	thoughts_has_plot_split <- get_country_mean_preds(d_child, m_children_has, type="thoughts", ages=age_range, split_urb_rural=TRUE)
	thoughts_has_plot_split$question <- "has"
	thoughts_has_plot_split$type <- "thoughts"

	feelings_has_plot_split <- get_country_mean_preds(d_child, m_children_has, type="feelings", ages=age_range, split_urb_rural=TRUE)
	feelings_has_plot_split$question <- "has"
	feelings_has_plot_split$type <- "feelings"

	thoughts_sim_plot_split <- get_country_mean_preds(filter(d_child, !is.na(response_similar)), m_children_sim, type="thoughts", ages=age_range, split_urb_rural=TRUE)
	thoughts_sim_plot_split$question <- "sim"
	thoughts_sim_plot_split$type <- "thoughts"

	feelings_sim_plot_split <- get_country_mean_preds(filter(d_child, !is.na(response_similar)), m_children_sim, type="feelings", ages=age_range, split_urb_rural=TRUE)
	feelings_sim_plot_split$question <- "sim"
	feelings_sim_plot_split$type <- "feelings"

	plot_tables <- rbind(feelings_has_plot, thoughts_has_plot, feelings_sim_plot, thoughts_sim_plot,
	                     feelings_has_plot_split, thoughts_has_plot_split, feelings_sim_plot_split, thoughts_sim_plot_split) %>%
		tidy_up()
	write_csv(plot_tables, "plot_tables.csv")
	plot_tables_split <- filter(plot_tables, Urb_rural != "combined")
	plot_tables <- filter(plot_tables, Urb_rural == "combined")
} else {
	# Load and slice up previously saved prediction tables to save time
	plot_tables <- read_csv("plot_tables.csv") %>%
		mutate(question = factor(question, c("Do animals have ...?", "Are animals' ... the same as humans'?"), ordered=T),
	               type = factor(type, c("Thoughts", "Feelings"), ordered=T))
	plot_tables_split <- filter(plot_tables, Urb_rural != "combined")
	plot_tables <- filter(plot_tables, Urb_rural == "combined")
	thoughts_has_plot <- filter(plot_tables, type=="Thoughts", str_starts(question, "Do"))
	thoughts_sim_plot <- filter(plot_tables, type=="Thoughts", str_starts(question, "Are"))
	feelings_has_plot <- filter(plot_tables, type=="Feelings", str_starts(question, "Do"))
	feelings_sim_plot <- filter(plot_tables, type=="Feelings", str_starts(question, "Are"))
	thoughts_has_plot_split <- filter(plot_tables_split, type=="Thoughts", str_starts(question, "Do"))
	thoughts_sim_plot_split <- filter(plot_tables_split, type=="Thoughts", str_starts(question, "Are"))
	feelings_has_plot_split <- filter(plot_tables_split, type=="Feelings", str_starts(question, "Do"))
	feelings_sim_plot_split <- filter(plot_tables_split, type=="Feelings", str_starts(question, "Are"))
}

# Function to make plot showing raw data points and model prediction curves
# (for positive response only)
line_and_dot_plot <- function(plot_table, raw_data, split=FALSE) {
	trub_15 <- as.vector(trubetskoy(n=15))
	trub_16 <- as.vector(trubetskoy(n=16))
	hacked_trub <- trub_15
	hacked_trub[14] <- trub_16[16]
	plot <- plot_table %>%
	ggplot() +
	geom_hline(yintercept=0.5, linetype="dotted")
        if(split) {
		plot <- plot + geom_jitter(aes(x=Age, y=p, group=Country, colour=Country, shape=Urb_rural, size=N), alpha=0.33, stroke=0, data=raw_data)
		plot <- plot + geom_line(aes(x=Age, y=pos_mean, group=paste(Country, Urb_rural), colour=Country, linetype=Urb_rural))
	} else {
		plot <- plot + geom_jitter(aes(x=Age, y=p, group=Country, colour=Country, size=N), alpha=0.33, stroke=0, data=raw_data)
		plot <- plot + geom_line(aes(x=Age, y=pos_mean, group=Country, colour=Country))
	}
	plot <- plot +
	ylim(c(0,1)) +
	ylab("Probability of positive answer") +
	xlab("Age (years)") +
	theme_light() +
	theme(plot.title = element_text(size=8),
	      axis.title = element_text(size=6),
	      axis.text = element_text(size=4),
	      legend.title = element_text(size=6),
	      legend.text = element_text(size=4),
	      legend.key.height = unit(3, "mm")
	) +
	scale_size(range=c(0.75, 2.25)) +
	scale_colour_manual(values=hacked_trub)
	return(plot)
}

# Use above function to create a faceted plot for all countries
p <- line_and_dot_plot(plot_tables, raw_unsplit) +
	facet_grid(question~type) +
	theme(strip.background = element_rect(fill="grey90"),
	      strip.text = element_text(colour="black", size=7))
ggsave("plots/countries_combined.png", plot=p, dpi=600, width=180, height=130, units="mm")

p <- line_and_dot_plot(plot_tables_split, raw, split=TRUE) +
	facet_grid(question~type) +
	theme(strip.background = element_rect(fill="grey90"),
	      strip.text = element_text(colour="black", size=7))
ggsave("plots/countries_combined_urb_rural.png", plot=p, dpi=600, width=180, height=130, units="mm")

# Function to make a stacked area plot showing relative probability of all
# three available responses
stack_pal <- c("#9FC8C8", "#1F6F6F", "#54A1A1")
country_faceted_stack_plot <- function(preds, title, legend=TRUE) {
	preds$Age <- AGE_MEAN + AGE_SD*preds$Age
	print(c(min(preds$Age), max(preds$Age)))
	intermed_name <- "Partially"
	if(str_starts(title, "Do")) {
		neg_name <- "No"
		pos_name <- "Yes"
	} else {
		neg_name <- "Different"
		pos_name <- "Same"
	}
	names(stack_pal) <- c(neg_name, pos_name, intermed_name)
	p <- ggplot(preds, aes(x=Age)) +
		geom_area(aes(y=pos_mean + part_mean + neg_mean, fill=pos_name), show.legend=legend) +
		geom_area(aes(y=part_mean + neg_mean, fill=intermed_name), show.legend=legend) +
		geom_area(aes(y=neg_mean, fill=neg_name), show.legend=legend) +
		ggtitle(title) +
		scale_x_continuous(expand=c(0, 0), limits=c(4,17), breaks=c(5, 10, 15)) +
		ylim(c(0, 1)) +
		xlab("Age (years)") +
		ylab("Probability of response") +
		scale_fill_manual(name="Response", values=stack_pal) +
	theme_light() +
	theme(axis.title = element_text(size=6),
	      axis.text = element_text(size=4),
	      strip.text = element_text(size=4, colour="black"),
	      strip.background = element_rect(fill="grey90"),
	      legend.title = element_text(size=6),
	      legend.text = element_text(size=4),
	      strip.text.x = element_text(margin = margin(b = 1, t=1)),
	      plot.title = element_text(size=8)
	) +
	facet_wrap(~Country_full, nrow=5)
	return(p)
}

# Create 4 plots (thoughts vs feelings x has vs similar) combining the two
# styles of plot created by the functions above
dummy_legend_has <- dummy_guide(
  labels = c("No", "Yes", "Partially"),
  fill   = c("#9FC8C8", "#54A1A1", "#1F6F6F"),
  colour = NA,
  title  = "Response",
  key = draw_key_polygon
)

dummy_legend_sim <- dummy_guide(
  labels = c("Different", "Same", "Partially"),
  fill   = c("#9FC8C8", "#54A1A1", "#1F6F6F"),
  colour = NA,
  title  = "Response",
  key = draw_key_polygon
)

re_form = ~(1 + type * Age | Country_full)

#### Excluding urb/rural effects:

## Thoughts has
nd <- expand_grid(type="thoughts", Country_full=unique(m_children_has$data$Country_full),
		  gender_thoughts=0, gender_feelings=0, Urb_rural = 0,
		  Age=age_range)
preds <- get_preds(m_children_has, recode_for_models(nd), re_formula=re_form)
stacked_country_plot_thoughts_has <- country_faceted_stack_plot(preds, "Do animals have thoughts?", FALSE)

## Thoughts similar
preds <- get_preds(m_children_sim, recode_for_models(nd), re_formula=re_form)
stacked_country_plot_thoughts_sim <- country_faceted_stack_plot(preds, "Are animals' thoughts the same as humans'?", FALSE)

## Feelings has
nd$type <- "feelings"
preds <- get_preds(m_children_has, recode_for_models(nd), re_formula=re_form)
stacked_country_plot_feelings_has <- country_faceted_stack_plot(preds, "Do animals have feelings?", TRUE)

## Feelings similar
preds <- get_preds(m_children_sim, recode_for_models(nd), re_formula=re_form)
stacked_country_plot_feelings_sim <- country_faceted_stack_plot(preds, "Are animals' feelings the same as humans'?", TRUE)

plotlist <- list(stacked_country_plot_thoughts_has,
                 stacked_country_plot_feelings_has,
                 stacked_country_plot_thoughts_sim,
                 stacked_country_plot_feelings_sim)
p <- wrap_plots(plotlist, nrow=2)
ggsave("plots/ordinal_combined.png", plot=p, dpi=600, width=180, height=130, units="mm")

#### For urban communities:

## Thoughts has
nd <- expand_grid(type="thoughts", Country_full=unique(m_children_has$data$Country_full),
		  gender_thoughts=0, gender_feelings=0, Urb_rural = "urban",
		  Age=age_range)
preds <- get_preds(m_children_has, recode_for_models(nd), re_formula=re_form)
urban_stacked_country_plot_thoughts_has <- country_faceted_stack_plot(preds, "Do animals have thoughts?", FALSE)

## Thoughts similar
preds <- get_preds(m_children_sim, recode_for_models(nd), re_formula=re_form)
urban_stacked_country_plot_thoughts_sim <- country_faceted_stack_plot(preds, "Are animals' thoughts the same as humans'?", FALSE)

## Feelings has
nd$type <- "feelings"
preds <- get_preds(m_children_has, recode_for_models(nd), re_formula=re_form)
urban_stacked_country_plot_feelings_has <- country_faceted_stack_plot(preds, "Do animals have feelings?", TRUE)

## Feelings similar
preds <- get_preds(m_children_sim, recode_for_models(nd), re_formula=re_form)
urban_stacked_country_plot_feelings_sim <- country_faceted_stack_plot(preds, "Are animals' feelings the same as humans'?", TRUE)

plotlist <- list(urban_stacked_country_plot_thoughts_has,
                 urban_stacked_country_plot_feelings_has,
                 urban_stacked_country_plot_thoughts_sim,
                 urban_stacked_country_plot_feelings_sim)
p <- wrap_plots(plotlist, nrow=2)
ggsave("plots/ordinal_combined_urban.png", plot=p, dpi=600, width=180, height=130, units="mm")

#### For rural communities:

## Thoughts has
nd <- expand_grid(type="thoughts", Country_full=unique(m_children_has$data$Country_full),
		  gender_thoughts=0, gender_feelings=0, Urb_rural = "rural",
		  Age=age_range)
preds <- get_preds(m_children_has, recode_for_models(nd), re_formula=re_form)
rural_stacked_country_plot_thoughts_has <- country_faceted_stack_plot(preds, "Do animals have thoughts?", FALSE)

## Thoughts similar
preds <- get_preds(m_children_sim, recode_for_models(nd), re_formula=re_form)
rural_stacked_country_plot_thoughts_sim <- country_faceted_stack_plot(preds, "Are animals' thoughts the same as humans'?", FALSE)

## Feelings has
nd$type <- "feelings"
preds <- get_preds(m_children_has, recode_for_models(nd), re_formula=re_form)
rural_stacked_country_plot_feelings_has <- country_faceted_stack_plot(preds, "Do animals have feelings?", TRUE)

## Feelings similar
preds <- get_preds(m_children_sim, recode_for_models(nd), re_formula=re_form)
rural_stacked_country_plot_feelings_sim <- country_faceted_stack_plot(preds, "Are animals' feelings the same as humans'?", TRUE)

plotlist <- list(rural_stacked_country_plot_thoughts_has,
                 rural_stacked_country_plot_feelings_has,
                 rural_stacked_country_plot_thoughts_sim,
                 rural_stacked_country_plot_feelings_sim)
p <- wrap_plots(plotlist, nrow=2)
ggsave("plots/ordinal_combined_rural.png", plot=p, dpi=600, width=180, height=130, units="mm")

### Plot adult responses
m_adults_has <- readRDS("m_adults_has.rds")
m_adults_sim <- readRDS("m_adults_sim.rds")

adult_stack_plot <- function(preds, title, xaxis_lab=FALSE) {
	preds <- rename(preds, Country = Country_full)
	ranked_countries <- preds %>%
		arrange(pos_mean) %>%
		pull(Country)
	preds$Country <- factor(preds$Country, ranked_countries, ordered=TRUE)
	intermed_name <- "Partially"
	if(str_starts(title, "Do")) {
		neg_name <- "No"
		pos_name <- "Yes"
	} else {
		neg_name <- "Different"
		pos_name <- "Same"
	}
	names(stack_pal) <- c(neg_name, pos_name, intermed_name)
	p <- ggplot(preds, aes(x=Country)) +
		geom_bar(aes(y=pos_mean + part_mean + neg_mean, fill=pos_name), stat="identity", position="stack") +
		geom_bar(aes(y=part_mean + neg_mean, fill=intermed_name), stat="identity", position="stack") +
		geom_bar(aes(y=neg_mean, fill=neg_name), position="stack", stat="identity") +
		ggtitle(title) +
		ylab("Probability of response") +
		scale_fill_manual(name="Response", values=stack_pal) +
	theme_light() +
	theme(axis.title = element_text(size=6),
	      axis.title.y = element_text(size=5),
	      axis.text = element_text(size=4),
	      strip.text = element_text(size=4, colour="black"),
	      strip.background = element_rect(fill="grey90"),
	      legend.title = element_text(size=6),
	      legend.text = element_text(size=4),
	      legend.key.size = unit(0.5, "cm"),
	      strip.text.x = element_text(margin = margin(b = 1, t=1)),
	      plot.title = element_text(size=8),
	      plot.margin = margin(1.5, 5.5, 1.5, 5.5, "points")
	)
	if(!xaxis_lab) {
		p <- p + theme(axis.title.x = element_blank())
	}
	return(p)
}

nd <- expand_grid(type="thoughts", Country_full=unique(m_adults_has$data$Country_full),
		  gender_thoughts=0, gender_feelings=0, Urb_rural = 0, Age=0)
preds <- get_preds(m_adults_has, recode_for_models(nd), re_formula=re_form)
stacked_adult_plot_thoughts_has <- adult_stack_plot(preds, "Do animals have thoughts?", FALSE)

preds <- get_preds(m_adults_sim, recode_for_models(nd), re_formula=re_form)
stacked_adult_plot_thoughts_sim <- adult_stack_plot(preds, "Are animals' thoughts the same as humans'?", FALSE)

nd$type="feelings"
preds <- get_preds(m_adults_has, recode_for_models(nd), re_formula=re_form)
stacked_adult_plot_feelings_has <- adult_stack_plot(preds, "Do animals have feelings?", FALSE)

preds <- get_preds(m_adults_sim, recode_for_models(nd), re_formula=re_form)
stacked_adult_plot_feelings_sim <- adult_stack_plot(preds, "Are animals' feelings the same as humans'?", TRUE)

plotlist <- list(stacked_adult_plot_thoughts_has,
                 stacked_adult_plot_feelings_has,
                 stacked_adult_plot_thoughts_sim,
                 stacked_adult_plot_feelings_sim)
p <- wrap_plots(plotlist, ncol=1)
ggsave("plots/adult_ordinal.png", plot=p, dpi=600, width=180, height=120, units="mm")