class TestBoltzTraP2Interpolater(unittest.TestCase): """Tests for interpolating a band structure using BoltzTraP2.""" def setUp(self): vr = Vasprun(os.path.join(gaas_files, "vasprun.xml.gz"), parse_projected_eigen=True) bs = vr.get_band_structure() num_electrons = vr.parameters["NELECT"] self.kpoints = np.array(vr.actual_kpoints) self.interpolater = Interpolater(bs, num_electrons, interpolate_projections=True, interpolation_factor=1) def test_initialisation(self): """Test coefficients and parameters are calculated correctly.""" self.interpolater.initialize() params = self.interpolater._parameters assert_array_equal(params[0][0], [[0, 0, 0]]) self.assertAlmostEqual(params[1][0][0], 9.14055614) self.assertAlmostEqual(params[2][0][0].real, -2.33144546) def test_get_energies(self): """Test getting the interpolated energy, velocity and effective mass.""" # test just getting energy energies = self.interpolater.get_energies(self.kpoints, 25, return_velocity=False, curvature=False) self.assertEqual(energies.shape, (138, )) self.assertAlmostEqual(energies[0], 3.852399483908641) # test energy + velocity energies, velocities = self.interpolater.get_energies( self.kpoints, 25, return_velocity=True, curvature=False) self.assertEqual(energies.shape, (138, )) self.assertAlmostEqual(energies[0], 3.852399483908641) self.assertEqual(velocities.shape, (138, 3)) self.assertAlmostEqual(velocities[10][0], 5.401166156893334e6, places=0) # test energy + effective_mass energies, effective_masses = self.interpolater.get_energies( self.kpoints, 25, return_velocity=False, curvature=True) self.assertEqual(energies.shape, (138, )) self.assertAlmostEqual(energies[0], 3.852399483908641) self.assertEqual(effective_masses.shape, (138, 3, 3)) self.assertAlmostEqual(effective_masses[10][0][0], 0.2242569613626494) # test energy + velocity + effective_mass energies, velocities, effective_masses = self.interpolater.get_energies( self.kpoints, 25, return_velocity=True, curvature=True) self.assertEqual(energies.shape, (138, )) self.assertAlmostEqual(energies[0], 3.852399483908641) self.assertEqual(velocities.shape, (138, 3)) self.assertAlmostEqual(velocities[10][0], 5.401166156893334e6, places=0) self.assertEqual(effective_masses.shape, (138, 3, 3)) self.assertAlmostEqual(effective_masses[10][0][0], 0.2242569613626494) def test_get_energies_scissor(self): """Test scissoring of band energies.""" # test valence band energies = self.interpolater.get_energies(self.kpoints, 25, return_velocity=False, curvature=False, scissor=1.0) self.assertEqual(energies.shape, (138, )) self.assertAlmostEqual(energies[0], 3.852399483908641 - 0.5) # test conduction band energies = self.interpolater.get_energies(self.kpoints, 33, return_velocity=False, curvature=False, scissor=1.0) self.assertEqual(energies.shape, (138, )) self.assertAlmostEqual(energies[0], 7.301700765 + 0.5) def test_get_energies_multiple_bands(self): """Test getting the interpolated data for multiple bands.""" # test just getting energy energies = self.interpolater.get_energies(self.kpoints, [25, 35], return_velocity=False, curvature=False) self.assertEqual(energies.shape, (2, 138)) self.assertAlmostEqual(energies[0][0], 3.852399483908641) self.assertAlmostEqual(energies[1][0], 9.594401616456384) # test energy + velocity energies, velocities = self.interpolater.get_energies( self.kpoints, [25, 35], return_velocity=True, curvature=False) self.assertEqual(energies.shape, (2, 138)) self.assertAlmostEqual(energies[0][0], 3.852399483908641) self.assertAlmostEqual(energies[1][0], 9.594401616456384) self.assertEqual(velocities.shape, (2, 138, 3)) self.assertAlmostEqual(velocities[0][10][0], 5.401166156893334e6, places=0) self.assertAlmostEqual(velocities[1][10][0], 1.1686720831758736e8, places=0) # test energy + effective_mass energies, effective_masses = self.interpolater.get_energies( self.kpoints, [25, 35], return_velocity=False, curvature=True) self.assertEqual(energies.shape, (2, 138)) self.assertAlmostEqual(energies[0][0], 3.852399483908641) self.assertAlmostEqual(energies[1][0], 9.594401616456384) self.assertEqual(effective_masses.shape, (2, 138, 3, 3)) self.assertAlmostEqual(effective_masses[0][10][0][0], 0.224256961362649) self.assertAlmostEqual(effective_masses[1][10][0][0], 0.009057103344700) # test energy + velocity + effective_mass energies, velocities, effective_masses = self.interpolater.get_energies( self.kpoints, [25, 35], return_velocity=True, curvature=True) self.assertEqual(energies.shape, (2, 138)) self.assertAlmostEqual(energies[0][0], 3.852399483908641) self.assertAlmostEqual(energies[1][0], 9.594401616456384) self.assertEqual(velocities.shape, (2, 138, 3)) self.assertAlmostEqual(velocities[0][10][0], 5.401166156893334e6, places=0) self.assertAlmostEqual(velocities[1][10][0], 1.1686720831758736e8, places=0) self.assertEqual(effective_masses.shape, (2, 138, 3, 3)) self.assertAlmostEqual(effective_masses[0][10][0][0], 0.224256961362649) self.assertAlmostEqual(effective_masses[1][10][0][0], 0.009057103344700) def test_get_energies_all_bands(self): # test all bands energies, velocities, effective_masses = self.interpolater.get_energies( self.kpoints, None, return_velocity=True, curvature=True) self.assertEqual(energies.shape, (96, 138)) self.assertAlmostEqual(energies[25][0], 3.852399483908641) self.assertAlmostEqual(energies[35][0], 9.594401616456384) self.assertEqual(velocities.shape, (96, 138, 3)) self.assertAlmostEqual(velocities[25][10][0], 5.401166156893334e6, places=0) self.assertAlmostEqual(velocities[35][10][0], 1.1686720831758736e8, places=0) self.assertEqual(effective_masses.shape, (96, 138, 3, 3)) self.assertAlmostEqual(effective_masses[25][10][0][0], 0.22425696136264) self.assertAlmostEqual(effective_masses[35][10][0][0], 0.00905710334470) def test_get_dos(self): """Test generating the interpolated DOS.""" dos = self.interpolater.get_dos([10, 10, 10], emin=-10, emax=10, width=0.075) self.assertEqual(dos.shape, (20000, 2)) self.assertEqual(dos[0][0], -10) self.assertAlmostEqual(dos[15000][1], 3.5362612128412807) def test_get_energies_symprec(self): # vr = Vasprun(os.path.join(tin_dioxide_files, 'vasprun.xml.gz'), # parse_projected_eigen=True) vr = Vasprun(os.path.join(pbs_files, "vasprun.xml.gz"), parse_projected_eigen=True) bs = vr.get_band_structure() num_electrons = vr.parameters["NELECT"] interpolater = Interpolater(bs, num_electrons, interpolate_projections=True, interpolation_factor=1) ir_kpoints, weights, kpoints, ir_kpoints_idx, ir_to_full_idx = get_kpoints( [13, 15, 29], vr.final_structure, boltztrap_ordering=True, return_full_kpoints=True, ) initialize_amset_logger() ( energies, velocities, curvature, projections, sym_info, ) = interpolater.get_energies( kpoints, None, return_velocity=True, atomic_units=True, return_curvature=True, return_projections=True, symprec=0.1, return_vel_outer_prod=True, return_kpoint_mapping=True, ) ( energies_no_sym, velocities_no_sym, curvature_no_sym, projections_no_sym, ) = interpolater.get_energies( kpoints, None, return_velocity=True, atomic_units=True, return_curvature=True, return_projections=True, return_vel_outer_prod=True, symprec=None, ) np.testing.assert_array_equal(ir_to_full_idx, sym_info["ir_to_full_idx"]) np.testing.assert_array_equal(ir_kpoints_idx, sym_info["ir_kpoints_idx"]) # print(velocities[Spin.up][5, :, :, -3:]) # print(velocities_no_sym[Spin.up][5, :, :, -3:]) # print(sym_info["ir_to_full_idx"][-10:]) np.testing.assert_array_almost_equal(energies[Spin.up], energies_no_sym[Spin.up], decimal=12) np.testing.assert_array_almost_equal(velocities[Spin.up], velocities_no_sym[Spin.up], decimal=12) np.testing.assert_array_almost_equal(curvature[Spin.up], curvature_no_sym[Spin.up], decimal=12) for l in projections[Spin.up]: np.testing.assert_array_almost_equal( projections[Spin.up][l], projections_no_sym[Spin.up][l]) def test_get_energies_interpolater(self): initialize_amset_logger() amset_data = self.interpolater.get_amset_data() energies, velocities, projections, sym_info = self.interpolater.get_energies( amset_data.kpoints, None, return_velocity=True, atomic_units=True, curvature=False, return_projections=True, symprec=0.1, return_vel_outer_prod=True, return_kpoint_mapping=True, ) np.testing.assert_array_almost_equal(energies[Spin.up], amset_data.energies[Spin.up]) np.testing.assert_array_almost_equal( velocities[Spin.up], amset_data.velocities_product[Spin.up]) for l in projections[Spin.up]: np.testing.assert_array_almost_equal( projections[Spin.up][l], amset_data._projections[Spin.up][l]) np.testing.assert_array_equal(sym_info["ir_kpoints_idx"], amset_data.ir_kpoints_idx) np.testing.assert_array_equal(sym_info["ir_to_full_idx"], amset_data.ir_to_full_kpoint_mapping)
class ElectronicStructurePlotter(object): def __init__( self, bandstructure: BandStructure, nelect: int, soc: bool = False, interpolation_factor=defaults["interpolation_factor"], print_log=defaults["print_log"], symprec=defaults["symprec"], energy_cutoff=None, ): if print_log: initialize_amset_logger( filename="amset_electronic_structure_plot.log") self.symprec = symprec self.interpolater = Interpolater( bandstructure, nelect, interpolation_factor=interpolation_factor, soc=soc) self.structure = bandstructure.structure self.energy_cutoff = energy_cutoff @classmethod def from_vasprun(cls, vasprun, **kwargs): from pymatgen.io.vasp import Vasprun if isinstance(vasprun, Path): vasprun = vasprun.as_posix() if isinstance(vasprun, str): vasprun_gz = vasprun + ".gz" if Path(vasprun).exists(): vasprun = Vasprun(vasprun) elif Path(vasprun_gz).exists(): vasprun = Vasprun(vasprun_gz) bandstructure = vasprun.get_band_structure() nelect = vasprun.parameters["NELECT"] soc = vasprun.parameters["LSORBIT"] return cls(bandstructure, nelect, soc=soc, **kwargs) def get_plot( self, plot_dos=True, plot_band_structure=True, kpath=None, line_density=100.0, dos_kpoints=50.0, dos_estep=defaults["dos_estep"], zero_to_efermi=True, vbm_cbm_marker=False, height=6, width=None, emin=None, emax=None, elabel="Energy (eV)", plt=None, dos_aspect=3, aspect=None, style=None, no_base_style=False, fonts=None, ): if width is None and plot_band_structure: width = 6 if self.energy_cutoff is None: self.energy_cutoff = max([abs(emin or 6), abs(emax or 6)]) if plot_band_structure: logger.info("Generating band structure...") bs_plotter = self.get_bs_plotter(line_density=line_density, kpath=kpath) if plot_dos: logger.info("Generating density of states...") dos_plotter = self.get_dos_plotter(dos_density=dos_kpoints, dos_estep=dos_estep) logger.info("Plotting...") common_kwargs = dict( zero_to_efermi=zero_to_efermi, width=width, height=height, plt=plt, style=style, no_base_style=no_base_style, fonts=fonts, ) if plot_band_structure: bs_kwargs = { "ymin": emin, "ymax": emax, "ylabel": elabel, "aspect": aspect, "vbm_cbm_marker": vbm_cbm_marker, } common_kwargs.update(bs_kwargs) if plot_dos: bs_dos_kwargs = { # "dos_label": "DOS", "dos_aspect": dos_aspect, "dos_plotter": dos_plotter, "plot_dos_legend": False, } common_kwargs.update(bs_dos_kwargs) return bs_plotter.get_plot(**common_kwargs) else: dos_kwargs = { "xmin": emin, "xmax": emax, "legend_on": False, "xlabel": elabel, "ylabel": "DOS", } common_kwargs.update(dos_kwargs) return dos_plotter.get_plot(common_kwargs) def get_bs_plotter(self, line_density=100, kpath=None): lm_bs = self.interpolater.get_line_mode_band_structure( line_density=line_density, kpath=kpath, symprec=self.symprec, energy_cutoff=self.energy_cutoff, ) return SBSPlotter(lm_bs) def get_dos_plotter(self, dos_density=50.0, dos_estep=defaults["dos_estep"]): dos = self.interpolater.get_dos( kpoint_mesh=dos_density, symprec=self.symprec, estep=dos_estep, energy_cutoff=self.energy_cutoff, ) return SDOSPlotter(dos, {})