from __future__ import annotations import functools import io import os from dataclasses import dataclass from datetime import datetime from typing import Union import biotite.structure as bs import biotite.structure.io.pdbx as pdbx from . import esmfold2_residue_constants as residue_constants # Define PathOrBuffer for the opensource version PathOrBuffer = Union[str, os.PathLike, io.StringIO] class NoProteinError(Exception): pass @dataclass class Residue: residue_number: int | None = None insertion_code: str = "" hetflag: bool = False @dataclass class MmcifHeader: release_date: datetime | None = None resolution: float | None = None structure_method: str = "UNKNOWN" class MmcifWrapper: def __init__(self, id: str | None = None): self.id: str = id or "" self.raw: pdbx.CIFFile | None = None self.structure: bs.AtomArray self.header: MmcifHeader = MmcifHeader() self.entities: dict[int, list[str]] = {} self.chain_to_seqres: dict[str, str] = {} self.seqres_to_structure: dict[str, dict[int, Residue]] = {} @classmethod def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper: obj = cls(id=id) obj._load(path) return obj def _load(self, path: PathOrBuffer, fileid: str | None = None): """Load mmCIF data from file.""" self.raw = pdbx.CIFFile.read(path) self._parse_structure() self._parse_header() self._parse_entities() self._parse_sequences() def _parse_structure(self): """Parse the atomic structure from mmCIF.""" try: structure = pdbx.get_structure(self.raw, model=1) if structure is None or not isinstance(structure, bs.AtomArray): raise NoProteinError("No structure found in mmCIF file") if len(structure) == 0: raise NoProteinError("Empty structure in mmCIF file") self.structure = structure except Exception as e: raise ValueError(f"Failed to parse structure: {e}") def _parse_header(self): """Parse header information from mmCIF.""" if not self.raw: return try: # Get the first (and usually only) block block = self.raw.block # Parse release date if "pdbx_database_status" in block: status_cat = block["pdbx_database_status"] if "recvd_initial_deposition_date" in status_cat: date_str = status_cat["recvd_initial_deposition_date"].as_item() if date_str and date_str != "?": try: self.header.release_date = datetime.strptime( date_str, "%Y-%m-%d" ) except ValueError: pass # Parse resolution if "refine" in block: refine_cat = block["refine"] if "ls_d_res_high" in refine_cat: res_str = refine_cat["ls_d_res_high"].as_item() if res_str and res_str != "?": try: self.header.resolution = float(res_str) except ValueError: pass # Parse structure method if "exptl" in block: exptl_cat = block["exptl"] if "method" in exptl_cat: method = exptl_cat["method"].as_item() if method and method != "?": self.header.structure_method = method.upper() except Exception: # If parsing fails, keep default values pass def _parse_entities(self): """Parse entity information and map to chains.""" if not self.raw: return try: block = self.raw.block # Parse entity information if "entity" in block: entity_cat = block["entity"] entity_ids = entity_cat["id"].as_array(str) entity_types = entity_cat["type"].as_array(str) # Initialize entities dict with all entities (not just polymers) for i, (entity_id, entity_type) in enumerate( zip(entity_ids, entity_types) ): self.entities[int(entity_id)] = [] # Map polymer chains to entities using entity_poly if "entity_poly" in block: poly_cat = block["entity_poly"] entity_ids = poly_cat["entity_id"].as_array(str) chain_lists = poly_cat["pdbx_strand_id"].as_array(str) for entity_id, chain_list in zip(entity_ids, chain_lists): entity_id = int(entity_id) # Chain list is comma-separated chains = [c.strip() for c in chain_list.split(",") if c.strip()] if entity_id in self.entities: self.entities[entity_id] = chains # Map non-polymer chains using struct_asym for entities not covered by entity_poly if "struct_asym" in block: asym_cat = block["struct_asym"] asym_ids = asym_cat["id"].as_array(str) entity_ids = asym_cat["entity_id"].as_array(str) for asym_id, entity_id in zip(asym_ids, entity_ids): entity_id = int(entity_id) # Only add if entity exists but has no chains yet (non-polymer entities) if entity_id in self.entities and not self.entities[entity_id]: self.entities[entity_id].append(asym_id) except Exception: # If parsing fails, try to infer from structure if ( self.structure and hasattr(self.structure, "chain_id") and self.structure.chain_id is not None and hasattr(self.structure.chain_id, "__iter__") ): chain_ids = list(set(self.structure.chain_id)) self.entities = {1: chain_ids} def _parse_sequences(self): """Parse sequence information from mmCIF.""" if not self.raw: return block = self.raw.block # Parse polymer sequences if "entity_poly" in block: poly_cat = block["entity_poly"] entity_ids = poly_cat["entity_id"].as_array(str) sequences = poly_cat["pdbx_seq_one_letter_code_can"].as_array(str) chain_lists = poly_cat["pdbx_strand_id"].as_array(str) for entity_id, sequence, chain_list in zip( entity_ids, sequences, chain_lists ): # Clean up sequence (remove whitespace and newlines) clean_seq = "".join(sequence.split()) chains = [c.strip() for c in chain_list.split(",") if c.strip()] for chain_id in chains: self.chain_to_seqres[chain_id] = clean_seq # Parse sequence to structure mapping if "pdbx_poly_seq_scheme" in block: seq_cat = block["pdbx_poly_seq_scheme"] asym_ids = seq_cat["asym_id"].as_array(str) # Internal chain IDs seq_positions = seq_cat["seq_id"].as_array(str) auth_seq_nums = seq_cat["auth_seq_num"].as_array(str) ins_codes = ( seq_cat["pdb_ins_code"].as_array(str) if "pdb_ins_code" in seq_cat else [""] * len(asym_ids) ) hetflags = ( seq_cat["hetflag"].as_array(str) if "hetflag" in seq_cat else ["N"] * len(asym_ids) ) # Get author chain IDs if available auth_chain_ids = ( seq_cat["pdb_strand_id"].as_array(str) if "pdb_strand_id" in seq_cat else asym_ids # Fallback to internal IDs ) # Build mapping from internal chain ID to author chain ID asym_to_auth_mapping = {} for asym_id, auth_id in zip(asym_ids, auth_chain_ids): asym_to_auth_mapping[asym_id] = auth_id # Group by internal chain ID first, then map to author chain ID chain_data = {} for asym_id, seq_pos, auth_seq, ins_code, hetflag in zip( asym_ids, seq_positions, auth_seq_nums, ins_codes, hetflags ): if asym_id not in chain_data: chain_data[asym_id] = {} try: seq_index = int(seq_pos) - 1 # Convert to 0-based indexing res_num = int(auth_seq) if auth_seq != "?" else None except ValueError: continue if res_num is not None: # Convert mmCIF "." and "?" to empty string clean_ins_code = "" if ins_code in [".", "?"] else ins_code else: clean_ins_code = "" res_num = None is_het = hetflag.upper() == "Y" # type: ignore chain_data[asym_id][seq_index] = Residue( residue_number=res_num, insertion_code=clean_ins_code, # type: ignore hetflag=is_het, ) # Handle cases where multiple residues have the same auth_seq_num # by adjusting residue numbers to be unique within each chain for asym_id, residue_data in chain_data.items(): # Check if there are duplicate residue numbers in this chain positions_with_same_num = {} for seq_idx, res_at_pos in residue_data.items(): if res_at_pos.residue_number is not None: res_num = res_at_pos.residue_number if res_num not in positions_with_same_num: positions_with_same_num[res_num] = [] positions_with_same_num[res_num].append(seq_idx) # Fix duplicate residue numbers by making them sequential for res_num, seq_indices in positions_with_same_num.items(): if len(seq_indices) > 1: # Multiple residues have the same residue number # Make them sequential starting from the original number seq_indices.sort() # Ensure consistent ordering for i, seq_idx in enumerate(seq_indices): original_pos = residue_data[seq_idx] new_pos = Residue( residue_number=res_num + i, insertion_code=original_pos.insertion_code, hetflag=original_pos.hetflag, ) residue_data[seq_idx] = new_pos # Create ordered mappings using author chain IDs for asym_id in chain_data: auth_chain_id = asym_to_auth_mapping.get(asym_id, asym_id) if auth_chain_id in self.chain_to_seqres: seq_len = len(self.chain_to_seqres[auth_chain_id]) ordered_mapping = {} for i in range(seq_len): if i in chain_data[asym_id]: ordered_mapping[i] = chain_data[asym_id][i] else: # Missing residue - no structure coordinates ordered_mapping[i] = Residue( residue_number=None, insertion_code="", hetflag=False ) self.seqres_to_structure[auth_chain_id] = ordered_mapping else: # Handle case where auth_chain_id is not in chain_to_seqres # This can happen if the chain is not a polymer or if there's a parsing issue # Create a basic mapping based on the chain_data if chain_data[asym_id]: # Sort by sequence index to create ordered mapping sorted_indices = sorted(chain_data[asym_id].keys()) ordered_mapping = {} for i, seq_idx in enumerate(sorted_indices): ordered_mapping[i] = chain_data[asym_id][seq_idx] self.seqres_to_structure[auth_chain_id] = ordered_mapping # Ensure all chains have complete mappings for chain_id in self.chain_to_seqres: if chain_id not in self.seqres_to_structure: seq_len = len(self.chain_to_seqres[chain_id]) self.seqres_to_structure[chain_id] = { i: Residue(residue_number=None, insertion_code="", hetflag=False) for i in range(seq_len) } else: # Fill in any missing indices seq_len = len(self.chain_to_seqres[chain_id]) mapping = self.seqres_to_structure[chain_id] for i in range(seq_len): if i not in mapping: mapping[i] = Residue( residue_number=None, insertion_code="", hetflag=False ) # Fallback: create basic mappings from structure for missing chains if ( self.structure and hasattr(self.structure, "chain_id") and self.structure.chain_id is not None and hasattr(self.structure.chain_id, "__iter__") ): for chain_id in set(self.structure.chain_id): if chain_id not in self.seqres_to_structure: chain_structure = self.structure[ self.structure.chain_id == chain_id ] if ( hasattr(chain_structure, "res_id") and chain_structure.res_id is not None and hasattr(chain_structure.res_id, "__iter__") ): residue_ids = list(set(chain_structure.res_id)) residue_ids.sort() self.seqres_to_structure[chain_id] = { i: Residue( residue_number=res_id, insertion_code="", hetflag=False ) for i, res_id in enumerate(residue_ids) } def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]: """Parse non-polymer coordinates from mmCIF block data.""" nonpoly_coords = {} # Get non-polymer entities from the mmCIF block assert self.raw is not None block = self.raw.block nonpoly_entities = set() # Find non-polymer entities if "entity" in block: entity_cat = block["entity"] entity_ids = entity_cat["id"].as_array(str) entity_types = entity_cat["type"].as_array(str) for entity_id, entity_type in zip(entity_ids, entity_types): if entity_type.upper() in ["NON-POLYMER", "WATER", "BRANCHED"]: nonpoly_entities.add(entity_id) # Map entities to chains for non-polymers entity_to_chains = {} if "pdbx_entity_nonpoly" in block: nonpoly_cat = block["pdbx_entity_nonpoly"] entity_ids = nonpoly_cat["entity_id"].as_array(str) comp_ids = nonpoly_cat["comp_id"].as_array(str) for entity_id, comp_id in zip(entity_ids, comp_ids): if entity_id in nonpoly_entities: entity_to_chains[entity_id] = comp_id # Get atom site information for non-polymers if "atom_site" in block: atom_cat = block["atom_site"] atom_chain_ids = atom_cat["label_asym_id"].as_array(str) atom_entity_ids = atom_cat["label_entity_id"].as_array(str) atom_comp_ids = atom_cat["label_comp_id"].as_array(str) # Group non-polymer atoms by entity and chain nonpoly_atom_groups = {} for i, (chain_id, entity_id, comp_id) in enumerate( zip(atom_chain_ids, atom_entity_ids, atom_comp_ids) ): if entity_id in nonpoly_entities: key = (comp_id, chain_id) if key not in nonpoly_atom_groups: nonpoly_atom_groups[key] = [] nonpoly_atom_groups[key].append(i) # Extract coordinates for each non-polymer group for (comp_id, chain_id), atom_indices in nonpoly_atom_groups.items(): # Match atoms by comparing chain_id and residue name structure_mask = (self.structure.chain_id == chain_id) & ( self.structure.res_name == comp_id ) if structure_mask.any(): nonpoly_array = self.structure[structure_mask] if ( isinstance(nonpoly_array, (bs.AtomArray, bs.AtomArrayStack)) and len(nonpoly_array) > 0 ): nonpoly_coords[(comp_id, chain_id)] = nonpoly_array return nonpoly_coords def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]: """Fallback method to extract heteroatoms directly from structure.""" nonpoly_coords = {} if not (self.structure and hasattr(self.structure, "chain_id")): return nonpoly_coords # Create set of standard residues from residue_constants standard_residues = set(residue_constants.resnames[:-1]) # Exclude 'UNK' standard_residues.update({"A", "C", "G", "T", "U"}) # Add nucleic acids if hasattr(self.structure, "chain_id") and self.structure.chain_id is not None: for chain_id in set(self.structure.chain_id): chain_structure = self.structure[self.structure.chain_id == chain_id] # Find non-standard residues if ( hasattr(chain_structure, "res_name") and chain_structure.res_name is not None and hasattr(chain_structure.res_name, "__iter__") ): for res_name in set(chain_structure.res_name): if res_name not in standard_residues: res_mask = (chain_structure.chain_id == chain_id) & ( chain_structure.res_name == res_name ) if res_mask.any() and isinstance( chain_structure, (bs.AtomArray, bs.AtomArrayStack) ): nonpoly_array = chain_structure[res_mask] nonpoly_coords[(res_name, chain_id)] = nonpoly_array return nonpoly_coords @functools.cached_property def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]: """ Extract non-polymer coordinates (ligands, cofactors, etc.) from mmCIF structure. Returns a dictionary mapping (nonpolymer_info, chain_id) tuples to AtomArrays. """ if not self.structure or not self.raw: return {} try: return self._parse_nonpoly_from_mmcif() except Exception: return self._parse_nonpoly_fallback()