qscat.viz

Rendering wavefunctions: a cached sparse projector onto an equidistant grid, complex-plane domain colouring, and static or animated 2-D field plots. Needs the plot extra (uv sync --all-packages --extra plot).

EquidistantProjector	Cached projector from a 2-D TensorGrid state to a uniform sampling grid.
WavefunctionArtist	Draw and update one domain-coloured wavefunction panel in a given axes.
complex_to_hsv	Complex array -> HSV array (shape z.shape + (3,)), all channels in [0, 1].
complex_to_rgb	Complex array -> RGB array (shape z.shape + (3,)), channels in [0, 1].
hsv_to_rgb	Vectorised HSV -> RGB (channels last), matching matplotlib's convention.
region_magnitudes	Per-point brightness scale that normalises each region separately.
energy_contour_levels	Assemble energy-relevant potential-contour levels (Hartree), sorted.
plot_wavefunction_2d	Project, domain-colour, and draw a 2-D state; save to path if given.
animate_wavefunction	Animate a sequence of states through one WavefunctionArtist.
animate_artists	Animate several panels (artist + its frames) together in one figure.

Wavefunction visualisation: project a 2-D FEM-DVR-ECS state to a uniform grid and domain-colour the complex field.

Two composable pieces, ported from eMoScat’s EquidistantProjector2d + display_wf.py:

• EquidistantProjector – cached sparse projection of a 2-D TensorGrid state onto an equidistant sampling grid (build once, apply per frame).

• complex_to_rgb / complex_to_hsv – domain colouring (phase -> hue, magnitude -> brightness), pure numpy.

• plot_wavefunction_2d – project + colour + draw (matplotlib, the optional qscat[plot] extra; imported lazily).

This is the building block for time-dependent wavefunction animations: propagate psi(t), project + colour each frame, stitch. matplotlib is NOT imported at module load, so import qscat.viz works without the plotting extra.

class qscat.viz.EquidistantProjector(tgrid, *, samples=400, extent=None)

Cached projector from a 2-D TensorGrid state to a uniform sampling grid.

Parameters:

• tgrid (TensorGrid) – A 2-D tensor grid (len(tgrid.grids) == 2); axis 0 is the first grid (e.g. electronic r), axis 1 the second (e.g. nuclear R).

• samples (int or tuple[int, int]) – Number of uniform samples per axis (a single int applies to both).

• extent (tuple[tuple[float, float], tuple[float, float]] or None) – ((a0, b0), (a1, b1)) real sampling ranges per axis; defaults to each grid’s full real region [x_min, R0] when None.

project(state)

Project a flat 2-D state to the uniform grid, shape (n0, n1) complex.

Parameters:

state (NDArray[complex128])

Return type:

NDArray[complex128]

project_values(field)

Project a nodal-VALUE field (e.g. a potential) on the same grid.

project interpolates a √w-scaled DVR state (a wavefunction). A field given as plain nodal values f(x_i) (e.g. a potential surface) has interpolant sum_i f(x_i) L_i(x) – no 1/√w – so it must be scaled by √w before going through the same operator. Returns shape (n0, n1).

Parameters:

field (NDArray[complex128])

Return type:

NDArray[complex128]

class qscat.viz.WavefunctionArtist(ax, projector, *, mag, inverse=False, title=None, xlabel='axis 1', ylabel='axis 0', contours=False, contour_field='magnitude', contour_color=None, contour_alpha=0.6, contour_linewidth=0.6, potential=None, potential_levels=None, potential_style=':', potential_color=None, potential_alpha=0.7, potential_linewidth=0.6, potential_labels=True, potential_label_fmt='%.3f', eps=None, v_init=0, energies=None)

Draw and update one domain-coloured wavefunction panel in a given axes.

The potential turning-surface overlay (if any) is STATIC – drawn once at construction. update(state) refreshes only the state-dependent layers (the domain-coloured image and the |psi| contours), so animation redraws the minimum per frame. See plot_wavefunction_2d for the parameter meanings; the style kwargs are identical, including the caveat that an array mag’s magnitude contour levels key off its maximum, not its per-point values.

Parameters:

• ax (Any)

• projector (EquidistantProjector)

• mag (float | npt.NDArray[np.float64])

• inverse (bool)

• title (str | None)

• xlabel (str)

• ylabel (str)

• contours (bool | int | Sequence[float])

• contour_field (Literal['magnitude', 'potential'])

• contour_color (str | None)

• contour_alpha (float)

• contour_linewidth (float)

• potential (Any)

• potential_levels (Sequence[float] | Literal['auto'] | None)

• potential_style (str)

• potential_color (str | None)

• potential_alpha (float)

• potential_linewidth (float)

• potential_labels (bool)

• potential_label_fmt (str)

• eps (npt.ArrayLike | None)

• v_init (int)

• energies (npt.ArrayLike | None)

update(state, *, phase=0.0)

Refresh the image + |psi| contours for a new state; return changed artists.

phase (radians) applies a global hue rotation psi -> e^{i*phase}*psi to the COLOURING only – used to view the phase relative to a reference energy (e.g. phase = E_ref * t removes the channel base-energy spin). Since it is a global phase it leaves |psi| (brightness, contours, the potential overlay) unchanged, and by linearity of the projector it is equivalent to phasing the state – applied here on the small projected field, so it costs nothing extra over the per-frame recolour.

Parameters:

• state (NDArray[complex128])

• phase (float)

Return type:

list[Any]

qscat.viz.animate_artists(fig, panels, *, outfile=None, fps=15, writer=None)

Animate several panels (artist + its frames) together in one figure.

Each panel is (artist, frames); all are advanced in lockstep (the frame count is the shortest panel’s). Use this for composed / side-by-side views: build a multi-axes figure, a WavefunctionArtist per axes, and pass them all.

Parameters:

• fig (Any)

• panels (Sequence[tuple[WavefunctionArtist, Sequence[NDArray[complex128]]]])

• outfile (str | PathLike[str] | None)

• fps (int)

• writer (Any)

Return type:

Any

qscat.viz.animate_wavefunction(projector, frames, *, mag, times=None, time_fmt='t = {:.1f}', phase_reference=0.0, outfile=None, fps=15, writer=None, ax=None, figsize=(8, 6), title=None, **style)

Animate a sequence of states through one WavefunctionArtist.

Parameters:

• projector (EquidistantProjector) – Cached projector for the states’ tensor grid.

• frames (iterable of ndarray) – The states psi(t) to animate, one per frame.

• mag (float or ndarray) – Fixed brightness scale across all frames. A scalar applies one scale to the whole field; an array (same shape as the projected field, e.g. from region_magnitudes) gives a per-point scale, held fixed across frames along with the potential overlay. As in plot_wavefunction_2d, magnitude contour levels still key off the array’s maximum, so weaker regions get no meaningful contour lines despite their own brightness scale.

• times (sequence of float, optional) – Per-frame times; when given the title shows time_fmt.format(t) and enable phase_reference.

• phase_reference (float, optional) – Channel base energy E_ref (Hartree). Each frame is coloured after a global phase e^{+i E_ref * times[i]}, i.e. the phase is shown RELATIVE to E_ref – removing the fast base-energy hue spin so the wavepacket’s relative motion in the channel is visible. Default 0 (no shift); requires times. |psi| (brightness/contours) is unaffected.

• outfile (path-like, optional) – Save target: .mp4 (ffmpeg) or .gif (pillow). If None, the FuncAnimation is returned unsaved.

• fps (optional) – Frame rate, an explicit matplotlib writer, an existing Axes (else a new figure), figure size, and a base title.

• writer (optional) – Frame rate, an explicit matplotlib writer, an existing Axes (else a new figure), figure size, and a base title.

• ax (optional) – Frame rate, an explicit matplotlib writer, an existing Axes (else a new figure), figure size, and a base title.

• figsize (optional) – Frame rate, an explicit matplotlib writer, an existing Axes (else a new figure), figure size, and a base title.

• title (optional) – Frame rate, an explicit matplotlib writer, an existing Axes (else a new figure), figure size, and a base title.

• **style – Forwarded to WavefunctionArtist (contours/potential/colours/…).

• time_fmt (str)

Return type:

matplotlib.animation.FuncAnimation

qscat.viz.complex_to_hsv(z, mag=1.0, *, inverse=False)

Complex array -> HSV array (shape z.shape + (3,)), all channels in [0, 1].

Parameters:

• z (ArrayLike)

• mag (float | NDArray[float64])

• inverse (bool)

Return type:

NDArray[float64]

qscat.viz.complex_to_rgb(z, mag=1.0, *, inverse=False)

Complex array -> RGB array (shape z.shape + (3,)), channels in [0, 1].

Parameters:

• z (ArrayLike)

• mag (float | NDArray[float64])

• inverse (bool)

Return type:

NDArray[float64]

qscat.viz.energy_contour_levels(*, eps=None, v_init=0, energies=None, include_thresholds=True, e_range=None, max_levels=12, min_spacing=None)

Assemble energy-relevant potential-contour levels (Hartree), sorted.

Parameters:

• eps (array_like, optional) – Vibrational energies. Their values become threshold levels (when include_thresholds), and eps[v_init] is the base for total energies.

• v_init (int, optional) – Initial vibrational level; eps[v_init] + E is the total energy.

• energies (array_like, optional) – Collision energies; each adds a total-energy level eps[v_init] + E (requires eps).

• include_thresholds (bool, optional) – Include the vibrational-threshold levels eps_v.

• e_range (tuple[float, float], optional) – Keep only levels within [lo, hi] (e.g. the shown potential’s range).

• max_levels (int, optional) – Cap the number of levels, keeping an evenly-spread subset.

• min_spacing (float, optional) – Drop levels closer than this to the previous kept level (de-clutter).

Returns:

Sorted, de-cluttered levels.

Return type:

list[float]

qscat.viz.hsv_to_rgb(hsv)

Vectorised HSV -> RGB (channels last), matching matplotlib’s convention.

Parameters:

hsv (NDArray[float64])

Return type:

NDArray[float64]

qscat.viz.plot_wavefunction_2d(projector, state, *, mag, path=None, inverse=False, title=None, xlabel='axis 1', ylabel='axis 0', ax=None, contours=False, contour_field='magnitude', potential=None, contour_color=None, contour_alpha=0.6, contour_linewidth=0.6, potential_levels=None, potential_style=':', potential_color=None, potential_alpha=0.7, potential_linewidth=0.6, potential_labels=True, potential_label_fmt='%.3f', eps=None, v_init=0, energies=None, phase=0.0)

Project, domain-colour, and draw a 2-D state; save to path if given.

A convenience wrapper: it makes a figure/axes if ax is None, builds a WavefunctionArtist, draws state once, and returns the image handle.

Parameters:

• projector (EquidistantProjector) – A cached projector for the state’s tensor grid.

• state (ndarray) – The flat 2-D complex state to render.

• mag (float or ndarray) – Magnitude mapping to full brightness (see complex_to_hsv). A scalar applies one brightness scale to the whole field; an array (same shape as the projected field, e.g. from region_magnitudes) gives a per-point scale, letting a weak region reach full brightness even when a strong region would otherwise saturate the whole colour map. Note: magnitude contour levels (contours=True, contour_field="magnitude") must be scalars, so with an array mag they key off its MAXIMUM (the brightest region) – weaker regions get no meaningful contour lines even though their brightness is still individually scaled.

• path (path-like, optional) – If given, the figure is saved here (PNG).

• inverse (bool, optional) – Light-background colour mapping (swaps saturation/value). Distinct from the pending print-mode value inversion (see complex_to_hsv).

• title (str, optional) – Axis title/labels; axis 0 is the projector’s first grid, axis 1 the second (e.g. ylabel="electronic r", xlabel="nuclear R").

• xlabel (str, optional) – Axis title/labels; axis 0 is the projector’s first grid, axis 1 the second (e.g. ylabel="electronic r", xlabel="nuclear R").

• ylabel (str, optional) – Axis title/labels; axis 0 is the projector’s first grid, axis 1 the second (e.g. ylabel="electronic r", xlabel="nuclear R").

• ax (matplotlib Axes, optional) – Draw into an existing Axes (composition); a new figure is made when None.

• contours (bool or int or sequence of float, optional) – |psi| contour overlay: False (default) none; True default levels; int that many; a sequence explicit levels.

• contour_field ({"magnitude", "potential"}, optional) – Primary contour source: |psi| (levels k*mag/5) or the potential field.

• potential (ndarray or callable, optional) – The potential field: nodal values on the SAME tensor grid (projected via EquidistantProjector.project_values) or a callable V(r, R). For the FULL 2-D PES pass model.surface (v0(R) + ell(ell+1)/2r^2 + v_int(r,R)), NOT v0 or interaction_diag alone.

• contour_color (optional) – |psi|-contour style. The colour defaults to whatever contrasts with the ground: white on the normal dark render, black under inverse. Pass a colour to override.

• contour_alpha (optional) – |psi|-contour style. The colour defaults to whatever contrasts with the ground: white on the normal dark render, black under inverse. Pass a colour to override.

• contour_linewidth (optional) – |psi|-contour style. The colour defaults to whatever contrasts with the ground: white on the normal dark render, black under inverse. Pass a colour to override.

• potential_levels (sequence of float or "auto", optional) – Enables the dotted potential overlay (in addition to |psi|) when this and potential are given. Explicit energies (Hartree) or "auto" (turning surfaces via energy_contour_levels from eps/energies).

• potential_style (optional) – Potential-overlay style; dotted at 0.7 alpha, in a grey that follows inverse (light on dark, dark on light).

• potential_color (optional) – Potential-overlay style; dotted at 0.7 alpha, in a grey that follows inverse (light on dark, dark on light).

• potential_alpha (optional) – Potential-overlay style; dotted at 0.7 alpha, in a grey that follows inverse (light on dark, dark on light).

• potential_linewidth (optional) – Potential-overlay style; dotted at 0.7 alpha, in a grey that follows inverse (light on dark, dark on light).

• potential_labels (bool, optional) – Inline-label each dotted line with its energy (potential_label_fmt).

• eps (optional) – For potential_levels="auto": vibrational energies, initial level, and collision energies (levels eps_v and eps[v_init] + E).

• v_init (optional) – For potential_levels="auto": vibrational energies, initial level, and collision energies (levels eps_v and eps[v_init] + E).

• energies (optional) – For potential_levels="auto": vibrational energies, initial level, and collision energies (levels eps_v and eps[v_init] + E).

• phase (float, optional) – Global hue-rotation (radians) applied to the colouring only (psi -> e^{i*phase}*psi); leaves |psi| unchanged. For animation, animate_wavefunction’s phase_reference sets phase = E_ref*t to view the phase relative to a channel base energy.

• potential_label_fmt (str)

Returns:

The drawn image handle.

Return type:

matplotlib.image.AxesImage

Raises:

ModuleNotFoundError – If matplotlib is not installed (uv sync –all-packages –extra plot).

qscat.viz.region_magnitudes(magnitude, *, axis, boundaries, percentile=99.5, floor=1e-300)

Per-point brightness scale that normalises each region separately.

A single global mag ties the whole field to its largest feature, so a region whose amplitude is orders of magnitude smaller renders black. This splits axis at boundaries and gives every point the scale of its own region, so each region is visible on its own terms.

The scale is a percentile rather than the maximum, so one outlying point cannot flatten the region it sits in.

Parameters:

• magnitude (array_like) – Non-negative field, e.g. np.abs(psi).

• axis (int) – Axis the split runs along.

• boundaries (sequence of int) – Strictly increasing split indices along axis, each in 1 .. n-1. An empty sequence means one region (the whole field).

• percentile (float) – Percentile of each region’s magnitudes mapped to full brightness.

• floor (float) – Lower clamp, so an all-zero region yields a positive scale rather than a divide-by-zero downstream.

Returns:

Same shape as magnitude; every point carries its region’s scale. Pass it straight to complex_to_rgb as mag.

Return type:

ndarray