Graph AD for Forces

[DhairyaLGandhi]

ref #111

x-ref: #111 (comment), #10

Okay, so we have the following working.

julia> using ChemistryFeaturization, Zygote, Xtals

julia> cr_path = normpath(joinpath(pathof(Xtals), "..", "..", "test", "data", "crystals"))
"/Users/dhairyagandhi/.julia/packages/Xtals/Kf4en/test/data/crystals"

julia> using ChemistryFeaturization.Utils.GraphBuilding

julia> c = Crystal(joinpath(cr_path, "IRMOF-1.cif"))
Name: /Users/dhairyagandhi/.julia/packages/Xtals/Kf4en/test/data/crystals/IRMOF-1.cif
Bravais unit cell of a crystal.
	Unit cell angles α = 90.000000 deg. β = 90.000000 deg. γ = 90.000000 deg.
	Unit cell dimensions a = 25.832000 Å. b = 25.832000 Å, c = 25.832000 Å
	Volume of unit cell: 17237.492730 ų

	# atoms = 424
	# charges = 0
	chemical formula: Dict(:Zn => 4, :H => 12, :O => 13, :C => 24)
	space Group: P1
	symmetry Operations:
		'x, y, z'

julia> gradient(c) do c
         w, s = build_graph2(c)
         sum(w)
       end
((name = nothing, box = (a = -31.83244099108285, b = -31.832440991082862, c = -31.832440991082862, α = -2.1514963693844834e-14, β = -2.8035019022818323e-14, γ = -3.751262522266638e-14, Ω = nothing, f_to_c = nothing, c_to_f = nothing, reciprocal_lattice = nothing), atoms = (n = nothing, species = nothing, coords = (xf = [-4.78645247969116 4.786452479691137 … 51.735173803203864 51.735173803203864; 4.786452479691156 -4.786452479691122 … 77.23782826208793 -77.23782826208932; 4.786452479691155 4.7864524796911665 … -77.23782826208932 77.23782826208792],)), charges = nothing, bonds = nothing, symmetry = nothing),)

which are the gradients for the crystal and graph AD 🎉

Couple of notes:

1. There were a few missing methods in Xtals to calculate distance, I added those Allow vectors of indices while computing distance SimonEnsemble/Xtals.jl#101 (and also included a copy in the PR)
2. replicate is written with a lot of mutation, and I figured we could write a simpler version, so I did. How do we want to handle this? We could propose this change to Xtals.jl (the performance is roughly same, but of course this allocates) (this implementation is roughly 3x faster)
3. I have refactored some of the indexing utilities and removed them from AD to save some time there. It means we can write it in a neater way if desired.
4. How do we want to test this? Do we have some cases in mind?
5. I might have messed up the rebase, so I'll have to move this to a different PR.

Also, we might have some performance on the table if we optimise this further going by the number of allocations.

julia> @btime gradient($c) do c
         w, s = build_graph(c)
         sum(w)
       end;
  56.376 ms (572641 allocations: 50.27 MiB)

[DhairyaLGandhi]

Another piece of good news is that there is very little involved for AD, for the most part we have refactored some functions and added a couple methods that would be good for Xtals.jl to have anyway.

[rkurchin]

It seems like some of the additional tests are failing on this at the moment?

[DhairyaLGandhi]

Yeah I'm not sure what's changed here. The weights_cutoff function is the same as before.

[rkurchin]

The way things are getting reorganized, this code will probably end up belonging in AtomGraphs (should be registered in a couple days) rather than here, JFYI