| 309 | |
| 310 | |
| 311 | def classify_chains(chains, residue_types): |
| 312 | nuc_residue_set = {"DA", "DC", "DT", "DG", "A", "C", "U", "G"} |
| 313 | chain_types = {} |
| 314 | |
| 315 | # Get unique chains and iterate over them |
| 316 | _, first_idx = np.unique(chains, return_index=True) |
| 317 | unique_chains = chains[np.sort(first_idx)] |
| 318 | |
| 319 | for chain in unique_chains: |
| 320 | # Find indices where the current chain is located |
| 321 | indices = np.where(chains == chain)[0] |
| 322 | # Get the residues for these indices |
| 323 | chain_residues = residue_types[indices] |
| 324 | # Count nucleic acid residues |
| 325 | nuc_count = sum(residue in nuc_residue_set for residue in chain_residues) |
| 326 | |
| 327 | # Determine if the chain is a nucleic acid or protein |
| 328 | chain_types[chain] = 'nucleic_acid' if nuc_count > 0 else 'protein' |
| 329 | |
| 330 | return chain_types |
| 331 | |
| 332 | |
| 333 | # Load residues from AlphaFold PDB or mmCIF file into lists; each residue is a dictionary |