Category Archives: Julia

ABC of A/B Testing

By: Blog by Bogumił Kamiński

Re-posted from: https://bkamins.github.io/julialang/2020/11/27/abt.html

Introduction

After my last post
on random number generation I think that it is good to add one more thing
that Chris Rackauckas pointed out to me as relevant to show a complete
picture. If you are working with pseudo random number generators it is good
to know that Julia has an excellent RandomNumbers.jl package that
offers a wide array of different random number generators to choose from.
So in this post I will use xoroshiro128** generator.

The post was tested under Julia 1.5.3 and packages DataFrames.jl 0.22.1,
Distributions.jl 0.23.8, FreqTables.jl 0.4.2, Pipe.jl 1.3.0,
RandomNumbers.jl 1.4.0, PyPlot.jl 2.9.9 and intends to present a pretty complete
work-flow of investigating a sample data science problem.

A/B testing

In the post we will investigate several simple options how A/B tests
can be performed.

In our scenario we assume that we have to choose from \(n\) alternatives.
Each alternative \(i\in[n]\) has an unknown probability of success \(p_i\).
Assume that e.g. we have \(n\) different ways to promote a product to a customer
and \(p_i\) is a probability that customer buys after seeing a promotion.

Initially we do not know which of the options is the best, i.e. has the highest
\(p_i\), but we are allowed to present them to customers and collect their
responses to learn the unknown probabilities.

The most simple way to run such an experiment is to present each option to the
customer with the same probability \(1/n\). Such an approach focuses purely on
exploration of the alternatives. However, intuitively we feel that it is
wasteful. If after some trials we see that some option is very bad and has
marginal chances of being the best then it is not worth showing it to the
customer. We are not only losing money (we present something that we know is not
very good), but additionally we will not learn much in terms of choosing the
best option.

Another approach would be to always show our customer an option that we
currently believe is best. Such a strategy focuses purely on exploitation.
Similarly to the uniform sampling described above, we can see that it does not
have to lead to good outcomes. We can get stuck with sub-optimal solution, as
we do not give ourselves a chance to learn that maybe some alternative that is
currently not considered the best is actually superior.

So how can we balance the exploration with exploitation? There are many
approaches that can be considered, but one of the simplest, yet very efficient,
is Thompson sampling. In this approach we do the following. We constantly
keep information about the distribution of our beliefs about the quality of each
alternative. Denote the random variable corresponding to the distribution of
these beliefs as \(P_i\) for alternative \(i\). Now in each step we
independently sample a value from each distribution \(P_i\) and choose to
measure the option that has the highest value of the sample. Now why does this
heuristic work? Note that we will tend to sample the alternatives that have a
high value of beliefs (so we do not waste time on sampling very bad options)
but at the same time we will also from time to time sample options that have
a high variance of \(P_i\) distribution (so the options about which we do not
know much; even if at the moment of sampling we believe they are not so good
our beliefs show us that actually they could turn out good if we collect more
evidence on them).

In the codes below I compare these three strategies of running A/B tests. I call
them respectively:

  1. Uniform: sample each alternative with probability \(1/n\);
  2. Greedy: sample the alternative with the highest \(E(P_i)\);
  3. Thompson: sample the alternative with the highest value of a single
    independent sample from \(P_i\).

Before we turn to coding one last thing should be discussed. How should we
represent our beliefs \(P_i\)? Fortunately here Bayesian statistics comes handy.
As our experiments are Bernoulli trials, we know that if we represent the
prior beliefs as Beta distributed then the posterior beliefs are also
going to follow the same distribution. Formally we say that the Beta
distribution is a conjugate prior of Bernoulli distribution. If our beliefs have
\(Beta(\alpha, \beta)\) distribution and we see a success then the distribution
of beliefs becomes \(Beta(\alpha+1, \beta)\). On the other hand if we see a
failure it becomes \(Beta(\alpha, \beta+1)\). As you can see it is very easy to
implement.

Setting up the stage

In a fresh Julia session we first load the required packages:

using DataFrames
using Distributions
using FreqTables
using Pipe
using RandomNumbers
using PyPlot
using LinearAlgebra

Next define a function that updates the beliefs about the quality of some
alternative:

update_belief(b::Beta, r::Bool) = Beta(b.α + r, b.β + !r)

Finally define the three decision rules that we want to compare: uniform,
greedy, and Thompson:

abstract type ABRule end

struct Greedy <: ABRule end
struct Thompson <: ABRule end
struct Unif <: ABRule end

function decide(options::Vector{<:Beta}, ::Unif)
    return rand(axes(options, 1))
end

function decide(options::Vector{<:Beta}, ::Greedy)
    m = mean.(options)
    mm = maximum(m)
    return rand(findall(==(mm), m))
end

function decide(options::Vector{<:Beta}, ::Thompson)
    m = rand.(options)
    mm = maximum(m)
    return rand(findall(==(mm), m))
end

One technical aspect of our implementations of decision making rule for greedy
and Thompson sampling is that we have to take into consideration that more than
one option can have the same evaluation. That is why we need first to find
which option has maximum mean with mm = maximum(m) and then randomly select
one of the options that are best with rand(findall(==(mm), m)).

Now we are ready to run the experiment.

Implementing the experimenter

Here is an implementation of the function that runs the experiments:

function coupled_run!(n::Int, steps::Int, rng::AbstractRNG,
                      results::DataFrame, runid::Int)
    truep = rand(rng, n)
    bestp = argmax(truep)
    streams = [rand(rng, Bernoulli(p), steps) for p in truep]

    for alg in (Greedy(), Unif(), Thompson())
        options = [Beta(1, 1) for i in 1:n]
        loc = fill(1, n)
        for i in 1:steps
            d = decide(options, alg)
            options[d] = update_belief(options[d], streams[d][loc[d]])
            loc[d] += 1
        end
        push!(results, (runid=runid, alg=string(typeof(alg)),
                        correct=mean(options[bestp]) == maximum(mean.(options)),
                        payoff=sum(x -> x.α - 1, options)))
    end
end

The coupler_run! takes five arguments: k is number of alternatives to
consider, steps is the length of A/B test experiment, rng is random number
generator we will use (remember that we want to test RandomNumbers.jl library),
results is a data frame that we will update and store the experiment results
in, finally runid is experiment number.

Let us comment on several crucial parts of the above code. Note that we sample
unknown true probabilities of success in the experiment with
truep = rand(rng, n). As they are uniformly distributed we set our initial
beliefs about options to also follow a uniform distribution in this line:
options = [Beta(1, 1) for i in 1:n].

The next notable thing is that we want the experiments for all algorithms we
consider to be coupled, that is: each alternative should generate the same
sequence of successes and failures for each option. Therefore, for simplicity,
we pre-populate the results of the trials in this line
streams = [rand(rng, Bernoulli(p), steps) for p in truep] and then just tack
in loc vector which random number for each option should be picked. Such an
approach will reduce noise when comparing the alternatives. (note that we could
have implemented it more efficiently, but I did not want to over-complicate the
code).

Finally note that we are collecting two informations:

  • correct: if what we think is best at the end of the experiment is best
    indeed.
  • payoff: how much payoff we have collected when running the experiment.

Note that both of them are desirable. Everyone probably agrees that it would
be best to have correct near to one (so that after running the experiment
we are confident that we know what is best). At the same time it is also
desirable to ensure that the experiment itself is not wasteful, so that we
collect as big payoff as possible.

We will compare the three considered A/B testing rules using these two criteria.

Analyzing the experiment results

First run the experiment (it should take under 20 seconds):

julia> results = DataFrame()
0×0 DataFrame

julia> rng = Xorshifts.Xoroshiro128StarStar(1234)
RandomNumbers.Xorshifts.Xoroshiro128StarStar(0x9e929e92000004d2, 0xda409a400013489a)

julia> @time foreach(runid -> coupled_run!(10, 500, rng, results, runid), 1:20_000)
 18.736959 seconds (143.71 M allocations: 8.944 GiB, 6.78% gc time)

julia> results
60000×4 DataFrame
   Row │ runid  alg       correct  payoff
       │ Int64  String    Bool     Float64
───────┼───────────────────────────────────
     1 │     1  Greedy      false    451.0
     2 │     1  Unif         true    354.0
     3 │     1  Thompson     true    486.0
     4 │     2  Greedy       true    468.0
     5 │     2  Unif         true    264.0
   ⋮   │   ⋮       ⋮         ⋮        ⋮
 59997 │ 19999  Thompson     true    464.0
 59998 │ 20000  Greedy      false    472.0
 59999 │ 20000  Unif         true    394.0
 60000 │ 20000  Thompson     true    462.0
                         59991 rows omitted

We have run the sampling 20,000 times in under 20 seconds. We have 60,000 rows
in our DataFrame as for each sampling we produced three rows corresponding
to three alternatives we consider. In our experiment we considered 10
alternatives and a budget of 500 samples for testing.

First we collect some initial statistics about the alternatives:

julia> @pipe groupby(results, :alg) |>
             combine(_, [:correct, :payoff] .=> mean)
3×3 DataFrame
 Row │ alg       correct_mean  payoff_mean
     │ String    Float64       Float64
─────┼─────────────────────────────────────
   1 │ Greedy         0.38145      405.997
   2 │ Unif           0.8031       250.394
   3 │ Thompson       0.87955      431.31

We see that Thompson sampling has both best probability of being correct and
the best cumulative payoff. Interestingly greedy strategy has low probability
of being correct, but relatively high payoff. The opposite holds for uniform
sampling strategy.

Now we investigate correct selection probability in more detail:

julia> correct_wide = unstack(results, :runid, :alg, :correct)
20000×4 DataFrame
   Row │ runid  Greedy  Unif   Thompson
       │ Int64  Bool?   Bool?  Bool?
───────┼────────────────────────────────
     1 │     1   false   true      true
     2 │     2    true   true      true
     3 │     3   false   true      true
     4 │     4   false   true      true
     5 │     5   false  false     false
   ⋮   │   ⋮      ⋮       ⋮       ⋮
 19997 │ 19997   false   true      true
 19998 │ 19998   false  false      true
 19999 │ 19999   false   true      true
 20000 │ 20000   false   true      true
                      19991 rows omitted

julia> proptable(correct_wide, :Greedy, :Unif)
2×2 Named Array{Float64,2}
Greedy ╲ Unif │   false     true
──────────────┼─────────────────
false         │ 0.14145   0.4771
true          │ 0.05545    0.326

julia> proptable(correct_wide, :Greedy, :Thompson)
2×2 Named Array{Float64,2}
Greedy ╲ Thompson │   false     true
──────────────────┼─────────────────
false             │  0.0907  0.52785
true              │ 0.02975   0.3517

julia> proptable(correct_wide, :Unif, :Thompson)
2×2 Named Array{Float64,2}
Unif ╲ Thompson │   false     true
────────────────┼─────────────────
false           │  0.0893   0.1076
true            │ 0.03115  0.77195

Observe that Thompson and uniform sampling give quite similar results most of
the time , with Thompson sampling having a slight edge when they disagree.
Greedy strategy is much worse. It is easy to understand these findings. Greedy
strategy is not doing enough exploration. Both Thompson and uniform sampling
do explore. The edge of Thompson sampling is that it does not waste time testing
very bad alternatives (so it can focus on the good ones and better discriminate
between them).

Similarly we can investigate which alternative gives the best payoff (remember
that we made sure to have the scenarios coupled):

julia> function find_best(alg, payoff)
           best_idxs = findall(==(maximum(payoff)), payoff)
           best_algs = alg[best_idxs]
           sort!(best_algs)
           return join(best_algs, '|')
       end
find_best (generic function with 1 method)

julia> alg_freq = @pipe groupby(results, :runid) |>
                        combine(_, [:alg, :payoff] => find_best => :best) |>
                        groupby(_, :best, sort=true) |>
                        combine(_, nrow)
3×2 DataFrame
 Row │ best             nrow
     │ String           Int64
─────┼────────────────────────
   1 │ Greedy            9736
   2 │ Greedy|Thompson    120
   3 │ Thompson         10144

We observe that uniform sampling for 20,000 simulations was never the best,
while Thompson sampling has a slight edge over greedy sampling in the
probability of being best (there is also a small probability of a tie).

As an exercise let us check, sticking to Bayesian approach, what is the
probability that Thompson sampling has higher probability of being best than
greedy one. Here we use the fact that the Dirichlet distribution is a
conjugate prior for the Categorical distribution, and set uniform a priori
beliefs:

julia> posterior = Dirichlet(alg_freq.nrow .+ 1)
Dirichlet{Float64}(alpha=[9737.0, 121.0, 10145.0])

julia> posterior_dis = rand(rng, posterior, 100_000)
3×100000 Array{Float64,2}:
 0.484744    0.486453    0.492359    …  0.494059   0.484432    0.482332
 0.00632059  0.00692795  0.00563844     0.0076129  0.00706873  0.00782947
 0.508936    0.506619    0.502002       0.498328   0.508499    0.509838

julia> mean(x -> x[1] < x[3], eachcol(posterior_dis))
0.99707

We can see that given the sample size of 20,000 this probability is greater than
99%.

However, we see that the considered fractions actually do not differ that much:

julia> transform(alg_freq, :nrow => (x -> x ./ sum(x)) => :freq)
3×3 DataFrame
 Row │ best             nrow   freq
     │ String           Int64  Float64
─────┼─────────────────────────────────
   1 │ Greedy            9736   0.4868
   2 │ Greedy|Thompson    120   0.006
   3 │ Thompson         10144   0.5072

and the difference in average payoffs we observed above was relatively higher.
Let us dig into this. First we plot a distribution of difference between payoffs
with Thompson and greedy sampling:

payoff_wide = unstack(results, :runid, :alg, :payoff)
Δpayoff = payoff_wide.Thompson .- payoff_wide.Greedy
hist(Δpayoff, 25)
xlabel("Δpayoff: Thompson-Greedy")

getting the following plot:

Delta of payoff

Now we see what is going on. If Thompson sampling is worse than greedy sampling
it is worse only slightly. On the other hand greedy sampling can be much worse
than Thompson sampling (and we understand when: this is a scenario when greedy
algorithm gets stuck at some sub-optimal alternative).

We also quickly check the distribution of estimator of the difference in payoffs
using Bayesian bootstrapping (to keep to the style we use in this post):

julia> dd = Dirichlet(length(Δpayoff), 1);
julia> boot = [dot(Δpayoff, rand(rng, dd)) for _ in 1:10_000];

julia> quantile(boot, 0:0.25:1)
5-element Array{Float64,1}:
 23.427965439713827
 24.64825057737705
 24.911299035010433
 25.17170629935465
 26.513561155964837

and we see, that, as expected the distribution of the difference is very
significantly greater than zero (we could probably easily detect this
significance even with a much smaller number of tests than 20,000).

Conclusion

That was a long post, so let me finish at this point. I hope you enjoyed it and
could also appreciate the data wrangling capabilities that JuliaData ecosystem
offers you. I also encourage you to run some more experiments on your own
to chcek how the considered A/B testing strategies behave under different
parameterizations.

And a usual caveat – performance geeks (are there any in the Julia community?)
can find several spots in my codes that I left sub-optimal to keep them simple.
Feel free to experiment with improving the execution speed of my examples.

E-graphs in Julia (Part I)

By: Hey There Buddo!

Re-posted from: https:/www.philipzucker.com/egraph-1/

E-graphs are the data structure that efficiently compute with equalities between terms. It can be used to prove equivalences for equational reasoning, but also to find equational rewrites, which is useful for program transformations or algebraic simplification.

I got a hot tip from Talia Ringer on twitter about a paper and library called egg. Egg implements E-graphs as a performant generic library in Rust. The paper explains the data structure, a efficiency twist about when to fix its invariants, and shows some compelling use case studies. The claimed speedups over a naive implementation of E-graphs are absolutely nuts (x1000?!).

E-graphs compactly contain all the terms that are generated by a set of equalities and make it fast to perform the operation of congruence closure. Congruence is the property that $a = b \rightarrow f(a) = f(b)$ for all functions $f$. Functions take equals to equals. Congruence closure is taking an equivalence relation and closing it under this congruence operation, saturating the relation.

I decided to try parallel tracks of attempting to rebuild a version in Julia vs trying to build bindings to the implementation in egg. It is very unlikely I’ll beat the performance of egg, but native language flexibility is awful nice. Regardless I figured it’s a good learning experience.

SMT solvers like Z3 use E-graphs to reason about equality of uninterpreted functions. Understanding the E-graph structure has given me some possible insight into where I went wrong. Equality reasoning like this has been on my mind since I failed getting off-the-shelf provers to prove equations produced by Catlab.jl. So that’s sort of where I’m trying to head with this.

The actual E-graph data structure looks something like if a Union-Find and a Hash-Cons has a baby, so let’s talk about about those a bit.

Union Find or Disjoint Set Data Structure

The union find data structure is a fast data structure for representing partitions of sets.

Roughly how it works is that every element points to a parent in the partition. When you want to know what partition your thing is in, you follow the parent pointers up to a canonical element. While doing this you trickle the pointers upwards so the query will be faster next time. Eventually all the pointers in the partition will point directly to the canonical element with no indirection.

To merge two sets, set the parent pointer of one canonical element to the other.

That’s the rough sketch. There’s more in the wikipedia article https://en.wikipedia.org/wiki/Disjoint-set_data_structure. Julia has an implementation in the DataStructures.jl library https://juliacollections.github.io/DataStructures.jl/latest/disjoint_sets/ which I’m just going to use.

To start to see how union find is useful for equations, let’s say we have a pile of atoms :a :b :c :d and some equations :a = :c, :b = :d, :d = :c. Question: is :a = :b? For this small example you can see that it is, but the disjoint set structure will do the same.

using DataStructures
s = DisjointSets(:a,:b,:c,:d)
union!(s, :a, :c)
union!(s, :b, :d)
union!(s, :c, :d)
in_same_set(s,:a,:c)
# >>> true

Now the question is, how can we extend this to compound terms that appear in syntax trees like :(g(f(a)))?
The trick is consider each subterm as an atom in it’s own right, :a, :(f(a)), :(g(f(a))) and place these things into the disjoint set. This trick is significantly aided by the technique of hash-consing, which brings the structure and performance differences between compound terms and atoms closer in alignment.

Hash-Consing

Hash-Consing is a programming technique to convert tree data structures into a DAG data structure where all identical subterms are shared rather than copied. This has the advantage of faster processing and decreased memory usage for trees with a lot of repetitive subterms and fast equality comparison since you can check the unique identity of a term immediately rather than recursing down it.
The way this is done by identifying each node with a unique id and maintaining a mapping of nodes to ids in a hash table.
Whenever you create new nodes, you use smart constructors that first check if the node is already defined in the hash table. If it is, this canonical node is returned, otherwise the node is added to the hash table.

You can recursively convert a tree into it’s hash consed from by destructing it and reconstructing it with the smart constructors.
The unique id may be the memory address where the node is stored or it may come from a counter being incremented every time you create a novel term.

The canonical reference on hash-consing is Type-Safe Modular Hash-Consing. Hash consing has been folklore for a very, very long time though. There is a reference in here to a paper by Ershov in 1958 for example. A really fun example in there is using HashConsing for a BDD implementation.

Another term sometimes used is “interned”. Julia’s Symbol is interned.

E-Graphs

When we’re talking about equations, it no longer makes sense to give each term an id as in the pure hash-cons. Instead we want to give each equivalence class of terms a unique id and maintain the equivalence classes in a union-find structure. As we put more equalities into the union-find structure and the equivalence classes grow, these canonical ids change under our feet, which is annoying but manageable.

The E-graph maintains:

  • a map from hash-consed terms to class ids
  • a map from class-ids to equivalence classes
  • a unionfind data structure on class ids

Each equivalence class needs to maintain

  • an array of terms in the class
  • an array of possible parent terms for efficiently propagating congruences

We will index into congruence classes with Id. Wrapping the Int64 in a struct should have no performance cost, and keeps things sane.

struct Id
    id::Int64
end

We need a structure to hold hash consed terms of our language. AutoHashEquals.jl makes it so hash and == use a structural definition.

using AutoHashEquals
@auto_hash_equals mutable struct Term
    head::Symbol
    args::Array{Id}
end

An equivalence class is a set of nodes. We also need to store parents of these nodes in order to make congruence closure and hash cons updating fast.

mutable struct EClass
    nodes::Array{Term}
    parents::Array{Tuple{Term,Id}}
end

Finally we have the actual EGraph data structure. It holds a unionfind that keeps equivalences of classes, a memo field to map Term to their equivalence class, a classes field to map equivalence class Id to the EClass. Maintaining all the appropriate invariants and connections between these pieces sucks.

using DataStructures

mutable struct EGraph
    unionfind::IntDisjointSets
    memo::Dict{Term,Id} # int32 UInt32?
    classes::Dict{Id,EClass} # Use array?
    dirty_unions::Array{Id}
end

# Build an empty EGraph
EGraph() = EGraph(IntDisjointSets(0), Dict(), Dict(), [])

Finding the root class or querying whether in_same_set for an Id just passes off the query to the union find.

# returns the canonical Id
find_root!(e::EGraph, id::Id) = Id(DataStructures.find_root!(e.unionfind, id.id))

in_same_class(e::EGraph, t1::Term, t2::Term) = in_same_set(e.unionfind, e.memo[t1] , e.memo[t2])

We can canonicalize terms so that the Id arguments only point to the root class.

function canonicalize!(e::EGraph, t::Term)
    t.args = [ find_root!(e, a) for a in t.args ]
end

We can lookup the class a term belongs to by canonicalizing and using the memo table

function find_class!(e::EGraph, t::Term)  # lookup
    canonicalize!(e, t) #t.args = [ find_root!(e, a) for a in t.args ]  # canonicalize the term
    if haskey(e.memo, t)
        id = e.memo[t]
        return find_root!(e, id)
    else
        return nothing
    end
end

We can add new terms by finding if they already exist in the hashcons. If so it, it does nothing.
But if not, we need to make a new EClass for this term to live in. This is the only way new EClasses are made.
It is important to maintain the hashcons correctly and canonicalize so that we don’t accidentally create a new EClass for a term that should already be in another EClass.

function Base.push!(e::EGraph, t::Term)
    id = find_class!(e,t) # also canonicalizes t
    if id == nothing # term not in egraph. Make new class and put in graph
        id = Id(push!(e.unionfind))
        cls = EClass( [t] , [] )
        for child in t.args # set parent pointers
            push!(e.classes[child].parents,  (t, id))
        end
        e.classes[id] = cls
        e.memo[t] = id
        return id
    else
        return id
    end
end

Alright. Here is where stuff gets hairy. We want to union two equivalence classes.

  1. Cananonize the Id.
  2. Check if already equivalent. If so, we’re done
  3. Ask unionfind to perform union
  4. Recanonize the arguments of nodes in memo
  5. Merge nodes and parents arrays of the two eclasses
  6. Mark equivalence class as dirty for rebuilding.
  7. Destroy the unneeded EClass in classes
function Base.union!(e::EGraph, id1::Id, id2::Id)
    id1 = find_root!(e, id1)
    id2 = find_root!(e, id2)
    # An invariant is that the EClass data should always keyed with the root of the union find
    # in the e.classes datastructure
    if id1 != id2 # if not already in same class
        id3 = Id(union!(e.unionfind, id1.id,id2.id)) # perform the union find
        if id3 == id1 # picked id1 as root
            to = id1
            from = id2
        elseif id3 == id2 # picked id2 as root
            to = id2
            from = id1
        else
            @assert false
        end
            
        push!(e.dirty_unions, id3) # id3 will need it's parents processed in rebuild!
        
        # we empty out the e.class[from] and put everything in e.class[to]
        for t in e.classes[from].nodes
            push!(e.classes[to].nodes, t) 
        end
        
        # recanonize all nodes in memo.
        for t in e.classes[to].nodes
            delete!(e.memo, t) # remove stale term
            canonicalize!(e, t) # update term in place in nodes array
            e.memo[t] = to # now this term is in the to eqauivalence class
        end
        
        # merge parents list
        for (p,id) in e.classes[from].parents
            push!(e.classes[to].parents, (p, find_root!(e,id)))
        end
        
        # destroy "from" from e.classes. It's information should now be
        # in e.classes[to] and it should never be accessed.
        delete!(e.classes, from) 
    end 
        
end

This code might be right. Lemme tell ya, I’ve had versions only subtly different that weren’t right. So I dunno. Buyer beware.

Finally the last big chunk is rebuild!. This function works through the dirty_unions array and fixes parent nodes and propagates any congruences.


function repair!(e::EGraph, id::Id)
    cls = e.classes[id]
    
    #  for every parent, update the hash cons. We need to repair that the term has possibly a wrong id in it
    for (t,t_id) in cls.parents
        delete!( e.memo, t) # the parent should be updated to use the updated class id
        canonicalize!(e, t) # canonicalize
        e.memo[t] = find_root!(e, t_id) # replace in hash cons  
    end
    
    # now we need to discover possible new congruence eqaulities in the parent nodes.
    # we do this by building up a parent hash to see if any new terms are generated.
    new_parents = Dict()
    for (t,t_id) in cls.parents
        canonicalize!(e,t) # canonicalize. Unnecessary by the above?
        if haskey(new_parents, t)
            union!(e, t_id, new_parents[t])
        end
        new_parents[t] = find_root!(e, t_id)
    end
    e.classes[id].parents = [ (p,id) for (p,id) in new_parents]
end

function rebuild!(e::EGraph)
    while length(e.dirty_unions) > 0
       todo = Set([ find_root!(e,id) for id in e.dirty_unions])
       e.dirty_unions = []
       for id in todo
            repair!(e, id)
        end
    end
end


Here’s some helper functions to build constants and terms

function constant!(e::EGraph, x::Symbol)
      t = Term(x, [])
      push!(e, t)
end

term!(e::EGraph, f::Symbol) = (args...) -> begin
    t = Term(f, collect(args))
    push!(e,  t)
end

and some other helper functions to help draw the EGraph, which is pretty crucial

using LightGraphs
using GraphPlot
using Colors

function graph(e::EGraph) 
    nverts = length(e.memo)
    g = SimpleDiGraph(nverts)
    vertmap = Dict([ t => n for (n, (t, id)) in enumerate(e.memo)])
    #=for (id, cls) in e.classes
        for t1 in cls.nodes
            for (t2,_) in cls.parents
                add_edge!(g, vertmap[t2], vertmap[t1])
            end
        end
    end =#
    for (n,(t,id)) in enumerate(e.memo)
        for n2 in t.args
            for t2 in e.classes[n2].nodes
                add_edge!(g, n, vertmap[t2])
            end
        end
    end
    nodelabel = [ t.head for (t, id) in e.memo]
    classmap = Dict([ (id,n)  for (n,id) in enumerate(Set([ id.id for (t, id) in e.memo]))])
    nodecolor = [classmap[id.id] for (t, id) in e.memo]
    
    return g, nodelabel, nodecolor
end

function graphplot(e::EGraph) 
    g, nodelabel,nodecolor = graph(e)


    # Generate n maximally distinguishable colors in LCHab space.
    nodefillc = distinguishable_colors(maximum(nodecolor), colorant"blue")
    gplot(g, nodelabel=nodelabel, nodefillc=nodefillc[nodecolor])
end

Example Usage

A simple toy problem is to use equalities like $f^n(a) == a$. What this does is put a loop in the EGraph.

e = EGraph()
a = constant!(e, :a)
f = term!(e, :f)
apply(n, f, x) = n == 0 ? x : apply(n-1,f,f(x))


union!(e, a, apply(6,f,a))
display(graphplot(e))

If you think about it, if we assert another equality $f^m(a)==a$, what we’re kind of doing is computing the $gcd(n,m)$. Here the gcd(6,9) = 3.

union!(e, a, apply(9,f,a))
rebuild!(e)
display(graphplot(e))

And then if additionally we assert a prime number of applications of f, we reduced the loop of nodes to a single equivalence class

union!(e, a, apply(11,f,a))
rebuild!(e)
display(graphplot(e))

We’ve only just begun. To do anything really fun, we need to implement pattern matching over the E-graph (EMatching).

We’ll see.

Bits and Bobbles

  • I should try just binding to egg. A combo of https://github.com/herbie-fp/egg-herbie and https://github.com/felipenoris/JuliaPackageWithRustDep.jl makes it seem like this shouldn’t be so bad. My implementation is trash compared to egg’s and I haven’t even done pattern matching yet
  • Factoring out the hashcons in the struct definition is like factoring out Fix in recursion schemes
  • Factoring out hash cons is like factoring out the fix for a recursive definition. When you do this, it also opens up the option of memoization, like in fib
  • Unification and union find. If you don’t congruence close, and instead merge recursively down terms, you get a kind of unification. Each equivalence class can only hold variables and terms of the same head for the unification to succeed. This can be checked after unioning everything. I think this is roughly what the MM algorithm does
  • Halfway House E-graph. For a loss of efficiency, you can store a simpler e graph structure that is really is basically a hash cons and union-find smashed together.
  • The models returns by Z3 for a bunch of equalities are basically a serialization of the E graph. Could this be a way to build a rewrite library? You might have to reimplement pattern matching.

References