






import math
import unittest






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				assert isinstance(lowerCamelCase_		,   lowerCamelCase_				) and (
				    number >= 0
				), "'number' must been an int and positive"

				if 1 < number < 4:
								# 2 and 3 are primes
								return True
				elif number < 2 or number % 2 == 0 or number % 3 == 0:
								# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
								return False

				# All primes number are in format of 6k +/- 1
				for i in range(5		,   int(math.sqrt(lowerCamelCase_				) + 1				)		,   6				):
								if number % i == 0 or number % (i + 2) == 0:
												return False
				return True
class       _UpperCAmelCase					(       unittest.TestCase	):

				"""simple docstring"""




				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								self.assertTrue(is_prime(2  )  )
								self.assertTrue(is_prime(3  )  )
								self.assertTrue(is_prime(5  )  )
								self.assertTrue(is_prime(7  )  )
								self.assertTrue(is_prime(11  )  )
								self.assertTrue(is_prime(13  )  )
								self.assertTrue(is_prime(17  )  )
								self.assertTrue(is_prime(19  )  )
								self.assertTrue(is_prime(23  )  )
								self.assertTrue(is_prime(29  )  )




				def 				lowercase__   (     self						:    Any  ):

								'''simple docstring'''
								with self.assertRaises(lowerCamelCase  ):
												is_prime(-19  )
								self.assertFalse(
								    is_prime(0  ),  '''Zero doesn\'t have any positive factors, primes must have exactly two.''',  )
								self.assertFalse(
								    is_prime(1  ),  '''One only has 1 positive factor, primes must have exactly two.''',  )
								self.assertFalse(is_prime(2 * 2  )  )
								self.assertFalse(is_prime(2 * 3  )  )
								self.assertFalse(is_prime(3 * 3  )  )
								self.assertFalse(is_prime(3 * 5  )  )
								self.assertFalse(is_prime(3 * 5 * 7  )  )


if __name__ == "__main__":
							unittest.main()












from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			[]
				lowercase__    =			[]
				lowercase__    =			[]

				for rt in rc.restypes:
								lowercase__    =			rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
								restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]				)
								lowercase__    =			{name: i for i, name in enumerate(lowerCamelCase_				)}
								restype_atomaa_to_atomaa_list.append(
								    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]				)

								restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]				)

				# Add dummy mapping for restype 'UNK'
				restype_atomaa_to_atomaa_list.append([0] * 14				)
				restype_atomaa_to_atomaa_list.append([0] * 37				)
				restype_atomaa_mask_list.append([0.0] * 14				)

				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.intaa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.intaa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.floataa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			protein['''aatype'''].to(torch.long				)

				# create the mapping for (residx, atom14) --> atom37, i.e. an array
				# with shape (num_res, 14) containing the atom37 indices for this protein
				lowercase__    =			restype_atomaa_to_atomaa[protein_aatype]
				lowercase__    =			restype_atomaa_mask[protein_aatype]

				lowercase__    =			residx_atomaa_mask
				lowercase__    =			residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				lowercase__    =			restype_atomaa_to_atomaa[protein_aatype]
				lowercase__    =			residx_atomaa_to_atomaa.long()

				# create the corresponding mask
				lowercase__    =			torch.zeros([21, 37]		,   dtype=torch.floataa		,   device=protein['''aatype'''].device				)
				for restype, restype_letter in enumerate(rc.restypes				):
								lowercase__    =			rc.restype_atoa[restype_letter]
								lowercase__    =			rc.residue_atoms[restype_name]
								for atom_name in atom_names:
												lowercase__    =			rc.atom_order[atom_name]
												lowercase__    =			1

				lowercase__    =			restype_atomaa_mask[protein_aatype]
				lowercase__    =			residx_atomaa_mask

				return protein






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			tree_map(lambda lowerCamelCase_				: torch.tensor(lowerCamelCase_		,   device=batch['''aatype'''].device				)		,   lowerCamelCase_		,   np.ndarray				)
				lowercase__    =			tensor_tree_map(lambda lowerCamelCase_				: np.array(lowerCamelCase_				)		,   make_atomaa_masks(lowerCamelCase_				)				)
				return out




