"""
block_matrix.py
=================
In-memory containers for a *single* block matrix:
* **BlockMatrix** - raw sparse blocks (H, D, S, …) grouped by element pair.
* **IrrepsBlockData** - same data after change-of-basis to irrep vectors.
Both dataclasses keep:
* `atoms` : list[str] (global order)
* `pair_edges` : dict[key] -> (2,E_ab) long tensor with global indices
* `lookup` : dict[(i,j)] -> (key, local_idx) for O(1) access
* reference to the shared `BlockIrrepMapper`
Conversion between the two is a one-liner via `.to_vectors()` / `.to_blocks()`.
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Dict, Tuple, TYPE_CHECKING
from typing import Sequence
if TYPE_CHECKING:
from core.basis_converter import OpenMXE3NNConverter
import torch
from core.block_irrep_mapper import BlockIrrepMapper
from core.orbital_irrep_config import OrbitalIrrepConfig
PairKey = str # canonical "A-B"
# --------------------------------------------------------------------------- #
[docs]
@dataclass
class BlockMatrix:
atoms: Tuple[str, ...]
atom_counts: Dict[str, int]
pair_blocks: Dict[PairKey, torch.Tensor] # (E_ab, d_A, d_B)
pair_edges: Dict[PairKey, torch.Tensor] # (5, E_ab)
lookup: Dict[Tuple[int, int], Tuple[PairKey, int]]
orbital_cfg: OrbitalIrrepConfig
basis: str # "openmx" | "e3nn"
_coverage_masks: dict | None = None
@property
def coverage_masks(self):
return (
self._coverage_masks
if self._coverage_masks is not None
else {
key: torch.ones_like(value, dtype=torch.bool)
for key, value in self.pair_blocks.items()
}
)
# --------------- convenience constructors ------------------------------ #
@classmethod
def empty(cls, atoms, orbital_cfg, basis="e3nn"):
from collections import Counter
return cls(tuple(atoms), Counter(atoms), {}, {}, {}, orbital_cfg, basis)
# --------------- dict-like access -------------------------------------- #
def __getitem__(self, item):
# item = (i, j) global indices, sum periodic images
if isinstance(item, tuple) and len(item) == 2:
print("Warning: BlockMatrix.__getitem__((i,j)) sums over periodic images.")
ret = None
for key in self.lookup:
# key is (sx, sy, sz, i, j)
if key[3] == item[0] and key[4] == item[1]:
k_key, k_idx = self.lookup[key]
block = self.pair_blocks[k_key][k_idx]
if ret is None:
ret = block.clone()
else:
ret += block
if ret is None:
raise KeyError(f"No blocks found for indices {item}")
return ret
# item = (sx, sy, sz, i, j) global indices
if isinstance(item, tuple) and len(item) == 5:
key, k = self.lookup[item]
return self.pair_blocks[key][k]
# item = "A-B"
if isinstance(item, str):
return self.pair_blocks[item]
raise KeyError("use (i,j), (sx,sy,sz,i,j) or 'A-B'")
def edges(self, key: PairKey) -> torch.Tensor:
return self.pair_edges[key]
def keys(self):
return self.pair_blocks.keys()
# ─────────────────────────────────────────────────────────────────────────
# helpers to access diagonal / off-diagonal blocks
# ─────────────────────────────────────────────────────────────────────────
[docs]
def diag(self) -> Dict[PairKey, torch.Tensor]:
"""
Return a **dict** ``key -> tensor`` that contains *only the blocks whose
global source/target atom are identical* (self-edges).
Off-diagonal blocks are omitted; keys that would become empty disappear.
"""
diag_dict: Dict[PairKey, torch.Tensor] = {}
for key, blk in self.pair_blocks.items():
el_a, el_b = key.split("-")
if el_a != el_b:
continue
n_atoms = self.atom_counts[el_a]
diag_dict[key] = blk[:n_atoms]
return diag_dict
[docs]
def offdiag(self) -> Dict[PairKey, torch.Tensor]:
"""
Same as :meth:`diag` but returns the **off-diagonal** blocks
( *i* ≠ *j* ).
"""
off_dict: Dict[PairKey, torch.Tensor] = {}
for key, blk in self.pair_blocks.items():
el_a, el_b = key.split("-")
if el_a == el_b:
n_atoms = self.atom_counts[el_a]
off_dict[key] = blk[n_atoms:]
else:
off_dict[key] = blk
return off_dict
# --------------- device handling --------------------------------------- #
def to(self, device, non_blocking: bool = False):
new_blocks = {
k: v.to(device, non_blocking=non_blocking)
for k, v in self.pair_blocks.items()
}
new_edges = {
k: v.to(device, non_blocking=non_blocking)
for k, v in self.pair_edges.items()
}
return BlockMatrix(
self.atoms,
self.atom_counts,
new_blocks,
new_edges,
self.lookup,
self.orbital_cfg,
self.basis,
_coverage_masks={
k: v.to(device, non_blocking=non_blocking)
for k, v in self.coverage_masks.items()
},
)
def pin_memory(self) -> "BlockMatrix":
pinned_blocks = {k: v.pin_memory() for k, v in self.pair_blocks.items()}
pinned_edges = {k: v.pin_memory() for k, v in self.pair_edges.items()}
return BlockMatrix(
self.atoms,
self.atom_counts,
pinned_blocks,
pinned_edges,
self.lookup,
self.orbital_cfg,
self.basis,
_coverage_masks={k: v.pin_memory() for k, v in self.coverage_masks.items()},
)
def detach(self) -> "BlockMatrix":
detached_blocks = {k: v.detach() for k, v in self.pair_blocks.items()}
detached_edges = {k: v.detach() for k, v in self.pair_edges.items()}
return BlockMatrix(
self.atoms,
self.atom_counts,
detached_blocks,
detached_edges,
self.lookup,
self.orbital_cfg,
self.basis,
_coverage_masks=self.coverage_masks,
)
# --------------- change-of-basis --------------------------------------- #
[docs]
def to_vectors(self, mapper: BlockIrrepMapper, *, layout=None) -> IrrepsBlockData:
"""Convert blocks to irreducible representation vectors using a mapper."""
pair_vec: Dict[PairKey, torch.Tensor] = {}
for key, blk in self.pair_blocks.items():
pair_vec[key] = mapper.blocks_to_vectors(key, blk, layout=layout)
return IrrepsBlockData(
self.atoms,
self.atom_counts,
pair_vec,
self.pair_edges,
self.lookup,
self.orbital_cfg,
self.basis,
layout=layout,
)
# ------------------------------------------------------------------ transpose
[docs]
def transpose(self) -> "BlockMatrix":
"""
Return a **new** snapshot representing the transposed matrix.
"""
# 1. Perform a simple transpose, flipping keys and edges.
transposed_blocks = {}
transposed_edges = {}
for key, blk in self.pair_blocks.items():
el_a, el_b = key.split("-")
new_key = f"{el_b}-{el_a}"
transposed_blocks[new_key] = blk.transpose(-1, -2).clone()
edges = self.pair_edges[key]
transposed_edges[new_key] = torch.stack(
[-edges[0], -edges[1], -edges[2], edges[4], edges[3]]
)
final_blocks = {}
final_edges = {}
# 2. For each key in the transposed matrix, restore original order if the key
# existed in the original matrix. Otherwise, create a canonical order.
for key, t_edges in transposed_edges.items():
if key in self.pair_edges:
# This key existed in the original matrix. We should match its edge order.
original_edges = self.pair_edges[key]
# Build a map from a transposed edge to its current index.
map_edge_to_idx = {
tuple(t_edges[:, i].tolist()): i for i in range(t_edges.shape[1])
}
# Create permutation by looking up original edges in the map.
try:
perm = torch.tensor(
[
map_edge_to_idx[tuple(original_edges[:, i].tolist())]
for i in range(original_edges.shape[1])
]
)
final_edges[key] = t_edges[:, perm]
final_blocks[key] = transposed_blocks[key][perm]
except KeyError:
raise Exception(
f"Key '{key}' not found in the transposed matrix; matrix not symmetric."
)
else:
raise Exception(
f"Key '{key}' not found in the original matrix; matrix not symmetric."
)
# Rebuild lookup
new_lookup = {}
for key, edges in final_edges.items():
for idx, (sx, sy, sz, i, j) in enumerate(edges.t().tolist()):
new_lookup[(sx, sy, sz, i, j)] = (key, idx)
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=final_blocks,
pair_edges=final_edges,
lookup=new_lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
def transpose_aligned(
self, reverse_alignment: Dict[str, Tuple[str, torch.Tensor]]
) -> "BlockMatrix":
"""
Return the transpose assuming reverse-edge permutations are already known.
``reverse_alignment[key] = (rev_key, perm)`` means that
``self.pair_blocks[rev_key][perm]`` matches the edge order of ``key`` after
applying the reverse-edge transform and transposing each block.
"""
final_blocks = {}
for key, blk in self.pair_blocks.items():
if key not in reverse_alignment:
raise ValueError(f"Missing reverse alignment for key '{key}'")
rev_key, perm = reverse_alignment[key]
if rev_key not in self.pair_blocks:
raise ValueError(
f"Reverse key '{rev_key}' missing in BlockMatrix (for key '{key}')"
)
rev_blocks = self.pair_blocks[rev_key]
if perm.device != rev_blocks.device:
perm = perm.to(rev_blocks.device)
final_blocks[key] = rev_blocks.index_select(0, perm).transpose(-1, -2)
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=final_blocks,
pair_edges=self.pair_edges,
lookup=self.lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
def __mul__(self, scalar: float) -> "BlockMatrix":
"""
Scalar multiplication of all blocks by a float or int.
"""
if not isinstance(scalar, (float, int)):
return NotImplemented
new_blocks = {k: v * scalar for k, v in self.pair_blocks.items()}
return self._replace_pair_blocks(new_blocks, basis=self.basis)
def __rmul__(self, scalar: float) -> "BlockMatrix":
"""
Reflected scalar multiplication so ``scalar * BlockMatrix`` works.
"""
return self.__mul__(scalar)
# ------------------------------------------------------------------ edge reordering
def reorder_edges(self, order_dict: Dict[str, torch.Tensor]) -> "BlockMatrix":
"""
Re-order edge *rows* for given keys. ``order_dict`` maps
``key -> permutation indices`` (1-D LongTensor of length ``E_key``).
Keys **not** in ``order_dict`` keep their original order.
"""
pair_blocks, pair_edges, lookup = {}, {}, {}
for key, blk in self.pair_blocks.items():
if key in order_dict:
idx = order_dict[key]
blk = blk[idx]
edges = self.pair_edges[key][:, idx]
else:
edges = self.pair_edges[key]
pair_blocks[key] = blk
pair_edges[key] = edges
for new_k, (sx, sy, sz, i, j) in enumerate(edges.t().tolist()):
lookup[(sx, sy, sz, i, j)] = (key, new_k)
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=pair_blocks,
pair_edges=pair_edges,
lookup=lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
_coverage_masks={
k: mask[order_dict[k]] if k in order_dict else mask
for k, mask in self.coverage_masks.items()
},
)
def relabel_atom_images(self, atom_image_shifts: torch.Tensor) -> "BlockMatrix":
"""
Relabel edge shifts to a different per-atom periodic image convention.
``atom_image_shifts[i]`` is the integer lattice shift that maps atom ``i``
from the old convention to the new one.
"""
if atom_image_shifts.ndim != 2 or atom_image_shifts.shape[1] != 3:
raise ValueError(
"atom_image_shifts must have shape (n_atoms, 3), got "
f"{tuple(atom_image_shifts.shape)}"
)
if atom_image_shifts.shape[0] != len(self.atoms):
raise ValueError(
"atom_image_shifts length must match number of atoms: "
f"{atom_image_shifts.shape[0]} != {len(self.atoms)}"
)
atom_image_shifts = atom_image_shifts.to(dtype=torch.long)
relabeled_blocks = {k: v for k, v in self.pair_blocks.items()}
relabeled_edges: Dict[PairKey, torch.Tensor] = {}
relabeled_lookup: Dict[Tuple[int, int], Tuple[PairKey, int]] = {}
for key, edges in self.pair_edges.items():
new_edges = edges.clone()
src = edges[3].to(dtype=torch.long)
dst = edges[4].to(dtype=torch.long)
shift_delta = atom_image_shifts[src] - atom_image_shifts[dst]
new_edges[:3] = new_edges[:3] + shift_delta.T
relabeled_edges[key] = new_edges
for idx, (sx, sy, sz, i, j) in enumerate(new_edges.t().tolist()):
relabeled_lookup[(sx, sy, sz, i, j)] = (key, idx)
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=relabeled_blocks,
pair_edges=relabeled_edges,
lookup=relabeled_lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
# ------------------------------------------------------------------ arithmetic
# private helper ------------------------------------------------------------
def _align_with(self, other: "BlockMatrix") -> Tuple["BlockMatrix", "BlockMatrix"]:
"""
Make sure the matrices have the same atoms, basis, orbital_cfg and edge order.
"""
if not isinstance(other, BlockMatrix):
raise TypeError("Operand must be BlockMatrix")
if self.atoms != other.atoms:
raise ValueError("Atoms differ; cannot add/subtract snapshots")
if self.basis != other.basis:
raise ValueError("Basis differs (openmx vs e3nn)")
if self.orbital_cfg.to_dict() != other.orbital_cfg.to_dict():
raise ValueError("OrbitalIrrepConfig differs")
if self.keys() != other.keys():
raise ValueError("Snapshots contain different element-pair keys")
reorder_dict = {}
for k in self.keys():
if not torch.equal(self.pair_edges[k], other.pair_edges[k]):
# If edge order differs, reorder self to match other.
map_edge_to_idx = {
tuple(self.pair_edges[k][:, i].tolist()): i
for i in range(self.pair_edges[k].shape[1])
}
try:
perm = torch.tensor(
[
map_edge_to_idx[tuple(other.pair_edges[k][:, i].tolist())]
for i in range(other.pair_edges[k].shape[1])
]
)
reorder_dict[k] = perm
except KeyError:
raise ValueError(f"Edge sets for key '{k}' differ between matrices")
if self.pair_blocks[k].shape != other.pair_blocks[k].shape:
raise ValueError(f"Shape mismatch for key '{k}'")
if reorder_dict:
reordered_self = self.reorder_edges(reorder_dict)
return reordered_self, other
else:
return self, other
# -------------- public dunder ops -----------------------------------------
def __add__(self, other):
if other == 0: # allow sum() with start=0
return self
a, b = self._align_with(other)
new_blocks = {k: a.pair_blocks[k] + b.pair_blocks[k] for k in a.pair_blocks}
return a._replace_pair_blocks(new_blocks, basis=a.basis)
def add_aligned(self, other: "BlockMatrix") -> "BlockMatrix":
"""
Fast addition assuming pair keys, edge order, and block shapes are already aligned.
"""
if not isinstance(other, BlockMatrix):
raise TypeError("Operand must be BlockMatrix")
if self.atoms != other.atoms:
raise ValueError("Atoms differ; cannot add aligned BlockMatrix objects")
if self.basis != other.basis:
raise ValueError("Basis differs (openmx vs e3nn)")
if self.orbital_cfg.to_dict() != other.orbital_cfg.to_dict():
raise ValueError("OrbitalIrrepConfig differs")
if self.keys() != other.keys():
raise ValueError("Snapshots contain different element-pair keys")
new_blocks = {}
for k in self.pair_blocks:
if self.pair_blocks[k].shape != other.pair_blocks[k].shape:
raise ValueError(f"Shape mismatch for key '{k}'")
new_blocks[k] = self.pair_blocks[k] + other.pair_blocks[k]
return self._replace_pair_blocks(new_blocks, basis=self.basis)
__radd__ = __add__ # commutative
def __neg__(self):
new_blocks = {k: -v for k, v in self.pair_blocks.items()}
return self._replace_pair_blocks(new_blocks, basis=self.basis)
def __sub__(self, other):
if other == 0:
return self
a, b = self._align_with(other)
new_blocks = {k: a.pair_blocks[k] - b.pair_blocks[k] for k in a.pair_blocks}
return a._replace_pair_blocks(new_blocks, basis=a.basis)
def sub_aligned(self, other: "BlockMatrix") -> "BlockMatrix":
"""
Fast subtraction assuming pair keys, edge order, and block shapes are already aligned.
"""
if not isinstance(other, BlockMatrix):
raise TypeError("Operand must be BlockMatrix")
if self.atoms != other.atoms:
raise ValueError(
"Atoms differ; cannot subtract aligned BlockMatrix objects"
)
if self.basis != other.basis:
raise ValueError("Basis differs (openmx vs e3nn)")
if self.orbital_cfg.to_dict() != other.orbital_cfg.to_dict():
raise ValueError("OrbitalIrrepConfig differs")
if self.keys() != other.keys():
raise ValueError("Snapshots contain different element-pair keys")
new_blocks = {}
for k in self.pair_blocks:
if self.pair_blocks[k].shape != other.pair_blocks[k].shape:
raise ValueError(f"Shape mismatch for key '{k}'")
new_blocks[k] = self.pair_blocks[k] - other.pair_blocks[k]
return self._replace_pair_blocks(new_blocks, basis=self.basis)
def conjugate(self):
return self._replace_pair_blocks(
{key: value.conj() for key, value in self.pair_blocks.items()},
basis=self.basis,
)
def adjoint(self):
return self.transpose().conjugate()
def symmetrize_aligned(
self, reverse_alignment: Dict[str, Tuple[str, torch.Tensor]]
) -> "BlockMatrix":
"""
Compute ``0.5 * (M + M^T)`` using precomputed reverse-edge permutations.
"""
return (
self.add_aligned(self.transpose_aligned(reverse_alignment).conjugate())
* 0.5
)
def __rsub__(self, other):
# allow (0 - BlockMatrix)
if other == 0:
return -self
return NotImplemented
# --------------------- dense helpers --------------------------------------
# global dim offsets ---------------------------------------------------------
def _physical_dimension(self, species):
dimensions = {
block.shape[-2 + position]
for pair, block in self.pair_blocks.items()
for position, atom in enumerate(pair.split("-"))
if atom == species
}
if len(dimensions) != 1:
raise ValueError(
f"Inconsistent or absent physical block dimension for {species}"
)
return dimensions.pop()
def _atom_offsets(self) -> Tuple[torch.Tensor, int]:
"""Return 1-D tensor of start indices per atom and total dimension."""
dims = torch.tensor(
[self._physical_dimension(el) for el in self.atoms],
dtype=torch.long,
)
offsets = torch.cumsum(
torch.cat([torch.tensor([0]), dims[:-1]]), dim=0
) # start indices
total_dim = int(offsets[-1] + dims[-1])
return offsets, total_dim
# --- public API -------------------------------------------------------------
[docs]
def to_dense(self) -> torch.Tensor:
"""
Assemble a full dense matrix of shape ``(Σ d_i, Σ d_i)`` where
``d_i`` is orbital dimension of atom *i*.
Summing all periodic images into the central cell.
"""
offsets, total = self._atom_offsets()
device = next(iter(self.pair_blocks.values())).device
dtype = next(iter(self.pair_blocks.values())).dtype
dense = torch.zeros(total, total, device=device, dtype=dtype)
for (_sx, _sy, _sz, i, j), (key, k) in self.lookup.items():
d_i, d_j = self.pair_blocks[key].shape[-2:]
r0 = int(offsets[i])
c0 = int(offsets[j])
dense[r0 : r0 + d_i, c0 : c0 + d_j] += self.pair_blocks[key][k]
return dense
# ------------------ serialisation ------------------------------------
def _to_payload(self) -> dict:
"""Plain python types + **CPU** tensors -> ready for torch.save."""
blocks_cpu = {k: v.detach().cpu() for k, v in self.pair_blocks.items()}
edges_cpu = {k: v.detach().cpu() for k, v in self.pair_edges.items()}
return {
"atoms": list(self.atoms),
"atom_counts": self.atom_counts,
"orbital_cfg": self.orbital_cfg.to_dict(),
"pair_blocks": blocks_cpu,
"pair_edges": edges_cpu,
"type": "block",
"basis": self.basis,
"coverage_masks": {
k: v.detach().cpu() for k, v in self.coverage_masks.items()
},
}
def save(self, path: str | "os.PathLike[str]") -> None:
import torch
torch.save(self._to_payload(), path)
# ------------------ alternate constructors ---------------------------
@classmethod
def load(cls, path, device="cpu") -> "BlockMatrix":
import torch
from core.orbital_irrep_config import OrbitalIrrepConfig
payload = torch.load(path, map_location="cpu", weights_only=False)
if payload.get("type") != "block":
raise ValueError("file does not contain block matrix")
orb_cfg = OrbitalIrrepConfig.from_dict(payload["orbital_cfg"])
pair_blocks = {k: v.to(device) for k, v in payload["pair_blocks"].items()}
pair_edges = {k: v.to(device) for k, v in payload["pair_edges"].items()}
# rebuild lookup
lookup = {}
for key, edges in pair_edges.items():
for idx, (sx, sy, sz, i, j) in enumerate(edges.t().tolist()):
lookup[(sx, sy, sz, i, j)] = (key, idx)
atoms = tuple(payload["atoms"])
from collections import Counter
atom_counts = payload.get("atom_counts", Counter(atoms))
return cls(
atoms,
atom_counts,
pair_blocks,
pair_edges,
lookup,
orb_cfg,
payload.get("basis", "openmx"),
_coverage_masks=payload.get("coverage_masks"),
).to(device)
# ------------- internal helper to clone with new blocks / basis ---------
def _replace_pair_blocks(
self, new_blocks: Dict[PairKey, torch.Tensor], *, basis: str
):
"""Return a shallow copy with *pair_blocks* replaced."""
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=new_blocks,
pair_edges=self.pair_edges,
lookup=self.lookup,
orbital_cfg=self.orbital_cfg,
basis=basis,
_coverage_masks=self._coverage_masks,
)
def _apply_edge_mask(
self,
mask_dict: Dict[PairKey, torch.Tensor | Sequence[bool]],
*,
drop_empty: bool = True,
) -> "BlockMatrix":
"""
Internal utility - return a **new** BlockMatrix where, for every
``key`` contained in ``mask_dict``, only the edges with a *True*
entry in the 1-D boolean mask are kept.
Keys **not** present in the dict are copied unchanged.
"""
pair_blocks, pair_edges, lookup = {}, {}, {}
for key, blk in self.pair_blocks.items():
if key in mask_dict:
mask = torch.as_tensor(
mask_dict[key], dtype=torch.bool, device=blk.device
)
if mask.ndim != 1 or mask.numel() != blk.shape[0]:
raise ValueError(f"Mask for '{key}' has wrong shape")
blk_kept = blk[mask]
edges_kept = self.pair_edges[key][:, mask]
else: # untouched
blk_kept = blk
edges_kept = self.pair_edges[key]
if blk_kept.shape[0] == 0 and drop_empty:
continue # drop key entirely
pair_blocks[key] = blk_kept
pair_edges[key] = edges_kept
for idx, (sx, sy, sz, i, j) in enumerate(edges_kept.t().tolist()):
lookup[(sx, sy, sz, i, j)] = (key, idx)
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=pair_blocks,
pair_edges=pair_edges,
lookup=lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
# ---------------- basis conversion wrappers ----------------------------
def to_e3nn(self, converter: "OpenMXE3NNConverter") -> "BlockMatrix":
matrix = converter.matrix_to_e3nn(self)
matrix.basis = "e3nn"
return matrix
def to_openmx(self, converter: "OpenMXE3NNConverter") -> "BlockMatrix":
matrix = converter.matrix_to_openmx(self)
matrix.basis = "openmx"
return matrix
def change_basis(self, d_dict: Dict[str, torch.Tensor]) -> "BlockMatrix":
"""
Change the basis of the block matrix using a dictionary of transformation matrices.
Each key in `d_dict` corresponds to an element symbol, and the value is a
transformation matrix that will be applied to the blocks associated with that element.
The transformation is applied as follows:
For a block corresponding to the pair (A-B):
.. math::
\\text{new\_block}_{AB} = d_{A} \\cdot \\text{block}_{AB} \\cdot d_{B}^T
where :math:`d_{A}` and :math:`d_{B}` are the transformation matrices for elements A and B,
respectively.
Parameters
----------
d_dict : Dict[str, torch.Tensor]
A dictionary mapping element symbols to transformation matrices.
Each matrix should have the shape (d_A, d_A) for element A.
Returns
-------
BlockMatrix
A new BlockMatrix instance with the transformed blocks.
Notes
-----
This method assumes that the transformation matrices in `d_dict` are square matrices
with dimensions matching the orbital dimensions of the respective elements.
"""
new_pair_blocks = {}
for key, blk in self.pair_blocks.items():
el_i, el_j = key.split("-")
new_pair_blocks[key] = d_dict[el_i] @ blk @ d_dict[el_j].T
# Return a new BlockMatrix with the updated blocks and unchanged edges
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=new_pair_blocks,
pair_edges=self.pair_edges,
lookup=self.lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
# ------------------------ alternate constructor -----------------------------
[docs]
@classmethod
def from_dense(
cls,
matrix: torch.Tensor,
orbital_cfg, # OrbitalIrrepConfig
atoms: Tuple[str, ...] | list[str],
*,
diagonal: bool = False,
sparsity_threshold: float = 0.0,
basis: str,
) -> "BlockMatrix":
"""
Build a block snapshot from a fully dense matrix *in global atom order*.
Primarily for tests / debugging.
"""
atoms = tuple(atoms)
from collections import Counter
atom_counts = Counter(atoms)
offsets = [0]
for el in atoms[:-1]:
d = orbital_cfg.block_dims(f"{el}-{el}")[0]
offsets.append(offsets[-1] + d)
offsets_t = torch.tensor(offsets, dtype=torch.long, device=matrix.device)
pair_blocks, pair_edges, lookup = {}, {}, {}
for i, el_i in enumerate(atoms):
d_i = orbital_cfg.block_dims(f"{el_i}-{el_i}")[0]
r0 = int(offsets_t[i])
for j, el_j in enumerate(atoms):
d_j = orbital_cfg.block_dims(f"{el_j}-{el_j}")[0]
c0 = int(offsets_t[j])
blk = matrix[r0 : r0 + d_i, c0 : c0 + d_j].clone()
key = f"{el_i}-{el_j}"
if key not in pair_blocks:
pair_blocks[key], pair_edges[key] = [], []
local_idx = len(pair_blocks[key])
pair_blocks[key].append(blk)
pair_edges[key].append([0, 0, 0, i, j])
lookup[(0, 0, 0, i, j)] = (key, local_idx)
print("Warning: from_dense currently assumes no periodic images!")
# stack
pair_blocks = {k: torch.stack(v) for k, v in pair_blocks.items()}
pair_edges = {
k: torch.tensor(v, dtype=torch.long).t() for k, v in pair_edges.items()
}
snapshot = cls(
atoms, atom_counts, pair_blocks, pair_edges, lookup, orbital_cfg, basis
)
if sparsity_threshold is not None:
snapshot = snapshot.sparsify(sparsity_threshold)
return snapshot
# ---------------------------------------------------------------- sparsify
[docs]
def sparsify(self, threshold: float) -> "BlockMatrix":
"""
Return a **new** snapshot in which only blocks whose root-mean-square
(RMS) magnitude is **strictly greater** than ``threshold`` are kept.
A block's RMS is computed as
``rms = sqrt( mean( block**2 ) )`` across its last two dimensions
(orbital rows & columns).
Parameters
----------
threshold
Cut-off applied *per block*. A Python float is accepted; it is
internally cast to the block's dtype and device.
Notes
-----
• Pruning is performed **per block** (first tensor dimension).
Keys for which *all* blocks are removed disappear entirely.
"""
if not self.pair_blocks:
return self # nothing to do
sample_blk = next(iter(self.pair_blocks.values()))
thr = torch.as_tensor(
threshold, dtype=sample_blk.dtype, device=sample_blk.device
)
new_blocks: Dict[PairKey, torch.Tensor] = {}
new_edges: Dict[PairKey, torch.Tensor] = {}
new_lookup: Dict[Tuple[int, int], Tuple[PairKey, int]] = {}
for key, blk in self.pair_blocks.items():
# (E, d_i, d_j) -> (E,)
rms = blk.pow(2).mean(dim=(-2, -1)).sqrt()
keep = rms > thr
if torch.any(keep):
blk_kept = blk[keep]
edges_kept = self.pair_edges[key][:, keep]
new_blocks[key] = blk_kept
new_edges[key] = edges_kept
# rebuild lookup for the surviving edges of this key
for local_idx, (sx, sy, sz, i, j) in enumerate(edges_kept.t().tolist()):
new_lookup[(sx, sy, sz, i, j)] = (key, local_idx)
# assemble the sparsified snapshot
return BlockMatrix(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_blocks=new_blocks,
pair_edges=new_edges,
lookup=new_lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
# ----------------------------------------------------------------- reload from payload
@classmethod
def from_payload(
cls, payload: dict, device: str | torch.device = "cpu"
) -> "BlockMatrix":
"""
Build :class:`BlockMatrix` from a dict previously produced by
:meth:`_to_payload`. Used internally by Snapshot.load().
"""
from core.orbital_irrep_config import OrbitalIrrepConfig # local import
orb_cfg = OrbitalIrrepConfig.from_dict(payload["orbital_cfg"])
pair_blocks = {k: v.to(device) for k, v in payload["pair_blocks"].items()}
pair_edges = {k: v.to(device) for k, v in payload["pair_edges"].items()}
# rebuild lookup table
lookup: Dict[Tuple[int, int], Tuple[str, int]] = {}
for key, edges in pair_edges.items():
for idx, (sx, sy, sz, i, j) in enumerate(edges.t().tolist()):
lookup[(sx, sy, sz, i, j)] = (key, idx)
atoms = tuple(payload["atoms"])
from collections import Counter
atom_counts = payload.get("atom_counts", Counter(atoms))
return cls(
atoms=atoms,
atom_counts=atom_counts,
pair_blocks=pair_blocks,
pair_edges=pair_edges,
lookup=lookup,
orbital_cfg=orb_cfg,
basis=payload.get("basis", "openmx"),
_coverage_masks={
k: v.to(device) for k, v in payload.get("coverage_masks", {}).items()
}
or None,
)
# ════════════════════════════════════════════════════════════════════════════
# ROTATION
# ═══════════════════════════════════════════════════════════════════════════
[docs]
def rotate(self, R: torch.Tensor) -> "BlockMatrix":
"""
Return a **new** :class:`BlockMatrix` whose *orbital reference frame*
has been rotated by the (3 x 3) matrix **R** (active rotation).
Notes
-----
* The operation is block-wise:
M'_(A,B) = U_A · M_(A,B) · U_Bᵀ
where ``U_A = irreps_A.D_from_matrix(R)`` is block-diagonal with one
Wigner-D copy per orbital.
"""
if R.shape != (3, 3):
raise ValueError("R must be a 3x3 rotation matrix")
device = next(iter(self.pair_blocks.values())).device
R = R.to(device=device, dtype=torch.float32)
# cache one U per element
U_cache: Dict[str, torch.Tensor] = {}
for el in self.orbital_cfg.elements():
irr = self.orbital_cfg.element_to_irreps[el]
U_cache[el] = irr.D_from_matrix(R) # (dim_el, dim_el)
# rotate every block ------------------------------------------------
new_blocks: Dict[PairKey, torch.Tensor] = {}
for key, blk in self.pair_blocks.items():
el_i, el_j = key.split("-")
U_i = U_cache[el_i]
U_j = U_cache[el_j]
# tensor contraction: (E,d_i,d_j)
blk_rot = U_i @ blk @ U_j.T
new_blocks[key] = blk_rot
# edge indices / lookup are unchanged
return self._replace_pair_blocks(new_blocks, basis=self.basis)
# --------------------------------------------------------------------------- #
[docs]
@dataclass
class IrrepsBlockData:
atoms: Tuple[str, ...]
atom_counts: Dict[str, int]
pair_vectors: Dict[PairKey, torch.Tensor] # (E_ab, n_vec_AB)
pair_edges: Dict[PairKey, torch.Tensor]
lookup: Dict[Tuple[int, int], Tuple[PairKey, int]]
orbital_cfg: OrbitalIrrepConfig
basis: str = "e3nn" # always "e3nn"
layout: object | None = None
packed_values: torch.Tensor | None = None
packed_edges: torch.Tensor | None = None
frames: torch.Tensor | None = None
operator: str | None = None
[docs]
@staticmethod
def pack_matrix(matrix, edges, layout, *, frames=None, operator=None):
"""Express physical blocks in a declared supervision frame and unit."""
rows = []
from utils.units import HARTREE_TO_EV
if matrix.basis != "openmx":
from core.basis_converter import OpenMXE3NNConverter
matrix = OpenMXE3NNConverter(matrix.orbital_cfg).matrix_to_openmx(matrix)
for index, edge in enumerate(edges.T.tolist()):
pair, row = matrix.lookup[tuple(edge)]
block = matrix.pair_blocks[pair][row]
atoms = tuple(pair.split("-"))
if frames is not None:
block = layout.rotation.rotate_blocks(
block,
frames[index],
layout.species_shells[atoms[0]],
layout.species_shells[atoms[1]],
)
if operator == "hamiltonian":
block = block * HARTREE_TO_EV
rows.append(layout.pack_blocks(atoms, block))
return torch.stack(rows)
[docs]
@classmethod
def from_packed(cls, values, x, mapper, layout, *, frame="global", operator=None):
"""Represent packed orbital channels through MANDALA's public container."""
from collections import Counter
species = tuple(layout.species_shells)
atoms = tuple(species[int(index)] for index in x["node_type_idx"])
edges = x.get("edge_keys")
if edges is None:
shifts = x.get(
"edge_shift", x["edge_index"].new_zeros((3, x["edge_index"].shape[1]))
)
edges = torch.cat((shifts, x["edge_index"]), dim=0)
pair_vectors, pair_edges, lookup = {}, {}, {}
for left in species:
for right in species:
selected = torch.tensor(
[
index
for index, edge in enumerate(edges.T)
if atoms[int(edge[3])] == left and atoms[int(edge[4])] == right
],
device=values.device,
dtype=torch.long,
)
if not len(selected):
continue
pair = f"{left}-{right}"
pair_vectors[pair] = layout.encode(values.index_select(0, selected))
pair_edges[pair] = edges.index_select(1, selected)
for local, global_index in enumerate(selected.tolist()):
lookup[tuple(edges[:, global_index].tolist())] = (pair, local)
frames = x.get("bond_frames") if frame == "bond" else None
if frame == "bond" and frames is None:
raise ValueError("Bond-frame output requires explicit frames")
return cls(
atoms,
Counter(atoms),
pair_vectors,
pair_edges,
lookup,
mapper.orbital_cfg,
"openmx",
layout,
values,
edges,
frames,
operator,
)
# -------- device -------- #
def _map_tensors(self, operation):
from dataclasses import replace
return replace(
self,
pair_vectors={
key: operation(value) for key, value in self.pair_vectors.items()
},
pair_edges={
key: operation(value) for key, value in self.pair_edges.items()
},
packed_values=(
operation(self.packed_values)
if self.packed_values is not None
else None
),
packed_edges=(
operation(self.packed_edges) if self.packed_edges is not None else None
),
frames=operation(self.frames) if self.frames is not None else None,
)
def to(self, device, non_blocking=False):
return self._map_tensors(
lambda value: value.to(device, non_blocking=non_blocking)
)
def pin_memory(self):
return self._map_tensors(lambda value: value.pin_memory())
def detach(self):
return self._map_tensors(lambda value: value.detach())
# -------- inverse change-of-basis ----- #
def to_blocks(self, mapper: BlockIrrepMapper) -> BlockMatrix:
pair_blk: Dict[PairKey, torch.Tensor] = {}
for key, vec in self.pair_vectors.items():
pair_blk[key] = mapper.vectors_to_blocks(key, vec, layout=self.layout)
if self.frames is not None:
global_rows = {
tuple(edge.tolist()): index
for index, edge in enumerate(self.packed_edges.T)
}
frames = torch.stack(
[
self.frames[global_rows[tuple(edge.tolist())]]
for edge in self.pair_edges[key].T
]
)
left, right = key.split("-")
pair_blk[key] = self.layout.rotation.rotate_blocks(
pair_blk[key],
frames.transpose(-1, -2),
self.layout.species_shells[left],
self.layout.species_shells[right],
)
if self.layout is not None and self.operator == "hamiltonian":
from utils.units import HARTREE_TO_EV
pair_blk[key] = pair_blk[key] / HARTREE_TO_EV
result = BlockMatrix(
self.atoms,
self.atom_counts,
pair_blk,
self.pair_edges,
self.lookup,
self.orbital_cfg,
self.basis,
)
if self.layout is not None:
result._coverage_masks = {
pair: self.layout.unpack_blocks(
tuple(pair.split("-")),
self.layout.packed_mask(
tuple(pair.split("-")), device=block.device
),
).expand_as(block)
for pair, block in pair_blk.items()
}
return result
# ─────────────────────────────────────────────────────────────────────────
# diag / offdiag access for vector form
# ─────────────────────────────────────────────────────────────────────────
def diag(self) -> Dict[PairKey, torch.Tensor]:
diag_dict: Dict[PairKey, torch.Tensor] = {}
for key, vec in self.pair_vectors.items():
el_a, el_b = key.split("-")
if el_a != el_b:
continue
n_atoms = self.atom_counts[el_a]
diag_dict[key] = vec[:n_atoms]
return diag_dict
def offdiag(self) -> Dict[PairKey, torch.Tensor]:
off_dict: Dict[PairKey, torch.Tensor] = {}
for key, vec in self.pair_vectors.items():
el_a, el_b = key.split("-")
if el_a == el_b:
n_atoms = self.atom_counts[el_a]
off_dict[key] = vec[n_atoms:]
else:
off_dict[key] = vec
return off_dict
# -------- indexing paralleling BlockMatrix -------- #
def __getitem__(self, item):
# item = (i, j) global indices, sum periodic images
if isinstance(item, tuple) and len(item) == 2:
print(
"Warning: IrrepsBlockData.__getitem__((i,j)) sums over periodic images."
)
item_i, item_j = item
ret = None
for key in self.lookup:
sx, sy, sz, i, j = key
if i == item_i and j == item_j:
k_key, k_idx = self.lookup[key]
vec = self.pair_vectors[k_key][k_idx]
if ret is None:
ret = vec.clone()
else:
ret += vec
if ret is None:
raise KeyError(f"No vectors found for indices {item}")
return ret
# item = (sx, sy, sz, i, j) global indices
if isinstance(item, tuple) and len(item) == 5:
key, k = self.lookup[item]
return self.pair_vectors[key][k]
# item = "A-B"
if isinstance(item, str):
return self.pair_vectors[item]
raise KeyError("use (i,j), (sx,sy,sz,i,j) or 'A-B'")
# ════════════════════════════════════════════════════════════════════════════
# ROTATION
# ═══════════════════════════════════════════════════════════════════════════
[docs]
def rotate(self, R: torch.Tensor, mapper: BlockIrrepMapper) -> "IrrepsBlockData":
"""
Return a **new** :class:`IrrepsBlockData` whose *orbital reference frame*
has been rotated by the (3 x 3) matrix **R** (active rotation).
Parameters
----------
R
The (3 x 3) rotation matrix.
mapper
The BlockIrrepMapper that provides the pair irreps for computing
Wigner-D matrices.
Notes
-----
* For irrep vectors representing the tensor product of orbitals from
atoms A and B, we apply the Wigner-D matrix for the output irreps:
v'_(A,B) = D_out(R) · v_(A,B)
where ``D_out = pair_irreps.D_from_matrix(R)`` acts on the tensor
product irreps.
"""
if R.shape != (3, 3):
raise ValueError("R must be a 3x3 rotation matrix")
device = next(iter(self.pair_vectors.values())).device
R = R.to(device=device, dtype=torch.float32)
# cache one Wigner-D matrix per pair type
D_cache: Dict[PairKey, torch.Tensor] = {}
for key in self.pair_vectors.keys():
pair_irreps = mapper.get_pair_irreps(key)
D_cache[key] = pair_irreps.D_from_matrix(R) # (n_vec, n_vec)
# rotate every vector ------------------------------------------------
new_vectors: Dict[PairKey, torch.Tensor] = {}
for key, vec in self.pair_vectors.items():
D = D_cache[key]
# tensor contraction: (E, n_vec)
vec_rot = vec @ D.T
new_vectors[key] = vec_rot
# edge indices / lookup are unchanged
return IrrepsBlockData(
atoms=self.atoms,
atom_counts=self.atom_counts,
pair_vectors=new_vectors,
pair_edges=self.pair_edges,
lookup=self.lookup,
orbital_cfg=self.orbital_cfg,
basis=self.basis,
)
# ------------------------------------------------------------------ serialisation
def _to_payload(self) -> dict:
"""
Convert to a CPU-resident, torch-savable python dict.
"""
vec_cpu = {k: v.detach().cpu() for k, v in self.pair_vectors.items()}
edges_cpu = {k: v.detach().cpu() for k, v in self.pair_edges.items()}
return {
"atoms": list(self.atoms),
"atom_counts": self.atom_counts,
"orbital_cfg": self.orbital_cfg.to_dict(),
"pair_vectors": vec_cpu,
"pair_edges": edges_cpu,
"type": "irrep",
"basis": self.basis,
"layout": (
self.layout.layout_signature() if self.layout is not None else None
),
"packed_values": (
self.packed_values.detach().cpu()
if self.packed_values is not None
else None
),
"packed_edges": (
self.packed_edges.detach().cpu()
if self.packed_edges is not None
else None
),
"frames": self.frames.detach().cpu() if self.frames is not None else None,
"operator": self.operator,
}
def save(self, path: str | "os.PathLike[str]") -> None:
import torch
torch.save(self._to_payload(), path)
# --------------------- alternate constructor ----------------------------- #
@classmethod
def load(cls, path, device="cpu") -> "IrrepsBlockData":
import torch
from core.orbital_irrep_config import OrbitalIrrepConfig
payload = torch.load(path, map_location="cpu", weights_only=False)
if payload.get("type") != "irrep":
raise ValueError("file does not contain irrep snapshot")
orb_cfg = OrbitalIrrepConfig.from_dict(payload["orbital_cfg"])
from core.operator_layout import ShellTensorCodec
pair_vec = {k: v.to(device) for k, v in payload["pair_vectors"].items()}
pair_edges = {k: v.to(device) for k, v in payload["pair_edges"].items()}
# rebuild lookup
lookup = {}
for key, edges in pair_edges.items():
for idx, (sx, sy, sz, i, j) in enumerate(edges.t().tolist()):
lookup[(sx, sy, sz, i, j)] = (key, idx)
atoms = tuple(payload["atoms"])
from collections import Counter
atom_counts = payload.get("atom_counts", Counter(atoms))
return cls(
atoms,
atom_counts,
pair_vec,
pair_edges,
lookup,
orb_cfg,
payload.get("basis", "e3nn"),
layout=(
ShellTensorCodec.from_signature(
payload["layout"],
device=device,
dtype=next(iter(pair_vec.values())).real.dtype,
)
if payload.get("layout") is not None
else None
),
packed_values=payload.get("packed_values"),
packed_edges=payload.get("packed_edges"),
frames=payload.get("frames"),
operator=payload.get("operator"),
).to(device)