# Path: alphafold/model/lddt.py
# def lddt(predicted_points,
#          true_points,
#          true_points_mask,
#          cutoff=15.,
#          per_residue=False):
#   """Measure (approximate) lDDT for a batch of coordinates.
# 
#   lDDT reference:
#   Mariani, V., Biasini, M., Barbato, A. & Schwede, T. lDDT: A local
#   superposition-free score for comparing protein structures and models using
#   distance difference tests. Bioinformatics 29, 2722–2728 (2013).
# 
#   lDDT is a measure of the difference between the true distance matrix and the
#   distance matrix of the predicted points.  The difference is computed only on
#   points closer than cutoff *in the true structure*.
# 
#   This function does not compute the exact lDDT value that the original paper
#   describes because it does not include terms for physical feasibility
#   (e.g. bond length violations). Therefore this is only an approximate
#   lDDT score.
# 
#   Args:
#     predicted_points: (batch, length, 3) array of predicted 3D points
#     true_points: (batch, length, 3) array of true 3D points
#     true_points_mask: (batch, length, 1) binary-valued float array.  This mask
#       should be 1 for points that exist in the true points.
#     cutoff: Maximum distance for a pair of points to be included
#     per_residue: If true, return score for each residue.  Note that the overall
#       lDDT is not exactly the mean of the per_residue lDDT's because some
#       residues have more contacts than others.
# 
#   Returns:
#     An (approximate, see above) lDDT score in the range 0-1.
#   """
# 
#   assert len(predicted_points.shape) == 3
#   assert predicted_points.shape[-1] == 3
#   assert true_points_mask.shape[-1] == 1
#   assert len(true_points_mask.shape) == 3
# 
#   # Compute true and predicted distance matrices.
#   dmat_true = jnp.sqrt(1e-10 + jnp.sum(
#       (true_points[:, :, None] - true_points[:, None, :])**2, axis=-1))
# 
#   dmat_predicted = jnp.sqrt(1e-10 + jnp.sum(
#       (predicted_points[:, :, None] -
#        predicted_points[:, None, :])**2, axis=-1))
# 
#   dists_to_score = (
#       (dmat_true < cutoff).astype(jnp.float32) * true_points_mask *
#       jnp.transpose(true_points_mask, [0, 2, 1]) *
#       (1. - jnp.eye(dmat_true.shape[1]))  # Exclude self-interaction.
#   )
# 
#   # Shift unscored distances to be far away.
#   dist_l1 = jnp.abs(dmat_true - dmat_predicted)
# 
#   # True lDDT uses a number of fixed bins.
#   # We ignore the physical plausibility correction to lDDT, though.
#   score = 0.25 * ((dist_l1 < 0.5).astype(jnp.float32) +
#                   (dist_l1 < 1.0).astype(jnp.float32) +
#                   (dist_l1 < 2.0).astype(jnp.float32) +
#                   (dist_l1 < 4.0).astype(jnp.float32))
# 
#   # Normalize over the appropriate axes.
#   reduce_axes = (-1,) if per_residue else (-2, -1)
#   norm = 1. / (1e-10 + jnp.sum(dists_to_score, axis=reduce_axes))
#   score = norm * (1e-10 + jnp.sum(dists_to_score * score, axis=reduce_axes))
# 
#   return score


from absl.testing import absltest
from absl.testing import parameterized
from alphafold.model import lddt
import numpy as np
       [1, 1]),
      ('one_a_dist',
       [[0, 0, 0], [5, 0, 0]],
       [[0, 0, 0], [6-1e-5, 0, 0]],
       [0.75, 0.75]),
      ('two_a_dist',
       [[0, 0, 0], [5, 0, 0]],
       [[0, 0, 0], [7-1e-5, 0, 0]],
       [0.5, 0.5]),
      ('four_a_dist',
       [[0, 0, 0], [5, 0, 0]],
       [[0, 0, 0], [9-1e-5, 0, 0]],
       [0.25, 0.25],),
      ('five_a_dist',
       [[0, 0, 0], [16-1e-5, 0, 0]],
       [[0, 0, 0], [11, 0, 0]],
       [0, 0]),
      ('no_pairs',
       [[0, 0, 0], [20, 0, 0]],
       [[0, 0, 0], [25-1e-5, 0, 0]],
       [1, 1]),
  )
  def test_lddt(
      self, predicted_pos, true_pos, exp_lddt):
    predicted_pos = np.array([predicted_pos], dtype=np.float32)
    true_points_mask = np.array([[[1]] * len(true_pos)], dtype=np.float32)
    true_pos = np.array([true_pos], dtype=np.float32)
    cutoff = 15.0
    per_residue = True


# Path: alphafold/model/lddt.py
# def lddt(predicted_points,
#          true_points,
#          true_points_mask,
#          cutoff=15.,
#          per_residue=False):
#   """Measure (approximate) lDDT for a batch of coordinates.
# 
#   lDDT reference:
#   Mariani, V., Biasini, M., Barbato, A. & Schwede, T. lDDT: A local
#   superposition-free score for comparing protein structures and models using
#   distance difference tests. Bioinformatics 29, 2722–2728 (2013).
# 
#   lDDT is a measure of the difference between the true distance matrix and the
#   distance matrix of the predicted points.  The difference is computed only on
#   points closer than cutoff *in the true structure*.
# 
#   This function does not compute the exact lDDT value that the original paper
#   describes because it does not include terms for physical feasibility
#   (e.g. bond length violations). Therefore this is only an approximate
#   lDDT score.
# 
#   Args:
#     predicted_points: (batch, length, 3) array of predicted 3D points
#     true_points: (batch, length, 3) array of true 3D points
#     true_points_mask: (batch, length, 1) binary-valued float array.  This mask
#       should be 1 for points that exist in the true points.
#     cutoff: Maximum distance for a pair of points to be included
#     per_residue: If true, return score for each residue.  Note that the overall
#       lDDT is not exactly the mean of the per_residue lDDT's because some
#       residues have more contacts than others.
# 
#   Returns:
#     An (approximate, see above) lDDT score in the range 0-1.
#   """
# 
#   assert len(predicted_points.shape) == 3
#   assert predicted_points.shape[-1] == 3
#   assert true_points_mask.shape[-1] == 1
#   assert len(true_points_mask.shape) == 3
# 
#   # Compute true and predicted distance matrices.
#   dmat_true = jnp.sqrt(1e-10 + jnp.sum(
#       (true_points[:, :, None] - true_points[:, None, :])**2, axis=-1))
# 
#   dmat_predicted = jnp.sqrt(1e-10 + jnp.sum(
#       (predicted_points[:, :, None] -
#        predicted_points[:, None, :])**2, axis=-1))
# 
#   dists_to_score = (
#       (dmat_true < cutoff).astype(jnp.float32) * true_points_mask *
#       jnp.transpose(true_points_mask, [0, 2, 1]) *
#       (1. - jnp.eye(dmat_true.shape[1]))  # Exclude self-interaction.
#   )
# 
#   # Shift unscored distances to be far away.
#   dist_l1 = jnp.abs(dmat_true - dmat_predicted)
# 
#   # True lDDT uses a number of fixed bins.
#   # We ignore the physical plausibility correction to lDDT, though.
#   score = 0.25 * ((dist_l1 < 0.5).astype(jnp.float32) +
#                   (dist_l1 < 1.0).astype(jnp.float32) +
#                   (dist_l1 < 2.0).astype(jnp.float32) +
#                   (dist_l1 < 4.0).astype(jnp.float32))
# 
#   # Normalize over the appropriate axes.
#   reduce_axes = (-1,) if per_residue else (-2, -1)
#   norm = 1. / (1e-10 + jnp.sum(dists_to_score, axis=reduce_axes))
#   score = norm * (1e-10 + jnp.sum(dists_to_score * score, axis=reduce_axes))
# 
#   return score


# Path: alphafold/model/lddt_test.py
from absl.testing import absltest
from absl.testing import parameterized
from alphafold.model import lddt
import numpy as np

       [1, 1]),
      ('one_a_dist',
       [[0, 0, 0], [5, 0, 0]],
       [[0, 0, 0], [6-1e-5, 0, 0]],
       [0.75, 0.75]),
      ('two_a_dist',
       [[0, 0, 0], [5, 0, 0]],
       [[0, 0, 0], [7-1e-5, 0, 0]],
       [0.5, 0.5]),
      ('four_a_dist',
       [[0, 0, 0], [5, 0, 0]],
       [[0, 0, 0], [9-1e-5, 0, 0]],
       [0.25, 0.25],),
      ('five_a_dist',
       [[0, 0, 0], [16-1e-5, 0, 0]],
       [[0, 0, 0], [11, 0, 0]],
       [0, 0]),
      ('no_pairs',
       [[0, 0, 0], [20, 0, 0]],
       [[0, 0, 0], [25-1e-5, 0, 0]],
       [1, 1]),
  )
  def test_lddt(
      self, predicted_pos, true_pos, exp_lddt):
    predicted_pos = np.array([predicted_pos], dtype=np.float32)
    true_points_mask = np.array([[[1]] * len(true_pos)], dtype=np.float32)
    true_pos = np.array([true_pos], dtype=np.float32)
    cutoff = 15.0
    per_residue = True


    result = lddt.lddt(