configs for max-k configurations (#28)

* use tropical in extended tropical

* done

Author: GiggleLiu

examples/IndependentSet.jl

@@ -92,7 +92,8 @@ independence_polynomial = solve(problem, GraphPolynomial(; method=:finitefield))
 # ### Configuration properties
 # ##### finding one maximum independent set (MIS)
 # One can use [`SingleConfigMax`](@ref) to find one of the solution with largest set size, and it has two implementations.
-# The unbounded (default) version uses [`ConfigSampler`](@ref) to sample one of the best solutions directly.
+# The unbounded (default) version uses a joint type of [`CountingTropical`](@ref) and [`ConfigSampler`](@ref) in computation,
+# where `CountingTropical` finds the maximum size and `ConfigSampler` samples one of the best solutions.
 # The bounded version uses the binary gradient back-propagation (see our paper) to compute the gradients.
 # It requires caching intermediate states, but is often faster (on CPU) because it can use [`TropicalGEMM`](https://github.com/TensorBFS/TropicalGEMM.jl) (see [Performance Tips](@ref)).
 max_config = solve(problem, SingleConfigMax(; bounded=false))[]
@@ -172,6 +173,18 @@ show_graph(graph; locs=locations, vertex_colors=
 spectrum = solve(problem, SizeMax(10))
 
 # It uses the [`ExtendedTropical`](@ref) as the tensor elements.
+# One can get sets with maximum `K` sizes, by combining [`ExtendedTropical`](@ref) and the algebra in the previous section for sampling one configuration.
+max5_configs = solve(problem, SingleConfigMax(5))[]
+
+imgs_max5 = ntuple(k->show_graph(graph;
+                    locs=locations, scale=0.25,
+                    vertex_colors=[iszero(max5_configs.orders[k].c.data[i]) ? "white" : "red"
+                    for i=1:nv(graph)]), 5);
+
+Compose.set_default_graphic_size(18cm, 4cm)
+
+Compose.compose(context(),
+     ntuple(k->(context((k-1)/5, 0.0, 1.2/5, 1.0), imgs_max5[k]), 5)...)
 
 # ## Open vertices and MIS tensor analysis
 # The following code computes the MIS tropical tensor (reference to be added) with open vertices 1, 2 and 3.

src/arithematics.jl

@@ -175,7 +175,7 @@ end
 
 Extended Tropical numbers with largest `K` orders keeped,
 or the [`TruncatedPoly`](@ref) without coefficients,
-`TO` is the element type of orders.
+`TO` is the element type of orders, usually [`Tropical`](@ref) numbers.
 This algebra maps
 
 * `+` to finding largest `K` values of union of two sets.
@@ -186,23 +186,23 @@ This algebra maps
 Example
 ------------------------------
 ```jldoctest; setup=(using GraphTensorNetworks)
-julia> x = ExtendedTropical{3}([1.0, 2, 3])
-ExtendedTropical{3, Float64}([1.0, 2.0, 3.0])
+julia> x = ExtendedTropical{3}(Tropical.([1.0, 2, 3]))
+ExtendedTropical{3, TropicalF64}(TropicalF64[1.0ₜ, 2.0ₜ, 3.0ₜ])
 
-julia> y = ExtendedTropical{3}([-Inf, 2, 5])
-ExtendedTropical{3, Float64}([-Inf, 2.0, 5.0])
+julia> y = ExtendedTropical{3}(Tropical.([-Inf, 2, 5]))
+ExtendedTropical{3, TropicalF64}(TropicalF64[-Infₜ, 2.0ₜ, 5.0ₜ])
 
 julia> x * y
-ExtendedTropical{3, Float64}([6.0, 7.0, 8.0])
+ExtendedTropical{3, TropicalF64}(TropicalF64[6.0ₜ, 7.0ₜ, 8.0ₜ])
 
 julia> x + y
-ExtendedTropical{3, Float64}([2.0, 3.0, 5.0])
+ExtendedTropical{3, TropicalF64}(TropicalF64[2.0ₜ, 3.0ₜ, 5.0ₜ])
 
 julia> one(x)
-ExtendedTropical{3, Float64}([-Inf, -Inf, 0.0])
+ExtendedTropical{3, TropicalF64}(TropicalF64[-Infₜ, -Infₜ, 0.0ₜ])
 
 julia> zero(x)
-ExtendedTropical{3, Float64}([-Inf, -Inf, -Inf])
+ExtendedTropical{3, TropicalF64}(TropicalF64[-Infₜ, -Infₜ, -Infₜ])
 ```
 """
 struct ExtendedTropical{K,TO} <: Number
@@ -223,10 +223,10 @@ 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]
+    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])]
+    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)
@@ -237,7 +237,7 @@ 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]))
+        push!(queue, (i, j, A[i]*B[j]))
     end
 end
 
@@ -281,14 +281,14 @@ 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))})
+        return one(ExtendedTropical{K,TO})
     else
-        return ExtendedTropical{K,TO}(a.orders .* b)
+        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(::Type{ExtendedTropical{K,TO}}) where {K,TO} = ExtendedTropical{K,TO}(fill(zero(TO), K))
+Base.one(::Type{ExtendedTropical{K,TO}}) where {K,TO} = ExtendedTropical{K,TO}(map(i->i==K ? one(TO) : zero(TO), 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})
 
@@ -703,6 +703,7 @@ for (F,TP) in [(:set_type, :ConfigEnumerator), (:sampler_type, :ConfigSampler),
         end
     end
 end
+sampler_type(::Type{ExtendedTropical{K,T}}, n::Int, nflavor::Int) where {K,T} = ExtendedTropical{K, sampler_type(T, n, nflavor)}
 
 # utilities for creating onehot vectors
 onehotv(::Type{ConfigEnumerator{N,S,C}}, i::Integer, v) where {N,S,C} = ConfigEnumerator([onehotv(StaticElementVector{N,S,C}, i, v)])
@@ -774,8 +775,7 @@ function _x(::Type{Tropical{TV}}; invert) where {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
+    return ExtendedTropical{K,TO}(map(i->i==K ? _x(TO; invert=invert) : zero(TO), 1:K))
 end
 
 # for finding all solutions
@@ -796,12 +796,14 @@ end
 function _onehotv(::Type{BS}, x, v) where {BS<:AbstractSetNumber}
     onehotv(BS, x, v)
 end
+function _onehotv(::Type{ExtendedTropical{K,TO}}, x, v) where {K,T,BS<:AbstractSetNumber,TO<:CountingTropical{T,BS}}
+    ExtendedTropical{K,TO}(map(i->i==K ? _onehotv(TO, x, v) : zero(TO), 1:K))
+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))
+post_invert_exponent(t::ExtendedTropical{K}) where K = ExtendedTropical{K}(map(i->inv(t.orders[i]), K:-1:1))

src/interfaces.jl

@@ -109,10 +109,10 @@ GraphPolynomial(; method::Symbol = :finitefield, kwargs...) = GraphPolynomial{me
 graph_polynomial_method(::GraphPolynomial{METHOD}) where METHOD = METHOD
 
 """
-    SingleConfigMax{BOUNDED} <: AbstractProperty
-    SingleConfigMax(; bounded=false)
+    SingleConfigMax{K, BOUNDED} <: AbstractProperty
+    SingleConfigMax(k::Int; bounded=false)
 
-Finding single best solution, e.g. for [`IndependentSet`](@ref) problem, it is one of the maximum independent sets.
+Finding single solution for largest-K sizes, e.g. for [`IndependentSet`](@ref) problem, it is one of the maximum independent sets.
 
 * The corresponding data type is [`CountingTropical{Float64,<:ConfigSampler}`](@ref) if `BOUNDED` is `false`, [`Tropical`](@ref) otherwise.
 * Weighted graph problems is supported.
@@ -122,14 +122,15 @@ Keyword Arguments
 ----------------------------
 * `bounded`, if it is true, use bounding trick (or boolean gradients) to reduce the working memory to store intermediate configurations.
 """
-struct SingleConfigMax{BOUNDED} <:AbstractProperty end
-SingleConfigMax(; bounded::Bool=false) = SingleConfigMax{bounded}()
+struct SingleConfigMax{K,BOUNDED} <:AbstractProperty end
+SingleConfigMax(k::Int=1; bounded::Bool=false) = SingleConfigMax{k, bounded}()
+max_k(::SingleConfigMax{K}) where K = K
 
 """
-    SingleConfigMin{BOUNDED} <: AbstractProperty
-    SingleConfigMin(; bounded=false)
+    SingleConfigMin{K, BOUNDED} <: AbstractProperty
+    SingleConfigMin(k::Int; bounded=false)
 
-Finding single "worst" solution.
+Finding single solution with smallest-K size.
 
 * The corresponding data type is inverted [`CountingTropical{Float64,<:ConfigSampler}`](@ref) if `BOUNDED` is `false`, inverted [`Tropical`](@ref) otherwise.
 * Weighted graph problems is supported.
@@ -139,8 +140,9 @@ Keyword Arguments
 ----------------------------
 * `bounded`, if it is true, use bounding trick (or boolean gradients) to reduce the working memory to store intermediate configurations.
 """
-struct SingleConfigMin{BOUNDED} <:AbstractProperty end
-SingleConfigMin(; bounded::Bool=false) = SingleConfigMin{bounded}()
+struct SingleConfigMin{K,BOUNDED} <:AbstractProperty end
+SingleConfigMin(k::Int=1; bounded::Bool=false) = SingleConfigMin{k,bounded}()
+min_k(::SingleConfigMin{K}) where K = K
 
 """
     ConfigsAll{TREESTORAGE} <:AbstractProperty
@@ -216,7 +218,7 @@ Positional Arguments
     * [`CountingAll`](@ref) for counting all configurations,
     * [`GraphPolynomial`](@ref) for evaluating the graph polynomial,
 
-    * [`SingleConfigMax`](@ref) for finding one maximum configuration,
+    * [`SingleConfigMax`](@ref) for finding one maximum configuration for each size,
     * [`ConfigsMax`](@ref) for enumerating configurations with largest-K sizes,
     * [`ConfigsMin`](@ref) for enumerating configurations with smallest-K sizes,
     * [`ConfigsAll`](@ref) for enumerating all configurations,
@@ -236,9 +238,9 @@ function solve(gp::GraphProblem, property::AbstractProperty; T=Float64, usecuda=
     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)
+        return contractx(gp, _x(ExtendedTropical{max_k(property), Tropical{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))
+        return post_invert_exponent.(contractx(gp, _x(ExtendedTropical{max_k(property), Tropical{T}}; invert=true); usecuda=usecuda))
     elseif property isa CountingAll
         return contractx(gp, one(T); usecuda=usecuda)
     elseif property isa CountingMax{1}
@@ -251,10 +253,14 @@ function solve(gp::GraphProblem, property::AbstractProperty; T=Float64, usecuda=
         return post_invert_exponent.(contractx(gp, pre_invert_exponent(TruncatedPoly(ntuple(i->i == min_k(property) ? one(T) : zero(T), min_k(property)), one(T))); usecuda=usecuda))
     elseif property isa GraphPolynomial
         return graph_polynomial(gp, Val(graph_polynomial_method(property)); usecuda=usecuda, T=T, property.kwargs...)
-    elseif property isa SingleConfigMax{false}
-        return solutions(gp, CountingTropical{T,T}; all=false, usecuda=usecuda, )
-    elseif property isa SingleConfigMin{false}
+    elseif property isa SingleConfigMax{1,false}
+        return solutions(gp, CountingTropical{T,T}; all=false, usecuda=usecuda)
+    elseif property isa (SingleConfigMax{K,false} where K)
+        return solutions(gp, ExtendedTropical{max_k(property),CountingTropical{T,T}}; all=false, usecuda=usecuda)
+    elseif property isa SingleConfigMin{1,false}
         return solutions(gp, CountingTropical{T,T}; all=false, usecuda=usecuda, invert=true)
+    elseif property isa (SingleConfigMin{K,false} where K)
+        return solutions(gp, ExtendedTropical{min_k(property),CountingTropical{T,T}}; all=false, usecuda=usecuda, invert=true)
     elseif property isa ConfigsMax{1,false}
         return solutions(gp, CountingTropical{T,T}; all=true, usecuda=usecuda, tree_storage=tree_storage(property))
     elseif property isa ConfigsMin{1,false}
@@ -265,10 +271,16 @@ function solve(gp::GraphProblem, property::AbstractProperty; T=Float64, usecuda=
         return solutions(gp, TruncatedPoly{min_k(property),T,T}; all=true, usecuda=usecuda, invert=true)
     elseif property isa ConfigsAll
         return solutions(gp, Real; all=true, usecuda=usecuda, tree_storage=tree_storage(property))
-    elseif property isa SingleConfigMax{true}
+    elseif property isa SingleConfigMax{1,true}
         return best_solutions(gp; all=false, usecuda=usecuda, T=T)
-    elseif property isa SingleConfigMin{true}
+    elseif property isa (SingleConfigMax{K,true} where K)
+        @warn "bounded `SingleConfigMax` property for `K != 1` is not implemented. Switching to the unbounded version."
+        return solve(gp, SingleConfigMax{max_k(property),false}(); T, usecuda)
+    elseif property isa SingleConfigMin{1,true}
         return best_solutions(gp; all=false, usecuda=usecuda, invert=true, T=T)
+    elseif property isa (SingleConfigMin{K,true} where K)
+        @warn "bounded `SingleConfigMin` property for `K != 1` is not implemented. Switching to the unbounded version."
+        return solve(gp, SingleConfigMin{min_k(property),false}(); T, usecuda)
     elseif property isa ConfigsMax{1,true}
         return best_solutions(gp; all=true, usecuda=usecuda, tree_storage=tree_storage(property), T=T)
     elseif property isa ConfigsMin{1,true}
@@ -396,10 +408,20 @@ Memory estimation in number of bytes to compute certain `property` of a `problem
 function estimate_memory(problem::GraphProblem, property::AbstractProperty; T=Float64)
     _estimate_memory(tensor_element_type(T, length(labels(problem)), nflavor(problem), property), problem)
 end
-function estimate_memory(problem::GraphProblem, ::Union{SingleConfigMax{true},SingleConfigMin{true}}; T=Float64)
+function estimate_memory(problem::GraphProblem, property::Union{SingleConfigMax{K,BOUNDED},SingleConfigMin{K,BOUNDED}}; T=Float64) where {K, BOUNDED}
     tc, sc, rw = timespacereadwrite_complexity(problem.code, _size_dict(problem))
     # caching all tensors is equivalent to counting the total number of writes
-    return ceil(Int, exp2(rw - 1)) * sizeof(Tropical{T})
+    if K == 1 && BOUNDED
+        return ceil(Int, exp2(rw - 1)) * sizeof(Tropical{T})
+    elseif K == 1 & !BOUNDED
+        n, nf = length(labels(problem)), nflavor(problem)
+        return peak_memory(problem.code, _size_dict(problem)) * (sizeof(tensor_element_type(T, n, nf, property)) * K)
+    else
+        # NOTE: the `K > 1` case does not respect bounding
+        n, nf = length(labels(problem)), nflavor(problem)
+        TT = tensor_element_type(T, n, nf, property)
+        return peak_memory(problem.code, _size_dict(problem)) * (sizeof(tensor_element_type(T, n, nf, SingleConfigMax{1,BOUNDED}())) * K + sizeof(TT))
+    end
 end
 function estimate_memory(problem::GraphProblem, ::GraphPolynomial{:polynomial}; T=Float64)
     # this is the upper bound
@@ -424,8 +446,7 @@ end
 
 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})),
+        (:(SizeMax{K}), :(ExtendedTropical{K,Tropical{T}})), (:(SizeMin{K}), :(ExtendedTropical{K,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})),
         (:(GraphPolynomial{:finitefield}), :(Mod{N,Int32} where N)), (:(GraphPolynomial{:fft}), :(Complex{T})), 
@@ -434,11 +455,17 @@ for (PROP, ET) in [
     @eval tensor_element_type(::Type{T}, n::Int, nflavor::Int, ::$PROP) where {T,K} = $ET
 end
 
-for (PROP, ET) in [(:(SingleConfigMax{false}), :(CountingTropical{T,T})), (:(SingleConfigMin{false}), :(CountingTropical{T,T}))]
-    @eval function tensor_element_type(::Type{T}, n::Int, nflavor::Int, ::$PROP) where {T}
-        sampler_type($ET, n, nflavor)
+function tensor_element_type(::Type{T}, n::Int, nflavor::Int, ::PROP) where {T, K, BOUNDED, PROP<:Union{SingleConfigMax{K,BOUNDED},SingleConfigMin{K,BOUNDED}}}
+    if K == 1 && BOUNDED
+        return Tropical{T}
+    elseif K == 1 && !BOUNDED
+        return sampler_type(CountingTropical{T,T}, n, nflavor)
+    else
+        # NOTE: the `K > 1` case does not respect bounding
+        return sampler_type(ExtendedTropical{K,CountingTropical{T,T}}, n, nflavor)
     end
 end
+
 for (PROP, ET) in [
         (:(ConfigsMax{1}), :(CountingTropical{T,T})), (:(ConfigsMin{1}), :(CountingTropical{T,T})),
         (:(ConfigsMax{K}), :(TruncatedPoly{K,T,T})), (:(ConfigsMin{K}), :(TruncatedPoly{K,T,T})),

src/networks/networks.jl

@@ -127,7 +127,7 @@ end
 add_labels!(tensors::AbstractVector{<:AbstractArray}, ixs, labels) = tensors
 
 const SetPolyNumbers{T} = Union{Polynomial{T}, TruncatedPoly{K,T} where K, CountingTropical{TV,T} where TV} where T<:AbstractSetNumber
-function add_labels!(tensors::AbstractVector{<:AbstractArray{T}}, ixs, labels) where T <: Union{AbstractSetNumber, SetPolyNumbers}
+function add_labels!(tensors::AbstractVector{<:AbstractArray{T}}, ixs, labels) where T <: Union{AbstractSetNumber, SetPolyNumbers, ExtendedTropical{K,T} where {K,T<:SetPolyNumbers}}
     for (t, ix) in zip(tensors, ixs)
         for (dim, l) in enumerate(ix)
             index = findfirst(==(l), labels)