def get_plotter(self, components=("xx", ), reim="reim", broad=0.00005, emin=0, emax=None, divs=500, **kwargs): """ Return an instance of the Spectrum plotter containing the different requested components Arguments: components: A list with the components of the dielectric tensor to plot. Can be either two indexes or a string like 'xx' to plot the (0,0) component reim: If 're' (im) is present in the string plots the real (imaginary) part of the dielectric tensor emin, emax: minimum and maximum energy in which to obtain the spectra divs: number of frequency samples between emin and emax """ directions_map = {"x": 0, "y": 1, "z": 2, 0: 0, 1: 1, 2: 2} reim_label = {"re": "Re", "im": "Im"} plotter = SpectrumPlotter() for component in components: i, j = [directions_map[direction] for direction in component] for fstr in ("re", "im"): if fstr in reim: label = r"%s{$\epsilon_{%s%s}$}" % ( reim_label[fstr], "xyz"[i], "xyz"[j], ) spectrum = self.get_spectrum(component, fstr, broad=broad, emin=emin, emax=emax, divs=divs) spectrum.XLABEL = r"Frequency (meV)" spectrum.YLABEL = r"$\epsilon(\omega)$" plotter.add_spectrum(label, spectrum) return plotter
class SpectrumPlotterTest(PymatgenTest): def setUp(self): self.xanes = XANES.from_dict(spect_data_dict) def test_get_plot(self): self.plotter = SpectrumPlotter(yshift=0.2) self.plotter.add_spectrum("Pd2O", self.xanes) xanes = self.xanes.copy() xanes.y += np.random.randn(len(xanes.y)) * 0.005 self.plotter.add_spectrum("Pd2O + noise", xanes) self.plotter.add_spectrum("Pd2O - replot", xanes, "k") plt = self.plotter.get_plot() self.plotter.save_plot("spectrum_plotter_test.eps") os.remove("spectrum_plotter_test.eps") def test_get_stacked_plot(self): self.plotter = SpectrumPlotter(yshift=0.2, stack=True) self.plotter.add_spectrum("Pd2O", self.xanes, "b") xanes = self.xanes.copy() xanes.y += np.random.randn(len(xanes.y)) * 0.005 self.plotter.add_spectrum("Pd2O + noise", xanes, "r") plt = self.plotter.get_plot()