Plot periodic table with Pymatgen and Matplotlib

Kyle
2019-03-08T12:12:24Z

periodic_table.py

| import re
|
| import matplotlib.patches as mpatches
| import matplotlib.pyplot as plt
| import numpy as np
| from matplotlib.collections import PatchCollection
| from pymatgen.core.periodic_table import Element
|
|
| class PeriodicTablePlotter(object):
| '''
| frame: str or dict, describe shapes of element cells.
| str: patch class name under matplotlib.patches.
| dict: {
| 'shape': patch class name,
| 'parms': width, height and other parms, needed for 'shape',
| 'kparms': kparms, supported by 'shape'
| }
| colors: dict, custom frame and background colors for elements.
| {
| 'frame': str, list
| 'background': str, list, {'depend_on': data_key [, 'cmp': color map, 'cbar': True]}
| }
| if you use a gradient, color will based on the data from Element().data or you provide.
| depend_on specify which data you want to use.
| label: dict or function(not supported yet).
| dict: {
| 'label1': [element symbols or atomic indexes],
| ...
| }
| or more detail:
| dict: {
| 'label1': {
| 'elements': [element symbols or atomic indexes],
| 'color': not necessary
| }
| ...
| }
| function: func(element symbol or atomic index) returns color and label
| text: text content and format within element cells (not supported yet).
| data: data beside pymatgen or you want overide for pymatgen
| {symbol: {k1: v1, k2:v2}, ...}
| '''
|
| def __init__(self, frame=None, colors=None, labels=None, text=None, data=None, font_family=None):
| self.frame = frame or {}
| if isinstance(self.frame, str):
| self.frame = {'shape': self.frame or None}
| # self.frame = _frame
| self.frame.setdefault('shape', 'Rectangle')
| self.frame.setdefault('_parms', self.frame.get('parms', [1, 1]))
| self.frame.setdefault('kparms', {})
| # _colors = {'frame': 'black', 'background': 'none'}
| # _colors.update(colors)
| # self.colors = _colors
| self.colors = colors or {}
| self.colors.setdefault('frame', 'black')
| self.colors.setdefault('background', 'none')
| if (not labels) or 'background' in labels or 'frame' in labels:
| self.labels = labels or {}
| else:
| self.labels = {'background': labels}
| self._text = text or None
| self._data = {
| Element.from_Z(i).symbol: {
| 'Symbol': Element.from_Z(i).symbol,
| **Element.from_Z(i).data
| }
| for i in range(1, 104)
| }
| if data:
| [self._data[k].update(data[k]) for k in data]
| self.font_family = font_family or None
|
| def get_frame(self):
| self.frame['parms'] = list(self.frame['_parms'])
| if self.frame['shape'] == 'FancyBboxPatch':
| self.frame['kparms'].setdefault('boxstyle', 'round, pad=0.1')
| pad = float([*re.findall('pad=(\d+[.]\d+)', self.frame['kparms']['boxstyle']), 0.3][0])
| self.frame['parms'][0:2] = list(np.array(self.frame['_parms'][0:2]) - 2 * pad)
| try:
| return getattr(mpatches, self.frame['shape'])
| except AttributeError:
| raise ValueError('Unsupported shape error')
|
| def get_xys(self):
| width_and_height = np.array(self.frame['_parms'][0:2])
| text_shift = self.frame.get('text_shift', np.array([0.5, 0.5]) * width_and_height)
| if self.frame['shape'] == 'FancyBboxPatch':
| text_shift -= 0.10 * width_and_height
|
| elements = [Element.from_Z(i) for i in range(1, 104)]
| group_and_rows = [np.array([el.group, el.row]) for el in elements]
| frame_xys = [np.array([1, -1]) * gr * width_and_height for gr in group_and_rows]
| text_xys = [text_shift + fxy for fxy in frame_xys]
| return frame_xys, text_xys
|
| def get_colors_and_legend(self):
| for v in self.colors.values():
| if isinstance(v, list) and len(v) != len(self._data):
| raise ValueError('Length of color list should be compatible with Elements.')
| default_color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
| default_edgecolor = ['black', self.colors['frame']][isinstance(self.colors['frame'], str)]
| default_facecolor = ['none', self.colors['background']][isinstance(self.colors['background'], str)]
|
| colors = {
| 'edgecolors':
| self.colors.get('frame', default_edgecolor),
| 'facecolors': [self.colors.get('background', default_facecolor), None][isinstance(
| self.colors['background'], dict)]
| }
| legend_handles = []
|
| for k, v in colors.items():
| labels = self.labels.get({'edgecolors': 'frame', 'facecolors': 'background'}[k], None)
| if labels:
| if not v:
| raise ValueError('Gradient mapped background can\'t be labelled.')
| if isinstance(colors[k], str):
| colors[k] = [colors[k]] * len(self._data)
| if isinstance(list(labels.values())[0], list):
| labels = {lk: {'elements': lv} for lk, lv in labels.items()}
| for i, l in enumerate(labels):
| label = labels[l]
| label.setdefault('label', l)
| indexes = [Element(s).Z - 1 for s in label['elements']]
| label.setdefault('color', default_color_cycle[i])
| for j in indexes:
| colors[k][j] = label['color']
| label.setdefault('frame_kparms', {})
| for lkk in ['edgecolor', 'facecolor']:
| label['frame_kparms'].setdefault(lkk, 'none')
| label['frame_kparms'][k[0:-1]] = label['color']
|
| patch_legend = mpatches.Patch(label=label['label'], **label['frame_kparms'])
| legend_handles.append(patch_legend)
|
| return colors, legend_handles
|
| def get_plot(self, title=None, figsize=None, legend_kparms={}):
| fig, ax = plt.subplots(figsize=figsize, frameon=False)
|
| frame_xys, text_xys = self.get_xys()
| colors, legend_handles = self.get_colors_and_legend()
|
| frame = self.get_frame()
| frames = [frame(xy, *self.frame['parms'], **self.frame['kparms']) for xy in frame_xys]
| collection = PatchCollection(frames, **colors)
| ax.add_collection(collection)
|
| text_kparms = {
| 'horizontalalignment': 'center',
| 'verticalalignment': 'center',
| 'family': self.font_family,
| 'size': 'x-large'
| }
| [
| ax.text(
| *text_xys[i],
| Element.from_Z(i + 1).symbol,
| **text_kparms) for i in range(len(self._data))
| ]
| ax.axis('equal')
| ax.axis('off')
| ax.get_xaxis().set_visible(False)
| ax.get_yaxis().set_visible(False)
|
| if legend_handles:
| legend_kparms.update({'handles': legend_handles})
| legend_kparms.setdefault('loc', 'upper center')
| legend_kparms.setdefault('bbox_to_anchor', (0.4, 1.0))
| legend_kparms.setdefault('prop', {'family': self.font_family, 'size': 'xx-large'})
| ax.legend(**legend_kparms)
| if title:
| plt.title(title, fontsize='xx-large')
| return plt
|
| def show(self, title=None, figsize=None):
| plt = self.get_plot(title=title, figsize=figsize)
| plt.show()