Towards graph AD for Forces

Project.toml

@@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 CSV = "0.7, 0.8"
@@ -35,7 +36,9 @@ SimpleWeightedGraphs = "1"
 julia = "1.4, 1.5, 1.6"
 
 [extras]
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test"]
+test = ["FiniteDifferences", "Test", "Zygote"]

src/utils/adjoints.jl

@@ -0,0 +1,60 @@
+using Zygote # , ChainRulesCore
+using Zygote: @adjoint
+using LinearAlgebra
+
+@adjoint function Base.Generator(f, iter)
+  ys, backs = Zygote.unzip([Zygote.pullback(f, x) for x in iter])
+  Base.Generator(f, iter), Δ -> begin
+    b(d::Dict) = [back(v)[1] for (back,v) in zip(backs,values(d))]
+    b(nt::NamedTuple{(:f, :iter)}) = [back(i)[1] for (back,i) in zip(backs,nt.iter)]
+    (nothing, b(Δ))
+  end
+end
+
+@adjoint function Base.Iterators.Zip(is)
+  Zip_pullback(Δ) = (Zygote.unzip(Δ),)
+  return Base.Iterators.Zip(is), Zip_pullback
+end
+
+@adjoint function Pair(a, b)
+  Pair(a, b), Δ -> (Δ, nothing)
+end
+
+@adjoint function Dict(g::Base.Generator)
+  ys, backs = Zygote.unzip([Zygote.pullback(g.f, args) for args in g.iter])
+  Dict(ys...), Δ -> begin
+    ∂d = first(backs)(Δ)[1]
+    d = Dict(ys...)
+    for (k,v) in pairs(d)
+      d[k] = _zero(v)
+    end
+    (merge(d, ∂d), )
+  end
+end
+
+_zero(x) = zero(x)
+_zero(::Nothing) = nothing
+
+@adjoint function _cutoff!(weight_mat, f, ijd,
+                           nb_counts, longest_dists;
+                           max_num_nbr = 12)
+  y, ld = _cutoff!(weight_mat, f, ijd,
+               nb_counts, longest_dists;
+               max_num_nbr = max_num_nbr)
+  function cutoff_pb((Δ,nt))
+    s = size(Δ)
+    Δ = vec(collect(Δ))
+    for (ix,(_,_,d)) in zip(eachindex(Δ), ijd)
+      y_, back_ = Zygote.pullback(f, d)
+      Δ[ix] *= back_(Δ[ix])[1]
+    end
+    (reshape(Δ, s), nothing,
+    collect(zip(fill(nothing, size(Δ,1)),
+                fill(nothing, size(Δ,1)),
+                Δ)),
+    nothing,
+    nothing)
+  end
+
+  (y,ld), cutoff_pb
+end

src/utils/graph_building.jl

@@ -7,6 +7,7 @@ export inverse_square, exp_decay
 using PyCall
 using ChemistryFeaturization
 using Serialization
+using Zygote
 
 # options for decay of bond weights with distance...
 # user can of course write their own as well
@@ -91,20 +92,32 @@ function weights_cutoff(is, js, dists; max_num_nbr = 12, dist_decay_func = inver
 
     # iterate over list of tuples to build edge weights...
     # note that neighbor list double counts so we only have to increment one counter per pair
-    weight_mat = zeros(Float32, num_atoms, num_atoms)
+    weight_mat = zeros(Float64, round(Int,num_atoms), round(Int,num_atoms))
+    weight_mat, longest_dists = _cutoff!(weight_mat,
+                                         dist_decay_func,
+                                         ijd,
+                                         nb_counts,
+                                         longest_dists)
+    # average across diagonal, just in case
+    weight_mat = 0.5 .* (weight_mat .+ weight_mat')
+end
+
+
+function _cutoff!(weight_mat, f, ijd,
+                  nb_counts, longest_dists; max_num_nbr = 12)
+
     for (i, j, d) in ijd
         # if we're under the max OR if it's at the same distance as the previous one
-        if nb_counts[i] < max_num_nbr || isapprox(longest_dists[i], d)
-            weight_mat[i, j] += dist_decay_func(d)
-            longest_dists[i] = d
-            nb_counts[i] += 1
+        if nb_counts[round(Int,i)] < max_num_nbr || isapprox(longest_dists[round(Int,i)], d)
+            weight_mat[round(Int,i), round(Int,j)] += f(d)
+            longest_dists[round(Int,i)] = d
+            nb_counts[round(Int,i)] += 1
         end
     end
-
-    # average across diagonal, just in case
-    weight_mat = 0.5 .* (weight_mat .+ weight_mat')
+    weight_mat, longest_dists
 end
 
+
 """
 Build graph using neighbors from faces of Voronoi polyedra and weights from areas. Based on the approach from https://github.com/ulissigroup/uncertainty_benchmarking/blob/aabb407807e35b5fd6ad06b14b440609ae09e6ef/BNN/data_pyro.py#L268
 """
@@ -135,4 +148,6 @@ end
 
 # TODO: graphs from SMILES via OpenSMILES.jl
 
+include("adjoints.jl")
+
 end

test/utils/GraphBuilding_tests.jl

@@ -1,5 +1,6 @@
 using Test
 using ChemistryFeaturization.Utils.GraphBuilding
+using Zygote, FiniteDifferences
 
 @testset "GraphBuilding" begin
     adj, els = build_graph(
@@ -39,3 +40,25 @@ using ChemistryFeaturization.Utils.GraphBuilding
     @test all(isapprox.(adj[2:5, 1], 1.0, atol = 1e-4))
     @test all(isapprox.(adj[3:2, 2], 0.375, atol = 1e-5))
 end
+
+@testset "Graph Building AD tests" begin
+
+  function test_fd(i, j, dist)
+      fd = grad(forward_fdm(2,1),
+                (i,j,dist) -> sum(GraphBuilding.weights_cutoff(i,j,dist)),
+                i, j, dist)
+
+      gs = gradient(i, j, dist) do i, j, dist
+          sum(GraphBuilding.weights_cutoff(i, j, dist))
+      end
+
+      @test gs[1] == fill(nothing, length(i))
+      @test gs[2] == fill(nothing, length(j))
+      @test gs[3] ≈ fd[3]
+  end
+
+  # test with non-overlapping indices
+  test_fd(collect(1:10), collect(1:10), Float64.(collect(1:10)))
+  # test with overlapping indices
+  test_fd(rand(1:10, 100), rand(1:10, 100), rand(100))
+end