Estimator evaluation

docs/src/km.md

@@ -24,6 +24,7 @@ using Greenwood's formula:
 
 ```@docs
 Survival.KaplanMeier
+Survival.KaplanMeier(t::Real)
 StatsAPI.fit(::Type{KaplanMeier}, ::Any, ::Any)
 StatsAPI.confint(::KaplanMeier)
 ```

docs/src/na.md

@@ -24,6 +24,7 @@ from ``n_i`` samples:
 
 ```@docs
 Survival.NelsonAalen
+Survival.NelsonAalen(t::Real)
 StatsAPI.fit(::Type{NelsonAalen}, ::Any, ::Any)
 StatsAPI.confint(::NelsonAalen)
 ```

src/kaplanmeier.jl

@@ -23,6 +23,26 @@ struct KaplanMeier{S,T} <: NonparametricEstimator
     stderr::Vector{S}
 end
 
+"""
+To computes the survival estimation at time `t` of a KaplanMeier fit you can use.
+
+```julia
+km = fit(KaplanMeier, ...)
+t = 5
+km(t) # evaluates the estimator at time 5
+```
+"""
+function (km::KaplanMeier)(t::Real)
+    time_points = km.events.time
+    survival = km.survival
+    if t < time_points[1]
+        return estimator_start(typeof(km))
+    else
+        id = findlast(x -> x <= t, time_points) 
+        return survival[id]
+    end
+end
+
 estimator_eltype(::Type{<:KaplanMeier{S}}) where {S} = S
 estimator_eltype(::Type{KaplanMeier}) = Float64
 

src/nelsonaalen.jl

@@ -23,6 +23,26 @@ struct NelsonAalen{S,T} <: NonparametricEstimator
     stderr::Vector{S}
 end
 
+"""
+To computes the cumulative hazard estimation at time `t` of a NelsonAalen fit you can use.
+
+```julia
+na = fit(NelsonAalen, ...)
+t = 5
+na(t) # evaluates the estimator at time 5
+```
+"""
+function (na::NelsonAalen)(t::Real)
+    time_points = na.events.time
+    cumulative_hazard = na.chaz
+    if t < time_points[1]
+        return estimator_start(typeof(na))
+    else
+        id = findlast(x -> x <= t, time_points) 
+        return cumulative_hazard[id]
+    end
+end
+
 estimator_eltype(::Type{<:NelsonAalen{S}}) where {S} = S
 estimator_eltype(::Type{NelsonAalen}) = Float64
 

test/runtests.jl

@@ -163,6 +163,14 @@ end
     @test_throws ArgumentError fit(KaplanMeier, EventTime{Int}[])
 
     @test_deprecated confint(km, 0.05)
+
+    event_times = km.events.time
+    @test km.(event_times) == km.survival
+    # check that times between events have the right value
+    between_events = [(event_times[i] + event_times[i+1]) / 2 for i in 1:(length(event_times)-1)]
+    @test km(0) == 1
+    @test km.(between_events) == km.survival[1:end-1]
+    @test km(event_times[end] + 1) == km.survival[end]
 end
 
 @testset "Nelson-Aalen" begin
@@ -204,6 +212,16 @@ end
     @test na.stderr ≈ na_f32.stderr
 
     @test_deprecated confint(na, 0.05)
+
+
+    event_times = na.events.time
+    @test na.(event_times) == na.chaz
+    # check that times between events have the right value
+    between_events = [(event_times[i] + event_times[i+1]) / 2 for i in 1:(length(event_times)-1)]
+    @test na(0) == 0
+    @test na.(between_events) == na.chaz[1:end-1]
+    @test na(event_times[end] + 1) == na.chaz[end]
+
 end