import numpy as np
[docs]
def site_offset_to_spos(site, offset, basis_spos):
""" Returns the scaled position of an atom at specified site and offset
relative to the basis in scaled coordinates"""
return offset + basis_spos[site]
[docs]
def spos_to_pos(spos, cell):
""" Returns the Cartesian coordinate given the scaled coordinate and cell
metric (cell vectors as rows)"""
return np.dot(spos, cell)
[docs]
def pos_to_spos(pos, cell):
""" Inverse of sps_to_pos"""
return np.linalg.solve(cell.T, pos)
[docs]
def spos_to_site_offset(spos, basis_spos, symprec):
""" Returns the site and offset of the atom at the specified scaled
coordinate given the scaled positions of the basis atoms"""
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, cell, basis_spos, symprec):
""" helper to map pos -> spos -> site/offset"""
spos = pos_to_spos(pos, cell)
return spos_to_site_offset(spos, basis_spos, symprec)
[docs]
def site_offset_to_pos(site, offset, cell, basis_spos):
""" helper to map site/offset -> spos -> pos"""
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 crustal"""
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):
return self._site
@property
def offset(self):
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):
return site_offset_to_pos(self._site, self._offset,
self._structure.cell,
self._structure.spos)
@property
def number(self):
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.Atoms class but is a bit more
carefull about pbc and scaled coordinates. It also returns hiphive.Atom
objects instead"""
def __init__(self, atoms, symprec=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):
return self._spos
@property
def cell(self):
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, symprec=None):
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
primitve lattice"""
def __init__(self, supercell, prim, symprec):
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))
[docs]
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, offset):
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)