Quickstart

This page takes the bundled eight-atom silicon calculation (OpenMX, PBE, Si7.0-s2p2d1 basis: two s, two p and one d shell, 13 orbitals per atom) from raw files to sparse operators, observables, irreducible representations and a first model prediction.

Operators and observables

from pathlib import Path

import torch

from core.block_irrep_mapper import BlockIrrepMapper
from core.sparse_math import trace_matmul_sparse_block_matrix
from data.snapshot import Snapshot
from net.common import Config
from utils.units import HARTREE_TO_EV

ROOT = Path(__file__).resolve().parents[2]
SILICON = ROOT / "data" / "small" / "Si"

# 1. Load an OpenMX calculation (HS.out + its output file) into the common
#    e3nn real-spherical-harmonic convention. All stored blocks are kept.
cfg = Config(verbosity=0)
snapshot = Snapshot.from_openmx(
    SILICON / "HS.out", SILICON / "Si.out", convention="e3nn", cfg=cfg
)
H, S, D = snapshot.hamiltonian, snapshot.overlap, snapshot.density
print("atoms:", len(H.atoms), "| species pairs:", sorted(H.keys()))
print("orbital irreps per species:", snapshot.density.orbital_cfg.to_dict())
print("Si-Si blocks (edges, rows, cols):", tuple(H["Si-Si"].shape))

# 2. Every block is labelled by its directed periodic edge (L, i, j).
edges = H.edges("Si-Si")  # rows: Lx, Ly, Lz, i, j
onsite = (edges[:3] == 0).all(0) & (edges[3] == edges[4])
print("onsite blocks:", int(onsite.sum()), "| total blocks:", edges.shape[1])
local = snapshot.filter_by_distance(8.0)
print("blocks within 8 A:", local.hamiltonian["Si-Si"].shape[0])

# 3. Observables are sparse trace contractions Tr(AB) = sum tr(A_ij^L B_ji^-L).
electrons = float(trace_matmul_sparse_block_matrix(D, S))
band_energy = float(trace_matmul_sparse_block_matrix(D, H)) * HARTREE_TO_EV
print(f"Tr(DS) = {electrons:.6f} electrons | Tr(DH) = {band_energy:.4f} eV")

# 4. Hermiticity of the stored real-space operator: H_ij^L = (H_ji^-L)^T.
asymmetry = float((H - H.adjoint()).to_dense().abs().max())
print(f"max |H - H^T| = {asymmetry:.2e} Hartree")

# 5. Change of basis into irreducible representations and back.
mapper = BlockIrrepMapper(snapshot.density.orbital_cfg, dtype=torch.float32)
vectors = H.to_vectors(mapper)
roundtrip = float((vectors.to_blocks(mapper)["Si-Si"] - H["Si-Si"]).abs().max())
print("Si-Si irreps:", mapper.get_pair_irreps("Si-Si"))
print(f"block -> irreps -> block error: {roundtrip:.1e}")

# 6. Generalized eigenvalue problem H(k) c = e S(k) c along a band path.
bands = snapshot.get_band_structure(path="GXMGR", npoints=60)
energies = bands.eigenvalues * HARTREE_TO_EV
print("band structure:", tuple(energies.shape), "(k-points, bands)")

PYTHONPATH=src:. python examples/quickstart/operators.py

atoms: 8 | species pairs: ['Si-Si']
orbital irreps per species: {'Si': ['2x0e', '2x1o', '1x2e']}
Si-Si blocks (edges, rows, cols): (1336, 13, 13)
onsite blocks: 8 | total blocks: 1336
blocks within 8 A: 792
Tr(DS) = 15.999713 electrons | Tr(DH) = -135.7970 eV
max |H - H^T| = 8.92e-11 Hartree
Si-Si irreps: 9x0e+12x1o+5x1e+4x2o+9x2e+4x3o+1x3e+1x4e
block -> irreps -> block error: 1.5e-08
band structure: (60, 104) (k-points, bands)

A few things are worth noticing:

  • A BlockMatrix stores one dense block per directed periodic edge \((i, j, \bm L)\) and groups them by ordered species pair ("Si-Si"). The 1336 blocks are all atom pairs whose orbitals overlap in the DFT calculation; 792 of them lie within 8 Å.

  • The density matrix is stored per spin channel, so for this non-spin-polarized calculation \(\Tr(DS)=N_e/2=16\) for 32 valence electrons.

  • Each \(13\times13\) block (169 numbers) is exactly the direct sum of the irreducible representations printed above. Their dimensions sum to 169; the operator theory explains where they come from.

  • The band structure is obtained from the generalized eigenvalue problem on the full overlap support; 104 bands are \(8\times13\) orbitals.

A model on the same structure

from pathlib import Path

import torch

from core.sparse_math import trace_matmul_sparse_block_matrix
from data.factory import DatasetFactory
from net.common import Config
from net.e3gnn import E3GNN

ROOT = Path(__file__).resolve().parents[2]
SILICON = ROOT / "data" / "small" / "Si"

# A deliberately small network: two message-passing layers, l <= 2 features.
cfg = Config(
    cutoff_radius=6.0,
    l_max=2,
    hidden_irreps="32x0e+8x0o+16x1o+8x1e+8x2e+4x2o+4x3o+4x3e+2x4e",
    n_radial=16,
    radial_layers=(32,),
    num_layers_gnn=2,
    node_update_message_agg="attention",
    matrix_targets=["hamiltonian", "overlap", "density"],
    verbosity=0,
)

# The factory parses the calculation, applies the cutoff, builds the periodic
# graph and returns the orbital mapper shared by data and model.
factory = DatasetFactory(cfg, convention="e3nn")
factory.add_snapshot(SILICON / "HS.out", SILICON / "Si.out", purpose="train")
train, _, mapper = factory.create()
x, y = train[0]

model = E3GNN(mapper=mapper, cfg=cfg)
print("parameters:", sum(p.numel() for p in model.parameters()))
print("pipeline stages:", list(model.pipeline.transforms))

# The forward pass returns irrep coefficients per atom pair ...
irreps = model(x)
print("Si-Si irrep vectors:", tuple(irreps["hamiltonian"]["Si-Si"].shape))

# ... and predict_matrices returns physical BlockMatrix operators.
with torch.no_grad():
    predicted = model.predict_matrices(x)
H, S, D = (predicted[name] for name in ("hamiltonian", "overlap", "density"))
print("Si-Si blocks:", tuple(H["Si-Si"].shape))
print("Tr(DS) of the untrained model:", float(trace_matmul_sparse_block_matrix(D, S)))

parameters: 514156
pipeline stages: ['node_encoding', 'edge_encoding', 'node_message_0', 'edge_message_0',
                  'node_message_1', 'edge_message_1', 'readout_hamiltonian',
                  'readout_overlap', 'readout_density']
Si-Si irrep vectors: (376, 169)
Si-Si blocks: (376, 13, 13)
Tr(DS) of the untrained model: -10.288435935974121

The flat Config compiles into the nine-stage FeaturePipeline listed above: two encoders, two node/edge message-passing layers and one readout per operator. model(x) returns IrrepsBlockData (irrep coefficients per atom pair); predict_matrices() converts them into physical BlockMatrix operators on the 376 edges of the 6 Å graph. All three operators of an untrained model are, of course, meaningless; Your first model trains one.

Where to go next