Bug fix: SpectrumPlotter.add_spectra

[minhsueh]

Original SpectrumPlotter in pymatgen.vis.plotters

def add_spectra(self, spectra_dict, key_sort_func=None):
    """
    Add a dictionary of doses, with an optional sorting function for the
    keys.

    Args:
        dos_dict: dict of {label: Dos}
        key_sort_func: function used to sort the dos_dict keys.
    """
    keys = sorted(spectra_dict, key=key_sort_func) if key_sort_func else list(spectra_dict)
    for label in keys:
        self.add_spectra(label, spectra_dict[label])

Issue 1:
The original docstring suggests using the Dos object, but attributes x and y were already overridden to energies and densities in Dos, which caused the error:

AttributeError: 'Dos' object has no attribute 'x'

Although the DosPlotter is available to plot DOS, I think this function shouldn't be deprecated because it can still parse Spectrum.

Suggestion 1:
Update the docstring to reflect Spectrum support.

Issue 2:
When I passed in Spectrum, it resulted in the error:

in add_spectra
    keys = sorted(spectra_dict, key=key_sort_func) if key_sort_func else list(spectra_dict)
TypeError: 'XAS' object is not callable

Reason: Recursive call within the function.

Suggestion 2:
Modify:

self.add_spectra(label, spectra_dict[label])

to:

self.add_spectrum(label, spectra_dict[label])

Note:

1. I used similar codes with pymatgen.tests.vis.test_plotters.py to test.

from __future__ import annotations

import json
import os

import matplotlib.pyplot as plt
import numpy as np
from monty.json import MontyDecoder

from pymatgen.analysis.xas.spectrum import XAS
from pymatgen.util.testing import TEST_FILES_DIR, PymatgenTest
from pymatgen.vis.plotters import SpectrumPlotter

from pymatgen.electronic_structure.dos import DOS

test_dir = f"{TEST_FILES_DIR}/spectrum_test"

with open(f"{test_dir}/LiCoO2_k_xanes.json") as fp:
    spect_data_dict = json.load(fp, cls=MontyDecoder)

xanes = XAS.from_dict(spect_data_dict)
plotter = SpectrumPlotter(yshift=0.2)

xanes_dict = dict()
xanes_dict["LiCoO2"] = xanes
xanes = xanes.copy()
xanes.y += np.random.randn(len(xanes.y)) * 0.005
xanes_dict["LiCoO2 + noise"] = xanes
xanes_dict["LiCoO2 - replot"] = xanes

plotter.add_spectra(xanes_dict)
ax = plotter.get_plot()
plotter.save_plot('test3.png', img_format='png')

2. This is my first time submitting a pull request at Pymatgen. Please let me know if anything is incorrectly formatted or needs improvement. Thank you :)

[janosh]

Thanks @minhsueh! Can you add a test for this method to prevent future regressions?

[minhsueh]

Hi @janosh, I have added a test in pymatgen.tests.vis.test_plotters

Also, I added a line to check the given object has attributes x and y in SpectrumPlotter.add_spectrum. This can directly point out the problem that the given object has no x and y.

assert hasattr(spectrum, "x") and hasattr(
            spectrum, "y"
        ), f"{type(spectrum)} has no attribute x and y, please check!"

Please let me know if anything is needed.

[janosh]

Thanks @minhsueh! 👍

PR labeled breaking because I dropped superfluous img_format="eps" kwarg.
Specify image format via filename extension instead.