Plotting

One-line plots of wavefunctions, potentials and eigenstates.

Every function draws into ax (the current axes if omitted), labels the axes and returns what it drew, so the plot can be refined afterwards:

so.plotting.plot1d(grid) # |psi|^2 with the potential in the background so.plotting.plot2d(grid, kind="phase", potential_levels=3) so.plotting.plot_eigenstates(grid, states, energies) so.plotting.plot_phase_space(*so.wigner(grid))

plot1d(grid, psi=None, ax=None, potential=True, scale_pot=None, **kwargs)

Plot \(\vert\psi\vert^2\) on a 1d grid, with the (scaled) real part of the potential as a shaded background.

Parameters:

Name	Type	Description	Default
grid	Grid1d		required
psi	array	Wavefunction; defaults to grid.psi.	None
ax	matplotlib Axes		None
potential	bool	Draw the potential.	True
scale_pot	float	Factor applied to the potential; by default its largest absolute value is scaled to the maximum of \(\vert\psi\vert^2\).	None
**kwargs		Passed on to ax.plot for the density.	{}

Returns:

Name	Type	Description
line	matplotlib Line2D

plot2d(grid, data=None, kind='density', ax=None, potential_levels=0, vmax=None, **kwargs)

Show a field on a 2d grid as an image with the right orientation and extent.

Parameters:

Name	Type	Description	Default
grid	Grid2d		required
data	array	Array indexed as data[i_x, i_y]; defaults to grid.psi.	None
kind	(density, abs, phase, real)	Show \(\vert\psi\vert^2\), \(\vert\psi\vert\), \(\psi\) with the phase as hue and \(\vert\psi\vert\) as brightness, or the real part of data (e.g. a potential).	"density"
ax	matplotlib Axes		None
potential_levels	int or sequence	Contour levels of the real part of the potential drawn on top; 0 or None for none.	0
vmax	float	Value for full color (or full brightness for "phase").	None
**kwargs		Passed on to ax.imshow, e.g. cmap.	{}

Returns:

Name	Type	Description
image	matplotlib AxesImage

plot_momentum(grid, psi=None, ax=None, **kwargs)

Plot the momentum distribution \(\vert\tilde\psi(k)\vert^2\), normalized to \(\int\vert\tilde\psi\vert^2\,dk = 1\).

In 1d this is a line plot, in 2d an image over \((k_x, k_y)\).

Parameters:

Name	Type	Description	Default
grid	Grid1d or Grid2d		required
psi	array	Wavefunction; defaults to grid.psi.	None
ax	matplotlib Axes		None
**kwargs		Passed on to ax.plot (1d) or ax.imshow (2d).	{}

Returns:

Type	Description
line or image

plot_eigenstates(grid, states, energies, ax=None, scale=None, potential=True)

Plot 1d eigenstates as \(\vert\varphi_n\vert^2 + E_n\) over the potential, the usual "ladder" picture.

Parameters:

Name	Type	Description	Default
grid	Grid1d		required
states	sequence of arrays		required
energies	sequence of float		required
ax	matplotlib Axes		None
scale	float	Factor applied to the densities; by default the largest density is scaled to 0.8 times the smallest level spacing.	None
potential	bool	Draw the real part of the potential.	True

Returns:

Name	Type	Description
lines	list of matplotlib Line2D

plot_phase_space(x, p, W, ax=None, vmax=None, **kwargs)

Show a phase space function such as the Wigner function as an image over \((x, p)\).

Functions with negative values are drawn with a diverging colormap centered at zero (red positive, blue negative), non-negative ones such as the Husimi function with a sequential one.

Parameters:

Name	Type	Description	Default
x	arrays	As returned by so.wigner() or so.husimi(): W[i, j] at (x[i], p[j]).	required
p	arrays	As returned by so.wigner() or so.husimi(): W[i, j] at (x[i], p[j]).	required
W	arrays	As returned by so.wigner() or so.husimi(): W[i, j] at (x[i], p[j]).	required
ax	matplotlib Axes		None
vmax	float	Value for full color; defaults to \(\max\vert W\vert\).	None
**kwargs		Passed on to ax.imshow, e.g. cmap.	{}

Returns:

Name	Type	Description
image	matplotlib AxesImage