Add native Cox posterior predictions and restricted mean survival times by wds15 · Pull Request #1964 · paul-buerkner/brms

wds15 · 2026-11-02T15:09:40Z

This PR adds native Cox response predictions through the existing brms
prediction framework, while keeping additional public API commitments small.

Predictive distributions: posterior_predict() now supports Cox
event-time draws, CDFs/PIT values, densities, and quantiles using the stored
baseline spline and posterior simplex coefficients, including stratified
baselines.
Restricted means: posterior_epred() and fitted() return restricted
mean survival time, E[min(T, horizon)], with an explicitly supplied horizon
rather than an unrestricted mean.
Tail policies: predictions beyond the spline boundary support
constant-hazard extrapolation ("constant", the default), loglinear
cumulative-hazard extrapolation ("loglinear"), or no extrapolation
("none"). Without extrapolation, negative draws encode right censoring at
the boundary.
Numerical stability: truncation and quantile inversion use
cumulative-hazard calculations to avoid loss of precision in extreme tails.
Compatibility: existing precomputed-baseline likelihood calculations
are retained. Native predictions require a compatible fit with its stored
spline; unsupported fits receive a refitting error.
Regression coverage: tests cover preparation, prediction outputs,
tails, truncation, and RMST.

Optional recipes in `demo/`

The optional utility layer lives in demo/cox.R, rather than the package
namespace, to avoid committing brms to additional public utility APIs.
The indexed, self-contained demo simulates censored data, fits a small Cox
model, and demonstrates:

Baseline hazard and cumulative hazard (cox_bhaz, cox_cumbhaz).
Conditional hazard, cumulative hazard, and survival (cox_hazard,
cox_cumhaz, cox_surv).
Full survival curves and memory-saving summaries (cox_surv_curve,
cox_surv_summary).
Explicit censoring of predictive draws, reverse Kaplan-Meier censoring-time
simulation, and decoding signed times (cox_censor, cox_censor_times,
cox_decode).
Native predictions, RMST, and the no-extrapolation censoring convention.

Run the demo after installing this branch:

demo("cox", package = "brms")

These are demonstration functions, not exported or supported package APIs.
They use implementation-dependent private helpers and prepared-object fields.
The survival recipes describe untruncated event times; native predictive
outputs and RMST respect model truncation. Reverse-KM examples use the optional
prodlim package.

Scope and limitations

The existing general prediction-output interface is reused, not introduced by
this PR. Cox-Snell residuals can be derived from log-survival predictions, but
censoring indicators must be retained explicitly for diagnostics.
Censoring-aware residual diagnostics and plotting remain external.

Censoring is not automatically simulated from the observed censoring pattern,
and ordinary summaries of signed censoring codes are not survival-time
summaries. Cox mixtures are outside the initial scope.

Retain canonical basis, simplex draws, strata and fitting-version metadata without disabling legacy likelihood paths. Handle singleton strata, response-specific data names and unchecked prediction responses. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Implement time-aware hazard, density, CDF, quantile and random outputs with stable truncation and probability-aware inversion. Preserve precomputed likelihood semantics, add censor decoding and reject incompatible bayesplot replications. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Expose baseline and conditional survival quantities, curves and streaming summaries. Compute restricted means with shared Gauss-Legendre panels and stable truncation, forwarding the explicit horizon through prediction wrappers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add reverse-KM censoring simulation, censoring-aware PIT, Cox-Snell residuals and PSIS-weighted LOO-PIT. Share likelihood/CDF preparation, reject incompatible precomputed weights and document the numerical feature stack without adding plots. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Read the loaded namespace version after requireNamespace succeeds instead of rereading installed package metadata. Preserve all checks and function signatures while validating the dependency actually executing in each process. Add a regression covering metadata-free version validation and rejection of an insufficient version. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Move response-time baseline matrices into dedicated Cox response preparation and retain existing prediction and likelihood behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Keep advanced predictive controls in posterior_predict, document the interface boundary, and cover prediction summaries and Cox plotting safeguards. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Keep shared Cox predictive outputs unchanged while removing the standalone calibration APIs and their dedicated documentation and tests. The original implementation is preserved locally on followup/cox-calibration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Keep native prediction and RMST unchanged while removing utility exports and standalone help. Preserve survival and censoring recipes in an indexed Cox demo and adapt regression coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

wds15 · 2026-11-02T15:13:45Z

Hi @paul-buerkner !

Here is the Cox work I was mentioning. I have asked AI to stay as close as possible to conventions in the repo. From the history you can see that the agent was first quite prolific and did change a lot of things.

The agent also wrote (useful) functions which deal with residuals under censoring and visualising the baseline hazard. I felt like the best way to deal with them is to move them into the demo folder... but we can delete them also - whatever you think is best. Adding support for censored residuals in one or the other way would be cool, but is certainly a bigger task of its own.

Let me know if I can help in any way the review.

Best,
Sebastian

paul-buerkner

Thank you very much! I have added a bunch of comments mostly related to code structure and consistency with existing code.

paul-buerkner · 2026-11-05T12:48:24Z

Please remove the demo files once everything else is good, because I would like not to merge them into master.

Alternatively, do a vignette. We can then choose to not add them to the CRAN submission. @VisruthSK can you remind me how to exclude vignettes from CRAN? I forgot how exactly this was done.

Could do either

params: EVAL: !r identical(Sys.getenv("NOT_CRAN"), "true")

to have a vignette that doesn't get executed, or we could use an "article", in vignettes/articles which will be built by pkgdown but ignored by CRAN (see usethis::use_article()).

paul-buerkner · 2026-11-05T12:49:14Z

@@ -0,0 +1,221 @@
+# Evaluate the fitted baseline without extrapolating its spline basis.


I would prefer there not to be family-specific R files. Where in the existing files would they best fit? In posterior_predict, or in distributions, for example?

Since this is largely an AI PR, you can tell AI to better respect the existing file structure of brms for this purpose.

paul-buerkner · 2026-11-05T12:50:52Z

+    bs <- bhaz_basis(bframe, data, args = bhaz$args)
+    y <- bframe$frame$resp$values
+    if (!check_response) {
+      # Use in-support placeholders for unchecked response-time matrices only.


what does that comment mean?

paul-buerkner · 2026-11-05T12:52:30Z

+      c(out) <- data_response_bhaz(x, data, check_response = check_response)
+    }
+  } else if (is.brmsframe(bframe)) {
+    data <- subset_data(data, bframe)


why is bframe$dpears looped through here? If it is a function dealing with responses, it the dpars don't matter but the whole brmsframe should hold all the relevant information. If this is not actually true, why not?

paul-buerkner · 2026-11-05T12:53:06Z


-# data for the baseline functions of Cox models
+# Response-time matrices are separate from predictor dimensions and priors so
+# response checking also applies to nonlinear and multivariate Cox models.


I don't quite understand this comment. Why is the refactor as a data_response_ function necessary?

paul-buerkner · 2026-11-05T13:10:34Z

+  # sbhaz matrices, one per stratum; Jgrbhaz selects each observation's stratum.
+  out <- list(
+    basis = bframe$frame$basis$bhaz$basis_matrix,
+    version = cox_version,


can we somehow avoid needing the cox_version argument?

paul-buerkner · 2026-11-05T13:12:59Z

@@ -0,0 +1,441 @@
+context("Native Cox predictions")


do we really needed that many tests? Can we somehow shorten the amount of tests?

paul-buerkner · 2026-11-05T13:13:52Z

 })

+test_that("predict preserves raw draws and summary controls", {
+  fit <- brms:::rename_pars(brms:::brmsfit_example3)


paul-buerkner · 2026-11-05T13:21:06Z

    ape,
    arm,
    statmod,
+    prodlim,


for which purpose is this package needed?

that is used in the demo code - which I will drop as per your preference.. the demo code can become an article at a later stage. Thus, prodlim can be dropped - prodlim is a much faster alternative to the survival package to get empirical Kaplan Meier curve estimates.

paul-buerkner · 2026-11-05T13:21:54Z


+* Add native Cox predictive draws, CDFs, densities, and quantiles with explicit
+tail extrapolation policies, and restricted mean survival time with a required
+horizon. The Cox demo includes optional survival and censoring utility


remember to remove the demo note. What you can do instead of demo is a vignette that we then do not include in the CRAN version such that we do not have issues with package size upon CRAN submission.

wds15 · 2026-11-06T08:44:08Z

Many thanks for you review. I feel sorry for the many glitches from AI I put up. I have already spend quite some effort into aligning AI with brms conventions... but obviously not enough. I am collecting these conventions for me right now which makes it more straightforward in the future to have an agent respect the conventions automatically. If that is of interest I can share these as well (separate issue and PR on that). This should improve compliance of the agents wrt to how you structured the code.

paul-buerkner · 2026-11-06T09:25:46Z

Thank you! The challenge for the agents is perhaps also that my code structure conventions are not the best in many places. I would need considerable time (and AI help) to refactor but then the danger might be that not even I understand the brms codebase anymore. A scenario I am trying to avoid for as long as possible.

wds15 · 2026-11-06T09:57:22Z

The agents need to follow what we want. Given the many violations (agent vs what you expect) we see that suggests that the agents were trained differently. This is where good agent instructions can really help. These will make an agent respect the conventions by teaching the agent your conventions. If some of those considerations are written up somewhere, then that is a starting point for such agent instructions.

Copilot and others added 13 commits October 29, 2026 12:51

Localize Cox response checks without extending predictor interfaces

3268bab

Move response-time baseline matrices into dedicated Cox response preparation and retain existing prediction and likelihood behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Restore predict.brmsfit interface to master

2a0cdf1

Keep advanced predictive controls in posterior_predict, document the interface boundary, and cover prediction summaries and Cox plotting safeguards. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Explain performance regression guard in dependency test

dd7e617

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Clarify Cox prediction and likelihood input paths

7f1c1c9

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Explain Cox preparation and numerical invariants

de3b72b

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

bump NEWS

91d5224

paul-buerkner requested changes Nov 5, 2026

View reviewed changes

		@@ -0,0 +1,221 @@
		# Evaluate the fitted baseline without extrapolating its spline basis.

Conversation

wds15 commented Nov 2, 2026

Optional recipes in demo/

Scope and limitations

Uh oh!

wds15 commented Nov 2, 2026

Uh oh!

paul-buerkner left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wds15 commented Nov 6, 2026

Uh oh!

paul-buerkner commented Nov 6, 2026

Uh oh!

wds15 commented Nov 6, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Optional recipes in `demo/`