wavepacket.plot

Various helper utilities for plotting.

For an introduction to plotting, see Plotting wave functions and density operators.

Classes

BasePlot1D

Base class for the 1D plotters.

SimplePlot1D

Simple plot of a one-dimensional density.

StackedPlot1D

Helper class to stack multiple plots on top of each other.

BaseContourPlot2D

Base class for 2D contour plots.

ContourPlot2D

Draws a two-dimensional contour plot and the reduced densities along each DOF.

StackedContourPlot2D

Draws multiple contour plots in one figure.

Package Contents

class wavepacket.plot.BasePlot1D(state: wavepacket.grid.State, potential: wavepacket.operator.OperatorBase | None = None, hamiltonian: wavepacket.operator.OperatorBase | None = None)

Bases: abc.ABC

Base class for the 1D plotters.

The class contains common code for presenting and guessing the various parameters of the plots (x limit, y limit, conversion factors), and how to draw on a given axes. Derived classes only need to generate and maintain the figure and axes object(s).

Attributes:
xlim: tuple[float, float]

The range (min, max) of the x-axis

ylim: tuple[float, float]

The range (min, max) of the y-axis of the plot

conversion_factor: float

The factor that converts from the density to energy units. Constant 1 if no potential is plotted.

xlim
abstractmethod plot(t: float, state: wavepacket.grid.State) matplotlib.pyplot.Axes

Plots a state, possibly together with the potential.

If a potential was supplied, it is also plotted, and the density shifted by the energy given as expectation value of the Hamiltonian. If the supplied potential is time-dependent, it is plotted at the given time.

Parameters:
t: float

The time at which the state applies.

state: wp.grid.State

The state whose density is plotted.

Returns:
plt.Axes

The Matplotlib axes object on which we plotted the state for possible further manipulation.

class wavepacket.plot.SimplePlot1D(state: wavepacket.grid.State, potential: wavepacket.operator.OperatorBase | None = None, hamiltonian: wavepacket.operator.OperatorBase | None = None)

Bases: BasePlot1D

Simple plot of a one-dimensional density.

This class creates a single plot of the density of a wave function or density operator, optionally together with the potential and the state’s energy. It can be used for a quick and dirty way of showing the dynamics of a simple quantum system.

Customization of the plot is limited, see BasePlot1D for the customizable attributes. The underlying grid must be one-dimensional.

Parameters:
state: wp.State

An example state for plotting; usually the initial state. This is only used to derive some reasonable defaults for the plots.

potential: wp.operator.OperatorBase, optional

The potential that is also plotted together with the state’s density. If no potential is given, only the density is plotted.

hamiltonian: wp.operator.OperatorBase, optional

The Hamiltonian of the system, usually the time-independent part. The plotted density is shifted in y (energy) direction by the expectation value of this Hamiltonian. If no Hamiltonian is given, the potential operator stands in for the Hamiltonian.

Attributes:
figure: matplotlib.pylot.Figure

The figure that we plot on.

plot(t: float, state: wavepacket.grid.State) matplotlib.pyplot.Axes

Plots a state, possibly together with the potential.

If a potential was supplied, it is also plotted, and the density shifted by the energy given as expectation value of the Hamiltonian. If the supplied potential is time-dependent, it is plotted at the given time.

Parameters:
t: float

The time at which the state applies.

state: wp.grid.State

The state whose density is plotted.

Returns:
plt.Axes

The Matplotlib axes object on which we plotted the state for possible further manipulation.

class wavepacket.plot.StackedPlot1D(num_plots: int, state: wavepacket.grid.State, potential: wavepacket.operator.OperatorBase | None = None, hamiltonian: wavepacket.operator.OperatorBase | None = None)

Bases: BasePlot1D

Helper class to stack multiple plots on top of each other.

This class does two things: It creates a Matplotlib figure with multiple axes stacked on top of each other, and it provides a plot function that conveniently plots the density of a state on in subsequent of these axes.

Customization of the plots is possible but limited for ease of use. The underlying grid must be one-dimensional. See the base class BasePlot1D for the attributes to tweak the plotting behavior.

This plot helper is probably most useful for Jupyter notebooks, where all created figures are implicitly plotted after execution of a code block, and where plot “animations” are difficult.

Parameters:
num_plots: int

The number of plots to stack. Should equal the number of calls to th plot function. If the class runs out of axes to plot onto, it continues plotting on the last axes.

state: wp.grid.State

An example state for plotting; usually the initial state. This is only used to derive some reasonable defaults for the plots.

potential: wp.operator.OperatorBase, optional

The potential that is also plotted together with the state’s density. If no potential is given, only the density is plotted.

hamiltonian: wp.operator.OperatorBase, optional

The Hamiltonian of the system, usually the time-independent part. The plotted density is shifted in y (energy) direction by the expectation value of this Hamiltonian. If no Hamiltonian is given, the potential operator stands in for the Hamiltonian.

Attributes:
figure: matplotlib.pylot.Figure

The figure that we plot on.

plot(t: float, state: wavepacket.grid.State) matplotlib.pyplot.Axes

Plots a state, possibly together with the potential.

If a potential was supplied, it is also plotted, and the density shifted by the energy given as expectation value of the Hamiltonian. If the supplied potential is time-dependent, it is plotted at the given time.

Parameters:
t: float

The time at which the state applies.

state: wp.grid.State

The state whose density is plotted.

Returns:
plt.Axes

The Matplotlib axes object on which we plotted the state for possible further manipulation.

class wavepacket.plot.BaseContourPlot2D(state: wavepacket.grid.State, potential: wavepacket.operator.OperatorBase | None = None)

Bases: abc.ABC

Base class for 2D contour plots.

This class contains some common code and attributes for the different plots. Derived classes are mainly concerned with layout and specific extensions.

Attributes:
xlim: tuple[float,float]

The range (min, max) of the x-axis.

ylim: tuple[float, float]

The range (min, max) of the y-axis.

contours: array_like[float]

Levels at which contour lines for the density should be drawn.

potential_contours: array_like[float]

Levels at which contour lines for the potential should be drawn.

xlim
ylim
contours
abstractmethod plot(t: float, state: wavepacket.grid.State) matplotlib.pyplot.Axes

Draws the contour plot and returns the axis of the contour plot.

If a potential was supplied, its contours are also drawn. Specific plots may also populate other axes, these are not normally accessible here. If the supplied potential is time-dependent, it is plotted at the given time.

Parameters:
t: float

The time at which the state applies.

state: wp.grid.State

The state whose density is plotted.

Returns:
plt.Axes

The Matplotlib axes object on which we plotted the state for possible further manipulation.

class wavepacket.plot.ContourPlot2D(state: wavepacket.grid.State, potential: wavepacket.operator.OperatorBase | None = None)

Bases: BaseContourPlot2D

Draws a two-dimensional contour plot and the reduced densities along each DOF.

This plot draws not one, but three axis: The actual contour plot, but also two additional axes were the reduced density along the first/second degree of freedom are drawn.

Customization is done almost exclusively through the attributes of the base class BaseContourPlot2D.

Parameters:
state: wp.grid.State

This state exclusively serves for deduction of parameters (contour levels etc.). It should be either somewhat representative, or you need to adjust the attributes later.

potential: wp.operator.OperatorBase | None, default=None

The potential to draw together with the contours of the state. May be time-dependent, in which case it is drawn at the corresponding point in time.

Attributes:
figure: plt.Figure

The figure of the plot. Can be used for global changes, such as resizing, saving etc.

max_marginals: tuple[float, float]

The upper range of the marginals, i.e., reduced densities. The default is to initialize them from the supplied state.

figure
max_marginals
plot(t: float, state: wavepacket.grid.State) matplotlib.pyplot.Axes

Draws the contour plot and returns the axis of the contour plot.

If a potential was supplied, its contours are also drawn. Specific plots may also populate other axes, these are not normally accessible here. If the supplied potential is time-dependent, it is plotted at the given time.

Parameters:
t: float

The time at which the state applies.

state: wp.grid.State

The state whose density is plotted.

Returns:
plt.Axes

The Matplotlib axes object on which we plotted the state for possible further manipulation.

class wavepacket.plot.StackedContourPlot2D(num_rows: int, num_cols: int, state: wavepacket.grid.State, potential: wavepacket.operator.OperatorBase | None = None)

Bases: BaseContourPlot2D

Draws multiple contour plots in one figure.

This plot class is mainly useful for static Jupyter notebooks and similar applications. You specify a number of contour plots, and after ech plot() command, you jump to the next plot. Similar to wavepacket.plot.StackedContourPlot2D, but we lay out the contour plots in a table pattern, not a single row.

See the base class wavepacket.plot.BaseContourPlot2D for settable attributes.

Parameters:
num_rows: int

How many rows of plots are laid out in the figure.

num_cols: int

How many columns of plots are laid out in the figure.

state: wavepacket.grid.State

A reference state from which plot settings are derived (e.g., good contour lines)

potential: wavepacket.operator.OperatorBase | None, default=None

You can optionally supply a potential, then the potential is also plotted with contour lines.

Attributes:
figure: matplotlib.figure.Figure

The handle to the figure on which we draw the plots.

figure
plot(t: float, state: wavepacket.grid.State) matplotlib.pyplot.Axes

Draws the contour plot and returns the axis of the contour plot.

If a potential was supplied, its contours are also drawn. Specific plots may also populate other axes, these are not normally accessible here. If the supplied potential is time-dependent, it is plotted at the given time.

Parameters:
t: float

The time at which the state applies.

state: wp.grid.State

The state whose density is plotted.

Returns:
plt.Axes

The Matplotlib axes object on which we plotted the state for possible further manipulation.