mindsponge.metrics.supervised_chi

mindsponge.metrics.supervised_chi(sequence_mask, aatype, sin_cos_true_chi, torsion_angle_mask, sin_cos_pred_chi, sin_cos_unnormalized_pred, chi_weight, angle_norm_weight, chi_pi_periodic)[source]

Computes loss for direct chi angle supervision. The torsion angles are represented by the sine and cosine value of the angle. This loss is composed of 2 items, the error of normalized predicted sine and cosine value, called chi angle difference loss; the other term is the difference between L2 norm of sine cosine value and 1, called angle norm loss. Jumper et al. (2021) Suppl. Alg. 27 “torsionAngleLoss”.

Parameters

  • sequence_mask (Tensor) – The mask tensor for sequence of shape \((N_{res},)\) with \(N_{res}\) the number of residues in protein.

  • aatype (Tensor) – The amino acid type tensor of shape \((N_{res},)\).

  • sin_cos_true_chi (Tensor) – Tensor of shape \((N_{res}, 14)\) which is the sine and cosine value of torsion angles. There are 7 torsion angles per residue, 3 for backbone and 4 for sidechain.

  • torsion_angle_mask (Tensor) – The binary mask for sidechain torsion angles of shape \((N_{res}, 4)\)

  • sin_cos_pred_chi (Tensor) – The predicted sine and cosine value (normalized) of torsion angles of shape \((N_{res}, 4, 2)\).

  • sin_cos_unnormalized_pred (Tensor) – The predicted sine and cosine value (unnormalized) of torsion angles of shape \((N_{recycle}, N_{res}, 7, 2)\) with \(N_{recycle}\) is the recycle number of FoldIteration in Structure module.

  • chi_weight (float) – The weight of chi angle difference loss term, constant.

  • angle_norm_weight (float) – The weight of angle norm loss term, constant.

  • chi_pi_periodic (Tensor) – Chi angles that are pi periodic: they can be rotated by a multiple of pi without affecting the structure. Constants of residues of shape \((21, 4)\), 20 types of amino acids + unknown.

Returns

  • loss (Tensor) - Supervised chi angle loss with shape \(()\) .

Supported Platforms:

Ascend GPU

Examples

>>> import numpy as np
>>> np.random.seed(0)
>>> from mindsponge.metrics import supervised_chi
>>> from mindsponge.common import residue_constants
>>> from mindspore import dtype as mstype
>>> from mindspore import Tensor
>>> sequence_mask = Tensor(np.random.rand(256, )).astype(mstype.float32)
>>> aatype = Tensor(np.random.randint(0, 21, (256,) )).astype(mstype.int32)
>>> sin_cos_true_chi = Tensor(np.random.rand(256, 4, 2)).astype(mstype.float32)
>>> torsion_angle_mask = Tensor(np.random.rand(256, 4)).astype(mstype.float32)
>>> sin_cos_pred_chi = Tensor(np.random.rand(256, 14)).astype(mstype.float32)
>>> sin_cos_unnormalized_pred = Tensor(np.random.rand(8, 256, 7, 2)).astype(mstype.float32)
>>> chi_weight = 0.1
>>> angle_norm_weight = 0.2
>>> chi_pi_periodic = Tensor(residue_constants.chi_pi_periodic).astype(mstype.float32)
>>> chi_loss = supervised_chi(sequence_mask, aatype, sin_cos_true_chi, torsion_angle_mask, sin_cos_pred_chi,
...                           sin_cos_unnormalized_pred, chi_weight, angle_norm_weight, chi_pi_periodic)
>>> print(chi_loss)
0.061829045