`rev_cumsum_exp!()` optimization

General Usage
5

You are right. Here is a full example, where I highlighted as a function the sensitive operation (eating 90% of my runtime). Its a mix between the code we were looking at here and the one from the previous thread, where i ended up using your solution.

using RDatasets
using LinearAlgebra
using BenchmarkTools
using StatsBase

# The dataset: 
const colon = let
    df = dataset("survival","colon")
    df = dropmissing(df, [:Nodes, :Differ])
    
    T = Float64.(df.Time)
    Δ = Int64.(df.Status)
    X = Float64.(hcat(df.Sex, df.Age, df.Obstruct,
    df.Perfor, df.Adhere, df.Nodes,
    df.Differ, df.Extent, df.Surg,
    df.Node4, df.EType))
    (T,Δ,X)
end

struct Cox
    X::Matrix{Float64}
    T::Vector{Float64}
    Δ::Vector{Bool}
    sX::Vector{Float64}
    G::Vector{Float64}
    η::Vector{Float64}
    B::Vector{Float64}
    I::Vector{Int64}
    J::Vector{Int64}
    _A::Vector{Float64}
    _B::Vector{Float64}
    function Cox(T,Δ,X)

        # Allocate: 
        n,m = size(X)
        G, B  = zeros(m), zeros(m)
        η, _A, _B  = zeros(n), zeros(n), zeros(n)

        # Precompute a few things: 
        # This should also be optimized, its taking a lot of time right now..
        o = sortperm(T)
        To = T[o]
        Δo = Bool.(Δ[o])
        Xo = X[o,:]
        sX = X'Δ
        Ro = StatsBase.competerank(To)
        I = Ro[Δo]

        # J is a form of ranks, filtered on Δ. Did not really found a function to do it. 
        J = zeros(Int64,n)
        for j in 1:n
            if Δo[j]
                for i in 1:n
                    J[i] += (To[i] >= To[j])
                end
            end
        end
        
        # Compute hessian bounds:
        for l in 1:m # for each dimension. 
            lastj = n+1
            Mₓ = Xo[end,l]
            mₓ = Xo[end,l]
            for i in n:-1:1
                for j in lastj-1:-1:Ro[i]
                    Mₓ = max(Xo[j,l], Mₓ)
                    mₓ = min(Xo[j,l], mₓ)
                end
                lastj = Ro[i]
                B[l] += (1/4) * Δ[i] * (Mₓ-mₓ)^2
            end
        end

        # Instantiate: 
        new(Xo, To, Δo, sX, G, η, B, I, J, _A, _B)
    end
end

function most_of_my_runtime!(A,B,η,I,J)
    # This is equivalent to : 
    #   A .= exp.(η)
    #   B .= 1 ./ reverse(cumsum(reverse(A)))
    #   A .*= cumsum(B[I])[J]
    # But 40% faster for the moment.

    aₖ = bₖ = cₖ = zero(eltype(A))
    n, lastj = length(η), 0
    local j

    @inbounds for k in n:-1:1
        aₖ = exp(η[k])
        A[k] = aₖ
        bₖ += aₖ
        B[k] = bₖ
    end
    @inbounds for k in 1:n
        for outer j in lastj+1:J[k]
            cₖ += inv(B[I[j]])
        end
        lastj = j
        A[k] *= cₖ
    end
end

function update!(β, M::Cox)
    mul!(M.η, M.X, β)
    most_of_my_runtime!(M._A, M._B, M.η, M.I, M.J)
    mul!(M.G, M.X', M._A)
    M.G .= M.sX .- M.G
    β .+= M.G ./ M.B 
    return nothing
end

function getβ(M::Cox; max_iter = 10000, tol = 1e-6)

    β = zeros(size(M.X, 2))
    βᵢ = similar(β)
    for i in 1:max_iter
        βᵢ .= β
        update!(β, M) 
        gap = norm(βᵢ - β)
        # println(gap)
        if gap < tol
            break
        end
    end
    return β
end

M = Cox(colon...)
getβ(M)
@profview [getβ(M) for i in 1:100]

The profview I have looks like:

image
image1199×510 29.3 KB

I thought about trying to intertwine the two loops, but i was not able to find a better version yet.