




"""simple docstring"""




from typing import Dict, Optional, Tuple

import torch







def       __UpperCAmelCase      (				__UpperCamelCase					):
   __lowercase	:       Union[str, Any]				=   boundaries[1] - boundaries[0]
   __lowercase	:       Optional[int]				=   boundaries + step / 2
   __lowercase	:       int				=   torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1					)]					,							dim=0					)
   return bin_centers







def       __UpperCAmelCase      (				__UpperCamelCase					,							__UpperCamelCase					,							):
   __lowercase	:       List[Any]				=   _calculate_bin_centers(__lowercase					)
   return (
       torch.sum(aligned_distance_error_probs * bin_centers					,							dim=-1					),
       bin_centers[-1],
   )







def       __UpperCAmelCase      (				__UpperCamelCase					,							__UpperCamelCase = 31					,							__UpperCamelCase = 64					,							**__UpperCamelCase					,							):
   __lowercase	:       List[Any]				=   torch.linspace(0					,							__lowercase					,							steps=(no_bins - 1)					,							device=logits.device					)

   __lowercase	:       Any				=   torch.nn.functional.softmax(__lowercase					,							dim=-1					)
   __lowercase	,__lowercase	:       List[Any]				=   _calculate_expected_aligned_error(
       alignment_confidence_breaks=__lowercase					,							aligned_distance_error_probs=__lowercase					,							)

   return {
       "aligned_confidence_probs": aligned_confidence_probs,
       "predicted_aligned_error": predicted_aligned_error,
       "max_predicted_aligned_error": max_predicted_aligned_error,
   }







def       __UpperCAmelCase      (				__UpperCamelCase					,							__UpperCamelCase = None					,							__UpperCamelCase = 31					,							__UpperCamelCase = 64					,							__UpperCamelCase = 1e-8					,							**__UpperCamelCase					,							):

   if residue_weights is None:
      __lowercase	:       Optional[Any]				=   logits.new_ones(logits.shape[-2]					)

   __lowercase	:       List[str]				=   torch.linspace(0					,							__lowercase					,							steps=(no_bins - 1)					,							device=logits.device					)

   __lowercase	:       Dict				=   _calculate_bin_centers(__lowercase					)
   torch.sum(__lowercase					)
   __lowercase	:       str				=   logits.shape[-2]
   __lowercase	:       str				=   max(__lowercase					,							19					)

   __lowercase	:       int				=   1.24 * (clipped_n - 15) ** (1.0 / 3) - 1.8

   __lowercase	:       Union[str, Any]				=   torch.nn.functional.softmax(__lowercase					,							dim=-1					)

   __lowercase	:       Union[str, Any]				=   1.0 / (1 + (bin_centers**2) / (da**2))
   __lowercase	:       List[Any]				=   torch.sum(probs * tm_per_bin					,							dim=-1					)

   __lowercase	:       Dict				=   residue_weights / (eps + residue_weights.sum())
   __lowercase	:       Optional[Any]				=   torch.sum(predicted_tm_term * normed_residue_mask					,							dim=-1					)

   __lowercase	:       Optional[Any]				=   per_alignment * residue_weights

   __lowercase	:       List[Any]				=   (weighted == torch.max(__lowercase					)).nonzero()[0]
   return per_alignment[tuple(__lowercase					)]









"""simple docstring"""




a_            =   1_0_0_0







def       __UpperCAmelCase      (				__UpperCamelCase					):
			return x * x







def       __UpperCAmelCase      (				__UpperCamelCase					,							__UpperCamelCase					,							__UpperCamelCase					,							__UpperCamelCase = 4					):
			assert callable(
			    __UpperCamelCase					), f"""the function(object) passed should be callable your input : {function}"""
			assert isinstance(__UpperCamelCase					,							(float, int)					), f"""a should be float or integer your input : {a}"""
			assert isinstance(function(__UpperCamelCase					)					,							(float, int)					), (
			    "the function should return integer or float return type of your function, "
			    f"""{type(__UpperCamelCase					)}"""
			)
			assert isinstance(__UpperCamelCase					,							(float, int)					), f"""b should be float or integer your input : {b}"""
			assert (
			    isinstance(__UpperCamelCase					,							__UpperCamelCase					) and precision > 0
			), f"""precision should be positive integer your input : {precision}"""

			# just applying the formula of simpson for approximate integration written in
			# mentioned article in first comment of this file and above this function

			__lowercase	:       Optional[Any]				=   (b - a) / N_STEPS
			__lowercase	:       Optional[int]				=   function(__UpperCamelCase					) + function(__UpperCamelCase					)

			for i in range(1					,							__UpperCamelCase					):
						__lowercase	:       Union[str, Any]				=   a + h * i
						result += function(__UpperCamelCase					) * (4 if i % 2 else 2)

			result *= h / 3
			return round(__UpperCamelCase					,							__UpperCamelCase					)


if __name__ == "__main__":
			import doctest

			doctest.testmod()



