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