fix edge cases in atomselect and bondguesser with empty molecules or single atoms

Author: stefdoerr

moleculekit/atomselect/analyze.py

@@ -70,6 +70,22 @@ def analyze(mol: Molecule, bonds, _profile=False):
     from moleculekit.periodictable import periodictable
     import numpy as np
 
+    if mol.numAtoms == 0:
+        empty_bool = np.zeros(0, dtype=bool)
+        empty_uint = np.zeros(0, dtype=np.uint32)
+        return {
+            "waters": empty_bool.copy(),
+            "lipids": empty_bool.copy(),
+            "ions": empty_bool.copy(),
+            "residues": empty_uint.copy(),
+            "protein_bb": empty_bool.copy(),
+            "nucleic_bb": empty_bool.copy(),
+            "protein": empty_bool.copy(),
+            "nucleic": empty_bool.copy(),
+            "fragments": empty_uint.copy(),
+            "sidechain": empty_bool.copy(),
+        }
+
     insertion = np.unique(mol.insertion, return_inverse=True)[1].astype(np.uint32)
     chain_id = np.unique(mol.chain, return_inverse=True)[1].astype(np.uint32)
     seg_id = np.unique(mol.segid, return_inverse=True)[1].astype(np.uint32)

moleculekit/atomselect/atomselect.py

@@ -268,6 +268,12 @@ def fn(x, y):
 
 
 def atomselect(mol, selection, bonds, _debug=False, _analysis=None, _return_ast=False):
+    if mol.numAtoms == 0:
+        mask = np.zeros(0, dtype=bool)
+        if _return_ast:
+            return mask, None
+        return mask
+
     if _analysis is None:
         _analysis = analyze(mol, bonds)
 

moleculekit/molecule.py

@@ -1125,6 +1125,9 @@ def _guessBonds(self, rdkit=False):
         """
         from moleculekit.bondguesser import guess_bonds, guess_bonds_rdkit
 
+        if self.numAtoms == 0:
+            return np.empty((0, 2), dtype=Molecule._dtypes["bonds"])
+
         if self.numFrames == 0:
             raise RuntimeError(
                 "No coordinates found in the molecule. Cannot guess bonds."

tests/test_atomselect.py

@@ -183,6 +183,89 @@ def _test_atomselect(pdbid, sel, _pdbmols):
             pickle.dump(results, f)
 
 
+def _test_empty_molecule():
+    mol = Molecule()
+    selections = [
+        "all",
+        "protein",
+        "nucleic",
+        "water",
+        "lipid",
+        "ion",
+        "backbone",
+        "sidechain",
+        "hydrogen",
+        "noh",
+        "name CA",
+        "resname ALA",
+        "resid 1",
+        "chain A",
+        "index 0",
+        "serial 1",
+        "element C",
+        "mass < 5",
+        "x < 6",
+        "beta >= 0",
+        "not protein",
+    ]
+    for sel in selections:
+        res = mol.atomselect(sel)
+        assert res.shape == (0,), f"Expected empty result for '{sel}', got shape {res.shape}"
+        assert res.dtype == bool, f"Expected bool dtype for '{sel}', got {res.dtype}"
+
+
+def _test_single_atom_molecule():
+    mol = Molecule()
+    mol.empty(1)
+    mol.record[:] = "ATOM"
+    mol.name[:] = "CA"
+    mol.resname[:] = "ALA"
+    mol.resid[:] = 1
+    mol.chain[:] = "A"
+    mol.element[:] = "C"
+    mol.coords = np.zeros((1, 3, 1), dtype=np.float32)
+
+    expected_true = [
+        "all",
+        "name CA",
+        "resname ALA",
+        "resid 1",
+        "chain A",
+        "element C",
+        "index 0",
+        "serial 1",
+        "noh",
+        "not nucleic",
+        "not water",
+        "x < 6",
+        "beta >= 0",
+    ]
+    for sel in expected_true:
+        res = mol.atomselect(sel)
+        assert res.shape == (1,), f"Expected shape (1,) for '{sel}', got {res.shape}"
+        assert res.dtype == bool
+        assert res[0], f"Expected True for '{sel}'"
+
+    expected_false = [
+        "nucleic",
+        "water",
+        "lipid",
+        "ion",
+        "hydrogen",
+        "name CB",
+        "resname GLY",
+        "resid 2",
+        "chain B",
+        "element N",
+        "index 1",
+    ]
+    for sel in expected_false:
+        res = mol.atomselect(sel)
+        assert res.shape == (1,), f"Expected shape (1,) for '{sel}', got {res.shape}"
+        assert res.dtype == bool
+        assert not res[0], f"Expected False for '{sel}'"
+
+
 def _test_numprop_list_equality():
     pdb = os.path.join(curr_dir, "test_atomselect", "test.pdb")
     mol = Molecule(pdb)

tests/test_bondguesser.py

@@ -44,6 +44,31 @@ def _test_bond_guessing(pdbid):
     assert np.array_equal(bonds, bondsref)
 
 
+def _test_zero_atoms():
+    from moleculekit.molecule import Molecule
+    from moleculekit.bondguesser import guess_bonds
+
+    mol = Molecule()
+    bonds = guess_bonds(mol)
+    assert bonds.shape == (0, 2)
+    assert bonds.dtype == np.uint32
+
+
+def _test_single_atom():
+    from moleculekit.molecule import Molecule
+    from moleculekit.bondguesser import guess_bonds
+
+    mol = Molecule()
+    mol.empty(1)
+    mol.element[:] = "C"
+    mol.name[:] = "CA"
+    mol.coords = np.zeros((1, 3, 1), dtype=np.float32)
+
+    bonds = guess_bonds(mol)
+    assert bonds.shape == (0, 2)
+    assert bonds.dtype == np.uint32
+
+
 def _test_solvated_bond_guessing():
     from moleculekit.molecule import Molecule, calculateUniqueBonds
     from moleculekit.bondguesser import guess_bonds