Molly.jl: AD with Enyzme returns 0 gradient

Specific Domains Modelling & Simulations
, ,
1

I am trying to use Enzyme.jl with Molly.jl to differentiate a MD Simulation with respect to some parameters. Unfortunately my program returns only gradients of magnitude 0.

( This has also been posted to Molly.jl: AD with Enyzme returns 0 gradient · Issue #156 · JuliaMolSim/Molly.jl · GitHub )

MWE

function AD_test_no_unit()
    data_dir = joinpath(dirname(pathof(Molly)), "..", "data") 
    ff_clean = MolecularForceField(
        joinpath(data_dir, "force_fields", "ff99SBildn.xml"),
        joinpath(data_dir, "force_fields", "tip3p_standard.xml"),
        joinpath(data_dir, "force_fields", "his.xml");
        units=false
    )

    sys_clean = System(
        joinpath(data_dir, "6mrr_nowater.pdb"),
        ff_clean;
        units=false,
        gpu=false,
    )

    temp = 298.0
    simulator = Langevin(
        dt=0.001,    #u"ps",
        temperature=temp,
        friction=1.0,    #u"ps^-1",
        coupling=MonteCarloBarostat(1.0, temp, sys_clean.boundary),   #1.0 u"bar"
    )

    # this bond object holds the 2 parameters of interest
    bond_params  = sys_clean.specific_inter_lists[1].inters[1]
    bond_k_orig  = bond_params.k        # k  = 363171.19999999995
    bond_r0_orig = bond_params.r0       # r0 = 0.101

    random_velocities!(sys_clean, temp)

    sys_dirty      = deepcopy(sys_clean)
    sys_dirty_copy = deepcopy(sys_clean)

    #get energy from a run using unaltered parameters
    neighbors = find_neighbors(sys_clean, sys_clean.neighbor_finder; n_threads=8)
    simulate!(sys_clean, simulator, 10)
    sys_clean_energy = total_energy(sys_clean, neighbors)       # energy = 8632.53

    bond_k_alt  = Float64(100)
    bond_r0_alt = Float64(10)

    # run with unaltered parameters. Expect ≈0 loss and ≈0 derivative
    # returns ((0.0, 0.0, nothing, nothing, nothing),)
    grads_orig = autodiff(Reverse, loss, Active, Active(bond_k_orig), Active(bond_r0_orig), Const(sys_dirty), Const(sys_clean_energy), Const(simulator), Const(neighbors))
    sys_dirty = deepcopy(sys_dirty_copy)
    loss_orig = loss(bond_k_orig, bond_r0_orig, sys_dirty, sys_clean_energy, simulator, neighbors) # loss = 12.44
    sys_dirty = deepcopy(sys_dirty_copy)


    # run with   altered parameters. Expect high loss and high derivative
    # returns ((0.0, 0.0, nothing, nothing, nothing),)
    grads_alt = autodiff(Reverse, loss, Active, Active(bond_k_alt), Active(bond_r0_alt),   Const(sys_dirty), Const(sys_clean_energy), Const(simulator), Const(neighbors))
    sys_dirty = deepcopy(sys_dirty_copy)
    loss_alt  = loss(bond_k_alt, bond_r0_alt, sys_dirty, sys_clean_energy, simulator, neighbors)  # loss = 4907.96
    
    return grads_orig, loss_orig, grads_alt, loss_alt
end

function loss(k, r0, sys_dirty, sys_clean_energy, simulator, neighbors)
    sys_dirty.specific_inter_lists[1].inters[1] = HarmonicBond{Float64, Float64}(k, r0)
    simulate!(sys_dirty, simulator, 10; n_threads=1, run_loggers=false)
    return abs(total_energy(sys_dirty, neighbors) - sys_clean_energy)  
end

AD_test_no_unit()

No matter which values for the parameters i choose, i always get the same output.
How would you make this work, or is there maybe a better way to get these gradients using Enzyme? (I’m trying Zygote.jl soon if i can’t get this to work)

Additional info

(BiochemicalAlgorithms) pkg> status Enzyme
Project BiochemicalAlgorithms v0.1.0
⌅ [7da242da] Enzyme v0.11.2
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated`

(BiochemicalAlgorithms) pkg> status Molly
Project BiochemicalAlgorithms v0.1.0
  [aa0f7f06] Molly v0.18.1

I’m new to Automatic Differentiation, Enzyme and Molly, so it’s possible i am doing something fundamentally wrong. I also tried out ForwardMode of Enzyme to no avail either. With larger files (eg “6mrr_equil.pdb”) i have been getting the TypeAnalysisDepthLimit warning, but not with this file (“6mrr_nowater.pdb”)

1 Like
2

That is an extremely old version of Enzyme (the latest is 0.11.10). It looks like Molly is version locked to the old one for some reason. @jgreener64 anythinf stopping from using the latest Enzyme?

2 Likes
3

Using Enzyme to take gradients through a whole simulation isn’t tested yet so I’m surprised that doesn’t hit a hard error.

Currently the recommended approach is to use Zygote in user code, see Differentiable simulation · Molly.jl for examples. Enzyme is used internally via chain rules for the more intensive bits. The plan is to switch to using Enzyme throughout once a couple more features are added, then your code should work. This is exciting as it will allow a complete internal rewrite of the simulators, since Enzyme allows mutation and hence better memory usage.

I’ll check the latest Enzyme tomorrow, but some issues have prevented updating compat to date.

3 Likes
4

Thanks for the reply :slight_smile: It’s good to know Enzyme won’t work for a full simulation, as this is what my Thesis is going to be about

5

Oh FYI the reason you’re getting zero above is because you’re marking temporary storage locations used in the derivative computation as constant (forcing Enzyme to think there is no propagation between it and input/outputs). Later Enzyme versions should throw errors more aggressively.

1 Like