Orbital-matrix losses

Block-sparse matrix losses compare predictions and references with the same stored atom-pair layout. Direct errors can include all stored blocks or select onsite or offsite blocks. For Hermitian half-storage, their normalization refers to the stored values; omitted conjugate partners are not counted again.

The orthogonalized loss instead transforms the matrix difference to chosen k-points and uses the overlap to express it in an orthonormal basis. This gives a different comparison from averaging errors over the stored orbital matrix elements.

Stored matrix elements

Each sample's selected matrix elements and additional components are averaged before reduction across samples.

block_sparse_mae_loss

block_sparse_mae_loss(
    input: BlockSparseOrbData,
    target: BlockSparseOrbData,
    *,
    blocks: _BlockSelection = "all",
    reduction: _Reduction = "mean",
) -> Tensor

Return the mean absolute error of each block-sparse orbital matrix.

\[ \ell = \frac{1}{N_c\sum_{q\in\mathcal B_s} N_{\mathrm{orb},i_q}N_{\mathrm{orb},j_q}} \sum_{q\in\mathcal B_s}\sum_{\mu,\nu,c} \left|\widehat X^q_{\mu\nu c} - X^q_{\mu\nu c}\right|. \]

The set \(\mathcal B_s\) contains all stored blocks by default. With blocks="onsite", it contains blocks whose row and column atoms are equal and whose cell shift is zero; blocks="offsite" selects the complement. For complex values, the absolute value is the complex modulus. Here, \(N_c\) is the number of additional value components. Each stored value has equal weight within one sample, while reduction="mean" gives every sample equal weight. For Hermitian half-storage, this loss measures the stored values; it does not reconstruct or double-count the omitted conjugate blocks. Each sample must contain at least one selected matrix element.

Note

This elementwise L1 error is generally not invariant under orbital-basis rotations.

Parameters:

input (BlockSparseOrbData) –

values_real: [sum(num_values), *extra_shape]. values_imag: Optional [sum(num_values), *extra_shape]. num_values: [B].
target (BlockSparseOrbData) –

Target data with the same stored block layout and field shapes as input.
blocks (_BlockSelection, default: 'all' ) –

Stored blocks included in the loss.
reduction (_Reduction, default: 'mean' ) –

"none" returns one value per sample; "mean" or "sum" reduces those sample losses.

block_sparse_mse_loss

block_sparse_mse_loss(
    input: BlockSparseOrbData,
    target: BlockSparseOrbData,
    *,
    blocks: _BlockSelection = "all",
    reduction: _Reduction = "mean",
) -> Tensor

Return the mean squared error of each block-sparse orbital matrix.

\[ \ell = \frac{1}{N_c\sum_{q\in\mathcal B_s} N_{\mathrm{orb},i_q}N_{\mathrm{orb},j_q}} \sum_{q\in\mathcal B_s}\sum_{\mu,\nu,c} \left|\widehat X^q_{\mu\nu c} - X^q_{\mu\nu c}\right|^2. \]

The set \(\mathcal B_s\) contains all stored blocks by default. With blocks="onsite", it contains blocks whose row and column atoms are equal and whose cell shift is zero; blocks="offsite" selects the complement. For complex values, the squared absolute value is the sum of the squared real and imaginary errors. Here, \(N_c\) is the number of additional value components. Each stored value has equal weight within one sample, while reduction="mean" gives every sample equal weight. For Hermitian half-storage, this loss measures the stored values; it does not reconstruct or double-count the omitted conjugate blocks. No factor of \(1/2\) is applied. Each sample must contain at least one selected matrix element.

Parameters:

input (BlockSparseOrbData) –

values_real: [sum(num_values), *extra_shape]. values_imag: Optional [sum(num_values), *extra_shape]. num_values: [B].
target (BlockSparseOrbData) –

Target data with the same stored block layout and field shapes as input.
blocks (_BlockSelection, default: 'all' ) –

Stored blocks included in the loss.
reduction (_Reduction, default: 'mean' ) –

"none" returns one value per sample; "mean" or "sum" reduces those sample losses.

Orthogonalized Bloch matrices

This calculation combines a Bloch transform with Cholesky orthogonalization. The overlap must be positive-definite at every selected k-point.

block_sparse_orthogonalized_mse_loss

block_sparse_orthogonalized_mse_loss(
    input: BlockSparseOrbData,
    target: BlockSparseOrbData,
    overlap: BlockSparseOrbData,
    orb_counts: Tensor,
    batch_ptr: Tensor,
    kpoints: Tensor,
    *,
    reduction: _Reduction = "mean",
) -> Tensor

Return the MSE between Cholesky-orthogonalized Bloch matrices.

\[ \ell = \frac{1}{N_k N_{\mathrm{orb}}^2 N_c} \sum_{\mathbf k\in\mathcal K}\sum_c \left\| U(\mathbf k)^{-\dagger} \Delta X_c(\mathbf k) U(\mathbf k)^{-1} \right\|_F^2, \qquad S(\mathbf k)=U(\mathbf k)^\dagger U(\mathbf k). \]

Here, \(\Delta X(\mathbf k)\) is the Bloch transform of input - target, and \(S(\mathbf k)\) is the Bloch transform of overlap. Thus the transformed difference equals the difference between input and target after both are mapped to the same orthonormal basis. The overlap must be positive-definite at every k-point; failed Cholesky factorization is reported by PyTorch. Samples are grouped by their exact orbital dimension, and the loss does not pad matrices.

Here, \(N_{\mathrm{orb}}\) is the sample's total number of orbitals, \(N_k=|\mathcal K|\), and \(N_c\) is the number of additional value components. Every dense matrix entry, k-point, and additional component has equal weight within one sample; reduction="mean" gives every sample equal weight. No factor of \(1/2\) is applied.

Parameters:

input (BlockSparseOrbData) –

values_real: [sum(num_values), *extra_shape]. values_imag: Optional [sum(num_values), *extra_shape]. num_values: [B].
target (BlockSparseOrbData) –

Target data with the same stored block layout and field shapes as input.
overlap (BlockSparseOrbData) –

Hermitian positive-definite overlap matrix with the same samples and no additional value axes.
orb_counts (Tensor) –

Number of orbitals on every atom, with shape [N_atom].
batch_ptr (Tensor) –

Sample boundaries on the atom axis, with shape [B+1].
kpoints (Tensor) –

One shared fractional k-point with shape [3] or multiple shared k-points with shape [N_kpoint, 3].
reduction (_Reduction, default: 'mean' ) –

"none" returns one value per sample; "mean" or "sum" reduces those sample losses.