# Indexing using multi-indices with fixed sum

**URL:** https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486
**Category:** General Usage
**Tags:** indexing
**Created:** [January 4, 2023, 8:11am UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486 "2023-01-04T08:11:02Z")
**Posts on this page:** 9
**Page:** 1

<div class="post-metadata">

### Author: ![sbarthelme](https://avatars.discourse-cdn.com/v4/letter/s/839c29/32.png) [@sbarthelme](https://discourse.julialang.org/u/sbarthelme)
#### Post date: [January 4, 2023, 8:11am UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/1 "2023-01-04T08:11:03Z")

</div>

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](https://interval.louisiana.edu/reliable-computing-journal/volume-17/reliable-computing-17-pp-11-21.pdf).  
The algorithm is naturally expressed as a loop that involves a multi-index \alpha:  
 ![casteljau](https://global.discourse-cdn.com/julialang/original/3X/d/f/df5f4bd57f6241a5600c1f052feb6883aadda969.png)  
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?

---

<div class="post-metadata">

### Author: ![tim.holy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tim.holy/32/52_2.png) [@tim.holy](https://discourse.julialang.org/u/tim.holy)
#### Post date: [January 4, 2023, 9:20am UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/2 "2023-01-04T09:20:04Z")

</div>

It seems like [Multidimensional algorithms and iteration](https://julialang.org/blog/2016/02/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))`.

---

<div class="post-metadata">

### Author: ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)
#### Post date: [January 4, 2023, 9:43am UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/3 "2023-01-04T09:43:49Z")

</div>

I think this is what you want?

```julia
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
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.

---

<div class="post-metadata">

### Author: ![Per](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/per/32/10387_2.png) [@Per](https://discourse.julialang.org/u/Per)
#### Post date: [January 4, 2023, 11:09am UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/4 "2023-01-04T11:09:00Z")

</div>

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

```julia
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
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)

```

---

<div class="post-metadata">

### Author: ![sbarthelme](https://avatars.discourse-cdn.com/v4/letter/s/839c29/32.png) [@sbarthelme](https://discourse.julialang.org/u/sbarthelme)
#### Post date: [January 4, 2023, 2:53pm UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/5 "2023-01-04T14:53:32Z")

</div>

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](https://en.wikipedia.org/wiki/Barycentric_coordinate_system). 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:

```plaintext

#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

```

---

<div class="post-metadata">

### Author: ![Per](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/per/32/10387_2.png) [@Per](https://discourse.julialang.org/u/Per)
#### Post date: [January 4, 2023, 9:16pm UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/6 "2023-01-04T21:16:31Z")

</div>

Just one small thing; you probably want

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

```

to avoid type instability.

---

<div class="post-metadata">

### Author: ![sbarthelme](https://avatars.discourse-cdn.com/v4/letter/s/839c29/32.png) [@sbarthelme](https://discourse.julialang.org/u/sbarthelme)
#### Post date: [January 5, 2023, 8:02am UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/7 "2023-01-05T08:02:37Z")

</div>

Noted, thanks!

---

<div class="post-metadata">

### Author: ![sbarthelme](https://avatars.discourse-cdn.com/v4/letter/s/839c29/32.png) [@sbarthelme](https://discourse.julialang.org/u/sbarthelme)
#### Post date: [January 5, 2023, 3:23pm UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/8 "2023-01-05T15:23:59Z")

</div>

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

```julia
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

```

---

<div class="post-metadata">

### Author: ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 5, 2023, 4:40pm UTC](https://discourse.julialang.org/t/indexing-using-multi-indices-with-fixed-sum/92486/9 "2023-01-05T16:40:46Z")

</div>

> [@sbarthelme](#):
>
> 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](https://interval.louisiana.edu/reliable-computing-journal/volume-17/reliable-computing-17-pp-11-21.pdf).

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](https://github.com/stevengj/FastChebInterp.jl/blob/d97fc423b9f8eff4e850f667871d6425e6d9cee1/src/eval.jl#L16-L56) in [FastChebInterp.jl](https://github.com/stevengj/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`).
