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