Coverage for /builds/ase/ase/ase/geometry/analysis.py : 70.72%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# flake8: noqa
2"""Tools for analyzing instances of :class:`~ase.Atoms`
3"""
5from typing import List, Optional
7import numpy as np
9from ase.neighborlist import build_neighbor_list, get_distance_matrix, get_distance_indices
10from ase.geometry.rdf import get_rdf, get_containing_cell_length
11from ase import Atoms
14__all__ = ['Analysis']
17def get_max_containing_cell_length(images: List[Atoms]):
18 i2diff = np.zeros(3)
19 for image in images:
20 np.maximum(get_containing_cell_length(image), i2diff, out=i2diff)
21 return i2diff
24def get_max_volume_estimate(images: List[Atoms]) -> float:
25 return np.product(get_max_containing_cell_length(images))
28class Analysis:
29 """Analysis class
31 Parameters for initialization:
33 images: :class:`~ase.Atoms` object or list of such
34 Images to analyze.
35 nl: None, :class:`~ase.neighborlist.NeighborList` object or list of such
36 Neighborlist(s) for the given images. One or nImages, depending if bonding
37 pattern changes or is constant. Using one Neigborlist greatly improves speed.
38 kwargs: options, dict
39 Arguments for constructing :class:`~ase.neighborlist.NeighborList` object if :data:`nl` is None.
41 The choice of ``bothways=True`` for the :class:`~ase.neighborlist.NeighborList` object
42 will not influence the amount of bonds/angles/dihedrals you get, all are reported
43 in both directions. Use the *unique*-labeled properties to get lists without
44 duplicates.
45 """
47 def __init__(self, images, nl=None, **kwargs):
48 self.images = images
50 if isinstance(nl, list):
51 assert len(nl) == self.nImages
52 self._nl = nl
53 elif nl is not None:
54 self._nl = [ nl ]
55 else:
56 self._nl = [ build_neighbor_list(self.images[0], **kwargs) ]
58 self._cache = {}
60 def _get_slice(self, imageIdx):
61 """Return a slice from user input.
62 Using *imageIdx* (can be integer or slice) the analyzed frames can be specified.
63 If *imageIdx* is None, all frames will be analyzed.
64 """
65 #get slice from imageIdx
66 if isinstance(imageIdx, int):
67 sl = slice(imageIdx, imageIdx+1)
68 elif isinstance(imageIdx, slice):
69 sl = imageIdx
70 elif imageIdx is None:
71 sl = slice(0, None)
72 else:
73 raise ValueError("Unsupported type for imageIdx in ase.geometry.analysis.Analysis._get_slice")
74 return sl
76 @property
77 def images(self):
78 """Images.
80 Set during initialization but can also be set later.
81 """
82 return self._images
84 @images.setter
85 def images(self, images):
86 """Set images"""
87 if isinstance(images, list):
88 self._images = images
89 else:
90 self._images = [ images ]
93 @images.deleter
94 def images(self):
95 """Delete images"""
96 self._images = None
98 @property
99 def nImages(self):
100 """Number of Images in this instance.
102 Cannot be set, is determined automatically.
103 """
104 return len(self.images)
106 @property
107 def nl(self):
108 """Neighbor Lists in this instance.
110 Set during initialization.
112 **No setter or deleter, only getter**
113 """
114 return self._nl
116 def _get_all_x(self, distance):
117 """Helper function to get bonds, angles, dihedrals"""
118 maxIter = self.nImages
119 if len(self.nl) == 1:
120 maxIter = 1
122 xList = []
123 for i in range(maxIter):
124 xList.append(get_distance_indices(self.distance_matrix[i], distance))
126 return xList
128 @property
129 def all_bonds(self):
130 """All Bonds.
132 A list with indices of bonded atoms for each neighborlist in *self*.
133 Atom i is connected to all atoms inside result[i]. Duplicates from PBCs are
134 removed. See also :data:`unique_bonds`.
136 **No setter or deleter, only getter**
137 """
138 if not 'allBonds' in self._cache:
139 self._cache['allBonds'] = self._get_all_x(1)
141 return self._cache['allBonds']
143 @property
144 def all_angles(self):
145 """All angles
147 A list with indices of atoms in angles for each neighborlist in *self*.
148 Atom i forms an angle to the atoms inside the tuples in result[i]:
149 i -- result[i][x][0] -- result[i][x][1]
150 where x is in range(number of angles from i). See also :data:`unique_angles`.
152 **No setter or deleter, only getter**
153 """
154 if not 'allAngles' in self._cache:
155 self._cache['allAngles'] = []
156 distList = self._get_all_x(2)
158 for imI in range(len(distList)):
159 self._cache['allAngles'].append([])
160 #iterate over second neighbors of all atoms
161 for iAtom, secNeighs in enumerate(distList[imI]):
162 self._cache['allAngles'][-1].append([])
163 if len(secNeighs) == 0:
164 continue
165 firstNeighs = self.all_bonds[imI][iAtom]
166 #iterate over second neighbors of iAtom
167 for kAtom in secNeighs:
168 relevantFirstNeighs = [ idx for idx in firstNeighs if kAtom in self.all_bonds[imI][idx] ]
169 #iterate over all atoms that are connected to iAtom and kAtom
170 for jAtom in relevantFirstNeighs:
171 self._cache['allAngles'][-1][-1].append((jAtom, kAtom))
173 return self._cache['allAngles']
175 @property
176 def all_dihedrals(self):
177 """All dihedrals
179 Returns a list with indices of atoms in dihedrals for each neighborlist in this instance.
180 Atom i forms a dihedral to the atoms inside the tuples in result[i]:
181 i -- result[i][x][0] -- result[i][x][1] -- result[i][x][2]
182 where x is in range(number of dihedrals from i). See also :data:`unique_dihedrals`.
184 **No setter or deleter, only getter**
185 """
186 if not 'allDihedrals' in self._cache:
187 self._cache['allDihedrals'] = []
188 distList = self._get_all_x(3)
190 for imI in range(len(distList)):
191 self._cache['allDihedrals'].append([])
192 for iAtom, thirdNeighs in enumerate(distList[imI]):
193 self._cache['allDihedrals'][-1].append([])
194 if len(thirdNeighs) == 0:
195 continue
196 anglesI = self.all_angles[imI][iAtom]
197 #iterate over third neighbors of iAtom
198 for lAtom in thirdNeighs:
199 secondNeighs = [ angle[-1] for angle in anglesI ]
200 firstNeighs = [ angle[0] for angle in anglesI ]
201 relevantSecondNeighs = [ idx for idx in secondNeighs if lAtom in self.all_bonds[imI][idx] ]
202 relevantFirstNeighs = [ firstNeighs[secondNeighs.index(idx)] for idx in relevantSecondNeighs ]
203 #iterate over all atoms that are connected to iAtom and lAtom
204 for jAtom, kAtom in zip(relevantFirstNeighs, relevantSecondNeighs):
205 #remove dihedrals in circles
206 tupl = (jAtom, kAtom, lAtom)
207 if len(set((iAtom, ) + tupl)) != 4:
208 continue
209 #avoid duplicates
210 elif tupl in self._cache['allDihedrals'][-1][-1]:
211 continue
212 elif iAtom in tupl:
213 raise RuntimeError("Something is wrong in analysis.all_dihedrals!")
214 self._cache['allDihedrals'][-1][-1].append((jAtom, kAtom, lAtom))
216 return self._cache['allDihedrals']
218 @property
219 def adjacency_matrix(self):
220 """The adjacency/connectivity matrix.
222 If not already done, build a list of adjacency matrices for all :data:`nl`.
224 **No setter or deleter, only getter**
225 """
227 if not 'adjacencyMatrix' in self._cache:
228 self._cache['adjacencyMatrix'] = []
229 for i in range(len(self.nl)):
230 self._cache['adjacencyMatrix'].append(self.nl[i].get_connectivity_matrix())
232 return self._cache['adjacencyMatrix']
234 @property
235 def distance_matrix(self):
236 """The distance matrix.
238 If not already done, build a list of distance matrices for all :data:`nl`. See
239 :meth:`ase.neighborlist.get_distance_matrix`.
241 **No setter or deleter, only getter**
242 """
244 if not 'distanceMatrix' in self._cache:
245 self._cache['distanceMatrix'] = []
246 for i in range(len(self.nl)):
247 self._cache['distanceMatrix'].append(get_distance_matrix(self.adjacency_matrix[i]))
249 return self._cache['distanceMatrix']
252 @property
253 def unique_bonds(self):
254 """Get Unique Bonds.
256 :data:`all_bonds` i-j without j-i. This is the upper triangle of the
257 connectivity matrix (i,j), `i < j`
259 """
260 bonds = []
261 for imI in range(len(self.all_bonds)):
262 bonds.append([])
263 for iAtom, bonded in enumerate(self.all_bonds[imI]):
264 bonds[-1].append([ jAtom for jAtom in bonded if jAtom > iAtom ])
266 return bonds
269 def _filter_unique(self, l):
270 """Helper function to filter for unique lists in a list
271 that also contains the reversed items.
272 """
273 r = []
274 #iterate over images
275 for imI in range(len(l)):
276 r.append([])
277 #iterate over atoms
278 for i, tuples in enumerate(l[imI]):
279 #add the ones where i is smaller than the last element
280 r[-1].append([ x for x in tuples if i < x[-1] ])
281 return r
283 def clear_cache(self):
284 """Delete all cached information."""
285 self._cache = {}
287 @property
288 def unique_angles(self):
289 """Get Unique Angles.
291 :data:`all_angles` i-j-k without k-j-i.
293 """
294 return self._filter_unique(self.all_angles)
296 @property
297 def unique_dihedrals(self):
298 """Get Unique Dihedrals.
300 :data:`all_dihedrals` i-j-k-l without l-k-j-i.
302 """
303 return self._filter_unique(self.all_dihedrals)
306 def _get_symbol_idxs(self, imI, sym):
307 """Get list of indices of element *sym*"""
308 if isinstance(imI, int):
309 return [ idx for idx in range(len(self.images[imI])) if self.images[imI][idx].symbol == sym ]
310 else:
311 return [ idx for idx in range(len(imI)) if imI[idx].symbol == sym ]
314 def _idxTuple2SymbolTuple(self, imI, tup):
315 """Converts a tuple of indices to their symbols"""
316 return ( self.images[imI][idx].symbol for idx in tup )
319 def get_bonds(self, A, B, unique=True):
320 """Get bonds from element A to element B.
322 Parameters:
324 A, B: str
325 Get Bonds between elements A and B
326 unique: bool
327 Return the bonds both ways or just one way (A-B and B-A or only A-B)
329 Returns:
331 return: list of lists of tuples
332 return[imageIdx][atomIdx][bondI], each tuple starts with atomIdx.
334 Use :func:`get_values` to convert the returned list to values.
335 """
336 r = []
337 for imI in range(len(self.all_bonds)):
338 r.append([])
339 aIdxs = self._get_symbol_idxs(imI, A)
340 if A != B:
341 bIdxs = self._get_symbol_idxs(imI, B)
342 for idx in aIdxs:
343 bonded = self.all_bonds[imI][idx]
344 if A == B:
345 r[-1].extend([ (idx, x) for x in bonded if ( x in aIdxs ) and ( x > idx ) ])
346 else:
347 r[-1].extend([ (idx, x) for x in bonded if x in bIdxs ])
349 if not unique:
350 r[-1] += [ x[::-1] for x in r[-1] ]
352 return r
355 def get_angles(self, A, B, C, unique=True):
356 """Get angles from given elements A-B-C.
358 Parameters:
360 A, B, C: str
361 Get Angles between elements A, B and C. **B will be the central atom**.
362 unique: bool
363 Return the angles both ways or just one way (A-B-C and C-B-A or only A-B-C)
365 Returns:
367 return: list of lists of tuples
368 return[imageIdx][atomIdx][angleI], each tuple starts with atomIdx.
370 Use :func:`get_values` to convert the returned list to values.
371 """
372 from itertools import product, combinations
373 r = []
374 for imI in range(len(self.all_angles)):
375 r.append([])
376 #Middle Atom is fixed
377 bIdxs = self._get_symbol_idxs(imI, B)
378 for bIdx in bIdxs:
379 bondedA = [ idx for idx in self.all_bonds[imI][bIdx] if self.images[imI][idx].symbol == A ]
380 if len(bondedA) == 0:
381 continue
383 if A != C:
384 bondedC = [ idx for idx in self.all_bonds[imI][bIdx] if self.images[imI][idx].symbol == C ]
385 if len(bondedC) == 0:
386 continue
388 if A == C:
389 extend = [ (x[0], bIdx, x[1]) for x in list(combinations(bondedA, 2)) ]
390 else:
391 extend = list(product(bondedA, [bIdx], bondedC))
393 if not unique:
394 extend += [ x[::-1] for x in extend ]
396 r[-1].extend(extend)
397 return r
400 def get_dihedrals(self, A, B, C, D, unique=True):
401 """Get dihedrals A-B-C-D.
403 Parameters:
405 A, B, C, D: str
406 Get Dihedralss between elements A, B, C and D. **B-C will be the central axis**.
407 unique: bool
408 Return the dihedrals both ways or just one way (A-B-C-D and D-C-B-A or only A-B-C-D)
410 Returns:
412 return: list of lists of tuples
413 return[imageIdx][atomIdx][dihedralI], each tuple starts with atomIdx.
415 Use :func:`get_values` to convert the returned list to values.
416 """
417 r = []
418 for imI in range(len(self.all_dihedrals)):
419 r.append([])
420 #get indices of elements
421 aIdxs = self._get_symbol_idxs(imI, A)
422 bIdxs = self._get_symbol_idxs(imI, B)
423 cIdxs = self._get_symbol_idxs(imI, C)
424 dIdxs = self._get_symbol_idxs(imI, D)
425 for aIdx in aIdxs:
426 dihedrals = [ (aIdx, ) + d for d in self.all_dihedrals[imI][aIdx] if ( d[0] in bIdxs ) and ( d[1] in cIdxs ) and ( d[2] in dIdxs ) ]
427 if not unique:
428 dihedrals += [ d[::-1] for d in dihedrals ]
429 r[-1].extend(dihedrals)
431 return r
434 def get_bond_value(self, imIdx, idxs, mic=True, **kwargs):
435 """Get bond length.
437 Parameters:
439 imIdx: int
440 Index of Image to get value from.
441 idxs: tuple or list of integers
442 Get distance between atoms idxs[0]-idxs[1].
443 mic: bool
444 Passed on to :func:`ase.Atoms.get_distance` for retrieving the value, defaults to True.
445 If the cell of the image is correctly set, there should be no reason to change this.
446 kwargs: options or dict
447 Passed on to :func:`ase.Atoms.get_distance`.
449 Returns:
451 return: float
452 Value returned by image.get_distance.
453 """
454 return self.images[imIdx].get_distance(idxs[0], idxs[1], mic=mic, **kwargs)
456 def get_angle_value(self, imIdx, idxs, mic=True, **kwargs):
457 """Get angle.
459 Parameters:
461 imIdx: int
462 Index of Image to get value from.
463 idxs: tuple or list of integers
464 Get angle between atoms idxs[0]-idxs[1]-idxs[2].
465 mic: bool
466 Passed on to :func:`ase.Atoms.get_angle` for retrieving the value, defaults to True.
467 If the cell of the image is correctly set, there should be no reason to change this.
468 kwargs: options or dict
469 Passed on to :func:`ase.Atoms.get_angle`.
471 Returns:
473 return: float
474 Value returned by image.get_angle.
475 """
476 return self.images[imIdx].get_angle(idxs[0], idxs[1], idxs[2], mic=True, **kwargs)
478 def get_dihedral_value(self, imIdx, idxs, mic=True, **kwargs):
479 """Get dihedral.
481 Parameters:
483 imIdx: int
484 Index of Image to get value from.
485 idxs: tuple or list of integers
486 Get angle between atoms idxs[0]-idxs[1]-idxs[2]-idxs[3].
487 mic: bool
488 Passed on to :func:`ase.Atoms.get_dihedral` for retrieving the value, defaults to True.
489 If the cell of the image is correctly set, there should be no reason to change this.
490 kwargs: options or dict
491 Passed on to :func:`ase.Atoms.get_dihedral`.
493 Returns:
495 return: float
496 Value returned by image.get_dihedral.
497 """
498 return self.images[imIdx].get_dihedral(idxs[0], idxs[1], idxs[2], idxs[3], mic=mic, **kwargs)
500 def get_values(self, inputList, imageIdx=None, mic=True, **kwargs):
501 """Get Bond/Angle/Dihedral values.
503 Parameters:
505 inputList: list of lists of tuples
506 Can be any list provided by :meth:`~ase.geometry.analysis.Analysis.get_bonds`,
507 :meth:`~ase.geometry.analysis.Analysis.get_angles` or
508 :meth:`~ase.geometry.analysis.Analysis.get_dihedrals`.
509 imageIdx: integer or slice
510 The images from :data:`images` to be analyzed. If None, all frames will be analyzed.
511 See :func:`~ase.geometry.analysis.Analysis._get_slice` for details.
512 mic: bool
513 Passed on to :class:`~ase.Atoms` for retrieving the values, defaults to True.
514 If the cells of the images are correctly set, there should be no reason to change this.
515 kwargs: options or dict
516 Passed on to the :class:`~ase.Atoms` classes functions for retrieving the values.
518 Returns:
520 return: list of lists of floats
521 return[imageIdx][valueIdx]. Has the same shape as the *inputList*, instead of each
522 tuple there is a float with the value this tuple yields.
524 The type of value requested is determined from the length of the tuple inputList[0][0].
525 The methods from the :class:`~ase.Atoms` class are used.
526 """
528 sl = self._get_slice(imageIdx)
530 #get method to call from length of inputList
531 if len(inputList[0][0]) == 2:
532 get = self.get_bond_value
533 elif len(inputList[0][0]) == 3:
534 get = self.get_angle_value
535 elif len(inputList[0][0]) == 4:
536 get = self.get_dihedral_value
537 else:
538 raise ValueError("inputList in ase.geometry.analysis.Analysis.get_values has a bad shape.")
540 #check if length of slice and inputList match
541 singleNL = False
542 if len(inputList) != len(self.images[sl]):
543 #only one nl for all images
544 if len(inputList) == 1 and len(self.nl) == 1:
545 singleNL = True
546 else:
547 raise RuntimeError("Length of inputList does not match length of \
548 images requested, but it also is not one item long.")
550 r = []
551 for inputIdx, image in enumerate(self.images[sl]):
552 imageIdx = self.images.index(image)
553 r.append([])
554 #always use first list from input if only a single neighborlist was used
555 if singleNL:
556 inputIdx = 0
557 for tupl in inputList[inputIdx]:
558 r[-1].append(get(imageIdx, tupl, mic=mic, **kwargs))
560 return r
562 def get_max_volume_estimate(self):
563 return get_max_volume_estimate(self.images)
565 def get_rdf(self, rmax, nbins, imageIdx=None, elements=None, return_dists=False,
566 volume: Optional[float] = None):
567 """Get RDF.
569 Wrapper for :meth:`ase.ga.utilities.get_rdf` with more selection possibilities.
571 Parameters:
573 rmax: float
574 Maximum distance of RDF.
575 nbins: int
576 Number of bins to divide RDF.
577 imageIdx: int/slice/None
578 Images to analyze, see :func:`_get_slice` for details.
579 elements: str/int/list/tuple
580 Make partial RDFs.
582 If elements is *None*, a full RDF is calculated. If elements is an *integer* or a *list/tuple
583 of integers*, only those atoms will contribute to the RDF (like a mask). If elements
584 is a *string* or a *list/tuple of strings*, only Atoms of those elements will contribute.
586 Returns:
588 return: list of lists / list of tuples of lists
589 If return_dists is True, the returned tuples contain (rdf, distances). Otherwise
590 only rdfs for each image are returned.
591 """
593 sl = self._get_slice(imageIdx)
595 ls_rdf = []
596 el = None
598 for image in self.images[sl]:
599 if elements is None:
600 tmp_image = image
601 #integers
602 elif isinstance(elements, int):
603 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
604 tmp_image.append(image[elements])
605 #strings
606 elif isinstance(elements, str):
607 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
608 for idx in self._get_symbol_idxs(image, elements):
609 tmp_image.append(image[idx])
610 #lists
611 elif isinstance(elements, list) or isinstance(elements, tuple):
612 #list of ints
613 if all(isinstance(x, int) for x in elements):
614 if len(elements) == 2:
615 #use builtin get_rdf mask
616 el = elements
617 tmp_image = image
618 else:
619 #create dummy image
620 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
621 for idx in elements:
622 tmp_image.append(image[idx])
623 #list of strings
624 elif all(isinstance(x, str) for x in elements):
625 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
626 for element in elements:
627 for idx in self._get_symbol_idxs(image, element):
628 tmp_image.append(image[idx])
629 else:
630 raise ValueError("Unsupported type of elements given in ase.geometry.analysis.Analysis.get_rdf!")
631 else:
632 raise ValueError("Unsupported type of elements given in ase.geometry.analysis.Analysis.get_rdf!")
634 rdf = get_rdf(tmp_image, rmax, nbins, elements=el, no_dists=(not return_dists),
635 volume=volume)
636 ls_rdf.append(rdf)
638 return ls_rdf