Extended tropical for computing spectrums (#24)

* update

* tests pass

* fix docs

* slow but works

Author: GiggleLiu

docs/src/ref.md

@@ -80,6 +80,7 @@ is_commutative_semiring
 ```@docs
 TropicalNumbers.Tropical
 TropicalNumbers.CountingTropical
+ExtendedTropical
 Mods.Mod
 Polynomials.Polynomial
 TruncatedPoly

src/GraphTensorNetworks.jl

@@ -18,7 +18,7 @@ export estimate_memory
 export StaticBitVector, StaticElementVector, @bv_str
 export is_commutative_semiring
 export Max2Poly, TruncatedPoly, Polynomial, Tropical, CountingTropical, StaticElementVector, Mod, ConfigEnumerator, onehotv, ConfigSampler, TreeConfigEnumerator
-export CountingTropicalF64, CountingTropicalF32, TropicalF64, TropicalF32
+export CountingTropicalF64, CountingTropicalF32, TropicalF64, TropicalF32, ExtendedTropical
 
 # Lower level APIs
 export AllConfigs, SingleConfig

src/arithematics.jl

@@ -79,6 +79,9 @@ function is_commutative_semiring(a::T, b::T, c::T) where T
     return true
 end
 
+######################## Truncated Polynomial ######################
+# TODO: store orders to support non-integer weights
+# (↑) Maybe not so nessesary, no use case for counting degeneracy when using floating point weights.
 """
     TruncatedPoly{K,T,TO} <: Number
     TruncatedPoly(coeffs::Tuple, maxorder)
@@ -165,6 +168,124 @@ function Base.show(io::IO, ::MIME"text/plain", x::TruncatedPoly{K}) where K
     end
 end
 
+############################ ExtendedTropical #####################
+"""
+    ExtendedTropical{K,TO} <: Number
+    ExtendedTropical(orders)
+
+Extended Tropical numbers with largest `K` orders keeped,
+or the [`TruncatedPoly`](@ref) without coefficients,
+`TO` is the element type of orders.
+
+Example
+------------------------------
+```jldoctest; setup=(using GraphTensorNetworks)
+julia> x = ExtendedTropical{3}([1.0, 2, 3])
+ExtendedTropical{3, Float64}([1.0, 2.0, 3.0])
+
+julia> y = ExtendedTropical{3}([-Inf, 2, 5])
+ExtendedTropical{3, Float64}([-Inf, 2.0, 5.0])
+
+julia> x * y
+ExtendedTropical{3, Float64}([6.0, 7.0, 8.0])
+
+julia> x + y
+ExtendedTropical{3, Float64}([2.0, 3.0, 5.0])
+
+julia> one(x)
+ExtendedTropical{3, Float64}([-Inf, -Inf, 0.0])
+
+julia> zero(x)
+ExtendedTropical{3, Float64}([-Inf, -Inf, -Inf])
+```
+"""
+struct ExtendedTropical{K,TO} <: Number
+    orders::Vector{TO}
+end
+function ExtendedTropical{K}(x::Vector{T}) where {T, K}
+    @assert length(x) == K
+    @assert issorted(x)
+    ExtendedTropical{K,T}(x)
+end
+Base.:(==)(a::ExtendedTropical{K}, b::ExtendedTropical{K}) where K = all(i->a.orders[i] == b.orders[i], 1:K)
+
+function Base.:*(a::ExtendedTropical{K,TO}, b::ExtendedTropical{K,TO}) where {K,TO}
+    res = Vector{TO}(undef, K)
+    return ExtendedTropical{K,TO}(sorted_sum_combination!(res, a.orders, b.orders))
+end
+
+function sorted_sum_combination!(res::AbstractVector{TO}, A::AbstractVector{TO}, B::AbstractVector{TO}) where TO
+    K = length(res)
+    @assert length(B) == length(A) == K
+    maxval = A[K] + B[K]
+    ptr = K
+    res[ptr] = maxval
+    queue = [(K,K-1,A[K]+B[K-1]), (K-1,K,A[K-1]+B[K])]
+    for k = 1:K-1
+        (i, j, res[K-k]) = _pop_max_sum!(queue)   # TODO: do not enumerate, use better data structures
+        _push_if_not_exists!(queue, i, j-1, A, B)
+        _push_if_not_exists!(queue, i-1, j, A, B)
+    end
+    return res
+end
+
+function _push_if_not_exists!(queue, i, j, A, B)
+    @inbounds if j>=1 && i>=1 && !any(x->x[1] >= i && x[2] >= j, queue)
+        push!(queue, (i, j, A[i] + B[j]))
+    end
+end
+
+function _pop_max_sum!(queue)
+    maxsum = first(queue)[3]
+    maxloc = 1
+    @inbounds for i=2:length(queue)
+        m = queue[i][3]
+        if m > maxsum
+            maxsum = m
+            maxloc = i
+        end
+    end
+    @inbounds data = queue[maxloc]
+    deleteat!(queue, maxloc)
+    return data
+end
+
+function Base.:+(a::ExtendedTropical{K,TO}, b::ExtendedTropical{K,TO}) where {K,TO}
+    res = Vector{TO}(undef, K)
+    ptr1, ptr2 = K, K
+    @inbounds va, vb = a.orders[ptr1], b.orders[ptr2]
+    @inbounds for i=K:-1:1
+        if va > vb
+            res[i] = va
+            if ptr1 != 1
+                ptr1 -= 1
+                va = a.orders[ptr1]
+            end
+        else
+            res[i] = vb
+            if ptr2 != 1
+                ptr2 -= 1
+                vb = b.orders[ptr2]
+            end
+        end
+    end
+    return ExtendedTropical{K,TO}(res)
+end
+
+Base.:^(a::ExtendedTropical, b::Integer) = Base.invoke(^, Tuple{ExtendedTropical, Real}, a, b)
+function Base.:^(a::ExtendedTropical{K,TO}, b::Real) where {K,TO}
+    if iszero(b)  # to avoid NaN
+        return one(ExtendedTropical{K,promote_type(TO,typeof(b))})
+    else
+        return ExtendedTropical{K,TO}(a.orders .* b)
+    end
+end
+
+Base.zero(::Type{ExtendedTropical{K,TO}}) where {K,TO} = ExtendedTropical{K,TO}(fill(zero(Tropical{TO}).n, K))
+Base.one(::Type{ExtendedTropical{K,TO}}) where {K,TO} = ExtendedTropical{K,TO}(map(i->i==K ? one(Tropical{TO}).n : zero(Tropical{TO}).n, 1:K))
+Base.zero(::ExtendedTropical{K,TO}) where {K,TO} = zero(ExtendedTropical{K,TO})
+Base.one(::ExtendedTropical{K,TO}) where {K,TO} = one(ExtendedTropical{K,TO})
+
 ############################ SET Numbers ##########################
 abstract type AbstractSetNumber end
 
@@ -501,7 +622,7 @@ function Base.copy(c::TreeConfigEnumerator{N,S,C}) where {N,S,C}
 end
 
 # Handle boolean, this is a patch for CUDA matmul
-for TYPE in [:AbstractSetNumber, :TruncatedPoly]
+for TYPE in [:AbstractSetNumber, :TruncatedPoly, :ExtendedTropical]
     @eval Base.:*(a::Bool, y::$TYPE) = a ? y : zero(y)
     @eval Base.:*(y::$TYPE, a::Bool) = a ? y : zero(y)
 end
@@ -550,9 +671,22 @@ function _x(::Type{Tropical{TV}}; invert) where {TV}
     ret = Tropical{TV}(one(TV))
     invert ? pre_invert_exponent(ret) : ret
 end
+function _x(::Type{ExtendedTropical{K,TO}}; invert) where {K,TO}
+    ret =ExtendedTropical{K,TO}(map(i->i==K ? one(TO) : zero(Tropical{TO}).n, 1:K))
+    invert ? pre_invert_exponent(ret) : ret
+end
 
 # for finding all solutions
 function _x(::Type{T}; invert) where {T<:AbstractSetNumber}
     ret = one(T)
     invert ? pre_invert_exponent(ret) : ret
 end
+
+# negate the exponents before entering the solver
+pre_invert_exponent(t::TruncatedPoly{K}) where K = TruncatedPoly(t.coeffs, -t.maxorder)
+pre_invert_exponent(t::TropicalNumbers.TropicalTypes) = inv(t)
+pre_invert_exponent(t::ExtendedTropical{K}) where K = ExtendedTropical{K}(map(i->i==K ? -t.orders[i] : t.orders[i], 1:K))
+# negate the exponents after entering the solver
+post_invert_exponent(t::TruncatedPoly{K}) where K = TruncatedPoly(ntuple(i->t.coeffs[K-i+1], K), -t.maxorder+(K-1))
+post_invert_exponent(t::TropicalNumbers.TropicalTypes) = inv(t)
+post_invert_exponent(t::ExtendedTropical{K}) where K = ExtendedTropical{K}(map(i->-t.orders[i], K:-1:1))

src/interfaces.jl

@@ -1,28 +1,33 @@
 abstract type AbstractProperty end
 
 """
-    SizeMax <: AbstractProperty
-    SizeMax()
+    SizeMax{K} <: AbstractProperty
+    SizeMax(k::Int)
 
-The maximum set size. e.g. the largest size of the [`IndependentSet`](@ref)  problem is also know as the independence number.
+The maximum-K set sizes. e.g. the largest size of the [`IndependentSet`](@ref)  problem is also know as the independence number.
 
-* The corresponding tensor element type is max-plus tropical number [`Tropical`](@ref).
+* The corresponding tensor element type are max-plus tropical number [`Tropical`](@ref) for `K == 1` and [`ExtendedTropical`](@ref) for `K > 1`.
 * It is compatible with weighted graph problems.
-* BLAS (on CPU) and GPU are supported,
+* BLAS (on CPU) and GPU are supported only for `K == 1`,
 """
-struct SizeMax <: AbstractProperty end
+struct SizeMax{K} <: AbstractProperty end
+SizeMax(k::Int=1) = SizeMax{k}()
+max_k(::SizeMax{K}) where K = K
 
 """
-    SizeMin <: AbstractProperty
-    SizeMin()
+    SizeMin{K} <: AbstractProperty
+    SizeMin(k::Int)
 
-The maximum set size. e.g. the smallest size ofthe [`MaximalIS`](@ref) problem is also known as the independent domination number.
+The minimum-K set sizes. e.g. the smallest size ofthe [`MaximalIS`](@ref) problem is also known as the independent domination number.
 
-* The corresponding tensor element type inverted max-plus tropical number [`Tropical`](@ref), which is equivalent to the min-plus tropical number.
+* The corresponding tensor element type are inverted max-plus tropical number [`Tropical`](@ref) for `K == 1` and inverted [`ExtendedTropical`](@ref) for `K > 1`.
+The inverted Tropical number emulates the min-plus tropical number.
 * It is compatible with weighted graph problems.
-* BLAS (on CPU) and GPU are supported,
+* BLAS (on CPU) and GPU are supported only for `K == 1`,
 """
-struct SizeMin <: AbstractProperty end
+struct SizeMin{K} <: AbstractProperty end
+SizeMin(k::Int=1) = SizeMin{k}()
+max_k(::SizeMin{K}) where K = K
 
 """
     CountingAll <: AbstractProperty
@@ -226,10 +231,14 @@ function solve(gp::GraphProblem, property::AbstractProperty; T=Float64, usecuda=
     if !_solvable(gp, property)
         throw(ArgumentError("Graph property `$(typeof(property))` is not computable for graph problem of type `$(typeof(gp))`."))
     end
-    if property isa SizeMax
+    if property isa SizeMax{1}
         return contractx(gp, _x(Tropical{T}; invert=false); usecuda=usecuda)
-    elseif property isa SizeMin
+    elseif property isa SizeMin{1}
         return post_invert_exponent.(contractx(gp, _x(Tropical{T}; invert=true); usecuda=usecuda))
+    elseif property isa SizeMax
+        return contractx(gp, _x(ExtendedTropical{max_k(property), T}; invert=false); usecuda=usecuda)
+    elseif property isa SizeMin
+        return post_invert_exponent.(contractx(gp, _x(ExtendedTropical{max_k(property), T}; invert=true); usecuda=usecuda))
     elseif property isa CountingAll
         return contractx(gp, one(T); usecuda=usecuda)
     elseif property isa CountingMax{1}
@@ -276,13 +285,6 @@ end
 # raise an error if the property for problem can not be computed
 _solvable(::Any, ::Any) = true
 
-# negate the exponents before entering the solver
-pre_invert_exponent(t::TruncatedPoly{K}) where K = TruncatedPoly(t.coeffs, -t.maxorder)
-pre_invert_exponent(t::TropicalNumbers.TropicalTypes) = inv(t)
-# negate the exponents after entering the solver
-post_invert_exponent(t::TruncatedPoly{K}) where K = TruncatedPoly(ntuple(i->t.coeffs[K-i+1], K), -t.maxorder+(K-1))
-post_invert_exponent(t::TropicalNumbers.TropicalTypes) = inv(t)
-
 """
     max_size(problem; usecuda=false)
 
@@ -403,6 +405,9 @@ function estimate_memory(problem::GraphProblem, ::GraphPolynomial{:polynomial};
     # this is the upper bound
     return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * length(labels(problem)))
 end
+function estimate_memory(problem::GraphProblem, ::Union{SizeMax{K},SizeMin{K}}; T=Float64) where K
+    return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * K)
+end
 
 function _size_dict(problem)
     lbs = labels(problem)
@@ -417,7 +422,9 @@ function _estimate_memory(::Type{ET}, problem::GraphProblem) where ET
     return peak_memory(problem.code, _size_dict(problem)) * sizeof(ET)
 end
 
-for (PROP, ET) in [(:SizeMax, :(Tropical{T})), (:SizeMin, :(Tropical{T})),
+for (PROP, ET) in [
+        (:(SizeMax{1}), :(Tropical{T})), (:(SizeMin{1}), :(Tropical{T})),
+        (:(SizeMax{K}), :(ExtendedTropical{K,T})), (:(SizeMin{K}), :(ExtendedTropical{K,T})),
         (:(SingleConfigMax{true}), :(Tropical{T})), (:(SingleConfigMin{true}), :(Tropical{T})),
         (:(CountingAll), :T), (:(CountingMax{1}), :(CountingTropical{T,T})), (:(CountingMin{1}), :(CountingTropical{T,T})),
         (:(CountingMax{K}), :(TruncatedPoly{K,T,T})), (:(CountingMin{K}), :(TruncatedPoly{K,T,T})),

test/arithematics.jl

@@ -34,6 +34,7 @@ end
                     (Max2Poly(1,2,3.0), Max2Poly(3,2,2.0), Max2Poly(4,7,1.0)),
                     (TruncatedPoly((1,2,3),3.0), TruncatedPoly((7,3,2),2.0), TruncatedPoly((1,4,7),1.0)),
                     (TropicalF64(5), TropicalF64(3), TropicalF64(-9)),
+                    (ExtendedTropical{2}([2.2,3.1]), ExtendedTropical{2}([-1.0, 4.0]), ExtendedTropical{2}([-Inf, 0.6])),
                     (CountingTropicalF64(5, 3), CountingTropicalF64(3, 9), CountingTropicalF64(-3, 2)),
                     (CountingTropical(5.0, ConfigSampler(StaticBitVector(rand(Bool, 10)))), CountingTropical(3.0, ConfigSampler(StaticBitVector(rand(Bool, 10)))), CountingTropical(-3.0, ConfigSampler(StaticBitVector(rand(Bool, 10))))),
                     (CountingTropical(5.0, ConfigEnumerator([StaticBitVector(rand(Bool, 10)) for j=1:3])), CountingTropical(3.0, ConfigEnumerator([StaticBitVector(rand(Bool, 10)) for j=1:4])), CountingTropical(-3.0, ConfigEnumerator([StaticBitVector(rand(Bool, 10)) for j=1:5]))),
@@ -97,6 +98,14 @@ end
     x = ConfigSampler(bv"00111")
     @test x ^ 0 == one(x)
     @test x ^ 2.0 == x
+
+    x = ExtendedTropical{3}([1.0, 2.0, 3.0])
+    @test x ^ 1 == x
+    @test x ^ 0 == one(x)
+    @test x ^ 1.0 == x
+    @test x ^ 0.0 == one(x)
+    @test x ^ 2 == ExtendedTropical{3}([2.0, 4.0, 6.0])
+    @test x ^ 2.0 == ExtendedTropical{3}([2.0, 4.0, 6.0])
 end
 
 @testset "push coverage" begin
@@ -113,4 +122,32 @@ end
     @test copy(y) !== y
     @test !iszero(y)
     println((x * x) * zero(x))
+end
+
+@testset "Truncated Tropical" begin
+    # +
+    a = ExtendedTropical{3}([1,2,3])
+    b = ExtendedTropical{3}([4,5,6])
+    c = ExtendedTropical{3}([0,1,2])
+    @test a + b == ExtendedTropical{3}([4,5,6])
+    @test b + a == ExtendedTropical{3}([4,5,6])
+    @test c + a == ExtendedTropical{3}([2,2,3])
+    @test a + c == ExtendedTropical{3}([2,2,3])
+
+    # *
+    function naive_mul(a, b)
+        K = length(a)
+        return sort!(vec([x+y for x in a, y in b]))[end-K+1:end]
+    end
+    d = ExtendedTropical{3}([0,1,20])
+    @test naive_mul(a.orders, b.orders) == (a * b).orders
+    @test naive_mul(b.orders, a.orders) == (b * a).orders
+    @test naive_mul(a.orders, d.orders) == (a * d).orders
+    @test naive_mul(d.orders, a.orders) == (d * a).orders
+    @test naive_mul(d.orders, d.orders) == (d * d).orders
+    for i=1:20
+        a = ExtendedTropical{100}(sort!(randn(100)))
+        b = ExtendedTropical{100}(sort!(randn(100)))
+        @test naive_mul(a.orders, b.orders) == (a * b).orders
+    end
 end

test/interfaces.jl

@@ -141,6 +141,11 @@ end
     res4b = solve(gp, ConfigsMax(1; bounded=true))[]
     @test res4 == res4b
     @test res4.c.data[] == res3.c.data
+
+    g = Graphs.smallgraph("petersen")
+    gp = IndependentSet(g; weights=fill(0.5, 10))
+    res5 = solve(gp, SizeMax(6))[]
+    @test res5.orders == [3.0,4,4,4,4,4] ./ 2
 end
 
 @testset "tree storage" begin
@@ -184,7 +189,7 @@ end
 @testset "memory estimation" begin
     gp = IndependentSet(smallgraph(:petersen))
     for property in [
-            SizeMax(), SizeMin(), CountingMax(), CountingMin(), CountingMax(2), CountingMin(2),
+            SizeMax(), SizeMin(), SizeMax(3), SizeMin(3), CountingMax(), CountingMin(), CountingMax(2), CountingMin(2),
             ConfigsMax(;bounded=true), ConfigsMin(;bounded=true), ConfigsMax(2;bounded=true), ConfigsMin(2;bounded=true), 
             ConfigsMax(;bounded=false), ConfigsMin(;bounded=false), ConfigsMax(2;bounded=false), ConfigsMin(2;bounded=false), SingleConfigMax(;bounded=false), SingleConfigMin(;bounded=false),
             CountingAll(), ConfigsAll(),

test/networks/MaximalIS.jl

@@ -23,9 +23,11 @@ end
 @testset "counting minimum maximal IS" begin
     g = smallgraph(:tutte)
     res = solve(MaximalIS(g), SizeMin())[]
+    res2 = solve(MaximalIS(g), SizeMin(10))[]
     poly = solve(MaximalIS(g), GraphPolynomial())[]
     @test poly == Polynomial([fill(0.0, 13)..., 2, 150, 7510, 71669, 66252, 14925, 571])
     @test res.n == 13
+    @test res2.orders == [13, 13, fill(14, 8)...]
     @test solve(MaximalIS(g), CountingMin())[].c == 2
     min2 = solve(MaximalIS(g), CountingMin(3))[]
     @test min2.maxorder == 15