# Online resource 5
# Title:A method for the taphonomic assessment of bone tools using 3D surface texture analysis of bone microtopography
# Authors: Naomi L. Martisius a,b*, Shannon P. McPherron b  , Ellen Schulz-Kornas c,d,b, Marie Soressi e,b, Teresa E. Steele a,b
# Affiliations: a Department of Anthropology, University of California, Davis, Davis, CA, USA; b Department of Human Evolution, Max Planck Institute for Evolutionary Anthropology, Leipzig, Germany; c Department of Cariology, Endodontology and Periodontology, University of Leipzig, Leipzig, Germany; d Max Planck Weizmann Center for Integrative Archaeology and Anthropology, Max Planck Institute for Evolutionary Anthropology, Leipzig, Germany; e Faculty of Archaeology, Leiden University, Leiden, The Netherlands
# * Corresponding author: Naomi L. Martisius (nlmartisius@ucdavis.edu)

# R script to read in 3D surface texture parameters, perform basic data management and transformations,
# setup data objects to pass to stan, and fit model M1 via rstan. 

library(StanHeaders)
library(ggplot2)
library(rstan)
library(loo)

# multithread options
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())


########### DATA INPUT, MANAGEMENT, AND TRANSFORMATION. ##############

d <- read.csv("ESM2.csv")

# Data frame containing only logged features for use below. 
d.manova <- log(d[,c("Sa", "Sal", "Spc", "Smr1", "IsT")])

# Integer-coded ID and Level, for modeling.
ID.int <- d$ID
ID.int.factor <- factor(ID.int)
ID.int <- as.numeric(ID.int.factor)
Source.int <- d$Source
Source.int.factor <- factor(Source.int)
Source.int <- as.numeric(Source.int.factor)

# Set-up data to pass to Stan. 
# Center data matrix at feature means.
col.means <- apply(d.manova, MARGIN=2, FUN="mean")
y.centered <- sweep(d.manova, MARGIN=2, STATS=col.means) 

##### MODEL M0: MULTIVARIATE T OBSERVATIONS WITH RANDOM EFFECTS FOR SPECIMEN AND SOURCE. ######
data_list <- list(
  K= dim(d.manova)[2],
  J= length(unique(d$ID)),
  M= length(unique(d$Source)),
  N= dim(d.manova)[1],
  y= as.matrix(y.centered), ## features centered at zero
  specimen= ID.int,
  lot= Source.int
)

m0.stan <- stan(
  file = "ESM3.stan", 
  model_name = "M0", 
  data=data_list, iter=6000, warmup=4000, chains=4,
  control = list(stepsize = 0.5, adapt_delta = 0.99, max_treedepth = 14))

# Extract the average Mahalanobis squared distances over realizations. 
# divide by K = number of features to scale for comparison with F distribution
Maha_sqd_scaled <- summary(m0.stan, pars=c("Maha_sqd"))$summary[,"mean"]/dim(d.manova)[2]

# Plot the distances -vs- F distribution quantiles.
# Extract the average mulivariate T degrees of freedom over samples. 
t.df <- summary(m0.stan, pars=c("DF_obs"))$summary[,"mean"]
# Plot the distances vs. F distribution (df1 = K, df2 = t.df) quantiles.
plot(qf(ppoints(Maha_sqd_scaled), df1=dim(d.manova)[2], df2=t.df), sort(Maha_sqd_scaled),
     xlab="F-distribution quantile", ylab="Squared, Scaled Mahalanobis distance", main="M0 Goodness of Fit",
     bty="n")
abline(a=0, b=1)

# Extract log-likelihoods of each observation for each posterior sample. 
logL.loo <- extract_log_lik(m0.stan)
waic(logL.loo)  
loo(logL.loo) 


##### MODEL M1: fixed effects for Specimen Type added to M0. #####

# Set up design matrix for fixed effects.
# Any one of the features can go on the left-hand side of "~" below.  
design.frame <- model.frame(Sa ~ -1 + Specimen.Type, data=d)

design.mat <- model.matrix(Sa ~ -1 + Specimen.Type, data=design.frame)
# Obtain singular values of the design matrix. Check to see that the design matrix is not singular. 
svd(design.mat)$d # The design matrix is non-singular. 

data_list <- list(
  P= dim(design.mat)[2], 
  K= dim(d.manova)[2],
  J= length(unique(d$ID)),
  M= length(unique(d$Source)),
  N= dim(d.manova)[1],
  y= as.matrix(y.centered), ## features centered at zero
  x= design.mat, 
  specimen= ID.int,
  lot= Source.int
)

m1.stan <- stan(
  file = "ESM4.stan",
  model_name = "M1", 
  data=data_list, iter=6000, warmup=4000, chains=4,
  control = list(stepsize = 0.5, adapt_delta = 0.99, max_treedepth = 14))

# Extract the average Mahalanobis squared distances over realizations. 
# divide by K = number of features to scale for comparison with F distribution
Maha_sqd_scaled <- summary(m1.stan, pars=c("Maha_sqd"))$summary[,"mean"]/dim(d.manova)[2]

# Plot the distances -vs- F distribution quantiles.
# Extract the average mulivariate T degrees of freedom over samples. 
t.df <- summary(m1.stan, pars=c("DF_obs"))$summary[,"mean"]
# Plot the distances vs. F distribution (df1 = K, df2 = t.df) quantiles.
plot(qf(ppoints(Maha_sqd_scaled), df1=dim(d.manova)[2], df2=t.df), sort(Maha_sqd_scaled),
     xlab="F-distribution quantile", ylab="Squared, Scaled Mahalanobis distance", main="M1 Goodness of Fit",
     bty="n")
abline(a=0, b=1)

# Extract log-likelihoods of each observation for each posterior sample. 
logL.loo <- extract_log_lik(m1.stan)
waic(logL.loo)  
loo(logL.loo)