Coverage for /builds/kinetik161/ase/ase/gui/images.py: 70.67%
283 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1import warnings
2from math import sqrt
4import numpy as np
6from ase import Atoms
7from ase.calculators.singlepoint import SinglePointCalculator
8from ase.constraints import FixAtoms
9from ase.data import covalent_radii
10from ase.geometry import find_mic
11from ase.gui.defaults import read_defaults
12from ase.gui.i18n import _
13from ase.io import read, string2index, write
16class Images:
17 def __init__(self, images=None):
18 self.covalent_radii = covalent_radii.copy()
19 self.config = read_defaults()
20 self.atom_scale = self.config['radii_scale']
21 if images is None:
22 images = [Atoms()]
23 self.initialize(images)
25 def __len__(self):
26 return len(self._images)
28 def __getitem__(self, index):
29 return self._images[index]
31 def __iter__(self):
32 return iter(self._images)
34 # XXXXXXX hack
35 # compatibility hacks while allowing variable number of atoms
36 def get_dynamic(self, atoms):
37 dynamic = np.ones(len(atoms), bool)
38 for constraint in atoms.constraints:
39 if isinstance(constraint, FixAtoms):
40 dynamic[constraint.index] = False
41 return dynamic
43 def set_dynamic(self, mask, value):
44 # Does not make much sense if different images have different
45 # atom counts. Attempts to apply mask to all images,
46 # to the extent possible.
47 for atoms in self:
48 dynamic = self.get_dynamic(atoms)
49 dynamic[mask[:len(atoms)]] = value
50 atoms.constraints = [c for c in atoms.constraints
51 if not isinstance(c, FixAtoms)]
52 atoms.constraints.append(FixAtoms(mask=~dynamic))
54 def scale_radii(self, scaling_factor):
55 self.covalent_radii *= scaling_factor
57 def get_energy(self, atoms):
58 try:
59 e = atoms.get_potential_energy() * self.repeat.prod()
60 except RuntimeError:
61 e = np.nan
62 return e
64 def get_forces(self, atoms):
65 try:
66 F = atoms.get_forces(apply_constraint=False)
67 except RuntimeError:
68 return None
69 else:
70 return F
72 def initialize(self, images, filenames=None):
73 nimages = len(images)
74 if filenames is None:
75 filenames = [None] * nimages
76 self.filenames = filenames
78 warning = False
80 self._images = []
82 # Whether length or chemical composition changes:
83 self.have_varying_species = False
84 for i, atoms in enumerate(images):
85 # copy atoms or not? Not copying allows back-editing,
86 # but copying actually forgets things like the attached
87 # calculator (might have forces/energies
88 self._images.append(atoms)
89 self.have_varying_species |= not np.array_equal(self[0].numbers,
90 atoms.numbers)
91 if hasattr(self, 'Q'):
92 assert False # XXX askhl fix quaternions
93 self.Q[i] = atoms.get_quaternions()
94 if (atoms.pbc != self[0].pbc).any():
95 warning = True
97 if warning:
98 import warnings
99 warnings.warn('Not all images have the same boundary conditions!')
101 self.maxnatoms = max(len(atoms) for atoms in self)
102 self.selected = np.zeros(self.maxnatoms, bool)
103 self.selected_ordered = []
104 self.visible = np.ones(self.maxnatoms, bool)
105 self.repeat = np.ones(3, int)
107 def get_radii(self, atoms):
108 radii = np.array([self.covalent_radii[z] for z in atoms.numbers])
109 radii *= self.atom_scale
110 return radii
112 def read(self, filenames, default_index=':', filetype=None):
113 if isinstance(default_index, str):
114 default_index = string2index(default_index)
116 images = []
117 names = []
118 for filename in filenames:
119 from ase.io.formats import parse_filename
121 if '@' in filename and 'postgres' not in filename or \
122 'postgres' in filename and filename.count('@') == 2:
123 actual_filename, index = parse_filename(filename, None)
124 else:
125 actual_filename, index = parse_filename(filename,
126 default_index)
128 # Read from stdin:
129 if filename == '-':
130 import sys
131 from io import BytesIO
132 buf = BytesIO(sys.stdin.buffer.read())
133 buf.seek(0)
134 filename = buf
135 filetype = 'traj'
137 imgs = read(filename, index, filetype)
138 if hasattr(imgs, 'iterimages'):
139 imgs = list(imgs.iterimages())
141 images.extend(imgs)
143 # Name each file as filename@index:
144 if isinstance(index, slice):
145 start = index.start or 0
146 step = index.step or 1
147 else:
148 start = index
149 step = 1
150 for i, img in enumerate(imgs):
151 if isinstance(start, int):
152 names.append('{}@{}'.format(
153 actual_filename, start + i * step))
154 else:
155 names.append(f'{actual_filename}@{start}')
157 self.initialize(images, names)
159 def repeat_results(self, atoms, repeat=None, oldprod=None):
160 """Return a dictionary which updates the magmoms, energy and forces
161 to the repeated amount of atoms.
162 """
163 def getresult(name, get_quantity):
164 # ase/io/trajectory.py line 170 does this by using
165 # the get_property(prop, atoms, allow_calculation=False)
166 # so that is an alternative option.
167 try:
168 if (not atoms.calc or
169 atoms.calc.calculation_required(atoms, [name])):
170 quantity = None
171 else:
172 quantity = get_quantity()
173 except Exception as err:
174 quantity = None
175 errmsg = ('An error occurred while retrieving {} '
176 'from the calculator: {}'.format(name, err))
177 warnings.warn(errmsg)
178 return quantity
180 if repeat is None:
181 repeat = self.repeat.prod()
182 if oldprod is None:
183 oldprod = self.repeat.prod()
185 results = {}
187 original_length = len(atoms) // oldprod
188 newprod = repeat.prod()
190 # Read the old properties
191 magmoms = getresult('magmoms', atoms.get_magnetic_moments)
192 magmom = getresult('magmom', atoms.get_magnetic_moment)
193 energy = getresult('energy', atoms.get_potential_energy)
194 forces = getresult('forces', atoms.get_forces)
196 # Update old properties to the repeated image
197 if magmoms is not None:
198 magmoms = np.tile(magmoms[:original_length], newprod)
199 results['magmoms'] = magmoms
201 if magmom is not None:
202 magmom = magmom * newprod / oldprod
203 results['magmom'] = magmom
205 if forces is not None:
206 forces = np.tile(forces[:original_length].T, newprod).T
207 results['forces'] = forces
209 if energy is not None:
210 energy = energy * newprod / oldprod
211 results['energy'] = energy
213 return results
215 def repeat_unit_cell(self):
216 for atoms in self:
217 # Get quantities taking into account current repeat():'
218 results = self.repeat_results(atoms, self.repeat.prod(),
219 oldprod=self.repeat.prod())
221 atoms.cell *= self.repeat.reshape((3, 1))
222 atoms.calc = SinglePointCalculator(atoms, **results)
223 self.repeat = np.ones(3, int)
225 def repeat_images(self, repeat):
226 from ase.constraints import FixAtoms
227 repeat = np.array(repeat)
228 oldprod = self.repeat.prod()
229 images = []
230 constraints_removed = False
232 for i, atoms in enumerate(self):
233 refcell = atoms.get_cell()
234 fa = []
235 for c in atoms._constraints:
236 if isinstance(c, FixAtoms):
237 fa.append(c)
238 else:
239 constraints_removed = True
240 atoms.set_constraint(fa)
242 # Update results dictionary to repeated atoms
243 results = self.repeat_results(atoms, repeat, oldprod)
245 del atoms[len(atoms) // oldprod:] # Original atoms
247 atoms *= repeat
248 atoms.cell = refcell
250 atoms.calc = SinglePointCalculator(atoms, **results)
252 images.append(atoms)
254 if constraints_removed:
255 from ase.gui.ui import showwarning, tk
257 # We must be able to show warning before the main GUI
258 # has been created. So we create a new window,
259 # then show the warning, then destroy the window.
260 tmpwindow = tk.Tk()
261 tmpwindow.withdraw() # Host window will never be shown
262 showwarning(_('Constraints discarded'),
263 _('Constraints other than FixAtoms '
264 'have been discarded.'))
265 tmpwindow.destroy()
267 self.initialize(images, filenames=self.filenames)
268 self.repeat = repeat
270 def center(self):
271 """Center each image in the existing unit cell, keeping the
272 cell constant."""
273 for atoms in self:
274 atoms.center()
276 def graph(self, expr):
277 """Routine to create the data in graphs, defined by the
278 string expr."""
279 import ase.units as units
280 code = compile(expr + ',', '<input>', 'eval')
282 nimages = len(self)
284 def d(n1, n2):
285 return sqrt(((R[n1] - R[n2])**2).sum())
287 def a(n1, n2, n3):
288 v1 = R[n1] - R[n2]
289 v2 = R[n3] - R[n2]
290 arg = np.vdot(v1, v2) / (sqrt((v1**2).sum() * (v2**2).sum()))
291 if arg > 1.0:
292 arg = 1.0
293 if arg < -1.0:
294 arg = -1.0
295 return 180.0 * np.arccos(arg) / np.pi
297 def dih(n1, n2, n3, n4):
298 # vector 0->1, 1->2, 2->3 and their normalized cross products:
299 a = R[n2] - R[n1]
300 b = R[n3] - R[n2]
301 c = R[n4] - R[n3]
302 bxa = np.cross(b, a)
303 bxa /= np.sqrt(np.vdot(bxa, bxa))
304 cxb = np.cross(c, b)
305 cxb /= np.sqrt(np.vdot(cxb, cxb))
306 angle = np.vdot(bxa, cxb)
307 # check for numerical trouble due to finite precision:
308 if angle < -1:
309 angle = -1
310 if angle > 1:
311 angle = 1
312 angle = np.arccos(angle)
313 if np.vdot(bxa, c) > 0:
314 angle = 2 * np.pi - angle
315 return angle * 180.0 / np.pi
317 # get number of mobile atoms for temperature calculation
318 E = np.array([self.get_energy(atoms) for atoms in self])
320 s = 0.0
322 # Namespace for eval:
323 ns = {'E': E,
324 'd': d, 'a': a, 'dih': dih}
326 data = []
327 for i in range(nimages):
328 ns['i'] = i
329 ns['s'] = s
330 ns['R'] = R = self[i].get_positions()
331 ns['V'] = self[i].get_velocities()
332 F = self.get_forces(self[i])
333 if F is not None:
334 ns['F'] = F
335 ns['A'] = self[i].get_cell()
336 ns['M'] = self[i].get_masses()
337 # XXX askhl verify:
338 dynamic = self.get_dynamic(self[i])
339 if F is not None:
340 ns['f'] = f = ((F * dynamic[:, None])**2).sum(1)**.5
341 ns['fmax'] = max(f)
342 ns['fave'] = f.mean()
343 ns['epot'] = epot = E[i]
344 ns['ekin'] = ekin = self[i].get_kinetic_energy()
345 ns['e'] = epot + ekin
346 ndynamic = dynamic.sum()
347 if ndynamic > 0:
348 ns['T'] = 2.0 * ekin / (3.0 * ndynamic * units.kB)
349 data = eval(code, ns)
350 if i == 0:
351 nvariables = len(data)
352 xy = np.empty((nvariables, nimages))
353 xy[:, i] = data
354 if i + 1 < nimages and not self.have_varying_species:
355 dR = find_mic(self[i + 1].positions - R, self[i].get_cell(),
356 self[i].get_pbc())[0]
357 s += sqrt((dR**2).sum())
358 return xy
360 def write(self, filename, rotations='', bbox=None,
361 **kwargs):
362 # XXX We should show the unit cell whenever there is one
363 indices = range(len(self))
364 p = filename.rfind('@')
365 if p != -1:
366 try:
367 slice = string2index(filename[p + 1:])
368 except ValueError:
369 pass
370 else:
371 indices = indices[slice]
372 filename = filename[:p]
373 if isinstance(indices, int):
374 indices = [indices]
376 images = [self.get_atoms(i) for i in indices]
377 if len(filename) > 4 and filename[-4:] in ['.eps', '.png', '.pov']:
378 write(filename, images,
379 rotation=rotations,
380 bbox=bbox, **kwargs)
381 else:
382 write(filename, images, **kwargs)
384 def get_atoms(self, frame, remove_hidden=False):
385 atoms = self[frame]
386 try:
387 E = atoms.get_potential_energy()
388 except RuntimeError:
389 E = None
390 try:
391 F = atoms.get_forces()
392 except RuntimeError:
393 F = None
395 # Remove hidden atoms if applicable
396 if remove_hidden:
397 atoms = atoms[self.visible]
398 if F is not None:
399 F = F[self.visible]
400 atoms.calc = SinglePointCalculator(atoms, energy=E, forces=F)
401 return atoms
403 def delete(self, i):
404 self.images.pop(i)
405 self.filenames.pop(i)
406 self.initialize(self.images, self.filenames)