# \`rev\_cumsum\_exp!()\` optimization

**URL:** <https://discourse.julialang.org/t/rev-cumsum-exp-optimization/128664>\
**Category:** General Usage\
**Created:** [May 3, 2025, 3:50pm UTC](https://discourse.julialang.org/t/rev-cumsum-exp-optimization/128664 "2025-05-03T15:50:49Z")\
**Posts on this page:** 1\
**Showing post:** 5

<div class="post-metadata">

**Author:** ![lrnv](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lrnv/32/19373_2.png) [@lrnv](https://discourse.julialang.org/u/lrnv)\
**Post date:** [May 4, 2025, 8:43am UTC](https://discourse.julialang.org/t/rev-cumsum-exp-optimization/128664/5 "2025-05-04T08:43:07Z")

</div>

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](https://discourse.julialang.org/t/a-weird-operation-to-optimize/128658/14), where i ended up using your solution.

```julia
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](https://global.discourse-cdn.com/julialang/original/3X/3/a/3adbd13a40bbaebbb66d21d883515725d9610d0d.png)

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

---

_[View the full topic](https://discourse.julialang.org/t/rev-cumsum-exp-optimization/128664)._