Add meta collection utility

[charleskawczynski]

We should add a meta collection utility, for improving algorithm transparency, tests and optimization monitoring. Here's a scratch pad of what I had locally:

using Revise;
import ClimaTimeSteppers as CTS
import OrdinaryDiffEq as ODE
import LinearAlgebra as LA

struct SchurComplementW{T}; p::T; end
Base.similar(w::SchurComplementW) = w

# T_imp!(Yₜ, Y, p, t) = collect_meta(p, t, T_imp!)
# T_exp!(Yₜ, Y, p, t) = collect_meta(p, t, T_exp!)
# Wfact!(Yₜ, Y, p, Δt, t) = collect_meta(p, t, Wfact!)
# implicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, implicit_tendency!)
# explicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, explicit_tendency!)
# dss!(Y, p, t) = collect_meta(p, t, dss!)
# lim!(U, p, t, u) = collect_meta(p, t, lim!)
# T_lim!(Yₜ, Y, p, t) = collect_meta(p, t, T_lim!)
# stage_callback!(Y, p, t) = collect_meta(p, t, stage_callback!)
# tgrad!(Yₜ, Y, p, t) = collect_meta(p, t, tgrad!)
# LA.ldiv!(x, W::SchurComplementW, b) = collect_meta(p, nothing, LA.ldiv!)

Wfact!(Yₜ, Y, p, Δt, t) = collect_meta(p, t, :Wfact!)
implicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, :implicit_tendency!)
explicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, :explicit_tendency!)
dss!(Y, p, t) = nothing
lim!(U, p, t, u) = nothing
T_lim!(Yₜ, Y, p, t) = nothing
stage_callback!(Y, p, t) = nothing
tgrad!(Yₜ, Y, p, t) = collect_meta(p, t, :tgrad!)
LA.ldiv!(x, W::SchurComplementW, b) = collect_meta(p, nothing, :ldiv!)

struct AlgoMeta{D, T}
    call_counter_dict::D
    func_call_order::T
    function AlgoMeta()
        func_call_order = []
        call_counter_dict = Dict{Symbol,Vector{Int}}()
        D = typeof(call_counter_dict)
        T = typeof(func_call_order)
        return new{D, T}(call_counter_dict, func_call_order)
    end
end

function collect_meta(p, t, f)
    (; meta) = p
    fun = Symbol(f)
    push!(meta.func_call_order, (fun, t))
    if haskey(meta.call_counter_dict, fun)
        meta.call_counter_dict[fun][1] += 1
    else
        meta.call_counter_dict[fun] = Int[0]
    end
end
FT = Float64
Y₀ = zeros(FT, 1)
p = (; meta = AlgoMeta())
jac_prototype=SchurComplementW(p)
func_args = (; jac_prototype, Wfact = Wfact!, tgrad = tgrad!)
split_tendency_func = CTS.ClimaODEFunction(;
    T_exp! = explicit_tendency!,
    T_imp! = ODE.ODEFunction(implicit_tendency!; func_args...),
    dss! = dss!,
    lim!,
    T_lim!,
    stage_callback!,
)
prob = ODE.ODEProblem(split_tendency_func, Y₀, (FT(0), FT(1)), p)
alg = CTS.ARS343(CTS.NewtonsMethod(; max_iters=3))
integrator = ODE.init(prob, alg; dt=FT(0.1), save_everystep = true)
CTS.step_u!(integrator);

function summarize(meta::AlgoMeta)
    println("------------- Function call order")
    for (func, t) in meta.func_call_order
        println("     $func called at time $t")
    end
    println("------------- N-total calls")
    for k in keys(meta.call_counter_dict)
        meta.call_counter_dict[k] == [0] && continue
        println("     $(meta.call_counter_dict[k]): $k")
    end
    println("-------------")
    return nothing
end

summarize(p.meta)

using Revise;
import ClimaTimeSteppers as CTS
import OrdinaryDiffEq as ODE
import LinearAlgebra as LA

struct SchurComplementW{T}; p::T; end
Base.similar(w::SchurComplementW) = w

Wfact!(Yₜ, Y, p, Δt, t) = collect_meta(p, t, :Wfact!)
implicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, :implicit_tendency!)
explicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, :explicit_tendency!)
dss!(Y, p, t) = nothing
lim!(U, p, t, u) = nothing
T_lim!(Yₜ, Y, p, t) = nothing
stage_callback!(Y, p, t) = nothing
tgrad!(Yₜ, Y, p, t) = collect_meta(p, t, :tgrad!)
LA.ldiv!(x, W::SchurComplementW, b) = collect_meta(p, nothing, :ldiv!)


struct AlgoMeta{D, T}
    call_counter_dict::D
    func_call_order::T
    function AlgoMeta()
        func_call_order = []
        call_counter_dict = Dict{Symbol,Vector{Int}}()
        D = typeof(call_counter_dict)
        T = typeof(func_call_order)
        return new{D, T}(call_counter_dict, func_call_order)
    end
end

function collect_meta(p, t, f)
    (; meta) = p
    fun = Symbol(f)
    push!(meta.func_call_order, (fun, t))
    if haskey(meta.call_counter_dict, fun)
        meta.call_counter_dict[fun][1] += 1
    else
        meta.call_counter_dict[fun] = Int[0]
    end
end
FT = Float64;
Y₀ = zeros(FT, 1);
p = (; meta = AlgoMeta());
jac_prototype=SchurComplementW(p);
tspan = (FT(0), FT(1));

jac_kwargs = (; jac_prototype, Wfact = Wfact!);
remaining_func = explicit_tendency!;
prob = ODE.SplitODEProblem(
    ODE.ODEFunction(
        implicit_tendency!;
        jac_kwargs...,
        tgrad = tgrad!,
    ),
    remaining_func,
    Y₀,
    tspan,
    p,
);

alg = CTS.ARS343(CTS.NewtonsMethod(; max_iters=3));
integrator = ODE.init(prob, alg; dt=FT(0.1), save_everystep = true);
CTS.step_u!(integrator);

function summarize(meta::AlgoMeta)
    println("------------- Function call order")
    for (func, t) in meta.func_call_order
        println("     $func called at time $t")
    end
    println("------------- N-total calls")
    for k in keys(meta.call_counter_dict)
        meta.call_counter_dict[k] == [0] && continue
        println("     $(meta.call_counter_dict[k]): $k")
    end
    println("-------------")
    return nothing
end

summarize(p.meta)

[charleskawczynski]

More local notes:

using Revise;
import ClimaTimeSteppers as CTS
import OrdinaryDiffEq as ODE
import LinearAlgebra as LA

struct SchurComplementW{T}; p::T; end
Base.similar(w::SchurComplementW) = w
T_imp!(Yₜ, Y, p, t) = collect_meta(p, t, T_imp!)
T_exp!(Yₜ, Y, p, t) = collect_meta(p, t, T_exp!)
Wfact!(Yₜ, Y, p, Δt, t) = collect_meta(p, t, Wfact!)
implicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, implicit_tendency!)
explicit_tendency!(Yₜ, Y, p, t) = collect_meta(p, t, explicit_tendency!)
tgrad!(Yₜ, Y, p, t) = collect_meta(p, t, tgrad!)
LA.ldiv!(x, W::SchurComplementW, b) = collect_meta(p, nothing, LA.ldiv!)

function collect_meta(p, t, f)
    (; algo_mapper) = p
    push!(algo_mapper[string(f)], t)
    algo_mapper["ctr"][1] += 1
end
Y₀ = zeros(Float64, 1)
FT = Float64
algo_mapper = Dict(
    "ctr" => [0],
    "T_imp!" => [],
    "T_exp!" => [],
    "Wfact!" => [],
    "implicit_tendency!" => [],
    "explicit_tendency!" => [],
    "tgrad!" => [],
    "ldiv!" => [],
)
p = (; algo_mapper)
jac_prototype=SchurComplementW(p)
func_args = (; jac_prototype, Wfact = Wfact!, tgrad = tgrad!)
split_tendency_func = CTS.ClimaODEFunction(;
    T_exp! = explicit_tendency!,
    T_imp! = ODE.ODEFunction(implicit_tendency!; func_args...),
)
prob = ODE.ODEProblem(split_tendency_func, Y₀, (FT(0), FT(1)), p)
alg = CTS.NewARS343(CTS.NewtonsMethod(; max_iters=1))
integrator = ODE.init(prob, alg; dt=FT(0.1), save_everystep = true)
CTS.step_u!(integrator)