import itertools import numpy as np from ost import mol from ost import conop # use cdist of scipy, fallback to (slower) numpy implementation if scipy is not # available try: from scipy.spatial.distance import cdist except: def cdist(p1, p2): x2 = np.sum(p1**2, axis=1) # (m) y2 = np.sum(p2**2, axis=1) # (n) xy = np.matmul(p1, p2.T) # (m, n) x2 = x2.reshape(-1, 1) return np.sqrt(x2 - 2*xy + y2) # (m, n) def blockwise_cdist(A, B, block_size=1000): """ Memory efficient cdist implementation that performs blockwise operations scipy cdist uses 64 bit floats (double) which can scratch at the upper memory end for most machines when number of positions become larger. E.g. ~4000 residues might for example have 35000 atom positions. That's Almost 10GB to hold all pairwise distances in 64bit floats. This function calls cdist blockwise and stores the results in a 32bit float matrix. This function is adapted from chatgpt output """ A = A.astype(np.float32) B = B.astype(np.float32) M, N = A.shape[0], B.shape[0] D = np.empty((M, N), dtype=np.float32) # Output in float32 to save memory for i in range(0, M, block_size): A_block = A[i:i+block_size] D[i:i+block_size, :] = cdist(A_block, B).astype(np.float32) return D class CustomCompound: """ Defines atoms for custom compounds LDDT requires the reference atoms of a compound which are typically extracted from a :class:`ost.conop.CompoundLib`. This lightweight container allows to handle arbitrary compounds which are not necessarily in the compound library. :param atom_names: Names of atoms of custom compound :type atom_names: :class:`list` of :class:`str` """ def __init__(self, atom_names): self.atom_names = atom_names @staticmethod def FromResidue(res): """ Construct custom compound from residue :param res: Residue from which reference atom names are extracted, hydrogen/deuterium atoms are filtered out :type res: :class:`ost.mol.ResidueView`/:class:`ost.mol.ResidueHandle` :returns: :class:`CustomCompound` """ at_names = [a.name for a in res.atoms if a.element not in ["H", "D"]] if len(at_names) != len(set(at_names)): raise RuntimeError("Duplicate atoms detected in CustomCompound") compound = CustomCompound(at_names) return compound class SymmetrySettings: """Container for symmetric compounds LDDT considers symmetries and selects the one resulting in the highest possible score. A symmetry is defined as a renaming operation on one or more atoms that leads to a chemically equivalent residue. Example would be OD1 and OD2 in ASP => renaming OD1 to OD2 and vice versa gives a chemically equivalent residue. Use :func:`AddSymmetricCompound` to define a symmetry which can then directly be accessed through the *symmetric_compounds* member. """ def __init__(self): self.symmetric_compounds = dict() def AddSymmetricCompound(self, name, symmetric_atoms): """Adds symmetry for compound with *name* :param name: Name of compound with symmetry :type name: :class:`str` :param symmetric_atoms: Pairs of atom names that define renaming operation, i.e. after applying all switches defined in the tuples, the resulting residue should be chemically equivalent. Atom names must refer to the PDB component dictionary. :type symmetric_atoms: :class:`list` of :class:`tuple` """ for pair in symmetric_atoms: if len(pair) != 2: raise RuntimeError("Expect pairs when defining symmetries") self.symmetric_compounds[name] = symmetric_atoms def GetDefaultSymmetrySettings(): """Constructs and returns :class:`SymmetrySettings` object for natural amino acids """ symmetry_settings = SymmetrySettings() # ASP symmetry_settings.AddSymmetricCompound("ASP", [("OD1", "OD2")]) # GLU symmetry_settings.AddSymmetricCompound("GLU", [("OE1", "OE2")]) # LEU symmetry_settings.AddSymmetricCompound("LEU", [("CD1", "CD2")]) # VAL symmetry_settings.AddSymmetricCompound("VAL", [("CG1", "CG2")]) # ARG symmetry_settings.AddSymmetricCompound("ARG", [("NH1", "NH2")]) # PHE symmetry_settings.AddSymmetricCompound( "PHE", [("CD1", "CD2"), ("CE1", "CE2")] ) # TYR symmetry_settings.AddSymmetricCompound( "TYR", [("CD1", "CD2"), ("CE1", "CE2")] ) # nucleotides nuc_names = ["A", "C", "G", "U", "DA", "DC", "DG", "DT"] for nuc_name in nuc_names: symmetry_settings.AddSymmetricCompound( nuc_name, [("OP1","OP2")] ) return symmetry_settings class lDDTScorer: """LDDT scorer object for a specific target Sets up everything to score models of that target. LDDT (local distance difference test) is defined as fraction of pairwise distances which exhibit a difference < threshold when considering target and model. In case of multiple thresholds, the average is returned. See V. Mariani, M. Biasini, A. Barbato, T. Schwede, lDDT : A local superposition-free score for comparing protein structures and models using distance difference tests, Bioinformatics, 2013 :param target: The target :type target: :class:`ost.mol.EntityHandle`/:class:`ost.mol.EntityView` :param compound_lib: Compound library from which a compound for each residue is extracted based on its name. Uses :func:`ost.conop.GetDefaultLib` if not given, raises if this returns no valid compound library. Atoms defined in the compound are searched in the residue and build the reference for scoring. If the residue has atoms with names ["A", "B", "C"] but the corresponding compound only has ["A", "B"], "A" and "B" are considered for scoring. If the residue has atoms ["A", "B"] but the compound has ["A", "B", "C"], "C" is considered missing and does not influence scoring, even if present in the model. :param custom_compounds: Custom compounds defining reference atoms. If given, *custom_compounds* take precedent over *compound_lib*. :type custom_compounds: :class:`dict` with residue names (:class:`str`) as key and :class:`CustomCompound` as value. :type compound_lib: :class:`ost.conop.CompoundLib` :param inclusion_radius: All pairwise distances < *inclusion_radius* are considered for scoring :type inclusion_radius: :class:`float` :param sequence_separation: Only pairwise distances between atoms of residues which are further apart than this threshold are considered. Residue distance is based on resnum. The default (0) considers all pairwise distances except intra-residue distances. :type sequence_separation: :class:`int` :param symmetry_settings: Define residues exhibiting internal symmetry, uses :func:`GetDefaultSymmetrySettings` if not given. :type symmetry_settings: :class:`SymmetrySettings` :param seqres_mapping: Mapping of model residues at the scoring stage happens with residue numbers defining their location in a reference sequence (SEQRES) using one based indexing. If the residue numbers in *target* don't correspond to that SEQRES, you can specify the mapping manually. You can provide a dictionary to specify a reference sequence (SEQRES) for one or more chain(s). Key: chain name, value: alignment (seq1: SEQRES, seq2: sequence of residues in chain). Example: The residues in a chain with name "A" have sequence "YEAH" and residue numbers [42,43,44,45]. You can provide an alignment with seq1 "``HELLYEAH``" and seq2 "``----YEAH``". "Y" gets assigned residue number 5, "E" gets assigned 6 and so on no matter what the original residue numbers were. :type seqres_mapping: :class:`dict` (key: :class:`str`, value: :class:`ost.seq.AlignmentHandle`) :param bb_only: Only consider atoms with name "CA" in case of amino acids and "C3'" for Nucleotides. this invalidates *compound_lib*. Raises if any residue in *target* is not `r.chem_class.IsPeptideLinking()` or `r.chem_class.IsNucleotideLinking()` :type bb_only: :class:`bool` :raises: :class:`RuntimeError` if *target* contains compound which is not in *compound_lib*, :class:`RuntimeError` if *symmetry_settings* specifies symmetric atoms that are not present in the according compound in *compound_lib*, :class:`RuntimeError` if *seqres_mapping* is not provided and *target* contains residue numbers with insertion codes or the residue numbers for each chain are not monotonically increasing, :class:`RuntimeError` if *seqres_mapping* is provided but an alignment is invalid (seq1 contains gaps, mismatch in seq1/seq2, seq2 does not match residues in corresponding chains). """ def __init__( self, target, compound_lib=None, custom_compounds=None, inclusion_radius=15, sequence_separation=0, symmetry_settings=None, seqres_mapping=dict(), bb_only=False ): self.target = target self.inclusion_radius = inclusion_radius self.sequence_separation = sequence_separation if compound_lib is None: compound_lib = conop.GetDefaultLib() if compound_lib is None: raise RuntimeError("No compound_lib given and conop.GetDefaultLib " "returns no valid compound library") self.compound_lib = compound_lib self.custom_compounds = custom_compounds if symmetry_settings is None: self.symmetry_settings = GetDefaultSymmetrySettings() else: self.symmetry_settings = symmetry_settings # whether to only consider atoms with name "CA" (amino acids) or C3' # (nucleotides), invalidates *compound_lib* self.bb_only=bb_only # names of heavy atoms of each unique compound present in *target* as # extracted from *compound_lib*, e.g. # self.compound_anames["GLY"] = ["N", "CA", "C", "O"] self.compound_anames = dict() # stores symmetry information for those compounds as defined in # *symmetry_settings* self.compound_symmetric_atoms = dict() # list of len(target.chains) containing all chain names in *target* self.chain_names = list() # list of len(target.residues) containing all compound names in *target* self.compound_names = list() # list of len(target.residues) defining start pos in internal reference # positions for each residue self.res_start_indices = list() # list of len(target.residues) defining residue numbers in target self.res_resnums = list() # list of len(target.chains) defining start pos in internal reference # positions for each chain self.chain_start_indices = list() # list of len(target.chains) defining start pos in self.compound_names # for each chain self.chain_res_start_indices = list() # maps residues in *target* to indices in # self.compound_names/self.res_start_indices. A residue gets identified # by a tuple (first element: chain name, second element: residue number, # residue number is either the actual residue number in *target* or # given by *seqres_mapping*) self.res_mapper = dict() # number of atoms as specified in compounds. not all are necessarily # covered by structure self.n_atoms = None # stores an index for each AtomHandle in *target* # (atom hashcode => index) self.atom_indices = dict() # store indices of all atoms that have symmetry properties self.symmetric_atoms = set() # the actual target positions in a numpy array of shape (self.n_atoms,3) self.positions = None # setup members defined above self._SetupEnv(self.compound_lib, self.custom_compounds, self.symmetry_settings, seqres_mapping, self.bb_only) # distance related members are lazily computed as they're affected # by different flavours of LDDT (e.g. LDDT including inter-chain # contacts or not etc.) # stores for each atom the other atoms within inclusion_radius self._ref_indices = None # the corresponding distances self._ref_distances = None # The following lists will be sparsely populated. We keep for each # symmetry related atom the distances towards all atoms which are NOT # affected by symmetry. So we can evaluate two symmetric versions # against the fixed stuff later on and select the better scoring one. self._sym_ref_indices = None self._sym_ref_distances = None # exactly the same as above but without interchain contacts # => single-chain (sc) self._ref_indices_sc = None self._ref_distances_sc = None self._sym_ref_indices_sc = None self._sym_ref_distances_sc = None # exactly the same as above but without intrachain contacts # => inter-chain (ic) self._ref_indices_ic = None self._ref_distances_ic = None self._sym_ref_indices_ic = None self._sym_ref_distances_ic = None # input parameter checking self._ProcessSequenceSeparation() @property def ref_indices(self): if self._ref_indices is None: self._ref_indices, self._ref_distances = \ lDDTScorer._SetupDistances(self.target, self.n_atoms, self.atom_indices, self.inclusion_radius) return self._ref_indices @property def ref_distances(self): if self._ref_distances is None: self._ref_indices, self._ref_distances = \ lDDTScorer._SetupDistances(self.target, self.n_atoms, self.atom_indices, self.inclusion_radius) return self._ref_distances @property def sym_ref_indices(self): if self._sym_ref_indices is None: self._sym_ref_indices, self._sym_ref_distances = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, self.ref_indices, self.ref_distances) return self._sym_ref_indices @property def sym_ref_distances(self): if self._sym_ref_distances is None: self._sym_ref_indices, self._sym_ref_distances = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, self.ref_indices, self.ref_distances) return self._sym_ref_distances @property def ref_indices_sc(self): if self._ref_indices_sc is None: self._ref_indices_sc, self._ref_distances_sc = \ lDDTScorer._SetupDistancesSC(self.n_atoms, self.chain_start_indices, self.ref_indices, self.ref_distances) return self._ref_indices_sc @property def ref_distances_sc(self): if self._ref_distances_sc is None: self._ref_indices_sc, self._ref_distances_sc = \ lDDTScorer._SetupDistancesSC(self.n_atoms, self.chain_start_indices, self.ref_indices, self.ref_distances) return self._ref_distances_sc @property def sym_ref_indices_sc(self): if self._sym_ref_indices_sc is None: self._sym_ref_indices_sc, self._sym_ref_distances_sc = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, self.ref_indices_sc, self.ref_distances_sc) return self._sym_ref_indices_sc @property def sym_ref_distances_sc(self): if self._sym_ref_distances_sc is None: self._sym_ref_indices_sc, self._sym_ref_distances_sc = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, self.ref_indices_sc, self.ref_distances_sc) return self._sym_ref_distances_sc @property def ref_indices_ic(self): if self._ref_indices_ic is None: self._ref_indices_ic, self._ref_distances_ic = \ lDDTScorer._SetupDistancesIC(self.n_atoms, self.chain_start_indices, self.ref_indices, self.ref_distances) return self._ref_indices_ic @property def ref_distances_ic(self): if self._ref_distances_ic is None: self._ref_indices_ic, self._ref_distances_ic = \ lDDTScorer._SetupDistancesIC(self.n_atoms, self.chain_start_indices, self.ref_indices, self.ref_distances) return self._ref_distances_ic @property def sym_ref_indices_ic(self): if self._sym_ref_indices_ic is None: self._sym_ref_indices_ic, self._sym_ref_distances_ic = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, self.ref_indices_ic, self.ref_distances_ic) return self._sym_ref_indices_ic @property def sym_ref_distances_ic(self): if self._sym_ref_distances_ic is None: self._sym_ref_indices_ic, self._sym_ref_distances_ic = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, self.ref_indices_ic, self.ref_distances_ic) return self._sym_ref_distances_ic def lDDT(self, model, thresholds = [0.5, 1.0, 2.0, 4.0], local_lddt_prop=None, local_contact_prop=None, chain_mapping=None, no_interchain=False, no_intrachain=False, penalize_extra_chains=False, residue_mapping=None, return_dist_test=False, check_resnames=True, add_mdl_contacts=False, interaction_data=None, set_atom_props=False): """Computes LDDT of *model* - globally and per-residue :param model: Model to be scored - models are preferably scored upon performing stereo-chemistry checks in order to punish for non-sensical irregularities. This must be done separately as a pre-processing step. Target contacts that are not covered by *model* are considered not conserved, thus decreasing LDDT score. This also includes missing model chains or model chains for which no mapping is provided in *chain_mapping*. :type model: :class:`ost.mol.EntityHandle`/:class:`ost.mol.EntityView` :param thresholds: Thresholds of distance differences to be considered as correct - see docs in constructor for more info. default: [0.5, 1.0, 2.0, 4.0] :type thresholds: :class:`list` of :class:`floats` :param local_lddt_prop: If set, per-residue scores will be assigned as generic float property of that name :type local_lddt_prop: :class:`str` :param local_contact_prop: If set, number of expected contacts as well as number of conserved contacts will be assigned as generic int property. Excected contacts will be set as _exp, conserved contacts as _cons. Values are summed over all thresholds. :type local_contact_prop: :class:`str` :param chain_mapping: Mapping of model chains (key) onto target chains (value). This is required if target or model have more than one chain. :type chain_mapping: :class:`dict` with :class:`str` as keys/values :param no_interchain: Whether to exclude interchain contacts :type no_interchain: :class:`bool` :param no_intrachain: Whether to exclude intrachain contacts (i.e. only consider interface related contacts) :type no_intrachain: :class:`bool` :param penalize_extra_chains: Whether to include a fixed penalty for additional chains in the model that are not mapped to the target. ONLY AFFECTS RETURNED GLOBAL SCORE. In detail: adds the number of intra-chain contacts of each extra chain to the expected contacts, thus adding a penalty. :type penalize_extra_chains: :class:`bool` :param residue_mapping: By default, residue mapping is based on residue numbers. That means, a model chain and the respective target chain map to the same underlying reference sequence (SEQRES). Alternatively, you can specify one or several alignment(s) between model and target chains by providing a dictionary. key: Name of chain in model (respective target chain is extracted from *chain_mapping*), value: Alignment with first sequence corresponding to target chain and second sequence to model chain. There is NO reference sequence involved, so the two sequences MUST exactly match the actual residues observed in the respective target/model chains (ATOMSEQ). :type residue_mapping: :class:`dict` with key: :class:`str`, value: :class:`ost.seq.AlignmentHandle` :param return_dist_test: Whether to additionally return the underlying per-residue data for the distance difference test. Adds five objects to the return tuple. First: Number of total contacts summed over all thresholds Second: Number of conserved contacts summed over all thresholds Third: list with length of scored residues. Contains indices referring to model.residues. Fourth: numpy array of size len(scored_residues) containing the number of total contacts, Fifth: numpy matrix of shape (len(scored_residues), len(thresholds)) specifying how many for each threshold are conserved. :param check_resnames: On by default. Enforces residue name matches between mapped model and target residues. :type check_resnames: :class:`bool` :param add_mdl_contacts: Adds model contacts - Only using contacts that are within a certain distance threshold in the target does not penalize for added model contacts. If set to True, this flag will also consider target contacts that are within the specified distance threshold in the model but not necessarily in the target. No contact will be added if the respective atom pair is not resolved in the target. :type add_mdl_contacts: :class:`bool` :param interaction_data: Pro param - don't use :type interaction_data: :class:`tuple` :param set_atom_props: If True, sets generic properties on a per atom level if *local_lddt_prop*/*local_contact_prop* are set as well. In other words: this is the only way you can get per-atom LDDT values. :type set_atom_props: :class:`bool` :returns: global and per-residue LDDT scores as a tuple - first element is global LDDT score (None if *target* has no contacts) and second element a list of per-residue scores with length len(*model*.residues). None is assigned to residues that are not covered by target. If a residue is covered but has no contacts in *target*, 0.0 is assigned. """ if chain_mapping is None: if len(self.chain_names) > 1 or len(model.chains) > 1: raise NotImplementedError("Must provide chain mapping if " "target or model have > 1 chains.") chain_mapping = {model.chains[0].GetName(): self.chain_names[0]} else: # check whether chains specified in mapping exist for model_chain, target_chain in chain_mapping.items(): if target_chain not in self.chain_names: raise RuntimeError(f"Target chain specified in " f"chain_mapping ({target_chain}) does " f"not exist. Target has chains: " f"{self.chain_names}") ch = model.FindChain(model_chain) if not ch.IsValid(): raise RuntimeError(f"Model chain specified in " f"chain_mapping ({model_chain}) does " f"not exist. Model has chains: " f"{[c.GetName() for c in model.chains]}") # data objects defining model data - see _ProcessModel for rough # description pos, res_ref_atom_indices, res_atom_indices, res_atom_hashes, \ res_indices, ref_res_indices, symmetries = \ self._ProcessModel(model, chain_mapping, residue_mapping = residue_mapping, nirvana_dist = self.inclusion_radius + max(thresholds), check_resnames = check_resnames) if no_interchain and no_intrachain: raise RuntimeError("no_interchain and no_intrachain flags are " "mutually exclusive") sym_ref_indices = None sym_ref_distances = None ref_indices = None ref_distances = None if interaction_data is None: if no_interchain: sym_ref_indices = self.sym_ref_indices_sc sym_ref_distances = self.sym_ref_distances_sc ref_indices = self.ref_indices_sc ref_distances = self.ref_distances_sc elif no_intrachain: sym_ref_indices = self.sym_ref_indices_ic sym_ref_distances = self.sym_ref_distances_ic ref_indices = self.ref_indices_ic ref_distances = self.ref_distances_ic else: sym_ref_indices = self.sym_ref_indices sym_ref_distances = self.sym_ref_distances ref_indices = self.ref_indices ref_distances = self.ref_distances if add_mdl_contacts: ref_indices, ref_distances = \ self._AddMdlContacts(model, res_atom_indices, res_atom_hashes, ref_indices, ref_distances, no_interchain, no_intrachain) # recompute symmetry related indices/distances sym_ref_indices, sym_ref_distances = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, ref_indices, ref_distances) else: sym_ref_indices, sym_ref_distances, ref_indices, ref_distances = \ interaction_data self._ResolveSymmetries(pos, thresholds, symmetries, sym_ref_indices, sym_ref_distances) atom_indices = list(itertools.chain.from_iterable(res_atom_indices)) per_atom_exp = np.asarray([self._GetNExp(i, ref_indices) for i in atom_indices], dtype=np.int32) per_res_exp = np.asarray([self._GetNExp(res_ref_atom_indices[idx], ref_indices) for idx in range(len(res_indices))], dtype=np.int32) per_atom_conserved = self._EvalAtoms(pos, atom_indices, thresholds, ref_indices, ref_distances) per_res_conserved = np.zeros((len(res_atom_indices), len(thresholds)), dtype=np.int32) start_idx = 0 for r_idx in range(len(res_atom_indices)): end_idx = start_idx + len(res_atom_indices[r_idx]) per_res_conserved[r_idx] = np.sum(per_atom_conserved[start_idx:end_idx,:], axis=0) start_idx = end_idx n_thresh = len(thresholds) # do per-residue scores per_res_lDDT = [None] * model.GetResidueCount() for idx in range(len(res_indices)): n_exp = n_thresh * per_res_exp[idx] if n_exp > 0: score = np.sum(per_res_conserved[idx,:]) / n_exp per_res_lDDT[res_indices[idx]] = score else: per_res_lDDT[res_indices[idx]] = 0.0 # do full model score n_distances = sum([len(x) for x in ref_indices]) if penalize_extra_chains: n_distances += self._GetExtraModelChainPenalty(model, chain_mapping) lDDT_tot = int(n_thresh * n_distances) lDDT_cons = int(np.sum(per_res_conserved)) lDDT = None if lDDT_tot > 0: lDDT = float(lDDT_cons) / lDDT_tot # set properties if necessary if local_lddt_prop: residues = model.residues for idx in res_indices: residues[idx].SetFloatProp(local_lddt_prop, per_res_lDDT[idx]) if local_contact_prop: residues = model.residues exp_prop = local_contact_prop + "_exp" conserved_prop = local_contact_prop + "_cons" for i, r_idx in enumerate(res_indices): residues[r_idx].SetIntProp(exp_prop, n_thresh * int(per_res_exp[i])) residues[r_idx].SetIntProp(conserved_prop, int(np.sum(per_res_conserved[i,:]))) if set_atom_props and (local_lddt_prop or local_contact_prop): atom_list = list() residues = model.residues for i, indices in enumerate(res_atom_indices): r = residues[res_indices[i]] r_idx = ref_res_indices[i] res_start_idx = self.res_start_indices[r_idx] anames = self.compound_anames[self.compound_names[r_idx]] for a_i in indices: a = r.FindAtom(anames[a_i - res_start_idx]) assert(a.IsValid()) atom_list.append(a) summed_per_atom_conserved = per_atom_conserved.sum(axis=1) if local_lddt_prop: # the only place where actually need to compute per-atom LDDT # scores for a_idx in range(len(atom_list)): if per_atom_exp[a_idx] != 0: tmp = summed_per_atom_conserved[a_idx] / per_atom_exp[a_idx] tmp = tmp / n_thresh atom_list[a_idx].SetFloatProp(local_lddt_prop, tmp) if local_contact_prop: conserved_prop = local_contact_prop + "_cons" exp_prop = local_contact_prop + "_exp" for a_idx in range(len(atom_list)): # do number of conserved contacts tmp = summed_per_atom_conserved[a_idx] atom_list[a_idx].SetIntProp(conserved_prop, tmp) # do number of expected contacts tmp = per_atom_exp[a_idx] * n_thresh atom_list[a_idx].SetIntProp(exp_prop, tmp) if return_dist_test: return lDDT, per_res_lDDT, lDDT_tot, lDDT_cons, res_indices, \ per_res_exp, per_res_conserved else: return lDDT, per_res_lDDT def DRMSD(self, model, dist_cap = 5, chain_mapping=None, no_interchain=False, no_intrachain=False, residue_mapping=None, check_resnames=True, add_mdl_contacts=False, interaction_data=None): """ DRMSD of *model* - globally and per-residue Very similar to LDDT as we operate on distance differences for all interatomic distances within the same inclusion radius as in LDDT. DRMSD is the distance rmsd, i.e. the RMSD of distance differences. Distance differences are capped at *dist_cap* which is also the default value for missing distances. :param model: Model to be scored - models are preferably scored upon performing stereo-chemistry checks in order to punish for non-sensical irregularities. This must be done separately as a pre-processing step. Target contacts that are not covered by *model* are considered not conserved, thus increasing DRMSD score. This also includes missing model chains or model chains for which no mapping is provided in *chain_mapping*. :type model: :class:`ost.mol.EntityHandle`/:class:`ost.mol.EntityView` :param dist_cap: Cap for distance differences. :type dist_cap: :class:`float` :param chain_mapping: Mapping of model chains (key) onto target chains (value). This is required if target or model have more than one chain. :type chain_mapping: :class:`dict` with :class:`str` as keys/values :param no_interchain: Whether to exclude interchain contacts :type no_interchain: :class:`bool` :param no_intrachain: Whether to exclude intrachain contacts (i.e. only consider interface related contacts) :type no_intrachain: :class:`bool` :param residue_mapping: By default, residue mapping is based on residue numbers. That means, a model chain and the respective target chain map to the same underlying reference sequence (SEQRES). Alternatively, you can specify one or several alignment(s) between model and target chains by providing a dictionary. key: Name of chain in model (respective target chain is extracted from *chain_mapping*), value: Alignment with first sequence corresponding to target chain and second sequence to model chain. There is NO reference sequence involved, so the two sequences MUST exactly match the actual residues observed in the respective target/model chains (ATOMSEQ). :type residue_mapping: :class:`dict` with key: :class:`str`, value: :class:`ost.seq.AlignmentHandle` :param check_resnames: On by default. Enforces residue name matches between mapped model and target residues. :type check_resnames: :class:`bool` :param add_mdl_contacts: Adds model contacts - Only using contacts that are within a certain distance threshold in the target does not penalize for added model contacts. If set to True, this flag will also consider target contacts that are within the specified distance threshold in the model but not necessarily in the target. No contact will be added if the respective atom pair is not resolved in the target. :type add_mdl_contacts: :class:`bool` :param interaction_data: Pro param - don't use :type interaction_data: :class:`tuple` :returns: global and per-residue DRMSD scores as a tuple - first element is global DRMSD score (None if *target* has no contacts) and second element a list of per-residue scores with length len(*model*.residues). None is assigned to residues that are not covered by target. If a residue is covered but has no contacts in *target*, None is assigned. """ if chain_mapping is None: if len(self.chain_names) > 1 or len(model.chains) > 1: raise NotImplementedError("Must provide chain mapping if " "target or model have > 1 chains.") chain_mapping = {model.chains[0].GetName(): self.chain_names[0]} else: # check whether chains specified in mapping exist for model_chain, target_chain in chain_mapping.items(): if target_chain not in self.chain_names: raise RuntimeError(f"Target chain specified in " f"chain_mapping ({target_chain}) does " f"not exist. Target has chains: " f"{self.chain_names}") ch = model.FindChain(model_chain) if not ch.IsValid(): raise RuntimeError(f"Model chain specified in " f"chain_mapping ({model_chain}) does " f"not exist. Model has chains: " f"{[c.GetName() for c in model.chains]}") # data objects defining model data - see _ProcessModel for rough # description pos, res_ref_atom_indices, res_atom_indices, res_atom_hashes, \ res_indices, ref_res_indices, symmetries = \ self._ProcessModel(model, chain_mapping, residue_mapping = residue_mapping, nirvana_dist = self.inclusion_radius + dist_cap, check_resnames = check_resnames) if no_interchain and no_intrachain: raise RuntimeError("no_interchain and no_intrachain flags are " "mutually exclusive") sym_ref_indices = None sym_ref_distances = None ref_indices = None ref_distances = None if interaction_data is None: if no_interchain: sym_ref_indices = self.sym_ref_indices_sc sym_ref_distances = self.sym_ref_distances_sc ref_indices = self.ref_indices_sc ref_distances = self.ref_distances_sc elif no_intrachain: sym_ref_indices = self.sym_ref_indices_ic sym_ref_distances = self.sym_ref_distances_ic ref_indices = self.ref_indices_ic ref_distances = self.ref_distances_ic else: sym_ref_indices = self.sym_ref_indices sym_ref_distances = self.sym_ref_distances ref_indices = self.ref_indices ref_distances = self.ref_distances if add_mdl_contacts: ref_indices, ref_distances = \ self._AddMdlContacts(model, res_atom_indices, res_atom_hashes, ref_indices, ref_distances, no_interchain, no_intrachain) # recompute symmetry related indices/distances sym_ref_indices, sym_ref_distances = \ lDDTScorer._NonSymDistances(self.n_atoms, self.symmetric_atoms, ref_indices, ref_distances) else: sym_ref_indices, sym_ref_distances, ref_indices, ref_distances = \ interaction_data self._ResolveSymmetriesSSD(pos, dist_cap, symmetries, sym_ref_indices, sym_ref_distances) atom_indices = list(itertools.chain.from_iterable(res_atom_indices)) per_atom_exp = np.asarray([self._GetNExp(i, ref_indices) for i in atom_indices], dtype=np.int32) per_res_exp = np.asarray([self._GetNExp(res_ref_atom_indices[idx], ref_indices) for idx in range(len(res_indices))], dtype=np.int32) per_atom_ssd = self._EvalAtomsSSD(pos, atom_indices, dist_cap, ref_indices, ref_distances) # do per residue scores start_idx = 0 per_res_drmsd = [None] * model.GetResidueCount() for r_idx in range(len(res_atom_indices)): end_idx = start_idx + len(res_atom_indices[r_idx]) n_tot = per_res_exp[r_idx] if n_tot > 0: ssd = np.sum(per_atom_ssd[start_idx:end_idx]) # add penalties from distances involving atoms that are not # present in the model n_missing = n_tot - np.sum(per_atom_exp[start_idx:end_idx]) ssd += n_missing*dist_cap*dist_cap per_res_drmsd[res_indices[r_idx]] = np.sqrt(ssd/n_tot) start_idx = end_idx # do full model score drmsd = None n_tot = sum([len(x) for x in ref_indices]) if n_tot > 0: ssd = np.sum(per_atom_ssd) # add penalties from distances involving atoms that are not # present in the model n_missing = n_tot - np.sum(per_atom_exp) ssd += (dist_cap*dist_cap*n_missing) drmsd = np.sqrt(ssd/n_tot) return drmsd, per_res_drmsd def GetNChainContacts(self, target_chain, no_interchain=False): """Returns number of contacts expected for a certain chain in *target* :param target_chain: Chain in *target* for which you want the number of expected contacts :type target_chain: :class:`str` :param no_interchain: Whether to exclude interchain contacts :type no_interchain: :class:`bool` :raises: :class:`RuntimeError` if specified chain doesnt exist """ if target_chain not in self.chain_names: raise RuntimeError(f"Specified chain name ({target_chain}) not in " f"target") ch_idx = self.chain_names.index(target_chain) s = self.chain_start_indices[ch_idx] e = self.n_atoms if ch_idx + 1 < len(self.chain_names): e = self.chain_start_indices[ch_idx+1] if no_interchain: return self._GetNExp(list(range(s, e)), self.ref_indices_sc) else: return self._GetNExp(list(range(s, e)), self.ref_indices) def _ProcessModel(self, model, chain_mapping, residue_mapping = None, nirvana_dist = 100, check_resnames = True): """ Helper that generates data structures from model """ # initialize positions with values far in nirvana. If a position is not # set, it should be far away from any position in model. max_pos = model.bounds.GetMax() max_coordinate = abs(max(max_pos[0], max_pos[1], max_pos[2])) max_coordinate += 42 * nirvana_dist pos = np.ones((self.n_atoms, 3), dtype=np.float32) * max_coordinate # for each scored residue in model a list of indices describing the # atoms from the reference that should be there res_ref_atom_indices = list() # for each scored residue in model a list of indices of atoms that are # actually there res_atom_indices = list() # and the respective hash codes # this is required if add_mdl_contacts is set to True res_atom_hashes = list() # indices of the scored residues res_indices = list() # respective residue indices in reference ref_res_indices = list() # Will contain one element per symmetry group symmetries = list() current_model_res_idx = -1 for ch in model.chains: model_ch_name = ch.GetName() if model_ch_name not in chain_mapping: current_model_res_idx += len(ch.residues) continue # additional model chain which is not mapped target_ch_name = chain_mapping[model_ch_name] rnums = self._GetChainRNums(ch, residue_mapping, model_ch_name, target_ch_name) for r, rnum in zip(ch.residues, rnums): current_model_res_idx += 1 res_mapper_key = (target_ch_name, rnum) if res_mapper_key not in self.res_mapper: continue r_idx = self.res_mapper[res_mapper_key] if check_resnames and r.name != self.compound_names[r_idx]: raise RuntimeError( f"Residue name mismatch for {r}, " f" expect {self.compound_names[r_idx]}" ) res_start_idx = self.res_start_indices[r_idx] rname = self.compound_names[r_idx] anames = self.compound_anames[rname] atoms = [r.FindAtom(aname) for aname in anames] res_ref_atom_indices.append( list(range(res_start_idx, res_start_idx + len(anames))) ) res_atom_indices.append(list()) res_atom_hashes.append(list()) res_indices.append(current_model_res_idx) ref_res_indices.append(r_idx) for a_idx, a in enumerate(atoms): if a.IsValid(): p = a.GetPos() pos[res_start_idx + a_idx][0] = p[0] pos[res_start_idx + a_idx][1] = p[1] pos[res_start_idx + a_idx][2] = p[2] res_atom_indices[-1].append(res_start_idx + a_idx) res_atom_hashes[-1].append(a.handle.GetHashCode()) if rname in self.compound_symmetric_atoms: sym_indices = list() for sym_tuple in self.compound_symmetric_atoms[rname]: a_one = atoms[sym_tuple[0]] a_two = atoms[sym_tuple[1]] if a_one.IsValid() and a_two.IsValid(): sym_indices.append( ( res_start_idx + sym_tuple[0], res_start_idx + sym_tuple[1], ) ) if len(sym_indices) > 0: symmetries.append(sym_indices) return (pos, res_ref_atom_indices, res_atom_indices, res_atom_hashes, res_indices, ref_res_indices, symmetries) def _GetExtraModelChainPenalty(self, model, chain_mapping): """Counts n distances in extra model chains to be added as penalty """ penalty = 0 for chain in model.chains: ch_name = chain.GetName() if ch_name not in chain_mapping: sm = self.symmetry_settings mdl_sel = model.Select(f"cname={mol.QueryQuoteName(ch_name)}") dummy_scorer = lDDTScorer(mdl_sel, self.compound_lib, symmetry_settings = sm, inclusion_radius = self.inclusion_radius, bb_only = self.bb_only) penalty += sum([len(x) for x in dummy_scorer.ref_indices]) return penalty def _GetChainRNums(self, ch, residue_mapping, model_ch_name, target_ch_name): """Map residues in model chain to target residues There are two options: one is simply using residue numbers, the other is a custom mapping as given in *residue_mapping* """ if residue_mapping and model_ch_name in residue_mapping: # extract residue numbers from target chain ch_idx = self.chain_names.index(target_ch_name) start_idx = self.chain_res_start_indices[ch_idx] if ch_idx < len(self.chain_names) - 1: end_idx = self.chain_res_start_indices[ch_idx+1] else: end_idx = len(self.compound_names) target_rnums = self.res_resnums[start_idx:end_idx] # get sequences from alignment and do consistency checks target_seq = residue_mapping[model_ch_name].GetSequence(0) model_seq = residue_mapping[model_ch_name].GetSequence(1) if len(target_seq.GetGaplessString()) != len(target_rnums): raise RuntimeError(f"Try to perform residue mapping for " f"model chain {model_ch_name} which " f"maps to {target_ch_name} in target. " f"Target sequence in alignment suggests " f"{len(target_seq.GetGaplessString())} " f"residues but {len(target_rnums)} are " f"expected.") if len(model_seq.GetGaplessString()) != len(ch.residues): raise RuntimeError(f"Try to perform residue mapping for " f"model chain {model_ch_name} which " f"maps to {target_ch_name} in target. " f"Model sequence in alignment suggests " f"{len(model_seq.GetGaplessString())} " f"residues but {len(ch.residues)} are " f"expected.") rnums = list() target_idx = -1 for col in residue_mapping[model_ch_name]: if col[0] != '-': target_idx += 1 # handle match if col[0] != '-' and col[1] != '-': rnums.append(target_rnums[target_idx]) # insertion in model adds None to rnum if col[0] == '-' and col[1] != '-': rnums.append(None) else: rnums = [r.GetNumber() for r in ch.residues] return rnums def _SetupEnv(self, compound_lib, custom_compounds, symmetry_settings, seqres_mapping, bb_only): """Sets target related lDDTScorer members defined in constructor No distance related members - see _SetupDistances """ residue_numbers = self._GetTargetResidueNumbers(self.target, seqres_mapping) current_idx = 0 positions = list() for chain in self.target.chains: ch_name = chain.GetName() self.chain_names.append(ch_name) self.chain_start_indices.append(current_idx) self.chain_res_start_indices.append(len(self.compound_names)) for r, rnum in zip(chain.residues, residue_numbers[ch_name]): if r.name not in self.compound_anames: # sets compound info in self.compound_anames and # self.compound_symmetric_atoms self._SetupCompound(r, compound_lib, custom_compounds, symmetry_settings, bb_only) self.res_start_indices.append(current_idx) self.res_mapper[(ch_name, rnum)] = len(self.compound_names) self.compound_names.append(r.name) self.res_resnums.append(rnum) atoms = [r.FindAtom(an) for an in self.compound_anames[r.name]] for a in atoms: if a.IsValid(): self.atom_indices[a.handle.GetHashCode()] = current_idx p = a.GetPos() positions.append(np.asarray([p[0], p[1], p[2]], dtype=np.float32)) else: positions.append(np.zeros(3, dtype=np.float32)) current_idx += 1 if r.name in self.compound_symmetric_atoms: for sym_tuple in self.compound_symmetric_atoms[r.name]: for a_idx in sym_tuple: a = atoms[a_idx] if a.IsValid(): hashcode = a.handle.GetHashCode() self.symmetric_atoms.add( self.atom_indices[hashcode] ) self.positions = np.vstack(positions) self.n_atoms = current_idx def _GetTargetResidueNumbers(self, target, seqres_mapping): """Returns residue numbers for each chain in target as dict They're either directly extracted from the raw residue number from the structure or from user provided alignments """ residue_numbers = dict() for ch in target.chains: ch_name = ch.GetName() rnums = list() if ch_name in seqres_mapping: seqres = seqres_mapping[ch_name].GetSequence(0).GetString() atomseq = seqres_mapping[ch_name].GetSequence(1).GetString() # SEQRES must not contain gaps if "-" in seqres: raise RuntimeError( "SEQRES in seqres_mapping must not " "contain gaps" ) atomseq_from_chain = [r.one_letter_code for r in ch.residues] if atomseq.replace("-", "") != atomseq_from_chain: raise RuntimeError( "ATOMSEQ in seqres_mapping must match " "raw sequence extracted from chain " "residues" ) rnum = 0 for seqres_olc, atomseq_olc in zip(seqres, atomseq): if seqres_olc != "-": rnum += 1 if atomseq_olc != "-": if seqres_olc != atomseq_olc: raise RuntimeError( f"Residue with number {rnum} in " f"chain {ch_name} has SEQRES " f"ATOMSEQ mismatch" ) rnums.append(mol.ResNum(rnum)) else: rnums = [r.GetNumber() for r in ch.residues] assert len(rnums) == len(ch.residues) residue_numbers[ch_name] = rnums return residue_numbers def _SetupCompound(self, r, compound_lib, custom_compounds, symmetry_settings, bb_only): """fill self.compound_anames/self.compound_symmetric_atoms """ if bb_only: # throw away compound_lib info if r.chem_class.IsPeptideLinking(): self.compound_anames[r.name] = ["CA"] elif r.chem_class.IsNucleotideLinking(): self.compound_anames[r.name] = ["C3'"] else: raise RuntimeError(f"Only support amino acids and nucleotides " f"if bb_only is True, failed with {str(r)}") self.compound_symmetric_atoms[r.name] = list() else: atom_names = list() symmetric_atoms = list() if custom_compounds is not None and r.GetName() in custom_compounds: atom_names = list(custom_compounds[r.GetName()].atom_names) else: compound = compound_lib.FindCompound(r.name) if compound is None: raise RuntimeError(f"no entry for {r} in compound_lib") for atom_spec in compound.GetAtomSpecs(): if atom_spec.element not in ["H", "D"]: atom_names.append(atom_spec.name) if r.name in symmetry_settings.symmetric_compounds: for pair in symmetry_settings.symmetric_compounds[r.name]: try: a = atom_names.index(pair[0]) b = atom_names.index(pair[1]) except: msg = f"Could not find symmetric atoms " msg += f"({pair[0]}, {pair[1]}) for {r.name} " msg += f"as specified in SymmetrySettings in " msg += f"compound from component dictionary. " msg += f"Atoms in compound: {atom_names}" raise RuntimeError(msg) symmetric_atoms.append((a, b)) self.compound_anames[r.name] = atom_names if len(symmetric_atoms) > 0: self.compound_symmetric_atoms[r.name] = symmetric_atoms def _AddMdlContacts(self, model, res_atom_indices, res_atom_hashes, ref_indices, ref_distances, no_interchain, no_intrachain): # buildup an index map for mdl atoms that are also present in target in_target = np.zeros(self.n_atoms, dtype=bool) for i in self.atom_indices.values(): in_target[i] = True mdl_atom_indices = dict() for at_indices, at_hashes in zip(res_atom_indices, res_atom_hashes): for i, h in zip(at_indices, at_hashes): if in_target[i]: mdl_atom_indices[h] = i # get contacts for mdl - the contacts are only from atom pairs that # are also present in target, as we only provide the respective # hashes in mdl_atom_indices mdl_ref_indices, mdl_ref_distances = \ lDDTScorer._SetupDistances(model, self.n_atoms, mdl_atom_indices, self.inclusion_radius) if no_interchain: mdl_ref_indices, mdl_ref_distances = \ lDDTScorer._SetupDistancesSC(self.n_atoms, self.chain_start_indices, mdl_ref_indices, mdl_ref_distances) if no_intrachain: mdl_ref_indices, mdl_ref_distances = \ lDDTScorer._SetupDistancesIC(self.n_atoms, self.chain_start_indices, mdl_ref_indices, mdl_ref_distances) # update ref_indices/ref_distances => add mdl contacts for i in range(self.n_atoms): mask = np.isin(mdl_ref_indices[i], ref_indices[i], assume_unique=True, invert=True) if np.sum(mask) > 0: added_mdl_indices = mdl_ref_indices[i][mask] ref_indices[i] = np.append(ref_indices[i], added_mdl_indices) # distances need to be recomputed from ref positions tmp = self.positions.take(added_mdl_indices, axis=0) np.subtract(tmp, self.positions[i][None, :], out=tmp) np.square(tmp, out=tmp) tmp = tmp.sum(axis=1) np.sqrt(tmp, out=tmp) # distances against all relevant atoms ref_distances[i] = np.append(ref_distances[i], tmp) return (ref_indices, ref_distances) @staticmethod def _SetupDistances(structure, n_atoms, atom_index_mapping, inclusion_radius): """Compute distance related members of lDDTScorer Brute force all vs all distance computation kills LDDT for large complexes. Instead of building some KD tree data structure, we make use of expected spatial proximity of atoms in the same chain. Distances are computed as follows: - process each chain individually - perform crude collision detection - process potentially interacting chain pairs - concatenate distances from all processing steps """ ref_indices = [np.asarray([], dtype=np.int32) for idx in range(n_atoms)] ref_distances = [np.asarray([], dtype=np.float32) for idx in range(n_atoms)] indices = [list() for _ in range(n_atoms)] distances = [list() for _ in range(n_atoms)] per_chain_pos = list() per_chain_indices = list() # Process individual chains for ch in structure.chains: pos_list = list() atom_indices = list() mask_start = list() mask_end = list() r_start_idx = 0 for r_idx, r in enumerate(ch.residues): n_valid_atoms = 0 for a in r.atoms: hash_code = a.handle.GetHashCode() if hash_code in atom_index_mapping: p = a.GetPos() pos_list.append(np.asarray([p[0], p[1], p[2]], dtype=np.float32)) atom_indices.append(atom_index_mapping[hash_code]) n_valid_atoms += 1 mask_start.extend([r_start_idx] * n_valid_atoms) mask_end.extend([r_start_idx + n_valid_atoms] * n_valid_atoms) r_start_idx += n_valid_atoms if len(pos_list) == 0: # nothing to do... continue pos = np.vstack(pos_list) atom_indices = np.asarray(atom_indices, dtype=np.int32) if atom_indices.shape[0] > 20000: dists = blockwise_cdist(pos, pos) else: dists = cdist(pos, pos) # apply masks far_away = 2 * inclusion_radius for idx in range(atom_indices.shape[0]): dists[idx, range(mask_start[idx], mask_end[idx])] = far_away # fish out and store close atoms within inclusion radius within_mask = dists < inclusion_radius for idx in range(atom_indices.shape[0]): indices_to_append = atom_indices[within_mask[idx,:]] if indices_to_append.shape[0] > 0: full_at_idx = atom_indices[idx] indices[full_at_idx].append(indices_to_append) distances[full_at_idx].append(dists[idx, within_mask[idx,:]]) dists = None per_chain_pos.append(pos) per_chain_indices.append(atom_indices) # perform crude collision detection min_pos = [p.min(0) for p in per_chain_pos] max_pos = [p.max(0) for p in per_chain_pos] chain_pairs = list() for idx_one in range(len(per_chain_pos)): for idx_two in range(idx_one + 1, len(per_chain_pos)): if np.max(min_pos[idx_one] - max_pos[idx_two]) > inclusion_radius: continue if np.max(min_pos[idx_two] - max_pos[idx_one]) > inclusion_radius: continue chain_pairs.append((idx_one, idx_two)) # process potentially interacting chains for pair in chain_pairs: if per_chain_pos[pair[0]].shape[0] > 20000 or per_chain_pos[pair[1]].shape[0] > 20000: dists = blockwise_cdist(per_chain_pos[pair[0]], per_chain_pos[pair[1]]) else: dists = cdist(per_chain_pos[pair[0]], per_chain_pos[pair[1]]) within = dists <= inclusion_radius # process pair[0] tmp = within.sum(axis=1) for idx in range(tmp.shape[0]): if tmp[idx] > 0: # even though not being a strict requirement, we perform an # insertion here such that the indices for each atom will be # sorted after the hstack operation at_idx = per_chain_indices[pair[0]][idx] indices_to_insert = per_chain_indices[pair[1]][within[idx,:]] distances_to_insert = dists[idx, within[idx, :]] insertion_idx = len(indices[at_idx]) for i in range(insertion_idx): if indices_to_insert[0] > indices[at_idx][i][0]: insertion_idx = i break indices[at_idx].insert(insertion_idx, indices_to_insert) distances[at_idx].insert(insertion_idx, distances_to_insert) # process pair[1] tmp = within.sum(axis=0) for idx in range(tmp.shape[0]): if tmp[idx] > 0: # even though not being a strict requirement, we perform an # insertion here such that the indices for each atom will be # sorted after the hstack operation at_idx = per_chain_indices[pair[1]][idx] indices_to_insert = per_chain_indices[pair[0]][within[:, idx]] distances_to_insert = dists[within[:, idx], idx] insertion_idx = len(indices[at_idx]) for i in range(insertion_idx): if indices_to_insert[0] > indices[at_idx][i][0]: insertion_idx = i break indices[at_idx].insert(insertion_idx, indices_to_insert) distances[at_idx].insert(insertion_idx, distances_to_insert) dists = None # concatenate distances from all processing steps for at_idx in range(n_atoms): if len(indices[at_idx]) > 0: ref_indices[at_idx] = np.hstack(indices[at_idx]) ref_distances[at_idx] = np.hstack(distances[at_idx]) return (ref_indices, ref_distances) @staticmethod def _SetupDistancesSC(n_atoms, chain_start_indices, ref_indices, ref_distances): """Select subset of contacts only covering intra-chain contacts """ # init ref_indices_sc = [np.asarray([], dtype=np.int32) for idx in range(n_atoms)] ref_distances_sc = [np.asarray([], dtype=np.float32) for idx in range(n_atoms)] n_chains = len(chain_start_indices) for ch_idx in range(n_chains): chain_s = chain_start_indices[ch_idx] chain_e = n_atoms if ch_idx + 1 < n_chains: chain_e = chain_start_indices[ch_idx+1] for i in range(chain_s, chain_e): if len(ref_indices[i]) > 0: intra_idx = np.where(np.logical_and(ref_indices[i]>=chain_s, ref_indices[i] 0: inter_idx = np.where(np.logical_or(ref_indices[i]=chain_e))[0] ref_indices_ic[i] = ref_indices[i][inter_idx] ref_distances_ic[i] = ref_distances[i][inter_idx] return (ref_indices_ic, ref_distances_ic) @staticmethod def _NonSymDistances(n_atoms, symmetric_atoms, ref_indices, ref_distances): """Transfer indices/distances of non-symmetric atoms and return """ sym_ref_indices = [np.asarray([], dtype=np.int32) for idx in range(n_atoms)] sym_ref_distances = [np.asarray([], dtype=np.float32) for idx in range(n_atoms)] for idx in symmetric_atoms: indices = list() distances = list() for i, d in zip(ref_indices[idx], ref_distances[idx]): if i not in symmetric_atoms: indices.append(i) distances.append(d) sym_ref_indices[idx] = indices sym_ref_distances[idx] = np.asarray(distances) return (sym_ref_indices, sym_ref_distances) def _EvalAtom(self, pos, atom_idx, thresholds, ref_indices, ref_distances): """Computes number of distance differences within given thresholds returns np.array with len(thresholds) elements """ a_p = pos[atom_idx, :] tmp = pos.take(ref_indices[atom_idx], axis=0) np.subtract(tmp, a_p[None, :], out=tmp) np.square(tmp, out=tmp) tmp = tmp.sum(axis=1) np.sqrt(tmp, out=tmp) # distances against all relevant atoms np.subtract(ref_distances[atom_idx], tmp, out=tmp) np.absolute(tmp, out=tmp) # absolute dist diffs return np.asarray([(tmp <= thresh).sum() for thresh in thresholds], dtype=np.int32) def _EvalAtoms( self, pos, atom_indices, thresholds, ref_indices, ref_distances ): """Calls _EvalAtom for several atoms and sums up the computed number of distance differences within given thresholds returns numpy matrix of shape (n_atoms, len(threshold)) """ conserved = np.zeros((len(atom_indices), len(thresholds)), dtype=np.int32) for a_idx, a in enumerate(atom_indices): conserved[a_idx, :] = self._EvalAtom(pos, a, thresholds, ref_indices, ref_distances) return conserved def _EvalResidues(self, pos, thresholds, res_atom_indices, ref_indices, ref_distances): """Calls _EvalAtoms for a bunch of residues residues are defined in *res_atom_indices* as lists of atom indices returns numpy matrix of shape (n_residues, len(thresholds)). """ conserved = np.zeros((len(res_atom_indices), len(thresholds)), dtype=np.int32) for rai_idx, rai in enumerate(res_atom_indices): conserved[rai_idx,:] = np.sum(self._EvalAtoms(pos, rai, thresholds, ref_indices, ref_distances), axis=0) return conserved def _ProcessSequenceSeparation(self): if self.sequence_separation != 0: raise NotImplementedError("Congratulations! You're the first one " "requesting a non-default " "sequence_separation in the new and " "awesome LDDT implementation. A crate of " "beer for Gabriel and he'll implement " "it.") def _GetNExp(self, atom_idx, ref_indices): """Returns number of close atoms around one or several atoms """ if isinstance(atom_idx, int): return len(ref_indices[atom_idx]) elif isinstance(atom_idx, list): return sum([len(ref_indices[idx]) for idx in atom_idx]) else: raise RuntimeError("invalid input type") def _ResolveSymmetries(self, pos, thresholds, symmetries, sym_ref_indices, sym_ref_distances): """Swaps symmetric positions in-place in order to maximize LDDT scores towards non-symmetric atoms. """ for sym in symmetries: atom_indices = list() for sym_tuple in sym: atom_indices += [sym_tuple[0], sym_tuple[1]] tot = self._GetNExp(atom_indices, sym_ref_indices) if tot == 0: continue # nothing to do # score as is sym_one_conserved = self._EvalAtoms( pos, atom_indices, thresholds, sym_ref_indices, sym_ref_distances, ) # switch positions and score again for pair in sym: pos[[pair[0], pair[1]]] = pos[[pair[1], pair[0]]] sym_two_conserved = self._EvalAtoms( pos, atom_indices, thresholds, sym_ref_indices, sym_ref_distances, ) sym_one_score = np.sum(sym_one_conserved) / (len(thresholds) * tot) sym_two_score = np.sum(sym_two_conserved) / (len(thresholds) * tot) if sym_one_score >= sym_two_score: # switch back, initial positions were better or equal # for the equal case: we still switch back to reproduce the old # LDDT behaviour for pair in sym: pos[[pair[0], pair[1]]] = pos[[pair[1], pair[0]]] def _EvalAtomSSD(self, pos, atom_idx, dist_cap, ref_indices, ref_distances): """ Computes summed squared distances distances are capped at dist_cap """ a_p = pos[atom_idx, :] tmp = pos.take(ref_indices[atom_idx], axis=0) np.subtract(tmp, a_p[None, :], out=tmp) np.square(tmp, out=tmp) tmp = tmp.sum(axis=1) np.sqrt(tmp, out=tmp) # distances against all relevant atoms np.subtract(ref_distances[atom_idx], tmp, out=tmp) # distance difference np.square(tmp, out=tmp) # squared distance difference squared_dist_cap = dist_cap*dist_cap tmp[tmp > squared_dist_cap] = squared_dist_cap return tmp.sum() def _EvalAtomsSSD( self, pos, atom_indices, dist_cap, ref_indices, ref_distances ): """Calls _EvalAtomSSD for several atoms """ return np.asarray([self._EvalAtomSSD(pos, a, dist_cap, ref_indices, ref_distances) for a in atom_indices], dtype=np.float32) def _ResolveSymmetriesSSD(self, pos, dist_cap, symmetries, sym_ref_indices, sym_ref_distances): """Swaps symmetric positions in-place in order to maximize summed squared distances towards non-symmetric atoms. """ for sym in symmetries: atom_indices = list() for sym_tuple in sym: atom_indices += [sym_tuple[0], sym_tuple[1]] tot = self._GetNExp(atom_indices, sym_ref_indices) if tot == 0: continue # nothing to do # score as is sym_one_ssd = self._EvalAtomsSSD( pos, atom_indices, dist_cap, sym_ref_indices, sym_ref_distances, ) # switch positions and score again for pair in sym: pos[[pair[0], pair[1]]] = pos[[pair[1], pair[0]]] sym_two_ssd = self._EvalAtomsSSD( pos, atom_indices, dist_cap, sym_ref_indices, sym_ref_distances, ) sym_one_score = np.sum(sym_one_ssd) sym_two_score = np.sum(sym_two_ssd) if sym_one_score < sym_two_score: # switch back, initial positions were better for pair in sym: pos[[pair[0], pair[1]]] = pos[[pair[1], pair[0]]]