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