Skip to content

qiu_classical_simulation.wave_optics.visualization

Plots of transverse fields.

Functions:

  • plot_wavefunction –

    Plot the magnitude and the phase of a field over its samples, side by side.

plot_wavefunction

plot_wavefunction(psi: ArrayLike, plot_size_scale: float = 1, normalize: bool = True, ylim: float | None = None) -> tuple[Figure, tuple[Axes, Axes]]

Plot the magnitude and the phase of a field over its samples, side by side.

Parameters:

  • psi (ArrayLike) –

    The amplitudes of the field.

  • plot_size_scale (float, default: 1 ) –

    The scale of the figure size.

  • normalize (bool, default: True ) –

    If True, plot the magnitude normalized to unit norm.

  • ylim (float | None, default: None ) –

    The upper limit of the magnitude axis, if given.

Returns:

  • tuple[Figure, tuple[Axes, Axes]] –

    The figure and its axes of the magnitude and the phase.

Source code in packages/qiu-classical-simulation/src/qiu_classical_simulation/wave_optics/visualization.py
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def plot_wavefunction(
    psi: npt.ArrayLike,
    plot_size_scale: float = 1,
    normalize: bool = True,
    ylim: float | None = None,
) -> tuple[Figure, tuple[Axes, Axes]]:
    """Plot the magnitude and the phase of a field over its samples, side by side.

    Args:
        psi: The amplitudes of the field.
        plot_size_scale: The scale of the figure size.
        normalize: If True, plot the magnitude normalized to unit norm.
        ylim: The upper limit of the magnitude axis, if given.

    Returns:
        The figure and its axes of the magnitude and the phase.
    """
    amplitudes = np.asarray(psi)
    magnitude = np.abs(amplitudes)
    if normalize:
        magnitude = magnitude / np.linalg.norm(magnitude)

    fig, (ax_magnitude, ax_phase) = plt.subplots(
        1, 2, figsize=(2 * 6.4 * plot_size_scale, 4.8 * plot_size_scale)
    )
    ax_magnitude.plot(magnitude)
    ax_phase.plot(np.angle(amplitudes))

    if ylim is not None:
        ax_magnitude.set_ylim(0, ylim)
    ax_phase.set_ylim(-np.pi - 0.3, np.pi + 0.3)
    ax_phase.set_yticks(_PHASE_TICKS, _PHASE_LABELS)

    ax_magnitude.set_xlabel(r"$|x\rangle$")
    ax_phase.set_xlabel(r"$|x\rangle$")
    ax_magnitude.set_ylabel(r"$|\psi(x)|$")
    ax_phase.set_ylabel(r"$\angle\psi(x)$")
    return fig, (ax_magnitude, ax_phase)