Allowing SHAKE/RATTLE constraints on GPU simulation

[Omar-Elrefaei]

As discussed I've adapted the constraint code to allow for GPU simulation.

testing with the protein simulation in `benchmarks.jl`

using Molly, DelimitedFiles, Test
data_dir = normpath("./data")
ff = MolecularForceField(
    joinpath(data_dir, "force_fields", "ff99SBildn.xml"),
    joinpath(data_dir, "force_fields", "tip3p_standard.xml"),
    joinpath(data_dir, "force_fields", "his.xml"),
)

# takes more than 40s
sys_init = System(joinpath(data_dir, "6mrr_equil.pdb"), ff; gpu=false)
sys_init_gpu = System(joinpath(data_dir, "6mrr_equil.pdb"), ff; gpu=true)

function protein(sys_init; is_constrained=false, nsteps=100)
    println("Constraints $is_constrained, GPU $(is_on_gpu(sys_init)), nsteps $nsteps")
    constraints = []
    bond_length = 0.0872u"nm"
    for i in 1171:3:15952
        push!(constraints, DistanceConstraint(i, i + 1, bond_length))
        push!(constraints, DistanceConstraint(i, i + 2, bond_length))
    end
    if is_constrained
        cons = (SHAKE_RATTLE([constraints...], length(sys_init), 1e-4u"nm", 1e-4u"nm^2 * ps^-1"),)
        sys_init = System(sys_init; constraints=cons)
    end

    sim = VelocityVerlet(dt=0.001u"ps")
    @time "Protein Simulation" simulate!(sys_init, sim, nsteps)
end

Note: the last example is not supported and crashes on master

julia> protein(sys_cpu, is_constrained=false);
Constraints false, GPU false 
Protein Simulation: 10.376156 seconds (4.97 k allocations: 623.307 MiB) 

julia> protein(sys_cpu, is_constrained=true); 
 Constraints true, GPU false 
Protein Simulation: 10.236894 seconds (5.17 k allocations: 660.505 MiB)

julia> protein(sys_gpu, is_constrained=false);
Constraints false, GPU true
Protein Simulation:  1.846926 seconds (129.68 k allocations: 198.061 MiB)

julia> protein(sys_gpu, is_constrained=true);
Constraints true, GPU true
Protein Simulation: 1.654842 seconds (242.76 k allocations: 200.673 MiB)

I'm submitting a draft PR, because there is a few concerns I was to attend to first before this is ready.

[codecov]

Codecov Report

Attention: Patch coverage is 45.71429% with 57 lines in your changes are missing coverage. Please review.

Project coverage is 71.81%. Comparing base (b769838) to head (6a7e014).

Files	Patch %	Lines
src/constraints/shake.jl	45.19%	57 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #175      +/-   ##
==========================================
- Coverage   72.55%   71.81%   -0.75%     
==========================================
  Files          36       36              
  Lines        5481     5542      +61     
==========================================
+ Hits         3977     3980       +3     
- Misses       1504     1562      +58     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[jgreener64]

Looks like a great start.

Would be good to change the existing test to run on GPU too.