Repository: biotite-dev/biotite

Issue:
altloc_id missing from get_structure() output
The last few recent commits have altered the behavior of get_structure() so that its output is missing the altloc_id annotation. I can't tell if this is expected behavior or a bug, however the docstring implies that its a bug.
```
    altloc : {'first', 'occupancy', 'all'}
        This parameter defines how *altloc* IDs are handled:
            - ``'first'`` - Use atoms that have the first *altloc* ID
              appearing in a residue.
            - ``'occupancy'`` - Use atoms that have the *altloc* ID
              with the highest occupancy for a residue.
            - ``'all'`` - Use all atoms.
              Note that this leads to duplicate atoms.
              When this option is chosen, the ``altloc_id`` annotation
              array is added to the returned structure.
```

So it seems that either its a bug or the docstring should get updated.

Minimal example to reproduce:

```
from biotite.database import rcsb
import biotite.structure.io.pdbx as pdbx

ciffile = rcsb.fetch("4zxb", "cif")
cif = pdbx.CIFFile.read(ciffile)
array = pdbx.get_structure(
    cif,
    model=1,
    altloc="all",
    use_author_fields=False,
)
array.get_annotation_categories()
```

current output from the latest commit is
```
['chain_id',
 'res_id',
 'ins_code',
 'res_name',
 'hetero',
 'atom_name',
 'element']
```
rolling back to `d0ae83d` gives
```
['chain_id',
 'res_id',
 'ins_code',
 'res_name',
 'hetero',
 'atom_name',
 'element',
 'altloc_id']
```

Discussion:
@padix-key: Good catch! In #781 I introduced this bug and #795 will fix it.@davegrays: awesome, thank you for the fix!

Solution patch:
diff --git a/src/biotite/structure/io/pdbx/convert.py b/src/biotite/structure/io/pdbx/convert.py
index 94b3a15c..84a76c9d 100644
--- a/src/biotite/structure/io/pdbx/convert.py
+++ b/src/biotite/structure/io/pdbx/convert.py
@@ -19,10 +19,16 @@ __all__ = [
 
 import itertools
 import warnings
+from collections import defaultdict
 import numpy as np
 from biotite.file import InvalidFileError
 from biotite.sequence.seqtypes import NucleotideSequence, ProteinSequence
-from biotite.structure.atoms import AtomArray, AtomArrayStack, repeat
+from biotite.structure.atoms import (
+    AtomArray,
+    AtomArrayStack,
+    concatenate,
+    repeat,
+)
 from biotite.structure.bonds import BondList, BondType, connect_via_residue_names
 from biotite.structure.box import (
     coord_to_fraction,
@@ -766,8 +772,10 @@ def _filter_altloc(array, atom_site, altloc):
     altloc_ids = atom_site.get("label_alt_id")
     occupancy = atom_site.get("occupancy")
 
-    # Filter altloc IDs and return
-    if altloc_ids is None or (altloc_ids.mask.array != MaskValue.PRESENT).all():
+    if altloc == "all":
+        array.set_annotation("altloc_id", altloc_ids.as_array(str))
+        return array, atom_site
+    elif altloc_ids is None or (altloc_ids.mask.array != MaskValue.PRESENT).all():
         # No altlocs in atom_site category
         return array, atom_site
     elif altloc == "occupancy" and occupancy is not None:
@@ -779,9 +787,6 @@ def _filter_altloc(array, atom_site, altloc):
     elif altloc == "first":
         mask = filter_first_altloc(array, altloc_ids.as_array(str))
         return array[..., mask], _filter(atom_site, mask)
-    elif altloc == "all":
-        array.set_annotation("altloc_id", altloc_ids.as_array(str))
-        return array, atom_site
     else:
         raise ValueError(f"'{altloc}' is not a valid 'altloc' option")
 
@@ -1710,7 +1715,7 @@ def get_assembly(
     )
 
     ### Get transformations and apply them to the affected asym IDs
-    assembly = None
+    chain_ops = defaultdict(list)
     for id, op_expr, asym_id_expr in zip(
         assembly_gen_category["assembly_id"].as_array(str),
         assembly_gen_category["oper_expression"].as_array(str),
@@ -1719,19 +1724,22 @@ def get_assembly(
         # Find the operation expressions for given assembly ID
         # We already asserted that the ID is actually present
         if id == assembly_id:
-            operations = _parse_operation_expression(op_expr)
-            asym_ids = asym_id_expr.split(",")
-            # Filter affected asym IDs
-            sub_structure = structure[..., np.isin(structure.label_asym_id, asym_ids)]
-            sub_assembly = _apply_transformations(
-                sub_structure, transformations, operations
-            )
-            # Merge the chains with asym IDs for this operation
-            # with chains from other operations
-            if assembly is None:
-                assembly = sub_assembly
-            else:
-                assembly += sub_assembly
+            for chain_id in asym_id_expr.split(","):
+                chain_ops[chain_id].extend(_parse_operation_expression(op_expr))
+
+    sub_assemblies = []
+    for asym_id, op_list in chain_ops.items():
+        sub_struct = structure[..., structure.label_asym_id == asym_id]
+        sub_assembly = _apply_transformations(sub_struct, transformations, op_list)
+        # Merge the chain's sub_assembly into the rest of the assembly
+        sub_assemblies.append(sub_assembly)
+    assembly = concatenate(sub_assemblies)
+
+    # Sort AtomArray or AtomArrayStack by 'sym_id'
+    max_sym_id = assembly.sym_id.max()
+    assembly = concatenate(
+        [assembly[..., assembly.sym_id == sym_id] for sym_id in range(max_sym_id + 1)]
+    )
 
     # Remove 'label_asym_id', if it was not included in the original
     # user-supplied 'extra_fields'


Test patch:
diff --git a/tests/structure/data/README.rst b/tests/structure/data/README.rst
index 2b9297d5..f299c641 100644
--- a/tests/structure/data/README.rst
+++ b/tests/structure/data/README.rst
@@ -5,6 +5,8 @@ Test structures
 1DIX: Structure contains insertion code
 1F2N: Structure contains biological assembly
 5ZNG: Structure contains biological assembly with multiple entries in _pdbx_struct_assembly_gen
+4ZXB: Multiple entries in _pdbx_struct_assembly_gen for the same assembly_id
+1NCB: Multiple entries in _pdbx_struct_assembly_gen for the same assembly_id and chains
 1GYA: Large multi-model structure
 2AXD: Multi-model structure with different number of atoms per model
 1IGY: Multi-chain structure
diff --git a/tests/structure/data/ids.txt b/tests/structure/data/ids.txt
index e905de2f..4e491e44 100644
--- a/tests/structure/data/ids.txt
+++ b/tests/structure/data/ids.txt
@@ -16,4 +16,6 @@
 5eil
 4p5j
 1crr
-7gsa
\ No newline at end of file
+7gsa
+4zxb
+1ncb
\ No newline at end of file
diff --git a/tests/structure/io/test_pdbx.py b/tests/structure/io/test_pdbx.py
index 97e3686b..74ba577a 100644
--- a/tests/structure/io/test_pdbx.py
+++ b/tests/structure/io/test_pdbx.py
@@ -148,6 +148,35 @@ def test_conversion(monkeypatch, tmpdir, format, path, model, find_matches_by_di
     assert test_atoms.coord.tolist() == ref_atoms.coord.tolist()
 
 
+@pytest.mark.parametrize(
+    "pdb_id",
+    [
+        "1aki",  # has no altloc IDs
+        "3o5r",  # has altloc IDs
+    ],
+)
+def test_filter_altloc(pdb_id):
+    """
+    Check if the different ``altloc`` options give the expected results.
+    """
+    pdbx_file = pdbx.BinaryCIFFile.read(join(data_dir("structure"), f"{pdb_id}.bcif"))
+    atoms = {}
+    for altloc in ["first", "occupancy", "all"]:
+        atoms[altloc] = pdbx.get_structure(pdbx_file, model=1, altloc=altloc)
+
+    # The 'altloc_id' annotation should only be present in the 'altloc=all' case
+    assert "altloc_id" in atoms["all"].get_annotation_categories()
+    assert "altloc_id" not in atoms["first"].get_annotation_categories()
+    assert "altloc_id" not in atoms["occupancy"].get_annotation_categories()
+    # Independent of which altloc atom is selected, only one atom must be selected...
+    assert atoms["occupancy"].array_length() == atoms["first"].array_length()
+    # ...with the exception of the 'altloc=all' case in which more atoms are selected
+    if np.any(atoms["all"].altloc_id != "."):
+        assert atoms["all"].array_length() > atoms["first"].array_length()
+    else:
+        assert atoms["all"].array_length() == atoms["first"].array_length()
+
+
 @pytest.mark.parametrize(
     "format, path",
     itertools.product(
@@ -493,26 +522,22 @@ def test_assembly_chain_count(format, pdb_id, model):
         ("1f2n", "4", 6),
         # Multiple combined operations
         ("1f2n", "6", 60),
+        # Multiple entries in pdbx_struct_assembly_gen for the same assembly_id
+        ("4zxb", "1", 2),
+        # Multiple entries in pdbx_struct_assembly_gen for the same assembly_id and chain
+        ("1ncb", "2", 8),
     ],
 )
 def test_assembly_sym_id(pdb_id, assembly_id, symmetric_unit_count):
     """
     Check if the :func:`get_assembly()` function returns the correct
-    symmetry ID annotation for a known example.
+    number of symmetry IDs for a known example.
     """
     pdbx_file = pdbx.BinaryCIFFile.read(join(data_dir("structure"), f"{pdb_id}.bcif"))
     assembly = pdbx.get_assembly(pdbx_file, assembly_id=assembly_id)
-    # 'unique_indices' contains the FIRST occurence of each unique value
-    unique_sym_ids, unique_indices = np.unique(assembly.sym_id, return_index=True)
-    # Sort by first occurrence
-    order = np.argsort(unique_indices)
-    unique_sym_ids = unique_sym_ids[order]
-    unique_indices = unique_indices[order]
-    assert unique_sym_ids.tolist() == list(range(symmetric_unit_count))
-    # Every asymmetric unit should have the same length,
-    # as each operation is applied to all atoms in the asymmetric unit
-    asym_lengths = np.diff(np.append(unique_indices, assembly.array_length()))
-    assert (asym_lengths == asym_lengths[0]).all()
+    assert sorted(np.unique(assembly.sym_id).tolist()) == list(
+        range(symmetric_unit_count)
+    )
 
 
 @pytest.mark.parametrize("model", [None, 1])
diff --git a/tests/structure/test_sequence.py b/tests/structure/test_sequence.py
index bcac9b9e..23b96677 100644
--- a/tests/structure/test_sequence.py
+++ b/tests/structure/test_sequence.py
@@ -4,8 +4,8 @@
 
 import glob
 from os.path import join
+import numpy as np
 import pytest
-import biotite.sequence as seq
 import biotite.sequence.align as align
 import biotite.structure as struc
 import biotite.structure.io.pdbx as pdbx
@@ -54,12 +54,20 @@ def _find_best_match(sequence, ref_sequences):
     for ref_sequence in ref_sequences.values():
         if not isinstance(sequence, type(ref_sequence)):
             continue
-        if isinstance(sequence, seq.ProteinSequence):
-            matrix = align.SubstitutionMatrix.std_protein_matrix()
-        else:
-            matrix = align.SubstitutionMatrix.std_nucleotide_matrix()
+        # We are not handling homology here, but residues may be missing due to
+        # experimental artifacts -> use simply uniform substitution matrix
+        matrix = align.SubstitutionMatrix(
+            sequence.alphabet,
+            sequence.alphabet,
+            np.eye(len(sequence.alphabet), dtype=int),
+        )
         alignment = align.align_optimal(
-            sequence, ref_sequence, matrix, terminal_penalty=False, max_number=1
+            sequence,
+            ref_sequence,
+            matrix,
+            gap_penalty=-1,
+            terminal_penalty=True,
+            max_number=1,
         )[0]
         # The 'shortest' identity is 1.0, if every residue in the
         # test sequence is aligned to an identical residue