"""Validate ARPES plotting utilities.
Extended Summary
----------------
The tests exercise the public plotting functions. They cover default and
custom rendering, shape validation, and existing axes. They also cover color
limits, color bars, and empty k-path labels.
"""
import chex
import equinox as eqx
import jax.numpy as jnp
import matplotlib
import pytest
from beartype import beartype
from beartype.typing import Any, Callable
from jaxtyping import Array, Float, jaxtyped
import diffpes
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from matplotlib.collections import PathCollection
from matplotlib.figure import Figure
from matplotlib.image import AxesImage
from diffpes.inout import (
apply_kpath_ticks,
list_band_scatter_presets,
plot_arpes_spectrum,
plot_arpes_with_kpath,
plot_band_scatter_preset,
plot_band_scatter_with_kpath,
)
from diffpes.types import (
ArpesSpectrum,
BandStructure,
KPathInfo,
OrbitalProjection,
make_arpes_spectrum,
make_band_structure,
make_kpath_info,
make_orbital_projection,
make_spin_orbital_projection,
)
def _unvalidated_kpath(
label_indices: list[int], labels: tuple[str, ...]
) -> KPathInfo:
"""Build malformed legacy k-path metadata for plotting edge tests."""
kpath: KPathInfo = KPathInfo(
num_kpoints=jnp.asarray(60, dtype=jnp.int32),
label_indices=jnp.asarray(label_indices, dtype=jnp.int32),
points_per_segment=jnp.asarray(0, dtype=jnp.int32),
segments=jnp.asarray(0, dtype=jnp.int32),
kpoints=None,
weights=None,
grid=None,
shift=None,
mode="Line-mode",
labels=labels,
comment="",
coordinate_mode="",
)
return kpath
@jaxtyped(typechecker=beartype)
def _make_spectrum(nk: int = 20, ne: int = 120) -> ArpesSpectrum:
"""Build a minimal ArpesSpectrum for plotting tests.
Creates a valid ArpesSpectrum with intensity shape (nk, ne) and
energy_axis length ne, so that plot functions receive consistent
test data without reading files.
Parameters
----------
nk : int, optional
Number of k-points (first dimension of intensity). Default 20.
ne : int, optional
Number of energy points (second dimension of intensity and
length of energy_axis). Default 120.
Returns
-------
spectrum : ArpesSpectrum
PyTree with intensity (nk, ne) and energy_axis (ne,).
"""
intensity: Float[Array, "nk ne"] = jnp.linspace(
0.0, 1.0, nk * ne, dtype=jnp.float64
).reshape(nk, ne)
energy_axis: Float[Array, " ne"] = jnp.linspace(-2.0, 0.5, ne)
spectrum: ArpesSpectrum = make_arpes_spectrum(
intensity=intensity,
energy_axis=energy_axis,
)
return spectrum
[docs]
class TestPlotArpesSpectrum(chex.TestCase):
"""Validate :func:`diffpes.inout.plot_arpes_spectrum`.
Covers default plotting, optional color limits and colorbar,
validation of spectrum array dimensions and shape compatibility,
and reuse of a user-provided matplotlib axis.
:see: :func:`~diffpes.inout.plot_arpes_spectrum`
"""
[docs]
def test_returns_expected_objects(self) -> None:
"""Plot with default options returns figure, axis, and image with correct shape and labels.
The test builds a spectrum and calls plot_arpes_spectrum with colorbar=False.
The test checks the transposed image shape ``(120, 20)``. It also
checks the default axis labels and title. The test closes the figure.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: diffpes.types.ArpesSpectrum
fig: Figure
ax: Axes
image: AxesImage
spectrum = _make_spectrum()
fig, ax, image = plot_arpes_spectrum(spectrum, colorbar=False)
chex.assert_equal(image.get_array().shape, (120, 20))
chex.assert_equal(ax.get_xlabel(), "k-point index")
chex.assert_equal(ax.get_ylabel(), "Energy (eV)")
chex.assert_equal(ax.get_title(), "Simulated ARPES Spectrum")
plt.close(fig)
[docs]
def test_with_clim_and_colorbar(self) -> None:
"""Verify color limits and a color bar with custom options.
The test calls plot_arpes_spectrum with colorbar=True and clim=(0.0, 0.5).
The test compares the image color limits with ``(0.0, 0.5)`` through
``get_clim``. This input covers the limits and color bar paths.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: diffpes.types.ArpesSpectrum
fig: Figure
ax: Axes
image: AxesImage
spectrum = _make_spectrum()
fig, ax, image = plot_arpes_spectrum(
spectrum, colorbar=True, clim=(0.0, 0.5)
)
chex.assert_equal(image.get_clim(), (0.0, 0.5))
plt.close(fig)
[docs]
def test_validation_rejects_wrong_intensity_ndim(self) -> None:
"""_prepare_plot_arrays raises ValueError when intensity is not 2D.
The test constructs an ArpesSpectrum with 1D intensity (bypassing the
factory's type checks) and calls plot_arpes_spectrum. Expects
a ValueError whose message indicates that spectrum.intensity
must have shape (K, E). This validates the dimension check
inside the plotting pipeline.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: ArpesSpectrum = eqx.tree_at(
lambda candidate: candidate.intensity,
_make_spectrum(nk=5, ne=10),
jnp.linspace(0.0, 1.0, 10),
)
with pytest.raises(
ValueError, match="Expected spectrum.intensity to have shape"
):
plot_arpes_spectrum(spectrum, colorbar=False)
[docs]
def test_validation_rejects_wrong_energy_axis_ndim(self) -> None:
"""_prepare_plot_arrays raises ValueError when energy_axis is not 1D.
The test constructs an ArpesSpectrum with 2D energy_axis and calls
plot_arpes_spectrum. Expects a ValueError whose message
indicates that spectrum.energy_axis must have shape (E,).
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: ArpesSpectrum = eqx.tree_at(
lambda candidate: candidate.energy_axis,
_make_spectrum(nk=5, ne=10),
jnp.zeros((10, 2)),
)
with pytest.raises(
ValueError, match="Expected spectrum.energy_axis to have shape"
):
plot_arpes_spectrum(spectrum, colorbar=False)
[docs]
def test_validation_rejects_shape_mismatch(self) -> None:
"""_prepare_plot_arrays raises ValueError when intensity and energy_axis lengths disagree.
The test uses intensity of shape (5, 10) and energy_axis of length 7, so
intensity.shape[1] != energy_axis.shape[0]. Expects a ValueError
with a message about incompatible shapes.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: ArpesSpectrum = eqx.tree_at(
lambda candidate: candidate.energy_axis,
_make_spectrum(nk=5, ne=10),
jnp.linspace(-1.0, 1.0, 7),
)
with pytest.raises(ValueError, match="Incompatible shapes"):
plot_arpes_spectrum(spectrum, colorbar=False)
[docs]
def test_uses_existing_axis(self) -> None:
"""Verify reuse of a given figure and axis.
The test creates a figure and axis with plt.subplots(), then passes ax to
``plot_arpes_spectrum``. The test verifies the identities of the
returned figure and axis. It also verifies the image on the given axis.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: diffpes.types.ArpesSpectrum
fig: Figure
ax: Axes
out_fig: Figure
out_ax: Axes
spectrum = _make_spectrum(nk=10, ne=40)
fig, ax = plt.subplots()
out_fig, out_ax, _ = plot_arpes_spectrum(
spectrum, ax=ax, colorbar=False
)
chex.assert_equal(out_fig is fig, True)
chex.assert_equal(out_ax is ax, True)
plt.close(fig)
[docs]
class TestApplyKpathTicks(chex.TestCase):
"""Validate :func:`diffpes.inout.apply_kpath_ticks`.
The tests apply symmetry-point ticks and labels to an axis. They cover
fewer labels than indices and the path without labels.
:see: :func:`~diffpes.inout.apply_kpath_ticks`
"""
[docs]
def test_sets_ticks_and_labels(self) -> None:
"""apply_kpath_ticks sets x-axis ticks and labels from KPathInfo.
The test builds a KPathInfo with four label indices and four labels
(G, M, K, G). The test applies ``apply_kpath_ticks`` to a new axis.
It compares the x-axis labels with the expected ordered labels.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
fig: Figure
ax: Axes
kpath: diffpes.types.KPathInfo
labels: list[str]
fig, ax = plt.subplots()
kpath = make_kpath_info(
num_kpoints=60,
label_indices=[0, 19, 39, 59],
segments=3,
labels=("G", "M", "K", "G"),
)
apply_kpath_ticks(ax, kpath)
labels = [tick.get_text() for tick in ax.get_xticklabels()]
chex.assert_equal(labels, ["G", "M", "K", "G"])
plt.close(fig)
[docs]
def test_handles_label_index_mismatch(self) -> None:
"""Verify truncation when labels are fewer than label indices.
The fixture provides four label indices but only two labels.
The test expects exactly the two available labels. This input covers
the truncation logic without an index error.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
fig: Figure
ax: Axes
kpath: diffpes.types.KPathInfo
labels: list[str]
fig, ax = plt.subplots()
kpath = _unvalidated_kpath([0, 19, 39, 59], ("G", "M"))
apply_kpath_ticks(ax, kpath)
labels = [tick.get_text() for tick in ax.get_xticklabels()]
chex.assert_equal(labels, ["G", "M"])
plt.close(fig)
[docs]
def test_empty_labels_returns_ax_unchanged(self) -> None:
"""Verify the unchanged axis when the k-path has no labels.
The test builds a KPathInfo with empty label_indices and empty labels so
that ``n_labels`` is zero. The test verifies the identity of the
returned axis. It also verifies the return path without tick changes.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
fig: Figure
ax: Axes
kpath: diffpes.types.KPathInfo
out: Array
fig, ax = plt.subplots()
kpath = _unvalidated_kpath([], ())
out = apply_kpath_ticks(ax, kpath)
chex.assert_equal(out is ax, True)
plt.close(fig)
[docs]
class TestPlotArpesWithKpath(chex.TestCase):
"""Validate :func:`diffpes.inout.plot_arpes_with_kpath`.
Covers the combined workflow of plotting an ARPES spectrum and
annotating the k-axis with symmetry labels from KPathInfo.
:see: :func:`~diffpes.inout.plot_arpes_with_kpath`
"""
[docs]
def test_combined_plot(self) -> None:
"""plot_arpes_with_kpath produces a spectrum image and applies k-path ticks and labels.
The test builds a spectrum and a KPathInfo with three symmetry points.
The test calls ``plot_arpes_with_kpath`` and checks the image shape.
It also compares the tick labels and the x-axis label. These checks
verify the spectrum plot and its k-path annotation.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: diffpes.types.ArpesSpectrum
kpath: diffpes.types.KPathInfo
fig: Figure
ax: Axes
image: AxesImage
labels: list[str]
spectrum = _make_spectrum()
kpath = make_kpath_info(
num_kpoints=20,
label_indices=[0, 9, 19],
segments=2,
labels=("G", "M", "K"),
)
fig, ax, image = plot_arpes_with_kpath(
spectrum=spectrum,
kpath=kpath,
colorbar=False,
)
chex.assert_equal(image.get_array().shape, (120, 20))
labels = [tick.get_text() for tick in ax.get_xticklabels()]
chex.assert_equal(labels, ["G", "M", "K"])
chex.assert_equal(ax.get_xlabel(), "Momentum (k)")
plt.close(fig)
@jaxtyped(typechecker=beartype)
def _make_band_and_projection(
nk: int = 12,
nb: int = 3,
na: int = 2,
) -> tuple[BandStructure, OrbitalProjection]:
"""Build minimal band/projection inputs for band-scatter tests."""
eigen: Float[Array, "nk nb"] = jnp.linspace(
-1.2, 0.8, nk * nb, dtype=jnp.float64
).reshape(nk, nb)
kx: Float[Array, " nk"] = jnp.linspace(0.0, 1.0, nk, dtype=jnp.float64)
kpoints: Float[Array, "nk 3"] = jnp.stack(
[kx, jnp.zeros_like(kx), jnp.zeros_like(kx)],
axis=1,
)
bands: BandStructure = make_band_structure(
eigenvalues=eigen,
kpoints=kpoints,
fermi_energy=0.15,
)
projections: Float[Array, "nk nb na 9"] = (
jnp.ones((nk, nb, na, 9), dtype=jnp.float64) * 0.05
)
projections = projections.at[..., 1:4].set(0.2)
spin: Float[Array, "nk nb na 6"] = jnp.zeros(
(nk, nb, na, 6), dtype=jnp.float64
)
spin = spin.at[: nk // 2, ..., 4].set(0.2)
spin = spin.at[nk // 2 :, ..., 5].set(0.3)
orbital_projection: OrbitalProjection = make_orbital_projection(
projections=projections,
spin=spin,
)
result: tuple[BandStructure, OrbitalProjection] = (
bands,
orbital_projection,
)
return result
[docs]
class TestListBandScatterPresets(chex.TestCase):
"""Validate :func:`~diffpes.inout.list_band_scatter_presets`.
Covers the stable public names for orbital, spin, and orbital-angular-
momentum scatter modes.
:see: :func:`~diffpes.inout.list_band_scatter_presets`
"""
[docs]
def test_returns_each_preset_family(self) -> None:
"""Return at least one name from each supported preset family.
The result must expose orbital, signed-spin, and OAM choices so callers
can build selection controls without reading private tables.
Notes
-----
The test calls the listing function once and checks representative public names
in the returned immutable tuple.
"""
presets: tuple[str, ...]
presets = list_band_scatter_presets()
assert "p" in presets
assert "spin_z" in presets
assert "oam_total" in presets
[docs]
class TestPlotBandScatterPreset(chex.TestCase):
"""Validate projected-band scatter plotting presets.
:see: :func:`~diffpes.inout.plot_band_scatter_preset`
"""
[docs]
def test_lists_presets(self) -> None:
"""list_band_scatter_presets returns known keys.
The test establishes the lists presets contract for plot band scatter with the
concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
presets: tuple[str, ...]
presets = list_band_scatter_presets()
assert "p" in presets
assert "d" in presets
assert "spin_z" in presets
assert "oam_total" in presets
[docs]
def test_orbital_preset_plot(self) -> None:
"""Verify the orbital preset scatter for each k-point and band.
The test establishes the orbital preset plot contract for plot band scatter
with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
bands, orb = _make_band_and_projection(nk=10, nb=4, na=2)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands,
orb_proj=orb,
preset="p",
colorbar=False,
)
chex.assert_equal(scatter.get_offsets().shape[0], 40)
chex.assert_equal(ax.get_xlabel(), "Momentum (k)")
chex.assert_equal(ax.get_ylabel(), "Energy (eV)")
plt.close(fig)
[docs]
def test_signed_spin_preset_with_colorbar(self) -> None:
"""Verify a signed-spin preset with a color bar.
The test establishes the signed spin preset with colorbar contract for plot
band scatter with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
fig: Figure
bands, orb = _make_band_and_projection(nk=8, nb=2, na=2)
fig, _, _ = plot_band_scatter_preset(
bands=bands,
orb_proj=orb,
preset="spin_z",
colorbar=True,
)
chex.assert_equal(len(fig.axes), 2)
plt.close(fig)
[docs]
def test_spin_preset_requires_spin_data(self) -> None:
"""Verify that spin presets require a spin field.
The test establishes the spin preset requires spin data contract for plot band
scatter with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
no_spin: diffpes.types.OrbitalProjection
bands, orb = _make_band_and_projection(nk=8, nb=2, na=1)
no_spin = make_orbital_projection(
projections=orb.projections, spin=None
)
with pytest.raises(ValueError, match="requires spin data"):
plot_band_scatter_preset(
bands=bands,
orb_proj=no_spin,
preset="spin_z",
)
[docs]
def test_band_scatter_with_kpath(self) -> None:
"""Verify symmetry labels on the projected-band scatter.
The test establishes the band scatter with kpath contract for plot band scatter
with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
kpath: diffpes.types.KPathInfo
fig: Figure
ax: Axes
scatter: PathCollection
labels: list[str]
bands, orb = _make_band_and_projection(nk=10, nb=3, na=2)
kpath = make_kpath_info(
num_kpoints=10,
label_indices=[0, 4, 9],
segments=2,
labels=("G", "M", "K"),
)
fig, ax, scatter = plot_band_scatter_with_kpath(
bands=bands,
orb_proj=orb,
kpath=kpath,
preset="d",
colorbar=False,
)
chex.assert_equal(scatter.get_offsets().shape[0], 30)
labels = [tick.get_text() for tick in ax.get_xticklabels()]
chex.assert_equal(labels, ["G", "M", "K"])
plt.close(fig)
[docs]
class TestPlotBandScatterWithKpath(chex.TestCase):
"""Validate :func:`~diffpes.inout.plot_band_scatter_with_kpath`.
Covers composition of projected-band marker weights with line-mode
symmetry labels on the shared momentum axis.
:see: :func:`~diffpes.inout.plot_band_scatter_with_kpath`
"""
[docs]
def test_applies_labels_to_projected_bands(self) -> None:
"""Apply all supplied symmetry labels to the scatter axis.
A ten-point, three-band fixture must produce thirty offsets and retain
the three requested high-symmetry labels in order.
Notes
-----
The test builds deterministic band and projection carriers, applies a three-label
``KPathInfo``, and checks the collection size and rendered tick text.
"""
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
kpath: diffpes.types.KPathInfo
fig: Figure
ax: Axes
scatter: PathCollection
labels: list[str]
bands, orb = _make_band_and_projection(nk=10, nb=3, na=2)
kpath = make_kpath_info(
num_kpoints=10,
label_indices=[0, 4, 9],
segments=2,
labels=("G", "M", "K"),
)
fig, ax, scatter = plot_band_scatter_with_kpath(
bands=bands,
orb_proj=orb,
kpath=kpath,
preset="d",
colorbar=False,
)
chex.assert_equal(scatter.get_offsets().shape[0], 30)
labels = [tick.get_text() for tick in ax.get_xticklabels()]
chex.assert_equal(labels, ["G", "M", "K"])
plt.close(fig)
[docs]
class TestPlotBandScatterEdgeCases(chex.TestCase):
"""Validate additional paths in the band-scatter plotting helpers.
The tests cover invalid array ranks and incompatible weight shapes.
They cover selections by atom, orbital, spin, and OAM. They also cover
unknown presets, missing OAM data, and an existing axis.
:see: :func:`~diffpes.inout.plot_band_scatter_preset`
"""
def _make_bands_1d(self, nk=4, nb=2):
"""Build BandStructure with 1D eigenvalues (bypassing factory)."""
valid_bands: BandStructure = make_band_structure(
eigenvalues=jnp.zeros((nk, nb), dtype=jnp.float64),
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
kpoint_weights=jnp.zeros(nk, dtype=jnp.float64),
fermi_energy=jnp.float64(0.0),
)
malformed_bands: BandStructure = eqx.tree_at(
lambda candidate: candidate.eigenvalues,
valid_bands,
jnp.zeros(nk * nb, dtype=jnp.float64),
)
return malformed_bands
def _make_orb_with_spin_and_oam(self, nk=4, nb=2, na=1):
"""Build OrbitalProjection with spin and OAM attached."""
proj: Array
spin: Array
oam: Array
proj = jnp.ones((nk, nb, na, 9), dtype=jnp.float64) * 0.1
spin = jnp.zeros((nk, nb, na, 6), dtype=jnp.float64)
spin = spin.at[..., 0].set(0.3)
spin = spin.at[..., 4].set(0.2)
oam = jnp.ones((nk, nb, na, 3), dtype=jnp.float64) * 0.05
return make_orbital_projection(projections=proj, spin=spin, oam=oam)
[docs]
def test_prepare_band_arrays_wrong_ndim_raises(self) -> None:
"""``_prepare_band_arrays`` with 1D eigenvalues raises ValueError (lines 449-450).
The test constructs a BandStructure with 1D eigenvalues (bypassing the
factory), then calls ``plot_band_scatter_preset``. Asserts a
``ValueError`` matching ``"shape (K, B)"``.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
bands = self._make_bands_1d(nk=4, nb=2)
proj = jnp.ones((4, 2, 1, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
with pytest.raises(ValueError, match="shape"):
plot_band_scatter_preset(bands=bands, orb_proj=orb, preset="p")
[docs]
def test_subset_atom_axis_with_indices(self) -> None:
"""Verify atom-axis selection with explicit atom indices.
The test plots the p preset with ``atom_indices=[0]``.
Thus, ``_subset_atom_axis`` receives an index array.
The test asserts the scatter point count equals nk * nb.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
na: int
eigen: Array
kpoints: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
nk, nb, na = 6, 2, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
kpoints = jnp.zeros((nk, 3), dtype=jnp.float64)
bands = make_band_structure(eigenvalues=eigen, kpoints=kpoints)
proj = jnp.ones((nk, nb, na, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands,
orb_proj=orb,
preset="p",
atom_indices=[0],
colorbar=False,
)
chex.assert_equal(scatter.get_offsets().shape[0], nk * nb)
plt.close(fig)
[docs]
def test_s_orbital_preset(self) -> None:
"""Verify the s-orbital preset branch.
The test calls ``plot_band_scatter_preset`` with ``preset='s'``. Asserts
that the scatter renders without error and the point count is
correct (exercises the ``ORBITAL_INDEX[key]`` branch).
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
nk, nb = 6, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb, 1, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="s", colorbar=False
)
chex.assert_equal(scatter.get_offsets().shape[0], nk * nb)
plt.close(fig)
[docs]
def test_total_preset(self) -> None:
"""Verify that the total preset sums all orbital channels.
The test calls ``plot_band_scatter_preset`` with ``preset='total'``.
The test asserts the scatter renders and the point count is correct
(exercises the ``elif key == 'total'`` branch at line 501-502).
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
nk, nb = 4, 3
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb, 1, 9), dtype=jnp.float64) * 0.05
orb = make_orbital_projection(projections=proj)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="total", colorbar=False
)
chex.assert_equal(scatter.get_offsets().shape[0], nk * nb)
plt.close(fig)
[docs]
def test_spin_channel_preset(self) -> None:
"""Verify the spin-channel preset branch.
The test calls ``plot_band_scatter_preset`` with ``preset='spin_z_up'``.
This exercises the ``if key in spin_channel`` branch (line 527-528).
The test asserts the scatter renders without error.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
spin: Array
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
nk, nb = 4, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb, 1, 9), dtype=jnp.float64) * 0.1
spin = jnp.ones((nk, nb, 1, 6), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj, spin=spin)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="spin_z_up", colorbar=False
)
chex.assert_equal(scatter.get_offsets().shape[0], nk * nb)
plt.close(fig)
[docs]
def test_oam_preset_with_oam_data(self) -> None:
"""Verify OAM selection and the component branch.
The test calls ``plot_band_scatter_preset`` with ``preset='oam_total'``
and an OrbitalProjection that has ``oam`` data. This exercises
line 541 (``oam_arr = _subset_atom_axis(...)``) and lines
554-556 (``if key in oam_component: weights = ...``).
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
nk, nb = 4, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
orb = self._make_orb_with_spin_and_oam(nk=nk, nb=nb, na=1)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="oam_total", colorbar=False
)
chex.assert_equal(scatter.get_offsets().shape[0], nk * nb)
plt.close(fig)
[docs]
def test_oam_abs_total_preset(self) -> None:
"""Verify the absolute-total OAM preset.
The test calls ``plot_band_scatter_preset`` with ``preset='oam_abs_total'``
and an OrbitalProjection with OAM data present. This exercises
the ``elif key == 'oam_abs_total'`` branch at lines 557-559,
which computes ``np.sum(np.abs(oam_arr[..., 2]), axis=2)`` and
sets ``signed = False``. Asserts the scatter renders without error.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
fig: Figure
ax: Axes
scatter: PathCollection
nk, nb = 4, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
orb = self._make_orb_with_spin_and_oam(nk=nk, nb=nb, na=1)
fig, ax, scatter = plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="oam_abs_total", colorbar=False
)
chex.assert_equal(scatter.get_offsets().shape[0], nk * nb)
plt.close(fig)
[docs]
def test_oam_preset_without_oam_data_raises(self) -> None:
"""Verify that an OAM preset requires OAM data.
The test calls ``plot_band_scatter_preset`` with ``preset='oam_p'`` but
provides an OrbitalProjection with ``oam=None``. Asserts
``ValueError`` matching ``"requires OAM data"``.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
nk, nb = 4, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb, 1, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
with pytest.raises(ValueError, match="requires OAM data"):
plot_band_scatter_preset(bands=bands, orb_proj=orb, preset="oam_p")
[docs]
def test_unknown_preset_raises(self) -> None:
"""Verify that an unknown preset raises ``ValueError``.
The test calls ``plot_band_scatter_preset`` with an unrecognized preset
name. Asserts ``ValueError`` matching ``"Unknown preset"``.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
nk, nb = 4, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb, 1, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
with pytest.raises(ValueError, match="Unknown preset"):
plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="not_a_real_preset"
)
[docs]
def test_weight_shape_mismatch_raises(self) -> None:
"""Verify rejection of weights with an incompatible shape.
The test creates incompatible band and projection shapes.
The ``"p"`` preset produces weights with shape ``(4, 3)`` instead of
``(4, 2)``. The test expects a ``ValueError`` that matches
``"Preset weights must have shape"``.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb_bands: int
nb_proj: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
nk = 4
nb_bands = 2
nb_proj = 3
eigen = jnp.linspace(
-1.0, 0.5, nk * nb_bands, dtype=jnp.float64
).reshape(nk, nb_bands)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb_proj, 1, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
with pytest.raises(ValueError, match="Preset weights must have shape"):
plot_band_scatter_preset(bands=bands, orb_proj=orb, preset="p")
[docs]
def test_uses_provided_ax(self) -> None:
"""Verify reuse of a given axis for a band scatter.
The test creates a figure and passes its axis to the plotting function.
It verifies the identity of the returned figure. This check covers
the existing-axis path instead of the new-figure path.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
nk: int
nb: int
eigen: Array
bands: diffpes.types.BandStructure
proj: Array
orb: diffpes.types.OrbitalProjection
fig0: Figure
ax0: Axes
out_fig: Figure
out_ax: Axes
nk, nb = 6, 2
eigen = jnp.linspace(-1.0, 0.5, nk * nb, dtype=jnp.float64).reshape(
nk, nb
)
bands = make_band_structure(
eigenvalues=eigen,
kpoints=jnp.zeros((nk, 3), dtype=jnp.float64),
)
proj = jnp.ones((nk, nb, 1, 9), dtype=jnp.float64) * 0.1
orb = make_orbital_projection(projections=proj)
fig0, ax0 = plt.subplots()
out_fig, out_ax, _ = plot_band_scatter_preset(
bands=bands, orb_proj=orb, preset="p", ax=ax0, colorbar=False
)
chex.assert_equal(out_fig is fig0, True)
plt.close(fig0)