Molly.jl: AD with Enyzme returns 0 gradient

[pushingPulling]

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 was also posted on https://discourse.julialang.org/t/molly-jl-ad-with-enyzme-returns-0-gradient/105835 , but i thought it wouldn't hurt to post it here too. )

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

[jgreener64]

See discussion at https://discourse.julialang.org/t/molly-jl-ad-with-enyzme-returns-0-gradient/105835.