Source code for quippy.atoms

# HQ X
# HQ X   quippy: Python interface to QUIP atomistic simulation library
# HQ X
# HQ X   Copyright James Kermode 2010
# HQ X
# HQ X   These portions of the source code are released under the GNU General
# HQ X   Public License, version 2,
# HQ X
# HQ X   If you would like to license the source code under different terms,
# HQ X   please contact James Kermode,
# HQ X
# HQ X   When using this software, please cite the following reference:
# HQ X
# HQ X
# HQ X

import os
import weakref
import copy
import sys
import logging
from math import pi, sqrt

import numpy as np

from quippy import _atoms
from quippy._atoms import *

from quippy.dictionary import (T_INTEGER_A, T_REAL_A, T_LOGICAL_A, T_CHAR_A,
                               T_INTEGER_A2, T_REAL_A2)
from quippy.table import TABLE_STRING_LENGTH
from quippy.oo_fortran import update_doc_string
from quippy.farray import frange, farray, fzeros, fvar, s2a
from quippy.dictmixin import DictMixin
from quippy.util import infer_format, parse_slice
from quippy.units import N_A, MASSCONVERT
from quippy.periodictable import ElementMass
from quippy.extendable_str import Extendable_str
from quippy import QUIPPY_TRUE, QUIPPY_FALSE
from quippy import available_modules, FortranDerivedTypes, get_fortran_indexing

atomslog = logging.getLogger('quippy.atoms')

__doc__ = (_atoms.__doc__ + """

This module defines the :class:`Atoms`, which stores and manipulates a
collection of atoms, as well as the :class:`Connection` class which stores
topology and neighbour lists, and the :class:`DomainDecomposition` class.


__all__ = _atoms.__all__ + ['NeighbourInfo', 'get_lattice_params_', 'get_lattice_params']

if 'ase' in available_modules:
    import ase
    from ase.spacegroup import Spacegroup
    # adsorbate_info as a property is deprecated in ase
    # since it shows up in `dir()` the methods in
    # oo_fortran will error when trying to access it as
    # a property. Just remove it so that it doesn't show up
    # at all as a property any more.
    if hasattr(ase.Atoms, 'adsorbate_info'):
        delattr(ase.Atoms, 'adsorbate_info')
    raise ImportError("mandatory dependency ase cannot be imported. Install using 'pip install ase'.")

if 'phonopy' in available_modules:
    from phonopy.structure.atoms import Atoms as PhonopyAtoms

get_lattice_params_ = get_lattice_params

[docs]def get_lattice_params(lattice): """ Wrapper around Fortran :func:`get_lattice_params_` Returns parameters of `lattice` as 6-tuple (a,b,c,alpha,beta,gamma). """ a,b,c,alpha,beta,gamma = (farray(0.), farray(0.), farray(0.), farray(0.), farray(0.), farray(0.)) _atoms.get_lattice_params(lattice,a,b,c,alpha,beta,gamma) return tuple(float(x) for x in (a,b,c,alpha,beta,gamma))
class NeighbourInfo(object): """Store information about a single neighbour of an atom""" __slots__ = ('j', 'distance', 'diff', 'cosines', 'shift') def __init__(self, j, distance, diff, cosines, shift): self.j = j self.distance = float(distance) self.diff = diff.copy() self.cosines = cosines.copy() self.shift = shift.copy() def __int__(self): return self.j def __repr__(self): return ('NeighbourInfo(j=%d,distance=%f,diff=%s,cosines=%s,shift=%s)' % (self.j, self.distance, self.diff, self.cosines, self.shift))
[docs]class Connection(_atoms.Connection): __doc__ = update_doc_string(_atoms.Connection.__doc__, """ The :class:`Connection` is a subclass of :class:`_atoms.Connection` which adds supports for iteration over all atoms, and indexing e.g. ``at.connect.neighbours[1]`` returns a list of the neighbours of the atom with index 1. When indexed with an integer from 1 to `at.n`, returns an array of :class:`NeighbourInfo` objects, each of which corresponds to a particular pair `(i,j)` and has attributes `j`, `distance`, `diff`, `cosines` and `shift`. If ``fortran_indexing`` is True, atom and neighbour indices start from 1; otherwise they are numbered from zero. If connectivity information has not already been calculated :meth:`calc_connect` will be called automatically. The code to loop over the neighbours of all atoms is quite idiomatic:: for i in at.indices: for neighb in at.connect[i]: print (neighb.j, neighb.distance, neighb.diff, neighb.cosines, neighb.shift) Note that this provides a more Pythonic interface to the atomic connectivity information than the wrapped Fortran functions :meth:`Atoms.n_neighbours` and :meth:`Atoms.neighbour`. """) # def __init__(self, n=None, nbuffer=None, pos=None, # lattice=None, g=None, origin=None, # extent=None, nn_guess=None, fill=None, # fpointer=None, finalise=True): # _atoms.Connection.__init__(self, n, nbuffer, pos, # lattice, g, origin, # extent, nn_guess, fill, # fpointer, finalise) def is_neighbour(self, i, j): return (i,j) in self.pairs()
[docs] def pairs(self): """Yield pairs of atoms (i,j) with i < j which are neighbours""" for i, neighbour_list in zip(self.parent.indices, self.iterneighbours()): for neighb in neighbour_list: if i < neighb.j: yield (i,neighb.j)
def __eq__(self, other): if not isinstance(other, self.__class__): return False # Neighbours are considered to be equal if *topology* matches, # not distances, displacement vectors and shifts. return sorted(self.pairs()) == sorted(other.pairs()) def __ne__(self, other): return not self.__eq__(other) def __iter__(self): return self.iterneighbours()
[docs] def iterneighbours(self): """Iterate over the neighbours of all atoms""" for i in self.parent.indices: yield self[i]
def __getitem__(self, i): if not self.initialised: if self is self.parent.hysteretic_connect: self.calc_connect_hysteretic(self.parent) else: self.calc_connect(self.parent) distance = farray(0.0) diff = fzeros(3) cosines = fzeros(3) shift = fzeros(3, dtype=np.int32) res = [] if not get_fortran_indexing(): i = i+1 # convert to 1-based indexing for n in frange(self.n_neighbours(i)): j = self.neighbour(self.parent, i, n, distance, diff, cosines, shift) if not get_fortran_indexing(): j = j-1 res.append(NeighbourInfo(j, distance, diff, cosines, shift)) if get_fortran_indexing(): res = farray(res) # to give 1-based indexing return res
[docs] def distances(self, Z1=None, Z2=None): """Distances between pairs of neighbours, optionally filtered by species (Z1,Z2)""" for i in self.parent.indices: for neighb in self[i]: if neighb.j > i: continue if Z1 is not None and Z2 is not None: if sorted((self.parent.z[i], self.parent.z[neighb.j])) == sorted((Z1, Z2)): yield neighb.distance else: yield neighb.distance
[docs] def get_neighbours(self, i): """ Return neighbours of atom i Return arrays of indices and offsets to neighbouring atoms. The positions of the neighbor atoms can be calculated like this:: indices, offsets = atoms.connect.get_neighbors(42) for i, offset in zip(indices, offsets): print atoms.positions[i] + dot(offset, atoms.get_cell()) Compatible with ase.calculators.neighborlist.NeighborList.get_neighbors(), providing that NeighborList is constructed with bothways=True and self_interaction=False. """ neighbours = self[i] indices = np.array([n.j for n in neighbours]) offsets = np.r_[[n.shift for n in neighbours]] return (indices, offsets)
[docs] def get_neighbors(self, i): """ Variant spelling of :meth:`get_neighbours` """ return self.get_neighbours(i)
class PropertiesWrapper(DictMixin): """Wrapper between quippy properties and ASE arrays""" def __init__(self, at): DictMixin.__init__(self) self.atref = weakref.ref(at) def __getitem__(self, key): at = self.atref() res =[at.name_map.get(key, key)].view(np.ndarray).T #if res.dtype == 'int32': # convert dtype int32 -> int # res = res.astype(int) # atomslog.debug('Making copy of arrays["%s"] since quippy/ASE dtypes incompatible' % key) return res def __setitem__(self, key, value): at = self.atref() atomslog.debug('Setting arrays["%s"]' % key) key = Atoms.name_map.get(key, key) value = np.array(value) if value.shape[0] != at.n: atomslog.debug(('Assignment to arrays["%s"]: changing size '+ ' from %d to %d') % (key, at.n, value.shape[0])) for p in at.remove_property(p) at.n = value.shape[0] at.nbuffer = at.n at.ndomain = at.n if value.dtype.kind != 'S': value = value.T if not at.has_property(key): at.add_property(key, value) else:[key][...] = value def __delitem__(self, key): at = self.atref() atomslog.debug('Deletion of arrays["%s"]' % key) key = at.name_map.get(key, key) at.remove_property(key) def keys(self): at = self.atref() return [at.rev_name_map.get(key, key) for key in]
[docs]class Atoms(_atoms.Atoms, ase.Atoms): __doc__ = update_doc_string(_atoms.Atoms.__doc__, """ The :class:`Atoms` class is a Pythonic wrapper over the auto-generated :class:`quippy._atoms.Atoms` class. Atoms object are usually constructed either by reading from an input file in one of the :ref:`fileformats`, or by using the structure creation functions in the :mod:`quippy.structures` or :mod:`ase.lattice` modules. For example to read from an :ref:`extendedxyz` file, use:: from quippy.atoms import Atoms atoms = Atoms('') Or, to create an 8-atom bulk diamond cubic cell of silicon:: from quippy.structures import diamond si_bulk = diamond(5.44, 14) The :class:`Atoms` class is inherited from the :class:`ase.atoms.Atoms` so has all the ASE Atoms attributes and methods. This means that quippy and ASE Atoms objects are fully interoperable.""", signature='Atoms([symbols, positions, numbers, tags, momenta, masses, magmoms, charges, scaled_positions, cell, pbc, constraint, calculator, info, n, lattice, properties, params, fixed_size, **read_args])') _cmp_skip_fields = ['own_this', 'ref_count', 'domain', 'connect', 'hysteretic_connect', 'source'] name_map = {'positions' : 'pos', 'numbers' : 'Z', 'initial_charges' : 'charge', 'initial_magmoms' : 'magmoms'} rev_name_map = dict(zip(name_map.values(), name_map.keys())) def __init__(self, symbols=None, positions=None, numbers=None, tags=None, momenta=None, masses=None, magmoms=None, charges=None, scaled_positions=None, cell=None, pbc=None, constraint=None, calculator=None, info=None, n=None, lattice=None, properties=None, params=None, fixed_size=None, set_species=True, fpointer=None, finalise=True, **readargs): # check for mutually exclusive options if cell is not None and lattice is not None: raise ValueError('only one of cell and lattice can be present') if n is None: n = 0 if cell is not None: lattice = cell.T if lattice is None: lattice = np.eye(3) from quippy import Dictionary if properties is not None and not isinstance(properties, Dictionary): properties = Dictionary(properties) if params is not None and not isinstance(params, Dictionary): params = Dictionary(params) _atoms.Atoms.__init__(self, n=n, lattice=lattice, properties=properties, params=params, fixed_size=fixed_size, fpointer=fpointer, finalise=finalise) self._ase_arrays = PropertiesWrapper(self) # If first argument is quippy.Atoms instance, copy data from it if isinstance(symbols, self.__class__): self.copy_from(symbols) symbols = None # Phonopy compatibility if 'phonopy' in available_modules: if symbols is not None and isinstance(symbols, type(PhonopyAtoms)): atoms = symbols symbols = atoms.get_chemical_symbols() cell = atoms.get_cell() positions = atoms.get_positions() masses = atoms.get_masses() # Try to read from first argument, if it's not ase.Atoms if symbols is not None and not isinstance(symbols, ase.Atoms): self.read_from(symbols, **readargs) symbols = None ## ASE compatibility remove_properties = [] if symbols is None and numbers is None: if self.has_property('z'): numbers = self.z.view(np.ndarray) else: numbers = [0]*len(self) remove_properties.append('Z') if symbols is None and positions is None: if self.has_property('pos'): positions = self.pos.view(np.ndarray).T else: remove_properties.append('pos') # Make sure argument to ase.Atoms constructor are consistent with # properties already present in this Atoms object if symbols is None and momenta is None and self.has_property('momenta'): momenta = self.get_momenta() if symbols is None and masses is None and self.has_property('masses'): masses = self.get_masses() if symbols is None and cell is None: cell = self.lattice.T.view(np.ndarray) if symbols is None and pbc is None: pbc = self.get_pbc() #print 'setting pbc to ', pbc if charges is None and self.has_property('charge'): charges = self.charge.view(np.ndarray) ase.Atoms.__init__(self, symbols, positions, numbers, tags, momenta, masses, magmoms, charges, scaled_positions, cell, pbc, constraint, calculator) # remove anything that ASE added that we don't want for p in remove_properties: self.remove_property(p) if isinstance(symbols, ase.Atoms): self.copy_from(symbols) ## end ASE compatibility if set_species and self.has_property('Z'): if not self.has_property('species'): self.add_property('species', ' '*TABLE_STRING_LENGTH) if self.n != 0 and not (self.z == 0).all(): mass = None if self.has_property('mass'): mass = self.mass.copy() self.set_atoms(self.z, mass=mass) # initialise species from z if info is not None: self.params.update(info) self._initialised = True # synonyms for backwards compatibility self.neighbours = self.connect self.hysteretic_neighbours = self.hysteretic_connect
[docs] def set_atomic_numbers(self, numbers, set_species=True): """Set atomic numbers and optionally also species property (default True)""" # override ase.Atoms.set_atomic_numbers() to keep QUIP Z and species in sync ase.Atoms.set_atomic_numbers(self, numbers) if set_species: if not self.has_property('species'): self.add_property('species', ' '*TABLE_STRING_LENGTH) if self.n != 0 and not (self.z == 0).all(): self.set_atoms(self.z) # set species from Z
[docs] def set_chemical_symbols(self, symbols, set_species=True): """Set chemical symbols - sets Z and optionally also species properties (default True)""" # override ase.Atoms.set_chemical_symbols() to keep QUIP Z and species in sync ase.Atoms.set_chemical_symbols(self, symbols) if set_species: if not self.has_property('species'): self.add_property('species', ' '*TABLE_STRING_LENGTH) if self.n != 0 and not (self.z == 0).all(): self.set_atoms(self.z) # set species from Z
[docs] def new_array(self, name, a, dtype=None, shape=None): # we overrride ase.Atoms.new_array() to allow "special" arrays # like "numbers", "positions" to be added more than once without # raising a RuntimeError if name in self.name_map and name in self.arrays: self.arrays[name] = a return ase.Atoms.new_array(self, name, a, dtype, shape)
[docs] def set_lattice(self, lattice, scale_positions=False): """Change the lattice vectors, keeping the inverse lattice vectors up to date. Optionally map the existing atoms into the new cell and recalculate connectivity (by default scale_positions=False).""" _atoms.Atoms.set_lattice(self, lattice, scale_positions)
def _get_cell(self): """Get ASE cell from QUIP lattice""" return self.lattice.view(np.ndarray).T def _set_cell(self, cell): """Set QUIP lattice from ASE cell""" self.set_lattice(cell.T, scale_positions=False) _cell = property(_get_cell, _set_cell) def _set_pbc(self, pbc): self.is_periodic = np.array(pbc).astype(int) def _get_pbc(self): return self.is_periodic.view(np.ndarray) == QUIPPY_TRUE _pbc = property(_get_pbc, _set_pbc) def _get_ase_arrays(self): """Provides access to ASE arrays, stored in QUIP properties dict""" return self._ase_arrays def _set_ase_arrays(self, value): """Set ASE arrays. Does not remove existing QUIP properties.""" self._ase_arrays.update(value) arrays = property(_get_ase_arrays, _set_ase_arrays) def _get_info(self): """ASE info dictionary Entries are actually stored in QUIP params dictionary.""" return self.params def _set_info(self, value): """Set ASE info dictionary. Entries are actually stored in QUIP params dictionary. Note that clearing doesn't empty params """ self.params.update(value) info = property(_get_info, _set_info) def _indices(self): """Return array of atoms indices If global ``fortran_indexing`` is True, returns FortranArray containing numbers 1..self.n. Otherwise, returns a standard numpuy array containing numbers in range 0..(self.n-1).""" if get_fortran_indexing(): return farray(list(frange(len(self)))) else: return np.array(list(range(len(self)))) indices = property(_indices)
[docs] def iteratoms(self): """Iterate over atoms, calling get_atom() for each one""" for i in self.indices: yield self.get_atom(i)
[docs] def equivalent(self, other): """Test for equivalence of two Atoms objects. Equivalence is less strong than equality. Equality (written `self == other`) requires all properties and parameters to be equal. Equivalence requires only that the number of atoms, positions, atomic numbers, unit cell and periodic boundary conditions match. .. note:: The quippy expression a.equivalent(b) has the same definition as a == b in ASE. This means that a quippy.Atoms instance can be compared with an ase.Atoms instance using this method. """ try: a = self.arrays b = other.arrays return (len(self) == len(other) and (a['positions'] == b['positions']).all() and (a['numbers'] == b['numbers']).all() and (self._cell == other.cell).all() and (self._pbc == other.pbc).all()) except AttributeError: return False
[docs] @classmethod def read(cls, source, format=None, **kwargs): """ Class method to read Atoms object from file `source` according to `format` If `format` is None, filetype is inferred from filename. Returns a new Atoms instance; to read into an existing Atoms object, use the read_from() method. If `source` corresponds to a known format then it used to construct an appropriate iterator from the :attr:`AtomsReaders` dictionary. See :ref:`fileformats` for a list of supported file formats. If `source` corresponds to an unknown format then it is expected to be an iterator returning :class:`Atoms` objects. """ if isinstance(source, basestring) and '@' in os.path.basename(source): source, frame = source.split('@') if source.endswith('.db'): source = source+'@'+frame format = 'db' else: frame = parse_slice(frame) if 'frame' in kwargs: raise ValueError("Conflicting frame references given: kwarg frame=%r and @-reference %r" % (kwargs['frame'], frame)) if not isinstance(frame, int): raise ValueError("Frame @-reference %r does not resolve to single frame" % frame) kwargs['frame'] = frame from import AtomsReaders filename, source, format = infer_format(source, format, AtomsReaders) opened = False if format in AtomsReaders: source = AtomsReaders[format](source, format=format, **kwargs) opened = True if isinstance(source, basestring): raise IOError("Don't know how to read from file '%s'" % source) if not hasattr(source, '__iter__'): raise IOError('Cannot read from %r - not an iterator' % source) at = iter(source).next() if not isinstance(at, cls): raise ValueError('Object %r read from %r is not Atoms instance' % (at, source)) if opened and hasattr(source, 'close'): source.close() if filename is not None: at.filename = filename return at
[docs] def write(self, dest=None, format=None, properties=None, prefix=None, **kwargs): """ Write this :class:`Atoms` object to `dest`. If `format` is absent it is inferred from the file extension or type of `dest`, as described for the :meth:`read` method. If `properties` is present, it should be a list of property names to include in the output file, e.g. `['species', 'pos']`. See :ref:`fileformats` for a list of supported file formats. """ if dest is None: # if filename is missing, save back to file from # which we loaded configuration if hasattr(self, 'filename'): dest = self.filename else: raise ValueError("No 'dest' and Atoms has no stored filename") from import AtomsWriters filename, dest, format = infer_format(dest, format, AtomsWriters) opened = filename is not None if format in AtomsWriters: dest = AtomsWriters[format](dest, **kwargs) if not hasattr(dest, 'write'): raise ValueError('Don\'t know how to write to "%s" in format "%s"' % (dest, format)) write_kwargs = {} if properties is not None: write_kwargs['properties'] = properties if prefix is not None: write_kwargs['prefix'] = prefix try: res = dest.write(self, **write_kwargs) except TypeError: raise ValueError('destination %r doesn\'t support arguments %r' % (dest, write_kwargs)) if opened and hasattr(dest, 'close'): dest.close() return res
[docs] def select(self, mask=None, list=None, orig_index=True): """Return a new :class:`Atoms` containing a subset of the atoms in this Atoms object One of either `mask` or `list` should be present. If `mask` is given it should be a rank one array of length `self.n`. In this case atoms corresponding to true values in `mask` will be included in the result. If `list` is present it should be an arry of list containing atom indices to include in the result. If `orig_index` is True (default), the new object will contain an ``orig_index`` property mapping the indices of the new atoms back to the original larger Atoms object. """ if mask is not None: mask = farray(mask) out = self.__class__(n=mask.sum(), lattice=self.lattice, properties={}, params={}), self, mask=mask, orig_index=orig_index) elif list is not None: list = farray(list) out = self.__class__(n=len(list), lattice=self.lattice), self, list=list, orig_index=orig_index) else: raise ValueError('Either mask or list must be present.') return out
[docs] def copy(self): """ Return a copy of this :class:`Atoms` object """ other = self.__class__(n=self.n, lattice=self.lattice,, params=self.params) # copy any normal (not Fortran) attributes for k, v in self.__dict__.iteritems(): if not k.startswith('_') and k not in other.__dict__: other.__dict__[k] = v # from _atoms.Atoms other.cutoff = self.cutoff other.cutoff_skin = self.cutoff_skin other.nneightol = self.nneightol # from ase.Atoms other.constraints = copy.deepcopy(self.constraints) return other
[docs] def copy_from(self, other): """Replace contents of this Atoms object with data from `other`.""" self.__class__.__del__(self) if isinstance(other, _atoms.Atoms): _atoms.Atoms.__init__(self, n=other.n, lattice=other.lattice,, params=other.params) self.cutoff = other.cutoff self.cutoff_skin = other.cutoff_skin self.nneightol = other.nneightol elif isinstance(other, ase.Atoms): _atoms.Atoms.__init__(self, n=0, lattice=np.eye(3)) ase.Atoms.__init__(self, other) # copy params/info dicts if hasattr(other, 'params'): self.params.update(other.params) if hasattr(other, 'info'): self.params.update( if 'nneightol' in self.nneightol =['nneightol'] if 'cutoff' in self.set_cutoff(['cutoff'],'cutoff_break')) if isinstance('spacegroup', None), Spacegroup): self.params['spacegroup'] =['spacegroup'].symbol # create extra properties for any non-standard arrays standard_ase_arrays = ['positions', 'numbers', 'masses', 'initial_charges', 'momenta', 'tags', 'initial_magmoms' ] for ase_name, value in other.arrays.iteritems(): quippy_name = self.name_map.get(ase_name, ase_name) if ase_name not in standard_ase_arrays: self.add_property(quippy_name, np.transpose(value)) self.constraints = copy.deepcopy(other.constraints) else: raise TypeError('can only copy from instances of quippy.Atoms or ase.Atoms') # copy any normal (not Fortran) attributes for k, v in other.__dict__.iteritems(): if not k.startswith('_') and k not in self.__dict__: self.__dict__[k] = v
[docs] def read_from(self, source, **readargs): """Replace contents of this Atoms object with Atoms read from `source`""" try: self.copy_from(source) except TypeError: tmp =, **readargs) self.shallow_copy_from(tmp)
# tmp goes out of scope here, but reference counting # prevents it from being free'd. def __getattr__(self, name): #print 'getattr', name #if name in if name == '_fpointer': raise AttributeError('Atoms object not initialised!') try: return[name] except KeyError: try: return self.params[name] except KeyError: raise AttributeError('Unknown Atoms attribute %s' % name) def __setattr__(self, name, value): #print 'setattr', name, value if not '_initialised' in self.__dict__: object.__setattr__(self, name, value) elif is not None and name in[name][...] = value elif self.params._fpointer is not None and name in self.params: if self.params.is_array(name): self.params[name][...] = value else: self.params[name] = value else: object.__setattr__(self, name, value)
[docs] def md5_hash(self, ndigits): """Hash an atoms object with a precision of ndigits decimal digits. Atomic numbers, lattice and fractional positions are fed to MD5 to form the hash.""" def rounded_string_rep(a, ndigits): return np.array2string(a, precision=ndigits, suppress_small=True).replace('-0. ', ' 0. ') # Compute fractional positions, round them to ndigits, then sort them # for hash stability flat_frac_pos =,self.pos).flatten() flat_frac_pos.sort() # md5 module deprecated in Python 2.5 and later try: import hashlib md5 = hashlib.md5 except ImportError: import md5 md5 = m = md5() m.update(rounded_string_rep(self.lattice.flatten(), ndigits)) m.update(str(self.z)) m.update(rounded_string_rep(flat_frac_pos, ndigits)) return m.hexdigest()
def __hash__(self): return hash(self.md5_hash(4)) #def __getitem__(self, i): # we override ase.Atoms.__getitem__ so we can raise # exception if we're using fortran indexing # if self.fortran_indexing: # raise RuntimeError('Atoms[i] inconsistent with fortran indexing') # return ase.Atoms.__getitem__(self, i)
[docs] def get_atom(self, i): """Return a dictionary containing the properties of the atom with index `i`. If fortran_indexing=True (the default), `i` should be in range 1..self.n, otherwise it should be in range 0..(self.n-1).""" if (get_fortran_indexing() and (i < 1 or i > self.n)) or \ (not get_fortran_indexing() and (i < 0 or i > self.n-1)): raise IndexError('Atoms index out of range') atom = {} atom['_index'] = i atom['atoms'] = self for k in v =[k][...,i] if isinstance(v,np.ndarray): if v.dtype.kind == 'S': v = ''.join(v).strip() elif v.shape == (): v = v.item() atom[k.lower()] = v return atom
[docs] def print_atom(self, i): """Pretty-print the properties of the atom with index `i`""" at = self.get_atom(i) title = 'Atom %d' % at['_index'] title = title + '\n' + '-'*len(title)+'\n\n' fields = ['%-15s = %s' % (k,at[k]) for k in sorted(at.keys()) if k not in ['_index', 'atoms']] print title+'\n'.join(fields)
[docs] def density(self): """Density in units of :math:`g/m^3`. If `mass` property exists, use that, otherwise we use `z` and ElementMass table.""" if self.has_property('mass'): mass = sum(self.mass)/MASSCONVERT/1.0e3 else: mass = sum(ElementMass[z] for z in self.z)/MASSCONVERT/1.0e3 return mass/(N_A*self.cell_volume()*1.0e-30)/1.0e3
[docs] def add_property(self, name, value, n_cols=None, overwrite=None, property_type=None): """ Add a new property to this Atoms object. `name` is the name of the new property and `value` should be either a scalar or an array representing the value, which should be either integer, real, logical or string. If a scalar is given for `value` it is copied to every element in the new property. `n_cols` can be specified to create a 2D property from a scalar initial value - the default is 1 which creates a 1D property. If an array is given for `value` it should either have shape (self.n,) for a 1D property or (n_cols,self.n) for a 2D property. In this case `n_cols` is inferred from the shape of the `value` and shouldn't be passed as an argument. If `property_type` is present, then no attempt is made to infer the type from `value`. This is necessary to resolve ambiguity between integer and logical types. If property with the same type is already present then no error occurs.If `overwrite` is true, the value will be overwritten with that given in `value`, otherwise the old value is retained. Here are some examples:: a = Atoms(n=10, lattice=10.0*fidentity(3)) a.add_property('mark', 1) # Scalar integer a.add_property('bool', False) # Scalar logical a.add_property('local_energy', 0.0) # Scalar real a.add_property('force', 0.0, n_cols=3) # Vector real a.add_property('label', '') # Scalar string a.add_property('count', [1,2,3,4,5,6,7,8,9,10]) # From list a.add_property('norm_pos', a.pos.norm()) # From 1D array a.add_property('pos', new_pos) # Overwrite positions with array new_pos # which should have shape (3,10) """ kwargs = {} if n_cols is not None: kwargs['n_cols'] = n_cols if overwrite is not None: kwargs['overwrite'] = overwrite if (isinstance(value, np.ndarray) and value.dtype.kind in ['O', 'S'] and value.shape != (len(self), TABLE_STRING_LENGTH)): value = s2a(value.astype('str'), TABLE_STRING_LENGTH).T if property_type is None: _atoms.Atoms.add_property(self, name, value, **kwargs) else: # override value_ref if property_type is specified new_property = not self.has_property(name) type_to_value_ref = { T_INTEGER_A : 0, T_REAL_A : 0.0, T_CHAR_A : " "*TABLE_STRING_LENGTH, T_LOGICAL_A : False, T_INTEGER_A2 : 0, T_REAL_A2: 0.0 } try: value_ref = type_to_value_ref[property_type] except KeyError: raise ValueError('Unknown property_type %d' % property_type) if (hasattr(value, 'shape') and len(value.shape) == 2 and property_type != T_CHAR_A and n_cols is None): kwargs['n_cols'] = value.shape[0] _atoms.Atoms.add_property(self, name, value_ref, **kwargs) if new_property or overwrite: getattr(self, name.lower())[:] = value
def __getstate__(self): return self.write('string') def __setstate__(self, state): self.read_from(state, format='string') def __reduce__(self): return (Atoms, (), self.__getstate__(), None, None)
[docs] def mem_estimate(self): """Estimate memory usage of this Atoms object, in bytes""" sizeof_table = 320 mem = sum([p.itemsize*p.size for p in]) if self.connect.initialised: c = self.connect mem += sizeof_table*self.n*2 # neighbour1 and neighbour2 tables mem += 32*c.n_neighbours_total() # neighbour data mem += c.cell_heads.size*c.cell_heads.itemsize # cell data return mem
[docs] def extend(self, other): """Extend atoms object by appending atoms from *other*.""" # modified version of ase.Atoms.extend() to work with QUIP data storage if isinstance(other, ase.Atom): other = self.__class__([other]) n1 = len(self) n2 = len(other) # first make a copy of self.arrays so that we can resize Atoms arrays = dict([ (key, value.copy()) for (key, value) in self.arrays.items() ]) atomslog.debug('old arrays %r' % arrays) for name, a1 in arrays.items(): a = np.zeros((n1 + n2,) + a1.shape[1:], a1.dtype) a[:n1] = a1 if name == 'masses': a2 = other.get_masses() else: a2 = other.arrays.get(name) if a2 is not None: a[n1:] = a2 self.arrays[name] = a for name, a2 in other.arrays.items(): if name in self.arrays: continue a = np.empty((n1 + n2,) + a2.shape[1:], a2.dtype) a[n1:] = a2 if name == 'masses': a[:n1] = self.get_masses()[:n1] else: a[:n1] = 0 self.set_array(name, a) atomslog.debug('new arrays %r' % self.arrays) return self
__iadd__ = extend def __imul__(self, m): """In-place repeat of atoms.""" # modified version of ase.Atoms.extend() to work with QUIP data storage if isinstance(m, int): m = (m, m, m) M = np.product(m) n = len(self) # first make a copy of self.arrays so that we can resize Atoms arrays = dict([ (key, value.copy()) for (key, value) in self.arrays.items() ]) for name, a in arrays.items(): self.arrays[name] = np.tile(a, (M,) + (1,) * (len(a.shape) - 1)) positions = self.arrays['positions'] i0 = 0 for m0 in range(m[0]): for m1 in range(m[1]): for m2 in range(m[2]): i1 = i0 + n positions[i0:i1] +=, m1, m2), self._cell) i0 = i1 if self.constraints is not None: self.constraints = [c.repeat(m, n) for c in self.constraints] self._cell = np.array([m[c] * self._cell[c] for c in range(3)]) return self
FortranDerivedTypes['type(atoms)'] = Atoms FortranDerivedTypes['type(connection)'] = Connection