# MIT License
#
# Copyright (c) 2025 Isacks-Co contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
import hashlib
import numpy as np
import logging
import os
from ase import Atoms
from ase.build import supercells
from ase.io import read
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation
from ase.neighborlist import NeighborList, natural_cutoffs
from .potential import Potential
logger = logging.getLogger(__name__)
[docs]
class AtomicStructure:
"""Wrapper for the ASE atoms object and a Potential.
This class provides a way to represent an atomic structure inluding its
potential. It provides methods to compute relevant instantaneous quantities as
well as some conveniance features for a better MD workflow.
Attributes:
_atoms (ASE Atoms): The underlying atoms object
_label (str): A unique hashlabel for the structure
"""
def __init__(self, atoms: Atoms, supercells = [1, 1, 1], special_label=None, label=None):
"""Initializes an AtomicStructure instance
Args:
atoms (ASE Atoms): The atoms object to create the structure from
special_label (str): Extra classification label if needed
label (str): Structure label, if None it will be autogenerated
"""
self._atoms = atoms.copy()
self.supercells = supercells
if atoms.calc != None:
self._atoms.calc = atoms.calc
else:
raise ValueError("Atoms object needs to have an attached calculator")
self._label = self._generateHashLabel(special_label) if label == None else label
# Write the label on top of sampledata if it exists.
if os.path.exists("sampledata.txt") and os.path.getsize("sampledata.txt") == 0:
with open(f"sampledata.txt", "w") as f:
f.write(f"{self._label}\n")
[docs]
@classmethod
def fromFile(cls, path, pbc=[True, True, True], supercells=[1, 1, 1], potential: Potential = None):
"""Creates a AtomicStructure object directly from a structure file
Args:
path (str): Path to the file containing the crystal structure
pbc (Bool): Whether to use periodic boundary conditions or not
supercells (list[int]): Number of supercells to use in each direction
potential (Potential): Interatomic potential to use
"""
atoms = read(path) * supercells
atoms.calc = potential.getASEPotentialCalculator()
if isinstance(pbc, bool):
atoms.pbc = [pbc, pbc, pbc]
else:
if len(pbc) == 3:
atoms.pbc = pbc
else:
raise ValueError("pbc needs to be a single boolean or a list of three booleans")
return cls(atoms, supercells)
# Standard dunder methods
def __len__(self):
"""Returns the number of atoms in the structure"""
return len(self._atoms)
def __str__(self):
"""Returns the string representation of the ASE atoms object"""
return f"{self._atoms}"
def __copy__(self):
"""Creates a copy of the AtomicStructure object"""
return AtomicStructure(self._atoms, label=self.label)
# Properties.
@property
def label(self):
"""Label of the given structure
Returns:
(str): label of the structure
"""
return self._label
@property
def potential(self):
"""Potential that is used
Returns:
(Potential): The structures potential
"""
return self._atoms.calc
@potential.setter
def potential(self, new_pot):
"""Set the potential calculator for the structure.
Args:
new_pot (Potential): Calculator to assign to the atoms object.
"""
self._atoms.calc = new_pot.getASEPotentialCalculator()
@property
def positions(self):
""" Positions of all atoms
Returns:
list[Atoms]:
"""
return self._atoms.get_positions()
@positions.setter
def positions(self, new_positions):
"""Set atomic positions.
Args:
new_positions (array-like): New atomic coordinates (Å).
"""
self._atoms.set_positions(new_positions)
@property
def velocities(self):
"""numpy.ndarray | None: Atomic velocities (Å/fs)."""
vel = self._atoms.get_velocities()
return vel.copy() if vel is not None else None
@velocities.setter
def velocities(self, new_vel):
"""Set atomic velocities.
Args:
new_vel (array-like): Velocities for each atom (Å/fs).
"""
self._atoms.set_velocities(new_vel)
@property
def forces(self):
"""numpy.ndarray: Forces acting on each atom (eV/Å).
Raises:
RuntimeError: If no potential calculator is defined.
"""
if self._atoms.calc is None:
raise RuntimeError("Cannot get forces: no potential assigned")
return self._atoms.get_forces()
@property
def stress(self):
"""numpy.ndarray: Stress tensor in Voigt notation (eV/ų).
Raises:
RuntimeError: If no potential calculator is defined.
"""
if self._atoms.calc is None:
raise RuntimeError("Cannot get forces: no potential assigned")
return self._atoms.get_stress(voigt=True)
@property
def masses(self):
"""numpy.ndarray: Atomic masses (amu)."""
return self._atoms.get_masses()
@property
def potential_energy(self):
"""float: Total potential energy of the structure (eV).
Raises:
RuntimeError: If no potential calculator is defined.
"""
if self._atoms.calc is None:
raise RuntimeError("Cannot get energy: no potential assigned")
return self._atoms.get_potential_energy()
@property
def kinetic_energy(self):
"""float: Total kinetic energy (eV).
Raises:
RuntimeError: If no potential calculator is defined.
"""
if self._atoms.calc is None:
raise RuntimeError("Cannot get energy: no potential assigned")
return self._atoms.get_kinetic_energy()
@property
def total_energy(self):
"""float: Total energy (potential + kinetic) in eV.
Raises:
RuntimeError: If no potential calculator is defined.
"""
if self._atoms.calc is None:
raise RuntimeError("Cannot get energy: no potential assigned")
return self._atoms.get_total_energy()
@property
def temperature(self):
"""float: Instantaneous temperature (K)."""
if self._atoms.calc is None:
raise RuntimeError("Cannot get energy: no potential assigned")
return self._atoms.get_temperature()
@property
def volume(self):
"""float: Volume of the cell (ų)."""
if self._atoms.calc is None:
raise RuntimeError("Cannot get energy: no potential assigned")
return self._atoms.get_volume()
@property
def cellpar(self):
"""numpy.ndarray: Cell parameters (a, b, c, α, β, γ)."""
return self._atoms.cell.cellpar().copy()
@property
def cell(self):
"""numpy.ndarray: 3×3 cell matrix defining the simulation cell."""
return self._atoms.cell.array.copy()
@cell.setter
def cell(self, new_cell):
"""Set the atomic cell and scale positions.
Args:
new_cell (array-like): New 3×3 cell matrix.
"""
self._atoms.set_cell(new_cell, scale_atoms=True)
@property
def lattice_constant(self):
"""float: Average in-plane lattice constant (Å)."""
cell = self.cell / self.supercells
a_len, b_len = np.linalg.norm(cell[0]), np.linalg.norm(cell[1])
return (a_len + b_len) / 2
@property
def get_cell(self):
"""numpy.ndarray: Return the cell matrix."""
return self._atoms.get_cell()
@property
def symbols(self):
"""list[str]: Chemical symbols of all atoms."""
return self._atoms.get_chemical_symbols()
@property
def get_atomic_numbers(self):
"""numpy.ndarray: Atomic numbers of all atoms."""
return self._atoms.get_atomic_numbers()
@property
def get_all_distances(self):
"""numpy.ndarray: Pairwise atomic distances (Å)."""
return self._atoms.get_all_distances()
[docs]
def cohesive_energy(self, potential_energy=None):
"""Compute the cohesive energy per atom.
Cohesive energy is defined as the difference between the sum of isolated
atomic energies and the bulk structure energy, divided by the number
of atoms:
E_coh = (Σ E_atom − E_bulk) / N
Args:
potential_energy (float, optional): Bulk structure potential energy (eV).
If not provided, it is computed automatically.
Returns:
float: Cohesive energy per atom (eV/atom).
"""
number = self.__len__()
e_atoms = 0
for symbol in self._atoms.get_chemical_symbols():
atom = Atoms(symbol, positions=[[0, 0, 0]], cell=[10, 10, 10], pbc=False)
atom.calc = self._atoms.calc
e_atoms += atom.get_potential_energy()
potential_energy = self.potential_energy
e_coh = (e_atoms - potential_energy) / number
return e_coh
@property
def internal_pressure(self):
"""
Compute internal pressure using atomic units internally.
Instantaneous: P = (1/3V) [ 2 N E_kin + sum_i r_i · f_i ]
Returns a mean over the instantaneous frames in the trajectory
Unit: ev/Å^3
Returns:
P (float): Internal pressure
"""
P = abs(-self.stress[0:3].sum()/3)
return P
[docs]
def setVelocitiesMB(self, temperature_K):
"""Assign velocities from a Maxwell–Boltzmann distribution.
The procedure removes net linear and angular momentum.
Args:
temperature_K (float): Target temperature in Kelvin.
"""
MaxwellBoltzmannDistribution(self._atoms, temperature_K=temperature_K, force_temp=True)
Stationary(self._atoms)
ZeroRotation(self._atoms)
[docs]
def computeMSD(self, orig_struct):
"""Compute mean-square displacement relative to a reference structure.
Args:
orig_struct (AtomicStructure): Reference configuration.
Returns:
float: Mean-square displacement (Ų).
Raises:
TypeError: If `orig_struct` is not an `AtomicStructure`.
"""
if not isinstance(orig_struct, AtomicStructure):
raise TypeError("orig_struct needs to be type AtomicStructure")
r_0 = orig_struct.positions
r_n = self.positions
return np.mean((r_0 - r_n) ** 2)
[docs]
def computeNearestNeighbour(self):
"""Compute the mean nearest-neighbor distance.
Returns:
float: Mean nearest-neighbor distance (Å).
Raises:
ValueError: If nearest neighbors cannot be found for an atom.
"""
INF = 1e9
NN_list = []
cutoff = natural_cutoffs(self._atoms)
neighbor_list = NeighborList(cutoff, bothways=True)
neighbor_list.update(self._atoms)
for current_atom in range(self._atoms.get_global_number_of_atoms()):
indices, offsets = neighbor_list.get_neighbors(current_atom)
nearest_distance = INF
for neighbor_index, offset in zip(indices[1:], offsets[1:]):
NN_vector = (
self._atoms.positions[neighbor_index]
+ offset @ self._atoms.get_cell()
- self._atoms.positions[current_atom]
)
distance = np.sqrt(NN_vector.dot(NN_vector))
if distance < nearest_distance:
nearest_distance = distance
if nearest_distance == INF:
raise ValueError(
f"No nearest neighbor found for atom {current_atom}"
)
NN_list.append(nearest_distance)
return np.mean(NN_list)
def _generateHashLabel(self, special_label):
"""Generate a semi-random hash label for the structure.
Args:
special_label (str | None): Custom label inserted into the hash.
Returns:
str: Hash label of the form
`<formula>_<special>_<hash>`.
"""
formula = self._atoms.get_chemical_formula()
data = (
self._atoms.numbers.tobytes()
+ self._atoms.positions.tobytes()
+ np.random.bytes(4)
)
short_hash = hashlib.md5(data).hexdigest()[:6]
if special_label is None:
label = f"{formula}_{short_hash}"
else:
label = f"{formula}_{special_label}_{short_hash}"
return label
[docs]
def getAtoms(self):
"""Unsafe helper: return the underlying ASE Atoms object.
Returns:
ase.Atoms: Raw Atoms object (use with caution).
"""
return self._atoms