1# Copyright 2008, 2009 CAMd
2# (see accompanying license files for details).
3
4"""Definition of the Atoms class.
5
6This module defines the central object in the ASE package: the Atoms
7object.
8"""
9import copy
10import numbers
11from math import cos, sin, pi
12
13import numpy as np
14
15import ase.units as units
16from ase.atom import Atom
17from ase.cell import Cell
18from ase.stress import voigt_6_to_full_3x3_stress, full_3x3_to_voigt_6_stress
19from ase.data import atomic_masses, atomic_masses_common
20from ase.geometry import (wrap_positions, find_mic, get_angles, get_distances,
21                          get_dihedrals)
22from ase.symbols import Symbols, symbols2numbers
23from ase.utils import deprecated
24
25
26class Atoms:
27    """Atoms object.
28
29    The Atoms object can represent an isolated molecule, or a
30    periodically repeated structure.  It has a unit cell and
31    there may be periodic boundary conditions along any of the three
32    unit cell axes.
33    Information about the atoms (atomic numbers and position) is
34    stored in ndarrays.  Optionally, there can be information about
35    tags, momenta, masses, magnetic moments and charges.
36
37    In order to calculate energies, forces and stresses, a calculator
38    object has to attached to the atoms object.
39
40    Parameters:
41
42    symbols: str (formula) or list of str
43        Can be a string formula, a list of symbols or a list of
44        Atom objects.  Examples: 'H2O', 'COPt12', ['H', 'H', 'O'],
45        [Atom('Ne', (x, y, z)), ...].
46    positions: list of xyz-positions
47        Atomic positions.  Anything that can be converted to an
48        ndarray of shape (n, 3) will do: [(x1,y1,z1), (x2,y2,z2),
49        ...].
50    scaled_positions: list of scaled-positions
51        Like positions, but given in units of the unit cell.
52        Can not be set at the same time as positions.
53    numbers: list of int
54        Atomic numbers (use only one of symbols/numbers).
55    tags: list of int
56        Special purpose tags.
57    momenta: list of xyz-momenta
58        Momenta for all atoms.
59    masses: list of float
60        Atomic masses in atomic units.
61    magmoms: list of float or list of xyz-values
62        Magnetic moments.  Can be either a single value for each atom
63        for collinear calculations or three numbers for each atom for
64        non-collinear calculations.
65    charges: list of float
66        Initial atomic charges.
67    cell: 3x3 matrix or length 3 or 6 vector
68        Unit cell vectors.  Can also be given as just three
69        numbers for orthorhombic cells, or 6 numbers, where
70        first three are lengths of unit cell vectors, and the
71        other three are angles between them (in degrees), in following order:
72        [len(a), len(b), len(c), angle(b,c), angle(a,c), angle(a,b)].
73        First vector will lie in x-direction, second in xy-plane,
74        and the third one in z-positive subspace.
75        Default value: [0, 0, 0].
76    celldisp: Vector
77        Unit cell displacement vector. To visualize a displaced cell
78        around the center of mass of a Systems of atoms. Default value
79        = (0,0,0)
80    pbc: one or three bool
81        Periodic boundary conditions flags.  Examples: True,
82        False, 0, 1, (1, 1, 0), (True, False, False).  Default
83        value: False.
84    constraint: constraint object(s)
85        Used for applying one or more constraints during structure
86        optimization.
87    calculator: calculator object
88        Used to attach a calculator for calculating energies and atomic
89        forces.
90    info: dict of key-value pairs
91        Dictionary of key-value pairs with additional information
92        about the system.  The following keys may be used by ase:
93
94          - spacegroup: Spacegroup instance
95          - unit_cell: 'conventional' | 'primitive' | int | 3 ints
96          - adsorbate_info: Information about special adsorption sites
97
98        Items in the info attribute survives copy and slicing and can
99        be stored in and retrieved from trajectory files given that the
100        key is a string, the value is JSON-compatible and, if the value is a
101        user-defined object, its base class is importable.  One should
102        not make any assumptions about the existence of keys.
103
104    Examples:
105
106    These three are equivalent:
107
108    >>> d = 1.104  # N2 bondlength
109    >>> a = Atoms('N2', [(0, 0, 0), (0, 0, d)])
110    >>> a = Atoms(numbers=[7, 7], positions=[(0, 0, 0), (0, 0, d)])
111    >>> a = Atoms([Atom('N', (0, 0, 0)), Atom('N', (0, 0, d))])
112
113    FCC gold:
114
115    >>> a = 4.05  # Gold lattice constant
116    >>> b = a / 2
117    >>> fcc = Atoms('Au',
118    ...             cell=[(0, b, b), (b, 0, b), (b, b, 0)],
119    ...             pbc=True)
120
121    Hydrogen wire:
122
123    >>> d = 0.9  # H-H distance
124    >>> h = Atoms('H', positions=[(0, 0, 0)],
125    ...           cell=(d, 0, 0),
126    ...           pbc=(1, 0, 0))
127    """
128
129    ase_objtype = 'atoms'  # For JSONability
130
131    def __init__(self, symbols=None,
132                 positions=None, numbers=None,
133                 tags=None, momenta=None, masses=None,
134                 magmoms=None, charges=None,
135                 scaled_positions=None,
136                 cell=None, pbc=None, celldisp=None,
137                 constraint=None,
138                 calculator=None,
139                 info=None,
140                 velocities=None):
141
142        self._cellobj = Cell.new()
143        self._pbc = np.zeros(3, bool)
144
145        atoms = None
146
147        if hasattr(symbols, 'get_positions'):
148            atoms = symbols
149            symbols = None
150        elif (isinstance(symbols, (list, tuple)) and
151              len(symbols) > 0 and isinstance(symbols[0], Atom)):
152            # Get data from a list or tuple of Atom objects:
153            data = [[atom.get_raw(name) for atom in symbols]
154                    for name in
155                    ['position', 'number', 'tag', 'momentum',
156                     'mass', 'magmom', 'charge']]
157            atoms = self.__class__(None, *data)
158            symbols = None
159
160        if atoms is not None:
161            # Get data from another Atoms object:
162            if scaled_positions is not None:
163                raise NotImplementedError
164            if symbols is None and numbers is None:
165                numbers = atoms.get_atomic_numbers()
166            if positions is None:
167                positions = atoms.get_positions()
168            if tags is None and atoms.has('tags'):
169                tags = atoms.get_tags()
170            if momenta is None and atoms.has('momenta'):
171                momenta = atoms.get_momenta()
172            if magmoms is None and atoms.has('initial_magmoms'):
173                magmoms = atoms.get_initial_magnetic_moments()
174            if masses is None and atoms.has('masses'):
175                masses = atoms.get_masses()
176            if charges is None and atoms.has('initial_charges'):
177                charges = atoms.get_initial_charges()
178            if cell is None:
179                cell = atoms.get_cell()
180            if celldisp is None:
181                celldisp = atoms.get_celldisp()
182            if pbc is None:
183                pbc = atoms.get_pbc()
184            if constraint is None:
185                constraint = [c.copy() for c in atoms.constraints]
186            if calculator is None:
187                calculator = atoms.calc
188            if info is None:
189                info = copy.deepcopy(atoms.info)
190
191        self.arrays = {}
192
193        if symbols is None:
194            if numbers is None:
195                if positions is not None:
196                    natoms = len(positions)
197                elif scaled_positions is not None:
198                    natoms = len(scaled_positions)
199                else:
200                    natoms = 0
201                numbers = np.zeros(natoms, int)
202            self.new_array('numbers', numbers, int)
203        else:
204            if numbers is not None:
205                raise TypeError(
206                    'Use only one of "symbols" and "numbers".')
207            else:
208                self.new_array('numbers', symbols2numbers(symbols), int)
209
210        if self.numbers.ndim != 1:
211            raise ValueError('"numbers" must be 1-dimensional.')
212
213        if cell is None:
214            cell = np.zeros((3, 3))
215        self.set_cell(cell)
216
217        if celldisp is None:
218            celldisp = np.zeros(shape=(3, 1))
219        self.set_celldisp(celldisp)
220
221        if positions is None:
222            if scaled_positions is None:
223                positions = np.zeros((len(self.arrays['numbers']), 3))
224            else:
225                assert self.cell.rank == 3
226                positions = np.dot(scaled_positions, self.cell)
227        else:
228            if scaled_positions is not None:
229                raise TypeError(
230                    'Use only one of "symbols" and "numbers".')
231        self.new_array('positions', positions, float, (3,))
232
233        self.set_constraint(constraint)
234        self.set_tags(default(tags, 0))
235        self.set_masses(default(masses, None))
236        self.set_initial_magnetic_moments(default(magmoms, 0.0))
237        self.set_initial_charges(default(charges, 0.0))
238        if pbc is None:
239            pbc = False
240        self.set_pbc(pbc)
241        self.set_momenta(default(momenta, (0.0, 0.0, 0.0)),
242                         apply_constraint=False)
243
244        if velocities is not None:
245            if momenta is None:
246                self.set_velocities(velocities)
247            else:
248                raise TypeError(
249                    'Use only one of "momenta" and "velocities".')
250
251        if info is None:
252            self.info = {}
253        else:
254            self.info = dict(info)
255
256        self.calc = calculator
257
258    @property
259    def symbols(self):
260        """Get chemical symbols as a :class:`ase.symbols.Symbols` object.
261
262        The object works like ``atoms.numbers`` except its values
263        are strings.  It supports in-place editing."""
264        return Symbols(self.numbers)
265
266    @symbols.setter
267    def symbols(self, obj):
268        new_symbols = Symbols.fromsymbols(obj)
269        self.numbers[:] = new_symbols.numbers
270
271    @deprecated(DeprecationWarning('Please use atoms.calc = calc'))
272    def set_calculator(self, calc=None):
273        """Attach calculator object.
274
275        Please use the equivalent atoms.calc = calc instead of this
276        method."""
277        self.calc = calc
278
279    @deprecated(DeprecationWarning('Please use atoms.calc'))
280    def get_calculator(self):
281        """Get currently attached calculator object.
282
283        Please use the equivalent atoms.calc instead of
284        atoms.get_calculator()."""
285        return self.calc
286
287    @property
288    def calc(self):
289        """Calculator object."""
290        return self._calc
291
292    @calc.setter
293    def calc(self, calc):
294        self._calc = calc
295        if hasattr(calc, 'set_atoms'):
296            calc.set_atoms(self)
297
298    @calc.deleter  # type: ignore
299    @deprecated(DeprecationWarning('Please use atoms.calc = None'))
300    def calc(self):
301        self._calc = None
302
303    @property  # type: ignore
304    @deprecated('Please use atoms.cell.rank instead')
305    def number_of_lattice_vectors(self):
306        """Number of (non-zero) lattice vectors."""
307        return self.cell.rank
308
309    def set_constraint(self, constraint=None):
310        """Apply one or more constrains.
311
312        The *constraint* argument must be one constraint object or a
313        list of constraint objects."""
314        if constraint is None:
315            self._constraints = []
316        else:
317            if isinstance(constraint, list):
318                self._constraints = constraint
319            elif isinstance(constraint, tuple):
320                self._constraints = list(constraint)
321            else:
322                self._constraints = [constraint]
323
324    def _get_constraints(self):
325        return self._constraints
326
327    def _del_constraints(self):
328        self._constraints = []
329
330    constraints = property(_get_constraints, set_constraint, _del_constraints,
331                           'Constraints of the atoms.')
332
333    def set_cell(self, cell, scale_atoms=False, apply_constraint=True):
334        """Set unit cell vectors.
335
336        Parameters:
337
338        cell: 3x3 matrix or length 3 or 6 vector
339            Unit cell.  A 3x3 matrix (the three unit cell vectors) or
340            just three numbers for an orthorhombic cell. Another option is
341            6 numbers, which describes unit cell with lengths of unit cell
342            vectors and with angles between them (in degrees), in following
343            order: [len(a), len(b), len(c), angle(b,c), angle(a,c),
344            angle(a,b)].  First vector will lie in x-direction, second in
345            xy-plane, and the third one in z-positive subspace.
346        scale_atoms: bool
347            Fix atomic positions or move atoms with the unit cell?
348            Default behavior is to *not* move the atoms (scale_atoms=False).
349        apply_constraint: bool
350            Whether to apply constraints to the given cell.
351
352        Examples:
353
354        Two equivalent ways to define an orthorhombic cell:
355
356        >>> atoms = Atoms('He')
357        >>> a, b, c = 7, 7.5, 8
358        >>> atoms.set_cell([a, b, c])
359        >>> atoms.set_cell([(a, 0, 0), (0, b, 0), (0, 0, c)])
360
361        FCC unit cell:
362
363        >>> atoms.set_cell([(0, b, b), (b, 0, b), (b, b, 0)])
364
365        Hexagonal unit cell:
366
367        >>> atoms.set_cell([a, a, c, 90, 90, 120])
368
369        Rhombohedral unit cell:
370
371        >>> alpha = 77
372        >>> atoms.set_cell([a, a, a, alpha, alpha, alpha])
373        """
374
375        # Override pbcs if and only if given a Cell object:
376        cell = Cell.new(cell)
377
378        # XXX not working well during initialize due to missing _constraints
379        if apply_constraint and hasattr(self, '_constraints'):
380            for constraint in self.constraints:
381                if hasattr(constraint, 'adjust_cell'):
382                    constraint.adjust_cell(self, cell)
383
384        if scale_atoms:
385            M = np.linalg.solve(self.cell.complete(), cell.complete())
386            self.positions[:] = np.dot(self.positions, M)
387
388        self.cell[:] = cell
389
390    def set_celldisp(self, celldisp):
391        """Set the unit cell displacement vectors."""
392        celldisp = np.array(celldisp, float)
393        self._celldisp = celldisp
394
395    def get_celldisp(self):
396        """Get the unit cell displacement vectors."""
397        return self._celldisp.copy()
398
399    def get_cell(self, complete=False):
400        """Get the three unit cell vectors as a `class`:ase.cell.Cell` object.
401
402        The Cell object resembles a 3x3 ndarray, and cell[i, j]
403        is the jth Cartesian coordinate of the ith cell vector."""
404        if complete:
405            cell = self.cell.complete()
406        else:
407            cell = self.cell.copy()
408
409        return cell
410
411    @deprecated('Please use atoms.cell.cellpar() instead')
412    def get_cell_lengths_and_angles(self):
413        """Get unit cell parameters. Sequence of 6 numbers.
414
415        First three are unit cell vector lengths and second three
416        are angles between them::
417
418            [len(a), len(b), len(c), angle(b,c), angle(a,c), angle(a,b)]
419
420        in degrees.
421        """
422        return self.cell.cellpar()
423
424    @deprecated('Please use atoms.cell.reciprocal()')
425    def get_reciprocal_cell(self):
426        """Get the three reciprocal lattice vectors as a 3x3 ndarray.
427
428        Note that the commonly used factor of 2 pi for Fourier
429        transforms is not included here."""
430
431        return self.cell.reciprocal()
432
433    @property
434    def pbc(self):
435        """Reference to pbc-flags for in-place manipulations."""
436        return self._pbc
437
438    @pbc.setter
439    def pbc(self, pbc):
440        self._pbc[:] = pbc
441
442    def set_pbc(self, pbc):
443        """Set periodic boundary condition flags."""
444        self.pbc = pbc
445
446    def get_pbc(self):
447        """Get periodic boundary condition flags."""
448        return self.pbc.copy()
449
450    def new_array(self, name, a, dtype=None, shape=None):
451        """Add new array.
452
453        If *shape* is not *None*, the shape of *a* will be checked."""
454
455        if dtype is not None:
456            a = np.array(a, dtype, order='C')
457            if len(a) == 0 and shape is not None:
458                a.shape = (-1,) + shape
459        else:
460            if not a.flags['C_CONTIGUOUS']:
461                a = np.ascontiguousarray(a)
462            else:
463                a = a.copy()
464
465        if name in self.arrays:
466            raise RuntimeError('Array {} already present'.format(name))
467
468        for b in self.arrays.values():
469            if len(a) != len(b):
470                raise ValueError('Array "%s" has wrong length: %d != %d.' %
471                                 (name, len(a), len(b)))
472            break
473
474        if shape is not None and a.shape[1:] != shape:
475            raise ValueError('Array "%s" has wrong shape %s != %s.' %
476                             (name, a.shape, (a.shape[0:1] + shape)))
477
478        self.arrays[name] = a
479
480    def get_array(self, name, copy=True):
481        """Get an array.
482
483        Returns a copy unless the optional argument copy is false.
484        """
485        if copy:
486            return self.arrays[name].copy()
487        else:
488            return self.arrays[name]
489
490    def set_array(self, name, a, dtype=None, shape=None):
491        """Update array.
492
493        If *shape* is not *None*, the shape of *a* will be checked.
494        If *a* is *None*, then the array is deleted."""
495
496        b = self.arrays.get(name)
497        if b is None:
498            if a is not None:
499                self.new_array(name, a, dtype, shape)
500        else:
501            if a is None:
502                del self.arrays[name]
503            else:
504                a = np.asarray(a)
505                if a.shape != b.shape:
506                    raise ValueError('Array "%s" has wrong shape %s != %s.' %
507                                     (name, a.shape, b.shape))
508                b[:] = a
509
510    def has(self, name):
511        """Check for existence of array.
512
513        name must be one of: 'tags', 'momenta', 'masses', 'initial_magmoms',
514        'initial_charges'."""
515        # XXX extend has to calculator properties
516        return name in self.arrays
517
518    def set_atomic_numbers(self, numbers):
519        """Set atomic numbers."""
520        self.set_array('numbers', numbers, int, ())
521
522    def get_atomic_numbers(self):
523        """Get integer array of atomic numbers."""
524        return self.arrays['numbers'].copy()
525
526    def get_chemical_symbols(self):
527        """Get list of chemical symbol strings.
528
529        Equivalent to ``list(atoms.symbols)``."""
530        return list(self.symbols)
531
532    def set_chemical_symbols(self, symbols):
533        """Set chemical symbols."""
534        self.set_array('numbers', symbols2numbers(symbols), int, ())
535
536    def get_chemical_formula(self, mode='hill', empirical=False):
537        """Get the chemical formula as a string based on the chemical symbols.
538
539        Parameters:
540
541        mode: str
542            There are four different modes available:
543
544            'all': The list of chemical symbols are contracted to a string,
545            e.g. ['C', 'H', 'H', 'H', 'O', 'H'] becomes 'CHHHOH'.
546
547            'reduce': The same as 'all' where repeated elements are contracted
548            to a single symbol and a number, e.g. 'CHHHOCHHH' is reduced to
549            'CH3OCH3'.
550
551            'hill': The list of chemical symbols are contracted to a string
552            following the Hill notation (alphabetical order with C and H
553            first), e.g. 'CHHHOCHHH' is reduced to 'C2H6O' and 'SOOHOHO' to
554            'H2O4S'. This is default.
555
556            'metal': The list of chemical symbols (alphabetical metals,
557            and alphabetical non-metals)
558
559        empirical, bool (optional, default=False)
560            Divide the symbol counts by their greatest common divisor to yield
561            an empirical formula. Only for mode `metal` and `hill`.
562        """
563        return self.symbols.get_chemical_formula(mode, empirical)
564
565    def set_tags(self, tags):
566        """Set tags for all atoms. If only one tag is supplied, it is
567        applied to all atoms."""
568        if isinstance(tags, int):
569            tags = [tags] * len(self)
570        self.set_array('tags', tags, int, ())
571
572    def get_tags(self):
573        """Get integer array of tags."""
574        if 'tags' in self.arrays:
575            return self.arrays['tags'].copy()
576        else:
577            return np.zeros(len(self), int)
578
579    def set_momenta(self, momenta, apply_constraint=True):
580        """Set momenta."""
581        if (apply_constraint and len(self.constraints) > 0 and
582           momenta is not None):
583            momenta = np.array(momenta)  # modify a copy
584            for constraint in self.constraints:
585                if hasattr(constraint, 'adjust_momenta'):
586                    constraint.adjust_momenta(self, momenta)
587        self.set_array('momenta', momenta, float, (3,))
588
589    def set_velocities(self, velocities):
590        """Set the momenta by specifying the velocities."""
591        self.set_momenta(self.get_masses()[:, np.newaxis] * velocities)
592
593    def get_momenta(self):
594        """Get array of momenta."""
595        if 'momenta' in self.arrays:
596            return self.arrays['momenta'].copy()
597        else:
598            return np.zeros((len(self), 3))
599
600    def set_masses(self, masses='defaults'):
601        """Set atomic masses in atomic mass units.
602
603        The array masses should contain a list of masses.  In case
604        the masses argument is not given or for those elements of the
605        masses list that are None, standard values are set."""
606
607        if isinstance(masses, str):
608            if masses == 'defaults':
609                masses = atomic_masses[self.arrays['numbers']]
610            elif masses == 'most_common':
611                masses = atomic_masses_common[self.arrays['numbers']]
612        elif masses is None:
613            pass
614        elif not isinstance(masses, np.ndarray):
615            masses = list(masses)
616            for i, mass in enumerate(masses):
617                if mass is None:
618                    masses[i] = atomic_masses[self.numbers[i]]
619        self.set_array('masses', masses, float, ())
620
621    def get_masses(self):
622        """Get array of masses in atomic mass units."""
623        if 'masses' in self.arrays:
624            return self.arrays['masses'].copy()
625        else:
626            return atomic_masses[self.arrays['numbers']]
627
628    def set_initial_magnetic_moments(self, magmoms=None):
629        """Set the initial magnetic moments.
630
631        Use either one or three numbers for every atom (collinear
632        or non-collinear spins)."""
633
634        if magmoms is None:
635            self.set_array('initial_magmoms', None)
636        else:
637            magmoms = np.asarray(magmoms)
638            self.set_array('initial_magmoms', magmoms, float,
639                           magmoms.shape[1:])
640
641    def get_initial_magnetic_moments(self):
642        """Get array of initial magnetic moments."""
643        if 'initial_magmoms' in self.arrays:
644            return self.arrays['initial_magmoms'].copy()
645        else:
646            return np.zeros(len(self))
647
648    def get_magnetic_moments(self):
649        """Get calculated local magnetic moments."""
650        if self._calc is None:
651            raise RuntimeError('Atoms object has no calculator.')
652        return self._calc.get_magnetic_moments(self)
653
654    def get_magnetic_moment(self):
655        """Get calculated total magnetic moment."""
656        if self._calc is None:
657            raise RuntimeError('Atoms object has no calculator.')
658        return self._calc.get_magnetic_moment(self)
659
660    def set_initial_charges(self, charges=None):
661        """Set the initial charges."""
662
663        if charges is None:
664            self.set_array('initial_charges', None)
665        else:
666            self.set_array('initial_charges', charges, float, ())
667
668    def get_initial_charges(self):
669        """Get array of initial charges."""
670        if 'initial_charges' in self.arrays:
671            return self.arrays['initial_charges'].copy()
672        else:
673            return np.zeros(len(self))
674
675    def get_charges(self):
676        """Get calculated charges."""
677        if self._calc is None:
678            raise RuntimeError('Atoms object has no calculator.')
679        try:
680            return self._calc.get_charges(self)
681        except AttributeError:
682            from ase.calculators.calculator import PropertyNotImplementedError
683            raise PropertyNotImplementedError
684
685    def set_positions(self, newpositions, apply_constraint=True):
686        """Set positions, honoring any constraints. To ignore constraints,
687        use *apply_constraint=False*."""
688        if self.constraints and apply_constraint:
689            newpositions = np.array(newpositions, float)
690            for constraint in self.constraints:
691                constraint.adjust_positions(self, newpositions)
692
693        self.set_array('positions', newpositions, shape=(3,))
694
695    def get_positions(self, wrap=False, **wrap_kw):
696        """Get array of positions.
697
698        Parameters:
699
700        wrap: bool
701            wrap atoms back to the cell before returning positions
702        wrap_kw: (keyword=value) pairs
703            optional keywords `pbc`, `center`, `pretty_translation`, `eps`,
704            see :func:`ase.geometry.wrap_positions`
705        """
706        if wrap:
707            if 'pbc' not in wrap_kw:
708                wrap_kw['pbc'] = self.pbc
709            return wrap_positions(self.positions, self.cell, **wrap_kw)
710        else:
711            return self.arrays['positions'].copy()
712
713    def get_potential_energy(self, force_consistent=False,
714                             apply_constraint=True):
715        """Calculate potential energy.
716
717        Ask the attached calculator to calculate the potential energy and
718        apply constraints.  Use *apply_constraint=False* to get the raw
719        forces.
720
721        When supported by the calculator, either the energy extrapolated
722        to zero Kelvin or the energy consistent with the forces (the free
723        energy) can be returned.
724        """
725        if self._calc is None:
726            raise RuntimeError('Atoms object has no calculator.')
727        if force_consistent:
728            energy = self._calc.get_potential_energy(
729                self, force_consistent=force_consistent)
730        else:
731            energy = self._calc.get_potential_energy(self)
732        if apply_constraint:
733            for constraint in self.constraints:
734                if hasattr(constraint, 'adjust_potential_energy'):
735                    energy += constraint.adjust_potential_energy(self)
736        return energy
737
738    def get_properties(self, properties):
739        """This method is experimental; currently for internal use."""
740        # XXX Something about constraints.
741        if self._calc is None:
742            raise RuntimeError('Atoms object has no calculator.')
743        return self._calc.calculate_properties(self, properties)
744
745    def get_potential_energies(self):
746        """Calculate the potential energies of all the atoms.
747
748        Only available with calculators supporting per-atom energies
749        (e.g. classical potentials).
750        """
751        if self._calc is None:
752            raise RuntimeError('Atoms object has no calculator.')
753        return self._calc.get_potential_energies(self)
754
755    def get_kinetic_energy(self):
756        """Get the kinetic energy."""
757        momenta = self.arrays.get('momenta')
758        if momenta is None:
759            return 0.0
760        return 0.5 * np.vdot(momenta, self.get_velocities())
761
762    def get_velocities(self):
763        """Get array of velocities."""
764        momenta = self.get_momenta()
765        masses = self.get_masses()
766        return momenta / masses[:, np.newaxis]
767
768    def get_total_energy(self):
769        """Get the total energy - potential plus kinetic energy."""
770        return self.get_potential_energy() + self.get_kinetic_energy()
771
772    def get_forces(self, apply_constraint=True, md=False):
773        """Calculate atomic forces.
774
775        Ask the attached calculator to calculate the forces and apply
776        constraints.  Use *apply_constraint=False* to get the raw
777        forces.
778
779        For molecular dynamics (md=True) we don't apply the constraint
780        to the forces but to the momenta. When holonomic constraints for
781        rigid linear triatomic molecules are present, ask the constraints
782        to redistribute the forces within each triple defined in the
783        constraints (required for molecular dynamics with this type of
784        constraints)."""
785
786        if self._calc is None:
787            raise RuntimeError('Atoms object has no calculator.')
788        forces = self._calc.get_forces(self)
789
790        if apply_constraint:
791            # We need a special md flag here because for MD we want
792            # to skip real constraints but include special "constraints"
793            # Like Hookean.
794            for constraint in self.constraints:
795                if md and hasattr(constraint, 'redistribute_forces_md'):
796                    constraint.redistribute_forces_md(self, forces)
797                if not md or hasattr(constraint, 'adjust_potential_energy'):
798                    constraint.adjust_forces(self, forces)
799        return forces
800
801    # Informs calculators (e.g. Asap) that ideal gas contribution is added here.
802    _ase_handles_dynamic_stress = True
803
804    def get_stress(self, voigt=True, apply_constraint=True,
805                   include_ideal_gas=False):
806        """Calculate stress tensor.
807
808        Returns an array of the six independent components of the
809        symmetric stress tensor, in the traditional Voigt order
810        (xx, yy, zz, yz, xz, xy) or as a 3x3 matrix.  Default is Voigt
811        order.
812
813        The ideal gas contribution to the stresses is added if the
814        atoms have momenta and ``include_ideal_gas`` is set to True.
815        """
816
817        if self._calc is None:
818            raise RuntimeError('Atoms object has no calculator.')
819
820        stress = self._calc.get_stress(self)
821        shape = stress.shape
822
823        if shape == (3, 3):
824            # Convert to the Voigt form before possibly applying
825            # constraints and adding the dynamic part of the stress
826            # (the "ideal gas contribution").
827            stress = full_3x3_to_voigt_6_stress(stress)
828        else:
829            assert shape == (6,)
830
831        if apply_constraint:
832            for constraint in self.constraints:
833                if hasattr(constraint, 'adjust_stress'):
834                    constraint.adjust_stress(self, stress)
835
836        # Add ideal gas contribution, if applicable
837        if include_ideal_gas and self.has('momenta'):
838            stresscomp = np.array([[0, 5, 4], [5, 1, 3], [4, 3, 2]])
839            p = self.get_momenta()
840            masses = self.get_masses()
841            invmass = 1.0 / masses
842            invvol = 1.0 / self.get_volume()
843            for alpha in range(3):
844                for beta in range(alpha, 3):
845                    stress[stresscomp[alpha, beta]] -= (
846                        p[:, alpha] * p[:, beta] * invmass).sum() * invvol
847
848        if voigt:
849            return stress
850        else:
851            return voigt_6_to_full_3x3_stress(stress)
852
853    def get_stresses(self, include_ideal_gas=False, voigt=True):
854        """Calculate the stress-tensor of all the atoms.
855
856        Only available with calculators supporting per-atom energies and
857        stresses (e.g. classical potentials).  Even for such calculators
858        there is a certain arbitrariness in defining per-atom stresses.
859
860        The ideal gas contribution to the stresses is added if the
861        atoms have momenta and ``include_ideal_gas`` is set to True.
862        """
863        if self._calc is None:
864            raise RuntimeError('Atoms object has no calculator.')
865        stresses = self._calc.get_stresses(self)
866
867        # make sure `stresses` are in voigt form
868        if np.shape(stresses)[1:] == (3, 3):
869            stresses_voigt = [full_3x3_to_voigt_6_stress(s) for s in stresses]
870            stresses = np.array(stresses_voigt)
871
872        # REMARK: The ideal gas contribution is intensive, i.e., the volume
873        # is divided out. We currently don't check if `stresses` are intensive
874        # as well, i.e., if `a.get_stresses.sum(axis=0) == a.get_stress()`.
875        # It might be good to check this here, but adds computational overhead.
876
877        if include_ideal_gas and self.has('momenta'):
878            stresscomp = np.array([[0, 5, 4], [5, 1, 3], [4, 3, 2]])
879            if hasattr(self._calc, 'get_atomic_volumes'):
880                invvol = 1.0 / self._calc.get_atomic_volumes()
881            else:
882                invvol = self.get_global_number_of_atoms() / self.get_volume()
883            p = self.get_momenta()
884            invmass = 1.0 / self.get_masses()
885            for alpha in range(3):
886                for beta in range(alpha, 3):
887                    stresses[:, stresscomp[alpha, beta]] -= (
888                        p[:, alpha] * p[:, beta] * invmass * invvol)
889        if voigt:
890            return stresses
891        else:
892            stresses_3x3 = [voigt_6_to_full_3x3_stress(s) for s in stresses]
893            return np.array(stresses_3x3)
894
895    def get_dipole_moment(self):
896        """Calculate the electric dipole moment for the atoms object.
897
898        Only available for calculators which has a get_dipole_moment()
899        method."""
900
901        if self._calc is None:
902            raise RuntimeError('Atoms object has no calculator.')
903        return self._calc.get_dipole_moment(self)
904
905    def copy(self):
906        """Return a copy."""
907        atoms = self.__class__(cell=self.cell, pbc=self.pbc, info=self.info,
908                               celldisp=self._celldisp.copy())
909
910        atoms.arrays = {}
911        for name, a in self.arrays.items():
912            atoms.arrays[name] = a.copy()
913        atoms.constraints = copy.deepcopy(self.constraints)
914        return atoms
915
916    def todict(self):
917        """For basic JSON (non-database) support."""
918        d = dict(self.arrays)
919        d['cell'] = np.asarray(self.cell)
920        d['pbc'] = self.pbc
921        if self._celldisp.any():
922            d['celldisp'] = self._celldisp
923        if self.constraints:
924            d['constraints'] = self.constraints
925        if self.info:
926            d['info'] = self.info
927        # Calculator...  trouble.
928        return d
929
930    @classmethod
931    def fromdict(cls, dct):
932        """Rebuild atoms object from dictionary representation (todict)."""
933        dct = dct.copy()
934        kw = {}
935        for name in ['numbers', 'positions', 'cell', 'pbc']:
936            kw[name] = dct.pop(name)
937
938        constraints = dct.pop('constraints', None)
939        if constraints:
940            from ase.constraints import dict2constraint
941            constraints = [dict2constraint(d) for d in constraints]
942
943        info = dct.pop('info', None)
944
945        atoms = cls(constraint=constraints,
946                    celldisp=dct.pop('celldisp', None),
947                    info=info, **kw)
948        natoms = len(atoms)
949
950        # Some arrays are named differently from the atoms __init__ keywords.
951        # Also, there may be custom arrays.  Hence we set them directly:
952        for name, arr in dct.items():
953            assert len(arr) == natoms, name
954            assert isinstance(arr, np.ndarray)
955            atoms.arrays[name] = arr
956        return atoms
957
958    def __len__(self):
959        return len(self.arrays['positions'])
960
961    def get_number_of_atoms(self):
962        """Deprecated, please do not use.
963
964        You probably want len(atoms).  Or if your atoms are distributed,
965        use (and see) get_global_number_of_atoms()."""
966        import warnings
967        warnings.warn('Use get_global_number_of_atoms() instead',
968                      np.VisibleDeprecationWarning)
969        return len(self)
970
971    def get_global_number_of_atoms(self):
972        """Returns the global number of atoms in a distributed-atoms parallel
973        simulation.
974
975        DO NOT USE UNLESS YOU KNOW WHAT YOU ARE DOING!
976
977        Equivalent to len(atoms) in the standard ASE Atoms class.  You should
978        normally use len(atoms) instead.  This function's only purpose is to
979        make compatibility between ASE and Asap easier to maintain by having a
980        few places in ASE use this function instead.  It is typically only
981        when counting the global number of degrees of freedom or in similar
982        situations.
983        """
984        return len(self)
985
986    def __repr__(self):
987        tokens = []
988
989        N = len(self)
990        if N <= 60:
991            symbols = self.get_chemical_formula('reduce')
992        else:
993            symbols = self.get_chemical_formula('hill')
994        tokens.append("symbols='{0}'".format(symbols))
995
996        if self.pbc.any() and not self.pbc.all():
997            tokens.append('pbc={0}'.format(self.pbc.tolist()))
998        else:
999            tokens.append('pbc={0}'.format(self.pbc[0]))
1000
1001        cell = self.cell
1002        if cell:
1003            if cell.orthorhombic:
1004                cell = cell.lengths().tolist()
1005            else:
1006                cell = cell.tolist()
1007            tokens.append('cell={0}'.format(cell))
1008
1009        for name in sorted(self.arrays):
1010            if name in ['numbers', 'positions']:
1011                continue
1012            tokens.append('{0}=...'.format(name))
1013
1014        if self.constraints:
1015            if len(self.constraints) == 1:
1016                constraint = self.constraints[0]
1017            else:
1018                constraint = self.constraints
1019            tokens.append('constraint={0}'.format(repr(constraint)))
1020
1021        if self._calc is not None:
1022            tokens.append('calculator={0}(...)'
1023                          .format(self._calc.__class__.__name__))
1024
1025        return '{0}({1})'.format(self.__class__.__name__, ', '.join(tokens))
1026
1027    def __add__(self, other):
1028        atoms = self.copy()
1029        atoms += other
1030        return atoms
1031
1032    def extend(self, other):
1033        """Extend atoms object by appending atoms from *other*."""
1034        if isinstance(other, Atom):
1035            other = self.__class__([other])
1036
1037        n1 = len(self)
1038        n2 = len(other)
1039
1040        for name, a1 in self.arrays.items():
1041            a = np.zeros((n1 + n2,) + a1.shape[1:], a1.dtype)
1042            a[:n1] = a1
1043            if name == 'masses':
1044                a2 = other.get_masses()
1045            else:
1046                a2 = other.arrays.get(name)
1047            if a2 is not None:
1048                a[n1:] = a2
1049            self.arrays[name] = a
1050
1051        for name, a2 in other.arrays.items():
1052            if name in self.arrays:
1053                continue
1054            a = np.empty((n1 + n2,) + a2.shape[1:], a2.dtype)
1055            a[n1:] = a2
1056            if name == 'masses':
1057                a[:n1] = self.get_masses()[:n1]
1058            else:
1059                a[:n1] = 0
1060
1061            self.set_array(name, a)
1062
1063    def __iadd__(self, other):
1064        self.extend(other)
1065        return self
1066
1067    def append(self, atom):
1068        """Append atom to end."""
1069        self.extend(self.__class__([atom]))
1070
1071    def __iter__(self):
1072        for i in range(len(self)):
1073            yield self[i]
1074
1075    def __getitem__(self, i):
1076        """Return a subset of the atoms.
1077
1078        i -- scalar integer, list of integers, or slice object
1079        describing which atoms to return.
1080
1081        If i is a scalar, return an Atom object. If i is a list or a
1082        slice, return an Atoms object with the same cell, pbc, and
1083        other associated info as the original Atoms object. The
1084        indices of the constraints will be shuffled so that they match
1085        the indexing in the subset returned.
1086
1087        """
1088
1089        if isinstance(i, numbers.Integral):
1090            natoms = len(self)
1091            if i < -natoms or i >= natoms:
1092                raise IndexError('Index out of range.')
1093
1094            return Atom(atoms=self, index=i)
1095        elif not isinstance(i, slice):
1096            i = np.array(i)
1097            # if i is a mask
1098            if i.dtype == bool:
1099                if len(i) != len(self):
1100                    raise IndexError('Length of mask {} must equal '
1101                                     'number of atoms {}'
1102                                     .format(len(i), len(self)))
1103                i = np.arange(len(self))[i]
1104
1105        import copy
1106
1107        conadd = []
1108        # Constraints need to be deepcopied, but only the relevant ones.
1109        for con in copy.deepcopy(self.constraints):
1110            try:
1111                con.index_shuffle(self, i)
1112            except (IndexError, NotImplementedError):
1113                pass
1114            else:
1115                conadd.append(con)
1116
1117        atoms = self.__class__(cell=self.cell, pbc=self.pbc, info=self.info,
1118                               # should be communicated to the slice as well
1119                               celldisp=self._celldisp)
1120        # TODO: Do we need to shuffle indices in adsorbate_info too?
1121
1122        atoms.arrays = {}
1123        for name, a in self.arrays.items():
1124            atoms.arrays[name] = a[i].copy()
1125
1126        atoms.constraints = conadd
1127        return atoms
1128
1129    def __delitem__(self, i):
1130        from ase.constraints import FixAtoms
1131        for c in self._constraints:
1132            if not isinstance(c, FixAtoms):
1133                raise RuntimeError('Remove constraint using set_constraint() '
1134                                   'before deleting atoms.')
1135
1136        if isinstance(i, list) and len(i) > 0:
1137            # Make sure a list of booleans will work correctly and not be
1138            # interpreted at 0 and 1 indices.
1139            i = np.array(i)
1140
1141        if len(self._constraints) > 0:
1142            n = len(self)
1143            i = np.arange(n)[i]
1144            if isinstance(i, int):
1145                i = [i]
1146            constraints = []
1147            for c in self._constraints:
1148                c = c.delete_atoms(i, n)
1149                if c is not None:
1150                    constraints.append(c)
1151            self.constraints = constraints
1152
1153        mask = np.ones(len(self), bool)
1154        mask[i] = False
1155        for name, a in self.arrays.items():
1156            self.arrays[name] = a[mask]
1157
1158    def pop(self, i=-1):
1159        """Remove and return atom at index *i* (default last)."""
1160        atom = self[i]
1161        atom.cut_reference_to_atoms()
1162        del self[i]
1163        return atom
1164
1165    def __imul__(self, m):
1166        """In-place repeat of atoms."""
1167        if isinstance(m, int):
1168            m = (m, m, m)
1169
1170        for x, vec in zip(m, self.cell):
1171            if x != 1 and not vec.any():
1172                raise ValueError('Cannot repeat along undefined lattice '
1173                                 'vector')
1174
1175        M = np.product(m)
1176        n = len(self)
1177
1178        for name, a in self.arrays.items():
1179            self.arrays[name] = np.tile(a, (M,) + (1,) * (len(a.shape) - 1))
1180
1181        positions = self.arrays['positions']
1182        i0 = 0
1183        for m0 in range(m[0]):
1184            for m1 in range(m[1]):
1185                for m2 in range(m[2]):
1186                    i1 = i0 + n
1187                    positions[i0:i1] += np.dot((m0, m1, m2), self.cell)
1188                    i0 = i1
1189
1190        if self.constraints is not None:
1191            self.constraints = [c.repeat(m, n) for c in self.constraints]
1192
1193        self.cell = np.array([m[c] * self.cell[c] for c in range(3)])
1194
1195        return self
1196
1197    def repeat(self, rep):
1198        """Create new repeated atoms object.
1199
1200        The *rep* argument should be a sequence of three positive
1201        integers like *(2,3,1)* or a single integer (*r*) equivalent
1202        to *(r,r,r)*."""
1203
1204        atoms = self.copy()
1205        atoms *= rep
1206        return atoms
1207
1208    def __mul__(self, rep):
1209        return self.repeat(rep)
1210
1211    def translate(self, displacement):
1212        """Translate atomic positions.
1213
1214        The displacement argument can be a float an xyz vector or an
1215        nx3 array (where n is the number of atoms)."""
1216
1217        self.arrays['positions'] += np.array(displacement)
1218
1219    def center(self, vacuum=None, axis=(0, 1, 2), about=None):
1220        """Center atoms in unit cell.
1221
1222        Centers the atoms in the unit cell, so there is the same
1223        amount of vacuum on all sides.
1224
1225        vacuum: float (default: None)
1226            If specified adjust the amount of vacuum when centering.
1227            If vacuum=10.0 there will thus be 10 Angstrom of vacuum
1228            on each side.
1229        axis: int or sequence of ints
1230            Axis or axes to act on.  Default: Act on all axes.
1231        about: float or array (default: None)
1232            If specified, center the atoms about <about>.
1233            I.e., about=(0., 0., 0.) (or just "about=0.", interpreted
1234            identically), to center about the origin.
1235        """
1236
1237        # Find the orientations of the faces of the unit cell
1238        cell = self.cell.complete()
1239        dirs = np.zeros_like(cell)
1240
1241        lengths = cell.lengths()
1242        for i in range(3):
1243            dirs[i] = np.cross(cell[i - 1], cell[i - 2])
1244            dirs[i] /= np.linalg.norm(dirs[i])
1245            if dirs[i] @ cell[i] < 0.0:
1246                dirs[i] *= -1
1247
1248        if isinstance(axis, int):
1249            axes = (axis,)
1250        else:
1251            axes = axis
1252
1253        # Now, decide how much each basis vector should be made longer
1254        pos = self.positions
1255        longer = np.zeros(3)
1256        shift = np.zeros(3)
1257        for i in axes:
1258            if len(pos):
1259                scalarprod = pos @ dirs[i]
1260                p0 = scalarprod.min()
1261                p1 = scalarprod.max()
1262            else:
1263                p0 = 0
1264                p1 = 0
1265            height = cell[i] @ dirs[i]
1266            if vacuum is not None:
1267                lng = (p1 - p0 + 2 * vacuum) - height
1268            else:
1269                lng = 0.0  # Do not change unit cell size!
1270            top = lng + height - p1
1271            shf = 0.5 * (top - p0)
1272            cosphi = cell[i] @ dirs[i] / lengths[i]
1273            longer[i] = lng / cosphi
1274            shift[i] = shf / cosphi
1275
1276        # Now, do it!
1277        translation = np.zeros(3)
1278        for i in axes:
1279            nowlen = lengths[i]
1280            if vacuum is not None:
1281                self.cell[i] = cell[i] * (1 + longer[i] / nowlen)
1282            translation += shift[i] * cell[i] / nowlen
1283
1284            # We calculated translations using the completed cell,
1285            # so directions without cell vectors will have been centered
1286            # along a "fake" vector of length 1.
1287            # Therefore, we adjust by -0.5:
1288            if not any(self.cell[i]):
1289                translation[i] -= 0.5
1290
1291        # Optionally, translate to center about a point in space.
1292        if about is not None:
1293            for vector in self.cell:
1294                translation -= vector / 2.0
1295            translation += about
1296
1297        self.positions += translation
1298
1299    def get_center_of_mass(self, scaled=False):
1300        """Get the center of mass.
1301
1302        If scaled=True the center of mass in scaled coordinates
1303        is returned."""
1304        masses = self.get_masses()
1305        com = masses @ self.positions / masses.sum()
1306        if scaled:
1307            return self.cell.scaled_positions(com)
1308        else:
1309            return com
1310
1311    def set_center_of_mass(self, com, scaled=False):
1312        """Set the center of mass.
1313
1314        If scaled=True the center of mass is expected in scaled coordinates.
1315        Constraints are considered for scaled=False.
1316        """
1317        old_com = self.get_center_of_mass(scaled=scaled)
1318        difference = old_com - com
1319        if scaled:
1320            self.set_scaled_positions(self.get_scaled_positions() + difference)
1321        else:
1322            self.set_positions(self.get_positions() + difference)
1323
1324    def get_moments_of_inertia(self, vectors=False):
1325        """Get the moments of inertia along the principal axes.
1326
1327        The three principal moments of inertia are computed from the
1328        eigenvalues of the symmetric inertial tensor. Periodic boundary
1329        conditions are ignored. Units of the moments of inertia are
1330        amu*angstrom**2.
1331        """
1332        com = self.get_center_of_mass()
1333        positions = self.get_positions()
1334        positions -= com  # translate center of mass to origin
1335        masses = self.get_masses()
1336
1337        # Initialize elements of the inertial tensor
1338        I11 = I22 = I33 = I12 = I13 = I23 = 0.0
1339        for i in range(len(self)):
1340            x, y, z = positions[i]
1341            m = masses[i]
1342
1343            I11 += m * (y ** 2 + z ** 2)
1344            I22 += m * (x ** 2 + z ** 2)
1345            I33 += m * (x ** 2 + y ** 2)
1346            I12 += -m * x * y
1347            I13 += -m * x * z
1348            I23 += -m * y * z
1349
1350        I = np.array([[I11, I12, I13],
1351                      [I12, I22, I23],
1352                      [I13, I23, I33]])
1353
1354        evals, evecs = np.linalg.eigh(I)
1355        if vectors:
1356            return evals, evecs.transpose()
1357        else:
1358            return evals
1359
1360    def get_angular_momentum(self):
1361        """Get total angular momentum with respect to the center of mass."""
1362        com = self.get_center_of_mass()
1363        positions = self.get_positions()
1364        positions -= com  # translate center of mass to origin
1365        return np.cross(positions, self.get_momenta()).sum(0)
1366
1367    def rotate(self, a, v, center=(0, 0, 0), rotate_cell=False):
1368        """Rotate atoms based on a vector and an angle, or two vectors.
1369
1370        Parameters:
1371
1372        a = None:
1373            Angle that the atoms is rotated around the vector 'v'. 'a'
1374            can also be a vector and then 'a' is rotated
1375            into 'v'.
1376
1377        v:
1378            Vector to rotate the atoms around. Vectors can be given as
1379            strings: 'x', '-x', 'y', ... .
1380
1381        center = (0, 0, 0):
1382            The center is kept fixed under the rotation. Use 'COM' to fix
1383            the center of mass, 'COP' to fix the center of positions or
1384            'COU' to fix the center of cell.
1385
1386        rotate_cell = False:
1387            If true the cell is also rotated.
1388
1389        Examples:
1390
1391        Rotate 90 degrees around the z-axis, so that the x-axis is
1392        rotated into the y-axis:
1393
1394        >>> atoms = Atoms()
1395        >>> atoms.rotate(90, 'z')
1396        >>> atoms.rotate(90, (0, 0, 1))
1397        >>> atoms.rotate(-90, '-z')
1398        >>> atoms.rotate('x', 'y')
1399        >>> atoms.rotate((1, 0, 0), (0, 1, 0))
1400        """
1401
1402        if not isinstance(a, numbers.Real):
1403            a, v = v, a
1404
1405        norm = np.linalg.norm
1406        v = string2vector(v)
1407
1408        normv = norm(v)
1409
1410        if normv == 0.0:
1411            raise ZeroDivisionError('Cannot rotate: norm(v) == 0')
1412
1413        if isinstance(a, numbers.Real):
1414            a *= pi / 180
1415            v /= normv
1416            c = cos(a)
1417            s = sin(a)
1418        else:
1419            v2 = string2vector(a)
1420            v /= normv
1421            normv2 = np.linalg.norm(v2)
1422            if normv2 == 0:
1423                raise ZeroDivisionError('Cannot rotate: norm(a) == 0')
1424            v2 /= norm(v2)
1425            c = np.dot(v, v2)
1426            v = np.cross(v, v2)
1427            s = norm(v)
1428            # In case *v* and *a* are parallel, np.cross(v, v2) vanish
1429            # and can't be used as a rotation axis. However, in this
1430            # case any rotation axis perpendicular to v2 will do.
1431            eps = 1e-7
1432            if s < eps:
1433                v = np.cross((0, 0, 1), v2)
1434                if norm(v) < eps:
1435                    v = np.cross((1, 0, 0), v2)
1436                assert norm(v) >= eps
1437            elif s > 0:
1438                v /= s
1439
1440        center = self._centering_as_array(center)
1441
1442        p = self.arrays['positions'] - center
1443        self.arrays['positions'][:] = (c * p -
1444                                       np.cross(p, s * v) +
1445                                       np.outer(np.dot(p, v), (1.0 - c) * v) +
1446                                       center)
1447        if rotate_cell:
1448            rotcell = self.get_cell()
1449            rotcell[:] = (c * rotcell -
1450                          np.cross(rotcell, s * v) +
1451                          np.outer(np.dot(rotcell, v), (1.0 - c) * v))
1452            self.set_cell(rotcell)
1453
1454    def _centering_as_array(self, center):
1455        if isinstance(center, str):
1456            if center.lower() == 'com':
1457                center = self.get_center_of_mass()
1458            elif center.lower() == 'cop':
1459                center = self.get_positions().mean(axis=0)
1460            elif center.lower() == 'cou':
1461                center = self.get_cell().sum(axis=0) / 2
1462            else:
1463                raise ValueError('Cannot interpret center')
1464        else:
1465            center = np.array(center, float)
1466        return center
1467
1468    def euler_rotate(self, phi=0.0, theta=0.0, psi=0.0, center=(0, 0, 0)):
1469        """Rotate atoms via Euler angles (in degrees).
1470
1471        See e.g http://mathworld.wolfram.com/EulerAngles.html for explanation.
1472
1473        Parameters:
1474
1475        center :
1476            The point to rotate about. A sequence of length 3 with the
1477            coordinates, or 'COM' to select the center of mass, 'COP' to
1478            select center of positions or 'COU' to select center of cell.
1479        phi :
1480            The 1st rotation angle around the z axis.
1481        theta :
1482            Rotation around the x axis.
1483        psi :
1484            2nd rotation around the z axis.
1485
1486        """
1487        center = self._centering_as_array(center)
1488
1489        phi *= pi / 180
1490        theta *= pi / 180
1491        psi *= pi / 180
1492
1493        # First move the molecule to the origin In contrast to MATLAB,
1494        # numpy broadcasts the smaller array to the larger row-wise,
1495        # so there is no need to play with the Kronecker product.
1496        rcoords = self.positions - center
1497        # First Euler rotation about z in matrix form
1498        D = np.array(((cos(phi), sin(phi), 0.),
1499                      (-sin(phi), cos(phi), 0.),
1500                      (0., 0., 1.)))
1501        # Second Euler rotation about x:
1502        C = np.array(((1., 0., 0.),
1503                      (0., cos(theta), sin(theta)),
1504                      (0., -sin(theta), cos(theta))))
1505        # Third Euler rotation, 2nd rotation about z:
1506        B = np.array(((cos(psi), sin(psi), 0.),
1507                      (-sin(psi), cos(psi), 0.),
1508                      (0., 0., 1.)))
1509        # Total Euler rotation
1510        A = np.dot(B, np.dot(C, D))
1511        # Do the rotation
1512        rcoords = np.dot(A, np.transpose(rcoords))
1513        # Move back to the rotation point
1514        self.positions = np.transpose(rcoords) + center
1515
1516    def get_dihedral(self, a0, a1, a2, a3, mic=False):
1517        """Calculate dihedral angle.
1518
1519        Calculate dihedral angle (in degrees) between the vectors a0->a1
1520        and a2->a3.
1521
1522        Use mic=True to use the Minimum Image Convention and calculate the
1523        angle across periodic boundaries.
1524        """
1525        return self.get_dihedrals([[a0, a1, a2, a3]], mic=mic)[0]
1526
1527    def get_dihedrals(self, indices, mic=False):
1528        """Calculate dihedral angles.
1529
1530        Calculate dihedral angles (in degrees) between the list of vectors
1531        a0->a1 and a2->a3, where a0, a1, a2 and a3 are in each row of indices.
1532
1533        Use mic=True to use the Minimum Image Convention and calculate the
1534        angles across periodic boundaries.
1535        """
1536        indices = np.array(indices)
1537        assert indices.shape[1] == 4
1538
1539        a0s = self.positions[indices[:, 0]]
1540        a1s = self.positions[indices[:, 1]]
1541        a2s = self.positions[indices[:, 2]]
1542        a3s = self.positions[indices[:, 3]]
1543
1544        # vectors 0->1, 1->2, 2->3
1545        v0 = a1s - a0s
1546        v1 = a2s - a1s
1547        v2 = a3s - a2s
1548
1549        cell = None
1550        pbc = None
1551
1552        if mic:
1553            cell = self.cell
1554            pbc = self.pbc
1555
1556        return get_dihedrals(v0, v1, v2, cell=cell, pbc=pbc)
1557
1558    def _masked_rotate(self, center, axis, diff, mask):
1559        # do rotation of subgroup by copying it to temporary atoms object
1560        # and then rotating that
1561        #
1562        # recursive object definition might not be the most elegant thing,
1563        # more generally useful might be a rotation function with a mask?
1564        group = self.__class__()
1565        for i in range(len(self)):
1566            if mask[i]:
1567                group += self[i]
1568        group.translate(-center)
1569        group.rotate(diff * 180 / pi, axis)
1570        group.translate(center)
1571        # set positions in original atoms object
1572        j = 0
1573        for i in range(len(self)):
1574            if mask[i]:
1575                self.positions[i] = group[j].position
1576                j += 1
1577
1578    def set_dihedral(self, a1, a2, a3, a4, angle,
1579                     mask=None, indices=None):
1580        """Set the dihedral angle (degrees) between vectors a1->a2 and
1581        a3->a4 by changing the atom indexed by a4.
1582
1583        If mask is not None, all the atoms described in mask
1584        (read: the entire subgroup) are moved. Alternatively to the mask,
1585        the indices of the atoms to be rotated can be supplied. If both
1586        *mask* and *indices* are given, *indices* overwrites *mask*.
1587
1588        **Important**: If *mask* or *indices* is given and does not contain
1589        *a4*, *a4* will NOT be moved. In most cases you therefore want
1590        to include *a4* in *mask*/*indices*.
1591
1592        Example: the following defines a very crude
1593        ethane-like molecule and twists one half of it by 30 degrees.
1594
1595        >>> atoms = Atoms('HHCCHH', [[-1, 1, 0], [-1, -1, 0], [0, 0, 0],
1596        ...                          [1, 0, 0], [2, 1, 0], [2, -1, 0]])
1597        >>> atoms.set_dihedral(1, 2, 3, 4, 210, mask=[0, 0, 0, 1, 1, 1])
1598        """
1599
1600        angle *= pi / 180
1601
1602        # if not provided, set mask to the last atom in the
1603        # dihedral description
1604        if mask is None and indices is None:
1605            mask = np.zeros(len(self))
1606            mask[a4] = 1
1607        elif indices is not None:
1608            mask = [index in indices for index in range(len(self))]
1609
1610        # compute necessary in dihedral change, from current value
1611        current = self.get_dihedral(a1, a2, a3, a4) * pi / 180
1612        diff = angle - current
1613        axis = self.positions[a3] - self.positions[a2]
1614        center = self.positions[a3]
1615        self._masked_rotate(center, axis, diff, mask)
1616
1617    def rotate_dihedral(self, a1, a2, a3, a4,
1618                        angle=None, mask=None, indices=None):
1619        """Rotate dihedral angle.
1620
1621        Same usage as in :meth:`ase.Atoms.set_dihedral`: Rotate a group by a
1622        predefined dihedral angle, starting from its current configuration.
1623        """
1624        start = self.get_dihedral(a1, a2, a3, a4)
1625        self.set_dihedral(a1, a2, a3, a4, angle + start, mask, indices)
1626
1627    def get_angle(self, a1, a2, a3, mic=False):
1628        """Get angle formed by three atoms.
1629
1630        Calculate angle in degrees between the vectors a2->a1 and
1631        a2->a3.
1632
1633        Use mic=True to use the Minimum Image Convention and calculate the
1634        angle across periodic boundaries.
1635        """
1636        return self.get_angles([[a1, a2, a3]], mic=mic)[0]
1637
1638    def get_angles(self, indices, mic=False):
1639        """Get angle formed by three atoms for multiple groupings.
1640
1641        Calculate angle in degrees between vectors between atoms a2->a1
1642        and a2->a3, where a1, a2, and a3 are in each row of indices.
1643
1644        Use mic=True to use the Minimum Image Convention and calculate
1645        the angle across periodic boundaries.
1646        """
1647        indices = np.array(indices)
1648        assert indices.shape[1] == 3
1649
1650        a1s = self.positions[indices[:, 0]]
1651        a2s = self.positions[indices[:, 1]]
1652        a3s = self.positions[indices[:, 2]]
1653
1654        v12 = a1s - a2s
1655        v32 = a3s - a2s
1656
1657        cell = None
1658        pbc = None
1659
1660        if mic:
1661            cell = self.cell
1662            pbc = self.pbc
1663
1664        return get_angles(v12, v32, cell=cell, pbc=pbc)
1665
1666    def set_angle(self, a1, a2=None, a3=None, angle=None, mask=None,
1667                  indices=None, add=False):
1668        """Set angle (in degrees) formed by three atoms.
1669
1670        Sets the angle between vectors *a2*->*a1* and *a2*->*a3*.
1671
1672        If *add* is `True`, the angle will be changed by the value given.
1673
1674        Same usage as in :meth:`ase.Atoms.set_dihedral`.
1675        If *mask* and *indices*
1676        are given, *indices* overwrites *mask*. If *mask* and *indices*
1677        are not set, only *a3* is moved."""
1678
1679        if any(a is None for a in [a2, a3, angle]):
1680            raise ValueError('a2, a3, and angle must not be None')
1681
1682        # If not provided, set mask to the last atom in the angle description
1683        if mask is None and indices is None:
1684            mask = np.zeros(len(self))
1685            mask[a3] = 1
1686        elif indices is not None:
1687            mask = [index in indices for index in range(len(self))]
1688
1689        if add:
1690            diff = angle
1691        else:
1692            # Compute necessary in angle change, from current value
1693            diff = angle - self.get_angle(a1, a2, a3)
1694
1695        diff *= pi / 180
1696        # Do rotation of subgroup by copying it to temporary atoms object and
1697        # then rotating that
1698        v10 = self.positions[a1] - self.positions[a2]
1699        v12 = self.positions[a3] - self.positions[a2]
1700        v10 /= np.linalg.norm(v10)
1701        v12 /= np.linalg.norm(v12)
1702        axis = np.cross(v10, v12)
1703        center = self.positions[a2]
1704        self._masked_rotate(center, axis, diff, mask)
1705
1706    def rattle(self, stdev=0.001, seed=None, rng=None):
1707        """Randomly displace atoms.
1708
1709        This method adds random displacements to the atomic positions,
1710        taking a possible constraint into account.  The random numbers are
1711        drawn from a normal distribution of standard deviation stdev.
1712
1713        For a parallel calculation, it is important to use the same
1714        seed on all processors!  """
1715
1716        if seed is not None and rng is not None:
1717            raise ValueError('Please do not provide both seed and rng.')
1718
1719        if rng is None:
1720            if seed is None:
1721                seed = 42
1722            rng = np.random.RandomState(seed)
1723        positions = self.arrays['positions']
1724        self.set_positions(positions +
1725                           rng.normal(scale=stdev, size=positions.shape))
1726
1727    def get_distance(self, a0, a1, mic=False, vector=False):
1728        """Return distance between two atoms.
1729
1730        Use mic=True to use the Minimum Image Convention.
1731        vector=True gives the distance vector (from a0 to a1).
1732        """
1733        return self.get_distances(a0, [a1], mic=mic, vector=vector)[0]
1734
1735    def get_distances(self, a, indices, mic=False, vector=False):
1736        """Return distances of atom No.i with a list of atoms.
1737
1738        Use mic=True to use the Minimum Image Convention.
1739        vector=True gives the distance vector (from a to self[indices]).
1740        """
1741        R = self.arrays['positions']
1742        p1 = [R[a]]
1743        p2 = R[indices]
1744
1745        cell = None
1746        pbc = None
1747
1748        if mic:
1749            cell = self.cell
1750            pbc = self.pbc
1751
1752        D, D_len = get_distances(p1, p2, cell=cell, pbc=pbc)
1753
1754        if vector:
1755            D.shape = (-1, 3)
1756            return D
1757        else:
1758            D_len.shape = (-1,)
1759            return D_len
1760
1761    def get_all_distances(self, mic=False, vector=False):
1762        """Return distances of all of the atoms with all of the atoms.
1763
1764        Use mic=True to use the Minimum Image Convention.
1765        """
1766        R = self.arrays['positions']
1767
1768        cell = None
1769        pbc = None
1770
1771        if mic:
1772            cell = self.cell
1773            pbc = self.pbc
1774
1775        D, D_len = get_distances(R, cell=cell, pbc=pbc)
1776
1777        if vector:
1778            return D
1779        else:
1780            return D_len
1781
1782    def set_distance(self, a0, a1, distance, fix=0.5, mic=False,
1783                     mask=None, indices=None, add=False, factor=False):
1784        """Set the distance between two atoms.
1785
1786        Set the distance between atoms *a0* and *a1* to *distance*.
1787        By default, the center of the two atoms will be fixed.  Use
1788        *fix=0* to fix the first atom, *fix=1* to fix the second
1789        atom and *fix=0.5* (default) to fix the center of the bond.
1790
1791        If *mask* or *indices* are set (*mask* overwrites *indices*),
1792        only the atoms defined there are moved
1793        (see :meth:`ase.Atoms.set_dihedral`).
1794
1795        When *add* is true, the distance is changed by the value given.
1796        In combination
1797        with *factor* True, the value given is a factor scaling the distance.
1798
1799        It is assumed that the atoms in *mask*/*indices* move together
1800        with *a1*. If *fix=1*, only *a0* will therefore be moved."""
1801
1802        if a0 % len(self) == a1 % len(self):
1803            raise ValueError('a0 and a1 must not be the same')
1804
1805        if add:
1806            oldDist = self.get_distance(a0, a1, mic=mic)
1807            if factor:
1808                newDist = oldDist * distance
1809            else:
1810                newDist = oldDist + distance
1811            self.set_distance(a0, a1, newDist, fix=fix, mic=mic,
1812                              mask=mask, indices=indices, add=False,
1813                              factor=False)
1814            return
1815
1816        R = self.arrays['positions']
1817        D = np.array([R[a1] - R[a0]])
1818
1819        if mic:
1820            D, D_len = find_mic(D, self.cell, self.pbc)
1821        else:
1822            D_len = np.array([np.sqrt((D**2).sum())])
1823        x = 1.0 - distance / D_len[0]
1824
1825        if mask is None and indices is None:
1826            indices = [a0, a1]
1827        elif mask:
1828            indices = [i for i in range(len(self)) if mask[i]]
1829
1830        for i in indices:
1831            if i == a0:
1832                R[a0] += (x * fix) * D[0]
1833            else:
1834                R[i] -= (x * (1.0 - fix)) * D[0]
1835
1836    def get_scaled_positions(self, wrap=True):
1837        """Get positions relative to unit cell.
1838
1839        If wrap is True, atoms outside the unit cell will be wrapped into
1840        the cell in those directions with periodic boundary conditions
1841        so that the scaled coordinates are between zero and one.
1842
1843        If any cell vectors are zero, the corresponding coordinates
1844        are evaluated as if the cell were completed using
1845        ``cell.complete()``.  This means coordinates will be Cartesian
1846        as long as the non-zero cell vectors span a Cartesian axis or
1847        plane."""
1848
1849        fractional = self.cell.scaled_positions(self.positions)
1850
1851        if wrap:
1852            for i, periodic in enumerate(self.pbc):
1853                if periodic:
1854                    # Yes, we need to do it twice.
1855                    # See the scaled_positions.py test.
1856                    fractional[:, i] %= 1.0
1857                    fractional[:, i] %= 1.0
1858
1859        return fractional
1860
1861    def set_scaled_positions(self, scaled):
1862        """Set positions relative to unit cell."""
1863        self.positions[:] = self.cell.cartesian_positions(scaled)
1864
1865    def wrap(self, **wrap_kw):
1866        """Wrap positions to unit cell.
1867
1868        Parameters:
1869
1870        wrap_kw: (keyword=value) pairs
1871            optional keywords `pbc`, `center`, `pretty_translation`, `eps`,
1872            see :func:`ase.geometry.wrap_positions`
1873        """
1874
1875        if 'pbc' not in wrap_kw:
1876            wrap_kw['pbc'] = self.pbc
1877
1878        self.positions[:] = self.get_positions(wrap=True, **wrap_kw)
1879
1880    def get_temperature(self):
1881        """Get the temperature in Kelvin."""
1882        dof = len(self) * 3
1883        for constraint in self._constraints:
1884            dof -= constraint.get_removed_dof(self)
1885        ekin = self.get_kinetic_energy()
1886        return 2 * ekin / (dof * units.kB)
1887
1888    def __eq__(self, other):
1889        """Check for identity of two atoms objects.
1890
1891        Identity means: same positions, atomic numbers, unit cell and
1892        periodic boundary conditions."""
1893        if not isinstance(other, Atoms):
1894            return False
1895        a = self.arrays
1896        b = other.arrays
1897        return (len(self) == len(other) and
1898                (a['positions'] == b['positions']).all() and
1899                (a['numbers'] == b['numbers']).all() and
1900                (self.cell == other.cell).all() and
1901                (self.pbc == other.pbc).all())
1902
1903    def __ne__(self, other):
1904        """Check if two atoms objects are not equal.
1905
1906        Any differences in positions, atomic numbers, unit cell or
1907        periodic boundary condtions make atoms objects not equal.
1908        """
1909        eq = self.__eq__(other)
1910        if eq is NotImplemented:
1911            return eq
1912        else:
1913            return not eq
1914
1915    # @deprecated('Please use atoms.cell.volume')
1916    # We kind of want to deprecate this, but the ValueError behaviour
1917    # might be desirable.  Should we do this?
1918    def get_volume(self):
1919        """Get volume of unit cell."""
1920        if self.cell.rank != 3:
1921            raise ValueError(
1922                'You have {0} lattice vectors: volume not defined'
1923                .format(self.cell.rank))
1924        return self.cell.volume
1925
1926    def _get_positions(self):
1927        """Return reference to positions-array for in-place manipulations."""
1928        return self.arrays['positions']
1929
1930    def _set_positions(self, pos):
1931        """Set positions directly, bypassing constraints."""
1932        self.arrays['positions'][:] = pos
1933
1934    positions = property(_get_positions, _set_positions,
1935                         doc='Attribute for direct ' +
1936                         'manipulation of the positions.')
1937
1938    def _get_atomic_numbers(self):
1939        """Return reference to atomic numbers for in-place
1940        manipulations."""
1941        return self.arrays['numbers']
1942
1943    numbers = property(_get_atomic_numbers, set_atomic_numbers,
1944                       doc='Attribute for direct ' +
1945                       'manipulation of the atomic numbers.')
1946
1947    @property
1948    def cell(self):
1949        """The :class:`ase.cell.Cell` for direct manipulation."""
1950        return self._cellobj
1951
1952    @cell.setter
1953    def cell(self, cell):
1954        cell = Cell.ascell(cell)
1955        self._cellobj[:] = cell
1956
1957    def write(self, filename, format=None, **kwargs):
1958        """Write atoms object to a file.
1959
1960        see ase.io.write for formats.
1961        kwargs are passed to ase.io.write.
1962        """
1963        from ase.io import write
1964        write(filename, self, format, **kwargs)
1965
1966    def iterimages(self):
1967        yield self
1968
1969    def edit(self):
1970        """Modify atoms interactively through ASE's GUI viewer.
1971
1972        Conflicts leading to undesirable behaviour might arise
1973        when matplotlib has been pre-imported with certain
1974        incompatible backends and while trying to use the
1975        plot feature inside the interactive GUI. To circumvent,
1976        please set matplotlib.use('gtk') before calling this
1977        method.
1978        """
1979        from ase.gui.images import Images
1980        from ase.gui.gui import GUI
1981        images = Images([self])
1982        gui = GUI(images)
1983        gui.run()
1984
1985
1986def string2vector(v):
1987    if isinstance(v, str):
1988        if v[0] == '-':
1989            return -string2vector(v[1:])
1990        w = np.zeros(3)
1991        w['xyz'.index(v)] = 1.0
1992        return w
1993    return np.array(v, float)
1994
1995
1996def default(data, dflt):
1997    """Helper function for setting default values."""
1998    if data is None:
1999        return None
2000    elif isinstance(data, (list, tuple)):
2001        newdata = []
2002        allnone = True
2003        for x in data:
2004            if x is None:
2005                newdata.append(dflt)
2006            else:
2007                newdata.append(x)
2008                allnone = False
2009        if allnone:
2010            return None
2011        return newdata
2012    else:
2013        return data
2014