Is it a bug or just something I need to do when use plot_slab function?

When I tried to plot the (111) slab of \rm UO_2, I set the oxidation states to check whether a surface has a net dipole moment, however, I need to remove the oxidation states of the slabs when trying to plot them by plot_slab function otherwise it will return a KeyError.

Things work well with the following codes:

from pymatgen.core import Structure
UO2 = Structure.from_file('UO2_mp-1597_primitive.cif')

from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
UO2 = SpacegroupAnalyzer(UO2).get_conventional_standard_structure()

UO2.add_oxidation_state_by_element({'U': 4, 'O': -2})

from pymatgen.core.surface import generate_all_slabs

all_slabs = generate_all_slabs(UO2, 2, 10, 10)
slabs_111 = [slab for slab in all_slabs if slab.miller_index==(1, 1, 1)]

from pymatgen.analysis.adsorption import plot_slab
from matplotlib import pyplot as plt
%matplotlib inline

fig = plt.figure()
for n, slab in enumerate(slabs_111):
    ax = fig.add_subplot(1, 2, n+1)
    ax.set_title(n+1)
    ax.set_xticks([])
    ax.set_yticks([])
    slab.remove_oxidation_states()
    plot_slab(slab, ax, adsorption_sites=False)

plt.show()

But if I remove the slab.remove_oxidation_states(), the python interpreter will return a KeyError:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-88-f31828b8ff1d> in <module>
     23     ax.set_yticks([])
     24     #slab.remove_oxidation_states()
---> 25     plot_slab(slab, ax, adsorption_sites=False)
     26 
     27 plt.show()

~/anaconda3/envs/ASE/lib/python3.9/site-packages/pymatgen/analysis/adsorption.py in plot_slab(slab, ax, scale, repeat, window, draw_unit_cell, decay, adsorption_sites, inverse)
    679         ax.add_patch(patches.Circle(coord[:2] - lattsum * (repeat // 2),
    680                                     r, color='w', zorder=2 * n))
--> 681         color = color_dict[sites[n].species_string]
    682         ax.add_patch(patches.Circle(coord[:2] - lattsum * (repeat // 2), r,
    683                                     facecolor=color, alpha=alphas[n],

KeyError: 'O2-'

So is this a small bug or something I need to do if I want to plot the slabs which have oxidation states set?

Hi Youmu,

This was also mentioned here. I reported the issue on github and I believe the fix is pretty straightforward.

Best,
Peter