Coverage for langbrainscore/interface/metric.py: 38%

1import typing 

2from abc import ABC, abstractmethod 

3 

4import numpy as np 

5from langbrainscore.interface.cacheable import _Cacheable 

6 

7 

8class _Metric(_Cacheable, ABC): 

9 # class _Metric(ABC): 

10 """ 

11 checks that two arrays are comparable for a given similarity metric, 

12 then applies that metric to those inputs and returns score(s) 

13 

14 Args: 

15 np.ndarray: X 

16 np.ndarray: Y 

17 

18 Returns: 

19 Typing.Union[np.ndarray,np.float]: score(s) 

20 

21 Raises: 

22 ValueError: X and Y must be 1D or 2D arrays. 

23 ValueError: X and Y must have the same number of samples. 

24 ValueError: for most metrics, X and Y must have same number of dimensions. 

25 

26 """ 

27 

28 def __init__(self): 

29 pass 

30 

31 def __call__( 

32 self, X: np.ndarray, Y: np.ndarray 

33 ) -> typing.Union[np.float, np.ndarray]: 

34 if X.ndim == 1: 

35 X = X.reshape(-1, 1) 

36 if Y.ndim == 1: 

37 Y = Y.reshape(-1, 1) 

38 if any(y.ndim != 2 for y in [X, Y]): 

39 raise ValueError("X and Y must be 1D or 2D arrays.") 

40 if X.shape[0] != Y.shape[0]: 

41 raise ValueError("X and Y must have the same number of samples.") 

42 if self.__class__.__name__ not in ("RSA", "CKA"): 

43 if X.shape[1] != Y.shape[1]: 

44 raise ValueError("X and Y must have the same number of dimensions.") 

45 

46 score = self._apply_metric(X, Y) 

47 if not isinstance(score, np.ndarray): 

48 return np.array(score).reshape(-1) 

49 return score 

50 

51 @abstractmethod 

52 def _apply_metric( 

53 self, X: np.ndarray, Y: np.ndarray 

54 ) -> typing.Union[np.float, np.ndarray]: 

55 raise NotImplementedError 

56 

57 

58class _VectorMetric(_Metric): 

59 """ 

60 subclass of _Metric that applies relevant vector similarity metric 

61 along each column of the input arrays. 

62 """ 

63 

64 def __init__(self, reduction=None): 

65 """ 

66 args: 

67 callable: reduction (can also be None or False) 

68 

69 raises: 

70 TypeError: if reduction argument is not callable. 

71 """ 

72 if reduction: 

73 if not callable(reduction): 

74 raise TypeError("Reduction argument must be callable.") 

75 self._reduction = reduction 

76 super().__init__() 

77 

78 def _apply_metric( 

79 self, X: np.ndarray, Y: np.ndarray 

80 ) -> typing.Union[np.float, np.ndarray]: 

81 """ 

82 internal function that applies scoring function along each array dimension 

83 and then optionally applies a reduction, e.g., np.mean 

84 

85 args: 

86 np.ndarray: X 

87 np.ndarray: Y 

88 

89 """ 

90 scores = np.zeros(X.shape[1]) 

91 for i in range(scores.size): 

92 x = X[:, i] 

93 y = Y[:, i] 

94 nan = np.isnan(x) | np.isnan(y) 

95 try: 

96 scores[i] = self._score(x[~nan], y[~nan]) 

97 except: 

98 scores[i] = np.nan 

99 if self._reduction: 

100 return self._reduction(scores) 

101 if len(scores) == 1: 

102 return scores[0] 

103 return scores 

104 

105 @abstractmethod 

106 def _score(self, X: np.ndarray, Y: np.ndarray) -> np.float: 

107 raise NotImplementedError 

108 

109 

110class _MatrixMetric(_Metric): 

111 """ 

112 interface for similarity metrics that operate over entire matrices, e.g., RSA 

113 """ 

114 

115 def __init__(self): 

116 super().__init__() 

117 

118 def _apply_metric(self, X: np.ndarray, Y: np.ndarray) -> np.float: 

119 score = self._score(X, Y) 

120 return score 

121 

122 @abstractmethod 

123 def _score(self, X: np.ndarray, Y: np.ndarray) -> np.float: 

124 raise NotImplementedError