Indexing using multi-indices with fixed sum

I’m trying to implement de Casteljau’s algorithm in a dimension-independent way. For the curious, it’s alg. 2.13 in this reference.
The algorithm is naturally expressed as a loop that involves a multi-index \alpha:
casteljau
Here the inner loop involves iterating over all multi-indices in dimension k that sum to d-l, where d is some degree.
For instance, with k=3, and d-l=2, the multi-indices are
\alpha = (3,0,0), \alpha = (2,1,0), \alpha = (1,2,0), …, \alpha=(0,0,3).
Doing this in arbitrary dimension involves generating all such multi-indices, and storing the computed coefficients in a space-efficient manner (ie, in a vector of doubles).
The problem is essentially the same as iterating over some ordering of the monomials with degree up to d-l, because the last coefficient of \alpha is determined by the first k-1 coefficients.
It’s quite easy to implement the alg. in dimension 1, and in dimension 2, but I’m struggling to find a nice Julian way of implementing this in arbitrary dimension. Any suggestions?

It seems like Multidimensional algorithms and iteration should be exactly what you’re looking for. I’d use i::Int for your first index and j::CartesianIndex{k-1} for the remainder, and compute the value of the i from sum(Tuple(j)).

1 Like

I think this is what you want?

function multind(::Val{N}, s) where {N}
	N == 1 && return [CartesianIndex(s)]
	mapreduce(vcat, 0:s) do i
		y = multind(Val(N-1), s-i)
		CartesianIndex.(y, i)
	end
end

which produces

julia> multind(Val(3), 3)
10-element Vector{CartesianIndex{3}}:
 CartesianIndex(3, 0, 0)
 CartesianIndex(2, 1, 0)
 CartesianIndex(1, 2, 0)
 CartesianIndex(0, 3, 0)
 CartesianIndex(2, 0, 1)
 CartesianIndex(1, 1, 1)
 CartesianIndex(0, 2, 1)
 CartesianIndex(1, 0, 2)
 CartesianIndex(0, 1, 2)
 CartesianIndex(0, 0, 3)

I’m not entirely sure if there’s a typo in the post, where you want the indices to sum to d-l=2, but instead they sum to 3 in the example. I assumed that you want the sum to be 3.

1 Like

It is also possible to do the same thing, but with iterators instead of vectors. (This avoids storing the indices in memory.)

function alpha_iterator(::Val{N}, s, t=()) where {N}
    N <= 1 && return ((s, t...),) # Iterator over a single Tuple
    Iterators.flatten(alpha_iterator(Val(N-1), s-i, (i, t...)) for i in 0:s)
end

which produces the same sequence:

julia> println.(alpha_iterator(Val(3),3));
(3, 0, 0)
(2, 1, 0)
(1, 2, 0)
(0, 3, 0)
(2, 0, 1)
(1, 1, 1)
(0, 2, 1)
(1, 0, 2)
(0, 1, 2)
(0, 0, 3)
1 Like

Thanks for your help everyone. All three solutions neatly address the problem of generating multi-indices in arbitrary dimension. I’m pasting some code below that evaluates simplicial Bernstein polynomials in all dimensions using Per’s iterator. I’m not satisfied with the way coefficients are stored (it wastes space by an increasing amount as dimension grows) but at least it is easy to code.
To recap: simplicial Bernstein polynomials are polynomials defined over a simplex (an interval in dimension 1, a triangle in dim. 2, etc.). They are naturally expressed in terms of barycentric coordinates. A polynomial expressed in Bernstein form reads
p(\lambda) = \sum_{|\alpha| = d} b_\alpha {d \choose \alpha} \lambda^\alpha
where b_\alpha is a coefficient, {d \choose \alpha} is a multinomial coefficient and \lambda is a barycentric coordinate.
The formula above can be used for direct evaluation, but de Casteljau’s algorithm is more numerically stable and generates other useful quantities.
The code below evaluates a Bernstein polynomial given its coefficients, see demo_1d and demo_2d:


#The coefficients of a degree-d simplical Bernstein polynomial
#in dimension k - 1 are stored in a k-1 dimensional array of size
#(d+1) × (d+1) × … × (d+1)
struct BernsteinPoly
    x :: Array{Float64}
end

#we index into the coefficients using a multi-index α of length k
to_ind = (α) -> CartesianIndex(α[1:(end-1)] .+ 1)
Base.getindex(bp :: BernsteinPoly,α  :: Tuple) = bp.x[to_ind(α)]
Base.getindex(bp :: BernsteinPoly,α  :: Vector) = bp[Tuple(α)]
Base.setindex!(bp :: BernsteinPoly,v,α  :: Tuple) = bp.x[to_ind(α)] = v

#initialise a zero polynomial 
function BernsteinPoly(deg::Int,dim::Int)
    x = zeros(fill(deg+1,dim)...)
    BernsteinPoly(x)
end

function dim(bp :: BernsteinPoly)
    ndims(bp.x)
end

Base.show(io::IO,bp :: BernsteinPoly) = print(io,"Bernstein polynomial of degree $(degree(bp)) in dimension $(dim(bp))")

function degree(bp :: BernsteinPoly)
    size(bp.x,1)-1
end

function multinomial(d,α)
    factorial(d) / prod(factorial.(α))
end

#iterator generating all the multi-indices α of size N
#that sum to s
#code by Per Rudquist on discourse
function alpha_iterator(::Val{N}, s, t=()) where {N}
    N <= 1 && return ((s, t...),) # Iterator over a single Tuple
    Iterators.flatten(alpha_iterator(Val(N-1), s-i, (i, t...)) for i in 0:s)
end


#evaluate Bernstein polynomial at λ,where λ is a vector of simplicial coordinates
#this uses the (unstable) naive summation method, i.e. the definition of simplicial Bernstein polynomials
function naive_eval(bp :: BernsteinPoly,λ)
    @assert length(λ) == dim(bp)+1
    val = 0.0
    d = degree(bp)
    for α in alpha_iterator(Val(dim(bp)+1),d)
        val += bp[α]*prod(λ.^α)*multinomial(d,α)
    end
    val
end

#evaluate Bernstein polynomial at λ using the de Casteljau algorithm.
#see e.g. Leroy, (2017) Convergence under Subdivision and Complexity of Polynomial Minimization in the Simplicial Bernstein Basis
function casteljau(bp :: BernsteinPoly,λ)
    @assert length(λ) == dim(bp)+1
    d = degree(bp)
    k = dim(bp)
    #initialise pyramid of Bernstein coefficients
    B = [bp;
         [BernsteinPoly(d-l,dim(bp)) for l in 1:d]];
    #index into coefficient matrix
    to_ind = (α) -> CartesianIndex(α[1:(end-1)] .+ 1)
    #Casteljau recursion
    for l in 1:d
        for α in alpha_iterator(Val(k+1),d-l)
            β = collect(α) #β is going to be the same as α except
            #along one dimension
            for p in 0:k #loop over dimension
                β[p+1]+=1
                tmp = B[l][β]
                tmp = λ[p+1]*tmp
                B[l+1][α] += tmp
                β[p+1] -=1
            end
        end
    end
    B[end][Tuple(fill(0,k+1))]
end

function demo_dim1()
    #We define a Bernstein polynomial of degree 2 in dim 1
    bp = BernsteinPoly([0,1,0])
    #all these three curves should match
    plot((x)-> naive_eval(bp,[x,1-x]),0,1)
    plot!((x)-> casteljau(bp,[x,1-x]),0,1)
    plot!((x)->2*x*(1-x),0,1)
end

function demo_dim2()
    #We define a Bernstein polynomial of degree 2 in dim 2
    bp = BernsteinPoly([0 2 1 ;
                        2 0 0;
                        1 0 0])
    xx = range(0,1,100)
    #these two plots should match
    p1=contour(xx,xx,(x,y)-> (x+y <= 1 ?  casteljau(bp,[x,y,1-(x+y)]) : 0.0))
    p2=contour(xx,xx,(x,y)-> (x+y <= 1 ?  naive_eval(bp,[x,y,1-(x+y)]) : 0.0))
    plot(p1,p2)
end


4 Likes

Just one small thing; you probably want

struct BernsteinPoly{N}
    x :: Array{Float64,N}
end

to avoid type instability.

Noted, thanks!

Ensuring type-stability took a bit of work. Here’s the updated code.

using Plots,LinearAlgebra

#The coefficients of a degree-d simplical Bernstein polynomial
#in dimension k - 1 are stored in a k-1 dimensional array of size
#(d+1) × (d+1) × … × (d+1)
struct BernsteinPoly{D} 
    x :: Array{Float64,D}
end

#we index into the coefficients using a multi-index α of length k
#for instance: BernsteinPoly([0.,.1,.4])[(1,1)]
#gives the coefficient corresponding to the Bernstein polynomial indexed by
#α=(1,1)
function Base.getindex(bp :: BernsteinPoly{D},α  :: Tuple) where D
    β = α[1:(end-1)].+ 1
    bp[CartesianIndex(β)]
end
Base.getindex(bp :: BernsteinPoly{D},α  :: Vector) where D = bp[Tuple(α)]
Base.getindex(bp :: BernsteinPoly{D},i  :: CartesianIndex) where D = bp.x[i] 
function Base.setindex!(bp :: BernsteinPoly{D},v :: Number,α  :: Tuple) where D
    β = α[1:(end-1)].+ 1
    bp.x[CartesianIndex(β)] = v
    bp
end


#initialise a zero polynomial 
function BernsteinPoly(deg::Int,dim::Int)
    x = zeros(fill(deg+1,dim)...)
    BernsteinPoly{dim}(x)
end

function BernsteinPoly(arr::AbstractArray)
    BernsteinPoly{ndims(arr)}(arr)
end

function dim(bp :: BernsteinPoly)
    ndims(bp.x)
end

Base.show(io::IO,bp :: BernsteinPoly) = print(io,"Bernstein polynomial of degree $(degree(bp)) in dimension $(dim(bp))")

function degree(bp :: BernsteinPoly)
    size(bp.x,1)-1
end

function multinomial(d,α)
    factorial(d) / prod(factorial.(α))
end

#iterator generating all the multi-indices α of size N
#that sum to s
#code by Per Rudquist on discourse
function alpha_iterator(::Val{N}, s, t=()) where {N}
    N <= 1 && return ((s, t...),) # Iterator over a single Tuple
    Iterators.flatten(alpha_iterator(Val(N-1), s-i, (i, t...)) for i in 0:s)
end


#evaluate Bernstein polynomial at λ,where λ is a vector of simplicial coordinates
#this uses the (unstable) naive summation method, i.e. the definition of simplicial Bernstein polynomials
function naive_eval(bp :: BernsteinPoly,λ)
    @assert length(λ) == dim(bp)+1
    val = 0.0
    d = degree(bp)
    for α in alpha_iterator(Val(dim(bp)+1),d)
        val += bp[α]*prod(λ.^α)*multinomial(d,α)
    end
    val
end

#evaluate Bernstein polynomial at λ using the de Casteljau algorithm.
#see e.g. Leroy, (2017) Convergence under Subdivision and Complexity of Polynomial Minimization in the Simplicial Bernstein Basis
function casteljau(bp :: BernsteinPoly{D},λ) where D
    @assert length(λ) == dim(bp)+1
    d = degree(bp)
    k = dim(bp)
    #initialise pyramid of Bernstein coefficients
    B = Vector{BernsteinPoly{D}}()
    push!(B,bp)
    for l in 1:d
        push!(B,BernsteinPoly(d-l,dim(bp)))
    end
    #index into coefficient matrix
    to_ind = (α) -> CartesianIndex(α[1:(end-1)] .+ 1)
    #Casteljau recursion
    for l in 1:d
        for α in alpha_iterator(Val(k+1),d-l)
            β = collect(α) #β is going to be the same as α except
                           #along one dimension
            acc = 0.0
            for p in 0:k #loop over dimension
                β[p+1]+=1
                bt = NTuple{D+1,Int64}(β) #convert to tuple
                acc += B[l][bt]*λ[p+1]
                β[p+1] -=1
            end
            B[l+1][α] = acc
        end
    end
    B[end].x[1] #top of pyramid has a single coefficient
end

function demo_dim1()
    #We define a Bernstein polynomial of degree 2 in dim 1
    bp = BernsteinPoly([0,1,0])
    #all these three curves should match
    plot((x)-> naive_eval(bp,[x,1-x]),0,1)
    plot!((x)-> casteljau(bp,[x,1-x]),0,1)
    plot!((x)->2*x*(1-x),0,1)
end

function demo_dim2()
    #We define a Bernstein polynomial of degree 1 in dim 2
    bp = BernsteinPoly([0 2 1 ;
                        2 0 0;
                        1 0 0])
    xx = range(0,1,100)
    p1=contour(xx,xx,(x,y)-> (x+y <= 1 ?  casteljau(bp,[x,y,1-(x+y)]) : 0.0))
    p2=contour(xx,xx,(x,y)-> (x+y <= 1 ?  naive_eval(bp,[x,y,1-(x+y)]) : 0.0))
    plot(p1,p2)
end

Note that there is nothing wrong with just doing explicit recursion over the dimension as long as the base case is large (e.g. a 1d polynomial evaluation) to amortize the recursion overhead, and it can be a very natural way to implement this sort of algorithm. That’s how I implemented a multidimensional Clenshaw recurrence in FastChebInterp.jl, for example.

A useful trick when doing explicit recurrence over the dimension to use linear indexing, exploiting the fact that Array is column-major (since you are storing the coefficients yourself you don’t need to handle AbstractArray).