# Tree \<: Collection? Looking for API feedback in DynamicExpressions
**URL:** https://discourse.julialang.org/t/tree-collection-looking-for-api-feedback-in-dynamicexpressions/98109
**Category:** General Usage
**Tags:** question, package, performance, symbolic-regression
**Created:** [April 30, 2023, 12:09am UTC](https://discourse.julialang.org/t/tree-collection-looking-for-api-feedback-in-dynamicexpressions/98109 "2023-04-30T00:09:11Z")
**Posts on this page:** 2
**Page:** 1
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [April 30, 2023, 12:09am UTC](https://discourse.julialang.org/t/tree-collection-looking-for-api-feedback-in-dynamicexpressions/98109/1 "2023-04-30T00:09:11Z")
## Intro
I am wondering if I could others’ opinions on a potential interface for traversing expressions in [DynamicExpressions.jl](https://github.com/SymbolicML/DynamicExpressions.jl).
To give you a sense of what the package does, the original announcement thread for this package is here: [[ANN] DynamicExpressions.jl: Fast evaluation of runtime-generated expressions without compilation](https://discourse.julialang.org/t/ann-dynamicexpressions-jl-fast-evaluation-of-runtime-generated-expressions-without-compilation/88983).
DynamicExpressions.jl is used in [PySR](http://github.com/MilesCranmer/PySR) and [SymbolicRegression.jl](https://github.com/MilesCranmer/SymbolicRegression.jl/) as the backend - it lets those packages evaluate arbitrary expressions without additional compilation.
* * *
## Tree structure
Expressions are currently implemented as binary trees with the following rigid structure. It is built for type stability, fast traversals, and low memory usage.
```julia-auto
mutable struct Node{T}
degree::Int # 0 for constant/variable, 1 for cos/sin, 2 for +/* etc.
constant::Bool # false if variable
val::Union{T,Nothing} # If is a constant, this stores the actual value
# ------------------- (possibly undefined below)
feature::Int # If is a variable (e.g., x in cos(x)), this stores the feature index.
op::Int # If operator, this is the index of the operator in operators.binary_operators, or operators.unary_operators
l::Node{T} # Left child node. Only defined for degree=1 or degree=2.
r::Node{T} # Right child node. Only defined for degree=2.
end
```
(in the future I will likely change this to `Node{T,D}` for some max degree `D`, but so far there hasn’t been a strong need)
## Problem
Following a discussion with @jling I realized it would be nice if there was an easier way to traverse expressions and compute various properties. For example, you might want to know the number of occurrences of a particular operator (e.g., for solving [this puzzle](https://discourse.julialang.org/t/subset-sum-problem-with-symbolicregressions-jl/97479)), or you might want to extract all the constants in the expression for use in an optimizer (i.e., used in [SymbolicRegression.jl](https://github.com/MilesCranmer/SymbolicRegression.jl/) for BFGS-optimization of a given expression).
For example, right now if you want to compute the number of nodes in an expression (i.e., constants+operators+variables), you could write the function (which performs a [depth-first search](https://en.wikipedia.org/wiki/Depth-first_search))
```julia-auto
function count_nodes(tree::Node)::Int
if tree.degree == 0
return 1
elseif tree.degree == 1
return 1 + count_nodes(tree.l)
else
return 1 + count_nodes(tree.l) + count_nodes(tree.r)
end
end
```
likewise if you wanted to compute the depth of a tree, you could write:
```julia-auto
function count_depth(tree::Node)::Int
if tree.degree == 0
return 1
elseif tree.degree == 1
return 1 + count_depth(tree.l)
else
return 1 + max(count_depth(tree.l), count_depth(tree.r))
end
end
```
These specific functions are already implemented in `DynamicExpressions.jl`, but I think there are many other properties users would want to compute which are not available in the library.
I would like it to be easier for users to write their own functions operating on these trees.
## Trees as a Collection
Here is the idea for an interface: **treat the tree like an ordered Julia collection**. By default, the “order” of nodes in the collection would be that of a depth-first search, but I think in the future you could specify other orderings.
With this interface, I implemented four primitive functions:
- `mapreduce` (aggregation is commutative),
- `tree_mapreduce` (aggregation is non-commutative; needs to know parent from child),
- `filter_and_map`, (filters nodes, then maps, then collects) and
- `any` (to lazily traverse a tree).
From these, most other [core functions on collections](https://docs.julialang.org/en/v1/base/collections/) are implemented: `filter`, `collect`, `map`, `count`, `sum`, `all`, `iterate`, `in`, `length`, and others.
For example, `count_nodes` becomes just:
```julia-auto
count_nodes(tree::Node) = count(_ -> 1, tree) # == length(tree)
```
which is just as fast as the version I gave above. `count_depth` becomes:
```julia-auto
count_depth(tree::Node) = tree_mapreduce(_ -> 1, (p, c...) -> p + max(c...), tree)
```
which also experiences no change in performance.
Likewise you could get a list of nodes with `collect(tree)`, or do things like extract all nodes in the tree which store constants with:
```julia-auto
filter(t -> t.degree==0 && t.constant, tree)
```
or perhaps count the occurrence of a particular binary operator with:
```julia-auto
count(t -> t.degree==2 && t.op == 1, tree)
```
or see if there are any negative constants in the expression with:
```julia-auto
any(t -> t.degree==0 && t.constant && t.val<0, tree)
```
or take the product of all constants with:
```julia-auto
mapreduce(*, tree) do t
t.degree == 0 && t.constant && return t.val
return one(T)
end
```
`iter` is implemented, so you can also do things like:
```julia-auto
for node in tree
if node.degree==0 && !node.constant && node.feature == 1
println("Found feature 1 in tree!")
end
end
```
which will iterate through the nodes depth-first order.
Since every node in a tree is a tree itself, you could have traversals within traversals, like:
```julia-auto
any(tree) do subtree
if subtree.degree == 2 && subtree.op == 1
return any(t -> t.degree==0&&!t.constant&&t.feature==3, subtree)
end
return false
end
```
which checks if there are any instances where feature 3 appears within the first binary operator.
I also have `eachindex`, `setindex!`, and `getindex` implemented, which refer to their order in the tree-traversal search. However I might delete those as they seem unintuitive (and are not fast anyways, as every call to `getindex` would result in a fresh traversal).
## Implementation
You can find the code implementing this here: [https://github.com/SymbolicML/DynamicExpressions.jl/blob/tree-map/src/tree\_map.jl](https://github.com/SymbolicML/DynamicExpressions.jl/blob/tree-map/src/tree_map.jl) (coding tips appreciated too, I am always eager to learn more about optimizing Julia code)
* * *
# Your Feedback
What do you think of this interface?
Would it be weird as a user if you could treat a tree structure like an ordered collection?
Would you prefer function names which are more explicit, such as `tree_mapreduce`, `tree_any`, etc.?
Would you rather be required to declare the type of ordering of nodes explicitly, like:
```julia-auto
for node in DepthFirst(tree) # New type DepthFirst{Node{T}}
# ...
end
```
or would it be fine as a default if this order is assumed? (Regardless, I may add `DepthFirst / BreadthFirst` types, with `DepthFirst` assumed if not given).
* * *
Thanks for any tips.
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [April 30, 2023, 4:31pm UTC](https://discourse.julialang.org/t/tree-collection-looking-for-api-feedback-in-dynamicexpressions/98109/2 "2023-04-30T16:31:18Z")
One other thought would be to only require the `DepthFirst(tree)` operation when performing operations specific to _ordered_ collections, like indexing. But operations which are order-invariant like `mapreduce` could take `tree` as an argument.