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
