using Dictionaries: getindices function cartesian_to_linear(dims::Tuple) return Dictionary(vec(Tuple.(CartesianIndices(dims))), 1:prod(dims)) end front(itr, n=1) = Iterators.take(itr, length(itr) - n) tail(itr) = Iterators.drop(itr, 1) # Tree utils function line_to_tree(line::Vector) if length(line) == 1 && line[1] isa Vector return line[1] end if length(line) <= 2 return line end return [line_to_tree(line[1:(end - 1)]), line[end]] end # Pad with last value to length or truncate to length. # If it is a single value (non-Vector), fill with # that value to the length. function extend_or_truncate(x::Vector, length::Int) l = length - Base.length(x) return l >= 0 ? [x; fill(last(x), l)] : x[1:length] end extend_or_truncate(x, length::Int) = extend_or_truncate([x], length) using StructWalk: StructWalk, WalkStyle, postwalk # Treat `AbstractArray` as leaves. struct AbstractArrayLeafStyle <: WalkStyle end StructWalk.children(::AbstractArrayLeafStyle, x::AbstractArray) = () function extend_or_truncate_columns(nt::NamedTuple, length::Int) return map(x -> extend_or_truncate(x, length), nt) end function extend_or_truncate_columns_recursive(nt::NamedTuple, length::Int) return postwalk(AbstractArrayLeafStyle(), nt) do x x isa NamedTuple && return x return extend_or_truncate(x, length) end end #ToDo: remove #nrows(nt::NamedTuple) = isempty(nt) ? 0 : length(first(nt)) function row(nt::NamedTuple, i::Int) isempty(nt) ? (return nt) : (return map(x -> x[i], nt)) end # Similar to `Tables.rowtable(x)` function rows(nt::NamedTuple, length::Int) return [row(nt, i) for i in 1:length] end function rows_recursive(nt::NamedTuple, length::Int) return postwalk(AbstractArrayLeafStyle(), nt) do x !(x isa NamedTuple) && return x return rows(x, length) end end function expand(nt::NamedTuple, length::Int) nt_padded = extend_or_truncate_columns_recursive(nt, length) return rows_recursive(nt_padded, length) end function interleave(a::Vector, b::Vector) ab = flatten(collect(zip(a, b))) if length(a) == length(b) return ab elseif length(a) == length(b) + 1 return append!(ab, [last(a)]) else error( "Trying to interleave vectors of length $(length(a)) and $(length(b)), not implemented.", ) end end function getindices_narrow_keytype(d::Dictionary, indices) return convert(typeof(d), getindices(d, indices)) end