Canonical response function analysis in Stan

crf_demo.R

library(tidyverse)
library(rstan)
rstan_options(auto_write = TRUE)

curve(dweibull(x,shape=1.75,scale=1e3),from=0,to=3e3)
one_response = dweibull(x=1:3e3,shape=1.75,scale=1e3)
one_response = one_response/max(one_response)

# create a signal with pulses at t=2e3, t=4e3 & t=5e3
obs = rep(0,8e3)

skip = 2e3
i = (skip+1):(skip+length(one_response))
obs[i] = obs[i] + one_response

skip = 4e3
i = (skip+1):(skip+length(one_response))
obs[i] = obs[i] + one_response

skip = 5e3
i = (skip+1):(skip+length(one_response))
obs[i] = obs[i] + one_response

plot(obs,type='l')

obs = obs+rnorm(length(obs),0,.1)
plot(obs,type='l')


# Generate the predictor matrix ----
pred_mat = matrix(0,nrow=length(obs),ncol=length(obs))
for(i in 1:ncol(pred_mat)){
	temp = i+length(one_response)-1
	if(length(obs)>=temp){
		pred_mat[i:temp,i] = one_response
	}else{
		pred_mat[i:length(obs),i] = one_response[1:(length(one_response)-(temp-length(obs)))]
	}
}

#keep only a subset of the columns for reduced compute
pred_mat = pred_mat[,round(seq(1,length(obs)-100,length.out = 50))]

plot(pred_mat[,1],type='l')
plot(pred_mat[,2],type='l')
plot(pred_mat[,ncol(pred_mat)-1],type='l')
plot(pred_mat[,ncol(pred_mat)],type='l')


mod = rstan::stan_model('pooled_regression.stan')
post = rstan::sampling(
	object = mod
	, data = list(
		nX = ncol(pred_mat)
		, nY = nrow(pred_mat)
		, X = pred_mat
		, Y = obs
	)
	, seed = 1
	, iter = 2e2
	, chains = 4
	, cores = 4
)
print(
	post
	, pars = c('intercept','noise','pool')
	, probs = c(.025,.975)
	, digits = 2
)
print(
	post
	, pars = 'coefs'
	, probs = c(.025,.975)
	, digits = 2
)
coefs = rstan::extract(post,'coefs')[[1]]
coefs = tibble::as_tibble(data.frame(coefs))
coefs = tidyr::gather(coefs, key, value)
coefs$key = as.numeric(gsub('X','',coefs$key))
coefs %>%
	dplyr::group_by(key) %>%
	dplyr::summarise(
		lo95 = quantile(value,.025)
		, hi95 = quantile(value,.975)
		, sig = (lo95>0)|(hi95<0)
		, lo50 = quantile(value,.25)
		, hi50 = quantile(value,.75)
		, med = quantile(value,.5)
	) %>%
	ggplot(
		mapping = aes(
			x = key
			, y = med
			, colour = sig
		)
	) +
	geom_errorbar(
		mapping = aes(
			ymin = lo95
			, ymax = hi95
		)
		, size = 1
		, width = 0
		, show.legend = F
	)+
	geom_errorbar(
		mapping = aes(
			ymin = lo50
			, ymax = hi50
		)
		, size = 2
		, width = 0
		, show.legend = F
	)+
	geom_line(
		show.legend = F
		, colour='black'
	)+
	geom_point(
		size = .5
		, show.legend = F
		, colour='black'
	)

pooled_regression.stan

data{
	int nX ;
	int nY ;
	matrix[nY,nX] X ;
	vector[nY] Y ;
}
parameters{
	real intercept ;
	vector[nX] coefs;
	real<lower=0> noise ;
	real<lower=0> pool ;
}
model{
	intercept ~ normal(0,1) ;
	noise ~ normal(0,1) ;
	pool ~ normal(0,1) ;
	coefs ~ cauchy(0,pool) ;
	Y ~ normal(X*coefs+intercept,noise) ;
}