check for cg beads, include simulated names for molecules and constit…

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.
 

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()