diff --git a/src/QuantEcon.jl b/src/QuantEcon.jl index 356f69a5..fe3d5ca1 100644 --- a/src/QuantEcon.jl +++ b/src/QuantEcon.jl @@ -97,7 +97,7 @@ export smooth, periodogram, ar_periodogram, # util - meshgrid, gridmake, gridmake!, ckron, is_stable, + meshgrid, gridmake, gridmake!, ckron, is_stable, num_compositions, simplex_grid, simplex_index, # robustlq RBLQ, diff --git a/src/util.jl b/src/util.jl index b8fcd3c9..c323bfb3 100644 --- a/src/util.jl +++ b/src/util.jl @@ -145,3 +145,128 @@ function is_stable(A::AbstractMatrix) end + +""" +The total number of m-part compositions of n, which is equal to + +(n + m - 1) choose (m - 1) + +##### Arguments + +- `m`::Int : Number of parts of composition +- `n`::Int : Integer to decompose + +##### Returns + +- ::Int - Total number of m-part compositions of n +""" +function num_compositions(m, n) + return binomial(n+m-1, m-1) +end + + +""" +Construct an array consisting of the integer points in the +(m-1)-dimensional simplex :math:`\{x \mid x_0 + \cdots + x_{m-1} = n +\}`, or equivalently, the m-part compositions of n, which are listed +in lexicographic order. The total number of the points (hence the +length of the output array) is L = (n+m-1)!/(n!*(m-1)!) (i.e., +(n+m-1) choose (m-1)). + +##### Arguments + +- `m`::Int : Dimension of each point. Must be a positive integer. +- `n`::Int : Number which the coordinates of each point sum to. Must + be a nonnegative integer. + +##### Returns +- `out`::Matrix{Int} : Array of shape (m, L) containing the integer + points in the simplex, aligned in lexicographic + order. + +##### Notes + +A grid of the (m-1)-dimensional *unit* simplex with n subdivisions +along each dimension can be obtained by `simplex_grid(m, n) / n`. + +##### Examples + +>>> simplex_grid(3, 4) + +3×15 Array{Int64,2}: + 0 0 0 0 0 1 1 1 1 2 2 2 3 3 4 + 0 1 2 3 4 0 1 2 3 0 1 2 0 1 0 + 4 3 2 1 0 3 2 1 0 2 1 0 1 0 0 + +##### References + +A. Nijenhuis and H. S. Wilf, Combinatorial Algorithms, Chapter 5, + Academic Press, 1978. +""" +function simplex_grid(m, n) + # Get number of elements in array and allocate + L = num_compositions(m, n) + out = Array{Int}(m, L) + + x = zeros(Int, m) + x[m] = n + + # Fill in first column + copy!(out, 1, x, 1, m) + + h = m + for i in 2:L + h -= 1 + + val = x[h+1] + x[h+1] = 0 + x[m] = val - 1 + x[h] += 1 + + copy!(out, 3*(i-1) + 1, x, 1, m) + + if val != 1 + h = m + end + end + + return out +end + + +""" +Return the index of the point x in the lexicographic order of the +integer points of the (m-1)-dimensional simplex :math:`\{x \mid x_0 ++ \cdots + x_{m-1} = n\}`. + +##### Arguments + +- `x`::Array{Int,1} : Integer point in the simplex, i.e., an array of + m nonnegative integers that sum to n. +- `m`::Int : Dimension of each point. Must be a positive integer. +- `n`::Int : Number which the coordinates of each point sum to. Must be a + nonnegative integer. + +##### Returns +- `idx`::Int : Index of x. + +""" +function simplex_index(x, m, n) + # If only one element then only one point in simplex + if m==1 + return 1 + end + + decumsum = reverse(cumsum(reverse(x[2:end]))) + idx = binomial(n+m-1, m-1) + for i in 1:m-1 + if decumsum[i] == 0 + break + end + + idx -= num_compositions(m - (i-1), decumsum[i]-1) + end + + return idx +end + diff --git a/test/runtests.jl b/test/runtests.jl index 0755250a..e8d0d6a0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -29,6 +29,7 @@ tests = [ "optimization", "interp", "sampler", + "util", ] diff --git a/test/test_util.jl b/test/test_util.jl index 24efb756..f9ad9f43 100644 --- a/test/test_util.jl +++ b/test/test_util.jl @@ -30,16 +30,35 @@ end @testset "fix" begin - @test 1 == @inferred fix(1.2) - @test [0, 2] == @inferred fix([0.9, 2.1]) - @test [0, 2] == @inferred fix([0, 2]) + @test 1 == @inferred QuantEcon.fix(1.2) + @test [0, 2] == @inferred QuantEcon.fix([0.9, 2.1]) + @test [0, 2] == @inferred QuantEcon.fix([0, 2]) out = [100, 100] - @inferred fix!([0, 2], out) + @inferred QuantEcon.fix!([0, 2], out) @test [0, 2] == out - @test [0 2; 2 0] == @inferred fix([0.5 2.9999; 3-2eps() -0.9]) + @test [0 2; 2 0] == @inferred QuantEcon.fix([0.5 2.9999; 3-2eps() -0.9]) end @test ([1 2; 1 2], [3 3; 4 4]) == @inferred meshgrid([1, 2], [3, 4]) + + @testset "simplex_tools" begin + grid_3_4 = [0 0 0 0 0 1 1 1 1 2 2 2 3 3 4 + 0 1 2 3 4 0 1 2 3 0 1 2 0 1 0 + 4 3 2 1 0 3 2 1 0 2 1 0 1 0 0] + points = [0 1 4 + 0 1 0 + 4 2 0] + idx = Array{Int}(3) + + for i in 1:3 + idx[i] = simplex_index(points[:, i], 3, 4) + end + + @test all(simplex_grid(3, 4) .== grid_3_4) + @test all(grid_3_4[:, idx] .== points) + @test size(grid_3_4, 2) == num_compositions(3, 4) + + end end