confusionmatrix.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219
  1. # Natural Language Toolkit: Confusion Matrices
  2. #
  3. # Copyright (C) 2001-2019 NLTK Project
  4. # Author: Edward Loper <edloper@gmail.com>
  5. # Steven Bird <stevenbird1@gmail.com>
  6. # URL: <http://nltk.org/>
  7. # For license information, see LICENSE.TXT
  8. from __future__ import print_function, unicode_literals
  9. from nltk.probability import FreqDist
  10. from nltk.compat import python_2_unicode_compatible
  11. @python_2_unicode_compatible
  12. class ConfusionMatrix(object):
  13. """
  14. The confusion matrix between a list of reference values and a
  15. corresponding list of test values. Entry *[r,t]* of this
  16. matrix is a count of the number of times that the reference value
  17. *r* corresponds to the test value *t*. E.g.:
  18. >>> from nltk.metrics import ConfusionMatrix
  19. >>> ref = 'DET NN VB DET JJ NN NN IN DET NN'.split()
  20. >>> test = 'DET VB VB DET NN NN NN IN DET NN'.split()
  21. >>> cm = ConfusionMatrix(ref, test)
  22. >>> print(cm['NN', 'NN'])
  23. 3
  24. Note that the diagonal entries *Ri=Tj* of this matrix
  25. corresponds to correct values; and the off-diagonal entries
  26. correspond to incorrect values.
  27. """
  28. def __init__(self, reference, test, sort_by_count=False):
  29. """
  30. Construct a new confusion matrix from a list of reference
  31. values and a corresponding list of test values.
  32. :type reference: list
  33. :param reference: An ordered list of reference values.
  34. :type test: list
  35. :param test: A list of values to compare against the
  36. corresponding reference values.
  37. :raise ValueError: If ``reference`` and ``length`` do not have
  38. the same length.
  39. """
  40. if len(reference) != len(test):
  41. raise ValueError('Lists must have the same length.')
  42. # Get a list of all values.
  43. if sort_by_count:
  44. ref_fdist = FreqDist(reference)
  45. test_fdist = FreqDist(test)
  46. def key(v):
  47. return -(ref_fdist[v] + test_fdist[v])
  48. values = sorted(set(reference + test), key=key)
  49. else:
  50. values = sorted(set(reference + test))
  51. # Construct a value->index dictionary
  52. indices = dict((val, i) for (i, val) in enumerate(values))
  53. # Make a confusion matrix table.
  54. confusion = [[0 for val in values] for val in values]
  55. max_conf = 0 # Maximum confusion
  56. for w, g in zip(reference, test):
  57. confusion[indices[w]][indices[g]] += 1
  58. max_conf = max(max_conf, confusion[indices[w]][indices[g]])
  59. #: A list of all values in ``reference`` or ``test``.
  60. self._values = values
  61. #: A dictionary mapping values in ``self._values`` to their indices.
  62. self._indices = indices
  63. #: The confusion matrix itself (as a list of lists of counts).
  64. self._confusion = confusion
  65. #: The greatest count in ``self._confusion`` (used for printing).
  66. self._max_conf = max_conf
  67. #: The total number of values in the confusion matrix.
  68. self._total = len(reference)
  69. #: The number of correct (on-diagonal) values in the matrix.
  70. self._correct = sum(confusion[i][i] for i in range(len(values)))
  71. def __getitem__(self, li_lj_tuple):
  72. """
  73. :return: The number of times that value ``li`` was expected and
  74. value ``lj`` was given.
  75. :rtype: int
  76. """
  77. (li, lj) = li_lj_tuple
  78. i = self._indices[li]
  79. j = self._indices[lj]
  80. return self._confusion[i][j]
  81. def __repr__(self):
  82. return '<ConfusionMatrix: %s/%s correct>' % (self._correct, self._total)
  83. def __str__(self):
  84. return self.pretty_format()
  85. def pretty_format(
  86. self,
  87. show_percents=False,
  88. values_in_chart=True,
  89. truncate=None,
  90. sort_by_count=False,
  91. ):
  92. """
  93. :return: A multi-line string representation of this confusion matrix.
  94. :type truncate: int
  95. :param truncate: If specified, then only show the specified
  96. number of values. Any sorting (e.g., sort_by_count)
  97. will be performed before truncation.
  98. :param sort_by_count: If true, then sort by the count of each
  99. label in the reference data. I.e., labels that occur more
  100. frequently in the reference label will be towards the left
  101. edge of the matrix, and labels that occur less frequently
  102. will be towards the right edge.
  103. @todo: add marginals?
  104. """
  105. confusion = self._confusion
  106. values = self._values
  107. if sort_by_count:
  108. values = sorted(
  109. values, key=lambda v: -sum(self._confusion[self._indices[v]])
  110. )
  111. if truncate:
  112. values = values[:truncate]
  113. if values_in_chart:
  114. value_strings = ["%s" % val for val in values]
  115. else:
  116. value_strings = [str(n + 1) for n in range(len(values))]
  117. # Construct a format string for row values
  118. valuelen = max(len(val) for val in value_strings)
  119. value_format = '%' + repr(valuelen) + 's | '
  120. # Construct a format string for matrix entries
  121. if show_percents:
  122. entrylen = 6
  123. entry_format = '%5.1f%%'
  124. zerostr = ' .'
  125. else:
  126. entrylen = len(repr(self._max_conf))
  127. entry_format = '%' + repr(entrylen) + 'd'
  128. zerostr = ' ' * (entrylen - 1) + '.'
  129. # Write the column values.
  130. s = ''
  131. for i in range(valuelen):
  132. s += (' ' * valuelen) + ' |'
  133. for val in value_strings:
  134. if i >= valuelen - len(val):
  135. s += val[i - valuelen + len(val)].rjust(entrylen + 1)
  136. else:
  137. s += ' ' * (entrylen + 1)
  138. s += ' |\n'
  139. # Write a dividing line
  140. s += '%s-+-%s+\n' % ('-' * valuelen, '-' * ((entrylen + 1) * len(values)))
  141. # Write the entries.
  142. for val, li in zip(value_strings, values):
  143. i = self._indices[li]
  144. s += value_format % val
  145. for lj in values:
  146. j = self._indices[lj]
  147. if confusion[i][j] == 0:
  148. s += zerostr
  149. elif show_percents:
  150. s += entry_format % (100.0 * confusion[i][j] / self._total)
  151. else:
  152. s += entry_format % confusion[i][j]
  153. if i == j:
  154. prevspace = s.rfind(' ')
  155. s = s[:prevspace] + '<' + s[prevspace + 1 :] + '>'
  156. else:
  157. s += ' '
  158. s += '|\n'
  159. # Write a dividing line
  160. s += '%s-+-%s+\n' % ('-' * valuelen, '-' * ((entrylen + 1) * len(values)))
  161. # Write a key
  162. s += '(row = reference; col = test)\n'
  163. if not values_in_chart:
  164. s += 'Value key:\n'
  165. for i, value in enumerate(values):
  166. s += '%6d: %s\n' % (i + 1, value)
  167. return s
  168. def key(self):
  169. values = self._values
  170. str = 'Value key:\n'
  171. indexlen = len(repr(len(values) - 1))
  172. key_format = ' %' + repr(indexlen) + 'd: %s\n'
  173. for i in range(len(values)):
  174. str += key_format % (i, values[i])
  175. return str
  176. def demo():
  177. reference = 'DET NN VB DET JJ NN NN IN DET NN'.split()
  178. test = 'DET VB VB DET NN NN NN IN DET NN'.split()
  179. print('Reference =', reference)
  180. print('Test =', test)
  181. print('Confusion matrix:')
  182. print(ConfusionMatrix(reference, test))
  183. print(ConfusionMatrix(reference, test).pretty_format(sort_by_count=True))
  184. if __name__ == '__main__':
  185. demo()