Update ResidueGroup/SegmentGroup with ids (PDBParser: force chainid to segid / Universe: set groups) (#4965)

* Fixes #4948 #2874
* Universe: add a function in the universe set_groups to update the topology and groups in the universe when atomwise resids/segids is given (more generalized method across different file type)
* PDBParser: add an argument force_chainids_to_segids in PDBReader to forcibly reading the chain ID into the segment ID (PDB specific method)
* update docs
* add tests
* update CHANGELOG

---------

Co-authored-by: Cédric Bouysset <bouysset.cedric@gmail.com>
Co-authored-by: Lily Wang <31115101+lilyminium@users.noreply.github.com>

Author: 3 people

package/CHANGELOG

@@ -15,7 +15,7 @@ The rules for this file:
 
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, orbeckst, BHM-Bob, TRY-ER, Abdulrahman-PROG, pbuslaev,
-         yuxuanzhuang
+         yuxuanzhuang, yuyuan871111
 
  * 2.10.0
 
@@ -34,6 +34,14 @@ Enhancements
  * Improve parsing of topology information from LAMMPS dump files to allow
    reading of mass, charge and element attributes. (Issue #3449, PR #4995)
  * Added timestep support for XDR writers and readers (Issue #4905, PR #4908)
+ * New feature in `Universe`: `set_groups`. This function allows updating 
+   `_topology`, `residues` (ResidueGroup), `segments` (SegmentGroup) in 
+   the `Universe` using user-provided atomwise `resids` and/or `segids`. 
+   (Issue #4948 #2874, PR #4965)   
+ * Adds an argument `force_chainids_to_segids` in `PDBParser`: this will 
+   force the use of ChainID as the segID when reading PDB (it is helpful 
+   if the segment IDs (segids) are partially missing in the PDB file). 
+   (Issue #4948 #2874, PR #4965)
 
 Changes
 

package/MDAnalysis/coordinates/PDB.py

@@ -244,6 +244,11 @@ class PDBReader(base.ReaderBase):
     :class:`PDBWriter`
     :class:`PDBReader`
 
+    If you would like to force the use of chainID as the segID when parsing PDB,
+    please use the keyword argument `force_chainids_to_segids=True`
+    (:class:`MDAnalysis.topology.PDBParser.PDBParser`).
+    This will prioritize the chain ID to the segment ID.
+
 
     .. versionchanged:: 0.11.0
        * Frames now 0-based instead of 1-based

package/MDAnalysis/core/universe.py

@@ -86,6 +86,7 @@
 from .topology import Topology
 from .topologyattrs import (
     AtomAttr, ResidueAttr, SegmentAttr,
+    Segindices, Segids, Resindices, Resids, Atomindices,
     BFACTOR_WARNING, _Connection
 )
 from .topologyobjects import TopologyObject
@@ -94,6 +95,105 @@
 logger = logging.getLogger("MDAnalysis.core.universe")
 
 
+def _update_topology_by_ids(universe, atomwise_resids, atomwise_segids):
+    """Update the topology of a Universe with new atomwise resids and segids.
+
+    Parameters
+    ----------
+    universe : Universe
+        The universe to update.
+    atomwise_resids : numpy.ndarray
+        The new atomwise residue indices.
+    atomwise_segids : numpy.ndarray
+        The new atomwise segment indices.
+    """
+    from ..topology.base import change_squash
+
+    # the original topology
+    top = universe._topology
+
+    # detect the atom level attributes (excluding residue level and above)
+    atom_attrindices = [
+        idx
+        for idx, each_attr in enumerate(top.attrs)
+        if (Residue not in each_attr.target_classes) and
+        (not isinstance(each_attr, Atomindices))
+    ]
+    attrs = [top.attrs[each_attr] for each_attr in atom_attrindices]
+
+    # create new residues level stuff
+    residue_attrindices = [
+        idx
+        for idx, each_attr in enumerate(top.attrs)
+        if (
+            (Residue in each_attr.target_classes) and
+            (Segment not in each_attr.target_classes) and
+            (not isinstance(each_attr, Resids)) and
+            (not isinstance(each_attr, Resindices))
+        )
+    ]  # residue level attributes except resids and resindices
+
+    res_criteria = [atomwise_resids, atomwise_segids] + [
+        getattr(universe.atoms, top.attrs[each_attr].attrname)
+        for each_attr in residue_attrindices
+        if top.attrs[each_attr].attrname != 'resnums'
+    ]
+
+    res_to_squash = [atomwise_resids, atomwise_segids] + [
+        getattr(universe.atoms, top.attrs[each_attr].attrname)
+        for each_attr in residue_attrindices
+    ]
+
+    residx, res_squashed = change_squash(res_criteria, res_to_squash)
+    resids = res_squashed[0]
+    res_squashed_segids = res_squashed[1]
+    n_residues = len(resids)
+    # all residue-level attributes except resids
+    res_squashed_res_attrs = res_squashed[2:]
+
+    res_attrs = [Resids(resids)] + [
+        # use the correspdoning type of the attribute with new sqaushed values
+        top.attrs[each_attr].__class__(res_squashed_res_attrs[idx])
+        for idx, each_attr in enumerate(residue_attrindices)
+    ]
+    attrs.extend(res_attrs)
+
+    # create new segment level stuff
+    segidx, (segids,) = change_squash((res_squashed_segids,), (res_squashed_segids,))
+    n_segments = len(segids)
+    attrs.append(Segids(segids))
+
+    # other attributes
+    other_attrs = [
+        each_attr
+        for each_attr in top.attrs
+        if (
+            (Segment in each_attr.target_classes) and
+            (not isinstance(each_attr, Segids)) and
+            (not isinstance(each_attr, Atomindices)) and
+            (not isinstance(each_attr, Resindices)) and
+            (not isinstance(each_attr, Segindices))
+        )
+    ]
+
+    # create new topology
+    top = Topology(
+        universe.atoms.n_atoms,
+        n_residues,
+        n_segments,
+        attrs=attrs,
+        atom_resindex=residx,
+        residue_segindex=segidx,
+    )
+
+    # add back other attributes
+    if len(other_attrs) > 0:
+        for each_otherattr in other_attrs:
+            top.add_TopologyAttr(each_otherattr)
+
+    # update the topology in the universe
+    universe._topology = top
+
 
 def _check_file_like(topology):
     if isstream(topology):
@@ -312,6 +412,12 @@ class Universe(object):
         functionality to treat independent trajectory files as a single virtual
         trajectory.
     **kwargs: extra arguments are passed to the topology parser.
+        For instance, when reading a PDB file
+        (:class:`PDBReader<MDAnalysis.coordinates.PDB>`,
+        :class:`PDBParser<MDAnalysis.topology.PDBParser>`), set
+        ``force_chainids_to_segids=True`` to make the universe use the
+        chainIDs (column 22) instead of the segmentIDs (column 73-76) as the
+        `segids` in the universe and select the corresponding SegmentGroup.
 
     Attributes
     ----------
@@ -380,6 +486,10 @@ class Universe(object):
         guessing masses and atom types after topology
         is read from a registered parser.
 
+    .. versionchanged:: 2.10.0
+        Added :meth: `~MDAnalysis.core.universe.Universe.set_groups`
+        API to set residues/segments based on the atomwise resids/segids.
+
     """
     def __init__(self, topology=None, *coordinates, all_coordinates=False,
                  format=None, topology_format=None, transformations=None,
@@ -1700,6 +1810,104 @@ def guess_TopologyAttrs(
             warnings.warn('Can not guess attributes '
                           'for universe with 0 atoms')
 
+    def set_groups(self, atomwise_resids=None, atomwise_segids=None):
+        """Set the groups (`ResidueGroup`, `SegmentGroup`) of the Universe
+        by atomwise resids/segids.
+
+        The `topology` will also be updated based on the provided `atomwise_resids`
+        and `atomwise_segids`. The original `resids` and `segids` will be stored
+        in attributes `atomwise_resids_orig` and/or `atomwise_segids_orig` if
+        they are modified.
+        See notes for the logic of the function.
+
+        Parameters
+        ----------
+        atomwise_resids:
+            A list of residue IDs to be set for the Universe. The length
+            of the list should be equal to the number of atoms in the Universe.
+            If `None`, the original resids will be used.
+
+        atomwise_segids:
+            A list of segment IDs to be set for the Universe. The length
+            of the list should be equal to the number of atoms in the Universe.
+            If `None`, the original segids will be used.
+
+        Raises
+        ------
+        AssertionError
+            If the length of the provided atomwise_resids or atomwise_segids
+            does not match the number of atoms in the Universe.
+
+        Notes
+        -----
+        First, the function will check if resids or segids is provided.
+        If both resids and segids are not provided (`None`), it will do nothing.
+        If only one of them is provided, it will use the original values for the
+        other one. If both are provided, it will use the provided values for
+        both resids and segids.
+        The function will then update the topology by a new generated topology
+        with new values of the resids and segids.
+        Finally, the corresponding new `ResidueGroup` and `SegmentGroup` will be
+        created by the updated topology.
+
+        Examples
+        --------
+        To set custom segment IDs for the segments of the Universe::
+
+            atomwise_segids = ['A', 'A', 'B', 'B']
+            u.set_groups(atomwise_segids=atomwise_segids)
+
+            # Now the Universe has two segments with segIDs 'A' and 'B'
+            u.segments
+            >>> <SegmentGroup with 2 segments>
+
+        .. versionadded:: 2.10.0
+        """
+        if (atomwise_resids is None) and (atomwise_segids is None):
+            warnings.warn("Not setting groups. Please provide atomwise_resids or "
+                          "atomwise_segids.")
+            return
+
+        # resids
+        if atomwise_resids is None:
+            atomwise_resids = self.atoms.resids
+
+        else:
+            # check the length of atomwise_resids
+            if len(atomwise_resids) != self.atoms.n_atoms:
+                raise ValueError(
+                    "The length of atomwise_resids should be the same as "
+                    "the number of atoms in the universe.")
+
+            self.atomwise_resids_orig = self.atoms.resids
+            logger.info("The new resids replaces the current one. "
+                        "The original resids is stored in "
+                        "atomwise_resids_orig.")
+
+        # segids
+        if atomwise_segids is None:
+            atomwise_segids = self.atoms.segids
+
+        else:
+            # check the length of atomwise_segids
+            if len(atomwise_segids) != self.atoms.n_atoms:
+                raise ValueError(
+                    "The length of atomwise_segids should be the same as "
+                    "the number of atoms in the universe.")
+
+            self.atomwise_segids_orig = self.atoms.segids
+            logger.info("The new resids replaces the current one. "
+                        "The original segids is stored in "
+                        "atomwise_segids_orig.")
+
+        atomwise_resids = np.array(atomwise_resids, dtype=int)
+        atomwise_segids = np.array(atomwise_segids, dtype=object)
+
+        _update_topology_by_ids(self,
+                                atomwise_resids=atomwise_resids,
+                                atomwise_segids=atomwise_segids)
+        _generate_from_topology(self)
+
 
 def Merge(*args):
     """Create a new new :class:`Universe` from one or more

package/MDAnalysis/topology/PDBParser.py

@@ -178,6 +178,11 @@ class PDBParser(TopologyReaderBase):
      - bonds
      - formalcharges
 
+    Note that `PDBParser` accepts an optional keyword argument
+    ``force_chainids_to_segids``. If set to ``True``, the chain IDs (even if
+    empty values are in the chain ID column in the file) will forcibly be used
+    instead of the segment IDs for creating segments.
+
     See Also
     --------
     :class:`MDAnalysis.coordinates.PDB.PDBReader`
@@ -207,7 +212,10 @@ class PDBParser(TopologyReaderBase):
         Removed type and mass guessing (attributes guessing takes place now
         through universe.guess_TopologyAttrs() API).
     .. versionchanged:: 2.10.0
-        segID is read from 73-76 instead of 67-76.
+        segID is read from 73-76 instead of 67-76 and added the
+        `force_chainids_to_segids` keyword argument. Some infos in logger will
+        be generated if the segids is not present or if the chainids are not
+        completely equal to segids.
     """
     format = ['PDB', 'ENT']
 
@@ -218,7 +226,7 @@ def parse(self, **kwargs):
         -------
         MDAnalysis Topology object
         """
-        top = self._parseatoms()
+        top = self._parseatoms(**kwargs)
 
         try:
             bonds = self._parsebonds(top.ids.values)
@@ -235,7 +243,7 @@ def parse(self, **kwargs):
 
         return top
 
-    def _parseatoms(self):
+    def _parseatoms(self, **kwargs):
         """Create the initial Topology object"""
         resid_prev = 0  # resid looping hack
 
@@ -325,6 +333,12 @@ def _parseatoms(self):
                         "found in the PDB file.")
             segids = chainids
 
+        # If force_chainids_to_segids is set, use chainids as segids
+        if kwargs.get("force_chainids_to_segids", False):
+            logger.info("force_chainids_to_segids is set. "
+                        "Using chain IDs as segment IDs.")
+            segids = chainids
+
         n_atoms = len(serials)
 
         attrs = []
@@ -403,11 +417,13 @@ def _parseatoms(self):
         n_residues = len(resids)
         attrs.append(Resnums(resnums))
         attrs.append(Resids(resids))
-        attrs.append(Resnums(resids.copy()))
         attrs.append(ICodes(icodes))
         attrs.append(Resnames(resnames))
 
-        if any(segids) and not any(val is None for val in segids):
+        if (
+            kwargs.get("force_chainids_to_segids", False) or
+            (any(segids) and not any(val is None for val in segids))
+        ):
             segidx, (segids,) = change_squash((segids,), (segids,))
             n_segments = len(segids)
             attrs.append(Segids(segids))