Neural Network Potentials about molly.jl

[jgreener64]

Some docs and examples for the experimental differentiable branch are up. Specific interactions now work.

https://juliamolsim.github.io/Molly.jl/latest/differentiable.html

[jgreener64]

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.

[bionicles]

They're working on a big chain rule update for Zygote which may indeed help. Thanks for fixing the docs

[jgreener64]

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.

[ChrisRackauckas]

It should be possible to setup DiffEq for timesteping and then use the adjoint method for the differentiation.

[jgreener64]

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()

[jgreener64]

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.