Easy access to non-hybrid topologies for endstates

[tlhr]

Is there an easy way to get access to a clean representation (non-hybrid) of the ligand trajectory at lambda[0] and lambda[-1]? Maybe I'm overcomplicating the indexing procedure, but actually it seems quite tedious.

In an ideal world I would envision something like an extra flag for the reader, e.g.:

u = mda.Universe('hybrid_system.pdb', 'simulation.nc', format=FEReader, state_id=0, hybrid=False)

This would make subsequent analyses a lot easier!

[tlhr]

I just realized that I was indeed overcomplicating the indexing, it seems that the lambda state is stored in the bfactor field (with 0.25, 0.5, 0.75 corresponding to state A, both states, and state B respectively). So getting a clean trajectory just involves something like this:

atom_mapping_b = np.array([atom.ix for atom in u.atoms if atom.bfactor in (0.5, 0.75)])
state_b = u.atoms[atom_mapping_b]

Leaving this open to see what the feelings on a flag are, otherwise this can be closed!

[IAlibay]

Ah yes, thanks for raising this @tlhr - the bfactor hack is.. well kind of a hack, we probably need to find a better long term solution to this problem and/or make sure it's well documented!

[nffaruk]

I wish I saw @tlhr's solution earlier, I found out about the bfactors from scouring the source, so +1 to including something in the readme. One remaining issue is that if you allow element change in the mapping, the atom type of State B will be the atom type of State A at the points of element change once you separate the hybrid topology, instead of the original atom type of State B.

My solution has been to load the transformation json (that you provide to quickrun) and extract the mapping info to recover the atom type:

import mdtraj as md
from rdkit import Chem

# load the separated topology trajectory of State B after dumping with mda/openfe_analysis
traj_B = md.load("traj_B.nc", top="traj_B.pdb")

with open("transform.json") as f:
    settings = json.load(f)

# Get atom types of original topologies
periodic_table = Chem.GetPeriodicTable()
compA_atoms = settings["mapping"]["componentA"]["atoms"]
compB_atoms = settings["mapping"]["componentB"]["atoms"]
symbolsA = [periodic_table.GetElementSymbol(a[0]) for a in compA_atoms]
symbolsB = [periodic_table.GetElementSymbol(a[0]) for a in compB_atoms]

# Replace where element change occurs in mapping
for i, idxs in enumerate(settings["mapping"]["componentA_to_componentB"].items()):    
    symA = symbolsA[int(idxs[0])]
    symB = symbolsB[int(idxs[1])]

    if symA != symB:
        elemB = md.element.get_by_symbol(symB)
        traj_B.top.atom(i).element = elemB
        traj_B.top.atom(i).name = symB

traj_B.save("traj_B.nc")
traj_B[0].save("traj_B.pdb")

An offset would be needed in the indexing if the protein is included in the trajectory