import itertools import os import pickle from pathlib import Path import random from typing import Dict, List import zipfile import numpy as np import torch from rdkit.Chem import Mol, AllChem from tqdm import tqdm from boltzgen.data import const from boltzgen.data.pad import pad_dim from boltzgen.model.loss.confidence import lddt_dist MOLDIR_ZIP_CACHE: Dict[ tuple[int, Path], zipfile.ZipFile ] = {} # moldir -> open zip file def _get_zipfile(moldir: Path) -> zipfile.ZipFile: """ Retrieve a cached ZipFile object for the given molecule directory zip file. If the ZipFile for the provided path and current process does not exist in the cache, it is created and stored. Otherwise, the cached instance is returned. Parameters ---------- moldir : Path Path to the .zip file containing molecule data. Returns ------- zipfile.ZipFile An open ZipFile object for reading molecule data. Notes ----- The cache is keyed by both process ID and zip file path to prevent issues with forked processes. """ pid = os.getpid() key = (pid, moldir) zf = MOLDIR_ZIP_CACHE.get(key) if zf is None: zf = zipfile.ZipFile(moldir, "r") MOLDIR_ZIP_CACHE[key] = zf return zf def load_molecules(moldir: str, molecules: List[str]) -> Dict[str, Mol]: """Load the given input data. Parameters ---------- moldir : str The path to the molecules directory or zip file. molecules : list[str] The molecules to load. Returns ------- dict[str, Mol] The loaded molecules. """ moldir = Path(moldir) loaded_mols: Dict[str, Mol] = {} if moldir.is_file() and moldir.suffix == ".zip": zip_file = _get_zipfile(moldir) for molecule in molecules: pkl_filename = f"{molecule}.pkl" with zip_file.open(pkl_filename, "r") as f: loaded_mols[molecule] = pickle.load(f) # noqa: S301 elif moldir.is_dir(): for molecule in molecules: path = moldir / f"{molecule}.pkl" with path.open("rb") as f: loaded_mols[molecule] = pickle.load(f) # noqa: S301 else: raise ValueError(f"Invalid moldir. Expected directory or zip file: {moldir}") return loaded_mols def load_canonicals(moldir: str) -> dict[str, Mol]: """Load the given input data. Parameters ---------- moldir : str The molecules to load. Returns ------- dict[str, Mol] The loaded molecules. """ return load_molecules(moldir, const.canonical_tokens) def load_all_molecules(moldir: str) -> dict[str, Mol]: """Load all molecular data from directory or zip file. Parameters ---------- moldir : str The path to the molecules directory or zip file. Returns ------- dict[str, Mol] The loaded molecules. """ moldir = Path(moldir) loaded_mols = {} if moldir.is_file() and moldir.suffix == ".zip": with zipfile.ZipFile(moldir, "r") as zip_file: pkl_files = [name for name in zip_file.namelist() if name.endswith(".pkl")] for pkl_filename in tqdm( pkl_files, desc="Loading molecules from zip", leave=False ): mol_name = Path(pkl_filename).stem with zip_file.open(pkl_filename) as f: loaded_mols[mol_name] = pickle.load(f) # noqa: S301 else: files = list(moldir.glob("*.pkl")) for path in tqdm( files, total=len(files), desc="Loading molecules", leave=False ): mol_name = path.stem with path.open("rb") as f: loaded_mols[mol_name] = pickle.load(f) # noqa: S301 return loaded_mols def get_symmetries(mols: dict[str, Mol]) -> Dict: # noqa: PLR0912 """Create a dictionary for the ligand symmetries. Parameters ---------- path : str The path to the ligand symmetries. Returns ------- Dict The ligand symmetries. """ symmetries = {} for key, mol in mols.items(): try: sym = pickle.loads(bytes.fromhex(mol.GetProp("symmetries"))) # noqa: S301 if mol.HasProp("pb_edge_index"): edge_index = pickle.loads( bytes.fromhex(mol.GetProp("pb_edge_index")) ).astype(np.int64) # noqa: S301 lower_bounds = pickle.loads( bytes.fromhex(mol.GetProp("pb_lower_bounds")) ) # noqa: S301 upper_bounds = pickle.loads( bytes.fromhex(mol.GetProp("pb_upper_bounds")) ) # noqa: S301 bond_mask = pickle.loads(bytes.fromhex(mol.GetProp("pb_bond_mask"))) # noqa: S301 angle_mask = pickle.loads(bytes.fromhex(mol.GetProp("pb_angle_mask"))) # noqa: S301 else: edge_index = np.empty((2, 0), dtype=np.int64) lower_bounds = np.array([], dtype=np.float32) upper_bounds = np.array([], dtype=np.float32) bond_mask = np.array([], dtype=np.float32) angle_mask = np.array([], dtype=np.float32) if mol.HasProp("chiral_atom_index"): chiral_atom_index = pickle.loads( bytes.fromhex(mol.GetProp("chiral_atom_index")) ).astype(np.int64) chiral_check_mask = pickle.loads( bytes.fromhex(mol.GetProp("chiral_check_mask")) ).astype(np.int64) chiral_atom_orientations = pickle.loads( bytes.fromhex(mol.GetProp("chiral_atom_orientations")) ) else: chiral_atom_index = np.empty((4, 0), dtype=np.int64) chiral_check_mask = np.array([], dtype=bool) chiral_atom_orientations = np.array([], dtype=bool) if mol.HasProp("stereo_bond_index"): stereo_bond_index = pickle.loads( bytes.fromhex(mol.GetProp("stereo_bond_index")) ).astype(np.int64) stereo_check_mask = pickle.loads( bytes.fromhex(mol.GetProp("stereo_check_mask")) ).astype(np.int64) stereo_bond_orientations = pickle.loads( bytes.fromhex(mol.GetProp("stereo_bond_orientations")) ) else: stereo_bond_index = np.empty((4, 0), dtype=np.int64) stereo_check_mask = np.array([], dtype=bool) stereo_bond_orientations = np.array([], dtype=bool) if mol.HasProp("aromatic_5_ring_index"): aromatic_5_ring_index = pickle.loads( bytes.fromhex(mol.GetProp("aromatic_5_ring_index")) ).astype(np.int64) else: aromatic_5_ring_index = np.empty((5, 0), dtype=np.int64) if mol.HasProp("aromatic_6_ring_index"): aromatic_6_ring_index = pickle.loads( bytes.fromhex(mol.GetProp("aromatic_6_ring_index")) ).astype(np.int64) else: aromatic_6_ring_index = np.empty((6, 0), dtype=np.int64) if mol.HasProp("planar_double_bond_index"): planar_double_bond_index = pickle.loads( bytes.fromhex(mol.GetProp("planar_double_bond_index")) ).astype(np.int64) else: planar_double_bond_index = np.empty((6, 0), dtype=np.int64) atom_names = [atom.GetProp("name") for atom in mol.GetAtoms()] symmetries[key] = ( sym, atom_names, edge_index, lower_bounds, upper_bounds, bond_mask, angle_mask, chiral_atom_index, chiral_check_mask, chiral_atom_orientations, stereo_bond_index, stereo_check_mask, stereo_bond_orientations, aromatic_5_ring_index, aromatic_6_ring_index, planar_double_bond_index, ) except Exception as e: # noqa: BLE001, PERF203, S110 pass return symmetries def compute_symmetry_idx_dictionary(data): # Compute the symmetry index dictionary total_count = 0 all_coords = [] for i, chain in enumerate(data.chains): chain.start_idx = total_count for j, token in enumerate(chain.tokens): token.start_idx = total_count - chain.start_idx all_coords.extend( [[atom.coords.x, atom.coords.y, atom.coords.z] for atom in token.atoms] ) total_count += len(token.atoms) return all_coords def get_current_idx_list(data): idx = [] for chain in data.chains: if chain.in_crop: for token in chain.tokens: if token.in_crop: idx.extend( [ chain.start_idx + token.start_idx + i for i in range(len(token.atoms)) ] ) return idx def all_different_after_swap(l): final = [s[-1] for s in l] return len(final) == len(set(final)) def minimum_lddt_symmetry_coords( coords: torch.Tensor, feats: dict, index_batch: int, ): all_coords = feats["all_coords"][index_batch].unsqueeze(0).to(coords) all_resolved_mask = ( feats["all_resolved_mask"][index_batch].to(coords).to(torch.bool) ) crop_to_all_atom_map = ( feats["crop_to_all_atom_map"][index_batch].to(coords).to(torch.long) ) chain_symmetries = feats["chain_swaps"][index_batch] amino_acids_symmetries = feats["amino_acids_symmetries"][index_batch] ligand_symmetries = feats["ligand_symmetries"][index_batch] dmat_predicted = torch.cdist( coords[:, : len(crop_to_all_atom_map)], coords[:, : len(crop_to_all_atom_map)] ) # Check best symmetry on chain swap best_true_coords = all_coords[:, crop_to_all_atom_map].clone() best_true_resolved_mask = all_resolved_mask[crop_to_all_atom_map].clone() best_lddt = -1.0 for c in chain_symmetries: true_all_coords = all_coords.clone() true_all_resolved_mask = all_resolved_mask.clone() for start1, end1, start2, end2, chainidx1, chainidx2 in c: true_all_coords[:, start1:end1] = all_coords[:, start2:end2] true_all_resolved_mask[start1:end1] = all_resolved_mask[start2:end2] true_coords = true_all_coords[:, crop_to_all_atom_map] true_resolved_mask = true_all_resolved_mask[crop_to_all_atom_map] dmat_true = torch.cdist(true_coords, true_coords) pair_mask = ( true_resolved_mask[:, None] * true_resolved_mask[None, :] * (1 - torch.eye(len(true_resolved_mask))).to(true_resolved_mask) ) lddt = lddt_dist( dmat_predicted, dmat_true, pair_mask, cutoff=15.0, per_atom=False )[0] lddt = lddt.item() if lddt > best_lddt and torch.sum(true_resolved_mask) > 3: best_lddt = lddt best_true_coords = true_coords best_true_resolved_mask = true_resolved_mask # atom symmetries (nucleic acid and protein residues), resolved greedily without recomputing alignment true_coords = best_true_coords.clone() true_resolved_mask = best_true_resolved_mask.clone() for symmetric_amino_or_lig in amino_acids_symmetries + ligand_symmetries: best_lddt_improvement = 0.0 indices = set() for c in symmetric_amino_or_lig: for i, j in c: indices.add(i) indices = sorted(list(indices)) indices = torch.from_numpy(np.asarray(indices)).to(true_coords.device).long() pred_coords_subset = coords[:, : len(crop_to_all_atom_map)][:, indices] sub_dmat_pred = torch.cdist( coords[:, : len(crop_to_all_atom_map)], pred_coords_subset ) for c in symmetric_amino_or_lig: # starting from greedy best, try to swap the atoms new_true_coords = true_coords.clone() new_true_resolved_mask = true_resolved_mask.clone() for i, j in c: new_true_coords[:, i] = true_coords[:, j] new_true_resolved_mask[i] = true_resolved_mask[j] true_coords_subset = true_coords[:, indices] new_true_coords_subset = new_true_coords[:, indices] sub_dmat_true = torch.cdist(true_coords, true_coords_subset) sub_dmat_new_true = torch.cdist(new_true_coords, new_true_coords_subset) sub_true_pair_lddt = ( true_resolved_mask[:, None] * true_resolved_mask[None, indices] ) sub_true_pair_lddt[indices] = ( sub_true_pair_lddt[indices] * (1 - torch.eye(len(indices))).to(sub_true_pair_lddt).bool() ) sub_new_true_pair_lddt = ( new_true_resolved_mask[:, None] * new_true_resolved_mask[None, indices] ) sub_new_true_pair_lddt[indices] = ( sub_new_true_pair_lddt[indices] * (1 - torch.eye(len(indices))).to(sub_true_pair_lddt).bool() ) lddt, total = lddt_dist( sub_dmat_pred, sub_dmat_true, sub_true_pair_lddt, cutoff=15.0, per_atom=False, ) new_lddt, new_total = lddt_dist( sub_dmat_pred, sub_dmat_new_true, sub_new_true_pair_lddt, cutoff=15.0, per_atom=False, ) lddt_improvement = new_lddt - lddt if lddt_improvement > best_lddt_improvement: best_true_coords = new_true_coords best_true_resolved_mask = new_true_resolved_mask best_lddt_improvement = lddt_improvement # greedily update best coordinates after each amino acid true_coords = best_true_coords.clone() true_resolved_mask = best_true_resolved_mask.clone() # Recomputing alignment true_coords = pad_dim(true_coords, 1, coords.shape[1] - true_coords.shape[1]) true_resolved_mask = pad_dim( true_resolved_mask, 0, coords.shape[1] - true_resolved_mask.shape[0], ) return true_coords, true_resolved_mask.unsqueeze(0) def compute_single_distogram_loss(pred, target, mask): # Compute the distogram loss errors = -1 * torch.sum( target * torch.nn.functional.log_softmax(pred, dim=-1), dim=-1, ) denom = 1e-5 + torch.sum(mask, dim=(-1, -2)) mean = errors * mask mean = torch.sum(mean, dim=-1) mean = mean / denom[..., None] batch_loss = torch.sum(mean, dim=-1) global_loss = torch.mean(batch_loss) return global_loss def minimum_lddt_symmetry_dist( pred_distogram: torch.Tensor, feats: dict, index_batch: int, ): # Note: for now only ligand symmetries are resolved disto_target = feats["disto_target"][index_batch] mask = feats["token_disto_mask"][index_batch] mask = mask[None, :] * mask[:, None] mask = mask * (1 - torch.eye(mask.shape[1])).to(disto_target) coords = feats["coords"][index_batch] ligand_symmetries = feats["ligand_symmetries"][index_batch] atom_to_token_map = feats["atom_to_token"][index_batch].argmax(dim=-1) # atom symmetries, resolved greedily without recomputing alignment for symmetric_amino_or_lig in ligand_symmetries: best_c, best_disto, best_loss_improvement = None, None, 0.0 for c in symmetric_amino_or_lig: # starting from greedy best, try to swap the atoms new_disto_target = disto_target.clone() indices = [] # fix the distogram by replacing first the columns then the rows disto_temp = new_disto_target.clone() for i, j in c: new_disto_target[:, atom_to_token_map[i]] = disto_temp[ :, atom_to_token_map[j] ] indices.append(atom_to_token_map[i].item()) disto_temp = new_disto_target.clone() for i, j in c: new_disto_target[atom_to_token_map[i], :] = disto_temp[ atom_to_token_map[j], : ] indices = ( torch.from_numpy(np.asarray(indices)).to(disto_target.device).long() ) pred_distogram_subset = pred_distogram[:, indices] disto_target_subset = disto_target[:, indices] new_disto_target_subset = new_disto_target[:, indices] mask_subset = mask[:, indices] loss = compute_single_distogram_loss( pred_distogram_subset, disto_target_subset, mask_subset ) new_loss = compute_single_distogram_loss( pred_distogram_subset, new_disto_target_subset, mask_subset ) loss_improvement = (loss - new_loss) * len(indices) if loss_improvement > best_loss_improvement: best_c = c best_disto = new_disto_target best_loss_improvement = loss_improvement # greedily update best coordinates after each ligand if best_loss_improvement > 0: disto_target = best_disto.clone() old_coords = coords.clone() for i, j in best_c: coords[:, i] = old_coords[:, j] # update features to be used in diffusion and in distogram loss feats["disto_target"][index_batch] = disto_target feats["coords"][index_batch] = coords return def compute_all_coords_mask(structure): # Compute all coords, crop mask and add start_idx to structure total_count = 0 all_coords = [] all_coords_crop_mask = [] all_resolved_mask = [] for i, chain in enumerate(structure.chains): chain.start_idx = total_count for j, token in enumerate(chain.tokens): token.start_idx = total_count - chain.start_idx all_coords.extend( [[atom.coords.x, atom.coords.y, atom.coords.z] for atom in token.atoms] ) all_coords_crop_mask.extend( [token.in_crop for _ in range(len(token.atoms))] ) all_resolved_mask.extend( [token.is_present for _ in range(len(token.atoms))] ) total_count += len(token.atoms) if len(all_coords_crop_mask) != len(all_resolved_mask): pass return all_coords, all_coords_crop_mask, all_resolved_mask def get_chain_symmetries(cropped, backbone_only, atom14, atom37, max_n_symmetries=100): # get all coordinates and resolved mask structure = cropped.structure all_coords = [] all_resolved_mask = [] original_atom_idx = [] chain_atom_idx = [] chain_atom_num = [] chain_in_crop = [] chain_asym_id = [] new_atom_idx = 0 for chain in structure.chains: atom_idx, atom_num = ( chain["atom_idx"], # Global index of first atom in the chain chain["atom_num"], # Number of atoms in the chain ) # compute coordinates and resolved mask # Whether each atom in the chain is actually resolved resolved_mask = structure.atoms["is_present"][atom_idx : atom_idx + atom_num] coords = structure.atoms["coords"][atom_idx : atom_idx + atom_num] in_crop = False for token in cropped.tokens: if token["asym_id"] == chain["asym_id"]: in_crop = True break all_coords.append(coords) all_resolved_mask.append(resolved_mask) original_atom_idx.append(atom_idx) chain_atom_idx.append(new_atom_idx) chain_atom_num.append(atom_num) chain_in_crop.append(in_crop) chain_asym_id.append(chain["asym_id"]) new_atom_idx += atom_num all_coords = np.concatenate(all_coords, axis=0) # Compute backmapping from token to all coords crop_to_all_atom_map = [] for token in cropped.tokens: chain_idx = chain_asym_id.index(token["asym_id"]) start = ( chain_atom_idx[chain_idx] - original_atom_idx[chain_idx] + token["atom_idx"] ) # add logic for backbone_only and atom14 if bool(token["is_standard"]) and bool(token["design_mask"]) and backbone_only: if token["mol_type"] == const.chain_type_ids["PROTEIN"]: atom_num = min(4, token["atom_num"]) elif token["mol_type"] == const.chain_type_ids["DNA"]: atom_num = min(11, token["atom_num"]) elif token["mol_type"] == const.chain_type_ids["RNA"]: atom_num = min(12, token["atom_num"]) elif ( bool(token["is_standard"]) and bool(token["design_mask"]) and (atom14 or atom37) ): if token["mol_type"] == const.chain_type_ids["PROTEIN"]: atom_num = 14 if atom14 else 37 elif token["mol_type"] == const.chain_type_ids["DNA"]: atom_num = 22 elif token["mol_type"] == const.chain_type_ids["RNA"]: atom_num = 23 else: atom_num = token["atom_num"] res_type = const.tokens[token["res_type"]] # special handling for atom14: real atoms map back with offset, fake atoms map back to the center rep atom if ( bool(token["is_standard"]) and bool(token["design_mask"]) and (atom14 or atom37) ): if atom14: local_array = np.arange(start, start + atom_num) placements = np.array(const.fake_atom_placements[res_type]) oxygen_offset = const.ref_atoms[res_type].index("O") nitrogen_offset = const.ref_atoms[res_type].index("N") local_array[placements == "O"] = start + oxygen_offset local_array[placements == "N"] = start + nitrogen_offset elif atom37: ca_offset = const.ref_atoms[res_type].index("CA") local_array = np.array( [ start + const.ref_atoms[res_type].index(name) if name in const.ref_atoms[res_type] else start + ca_offset for name in const.atom_types ] ) else: real_range = np.arange(start, start + token["atom_num"]) # need to add in offset for the rep atom rep_atom_offset = token["center_idx"] - token["atom_idx"] fake_range = np.array( [start + rep_atom_offset] * (atom_num - token["atom_num"]) ) local_array = np.concatenate([real_range, fake_range]) # for backbone_only and without backbone_only or atom14 else: local_array = np.arange(start, start + atom_num) crop_to_all_atom_map.append(local_array) crop_to_all_atom_map = np.concatenate(crop_to_all_atom_map, axis=0) # Compute the connections edge index for covalent bonds all_atom_to_crop_map = np.zeros(all_coords.shape[0], dtype=np.int64) all_atom_to_crop_map[crop_to_all_atom_map.astype(np.int64)] = np.arange( crop_to_all_atom_map.shape[0] ) connections_edge_index = [] crop_atom_set = set(crop_to_all_atom_map.astype(np.int64)) for connection in structure.bonds: if (connection["chain_1"] == connection["chain_2"]) and ( connection["res_1"] == connection["res_2"] ): continue atom_1, atom_2 = connection["atom_1"], connection["atom_2"] # Only include bonds where BOTH atoms are in the crop if atom_1 in crop_atom_set and atom_2 in crop_atom_set: connections_edge_index.append([atom_1, atom_2]) if len(connections_edge_index) > 0: connections_edge_index = np.array(connections_edge_index, dtype=np.int64).T connections_edge_index = all_atom_to_crop_map[connections_edge_index] else: connections_edge_index = np.empty((2, 0)) # Compute the symmetries between chains symmetries = [] swaps = [] for i, chain in enumerate(structure.chains): start = chain_atom_idx[i] end = start + chain_atom_num[i] if chain_in_crop[i]: possible_swaps = [] for j, chain2 in enumerate(structure.chains): start2 = chain_atom_idx[j] end2 = start2 + chain_atom_num[j] if ( chain["entity_id"] == chain2["entity_id"] and end - start == end2 - start2 ): possible_swaps.append((start, end, start2, end2, i, j)) swaps.append(possible_swaps) found = False for symmetry_idx, symmetry in enumerate(symmetries): j = symmetry[0][0] chain2 = structure.chains[j] start2 = chain_atom_idx[j] end2 = start2 + chain_atom_num[j] if ( chain["entity_id"] == chain2["entity_id"] and end - start == end2 - start2 ): symmetries[symmetry_idx].append( (i, start, end, chain_in_crop[i], chain["mol_type"]) ) found = True if not found: symmetries.append([(i, start, end, chain_in_crop[i], chain["mol_type"])]) combinations = itertools.product(*swaps) # to avoid combinatorial explosion, bound the number of combinations even considered combinations = list(itertools.islice(combinations, max_n_symmetries * 10)) # filter for all chains getting a different assignment combinations = [c for c in combinations if all_different_after_swap(c)] if len(combinations) > max_n_symmetries: combinations = random.sample(combinations, max_n_symmetries) if len(combinations) == 0: combinations.append([]) for i in range(len(symmetries) - 1, -1, -1): if not any(chain[3] for chain in symmetries[i]): symmetries.pop(i) features = {} features["all_coords"] = torch.Tensor(all_coords) # axis=1 with ensemble features["all_resolved_mask"] = torch.Tensor( np.concatenate(all_resolved_mask, axis=0) ) features["crop_to_all_atom_map"] = torch.Tensor(crop_to_all_atom_map) features["chain_symmetries"] = symmetries features["connections_edge_index"] = torch.tensor(connections_edge_index) features["chain_swaps"] = combinations return features def get_amino_acids_symmetries(cropped, backbone_only, atom14, atom37): # Compute standard amino-acids symmetries swaps = [] start_index_crop = 0 for token in cropped.tokens: symmetries = const.ref_symmetries.get(const.tokens[token["res_type"]], []) # add atom_num logic for backbone_only and atom14 if bool(token["is_standard"]) and bool(token["design_mask"]) and backbone_only: if token["mol_type"] == const.chain_type_ids["PROTEIN"]: atom_num = min(4, token["atom_num"]) elif token["mol_type"] == const.chain_type_ids["DNA"]: atom_num = min(11, token["atom_num"]) elif token["mol_type"] == const.chain_type_ids["RNA"]: atom_num = min(12, token["atom_num"]) elif ( bool(token["is_standard"]) and bool(token["design_mask"]) and (atom14 or atom37) ): if token["mol_type"] == const.chain_type_ids["PROTEIN"]: atom_num = 14 if atom14 else 37 elif token["mol_type"] == const.chain_type_ids["DNA"]: atom_num = 22 elif token["mol_type"] == const.chain_type_ids["RNA"]: atom_num = 23 else: atom_num = token["atom_num"] if len(symmetries) > 0: residue_swaps = [] for sym in symmetries: sym_new_idx = [ (i + start_index_crop, j + start_index_crop) for i, j in sym ] residue_swaps.append(sym_new_idx) swaps.append(residue_swaps) # start_index_crop += token["atom_num"] start_index_crop += atom_num features = {"amino_acids_symmetries": swaps} return features def slice_valid_index(index, ccd_to_valid_id_array, args=None): index = ccd_to_valid_id_array[index] valid_index_mask = (~np.isnan(index)).all(axis=0) index = index[:, valid_index_mask] if args is None: return index args = (arg[valid_index_mask] for arg in args) return index, args def get_ligand_symmetries(cropped, symmetries, return_physical_metrics=False): # Compute ligand and non-standard amino-acids symmetries structure = cropped.structure added_molecules = {} index_mols = [] atom_count = 0 for token in cropped.tokens: # check if molecule is already added by identifying it through asym_id and res_idx atom_count += token["atom_num"] mol_id = (token["asym_id"], token["res_idx"]) if mol_id in added_molecules: added_molecules[mol_id] += token["atom_num"] continue added_molecules[mol_id] = token["atom_num"] # get the molecule type and indices residue_idx = token["res_idx"] + structure.chains[token["asym_id"]]["res_idx"] mol_name = structure.residues[residue_idx]["name"] atom_idx = structure.residues[residue_idx]["atom_idx"] mol_atom_names = structure.atoms[ atom_idx : atom_idx + structure.residues[residue_idx]["atom_num"] ]["name"] if mol_name not in const.ref_symmetries: index_mols.append( (mol_name, atom_count - token["atom_num"], mol_id, mol_atom_names) ) # for each molecule, get the symmetries molecule_symmetries = [] all_edge_index = [] all_lower_bounds, all_upper_bounds = [], [] all_bond_mask, all_angle_mask = [], [] all_chiral_atom_index, all_chiral_check_mask, all_chiral_atom_orientations = ( [], [], [], ) all_stereo_bond_index, all_stereo_check_mask, all_stereo_bond_orientations = ( [], [], [], ) ( all_aromatic_5_ring_index, all_aromatic_6_ring_index, all_planar_double_bond_index, ) = ( [], [], [], ) for mol_name, start_mol, mol_id, mol_atom_names in index_mols: if not mol_name in symmetries: continue else: swaps = [] ( syms_ccd, mol_atom_names_ccd, edge_index, lower_bounds, upper_bounds, bond_mask, angle_mask, chiral_atom_index, chiral_check_mask, chiral_atom_orientations, stereo_bond_index, stereo_check_mask, stereo_bond_orientations, aromatic_5_ring_index, aromatic_6_ring_index, planar_double_bond_index, ) = symmetries[mol_name] # Get indices of mol_atom_names_ccd that are in mol_atom_names ccd_to_valid_ids = { mol_atom_names_ccd.index(name): i for i, name in enumerate(mol_atom_names) } ccd_to_valid_id_array = np.array( [ float("nan") if i not in ccd_to_valid_ids else ccd_to_valid_ids[i] for i in range(len(mol_atom_names_ccd)) ] ) ccd_valid_ids = set(ccd_to_valid_ids.keys()) syms = [] # Get syms for sym_ccd in syms_ccd: sym_dict = {} bool_add = True for i, j in enumerate(sym_ccd): if i in ccd_valid_ids: if j in ccd_valid_ids: i_true = ccd_to_valid_ids[i] j_true = ccd_to_valid_ids[j] sym_dict[i_true] = j_true else: bool_add = False break if bool_add: syms.append([sym_dict[i] for i in range(len(ccd_valid_ids))]) for sym in syms: if len(sym) != added_molecules[mol_id]: raise Exception( f"Symmetry length mismatch {len(sym)} {added_molecules[mol_id]}" ) # assert ( # len(sym) == added_molecules[mol_id] # ), f"Symmetry length mismatch {len(sym)} {added_molecules[mol_id]}" sym_new_idx = [] for i, j in enumerate(sym): if i != int(j): sym_new_idx.append((i + start_mol, int(j) + start_mol)) if len(sym_new_idx) > 0: swaps.append(sym_new_idx) if len(swaps) > 0: molecule_symmetries.append(swaps) if return_physical_metrics: edge_index, (lower_bounds, upper_bounds, bond_mask, angle_mask) = ( slice_valid_index( edge_index, ccd_to_valid_id_array, (lower_bounds, upper_bounds, bond_mask, angle_mask), ) ) all_edge_index.append(edge_index + start_mol) all_lower_bounds.append(lower_bounds) all_upper_bounds.append(upper_bounds) all_bond_mask.append(bond_mask) all_angle_mask.append(angle_mask) chiral_atom_index, (chiral_check_mask, chiral_atom_orientations) = ( slice_valid_index( chiral_atom_index, ccd_to_valid_id_array, (chiral_check_mask, chiral_atom_orientations), ) ) all_chiral_atom_index.append(chiral_atom_index + start_mol) all_chiral_check_mask.append(chiral_check_mask) all_chiral_atom_orientations.append(chiral_atom_orientations) stereo_bond_index, (stereo_check_mask, stereo_bond_orientations) = ( slice_valid_index( stereo_bond_index, ccd_to_valid_id_array, (stereo_check_mask, stereo_bond_orientations), ) ) all_stereo_bond_index.append(stereo_bond_index + start_mol) all_stereo_check_mask.append(stereo_check_mask) all_stereo_bond_orientations.append(stereo_bond_orientations) aromatic_5_ring_index = slice_valid_index( aromatic_5_ring_index, ccd_to_valid_id_array ) aromatic_6_ring_index = slice_valid_index( aromatic_6_ring_index, ccd_to_valid_id_array ) planar_double_bond_index = slice_valid_index( planar_double_bond_index, ccd_to_valid_id_array ) all_aromatic_5_ring_index.append(aromatic_5_ring_index + start_mol) all_aromatic_6_ring_index.append(aromatic_6_ring_index + start_mol) all_planar_double_bond_index.append( planar_double_bond_index + start_mol ) if return_physical_metrics: if len(all_edge_index) > 0: all_edge_index = np.concatenate(all_edge_index, axis=1) all_lower_bounds = np.concatenate(all_lower_bounds, axis=0) all_upper_bounds = np.concatenate(all_upper_bounds, axis=0) all_bond_mask = np.concatenate(all_bond_mask, axis=0) all_angle_mask = np.concatenate(all_angle_mask, axis=0) all_chiral_atom_index = np.concatenate(all_chiral_atom_index, axis=1) all_chiral_check_mask = np.concatenate(all_chiral_check_mask, axis=0) all_chiral_atom_orientations = np.concatenate( all_chiral_atom_orientations, axis=0 ) all_stereo_bond_index = np.concatenate(all_stereo_bond_index, axis=1) all_stereo_check_mask = np.concatenate(all_stereo_check_mask, axis=0) all_stereo_bond_orientations = np.concatenate( all_stereo_bond_orientations, axis=0 ) all_aromatic_5_ring_index = np.concatenate( all_aromatic_5_ring_index, axis=1 ) all_aromatic_6_ring_index = np.concatenate( all_aromatic_6_ring_index, axis=1 ) all_planar_double_bond_index = np.empty( (6, 0), dtype=np.int64 ) else: all_edge_index = np.empty((2, 0), dtype=np.int64) all_lower_bounds = np.array([], dtype=np.float32) all_upper_bounds = np.array([], dtype=np.float32) all_bond_mask = np.array([], dtype=bool) all_angle_mask = np.array([], dtype=bool) all_chiral_atom_index = np.empty((4, 0), dtype=np.int64) all_chiral_check_mask = np.array([], dtype=bool) all_chiral_atom_orientations = np.array([], dtype=bool) all_stereo_bond_index = np.empty((4, 0), dtype=np.int64) all_stereo_check_mask = np.array([], dtype=bool) all_stereo_bond_orientations = np.array([], dtype=bool) all_aromatic_5_ring_index = np.empty((5, 0), dtype=np.int64) all_aromatic_6_ring_index = np.empty((6, 0), dtype=np.int64) all_planar_double_bond_index = np.empty((6, 0), dtype=np.int64) features = { "ligand_symmetries": molecule_symmetries, "ligand_edge_index": torch.tensor(all_edge_index).long(), "ligand_edge_lower_bounds": torch.tensor(all_lower_bounds), "ligand_edge_upper_bounds": torch.tensor(all_upper_bounds), "ligand_edge_bond_mask": torch.tensor(all_bond_mask), "ligand_edge_angle_mask": torch.tensor(all_angle_mask), "ligand_chiral_atom_index": torch.tensor(all_chiral_atom_index).long(), "ligand_chiral_check_mask": torch.tensor(all_chiral_check_mask), "ligand_chiral_atom_orientations": torch.tensor( all_chiral_atom_orientations ), "ligand_stereo_bond_index": torch.tensor(all_stereo_bond_index).long(), "ligand_stereo_check_mask": torch.tensor(all_stereo_check_mask), "ligand_stereo_bond_orientations": torch.tensor( all_stereo_bond_orientations ), "ligand_aromatic_5_ring_index": torch.tensor( all_aromatic_5_ring_index ).long(), "ligand_aromatic_6_ring_index": torch.tensor( all_aromatic_6_ring_index ).long(), "ligand_planar_double_bond_index": torch.tensor( all_planar_double_bond_index ).long(), } else: features = { "ligand_symmetries": molecule_symmetries, } return features def mol_from_smile(smiles_list: list[str]) -> list[Mol]: """Load the smiles. Parameters ---------- smiles : str Path to the smiles file. Returns ------- List[Mol] """ mols_processed = [] failed = 0 for smile in tqdm(smiles_list): try: mol = AllChem.MolFromSmiles(smile) mol = AllChem.AddHs(mol) mol.SetProp("MOL_NAME", "RANDO") mol.SetProp("SMILES", smile) canonical_order = AllChem.CanonicalRankAtoms(mol) for atom, can_idx in zip(mol.GetAtoms(), canonical_order): atom.SetProp("name", atom.GetSymbol().upper() + str(can_idx + 1)) mols_processed.append(mol) except Exception as e: # noqa: PERF203, BLE001 failed += 1 print(f"Failed to process {failed} molecules.") return mols_processed def compute_3d(mol: Mol, num_conformers: int = 10) -> bool: """Generate 3D coordinates using EKTDG method. Taken from `pdbeccdutils.core.component.Component`. Parameters ---------- mol: Mol The RDKit molecule to process num_conformers: int, optional Number of conformers to generate, defaults to 10 Returns ------- bool Whether computation was successful. """ try: # Generate multiple conformers params = AllChem.ETKDGv3() params.numThreads = 0 params.timeout = 60 conf_ids = AllChem.EmbedMultipleConfs( mol, numConfs=num_conformers, params=params ) if len(conf_ids) < num_conformers: params.useRandomCoords = True conf_ids = AllChem.EmbedMultipleConfs( mol, numConfs=num_conformers, params=params ) if len(conf_ids) < 1: return False # Optimize each conformer conf_ids_final = [] for conf_id in conf_ids: try: AllChem.UFFOptimizeMolecule(mol, confId=conf_id, maxIters=1000) except: # noqa: E722 # Remove failed conformer mol.RemoveConformer(conf_id) continue conf_ids_final.append(conf_id) if len(conf_ids_final) < 1: return False except: # noqa: E722 # Handle force field or sanitization issues return False return True def get_moldir_atomic_numbers(moldir: str) -> list[int]: counter = set() mols = load_all_molecules(moldir) z_set = set() for mol in mols.values(): z_set.update(a.GetAtomicNum() for a in mol.GetAtoms()) z_set.discard(0) whitelist = sorted(z_set) return whitelist