diff options
| -rw-r--r-- | schroedinger/schrodinger_plot.py | 19 |
1 files changed, 10 insertions, 9 deletions
diff --git a/schroedinger/schrodinger_plot.py b/schroedinger/schrodinger_plot.py index 48032d4..b9ad14f 100644 --- a/schroedinger/schrodinger_plot.py +++ b/schroedinger/schrodinger_plot.py @@ -30,8 +30,8 @@ args = parser.parse_args() def plot_potential(ax: plt.Axes, solution_path: str): potential_data = np.loadtxt(f'{solution_path}/potential.dat') x = potential_data[:, 0] - V = potential_data[:, 1] - ax.plot(x, V, c='black', ls='-', marker='') + v = potential_data[:, 1] + ax.plot(x, v, c='black', ls='-', marker='') def plot_wavefuncs(ax: plt.Axes, solution_path: str, wavefunction_scale: float = 1.0): @@ -51,8 +51,10 @@ def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str): expected_values = expvalues_data[:, 0] uncertainties = expvalues_data[:, 1] + uncertainty_max = uncertainties.max() - x_lim = ax1.get_xlim() + x1_lim = ax1.get_xlim() + x2_lim = (0, uncertainty_max*1.1) y_lim = ax1.get_ylim() for index in range(expvalues_data.shape[0]): energy = energies[index] @@ -60,14 +62,13 @@ def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str): uncertainty = uncertainties[index] ax1.plot(expected_value, energy, marker='x', ls='', c='green', zorder=101) - ax1.hlines(energy, *x_lim, colors='gray', alpha=0.5) + ax1.hlines(energy, *x1_lim, colors='gray', alpha=0.5) - ax2.hlines(energy, *x_lim, colors='gray', alpha=0.5) + ax2.hlines(energy, *x2_lim, colors='gray', alpha=0.5) ax2.plot(uncertainty, energy, color='orange', marker='+', ls='', markersize=10) - uncertainty_max = uncertainties.max() - ax1.set(xlim=x_lim, ylim=y_lim) - ax2.set(xlim=(0, uncertainty_max*1.1), ylim=y_lim) + ax1.set(xlim=x1_lim, ylim=y_lim) + ax2.set(xlim=x2_lim, ylim=y_lim) def plot_solution(solution_path: str = args.solution, pdf_path: str = args.pdf, show_plot: bool = args.show, @@ -94,7 +95,7 @@ def plot_solution(solution_path: str = args.solution, pdf_path: str = args.pdf, def main(): - plot_solution('test', 'test') + plot_solution('test', 'test', x_lim=(-5, 5), energy_lim=(0, 3.5), wavefunction_scale=1.0) if __name__ == '__main__': |