Add summary printing to show(::IO, ::Problem) (#650)

Co-authored-by: Eric Hanson <5846501+ericphanson@users.noreply.github.com>

Author: odow

src/constraints/GenericConstraint.jl

@@ -17,6 +17,8 @@ function GenericConstraint{S}(child::AbstractExpr) where {S<:MOI.AbstractSet}
     return GenericConstraint(child, set_with_size(S, size(child)))
 end
 
+iscomplex(c::GenericConstraint) = iscomplex(c.child)
+
 function set_with_size(
     ::Type{S},
     sz::Tuple{Int,Int},

src/constraints/GeometricMeanHypoCone.jl

@@ -115,6 +115,12 @@ mutable struct GeometricMeanHypoConeConstraint <: Constraint
     end
 end
 
+function iscomplex(c::GeometricMeanHypoConeConstraint)
+    return iscomplex(c.T) || iscomplex(c.cone)
+end
+
+iscomplex(c::GeometricMeanHypoCone) = iscomplex(c.A) || iscomplex(c.B)
+
 function head(io::IO, ::GeometricMeanHypoConeConstraint)
     return print(io, "∈(GeometricMeanHypoCone)")
 end
@@ -128,7 +134,7 @@ function AbstractTrees.children(constraint::GeometricMeanHypoConeConstraint)
         constraint.T,
         constraint.cone.A,
         constraint.cone.B,
-        "t=$(constraint.cone.t)",
+        constraint.cone.t,
     )
 end
 

src/constraints/RelativeEntropyEpiCone.jl

@@ -106,6 +106,12 @@ mutable struct RelativeEntropyEpiConeConstraint <: Constraint
     end
 end
 
+function iscomplex(c::RelativeEntropyEpiConeConstraint)
+    return iscomplex(c.τ) || iscomplex(c.cone)
+end
+
+iscomplex(c::RelativeEntropyEpiCone) = iscomplex(c.X) || iscomplex(c.Y)
+
 function head(io::IO, ::RelativeEntropyEpiConeConstraint)
     return print(io, "∈(RelativeEntropyEpiCone)")
 end

src/problems.jl

@@ -41,9 +41,19 @@ function Base.getproperty(p::Problem, s::Symbol)
     return getfield(p, s)
 end
 
-dual_status(p::Problem) = MOI.get(p.model, MOI.DualStatus())
+function dual_status(p::Problem)
+    if p.model === nothing
+        return MOI.NO_SOLUTION
+    end
+    return MOI.get(p.model, MOI.DualStatus())
+end
 
-primal_status(p::Problem) = MOI.get(p.model, MOI.PrimalStatus())
+function primal_status(p::Problem)
+    if p.model === nothing
+        return MOI.NO_SOLUTION
+    end
+    return MOI.get(p.model, MOI.PrimalStatus())
+end
 
 termination_status(p::Problem) = MOI.get(p.model, MOI.TerminationStatus())
 

src/utilities/show.jl

@@ -197,7 +197,7 @@ end
 AbstractTrees.children(p::ProblemObjectiveRoot) = (p.objective,)
 
 function AbstractTrees.printnode(io::IO, p::ProblemObjectiveRoot)
-    return print(io, string(p.head))
+    return print(io, "  ", string(p.head))
 end
 
 struct ProblemConstraintsRoot
@@ -207,7 +207,7 @@ end
 AbstractTrees.children(p::ProblemConstraintsRoot) = p.constraints
 
 function AbstractTrees.printnode(io::IO, p::ProblemConstraintsRoot)
-    return print(io, "subject to")
+    return print(io, "  subject to")
 end
 
 function TreePrint.print_tree(io::IO, p::Problem, args...; kwargs...)
@@ -234,17 +234,40 @@ function TreePrint.print_tree(io::IO, p::Problem, args...; kwargs...)
     return
 end
 
+function _to_underscore(n::Integer)
+    x = Iterators.partition(digits(n), 3)
+    return join(reverse(join.(reverse.(x))), '_')
+end
+
+function _str_with_elements(x, y)
+    xs, ys = _to_underscore(x), _to_underscore(y)
+    return "$xs ($ys scalar elements)"
+end
+
 function Base.show(io::IO, p::Problem)
-    TreePrint.print_tree(io, p, MAXDEPTH[], MAXWIDTH[])
-    if p.status == MOI.OPTIMIZE_NOT_CALLED
-        print(io, "\nstatus: `solve!` not called yet")
-    else
-        print(io, "\ntermination status: $(p.status)")
-        print(io, "\nprimal status: $(primal_status(p))")
-        print(io, "\ndual status: $(dual_status(p))")
-    end
-    if p.status == "solved"
-        print(io, " with optimal value of $(round(p.optval, digits=4))")
+    # Print problem statistics
+    counts = Counts(p)
+    println(io, "Problem statistics")
+    var_str = _str_with_elements(counts.n_variables, counts.n_scalar_variables)
+    println(io, "  number of variables    : ", var_str)
+    con_str =
+        _str_with_elements(counts.n_constraints, counts.n_scalar_constraints)
+    println(io, "  number of constraints  : ", con_str)
+    println(io, "  number of coefficients : ", counts.n_nonzeros)
+    println(io, "  number of atoms        : ", counts.n_atoms)
+    println(io)
+    # Print solution summary
+    println(io, "Solution summary")
+    println(io, "  termination status : ", p.status)
+    println(io, "  primal status      : ", primal_status(p))
+    println(io, "  dual status        : ", dual_status(p))
+    if p.status == MOI.OPTIMAL && p.optval !== nothing
+        println(io, "  objective value    : ", round(p.optval; digits = 4))
     end
+    println(io)
+    # Expression tree
+    println(io, "Expression graph")
+    # We offset the depth by 1 so things are printed with a three-space offset.
+    TreePrint.print_tree(io, p, MAXDEPTH[] + 1, MAXWIDTH[]; depth = 1)
     return
 end

src/utilities/tree_interface.jl

@@ -38,3 +38,97 @@ function AbstractTrees.printnode(io::IO, node::Constant)
 end
 
 AbstractTrees.printnode(io::IO, node::AbstractVariable) = summary(io, node)
+
+mutable struct Counts
+    seen::Base.IdSet
+    n_variables::Int
+    n_scalar_variables::Int
+    n_constraints::Int
+    n_scalar_constraints::Int
+    n_atoms::Int
+    n_nonzeros::Int
+    Counts() = new(Base.IdSet(), 0, 0, 0, 0, 0, 0)
+end
+
+function Counts(p::Problem)
+    counts = Counts()
+    _add_to_problem_count(counts, p)
+    counts.n_atoms -= 1  # Don't count p as an atom
+    return counts
+end
+
+function _add_to_problem_count(counts::Counts, node::AbstractExpr)
+    if node in counts.seen
+        return
+    end
+    for n in AbstractTrees.PreOrderDFS(node)
+        if n !== node && !(n in counts.seen)
+            _add_to_problem_count(counts, n)
+        end
+    end
+    counts.n_atoms += 1
+    push!(counts.seen, node)
+    return
+end
+
+function _add_to_problem_count(counts::Counts, node::AbstractVariable)
+    if node in counts.seen
+        return
+    end
+    counts.n_variables += 1
+    counts.n_scalar_variables += (iscomplex(node) ? 2 : 1) * length(node)
+    push!(counts.seen, node)
+    for c in get_constraints(node)
+        _add_to_problem_count(counts, c)
+    end
+    return
+end
+
+_add_to_problem_count(::Counts, ::Vector{Constraint}) = nothing
+
+_add_to_problem_count(::Counts, ::Nothing) = nothing
+
+function _add_to_problem_count(counts::Counts, node::GenericConstraint)
+    counts.n_constraints += 1
+    counts.n_scalar_constraints +=
+        (iscomplex(node) ? 2 : 1) * MOI.dimension(node.set)
+    _add_to_problem_count(counts, node.child)
+    return
+end
+
+function _add_to_problem_count(counts::Counts, node::Constraint)
+    counts.n_constraints += 1
+    # TODO(odow): we don't know the general case
+    return
+end
+
+function _add_to_problem_count(counts::Counts, node::Constant)
+    _add_to_problem_count(counts, node.value)
+    push!(counts.seen, node)
+    return
+end
+
+function _add_to_problem_count(counts::Counts, node::ComplexConstant)
+    _add_to_problem_count(counts, node.real_constant)
+    _add_to_problem_count(counts, node.imag_constant)
+    push!(counts.seen, node)
+    return
+end
+
+function _add_to_problem_count(counts::Counts, node::Number)
+    counts.n_nonzeros += iscomplex(node) ? 2 : 1
+    return
+end
+
+function _add_to_problem_count(counts::Counts, node::AbstractArray{<:Number})
+    counts.n_nonzeros += length(node)
+    return
+end
+
+function _add_to_problem_count(
+    counts::Counts,
+    node::SparseArrays.SparseMatrixCSC,
+)
+    counts.n_nonzeros += SparseArrays.nnz(node)
+    return
+end

test/test_constraints.jl

@@ -348,6 +348,28 @@ function test_GeometricMeanEpiConeSquare()
     return
 end
 
+### constraints/RelativeEntropyEpiConeConstraint
+
+function test_constraints_RelativeEntropyEpiConeConstraint()
+    t = Variable(2, 2)
+    X = Variable(2, 2)
+    Y = Variable(2, 2)
+    c = t in RelativeEntropyEpiCone(X, Y)
+    @test !Convex.iscomplex(c)
+    return
+end
+
+### constraints/GeometricMeanHypoConeConstraint
+
+function test_constraints_GeometricMeanHypoConeConstraint()
+    T = Variable(2, 2)
+    A = Variable(2, 2)
+    B = Variable(2, 2)
+    c = T in GeometricMeanHypoCone(A, B, 1 // 2)
+    @test !Convex.iscomplex(c)
+    return
+end
+
 end  # TestConstraints
 
 TestConstraints.runtests()

test/test_problem_depot.jl

@@ -51,12 +51,15 @@ function test_Clarabel_warmstarts()
     Convex.ProblemDepot.run_tests(;
         exclude = [r"mip", r"sdp_quantum_relative_entropy3_lowrank"],
     ) do p
-        return solve!(
+        ret = solve!(
             p,
             Clarabel.Optimizer;
             silent_solver = true,
             warmstart = true,
         )
+        # verify post-solve printing does not error
+        @test sprint(show, p) isa AbstractString
+        return ret
     end
     return
 end