diff --git a/src/auxiliary.jl b/src/auxiliary.jl index 3d7f0ec5..ac0732de 100644 --- a/src/auxiliary.jl +++ b/src/auxiliary.jl @@ -32,76 +32,128 @@ end _kron_promote(A::AbstractArray, B::AbstractArray, a, b) = _kron(A, B) # Manhattan based distance enumeration: I is supposed to be one-based index -# TODO: is there any way to make this faster? +# Points are zero-based coordinates `x` with `0 ≤ x[i] < sz[i]`, ordered first by their +# Manhattan distance `d = sum(x)` to the origin, and then lexicographically. -# forward mapping from multidimensional to single Manhattan index -@inline function num_manhattan_points(d::Int, sz::Dims{1}) - return Int(sz[1] > d) -end -@inline function num_manhattan_points(d::Int, sz::Dims{N}) where {N} - d == 0 && return 1 +# Number of points `x` with `0 ≤ x[i] < sz[i]` and Manhattan distance `sum(x) < d`. +# +# Without the upper bounds, the number of points `x ≥ 0` with `sum(x) ≤ d - 1` is +# `binomial(d - 1 + N, N)` (stars and bars). The points that violate the upper bound in +# (at least) a subset `S` of the dimensions, i.e. with `x[i] ≥ sz[i]` for `i ∈ S`, are +# counted in the same way after shifting `x[i] → x[i] - sz[i]`, which gives +# `binomial(d - 1 - s + N, N)` with `s = sum(sz[i] for i in S)`. By inclusion-exclusion, +# the number of points that satisfy all bounds is then +# ∑_S (-1)^|S| binomial(d - 1 - s + N, N) +# where the subsets `S` are enumerated as the bits of `mask`, and terms with `s ≥ d` vanish. +@inline function num_manhattan_points_below(d::Int, sz::Dims{N}) where {N} + d ≤ 0 && return 0 num = 0 - for i in 1:min(sz[1], d + 1) - num += num_manhattan_points(d - i + 1, Base.tail(sz)) + for mask in 0:((1 << N) - 1) + s = 0 + sign = 1 + for i in 1:N + if isodd(mask >> (i - 1)) + s += min(sz[i], d) # clamp to avoid overflow, terms with s ≥ d vanish anyway + sign = -sign + end + end + s < d && (num += sign * _binomial(d - 1 - s + N, Val(N))) end return num end +# binomial(n, N) for n ≥ N, unrolled for small static N +# (this is faster than a lookup table, since the divisions by constants are cheap) +function _binomial(n::Int, ::Val{N}) where {N} + r = 1 + for k in 1:N + r = (r * (n - N + k)) ÷ k + end + return r +end + +# forward mapping from multidimensional to single Manhattan index @inline localoffset(d, I::Tuple{Int}, sz::Tuple{Int}) = 0 @inline function localoffset(d, I::NTuple{N, Int}, sz::NTuple{N, Int}) where {N} - offset = 0 - for i in 1:(I[1] - 1) - offset += num_manhattan_points(d - i + 1, Base.tail(sz)) - end - offset += localoffset(d - I[1] + 1, Base.tail(I), Base.tail(sz)) - return offset + # number of points at distance `d` with a smaller first coordinate + sz′ = Base.tail(sz) + offset = num_manhattan_points_below(d + 1, sz′) - num_manhattan_points_below(d - I[1] + 2, sz′) + return offset + localoffset(d - I[1] + 1, Base.tail(I), sz′) end to_manhattan_index(::Tuple{}, ::Tuple{}) = 1 to_manhattan_index(I::Tuple{Int}, ::Tuple{Int}) = I[1] function to_manhattan_index(I::Dims{N}, sz::Dims{N}) where {N} d = sum(I) - N # Manhattan distance to origin for one-based indices - d == 0 && return 1 - index = 1 - # count all the points with smaller Manhatten distance - for k in 0:(d - 1) - index += num_manhattan_points(k, sz) - end - return index + localoffset(d, I, sz) + return 1 + num_manhattan_points_below(d, sz) + localoffset(d, I, sz) end # inverse mapping -@inline invertlocaloffset(d, offset, sz::Tuple{Int}) = (d + 1,) +invertlocaloffset(d, offset, sz::Tuple{Int}) = (d + 1,) @inline function invertlocaloffset(d, offset, sz::Tuple{Int, Int, Vararg{Int}}) - i₁ = 1 - while i₁ < sz[1] - jump = num_manhattan_points(d - i₁ + 1, Base.tail(sz)) - if offset < jump + # find the largest first coordinate `x₁` such that the number of points at distance `d` + # with a smaller first coordinate does not exceed `offset` + sz′ = Base.tail(sz) + nd = num_manhattan_points_below(d + 1, sz′) + lo, hi = 0, min(sz[1] - 1, d) + while lo < hi + mid = (lo + hi + 1) >> 1 + if nd - num_manhattan_points_below(d - mid + 1, sz′) ≤ offset + lo = mid + else + hi = mid - 1 + end + end + offset -= nd - num_manhattan_points_below(d - lo + 1, sz′) + return (lo + 1, invertlocaloffset(d - lo, offset, sz′)...) +end + +# successor of a multidimensional index in Manhattan order, assuming it exists +function next_manhattan_index(I::Dims{N}, sz::Dims{N}) where {N} + # find the last position that can be incremented while lowering the distance of the + # positions after it, and otherwise move to the next layer + r = 0 # Manhattan distance of the positions after k + k = 0 + @inbounds for j in N:-1:1 + if r > 0 && I[j] < sz[j] + k = j break end - offset -= jump - i₁ += 1 + r += I[j] - 1 + end + if k > 0 + I = @inbounds Base.setindex(I, I[k] + 1, k) + r -= 1 + else + r += 1 end - return (i₁, invertlocaloffset(d - i₁ + 1, offset, Base.tail(sz))...) + # the smallest index with distance `r` for the positions after k is obtained by + # filling them up from the back + @inbounds for j in N:-1:(k + 1) + x = min(r, sz[j] - 1) + I = Base.setindex(I, x + 1, j) + r -= x + end + return I end manhattan_to_multidimensional_index(index::Int, ::Dims{0}) = () manhattan_to_multidimensional_index(index::Int, ::Dims{1}) = (index,) function manhattan_to_multidimensional_index(index::Int, sz::Dims{N}) where {N} - # find the layer of the Manhattan distance - index == 1 && return ntuple(one, Val(N)) - offset = 1 - d = 1 - while true - currentlayer = num_manhattan_points(d, sz) - if index <= offset + currentlayer - break + # find the layer of the Manhattan distance, i.e. the smallest `d` such that there are at + # least `index` points with distance `≤ d`, by exponential followed by binary search + lo, hi = 0, 1 + while num_manhattan_points_below(hi + 1, sz) < index + lo, hi = hi + 1, 2hi + end + while lo < hi + mid = (lo + hi) >> 1 + if num_manhattan_points_below(mid + 1, sz) < index + lo = mid + 1 + else + hi = mid end - d += 1 - offset += currentlayer end - # find the position within the layer - index -= offset - return invertlocaloffset(d, index - 1, sz) + return invertlocaloffset(lo, index - 1 - num_manhattan_points_below(lo, sz), sz) end diff --git a/src/product.jl b/src/product.jl index 6848b617..bafca661 100644 --- a/src/product.jl +++ b/src/product.jl @@ -47,9 +47,14 @@ function findindex(P::SectorValues{I}, c::I) where {I <: ProductSector} return to_manhattan_index(findindex.(values.(_sectors(I)), Tuple(c)), _size(P)) end -function Base.iterate(P::SectorValues{I}, i = 1) where {I <: ProductSector} - Base.IteratorSize(P) != Base.IsInfinite() && i > length(P) && return nothing - return getindex(P, i), i + 1 +function Base.iterate(P::SectorValues{I}) where {I <: ProductSector} + inds = map(Returns(1), _size(P)) + return I(getindex.(values.(_sectors(I)), inds)), (1, inds) +end +function Base.iterate(P::SectorValues{I}, (i, inds)) where {I <: ProductSector} + Base.IteratorSize(P) != Base.IsInfinite() && i ≥ length(P) && return nothing + inds′ = next_manhattan_index(inds, _size(P)) + return I(getindex.(values.(_sectors(I)), inds′)), (i + 1, inds′) end function Base.convert(::Type{ProductSector{T}}, t::Tuple) where {T <: SectorTuple}