Source code for sqsgenerator.core.structure

import numbers
import itertools
import numpy as np
import typing as T
from operator import attrgetter as attr
from sqsgenerator.core.core import Structure as Structure_


[docs]class Structure(Structure_): """ Structure class used to store structural information. This class is used by the core extension and is a wrapper around the extension internal ``sqsgenerator.core.core.Structure`` class. The class is designed to be array like, therefore it does not provide any setter functions. Internally this class uses **fractional coordinates** to represent structural information. """ @classmethod def from_extension_class(cls, o: Structure_): return cls(o.lattice.copy(), o.frac_coords.copy(), list(map(attr('symbol'), o.species)), o.pbc) def __init__(self, lattice: np.ndarray, frac_coords: np.ndarray, symbols: T.List[str], pbc: T.Tuple[bool, bool, bool] = (True, True, True)): """ Constructs a new structure from raw data :param lattice: the (3x3) lattice matrix. The three **rows** are interpreted as the lattice vectors **a**, **b** and **c**. :type lattice: np.ndarray :param frac_coords: the **fractional coordinates** of the lattice positions as (3xN) array. Each **row** will be treated as a lattice position :type frac_coords: np.ndarray :param symbols: a list of strings of length N specifying the atomic species which occupy the lattice positions :type symbols: List[str] :param pbc: the coordinate axes for which **periodic boundary conditions** should be applied. **Do not pass a value here**. This feature is not yet implemented (default is ``(True, True, True)``) :type pbc: Tuple[bool, bool, bool] :raises ValueError: if length of {frac_coords} and length of {symbols} do not match """ super(Structure, self).__init__(lattice, frac_coords, symbols, pbc) self._symbols = np.array(list(map(attr('symbol'), self.species)), dtype='<U3') # Uut is the longest symbol self._numbers = np.fromiter(map(attr('Z'), self.species), dtype=int, count=self.num_atoms) self._unique_species = set(np.unique(self._symbols)) @property def symbols(self) -> np.ndarray: """ A ``numpy.ndarray`` storing the symbols of the atomic species. E.g "*Fe*", "*Cr*", "*Ni*" :return: the array of symbols :rtype: numpy.ndarray """ return self._symbols @property def numbers(self) -> np.ndarray: """ The ordinal numbers of the atoms sitting on the lattice positions :return: array of ordinal numbers :rtype: numpy.ndarray """ return self._numbers @property def num_unique_species(self) -> int: """ The number of unique elements in the structure object :return: the number of elements :rtype: int """ return len(self._unique_species) @property def unique_species(self) -> T.Set[str]: """ A set of symbols of the occurring species :return: a set containing the symbol of the occurring species :rtype: Set[str] """ return self._unique_species def without_vacancies(self): """ Removes vacancies from the structure. Removes all lattice sites which are occupied by the "0" species :return: a structure with no vacancies :rtype: Structure """ return self[~(self.numbers == 0)] def __len__(self): return self.num_atoms def __repr__(self): def group_symbols(): for species, same in itertools.groupby(self._symbols): num_same = len(list(same)) yield species if num_same == 1 else f'{species}{num_same}' formula = ''.join(group_symbols()) return f'Structure({formula}, len={self.num_atoms})' def __eq__(self, other): if not isinstance(other, Structure): return False same_species = all(this_num == other_num for this_num, other_num in zip(self.numbers.flat, other.numbers.flat)) same_lattice = np.allclose(self.lattice, other.lattice) same_coords = np.allclose(self.frac_coords, other.frac_coords) return all((same_species, same_lattice, same_coords))
[docs] def sorted(self): """ Creates a new structure, where the lattice positions are ordered by the ordinal numbers of the occupying species :return: the sorted Structure :rtype: Structure """ return self.from_extension_class(super(Structure, self).sorted())
def rearranged(self, order): return self.from_extension_class(super(Structure, self).rearranged(order)) def __getitem__(self, item): if isinstance(item, numbers.Integral): if item < -self.num_atoms or item >= self.num_atoms: raise IndexError('Index out of range.') return self.species[item] if isinstance(item, (list, tuple, np.ndarray)): indices = np.array(item) if indices.dtype == bool: indices = np.argwhere(indices).flatten() elif indices.dtype != int: raise TypeError('Only integer numbers cann be used for slicing') # boolean mask slice elif isinstance(item, slice): indices = np.arange(self.num_atoms)[item] else: raise TypeError(f'Structure indices must not be of type {type(item)}') return Structure(self.lattice, self.frac_coords[indices], self.symbols[indices])
[docs] def to_dict(self) -> dict: """ Serializes the object into JSON/YAML serializable dictionary :return: the JSON/YAML serializable dictionary :rtype: Dict[str, Any] """ return structure_to_dict(self)
[docs] def slice_with_species(self, species: T.Iterable[str], which: T.Optional[T.Iterable[int]] = None): """ Creates a new structure containing the lattice positions specified by {which}. The new structure is occupied by the atomic elements specified in {species}. In case {which} is ``None`` all lattice positions are assumed to be occupied with a new {species}. The new structure will only contain lattice positions specified in {which}. :param species: the atomic species specified by their symbols :type species: Iterable[str] :param which: the indices of the lattice positions to choose (default is ``None``) :type which: Optional[Iterable[int]] :return: the (subset) structure with new species :rtype: Structure :raises ValueError: if length of which is < 1 or length of {which} and {species} does not match """ which = which or tuple(range(self.num_atoms)) return self.with_species(species, which=which)[which]
[docs] def with_species(self, species, which=None): """ Creates a new structure containing the lattice positions specified by {which}. The new structure is occupied by the atomic elements specified in {species}. In case {which} is ``None`` all lattice positions are assumed to be occupied with a new {species}. The new structure will only all lattice positions of the current structure, while on the positions specified by {which} are occupied with {species} :param species: the atomic species specified by their symbols :type species: Iterable[str] :param which: the indices of the lattice positions to choose (default is ``None``) :type which: Optional[Iterable[int]] :return: the structure with new species :rtype: Structure :raises ValueError: if length of which is < 1 or length of {which} and {species} does not match """ which = which or tuple(range(len(self))) species = list(species) if len(species) < 1: raise ValueError('Cannot create an empty structure') if len(which) != len(species): raise ValueError('Number of species does not match the number of specified atoms') new_symbols = self.symbols.copy() new_symbols[np.array(which)] = species return Structure(self.lattice, self.frac_coords, new_symbols.tolist())
def structure_to_dict(structure: Structure): structure = structure.without_vacancies() return dict( lattice=structure.lattice.tolist(), coords=structure.frac_coords.tolist(), species=structure.symbols.tolist() )
[docs]def make_supercell(structure: Structure, sa: int = 1, sb: int = 1, sc: int = 1) -> Structure: """ Creates a supercell of structure, which is repeated {sa}, {sb} and {sc} times :param structure: the structure to replicate :type structure: Structure :param sa: number of repetitions in directions of the first lattice vector (default is ``1``) :type sa: int :param sb: number of repetitions in directions of the second lattice vector (default is ``1``) :type sb: int :param sc: number of repetitions in directions of the third lattice vector (default is ``1``) :type sc: int :return: the supercell structure :rtype: Structure """ sizes = (sa, sb, sc) num_cells = np.prod(sizes) num_atoms_supercell = structure.num_atoms * num_cells scale = np.diag(sizes).astype(float) iscale = np.linalg.inv(scale) supercell_lattice = scale @ structure.lattice scaled_fc = structure.frac_coords @ iscale.T def make_translation_vector(a, b, c): return np.tile(np.array([a, b, c]) @ iscale, [structure.num_atoms, 1]) supercell_coords = np.vstack([ scaled_fc + make_translation_vector(*shift) for shift in itertools.product(*map(range, sizes)) ]) supercell_species = np.tile(structure.symbols, num_cells).tolist() assert supercell_coords.shape == (num_atoms_supercell, 3) assert len(supercell_species) == num_atoms_supercell structure_supercell = Structure(supercell_lattice, supercell_coords, supercell_species, (True, True, True)) return structure_supercell