[ANN] ConstraintTrees.jl -- tidy constrained systems for JuMP

Package Announcements
, ,
1

What is it about?

We’ve lately spent quite some time on developing constraint-based modeling software that works with JuMP (first COBREXA.jl, then FBCModelTests.jl, and lately DifferentialMetabolism.jl).

A common observation we made during the development of these packages is that constraint-based models typically consist of multiple interacting subsystems and stitching these together always became quite hacky, even with the nice JuMP frontend.

To clean that up, we needed a few new features:

  • Easy creation of constraints that connect multiple parts of the system. E.g., in metabolic modeling that could be tying the rate of a reaction to the amount of catalytic material produced by a cell.
  • A bit of polymorphism in what are variables and what are “derived values”. In particular, previously our code had to be specialized for models where reaction fluxes were represented as variables, for models where the fluxes were “split” into positive and negative contributions, and for models where fluxes were implicitly derived. We were not very happy about the resultant complexity.
  • Minimization of the burden of allocating the variables – In JuMP, @variables macro for shared variables acts basically as malloc() for shared memory, and suffers from same problems. Moreover, for users it was sufficiently challenging to make sure there are no naming & identifier conflicts between the set of constraints, and managing another set of identifiers connected to the constraints via a completely random bipartite system has proven to be daunting.

So we made ConstraintTrees.jl as a small stepping stone for bridging the organization gap between the complicated constraint systems and JuMP. Mainly, it takes care about all variable allocations for you, and provides a way to hierarchically organize your model contents in a directory-like structure so that you do not need to manually manage any global identifier pool.

What does it do?

In short, you get a dictionary-like structure called a ConstraintTree where you can create and store constraints; later you can pass these to JuMP for solving.

Let’s solve a quadratic optimization problem as a demo (this here is a shortened version, the full code is in the documentation): Let’s have an ellipse and a line, and find the point with the shortest distance between them.

First, let’s make the point; that one is literally just 2 variables:

import ConstraintTrees as C
point = C.variables(keys=[:x, :y])

The point is a simple ConstraintTree; we can explore it using dot notation:

point.x

This prints out a “constraint”, which is a “value” possibly (not necessarily) bounded into a closed interval. The values are either linear or quadratic combinations, possibly affine, of some variables (in this case each value on point is a trivial “combination” of one variable).

Now we can label our point using ^, e.g. call it A:

:A ^ point

We can also add other constraints to the point using * (an algebraic version of “and”, as with Julian string concatenation); for example here we force the point to rest on a line:

point_on_a_line = (
    :coords             ^ point *
    :on_line_constraint ^ C.Constraint(
        point.x.value - 2 * point.y.value,
        0.0,  # equality constraint
    )
)

…or we can restrict it to lie in an ellipse:

point_in_an_ellipse = (
    :coords ^ point *
    :in_ellipse_constraint ^ C.Constraint(
        C.squared(0.5 * point.x.value) + C.squared(point.y.value - 10.0),
        (-Inf, 1.0),  # interval constraint
    )
)

With that, we can make a system that has both the ellipse-point and a line-point. For that, we use + instead of * to combine the systems: + means combining them as independent ones, making sure the point on the ellipse is completely unrelated to the one on the line:

points = :line_point ^ point_on_a_line + :ellipse_point ^ point_in_an_ellipse

Notably, the system now has a total of 4 variables instead of just 2. This completes the “algebra of constraints” with operations^, * and +.

Finally, we might like finding the actual 2 point coordinates that are closest between the line and an ellipse. So, we can attach a suitable objective representation to our optimization problem:

points *= :squared_distance ^ C.Constraint(
    C.squared(points.line_point.x.value - points.ellipse_point.x.value) +
    C.squared(points.line_point.y.value - points.ellipse_point.y.value)
) # not an actual constraint, just kindof a "named value".

We can now use some prepared scriptage to get this solved in JuMP; and get another tree that contains all values of the constraints in it, including the coordinates of both solved points and the distance:

function optimize(...)
    #= copypaste a suitable implementation from the docs =#
end

import Clarabel
solution = C.ValueTree(
    points,
    optimize(
        points,
        objective = -points.squared_distance.value,
        optimizer = Clarabel.Optimizer,
    ),
)

The solution may be explored quite easily:

julia> solution
ConstraintTrees.Tree{Float64} with 3 elements:
  :ellipse_point    => ConstraintTrees.Tree{Float64}(#= 2 elements =#)
  :line_point       => ConstraintTrees.Tree{Float64}(#= 2 elements =#)
  :squared_distance => 58.9726

julia> solution.line_point.coords
ConstraintTrees.Tree{Float64} with 2 elements:
  :x => 4.84863
  :y => 2.42432

julia> solution.ellipse_point.coords
ConstraintTrees.Tree{Float64} with 2 elements:
  :x => 1.41431
  :y => 9.29294

julia> # the constrained value should be in the requested interval:
julia> solution.ellipse_point.in_ellipse_constraint
1.0

julia> sqrt(solution.squared_distance)
7.679360851160642

What do we plan to do with this?

The approach has proven to be quite nice for organizing seriously big constraint systems. We were able to simplify and generally remove a LOT of code already, and everything seems to scale and work just right. Some actual domain-specific example use can be seen in the other documentation example.

Enjoy the optimizations!
-mk

PS. We’re going to be grateful for any suggestions here from the JuMP community. Perhaps there’s another field where large constraint systems dominate the complexity that could benefit from such an approach?

10 Likes