Coverage for /builds/ase/ase/ase/geometry/analysis.py: 70.61%
262 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3# flake8: noqa
4"""Tools for analyzing instances of :class:`~ase.Atoms`
5"""
7from typing import List, Optional
9import numpy as np
11from ase import Atoms
12from ase.geometry.rdf import get_containing_cell_length, get_rdf
13from ase.neighborlist import (build_neighbor_list, get_distance_indices,
14 get_distance_matrix)
16__all__ = ['Analysis']
19def get_max_containing_cell_length(images: List[Atoms]):
20 i2diff = np.zeros(3)
21 for image in images:
22 np.maximum(get_containing_cell_length(image), i2diff, out=i2diff)
23 return i2diff
26def get_max_volume_estimate(images: List[Atoms]) -> float:
27 return np.prod(get_max_containing_cell_length(images))
30class Analysis:
31 """Analysis class
33 Parameters for initialization:
35 images: :class:`~ase.Atoms` object or list of such
36 Images to analyze.
37 nl: None, :class:`~ase.neighborlist.NeighborList` object or list of such
38 Neighborlist(s) for the given images. One or nImages, depending if bonding
39 pattern changes or is constant. Using one Neigborlist greatly improves speed.
40 kwargs: options, dict
41 Arguments for constructing :class:`~ase.neighborlist.NeighborList` object if :data:`nl` is None.
43 The choice of ``bothways=True`` for the :class:`~ase.neighborlist.NeighborList` object
44 will not influence the amount of bonds/angles/dihedrals you get, all are reported
45 in both directions. Use the *unique*-labeled properties to get lists without
46 duplicates.
47 """
49 def __init__(self, images, nl=None, **kwargs):
50 self.images = images
52 if isinstance(nl, list):
53 assert len(nl) == self.nImages
54 self._nl = nl
55 elif nl is not None:
56 self._nl = [nl]
57 else:
58 self._nl = [build_neighbor_list(self.images[0], **kwargs)]
60 self._cache = {}
62 def _get_slice(self, imageIdx):
63 """Return a slice from user input.
64 Using *imageIdx* (can be integer or slice) the analyzed frames can be specified.
65 If *imageIdx* is None, all frames will be analyzed.
66 """
67 # get slice from imageIdx
68 if isinstance(imageIdx, int):
69 sl = slice(imageIdx, imageIdx + 1)
70 elif isinstance(imageIdx, slice):
71 sl = imageIdx
72 elif imageIdx is None:
73 sl = slice(0, None)
74 else:
75 raise ValueError(
76 "Unsupported type for imageIdx in ase.geometry.analysis.Analysis._get_slice")
77 return sl
79 @property
80 def images(self):
81 """Images.
83 Set during initialization but can also be set later.
84 """
85 return self._images
87 @images.setter
88 def images(self, images):
89 """Set images"""
90 if isinstance(images, list):
91 self._images = images
92 else:
93 self._images = [images]
95 @images.deleter
96 def images(self):
97 """Delete images"""
98 self._images = None
100 @property
101 def nImages(self):
102 """Number of Images in this instance.
104 Cannot be set, is determined automatically.
105 """
106 return len(self.images)
108 @property
109 def nl(self):
110 """Neighbor Lists in this instance.
112 Set during initialization.
114 **No setter or deleter, only getter**
115 """
116 return self._nl
118 def _get_all_x(self, distance):
119 """Helper function to get bonds, angles, dihedrals"""
120 maxIter = self.nImages
121 if len(self.nl) == 1:
122 maxIter = 1
124 xList = []
125 for i in range(maxIter):
126 xList.append(get_distance_indices(
127 self.distance_matrix[i], distance))
129 return xList
131 @property
132 def all_bonds(self):
133 """All Bonds.
135 A list with indices of bonded atoms for each neighborlist in *self*.
136 Atom i is connected to all atoms inside result[i]. Duplicates from PBCs are
137 removed. See also :data:`unique_bonds`.
139 **No setter or deleter, only getter**
140 """
141 if 'allBonds' not in self._cache:
142 self._cache['allBonds'] = self._get_all_x(1)
144 return self._cache['allBonds']
146 @property
147 def all_angles(self):
148 """All angles
150 A list with indices of atoms in angles for each neighborlist in *self*.
151 Atom i forms an angle to the atoms inside the tuples in result[i]:
152 i -- result[i][x][0] -- result[i][x][1]
153 where x is in range(number of angles from i). See also :data:`unique_angles`.
155 **No setter or deleter, only getter**
156 """
157 if 'allAngles' not in self._cache:
158 self._cache['allAngles'] = []
159 distList = self._get_all_x(2)
161 for imI in range(len(distList)):
162 self._cache['allAngles'].append([])
163 # iterate over second neighbors of all atoms
164 for iAtom, secNeighs in enumerate(distList[imI]):
165 self._cache['allAngles'][-1].append([])
166 if len(secNeighs) == 0:
167 continue
168 firstNeighs = self.all_bonds[imI][iAtom]
169 # iterate over second neighbors of iAtom
170 for kAtom in secNeighs:
171 relevantFirstNeighs = [
172 idx for idx in firstNeighs if kAtom in self.all_bonds[imI][idx]]
173 # iterate over all atoms that are connected to iAtom and kAtom
174 for jAtom in relevantFirstNeighs:
175 self._cache['allAngles'][-1][-1].append(
176 (jAtom, kAtom))
178 return self._cache['allAngles']
180 @property
181 def all_dihedrals(self):
182 """All dihedrals
184 Returns a list with indices of atoms in dihedrals for each neighborlist in this instance.
185 Atom i forms a dihedral to the atoms inside the tuples in result[i]:
186 i -- result[i][x][0] -- result[i][x][1] -- result[i][x][2]
187 where x is in range(number of dihedrals from i). See also :data:`unique_dihedrals`.
189 **No setter or deleter, only getter**
190 """
191 if 'allDihedrals' not in self._cache:
192 self._cache['allDihedrals'] = []
193 distList = self._get_all_x(3)
195 for imI in range(len(distList)):
196 self._cache['allDihedrals'].append([])
197 for iAtom, thirdNeighs in enumerate(distList[imI]):
198 self._cache['allDihedrals'][-1].append([])
199 if len(thirdNeighs) == 0:
200 continue
201 anglesI = self.all_angles[imI][iAtom]
202 # iterate over third neighbors of iAtom
203 for lAtom in thirdNeighs:
204 secondNeighs = [angle[-1] for angle in anglesI]
205 firstNeighs = [angle[0] for angle in anglesI]
206 relevantSecondNeighs = [
207 idx for idx in secondNeighs if lAtom in self.all_bonds[imI][idx]]
208 relevantFirstNeighs = [
209 firstNeighs[secondNeighs.index(idx)] for idx in relevantSecondNeighs]
210 # iterate over all atoms that are connected to iAtom and lAtom
211 for jAtom, kAtom in zip(relevantFirstNeighs, relevantSecondNeighs):
212 # remove dihedrals in circles
213 tupl = (jAtom, kAtom, lAtom)
214 if len(set((iAtom, ) + tupl)) != 4:
215 continue
216 # avoid duplicates
217 elif tupl in self._cache['allDihedrals'][-1][-1]:
218 continue
219 elif iAtom in tupl:
220 raise RuntimeError(
221 "Something is wrong in analysis.all_dihedrals!")
222 self._cache['allDihedrals'][-1][-1].append(
223 (jAtom, kAtom, lAtom))
225 return self._cache['allDihedrals']
227 @property
228 def adjacency_matrix(self):
229 """The adjacency/connectivity matrix.
231 If not already done, build a list of adjacency matrices for all :data:`nl`.
233 **No setter or deleter, only getter**
234 """
236 if 'adjacencyMatrix' not in self._cache:
237 self._cache['adjacencyMatrix'] = []
238 for i in range(len(self.nl)):
239 self._cache['adjacencyMatrix'].append(
240 self.nl[i].get_connectivity_matrix())
242 return self._cache['adjacencyMatrix']
244 @property
245 def distance_matrix(self):
246 """The distance matrix.
248 If not already done, build a list of distance matrices for all :data:`nl`. See
249 :meth:`ase.neighborlist.get_distance_matrix`.
251 **No setter or deleter, only getter**
252 """
254 if 'distanceMatrix' not in self._cache:
255 self._cache['distanceMatrix'] = []
256 for i in range(len(self.nl)):
257 self._cache['distanceMatrix'].append(
258 get_distance_matrix(self.adjacency_matrix[i]))
260 return self._cache['distanceMatrix']
262 @property
263 def unique_bonds(self):
264 """Get Unique Bonds.
266 :data:`all_bonds` i-j without j-i. This is the upper triangle of the
267 connectivity matrix (i,j), `i < j`
269 """
270 bonds = []
271 for imI in range(len(self.all_bonds)):
272 bonds.append([])
273 for iAtom, bonded in enumerate(self.all_bonds[imI]):
274 bonds[-1].append([jAtom for jAtom in bonded if jAtom > iAtom])
276 return bonds
278 def _filter_unique(self, l):
279 """Helper function to filter for unique lists in a list
280 that also contains the reversed items.
281 """
282 r = []
283 # iterate over images
284 for imI in range(len(l)):
285 r.append([])
286 # iterate over atoms
287 for i, tuples in enumerate(l[imI]):
288 # add the ones where i is smaller than the last element
289 r[-1].append([x for x in tuples if i < x[-1]])
290 return r
292 def clear_cache(self):
293 """Delete all cached information."""
294 self._cache = {}
296 @property
297 def unique_angles(self):
298 """Get Unique Angles.
300 :data:`all_angles` i-j-k without k-j-i.
302 """
303 return self._filter_unique(self.all_angles)
305 @property
306 def unique_dihedrals(self):
307 """Get Unique Dihedrals.
309 :data:`all_dihedrals` i-j-k-l without l-k-j-i.
311 """
312 return self._filter_unique(self.all_dihedrals)
314 def _get_symbol_idxs(self, imI, sym):
315 """Get list of indices of element *sym*"""
316 if isinstance(imI, int):
317 return [idx for idx in range(len(self.images[imI])) if self.images[imI][idx].symbol == sym]
318 else:
319 return [idx for idx in range(len(imI)) if imI[idx].symbol == sym]
321 def _idxTuple2SymbolTuple(self, imI, tup):
322 """Converts a tuple of indices to their symbols"""
323 return (self.images[imI][idx].symbol for idx in tup)
325 def get_bonds(self, A, B, unique=True):
326 """Get bonds from element A to element B.
328 Parameters:
330 A, B: str
331 Get Bonds between elements A and B
332 unique: bool
333 Return the bonds both ways or just one way (A-B and B-A or only A-B)
335 Returns:
337 return: list of lists of tuples
338 return[imageIdx][atomIdx][bondI], each tuple starts with atomIdx.
340 Use :func:`get_values` to convert the returned list to values.
341 """
342 r = []
343 for imI in range(len(self.all_bonds)):
344 r.append([])
345 aIdxs = self._get_symbol_idxs(imI, A)
346 if A != B:
347 bIdxs = self._get_symbol_idxs(imI, B)
348 for idx in aIdxs:
349 bonded = self.all_bonds[imI][idx]
350 if A == B:
351 r[-1].extend([(idx, x)
352 for x in bonded if (x in aIdxs) and (x > idx)])
353 else:
354 r[-1].extend([(idx, x) for x in bonded if x in bIdxs])
356 if not unique:
357 r[-1] += [x[::-1] for x in r[-1]]
359 return r
361 def get_angles(self, A, B, C, unique=True):
362 """Get angles from given elements A-B-C.
364 Parameters:
366 A, B, C: str
367 Get Angles between elements A, B and C. **B will be the central atom**.
368 unique: bool
369 Return the angles both ways or just one way (A-B-C and C-B-A or only A-B-C)
371 Returns:
373 return: list of lists of tuples
374 return[imageIdx][atomIdx][angleI], each tuple starts with atomIdx.
376 Use :func:`get_values` to convert the returned list to values.
377 """
378 from itertools import combinations, product
379 r = []
380 for imI in range(len(self.all_angles)):
381 r.append([])
382 # Middle Atom is fixed
383 bIdxs = self._get_symbol_idxs(imI, B)
384 for bIdx in bIdxs:
385 bondedA = [idx for idx in self.all_bonds[imI]
386 [bIdx] if self.images[imI][idx].symbol == A]
387 if len(bondedA) == 0:
388 continue
390 if A != C:
391 bondedC = [idx for idx in self.all_bonds[imI]
392 [bIdx] if self.images[imI][idx].symbol == C]
393 if len(bondedC) == 0:
394 continue
396 if A == C:
397 extend = [(x[0], bIdx, x[1])
398 for x in list(combinations(bondedA, 2))]
399 else:
400 extend = list(product(bondedA, [bIdx], bondedC))
402 if not unique:
403 extend += [x[::-1] for x in extend]
405 r[-1].extend(extend)
406 return r
408 def get_dihedrals(self, A, B, C, D, unique=True):
409 """Get dihedrals A-B-C-D.
411 Parameters:
413 A, B, C, D: str
414 Get Dihedralss between elements A, B, C and D. **B-C will be the central axis**.
415 unique: bool
416 Return the dihedrals both ways or just one way (A-B-C-D and D-C-B-A or only A-B-C-D)
418 Returns:
420 return: list of lists of tuples
421 return[imageIdx][atomIdx][dihedralI], each tuple starts with atomIdx.
423 Use :func:`get_values` to convert the returned list to values.
424 """
425 r = []
426 for imI in range(len(self.all_dihedrals)):
427 r.append([])
428 # get indices of elements
429 aIdxs = self._get_symbol_idxs(imI, A)
430 bIdxs = self._get_symbol_idxs(imI, B)
431 cIdxs = self._get_symbol_idxs(imI, C)
432 dIdxs = self._get_symbol_idxs(imI, D)
433 for aIdx in aIdxs:
434 dihedrals = [(aIdx, ) + d for d in self.all_dihedrals[imI][aIdx]
435 if (d[0] in bIdxs) and (d[1] in cIdxs) and (d[2] in dIdxs)]
436 if not unique:
437 dihedrals += [d[::-1] for d in dihedrals]
438 r[-1].extend(dihedrals)
440 return r
442 def get_bond_value(self, imIdx, idxs, mic=True, **kwargs):
443 """Get bond length.
445 Parameters:
447 imIdx: int
448 Index of Image to get value from.
449 idxs: tuple or list of integers
450 Get distance between atoms idxs[0]-idxs[1].
451 mic: bool
452 Passed on to :func:`ase.Atoms.get_distance` for retrieving the value, defaults to True.
453 If the cell of the image is correctly set, there should be no reason to change this.
454 kwargs: options or dict
455 Passed on to :func:`ase.Atoms.get_distance`.
457 Returns:
459 return: float
460 Value returned by image.get_distance.
461 """
462 return self.images[imIdx].get_distance(idxs[0], idxs[1], mic=mic, **kwargs)
464 def get_angle_value(self, imIdx, idxs, mic=True, **kwargs):
465 """Get angle.
467 Parameters:
469 imIdx: int
470 Index of Image to get value from.
471 idxs: tuple or list of integers
472 Get angle between atoms idxs[0]-idxs[1]-idxs[2].
473 mic: bool
474 Passed on to :func:`ase.Atoms.get_angle` for retrieving the value, defaults to True.
475 If the cell of the image is correctly set, there should be no reason to change this.
476 kwargs: options or dict
477 Passed on to :func:`ase.Atoms.get_angle`.
479 Returns:
481 return: float
482 Value returned by image.get_angle.
483 """
484 return self.images[imIdx].get_angle(idxs[0], idxs[1], idxs[2], mic=True, **kwargs)
486 def get_dihedral_value(self, imIdx, idxs, mic=True, **kwargs):
487 """Get dihedral.
489 Parameters:
491 imIdx: int
492 Index of Image to get value from.
493 idxs: tuple or list of integers
494 Get angle between atoms idxs[0]-idxs[1]-idxs[2]-idxs[3].
495 mic: bool
496 Passed on to :func:`ase.Atoms.get_dihedral` for retrieving the value, defaults to True.
497 If the cell of the image is correctly set, there should be no reason to change this.
498 kwargs: options or dict
499 Passed on to :func:`ase.Atoms.get_dihedral`.
501 Returns:
503 return: float
504 Value returned by image.get_dihedral.
505 """
506 return self.images[imIdx].get_dihedral(idxs[0], idxs[1], idxs[2], idxs[3], mic=mic, **kwargs)
508 def get_values(self, inputList, imageIdx=None, mic=True, **kwargs):
509 """Get Bond/Angle/Dihedral values.
511 Parameters:
513 inputList: list of lists of tuples
514 Can be any list provided by :meth:`~ase.geometry.analysis.Analysis.get_bonds`,
515 :meth:`~ase.geometry.analysis.Analysis.get_angles` or
516 :meth:`~ase.geometry.analysis.Analysis.get_dihedrals`.
517 imageIdx: integer or slice
518 The images from :data:`images` to be analyzed. If None, all frames will be analyzed.
519 See :func:`~ase.geometry.analysis.Analysis._get_slice` for details.
520 mic: bool
521 Passed on to :class:`~ase.Atoms` for retrieving the values, defaults to True.
522 If the cells of the images are correctly set, there should be no reason to change this.
523 kwargs: options or dict
524 Passed on to the :class:`~ase.Atoms` classes functions for retrieving the values.
526 Returns:
528 return: list of lists of floats
529 return[imageIdx][valueIdx]. Has the same shape as the *inputList*, instead of each
530 tuple there is a float with the value this tuple yields.
532 The type of value requested is determined from the length of the tuple inputList[0][0].
533 The methods from the :class:`~ase.Atoms` class are used.
534 """
536 sl = self._get_slice(imageIdx)
538 # get method to call from length of inputList
539 if len(inputList[0][0]) == 2:
540 get = self.get_bond_value
541 elif len(inputList[0][0]) == 3:
542 get = self.get_angle_value
543 elif len(inputList[0][0]) == 4:
544 get = self.get_dihedral_value
545 else:
546 raise ValueError(
547 "inputList in ase.geometry.analysis.Analysis.get_values has a bad shape.")
549 # check if length of slice and inputList match
550 singleNL = False
551 if len(inputList) != len(self.images[sl]):
552 # only one nl for all images
553 if len(inputList) == 1 and len(self.nl) == 1:
554 singleNL = True
555 else:
556 raise RuntimeError("Length of inputList does not match length of \
557 images requested, but it also is not one item long.")
559 r = []
560 for inputIdx, image in enumerate(self.images[sl]):
561 imageIdx = self.images.index(image)
562 r.append([])
563 # always use first list from input if only a single neighborlist was used
564 if singleNL:
565 inputIdx = 0
566 for tupl in inputList[inputIdx]:
567 r[-1].append(get(imageIdx, tupl, mic=mic, **kwargs))
569 return r
571 def get_max_volume_estimate(self):
572 return get_max_volume_estimate(self.images)
574 def get_rdf(self, rmax, nbins, imageIdx=None, elements=None, return_dists=False,
575 volume: Optional[float] = None):
576 """Get RDF.
578 Wrapper for :meth:`ase.ga.utilities.get_rdf` with more selection possibilities.
580 Parameters:
582 rmax: float
583 Maximum distance of RDF.
584 nbins: int
585 Number of bins to divide RDF.
586 imageIdx: int/slice/None
587 Images to analyze, see :func:`_get_slice` for details.
588 elements: str/int/list/tuple
589 Make partial RDFs.
591 If elements is *None*, a full RDF is calculated. If elements is an *integer* or a *list/tuple
592 of integers*, only those atoms will contribute to the RDF (like a mask). If elements
593 is a *string* or a *list/tuple of strings*, only Atoms of those elements will contribute.
595 Returns:
597 return: list of lists / list of tuples of lists
598 If return_dists is True, the returned tuples contain (rdf, distances). Otherwise
599 only rdfs for each image are returned.
600 """
602 sl = self._get_slice(imageIdx)
604 ls_rdf = []
605 el = None
607 for image in self.images[sl]:
608 if elements is None:
609 tmp_image = image
610 # integers
611 elif isinstance(elements, int):
612 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
613 tmp_image.append(image[elements])
614 # strings
615 elif isinstance(elements, str):
616 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
617 for idx in self._get_symbol_idxs(image, elements):
618 tmp_image.append(image[idx])
619 # lists
620 elif isinstance(elements, (list, tuple)):
621 # list of ints
622 if all(isinstance(x, int) for x in elements):
623 if len(elements) == 2:
624 # use builtin get_rdf mask
625 el = elements
626 tmp_image = image
627 else:
628 # create dummy image
629 tmp_image = Atoms(
630 cell=image.get_cell(), pbc=image.get_pbc())
631 for idx in elements:
632 tmp_image.append(image[idx])
633 # list of strings
634 elif all(isinstance(x, str) for x in elements):
635 tmp_image = Atoms(cell=image.get_cell(),
636 pbc=image.get_pbc())
637 for element in elements:
638 for idx in self._get_symbol_idxs(image, element):
639 tmp_image.append(image[idx])
640 else:
641 raise ValueError(
642 "Unsupported type of elements given in ase.geometry.analysis.Analysis.get_rdf!")
643 else:
644 raise ValueError(
645 "Unsupported type of elements given in ase.geometry.analysis.Analysis.get_rdf!")
647 rdf = get_rdf(tmp_image, rmax, nbins, elements=el, no_dists=(not return_dists),
648 volume=volume)
649 ls_rdf.append(rdf)
651 return ls_rdf