R/weightitMSM.R

In WeightIt: Weighting for Covariate Balance in Observational Studies

#' Generate Balancing Weights for Longitudinal Treatments
#'
#' @description
#' `weightitMSM()` allows for the easy generation of balancing
#' weights for marginal structural models for time-varying treatments using a
#' variety of available methods for binary, continuous, and multi-category
#' treatments, as well as censoring. Some of these methods exist in other packages, which [weightit()]
#' calls; these packages must be installed to use the desired method.
#'
#' @inheritParams weightit
#' @param formula.list a list of formulas corresponding to each time point with
#'   the time-specific treatment variable on the left hand side and
#'   pre-treatment covariates to be balanced on the right hand side. The
#'   formulas must be in temporal order, and must contain all covariates to be
#'   balanced at that time point (i.e., treatments and covariates featured in
#'   early formulas should appear in later ones). Interactions and functions of
#'   covariates are allowed. As in [weightit()], a formula may have an empty
#'   right hand side (e.g., `A_1 ~ 1`), which requests a marginal model at that
#'   time point; see *Empty model formulas* in Details at [weightit()].
#' @param data an optional data set in the form of a data frame that contains
#'   the variables in the formulas in `formula.list`. This must be a wide data
#'   set with exactly one row per unit.
#' @param method a string of length 1 containing the name of the method that
#'   will be used to estimate weights. See [weightit()] for allowable options.
#'   The default is `"glm"`, which estimates the weights using generalized
#'   linear models.
#' @param stabilize `logical`; whether or not to stabilize the weights.
#'   Stabilizing the weights involves fitting a model predicting treatment at
#'   each time point from treatment status at prior time points. If `TRUE`, a
#'   fully saturated model will be fit (i.e., all interactions between all
#'   treatments up to each time point), essentially using the observed treatment
#'   probabilities in the numerator (for binary and multi-category treatments).
#'   This may yield an error if some combinations are not observed. Default is
#'   `FALSE`. To manually specify stabilization model formulas, e.g., to specify
#'   non-saturated models, use `num.formula`. With many time points, saturated
#'   models may be time-consuming or impossible to fit.
#' @param num.formula an optional one-sided formula with the stabilization
#'   factors (other than the previous treatments) on the right hand side, which
#'   adds, for each time point, the stabilization factors to a model saturated
#'   with previous treatments. See Cole & Hernán (2008) for a discussion of how
#'   to specify this model; including stabilization factors can change the
#'   estimand without proper adjustment, and should be done with caution. Can
#'   also be a list of one-sided formulas, one for each entry of `formula.list`,
#'   including any censoring entries. Unless you
#'   know what you are doing, we recommend setting `stabilize = TRUE` and
#'   ignoring `num.formula`.
#' @param include.obj `logical`; whether to include in the output a list of the fit objects
#'   created in the process of estimating the weights at each time point. For
#'   example, with `method = "glm"`, a list of the `glm` objects containing the
#'   propensity score models at each time point will be included. See the help
#'   pages for each method for information on what object will be included if
#'   `TRUE`.
#' @param is.MSM.method `logical`; whether the method estimates weights for multiple time
#'   points all at once (`TRUE`) or by estimating weights at each time point and
#'   then multiplying them together (`FALSE`). This is only relevant for
#'   user-specified functions.
#' @param weightit.force `logical`; several methods are not valid for estimating weights
#'   with longitudinal treatments, and will produce an error message if
#'   attempted. Set to `TRUE` to bypass this error message.
#' @param ... other arguments that control aspects of fitting that are not covered by the above arguments. See Details at [weightit()].
#'
#' @returns
#' A `weightitMSM` object with the following elements:
#' \item{weights}{The estimated weights, one for each unit.}
#' \item{treat.list}{A list of the values of the time-varying treatment variables.}
#' \item{covs.list}{A list of the covariates used in the fitting at each time point. Only includes the raw covariates, which may have been altered in the fitting process.}
#' \item{estimand}{"ATE", currently the only estimand for MSMs with binary or multi-category treatments.}
#' \item{method}{The weight estimation method specified.}
#' \item{s.weights}{The provided sampling weights.}
#' \item{by}{A data.frame containing the `by` variable when specified.}
#' \item{stabilization}{The stabilization factors, if any.}
#'
#' When censoring is modeled (i.e., when any entry of `formula.list` has its left
#' side wrapped in [.cens()]), `treat.list` and `covs.list` describe the *treatment*
#' models only, while `formula.list` is kept exactly as supplied, markers included,
#' so that [update()] round-trips. The following additional components describe the
#' censoring models:
#' \item{cens.list}{A list of the values of the censoring indicators, one entry per
#' censoring time point. Each is 0 for units still under observation and 1 for units
#' censored at that time point, and `NA` for units censored earlier.}
#' \item{cens.covs.list}{A list of the covariates used to fit each censoring model.
#' As with `covs.list`, only the raw covariates are included.}
#' \item{cens.formula.list}{A list of the censoring model formulas, with the
#' [.cens()] marker retained on the left side.}
#' \item{cens.time}{The positions of the censoring models within `formula.list`, so
#' that `formula.list[cens.time]` recovers them and their timing relative to the
#' treatment models can be determined.}
#' \item{at.risk}{A logical matrix with one row per unit and one column per entry of
#' `formula.list`, named for the treatment or censoring variable modeled at that
#' entry. Each column records which units were still under observation when that
#' model was fit, i.e., which units contributed to it. Useful for assessing balance;
#' see [.cens()].}
#'
#' Censored units have a final weight of exactly 0, so `weights` is 0 for any unit
#' censored at any time point.
#'
#' When `keep.mparts` is `TRUE` (the default) and the chosen method is
#' compatible with M-estimation, the components related to M-estimation for use
#' in [glm_weightit()] are stored in the `"Mparts.list"` attribute. When `by` is
#' specified, the per-stratum components are combined into `"Mparts.list"` so
#' that the standard errors produced by [glm_weightit()] are asymptotically
#' equivalent to those from estimating the weights from models in which the `by`
#' variable is fully interacted with all the covariates at every time point. (For
#' `method = "cbps"`, this requires `is.MSM.method = FALSE`, i.e., estimating a
#' separate model at each time point, as M-estimation is not supported for the
#' single-model MSM version of CBPS.)
#'
#' @details
#'
#'
#' In general, `weightitMSM()` works by separating the estimation of weights
#' into separate procedures for each time period based on the formulas provided.
#' For each formula, `weightitMSM()` simply applies `weightit()` to that formula,
#' collects the weights for each time period, and multiplies them together to
#' arrive at longitudinal balancing weights.
#'
#' Each formula should contain all the covariates to be balanced on. For
#' example, the formula corresponding to the second time period should contain
#' all the baseline covariates, the treatment variable at the first time period,
#' and the time-varying covariates that took on values after the first treatment
#' and before the second. Currently only "wide" data sets are supported, where each
#' unit is represented by exactly one row that contains its covariate and
#' treatment history encoded in separate variables. You can use [reshape()] or
#' other functions to transform your data into this format; see example below.
#'
#' ## Censoring weights (IPCW)
#'
#' Censoring can be modeled by including entries in `formula.list` whose left side
#' is wrapped in [.cens()], placed in temporal order among the treatment
#' models. For example,
#'
#' ```
#' weightitMSM(list(A_1 ~ X1_0 + X2_0,
#'                  A_2 ~ X1_1 + X2_1 + A_1,
#'                  .cens(C_2) ~ X1_1 + X2_1 + A_1 + A_2,
#'                  A_3 ~ X1_2 + X2_2 + A_2),
#'             data = d, method = "glm")
#' ```
#'
#' models censoring occurring after the second treatment. Each censoring indicator
#' must be 0 for units still under observation and 1 for units censored at that time
#' point. See [.cens()] for details of what the resulting weights
#' estimate.
#'
#' Every model, treatment or censoring, is fit only among the units still under
#' observation when it is reached, and the resulting weights are multiplied together
#' across time points as usual. A unit censored at any time point therefore has a
#' final weight of exactly 0. Because such units drop out, missing values are
#' permitted in the treatments and covariates that follow their censoring; missing
#' values among units still under observation remain an error. `at.risk` in the
#' output has one column per time point recording which units were under
#' observation when that model was fit.
#'
#' Censoring time points are stabilized in exactly the same way as treatment time
#' points: with `stabilize = TRUE`, the numerator of a censoring weight is a model
#' for that censoring indicator given the preceding treatments (or a marginal model
#' when no treatment precedes it), and `num.formula` adds stabilization factors to
#' it as it does for a treatment. When `num.formula` is supplied as a list, it must
#' have one entry per entry of `formula.list`, censoring entries included. The
#' numerator of a censoring weight is itself a censoring model, so the stabilized
#' weight is \eqn{P(C = 0 | \cdot) / P(C = 0 | X)} for the units still under
#' observation and remains exactly 0 for those censored.
#'
#' The right side of a censoring formula may be empty, as in `.cens(C_2) ~ 1`, which
#' requests a marginal censoring model that assumes censoring at that time point is
#' independent of the covariates; its contribution to the product is \eqn{1/P(C = 0)}
#' for the units still under observation and 0 for those censored there. Everything
#' else is unaffected: the risk sets, the missing values permitted after censoring,
#' `stabilize`, `by`, and M-estimation all work as they do for a
#' covariate-dependent censoring model, and empty and non-empty censoring formulas
#' can be mixed freely. Each time point is fit separately, so only the empty ones
#' take the intercept-only shortcut described in *Empty model formulas* in Details at
#' [weightit()]; when `is.MSM.method = TRUE` there is no shortcut to take, because a
#' single set of weights is estimated for all time points at once, and a time point
#' with no covariates instead contributes only its intercept balance condition.
#'
#' @seealso
#' [weightit()] for information on the allowable methods
#'
#' [summary.weightitMSM()] for summarizing the weights
#'
#' @references
#' Cole, S. R., & Hernán, M. A. (2008). Constructing Inverse Probability Weights for Marginal Structural Models. *American Journal of Epidemiology*, 168(6), 656–664. \doi{10.1093/aje/kwn164}
#'
#' @examples
#' data("msmdata")
#' (W1 <- weightitMSM(list(A_1 ~ X1_0 + X2_0,
#'                         A_2 ~ X1_1 + X2_1 +
#'                           A_1 + X1_0 + X2_0,
#'                         A_3 ~ X1_2 + X2_2 +
#'                           A_2 + X1_1 + X2_1 +
#'                           A_1 + X1_0 + X2_0),
#'                    data = msmdata,
#'                    method = "glm"))
#' summary(W1)
#' cobalt::bal.tab(W1)
#'
#' # Using stabilization factors
#' W2 <- weightitMSM(list(A_1 ~ X1_0 + X2_0,
#'                         A_2 ~ X1_1 + X2_1 +
#'                           A_1 + X1_0 + X2_0,
#'                         A_3 ~ X1_2 + X2_2 +
#'                           A_2 + X1_1 + X2_1 +
#'                           A_1 + X1_0 + X2_0),
#'                    data = msmdata,
#'                    method = "glm",
#'                    stabilize = TRUE,
#'                    num.formula = list(~ 1,
#'                                       ~ A_1,
#'                                       ~ A_1 + A_2))
#'
#' # Same as above but with fully saturated stabilization factors
#' # (i.e., making the last entry in 'num.formula' A_1*A_2)
#' W3 <- weightitMSM(list(A_1 ~ X1_0 + X2_0,
#'                         A_2 ~ X1_1 + X2_1 +
#'                           A_1 + X1_0 + X2_0,
#'                         A_3 ~ X1_2 + X2_2 +
#'                           A_2 + X1_1 + X2_1 +
#'                           A_1 + X1_0 + X2_0),
#'                    data = msmdata,
#'                    method = "glm",
#'                    stabilize = TRUE)

#' @export
weightitMSM <- function(formula.list, data = NULL, method = "glm",
                        stabilize = FALSE, by = NULL,
                        s.weights = NULL, num.formula = NULL, missing = NULL, verbose = FALSE,
                        include.obj = FALSE, keep.mparts = TRUE,
                        is.MSM.method, weightit.force = FALSE, ...) {

  call <- match.call()

  ## Checks and processing ----

  #Checks

  ##Process method
  .check_acceptable_method(method, msm = TRUE, force = weightit.force)

  if (is_null(method)) {
    method <- NULL
    is.MSM.method <- TRUE
  }
  else if (is.character(method)) {
    method <- .method_to_proper_method(method)
    attr(method, "name") <- method
    if (missing(is.MSM.method)) is.MSM.method <- NULL
    is.MSM.method <- .process_MSM_method(is.MSM.method, method)
  }
  else { #function
    method.name <- paste(deparse(substitute(method)))
    .check_user_method(method)
    if (missing(is.MSM.method)) is.MSM.method <- NULL
    is.MSM.method <- .process_MSM_method(is.MSM.method, method)
    attr(method, "name") <- method.name
  }

  ##Process by
  if (is_not_null(...get("exact"))) {
    arg::wrn("{.arg by} has replaced {.arg exact} in the {.fun weightit} syntax, but {.arg exact} will always work")
    by <- ...get("exact")
    by.arg <- "exact"
  }
  else {
    by.arg <- "by"
  }

  reported.covs.list <- simple.covs.list <- covs.list <- treat.list <- w.list <- ps.list <-
    stabout <- sw.list <- Mparts.list <- stab.Mparts.list <- na.list <-
    atrisk.list <- make_list(length(formula.list))

  if (is_null(formula.list) || !is.list(formula.list) ||
      !all_apply(formula.list, rlang::is_formula, lhs = TRUE)) {
    arg::err("{.arg formula.list} must be a list of formulas")
  }

  for (i in seq_along(formula.list)) {

    #Process treat and covs from formula and data. An empty right side (e.g.,
    #`A_1 ~ 1` or `.cens(C_2) ~ 1`) yields a zero-column `model.covs` and is
    #allowed, as it is in `weightit()`: the corresponding model is marginal. This
    #is what makes it possible to model censoring that does not depend on
    #covariates while keeping the rest of the censoring machinery (risk sets,
    #tolerated NAs after censoring, M-estimation).
    t.c <- get_covs_and_treat_from_formula2(formula.list[[i]], data)
    simple.covs.list[[i]] <- t.c[["simple.covs"]]
    reported.covs.list[[i]] <- t.c[["reported.covs"]]

    covs.list[[i]] <- t.c[["model.covs"]]
    treat.list[[i]] <- t.c[["treat"]]

    if (is_null(treat.list[[i]])) {
      arg::err("no treatment variable was specified in the {ordinal(i)} formula")
    }

    n <- length(treat.list[[i]])

    if (nrow(covs.list[[i]]) != n) {
      arg::err("the treatment and covariates must have the same number of units")
    }

    treat.list[[i]] <- as.treat(treat.list[[i]], process = TRUE)

    treat.name <- .attr(treat.list[[i]], "treat.name")

    #Non-finite values are never allowed. Missing values are allowed only when a
    #censoring model is present, because a unit censored at an earlier time point
    #legitimately has no later treatment or censoring indicator. Those are checked
    #against the risk sets after this loop, once the risk sets are known.
    na.list[[i]] <- is.na(treat.list[[i]])

    if (!all(is.finite(treat.list[[i]][!na.list[[i]]]))) {
      arg::err(c("No non-finite values are allowed in the treatment variable.",
                 "i" = "Non-finite values found in {.var treat.name}"))
    }

    names(treat.list)[i] <- treat.name
    names(reported.covs.list)[i] <- treat.name

    if (!is.MSM.method) {
      .check_method_treat.type(method, get_treat_type(treat.list[[i]]))
    }

    #By is processed each for each time to check, but only last time is used for by.factor.
    processed.by <- .process_by(by, data = data,
                                treat = treat.list[[i]],
                                treat.name = treat.name,
                                by.arg = by.arg)
  }

  #Censoring models, marked by wrapping the censoring indicator in `.cens()` on the
  #LHS of a formula in `formula.list` (e.g., `.cens(C) ~ x1 + x2`). They are
  #interleaved with the treatment models in temporal order and folded into the
  #final product as inverse probability of censoring weights.
  is.cens <- vapply(treat.list, function(t) {
    identical(get_treat_type(t), "censoring")
  }, logical(1L))

  if (any(is.cens)) {
    #`method = "cbps"` is the only built-in method that estimates the weights for
    #all time points at once, and it supports censoring; a user-defined MSM method
    #has no way to receive the risk sets.
    #`==` rather than `identical()`: `method` carries a `"name"` attribute
    if (is.MSM.method && !(is.character(method) && isTRUE(method == "cbps"))) {
      arg::err(c("censoring models (specified with {.fun .cens} in {.arg formula.list}) cannot be used with this method when it estimates the weights for all time points simultaneously.",
                 "i" = "Set {.code is.MSM.method = FALSE} to estimate the weights separately at each time point."))
    }

    if (all(is.cens)) {
      arg::err("{.arg formula.list} must contain at least one treatment model")
    }
  }

  #Units still under observation when each model is fit. A censoring model is fit on
  #the units at risk just before that censoring event removes any of them.
  .ar <- rep.int(TRUE, n)

  for (i in seq_along(formula.list)) {
    atrisk.list[[i]] <- .ar

    if (is.cens[i]) {
      ind <- .make_cens_treat(treat.list[[i]])

      #The !is.na() guard matters: NA == 1 would propagate NA into the risk set
      .ar[.ar & !is.na(ind) & ind == 1] <- FALSE
    }
  }

  #Missing values in a treatment or censoring indicator are tolerated only for
  #units already censored; units still under observation must be fully observed.
  if (any(is.cens)) {
    for (i in seq_along(formula.list)) {
      if (any(na.list[[i]] & atrisk.list[[i]])) {
        arg::err(c("No missing values are allowed among the units still under observation.",
                   "i" = "Missing values found in {.var {names(treat.list)[i]}} among units not yet censored"))
      }
    }

    for (i in which(is.cens)) {
      ind <- .make_cens_treat(treat.list[[i]])[atrisk.list[[i]]]

      if (all(ind == 1)) {
        arg::err("all units still under observation are censored at {.var {names(treat.list)[i]}}, so no censoring weights can be estimated")
      }

      if (!any(ind == 1)) {
        arg::msg("no units are censored at {.var {names(treat.list)[i]}}; the corresponding censoring weights are all {.val {1}}")
      }
    }
  }
  else if (any_apply(na.list, any)) {
    bad <- names(treat.list)[vapply(na.list, any, logical(1L))]

    arg::err(c("No missing values are allowed in the treatment variable.",
               "i" = "Missing values found in {.var {bad}}"))
  }

  #Process missing
  missing <- {
    if (is_null(method) ||
        !any_apply(seq_along(formula.list), function(i) {
          #Covariates that are missing only for units already censored must not
          #trigger the missingness machinery, since those units are never used.
          anyNA(reported.covs.list[[i]][atrisk.list[[i]], , drop = FALSE])
        })) {
      ""
    }
    else .process_missing(missing, method)
  }

  #Process s.weights
  s.weights <- .process.s.weights(s.weights, data)

  if (is_null(s.weights)) s.weights <- rep.int(1, n)
  else .check_method_s.weights(method, s.weights)

  if (is_null(method)) {
    num.formula <- NULL
    stabilize <- FALSE
  }
  else if (is_not_null(num.formula)) {
    if (!isTRUE(stabilize)) {
      arg::msg("setting {.arg stabilize} to {.val {TRUE}} based on {.arg num.formula} input")
    }
    stabilize <- TRUE
  }

  if (stabilize) {
    if (!is.function(method) && !.weightit_methods[[method]]$stabilize_ok) {
      arg::wrn("{.arg stabilize} cannot be used with {(.method_to_phrase(method))} and will be ignored")
      stabilize <- FALSE
      num.formula <- NULL
    }
    else if (is_not_null(num.formula)) {
      #Censoring time points are stabilized like any other, so a list of numerator
      #formulas has one entry per entry of `formula.list`, censoring included.
      .check_num.formula(num.formula, data, env = parent.frame(),
                         formula.list = formula.list)
    }
  }

  #Process moments and int
  m.i.q <- .process_moments_int_quantile(method = method, ...)

  A <- list(...)
  A["s.weights"] <- list(s.weights)
  A["by.factor"] <- list(.attr(processed.by, "by.factor"))
  A["method"] <- list(method)
  A[c("moments", "int", "quantile")] <- m.i.q[c("moments", "int", "quantile")]
  A["subclass"] <- list(numeric())
  A["missing"] <- list(missing)
  A["verbose"] <- list(verbose)
  A["include.obj"] <- list(include.obj)

  if (is.MSM.method) {
    #Returns weights (w)
    A["covs.list"] <- list(covs.list)
    A["treat.list"] <- list(treat.list)
    A["stabilize"] <- list(stabilize)

    #Only passed when censoring is present, so that without it the fitting
    #function takes exactly its original code path
    A["atrisk.list"] <- list(if (any(is.cens)) atrisk.list)

    obj <- do.call("weightitMSM.fit", A)

    w <- obj[["weights"]]
    stabout <- NULL
    obj.list <- obj[["fit.obj"]]
    #Wrap into the same list-of-lists shape used by the per-time-point path below
    #so the shared combine step can flatten uniformly. weightitMSM.fit() returns
    #"Mparts.list" when `by` has >1 level (per-stratum parts) and "Mparts"
    #otherwise.
    Mparts.list <- list(clear_null(.attr(obj, "Mparts.list") %or% list(.attr(obj, "Mparts"))))
  }
  else {
    if (is_not_null(A[["link"]])) {
      if (length(A[["link"]]) == 1L) {
        A[["link"]] <- rep.int(A[["link"]], length(formula.list))
      }
      else if (length(A[["link"]]) != length(formula.list)) {
        arg::err("the argument to {.arg link} must have length {.or {unique(c(1, length(formula.list)))}}")
      }
    }

    obj.list <- make_list(length(formula.list))

    A["estimand"] <- list("ATE")
    A["focal"] <- list(character())
    A["stabilize"] <- list(FALSE)
    A["ps"] <- list(numeric())
    A["is.MSM.method"] <- list(FALSE)

    #Prior treatments, used as stabilization predictors; censoring indicators are
    #never included among them. Whether this is empty also stands in for "is this
    #the first time point?" when building the stabilization formulas below: it is
    #equivalent for the treatment models and is the behavior wanted at a censoring
    #time point that precedes every treatment.
    prior.treat.names <- character()

    for (i in seq_along(formula.list)) {
      A_i <- A
      if (length(A[["link"]]) == length(formula.list)) {
        A_i["link"] <- list(A[["link"]][[i]])
      }

      at.risk <- atrisk.list[[i]]

      A_i["covs"] <- list(covs.list[[i]])
      A_i["treat"] <- list(treat.list[[i]])
      A_i[".data"] <- list(data)
      A_i[".covs"] <- list(reported.covs.list[[i]])

      #Only units still under observation are used to fit the model. `NULL` rather
      #than an all-TRUE vector when there is no censoring, so that `weightit.fit()`
      #takes exactly its original code path.
      A_i["subset"] <- list(if (any(is.cens)) at.risk else NULL)

      ## Running models ----

      #Returns weights (w) and propensity score (ps)
      obj <- do.call("weightit.fit", A_i)

      #`weightit.fit()` leaves the weights of units outside `subset` as NA. Those
      #units contribute nothing at this time point, so their factor is 1. Censored
      #units already have a weight of exactly 0 from the censoring method itself.
      w_i <- obj[["weights"]]
      w_i[!at.risk] <- 1

      w.list[i] <- list(w_i)
      ps.list[i] <- list(obj[["ps"]])
      obj.list[i] <- list(obj[["fit.obj"]])
      #A list of parts for this time point: one per `by` group (already expanded
      #to full-sample size by weightit.fit()) when by has >1 level, otherwise the
      #single Mparts. Empty when the method supplies no Mparts.
      Mparts.list[[i]] <- clear_null(.attr(obj, "Mparts.list") %or% list(.attr(obj, "Mparts")))

      if (stabilize) {
        #Process stabilization formulas and get stab weights. Censoring time points
        #are stabilized exactly like treatment time points, and by a censoring model
        #of their own. The marker is stripped from the left side only so that
        #`get_covs_and_treat_from_formula2()` is asked for nothing but the numerator
        #covariates; the treatment passed to `weightit.fit()` below is still the
        #censoring-tagged indicator, so the numerator is fit by the same `.cens`
        #method as the denominator.
        f_i <- {
          if (is.cens[i]) .uncens_formula(formula.list[[i]])
          else formula.list[[i]]
        }

        if (rlang::is_formula(num.formula)) {
          if (is_null(prior.treat.names)) {
            stab.f <- update(f_i, num.formula)
          }
          else {
            stab.f <- update.formula(as.formula(paste(names(treat.list)[i], "~",
                                                      paste(prior.treat.names,
                                                            collapse = " * "))),
                                     as.formula(paste("~", deparse1(rlang::f_rhs(num.formula)), "+ .")))
          }
        }
        else if (is.list(num.formula)) {
          stab.f <- update(f_i, num.formula[[i]])
        }
        else {
          if (is_null(prior.treat.names)) {
            stab.f <- update(f_i, ". ~ 1")
          }
          else {
            stab.f <- update(f_i,
                             sprintf(". ~ %s", paste(prior.treat.names,
                                                     collapse = " * ")))
          }
        }

        stab.t.c_i <- get_covs_and_treat_from_formula2(stab.f, data)

        A_i["covs"] <- stab.t.c_i["model.covs"]
        A_i["method"] <- list("glm")
        A_i["moments"] <- list(integer())
        A_i["int"] <- list(FALSE)
        A_i["quantile"] <- list(list())

        sw_obj <- do.call("weightit.fit", A_i)

        #A censoring numerator is itself a censoring model, so its weights are 0 for
        #the units censored here; `.num_stab_weights()` neutralizes those so the
        #reciprocal is finite. The stabilized factor is then P(C = 0 | V) for the
        #units still under observation, and the denominator keeps those censored here
        #at exactly 0.
        sw_i <- 1 / .num_stab_weights(sw_obj, censoring = is.cens[i])
        sw_i[!at.risk] <- 1

        sw.list[[i]] <- sw_i
        stabout[[i]] <- stab.f[-2L]

        #Invert each numerator part (one per `by` group when by has >1 level).
        stab.Mparts.list[[i]] <- lapply(
          clear_null(.attr(sw_obj, "Mparts.list") %or% list(.attr(sw_obj, "Mparts"))),
          .invert_num_Mpart, censoring = is.cens[i])
      }

      if (!is.cens[i]) {
        prior.treat.names <- c(prior.treat.names, names(treat.list)[i])
      }
    }

    w <- Reduce("*", w.list, init = 1)

    if (stabilize) {
      #`clear_null()` is load-bearing, not cosmetic: a NULL entry would collapse `w`
      #to length 0 (Reduce("*", list(NULL, x), init = y) is numeric(0)).
      w <-  Reduce("*", clear_null(sw.list), init = w)

      unique.stabout <- unique(clear_null(stabout))

      if (length(unique.stabout) <= 1L) {
        stabout <- unique.stabout
      }
    }
    else {
      stabout <- NULL
    }

    if (include.obj) {
      #`treat.list` is still full length here, so censoring fit objects are named
      #after their censoring indicator; it is subset in the output below.
      names(obj.list) <- names(treat.list)
    }
  }

  if (is_not_null(method) && all_the_same(w)) {
    arg::wrn("all weights are {.val w[1L]}, possibly indicating an estimation failure")
  }

  ## Assemble output object----
  #`treat.list` and `covs.list` describe treatments only, so that printing, balance
  #assessment, and stabilization operate on treatments alone; the censoring models
  #are stored separately. `formula.list` is kept as supplied, `.cens()` markers
  #included, so that `update()` round-trips.
  out <- list(weights = w,
              treat.list = treat.list[!is.cens],
              covs.list = simple.covs.list[!is.cens],
              estimand = "ATE",
              method = method,
              s.weights = s.weights,
              by = processed.by,
              call = call,
              formula.list = formula.list,
              cens.list = if (any(is.cens)) treat.list[is.cens],
              cens.covs.list = if (any(is.cens)) simple.covs.list[is.cens],
              cens.formula.list = if (any(is.cens)) formula.list[is.cens],
              cens.time = if (any(is.cens)) which(is.cens),
              at.risk = if (any(is.cens)) {
                #One column per model, treatment or censoring, giving the units
                #still under observation when that model was fit. Needed to assess
                #balance, since `bal.tab()` cannot handle the NA treatments that
                #censoring implies; subsetting to a column gives exactly the sample
                #the corresponding model was fit on.
                do.call("cbind", atrisk.list) |>
                  `colnames<-`(names(treat.list))
              },
              stabilization = stabout,
              missing = if (nzchar(missing)) missing else NULL,
              env = parent.frame(),
              obj = obj.list
  )

  out <- clear_null(out)

  if (keep.mparts && all(lengths(Mparts.list) > 0L)) {
    #Each slot holds a list of parts (per time point, and per `by` group within a
    #time point); flatten one level into a single stacked Mparts.list.
    attr(out, "Mparts.list") <- clear_null(c(do.call("c", Mparts.list),
                                             do.call("c", stab.Mparts.list)))
  }

  class(out) <- c("weightitMSM", "weightit")

  out
}

#' @exportS3Method print weightitMSM
print.weightitMSM <- function(x, ...) {
  treat.types <- vapply(x[["treat.list"]], get_treat_type, character(1L))

  cat(sprintf("A %s object\n", .it(class(x)[1L])))

  if (is_not_null(x[["method"]])) {
    method_name <- {
      if (is_not_null(.attr(x[["method"]], "name"))) add_quotes(.attr(x[["method"]], "name"))
      else if (is.character(x[["method"]])) add_quotes(x[["method"]])
      else "user-defined"
    }

    method_note <- {
      if (is_not_null(.attr(x[["method"]], "package")))
        sprintf(" (converted from %s)", .it(.attr(x[["method"]], "package")))
      else if (is_not_null(x[["method"]]))
        sprintf(" (%s)", .method_to_phrase(x[["method"]]))
      else
        ""
    }

    cat(sprintf(" - method: %s%s\n",
                method_name,
                method_note))
  }
  else if (all_the_same(x[["weights"]])) {
    cat(" - method: no weighting\n")
  }

  cat(sprintf(" - number of obs.: %s\n",
              nobs(x)))

  cat(sprintf(" - sampling weights: %s\n",
              if (is_null(x[["s.weights"]]) || all_the_same(x[["s.weights"]])) "none" else "present"))

  cat(sprintf(" - number of time points: %s (%s)\n",
              length(x[["treat.list"]]),
              word_list(names(x[["treat.list"]]), and.or = FALSE)))

  cat(" - treatment:\n")
  for (i in seq_along(x[["treat.list"]])) {
    cat(sprintf("    + time %s: %s\n",
                i,
                switch(treat.types[i],
                       continuous = "continuous",
                       `multi-category` =,
                       multinomial = sprintf("%s-category (%s)",
                                             nunique(x[["treat.list"]][[i]]),
                                             word_list(levels(x[["treat.list"]][[i]]), and.or = FALSE)),
                       binary = "2-category")))
  }

  if (is_not_null(x[["cens.list"]])) {
    cat(" - censoring (IPCW):\n")
    for (i in seq_along(x[["cens.list"]])) {
      cat(sprintf("    + %s: %s of %s units censored\n",
                  names(x[["cens.list"]])[i],
                  sum(.make_cens_treat(x[["cens.list"]][[i]]) == 1, na.rm = TRUE),
                  nobs(x)))
    }
  }

  if (is_not_null(x[["cens.covs.list"]])) {
    cat(" - censoring covariates:\n")
    for (i in seq_along(x[["cens.covs.list"]])) {
      cat(sprintf("    + %s: %s\n",
                  names(x[["cens.list"]])[i],
                  if (is_null(x[["cens.covs.list"]][[i]])) "(none)"
                  else word_list(names(x[["cens.covs.list"]][[i]]), and.or = FALSE)))
    }
  }

  if (is_not_null(x[["covs.list"]])) {
    cat(" - covariates:\n")
    for (i in seq_along(x[["covs.list"]])) {
      if (i == 1L) {
        cat(sprintf("    + baseline: %s\n",
                    if (is_null(x$covs.list[[i]])) "(none)"
                    else word_list(names(x$covs.list[[i]]), and.or = FALSE)))
      }
      else {
        cat(sprintf("    + after time %s: %s\n",
                    i - 1L,
                    if (is_null(x$covs.list[[i]])) "(none)"
                    else word_list(names(x$covs.list[[i]]), and.or = FALSE)))
      }
    }
  }

  if (is_not_null(x[["missing"]]) && !identical(x[["missing"]], "")) {
    cat(sprintf(" - missingness method: %s\n",
                .missing_to_phrase(x[["missing"]])))
  }

  if (is_not_null(x[["by"]])) {
    cat(sprintf(" - by: %s\n",
                word_list(names(x[["by"]]), and.or = FALSE)))
  }

  if (is_not_null(x$stabilization)) {
    cat(" - stabilized")
    if (any_apply(x$stabilization, function(s) is_not_null(get_varnames(s)))) {
      cat(paste0("; stabilization factors:\n",
                 if (length(x$stabilization) == 1L) {
                   sprintf("      %s", word_list(.attr(terms(x[["stabilization"]][[1L]]), "term.labels"),
                                                 and.or = FALSE))
                 }
                 else {
                   paste(vapply(seq_along(x$stabilization), function(i) {
                     if (i == 1L) {
                       sprintf("    + baseline: %s",
                               if (is_null(.attr(terms(x[["stabilization"]][[i]]), "term.labels"))) "(none)"
                               else word_list(.attr(terms(x[["stabilization"]][[i]]), "term.labels"), and.or = FALSE))
                     }
                     else {
                       sprintf("    + after time %s: %s",
                               i - 1L,
                               word_list(.attr(terms(x[["stabilization"]][[i]]), "term.labels"), and.or = FALSE))
                     }
                   }, character(1L)), collapse = "\n")
                 }))
    }
  }

  #trim
  if (is_not_null(.attr(x, "trim"))) {
    trim.at <- .attr(x, "trim")[["at"]]
    trim.lower <- .attr(x, "trim")[["lower"]]
    trim.drop <- .attr(x, "trim")[["drop"]]
  }
  else if (is_not_null(.attr(x[["weights"]], "trim"))) {
    trim.at <- .attr(x[["weights"]], "trim")
    trim.lower <- .attr(x[["weights"]], "trim.lower")
    trim.drop <- FALSE
  }
  else {
    trim.at <- NULL
  }

  if (is_not_null(trim.at) && is_number(trim.at)) {
    if (trim.at < 1) {
      if (trim.lower) {
        trim.at <- c(1 - trim.at, trim.at)
      }

      cat(sprintf(" - weights trimmed at %s%s\n",
                  word_list(paste0(round(100 * trim.at, 2L), "%")),
                  if (trim.drop) " and units dropped" else ""))
    }
    else {
      cat(sprintf(" - weights trimmed at the %s %s%s\n",
                  if (trim.lower) "top and bottom" else "top",
                  trim.at,
                  if (trim.drop) " and units dropped" else ""))
    }
  }

  invisible(x)
}