structure_tools

  1"""
  2[structure_tools](https://github.com/isblab/af_pipeline/tree/main/af_pipeline/tools/structure_tools.py)
  3==================================
  4- This module provides utility functions and classes to work with structure files (PDB or CIF).
  5- Uses `Biopython` for structure manipulation and parsing.
  6- The structure object refers to `Bio.PDB.Structure.Structure`.
  7- The residue object refers to `Bio.PDB.Residue.Residue`.<br />
  8  Hence, the "residue" term is used for amino acids, nucleotides, ions, ligands.
  9"""
 10from textwrap import dedent
 11import warnings
 12import Bio
 13import Bio.PDB
 14import Bio.PDB.Structure
 15import Bio.PDB.Residue
 16import Bio.PDB.Atom
 17import numpy as np
 18from Bio.PDB.Chain import Chain
 19from Bio.PDB.mmcifio import MMCIFIO
 20from Bio.PDB.PDBIO import Select, PDBIO
 21from collections import defaultdict
 22from typing import Dict, Any, overload
 23from af_pipeline.constants.af_constants import (
 24    PROTEIN_ENTITIES,
 25    DNA_ENTITIES,
 26    RNA_ENTITIES,
 27    ION,
 28    PURINES,
 29    PYRIMIDINES,
 30    ALLOWED_PTMS,
 31    ALLOWED_DNA_MODS,
 32    ALLOWED_RNA_MODS,
 33    ALLOWED_LIGANDS,
 34    ONLY_CA_RESIDUES,
 35    AtomDecoration,
 36    EntityType,
 37    FileFormat,
 38    InteractionMapType,
 39    ResidueDecoration,
 40    ResidueMapDepth,
 41    VALID_INTERACTION_MAP_TYPES,
 42    ResidueMapKeys,
 43    TokenLevel,
 44)
 45
 46def get_interaction_map(
 47    coords1: np.ndarray,
 48    coords2: np.ndarray,
 49    contact_threshold: float,
 50    map_type: InteractionMapType,
 51) -> np.ndarray:
 52    """ Create an interaction map, given the input coordinates.
 53
 54    Returns a distance map or a contact map, based on the map_type specified.
 55
 56    ## Arguments:
 57
 58    - **coords1 (np.ndarray)**:<br />
 59        Coordinates of shape (N, 3) for the first set of atoms/residues.
 60
 61    - **coords2 (np.ndarray)**:<br />
 62        Coordinates of shape (M, 3) for the second set of atoms/residues.
 63
 64    - **contact_threshold (float)**:<br />
 65        Distance threshold to define a contact.
 66        Only used if `map_type` is `InteractionMapType.CONTACT`.
 67
 68    - **map_type (InteractionMapType)**:<br />
 69        Type of interaction map to create.
 70        Can be either `InteractionMapType.DISTANCE` or `InteractionMapType.CONTACT`.
 71
 72         - `InteractionMapType.DISTANCE`:
 73            Returns a distance map of shape (N, M)
 74            where each element (i, j) is the distance between the i-th atom/residue
 75            in `coords1` and the j-th atom/residue in `coords2`.
 76
 77         - `InteractionMapType.CONTACT`:
 78            Returns a binary contact map of shape (N, M)
 79            where each element (i, j) is 1 if the distance between the i-th atom/residue
 80            in `coords1` and the j-th atom/residue in `coords2` is less than or equal to
 81            `contact_threshold`, and 0 otherwise.
 82
 83    ## Returns:
 84
 85    - **np.ndarray**:<br />
 86        The interaction map as a numpy array.
 87    """
 88
 89    if map_type not in VALID_INTERACTION_MAP_TYPES:
 90        raise ValueError(
 91            f"Invalid map_type specified."\
 92            f"Expected one of {VALID_INTERACTION_MAP_TYPES}, got {map_type}"
 93        )
 94
 95    distance_map = get_distance_map(coords1, coords2)
 96
 97    if map_type == InteractionMapType.DISTANCE:
 98        return distance_map
 99
100    elif map_type == InteractionMapType.CONTACT:
101        contact_map = get_contact_map(distance_map, contact_threshold)
102        return contact_map
103
104def get_distance_map(coords1: np.ndarray, coords2: np.ndarray):
105    """ Create an all-v-all distance map.
106
107    Returns a matrix of distances between all pairs of atoms/residues in the
108    two sets of coordinates.
109
110    ## Arguments:
111
112    - **coords1 (np.ndarray)**:<br />
113        Coordinates of shape (N, 3) for the first set of atoms/residues.
114
115    - **coords2 (np.ndarray)**:<br />
116        Coordinates of shape (M, 3) for the second set of atoms/residues.
117
118    ## Returns:
119
120    - **np.ndarray**:<br />
121        A matrix of shape (N, M) where each element (i, j) is the Euclidean distance
122        between the i-th atom/residue in `coords1` and the j-th atom/residue in `coords2`.
123    """
124
125    from scipy.spatial import distance_matrix
126
127    distance_map = distance_matrix(coords1, coords2)
128
129    return distance_map
130
131def get_contact_map(distance_map: np.ndarray, contact_threshold: float):
132    """ Given the distance map, create a binary contact map by thresholding distances.
133
134    Returns a binary matrix, where 1 indicates a contact and 0 indicates no contact.
135
136    ## Arguments:
137
138    - **distance_map (np.ndarray)**:<br />
139        A matrix of shape (N, M) where each element (i, j) is the Euclidean distance
140        between the i-th atom/residue in `coords1` and the j-th atom/residue in `coords2`.
141
142    - **contact_threshold (float)**:<br />
143        Distance threshold to define a contact. If the distance between two atoms/residues
144        is less than or equal to this threshold, they are considered to be in contact.
145
146    ## Returns:
147
148    - **np.ndarray**:<br />
149        A binary matrix of shape (N, M) where each element (i, j) is 1 if the distance
150        between the i-th atom/residue in `coords1` and the j-th atom/residue in `coords2`
151        is less than or equal to `contact_threshold`, and 0 otherwise.
152    """
153
154    contact_map = np.where(distance_map <= contact_threshold, 1, 0)
155
156    return contact_map
157
158def get_token_level(residue: Bio.PDB.Residue.Residue) -> TokenLevel:
159    """ Get the token level of the residue.
160
161    Assuming the `residue` object is decorated with the appropriate attributes.
162    - `tokenLevel`: `TokenLevel.ATOM` if the residue has per-atom token, `TokenLevel.RESIDUE` otherwise.
163
164    Arguments:
165
166    - **residue (Bio.PDB.Residue.Residue)**:<br />
167        Biopython residue object.
168
169    Returns:
170
171    - **token_level (TokenLevel)**:<br />
172        Token level of the residue.
173    """
174
175    token_level = residue.xtra.get(
176        ResidueDecoration.TOKEN_LEVEL, TokenLevel.RESIDUE
177    )
178
179    return token_level
180
181def has_modifications(structure: Bio.PDB.Structure.Structure) -> bool:
182    """ Check if the structure has any modified residues.
183
184    Assuming the `residue` objects are decorated with the appropriate attributes.
185    - `is_modified`: `True` if the `residue` is modified.
186
187    Arguments:
188
189    - **structure (Bio.PDB.Structure.Structure)**:<br />
190        Biopython structure object.
191
192    Returns:
193
194    - **condition (bool)**:<br />
195        `True` if the structure has any modified residues, `False` otherwise.
196    """
197
198    for residue in structure.get_residues():
199        if residue.xtra.get(ResidueDecoration.IS_MODIFIED, False):
200            return True
201
202    return False
203
204def save_structure_obj(
205    structure: Bio.PDB.Structure.Structure,
206    out_file: str,
207    res_select_obj: Select = Select(),
208    save_type: str = FileFormat.CIF,
209    preserve_header_footer = False,
210):
211    """Save the selection in Biopython structure object as a PDB or CIF file.
212
213    Arguments:
214
215    - **structure (Bio.PDB.Structure.Structure)**:<br />
216        Biopython structure object to save.
217
218    - **out_file (str)**:<br />
219        Path to the output file where the structure will be saved.
220
221    - **res_select_obj (Bio.PDB.Select, optional)**:<br />
222        Biopython Select object to filter the residues to save.
223        Defaults to `Select()` which saves all residues.
224
225    - **save_type (str, optional)**:<br />
226        Type of file to save the structure as.
227        Can be either "pdb" or "cif".
228
229    - **preserve_header_footer (bool, optional)**:<br />
230        If `True`, the header and footer information from the structure
231        will be preserved in the saved file.
232        > [!NOTE]
233        > The header and footer information can only be preserved for CIF files.
234    """
235
236    if save_type == FileFormat.PDB:
237
238        io = PDBIO()
239        io.set_structure(structure)
240        io.save(out_file, res_select_obj)
241
242        if preserve_header_footer:
243            warnings.warn(
244                """
245
246                PDB files do not support headers and footers.
247                Saving without headers and footers.
248                """
249            )
250
251    elif save_type == FileFormat.CIF:
252
253        io = MMCIFIO()
254        io.set_structure(structure)
255        io.save(out_file, res_select_obj)
256
257        if preserve_header_footer:
258
259            header_footer = structure.xtra.get("header_footer", {})
260
261            if {"header", "footer"}.issubset(set(header_footer.keys())):
262
263                with open(out_file, "r") as f:
264                    lines = f.readlines()
265
266                lines.insert(0, "\n")
267                lines.append("\n")
268
269                for header_line in header_footer["header"][::-1]:
270                    lines.insert(0, f"{header_line}")
271
272                for footer_line in header_footer["footer"][::-1]:
273                    lines.append(f"{footer_line}")
274
275                with open(out_file, "w") as f:
276                    for line in lines:
277                        f.write(line)
278
279            else:
280                warnings.warn(
281                    """
282
283                    No header or footer information found in the structure.
284                    Saving without headers and footers.
285                    """
286                )
287
288def add_header_footer(
289    structure: Bio.PDB.Structure.Structure,
290    structure_file_path: str,
291) -> Bio.PDB.Structure.Structure:
292    """Add the header and footer information to the structure object.
293
294    While parsing the `CIF` file using Biopython `MMCIFParser`,
295    the header and footer information information is not retained.<br />
296    This function extracts the header and footer information from the
297    structure file and adds it to the structure object. \n
298
299    The information is stored in the `xtra` attribute of the structure
300    object under the key `header_footer`. \n
301
302    > [!NOTE]
303    > - `PDB` format does not have header and footer information.
304    > - Using this function leads to parsing of the structure file twice. Once
305    >   to get the structure object and once to extract the header and footer information.
306
307    Arguments:
308
309    - **structure (Bio.PDB.Structure.Structure)**:<br />
310        Biopython structure object.
311
312    - **structure_file_path (str)**:<br />
313        Path to the structure file.
314
315    Returns:
316
317    - **structure (Bio.PDB.Structure.Structure)**:<br />
318        Biopython Structure object with `header_footer` added in `xtra`.
319    """
320
321    with open(structure_file_path, "r") as f:
322        lines = f.readlines()
323
324    header_info = []
325    header_section = ""
326
327    for line in lines:
328        header_section += line
329
330        if line.startswith("#"):
331            header_info.append(header_section)
332            header_section = ""
333
334        if line.startswith("_atom_site"):
335            break
336
337    footer_info = []
338    footer_section = ""
339
340    for line in lines[::-1]:
341        footer_section = line + footer_section
342
343        if line.startswith("#"):
344            footer_info.append(footer_section)
345            footer_section = ""
346
347        if line.startswith("ATOM") or line.startswith("HETATM"):
348            break
349
350    structure.xtra["header_footer"] = {
351        "header": header_info,
352        "footer": footer_info,
353    }
354
355    return structure
356
357def decorate_residue(
358    residue: Bio.PDB.Residue.Residue,
359    xtra_field: str | None = None,
360    xtra_value: Any = None,
361):
362    """Decorate the residue.
363
364    This function adds `entityType` attribute to the residue's `xtra`
365    attribute depending on the type of entity.
366
367    Allowed entity types are:
368    ```python
369    - "proteinChain"
370    - "dnaSequence"
371    - "rnaSequence"
372    - "ligand"
373    - "ion"
374    - None #(if the residue does not belong to any of the above types)
375    ```
376
377    It also adds the following flags to the residue's `xtra` wherever applicable -
378    ```python
379    - "is_modified" # boolean indicating if the residue is modified
380    - "is_ca_only" # boolean indicating if the residue has only CA atom
381    - "is_purine" # boolean indicating if the nucleotide is purine
382    - "is_pyrimidine" # boolean indicating if the nucleotide is pyrimidine
383    - "is_ion" # boolean indicating if the residue is an ion.
384    - "is_ligand" # boolean indicating if the residue is a ligand.
385    ```
386
387    Arguments:
388
389    - **residue (Bio.PDB.Residue.Residue)**:<br />
390        Biopython residue object.
391    """
392
393    symbol = residue.get_resname()
394
395    if symbol in PROTEIN_ENTITIES:
396        residue.xtra[ResidueDecoration.ENTITY_TYPE] = EntityType.PROTEIN_CHAIN
397
398        if symbol in ALLOWED_PTMS:
399            residue.xtra[ResidueDecoration.IS_MODIFIED] = True
400
401        if symbol in ONLY_CA_RESIDUES:
402            residue.xtra[ResidueDecoration.IS_CA_ONLY] = True
403
404    elif symbol in DNA_ENTITIES:
405        residue.xtra[ResidueDecoration.ENTITY_TYPE] = EntityType.DNA_SEQUENCE
406
407        if symbol in ALLOWED_DNA_MODS:
408            residue.xtra[ResidueDecoration.IS_MODIFIED] = True
409    elif symbol in RNA_ENTITIES:
410        residue.xtra[ResidueDecoration.ENTITY_TYPE] = EntityType.RNA_SEQUENCE
411
412        if symbol in ALLOWED_RNA_MODS:
413            residue.xtra[ResidueDecoration.IS_MODIFIED] = True
414
415    elif symbol in ALLOWED_LIGANDS:
416        residue.xtra[ResidueDecoration.ENTITY_TYPE] = EntityType.LIGAND
417
418    elif symbol in ION:
419        residue.xtra[ResidueDecoration.ENTITY_TYPE] = EntityType.ION
420
421    else:
422        warnings.warn(dedent(f"""
423            The residue "{symbol}" does not belong to any known entity types.
424            Setting 'entityType' to None.""")
425        )
426        residue.xtra[ResidueDecoration.ENTITY_TYPE] = None
427
428    if symbol in PURINES:
429        residue.xtra[ResidueDecoration.IS_PURINE] = True
430
431    elif symbol in PYRIMIDINES:
432        residue.xtra[ResidueDecoration.IS_PYRIMIDINE] = True
433
434    if xtra_field is not None and xtra_value is not None:
435        if xtra_field in residue.xtra:
436            warnings.warn(dedent(f"""
437                The field '{xtra_field}' already exists in the residue's xtra.
438                Overwriting the value.""")
439            )
440        residue.xtra[xtra_field] = xtra_value
441
442    condition = (
443        residue.xtra.get(ResidueDecoration.IS_MODIFIED)
444        or residue.xtra.get(ResidueDecoration.ENTITY_TYPE) in [EntityType.LIGAND, None]
445    )
446    _get_token_level = {True: TokenLevel.ATOM, False: TokenLevel.RESIDUE}
447    residue.xtra[ResidueDecoration.TOKEN_LEVEL] = _get_token_level[condition]
448
449def decorate_atom(
450    atom: Bio.PDB.Atom.Atom,
451    xtra_field: str | None = None,
452    xtra_value: Any = None,
453):
454    """ Decorate the atom.
455
456    This function adds `is_representative` attribute to the atom's `xtra`
457
458    Arguments:
459
460    - **atom (Bio.PDB.Atom.Atom)**:<br />
461        Biopython atom object.
462
463    - **xtra_field (str | None, optional)**:<br />
464        field to add to the atom's `xtra`.
465
466    - **xtra_value (Any, optional)**:<br />
467        value of the field to add to the atom's `xtra`.
468    """
469
470    symbol = atom.get_name()
471    residue = atom.get_parent()
472
473    if not isinstance(residue, Bio.PDB.Residue.Residue):
474        raise TypeError(
475            f"Expected a `Bio.PDB.Residue.Residue` object, got {type(residue)}"
476        )
477
478    residue = residue.get_resname()
479
480    residue_attrs = (
481        residue in PROTEIN_ENTITIES,
482        residue in ONLY_CA_RESIDUES,
483        residue in PURINES,
484        residue in PYRIMIDINES,
485        symbol
486    )
487
488    atom_xtra_dict = {
489        (True, False, False, False, "CB"): True,
490        (True, True, False, False, "CA"): True,
491        (False, False, True, False, "C4"): True,
492        (False, False, False, True, "C2"): True,
493    }
494
495    atom.xtra[AtomDecoration.IS_REPRESENTATIVE] = atom_xtra_dict.get(
496        residue_attrs, False
497    )
498
499    if (
500        xtra_field is not None
501        and xtra_value is not None
502        and xtra_field in atom.xtra
503    ):
504        warnings.warn(dedent(f"""
505            The field '{xtra_field}' already exists in the atom's xtra.
506            Overwriting the value.""")
507        )
508        atom.xtra[xtra_field] = xtra_value
509
510
511class RenumberResidues:
512    """Class to renumber the residues based on the offset.
513
514    If the input sequence to the AlphFold does not start from the first residue,
515    the residue numbering in the predicted structure will be different from
516    the original (UniProt in case of proteins) numbering.
517
518    This class renumbers the residues in the structure based on the provided `offset`.
519
520    > [!NOTE]
521    > Although written with the purpose for AlphFold predicted structures,
522    > this class can be used to renumber residues in any structure file.
523    """
524
525    offset: dict
526    """ Offset describing start and end residue number for each chain in
527    the predicted structure.\n
528    example: `{'A': [1, 100], 'B': [101, 200]}`."""
529
530    def __init__(self, offset: dict = {}):
531
532        self.offset = offset
533
534    def renumber_structure(
535        self,
536        structure: Bio.PDB.Structure.Structure,
537    ):
538        """Renumber the residues in the structure based on the offset.
539
540        Arguments:
541
542        - **structure (Bio.PDB.Structure.Structure)**:<br />
543            Biopython structure object.
544
545        Returns:
546
547        - **structure (Bio.PDB.Structure.Structure)**:<br />
548            Biopython structure object with renumbered residues.
549        """
550
551        for model in structure:
552            for residue in model.get_residues():
553                chain_id = residue.parent.id
554                h, num, ins = residue.id
555
556                num = self.renumber_chain_res_num(
557                    chain_res_num=num,
558                    chain_id=chain_id,
559                )
560
561                residue.id = (h, num, ins)
562
563        return structure
564
565    def original_chain_res_num(
566        self,
567        chain_res_num: int,
568        chain_id: str,
569    ):
570        """Get the original residue number based on the offset.
571
572        Inverse of `renumber_chain_res_num`.
573
574        Arguments:
575
576        - **chain_res_num (int)**:<br />
577            Residue index (1-indexed) within the chain in the predicted structure.
578
579        - **chain_id (str)**:<br />
580            Chain ID of the residue.
581
582        Returns:
583
584        - **chain_res_num (int)**:<br />
585            Original residue number
586        """
587
588        if chain_id in self.offset:
589            chain_res_num -= (self.offset[chain_id][0] - 1)
590
591        return chain_res_num
592
593    def renumber_chain_res_num(
594        self,
595        chain_res_num: int,
596        chain_id: str,
597    ):
598        """Renumber the residue number based on the offset.
599
600        Given a residue index (1-indexed) in a chain, this function renumbers it
601        based on the offset provided for that chain.
602
603        Arguments:
604
605        - **chain_res_num (int)**:<br />
606            Residue index (1-indexed) within the chain in the predicted structure.
607
608        - **chain_id (str)**:<br />
609            Chain ID of the residue.
610
611        Returns:
612
613        - **chain_res_num (int)**:<br />
614            Renumbered residue number
615        """
616
617        if chain_id in self.offset:
618            chain_res_num += self.offset[chain_id][0] - 1
619
620        return chain_res_num
621
622    def renumber_region_of_interest(
623        self,
624        region_of_interest: Dict[str, list],
625    ):
626        """Offset the region of interest to the AF2/3 numbering.
627
628        Region of interest is defined by the user is as per the original numbering
629        (UniProt in case of proteins). \n
630        However, if the prediction is done on a fragment of the protein, the
631        residue numbering in the predicted structure will be different compared to
632        the one provided by the user. \n
633        This function renumbers the region of interest to the numbering of the
634        predicted structure based on the `offset`. \n
635
636        Arguments:
637
638        - **region_of_interest (dict)**:
639            Dictionary containing the region of interest for each chain.
640
641        Returns:
642
643        - **renumbered_region_of_interest (dict)**:
644            Dictionary containing the renumbered region of interest for
645            each chain.
646        """
647
648        renumbered_region_of_interest = {}
649
650        for chain_id in region_of_interest:
651
652            start, end = region_of_interest[chain_id]
653
654            if chain_id in self.offset:
655
656                start = start - (self.offset[chain_id][0] - 1)
657                end = end - (self.offset[chain_id][0] - 1)
658
659            renumbered_region_of_interest[chain_id] = [start, end]
660
661        return renumbered_region_of_interest
662
663    def residue_map(
664        self,
665        token_chain_ids: list,
666        token_res_ids: list,
667        token_atom_names: list,
668        depth: ResidueMapDepth = ResidueMapDepth.ATOM,
669    ):
670        """Create a map of residue indices to residue numbers and vice-versa.
671
672        `res_idx` is essentially the token index. \n
673        `res_num` is the residue number in the chain which will be same as UniProt
674        numbering if the `offset` is provided correctly or full-length sequences
675        were used in the prediction. \n
676        `offset` informs what is the starting residue number for each chain.
677
678        For example, if the input is as follows: \n
679        ```python
680        token_chain_ids = ['A', 'A', 'A', 'B', 'B']
681        token_res_ids = [1, 1, 2, 1, 2]
682        token_atom_names = ['N', 'CA', 'N', 'N', 'CA']
683        ```
684        The output will be: \n
685        ```python
686        idx_to_num = {
687            0: {'chain_id': 'A', 'res_num': 1, 'atom_name': 'N'},
688            1: {'chain_id': 'A', 'res_num': 1, 'atom_name': 'CA'},
689            2: {'chain_id': 'A', 'res_num': 2, 'atom_name': 'N'},
690            3: {'chain_id': 'B', 'res_num': 1, 'atom_name': 'N'},
691            4: {'chain_id': 'B', 'res_num': 2, 'atom_name': 'CA'},
692        }
693        num_to_idx = {
694            'A': {
695                1: {'N': 0, 'CA': 1},
696                2: {'N': 2},
697            },
698            'B': {
699                1: {'N': 3},
700                2: {'CA': 4},
701            },
702        }
703        ```
704
705        If the offset is provided as: \n
706        ```python
707        offset = {
708            'A': [10, 11],
709            'B': [30, 31],
710        }
711        ```
712        The output will be: \n
713        ```python
714        idx_to_num = {
715            0: {'chain_id': 'A', 'res_num': 10, 'atom_name': 'N'},
716            1: {'chain_id': 'A', 'res_num': 10, 'atom_name': 'CA'},
717            2: {'chain_id': 'A', 'res_num': 11, 'atom_name': 'N'},
718            3: {'chain_id': 'B', 'res_num': 30, 'atom_name': 'N'},
719            4: {'chain_id': 'B', 'res_num': 31, 'atom_name': 'CA'},
720        }
721        num_to_idx = {
722            'A': {
723                10: {'N': 0, 'CA': 1},
724                11: {'N': 2},
725            },
726            'B': {
727                30: {'N': 3},
728                31: {'CA': 4},
729            },
730        }
731        ```
732
733        Arguments:
734
735        - **token_chain_ids (list)**:<br />
736            Token chain IDs. Same as provided in the AF3 JSON output.
737
738        - **token_res_ids (list)**:<br />
739            Token residue IDs. Same as provided in the AF3 JSON output.
740
741        - **token_atom_names (list)**:<br />
742            Token atom names.
743
744        Returns:
745
746        - **(tuple)**:<br />
747
748            - `idx_to_num (Dict)`:
749                Dictionary mapping token indices to token numbers.
750
751            - `num_to_idx (Dict)`:
752                Dictionary mapping token numbers to token indices.
753        """
754
755        idx_to_num = {}
756        num_to_idx = defaultdict(dict)
757
758        for token_idx, (chain_id, token_num, atom_name) in enumerate(
759            zip(token_chain_ids, token_res_ids, token_atom_names)
760        ):
761
762            token_num = self.renumber_chain_res_num(
763                chain_res_num=token_num,
764                chain_id=chain_id,
765            )
766
767            idx_to_num[token_idx] = {
768                ResidueMapKeys.CHAIN_ID: chain_id,
769                ResidueMapKeys.TOKEN_NUM: token_num,
770            }
771
772            if depth == ResidueMapDepth.ATOM:
773                idx_to_num[token_idx][ResidueMapKeys.ATOM_NAME] = atom_name
774
775            if token_num not in num_to_idx[chain_id]:
776                if depth == ResidueMapDepth.ATOM:
777                    num_to_idx[chain_id][token_num] = {
778                        atom_name: token_idx
779                    }
780                elif depth == ResidueMapDepth.RESIDUE:
781                    num_to_idx[chain_id][token_num] = token_idx
782            else:
783                if depth == ResidueMapDepth.ATOM:
784                    num_to_idx[chain_id][token_num][atom_name] = token_idx
785                elif depth == ResidueMapDepth.RESIDUE:
786                    warnings.warn(
787                        f"""
788
789                        Multiple tokens found for residue number {token_num}
790                        in chain {chain_id}. Overwriting the previous token index.
791                        """
792                    )
793                    num_to_idx[chain_id][token_num] = token_idx
794
795        return idx_to_num, num_to_idx
796
797
798class ResidueSelect(Select):
799    """Class to select residues in the structure based on the input dictionary.
800
801    This is achieved by overriding the `Bio.PDB.Select.accept_residue` method
802    from Biopython.
803    """
804
805    confident_residues: Dict
806    """ Dictionary containing the chain ID as key and a list of residue
807    numbers as value.<br />
808    e.g. `{"A": [1, 2, 3], "B": [4, 5, 6]}`
809    """
810
811    def __init__(
812        self,
813        confident_residues: Dict
814    ):
815        self.confident_residues = confident_residues
816
817    @overload
818    def accept_residue(
819        self,
820        residue: Bio.PDB.Residue.Residue
821    ) -> bool:
822        ...
823
824    @overload
825    def accept_residue(self, residue):
826        """Overload this to reject residues for output."""
827        return 1
828
829    def accept_residue(
830        self,
831        residue: Bio.PDB.Residue.Residue
832    ) -> bool:
833        """Accept the residue if it's in `confident_residues`.
834
835        Arguments:
836
837        - **residue (Bio.PDB.Residue.Residue)**:<br />
838            Biopython residue object.
839
840        Returns:
841
842        - **(bool)**:<br />
843            `True` if the residue is in the `confident_residues`.
844        """
845
846        chain = residue.parent
847        if not isinstance(chain, Chain):
848            raise TypeError(
849                f"Expected a Bio.PDB.Chain.Chain object, got {type(chain)}"
850            )
851        chain_id = chain.id
852
853        return residue.id[1] in self.confident_residues.get(chain_id, [])