From 3526d7a9ddbaf33d7aa266dd0d24760eb4239ade Mon Sep 17 00:00:00 2001 From: PJ Date: Tue, 18 Aug 2026 20:12:59 +0530 Subject: [PATCH] refactor(analyze): open the log-rank up to a weight on the risk set The log-rank is one member of a family that differs only in how much each event time counts. Nothing else changes: the counts it reports stay counts whatever the weight, and the published-dataset results are unmoved. --- cmd/internal-tools/analyze/survival.go | 30 ++++++++++++++------- cmd/internal-tools/analyze/survival_test.go | 19 +++++++++++++ 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/cmd/internal-tools/analyze/survival.go b/cmd/internal-tools/analyze/survival.go index 28ac2b0..5330700 100644 --- a/cmd/internal-tools/analyze/survival.go +++ b/cmd/internal-tools/analyze/survival.go @@ -88,13 +88,21 @@ type logRankResult struct { PValue float64 `json:"p_value"` } -// logRank is the k-sample Mantel-Haenszel log-rank test. At every distinct -// event time it contrasts observed with expected events under the null of equal -// hazards, then combines the k-1 independent differences through their -// covariance matrix: chi-square = U' V^-1 U on k-1 degrees of freedom. -// Mantel (1966); Peto and Peto (1972); Klein and Moeschberger, Survival -// Analysis, 2nd ed., section 7.3. +// logRank is the k-sample Mantel-Haenszel log-rank test, the member of the +// weighted family that counts every event time alike. Mantel (1966); Peto and +// Peto (1972). func logRank(names []string, groups [][]observation) logRankResult { + return weightedLogRank(names, groups, func(atRisk float64) float64 { return 1 }) +} + +// weightedLogRank is the family the log-rank belongs to. At every distinct event +// time it contrasts observed with expected events under the null of equal +// hazards, weights that difference by weight(atRisk), and combines the k-1 +// independent weighted differences through their covariance matrix: +// chi-square = U' V^-1 U on k-1 degrees of freedom. Observed and Expected stay +// event counts whatever the weight, because a weighted count is not one. +// Klein and Moeschberger, Survival Analysis, 2nd ed., section 7.3. +func weightedLogRank(names []string, groups [][]observation, weight func(atRisk float64) float64) logRankResult { var keptNames []string var kept [][]observation for index, group := range groups { @@ -130,6 +138,7 @@ func logRank(names []string, groups [][]observation) logRankResult { } atRisk := make([]float64, count) deaths := make([]float64, count) + weightedDifference := make([]float64, count) for _, steps := range distinctSteps(pooled) { totalAtRisk, totalDeaths := 0.0, 0.0 @@ -149,14 +158,17 @@ func logRank(names []string, groups [][]observation) logRankResult { if totalDeaths == 0 { continue } + weightAtStep := weight(totalAtRisk) for index := range groups { + expected := totalDeaths * atRisk[index] / totalAtRisk result.Observed[index] += deaths[index] - result.Expected[index] += totalDeaths * atRisk[index] / totalAtRisk + result.Expected[index] += expected + weightedDifference[index] += weightAtStep * (deaths[index] - expected) } if totalAtRisk <= 1 { continue } - scale := totalDeaths * (totalAtRisk - totalDeaths) / (totalAtRisk - 1) + scale := weightAtStep * weightAtStep * totalDeaths * (totalAtRisk - totalDeaths) / (totalAtRisk - 1) for row := range groups { share := atRisk[row] / totalAtRisk covariance[row][row] += scale * share * (1 - share) @@ -174,7 +186,7 @@ func logRank(names []string, groups [][]observation) logRankResult { for row := 0; row < count-1; row++ { reduced[row] = make([]float64, count-1) copy(reduced[row], covariance[row][:count-1]) - difference[row] = result.Observed[row] - result.Expected[row] + difference[row] = weightedDifference[row] } solution, ok := solveLinearSystem(reduced, difference) if !ok { diff --git a/cmd/internal-tools/analyze/survival_test.go b/cmd/internal-tools/analyze/survival_test.go index 78e06cf..87d928e 100644 --- a/cmd/internal-tools/analyze/survival_test.go +++ b/cmd/internal-tools/analyze/survival_test.go @@ -302,3 +302,22 @@ func TestLogRank_NoEventsAnywhere(t *testing.T) { t.Errorf("chi-square %v p-value %v, want 0 and 1", result.ChiSquare, result.PValue) } } + +// The weighted family reports event counts, not weighted ones: the report +// prints observed against expected as counts of violations, and a weight that +// changes the statistic must leave those alone. +func TestWeightedLogRank_WeightsTheStatisticAndNotTheCounts(t *testing.T) { + names := []string{"6-mp", "placebo"} + groups := [][]observation{gehanSixMercaptopurine, gehanPlacebo} + plain := logRank(names, groups) + weighted := weightedLogRank(names, groups, atRiskWeight) + + if !slices.Equal(plain.Observed, weighted.Observed) || !slices.Equal(plain.Expected, weighted.Expected) { + t.Errorf("weighted observed %v expected %v, want the counts %v and %v", + weighted.Observed, weighted.Expected, plain.Observed, plain.Expected) + } + if math.Abs(weighted.ChiSquare-plain.ChiSquare) < 1e-9 { + t.Errorf("chi-square %v under Gehan's weight and %v unweighted, want the weight to reach the statistic", + weighted.ChiSquare, plain.ChiSquare) + } +}