From 926be60f4249cbb33bfc917d765d3ba250470b26 Mon Sep 17 00:00:00 2001 From: samuelmurail Date: Fri, 5 Jul 2024 00:37:37 +0200 Subject: [PATCH 1/2] Fix iptm chain --- alphafold/common/confidence.py | 77 ++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/alphafold/common/confidence.py b/alphafold/common/confidence.py index e5ea99af0..0516c4bff 100644 --- a/alphafold/common/confidence.py +++ b/alphafold/common/confidence.py @@ -16,6 +16,7 @@ import jax.numpy as jnp import jax +from typing import Any, Mapping, Optional, Union, List, Dict import numpy as np from alphafold.common import residue_constants import scipy.special @@ -168,6 +169,75 @@ def predicted_tm_score(logits, breaks, residue_weights = None, return (per_alignment * residue_weights).max() + +def predicted_tm_score_new(logits, breaks, residue_weights = None, + asym_id = None, use_jnp=False): + """Computes predicted TM alignment or predicted interface TM alignment score. + + Args: + logits: [num_res, num_res, num_bins] the logits output from + PredictedAlignedErrorHead. + breaks: [num_bins] the error bins. + residue_weights: [num_res] the per residue weights to use for the + expectation. + asym_id: [num_res] the asymmetric unit ID - the chain ID. Only needed for + ipTM calculation. + + Returns: + ptm_score: The predicted TM alignment or the predicted iTM score. + """ + if use_jnp: + _np, _softmax = jnp, jax.nn.softmax + else: + _np, _softmax = np, scipy.special.softmax + + # residue_weights has to be in [0, 1], but can be floating-point, i.e. the + # exp. resolved head's probability. + if residue_weights is None: + residue_weights = _np.ones(logits.shape[0]) + + bin_centers = _calculate_bin_centers(breaks, use_jnp=use_jnp) + num_res = residue_weights.shape[0] + + # Clip num_res to avoid negative/undefined d0. + clipped_num_res = _np.maximum(residue_weights.sum(), 19) + + # Compute d_0(num_res) as defined by TM-score, eqn. (5) in Yang & Skolnick + # "Scoring function for automated assessment of protein structure template + # quality", 2004: http://zhanglab.ccmb.med.umich.edu/papers/2004_3.pdf + d0 = 1.24 * (clipped_num_res - 15) ** (1./3) - 1.8 + + # Convert logits to probs. + probs = _softmax(logits, axis=-1) + + # TM-Score term for every bin. + tm_per_bin = 1. / (1 + _np.square(bin_centers) / _np.square(d0)) + # E_distances tm(distance). + predicted_tm_term = (probs * tm_per_bin).sum(-1) + + ## TO SOLVE, HOW TO GET CHAIN NUMBER !!! + + chain_num = jnp.max(asym_id) + 1 + + def get_cross_iptm(i, j): + pair_mask = jnp.logical_and(i * jnp.ones((num_res))[:, None] == asym_id[None, :] , j*jnp.ones((num_res))[None, :] == asym_id[:, None]) + chain_chain_predicted_tm_term = predicted_tm_term * pair_mask + pair_residue_weights = pair_mask * (residue_weights[None, :] * residue_weights[:, None]) + normed_residue_mask = pair_residue_weights / (1e-8 + pair_residue_weights.sum(-1, keepdims=True)) + per_alignment = (chain_chain_predicted_tm_term * normed_residue_mask).sum(-1) + return (per_alignment * residue_weights).max() + + iptm_matrix_list = [] + + for i in jnp.arange(chain_num): + local_list = [] + for j in jnp.arange(chain_num): + local_list.append(get_cross_iptm(i, j)) + iptm_matrix_list.append(local_list) + + return(iptm_matrix_list) + + def get_confidence_metrics(prediction_result, mask, rank_by = "plddt", use_jnp=False): """Post processes prediction_result to get confidence metrics.""" confidence_metrics = {} @@ -195,6 +265,13 @@ def get_confidence_metrics(prediction_result, mask, rank_by = "plddt", use_jnp=F residue_weights=mask, asym_id=prediction_result['predicted_aligned_error']['asym_id'], use_jnp=use_jnp) + confidence_metrics['new_iptm'] = predicted_tm_score_new( + logits=prediction_result['predicted_aligned_error']['logits'], + breaks=prediction_result['predicted_aligned_error']['breaks'], + residue_weights=mask, + asym_id=prediction_result['predicted_aligned_error']['asym_id'], + use_jnp=use_jnp + ) # compute mean_score if rank_by == "multimer": From fcfbe2b4ead4a30fcad3806a2ffcb44ae1d3e907 Mon Sep 17 00:00:00 2001 From: samuelmurail Date: Fri, 5 Jul 2024 11:23:03 +0200 Subject: [PATCH 2/2] Add chain_num as argument --- alphafold/common/confidence.py | 22 ++++++++++++---------- alphafold/model/modules_multimer.py | 8 +++++++- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/alphafold/common/confidence.py b/alphafold/common/confidence.py index 0516c4bff..e4979eccb 100644 --- a/alphafold/common/confidence.py +++ b/alphafold/common/confidence.py @@ -170,8 +170,8 @@ def predicted_tm_score(logits, breaks, residue_weights = None, return (per_alignment * residue_weights).max() -def predicted_tm_score_new(logits, breaks, residue_weights = None, - asym_id = None, use_jnp=False): +def predicted_tm_score_chain(logits, breaks, residue_weights = None, + asym_id = None, use_jnp=False, chain_num=None): """Computes predicted TM alignment or predicted interface TM alignment score. Args: @@ -191,6 +191,9 @@ def predicted_tm_score_new(logits, breaks, residue_weights = None, else: _np, _softmax = np, scipy.special.softmax + if chain_num is None: + chain_num = 1 + # residue_weights has to be in [0, 1], but can be floating-point, i.e. the # exp. resolved head's probability. if residue_weights is None: @@ -215,10 +218,8 @@ def predicted_tm_score_new(logits, breaks, residue_weights = None, # E_distances tm(distance). predicted_tm_term = (probs * tm_per_bin).sum(-1) - ## TO SOLVE, HOW TO GET CHAIN NUMBER !!! - - chain_num = jnp.max(asym_id) + 1 - + # jax.debug.print('residue weights = {x}',x=residue_weights) + def get_cross_iptm(i, j): pair_mask = jnp.logical_and(i * jnp.ones((num_res))[:, None] == asym_id[None, :] , j*jnp.ones((num_res))[None, :] == asym_id[:, None]) chain_chain_predicted_tm_term = predicted_tm_term * pair_mask @@ -235,10 +236,10 @@ def get_cross_iptm(i, j): local_list.append(get_cross_iptm(i, j)) iptm_matrix_list.append(local_list) - return(iptm_matrix_list) + return(iptm_matrix_list, predicted_tm_term) -def get_confidence_metrics(prediction_result, mask, rank_by = "plddt", use_jnp=False): +def get_confidence_metrics(prediction_result, mask, rank_by = "plddt", use_jnp=False, chain_num=None): """Post processes prediction_result to get confidence metrics.""" confidence_metrics = {} plddt = compute_plddt(prediction_result['predicted_lddt']['logits'], use_jnp=use_jnp) @@ -265,12 +266,13 @@ def get_confidence_metrics(prediction_result, mask, rank_by = "plddt", use_jnp=F residue_weights=mask, asym_id=prediction_result['predicted_aligned_error']['asym_id'], use_jnp=use_jnp) - confidence_metrics['new_iptm'] = predicted_tm_score_new( + confidence_metrics['chain_iptm'], confidence_metrics['ptm_matrix'] = predicted_tm_score_chain( logits=prediction_result['predicted_aligned_error']['logits'], breaks=prediction_result['predicted_aligned_error']['breaks'], residue_weights=mask, asym_id=prediction_result['predicted_aligned_error']['asym_id'], - use_jnp=use_jnp + use_jnp=use_jnp, + chain_num=chain_num, ) # compute mean_score diff --git a/alphafold/model/modules_multimer.py b/alphafold/model/modules_multimer.py index 5071c7f9f..afc5c0de7 100644 --- a/alphafold/model/modules_multimer.py +++ b/alphafold/model/modules_multimer.py @@ -461,12 +461,18 @@ def apply_network(prev, safe_key): if not return_representations: del ret['representations'] + # Extract chain NUM # MODIFIED + chain_num = c.embeddings_and_evoformer.max_relative_chain + 1 # MODIFIED + # add confidence metrics + ret.update(confidence.get_confidence_metrics( prediction_result=ret, mask=batch["seq_mask"], rank_by=self.config.rank_by, - use_jnp=True)) + use_jnp=True, + chain_num=chain_num # MODIFIED + )) ret["tol"] = confidence.compute_tol( prev["prev_pos"],