Open bionicles opened 4 years ago
Would it please be possible to add an example of neural network as the simulator?
I intend to put some examples of making custom potentials in the docs soon. Having a call to a neural network would work just fine.
However training a neural network by autodifferentiating back through the simulation is more problematic; I was playing around with it last night and running into problems because currently we mutate a lot of arrays.
Hopefully mutation support in Zygote will improve soon and then it will be possible. I did also try ReverseDiff but ran into some typing problems.
Long term it is a priority as it was one of the motivations of making this package.
Also, how can we make a nice Gif?
Now the link to the docs works you can find that information there.
They're working on a big chain rule update for Zygote which may indeed help. Thanks for fixing the docs
Yeah there is lots of exciting development going on.
I did manage to refactor the code to remove mutation and get autodiff working with Zygote, but there was a performance hit. I'll keep looking at that.
It should be possible to setup DiffEq for timesteping and then use the adjoint method for the differentiation.
I have a prototype working on the differentiable
branch. The following code dump repeatedly runs a simulation of 50 atoms for 500 steps and optimises the Lennard-Jones σ value to match a desired mean minimum separation of atoms at the end of the simulation. It uses a neighbour list and PBCs, and gets a low loss in 25 epochs.
EDIT: this is out of date now, see link to docs below.
using Molly
using Zygote
function meanminseparation(final_coords, box_size)
n_atoms = length(final_coords)
sum_dists = 0.0
for i in 1:n_atoms
min_dist = 100.0
for j in 1:n_atoms
i == j && continue
dist = sqrt(square_distance(i, j, final_coords, box_size))
min_dist = min(dist, min_dist)
end
sum_dists += min_dist
end
return sum_dists / n_atoms
end
dist_true = 1.0
σtrue = dist_true / (2 ^ (1 / 6))
n_atoms = 50
mass = 10.0
box_size = 5.0
coords = [box_size .* rand(SVector{3}) for i in 1:n_atoms]
temperature = 0.1
velocities = [velocity(mass, temperature) .* 0.0 for i in 1:n_atoms]
general_inters = Dict{String, GeneralInteraction}("LJ" => LennardJones(true))
neighbour_finder = DistanceNeighbourFinder(trues(n_atoms, n_atoms), 20, 2.0)
function loss(σ)
s = Simulation{typeof(coords)}(
VelocityVerlet(),
[Atom("", "", 0, "", 0.0, mass, σ, 0.2) for i in 1:n_atoms],
Dict{String, Vector{SpecificInteraction}}(),
general_inters,
coords,
velocities,
temperature,
box_size,
Tuple{Int, Int}[],
neighbour_finder,
NoThermostat(),
Logger[],
0.05,
10,
[0]
)
mms_start = meanminseparation(coords, box_size)
c = simulate!(s, 500, parallel=false)
mms_end = meanminseparation(c, box_size)
l = abs(mms_end - dist_true)
println("σ ", round(σ, digits=3))
println("mean min sep expected ", round(σ * (2 ^ (1 / 6)), digits=3))
println("mean min sep start ", round(mms_start, digits=3))
println("mean min sep end ", round(mms_end, digits=3))
println("loss ", round(l, digits=3))
return l
end
grad = gradient(loss, σtrue)[1]
# Simple training loop
function train()
σlearn = 0.8 / (2 ^ (1 / 6))
for epoch_n in 1:25
println("Epoch ", epoch_n)
grad = gradient(loss, σlearn)[1]
σlearn -= grad * 1e-2
println()
end
return σlearn
end
σlearn = train()
The integration steps were okay to implement - the harder bit was changing all the devectorised code in the force calculation that made it fast before. I have found there is a performance/memory hit in making it Zygote-friendly.
The next step is to work out how to code specific interactions in this scheme, e.g. a covalent bond between two specific atoms. I was hoping I could use a sparse matrix constructor to turn the results of a broadcast over bonds into a standard dense vector, but I haven't got that working with Zygote yet.
Also I'd like to look at using forward-mode autodiff as the memory requirements don't scale so badly with the length of the simulation.
Some docs and examples for the experimental differentiable
branch are up. Specific interactions now work.
https://juliamolsim.github.io/Molly.jl/latest/differentiable.html
Would it please be possible to add an example of neural network as the simulator?
Also, how can we make a nice Gif?