def structure_energy_map_figure_2D_2(self, fig, n, padding=2): atom_positions = self.carbon_data.getStructure(n) structure_energy = self.carbon_data.data_energies[n] atom_energies = self.get_energies_of_structure(n) e_min = min(atom_energies) e_max = max(atom_energies) sizes = self.calc_sizes_using_z_depth(atom_positions) ax = fig.add_subplot(111) E_tot = np.sum(atom_energies) DeltaE = structure_energy - E_tot plt.title("Structure #" + str(n) + r", $E^{{NN}}$={:.2f}eV".format(E_tot) + ", ΔE={:.2f}eV".format(DeltaE)) cc = CircleCollection(sizes=sizes, offsets=atom_positions[:, 0:2], transOffset=ax.transData) cc.set_array(atom_energies) ax.add_collection(cc) ax.autoscale_view() ax.axis("off") cbar = plt.colorbar(cc) cbar.set_label('Predicted energy [eV]') x = atom_positions[:, 0] y = atom_positions[:, 1] ax.set_xlim(np.min(x) - padding, np.max(x) + padding) ax.set_ylim(np.min(y) - padding, np.max(y) + padding) return fig
def plot_rydberg_snap(snap, ax=None): """ Plots a single-channel Rydberg snapshot, optionally into the given `ax`. """ from matplotlib.collections import CircleCollection W, H = snap.shape[-1], snap.shape[-2] plt.rc('axes', linewidth=6) radius = 200 rads = radius * np.ones_like(snap.flatten()) lattice = np.stack(np.meshgrid(np.arange(W), np.arange(H)), axis=-1).reshape(-1, 2) if ax is None: fig, ax = plt.subplots(figsize=(5, 5)) circs = CircleCollection( rads, offsets=lattice, cmap='Purples', transOffset=ax.transData, edgecolors='black', linewidths=3, ) circs.set_array(snap.flatten()) ax.add_collection(circs) ax.autoscale_view() ax.set_aspect(1.0) ax.set_facecolor('xkcd:grey') ax.tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, right=False, labelleft=False) return ax