implement variable squashing

Project.toml

@@ -1,7 +1,7 @@
 name = "ConstraintTrees"
 uuid = "5515826b-29c3-47a5-8849-8513ac836620"
 authors = ["The developers of ConstraintTrees.jl"]
-version = "1.2.0"
+version = "1.3.0"
 
 [deps]
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"

docs/src/4-functional-tree-processing.jl

@@ -42,9 +42,10 @@
 # - [`variables_for`](@ref ConstraintTrees.variables_for) allocates a variable
 #   for each constraint in the tree and allows the user to specify bounds
 #
-# Additionally, all these have their "indexed" variant which allows you to know
-# the path where the tree elements are being merged. The path is passed to the
-# handling function as a tuple of symbols. The variants are prefixed with `i`:
+# Additionally, all these have their "indexed" variant which allows the user to
+# know the path where the tree elements are being merged. The path is passed to
+# the handling function as a tuple of symbols. The variants are prefixed with
+# `i`:
 #
 # - [`imap`](@ref ConstraintTrees.imap)
 # - [`imapreduce`](@ref ConstraintTrees.ireduce) (here the path refers to the
@@ -202,7 +203,8 @@ tz.y
 @test isapprox(tz.y.x, 0.99) #src
 @test isapprox(tz.y.y, 0.78) #src
 
-# We also have the indexed variants; for example this allows us to only merge the `x` elements in points:
+# We also have the indexed variants; for example this allows us to only merge
+# the `x` elements in points:
 
 tx = C.imerge(t1, t2, Float64) do path, x, y
     last(path) == :x || return missing
@@ -313,3 +315,73 @@ end;
 
 # To prevent uncertainty, both functions always traverse the keys in sorted
 # order.
+
+# ## Removing constraints with `filter`
+#
+# In many cases it is beneficial to simplify the constraint system by
+# systematically removing constraints. [`filter`](@ref ConstraintTrees.filter)
+# and [`ifilter`](@ref ConstraintTrees.ifilter) run a function on all subtrees
+# and leaves (usually the leaves are [`Constraint`](@ref
+# ConstraintTrees.Constraint)s), and only retain these where the function
+# returns `true`.
+#
+# For example, this removes all constraints named `y`:
+
+filtered = C.ifilter(x) do ix, c
+    return c isa C.ConstraintTree || last(ix) != :y
+end
+
+filtered.z
+
+@test !haskey(filtered.x, :y) #src
+
+# Functions [`filter_leaves`](@ref ConstraintTrees.filter_leaves) and
+# [`ifilter_leaves`](@ref ConstraintTrees.ifilter_leaves) act similarly but
+# automatically assume that the directory structure is going to stay intact,
+# freeing the user from having to handle the subdirectories.
+#
+# The above example thus simplifies to:
+
+filtered = C.ifilter_leaves(x) do ix, c
+    last(ix) != :y
+end
+
+filtered.z
+
+@test !haskey(filtered.x, :y) #src
+
+# We can also remove whole variable ranges:
+
+filtered = C.filter_leaves(x) do c
+    all(>=(4), c.value.idxs)
+end
+
+# Notably, these operations leave the tree with a slightly sub-optimal state,
+# as there are indexes allocated for variables that are no longer used!
+
+C.var_count(filtered)
+
+# To fix the issue, it is possible to "squash" the variable indexes using
+# [`prune_variables`](@ref ConstraintTrees.prune_variables):
+
+pruned = C.prune_variables(filtered)
+
+C.var_count(pruned)
+
+@test C.var_count(pruned) == 3 #src
+
+# Note that given the renumbering, constraint trees are no longer compatible
+# after pruning, and should not be combined with `*`. As an anti-example, one
+# might be interested in pruning the variable values before joining them in to
+# larger constraint tree, e.g. to simplify larger quadratic values:
+
+pruned_qv = C.prune_variables(x.y.x.value * x.z.y.value)
+
+# Despite the variable count decreased, the value now corresponds to a
+# completely different value in the original tree! Compare:
+
+(pruned_qv, x.x.x.value * x.x.y.value)
+
+@test C.var_count(pruned_qv) == 2 #src
+@test pruned_qv.idxs == (x.x.x.value * x.x.y.value).idxs #src
+@test pruned_qv.weights == (x.x.x.value * x.x.y.value).weights #src

src/constraint_tree.jl

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import DataStructures: SortedDict
+import DataStructures: SortedDict, SortedSet
 
 """
 $(TYPEDEF)
@@ -98,7 +98,7 @@ $(TYPEDSIGNATURES)
 
 Find the expected count of variables in a [`ConstraintTree`](@ref).
 """
-var_count(x::ConstraintTree) = isempty(elems(x)) ? 0 : maximum(var_count.(values(elems(x))))
+var_count(x::ConstraintTree) = isempty(elems(x)) ? 0 : maximum(var_count.(values(x)))
 
 """
 $(TYPEDSIGNATURES)
@@ -159,6 +159,65 @@ incr_var_idxs(x::QuadraticValue, incr::Int) = QuadraticValue(
     weights = x.weights,
 )
 
+"""
+$(TYPEDSIGNATURES)
+
+Push all variable indexes found in `x` to the `out` container.
+
+(The container needs to support the standard `push!`.)
+"""
+collect_variables!(x::Constraint, out) = collect_variables!(x.value, out)
+collect_variables!(x::LinearValue, out) =
+    for idx in x.idxs
+        push!(out, idx)
+    end
+collect_variables!(x::QuadraticValue, out) =
+    for (idx, idy) in x.idxs
+        push!(out, idx, idy)
+    end
+collect_variables!(x::ConstraintTree, out) = collect_variables!.(values(x), Ref(out))
+
+"""
+$(TYPEDSIGNATURES)
+
+Prune the unused variable indexes from an object `x` (such as a
+[`ConstraintTree`](@ref)).
+
+This first runs [`collect_variables!`](@ref) to determine the actual used
+variables, then calls [`renumber_variables`](@ref) to create a renumbered
+object.
+"""
+function prune_variables(x)
+    vars = SortedSet{Int}()
+    collect_variables!(x, vars)
+    push!(vars, 0)
+    vv = collect(vars)
+    @assert vv[1] == 0 "variable indexes are broken"
+    return renumber_variables(x, SortedDict(vv .=> 0:length(vv)-1))
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Renumber all variables in an object (such as [`ConstraintTree`](@ref)). The new
+variable indexes are taken from the `mapping` parameter at the index of the old
+variable's index.
+
+This does not run any consistency checks on the result; the `mapping` must
+therefore be monotonically increasing, and the zero index must map to itself,
+otherwise invalid [`Value`](@ref)s will be produced.
+"""
+renumber_variables(x::ConstraintTree, mapping) =
+    ConstraintTree(k => renumber_variables(v, mapping) for (k, v) in x)
+renumber_variables(x::Constraint, mapping) =
+    Constraint(renumber_variables(x.value, mapping), x.bound)
+renumber_variables(x::LinearValue, mapping) =
+    LinearValue(idxs = [mapping[idx] for idx in x.idxs], weights = x.weights)
+renumber_variables(x::QuadraticValue, mapping) = QuadraticValue(
+    idxs = [(mapping[idx], mapping[idy]) for (idx, idy) in x.idxs],
+    weights = x.weights,
+)
+
 #
 # Algebraic construction
 #

src/tree.jl

@@ -156,6 +156,61 @@ end
 """
 $(TYPEDSIGNATURES)
 
+Filter all branches and leaves in a tree, leaving only the ones where `f`
+returns `true`.
+
+Note that the branches are passed to `f` as well. Use [`filter_leaves`](@ref)
+to only work with the leaf values.
+"""
+function filter(f, x::Tree{T}) where {T}
+    go(x::Tree) = Tree{T}(k => go(v) for (k, v) in x if f(v))
+    go(x) = x
+
+    go(x)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Like [`filter`](@ref) but the filtering predicate function also receives the
+"path" in the tree.
+"""
+function ifilter(f, x::Tree{T}) where {T}
+    go(ix, x::Tree{T}) =
+        Tree{T}(k => go(tuple(ix..., k), v) for (k, v) in x if f(tuple(ix..., k), v))
+    go(ix, x) = x
+
+    go((), x)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Like [`filter`](@ref) but the filtering predicate function `f` only receives
+the leaf values (i.e., no intermediate sub-trees).
+
+In turn, the result will retain the whole subtree structure (even if empty).
+"""
+function filter_leaves(f, x::Tree{T}) where {T}
+    flt(x::Tree{T}) = true
+    flt(x) = f(x)
+
+    filter(flt, x)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Combination of [`ifilter`](@ref) and [`filter_leaves`](@ref).
+"""
+ifilter_leaves(f, x::Tree{T}) where {T} =
+    let flt(_, x::Tree{T}) = true, flt(i, x) = f(i, x)
+        ifilter(flt, x)
+    end
+
+"""
+$(TYPEDSIGNATURES)
+
 Like [`map`](@ref), but discards the results, thus relying only on the side
 effects of `f`.