JuliaMolSim / Molly.jl

Molecular simulation in Julia
Other
371 stars 51 forks source link

Allowing SHAKE/RATTLE constraints on GPU simulation #175

Open Omar-Elrefaei opened 2 months ago

Omar-Elrefaei commented 2 months ago

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

testing with the protein simulation in `benchmarks.jl` ```julia 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[bot] commented 2 months ago

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 :warning:
Additional details and impacted files ```diff @@ 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 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

jgreener64 commented 2 months ago

Looks like a great start.

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