cleanup wip uncommit

src/Dice.jl

@@ -33,7 +33,6 @@ end
 
 include("autodiff/adnode.jl")
 include("autodiff/core.jl")
-include("autodiff/train.jl")
 include("dist/dist.jl")
 include("inference/inference.jl")
 include("analysis/analysis.jl")

src/inference/inference.jl

@@ -139,11 +139,6 @@ include("cudd/wmc.jl")
 # - pr(::Dist, evidence=..., errors=...)
 include("pr.jl")
 
-# Exposes functionality for changing the probabilities of flip_for's
-# to maximize a list of (possibly conditional) bools
-# Notable exports:
-# - train_group_probs!(::Vector{<:AnyBool}))
-# - train_group_probs!(::Vector{<:Tuple{<:AnyBool, <:AnyBool}})
 include("train_pr.jl")
 include("train_pr_losses.jl")
 

src/inference/train_pr.jl

@@ -1,5 +1,5 @@
 # The bridge between autodiff and cudd
-export step_vars!, train_pr!, LogPr, compute_mixed
+export LogPr, compute_mixed, train!
 using DataStructures: Queue
 
 mutable struct LogPr <: ADNode
@@ -30,6 +30,7 @@ function expand_logprs(l::LogPrExpander, root::ADNode)::ADNode
     foldup(root, fl, fi, ADNode, l.cache)
 end
 
+# Within root's LogPrs there are Dist{Bool} DAGs. Collect minimal roots all DAGs.
 function bool_roots(root::ADNode)
     # TODO: have non-root removal be done in src/inference/cudd/compile.jl
     seen_adnodes = Dict{ADNode, Nothing}()
@@ -52,56 +53,36 @@ function bool_roots(root::ADNode)
     setdiff(keys(seen_bools), non_roots)
 end
 
-# Find the log-probabilities and the log-probability gradient of a BDD
-function add_scaled_dict!(
-    x::AbstractDict{<:Any, <:Real},
-    y::AbstractDict{<:Any, <:Real},
-    s::Real
-)
-    for (k, v) in y
-        x[k] += v * s
-    end
-end
-
-function step_pr!(
+function compute_mixed(
     var_vals::Valuation,
-    loss::ADNode,
-    learning_rate::Real
+    x::ADNode
 )
-    l = LogPrExpander(WMC(BDDCompiler(bool_roots(loss))))
-    loss = expand_logprs(l, loss)
-    vals, derivs = differentiate(var_vals, Derivs(loss => 1))
-
-    # update vars
-    for (adnode, d) in derivs
-        if adnode isa Var
-            var_vals[adnode] -= d * learning_rate
-        end
-    end
-
-    vals[loss]
+    l = LogPrExpander(WMC(BDDCompiler(bool_roots(x))))
+    x = expand_logprs(l, x)
+    compute(var_vals, [x])[x]
 end
 
-# Train group_to_psp to such that generate() approximates dataset's distribution
-function train_pr!(
+function train!(
     var_vals::Valuation,
     loss::ADNode;
     epochs::Integer,
     learning_rate::Real,
 )
     losses = []
     for _ in 1:epochs
-        push!(losses, step_pr!(var_vals, loss, learning_rate))
+        l = LogPrExpander(WMC(BDDCompiler(bool_roots(loss))))
+        loss = expand_logprs(l, loss)
+        vals, derivs = differentiate(var_vals, Derivs(loss => 1))
+
+        # update vars
+        for (adnode, d) in derivs
+            if adnode isa Var
+                var_vals[adnode] -= d * learning_rate
+            end
+        end
+
+        push!(losses, vals[loss])
     end
     push!(losses, compute_mixed(var_vals, loss))
     losses
 end
-
-function compute_mixed(
-    var_vals::Valuation,
-    x::ADNode
-)
-    l = LogPrExpander(WMC(BDDCompiler(bool_roots(x))))
-    x = expand_logprs(l, x)
-    compute(var_vals, [x])[x]
-end

src/inference/train_pr_losses.jl

@@ -1,4 +1,3 @@
-
 export BoolToMax, mle_loss, kl_divergence
 
 struct BoolToMax

test/autodiff_test.jl

@@ -16,63 +16,58 @@ end
 @testset "test GD" begin
     x, y = Var("x"), Var("y")
 
-    params = Valuation(x => 0)
+    var_vals = Valuation(x => 0)
     e = 7 - (x - 5)^2
-    for _ in 1:100
-        step_maximize!(params, [e], 0.1)
-    end
-    vals = compute(params, [e])
-    @test vals[e] ≈ 7
-    @test vals[x] ≈ 5
+    train!(var_vals, -e, epochs=100, learning_rate=0.1)
+    @test compute_mixed(var_vals, e) ≈ 7
+    @test compute_mixed(var_vals, x) ≈ 5
 
-    var_vals = Valuation(x => 7, y => 1)
-    e = -(x*y-5)^2
-    for i in 1:100
-        step_maximize!(var_vals, [e], 0.01)
-    end
-    vals = compute(var_vals, [e])
-    @test abs(vals[e]) < 0.001
-    @test vals[x] * vals[y] ≈ 5
+    # var_vals = Valuation(x => 7, y => 1)
+    # e = -(x*y-5)^2
+    # train!(var_vals, e, 100, 0.01)
+    # vals = compute(var_vals, [e])
+    # @test abs(vals[e]) < 0.001
+    # @test vals[x] * vals[y] ≈ 5
 
-    psp = Var("psp") # pre-sigmoid probability
-    var_vals = Valuation(psp => 0)
-    p = sigmoid(psp)
-    # maximize logpr of flip(p) & flip(p) & !flip(p)
-    e = log(p * p * (1 - p))
-    for i in 1:500
-        step_maximize!(var_vals, [e], 0.1)
-    end
-    @test compute(var_vals, [p])[p] ≈ 2/3
+    # psp = Var("psp") # pre-sigmoid probability
+    # var_vals = Valuation(psp => 0)
+    # p = sigmoid(psp)
+    # # maximize logpr of flip(p) & flip(p) & !flip(p)
+    # e = log(p * p * (1 - p))
+    # for i in 1:500
+    #     step_maximize!(var_vals, [e], 0.1)
+    # end
+    # @test compute(var_vals, [p])[p] ≈ 2/3
 end
 
-@testset "matrices" begin
-    # Helper functions for matrices used as vectors
-    x(v) = v[1,1]
-    y(v) = v[2,1]
-    distance(u, v) = (x(u) - x(v))^2 + (y(u) - y(v))^2
-    to_matrix(v::Vector) = reshape(v, :, 1)
+# @testset "matrices" begin
+#     # Helper functions for matrices used as vectors
+#     x(v) = v[1,1]
+#     y(v) = v[2,1]
+#     distance(u, v) = (x(u) - x(v))^2 + (y(u) - y(v))^2
+#     to_matrix(v::Vector) = reshape(v, :, 1)
 
-    # Rotate [1, 2] by what angle to get closest to [-3, -3]?
-    θ = Var("θ")
-    var_vals = Valuation(θ => 0)
-    rotation_matrix = ADMatrix([[cos(θ) -sin(θ)]; [sin(θ) cos(θ)]])
-    rotated_vec = rotation_matrix * to_matrix([1, 2])
-    target_vec = to_matrix([-3, -3])
-    for _ in 1:2000
-        step_maximize!(var_vals, [-distance(rotated_vec, target_vec)], 0.003)
-    end
-    @test var_vals[θ] ≈ 5/8 * 2π - atan(2)
+#     # Rotate [1, 2] by what angle to get closest to [-3, -3]?
+#     θ = Var("θ")
+#     var_vals = Valuation(θ => 0)
+#     rotation_matrix = ADMatrix([[cos(θ) -sin(θ)]; [sin(θ) cos(θ)]])
+#     rotated_vec = rotation_matrix * to_matrix([1, 2])
+#     target_vec = to_matrix([-3, -3])
+#     for _ in 1:2000
+#         step_maximize!(var_vals, [-distance(rotated_vec, target_vec)], 0.003)
+#     end
+#     @test var_vals[θ] ≈ 5/8 * 2π - atan(2)
 
-    # Variables can also be matrices!
-    # Transform by [1, 2] by what matrix to get closest to [-3, -3]?
-    A = Var("A")
-    var_vals = Valuation(A => [[1 0]; [0 1]])
-    v = to_matrix([1, 2])
-    v′ = A * v
-    target_vec = to_matrix([-3, -3])
-    for _ in 1:2000
-        step_maximize!(var_vals, [-distance(v′, target_vec)], 0.003)
-    end
+#     # Variables can also be matrices!
+#     # Transform by [1, 2] by what matrix to get closest to [-3, -3]?
+#     A = Var("A")
+#     var_vals = Valuation(A => [[1 0]; [0 1]])
+#     v = to_matrix([1, 2])
+#     v′ = A * v
+#     target_vec = to_matrix([-3, -3])
+#     for _ in 1:2000
+#         step_maximize!(var_vals, [-distance(v′, target_vec)], 0.003)
+#     end
 
-    @test var_vals[A] * v ≈ [-3, -3]
-end
+#     @test var_vals[A] * v ≈ [-3, -3]
+# end