From ba18b60d0ca9cacb8def6311880e2a2ce9724f1e Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Thu, 6 Aug 2026 13:26:31 +0100 Subject: [PATCH] check for cg beads, include simulated names for molecules and constituent particles --- biosim_extractor/mdanalysis/toptraj.py | 75 +++++++++++++++++--------- tests/test_mdanalysis/test_toptraj.py | 3 ++ 2 files changed, 54 insertions(+), 24 deletions(-) diff --git a/biosim_extractor/mdanalysis/toptraj.py b/biosim_extractor/mdanalysis/toptraj.py index c1a3e43..55efe07 100644 --- a/biosim_extractor/mdanalysis/toptraj.py +++ b/biosim_extractor/mdanalysis/toptraj.py @@ -12,7 +12,7 @@ from MDAnalysis.exceptions import NoDataError from rdkit import Chem -alternative_residue_names = { +ALTERNATIVE_RESIDUE_NAMES = { "ALAD": "ALA", "ASH": "ASP", "ASPH": "ASP", @@ -27,6 +27,28 @@ "LYN": "LYS", } +MARTINI_BEAD_NAMES = [ + "BB", + "SC1", + "SC2", + "SC3", + "SC4", + "SC5", + "NC3", + "PO4", + "GL1", + "GL2", + "C1A", + "D2A", + "C3A", + "C4A", + "C1B", + "C2B", + "C3B", + "C4B", +] +BEAD_SEL_STRING = "name " + " ".join(MARTINI_BEAD_NAMES) + def get_protein_sequence(fragment): """ @@ -47,7 +69,7 @@ def get_protein_sequence(fragment): if len(protein_atoms) > 0: protein_residues = protein_atoms.residues # Handle alternative residue names - for old_name, new_name in alternative_residue_names.items(): + for old_name, new_name in ALTERNATIVE_RESIDUE_NAMES.items(): for residue in protein_residues: if residue.resname == old_name: residue.resname = new_name @@ -201,6 +223,9 @@ def safe_extract(func): "fixed_charges": lambda u: safe_extract( lambda: getattr(u.atoms, "charges", None) is not None ), + "coarse_grained": lambda u: safe_extract( + lambda: len(u.select_atoms(BEAD_SEL_STRING)) > 0 + ), } @@ -213,6 +238,16 @@ def safe_extract(func): "molecular_weight": lambda fragment: float(sum(fragment.masses)) if hasattr(fragment, "masses") else None, + "simulated_particle_names": lambda fragment: ( + ", ".join(list(fragment.atoms.names)) + if len(fragment.resnames) > 0 and len(set(fragment.resnames)) == 1 + else None + ), + "simulated_molecule_name": lambda fragment: ( + fragment.resnames[0] + if len(fragment.resnames) > 0 and len(set(fragment.resnames)) == 1 + else None + ), } @@ -252,6 +287,7 @@ def __init__(self, toppath, trajpath): self.lines = [] self.data = {} self.molecule_types = {} + self.cg = False # ========================= # MAIN ENTRY @@ -288,6 +324,10 @@ def _extract_molecules(self): except NoDataError: fragments = [] + # check if system is coarse-grained + if len(self.u.select_atoms(BEAD_SEL_STRING)) > 0: + self.cg = True + if not fragments: self.molecule_types = {} self.data["molecule_ids"] = {} @@ -303,7 +343,8 @@ def _extract_molecules(self): representatives = {} for fragment in fragments: signature = (tuple(fragment.resnames), tuple(fragment.atoms.names)) - if signature not in representatives: + molecular_weight = float(sum(fragment.masses)) + if signature not in representatives and molecular_weight != 0: representatives[signature] = { "count": molecule_types[signature], "fragment": fragment, # Store the actual AtomGroup @@ -337,10 +378,13 @@ def _find_molecule_IDs(self): molecule_ids[i][key] = func if len(mol_atoms) == 1: - atom = fragment.atoms[0] - molecule_ids[i]["molecular_formula"] = getattr( - atom, "element", atom.name.strip("0123456789") - ) + if not self.cg: + atom = fragment.atoms[0] + molecule_ids[i]["molecular_formula"] = getattr( + atom, "element", atom.name.strip("0123456789") + ) + else: + continue elif len(set(mol_residues)) == 1 and len(mol_atoms) > 1: try: # Attempt RDKit conversion first @@ -370,23 +414,6 @@ def _find_molecule_IDs(self): return self.data -# # ========================= -# # USAGE -# # ========================= -# if __name__ == "__main__": -# topology = sys.argv[1] -# trajectory = sys.argv[2] -# parser = TopTrajParser(topology, trajectory) -# result = parser.parse() - -# print(json.dumps(result, indent=2)) - - -# ========================= -# ENTRY POINT -# ========================= - - def parse_args(): """Parse command-line arguments. diff --git a/tests/test_mdanalysis/test_toptraj.py b/tests/test_mdanalysis/test_toptraj.py index 18f7c1f..95721aa 100644 --- a/tests/test_mdanalysis/test_toptraj.py +++ b/tests/test_mdanalysis/test_toptraj.py @@ -38,6 +38,7 @@ def test_parse_handles_non_callable_toptraj_extract(mock_universe, monkeypatch): def test_find_molecule_ids_atom_formula_fallback(monkeypatch): parser = object.__new__(toptraj.TopTrajParser) parser.data = {} + parser.cg = False atom = SimpleNamespace(name="C12") # no .element -> fallback strips digits frag = MagicMock() frag.atoms = MagicMock() @@ -290,10 +291,12 @@ def test_toptrajparser_find_molecule_ids_peptide_and_atom(mock_universe): atom_fragment.residues = make_mock_residues(1) atom_fragment.atoms[0].element = "H" atom_fragment.atoms[0].name = "H1" + atom_fragment.masses = [1.0] # Peptide fragment peptide_fragment = MagicMock() peptide_fragment.atoms = make_mock_atoms(2, names=["N", "CA"]) peptide_fragment.residues = make_mock_residues(2, resnames=["ALA", "GLY"]) + peptide_fragment.masses = [12.0, 14.0] peptide_fragment.convert_to.return_value = MagicMock() mock_u.atoms.fragments = [atom_fragment, peptide_fragment] mock_trajectory = MagicMock()