Generating binary trees

[MilesCranmer]

Hi @Seelengrab,

I'm looking to integrate Supposition.jl into the tests of DynamicExpressions.jl and SymbolicRegression.jl and was wondering if I could ask if I am going the right direction. I am attempting to build a generator for binary trees up to a certain max depth.

Here is what I have so far:

using DynamicExpressions: OperatorEnum, Node
using Supposition

# Wish to define trees over these nodes:
operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(abs, cos))

# Element type:
T = Float64
# Can use x1 through x5:
features = 1:5

# Constant nodes, like `1.0`
val_gen = Data.Floats{T}(; nans=true, infs=true)
val_node_gen = map(val -> Node{T}(; val), val_gen)

# Feature nodes, like `x1`
feature_gen = Data.SampledFrom(features)
feature_node_gen = map(feature -> Node{T}(; feature), feature_gen)

# Randomly choose a possible branch structure;
operator_gen = Data.SampledFrom(
    vcat(
        map(op -> (; op, degree=1), eachindex(operators.unaops)),
        map(op -> (; op, degree=2), eachindex(operators.binops)),
    ),
)

leaf_node_gen = val_node_gen | feature_node_gen

branch_gen = @composed function _branch_gen(
    l=leaf_node_gen, r=leaf_node_gen, operator=operator_gen
)
    if operator.degree == 1
        return Node{T}(; operator.op, l)
    else # degree == 2
        return Node{T}(; operator.op, l, r)
    end
end
# ^Generates short trees, like `cos(x5)` or `x1 - 0.3232`

Now, I wish to use branch_gen to recursively generate trees of arbitrary size. Here's my attempt:

Data.recursive(branch_gen; max_layers=3) do root
    leafs = filter(t -> t.degree == 0, root)
    leafs_gen = Data.SampledFrom(leafs)

    # ?
    # Want to call `branch_gen` here with the new leafs.
end

The filter(t -> t.degree == 0, root) is a utility in DynamicExpressions.jl to pick out all the leafs of a tree. Then, I put this in a data generator with SampledFrom.

However, I'm not sure the next step. What I want to do now is call branch_gen to get a new branch node, and then set the leaf node equal to that branch node, thus expanding the tree. But I'm not sure how to pass leafs_gen to branch_gen and I can't seem to create closures (#45 and #46).

Any tips? Alternative to the above approach, I'd love to hear any tips on generating binary trees like https://github.com/SymbolicML/DynamicExpressions.jl/blob/cab1143ae06b147cd81482a783912d9deb9247d6/src/Node.jl#L73-L135. (I couldn't seem to transfer the JSON example to this)

Thanks again for the nice package!!
Best,
Miles

[Seelengrab]

I think what's tripping you up is that you're thinking of the object passed to recursive as the root, not a child. Note the definition of jsonwrap from the docs:

function jsonwrap(child)
    vecs = Data.Vectors(child)
    dicts = Data.Dicts(strings, child)
    vecs | dicts
end

jsonwrap takes a Possibility for generating children (child), builds a Vector/Dict of them, and finally returns a Possibility for either. Translated to your example, you'll want to do something like

Data.recursive(branch_gen; max_layers=3) do childgen
    branches = Data.Vectors(childgen; min_size=2, max_size=2) # this really should be something that generates tuples...
    map(branches) do v
        Node{T}(; operator.op, v...)
    end
end

i.e. build a Possibility for new parents that's composing the generated children. In the simplest case, that truly can be just maping the function _branch_gen that's defined by @composed.

I agree that #45 would be the most appropriate instead of Data.Vectors! That would then also allow generating the operator for the new root node from that last map properly; as-is this code doesn't work because of that.

The crucial thing is that the function passed to recursive works on Possibility, and not on the objects produced by those Possibility.

[MilesCranmer]

(Although for posterity I think that operator_gen would need to be either binary_operator_gen or unary_operator_gen so you sample the right space of functions)

[Seelengrab]

Are there are any types of trees that this wouldn't be able to sample?

That's a good question, I haven't tried! My gut says that this might lead to only balanced trees, i.e. where each child is of the same height. If that's the case, it could be fixed like so, giving the option to not generate a new node at all:

Data.recursive(branch_gen; max_layers=3) do childgen
    a = map(operator_gen, childgen, childgen) do op, l, r
        Node{T}(; op, l, r)
    end
    b = map(operator_gen, childgen) do op, l
        Node{T}(; op, l)
    end
    a | b | childgen
end

This gives a 1/3 chance for each to occur. You could probably tweak that distribution by making use of WeightedSample:

Data.recursive(branch_gen; max_layers=3) do childgen
    a = map(operator_gen, childgen, childgen) do op, l, r
        Node{T}(; op, l, r)
    end
    b = map(operator_gen, childgen) do op, l
        Node{T}(; op, l)
    end
    ws = Data.WeightedSample((childgen, a, b), Float64[1, 3, 3])
    map(produce!, ws)
end

which should end up tweaking the probability for simply not generating a layer in a given branch to 1/7 per layer.

[MilesCranmer]

With the currently available method of

Data.recursive(branch_gen; max_layers=3) do childgen
    branches = Data.Vectors(childgen; min_size=2, max_size=2) # this really should be something that generates tuples...
    map(branches) do v
        Node{T}(; operator.op, v...)
    end
end

I'm not sure how to generate operators at each node? Each node should have a random operator index.

[MilesCranmer]

Oh maybe

tree_gen = Data.recursive(leaf_node_gen; max_layers=3) do childgen
    deg0_node_gen = childgen
    deg1_node_gen = map(
        Data.Dicts(
            unaop_gen,
            Data.Vectors(childgen; min_size=1, max_size=1);
            min_size=1,
            max_size=1
        )
    ) do cs
        (op, nodes) = only(cs)
        op
        l = only(nodes)
        Node{T}(; op, l)
    end
    deg2_node_gen = map(
        Data.Dicts(
            binop_gen,
            Data.Vectors(childgen; min_size=2, max_size=2);
            min_size=1,
            max_size=1
        )
    ) do cs
        (op, nodes) = only(cs)
        l = first(nodes)::Node
        r = last(nodes)::Node
        Node{T}(; op, l, r)
    end
    deg0_node_gen | deg1_node_gen | deg2_node_gen
end

Is the way to do it

[Seelengrab]

That's also an option, yeah - just much uglier than using #45! It should be up now on feat/multimap, so let me know if that implementation works for you.