plot_chess_masters.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. """
  2. =============
  3. Chess Masters
  4. =============
  5. An example of the MultiDiGraph class.
  6. The function `chess_pgn_graph` reads a collection of chess matches stored in
  7. the specified PGN file (PGN ="Portable Game Notation"). Here the (compressed)
  8. default file::
  9. chess_masters_WCC.pgn.bz2
  10. contains all 685 World Chess Championship matches from 1886--1985.
  11. (data from http://chessproblem.my-free-games.com/chess/games/Download-PGN.php)
  12. The `chess_pgn_graph()` function returns a `MultiDiGraph` with multiple edges.
  13. Each node is the last name of a chess master. Each edge is directed from white
  14. to black and contains selected game info.
  15. The key statement in `chess_pgn_graph` below is::
  16. G.add_edge(white, black, game_info)
  17. where `game_info` is a `dict` describing each game.
  18. """
  19. import matplotlib.pyplot as plt
  20. import networkx as nx
  21. # tag names specifying what game info should be
  22. # stored in the dict on each digraph edge
  23. game_details = ["Event", "Date", "Result", "ECO", "Site"]
  24. def chess_pgn_graph(pgn_file="chess_masters_WCC.pgn.bz2"):
  25. """Read chess games in pgn format in pgn_file.
  26. Filenames ending in .bz2 will be uncompressed.
  27. Return the MultiDiGraph of players connected by a chess game.
  28. Edges contain game data in a dict.
  29. """
  30. import bz2
  31. G = nx.MultiDiGraph()
  32. game = {}
  33. with bz2.BZ2File(pgn_file) as datafile:
  34. lines = [line.decode().rstrip("\r\n") for line in datafile]
  35. for line in lines:
  36. if line.startswith("["):
  37. tag, value = line[1:-1].split(" ", 1)
  38. game[str(tag)] = value.strip('"')
  39. else:
  40. # empty line after tag set indicates
  41. # we finished reading game info
  42. if game:
  43. white = game.pop("White")
  44. black = game.pop("Black")
  45. G.add_edge(white, black, **game)
  46. game = {}
  47. return G
  48. G = chess_pgn_graph()
  49. print(
  50. f"Loaded {G.number_of_edges()} chess games between {G.number_of_nodes()} players\n"
  51. )
  52. # identify connected components of the undirected version
  53. H = G.to_undirected()
  54. Gcc = [H.subgraph(c) for c in nx.connected_components(H)]
  55. if len(Gcc) > 1:
  56. print(f"Note the disconnected component consisting of:\n{Gcc[1].nodes()}")
  57. # find all games with B97 opening (as described in ECO)
  58. openings = {game_info["ECO"] for (white, black, game_info) in G.edges(data=True)}
  59. print(f"\nFrom a total of {len(openings)} different openings,")
  60. print("the following games used the Sicilian opening")
  61. print('with the Najdorff 7...Qb6 "Poisoned Pawn" variation.\n')
  62. for white, black, game_info in G.edges(data=True):
  63. if game_info["ECO"] == "B97":
  64. summary = f"{white} vs {black}\n"
  65. for k, v in game_info.items():
  66. summary += f" {k}: {v}\n"
  67. summary += "\n"
  68. print(summary)
  69. # make new undirected graph H without multi-edges
  70. H = nx.Graph(G)
  71. # edge width is proportional number of games played
  72. edgewidth = [len(G.get_edge_data(u, v)) for u, v in H.edges()]
  73. # node size is proportional to number of games won
  74. wins = dict.fromkeys(G.nodes(), 0.0)
  75. for u, v, d in G.edges(data=True):
  76. r = d["Result"].split("-")
  77. if r[0] == "1":
  78. wins[u] += 1.0
  79. elif r[0] == "1/2":
  80. wins[u] += 0.5
  81. wins[v] += 0.5
  82. else:
  83. wins[v] += 1.0
  84. nodesize = [wins[v] * 50 for v in H]
  85. # Generate layout for visualization
  86. pos = nx.kamada_kawai_layout(H)
  87. # Manual tweaking to limit node label overlap in the visualization
  88. pos["Reshevsky, Samuel H"] += (0.05, -0.10)
  89. pos["Botvinnik, Mikhail M"] += (0.03, -0.06)
  90. pos["Smyslov, Vassily V"] += (0.05, -0.03)
  91. fig, ax = plt.subplots(figsize=(12, 12))
  92. # Visualize graph components
  93. nx.draw_networkx_edges(H, pos, alpha=0.3, width=edgewidth, edge_color="m")
  94. nx.draw_networkx_nodes(H, pos, node_size=nodesize, node_color="#210070", alpha=0.9)
  95. label_options = {"ec": "k", "fc": "white", "alpha": 0.7}
  96. nx.draw_networkx_labels(H, pos, font_size=14, bbox=label_options)
  97. # Title/legend
  98. font = {"fontname": "Helvetica", "color": "k", "fontweight": "bold", "fontsize": 14}
  99. ax.set_title("World Chess Championship Games: 1886 - 1985", font)
  100. # Change font color for legend
  101. font["color"] = "r"
  102. ax.text(
  103. 0.80,
  104. 0.10,
  105. "edge width = # games played",
  106. horizontalalignment="center",
  107. transform=ax.transAxes,
  108. fontdict=font,
  109. )
  110. ax.text(
  111. 0.80,
  112. 0.06,
  113. "node size = # games won",
  114. horizontalalignment="center",
  115. transform=ax.transAxes,
  116. fontdict=font,
  117. )
  118. # Resize figure for label readability
  119. ax.margins(0.1, 0.05)
  120. fig.tight_layout()
  121. plt.axis("off")
  122. plt.show()