Inference

The inference package provides utilities for making predictions with trained models.

Predictor

class aenet.torch_training.inference.Predictor(model, descriptor, normalizer: NormalizationManager, atomic_energies: Dict[str, float] | None = None, device: torch.device = torch.device, dtype: torch.dtype = torch.float64)

Bases: object

Handles prediction for trained models.

Parameters:

• model (nn.Module) – Trained model (EnergyModelAdapter).

• descriptor (ChebyshevDescriptor) – Descriptor for featurization.

• normalizer (NormalizationManager) – Normalization manager with trained statistics.

• atomic_energies (dict, optional) – Atomic reference energies per species. These are added back to model predictions to get total energies. Defaults to zeros if not provided.

• device (torch.device) – Device for computation.

• dtype (torch.dtype) – Data type for tensors.

cohesive_energy(structure: Structure, atomic_energies: Dict[str, float] | None = None) → float

Compute cohesive energy from a structure with total energy.

Parameters:

• structure (Structure) – Structure containing total energy and species list.

• atomic_energies (dict, optional) – Per-species atomic reference energies. If None, uses predictor’s stored E_atomic.

Returns:

Cohesive energy (total - sum of atomic reference energies).

Return type:

float

Raises:

ValueError – If no atomic energies are available.

predict(structures: List[Structure], eval_forces: bool = False, return_atom_energies: bool = False, track_timing: bool = False, batch_size: int | None = None, num_workers: int | None = None, prefetch_factor: int | None = None, persistent_workers: bool | None = None) → Tuple[List[float], List[torch.Tensor] | None, List[torch.Tensor] | None, Dict[str, List[float]] | None]

Predict energies (and optionally forces) for structures.

Parameters:

• structures (list of Structure) – Structures to predict on.

• eval_forces (bool) – Whether to also predict forces.

• return_atom_energies (bool) – Whether to return per-atom energies.

• track_timing (bool) – Whether to track and return timing information.

Returns:

• energies (list of float) – Predicted energies (total energies if energy_target=’cohesive’, otherwise model outputs directly).

• forces (list of Tensor or None) – Predicted forces (N_i, 3) per structure if requested, otherwise None.

• atom_energies (list of Tensor or None) – Per-atom energies (N_i,) per structure if requested, otherwise None.

• timing (dict or None) – Timing information per structure if requested, otherwise None. Keys: ‘featurization’, ‘energy_eval’, ‘force_eval’, ‘total’

predict_dataset(dataset, eval_forces: bool = False, return_atom_energies: bool = False, track_timing: bool = False, batch_size: int | None = None, num_workers: int | None = None, prefetch_factor: int | None = None, persistent_workers: bool | None = None) → Tuple[List[float], List[torch.Tensor] | None, List[torch.Tensor] | None, Dict[str, List[float]] | None, Dict[str, Any]]

Predict energies from a dataset that already yields featurized samples.

This path is intended for PyTorch-only dataset-backed inference where reusing cached features is desirable.