ChemicalSpeciesToAtomTypeMapper

class torch_sim.models.nequip_framework.ChemicalSpeciesToAtomTypeMapper(chemical_symbols)

Bases: object

Maps atomic numbers to model-specific atom type indices.

This class provides functionality to map atomic numbers to the corresponding atom type indices used by the model. It handles cases where the model’s internal representation of atom types may differ from conventional chemical species, such as when modeling different charge states of the same element.

The mapping is created using a lookup table that converts atomic numbers to zero-based indices based on the provided list of chemical symbols. The order of chemical symbols must match the order of atom types expected by the model.

NOTE: This is adapted from the NequIP package.

Variables:

lookup_table (torch.Tensor) – Tensor mapping atomic numbers to model type indices. Contains -1 for unmapped atomic numbers.

Parameters:

chemical_symbols (list[str]) – List of chemical symbols in the order matching the model’s internal type ordering. Each symbol must be a valid chemical element symbol.

Raises:

AssertionError – If an invalid chemical symbol is provided.