--- title: "jmpost Quickstart" package: jmpost format: html: html-math-method: mathml vignette: > %\VignetteIndexEntry{jmpost Quickstart} %\VignetteEncoding{UTF-8} %\VignetteEngine{quarto::html} --- ```{r} #| label: setup #| include: false knitr::opts_chunk$set( collapse = TRUE, comment = "#>" ) library(jmpost) ``` # Introduction The `jmpost` package is used to fit joint models of time-to-event endpoints with an association term parameterised from a longitudinal submodel. The package was originally designed for the context of fitting tumour growth inhibition models to inform overall survival but is not limited to this application. It is important to note that `jmpost` is just a wrapper on top of the Stan programming language and leans heavily on the wider Stan ecosystem. For example no model diagnostics are provided by `jmpost` with users being recommended to use packages like `bayesplot` for this functionality. This vignette provides a basic outline of how to fit a joint model using `jmpost`. For more detailed information on the package and its functions please refer to the package documentation. # Model Specification Each `JointModel` needs to be specified with three parts: 1. `longitudinal`: The model for the longitudinal outcomes. 1. `survival`: The model for the survival outcomes. 1. `link`: The link that specifies how the `longitudinal` model parameters enter the `survival` model. ## Default Options Let's first specify a very simple joint model with: 1. A random slope model for the longitudinal outcome. 1. A Weibull proportional hazards model for the survival outcome. 1. The link between the two models is that the random slope from the longitudinal model enters as a product with a link coefficient into the linear predictor of the survival model. ```{r} #| label: simple_model simple_model <- JointModel( longitudinal = LongitudinalRandomSlope(), survival = SurvivalWeibullPH(), link = linkDSLD() ) ``` Note that here we use all the default options for the two models and the link, in particular the prior distributions and the initial values in the MCMC chain for the parameters are automatically chosen. We can see this from the arguments of the constructors (or from the help page): ```{r} #| label: constructor_args args(LongitudinalRandomSlope) ``` So here we see that the Longitudinal Random Slope model has 4 parameters that we can define a prior for. ## Specifying Priors We can alternatively also specify the prior distributions for the parameters manually. This is important in practice to obtain a meaningful model specification and hence converging MCMC chains that allow to estimate the posterior distributions. For the random slope model for the longitudinal outcome, we can e.g. say: ```{r} #| label: random_slope_prior_inits random_slope_model <- LongitudinalRandomSlope( intercept = prior_normal(40, 5), slope_mu = prior_normal(10, 2) ) ``` This sets the prior for the `intercept` to be a $N(40, 5)$ distribution and the prior for the `slope_mu` parameter to be a $N(10, 2)$ distribution. For the survival models, we can set different independent normal priors for the coefficients of the covariates in the linear predictor using the `prior_normal_vector()` specification. Here we can give a vector of the same length as the number of design matrix columns for both the means or standard deviations of the normal distributions, or we can give a single value that will be repeated for all covariates. For example, say we have 3 design matrix columns, i.e. coefficients, in the linear predictor of the survival model, then we can set the priors for these coefficients as follows: ```{r} #| label: survival_prior_inits survival_model <- SurvivalWeibullPH( beta = prior_normal_vector(mus = c(-1, 0, 5), sigmas = c(1, 1, 10)) ) survival_model ``` In order to know which position corresponds to which design matrix column and thus coefficient, we can use the `model.matrix()` function on the `DataSurvival` object. This is explained in the section "Formatting Data" below in detail. ### Priors on a transformed scale Sometimes prior information is available on a different scale from the model parameter. Use `transform` to specify the mapping **from the model parameter to the prior density scale**. For example, `sigma` is a residual standard deviation, so a prior on its variance can be specified as follows: ```{r} #| label: variance_scale_prior variance_prior_model <- LongitudinalRandomSlope( sigma = prior_normal(100, 100, transform = prior_transform_square()) ) ``` Here `sigma^2` has a normal distribution restricted to positive values. `jmpost` includes the Jacobian and maps model-imposed bounds to the variance scale when normalising the truncated density. The square transformation uses the positive branch and restricts the declared parameter to non-negative values. For log-scale intercepts, use `prior_transform_exp()`: for example, `prior_normal(60, 60, transform = prior_transform_exp())` places the normal prior on the exponential of the intercept. Use `prior_transform_exp(scale = 2)` for a prior on `exp(2 * intercept)`, such as a variance when the intercept is a log standard deviation. These transformations change the prior; they are independent of covariate predictor parametrizations. The same specification can be written with a wrapper: ```{r} #| label: transformed_prior_wrapper variance_prior <- prior_transform( prior_normal(100, 100), prior_transform_square() ) ``` Both forms support scalar distribution helpers, including `prior_normal()` and `prior_invgamma()`. A scalar distribution is applied independently to each component of a vector parameter. Vector-specific helpers such as `prior_normal_vector()` and the horseshoe prior do not accept transformations. Initial values are drawn on the density scale within the mapped bounds, inverse-transformed and then shrunk on the model parameter scale. ### Flat priors Use `prior_flat()` to add no prior density term on the declared parameter scale. Its `init` distribution only supplies initial values: ```{r} #| label: flat_coefficient_prior flat_coefficient_model <- SurvivalWeibullPH( beta = prior_flat(init = prior_normal(0, 1)) ) ``` Model-imposed bounds still apply. An unbounded flat prior is improper, and a proper posterior is not guaranteed. A flat prior is specific to its scale: flat on a log coefficient does not imply flat on its exponential. ## Separate Models It is also possible to not link the longitudinal and the survival models, by using the special `linkNone()` link specification. For example, ```{r} #| label: simple_model_no_link simple_model_no_link <- JointModel( longitudinal = LongitudinalRandomSlope(), survival = SurvivalWeibullPH(), link = linkNone() ) ``` would allow to fit the two models separately, but in the same MCMC chain. ## Single Models It is also possible to specify only the longitudinal or the survival model. Then these can be estimated on their own with separate MCMC chains. ```{r} #| label: single_models single_longitudinal <- JointModel( longitudinal = LongitudinalRandomSlope() ) single_survival <- JointModel( survival = SurvivalWeibullPH() ) ``` # Data Preparation Before we can fit the models, we need to prepare the data in the right format. ## Simulating Data Here we start from a simulated data set. - We assign 50 subjects each to the two treatment arms. - We use a time grid from 1 to 2000, e.g. specifying the days after randomization. - We use an exponentially distributed censoring time with mean of 9000 days. - We use a categorical covariate with three levels A, B and C in the overall survival model, drawn uniformly from the three levels. (Note that this is hardcoded at the moment, so the levels need to be A, B, C.) - We use another continuous covariate in the overall survival model generated from a standard normal distribution, with coefficient 0.3. - For the longitudinal outcome, we draw the values from a random slope model with the given parameters. - For the survival outcome, we draw the true value from a Weibull model. Note that it is fairly easy to put here another choice, you just need to specify a function of `time` returning the log baseline hazard under the given survival model. So let's run the code for that: ```{r} #| label: simulate_data set.seed(129) sim_data <- SimJointData( design = list( SimGroup(50, "Arm-A", "Study-X"), SimGroup(50, "Arm-B", "Study-X") ), longitudinal = SimLongitudinalRandomSlope( times = c(1, 20, 50, 100, 150, 200, 250, 300), intercept = 30, slope_mu = c(1, 2), slope_sigma = 0.2, sigma = 20, link_dsld = 0.1 ), survival = SimSurvivalWeibullPH( lambda = 1 / 300, gamma = 0.97, time_max = 2000, time_step = 1, lambda_cen = 1 / 9000, beta_cat = c( "A" = 0, "B" = -0.1, "C" = 0.5 ), beta_cont = 0.3 ) ) ``` We might get a message here that a few subjects did not die before the day 2000, but this is not of concern. Basically it just gives us a feeling of how many survival events are included in the data set. ## Formatting Data Next we bring the data into the right format. We start with extracting data into individual data sets, and then reducing the longitudinal data to specific time points. ```{r} #| label: individual_data os_data <- sim_data@survival long_data <- sim_data@longitudinal ``` Let's have a quick look at the format: The survival data has: - subject ID - time point - continuous covariate value - categorical covariate level - event indicator (1 for observed, 0 for censored) - study ID - treatment arm ```{r} #| label: os_data_structure head(os_data) ``` The longitudinal data has: - subject ID - time point - sum of longest diameters (SLD) - study ID - treatment arm - observation flag ```{r} #| label: lm_data_structure head(long_data) ``` Finally, we wrap these in the data formatting functions. Here the mapping of the column names to the required variables happens. This means that in our applications we don't have to use the same variable names as seen above, but we can use custom names and then apply the mapping here. ```{r} #| label: data_joint joint_data <- DataJoint( subject = DataSubject( data = os_data, subject = "subject", arm = "arm", study = "study" ), survival = DataSurvival( data = os_data, formula = Surv(time, event) ~ cov_cat + cov_cont ), longitudinal = DataLongitudinal( data = long_data, formula = sld ~ time, threshold = 5 ) ) ``` Note that we can also actively look into the design matrix of the `DataSurvival` object. This can be helpful if we want to know which coefficients are at which position, e.g. for setting the right prior means and standard deviations with `prior_normal_vector()`. It works like this: ```{r} #| label: design_matrix head(model.matrix(joint_data@survival)) ``` So here we see that the first coefficient corresponds to category `B` and the second coefficient corresponds to category `C` of the categorical covariate `cov_cat`, while the third coefficient corresponds to the continuous covariate `cov_cont`. # Model Fitting Now let's have a look how we can fit the (joint) models. ## Debugging Stan Code It is always possible to read out the Stan code that is contained in the `JointModel` object, using `write_stan()`: ```{r} #| label: debug_stan tmp <- tempfile() write_stan(simple_model, destination = tmp) first_part <- head(readLines(tmp), 10) cat(paste(first_part, collapse = "\n")) ``` ## Sampling Parameters Finally, `sampleStanModel()` is kicking off the MCMC sampler via `cmdstanr` in the backend. Note that in practice you would increase the number of warm-up and sampling iterations. ```{r} #| label: sampling mcmc_results <- sampleStanModel( simple_model, data = joint_data, iter_sampling = 500, iter_warmup = 500, chains = 1, parallel_chains = 1 ) ``` ## Convergence checks After the sampling finishes, we can inspect the parameter distributions. This is using the `cmdstanr` functions, because the `results` element is of class `CmdStanMCMC`. ```{r} #| label: inspect_distributions vars <- c( "lm_rs_intercept", "lm_rs_slope_mu", "lm_rs_slope_sigma", "lm_rs_sigma", "link_dsld", "sm_weibull_ph_lambda", "sm_weibull_ph_gamma", "beta_os_cov" ) cmdstanr::as.CmdStanMCMC(mcmc_results)$summary(vars) ``` We can already see here from the `rhat` statistic whether the MCMC sampler converged - values close to 1 indicate convergence, while values larger than 1 indicate divergence. In general, convergence is sensitive to the choice of: - Priors - Initial values - Sufficient warm-up iterations If you don't achieve convergence, then play around with different choices of the above. # Plotting We can now proceed towards investigating the results of the MCMC chain with plots. ## Trace plots Using the `draws()` method on the `results` element, we can extract the samples in a format that is understood by `bayesplot`. We can e.g. look at some simple trace plots: ```{r} #| fig-width: 6 vars_draws <- cmdstanr::as.CmdStanMCMC(mcmc_results)$draws(vars) library(bayesplot) mcmc_trace(vars_draws) ``` ## Longitudinal fit plots Using the `longitudinal()` method we can extract the longitudinal fit samples from the result, and then plot them for all subjects or those that we are interested in. For illustration, we will plot here the first 10 subjects: ```{r} #| fig-width: 6 selected_subjects <- head(os_data$subject, 10) long_quantities <- LongitudinalQuantities( mcmc_results, grid = GridFixed( subjects = selected_subjects ) ) as.data.frame(long_quantities) |> head() summary(long_quantities) |> head() autoplot(long_quantities) ``` ## Survival fit plots And using the `SurvivalQuantities()` method we can do the same for the estimated survival functions. ```{r} #| fig-width: 6 surv_quantities <- SurvivalQuantities( mcmc_results, grid = GridFixed( subjects = selected_subjects ) ) as.data.frame(surv_quantities) |> head() summary(surv_quantities) |> head() autoplot(surv_quantities) ``` We can also aggregate the estimated survival curves from groups of subjects, using the corresponding method. ```{r} #| warning: false #| fig-width: 6 surv_quantities <- SurvivalQuantities( mcmc_results, grid = GridGrouped( groups = split(os_data$subject, os_data$arm) ) ) autoplot(surv_quantities, add_km = TRUE) ``` ## Predicting Survival Quantities for Hypothetical Subjects The `SurvivalQuantities()` method can also be used to predict survival quantities for hypothetical subjects with any arbitrary covariate and longitudinal model parameter values. This is done via passing a `GridPrediction()` object to the method where the `newdata` and `params` arguments are used to specify the desired values for the covariates and longitudinal model parameters respectively. For example below we use these functions to predict the survival distribution for a hypothetical subject: ```{r} surv_quantities <- SurvivalQuantities( mcmc_results, grid = GridPrediction( times = seq(1, 800, by = 50), newdata = dplyr::tibble( cov_cat = c("A", "B"), cov_cont = c(1.2, 2) ), params = list( intercept = 40, slope = 0.05 ) ) ) autoplot(surv_quantities) ``` The exact names that are required for the `params` argument (i.e. the longitudinal model parameters) can be found by running `getPredictionNames()` on the longitudinal model object. For example: ```{r} getPredictionNames(LongitudinalGSF()) ``` ## Brier Score The `brierScore()` method can be used to extract the Brier Scores (predictive performance measure) from our `SurvivalQuantities` object. ```{r} sq <- SurvivalQuantities( mcmc_results, grid = GridFixed(times = c(1, 50, 100, 400, 800)), type = "surv" ) brierScore(sq) ``` ## Initial Values By default `jmpost` will set the initial values for all parameters to be a random value drawn from the prior distribution that has been shrunk towards the mean of said distribution e.g. for a `prior_normal(4, 2)` the initial value for each chain will be: ``` 4 * shrinkage_factor + rnorm(1, 4, 2) * (1 - shrinkage_factor) ``` Note that the shrinkage factor is set to 0.5 by default and can be changed via the `jmpost.prior_shrinkage` option e.g. ``` options("jmpost.prior_shrinkage" = 0.7) ``` If you wish to manually specify the initial values you can do so via the `initialValues()` function. For example: ```r joint_model <- JointModel( longitudinal = LongitudinalRandomSlope(), survival = SurvivalExponential(), link = linkNone() ) initial_values <- initialValues(joint_model, n_chains = 2) initial_values[[1]]$lm_rs_intercept <- 0.2 initial_values[[2]]$lm_rs_intercept <- 0.3 mcmc_results <- sampleStanModel( joint_model, data = DataJoint(...), init = initial_values, chains = 2 ) ``` Note the following: - `initialValues()` will return a list of lists where each sublist contains the initial values for the corresponding chain index. - `initialValues()` by default just returns 1 initial value for each parameter; the `sampleStanModel()` function will then broadcast this number as many times as required e.g. if you have 3 covariates in your survival model then the initial value for the Beta coeficient will be repeated 3 times. If however you want to specify individual initial values for each covariate you can do so by passing in a vector of the same length as the number of covariates. - `initialValues()` will not check if the proposed values are valid for constrained parameters. That is, if you are using a `prior_cauchy(0, 1)` prior for a parameter that should be `>0` then you will need to manually set the initial value (as described above) to ensure that it is a valid value. - For constrained parameters (e.g. variance parameters that must be $> 0$) `initialValues()` will continuously sample and discard initial values until it generates one that meet the constraints. If after 100 attempts no valid initial value has been found then it will throw an error. # Covariates in Longitudinal Models The covariate longitudinal models allow subject-level baseline covariates to affect the parameters of a longitudinal trajectory. They are provided by `LongitudinalRandomSlopeCov()`, `LongitudinalSteinFojoCov()`, `LongitudinalGSFCov()`, and `LongitudinalClaretBrunoCov()`. Each constructor uses one-sided formulas to specify the covariates for each model parameter. The variables in these formulas must be columns in the data passed to `DataSubject()`. The intercept column implied by a formula is removed because each predictor has a separate `*_intercept_prior`; the corresponding `*_coefficients_prior` controls the remaining design-matrix coefficients. For example, `~ arm + cov_cont` models differences from the reference arm together with a linear effect of `cov_cont`. Predictors can use `"linear"`, `"proportional"`, `"exponential"`, or `"log-linear"` parametrizations. The defaults use linear predictors for location parameters and log-linear predictors for scale parameters, keeping the latter positive. Covariates may differ between parameters and may overlap with, or differ from, the covariates in the survival model. The following short example uses the data prepared above. It lets the continuous baseline covariate affect the subject intercept, treatment arm affect the mean slope, and uses a common slope standard deviation. The usual `DataJoint()` and `sampleStanModel()` workflow is otherwise unchanged. ```{r} #| label: covariate-longitudinal-model covariate_model <- JointModel( longitudinal = LongitudinalRandomSlopeCov( mu_formula = ~cov_cont, slope_mu_formula = ~arm, slope_sigma_formula = ~1, mu_intercept_prior = prior_normal(30, 10), mu_coefficients_prior = prior_normal(0, 3), slope_mu_intercept_prior = prior_normal(1, 2), slope_mu_coefficients_prior = prior_normal(0, 2), slope_sigma_intercept_prior = prior_normal(log(0.2), 0.5), sigma = prior_lognormal(log(20), 0.5) ) ) set.seed(481) covariate_results <- sampleStanModel( covariate_model, data = joint_data, iter_sampling = 100, iter_warmup = 100, chains = 1, parallel_chains = 1, refresh = 0 ) cmdstanr::as.CmdStanMCMC(covariate_results)$summary( c( "lm_rsc_mu_intercept", "lm_rsc_mu_coefficients", "lm_rsc_slope_mu_intercept", "lm_rsc_slope_mu_coefficients" ) ) ``` The coefficient names and their order follow the non-intercept columns of the formula's design matrix. In a real analysis, use informative priors where possible, run longer chains, and assess convergence as described above.