Add option to enforce PBC when retrieving coords

[SimonBoothroyd]

Description

This PR adds an option to the HREMD config that enables enforcing PBC when writting to trajectory files.

Supersedes #3

Status

• Ready to go

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.94%. Comparing base (8ba6628) to head (7b23742).

Additional details and impacted files

@@           Coverage Diff           @@
##             main       #8   +/-   ##
=======================================
  Coverage   95.94%   95.94%           
=======================================
  Files          46       46           
  Lines        3008     3010    +2     
=======================================
+ Hits         2886     2888    +2     
  Misses        122      122           

Flag	Coverage Δ
unittests	95.94% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

[tomschultz-roivant]

Is there a reason you're not using the config object like you are above in hremd.py?

[SimonBoothroyd]

hmm yeah this function doesn't have a solo config, but rather a list of configs for each step, so there wasn't really a single good place to put this in a config or without changing the overall func signature that i could see.

[tomschultz-roivant]

Also are there tests that should be added for this?

[SimonBoothroyd]

because this is a more visual thing, i did a test locally by running:

import pathlib

import openmm.app
import openmm.unit
import parmed

import femto.md.config
import femto.md.hremd
import femto.md.simulate

topology = parmed.load_file("argon_start.pdb", structure=True)

system = openmm.System()
for i in range(100):
    system.addParticle(39.948 * openmm.unit.amu)
lj = openmm.NonbondedForce()
lj.setNonbondedMethod(openmm.NonbondedForce.CutoffPeriodic)
lj.setCutoffDistance(8.5 * openmm.unit.angstrom)
for i in range(100):
    lj.addParticle(
        0.0, 3.405 * openmm.unit.angstrom, 0.238 * openmm.unit.kilocalories_per_mole
    )
system.addForce(lj)

integrator = openmm.LangevinIntegrator(
    90.0 * openmm.unit.kelvin,
    1.0 / openmm.unit.picosecond,
    0.002 * openmm.unit.picoseconds,
)
simulation = openmm.app.Simulation(topology.topology, system, integrator)

for wrap in [False, True]:
    simulation.context.setPeriodicBoxVectors(*topology.box_vectors)
    simulation.context.setPositions(topology.positions)

    femto.md.hremd.run_hremd(
        simulation,
        [{}],
        femto.md.config.HREMD(
            n_warmup_steps=0,
            n_cycles=1500,
            n_steps_per_cycle=10,
            trajectory_interval=100,
            trajectory_enforce_pbc=wrap,
        ),
        output_dir=pathlib.Path(f"wrap-{wrap}"),
    )

for wrap in [False, True]:
    final = femto.md.simulate.simulate_state(
        system,
        topology,
        {},
        [
            femto.md.config.Minimization(),
            femto.md.config.Simulation(
                integrator=femto.md.config.LangevinIntegrator(),
                temperature=90.0 * openmm.unit.kelvin,
                pressure=None,
                n_steps=15000,
            ),
        ],
        femto.md.constants.OpenMMPlatform.CPU,
        enforce_pbc=wrap,
    )
    openmm.app.PDBFile.writeFile(
        topology.topology, final.getPositions(), f"final-{wrap}.pdb"
    )

and visualising the final structure / trajectories and everything looked good!

i can try and think of a smoke test if you'd prefer?