"""
E3GNN - PyTorch-Lightning implementation
----------------------------------------
* encoders.py -> NodeEncoder / EdgeEncoder
* layers.py -> MessageBlock (small & large graphs)
* heads.py -> DeepHead (H, S, D)
* sparse_math.trace_* -> Energy / electron count losses
"""
from __future__ import annotations
from typing import Dict, Any
import torch
import pytorch_lightning as pl
from collections import OrderedDict
from e3nn.o3 import Irreps
import time
from core.block_irrep_mapper import BlockIrrepMapper
from data.snapshot import Snapshot
from data.block_matrix import BlockMatrix, IrrepsBlockData
from data.envelope import scale_block_matrix_by_edge_values
from data.graph_features import compute_edge_geometry_from_static_edges
from net.common import Config, resolve_hidden_irreps, validate_training_config
from net.irrep_tools import (
build_irrep_projector_cache,
build_irrep_block_matrix_cache,
compute_irrep_metrics,
compute_hamiltonian_mae_contributions,
get_all_irreps,
project_irrep_vectors_to_blocks,
)
from net.observable_metrics import (
add_observable_metrics,
build_observable_predictions,
observable_loss,
rescale_density_prediction_to_num_electrons,
truncate_pred_block_matrix_to_target_prefix,
validate_observable_config,
)
from net.spectral_loss import compute_spectral_eigenvalue_loss
from net.objectives import build_objective_terms
from net.checkpoint_compat import strip_mapper_keys
from utils.summary import print_model_summary
from utils.units import HARTREE_TO_EV
[docs]
class E3GNN(pl.LightningModule):
"""
Full network **and** training logic.
Forward returns a dict with three `IrrepsBlockData` entries
(hamiltonian / overlap / density). Losses are computed inside
`training_step`.
"""
def __init__(
self,
mapper: BlockIrrepMapper,
cfg: Config,
*,
legacy_separate_shifted_self: bool = False,
):
super().__init__()
# Older slotted Config pickles predate these optional fields. Unpickling
# does not run the dataclass constructor, so initialize missing slots
# before Lightning inspects or serializes the restored configuration.
defaults = Config()
for field in Config.__dataclass_fields__:
if not hasattr(cfg, field):
setattr(cfg, field, getattr(defaults, field))
if not legacy_separate_shifted_self:
validate_training_config(cfg)
self.legacy_separate_shifted_self = bool(
legacy_separate_shifted_self and cfg.separate_shifted_self
)
self.save_hyperparameters(ignore=["mapper", "legacy_separate_shifted_self"])
# Structure-only restoration must not depend on reference .out files.
self.save_hyperparameters({"orbital_cfg": mapper.orbital_cfg.to_dict()})
self.cfg = cfg
# The mapper is now an nn.Module and will be moved to the correct device
# automatically by PyTorch Lightning.
self.mapper: BlockIrrepMapper = mapper
self._objective_terms = build_objective_terms(cfg.objective)
if cfg.pipeline is None:
from net.component_defaults import default_graph
cfg.pipeline = default_graph(cfg, len(mapper.orbital_cfg.elements()))
allow_zero_observable_control = bool(
getattr(cfg, "allow_zero_observable_loss_control", False)
)
if (
cfg.train_on_energy
and cfg.loss_coef_observables == 0.0
and not allow_zero_observable_control
):
raise ValueError(
"If training on energy, loss_coef_observables must be nonzero unless "
"allow_zero_observable_loss_control=True."
)
if (
cfg.train_on_num_electrons
and cfg.loss_coef_observables == 0.0
and not allow_zero_observable_control
):
raise ValueError(
"If training on number of electrons, loss_coef_observables must be "
"nonzero unless allow_zero_observable_loss_control=True."
)
if cfg.train_on_forces and cfg.loss_coef_forces == 0.0:
raise ValueError("If training on forces, loss_coef_forces must be nonzero.")
if cfg.train_on_stress and cfg.loss_coef_stress == 0.0:
raise ValueError("If training on stress, loss_coef_stress must be nonzero.")
if cfg.partial_train not in (None, "diag", "shifted_self", "offdiag"):
raise ValueError(
"partial_train must be one of None, 'diag', 'shifted_self', or 'offdiag'."
)
if cfg.train_on_irrep_parts and cfg.train_target != "matrix":
raise ValueError(
"train_on_irrep_parts expects matrix-space supervision and should only be used with train_target='matrix'."
)
if cfg.train_target not in {"matrix", "irreps"}:
raise ValueError("train_target must be one of 'matrix' or 'irreps'.")
if cfg.matrix_loss_reduction not in {"pair_sum", "pair_mean", "elements"}:
raise ValueError(
"matrix_loss_reduction must be pair_sum, pair_mean or elements"
)
if cfg.matrix_loss_reduction != "pair_sum":
if cfg.train_target != "matrix" or cfg.train_on_irrep_parts:
raise ValueError(
"matrix_loss_reduction requires orbital_blocks matrix-space "
"supervision, not irrep objectives"
)
if cfg.data.get("representation") == "packed_shells":
raise ValueError(
"matrix_loss_reduction applies to orbital_blocks; packed_shells "
"already uses its declared masked-entry objective"
)
if str(cfg.loss_weighting_mode).lower() != "off":
raise ValueError(
"matrix_loss_reduction pair_mean/elements requires "
"loss_weighting_mode='off'"
)
validate_observable_config(cfg)
observable_loss_kind = str(getattr(cfg, "observable_loss_kind", "mse")).lower()
if observable_loss_kind not in {"mse", "mae"}:
raise ValueError("observable_loss_kind must be one of 'mse' or 'mae'.")
spectral_loss_kind = str(getattr(cfg, "spectral_loss_kind", "huber")).lower()
if spectral_loss_kind not in {"huber", "mse", "mae"}:
raise ValueError(
"spectral_loss_kind must be one of 'huber', 'mse', or 'mae'."
)
envelope_mode = str(getattr(cfg, "hamiltonian_envelope_mode", "off")).lower()
if envelope_mode not in {"off", "normalize_target", "multiply_prediction"}:
raise ValueError(
"hamiltonian_envelope_mode must be one of 'off', "
"'normalize_target', or 'multiply_prediction'."
)
if bool(getattr(cfg, "pair_conditioned_radial_mlp", False)) and bool(
getattr(cfg, "train_on_irrep_parts", False)
):
raise ValueError(
"pair_conditioned_radial_mlp is not supported together with "
"train_on_irrep_parts."
)
if bool(getattr(cfg, "train_on_irrep_parts", False)) and envelope_mode != "off":
raise ValueError(
"hamiltonian_envelope_mode requires matrix-space training; "
"it cannot be used with train_on_irrep_parts."
)
if cfg.train_target == "irreps" and envelope_mode != "off":
raise ValueError(
"hamiltonian_envelope_mode requires matrix-space supervision and "
"cannot be used with train_target='irreps'."
)
if (
cfg.train_target == "irreps"
and str(cfg.loss_weighting_mode).lower() != "off"
):
raise ValueError(
"loss_weighting_mode requires matrix-edge supervision and cannot "
"be used with train_target='irreps'."
)
if bool(
getattr(cfg, "spectral_loss_enabled", False)
) and "hamiltonian" not in set(cfg.matrix_targets):
raise ValueError(
"spectral_loss_enabled requires 'hamiltonian' to be present in matrix_targets."
)
# ---------- shared irreps ---------------------------------------
self.hidden_irreps: Irreps = resolve_hidden_irreps(self.cfg)
self.neck_irreps: Irreps = resolve_hidden_irreps(self.cfg)
self.sh_irreps: Irreps = Irreps.spherical_harmonics(self.cfg.l_max)
if cfg.data.get("representation") == "packed_shells" and not cfg.pipeline.get(
"operator_layout"
):
shells = {
species: tuple(ir.l for mul, ir in irreps for _ in range(mul))
for species, irreps in mapper.orbital_cfg.element_to_irreps.items()
}
with torch.random.fork_rng(devices=[]):
resolved, _ = Config.resolve(
cfg.to_mapping(),
species_shells=shells,
spinful=cfg.layout.get("spinful", False),
)
cfg.pipeline = resolved["pipeline"]
from net.pipeline import FeaturePipeline
self.pipeline = FeaturePipeline(
cfg.pipeline, num_species=len(mapper.orbital_cfg.elements()), mapper=mapper
)
for operator, slot in self.pipeline.outputs.items():
spec = self.pipeline.feature_specs[slot]
if (
spec.domain != "edge"
or spec.representation not in {"matrix", "block_vectors"}
or spec.support != "output"
):
raise ValueError(
f"Operator {operator!r} must select matrix features on the output-edge support"
)
signature = self.pipeline.operator_layout
self.output_layout = (
self.mapper.resolve_layout(signature) if signature is not None else None
)
self.all_irreps = get_all_irreps(self.mapper, self.output_layout)
self.irrep_projectors = build_irrep_projector_cache(
self.mapper, self.all_irreps, layout=self.output_layout
)
self._nan_loss_detected = False
self._spectral_loss_debug_printed = False
self._apply_init_weights_factor()
self._apply_trainable_parameter_freeze()
# Print model summary if verbosity >= 1
print_model_summary(self, verbosity=self.cfg.verbosity)
@classmethod
def from_checkpoint_config(cls, mapper: BlockIrrepMapper, cfg: Config) -> "E3GNN":
"""Construct a saved architecture for strict state-dict loading/inference.
This explicit legacy boundary preserves shifted-self module names and
routing. It does not authorize training a retired architecture.
"""
return cls(mapper, cfg, legacy_separate_shifted_self=True)
@classmethod
def load_from_checkpoint(cls, checkpoint_path, *args, **kwargs):
"""Restore historical heads automatically from the stored configuration."""
kwargs["legacy_separate_shifted_self"] = True
kwargs.setdefault("weights_only", False)
return super().load_from_checkpoint(checkpoint_path, *args, **kwargs)
def on_fit_start(self):
validate_training_config(self.cfg)
@classmethod
def load_from_checkpoint_with_mapper(
cls,
checkpoint_path: str,
mapper: BlockIrrepMapper,
map_location: str | torch.device | None = "cpu",
strict: bool = True,
**kwargs,
) -> "E3GNN":
# Native checkpoints contain Config objects as well as tensor weights.
# Keep their trusted-payload loading semantics across PyTorch >= 2.6.
kwargs.setdefault("weights_only", False)
return cls.load_from_checkpoint(
checkpoint_path,
mapper=mapper,
map_location=map_location,
strict=strict,
**kwargs,
)
[docs]
def fit_training_data(self, dataset):
"""Fit declared train-only transforms once, using the dataset contract."""
if getattr(self, "_statistics_prepared", False):
return
for component in self.modules():
if component is self:
continue
fit = getattr(component, "fit_training_data", None)
if fit is not None:
fit(
dataset.samples,
layout=self.output_layout,
mapper=self.mapper,
target_frame=self.cfg.data.get("target_frame", "global"),
)
self._statistics_prepared = True
def load_state_dict(self, state_dict: dict[str, Any], strict: bool = True):
# Translate historical parameter paths at the checkpoint boundary only.
# Every newly constructed model uses the common component graph.
translated = {}
for key, value in strip_mapper_keys(state_dict).items():
if key.startswith("node_enc."):
key = "pipeline.transforms.node_encoding." + key.removeprefix(
"node_enc."
)
elif key.startswith("edge_enc."):
key = "pipeline.transforms.edge_encoding." + key.removeprefix(
"edge_enc."
)
elif key.startswith("mp_blocks."):
_, index, suffix = key.split(".", 2)
component, suffix = suffix.split(".", 1)
domain = {"node_upd": "node", "edge_upd": "edge"}[component]
key = f"pipeline.transforms.{domain}_message_{index}.{suffix}"
elif key.startswith("heads."):
_, operator, suffix = key.split(".", 2)
key = f"pipeline.transforms.readout_{operator}.{suffix}"
translated[key] = value
# Frozen basis buffers did not exist in older public checkpoints.
# Reconstruct only these constants; learned parameters remain strict.
learned = dict(self.named_parameters())
for key, value in super().state_dict().items():
if (
key.startswith("output_layout.")
and key not in learned
and key not in translated
):
translated[key] = value
return super().load_state_dict(translated, strict=strict)
def _apply_init_weights_factor(self) -> None:
factor = float(self.cfg.init_weights_factor)
if factor == 1.0:
return
with torch.no_grad():
for param in self.parameters():
if param.is_floating_point():
param.mul_(factor)
def _apply_trainable_parameter_freeze(self) -> None:
if not bool(getattr(self.cfg, "freeze_backbone_train_heads_only", False)):
return
for name, module in self.pipeline.transforms.items():
if self.pipeline.roles[name] != "readout":
for param in module.parameters():
param.requires_grad = False
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
frozen = sum(p.numel() for p in self.parameters() if not p.requires_grad)
print(
"--- Head-only fine-tuning enabled: "
f"trainable_params={trainable}, frozen_params={frozen} ---"
)
# ------------------------ util helpers -----------------------------
@staticmethod
def _mse(pred, target):
if pred.shape != target.shape:
raise ValueError(
f"Shape mismatch in block loss: {pred.shape=}, {target.shape=}"
)
return torch.mean(torch.abs(pred - target) ** 2)
@staticmethod
def _mae(pred, target):
if pred.shape != target.shape:
raise ValueError(
f"Shape mismatch in block loss: {pred.shape=}, {target.shape=}"
)
return torch.mean(torch.abs(pred - target))
def _wrap_head_output(
self,
raw: Dict[str, torch.Tensor],
x: Dict[str, Any],
) -> IrrepsBlockData:
"""
Convert DeepHead raw dict -> IrrepsBlockData with mapper.
"""
irreps_blocks = IrrepsBlockData(
atoms=x["atoms_tuple"],
atom_counts=x["atom_counts"],
pair_vectors=raw,
pair_edges=x["pred_pair_edges_static"],
lookup=x["pred_lookup_static"],
orbital_cfg=self.mapper.orbital_cfg,
)
return irreps_blocks
def _matrix_envelope_mode(self) -> str:
return str(getattr(self.cfg, "hamiltonian_envelope_mode", "off")).lower()
def _physicalize_predicted_matrix(
self,
*,
name: str,
pred_matrix: BlockMatrix,
x: Dict[str, Any],
) -> BlockMatrix:
mode = self._matrix_envelope_mode()
if name not in {"hamiltonian", "overlap"} or mode == "off":
return pred_matrix
if "edge_envelope" not in x:
raise ValueError(
"hamiltonian_envelope_mode requires edge_envelope in the batch."
)
if "edge_partitions" not in x:
raise ValueError(
"hamiltonian_envelope_mode requires edge_partitions in the batch."
)
if mode in {"multiply_prediction", "normalize_target"}:
return scale_block_matrix_by_edge_values(
pred_matrix,
x["edge_envelope"],
x["edge_partitions"],
)
raise ValueError(
"hamiltonian_envelope_mode must be one of 'off', "
"'normalize_target', or 'multiply_prediction'."
)
def predicted_irreps_to_block_matrices(
self,
predictions: Dict[str, IrrepsBlockData],
x: Dict[str, Any],
*,
physical: bool = True,
) -> Dict[str, BlockMatrix]:
matrices = {
name: pred.to_blocks(self.mapper) for name, pred in predictions.items()
}
if not physical:
return matrices
return {
name: self._physicalize_predicted_matrix(
name=name,
pred_matrix=matrix,
x=x,
)
for name, matrix in matrices.items()
}
[docs]
def predict_matrices(
self,
x: Dict[str, Any],
*,
physical: bool = True,
) -> Dict[str, BlockMatrix]:
"""Run inference and return predicted sparse block matrices.
This is a convenience wrapper around ``forward`` and
``predicted_irreps_to_block_matrices``. It does not move ``x`` between
devices; callers should prepare the input on the model's device.
"""
was_training = self.training
self.eval()
try:
with torch.inference_mode():
from data.sample import move_tree
x = move_tree(x, self.device, self.dtype)
predictions = self(x, physical_derivatives=True)
return self.predicted_irreps_to_block_matrices(
predictions,
x,
physical=physical,
)
finally:
self.train(was_training)
def _spectral_loss_enabled(self) -> bool:
return bool(getattr(self.cfg, "spectral_loss_enabled", False))
def _compute_spectral_loss(
self,
*,
x: Dict[str, Any],
physical_pred_hamiltonian: BlockMatrix,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
required = {
"spectral_kpoints_abs",
"spectral_shifts",
"spectral_gt_overlap_k",
"spectral_gt_eigs_ev",
"spectral_gt_fermi_ev",
"spectral_window_weights",
}
missing = sorted(key for key in required if key not in x)
if missing:
raise ValueError(
"spectral_loss_enabled requires cached spectral payload in the batch; "
f"missing keys: {missing}"
)
if not self._spectral_loss_debug_printed:
print(
"--- Running first spectral loss evaluation: "
f"kpoints={int(x['spectral_kpoints_abs'].shape[0])}, "
f"window_weight_sum={float(x['spectral_window_weights'].sum().item()):.1f}, "
f"loss_kind={str(self.cfg.spectral_loss_kind)}, "
f"huber_delta_ev={float(self.cfg.spectral_loss_huber_delta_ev):.3f} ---",
flush=True,
)
loss, stats = compute_spectral_eigenvalue_loss(
pred_hamiltonian=physical_pred_hamiltonian,
spectral_payload=x,
box=x["box"],
loss_kind=str(self.cfg.spectral_loss_kind),
huber_delta_ev=float(self.cfg.spectral_loss_huber_delta_ev),
overlap_psd_cleanup=bool(
getattr(self.cfg, "spectral_loss_overlap_psd_cleanup", True)
),
overlap_jitter=bool(
getattr(self.cfg, "spectral_loss_overlap_jitter", False)
),
)
if not self._spectral_loss_debug_printed:
print(
"--- First spectral loss evaluation finished: "
f"loss={float(loss.detach().item()):.6f}, "
f"spectral_mae_ev={float(stats['spectral_mae_ev'].detach().item()):.6f} ---",
flush=True,
)
self._spectral_loss_debug_printed = True
return loss, stats
def _apply_matrix_envelope_mode(
self,
*,
name: str,
pred_matrix: BlockMatrix,
target_matrix: BlockMatrix,
x: Dict[str, Any],
) -> tuple[BlockMatrix, BlockMatrix, BlockMatrix]:
mode = self._matrix_envelope_mode()
if name not in {"hamiltonian", "overlap"} or mode == "off":
return pred_matrix, target_matrix, pred_matrix
if "edge_envelope" not in x:
raise ValueError(
"hamiltonian_envelope_mode requires edge_envelope in the batch."
)
if "edge_partitions" not in x:
raise ValueError(
"hamiltonian_envelope_mode requires edge_partitions in the batch."
)
edge_envelope = x["edge_envelope"]
edge_partitions = x["edge_partitions"]
eps = float(getattr(self.cfg, "hamiltonian_envelope_eps", 1.0e-12))
if mode == "multiply_prediction":
scaled_pred = scale_block_matrix_by_edge_values(
pred_matrix,
edge_envelope,
edge_partitions,
)
return scaled_pred, target_matrix, scaled_pred
if mode == "normalize_target":
physical_pred = scale_block_matrix_by_edge_values(
pred_matrix,
edge_envelope,
edge_partitions,
)
normalized_target = scale_block_matrix_by_edge_values(
target_matrix,
edge_envelope,
edge_partitions,
inverse=True,
eps=eps,
allow_prefix_trim=True,
)
return pred_matrix, normalized_target, physical_pred
raise ValueError(
"hamiltonian_envelope_mode must be one of 'off', "
"'normalize_target', or 'multiply_prediction'."
)
def _edge_loss_weights(self, x: Dict[str, Any]) -> torch.Tensor | None:
mode = str(getattr(self.cfg, "loss_weighting_mode", "off")).lower()
if mode == "off":
return None
if "edge_partitions" not in x:
raise ValueError(
"loss_weighting_mode requires edge_partitions in the batch."
)
if mode in {"envelope_inverse_sqrt_clipped", "envelope_inverse_clipped"}:
if "edge_envelope" not in x:
raise ValueError(
"loss_weighting_mode based on envelope requires edge_envelope in the batch."
)
base = x["edge_envelope"]
eps = float(getattr(self.cfg, "hamiltonian_envelope_eps", 1.0e-12))
if mode == "envelope_inverse_sqrt_clipped":
weights = torch.rsqrt(base.clamp_min(eps))
else:
weights = base.clamp_min(eps).reciprocal()
elif mode == "distance_short_range_bias":
if "edge_length" not in x:
raise ValueError(
"distance_short_range_bias loss weighting requires edge_length in the batch."
)
weights = torch.rsqrt(1.0 + x["edge_length"].clamp_min(0.0))
else:
raise ValueError(
"loss_weighting_mode must be one of 'off', "
"'distance_short_range_bias', "
"'envelope_inverse_sqrt_clipped', or 'envelope_inverse_clipped'."
)
min_w = float(getattr(self.cfg, "loss_weight_min", 0.0))
max_w = float(getattr(self.cfg, "loss_weight_max", 1.0))
if max_w < min_w:
raise ValueError("loss_weight_max must be >= loss_weight_min.")
return weights.clamp(min=min_w, max=max_w)
def _weighted_edge_block_losses(
self,
*,
preds: torch.Tensor,
targets: torch.Tensor,
edge_weights: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
if preds.shape != targets.shape:
raise ValueError(
"Predicted and target block tensors must have identical shapes, got "
f"{tuple(preds.shape)} and {tuple(targets.shape)}"
)
if preds.ndim != 3:
raise ValueError(
f"Block tensors must have shape (E, D1, D2), got {tuple(preds.shape)}"
)
diff = preds - targets
mixture = self._objective_l1_fraction()
mse_diff = diff if mixture < 1.0 else diff.detach()
mae_diff = diff if mixture > 0.0 else diff.detach()
per_edge_mse = torch.mean(mse_diff.abs().square(), dim=(1, 2))
per_edge_mae = torch.mean(torch.abs(mae_diff), dim=(1, 2))
if edge_weights is None:
return per_edge_mse.mean(), per_edge_mae.mean()
if edge_weights.ndim != 1 or edge_weights.shape[0] != per_edge_mse.shape[0]:
raise ValueError(
"Edge weights must be a 1D tensor matching the number of selected "
f"edges, got weights={tuple(edge_weights.shape)} and "
f"selected_edges={int(per_edge_mse.shape[0])}"
)
edge_weights = edge_weights.to(device=preds.device, dtype=preds.dtype)
weight_sum = edge_weights.sum()
weighted_mse = torch.sum(edge_weights * per_edge_mse) / weight_sum
weighted_mae = torch.sum(edge_weights * per_edge_mae) / weight_sum
return weighted_mse, weighted_mae
def _objective_l1_fraction(self) -> float:
"""Return the actual objective mixture used by loss aggregation."""
kind = (self.cfg.objective or {}).get("kind")
if kind == "mae":
return 1.0
if kind == "mse":
return 0.0
return float(self.cfg.loss_l1_fraction)
def _observed_block_losses(self, preds, targets, mask):
"""Differentiable unweighted means and count for observed block entries."""
if preds.shape != targets.shape or mask.shape != targets.shape:
raise ValueError("Observed block predictions, targets and mask must align")
if mask.dtype != torch.bool:
raise ValueError("Observed block coverage mask must be boolean")
difference = (preds - targets)[mask]
if not difference.numel():
return None
if not torch.isfinite(difference).all():
raise ValueError("Nonfinite observed matrix objective error")
mixture = self._objective_l1_fraction()
mse_difference = difference if mixture < 1.0 else difference.detach()
mae_difference = difference if mixture > 0.0 else difference.detach()
return (
mse_difference.abs().square().mean(),
mae_difference.abs().mean(),
difference.numel(),
)
def _should_recompute_edge_features(self, x: Dict[str, Any]) -> bool:
if "positions" in x:
reference = x.get("_geometry_reference", {})
if any(
key not in reference or not torch.equal(value, reference[key].to(value))
for key in ("positions", "box")
if (value := x.get(key)) is not None
):
return True
if not self.cfg.precompute_edge_features:
return True
required = {slot for _, reads, _ in self.pipeline.bindings for slot in reads}
if ({"edge_sh", "edge_length_emb"} & required) - x.keys():
return True
if "positions" in x and x["positions"].requires_grad:
return True
if x.get("box") is not None and x["box"].requires_grad:
return True
return False
def _geometry_num_l(self):
import math
spec = self.pipeline.feature_specs.get("sub_edge_ang")
width = spec.dim if spec is not None else None
declared = self.cfg.data.get("num_l")
if width is not None:
degree = math.isqrt(width)
if degree * degree != width or (
declared is not None and declared != degree
):
raise ValueError(
"Local harmonic configuration disagrees with the typed angular input"
)
return degree
return declared if declared is not None else 5
def _populate_edge_features(self, x: Dict[str, Any]) -> None:
from data.graph_features import evaluate_graph_geometry
reads = {slot for _, slots, _ in self.pipeline.bindings for slot in slots}
evaluate_graph_geometry(
x,
local_coordinates=bool(
reads & {"sub_edge_ang", "sub_index", "sub_atom_idx", "sub_edge_idx"}
),
num_l=self._geometry_num_l(),
require_frames="bond_frames" in reads,
recompute_geometry=True,
derive_frames=self.cfg.data.get("frame_source", "derive") == "derive",
neighborhood_support=self.cfg.data.get("neighborhood_support", "output"),
)
if getattr(self.cfg, "hamiltonian_envelope_path", None) is not None:
from data.envelope import load_pair_envelope_table, build_edge_envelope
order = self.cfg.species.get("order") or (
tuple(self.output_layout.species_shells)
if self.output_layout is not None
else self.mapper.orbital_cfg.elements()
)
pairs = [f"{a}-{b}" for a in order for b in order]
table = load_pair_envelope_table(
self.cfg.hamiltonian_envelope_path,
pair_order=pairs,
dtype=x["edge_lengths"].dtype,
device=x["edge_lengths"].device,
)
x["edge_envelope"] = build_edge_envelope(
x["edge_lengths"], x["edge_type_idx"], table
)
if self.cfg.pair_distance_normalization == "pair_r0":
x["edge_r0"] = table.r0[x["edge_type_idx"]]
if not ({"edge_sh", "edge_length_emb"} & reads):
return
radial_lengths = None
radial_basis_end = None
if (
str(getattr(self.cfg, "pair_distance_normalization", "off")).lower()
== "pair_r0"
):
if "edge_length" not in x or "edge_r0" not in x:
raise ValueError(
"pair_distance_normalization='pair_r0' requires edge_length "
"and edge_r0 in the batch."
)
radial_lengths = x["edge_length"] / x["edge_r0"].clamp_min(1e-12)
radial_basis_end = 1.0
edge_length_emb, edge_sh, _ = compute_edge_geometry_from_static_edges(
positions=x["positions"],
box=x["box"],
edge_index=x["edge_index"],
edge_shift=x["edge_shift"],
sh_irreps=self.sh_irreps,
cutoff_radius=self.cfg.cutoff_radius,
n_radial=self.cfg.n_radial,
radial_embedding_scale=self.cfg.radial_embedding_scale,
radial_lengths=radial_lengths,
radial_basis_end=radial_basis_end,
edge_batch=x.get("edge_batch"),
)
x["edge_length_emb"] = edge_length_emb
x["edge_sh"] = edge_sh
def _partial_train_mask(self, edges_5d: torch.Tensor) -> torch.Tensor | None:
if self.cfg.partial_train is None:
return None
sx, sy, sz, i, j = edges_5d
is_diag = (i == j) & (sx == 0) & (sy == 0) & (sz == 0)
if self.cfg.partial_train == "diag":
return is_diag
if self.cfg.partial_train == "shifted_self":
# Only saved configurations can request this retired selector.
return (i == j) & ~((sx == 0) & (sy == 0) & (sz == 0))
if self.cfg.partial_train == "offdiag":
if self.legacy_separate_shifted_self:
return i != j
return ~is_diag
raise RuntimeError(f"Unsupported partial_train value: {self.cfg.partial_train}")
def _stop_training_on_nonfinite_loss(
self,
loss: torch.Tensor,
*,
stage: str,
batch_idx: int,
phase: str,
) -> bool:
if stage != "train":
return False
if torch.isfinite(loss).all():
return False
self._nan_loss_detected = True
print(
"--- Non-finite train loss detected "
f"(phase={phase}, batch_idx={batch_idx}); "
"stopping before optimizer step and finishing evaluation. ---",
flush=True,
)
trainer = getattr(self, "_trainer", None)
if trainer is not None:
try:
trainer.should_stop = True
except Exception:
pass
return True
def _compute_irrep_part_losses(
self,
pred: IrrepsBlockData,
target: IrrepsBlockData,
) -> tuple[torch.Tensor, torch.Tensor, dict[str, dict[str, torch.Tensor]]]:
total_mse = torch.tensor(0.0, device=self.device)
total_mae = torch.tensor(0.0, device=self.device)
per_irrep: dict[str, dict[str, torch.Tensor]] = {}
for irrep in self.all_irreps:
irrep_key = str(irrep)
irrep_mse = torch.tensor(0.0, device=self.device)
irrep_mae = torch.tensor(0.0, device=self.device)
saw_any = False
for pair_key, projector in self.irrep_projectors[irrep_key].items():
if (
pair_key not in pred.pair_vectors
or pair_key not in target.pair_vectors
):
continue
pred_blocks = project_irrep_vectors_to_blocks(
pred.pair_vectors[pair_key], projector, self.mapper
)
target_blocks = project_irrep_vectors_to_blocks(
target.pair_vectors[pair_key], projector, self.mapper
)
target_n = target_blocks.shape[0]
if target_n <= 0:
continue
if pred_blocks.shape[0] < target_n:
raise ValueError(
f"Predicted irrep blocks for key {pair_key} and irrep {irrep_key} "
f"are too short: pred_len={pred_blocks.shape[0]} target_len={target_n}"
)
pred_edges = pred.pair_edges[pair_key][:, :target_n]
target_edges = target.pair_edges[pair_key][:, :target_n]
if self.cfg.safety_checks:
assert torch.equal(
pred_edges, target_edges
), f"Edge mismatch in irrep-part loss for key {pair_key}, irrep {irrep_key}."
partial_mask = self._partial_train_mask(target_edges)
pred_selected = pred_blocks[:target_n]
target_selected = target_blocks[:target_n]
if partial_mask is not None:
pred_selected = pred_selected[partial_mask]
target_selected = target_selected[partial_mask]
if pred_selected.shape[0] == 0:
continue
irrep_mse = irrep_mse + self._mse(pred_selected, target_selected)
irrep_mae = irrep_mae + self._mae(pred_selected, target_selected)
saw_any = True
if saw_any:
per_irrep[irrep_key] = {"mse": irrep_mse, "mae": irrep_mae}
total_mse = total_mse + irrep_mse
total_mae = total_mae + irrep_mae
return total_mse, total_mae, per_irrep
def forward(self, x: Dict[str, Any], *, physical_derivatives: bool = False):
# Forward-local values must not leave stale autograd graphs in a dataset.
x = dict(x)
from data.graph_topology import prepare_graph_topology
prepare_graph_topology(x, self.cfg.data.get("supports"))
from data.graph_features import complete_graph_metadata
complete_graph_metadata(
x,
self.mapper,
species_order=self.cfg.species.get("order")
or (
tuple(self.output_layout.species_shells)
if self.output_layout is not None
else None
),
separate_shifted_self=self.legacy_separate_shifted_self,
)
from data.graph_features import evaluate_graph_geometry
reads = {slot for _, slots, _ in self.pipeline.bindings for slot in slots}
geometric = {
"edge_vectors",
"edge_lengths",
"message_vectors",
"message_lengths",
"message_edge_index",
"message_to_output",
"selfloop_edge",
"bond_frames",
"sub_edge_ang",
"sub_index",
}
if (reads & geometric) - x.keys():
evaluate_graph_geometry(
x,
local_coordinates=bool(
reads
& {"sub_edge_ang", "sub_index", "sub_atom_idx", "sub_edge_idx"}
),
num_l=self._geometry_num_l(),
require_frames="bond_frames" in reads,
neighborhood_support=self.cfg.data.get(
"neighborhood_support", "output"
),
)
activation_mags = OrderedDict() if self.cfg.log_activation_mag else None
if self._should_recompute_edge_features(x):
self._populate_edge_features(x)
outputs = self.pipeline(
x,
physical_derivatives=physical_derivatives,
activation_mags=activation_mags,
)
predictions = {
name: (
self._wrap_head_output(value, x)
if self.pipeline.feature_specs[
self.pipeline.outputs[name]
].representation
== "block_vectors"
else IrrepsBlockData.from_packed(
value,
x,
self.mapper,
self.output_layout,
frame=self.pipeline.feature_specs[
self.pipeline.outputs[name]
].frame,
operator=name,
)
)
for name, value in outputs.items()
}
# An original block-vector readout can feed the same shell layout as any
# other readout. Convert coordinates explicitly at this boundary.
if self.output_layout is not None:
for name, prediction in list(predictions.items()):
if prediction.layout is not None:
continue
edges = torch.cat((x["edge_shift"], x["edge_index"]))
values = IrrepsBlockData.pack_matrix(
prediction.to_blocks(self.mapper),
edges,
self.output_layout,
operator=name,
)
predictions[name] = IrrepsBlockData.from_packed(
values, x, self.mapper, self.output_layout, operator=name
)
self._last_activation_mags = activation_mags or {}
return predictions
@property
def node_enc(self):
return self.pipeline.transforms["node_encoding"]
@property
def edge_enc(self):
return self.pipeline.transforms["edge_encoding"]
@property
def mp_blocks(self):
from types import SimpleNamespace
return [
SimpleNamespace(
node_upd=self.pipeline.transforms[f"node_message_{index}"],
edge_upd=self.pipeline.transforms[f"edge_message_{index}"],
)
for index in range(self.cfg.num_layers_gnn)
]
@property
def heads(self):
return {
name.removeprefix("readout_"): module
for name, module in self.pipeline.transforms.items()
if self.pipeline.roles[name] == "readout"
}
# ==================== Lightning steps ====================================
def _graph_step(self, batch, batch_idx, stage: str, predictions=None):
"""
Shared logic for training and validation steps.
Logs metrics prefixed with stage ('train' or 'val').
"""
x, y = batch
supervision = {
name: record
for name, record in y.items()
if isinstance(record, dict) and "values" in record
}
if supervision and self.cfg.matrix_loss_reduction != "pair_sum":
raise ValueError(
"matrix_loss_reduction applies to orbital_blocks, not packed supervision"
)
y = dict(y)
for name, record in supervision.items():
y[name] = IrrepsBlockData.from_packed(
record["values"],
x,
self.mapper,
self.output_layout,
frame=self.pipeline.feature_specs[self.pipeline.outputs[name]].frame,
operator=name,
).to_blocks(self.mapper)
if self._objective_terms:
# Plugins must see which reference entries were observed,
# including missing values zero-filled by the packed loader.
edges = x.get("edge_keys")
if edges is None:
edges = torch.cat((x["edge_shift"], x["edge_index"]))
rows = {tuple(edge.tolist()): i for i, edge in enumerate(edges.T)}
y[name]._coverage_masks = {
pair: self.output_layout.unpack_blocks(
tuple(pair.split("-")),
record["mask"][
[rows[tuple(edge.tolist())] for edge in pair_edges.T]
],
)
for pair, pair_edges in y[name].pair_edges.items()
}
metrics = {}
record_timings = bool(self.cfg.benchmark)
t_fwd_start = time.perf_counter() if record_timings else 0.0
preds_irreps = (
predictions
if predictions is not None
else self(
x,
physical_derivatives=bool(
self.cfg.train_on_forces or self.cfg.train_on_stress
),
)
)
t_fwd_end = time.perf_counter() if record_timings else 0.0
# --- block mapping timing ---------------------------------------
t_map_start = time.perf_counter() if record_timings else 0.0
preds_matrix = {
name: preds_irreps[name].to_blocks(self.mapper)
for name in self.cfg.matrix_targets
}
t_map_end = time.perf_counter() if record_timings else 0.0
# --- Matrix Loss Calculation ------------------------------------
combined_matrix_losses = {}
combined_pair_losses: dict[str, dict[str, torch.Tensor]] = {}
physical_pred_matrices: dict[str, BlockMatrix] = {}
target_matrices: dict[str, BlockMatrix] = {
name: (
y[name].to_blocks(self.mapper)
if isinstance(y[name], IrrepsBlockData)
else y[name]
)
for name in ("hamiltonian", "overlap", "density")
if name in y
}
from core.basis_converter import OpenMXE3NNConverter
converter = OpenMXE3NNConverter(self.mapper.orbital_cfg)
for name, matrix in target_matrices.items():
basis = preds_matrix.get(name, next(iter(preds_matrix.values()))).basis
target_matrices[name] = (
converter.matrix_to_openmx(matrix)
if basis == "openmx"
else converter.matrix_to_e3nn(matrix)
)
hamiltonian_mae_contribs: dict[str, object] | None = None
irrep_block_cache_by_name: dict[
str, tuple[dict[str, BlockMatrix], dict[str, BlockMatrix]]
] = {}
allow_train_metrics = stage != "train" or bool(self.cfg.log_train_metrics)
edge_loss_weights = self._edge_loss_weights(x)
# num_atoms = x["positions"].shape[0]
def _compute_physical_matrix_metrics(
pred_matrix: BlockMatrix,
target_matrix: BlockMatrix,
matrix_name: str,
) -> tuple[torch.Tensor, torch.Tensor]:
mse_val = torch.tensor(0.0, device=self.device)
mae_val = torch.tensor(0.0, device=self.device)
for key in target_matrix.pair_blocks.keys():
if key not in pred_matrix.pair_blocks:
raise ValueError(f"Key {key} not found in predicted items.")
preds = pred_matrix.pair_blocks[key]
targets = target_matrix.pair_blocks[key]
target_n = targets.shape[0]
if target_n <= 0:
continue
if preds.shape[0] < target_n:
raise ValueError(
f"Predicted blocks for matrix {matrix_name}, key {key} are too short: "
f"pred_len={preds.shape[0]} target_len={target_n}"
)
preds = preds[:target_n]
targets = targets[:target_n]
pred_edges = pred_matrix.pair_edges[key][:, :target_n]
target_edges = target_matrix.pair_edges[key][:, :target_n]
if self.cfg.safety_checks:
assert torch.equal(
pred_edges, target_edges
), f"Edge mismatch in metric block loss for matrix {matrix_name}, key {key}."
partial_mask = self._partial_train_mask(target_edges)
if partial_mask is not None:
preds = preds[partial_mask]
targets = targets[partial_mask]
if preds.shape[0] == 0:
continue
if self.cfg.matrix_loss_reduction != "pair_sum":
mask = target_matrix.coverage_masks.get(key)
mask = (
torch.ones_like(
target_matrix.pair_blocks[key], dtype=torch.bool
)
if mask is None
else mask
)[:target_n]
if partial_mask is not None:
mask = mask[partial_mask]
diff = (preds - targets)[mask].detach()
if not diff.numel():
continue
else:
diff = (preds - targets).detach()
mse_val += torch.mean(diff.abs().square())
mae_val += torch.mean(torch.abs(diff))
return mse_val, mae_val
def _symmetrize_aligned(
matrix: BlockMatrix,
alignment: dict[str, tuple[str, torch.Tensor]] | None,
) -> BlockMatrix:
if alignment is None:
raise ValueError(
"Symmetrization requires precomputed reverse alignment."
)
transposed = matrix.transpose_aligned(alignment).conjugate()
return matrix.add_aligned(transposed) * 0.5
for name in self.cfg.matrix_targets:
per_irrep_metrics: dict[str, dict[str, torch.Tensor]] = {}
pred_irrep_blocks = None
target_irrep_blocks = None
pair_losses: dict[str, torch.Tensor] = {}
raw_pred_matrix = preds_matrix[name]
raw_target = y[name]
raw_target_matrix = target_matrices[name]
loss_pred_matrix, loss_target_matrix, metric_pred_matrix = (
self._apply_matrix_envelope_mode(
name=name,
pred_matrix=raw_pred_matrix,
target_matrix=raw_target_matrix,
x=x,
)
)
physical_pred_matrices[name] = metric_pred_matrix
raw_metric_pred_matrix = metric_pred_matrix
matrix_edge_weights = (
edge_loss_weights if name in {"hamiltonian", "overlap"} else None
)
need_irrep_cache = allow_train_metrics and (
self.cfg.log_per_irrep_metrics
or (
name == "hamiltonian"
and (
self.cfg.log_hamiltonian_irrep_contrib_metrics
or self.cfg.log_hamiltonian_pair_contrib_metrics
)
)
)
need_hamiltonian_contribs = allow_train_metrics and (
name == "hamiltonian"
and (
self.cfg.log_hamiltonian_irrep_contrib_metrics
or self.cfg.log_hamiltonian_pair_contrib_metrics
)
)
if name in supervision and (
self.cfg.train_on_irrep_parts or self.cfg.train_target == "irreps"
):
record = supervision[name]
expected = torch.stack(
[
self.output_layout.packed_mask(
(
raw_target_matrix.atoms[int(edge[3])],
raw_target_matrix.atoms[int(edge[4])],
),
device=record["mask"].device,
)
for edge in preds_irreps[name].packed_edges.T
]
)
if not torch.equal(record["mask"], expected):
raise ValueError(
"Irrep-part objectives require complete shell blocks; use a masked matrix objective for partial shells"
)
loss_mse_val, loss_mae_val, per_irrep_losses = (
self._compute_irrep_part_losses(
loss_pred_matrix.to_vectors(
self.mapper, layout=self.output_layout
),
loss_target_matrix.to_vectors(
self.mapper, layout=self.output_layout
),
)
)
metric_mse_val, metric_mae_val = _compute_physical_matrix_metrics(
metric_pred_matrix, raw_target_matrix, name
)
elif name in supervision:
from net.training_policies import masked_error
record = supervision[name]
prediction_record = preds_irreps[name]
p, t = loss_pred_matrix, loss_target_matrix
symmetry = self.cfg.symmetrize_output or (
stage != "train" and self.cfg.symmetrize_validation_output
)
if symmetry:
p = _symmetrize_aligned(
p,
x.get("pred_reverse_alignment", x.get("pred_trace_alignment")),
)
metric_pred_matrix = _symmetrize_aligned(
metric_pred_matrix,
x.get("pred_reverse_alignment", x.get("pred_trace_alignment")),
)
transform = symmetry or self._matrix_envelope_mode() != "off"
frame = prediction_record.frames
values = (
IrrepsBlockData.pack_matrix(
p,
prediction_record.packed_edges,
self.output_layout,
frames=frame,
operator=name,
)
if transform
else prediction_record.packed_values
)
targets = (
IrrepsBlockData.pack_matrix(
t,
prediction_record.packed_edges,
self.output_layout,
frames=frame,
operator=name,
)
if transform
else record["values"]
)
mask = record["mask"].clone()
edge_mask = self._partial_train_mask(prediction_record.packed_edges)
if edge_mask is not None:
mask &= edge_mask.reshape((-1,) + (1,) * (mask.ndim - 1))
options = self.cfg.objective or {}
convention = options.get("complex_reduction", "magnitude")
loss_mse_val = masked_error(
values,
targets,
mask,
kind="mse",
complex_reduction=convention,
weights=matrix_edge_weights,
)
loss_mae_val = masked_error(
values,
targets,
mask,
kind="mae",
complex_reduction=convention,
weights=matrix_edge_weights,
)
unit = HARTREE_TO_EV if name == "hamiltonian" else 1.0
metric_values = (
IrrepsBlockData.pack_matrix(
metric_pred_matrix,
prediction_record.packed_edges,
self.output_layout,
frames=frame,
operator=name,
)
if transform
else values
)
metric_mse_val = (
masked_error(
metric_values,
record["values"],
mask,
kind="mse",
complex_reduction=convention,
).detach()
/ unit**2
)
metric_mae_val = (
masked_error(
metric_values,
record["values"],
mask,
kind="mae",
complex_reduction=convention,
).detach()
/ unit
)
physical_pred_matrices[name] = metric_pred_matrix
elif self.cfg.train_on_irrep_parts:
if self._matrix_envelope_mode() != "off":
raise ValueError(
"hamiltonian_envelope_mode is not supported with train_on_irrep_parts."
)
target_irreps = raw_target_matrix.to_vectors(
self.mapper, layout=self.output_layout
)
prediction_irreps = (
raw_pred_matrix.to_vectors(self.mapper, layout=self.output_layout)
if self.output_layout is not None
else preds_irreps[name]
)
loss_mse_val, loss_mae_val, per_irrep_metrics = (
self._compute_irrep_part_losses(prediction_irreps, target_irreps)
)
metric_mse_val = loss_mse_val
metric_mae_val = loss_mae_val
elif self.cfg.train_target == "irreps":
prediction_irreps = preds_irreps[name]
if self.output_layout is not None:
prediction_irreps = raw_pred_matrix.to_vectors(
self.mapper, layout=self.output_layout
)
raw_target = raw_target_matrix.to_vectors(
self.mapper, layout=self.output_layout
)
loss_mse_val = torch.tensor(0.0, device=self.device)
loss_mae_val = torch.tensor(0.0, device=self.device)
for key, targets in raw_target.pair_vectors.items():
if key not in prediction_irreps.pair_vectors:
raise ValueError(
f"Key {key} not found in predicted irreps for {name}."
)
preds = prediction_irreps.pair_vectors[key]
target_n = targets.shape[0]
if preds.shape[0] < target_n:
raise ValueError(
f"Predicted vectors for matrix {name}, key {key} are too "
f"short: pred_len={preds.shape[0]} target_len={target_n}"
)
pred_edges = prediction_irreps.pair_edges[key][:, :target_n]
target_edges = raw_target.pair_edges[key][:, :target_n]
if self.cfg.safety_checks:
assert torch.equal(
pred_edges, target_edges
), f"Edge mismatch in irrep loss for matrix {name}, key {key}."
partial_mask = self._partial_train_mask(target_edges)
preds = preds[:target_n]
if partial_mask is not None:
preds = preds[partial_mask]
targets = targets[partial_mask]
if preds.shape[0] == 0:
continue
pair_mse = self._mse(preds, targets)
pair_mae = self._mae(preds, targets)
loss_mse_val += pair_mse
loss_mae_val += pair_mae
pair_losses[key] = (
1 - self._objective_l1_fraction()
) * pair_mse + self._objective_l1_fraction() * pair_mae
if (
stage != "train" and self.cfg.symmetrize_validation_output
) or self.cfg.symmetrize_output:
metric_pred_matrix = _symmetrize_aligned(
metric_pred_matrix,
x.get(
"pred_reverse_alignment",
x.get("pred_trace_alignment"),
),
)
metric_mse_val, metric_mae_val = _compute_physical_matrix_metrics(
metric_pred_matrix, raw_target_matrix, name
)
else:
p = loss_pred_matrix
t = loss_target_matrix
loss_and_metric_share_prediction = metric_pred_matrix is p
should_symmetrize = self.cfg.symmetrize_output or (
stage != "train" and self.cfg.symmetrize_validation_output
)
if should_symmetrize:
p = _symmetrize_aligned(
p,
x.get("pred_reverse_alignment", x.get("pred_trace_alignment")),
)
if self._matrix_envelope_mode() == "normalize_target":
t = _symmetrize_aligned(
t,
x.get(
"target_reverse_alignment",
x.get("pred_trace_alignment"),
),
)
if should_symmetrize:
if loss_and_metric_share_prediction:
metric_pred_matrix = p
else:
metric_pred_matrix = _symmetrize_aligned(
metric_pred_matrix,
x.get(
"pred_reverse_alignment",
x.get("pred_trace_alignment"),
),
)
p_items, t_items = p.pair_blocks, t.pair_blocks
loss_mse_val = torch.tensor(0.0, device=self.device)
loss_mae_val = torch.tensor(0.0, device=self.device)
reduction = self.cfg.matrix_loss_reduction
diagnostic_mse_val = torch.tensor(0.0, device=self.device)
diagnostic_mae_val = torch.tensor(0.0, device=self.device)
reduction_count = 0
coverage = t.coverage_masks if reduction != "pair_sum" else {}
# Vectorized loss calculation
for key in t_items.keys():
if key not in p_items.keys():
raise ValueError(f"Key {key} not found in predicted items.")
preds = p_items[key]
targets = t_items[key]
target_n = targets.shape[0]
if target_n <= 0:
continue
if preds.shape[0] < target_n:
raise ValueError(
f"Predicted blocks for matrix {name}, key {key} are too short: "
f"pred_len={preds.shape[0]} target_len={target_n}"
)
preds = preds[:target_n]
targets = targets[:target_n]
pred_edges = p.pair_edges[key][:, :target_n]
target_edges = t.pair_edges[key][:, :target_n]
if self.cfg.safety_checks:
assert torch.equal(
pred_edges, target_edges
), f"Edge mismatch in block loss for matrix {name}, key {key}."
partial_mask = self._partial_train_mask(target_edges)
if partial_mask is not None:
preds = preds[partial_mask]
targets = targets[partial_mask]
if preds.shape[0] == 0:
continue
if reduction != "pair_sum":
mask = coverage.get(key)
mask = (
torch.ones_like(t_items[key], dtype=torch.bool)
if mask is None
else mask
)[:target_n]
if partial_mask is not None:
mask = mask[partial_mask]
observed = self._observed_block_losses(preds, targets, mask)
if observed is None:
continue
weighted_pair_mse, weighted_pair_mae, entry_count = observed
diagnostic_mse_val += weighted_pair_mse.detach()
diagnostic_mae_val += weighted_pair_mae.detach()
factor = entry_count if reduction == "elements" else 1
loss_mse_val += weighted_pair_mse * factor
loss_mae_val += weighted_pair_mae * factor
reduction_count += factor
pair_losses[key] = factor * (
(1 - self._objective_l1_fraction()) * weighted_pair_mse
+ self._objective_l1_fraction() * weighted_pair_mae
)
continue
edge_weights_selected = None
if matrix_edge_weights is not None:
selected_global_idx = x["edge_partitions"][key]["global_idx"][
:target_n
]
if partial_mask is not None:
selected_global_idx = selected_global_idx[partial_mask]
edge_weights_selected = matrix_edge_weights.index_select(
0, selected_global_idx
)
weighted_pair_mse, weighted_pair_mae = (
self._weighted_edge_block_losses(
preds=preds,
targets=targets,
edge_weights=edge_weights_selected,
)
)
loss_mse_val += weighted_pair_mse
loss_mae_val += weighted_pair_mae
pair_losses[key] = (
1 - self._objective_l1_fraction()
) * weighted_pair_mse + self._objective_l1_fraction() * weighted_pair_mae
if reduction != "pair_sum":
if reduction_count == 0:
raise ValueError("No observed matrix entries for objective")
loss_mse_val /= reduction_count
loss_mae_val /= reduction_count
pair_losses = {
key: value / reduction_count
for key, value in pair_losses.items()
}
if (
metric_pred_matrix is p
and t is raw_target_matrix
and matrix_edge_weights is None
):
# Historical diagnostics are sums of pair means even when
# the training objective has a different denominator.
metric_mse_val = (
loss_mse_val.detach()
if reduction == "pair_sum"
else diagnostic_mse_val
)
metric_mae_val = (
loss_mae_val.detach()
if reduction == "pair_sum"
else diagnostic_mae_val
)
else:
metric_mse_val, metric_mae_val = _compute_physical_matrix_metrics(
metric_pred_matrix,
raw_target_matrix,
name,
)
combined_pair_losses[name] = pair_losses
if need_irrep_cache:
pred_irrep_blocks = build_irrep_block_matrix_cache(
metric_pred_matrix,
self.mapper,
self.all_irreps,
layout=self.output_layout,
)
target_irrep_blocks = build_irrep_block_matrix_cache(
raw_target_matrix,
self.mapper,
self.all_irreps,
layout=self.output_layout,
)
irrep_block_cache_by_name[name] = (
pred_irrep_blocks,
target_irrep_blocks,
)
if need_hamiltonian_contribs:
hamiltonian_mae_contribs = compute_hamiltonian_mae_contributions(
metric_pred_matrix,
raw_target_matrix,
self.mapper,
all_irreps=self.all_irreps,
compute_irrep_sums=self.cfg.log_hamiltonian_irrep_contrib_metrics,
compute_pair_sums=self.cfg.log_hamiltonian_pair_contrib_metrics,
require_exact_prefix=bool(self.cfg.require_exact_edge_match),
pred_irrep_blocks=pred_irrep_blocks,
target_irrep_blocks=target_irrep_blocks,
)
# Store for combined loss BEFORE unit conversion
mse_for_loss = loss_mse_val
mae_for_loss = loss_mae_val
mse_val = metric_mse_val
mae_val = metric_mae_val
if name == "hamiltonian":
# Convert to eV^2 and eV for logging only
mse_val = metric_mse_val * (HARTREE_TO_EV**2)
if hamiltonian_mae_contribs is not None:
denom = int(hamiltonian_mae_contribs["total_count"])
if denom > 0:
mae_val = (
float(hamiltonian_mae_contribs["total_abs_sum"]) / denom
) * HARTREE_TO_EV
else:
mae_val = metric_mae_val * HARTREE_TO_EV
metrics[f"{stage}/{name}_mae"] = mae_val
metrics[f"{stage}/{name}_mse"] = mse_val
# Explicit physical-entry reduction, independent of diagnostic flags.
# Preserve the historical objectives/monitors of existing runs.
from net.operator_metrics import matrix_error_metrics
for metric, value in matrix_error_metrics(
raw_metric_pred_matrix,
raw_target_matrix,
scale=HARTREE_TO_EV if name == "hamiltonian" else 1.0,
edge_selector=self._partial_train_mask,
).items():
metrics[f"{stage}/{name}_{metric}"] = value
if metric_pred_matrix is not raw_metric_pred_matrix:
for metric, value in matrix_error_metrics(
metric_pred_matrix,
raw_target_matrix,
scale=HARTREE_TO_EV if name == "hamiltonian" else 1.0,
edge_selector=self._partial_train_mask,
).items():
metrics[f"{stage}/{name}_{metric}_projected"] = value
if self.cfg.log_per_irrep_metrics or self.cfg.train_on_irrep_parts:
for irrep_key, irrep_vals in per_irrep_metrics.items():
irrep_mse = irrep_vals["mse"]
irrep_mae = irrep_vals["mae"]
irrep_loss = (
1 - self._objective_l1_fraction()
) * irrep_mse + self._objective_l1_fraction() * irrep_mae
if name == "hamiltonian":
metrics[f"{stage}/{name}_irrep_{irrep_key}_mae"] = (
irrep_mae * HARTREE_TO_EV
)
metrics[f"{stage}/{name}_irrep_{irrep_key}_mse"] = irrep_mse * (
HARTREE_TO_EV**2
)
else:
metrics[f"{stage}/{name}_irrep_{irrep_key}_mae"] = irrep_mae
metrics[f"{stage}/{name}_irrep_{irrep_key}_mse"] = irrep_mse
if self.cfg.train_on_irrep_parts:
metrics[f"{stage}/{name}_irrep_{irrep_key}_loss"] = irrep_loss
options = self.cfg.objective or {}
mixture = self._objective_l1_fraction()
combined_matrix_losses[name] = options.get("weights", {}).get(name, 1.0) * (
(1 - mixture) * mse_for_loss + mixture * mae_for_loss
)
loss_matrix = sum(combined_matrix_losses.values())
# --- Observable Evaluation --------------------------------------
t_obs_start = time.perf_counter() if record_timings else 0.0
loss_E_weighted = torch.tensor(0.0, device=self.device)
loss_N_weighted = torch.tensor(0.0, device=self.device)
loss_F_weighted = torch.tensor(0.0, device=self.device)
loss_stress_weighted = torch.tensor(0.0, device=self.device)
loss_spectral_weighted = torch.tensor(0.0, device=self.device)
observable_preds_matrix = {}
observable_trace_alignment = x["pred_trace_alignment"]
H_true = None
D_true = None
S_true = None
needs_gt_observables = (
self.cfg.enable_energy
or self.cfg.enable_num_electrons
or self.cfg.log_partial_gt_observables
or self.cfg.train_observables_on_gt
)
if needs_gt_observables and {"hamiltonian", "density", "overlap"}.issubset(
target_matrices
):
H_true = target_matrices["hamiltonian"]
D_true = target_matrices["density"]
S_true = target_matrices["overlap"]
E_true = y.get("energy")
N_true = y.get("num_electrons")
for name, pred_matrix in physical_pred_matrices.items():
target_matrix = target_matrices.get(name)
if target_matrix is None:
continue
observable_preds_matrix[name] = truncate_pred_block_matrix_to_target_prefix(
pred_matrix,
target_matrix,
)
observable_values = build_observable_predictions(
observable_preds_matrix,
trace_alignment=observable_trace_alignment,
H_true=H_true,
D_true=D_true,
S_true=S_true,
)
observable_metric_preds_matrix = observable_preds_matrix
observable_metric_values = observable_values
if bool(getattr(self.cfg, "rescale_density_to_num_electrons", False)):
observable_metric_preds_matrix, num_electrons_pre_rescale = (
rescale_density_prediction_to_num_electrons(
observable_preds_matrix,
trace_alignment=observable_trace_alignment,
num_electrons_target=N_true,
overlap_true=S_true,
)
)
if num_electrons_pre_rescale is not None:
metrics[f"{stage}/num_electrons_mae_pre_rescale"] = torch.mean(
torch.abs(num_electrons_pre_rescale - N_true)
)
observable_metric_values = build_observable_predictions(
observable_metric_preds_matrix,
trace_alignment=observable_trace_alignment,
H_true=H_true,
D_true=D_true,
S_true=S_true,
)
add_observable_metrics(
metrics,
stage=stage,
cfg=self.cfg,
observable_values=observable_metric_values,
energy_target=E_true,
num_electrons_target=N_true,
)
# Training losses always use the original, unscaled predictions.
loss_E_weighted, loss_N_weighted = observable_loss(
cfg=self.cfg,
observable_values=observable_values,
energy_target=E_true,
num_electrons_target=N_true,
mse_fn=self._mse,
device=self.device,
)
energy_loss_active = bool(
self.cfg.train_on_energy
and float(self.cfg.loss_coef_observables) != 0.0
and E_true is not None
and (
(
self.cfg.train_observables_on_gt
and any(key.startswith("energy_gt_") for key in observable_values)
)
or (
not self.cfg.train_observables_on_gt
and "energy" in observable_values
)
)
)
num_electrons_loss_active = bool(
self.cfg.train_on_num_electrons
and float(self.cfg.loss_coef_observables) != 0.0
and N_true is not None
and (
(
self.cfg.train_observables_on_gt
and any(
key.startswith("num_electrons_gt_") for key in observable_values
)
)
or (
not self.cfg.train_observables_on_gt
and "num_electrons" in observable_values
)
)
)
force_loss_active = bool(
self.cfg.enable_forces
and self.cfg.train_on_forces
and float(self.cfg.loss_coef_forces) != 0.0
and y.get("forces") is not None
)
if self.cfg.enable_forces and y.get("forces") is not None:
forces_pred = self.get_forces(preds_irreps, x["positions"], x["box"], x=x)
forces_true = y["forces"]
metrics[f"{stage}/forces_mae"] = torch.mean(
torch.abs(forces_pred - forces_true)
)
metrics[f"{stage}/forces_mse"] = self._mse(forces_pred, forces_true)
if self.cfg.train_on_forces:
loss_F_weighted = self.cfg.loss_coef_forces * self._mse(
forces_pred, forces_true
)
stress_loss_active = bool(
self.cfg.enable_stress
and self.cfg.train_on_stress
and self.cfg.loss_coef_stress != 0.0
and y.get("stress") is not None
)
if self.cfg.enable_stress and y.get("stress") is not None:
stress_pred = self.get_stress(preds_irreps, x["positions"], x["box"], x=x)
metrics[f"{stage}/stress_mae"] = self._mae(stress_pred, y["stress"])
metrics[f"{stage}/stress_mse"] = self._mse(stress_pred, y["stress"])
if stress_loss_active:
loss_stress_weighted = (
self.cfg.loss_coef_stress * metrics[f"{stage}/stress_mse"]
)
metrics[f"{stage}/loss_stress_weighted"] = loss_stress_weighted
if allow_train_metrics and self.cfg.log_per_irrep_metrics:
for name in self.cfg.matrix_targets:
if name not in preds_irreps or name not in y:
continue
cache = irrep_block_cache_by_name.get(name)
if cache is None:
cache = (
build_irrep_block_matrix_cache(
preds_matrix[name],
self.mapper,
self.all_irreps,
layout=self.output_layout,
),
build_irrep_block_matrix_cache(
target_matrices[name],
self.mapper,
self.all_irreps,
layout=self.output_layout,
),
)
irrep_block_cache_by_name[name] = cache
pred_irrep_blocks, target_irrep_blocks = cache
irrep_metrics = compute_irrep_metrics(
preds_matrix[name],
target_matrices[name],
self.all_irreps,
self.mapper,
pred_irrep_blocks=pred_irrep_blocks,
target_irrep_blocks=target_irrep_blocks,
)
for key, value in irrep_metrics.items():
metrics[f"{stage}/{name}_irrep_{key}"] = value
t_obs_end = time.perf_counter() if record_timings else 0.0
spectral_loss_active = False
if (
self._spectral_loss_enabled()
and "hamiltonian" in observable_preds_matrix
and (stage != "train" or bool(getattr(self.cfg, "train_on_spectral", True)))
):
spectral_loss, spectral_stats = self._compute_spectral_loss(
x=x,
physical_pred_hamiltonian=observable_preds_matrix["hamiltonian"],
)
if bool(getattr(self.cfg, "train_on_spectral", True)):
loss_spectral_weighted = (
float(self.cfg.spectral_loss_coef) * spectral_loss
)
spectral_loss_active = float(self.cfg.spectral_loss_coef) != 0.0
metrics[f"{stage}/loss_spectral"] = spectral_loss
metrics[f"{stage}/loss_spectral_weighted"] = loss_spectral_weighted
metrics[f"{stage}/spectral_mae_ev"] = spectral_stats["spectral_mae_ev"]
metrics[f"{stage}/spectral_weight_sum"] = spectral_stats[
"spectral_weight_sum"
]
# --- Total Loss Aggregation ---
loss = (
loss_matrix
+ loss_E_weighted
+ loss_N_weighted
+ loss_F_weighted
+ loss_stress_weighted
+ loss_spectral_weighted
)
for term in self._objective_terms:
term_loss = term(
predictions=physical_pred_matrices, targets=target_matrices, inputs=x
)
weighted = term.weight * term_loss
loss = loss + weighted
metrics[f"{stage}/loss_term_{term.name}"] = term_loss
metrics[f"{stage}/loss_term_{term.name}_weighted"] = weighted
# L1 and L2 regularization
if self.cfg.l1_reg_coef > 0:
l1_reg = sum(p.abs().sum() for p in self.parameters())
loss += self.cfg.l1_reg_coef * l1_reg
metrics[f"{stage}/loss_l1_reg"] = l1_reg
if self.cfg.l2_reg_coef > 0:
l2_reg = sum(p.pow(2).sum() for p in self.parameters())
loss += self.cfg.l2_reg_coef * l2_reg
metrics[f"{stage}/loss_l2_reg"] = l2_reg
if self._stop_training_on_nonfinite_loss(
loss,
stage=stage,
batch_idx=batch_idx,
phase="total",
):
return None
# --- Logging ------------------------------------------------------
metrics[f"{stage}/loss"] = loss
metrics[f"{stage}/loss_total"] = loss
metrics[f"{stage}/loss_matrix_total"] = loss_matrix
metrics[f"{stage}/loss_block_total"] = loss_matrix
if energy_loss_active:
metrics[f"{stage}/loss_energy"] = loss_E_weighted / float(
self.cfg.loss_coef_observables
)
metrics[f"{stage}/loss_energy_weighted"] = loss_E_weighted
if num_electrons_loss_active:
metrics[f"{stage}/loss_num_electrons"] = loss_N_weighted / float(
self.cfg.loss_coef_observables
)
metrics[f"{stage}/loss_num_electrons_weighted"] = loss_N_weighted
if force_loss_active:
metrics[f"{stage}/loss_forces"] = loss_F_weighted
metrics[f"{stage}/loss_forces_weighted"] = loss_F_weighted
if spectral_loss_active:
metrics[f"{stage}/loss_spectral_total"] = loss_spectral_weighted
for name, loss_val in combined_matrix_losses.items():
metrics[f"{stage}/loss_block_{name}"] = loss_val
if self.cfg.log_per_pair_loss_metrics:
for matrix_name, pair_losses in combined_pair_losses.items():
for pair_key, pair_loss in pair_losses.items():
metrics[f"{stage}/loss_block_{matrix_name}_{pair_key}"] = pair_loss
if (
hamiltonian_mae_contribs is not None
and int(hamiltonian_mae_contribs["total_count"]) > 0
):
denom = int(hamiltonian_mae_contribs["total_count"])
if self.cfg.log_hamiltonian_irrep_contrib_metrics:
for irrep_key, abs_sum in hamiltonian_mae_contribs[
"irrep_abs_sums"
].items():
metrics[f"{stage}/hamiltonian_mae_{irrep_key}"] = (
abs_sum / denom
) * HARTREE_TO_EV
if self.cfg.log_hamiltonian_pair_contrib_metrics:
for pair_key, abs_sum in hamiltonian_mae_contribs[
"pair_abs_sums"
].items():
metrics[f"{stage}/hamiltonian_mae_{pair_key.replace('-', '_')}"] = (
abs_sum / denom
) * HARTREE_TO_EV
if stage == "train":
fraction_denom = loss.clamp_min(1e-12)
for name, val in combined_matrix_losses.items():
metrics[f"frac/loss_{name}"] = val / fraction_denom
if self.cfg.log_per_pair_loss_metrics:
for matrix_name, pair_losses in combined_pair_losses.items():
for pair_key, pair_loss in pair_losses.items():
metrics[f"frac/loss_{matrix_name}_{pair_key}"] = (
pair_loss / fraction_denom
)
if energy_loss_active:
metrics["frac/loss_energy"] = loss_E_weighted / fraction_denom
if num_electrons_loss_active:
metrics["frac/loss_num_electrons"] = loss_N_weighted / fraction_denom
if force_loss_active:
metrics["frac/loss_forces"] = loss_F_weighted / fraction_denom
if spectral_loss_active:
metrics["frac/loss_spectral"] = loss_spectral_weighted / fraction_denom
# record per-batch timings for callback
if record_timings:
self._last_batch_times = {
"forward": t_fwd_end - t_fwd_start,
"map": t_map_end - t_map_start,
"obs": t_obs_end - t_obs_start,
}
if stage == "val" and bool(
getattr(self, "_capture_artifact_validation", False)
):
self._artifact_validation_payloads.append((x, y, preds_irreps))
weight = (
next(iter(supervision.values()))["values"].shape[0]
if supervision
and (self.cfg.objective or {}).get("epoch_weight", "edges") == "edges"
else (
int(next(iter(supervision.values()))["mask"].sum())
if supervision
else 1
)
)
return loss, metrics, weight
def _split_batch_predictions(self, predictions, graphs, union):
result = []
node_offset, edge_offset = 0, 0
for graph in graphs:
x = dict(graph[0])
from data.graph_features import complete_graph_metadata
complete_graph_metadata(
x,
self.mapper,
species_order=self.cfg.species.get("order")
or (
tuple(self.output_layout.species_shells)
if self.output_layout is not None
else None
),
separate_shifted_self=self.legacy_separate_shifted_self,
)
n, e = len(x["node_type_idx"]), x["edge_index"].shape[1]
for key in ("bond_frames", "edge_envelope", "edge_r0"):
if key in union:
x[key] = union[key][edge_offset : edge_offset + e]
records = {}
for name, prediction in predictions.items():
if prediction.packed_values is not None:
records[name] = IrrepsBlockData.from_packed(
prediction.packed_values[edge_offset : edge_offset + e],
x,
self.mapper,
self.output_layout,
frame=self.pipeline.feature_specs[
self.pipeline.outputs[name]
].frame,
operator=name,
)
else:
vectors, edges, lookup = {}, {}, {}
for pair, values in prediction.pair_vectors.items():
indices = prediction.pair_edges[pair]
selected = (indices[3] >= node_offset) & (
indices[3] < node_offset + n
)
if not selected.any():
continue
vectors[pair] = values[selected]
edges[pair] = indices[:, selected].clone()
edges[pair][3:] -= node_offset
lookup.update(
{
tuple(edge): (pair, row)
for row, edge in enumerate(edges[pair].T.tolist())
}
)
records[name] = IrrepsBlockData(
tuple(x["atoms"]),
x["atom_counts"],
vectors,
edges,
lookup,
self.mapper.orbital_cfg,
prediction.basis,
)
result.append(((x, graph[1]), records))
node_offset += n
edge_offset += e
return result
def _shared_step(self, batch, batch_idx, stage):
from data.sample import GraphBatch, union_graph_inputs
graphs = batch.samples if isinstance(batch, GraphBatch) else [batch]
with torch.set_grad_enabled(
torch.is_grad_enabled() or self.cfg.enable_forces or self.cfg.enable_stress
):
if len(graphs) > 1:
from data.graph_topology import prepare_graph_topology
graphs = [
(
prepare_graph_topology(
dict(graph[0]), self.cfg.data.get("supports")
),
graph[1],
)
for graph in graphs
]
union = union_graph_inputs(graphs, layouts=self.pipeline.input_layouts)
predictions = self(
union,
physical_derivatives=bool(
self.cfg.train_on_forces or self.cfg.train_on_stress
),
)
separated = self._split_batch_predictions(predictions, graphs, union)
results = [
self._graph_step(graph, batch_idx, stage, predictions=prediction)
for graph, prediction in separated
]
else:
results = [self._graph_step(graphs[0], batch_idx, stage)]
if any(result is None for result in results):
return None
keys = set.intersection(*(set(result[1]) for result in results))
metrics = {
key: sum(result[1][key] for result in results) / len(results)
for key in keys
}
matrix_loss = 0.0
for name in self.cfg.matrix_targets:
weights = []
for graph in graphs:
record = graph[1][name]
if isinstance(record, dict) and "mask" in record:
mask = record["mask"]
edges = torch.cat((graph[0]["edge_shift"], graph[0]["edge_index"]))
selected = self._partial_train_mask(edges)
weights.append(
int((mask if selected is None else mask[selected]).sum())
)
else:
weights.append(1)
if sum(weights) == 0:
raise ValueError(f"Batch has no supervised entries for {name}")
key = f"{stage}/loss_block_{name}"
matrix_loss = matrix_loss + sum(
result[1][key] * weight / sum(weights)
for result, weight in zip(results, weights)
)
for metric in (key, f"{stage}/{name}_mae", f"{stage}/{name}_mse"):
if metric in metrics:
metrics[metric] = sum(
result[1][metric] * weight / sum(weights)
for result, weight in zip(results, weights)
)
# Matrix objectives retain the source's reduction over masked entries.
# Observable losses and regularization reduce over independent cells.
loss = matrix_loss + sum(
result[0] - result[1][f"{stage}/loss_matrix_total"] for result in results
) / len(results)
metrics.update(
{
f"{stage}/loss": loss,
f"{stage}/loss_total": loss,
f"{stage}/loss_matrix_total": matrix_loss,
f"{stage}/loss_block_total": matrix_loss,
}
)
entry_metrics = {
key: value for key, value in metrics.items() if "_elementwise" in key
}
self.log_dict(
{key: value for key, value in metrics.items() if key not in entry_metrics},
prog_bar=True,
on_step=self.cfg.log_on_step,
on_epoch=self.cfg.log_on_epoch,
batch_size=sum(result[2] for result in results),
)
self.log_dict(
entry_metrics,
prog_bar=False,
on_step=self.cfg.log_on_step,
on_epoch=self.cfg.log_on_epoch,
batch_size=len(results),
)
return loss
def transfer_batch_to_device(self, batch, device, dataloader_idx):
from data.sample import GraphSample, GraphBatch, move_tree
return (
batch.to(device)
if isinstance(batch, (GraphSample, GraphBatch))
else move_tree(batch, device)
)
def training_step(self, batch, batch_idx):
return self._shared_step(batch, batch_idx, stage="train")
def validation_step(self, batch, batch_idx):
return self._shared_step(batch, batch_idx, stage="val")
def test_step(self, batch, batch_idx):
return self._shared_step(batch, batch_idx, stage="test")
def on_train_epoch_start(self):
"""Log learning rate at the beginning of each training epoch."""
if self.trainer.sanity_checking:
return
self._grad_norm_logged_this_epoch = False
# Get the first optimizer
optimizer = self.optimizers()
if not isinstance(optimizer, list):
optimizer = [optimizer]
lr = optimizer[0].param_groups[0]["lr"]
self.log("lr", lr, on_step=False, on_epoch=True, prog_bar=True, logger=True)
def on_before_optimizer_step(self, optimizer) -> None:
if not bool(getattr(self.cfg, "log_grad_norm", False)):
return
if bool(getattr(self, "_grad_norm_logged_this_epoch", False)):
return
total_norm_sq = torch.tensor(0.0, device=self.device)
saw_grad = False
for param in self.parameters():
if param.grad is None:
continue
grad = param.grad.detach()
total_norm_sq = total_norm_sq + torch.sum(grad * grad)
saw_grad = True
if saw_grad:
self._grad_norm_logged_this_epoch = True
self.log(
"grad_norm",
torch.sqrt(total_norm_sq),
on_step=False,
on_epoch=True,
prog_bar=False,
logger=True,
batch_size=1,
)
# ------------------------------------------------------------------ optimiser
def configure_optimizers(self):
trainable_params = [p for p in self.parameters() if p.requires_grad]
if not trainable_params:
raise ValueError("No trainable parameters remain after freeze settings.")
from net.training_policies import make_optimizer, make_scheduler
options = self.cfg.optimization
if options is None:
options = {"optimizer": {"kind": "adamw", "lr": self.cfg.lr}}
if self.cfg.use_lr_scheduler:
options["scheduler"] = {
"kind": "plateau",
"ema_alpha": 1.0,
"factor": self.cfg.lr_scheduler_factor,
"patience": self.cfg.lr_scheduler_patience,
"min_lr": self.cfg.lr_scheduler_min_lr,
"monitor": self.cfg.lr_scheduler_target,
}
optimizer = make_optimizer(trainable_params, options["optimizer"])
scheduler = make_scheduler(optimizer, options.get("scheduler"))
return (
optimizer
if scheduler is None
else {"optimizer": optimizer, "lr_scheduler": scheduler}
)
def lr_scheduler_step(self, scheduler, metric):
"""Advance the declared scheduler before checking the run's LR stop rule."""
if metric is None:
scheduler.step()
else:
scheduler.step(metric)
if self.cfg.optimization is not None:
minimum = self.cfg.optimization.get("stop_lr")
if (
minimum is not None
and max(group["lr"] for group in scheduler.optimizer.param_groups)
<= minimum
):
self.trainer.should_stop = True
# ------------------------------------------------------------------ force prediction
def predictions_to_snapshot(
self,
predictions: Dict[str, IrrepsBlockData],
positions: torch.Tensor,
box: torch.Tensor,
x: Dict[str, Any] | None = None,
) -> "Snapshot":
required = {"hamiltonian", "overlap", "density"}
missing = sorted(required.difference(predictions))
if missing:
raise ValueError(
"Energy-derived force/stress evaluation requires predictions for "
f"hamiltonian, overlap, and density; missing: {', '.join(missing)}."
)
if x is None and self._matrix_envelope_mode() != "off":
raise ValueError(
"predictions_to_snapshot requires the input batch `x` when "
"hamiltonian_envelope_mode is enabled."
)
block_matrices = self.predicted_irreps_to_block_matrices(
predictions,
x if x is not None else {},
physical=True,
)
return Snapshot(
hamiltonian=(
block_matrices["hamiltonian"]
if "hamiltonian" in block_matrices
else None
),
overlap=(
block_matrices["overlap"] if "overlap" in block_matrices else None
),
density=(
block_matrices["density"] if "density" in block_matrices else None
),
positions=positions,
box=box,
)
def get_forces(
self,
predictions: Dict[str, IrrepsBlockData],
positions: torch.Tensor,
box: torch.Tensor,
*,
x=None,
) -> torch.Tensor:
"""
Compute forces from predictions and positions.
Args:
predictions: Output from the forward pass, containing hamiltonian,
overlap, and density matrices.
positions: Atomic positions (N, 3).
box: Lattice box matrix (3, 3).
Returns:
Forces as a tensor of shape (N, 3).
Comments:
- Forces are computed as -∂E/∂r, where E is the energy from the hamiltonian.
- No Pulay correction needed
"""
snapshot = self.predictions_to_snapshot(predictions, positions, box, x=x)
energy = snapshot.get_energy()
grad_pos = torch.autograd.grad(
energy,
positions,
create_graph=self.cfg.train_on_forces, # needed for second derivatives
retain_graph=True,
)[0]
return -grad_pos
def get_box_grad(
self,
predictions: Dict[str, IrrepsBlockData],
positions: torch.Tensor,
box: torch.Tensor,
*,
x=None,
) -> torch.Tensor:
"""
Compute gradient of energy with respect to box.
Args:
predictions: Output from the forward pass, containing hamiltonian,
overlap, and density matrices.
positions: Atomic positions (N, 3).
box: Lattice box matrix (3, 3).
Returns:
Gradient of energy with respect to box as a tensor of shape (3, 3).
"""
snapshot = self.predictions_to_snapshot(predictions, positions, box, x=x)
energy = snapshot.get_energy()
# 1. get dE/dh
grad_box = torch.autograd.grad(
energy,
box,
create_graph=self.cfg.train_on_stress, # needed for second derivatives
retain_graph=True,
)[0]
return grad_box
def get_stress(
self,
predictions: Dict[str, IrrepsBlockData],
positions: torch.Tensor,
box: torch.Tensor,
symmetrize: bool = True,
*,
x=None,
) -> torch.Tensor:
"""
Compute stress tensor from energy and box.
Args:
energy: Scalar energy value.
box: Lattice box matrix (3, 3).
Returns:
Stress tensor as a (3, 3) tensor.
Comments:
- Stress is defined as σ_{αβ} = (1/Ω) ∑_γ h_{γα} (dE/dh_{γβ})
- No Pulay correction needed
"""
# 1. get dE/dh
grad_box = self.get_box_grad(predictions, positions, box, x=x)
# 2. compute volume
volume = torch.det(box)
# 3. compute stress tensor σ_{αβ} = (1/Ω) ∑_γ h_{γα} (dE/dh_{γβ})
# stress = (1/Ω) * box^T @ grad_box
stress = torch.matmul(box.t(), grad_box) / volume
# 4. optionally symmetrize: σ -> (σ+σ^T)/2
if symmetrize:
stress = 0.5 * (stress + stress.transpose(-1, -2))
return stress
def predict_forces(self, x: Dict[str, Any]) -> torch.Tensor:
x = dict(x)
if not x["positions"].requires_grad:
x["positions"] = x["positions"].clone().detach().requires_grad_(True)
predictions = self(x, physical_derivatives=True)
return self.get_forces(predictions, x["positions"], x["box"], x=x)
def predict_stress(self, x: Dict[str, Any]) -> torch.Tensor:
x = dict(x)
if x.get("box") is not None and not x["box"].requires_grad:
x["box"] = x["box"].clone().detach().requires_grad_(True)
predictions = self(x, physical_derivatives=True)
return self.get_stress(predictions, x["positions"], x["box"], x=x)