Coverage for /builds/kinetik161/ase/ase/geometry/analysis.py: 70.72%
263 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1# flake8: noqa
2"""Tools for analyzing instances of :class:`~ase.Atoms`
3"""
5from typing import List, Optional
7import numpy as np
9from ase import Atoms
10from ase.geometry.rdf import get_containing_cell_length, get_rdf
11from ase.neighborlist import (build_neighbor_list, get_distance_indices,
12 get_distance_matrix)
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.prod(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(
74 "Unsupported type for imageIdx in ase.geometry.analysis.Analysis._get_slice")
75 return sl
77 @property
78 def images(self):
79 """Images.
81 Set during initialization but can also be set later.
82 """
83 return self._images
85 @images.setter
86 def images(self, images):
87 """Set images"""
88 if isinstance(images, list):
89 self._images = images
90 else:
91 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(
125 self.distance_matrix[i], distance))
127 return xList
129 @property
130 def all_bonds(self):
131 """All Bonds.
133 A list with indices of bonded atoms for each neighborlist in *self*.
134 Atom i is connected to all atoms inside result[i]. Duplicates from PBCs are
135 removed. See also :data:`unique_bonds`.
137 **No setter or deleter, only getter**
138 """
139 if not 'allBonds' in self._cache:
140 self._cache['allBonds'] = self._get_all_x(1)
142 return self._cache['allBonds']
144 @property
145 def all_angles(self):
146 """All angles
148 A list with indices of atoms in angles for each neighborlist in *self*.
149 Atom i forms an angle to the atoms inside the tuples in result[i]:
150 i -- result[i][x][0] -- result[i][x][1]
151 where x is in range(number of angles from i). See also :data:`unique_angles`.
153 **No setter or deleter, only getter**
154 """
155 if not 'allAngles' in self._cache:
156 self._cache['allAngles'] = []
157 distList = self._get_all_x(2)
159 for imI in range(len(distList)):
160 self._cache['allAngles'].append([])
161 # iterate over second neighbors of all atoms
162 for iAtom, secNeighs in enumerate(distList[imI]):
163 self._cache['allAngles'][-1].append([])
164 if len(secNeighs) == 0:
165 continue
166 firstNeighs = self.all_bonds[imI][iAtom]
167 # iterate over second neighbors of iAtom
168 for kAtom in secNeighs:
169 relevantFirstNeighs = [
170 idx for idx in firstNeighs if kAtom in self.all_bonds[imI][idx]]
171 # iterate over all atoms that are connected to iAtom and kAtom
172 for jAtom in relevantFirstNeighs:
173 self._cache['allAngles'][-1][-1].append(
174 (jAtom, kAtom))
176 return self._cache['allAngles']
178 @property
179 def all_dihedrals(self):
180 """All dihedrals
182 Returns a list with indices of atoms in dihedrals for each neighborlist in this instance.
183 Atom i forms a dihedral to the atoms inside the tuples in result[i]:
184 i -- result[i][x][0] -- result[i][x][1] -- result[i][x][2]
185 where x is in range(number of dihedrals from i). See also :data:`unique_dihedrals`.
187 **No setter or deleter, only getter**
188 """
189 if not 'allDihedrals' in self._cache:
190 self._cache['allDihedrals'] = []
191 distList = self._get_all_x(3)
193 for imI in range(len(distList)):
194 self._cache['allDihedrals'].append([])
195 for iAtom, thirdNeighs in enumerate(distList[imI]):
196 self._cache['allDihedrals'][-1].append([])
197 if len(thirdNeighs) == 0:
198 continue
199 anglesI = self.all_angles[imI][iAtom]
200 # iterate over third neighbors of iAtom
201 for lAtom in thirdNeighs:
202 secondNeighs = [angle[-1] for angle in anglesI]
203 firstNeighs = [angle[0] for angle in anglesI]
204 relevantSecondNeighs = [
205 idx for idx in secondNeighs if lAtom in self.all_bonds[imI][idx]]
206 relevantFirstNeighs = [
207 firstNeighs[secondNeighs.index(idx)] for idx in relevantSecondNeighs]
208 # iterate over all atoms that are connected to iAtom and lAtom
209 for jAtom, kAtom in zip(relevantFirstNeighs, relevantSecondNeighs):
210 # remove dihedrals in circles
211 tupl = (jAtom, kAtom, lAtom)
212 if len(set((iAtom, ) + tupl)) != 4:
213 continue
214 # avoid duplicates
215 elif tupl in self._cache['allDihedrals'][-1][-1]:
216 continue
217 elif iAtom in tupl:
218 raise RuntimeError(
219 "Something is wrong in analysis.all_dihedrals!")
220 self._cache['allDihedrals'][-1][-1].append(
221 (jAtom, kAtom, lAtom))
223 return self._cache['allDihedrals']
225 @property
226 def adjacency_matrix(self):
227 """The adjacency/connectivity matrix.
229 If not already done, build a list of adjacency matrices for all :data:`nl`.
231 **No setter or deleter, only getter**
232 """
234 if not 'adjacencyMatrix' in self._cache:
235 self._cache['adjacencyMatrix'] = []
236 for i in range(len(self.nl)):
237 self._cache['adjacencyMatrix'].append(
238 self.nl[i].get_connectivity_matrix())
240 return self._cache['adjacencyMatrix']
242 @property
243 def distance_matrix(self):
244 """The distance matrix.
246 If not already done, build a list of distance matrices for all :data:`nl`. See
247 :meth:`ase.neighborlist.get_distance_matrix`.
249 **No setter or deleter, only getter**
250 """
252 if not 'distanceMatrix' in self._cache:
253 self._cache['distanceMatrix'] = []
254 for i in range(len(self.nl)):
255 self._cache['distanceMatrix'].append(
256 get_distance_matrix(self.adjacency_matrix[i]))
258 return self._cache['distanceMatrix']
260 @property
261 def unique_bonds(self):
262 """Get Unique Bonds.
264 :data:`all_bonds` i-j without j-i. This is the upper triangle of the
265 connectivity matrix (i,j), `i < j`
267 """
268 bonds = []
269 for imI in range(len(self.all_bonds)):
270 bonds.append([])
271 for iAtom, bonded in enumerate(self.all_bonds[imI]):
272 bonds[-1].append([jAtom for jAtom in bonded if jAtom > iAtom])
274 return bonds
276 def _filter_unique(self, l):
277 """Helper function to filter for unique lists in a list
278 that also contains the reversed items.
279 """
280 r = []
281 # iterate over images
282 for imI in range(len(l)):
283 r.append([])
284 # iterate over atoms
285 for i, tuples in enumerate(l[imI]):
286 # add the ones where i is smaller than the last element
287 r[-1].append([x for x in tuples if i < x[-1]])
288 return r
290 def clear_cache(self):
291 """Delete all cached information."""
292 self._cache = {}
294 @property
295 def unique_angles(self):
296 """Get Unique Angles.
298 :data:`all_angles` i-j-k without k-j-i.
300 """
301 return self._filter_unique(self.all_angles)
303 @property
304 def unique_dihedrals(self):
305 """Get Unique Dihedrals.
307 :data:`all_dihedrals` i-j-k-l without l-k-j-i.
309 """
310 return self._filter_unique(self.all_dihedrals)
312 def _get_symbol_idxs(self, imI, sym):
313 """Get list of indices of element *sym*"""
314 if isinstance(imI, int):
315 return [idx for idx in range(len(self.images[imI])) if self.images[imI][idx].symbol == sym]
316 else:
317 return [idx for idx in range(len(imI)) if imI[idx].symbol == sym]
319 def _idxTuple2SymbolTuple(self, imI, tup):
320 """Converts a tuple of indices to their symbols"""
321 return (self.images[imI][idx].symbol for idx in tup)
323 def get_bonds(self, A, B, unique=True):
324 """Get bonds from element A to element B.
326 Parameters:
328 A, B: str
329 Get Bonds between elements A and B
330 unique: bool
331 Return the bonds both ways or just one way (A-B and B-A or only A-B)
333 Returns:
335 return: list of lists of tuples
336 return[imageIdx][atomIdx][bondI], each tuple starts with atomIdx.
338 Use :func:`get_values` to convert the returned list to values.
339 """
340 r = []
341 for imI in range(len(self.all_bonds)):
342 r.append([])
343 aIdxs = self._get_symbol_idxs(imI, A)
344 if A != B:
345 bIdxs = self._get_symbol_idxs(imI, B)
346 for idx in aIdxs:
347 bonded = self.all_bonds[imI][idx]
348 if A == B:
349 r[-1].extend([(idx, x)
350 for x in bonded if (x in aIdxs) and (x > idx)])
351 else:
352 r[-1].extend([(idx, x) for x in bonded if x in bIdxs])
354 if not unique:
355 r[-1] += [x[::-1] for x in r[-1]]
357 return r
359 def get_angles(self, A, B, C, unique=True):
360 """Get angles from given elements A-B-C.
362 Parameters:
364 A, B, C: str
365 Get Angles between elements A, B and C. **B will be the central atom**.
366 unique: bool
367 Return the angles both ways or just one way (A-B-C and C-B-A or only A-B-C)
369 Returns:
371 return: list of lists of tuples
372 return[imageIdx][atomIdx][angleI], each tuple starts with atomIdx.
374 Use :func:`get_values` to convert the returned list to values.
375 """
376 from itertools import combinations, product
377 r = []
378 for imI in range(len(self.all_angles)):
379 r.append([])
380 # Middle Atom is fixed
381 bIdxs = self._get_symbol_idxs(imI, B)
382 for bIdx in bIdxs:
383 bondedA = [idx for idx in self.all_bonds[imI]
384 [bIdx] if self.images[imI][idx].symbol == A]
385 if len(bondedA) == 0:
386 continue
388 if A != C:
389 bondedC = [idx for idx in self.all_bonds[imI]
390 [bIdx] if self.images[imI][idx].symbol == C]
391 if len(bondedC) == 0:
392 continue
394 if A == C:
395 extend = [(x[0], bIdx, x[1])
396 for x in list(combinations(bondedA, 2))]
397 else:
398 extend = list(product(bondedA, [bIdx], bondedC))
400 if not unique:
401 extend += [x[::-1] for x in extend]
403 r[-1].extend(extend)
404 return r
406 def get_dihedrals(self, A, B, C, D, unique=True):
407 """Get dihedrals A-B-C-D.
409 Parameters:
411 A, B, C, D: str
412 Get Dihedralss between elements A, B, C and D. **B-C will be the central axis**.
413 unique: bool
414 Return the dihedrals both ways or just one way (A-B-C-D and D-C-B-A or only A-B-C-D)
416 Returns:
418 return: list of lists of tuples
419 return[imageIdx][atomIdx][dihedralI], each tuple starts with atomIdx.
421 Use :func:`get_values` to convert the returned list to values.
422 """
423 r = []
424 for imI in range(len(self.all_dihedrals)):
425 r.append([])
426 # get indices of elements
427 aIdxs = self._get_symbol_idxs(imI, A)
428 bIdxs = self._get_symbol_idxs(imI, B)
429 cIdxs = self._get_symbol_idxs(imI, C)
430 dIdxs = self._get_symbol_idxs(imI, D)
431 for aIdx in aIdxs:
432 dihedrals = [(aIdx, ) + d for d in self.all_dihedrals[imI][aIdx]
433 if (d[0] in bIdxs) and (d[1] in cIdxs) and (d[2] in dIdxs)]
434 if not unique:
435 dihedrals += [d[::-1] for d in dihedrals]
436 r[-1].extend(dihedrals)
438 return r
440 def get_bond_value(self, imIdx, idxs, mic=True, **kwargs):
441 """Get bond length.
443 Parameters:
445 imIdx: int
446 Index of Image to get value from.
447 idxs: tuple or list of integers
448 Get distance between atoms idxs[0]-idxs[1].
449 mic: bool
450 Passed on to :func:`ase.Atoms.get_distance` for retrieving the value, defaults to True.
451 If the cell of the image is correctly set, there should be no reason to change this.
452 kwargs: options or dict
453 Passed on to :func:`ase.Atoms.get_distance`.
455 Returns:
457 return: float
458 Value returned by image.get_distance.
459 """
460 return self.images[imIdx].get_distance(idxs[0], idxs[1], mic=mic, **kwargs)
462 def get_angle_value(self, imIdx, idxs, mic=True, **kwargs):
463 """Get angle.
465 Parameters:
467 imIdx: int
468 Index of Image to get value from.
469 idxs: tuple or list of integers
470 Get angle between atoms idxs[0]-idxs[1]-idxs[2].
471 mic: bool
472 Passed on to :func:`ase.Atoms.get_angle` for retrieving the value, defaults to True.
473 If the cell of the image is correctly set, there should be no reason to change this.
474 kwargs: options or dict
475 Passed on to :func:`ase.Atoms.get_angle`.
477 Returns:
479 return: float
480 Value returned by image.get_angle.
481 """
482 return self.images[imIdx].get_angle(idxs[0], idxs[1], idxs[2], mic=True, **kwargs)
484 def get_dihedral_value(self, imIdx, idxs, mic=True, **kwargs):
485 """Get dihedral.
487 Parameters:
489 imIdx: int
490 Index of Image to get value from.
491 idxs: tuple or list of integers
492 Get angle between atoms idxs[0]-idxs[1]-idxs[2]-idxs[3].
493 mic: bool
494 Passed on to :func:`ase.Atoms.get_dihedral` for retrieving the value, defaults to True.
495 If the cell of the image is correctly set, there should be no reason to change this.
496 kwargs: options or dict
497 Passed on to :func:`ase.Atoms.get_dihedral`.
499 Returns:
501 return: float
502 Value returned by image.get_dihedral.
503 """
504 return self.images[imIdx].get_dihedral(idxs[0], idxs[1], idxs[2], idxs[3], mic=mic, **kwargs)
506 def get_values(self, inputList, imageIdx=None, mic=True, **kwargs):
507 """Get Bond/Angle/Dihedral values.
509 Parameters:
511 inputList: list of lists of tuples
512 Can be any list provided by :meth:`~ase.geometry.analysis.Analysis.get_bonds`,
513 :meth:`~ase.geometry.analysis.Analysis.get_angles` or
514 :meth:`~ase.geometry.analysis.Analysis.get_dihedrals`.
515 imageIdx: integer or slice
516 The images from :data:`images` to be analyzed. If None, all frames will be analyzed.
517 See :func:`~ase.geometry.analysis.Analysis._get_slice` for details.
518 mic: bool
519 Passed on to :class:`~ase.Atoms` for retrieving the values, defaults to True.
520 If the cells of the images are correctly set, there should be no reason to change this.
521 kwargs: options or dict
522 Passed on to the :class:`~ase.Atoms` classes functions for retrieving the values.
524 Returns:
526 return: list of lists of floats
527 return[imageIdx][valueIdx]. Has the same shape as the *inputList*, instead of each
528 tuple there is a float with the value this tuple yields.
530 The type of value requested is determined from the length of the tuple inputList[0][0].
531 The methods from the :class:`~ase.Atoms` class are used.
532 """
534 sl = self._get_slice(imageIdx)
536 # get method to call from length of inputList
537 if len(inputList[0][0]) == 2:
538 get = self.get_bond_value
539 elif len(inputList[0][0]) == 3:
540 get = self.get_angle_value
541 elif len(inputList[0][0]) == 4:
542 get = self.get_dihedral_value
543 else:
544 raise ValueError(
545 "inputList in ase.geometry.analysis.Analysis.get_values has a bad shape.")
547 # check if length of slice and inputList match
548 singleNL = False
549 if len(inputList) != len(self.images[sl]):
550 # only one nl for all images
551 if len(inputList) == 1 and len(self.nl) == 1:
552 singleNL = True
553 else:
554 raise RuntimeError("Length of inputList does not match length of \
555 images requested, but it also is not one item long.")
557 r = []
558 for inputIdx, image in enumerate(self.images[sl]):
559 imageIdx = self.images.index(image)
560 r.append([])
561 # always use first list from input if only a single neighborlist was used
562 if singleNL:
563 inputIdx = 0
564 for tupl in inputList[inputIdx]:
565 r[-1].append(get(imageIdx, tupl, mic=mic, **kwargs))
567 return r
569 def get_max_volume_estimate(self):
570 return get_max_volume_estimate(self.images)
572 def get_rdf(self, rmax, nbins, imageIdx=None, elements=None, return_dists=False,
573 volume: Optional[float] = None):
574 """Get RDF.
576 Wrapper for :meth:`ase.ga.utilities.get_rdf` with more selection possibilities.
578 Parameters:
580 rmax: float
581 Maximum distance of RDF.
582 nbins: int
583 Number of bins to divide RDF.
584 imageIdx: int/slice/None
585 Images to analyze, see :func:`_get_slice` for details.
586 elements: str/int/list/tuple
587 Make partial RDFs.
589 If elements is *None*, a full RDF is calculated. If elements is an *integer* or a *list/tuple
590 of integers*, only those atoms will contribute to the RDF (like a mask). If elements
591 is a *string* or a *list/tuple of strings*, only Atoms of those elements will contribute.
593 Returns:
595 return: list of lists / list of tuples of lists
596 If return_dists is True, the returned tuples contain (rdf, distances). Otherwise
597 only rdfs for each image are returned.
598 """
600 sl = self._get_slice(imageIdx)
602 ls_rdf = []
603 el = None
605 for image in self.images[sl]:
606 if elements is None:
607 tmp_image = image
608 # integers
609 elif isinstance(elements, int):
610 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
611 tmp_image.append(image[elements])
612 # strings
613 elif isinstance(elements, str):
614 tmp_image = Atoms(cell=image.get_cell(), pbc=image.get_pbc())
615 for idx in self._get_symbol_idxs(image, elements):
616 tmp_image.append(image[idx])
617 # lists
618 elif isinstance(elements, (list, tuple)):
619 # list of ints
620 if all(isinstance(x, int) for x in elements):
621 if len(elements) == 2:
622 # use builtin get_rdf mask
623 el = elements
624 tmp_image = image
625 else:
626 # create dummy image
627 tmp_image = Atoms(
628 cell=image.get_cell(), pbc=image.get_pbc())
629 for idx in elements:
630 tmp_image.append(image[idx])
631 # list of strings
632 elif all(isinstance(x, str) for x in elements):
633 tmp_image = Atoms(cell=image.get_cell(),
634 pbc=image.get_pbc())
635 for element in elements:
636 for idx in self._get_symbol_idxs(image, element):
637 tmp_image.append(image[idx])
638 else:
639 raise ValueError(
640 "Unsupported type of elements given in ase.geometry.analysis.Analysis.get_rdf!")
641 else:
642 raise ValueError(
643 "Unsupported type of elements given in ase.geometry.analysis.Analysis.get_rdf!")
645 rdf = get_rdf(tmp_image, rmax, nbins, elements=el, no_dists=(not return_dists),
646 volume=volume)
647 ls_rdf.append(rdf)
649 return ls_rdf