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