From d61d6b1fa4e023d3fda6af1432c4665f9fa05731 Mon Sep 17 00:00:00 2001 From: Julian Teske Date: Wed, 18 May 2022 15:56:40 +0200 Subject: [PATCH] return the figure and axis in the energy spectrum plotting --- qopt/energy_spectrum.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/qopt/energy_spectrum.py b/qopt/energy_spectrum.py index 9494b55..7da17fe 100644 --- a/qopt/energy_spectrum.py +++ b/qopt/energy_spectrum.py @@ -76,7 +76,8 @@ def vector_color_map(vectors: np.array): def plot_energy_spectrum(hamiltonian: List[OperatorMatrix], x_val: np.array, - x_label: str): + x_label: str, + ax=None): """ Calculates and plots the energy spectra of hamilton operators. @@ -94,6 +95,9 @@ def plot_energy_spectrum(hamiltonian: List[OperatorMatrix], x_label: str Label of the x-axis. + ax: matplotlib pyplot axes + Instance of axes to plot the data in. Defaults to None. + """ d = hamiltonian[0].shape[0] eigenvalues = np.empty((len(hamiltonian), d)) @@ -103,9 +107,10 @@ def plot_energy_spectrum(hamiltonian: List[OperatorMatrix], eigenvalues[i, :] = eig_val eigenvectors[i, :, :] = np.abs(eig_vec) - fig, ax = plt.subplots() + if ax is None: + _, ax = plt.subplots() for i in range(d): ax.scatter(x=x_val, y=eigenvalues[:, i], c=vector_color_map(eigenvectors[:, :, i])) ax.set_xlabel(x_label) - return fig, ax + return ax