def max_cutoff(self):
return get_max_cutoff(self.cell)
def __init__(
self,
cell: "ndarray",
species: Union[List[str], List[int]],
positions: "ndarray",
mass_dict: dict = None,
prev_positions: "ndarray" = None,
species_labels: List[str] = None,
forces=None,
stds=None,
energy: float = None,
):
# Define cell (each row is a Bravais lattice vector).
self.cell = np.array(cell)
# Compute the max cutoff compatible with a 3x3x3 supercell of the
# structure.
self.max_cutoff = get_max_cutoff(self.cell)
# Set positions.
self.positions = np.array(positions)
# If species are strings, convert species to integers by atomic number
if species_labels is None:
self.species_labels = species
else:
self.species_labels = species_labels
self.coded_species = np.array([element_to_Z(spec) for spec in species])
self.nat = len(species)
# Default: atoms have no velocity
if prev_positions is None:
self.prev_positions = np.copy(self.positions)
else:
assert len(positions) == len(
prev_positions
), "Previous positions and positions are not same length"
self.prev_positions = prev_positions
# Set forces, energies, and stresses and their uncertainties.
if forces is not None:
self.forces = np.array(forces)
else:
self.forces = np.zeros((len(positions), 3))
if stds is not None:
self.stds = np.array(stds)
else:
self.stds = np.zeros((len(positions), 3))
self.energy = energy
self.local_energies = None
self.local_energy_stds = None
self.partial_stresses = None
self.partial_stress_stds = None
self.stress = None
self.stress_stds = None
# Potential energy attribute needed to mirror ASE atoms object.
self.potential_energy = None
self.mass_dict = mass_dict
# Convert from elements to atomic numbers in mass dict
if mass_dict is not None:
keys = list(mass_dict.keys())
for elt in keys:
if isinstance(elt, str):
mass_dict[element_to_Z(elt)] = mass_dict[elt]
if elt.isnumeric():
mass_dict[int(elt)] = mass_dict[elt]