Applying discrete effects about mvgam HOT 4 CLOSED

library(mvgam)
library(marginaleffects)
library(dplyr)
library(ggplot2); theme_set(theme_bw())

# Simulate data
t <- seq(from = 0, to = 20, length.out = 1000)
v1 <- rep(1, 500)
v2 <- rep(2, 500)
noise <- rnorm(1000, mean = 0.5, sd = 1)
values <- c(v1,v2) + noise
a <- rep("a", 500)
b <- rep("b", 500)
df <- data.frame(
  time = t,
  value = values,
  # Don't use the name 'group' as it causes problems for marginaleffects
  grp = c(a,b)
) |> 
  mutate(value = if_else(time > 15, value - ((time-14)^1.2), value)) |>
  # Add a time since the start of condition 'b' as a new variable
  mutate(time_since_b = c(rep(0, 500),
                          1:500))
df <- bind_rows(df, df, df, df, df) |> 
  mutate(time = row_number(),
         season = as.integer(time / 1000),
         series = factor("A"))

# Fit a model that estimates the nonlinear, delayed impact of entering
# condition 'b'
mod <- mvgam(value ~ grp + 
               s(time_since_b, k = 15),
             data = df,
             family = gaussian(),
             # Meanfield Variation inference should be just fine for this model,
             # and will be lightning fast
             algorithm = 'meanfield',
             backend = 'cmdstanr')
#> Compiling Stan program using cmdstanr
#> 
#> Start sampling
#> ------------------------------------------------------------ 
#> EXPERIMENTAL ALGORITHM: 
#>   This procedure has not been thoroughly tested and may be unstable 
#>   or buggy. The interface is subject to change. 
#> ------------------------------------------------------------ 
#> Gradient evaluation took 0.000375 seconds 
#> 1000 transitions using 10 leapfrog steps per transition would take 3.75 seconds. 
#> Adjust your expectations accordingly! 
#> Begin eta adaptation. 
#> Iteration:   1 / 250 [  0%]  (Adaptation) 
#> Iteration:  50 / 250 [ 20%]  (Adaptation) 
#> Iteration: 100 / 250 [ 40%]  (Adaptation) 
#> Iteration: 150 / 250 [ 60%]  (Adaptation) 
#> Iteration: 200 / 250 [ 80%]  (Adaptation) 
#> Success! Found best value [eta = 1] earlier than expected. 
#> Begin stochastic gradient ascent. 
#>   iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
#>    100      -623427.781             1.000            1.000 
#>    200      -386878.079             0.806            1.000 
#>    300      -320377.370             0.606            0.611 
#>    400        -9371.998             8.751            1.000 
#>    500        -7400.788             7.054            0.611 
#>    600        -7129.433             5.885            0.611 
#>    700        -7051.640             5.046            0.266 
#>    800        -7060.191             4.415            0.266 
#>    900        -7063.776             3.925            0.208 
#>   1000        -7053.622             3.532            0.208 
#>   1100        -7019.789             3.433            0.038   MAY BE DIVERGING... INSPECT ELBO 
#>   1200        -7045.633             3.372            0.011   MAY BE DIVERGING... INSPECT ELBO 
#>   1300        -7126.481             3.352            0.011   MAY BE DIVERGING... INSPECT ELBO 
#>   1400        -7040.557             0.035            0.011 
#>   1500        -7017.499             0.009            0.005   MEAN ELBO CONVERGED   MEDIAN ELBO CONVERGED 
#> Drawing a sample of size 500 from the approximate posterior...  
#> COMPLETED. 
#> Finished in  9.3 seconds.
summary(mod, include_betas = FALSE)
#> GAM formula:
#> value ~ grp + s(time_since_b, k = 15)
#> 
#> Family:
#> gaussian
#> 
#> Link function:
#> identity
#> 
#> Trend model:
#> None
#> 
#> N series:
#> 1 
#> 
#> N timepoints:
#> 5000 
#> 
#> Status:
#> Fitted using Stan 
#> 1 chains, each with iter = 500; warmup = ; thin = 1 
#> Total post-warmup draws = 500
#> 
#> 
#> Observation error parameter estimates:
#>              2.5%  50% 97.5% Rhat n.eff
#> sigma_obs[1] 0.95 0.98     1  NaN   NaN
#> 
#> GAM coefficient (beta) estimates:
#>             2.5%  50% 97.5% Rhat n.eff
#> (Intercept) 0.34 0.37  0.39  NaN   NaN
#> grpb        0.96 1.00  1.00  NaN   NaN
#> 
#> Approximate significance of GAM smooths:
#>                  edf Ref.df Chi.sq p-value    
#> s(time_since_b) 9.94     14  17780  <2e-16 ***
#> ---
#> Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1
#> 
#> Posterior approximation used: no diagnostics to compute

# Plot the delayed effect of entering the 'b' condition
plot_predictions(mod, by = c("time_since_b"), 
                 newdata = datagrid(time_since_b = unique,
                                    group = 'b'),
                 points = 0.2)

# Plot an estimate of the structural change when first entering condition 'b'
plot_predictions(mod, by = 'grp',
                 newdata = datagrid(time_since_b = 0,
                                    group = unique))

# Plot full predictions
fc <- hindcast(mod)
plot(fc)