Retrieving exact topology used for an endpoint in a relative transformation

[slochower]

It is sometimes useful to take an endpoint in a Perses calculation and continue to run vanilla MD. This is especially true if one of the endpoints is causing trouble.

Helpfully, I can see that recent versions of Perses write out the topologies to models/ and the serialized system and state to nan-error-logs/ (I'm not sure if this latter directory is only populated with irksome endpoints, but it's here now and it is appreciated.)

I have no issues loading the system and state, setting up an integrator, and running more steps. But it can often be useful to visualize the trajectory and to do this, I'll have to load a topology too, so that I can use (for example) simulation.reporters.append(PDBReporter(...)). The files models/out_complex_new.pdb and models/out_complex_old.pdb look like they contain the endpoints of the old to new transformation. But if I try to do something like...

system = deserialize_xml("nan-error-logs/iteration0-replica0-state0-system")

pdb = PDBFile("models/out_complex_old.pdb")
simulation = Simulation(pdb.topology, system, ...)
simulation.context.setPositions(pdb.positions)

Then I can't get the number of positions to match between the PDB file and the system. It's close but not exact. Is there something else that has been added to the Perses topology that is not in those PDB files?

[mikemhenry]

@zhang-ivy or @ijpulidos might be able to answer your question, things can get tricky if solvent is included or not as well if models/out_complex_old.pdb is alchemical or not

[mikemhenry]

Also as an FYI, we are adapting perses to run with OpenFE https://github.com/OpenFreeEnergy/openfe and new development is going on here: https://github.com/choderalab/feflow

If your issue does turn out to be a bug, I am happy to accept a PR/fix it and cut a new perses release :)

[slochower]

The potential bug aspect, which I didn't mention in the OP, is that some edges end up crashing with NAN coordinates but seem fine if I (outside of Perses) load the system/state XML as above and minimize before continuing with the integrator. So I'm wondering if minimization is not actually happening through the Perses workflow (by setting, e.g. n_equilibration_iterations: 1000 in the YAML) because the edges crash immediately with the NAN coordinates.

Do you have any additional documentation or examples for the feflow port yet?

[zhang-ivy]

Then I can't get the number of positions to match between the PDB file and the system.

I'm guessing this is because the out_complex_old.pdb has the old topology/positions, whereas the system file contains the hybrid system, but @ijpulidos would know better, since he's the one who wrote the new models/ feature

For now, if you have access to the HybridTopologyFactory object (this should be one of the files generated when you run perses), you can pull out all the relevant topologies/positions/systems with:

# Positions
old_positions = htf.old_positions(hybrid_positions)
new_positions = htf.new_positions(hybrid_positions)
hybrid_positions = htf.hybrid_positions

# Topologies
old_topology = htf._topology_proposal.old_topology # openmm topology
new_topology = htf._topology_proposal.new_topology # openmm topology
hybrid_topology = htf.hybrid_topology # mdtraj topology
hybrid_topology = htf.omm_hybrid_topology # openmm topology

# Systems
old_system = htf._old_system
new_system = htf._new_system
hybrid_system = htf._hybrid_system

@ijpulidos seems like it will be important to make sure all of these objects are dumped into models/ -- I'm not sure what's currently dumped there now

[zhang-ivy]

So I'm wondering if minimization is not actually happening through the Perses workflow (by setting, e.g. n_equilibration_iterations: 1000 in the YAML) because the edges crash immediately with the NAN coordinates.

Are you running repex? Minimization does happen automatically with the repex workflow, but not with the neq switching workflow. That said, the minimization that happens in the repex workflow involves only 100 steps: https://github.com/choderalab/perses/blob/main/perses/samplers/multistate.py#L33 I would recommend changing this to 0 as this will allow minimization to proceed until the default energy tolerance is reached (which is the default openmm scheme: http://docs.openmm.org/development/api-python/generated/openmm.openmm.LocalEnergyMinimizer.html)

[ijpulidos]

I'm guessing this is because the out_complex_old.pdb has the old topology/positions, whereas the system file contains the hybrid system, but @ijpulidos would know better, since he's the one who wrote the new models/ feature

@zhang-ivy You are correct, these PDB files are generated before any hybrid topology is being constructed. They are just generated as convenient points to know what the old/new topologies look like. They are not compatible to the serialized system when a NaN is detected.

@ijpulidos seems like it will be important to make sure all of these objects are dumped into models/ -- I'm not sure what's currently dumped there now

Yes, I agree, we probably need a bit more consistency in the outputs, even though it's not a bug as far as I can tell. At the time it was convenient to fulfill another request (I cannot remember exactly what was it but I can dig it up if needed).

[slochower]

Are you running repex? Minimization does happen automatically with the repex workflow, but not with the neq switching workflow. That said, the minimization that happens in the repex workflow involves only 100 steps: https://github.com/choderalab/perses/blob/main/perses/samplers/multistate.py#L33 I would recommend changing this to 0 as this will allow minimization to proceed until the default energy tolerance is reached (which is the default openmm scheme: http://docs.openmm.org/development/api-python/generated/openmm.openmm.LocalEnergyMinimizer.html)

Yes, repex. Wow, okay, I definitely missed that and thought that equilibration was controlled by the YAML key n_equilibration_iterations. Let me give this a try.

@zhang-ivy You are correct, these PDB files are generated before any hybrid topology is being constructed. They are just generated as convenient points to know what the old/new topologies look like. They are not compatible to the serialized system when a NaN is detected.

Oh, this makes sense.

For now, if you have access to the HybridTopologyFactory object (this should be one of the files generated when you run perses), you can pull out all the relevant topologies/positions/systems with:

Okay, I see. I was able to get this running using both the System and the topology from the HybridTopologyFactory and not using any of the other XML or PDB files. Is something like this what you had in mind?

from openeye import oechem  # Necessary for the pickle...
htf = np.load(args.factory, allow_pickle=True)

# Positions
htf = htf["arr_0"].flatten()[0]["complex"]
hybrid_positions = htf._compute_hybrid_positions()
old_positions = htf.old_positions(hybrid_positions)

# Topologies
old_topology = htf._topology_proposal.old_topology # openmm topology

# Systems
old_system = htf._old_system

_int = integrators.LangevinIntegrator(
    temperature=300 * unit.kelvin, timestep=4 * unit.femtoseconds
)

simulation = Simulation(old_topology, old_system, _int)
simulation.context.setPositions(old_positions)

print(f"Initial energy:" f" {simulation.context.getState(getEnergy=True).getPotentialEnergy()}")
simulation.minimizeEnergy()
print(f"Minimized energy:" f" "
      f"{simulation.context.getState(getEnergy=True).getPotentialEnergy()}")

simulation.reporters.append(PDBReporter('output.pdb', 100))
simulation.reporters.append(StateDataReporter(stdout, 100, step=True,
        potentialEnergy=True, temperature=True))
simulation.step(1000)

[ijpulidos]

@slochower That seems totally reasonable to me, yes, getting the topology/systems from the HTF object rather than the serialized pdb files. Are your endstate simulations running fine with the example script you posted?

[slochower]

Okay, fantastic -- thanks @ijpulidos @zhang-ivy. So far, following @zhang-ivy 's suggestion of setting minimisation_steps to 0 seems to be working!

Can someone point me to what n_equilibration_iterations in the YAML actually controls for repex runs? It looks like it gets propagated here: https://github.com/choderalab/perses/blob/f3aacde3588e5b90e67a21052d914a033d3c087b/perses/app/setup_relative_calculation.py#L846 And then used: https://github.com/choderalab/perses/blob/f3aacde3588e5b90e67a21052d914a033d3c087b/perses/app/setup_relative_calculation.py#L881-L887 Is this actually running equilibration on lambda = 0 and lambda = 1?

[zhang-ivy]

Yay, glad its working!

The second code snippet you shared is for setting up neq switching calculations. n_equilibration_iterations is used here for the repex: https://github.com/choderalab/perses/blob/f3aacde3588e5b90e67a21052d914a033d3c087b/perses/app/setup_relative_calculation.py#L980

hss_run is an instance of the perses HybridRepexSampler, which is a subclass of the openmmtools MultiStateSampler. The equilibrate function is defined here: https://github.com/choderalab/openmmtools/blob/main/openmmtools/multistate/multistatesampler.py#L654

It does look like equilibration is run for all replicas/thermodynamic states.

[zhang-ivy]

Also, we're super glad that you are using perses -- huge apologies that there aren't any docs. Thanks for your patience / happy to help clarify things whenever we can!