diffpes.tightb

Native tight-binding model construction, Slater-Koster coupling, degeneracy-safe diagonalization, and ARPES-side adapters for external electronic-structure sources.

Provide native tight-binding tools and ARPES-side adapters.

Extended Summary

The native tight-binding layer provides model construction and Slater-Koster coupling. It adds spin-orbit coupling, slabs, and degeneracy-safe diagonalization as the plan series progresses. It also consumes DiagonalizedBands from other electronic-structure sources.

This module retains:

  • ARPES-side adapters that stay here permanently: vasp_to_diagonalized, eigenvector_orbital_weights, orbital_coefficients.

  • Current fixtures, superseded by plan 04: build_hamiltonian_k, diagonalize_single_k, diagonalize_tb.

The following list describes the submodules:

  • diagonalize

    Diagonalize bands and adapt VASP outputs.

  • hamiltonian

    Build tight-binding Hamiltonians in JAX.

  • kspace

    Build differentiable paths and fixed-shape rasters in k-space.

  • projections

    Convert eigenvectors to orbital weights.

Routine Listings

build_hamiltonian_k()

Build the Bloch Hamiltonian H(k) at a single k-point.

build_arpes_kmesh()

Build a fixed-kz ARPES raster in fractional coordinates.

build_bz_mesh()

Build a fixed-shape reciprocal mesh and its first-zone mask.

build_kmesh_hv()

Build a photon-energy raster in fractional coordinates.

build_kpath()

Build a labeled path between k-space anchors.

diagonalize_single_k()

Diagonalize H(k) at a single k-point.

diagonalize_tb()

Diagonalize a TB model at all k-points.

eigenvector_orbital_weights()

Compute orbital weights from eigenvectors.

first_bz_mask()

Mark Cartesian points inside the first Brillouin zone.

kpath_arc_length()

Compute cumulative Cartesian distance along a k-path.

kpoints_cart_to_frac()

Convert Cartesian momenta to fractional k-points.

kpoints_frac_to_cart()

Convert fractional k-points to Cartesian momenta.

orbital_coefficients()

Return the raw complex orbital coefficients.

vasp_to_diagonalized()

Convert VASP BandStructure + OrbitalProjection to DiagonalizedBands.

diffpes.tightb.diagonalize_single_k(H_k: Complex[Array, 'O O']) → tuple[Float[Array, 'O'], Complex[Array, 'O O']][source]

Diagonalize H(k) at a single k-point.

The function calls the standard LAPACK-style Hermitian eigensolver jnp.linalg.eigh. The Hamiltonian construction guarantees Hermitian symmetry. Therefore, eigh produces real eigenvalues and an orthonormal eigenvector basis.

jnp.linalg.eigh returns eigenvalues in ascending order and eigenvectors as columns of the returned matrix: that is, eigenvectors[:, i] is the eigenvector corresponding to eigenvalues[i]. This column-eigenvector convention is the LAPACK/NumPy/JAX standard but differs from some physics textbooks that store eigenvectors as rows. The caller (diagonalize_tb) transposes to a band-major layout after vmapping.

This function contains one call to eigh. Tests can isolate or replace this function with a custom differentiable eigensolver.

See:

TestDiagonalizeSingleK

Parameters:

H_k (Complex[Array, 'O O']) – Hermitian Hamiltonian matrix.

Return type:

tuple[Float[Array, 'O'], Complex[Array, 'O O']]

Returns:

  • eigenvalues (Float[Array, " O"]) – Eigenvalues in ascending order.

  • eigenvectors (Complex[Array, "O O"]) – Eigenvector columns (eigenvectors[:, i] is the i-th).

Notes

JAX provides analytical gradients through eigh via implicit differentiation of the eigenvalue equation, so jax.grad(lambda p: eigenvalues(p).sum()) works without additional configuration. Degenerate eigenvalues can cause numerical instability in the backward pass; this is a known JAX limitation.

diffpes.tightb.diagonalize_tb(tb_model: TBModel, kpoints: Float[Array, 'K 3']) → DiagonalizedBands[source]

Diagonalize a TB model at all k-points.

The function builds H(k) for each k-point and calls jnp.linalg.eigh. jax.vmap vectorizes both operations. JAX differentiates them with respect to tb_model.hopping_params.

The internal _build_and_diag closure captures the model parameters. It maps one k-point to its eigenvalues and eigenvectors. jax.vmap applies this closure along the leading k-point axis. One vectorized call computes the complete band structure without Python loops.

After vmapping, the eigenvector array has shape (K, O, O) in the column-eigenvector convention of eigh. Thus, evecs[k, :, b] is band b at k-point k). A transpose (0, 2, 1) converts this to the band-major convention (K, B, O) where evecs[k, b, :] gives the orbital coefficients of band b at k-point k. This layout matches the DiagonalizedBands.eigenvectors contract used by the rest of the ARPYES pipeline (projections, matrix elements, spectral function).

The function sets the Fermi energy to 0.0 because bare tight-binding models have no absolute energy reference.

See:

TestDiagonalizeTB

Parameters:

  • tb_model (TBModel) – Tight-binding model.

  • kpoints (Float[Array, 'K 3']) – k-points in fractional coordinates.

Returns:

bands – Diagonalized electronic structure.

Return type:

DiagonalizedBands

Notes

Because jax.vmap traces the function once and broadcasts over the batch dimension, the number of k-points does not affect compilation time – only runtime. The full forward + backward pass (eigenvalues and their gradients with respect to hopping parameters) is differentiable end-to-end.

diffpes.tightb.vasp_to_diagonalized(bands: BandStructure, orb_proj: OrbitalProjection, orbital_basis: OrbitalBasis, phase_loss: Literal['warn', 'ignore', 'error'] = 'warn') → DiagonalizedBands[source]

Convert VASP BandStructure + OrbitalProjection to DiagonalizedBands.

VASP PROCAR gives |c_{k,b,orb}|^2, not the complex coefficients. This adapter uses sqrt(|c|^2) with positive sign as an approximation. Phase information is lost.

The adapter sums the orbital projections over atoms. It then maps them to the orbital basis order.

VASP’s PROCAR file provides site- and orbital-projected squared moduli |c_{k,b,atom,orb}|^2 for each eigenstate. The VASP projection discards the complex phase of each coefficient. Therefore, PROCAR data cannot recover the true complex eigenvectors. This adapter takes sqrt(|c|^2) with a positive sign. The approximation produces real, nonnegative coefficients.

The approximation is exact for an observable that depends only on the coefficient modulus. It introduces errors when an observable depends on relative orbital phases. Photoemission interference terms have this phase dependence.

Orbital mapping. VASP stores projections in a fixed 9-orbital ordering for the s, p, and d channels:

[s, py, pz, px, dxy, dyz, dz2, dxz, dx2 - y2]

This order differs from some standard orders, such as the (l, m) convention with m running from -l to +l). The function maps from (l, m) quantum numbers in the OrbitalBasis to VASP’s 9-orbital column index via a lookup table.

Sum the atoms. Before the orbital mapping, the adapter sums the projections over atom axis 2, which changes (K, B, A, 9) -> (K, B, 9). This is correct when the OrbitalBasis describes a single composite orbital per (l, m) channel rather than per-atom resolution.

Normalize the vectors. After the square root, the adapter normalizes each k-point and band eigenvector. Therefore, sum_orb |c_{k,b,orb}|^2 = 1. A safe division selects a zero vector and zero gradient when all projections equal zero.

See:

TestVaspToDiagonalized

Parameters:

  • bands (BandStructure) – VASP eigenvalues and k-points.

  • orb_proj (OrbitalProjection) – VASP orbital projections of shape (K, B, A, 9).

  • orbital_basis (OrbitalBasis) – Quantum number metadata defining which VASP orbital indices to use.

  • phase_loss (Literal['warn', 'ignore', 'error'], default: "warn") – Policy for handling lost phase information: - "warn" (default): emit a runtime warning. - "ignore": proceed without a warning. - "error": raise ValueError and abort.

Returns:

diag – Approximate diagonalized bands.

Return type:

DiagonalizedBands

Raises:

ValueError – If phase_loss is invalid, is "error", or the orbital basis contains a channel outside the VASP nine-orbital set.

Notes

The VASP 9-orbital ordering used here covers s, p, and d channels only. The adapter does not support f-orbital projections and raises ValueError for them. A PROCAR file from LORBIT=11 or higher contains all nine channels. LORBIT=10 can omit the decomposition by m and does not support this adapter.

The adapter converts the resulting eigenvectors to complex128 for the DiagonalizedBands type. Their imaginary parts remain zero.

diffpes.tightb.build_hamiltonian_k(k: Float[Array, '3'], hopping_params: Float[Array, 'H'], hopping_indices: tuple, n_orbitals: int, lattice_vectors: Float[Array, '3 3']) → Complex[Array, 'O O'][source]

Build the Bloch Hamiltonian H(k) at a single k-point.

\[H_{ij}(\mathbf{k}) = \sum_{\mathbf{R}} t_{ij,\mathbf{R}} \exp(i \mathbf{k} \cdot \mathbf{R})\]

where the sum runs over lattice vectors R defined by the hopping indices. The function makes the result Hermitian: H = (H_raw + H_raw^dag) / 2.

The function computes the Bloch sum entirely in fractional coordinates. k uses reciprocal lattice units, and each hopping R uses direct lattice units. Therefore, k_frac . R_frac includes the metric. The Cartesian phase equals exp(2 pi i k_frac . R_frac). This stage does not need the lattice matrix. The model retains it for Cartesian conversions.

The function iterates over each hopping in Python because the hopping count is small and static. For each hopping, it computes the Bloch phase. It adds the corresponding amplitude to H[orb_i, orb_j].

After the accumulation, the function applies (H + H^dag) / 2. A user-supplied hopping list can contain only the upper triangle. This operation fills the lower triangle symmetrically. It also corrects floating-point rounding that can prevent jnp.linalg.eigh from working.

See:

TestBuildHamiltonianK

Parameters:

  • k (Float[Array, '3']) – k-point in fractional coordinates.

  • hopping_params (Float[Array, 'H']) – Hopping amplitudes (differentiable).

  • hopping_indices (tuple) – (orb_i, orb_j, (R_x, R_y, R_z)) per hopping.

  • n_orbitals (int) – Number of orbitals in the unit cell.

  • lattice_vectors (Float[Array, '3 3']) – Real-space lattice vectors (rows).

Returns:

H_k – Hermitian Hamiltonian matrix.

Return type:

Complex[Array, 'O O']

Notes

The hopping amplitudes hopping_params are plain real floats. Complex spin-orbit hoppings require complex128 values and no automatic Hermitian completion.

The @jaxtyped and @beartype decorators detect shape and dtype errors when the caller invokes the function.

diffpes.tightb.build_arpes_kmesh(kx_axis_inv_ang: Float[Array, 'n_kx'], ky_axis_inv_ang: Float[Array, 'n_ky'], kz_inv_ang: float | Float[Array, ''], sample_azimuth: float | Float[Array, ''], geometry: CrystalGeometry) → KGrid[source]

Build a fixed-kz ARPES raster in fractional coordinates.

The input axes describe the laboratory frame. A negative sample azimuth rotates that raster into the crystal frame before fractional conversion.

See:

TestBuildArpesKmesh

Implementation Logic

  1. Build the laboratory raster:

    lab_x, lab_y = jnp.meshgrid(kx_axis, ky_axis)

    The result uses rows for y and columns for x.

  2. Apply the sample rotation:

    sample_x = cosine * lab_x + sine * lab_y

    This expression applies an active rotation by negative azimuth.

  3. Convert and return the grid:

    return make_kgrid(kpoints_frac, mesh_shape, kz)

    The fixed kz value remains a traced carrier leaf.

param kx_axis_inv_ang:

Laboratory x momenta in 1/Angstrom.

type kx_axis_inv_ang:

Float[Array, 'n_kx']

param ky_axis_inv_ang:

Laboratory y momenta in 1/Angstrom.

type ky_axis_inv_ang:

Float[Array, 'n_ky']

param kz_inv_ang:

Fixed out-of-plane momentum in 1/Angstrom.

type kz_inv_ang:

Union[float, Float[Array, '']]

param sample_azimuth:

Sample rotation about the surface normal in radians.

type sample_azimuth:

Union[float, Float[Array, '']]

param geometry:

Crystal geometry for Cartesian-to-fractional conversion.

type geometry:

CrystalGeometry

returns:

kgrid – Fractional raster with static shape (n_ky, n_kx).

rtype:

KGrid

Notes

Gradients flow through both axes, kz, the azimuth, and the lattice. The raster shape depends only on the static input shapes.

diffpes.tightb.build_bz_mesh(geometry: CrystalGeometry, n_per_axis: int, shell_radius: int = 2) → tuple[KGrid, Bool[Array, 'n_k']][source]

Build a fixed-shape reciprocal mesh and its first-zone mask.

The mesh samples the fractional cube from -1 to 1. A basis-derived bound must prove that this cube contains the complete first zone.

See:

TestBuildBzMesh

Implementation Logic

  1. Build the fractional cube:

    mesh_x, mesh_y, mesh_z = jnp.meshgrid(axis, axis, axis)

    Static axis lengths give the flattened point array a fixed shape.

  2. Compute zone membership:

    mask = first_bz_mask(kpoints_cart, geometry)

    The mask preserves the fixed point count during compiled execution.

  3. Return the grid and mask:

    return kgrid, mask

    The grid stores the three-dimensional mesh as rows of z coordinates.

param geometry:

Crystal geometry for reciprocal conversion.

type geometry:

CrystalGeometry

param n_per_axis:

Number of samples on each fractional axis. This value is static. A change causes retracing.

type n_per_axis:

int

param shell_radius:

Static reciprocal-coefficient radius for the first-zone test. Increase this value for a skew, anisotropic, or unreduced basis. Default is 2.

type shell_radius:

int, default: 2

rtype:

tuple[KGrid, Bool[Array, 'n_k']]

returns:

  • kgrid (KGrid) – Fractional mesh with n_per_axis**3 flattened points.

  • mask (Bool[Array, " n_k"]) – First-zone membership for every mesh point.

raises ValueError:

If n_per_axis is less than two or shell_radius is less than one.

raises EquinoxRuntimeError:

If the reciprocal basis does not prove that the fractional cube contains the first zone. The function also raises if the shell is not provably complete for the provisional first-zone points.

Notes

The static mesh shape is (n_per_axis**2, n_per_axis). Each row follows the z axis for one pair of x and y coordinates.

The coverage guard uses the inequalities abs(k dot b_i) <= norm(b_i)**2 / 2 that every first-zone point obeys. It maps their coordinate-wise bounds through the reciprocal Gram matrix. A failed conservative bound asks the caller for a reduced reciprocal basis instead of returning a partial zone.

diffpes.tightb.build_kmesh_hv(kpar_axis_inv_ang: Float[Array, 'n_kpar'], photon_energies_ev: Float[Array, 'n_hv'], work_function_ev: float | Float[Array, ''], inner_potential_ev: float | Float[Array, ''], sample_azimuth: float | Float[Array, ''], kpar_direction: Float[Array, '2'], geometry: CrystalGeometry) → KGrid[source]

Build a photon-energy raster in fractional coordinates.

Each row contains one photon energy and all requested parallel momenta. The free-electron final-state model supplies the row-dependent kz.

See:

TestBuildKmeshHv

Implementation Logic

  1. Compute each out-of-plane row:

    kz_rows = jax.vmap(kz_for_energy)(photon_energies_ev)

    One vectorized map preserves the dense photon-energy axis.

  2. Build and rotate the Cartesian raster:

    sample_x = cosine * lab_x + sine * lab_y

    The traced azimuth maps the laboratory direction into the sample.

  3. Convert and return the grid:

    return make_kgrid(..., photon_energy_axis_ev=photon_energies_ev)

    The carrier records one photon energy for each raster row.

param kpar_axis_inv_ang:

Signed parallel momenta in 1/Angstrom.

type kpar_axis_inv_ang:

Float[Array, 'n_kpar']

param photon_energies_ev:

Photon energies in eV.

type photon_energies_ev:

Float[Array, 'n_hv']

param work_function_ev:

Work function in eV.

type work_function_ev:

Union[float, Float[Array, '']]

param inner_potential_ev:

Inner potential in eV.

type inner_potential_ev:

Union[float, Float[Array, '']]

param sample_azimuth:

Sample rotation about the surface normal in radians.

type sample_azimuth:

Union[float, Float[Array, '']]

param kpar_direction:

Unit direction in the laboratory surface plane.

type kpar_direction:

Float[Array, '2']

param geometry:

Crystal geometry for Cartesian-to-fractional conversion.

type geometry:

CrystalGeometry

returns:

kgrid – Fractional raster with static shape (n_hv, n_kpar).

rtype:

KGrid

raises EquinoxRuntimeError:

If the direction is not finite and unit length. The function also rejects a raster that contains an evanescent channel.

Notes

The KGrid carrier stores no scalar kz because each row has a different value. The third coordinate of each point stores that value.

diffpes.tightb.build_kpath(anchors: Float[Array, 'n_anchor 3'], geometry: CrystalGeometry, n_per_segment: int, labels: tuple[str, ...], anchor_units: str = 'fractional') → KPath[source]

Build a labeled path between k-space anchors.

Each segment contains both endpoints. Adjacent segments therefore repeat their shared anchor, which matches the Chinook klib.kpath convention.

See:

TestBuildKpath

Implementation Logic

  1. Convert the anchors:

    fractional_anchors = kpoints_cart_to_frac(anchors, geometry)

    This branch applies only to anchors with absolute units.

  2. Interpolate all segments:

    segments = starts + fractions * (ends - starts)

    A broadcast creates every segment with a fixed output shape.

  3. Create the path carrier:

    return make_kpath(...)

    Static indices identify each anchor in the flattened point array.

param anchors:

Path anchors in the coordinates selected by anchor_units.

type anchors:

Float[Array, 'n_anchor 3']

param geometry:

Crystal geometry for absolute-to-fractional conversion.

type geometry:

CrystalGeometry

param n_per_segment:

Points in each segment, including both endpoints. This value is static. A change causes retracing.

type n_per_segment:

int

param labels:

One label for each anchor. This value is static. A change causes retracing.

type labels:

tuple[str, ...]

param anchor_units:

Static coordinate selector. Use "fractional" or "absolute". A change causes retracing. Default is "fractional".

type anchor_units:

str, default: "fractional"

returns:

kpath – Fractional path with (n_anchor - 1) * n_per_segment points.

rtype:

KPath

raises ValueError:

If the path has fewer than two anchors or invalid static settings.

raises EquinoxRuntimeError:

If an anchor is non-finite.

Notes

"absolute" means Cartesian momentum in 1/Angstrom. The interpolation remains differentiable with respect to anchors and the lattice.

diffpes.tightb.first_bz_mask(kpoints_cart: Float[Array, 'n_k 3'], geometry: CrystalGeometry, shell_radius: int = 2) → Bool[Array, 'n_k'][source]

Mark Cartesian points inside the first Brillouin zone.

The Wigner-Seitz test compares the origin distance with reciprocal-lattice points in a static integer shell. A singular-value bound proves that no point outside the shell can change a retained membership result.

See:

TestFirstBzMask

Parameters:

  • kpoints_cart (Float[Array, 'n_k 3']) – Cartesian momenta in 1/Angstrom.

  • geometry (CrystalGeometry) – Crystal geometry with reciprocal vectors in 1/Angstrom.

  • shell_radius (int, default: 2) – Static maximum absolute reciprocal coefficient. Increase this value for a skew, anisotropic, or unreduced basis. Default is 2.

Returns:

mask – True for points inside or on the first-zone boundary.

Return type:

Bool[Array, 'n_k']

Raises:

  • ValueError – If shell_radius is less than one.

  • EquinoxRuntimeError – If the static shell is not provably complete for the supplied points and reciprocal basis.

Notes

The comparisons use squared distances and include ties. The Boolean mask has no boundary derivative. Consumers must not differentiate through its discrete membership changes.

For any unseen reciprocal vector G = n @ B, its norm is at least sigma_min(B) * norm(n). A vector can beat the origin only when norm(G) <= 2 * norm(k). The function checks this sufficient bound for every point that survives the requested shell. It raises instead of returning an uncertified false positive when the bound is inconclusive.

diffpes.tightb.kpath_arc_length(kpath: KPath, geometry: CrystalGeometry) → Float[Array, 'n_k'][source]

Compute cumulative Cartesian distance along a k-path.

The plotting coordinate measures each segment after reciprocal-lattice conversion. Repeated junction points add zero distance.

See:

TestKpathArcLength

Parameters:

  • kpath (KPath) – Fractional path and its static plotting metadata.

  • geometry (CrystalGeometry) – Crystal geometry with reciprocal vectors in 1/Angstrom.

Returns:

arc_length – Cumulative Cartesian path distance in 1/Angstrom.

Return type:

Float[Array, 'n_k']

Notes

diffpes.maths.safe_norm() gives a zero gradient for a repeated junction. All nonzero segment lengths retain their usual gradients.

diffpes.tightb.kpoints_cart_to_frac(kpoints_cart: Float[Array, 'n_k 3'], geometry: CrystalGeometry) → Float[Array, 'n_k 3'][source]

Convert Cartesian momenta to fractional k-points.

The function uses the closed reciprocal identity instead of a numerical linear solve.

See:

TestKpointsCartToFrac

Parameters:

  • kpoints_cart (Float[Array, 'n_k 3']) – Cartesian momenta in 1/Angstrom.

  • geometry (CrystalGeometry) – Crystal geometry with real-space vectors in Angstrom.

Returns:

kpoints_frac – Fractional k-points.

Return type:

Float[Array, 'n_k 3']

Notes

Reciprocal and real-space lattice rows satisfy \(B A^T=2\pi I\). Thus, the exact expression is kpoints_cart @ geometry.lattice.T / (2 * pi).

diffpes.tightb.kpoints_frac_to_cart(kpoints_frac: Float[Array, 'n_k 3'], geometry: CrystalGeometry) → Float[Array, 'n_k 3'][source]

Convert fractional k-points to Cartesian momenta.

The row-vector convention applies the reciprocal lattice once. The reciprocal lattice includes the factor of \(2\pi\).

See:

TestKpointsFracToCart

Parameters:

  • kpoints_frac (Float[Array, 'n_k 3']) – Fractional k-points.

  • geometry (CrystalGeometry) – Crystal geometry with reciprocal vectors in 1/Angstrom.

Returns:

kpoints_cart – Cartesian momenta in 1/Angstrom.

Return type:

Float[Array, 'n_k 3']

Notes

The function computes kpoints_frac @ geometry.reciprocal. Gradients flow through the coordinates and all applicable lattice entries.

diffpes.tightb.eigenvector_orbital_weights(eigenvectors: Complex[Array, 'K B O']) → Float[Array, 'K B O'][source]

Compute orbital weights from eigenvectors.

For each eigenstate |psi_{k,b}> expanded in the orbital basis

\[|\psi_{k,b}\rangle = \sum_o c_{k,b,o} |o\rangle,\]

the orbital weight of orbital o in band b at k-point k is the squared modulus of the expansion coefficient:

\[w_{k,b,o} = |c_{k,b,o}|^2.\]

This is the probability of finding the electron in orbital o given that it occupies eigenstate (k, b). By construction, normalized eigenvectors give weights that sum to 1 over orbitals for each (k, b) pair.

Fat-band plots and orbital-resolved DOS use orbital weights. Photoemission matrix element computations also start from these weights. These computations also need the complex coefficients; see orbital_coefficients.

See:

TestEigenvectorOrbitalWeights

Parameters:

eigenvectors (Complex[Array, 'K B O']) – Complex orbital coefficients c_{k,b,orb}.

Returns:

weights – |c_{k,b,orb}|^2 per orbital.

Return type:

Float[Array, 'K B O']

Notes

The implementation uses jnp.abs(eigenvectors) ** 2 rather than (eigenvectors * eigenvectors.conj()).real for clarity. Both expressions are mathematically identical for complex arrays and produce the same JAX trace; the former is marginally more readable.

diffpes.tightb.orbital_coefficients(eigenvectors: Complex[Array, 'K B O']) → Complex[Array, 'K B O'][source]

Return the raw complex orbital coefficients.

A full matrix element computation needs the complex coefficients c_{k,b,orb}, not only |c|^2.

This is an identity function: it returns its input unchanged. Its purpose gives a clear name to each call site. A pipeline can need both eigenvector_orbital_weights and the raw coefficients. Callers can write:

weights = eigenvector_orbital_weights(evecs)
coeffs = orbital_coefficients(evecs)

Thus, each downstream path clearly identifies whether it uses only magnitudes or the full phase information.

In the Chinook-style matrix element computation, the function multiplies each complex coefficient by the applicable one-electron dipole matrix element. It then adds the products coherently. Interference between orbital channels depends on the relative phases of these coefficients. Therefore, this computation needs the complex values, not only |c|^2.

When vasp_to_diagonalized supplies the eigenvectors, the adapter loses the phases. The adapter then returns real, nonnegative coefficients. In that regime, the coherent interference terms are approximate.

See:

TestOrbitalCoefficients

Parameters:

eigenvectors (Complex[Array, 'K B O']) – Complex orbital coefficients.

Returns:

coefficients – Same as input (identity, for API clarity).

Return type:

Complex[Array, 'K B O']

Notes

Because this is an identity, jax.grad through orbital_coefficients adds zero overhead – the function compiles away entirely.

K-space builders

Build differentiable paths and fixed-shape rasters in k-space.

Extended Summary

This module converts fractional and Cartesian k-points through one reciprocal lattice contract. It also builds paths, reciprocal-space meshes, and ARPES rasters without data-dependent output shapes.

Routine Listings

kpoints_frac_to_cart()

Convert fractional k-points to Cartesian momenta.

kpoints_cart_to_frac()

Convert Cartesian momenta to fractional k-points.

build_kpath()

Build a labeled path between k-space anchors.

kpath_arc_length()

Compute cumulative Cartesian distance along a k-path.

first_bz_mask()

Mark Cartesian points inside the first Brillouin zone.

build_bz_mesh()

Build a fixed-shape reciprocal mesh and its first-zone mask.

build_arpes_kmesh()

Build a fixed-kz ARPES raster in fractional coordinates.

build_kmesh_hv()

Build a photon-energy raster in fractional coordinates.

Notes

Every conversion uses reciprocal vectors as matrix rows. The first-zone operation returns a mask because a gather creates a data-dependent shape.

diffpes.tightb.kspace.kpoints_frac_to_cart(kpoints_frac: Float[Array, 'n_k 3'], geometry: CrystalGeometry) → Float[Array, 'n_k 3'][source]

Convert fractional k-points to Cartesian momenta.

The row-vector convention applies the reciprocal lattice once. The reciprocal lattice includes the factor of \(2\pi\).

See:

TestKpointsFracToCart

Parameters:

  • kpoints_frac (Float[Array, 'n_k 3']) – Fractional k-points.

  • geometry (CrystalGeometry) – Crystal geometry with reciprocal vectors in 1/Angstrom.

Returns:

kpoints_cart – Cartesian momenta in 1/Angstrom.

Return type:

Float[Array, 'n_k 3']

Notes

The function computes kpoints_frac @ geometry.reciprocal. Gradients flow through the coordinates and all applicable lattice entries.

diffpes.tightb.kspace.kpoints_cart_to_frac(kpoints_cart: Float[Array, 'n_k 3'], geometry: CrystalGeometry) → Float[Array, 'n_k 3'][source]

Convert Cartesian momenta to fractional k-points.

The function uses the closed reciprocal identity instead of a numerical linear solve.

See:

TestKpointsCartToFrac

Parameters:

  • kpoints_cart (Float[Array, 'n_k 3']) – Cartesian momenta in 1/Angstrom.

  • geometry (CrystalGeometry) – Crystal geometry with real-space vectors in Angstrom.

Returns:

kpoints_frac – Fractional k-points.

Return type:

Float[Array, 'n_k 3']

Notes

Reciprocal and real-space lattice rows satisfy \(B A^T=2\pi I\). Thus, the exact expression is kpoints_cart @ geometry.lattice.T / (2 * pi).

diffpes.tightb.kspace.build_kpath(anchors: Float[Array, 'n_anchor 3'], geometry: CrystalGeometry, n_per_segment: int, labels: tuple[str, ...], anchor_units: str = 'fractional') → KPath[source]

Build a labeled path between k-space anchors.

Each segment contains both endpoints. Adjacent segments therefore repeat their shared anchor, which matches the Chinook klib.kpath convention.

See:

TestBuildKpath

Implementation Logic

  1. Convert the anchors:

    fractional_anchors = kpoints_cart_to_frac(anchors, geometry)

    This branch applies only to anchors with absolute units.

  2. Interpolate all segments:

    segments = starts + fractions * (ends - starts)

    A broadcast creates every segment with a fixed output shape.

  3. Create the path carrier:

    return make_kpath(...)

    Static indices identify each anchor in the flattened point array.

param anchors:

Path anchors in the coordinates selected by anchor_units.

type anchors:

Float[Array, 'n_anchor 3']

param geometry:

Crystal geometry for absolute-to-fractional conversion.

type geometry:

CrystalGeometry

param n_per_segment:

Points in each segment, including both endpoints. This value is static. A change causes retracing.

type n_per_segment:

int

param labels:

One label for each anchor. This value is static. A change causes retracing.

type labels:

tuple[str, ...]

param anchor_units:

Static coordinate selector. Use "fractional" or "absolute". A change causes retracing. Default is "fractional".

type anchor_units:

str, default: "fractional"

returns:

kpath – Fractional path with (n_anchor - 1) * n_per_segment points.

rtype:

KPath

raises ValueError:

If the path has fewer than two anchors or invalid static settings.

raises EquinoxRuntimeError:

If an anchor is non-finite.

Notes

"absolute" means Cartesian momentum in 1/Angstrom. The interpolation remains differentiable with respect to anchors and the lattice.

diffpes.tightb.kspace.kpath_arc_length(kpath: KPath, geometry: CrystalGeometry) → Float[Array, 'n_k'][source]

Compute cumulative Cartesian distance along a k-path.

The plotting coordinate measures each segment after reciprocal-lattice conversion. Repeated junction points add zero distance.

See:

TestKpathArcLength

Parameters:

  • kpath (KPath) – Fractional path and its static plotting metadata.

  • geometry (CrystalGeometry) – Crystal geometry with reciprocal vectors in 1/Angstrom.

Returns:

arc_length – Cumulative Cartesian path distance in 1/Angstrom.

Return type:

Float[Array, 'n_k']

Notes

diffpes.maths.safe_norm() gives a zero gradient for a repeated junction. All nonzero segment lengths retain their usual gradients.

diffpes.tightb.kspace.first_bz_mask(kpoints_cart: Float[Array, 'n_k 3'], geometry: CrystalGeometry, shell_radius: int = 2) → Bool[Array, 'n_k'][source]

Mark Cartesian points inside the first Brillouin zone.

The Wigner-Seitz test compares the origin distance with reciprocal-lattice points in a static integer shell. A singular-value bound proves that no point outside the shell can change a retained membership result.

See:

TestFirstBzMask

Parameters:

  • kpoints_cart (Float[Array, 'n_k 3']) – Cartesian momenta in 1/Angstrom.

  • geometry (CrystalGeometry) – Crystal geometry with reciprocal vectors in 1/Angstrom.

  • shell_radius (int, default: 2) – Static maximum absolute reciprocal coefficient. Increase this value for a skew, anisotropic, or unreduced basis. Default is 2.

Returns:

mask – True for points inside or on the first-zone boundary.

Return type:

Bool[Array, 'n_k']

Raises:

  • ValueError – If shell_radius is less than one.

  • EquinoxRuntimeError – If the static shell is not provably complete for the supplied points and reciprocal basis.

Notes

The comparisons use squared distances and include ties. The Boolean mask has no boundary derivative. Consumers must not differentiate through its discrete membership changes.

For any unseen reciprocal vector G = n @ B, its norm is at least sigma_min(B) * norm(n). A vector can beat the origin only when norm(G) <= 2 * norm(k). The function checks this sufficient bound for every point that survives the requested shell. It raises instead of returning an uncertified false positive when the bound is inconclusive.

diffpes.tightb.kspace.build_bz_mesh(geometry: CrystalGeometry, n_per_axis: int, shell_radius: int = 2) → tuple[KGrid, Bool[Array, 'n_k']][source]

Build a fixed-shape reciprocal mesh and its first-zone mask.

The mesh samples the fractional cube from -1 to 1. A basis-derived bound must prove that this cube contains the complete first zone.

See:

TestBuildBzMesh

Implementation Logic

  1. Build the fractional cube:

    mesh_x, mesh_y, mesh_z = jnp.meshgrid(axis, axis, axis)

    Static axis lengths give the flattened point array a fixed shape.

  2. Compute zone membership:

    mask = first_bz_mask(kpoints_cart, geometry)

    The mask preserves the fixed point count during compiled execution.

  3. Return the grid and mask:

    return kgrid, mask

    The grid stores the three-dimensional mesh as rows of z coordinates.

param geometry:

Crystal geometry for reciprocal conversion.

type geometry:

CrystalGeometry

param n_per_axis:

Number of samples on each fractional axis. This value is static. A change causes retracing.

type n_per_axis:

int

param shell_radius:

Static reciprocal-coefficient radius for the first-zone test. Increase this value for a skew, anisotropic, or unreduced basis. Default is 2.

type shell_radius:

int, default: 2

rtype:

tuple[KGrid, Bool[Array, 'n_k']]

returns:

  • kgrid (KGrid) – Fractional mesh with n_per_axis**3 flattened points.

  • mask (Bool[Array, " n_k"]) – First-zone membership for every mesh point.

raises ValueError:

If n_per_axis is less than two or shell_radius is less than one.

raises EquinoxRuntimeError:

If the reciprocal basis does not prove that the fractional cube contains the first zone. The function also raises if the shell is not provably complete for the provisional first-zone points.

Notes

The static mesh shape is (n_per_axis**2, n_per_axis). Each row follows the z axis for one pair of x and y coordinates.

The coverage guard uses the inequalities abs(k dot b_i) <= norm(b_i)**2 / 2 that every first-zone point obeys. It maps their coordinate-wise bounds through the reciprocal Gram matrix. A failed conservative bound asks the caller for a reduced reciprocal basis instead of returning a partial zone.

diffpes.tightb.kspace.build_arpes_kmesh(kx_axis_inv_ang: Float[Array, 'n_kx'], ky_axis_inv_ang: Float[Array, 'n_ky'], kz_inv_ang: float | Float[Array, ''], sample_azimuth: float | Float[Array, ''], geometry: CrystalGeometry) → KGrid[source]

Build a fixed-kz ARPES raster in fractional coordinates.

The input axes describe the laboratory frame. A negative sample azimuth rotates that raster into the crystal frame before fractional conversion.

See:

TestBuildArpesKmesh

Implementation Logic

  1. Build the laboratory raster:

    lab_x, lab_y = jnp.meshgrid(kx_axis, ky_axis)

    The result uses rows for y and columns for x.

  2. Apply the sample rotation:

    sample_x = cosine * lab_x + sine * lab_y

    This expression applies an active rotation by negative azimuth.

  3. Convert and return the grid:

    return make_kgrid(kpoints_frac, mesh_shape, kz)

    The fixed kz value remains a traced carrier leaf.

param kx_axis_inv_ang:

Laboratory x momenta in 1/Angstrom.

type kx_axis_inv_ang:

Float[Array, 'n_kx']

param ky_axis_inv_ang:

Laboratory y momenta in 1/Angstrom.

type ky_axis_inv_ang:

Float[Array, 'n_ky']

param kz_inv_ang:

Fixed out-of-plane momentum in 1/Angstrom.

type kz_inv_ang:

Union[float, Float[Array, '']]

param sample_azimuth:

Sample rotation about the surface normal in radians.

type sample_azimuth:

Union[float, Float[Array, '']]

param geometry:

Crystal geometry for Cartesian-to-fractional conversion.

type geometry:

CrystalGeometry

returns:

kgrid – Fractional raster with static shape (n_ky, n_kx).

rtype:

KGrid

Notes

Gradients flow through both axes, kz, the azimuth, and the lattice. The raster shape depends only on the static input shapes.

diffpes.tightb.kspace.build_kmesh_hv(kpar_axis_inv_ang: Float[Array, 'n_kpar'], photon_energies_ev: Float[Array, 'n_hv'], work_function_ev: float | Float[Array, ''], inner_potential_ev: float | Float[Array, ''], sample_azimuth: float | Float[Array, ''], kpar_direction: Float[Array, '2'], geometry: CrystalGeometry) → KGrid[source]

Build a photon-energy raster in fractional coordinates.

Each row contains one photon energy and all requested parallel momenta. The free-electron final-state model supplies the row-dependent kz.

See:

TestBuildKmeshHv

Implementation Logic

  1. Compute each out-of-plane row:

    kz_rows = jax.vmap(kz_for_energy)(photon_energies_ev)

    One vectorized map preserves the dense photon-energy axis.

  2. Build and rotate the Cartesian raster:

    sample_x = cosine * lab_x + sine * lab_y

    The traced azimuth maps the laboratory direction into the sample.

  3. Convert and return the grid:

    return make_kgrid(..., photon_energy_axis_ev=photon_energies_ev)

    The carrier records one photon energy for each raster row.

param kpar_axis_inv_ang:

Signed parallel momenta in 1/Angstrom.

type kpar_axis_inv_ang:

Float[Array, 'n_kpar']

param photon_energies_ev:

Photon energies in eV.

type photon_energies_ev:

Float[Array, 'n_hv']

param work_function_ev:

Work function in eV.

type work_function_ev:

Union[float, Float[Array, '']]

param inner_potential_ev:

Inner potential in eV.

type inner_potential_ev:

Union[float, Float[Array, '']]

param sample_azimuth:

Sample rotation about the surface normal in radians.

type sample_azimuth:

Union[float, Float[Array, '']]

param kpar_direction:

Unit direction in the laboratory surface plane.

type kpar_direction:

Float[Array, '2']

param geometry:

Crystal geometry for Cartesian-to-fractional conversion.

type geometry:

CrystalGeometry

returns:

kgrid – Fractional raster with static shape (n_hv, n_kpar).

rtype:

KGrid

raises EquinoxRuntimeError:

If the direction is not finite and unit length. The function also rejects a raster that contains an evanescent channel.

Notes

The KGrid carrier stores no scalar kz because each row has a different value. The third coordinate of each point stores that value.