Change to always return warmup period when save_warmup is TRUE

[Cole-Monnahan-NOAA]

You previously commited this change to fix a bug I reported in this issue. The main issue was fixed but it introduced a new bug.

When the user sets save_warmup=TRUE the diagnostic output gets trimmed, which is correct for checking diagnostics, but it then returns that. So downstream the user cannot look at the sampler parameters during the warmup. This e.g., breaks a link to something like ShinyStan. It is also unexpected from a user perspective that save_warmup does not control whether the warmup samples are saved.

This PR just moves the filtering of warmup samples when calling check_hmc_diagnostics.

It passes tests locally. A quick reprex from the README shows the behavior:

fit1 <- stan_sample(loglik_fun, inits, additional_args = list(y),
                   lower = c(-Inf, 0), # Enforce a positivity constraint for SD
                   num_chains = 1, seed = 1234, save_warmup=TRUE)
fit1@diagnostics |> nrow()

fit2 <- stan_sample(loglik_fun, inits, additional_args = list(y),
                    lower = c(-Inf, 0), # Enforce a positivity constraint for SD
                    num_chains = 1, seed = 1234, save_warmup=FALSE)
fit2@diagnostics |> nrow()

Where the first has 2000 rows and the second 1000 rows.

[andrjohns]

Ahhh great catch and many thanks for the PR!