Switch from Zygote to Enzyme

.github/workflows/CI.yml

@@ -26,8 +26,8 @@ jobs:
         arch:
           - x64
         test:
-          - NotZygote
-          - Zygote
+          - NotGradients
+          - Gradients
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

Project.toml

@@ -10,19 +10,14 @@ AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1"
 BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864"
-ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Chemfiles = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EzXML = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
-KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
@@ -34,17 +29,19 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
-UnitfulChainRules = "f31437dd-25a7-4345-875f-756556e6935d"
 UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 
 [extensions]
+MollyEnzymeExt = "Enzyme"
 MollyGLMakieExt = ["GLMakie", "Colors"]
+MollyKernelDensityExt = "KernelDensity"
 MollyPythonCallExt = "PythonCall"
 
 [compat]
@@ -54,18 +51,15 @@ AtomsCalculators = "0.2"
 BioStructures = "4"
 CUDA = "4.2, 5"
 CellListMap = "0.8.11, 0.9"
-ChainRules = "1.44"
-ChainRulesCore = "1"
 Chemfiles = "0.10.3"
-Colors = "0.11, 0.12"
+Colors = "0.11, 0.12, 0.13"
 Combinatorics = "1"
 DataStructures = "0.18"
 Distances = "0.10"
 Distributions = "0.23, 0.24, 0.25"
-Enzyme = "0.11.15, 0.12"
+Enzyme = "0.13"
 EzXML = "1"
 FLoops = "0.2"
-ForwardDiff = "0.10.35"
 GLMakie = "0.8, 0.9, 0.10"
 Graphs = "1.8"
 KernelDensity = "0.5, 0.6"
@@ -81,7 +75,5 @@ StaticArrays = "1.8.2"
 Statistics = "1.9"
 Unitful = "1"
 UnitfulAtomic = "1"
-UnitfulChainRules = "0.1.2"
 UnsafeAtomicsLLVM = "0.1, 0.2"
-Zygote = "0.6.67"
 julia = "1.9"

benchmark/benchmarks.jl

@@ -44,10 +44,10 @@ coords = [c1, c2]
 dr = vector(c1, c2, boundary)
 b1 = HarmonicBond(k=100_000.0u"kJ * mol^-1 * nm^-2", r0=0.6u"nm")
 
-SUITE["interactions"]["LennardJones force" ] = @benchmarkable force($(LennardJones()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary))
-SUITE["interactions"]["LennardJones energy"] = @benchmarkable potential_energy($(LennardJones()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary))
-SUITE["interactions"]["Coulomb force"      ] = @benchmarkable force($(Coulomb()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary))
-SUITE["interactions"]["Coulomb energy"     ] = @benchmarkable potential_energy($(Coulomb()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary))
+SUITE["interactions"]["LennardJones force" ] = @benchmarkable force($(LennardJones()), $(dr), $(a1), $(a1))
+SUITE["interactions"]["LennardJones energy"] = @benchmarkable potential_energy($(LennardJones()), $(dr), $(a1), $(a1))
+SUITE["interactions"]["Coulomb force"      ] = @benchmarkable force($(Coulomb()), $(dr), $(a1), $(a1))
+SUITE["interactions"]["Coulomb energy"     ] = @benchmarkable potential_energy($(Coulomb()), $(dr), $(a1), $(a1))
 SUITE["interactions"]["HarmonicBond force" ] = @benchmarkable force($(b1), $(c1), $(c2), $(boundary))
 SUITE["interactions"]["HarmonicBond energy"] = @benchmarkable potential_energy($(b1), $(c1), $(c2), $(boundary))
 
@@ -90,14 +90,14 @@ function test_sim(nl::Bool, parallel::Bool, f32::Bool, gpu::Bool)
     end
 
     if gpu
-        coords = CuArray(deepcopy(f32 ? starting_coords_f32 : starting_coords))
-        velocities = CuArray(deepcopy(f32 ? starting_velocities_f32 : starting_velocities))
-        atoms = CuArray([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
+        coords = CuArray(copy(f32 ? starting_coords_f32 : starting_coords))
+        velocities = CuArray(copy(f32 ? starting_velocities_f32 : starting_velocities))
+        atoms = CuArray([Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
                               ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms])
     else
-        coords = deepcopy(f32 ? starting_coords_f32 : starting_coords)
-        velocities = deepcopy(f32 ? starting_velocities_f32 : starting_velocities)
-        atoms = [Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
+        coords = copy(f32 ? starting_coords_f32 : starting_coords)
+        velocities = copy(f32 ? starting_velocities_f32 : starting_velocities)
+        atoms = [Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
                       ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms]
     end
 

docs/make.jl

@@ -7,6 +7,7 @@ makedocs(
         prettyurls=(get(ENV, "CI", nothing) == "true"),
         size_threshold_ignore=["api.md"],
     ),
+    modules=[Molly],
     pages=[
         "Home"                      => "index.md",
         "Documentation"             => "documentation.md",

docs/src/api.md

@@ -2,17 +2,23 @@
 
 The API reference can be found here.
 
-Molly also re-exports [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) and [Unitful.jl](https://github.com/PainterQubits/Unitful.jl), making the likes of `SVector` and `1.0u"nm"` available when you call `using Molly`.
+Molly re-exports [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) and [Unitful.jl](https://github.com/PainterQubits/Unitful.jl), making the likes of `SVector` and `1.0u"nm"` available when you call `using Molly`.
 
-The [`visualize`](@ref) function is in a package extension and is only available once you have called `using GLMakie`.
-The [`ASECalculator`](@ref) code is in a package extension and is only available once you have called `using PythonCall`.
+Package extensions are used in order to reduce the number of dependencies:
+- To use [`visualize`](@ref), call `using GLMakie`.
+- To use [`ASECalculator`](@ref), call `using PythonCall`.
+- To use [`rdf`](@ref), call `using KernelDensity`.
+
+## Exported names
 
 ```@index
-Order   = [:module, :type, :constant, :function, :macro]
+Order = [:module, :type, :constant, :function, :macro]
 ```
 
+## Docstrings
+
 ```@autodocs
 Modules = [Molly]
 Private = false
-Order   = [:module, :type, :constant, :function, :macro]
+Order = [:module, :type, :constant, :function, :macro]
 ```

docs/src/development.md

@@ -8,7 +8,7 @@ Various environmental variables can be set to modify the tests:
 - `VISTESTS` determines whether to run the [GLMakie.jl](https://github.com/JuliaPlots/Makie.jl) plotting tests which will error on remote systems where a display is not available, default `VISTESTS=1`.
 - `GPUTESTS` determines whether to run the GPU tests, default `GPUTESTS=1`.
 - `DEVICE` determines which GPU to run the GPU tests on, default `DEVICE=0`.
-- `GROUP` can be used to run a subset of the tests, options `All`/`Protein`/`Zygote`/`NotZygote`, default `GROUP=All`.
+- `GROUP` can be used to run a subset of the tests, options `All`/`Protein`/`Gradients`/`NotGradients`, default `GROUP=All`.
 The CI run does not carry out all tests - for example the GPU tests are not run - and this is reflected in the code coverage.
 
 ## Benchmarks

docs/src/differentiable.md

@@ -7,8 +7,8 @@ In the last few years, the deep learning revolution has broadened to include the
 The concept of using automatic differentiation (AD) to obtain exact gradients through physical simulations has many interesting applications, including parameterising force fields and training neural networks to describe atomic potentials.
 
 There are some projects that explore differentiable molecular simulations - see [Related software](@ref).
-However Julia provides a strong suite of AD tools, with [Zygote.jl](https://github.com/FluxML/Zygote.jl) and [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) allowing source-to-source transformations for much of the language.
-With Molly you can use the power of Zygote and Enzyme to obtain gradients through molecular simulations, even in the presence of complex interactions such as implicit solvation and stochasticity such as Langevin dynamics or the Andersen thermostat.
+However Julia provides a strong suite of AD tools, with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) allowing source-to-source transformations for much of the language.
+With Molly you can use the power of Enzyme to obtain gradients through molecular simulations, even in the presence of complex interactions such as implicit solvation and stochasticity such as Langevin dynamics or the Andersen thermostat.
 Reverse mode AD can be used on the CPU with multithreading and on the GPU; performance is typically within an order of magnitude of the primal run.
 Forward mode AD can also be used on the CPU.
 Pairwise, specific and general interactions work, along with neighbor lists, and the same abstractions for running simulations are used as in the main package.
@@ -44,7 +44,6 @@ end
 Now we can set up and run the simulation in a similar way to that described in the [Molly documentation](@ref).
 The difference is that we wrap the simulation in a `loss` function.
 This returns a single value that we want to obtain gradients with respect to, in this case the difference between the value of the above function at the end of the simulation and a target distance.
-The `Zygote.ignore()` block allows us to ignore code for the purposes of obtaining gradients; you could add the [`visualize`](@ref) function there for example.
 ```julia
 using Zygote
 using Format
@@ -66,8 +65,6 @@ neighbor_finder = DistanceNeighborFinder(
 lj = LennardJones(
     cutoff=DistanceCutoff(1.5),
     use_neighbors=true,
-    force_units=NoUnits,
-    energy_units=NoUnits,
 )
 pairwise_inters = (lj,)
 coords = place_atoms(n_atoms, boundary; min_dist=0.6)
@@ -78,8 +75,8 @@ simulator = VelocityVerlet(
 )
 
 function loss(σ, coords, velocities)
-    atoms = [Atom(0, 0.0, atom_mass, σ, 0.2, false) for i in 1:n_atoms]
-    loggers = (coords=CoordinateLogger(Float64, 10),)
+    atoms = [Atom(0, 0, atom_mass, 0.0, σ, 0.2) for i in 1:n_atoms]
+    loggers = (coords=CoordinatesLogger(Float64, 10),)
 
     sys = System(
         atoms=atoms,
@@ -98,10 +95,8 @@ function loss(σ, coords, velocities)
     mms_end = mean_min_separation(Array(sys.coords), boundary)
     loss_val = abs(mms_end - dist_true)
 
-    Zygote.ignore() do
-        printfmt("σ {:6.3f}  |  Mean min sep expected {:6.3f}  |  Mean min sep end {:6.3f}  |  Loss {:6.3f}  |  ",
-                  σ, σ * (2 ^ (1 / 6)), mms_end, loss_val)
-    end
+    printfmt("σ {:6.3f}  |  Mean min sep expected {:6.3f}  |  Mean min sep end {:6.3f}  |  Loss {:6.3f}  |  ",
+             σ, σ * (2 ^ (1 / 6)), mms_end, loss_val)
 
     return loss_val
 end
@@ -186,8 +181,8 @@ simulator = VelocityVerlet(
 )
 
 function loss(θ)
-    atoms = [Atom(0, 0.0, atom_mass, 0.0, 0.0, false) for i in 1:n_atoms]
-    loggers = (coords=CoordinateLogger(Float64, 2),)
+    atoms = [Atom(0, 0, atom_mass, 0.0, 0.0, 0.0) for i in 1:n_atoms]
+    loggers = (coords=CoordinatesLogger(Float64, 2),)
     specific_inter_lists = (
         InteractionList2Atoms(
             [1, 2, 4, 5],
@@ -204,9 +199,9 @@ function loss(θ)
 
     sys = System(
         atoms=atoms,
-        coords=deepcopy(coords),
+        coords=copy(coords),
         boundary=boundary,
-        velocities=deepcopy(velocities),
+        velocities=copy(velocities),
         specific_inter_lists=specific_inter_lists,
         loggers=loggers,
         force_units=NoUnits,
@@ -220,10 +215,8 @@ function loss(θ)
     dist_end = 0.5 * (d1 + d2)
     loss_val = abs(dist_end - dist_true)
 
-    Zygote.ignore() do
-        printfmt("θ {:5.1f}°  |  Final dist {:4.2f}  |  Loss {:5.3f}  |  ",
-                 rad2deg(θ), dist_end, loss_val)
-    end
+    printfmt("θ {:5.1f}°  |  Final dist {:4.2f}  |  Loss {:5.3f}  |  ",
+             rad2deg(θ), dist_end, loss_val)
 
     return loss_val
 end
@@ -278,7 +271,7 @@ The plot of these shows that the gradient has the expected sign either side of t
 
 ## Neural network potentials
 
-Since gradients can be computed with Zygote, [Flux](https://fluxml.ai) models can also be incorporated into simulations.
+[Flux](https://fluxml.ai) models can also be incorporated into simulations.
 Here we show a neural network in the force function, though they can also be used in other parts of the simulation.
 This example also shows how gradients for multiple parameters can be obtained, in this case the parameters of the neural network.
 The jump from single to multiple parameters is important because single parameters can be optimised using finite differencing, whereas differentiable simulation is well-placed to optimise many parameters simultaneously.
@@ -328,15 +321,15 @@ simulator = VelocityVerlet(
 )
 
 function loss()
-    atoms = [Atom(0, 0.0f0, mass, 0.0f0, 0.0f0, false) for i in 1:n_atoms]
-    loggers = (coords=CoordinateLogger(Float32, 10),)
+    atoms = [Atom(0, 0, mass, 0.0f0, 0.0f0, 0.0f0) for i in 1:n_atoms]
+    loggers = (coords=CoordinatesLogger(Float32, 10),)
     general_inters = (NNBonds(),)
 
     sys = System(
         atoms=atoms,
-        coords=deepcopy(coords),
+        coords=copy(coords),
         boundary=boundary,
-        velocities=deepcopy(velocities),
+        velocities=copy(velocities),
         general_inters=general_inters,
         loggers=loggers,
         force_units=NoUnits,
@@ -350,10 +343,8 @@ function loss()
                 norm(vector(sys.coords[3], sys.coords[1], boundary))) / 3
     loss_val = abs(dist_end - dist_true)
 
-    Zygote.ignore() do
-        printfmt("Dist end {:6.3f}  |  Loss {:6.3f}\n", dist_end, loss_val)
-        visualize(sys.loggers.coords, boundary, "sim.mp4"; show_boundary=false)
-    end
+    printfmt("Dist end {:6.3f}  |  Loss {:6.3f}\n", dist_end, loss_val)
+    visualize(sys.loggers.coords, boundary, "sim.mp4"; show_boundary=false)
 
     return loss_val
 end