Composable dataset statistics and normalization moments.
Bases: Protocol
Streaming computation that finalizes into one statistics result.
update(batch: Batch) -> None
Accumulate sufficient state from one batch.
finalize() -> ResultT
Construct the result from all accumulated batches.
dataclass
BlockSparseStats(
onsite_means: Tensor,
onsite_scales: Tensor,
offsite_means: Tensor,
offsite_scales: Tensor,
)
Canonical coupled-path moments of Hermitian block-sparse data.
state_dict() -> dict[str, Tensor]
Return a tensor-only state suitable for torch.save.
classmethod
from_state_dict(state: Mapping[str, Tensor]) -> BlockSparseStats
Reconstruct statistics using the stored tensors.
BlockSparseStatsAccumulator(
basis_set: BasisSet,
cg_transform: _BlockSparseCGTransform,
*,
target_name: str,
basis_role: str,
)
Accumulate canonical coupled-path moments of block-sparse data.
Physical blocks are first changed to coupled coordinates by the configured CG transform. Onsite means and scales are independent from offsite means and scales. A block is onsite exactly when its two atom indices are equal and its cell shift is zero. Hermitian reverse species and shell pairs share one path. Scalar paths use their centered population standard deviation; all higher-\(L\) paths keep zero mean and use component RMS. Every explicitly stored shell-pair occurrence has equal weight, and a path with zero variance receives unit scale.
update(batch: Batch) -> None
Accumulate moments from one batch containing the configured target.
dataclass
ConnectivityStats(total_edges: int, total_atoms: int)
Atom-weighted statistics of directed model connectivity.
property
avg_num_neighbors: float
Return the atom-weighted mean directed neighbor count.
state_dict() -> dict[str, Tensor]
Return a tensor-only state suitable for torch.save.
classmethod
from_state_dict(state: Mapping[str, Tensor]) -> ConnectivityStats
Reconstruct statistics from state_dict().
ConnectivityStatsAccumulator()
Accumulate atom-weighted statistics of directed model connectivity.
update(batch: Batch) -> None
Accumulate connectivity counts from one connected batch.
dataclass
OrbVectorStats(shell_means: tuple[Tensor, ...], shell_scales: tuple[Tensor, ...])
Equivariant shell moments of an orbital vector.
state_dict() -> dict[str, Tensor]
Return a tensor-only state suitable for torch.save.
classmethod
from_state_dict(state: Mapping[str, Tensor]) -> OrbVectorStats
Reconstruct statistics as views into the stored tensors.
OrbVectorStatsAccumulator(basis_set: BasisSet, *, target_name: str, basis_role: str)
Accumulate equivariant normalization moments for an orbital vector.
Scalar shells use their centered population standard deviation. Higher angular momenta keep zero mean and use component RMS. Every atom-shell occurrence has equal weight, and a path with zero variance receives unit scale.
update(batch: Batch) -> None
Accumulate moments from one batch containing the configured target.
accumulate_stats(
batches: Iterable[Batch], *accumulators: StatsAccumulator[object]
) -> None
Update multiple statistics accumulators in one pass over batches.