Faster Manhattan indexing for ProductSector

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

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}