Coverage for langbrainscore/metrics/metric.py: 45%

1import typing 

2 

3import numpy as np 

4import xarray as xr 

5from langbrainscore.interface import _MatrixMetric, _Metric, _VectorMetric 

6from scipy.stats import kendalltau, pearsonr, spearmanr 

7from sklearn.metrics import accuracy_score, mean_squared_error, pairwise_distances 

8 

9 

10# class Metric: 

11# """ 

12# wrapper for metric classes that confirms they instantiate the proper interface 

13# and coordinates their execution over the contents of supplied xarrays 

14# """ 

15 

16# def __init__(self, metric: typing.Union[_Metric, str], **kwargs) -> "Metric": 

17# assert issubclass(metric, _Metric) 

18# self._metric = metric(**kwargs) 

19 

20# def __call__(self, X: xr.DataArray, Y: xr.DataArray) -> np.ndarray: 

21# """ 

22# args: 

23# xr.DataArray: X 

24# xr.DataArray: Y 

25 

26# returns: 

27# score of specified metric applied to X and Y 

28# """ 

29# score = self._metric(X.values, Y.values) 

30# if not isinstance(score, np.ndarray): 

31# return np.array(score).reshape(-1) 

32# return score 

33 

34 

35class PearsonR(_VectorMetric): 

36 @staticmethod 

37 def _score(x: np.ndarray, y: np.ndarray) -> np.float: 

38 r, p = pearsonr(x, y) 

39 return r 

40 

41 

42class SpearmanRho(_VectorMetric): 

43 @staticmethod 

44 def _score(x: np.ndarray, y: np.ndarray) -> np.float: 

45 rho, p = spearmanr(x, y) 

46 return rho 

47 

48 

49class KendallTau(_VectorMetric): 

50 @staticmethod 

51 def _score(x: np.ndarray, y: np.ndarray) -> np.float: 

52 tau, p = kendalltau(x, y) 

53 return tau 

54 

55 

56class FisherCorr(_VectorMetric): 

57 @staticmethod 

58 def _score(x: np.ndarray, y: np.ndarray) -> np.float: 

59 r, p = pearsonr(x, y) 

60 corr = np.arctanh(r) 

61 return corr 

62 

63 

64class RMSE(_VectorMetric): 

65 @staticmethod 

66 def _score(x: np.ndarray, y: np.ndarray) -> np.float: 

67 loss = mean_squared_error(x, y, squared=False) 

68 return loss 

69 

70 

71class ClassificationAccuracy(_VectorMetric): 

72 @staticmethod 

73 def _score(x: np.ndarray, y: np.ndarray) -> np.float: 

74 score = accuracy_score(x, y, normalize=True) 

75 return score 

76 

77 

78class RSA(_MatrixMetric): 

79 """ 

80 evaluates representational similarity between two matrices for a given 

81 distance measure and vector comparison metric 

82 """ 

83 

84 def __init__(self, distance="correlation", comparison=PearsonR()): 

85 """ 

86 args: 

87 string: distance (anything accepted by sklearn.metrics.pairwise_distances) 

88 _VectorMetric: comparison 

89 """ 

90 self._distance = distance 

91 self._comparison = comparison 

92 super().__init__() 

93 

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

95 X_rdm = pairwise_distances(X, metric=self._distance) 

96 Y_rdm = pairwise_distances(Y, metric=self._distance) 

97 if any([m.shape[1] == 1 for m in (X, Y)]): # can't calc 1D corr dists 

98 X_rdm[np.isnan(X_rdm)] = 0 

99 Y_rdm[np.isnan(Y_rdm)] = 0 

100 indices = np.triu_indices(X_rdm.shape[0], k=1) 

101 score = self._comparison(X_rdm[indices], Y_rdm[indices]) 

102 return score 

103 

104 

105# inspired by https://github.com/yuanli2333/CKA-Centered-Kernel-Alignment/blob/master/CKA.py 

106class CKA(_MatrixMetric): 

107 """ 

108 evaluates centered kernel alignment distance between two matrices 

109 currently only implements linear kernel 

110 """ 

111 

112 def __init__(self): 

113 super().__init__() 

114 

115 @staticmethod 

116 def _center(K): 

117 N = K.shape[0] 

118 U = np.ones([N, N]) 

119 I = np.eye(N) 

120 H = I - U / N 

121 centered = H @ K @ H 

122 return centered 

123 

124 def _HSIC(self, A, B): 

125 L_A = A @ A.T 

126 L_B = B @ B.T 

127 HSIC = np.sum(self._center(L_A) * self._center(L_B)) 

128 return HSIC 

129 

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

131 HSIC_XY = self._HSIC(X, Y) 

132 HSIC_XX = self._HSIC(X, X) 

133 HSIC_YY = self._HSIC(Y, Y) 

134 score = HSIC_XY / (np.sqrt(HSIC_XX) * np.sqrt(HSIC_YY)) 

135 return score