Source code for data.structure_inference
from __future__ import annotations
from collections import Counter
import math
from pathlib import Path
from typing import Any
import torch
import torch.nn.functional as F
from ase.io import read as ase_read
from e3nn.o3 import Irreps
from core.block_irrep_mapper import BlockIrrepMapper
from core.orbital_irrep_config import OrbitalIrrepConfig
from data.edge_alignment import (
build_prediction_edge_metadata,
strict_reverse_edge_check,
)
from data.envelope import (
build_edge_envelope,
build_edge_r0_lookup,
load_pair_envelope_table,
)
from data.graph_features import compute_graph_features
from data.openmx_info_parser import parse_info_out
from net.common import Config, get_torch_dtype
def load_orbital_cfg_from_reference_info(
info_path: str | Path,
*,
dtype: torch.dtype = torch.float32,
) -> OrbitalIrrepConfig:
info = parse_info_out(info_path, dtype=dtype)
return OrbitalIrrepConfig.from_dict(info.orbital_set)
def load_structure_from_cif(
cif_path: str | Path,
*,
dtype: torch.dtype = torch.float32,
device: str | torch.device = "cpu",
) -> tuple[tuple[str, ...], torch.Tensor, torch.Tensor]:
atoms = ase_read(str(cif_path))
if atoms.cell is None or atoms.cell.rank < 3:
raise ValueError(
f"CIF structure does not define a full periodic cell: {cif_path}"
)
atoms_tuple = tuple(str(sym) for sym in atoms.get_chemical_symbols())
positions = torch.tensor(atoms.get_positions(), dtype=dtype, device=device)
box = torch.tensor(atoms.cell.array, dtype=dtype, device=device)
return atoms_tuple, positions, box
[docs]
def build_model_input_from_structure(
*,
atoms: tuple[str, ...],
positions: torch.Tensor,
box: torch.Tensor | None,
cfg: Config,
mapper: BlockIrrepMapper,
) -> dict[str, Any]:
# Structure-only inputs need an explicit geometric support when training
# used a stored DFT graph. This is an input policy, independent of the model.
data = getattr(cfg, "data", {})
graph_source = data.get(
"graph_source",
(
"stored"
if data.get("representation") == "packed_shells"
or cfg.graph_source == "target_edges"
else "radius"
),
)
if graph_source == "stored":
raise ValueError(
"Structure-only inference requires data.graph_source=radius and an explicit graph_radius; use native input to preserve stored support"
)
if data.get("graph_source") == "radius":
from dataclasses import replace
radius = data.get("graph_radius")
if radius is None or not math.isfinite(float(radius)) or radius <= 0:
raise ValueError("Radius graph requires a finite positive graph_radius")
cfg = replace(cfg, cutoff_radius=float(radius))
species = tuple(cfg.species.get("order", mapper.orbital_cfg.elements()))
edge_types = [f"{left}-{right}" for left in species for right in species]
edge_type2idx = {pair: index for index, pair in enumerate(edge_types)}
dtype = get_torch_dtype(cfg.dtype)
sh_irreps = Irreps.spherical_harmonics(cfg.l_max)
envelope_mode = str(getattr(cfg, "hamiltonian_envelope_mode", "off")).lower()
pair_distance_normalization = str(
getattr(cfg, "pair_distance_normalization", "off")
).lower()
loss_weighting_mode = str(getattr(cfg, "loss_weighting_mode", "off")).lower()
envelope_weight_modes = {
"envelope_inverse_sqrt_clipped",
"envelope_inverse_clipped",
}
need_envelope_table = (
getattr(cfg, "hamiltonian_envelope_path", None) is not None
or envelope_mode != "off"
or pair_distance_normalization != "off"
or loss_weighting_mode in envelope_weight_modes
)
envelope_table = None
edge_type_r0 = None
if need_envelope_table:
envelope_path = getattr(cfg, "hamiltonian_envelope_path", None)
if envelope_path is None:
raise ValueError(
"hamiltonian_envelope_path is required when envelope-based "
"evaluation features are enabled."
)
envelope_table = load_pair_envelope_table(
envelope_path,
pair_order=edge_types,
dtype=dtype,
device=positions.device,
)
if pair_distance_normalization == "pair_r0":
edge_type_r0 = envelope_table.r0.clone().detach()
(
edge_index,
edge_shift,
edge_type_idx,
edge_length_emb,
edge_sh,
num_self_edges,
edge_lengths,
) = compute_graph_features(
positions=positions,
box=box,
atoms=atoms,
cfg=cfg,
sh_irreps=sh_irreps,
edge_type2idx=edge_type2idx,
edge_type_r0=edge_type_r0,
)
strict_reverse_edge_check(edge_index, edge_shift, edge_set_name="graph")
elem2idx = {el: i for i, el in enumerate(species)}
try:
node_type_idx = torch.tensor(
[elem2idx[el] for el in atoms],
dtype=torch.long,
device=positions.device,
)
except KeyError as exc:
raise ValueError(
f"Structure contains element {exc.args[0]!r} that is missing from the orbital config."
) from exc
num_species = len(mapper.orbital_cfg.elements())
node_one_hot = F.one_hot(node_type_idx, num_classes=num_species).to(dtype=dtype)
src_type = node_type_idx[edge_index[0]]
dst_type = node_type_idx[edge_index[1]]
edge_one_hot = F.one_hot(
src_type * num_species + dst_type,
num_classes=num_species * num_species,
).to(dtype=dtype)
pred_metadata = build_prediction_edge_metadata(
edge_index=edge_index,
edge_shift=edge_shift,
edge_type_idx=edge_type_idx,
atoms=atoms,
edge_types=edge_types,
edge_type2idx=edge_type2idx,
separate_shifted_self=bool(cfg.separate_shifted_self),
)
x: dict[str, Any] = {
"node_type_idx": node_type_idx,
"node_one_hot": node_one_hot,
"positions": positions,
"box": box,
"atoms": atoms,
"atoms_tuple": atoms,
"atom_counts": dict(Counter(atoms)),
"edge_index": edge_index,
"edge_shift": edge_shift,
"edge_type_idx": edge_type_idx,
"edge_one_hot": edge_one_hot,
"edge_length": edge_lengths.to(dtype=dtype),
"num_self_edges": num_self_edges,
**pred_metadata,
}
if envelope_table is not None:
x["edge_envelope"] = build_edge_envelope(
edge_lengths=edge_lengths,
edge_type_idx=edge_type_idx,
envelope_table=envelope_table,
).to(dtype=dtype)
if envelope_table.r0_parameter_index is not None:
x["edge_r0"] = build_edge_r0_lookup(
edge_type_idx=edge_type_idx,
envelope_table=envelope_table,
).to(dtype=dtype)
if cfg.precompute_edge_features:
x["edge_length_emb"] = edge_length_emb
x["edge_sh"] = edge_sh
return x