Quick and dirty plotting functions

[Raumschifffan]

Not sure if this is the right place for this, but if anyone needs to plot the Kaplan-Meier-Estimates with Plots.jl, feel free to use the script below.

using Survival
using Plots

for f in [:plot, :plot!]
    @eval begin
        function Plots.$f(kme::KaplanMeier{S, T}; full=false, conf=false, options...) where {S, T}
            x = kme.events.time
            y = kme.survival
            options = Dict{Symbol, Any}(options)

            if conf
                lower, upper = collect(zip(Survival.confint(kme)...))
                upper = replace(upper .- y, NaN=>zero(S))
                lower = replace(y .- lower, NaN=>zero(S))
                options[:ribbon] = (upper, lower)
            end
            if full 
                x = [zero(T); x]
                y = [one(S); y]
                options[:ylims] = (zero(S), one(S))
                if conf
                    upper, lower = options[:ribbon]
                    options[:ribbon] = [zero(S); upper], [zero(S); lower]
                end
            end

            return $f(x, y; linewidth=2, options...)
        end

        function Plots.$f(kmes::Vector{KaplanMeier}; generators=Dict(), options...)
            $f()
            options = Dict{Symbol, Any}(options)
            for (i, kme) in enumerate(kmes)
                for (key, generator) in generators
                    options[key] = generator(i)
                end
                plot!(kme; options...)
            end
            return plot!(show=true)
        end
    end
end

To use like:

x = groupby(df, :sex)
y = [fit(KaplanMeier, group.time, group.event) for group in x]
generators = Dict(
    :label => (i) -> "$(x[i].sex[1])",
    :linewidth => (i) -> 2^i
)
plot(y, full=true, conf=true, generators=generators)