Estimating State Occupancy with Aalen-Johansen

Why estimate state occupancy?

In a multi-state study, a central question is often not just whether an event has happened, but which state people occupy over time. For example, after treatment a patient may remain alive without relapse, relapse, or die. The probability of occupying each state at a given time describes the distribution of the original cohort across those states. These probabilities are useful summaries of the whole transition process, while preserving the fact that a person can move through more than one state.

The Aalen–Johansen estimator is the multi-state analogue of the Kaplan–Meier estimator. It combines estimated transition hazards across event times through a product of transition matrices. Here the emphasis is on state-occupation probabilities for a cohort starting in a specified state, estimated from observed histories with mstate. This is distinct from producing a collection of covariate-specific predictions.

Definition: Let \(X(t)\) denote the state occupied at time \(t\), and suppose individuals begin in state \(r\). The state-occupation probability for state \(j\) at time \(t\) is \(P(X(t)=j\mid X(0)=r)\). The Aalen–Johansen estimator estimates these probabilities by multiplying estimated transition-probability matrices over observed transition times.

The product-integral calculation

Consider a finite set of states and an allowed transition structure. At event time \(u_k\), let \(d_{hj}(u_k)\) be the number of observed transitions from state \(h\) to state \(j\), and let \(Y_h(u_k)\) be the number at risk in state \(h\) just before that time. In the basic nonparametric construction, the estimated transition probability from \(h\) to \(j\) at that event time is \(d_{hj}(u_k)/Y_h(u_k)\), for \(j\ne h\). The diagonal entry is chosen so each row sums to one:

\widehat A_{hj}(u_k)= \begin{cases} d_{hj}(u_k)/Y_h(u_k), & j\ne h,\\[4pt] 1-\displaystyle\sum_{\ell\ne h}d_{h\ell}(u_k)/Y_h(u_k), & j=h. \end{cases}

The estimated occupancy row vector at time \(t\), starting from state \(r\), is the initial row vector multiplied by the event-time matrices up to \(t\). Events from different source states can occur at the same time: each source-state row is updated using its own risk set. A state with no outgoing transition at an event time has an unchanged row. In particular, an absorbing state such as death remains absorbing.

\widehat{\boldsymbol p}_r(t) =\boldsymbol e_r^{\mathsf T} \prod_{u_k\le t}\widehat A(u_k),

where \(\boldsymbol e_r^{\mathsf T}\) places probability one in the starting state and zero in the others. The product is taken in chronological order from earliest to latest event time. This construction estimates occupancy at each event time; between event times, the estimate stays constant unless a different time-grid convention is used.

Worked Example: Hand-calculating Occupancy in an Illness–Death Model

Suppose five hypothetical patients begin alive without relapse. The allowed transitions are alive without relapse to relapse, alive without relapse to death, and relapse to death. At time 1, one patient relapses and another dies without relapse. At time 2, one of the three remaining patients without relapse dies, and the patient who relapsed dies. Assume no other transitions occur before time 2.

At time 1, five patients are at risk in the initial state. The estimated probabilities of moving from that state to relapse and death are each \(1/5=0.2\), leaving \(1-0.2-0.2=0.6\) in the initial state. The patient who has already entered the relapse state remains there at this event time, and the death state is absorbing. Starting from the initial state, occupancy just after time 1 is therefore:

(1,0,0) \begin{pmatrix} 0.6&0.2&0.2\\ 0&1&0\\ 0&0&1 \end{pmatrix} =(0.6,0.2,0.2).

At time 2, three patients remain at risk in the initial state. One dies, so the probability of remaining in that state is \(1-1/3=2/3\), and the probability of moving to death is \(1/3\). The one patient at risk in the relapse state dies, so that row moves entirely to death. Multiplying the current occupancy vector by the time-2 matrix gives:

(0.6,0.2,0.2) \begin{pmatrix} 2/3&0&1/3\\ 0&0&1\\ 0&0&1 \end{pmatrix} =(0.4,0,0.6).

Thus the estimated probabilities at time 2 are 0.4 alive without relapse, 0 in relapse, and 0.6 dead. They are probabilities for the original cohort at that time, not separate risks calculated only among people who remain alive.

Assumptions and data structure

An Aalen–Johansen estimate is meaningful only when the recorded histories represent the states and transitions of interest. Each transition must be encoded according to a defined transition structure, and event times must use a consistent time origin and scale. In the software workflow below, a stratified Cox model with no covariates estimates a separate baseline cumulative hazard for each allowed transition; its product integral yields the nonparametric Aalen–Johansen occupancy estimate.

Correct transition coding: each observed move is assigned to an allowed transition, and the event indicator distinguishes a transition from censoring.
Appropriate risk sets: a person contributes to the risk set for transitions out of the state they currently occupy, not transitions out of every state.
Independent censoring: within the analysis being made, censoring must not selectively remove people with different future transition patterns, conditional on the information used to define the analysis.
Model interpretation: interpreting the product as transition probabilities conditional on the current state uses a Markov formulation. If transition behavior depends on time spent in a state or earlier history, the state representation may need to include that duration or history.
Complete observation of relevant moves: unobserved intermediate transitions can invalidate the encoded histories and their risk sets.

The first four conditions are not made true by software. In particular, the Aalen–Johansen calculation does not correct informative loss to follow-up or a misspecified transition structure. If people enter the study after the chosen time origin, delayed entry also needs to be represented appropriately rather than silently treated as ordinary follow-up from time zero.

Worked mstate workflow

The usual mstate workflow is to encode the transition structure, expand the observed histories into transition-specific intervals with msprep, fit a no-covariate stratified Cox model, obtain transition-specific cumulative hazards with msfit, and calculate the product integral. The result for a chosen initial state is a set of state-occupation probabilities over time. The next example uses hypothetical data chosen to match the hand calculation above.

Worked Example: Estimating Occupancy with mstate

Use the five hypothetical histories from the preceding example. Patient 1 relapses at time 1 and dies at time 2. Patient 2 dies without relapse at time 1. Patient 5 dies without relapse at time 2. Patients 3 and 4 are censored at time 3 without relapse. The following R commands illustrate the workflow; the time and status columns are all hypothetical.

First define the illness–death structure and enter one row per patient. For patients without a relapse, the relapse time is set to the end of observed follow-up, with the relapse status set to zero. The death status is one only for observed deaths.

Use trans = mstate::transMat(list(c(2,3), c(3), c()), names = c("No relapse", "Relapse", "Death")). Then create a wide data frame with columns id = 1:5, rel = c(1,1,3,3,2), relstat = c(1,0,0,0,0), death = c(2,1,3,3,2), and deathstat = c(1,1,0,0,1). Expand the histories with mstate::msprep(time = c(NA, "rel", "death"), status = c(NA, "relstat", "deathstat"), data = wide, trans = trans, id = "id").

Check the expanded records before fitting: each observed interval should begin in the state occupied at its start, end at the appropriate transition time or censoring time, and carry the corresponding transition number and status. Then fit separate baseline hazards by transition with no covariates. In R, use survival::coxph(survival::Surv(Tstart, Tstop, status) ~ strata(trans), data = msdata, method = "breslow"). The no-covariate model is intentional: it targets the cohort-level, nonparametric transition estimates rather than a covariate-specific prediction.

Obtain the transition-specific cumulative hazards with mstate::msfit(fit, newdata = data.frame(trans = 1:3), trans = trans, variance = FALSE). Then compute the occupancy estimates from state 1 at time zero with mstate::probtrans(msf, predt = 0). The output for the initial state is pt[[1]]; its columns named pstate1, pstate2, and pstate3 give the estimated probabilities for no relapse, relapse, and death. At times 1 and 2, the estimates should agree with the hand calculation: at time 1 they are 0.6, 0.2, and 0.2; at time 2 they are 0.4, 0, and 0.6.

The example assumes independent censoring and correctly recorded transitions. Here the two censorings occur after the event times used in the calculation, so they do not change the risk sets at times 1 or 2. In an actual analysis, verify this from the expanded data and study design; the fact that a short hypothetical data set has convenient censoring does not justify that assumption in a real application.

Checking that probabilities sum to one

At every time, state occupancy probabilities for a complete set of mutually exclusive states should sum to one, up to numerical rounding. This is both a basic probability check and a useful implementation diagnostic. Because every estimated transition matrix has rows summing to one, multiplying an occupancy vector that sums to one by such a matrix preserves that sum:

\sum_j \widehat p_j(u_k) =\sum_h \widehat p_h(u_k^-)\sum_j\widehat A_{hj}(u_k) =\sum_h \widehat p_h(u_k^-).

The equality justifies the check: if the starting vector sums to one and every event-time transition matrix is stochastic, the resulting occupancy vector also sums to one. For the mstate output, extract all columns named pstate followed by a state number and calculate the row sums over time. For example, in R use state_cols = grep("^pstate", names(pt[[1]]), value = TRUE), then occ_sum = rowSums(pt[[1]][, state_cols]). Inspect the minimum and maximum of occ_sum and compare each to one using a small numerical tolerance, such as 1e-8.

Worked Example: Checking the Probability Total

Suppose the three estimated probabilities at a reported time are 0.4000, 0.0000, and 0.6000. Their sum is \(0.4000+0.0000+0.6000=1.0000\), as expected. If software instead returns 0.4000, 0.0000, and 0.5900, the sum is 0.9900, which is not a rounding discrepancy at four decimal places. Do not rescale the values to force the total to one: first check whether all states were retained, whether the selected output corresponds to one initial state, and whether the transition structure and risk sets were specified correctly.

Common failures and alternatives

A sum below one often points to an omitted state, a truncated output table, or a mismatch between the transition matrix and the fitted hazards. A sum above one can suggest duplicate transition coding or an incorrect calculation that adds probabilities from different starting states. A state probability can also remain unexpectedly constant because the fitted data contain no observed outgoing transitions from that state. These patterns prompt investigation; they are not repaired simply by changing the plotted output.

The Aalen–Johansen approach is preferable to treating each transition as an isolated survival analysis when the target is the joint distribution across states. Separate Kaplan–Meier curves can fail to account for competing exits from a state, and independently estimated state curves need not add to one. For a simple competing-risks outcome, cumulative incidence methods provide the corresponding event-type probabilities while accounting for competing events. In a richer multi-state process, Aalen–Johansen uses the full transition structure and preserves the probability total when the model and data are coherent.

The estimate is not a cure for sparse data: late-time probabilities may be based on few people at risk and can be unstable. A total of one does not establish that censoring is independent, that all transitions were observed, or that the chosen state model is adequate. Also, occupancy estimates for a cohort starting in state 1 should not be read as predictions for a patient with particular covariates. Those are different targets and require an explicitly specified prediction procedure.

Key takeaway: The Aalen–Johansen estimator updates state occupancy by multiplying transition-probability matrices over time. With a complete set of states and a coherent transition model, probabilities across states sum to one at each time. Check that identity, but also inspect the data structure and assumptions that make the estimates interpretable.

Check Your Understanding

Work through these questions, then compare with the brief solutions.

In a three-state illness–death model, why is the probability of remaining in a state the diagonal entry of an event-time transition matrix?
At an event time, a source state has 10 people at risk, with one transition to each of two different states. What are the two off-diagonal probabilities and the diagonal probability for that source-state row?
In the worked example, what are the state-occupation probabilities just after time 1, and why do they sum to one?
A software output gives state probabilities summing to 0.97 at a late time. Give two checks to perform before interpreting the estimates.
Why can a correct probability total not demonstrate that censoring is independent?

Solutions. (1) The row must sum to one, so the diagonal is one minus the sum of the probabilities of leaving the state. (2) The off-diagonal entries are \(1/10=0.1\) each and the diagonal is \(1-0.1-0.1=0.8\). (3) They are 0.6, 0.2, and 0.2 for no relapse, relapse, and death; their sum is one because the event-time transition matrix has rows summing to one. (4) Check that all state-probability columns are included and that the fitted transition structure, expanded records, and risk sets are correct. (5) The sum-to-one identity follows from the transition-matrix calculation; it does not test whether censoring selectively removed people with different future transition patterns.