Splines to compare treatments

Say I have some time-series data for 2 groups: 1 group received placebo and the other received treatment. Can the effect of treatment be modeled with splines similar to a linear model. For example:

# DATA
n_A <- 10 # Number of rats in Group A
n_B <- 10 # Number of rats in Group B

m_NREMS <- 4 # number of measurements for NREMS.  Here, we'll assume that the rhythyms of NREMS express over 6-h blocks, giving us 4 blocks.

NREMS_A <- data.frame("Blk0" = rep(0, n_A),
                      "Blk1" = rgamma(n = n_A, shape=144, rate=6/5), # 2h NREMS
                      "Blk2" = rgamma(n=n_A, shape=324, rate=9/5), # 3h NREMS
                      "Blk3" = rgamma(n=n_A, shape=900, rate=3), # 5h NREMS
                      "Blk4" = rgamma(n=n_A, shape=576, rate=12/5)) # 4h NREMS

NREMS_B <- data.frame("Blk0" = rep(0, n_B),
                      "Blk1" = rgamma(n = n_B, shape=144, rate=6/5) + rnorm(n=n_B, mean=30, sd=6), # Similar to NREMS_A, but with the effect of S added
                      "Blk2" = rgamma(n=n_B, shape=324, rate=9/5) + rnorm(n=n_B, mean=20, sd=5),
                      "Blk3" = rgamma(n=n_B, shape=900, rate=3) + rnorm(n=n_B, mean=10, sd=4),
                      "Blk4" = rgamma(n=n_B, shape=576, rate=12/5) + rnorm(n=n_B, mean=5, sd=3))

NREMS_A_cuml <- NREMS_A
for(i in 2:ncol(NREMS_A)) {
  NREMS_A_cuml[,i] <- NREMS_A_cuml[,i] + NREMS_A_cuml[,i-1]
}

NREMS_B_cuml <- NREMS_B
for(i in 2:ncol(NREMS_B)) {
  NREMS_B_cuml[,i] <- NREMS_B_cuml[,i] + NREMS_B_cuml[,i-1]
}

NREMS_All <- data.frame("NREMS" = c(as.numeric(unlist(NREMS_A_cuml)), as.numeric(unlist(NREMS_B_cuml))),
                        "Group" = c(rep("A", (m_NREMS+1)*n_A), c(rep("B", (m_NREMS+1)*n_B))),
                        "Block" = c(rep(0:4, each=n_A), c(rep(0:4, each=n_B))))
NREMS_All$treatment <- ifelse(test = NREMS_All$Group=="A",
                              yes = FALSE,
                              no = TRUE)
NREMS_All$minute <- NREMS_All$Block*360

# SPLINES
num_knots <- 4
knot_list <- quantile(NREMS_All$minute, probs=seq(from=0, to=1, length.out=num_knots))
knot_degree <- 3

B <- bs(NREMS_All$minute,
        knots=knot_list[-c(1, num_knots)],
        degree=knot_degree,
        intercept=TRUE)

# MODEL
NREMS_model_1 <- quap(
  alist(
    NREMS ~ dnorm(mu, sigma),
      mu <- a +
            B %*% w[treatment],
        a ~ dnorm(180, 10),
        w[treatment] ~ dnorm(1, 1),
      sigma ~ dexp(1)
  ), data=list(NREMS=NREMS_All$NREMS,
               treatment=as.integer(NREMS_All$treatment),
               B=B),
     start=list(w=rep(0, ncol(B)))
)

I think that the problem is that B is a matrix and w is a vector such that w[treatment] is read as an index on the vector w, but I want it to be work like it does in the random effects models.

Any help much appreciated.

rmcelreath / stat_rethinking_2023

Splines to compare treatments #16