Coverage for NeuralTSNE/NeuralTSNE/TSNE/tests/test_helpers.py: 96%

1from unittest.mock import MagicMock, patch 

2 

3import numpy as np 

4import pytest 

5import torch 

6 

7from NeuralTSNE.TSNE.Helpers import Hbeta, x2p, x2p_job 

8 

9 

10@pytest.mark.parametrize( 

11 "D, beta, expected_H, expected_P", 

12 [ 

13 ( 

14 torch.tensor([[0.0, 2.0], [2.0, 0.0]]), 

15 0.5, 

16 1.2754, 

17 torch.tensor([[0.3655, 0.1345], [0.1345, 0.3655]]), 

18 ), 

19 ( 

20 torch.tensor([[0.0, 1.0, 2.0], [1.0, 0.0, 3.0], [2.0, 3.0, 0.0]]), 

21 0.7, 

22 1.9558, 

23 torch.tensor( 

24 [ 

25 [0.2114, 0.1050, 0.0521], 

26 [0.1050, 0.2114, 0.0259], 

27 [0.0521, 0.0259, 0.2114], 

28 ] 

29 ), 

30 ), 

31 ], 

32) 

33def test_Hbeta(D, beta, expected_H, expected_P): 

34 H, P = Hbeta(D, beta) 

35 assert torch.isclose(H, torch.tensor(expected_H), rtol=1e-3) 

36 assert torch.allclose(P, expected_P, rtol=1e-3) 

37 

38 

39@pytest.mark.parametrize("i", [7, 1, 21]) 

40@pytest.mark.parametrize("perplexity", [10, 50, 1000]) 

41@pytest.mark.parametrize("tolerance", [0.1, 0.01, 0.001, 1e-6]) 

42@pytest.mark.parametrize("max_iterations", [100, 200, 50]) 

43@pytest.mark.parametrize( 

44 "D", 

45 [ 

46 torch.tensor([17.0, 89.0, 123.0, 40.0, 67.0]), 

47 torch.tensor([0.6, 0.2311, 0.456, 0.01, 1.53]), 

48 ], 

49) # TODO: CHECK IF THIS IS CORRECT 

50def test_x2p_job( 

51 i: int, 

52 perplexity: int, 

53 D: torch.Tensor, 

54 tolerance: float, 

55 max_iterations: int, 

56): 

57 logU = torch.tensor([np.log(perplexity)], dtype=torch.float32) 

58 data = i, D, logU 

59 result = x2p_job(data, tolerance, max_iterations) 

60 i, P, Hdiff, iterations = result 

61 assert i == i 

62 if iterations != max_iterations: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 assert torch.allclose(Hdiff, torch.zeros_like(Hdiff), rtol=tolerance) 

64 else: 

65 estimated_tolerance = 10 ** torch.ceil(torch.log10(torch.abs(Hdiff)))[0] 

66 assert torch.isclose(torch.zeros_like(Hdiff), Hdiff, atol=estimated_tolerance) 

67 

68 

69@pytest.mark.parametrize( 

70 "X, perplexity, tolerance, expected_shape", 

71 [ 

72 [torch.tensor([[1.0, 2.0], [3.0, 4.0]]), 2, 0.1, (2, 2)], 

73 [ 

74 torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), 

75 3, 

76 0.01, 

77 (3, 3), 

78 ], 

79 [ 

80 torch.tensor( 

81 [ 

82 [1.0, 6.0, 4.0, 2.0, 5.0], 

83 [7.0, 2.0, 6.0, 7.0, 3.0], 

84 [8.0, 2.0, 1.0, 7.0, 9.0], 

85 ] 

86 ), 

87 5, 

88 0.001, 

89 (3, 3), 

90 ], 

91 ], 

92) 

93@patch("NeuralTSNE.TSNE.Helpers.helpers.x2p_job") 

94def test_x2p( 

95 mock_x2p_job: MagicMock, 

96 X: torch.Tensor, 

97 perplexity: int, 

98 tolerance: float, 

99 expected_shape: tuple, 

100): 

101 log_perplexity = torch.tensor([np.log(perplexity)], dtype=torch.float32) 

102 pair_squared_distances = torch.cdist(X, X, p=2) ** 2 

103 

104 pairwise_distances = pair_squared_distances[pair_squared_distances != 0].reshape( 

105 pair_squared_distances.shape[0], -1 

106 ) 

107 

108 returned_P = [ 

109 (i, torch.ones(pairwise_distances.shape[1])) for i in range(X.shape[0]) 

110 ] 

111 

112 mock_x2p_job.side_effect = returned_P 

113 

114 P = x2p(X, perplexity, tolerance) 

115 assert P.shape == expected_shape 

116 assert mock_x2p_job.call_count == X.shape[0] 

117 mock_x2p_calls = mock_x2p_job.call_args_list 

118 

119 P_diag = P.diag() 

120 assert torch.allclose(P_diag, torch.zeros_like(P_diag), rtol=1e-5) 

121 P_no_diag = P[P != 0].reshape(P.shape[0], -1) 

122 assert torch.allclose(P_no_diag, torch.ones_like(P_no_diag), rtol=1e-5) 

123 

124 for k, call_args in enumerate(mock_x2p_calls): 

125 data, tol = call_args[0] 

126 i, D, logU = data 

127 

128 assert i == k 

129 assert torch.allclose(D, pairwise_distances[k], rtol=1e-5) 

130 assert torch.isclose(logU, log_perplexity, rtol=1e-5) 

131 assert tol == tolerance