import numpy as np
from ase import Atoms
[docs]
def site_offset_to_spos(
site: int,
offset: np.ndarray,
basis_spos: list[np.ndarray],
) -> np.ndarray:
""" Returns the scaled position of an atom at the specified site and offset
relative to the basis in scaled coordinates.
Parameters
----------
site
Site index.
offset
Offset in scaled coordinates (should be integers).
basis_spos
Positions of atoms in basis in scaled coordinates.
"""
return offset + basis_spos[site]
[docs]
def spos_to_pos(spos: np.ndarray, cell: np.ndarray) -> np.ndarray:
""" Returns the Cartesian coordinate given the scaled coordinate and cell
metric (cell vectors as rows).
Parameters
----------
spos
Position in scaled coordinates.
cell
Cell metric.
"""
return np.dot(spos, cell)
[docs]
def pos_to_spos(pos: np.ndarray, cell: np.ndarray) -> np.ndarray:
"""Returns the scaled coordinate given the Cartesian coordinate and the cell metric.
Inverse of :func:`spos_to_pos`.
Parameters
----------
pos
Position in Cartesian coordinates.
cell
Cell metric.
"""
return np.linalg.solve(cell.T, pos)
[docs]
def spos_to_site_offset(
spos: np.ndarray,
basis_spos: list[np.ndarray],
symprec: float,
) -> np.ndarray:
""" Returns the site and offset of the atom at the specified scaled
coordinate given the scaled positions of the basis atoms.
Parameters
----------
spos
Position in scaled coordinates.
basis_spos
Positions of atoms in basis in scaled coordinates.
symprec
Tolerance imposed when rounding the offset.
"""
for site, sp in enumerate(basis_spos):
offset = spos - sp
rounded_offset = offset.round(0).astype(np.int64)
# TODO: fix tolerance (symprec should be w.r.t. cart. coord.)
if np.allclose(rounded_offset, offset, rtol=0, atol=symprec):
return site, rounded_offset
raise Exception('spos {} not compatible with basis {} using symprec {}'
.format(spos, basis_spos, symprec))
[docs]
def pos_to_site_offset(
pos: np.ndarray,
cell: np.ndarray,
basis_spos: list[np.ndarray],
symprec: float,
) -> np.ndarray:
"""Returns offset given the position in Cartesian coordinates and the cell metric.
Parameters
----------
pos
Position in Cartesian coordinates.
cell
Cell metric.
basis_spos
Positions of atoms in basis in scaled coordinates.
symprec
Tolerance imposed when rounding the offset.
"""
spos = pos_to_spos(pos, cell)
return spos_to_site_offset(spos, basis_spos, symprec)
[docs]
def site_offset_to_pos(
site: int,
offset: np.ndarray,
cell: np.ndarray,
basis_spos: list[np.ndarray],
) -> np.ndarray:
"""Returns the position in Cartesian coordinates given a site index,
an offset, the cell metric, and the position in the basis.
Parameters
----------
site
Site index.
offset
Offset in scaled coordinates (should be integers).
cell
Cell metric.
basis_spos
Positions of atoms in basis in scaled coordinates.
"""
spos = site_offset_to_spos(site, offset, basis_spos)
return spos_to_pos(spos, cell)
[docs]
class BaseAtom:
""" This class represents an atom placed in an infinite crystal.
Attributes
----------
site : int
Site index.
offset : list[int]
Offset in scaled coordinates.
"""
def __init__(self, site, offset):
assert type(site) is int, type(site)
assert len(offset) == 3, len(offset)
assert (all(type(i) is int for i in offset) or
all(type(i) is np.int64 for i in offset)), type(offset[0])
self._site = site
self._offset = np.array(offset)
@property
def site(self) -> int:
""" Site index. """
return self._site
@property
def offset(self) -> np.ndarray:
""" Offset in scaled coordinates. """
return self._offset
[docs]
def astype(self, dtype):
""" Useful arguments: list, tuple, np.int64"""
return dtype((self._site, *self._offset))
[docs]
class Atom(BaseAtom):
""" This class represents a crystal atom in a given structure.
"""
def __init__(self, *args, **kwargs):
self._structure = kwargs.pop('structure', None)
super().__init__(*args, **kwargs)
@property
def pos(self) -> np.ndarray:
return site_offset_to_pos(self._site, self._offset,
self._structure.cell,
self._structure.spos)
@property
def number(self) -> int:
return self._structure.numbers[self._site]
[docs]
class SupercellAtom(Atom):
""" Represents an atom in a supercell but site and offset given by an
underlying primitve cell."""
def __init__(self, *args, **kwargs):
self._index = kwargs.pop('index')
assert type(self._index) is int
super().__init__(*args, **kwargs)
@property
def index(self):
return self._index
[docs]
class Structure:
""" This class essentially wraps the ASE :class:`Atoms` class but is a bit more
careful with respect to periodic boundary conditions and scaled coordinates.
Note that it returns :class:`hiphive.Atom` objects when queried for individual atoms.
Parameters
----------
atoms
Atomic structure.
symprec
Tolerance imposed when rounding scaled coordinates.
"""
def __init__(self, atoms: Atoms, symprec: float = 1e-6):
spos = atoms.get_scaled_positions(wrap=False)
for sp in spos.flat:
if not (-symprec < sp < (1 - symprec)):
raise ValueError('bad spos {}'.format(sp))
self._spos = spos
self._cell = atoms.cell
self._numbers = atoms.numbers
def __len__(self):
return len(self._spos)
@property
def spos(self) -> np.ndarray:
""" Scaled coordinates. """
return self._spos
@property
def cell(self):
""" Cell metric. """
return self._cell
def __getitem__(self, index):
if index >= len(self):
raise IndexError('Structure contains {} atoms'.format(len(self)))
return Atom(index, (0, 0, 0), structure=self)
[docs]
def atom_from_pos(self, pos: np.ndarray, symprec: float = None) -> Atom:
""" Returns atom given a position in Cartesian coordinates.
Parameters
----------
pos
Position in Cartesian coordinates.
symprec
Tolerance imposed when rounding scaled coordinates.
"""
if symprec is None:
symprec = self._symprec
site, offset = pos_to_site_offset(pos, self._cell, self._spos, symprec)
return Atom(site, offset, structure=self)
[docs]
class Supercell:
""" This class tries to represent atoms in a supercell as positioned on the
primitive lattice.
Parameters
----------
supercell
Supercell structure.
prim
Primitive structure.
symprec
Tolerance imposed when rounding scaled coordinates.
"""
def __init__(self, supercell: Atoms, prim: Atoms, symprec: float):
self._supercell = Structure(supercell)
self._prim = Structure(prim)
self._symprec = symprec
self._map = list()
self._inverse_map_lookup = dict()
self._create_map()
def _create_map(self):
for atom in self._supercell:
atom = self._prim.atom_from_pos(atom.pos, self._symprec)
self._map.append(atom.astype(tuple))
def wrap_atom(self, atom):
atom = Atom(atom.site, atom.offset, structure=self._prim)
tup = atom.astype(tuple)
index = self._inverse_map_lookup.get(tup, None)
if index is None:
atom = self._supercell.atom_from_pos(atom.pos, self._symprec)
index = atom.site
self._inverse_map_lookup[tup] = index
return self[index]
[docs]
def index(self, site: int, offset: np.ndarray) -> int:
"""" Returns index of atom given a site index and an offset.
Parameters
----------
site
Site index.
offset
Offset in scaled coordinates (should be integers).
"""
atom = self.wrap_atom(BaseAtom(site, offset))
return atom.index
def __getitem__(self, index):
tup = self._map[index]
return SupercellAtom(tup[0], tup[1:], structure=self._prim,
index=index)
def __len__(self):
return len(self._supercell)
