diff --git a/vermouth/pdb/pdb.py b/vermouth/pdb/pdb.py index c7bec80b6..433c04ef3 100644 --- a/vermouth/pdb/pdb.py +++ b/vermouth/pdb/pdb.py @@ -61,6 +61,7 @@ def __init__(self, exclude=('SOL',), ignh=False, modelidx=1): self.ignh = ignh self.modelidx = modelidx self._skipahead = False + self.cryst = {} def dispatch(self, line): """ @@ -151,7 +152,6 @@ def _skip(line, lineno=0): site = _skip # CRYSTALLOGRAPHIC AND COORDINATE TRANSFORMATION SECTION - cryst1 = _skip origx1 = _skip origx2 = _skip origx3 = _skip @@ -265,6 +265,38 @@ def _atom(self, line, lineno=0): atom = _atom hetatm = _atom + def cryst1(self, line, lineno=0): + """ + Parse the CRYST1 record. Crystal structure information are stored with + the parser object and may be extracted later. + """ + fields = [ + ('', str, 6), + ('a', float, 9), + ('b', float, 9), + ('c', float, 9), + ('alpha', float, 7), + ('beta', float, 7), + ('gamma', float, 7), + ('space_group', str, 11), + ('z_value', int, 4), + ] + start = 0 + field_slices = [] + for name, type_, width in fields: + if name == "space_group": + start+=1 + if name: + field_slices.append((name, type_, slice(start, start + width))) + start += width + + for name, type_, slice_ in field_slices: + value = line[slice_].strip() + if value: + self.cryst[name] = type_(value) + else: + LOGGER.warning(f"CRYST1 directive incomplete. Missing entry for {name}.") + def model(self, line, lineno=0): """ Parse a MODEL record. If the model is not the same as :attr:`modelidx`, diff --git a/vermouth/tests/pdb/test_read_pdb.py b/vermouth/tests/pdb/test_read_pdb.py index 006609af3..4088cb7ea 100644 --- a/vermouth/tests/pdb/test_read_pdb.py +++ b/vermouth/tests/pdb/test_read_pdb.py @@ -155,7 +155,6 @@ def test_single_model(pdbstr, ignh, nnodesnedges): assert len(mol.nodes) == nnodes assert len(mol.edges) == nedges - @pytest.mark.parametrize('ignh', [True, False]) @pytest.mark.parametrize('modelidx', range(1, 16)) def test_integrative(ignh, modelidx): @@ -233,3 +232,34 @@ def test_atom_attributes(): assert np.allclose(n_attrs[n_idx][attr], mol.nodes[n_idx][attr]) else: assert n_attrs[n_idx][attr] == mol.nodes[n_idx][attr] + + +@pytest.mark.parametrize('pdbstr, cryst_dict', ( + # complete directive + ('''CRYST1 77.987 77.987 77.987 90.00 90.00 90.00 P 1 1 + MODEL 1 + ATOM 1 EO PEO 0 74.550 37.470 22.790 1.00 0.00 + ATOM 2 EO PEO 1 77.020 38.150 25.000 1.00 0.00 + ATOM 3 EO PEO 2 76.390 37.180 28.130 1.00 0.00 + ATOM 4 EO PEO 3 75.430 37.920 31.450 1.00 0.00 + ''', + {"a": 77.987, "b": 77.987, "c": 77.987, + "alpha": 90.0, "beta": 90.0, "gamma": 90, + "space_group": "P 1", "z_value": 1} + ), + # incomplete directive + ('''CRYST1 77.987 77.987 77.987 + MODEL 1 + ATOM 1 EO PEO 0 74.550 37.470 22.790 1.00 0.00 + ATOM 2 EO PEO 1 77.020 38.150 25.000 1.00 0.00 + ATOM 3 EO PEO 2 76.390 37.180 28.130 1.00 0.00 + ATOM 4 EO PEO 3 75.430 37.920 31.450 1.00 0.00 + ''', + {"a": 77.987, "b": 77.987, "c": 77.987,} + ))) +def test_cryst1(caplog, pdbstr, cryst_dict): + parser = PDBParser() + mols = list(parser.parse(pdbstr.splitlines())) + assert parser.cryst == cryst_dict + if len(cryst_dict) < 8: + assert any(rec.levelname == 'WARNING' for rec in caplog.records)