Empirical Survival Function with Interval Censored and Truncated Data

truncEDF.jl

using Distributions
using Random
using Plots

function p_up!(p,n,m,j1,j2,dj)
    bj = zeros(m)
    for i in 1:n
        q = sum(p[j1[i]:j2[i]])
        if q < 1.0
        b = (1.0-q)/q
        ind = setdiff(1:m,j1[i]:j2[i])
        r = vec(p[ind])
        r ./= sum(r)
        bj[ind] += b*r
        end
    end
    num = dj + bj
    copy!(p, num ./ sum(num))
end

function d_up(p,j1,j2,n,m)
    dj = zeros(m)
    for i in 1:n
        q = p[j1[i]:j2[i]]
        q ./= sum(q)
        dj[j1[i]:j2[i]] += q
    end
    return dj
end

function acount(L,R,breaks,n,m)
    afirst = zeros(Int,n)
    alast = zeros(Int,n)
    for i in 1:n
        alpha = L[i] .<= breaks[1:(m-1)] .&& breaks[2:m] .<= R[i]
        afirst[i] = findfirst(alpha)
        alast[i] =  findlast(alpha)
    end
    return afirst, alast
end

function bcount(U,breaks,n,m)
    bfirst = zeros(Int,n)
    blast = zeros(Int,n)
    for i in 1:n
        beta = breaks .<= U[i]
        bfirst[i] = findfirst(beta)
        blast[i] =  findlast(beta)
    end
    return bfirst, blast
end

function setbreaks(trow, O=0.0)
    tj = sort(unique(max.(O, trow)))
    return tj
end

function make_ic(rng::AbstractRNG, y::Vector, at::Vector, tau::Vector)
    n = length(y)
    E_R = zeros(n)
    E_L = zeros(n)
    S = zeros(n)
    for i in 1:n
        ue = rand(rng)
        E_R[i] = at[i] + tau[i] * ue
        E_L[i] = max(0.0, at[i] - tau[i] * (1.0 - ue))
        S[i] = y[i] + at[i]
    end
    return E_L, E_R, S
end


function make_ic(rng::AbstractRNG, y::Number, at::Number, tau::Number)
    ue = rand(rng)
    E_R = at + tau * ue
    E_L = max(0.0, at - tau * (1.0 - ue))
    S = y + at
    return E_L, E_R, S
end

function maketic(rng, dist, Tmax,  N)
    L = zeros(N)
    R = zeros(N)
    S = zeros(N)
    for i in 1:N
        y = rand(rng,dist)
        at = rand(rng, Uniform(0.0, Tmax))
        tau = rand(rng, Exponential(1.0))
        L[i], R[i], S[i] = make_ic(rng,y,at,tau)
    end
    d = S .<= Tmax
    return L[d], R[d], S[d]
end

function SurvCT(L, R, S, Tmax, iter=100, tol=1e-10)
    tj = setbreaks([S-L;S-R])
    n = length(L)
    m = length(tj)
    aind = acount(S-R, S-L,tj,n,m)
    p = (1.0/m)*ones(m)
    bind = bcount(Tmax.-S,tj,n,m)
    dj = d_up(p,aind[1],aind[2],n,m)
    p_up!(p, n, m, bind[1], bind[2],dj)
    con = false
    count = 0
    p2 = p #こういうとこcopy!()とか使ったほうがいいですか？
    for it in 1:iter
        count += 1
        dj = d_up(p,aind[1],aind[2],n,m)
        p_up!(p, n, m, bind[1], bind[2], dj)
        con = all(abs.(p2-p) .< tol)
        p = p2
    if con || any(isnan.(p))
        break
    end
    end
    return tj, 1.0 .- cumsum(p), con, count
end

Random.seed!(1234)
rng = MersenneTwister()
dat = maketic(rng, Weibull(1.5, 10.0), 100.0, 100)
fit1 = SurvCT(dat[1], dat[2], dat[3], 100.0)
plot(fit1[1], fit1[2], linetype=:steppost, legend=false)
plot!(x -> ccdf(Weibull(1.5, 10.0),x))

dat = maketic(rng, Weibull(0.5, 10.0), 100.0, 100)
fit2 = SurvCT(dat[1], dat[2], dat[3], 100.0)
plot(fit2[1], fit2[2], linetype=:steppost, legend=false)
plot!(x -> ccdf(Weibull(0.5, 10.0),x))

function simit(N, it, dist)
    res = []
    for i in 1:it
        dat = maketic(rng, dist, 100.0, N)
        push!(res, SurvCT(dat[1], dat[2], dat[3], 100.0))
    end
    return res
end

function drawit(res, truedist)
    p = plot(res[1][1], res[1][2], legend=false, linetype=:steppost, color=:gray, alpha=0.1)
    for i in 2:length(res)
    plot!(p, res[i][1], res[i][2], inetype=:steppost, color=:gray, alpha=0.1)
    end
    plot!(p, x -> ccdf(truedist,x),color=:firebrick)
end

res1 = simit(500, 100, Weibull(1.5, 10.0))
res2 = simit(500, 100, Gamma(0.5, 20.0))
res3 = simit(500, 100, Exponential(10.0))
mlnorm = MixtureModel(LogNormal[LogNormal(1.0, 0.5), LogNormal(2.0, 0.1)], [0.3, 0.7])
res4 = simit(500, 100, mlnorm)

drawit(res1,Weibull(1.5, 10.0))
drawit(res2,Gamma(0.5, 20.0))
drawit(res3,Exponential(10.0))
drawit(res4,mlnorm)