Repository: biotite-dev/biotite

Issue:
Add function to parse secondary structure elements from PDBx
As MMTF is now deprecated (#525) another method is desired to parse the secondary structure from PDB. PDBx files contain annotations of helices and sheets in the `struct_conf` and `struct_sheet_range` categories respectively.

Hence, a `pdbx.get_sse()` function could be implemented.

Hints from issue comments:
@ceziegler: Is there a preferred format for the output?@padix-key: I think a dictionary mapping the chain ID to the secondary structure elements (SSE) would be good. I think the SSE should be an ndarray
- either using letters like in https://www.biotite-python.org/apidoc/biotite.structure.annotate_sse.html
- or `IntEnum` values of a new `IntEnum` class, as e.g in https://www.biotite-python.org/apidoc/biotite.structure.BondType.html

Residues without defined secondary structure could get a `"c"` or some `COIL` value from the enum value assigned@padix-key: Thanks by the way for tackling this feature! I assigned the issue to you.@padix-key: Hi, are you still planning to implement this feature? If not, this is also fine, but I would free the issue up again.@ceziegler: Hi yes, I'll finish this today! @padix-key: Implemented by #710

Solution patch:
diff --git a/doc/apidoc.json b/doc/apidoc.json
index 39bc6bb5..432cd6c6 100644
--- a/doc/apidoc.json
+++ b/doc/apidoc.json
@@ -418,7 +418,8 @@
             "get_component",
             "set_component",
             "list_assemblies",
-            "get_assembly"
+            "get_assembly",
+            "get_sse"
         ],
         "CIF format" : [
             "CIFFile",
diff --git a/src/biotite/structure/io/pdbx/convert.py b/src/biotite/structure/io/pdbx/convert.py
index 048bfdc9..1f9dcd4d 100644
--- a/src/biotite/structure/io/pdbx/convert.py
+++ b/src/biotite/structure/io/pdbx/convert.py
@@ -3,7 +3,7 @@
 # information.
 
 __name__ = "biotite.structure.io.pdbx"
-__author__ = "Fabrice Allain, Patrick Kunzmann"
+__author__ = "Fabrice Allain, Patrick Kunzmann, Cheyenne Ziegler"
 __all__ = [
     "get_sequence",
     "get_model_count",
@@ -13,6 +13,7 @@ __all__ = [
     "set_component",
     "list_assemblies",
     "get_assembly",
+    "get_sse",
 ]
 
 import itertools
@@ -162,8 +163,8 @@ def get_sequence(pdbx_file, data_block=None):
     -------
     sequence_dict : Dictionary of Sequences
         Dictionary keys are derived from ``entity_poly.pdbx_strand_id``
-        (often equivalent to chain_id and atom_site.auth_asym_id
-        in most cases). Dictionary values are sequences.
+        (equivalent to ``atom_site.auth_asym_id``).
+        Dictionary values are sequences.
 
     Notes
     -----
@@ -219,9 +220,7 @@ def get_model_count(pdbx_file, data_block=None):
         The number of models.
     """
     block = _get_block(pdbx_file, data_block)
-    return len(
-        _get_model_starts(block["atom_site"]["pdbx_PDB_model_num"].as_array(np.int32))
-    )
+    return len(np.unique((block["atom_site"]["pdbx_PDB_model_num"].as_array(np.int32))))
 
 
 def get_structure(
@@ -322,13 +321,12 @@ def get_structure(
         raise InvalidFileError("Missing 'atom_site' category in file")
 
     models = atom_site["pdbx_PDB_model_num"].as_array(np.int32)
-    model_starts = _get_model_starts(models)
-    model_count = len(model_starts)
+    model_count = len(np.unique(models))
     atom_count = len(models)
 
     if model is None:
         # For a stack, the annotations are derived from the first model
-        model_atom_site = _filter_model(atom_site, model_starts, 1)
+        model_atom_site = _filter_model(atom_site, 1)
         # Any field of the category would work here to get the length
         model_length = model_atom_site.row_count
         atoms = AtomArrayStack(model_count, model_length)
@@ -374,7 +372,7 @@ def get_structure(
                 f"the given model {model} does not exist"
             )
 
-        model_atom_site = _filter_model(atom_site, model_starts, model)
+        model_atom_site = _filter_model(atom_site, model)
         # Any field of the category would work here to get the length
         model_length = model_atom_site.row_count
         atoms = AtomArray(model_length)
@@ -715,21 +713,26 @@ def _filter_altloc(array, atom_site, altloc):
         raise ValueError(f"'{altloc}' is not a valid 'altloc' option")
 
 
-def _get_model_starts(model_array):
-    """
-    Get the start index for each model in the arrays of the
-    ``atom_site`` category.
-    """
-    _, indices = np.unique(model_array, return_index=True)
-    indices.sort()
-    return indices
-
-
-def _filter_model(atom_site, model_starts, model):
+def _filter_model(atom_site, model):
     """
     Reduce the ``atom_site`` category to the values for the given
     model.
+
+    Parameters
+    ----------
+    atom_site : CIFCategory or BinaryCIFCategory
+        ``atom_site`` category containing all models.
+    model : int
+        The model to be selected.
+
+    Returns
+    -------
+    atom_site : CIFCategory or BinaryCIFCategory
+        The ``atom_site`` category containing only the selected model.
     """
+    models = atom_site["pdbx_PDB_model_num"].as_array(np.int32)
+    _, model_starts = np.unique(models, return_index=True)
+    model_starts.sort()
     # Append exclusive stop
     model_starts = np.append(model_starts, [atom_site.row_count])
     # Indexing starts at 0, but model number starts at 1
@@ -1467,6 +1470,7 @@ def list_assemblies(pdbx_file, data_block=None):
 
     Examples
     --------
+
     >>> import os.path
     >>> file = CIFFile.read(os.path.join(path_to_structures, "1f2n.cif"))
     >>> assembly_ids = list_assemblies(file)
@@ -1765,3 +1769,117 @@ def _convert_string_to_sequence(string, stype):
         return None
     else:
         raise InvalidFileError("mmCIF _entity_poly.type unsupported type: " + stype)
+
+
+def get_sse(pdbx_file, data_block=None, match_model=None):
+    """
+    Get the secondary structure from a PDBx file.
+
+    Parameters
+    ----------
+    pdbx_file : CIFFile or CIFBlock or BinaryCIFFile or BinaryCIFBlock
+        The file object.
+        The following categories are required:
+
+        - ``entity_poly``
+        - ``struct_conf`` (if alpha-helices are present)
+        - ``struct_sheet_range`` (if beta-strands are present)
+        - ``atom_site`` (if `match_model` is set)
+
+    data_block : str, optional
+        The name of the data block.
+        Default is the first (and most times only) data block of the
+        file.
+        If the data block object is passed directly to `pdbx_file`,
+        this parameter is ignored.
+    match_model : None, optional
+        If a model number is given, only secondary structure elements for residues are
+        kept, that are resolved in the given model.
+        This means secondary structure elements for residues that would not appear
+        in a corresponding :class:`AtomArray` from :func:`get_structure()` are removed.
+        By default, all residues in the sequence are kept.
+
+    Returns
+    -------
+    sse_dict : dict of str -> ndarray, dtype=str
+        The dictionary maps the chain ID (derived from ``auth_asym_id``) to the
+        secondary structure of the respective chain.
+
+        - ``"a"``: alpha-helix
+        - ``"b"``: beta-strand
+        - ``"c"``: coil or not an amino acid
+
+        Each secondary structure element corresponds to the ``label_seq_id`` of the
+        ``atom_site`` category.
+        This means that the 0-th position of the array corresponds to the residue
+        in ``atom_site`` with ``label_seq_id`` ``1``.
+
+    Examples
+    --------
+
+    >>> import os.path
+    >>> file = CIFFile.read(os.path.join(path_to_structures, "1aki.cif"))
+    >>> sse = get_sse(file, match_model=1)
+    >>> print(sse)
+    {'A': array(['c', 'c', 'c', 'c', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a',
+                 'a', 'c', 'c', 'c', 'c', 'c', 'a', 'a', 'a', 'c', 'c', 'a', 'a',
+                 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'c', 'c', 'c',
+                 'c', 'c', 'c', 'b', 'b', 'b', 'c', 'c', 'c', 'c', 'c', 'b', 'b',
+                 'b', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c',
+                 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c', 'c',
+                 'c', 'a', 'a', 'a', 'a', 'a', 'c', 'c', 'c', 'c', 'a', 'a', 'a',
+                 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'c', 'c', 'a',
+                 'a', 'a', 'a', 'c', 'a', 'a', 'a', 'a', 'a', 'a', 'c', 'c', 'c',
+                 'c', 'c', 'a', 'a', 'a', 'a', 'c', 'c', 'c', 'c', 'c', 'c'],
+                 dtype='<U1')}
+
+    If only secondary structure elements for resolved residues are requested, the length
+    of the returned array matches the number of peptide residues in the structure.
+
+    >>> file = CIFFile.read(os.path.join(path_to_structures, "3o5r.cif"))
+    >>> print(len(get_sse(file, match_model=1)["A"]))
+    128
+    >>> atoms = get_structure(file, model=1)
+    >>> atoms = atoms[filter_amino_acids(atoms) & (atoms.chain_id == "A")]
+    >>> print(get_residue_count(atoms))
+    128
+    """
+    block = _get_block(pdbx_file, data_block)
+
+    # Init all chains with "c" for coil
+    sse_dict = {
+        chain_id: np.repeat("c", len(sequence))
+        for chain_id, sequence in get_sequence(block).items()
+    }
+
+    # Populate SSE arrays with helices and strands
+    for sse_symbol, category_name in [
+        ("a", "struct_conf"),
+        ("b", "struct_sheet_range"),
+    ]:
+        if category_name in block:
+            category = block[category_name]
+            chains = category["beg_auth_asym_id"].as_array(str)
+            start_positions = category["beg_label_seq_id"].as_array(int)
+            end_positions = category["end_label_seq_id"].as_array(int)
+
+            # set alpha helix positions
+            for chain, start, end in zip(chains, start_positions, end_positions):
+                # Translate the 1-based positions from PDBx into 0-based array indices
+                sse_dict[chain][start - 1 : end] = sse_symbol
+
+    if match_model is not None:
+        model_atom_site = _filter_model(block["atom_site"], match_model)
+        chain_ids = model_atom_site["auth_asym_id"].as_array(str)
+        res_ids = model_atom_site["label_seq_id"].as_array(int, masked_value=-1)
+        # Filter out masked residues, i.e. residues not part of a chain
+        mask = res_ids != -1
+        chain_ids = chain_ids[mask]
+        res_ids = res_ids[mask]
+        for chain_id, sse in sse_dict.items():
+            res_ids_in_chain = res_ids[chain_ids == chain_id]
+            # Transform from 1-based residue ID to 0-based index
+            indices = np.unique(res_ids_in_chain) - 1
+            sse_dict[chain_id] = sse[indices]
+
+    return sse_dict


Test patch:
diff --git a/tests/structure/io/test_pdbx.py b/tests/structure/io/test_pdbx.py
index 2b80ecfd..ceda90de 100644
--- a/tests/structure/io/test_pdbx.py
+++ b/tests/structure/io/test_pdbx.py
@@ -577,6 +577,58 @@ def test_get_sequence(format):
     assert type(sequences_2["A"]) is seq.NucleotideSequence
 
 
+def test_get_sse():
+    """
+    Check if the secondary structure elements are returned for a short structure
+    where the reference secondary structure elements are taken from Mol*
+    (https://www.rcsb.org/3d-view/1AKI).
+    """
+    TOTAL_LENGTH = 129
+    # Ranges are given in terms of residue IDs, and the end is inclusive
+    HELIX_RANGES = [
+        (5, 14),
+        (20, 22),
+        (25, 36),
+        (80, 84),
+        (89, 101),
+        (104, 107),
+        (109, 114),
+        (120, 123),
+    ]
+    SHEET_RANGES = [
+        (43, 45),
+        (51, 53),
+    ]
+
+    ref_sse = np.full(TOTAL_LENGTH, "c", dtype="U1")
+    for helix_range in HELIX_RANGES:
+        # Conver to zero-based indexing
+        ref_sse[helix_range[0] - 1 : helix_range[1]] = "a"
+    for sheet_range in SHEET_RANGES:
+        ref_sse[sheet_range[0] - 1 : sheet_range[1]] = "b"
+
+    pdbx_file = pdbx.BinaryCIFFile.read(join(data_dir("structure"), "1aki.bcif"))
+    test_sse = pdbx.get_sse(pdbx_file)["A"]
+
+    assert test_sse.tolist() == ref_sse.tolist()
+
+
+@pytest.mark.parametrize("path", glob.glob(join(data_dir("structure"), "*.bcif")))
+def test_get_sse_length(path):
+    """
+    If `match_model` is set in :func:`get_sse()`, the length of the returned array
+    must match the number of residues in the structure.
+    """
+    pdbx_file = pdbx.BinaryCIFFile.read(path)
+    atoms = pdbx.get_structure(pdbx_file, model=1)
+    atoms = atoms[struc.filter_amino_acids(atoms)]
+    sse = pdbx.get_sse(pdbx_file, match_model=1)
+
+    for chain_id in np.unique(atoms.chain_id):
+        chain = atoms[atoms.chain_id == chain_id]
+        assert len(sse[chain_id]) == struc.get_residue_count(chain)
+
+
 def test_bcif_encoding():
     """
     Check if encoding and subsequent decoding data in a BinaryCIF file