tschuprows.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193
  1. # Copyright The Lightning team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import itertools
  15. from typing import Optional
  16. import torch
  17. from torch import Tensor
  18. from typing_extensions import Literal
  19. from torchmetrics.functional.classification.confusion_matrix import _multiclass_confusion_matrix_update
  20. from torchmetrics.functional.nominal.utils import (
  21. _compute_bias_corrected_values,
  22. _compute_chi_squared,
  23. _drop_empty_rows_and_cols,
  24. _handle_nan_in_data,
  25. _nominal_input_validation,
  26. _unable_to_use_bias_correction_warning,
  27. )
  28. def _tschuprows_t_update(
  29. preds: Tensor,
  30. target: Tensor,
  31. num_classes: int,
  32. nan_strategy: Literal["replace", "drop"] = "replace",
  33. nan_replace_value: Optional[float] = 0.0,
  34. ) -> Tensor:
  35. """Compute the bins to update the confusion matrix with for Tschuprow's T calculation.
  36. Args:
  37. preds: 1D or 2D tensor of categorical (nominal) data
  38. target: 1D or 2D tensor of categorical (nominal) data
  39. num_classes: Integer specifying the number of classes
  40. nan_strategy: Indication of whether to replace or drop ``NaN`` values
  41. nan_replace_value: Value to replace ``NaN`s when ``nan_strategy = 'replace```
  42. Returns:
  43. Non-reduced confusion matrix
  44. """
  45. preds = preds.argmax(1) if preds.ndim == 2 else preds
  46. target = target.argmax(1) if target.ndim == 2 else target
  47. preds, target = _handle_nan_in_data(preds, target, nan_strategy, nan_replace_value)
  48. return _multiclass_confusion_matrix_update(preds, target, num_classes)
  49. def _tschuprows_t_compute(confmat: Tensor, bias_correction: bool) -> Tensor:
  50. """Compute Tschuprow's T statistic based on a pre-computed confusion matrix.
  51. Args:
  52. confmat: Confusion matrix for observed data
  53. bias_correction: Indication of whether to use bias correction.
  54. Returns:
  55. Tschuprow's T statistic
  56. """
  57. confmat = _drop_empty_rows_and_cols(confmat)
  58. cm_sum = confmat.sum()
  59. chi_squared = _compute_chi_squared(confmat, bias_correction)
  60. phi_squared = chi_squared / cm_sum
  61. num_rows, num_cols = confmat.shape
  62. if bias_correction:
  63. phi_squared_corrected, rows_corrected, cols_corrected = _compute_bias_corrected_values(
  64. phi_squared, num_rows, num_cols, cm_sum
  65. )
  66. if torch.min(rows_corrected, cols_corrected) == 1:
  67. _unable_to_use_bias_correction_warning(metric_name="Tschuprow's T")
  68. return torch.tensor(float("nan"), device=confmat.device)
  69. tschuprows_t_value = torch.sqrt(phi_squared_corrected / torch.sqrt((rows_corrected - 1) * (cols_corrected - 1)))
  70. else:
  71. n_rows_tensor = torch.tensor(num_rows, device=phi_squared.device)
  72. n_cols_tensor = torch.tensor(num_cols, device=phi_squared.device)
  73. tschuprows_t_value = torch.sqrt(phi_squared / torch.sqrt((n_rows_tensor - 1) * (n_cols_tensor - 1)))
  74. return tschuprows_t_value.clamp(0.0, 1.0)
  75. def tschuprows_t(
  76. preds: Tensor,
  77. target: Tensor,
  78. bias_correction: bool = True,
  79. nan_strategy: Literal["replace", "drop"] = "replace",
  80. nan_replace_value: Optional[float] = 0.0,
  81. ) -> Tensor:
  82. r"""Compute `Tschuprow's T`_ statistic measuring the association between two categorical (nominal) data series.
  83. .. math::
  84. T = \sqrt{\frac{\chi^2 / n}{\sqrt{(r - 1) * (k - 1)}}}
  85. where
  86. .. math::
  87. \chi^2 = \sum_{i,j} \ frac{\left(n_{ij} - \frac{n_{i.} n_{.j}}{n}\right)^2}{\frac{n_{i.} n_{.j}}{n}}
  88. where :math:`n_{ij}` denotes the number of times the values :math:`(A_i, B_j)` are observed with :math:`A_i, B_j`
  89. represent frequencies of values in ``preds`` and ``target``, respectively.
  90. Tschuprow's T is a symmetric coefficient, i.e. :math:`T(preds, target) = T(target, preds)`.
  91. The output values lies in [0, 1] with 1 meaning the perfect association.
  92. Args:
  93. preds: 1D or 2D tensor of categorical (nominal) data:
  94. - 1D shape: (batch_size,)
  95. - 2D shape: (batch_size, num_classes)
  96. target: 1D or 2D tensor of categorical (nominal) data:
  97. - 1D shape: (batch_size,)
  98. - 2D shape: (batch_size, num_classes)
  99. bias_correction: Indication of whether to use bias correction.
  100. nan_strategy: Indication of whether to replace or drop ``NaN`` values
  101. nan_replace_value: Value to replace ``NaN``s when ``nan_strategy = 'replace'``
  102. Returns:
  103. Tschuprow's T statistic
  104. Example:
  105. >>> from torch import randint, round
  106. >>> from torchmetrics.functional.nominal import tschuprows_t
  107. >>> preds = randint(0, 4, (100,))
  108. >>> target = round(preds + torch.randn(100)).clamp(0, 4)
  109. >>> tschuprows_t(preds, target)
  110. tensor(0.4930)
  111. """
  112. _nominal_input_validation(nan_strategy, nan_replace_value)
  113. num_classes = len(torch.cat([preds, target]).unique())
  114. confmat = _tschuprows_t_update(preds, target, num_classes, nan_strategy, nan_replace_value)
  115. return _tschuprows_t_compute(confmat, bias_correction)
  116. def tschuprows_t_matrix(
  117. matrix: Tensor,
  118. bias_correction: bool = True,
  119. nan_strategy: Literal["replace", "drop"] = "replace",
  120. nan_replace_value: Optional[float] = 0.0,
  121. ) -> Tensor:
  122. r"""Compute `Tschuprow's T`_ statistic between a set of multiple variables.
  123. This can serve as a convenient tool to compute Tschuprow's T statistic for analyses of correlation between
  124. categorical variables in your dataset.
  125. Args:
  126. matrix: A tensor of categorical (nominal) data, where:
  127. - rows represent a number of data points
  128. - columns represent a number of categorical (nominal) features
  129. bias_correction: Indication of whether to use bias correction.
  130. nan_strategy: Indication of whether to replace or drop ``NaN`` values
  131. nan_replace_value: Value to replace ``NaN``s when ``nan_strategy = 'replace'``
  132. Returns:
  133. Tschuprow's T statistic for a dataset of categorical variables
  134. Example:
  135. >>> from torch import randint
  136. >>> from torchmetrics.functional.nominal import tschuprows_t_matrix
  137. >>> matrix = randint(0, 4, (200, 5))
  138. >>> tschuprows_t_matrix(matrix)
  139. tensor([[1.0000, 0.0637, 0.0000, 0.0542, 0.1337],
  140. [0.0637, 1.0000, 0.0000, 0.0000, 0.0000],
  141. [0.0000, 0.0000, 1.0000, 0.0000, 0.0649],
  142. [0.0542, 0.0000, 0.0000, 1.0000, 0.1100],
  143. [0.1337, 0.0000, 0.0649, 0.1100, 1.0000]])
  144. """
  145. _nominal_input_validation(nan_strategy, nan_replace_value)
  146. num_variables = matrix.shape[1]
  147. tschuprows_t_matrix_value = torch.ones(num_variables, num_variables, device=matrix.device)
  148. for i, j in itertools.combinations(range(num_variables), 2):
  149. x, y = matrix[:, i], matrix[:, j]
  150. num_classes = len(torch.cat([x, y]).unique())
  151. confmat = _tschuprows_t_update(x, y, num_classes, nan_strategy, nan_replace_value)
  152. tschuprows_t_matrix_value[i, j] = tschuprows_t_matrix_value[j, i] = _tschuprows_t_compute(
  153. confmat, bias_correction
  154. )
  155. return tschuprows_t_matrix_value