Replace _get_idx w/ CartesianIndices/LinearIndices

[charleskawczynski]

I went to add some topology unit tests, and realized that there are no unit tests/documentation for _get_idx. It turns out that we can simply replace most of the methods with two new simpler functions:

Base.@propagate_inbounds cart_ind(n, i::Integer) = @inbounds CartesianIndices(map(x->Base.OneTo(x), n))[i]
Base.@propagate_inbounds linear_ind(n, ci::CartesianIndex) = @inbounds LinearIndices(map(x->Base.OneTo(x), n))[ci]

Which is actually safer in that it will throw bounds errors if we pass in bad indices, and it's also faster on average:

#=
julia --project
using Revise; include(joinpath("test", "Topologies", "benchmark_get_idx.jl"))
=#
using ClimaCore.Topologies: _get_idx, _get_idx_metric
using Test
using BenchmarkTools

Base.@propagate_inbounds cart_ind(n, i::Integer) = @inbounds CartesianIndices(map(x->Base.OneTo(x), n))[i]
Base.@propagate_inbounds linear_ind(n, ci::CartesianIndex) = @inbounds LinearIndices(map(x->Base.OneTo(x), n))[ci]
get_idx(sizet, idx) = @inbounds _get_idx(sizet, idx)

@testset "_get_idx" begin
    # _get_idx(sizet::NTuple{3, Int}, loc::Int)
    for sizet in (
        # (1,1,1),
        (3,4,5),
        # (5,4,3),
        )
        for idx in 1:10:prod(sizet)
            t1 = @belapsed cart_ind($sizet, $idx)
            t2 = @belapsed get_idx($sizet, $idx)
            @test t1 ≤ t2
        end
    end

    # _get_idx(sizet::NTuple{5, Int}, loc::NTuple{5, Int})
    for sizet in (
        # (1,1,1,1),
        (3,4,5,6),
        # (6,5,4,3),
        )
        for idx in 1:100:prod(sizet)
            t1 = @belapsed cart_ind($sizet, $idx)
            t2 = @belapsed get_idx($sizet, $idx)
            @test t1 ≤ t2
            ci = cart_ind(sizet, idx)
            t1 = @belapsed linear_ind($sizet, $ci)
            t2 = @belapsed get_idx($sizet, $(ci.I))
            @test t1 ≤ t2
        end
    end

    # _get_idx(sizet::NTuple{4, Int}, loc::NTuple{4, Int})
    for sizet in (
        # (1,1,1,1,1),
        (3,4,5,6,6),
        # (6,5,4,3,3),
        )
        for idx in 1:200:prod(sizet)
            ci = cart_ind(sizet, idx)
            t2 = @belapsed get_idx($sizet, $(ci.I))
            t1 = @belapsed linear_ind($sizet, $ci)
            @test t1 ≤ t2
        end
    end

end

This seems like a win on all fronts: new code is documented (in base), safer, faster, and fewer lines of code.

If this doesn't workout, I'll revert and add the unit test for _get_idx:

-#=
julia --project
using Revise; include(joinpath("test", "Topologies", "unit_get_idx.jl"))
=#
using ClimaCore.Topologies: _get_idx, _get_idx_metric
using Test

Base.@propagate_inbounds cart_ind(n, i::Integer) = CartesianIndices(map(x->Base.OneTo(x), n))[i]
Base.@propagate_inbounds linear_ind(n, ci::CartesianIndex) = LinearIndices(map(x->Base.OneTo(x), n))[ci]

# sizet_metric = (nlevels, Nq, Nq, nmetric, nelems)
# Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
# function _get_idx_metric(sizet::NTuple{5, Int}, loc::NTuple{4, Int})
#     nmetric = sizet[4]
#     (i11, i12, i21, i22) = nmetric == 4 ? (1, 2, 3, 4) : (1, 2, 4, 5)
#     (level, i, j, elem) = loc
#     return (
#         _get_idx(sizet, (level, i, j, i11, elem)),
#         _get_idx(sizet, (level, i, j, i12, elem)),
#         _get_idx(sizet, (level, i, j, i21, elem)),
#         _get_idx(sizet, (level, i, j, i22, elem)),
#     )
#     return nothing
# end

Base.@propagate_inbounds function metric_ind(sizet::NTuple{5, Int}, loc::NTuple{4, Int})
    @inbounds begin
        nmetric = sizet[4]
        (i21, i22) = nmetric == 4 ? (3, 4) : (4, 5)
        (level, i, j, elem) = loc
        inds = (
            linear_ind(sizet, CartesianIndex(level, i, j, 1, elem)),
            linear_ind(sizet, CartesianIndex(level, i, j, 2, elem)),
            linear_ind(sizet, CartesianIndex(level, i, j, i21, elem)),
            linear_ind(sizet, CartesianIndex(level, i, j, i22, elem)),
        )
        return inds
    end
end


@testset "_get_idx" begin
    # _get_idx(sizet::NTuple{3, Int}, loc::Int)
    for sizet in (
        (1,1,1),
        (3,4,5),
        (5,4,3),
        )
        for idx in 1:prod(sizet)
            @test _get_idx(sizet, idx) == cart_ind(sizet, idx).I
        end
        @test_throws BoundsError cart_ind(sizet, prod(sizet)+1).I
        @test_throws BoundsError cart_ind(sizet, 0).I
        @test_throws BoundsError cart_ind(sizet, -1).I
    end

    # _get_idx(sizet::NTuple{5, Int}, loc::NTuple{5, Int})
    for sizet in (
        (1,1,1,1),
        (3,4,5,6),
        (6,5,4,3),
        )
        for idx in 1:prod(sizet)
            @test _get_idx(sizet, idx) == cart_ind(sizet, idx).I
            ci = cart_ind(sizet, idx)
            @test _get_idx(sizet, ci.I) == linear_ind(sizet, ci)
        end
        @test_throws BoundsError cart_ind(sizet, prod(sizet)+1).I
        @test_throws BoundsError cart_ind(sizet, 0).I
        @test_throws BoundsError cart_ind(sizet, -1).I
    end

    # _get_idx(sizet::NTuple{4, Int}, loc::NTuple{4, Int})
    for sizet in (
        (1,1,1,1,1),
        (3,4,5,6,6),
        (6,5,4,3,3),
        )
        for idx in 1:prod(sizet)
            ci = cart_ind(sizet, idx)
            @test _get_idx(sizet, ci.I) == linear_ind(sizet, ci)
        end
        @test_throws BoundsError cart_ind(sizet, prod(sizet)+1).I
        @test_throws BoundsError cart_ind(sizet, 0).I
        @test_throws BoundsError cart_ind(sizet, -1).I
    end

    # sizet=(1,1,1,4,1); loc=(1,1,2,4); @test _get_idx_metric(sizet, loc) == metric_ind(sizet, loc)
    # sizet=(3,4,5,6,6); loc=(3,4,5,6); @test _get_idx_metric(sizet, loc) == metric_ind(sizet, loc)
    # sizet=(3,4,5,4,6); loc=(3,4,5,4); @test _get_idx_metric(sizet, loc) == metric_ind(sizet, loc)
    # sizet=(6,5,4,3,3); loc=(6,5,4,3); @test _get_idx_metric(sizet, loc) == metric_ind(sizet, loc)
end