Coverage for /builds/ase/ase/ase/io/utils.py: 89.09%
330 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
2from itertools import islice
3from typing import IO
5import numpy as np
7from ase.data import atomic_numbers, covalent_radii
8from ase.data.colors import jmol_colors as default_colors
9from ase.io.formats import string2index
10from ase.utils import irotate, rotate
13def normalize(a):
14 return np.array(a) / np.linalg.norm(a)
17def complete_camera_vectors(look=None, up=None, right=None):
18 """Creates the camera (or look) basis vectors from user input and
19 will autocomplete missing vector or non-orthogonal vectors using dot
20 products. The look direction will be maintained, up direction has higher
21 priority than right direction"""
23 # ensure good input
24 if look is not None:
25 assert len(look) == 3
26 l = np.array(look)
28 if up is not None:
29 assert len(up) == 3
30 u = np.array(up)
32 if right is not None:
33 assert len(right) == 3
34 r = np.array(right)
36 if look is not None and up is not None:
37 r = normalize(np.cross(l, u))
38 u = normalize(np.cross(r, l)) # ensures complete perpendicularity
39 l = normalize(np.cross(u, r))
40 elif look is not None and right is not None:
41 u = normalize(np.cross(r, l))
42 r = normalize(np.cross(l, u)) # ensures complete perpendicularity
43 l = normalize(np.cross(u, r))
44 elif up is not None and right is not None:
45 l = normalize(np.cross(u, r))
46 r = normalize(np.cross(l, u)) # ensures complete perpendicularity
47 u = normalize(np.cross(r, u))
48 else:
49 raise ValueError('''At least two camera vectors of <look>, <up>,
50 or <right> must be specified''')
51 return l, u, r
54def get_cell_vertex_points(cell, disp=(0.0, 0.0, 0.0)):
55 """Returns 8x3 list of the cell vertex coordinates"""
56 cell_vertices = np.empty((2, 2, 2, 3))
57 displacement = np.array(disp)
58 for c1 in range(2):
59 for c2 in range(2):
60 for c3 in range(2):
61 cell_vertices[c1, c2, c3] = [c1, c2, c3] @ cell + displacement
62 cell_vertices.shape = (8, 3)
63 return cell_vertices
66def update_line_order_for_atoms(L, T, D, atoms, radii):
67 # why/how does this happen before the camera rotation???
68 R = atoms.get_positions()
69 r2 = radii**2
70 for n in range(len(L)):
71 d = D[T[n]]
72 if ((((R - L[n] - d)**2).sum(1) < r2) &
73 (((R - L[n] + d)**2).sum(1) < r2)).any():
74 T[n] = -1
75 return T
78def combine_bboxes(bbox_a, bbox_b):
79 """Combines bboxes using their extrema"""
80 bbox_low = np.minimum(bbox_a[0], bbox_b[0])
81 bbox_high = np.maximum(bbox_a[1], bbox_b[1])
82 return np.array([bbox_low, bbox_high])
85def has_cell(atoms):
86 return atoms.cell.rank > 0
89HIDE = 0
90SHOW_CELL = 1
91SHOW_CELL_AND_FIT_TO_ALL = 2
92SHOW_CELL_AND_FIT_TO_CELL = 3
95class PlottingVariables:
96 # removed writer - self
97 def __init__(self, atoms, rotation='', show_unit_cell=2,
98 radii=None, bbox=None, colors=None, scale=20,
99 maxwidth=500, extra_offset=(0., 0.),
100 auto_bbox_size=1.05,
101 auto_image_plane_z='front_all',
102 ):
104 assert show_unit_cell in (0, 1, 2, 3)
105 """Handles camera/paper space transformations used for rendering, 2D
106 plots, ...and a few legacy features. after camera rotations, the image
107 plane is set to the front of structure.
109 atoms: Atoms object
110 The Atoms object to render/plot.
112 rotation: string or 3x3 matrix
113 Controls camera rotation. Can be a string with euler angles in
114 degrees like '45x, 90y, 0z' or a rotation matrix.
115 (defaults to '0x, 0y, 0z')
117 show_unit_cell: int 0, 1, 2, or 3
118 0 cell is not shown, 1 cell is shown, 2 cell is shown and bounding
119 box is computed to fit atoms and cell, 3 bounding box is fixed to
120 cell only. (default 2)
122 radii: list of floats
123 a list of atomic radii for the atoms. (default None)
125 bbox: list of four floats
126 Allows explicit control of the image plane bounding box in the form
127 (xlo, ylo, xhi, yhi) where x and y are the horizontal and vertical
128 axes of the image plane. The units are in atomic coordinates without
129 the paperspace scale factor. (defaults to None the automatic
130 bounding box is used)
132 colors : a list of RGB color triples
133 a list of the RGB color triples for each atom. (default None, uses
134 Jmol colors)
136 scale: float
137 The ratio between the image plane units and atomic units, e.g.
138 Angstroms per cm. (default 20.0)
140 maxwidth: float
141 Limits the width of the image plane. (why?) Uses paperspace units.
142 (default 500)
144 extra_offset: (float, float)
145 Translates the image center in the image plane by (x,y) where x and
146 y are the horizontal and vertical shift distances, respectively.
147 (default (0.0, 0.0)) should only be used for small tweaks to the
148 automatically fit image plane
150 auto_bbox_size: float
151 Controls the padding given to the bounding box in the image plane.
152 With auto_bbox_size=1.0 the structure touches the edges of the
153 image. auto_bbox_size>1.0 gives whitespace padding. (default 1.05)
155 auto_image_plane_z: string ('front_all', 'front_auto', 'legacy')
156 After a camera rotation, controls where to put camera image plane
157 relative to the atoms and cell. 'front_all' puts everything in front
158 of the camera. 'front_auto' sets the image plane location to
159 respect the show_unit_cell option so that the atoms or cell can be
160 ignored when setting the image plane. 'legacy' leaves the image
161 plane passing through the origin for backwards compatibility.
162 (default: 'front_all')
163 """
165 self.show_unit_cell = show_unit_cell
166 self.numbers = atoms.get_atomic_numbers()
167 self.maxwidth = maxwidth
168 self.atoms = atoms
169 # not used in PlottingVariables, keeping for legacy
170 self.natoms = len(atoms)
172 self.auto_bbox_size = auto_bbox_size
173 self.auto_image_plane_z = auto_image_plane_z
174 self.offset = np.zeros(3)
175 self.extra_offset = np.array(extra_offset)
177 self.constraints = atoms.constraints
178 # extension for partial occupancies
179 self.frac_occ = False
180 self.tags = None
181 self.occs = None
183 if 'occupancy' in atoms.info:
184 self.occs = atoms.info['occupancy']
185 self.tags = atoms.get_tags()
186 self.frac_occ = True
188 # colors
189 self.colors = colors
190 if colors is None:
191 ncolors = len(default_colors)
192 self.colors = default_colors[self.numbers.clip(max=ncolors - 1)]
194 # radius
195 if radii is None:
196 radii = covalent_radii[self.numbers]
197 elif isinstance(radii, float):
198 radii = covalent_radii[self.numbers] * radii
199 else:
200 radii = np.array(radii)
202 self.radii = radii # radius in Angstroms
203 self.scale = scale # Angstroms per cm
205 self.set_rotation(rotation)
206 self.update_image_plane_offset_and_size_from_structure(bbox=bbox)
208 def to_dict(self):
209 out = {
210 'bbox': self.get_bbox(),
211 'rotation': self.rotation,
212 'scale': self.scale,
213 'colors': self.colors}
214 return out
216 @property
217 def d(self):
218 # XXX hopefully this can be deprecated someday.
219 """Returns paperspace diameters for scale and radii lists"""
220 return 2 * self.scale * self.radii
222 def set_rotation(self, rotation):
223 if rotation is not None:
224 if isinstance(rotation, str):
225 rotation = rotate(rotation)
226 self.rotation = rotation
227 self.update_patch_and_line_vars()
229 def update_image_plane_offset_and_size_from_structure(self, bbox=None):
230 """Updates image size to fit structure according to show_unit_cell
231 if bbox=None. Otherwise, sets the image size from bbox. bbox is in the
232 image plane. Note that bbox format is (xlo, ylo, xhi, yhi) for
233 compatibility reasons the internal functions use (2,3)"""
235 # zero out the offset so it's not involved in the
236 # to_image_plane_positions() calculations which are used to calcucate
237 # the offset
238 self.offset = np.zeros(3)
240 # computing the bboxes in self.atoms here makes it easier to follow the
241 # various options selection/choices later
242 bbox_atoms = self.get_bbox_from_atoms(self.atoms, self.d / 2)
243 if has_cell(self.atoms):
244 cell = self.atoms.get_cell()
245 disp = self.atoms.get_celldisp().flatten()
246 bbox_cell = self.get_bbox_from_cell(cell, disp)
247 bbox_combined = combine_bboxes(bbox_atoms, bbox_cell)
248 else:
249 bbox_combined = bbox_atoms
251 # bbox_auto is the bbox that matches the show_unit_cell option
252 if has_cell(self.atoms) and self.show_unit_cell in (
253 SHOW_CELL_AND_FIT_TO_ALL, SHOW_CELL_AND_FIT_TO_CELL):
255 if self.show_unit_cell == SHOW_CELL_AND_FIT_TO_ALL:
256 bbox_auto = bbox_combined
257 else:
258 bbox_auto = bbox_cell
259 else:
260 bbox_auto = bbox_atoms
262 #
263 if bbox is None:
264 middle = (bbox_auto[0] + bbox_auto[1]) / 2
265 im_size = self.auto_bbox_size * (bbox_auto[1] - bbox_auto[0])
266 # should auto_bbox_size pad the z_heght via offset?
268 if im_size[0] > self.maxwidth:
269 rescale_factor = self.maxwidth / im_size[0]
270 im_size *= rescale_factor
271 self.scale *= rescale_factor
272 offset = middle - im_size / 2
273 else:
274 width = (bbox[2] - bbox[0]) * self.scale
275 height = (bbox[3] - bbox[1]) * self.scale
277 im_size = np.array([width, height, 0])
278 offset = np.array([bbox[0], bbox[1], 0]) * self.scale
280 # this section shifts the image plane up and down parallel to the look
281 # direction to match the legacy option, or to force it allways touch the
282 # front most objects regardless of the show_unit_cell setting
283 if self.auto_image_plane_z == 'front_all':
284 offset[2] = bbox_combined[1, 2] # highest z in image orientation
285 elif self.auto_image_plane_z == 'legacy':
286 offset[2] = 0
287 elif self.auto_image_plane_z == 'front_auto':
288 offset[2] = bbox_auto[1, 2]
289 else:
290 raise ValueError(
291 f'bad image plane setting {self.auto_image_plane_z!r}')
293 # since we are moving the origin in the image plane (camera coordinates)
294 self.offset += offset
296 # Previously, the picture size changed with extra_offset, This is very
297 # counter intuitive and seems like a bug. Leaving it commented out in
298 # case someone relying on this likely bug needs to revert it.
299 self.w = im_size[0] # + self.extra_offset[0]
300 self.h = im_size[1] # + self.extra_offset[1]
302 # allows extra_offset to be 2D or 3D
303 for i in range(len(self.extra_offset)):
304 self.offset[i] -= self.extra_offset[i]
306 # we have to update the arcane stuff after every camera update.
307 self.update_patch_and_line_vars()
309 def center_camera_on_position(self, pos, scaled_position=False):
310 if scaled_position:
311 pos = pos @ self.atoms.cell
312 im_pos = self.to_image_plane_positions(pos)
313 cam_pos = self.to_image_plane_positions(self.get_image_plane_center())
314 in_plane_shift = im_pos - cam_pos
315 self.offset[0:2] += in_plane_shift[0:2]
316 self.update_patch_and_line_vars()
318 def get_bbox(self):
319 xlo = self.offset[0]
320 ylo = self.offset[1]
321 xhi = xlo + self.w
322 yhi = ylo + self.h
323 return np.array([xlo, ylo, xhi, yhi]) / self.scale
325 def set_rotation_from_camera_directions(self,
326 look=None, up=None, right=None,
327 scaled_position=False):
329 if scaled_position:
330 if look is not None:
331 look = look @ self.atoms.cell
332 if right is not None:
333 right = right @ self.atoms.cell
334 if up is not None:
335 up = up @ self.atoms.cell
337 look, up, right = complete_camera_vectors(look, up, right)
339 rotation = np.zeros((3, 3))
340 rotation[:, 0] = right
341 rotation[:, 1] = up
342 rotation[:, 2] = -look
343 self.rotation = rotation
344 self.update_patch_and_line_vars()
346 def get_rotation_angles(self):
347 """Gets the rotation angles from the rotation matrix in the current
348 PlottingVariables object"""
349 return irotate(self.rotation)
351 def get_rotation_angles_string(self, digits=5):
352 fmt = '%.{:d}f'.format(digits)
353 angles = self.get_rotation_angles()
354 outstring = (fmt + 'x, ' + fmt + 'y, ' + fmt + 'z') % (angles)
355 return outstring
357 def update_patch_and_line_vars(self):
358 """Updates all the line and path stuff that is still inobvious, this
359 function should be deprecated if nobody can understand why it's features
360 exist."""
361 cell = self.atoms.get_cell()
362 disp = self.atoms.get_celldisp().flatten()
363 positions = self.atoms.get_positions()
365 if self.show_unit_cell in (
366 SHOW_CELL, SHOW_CELL_AND_FIT_TO_ALL, SHOW_CELL_AND_FIT_TO_CELL):
368 L, T, D = cell_to_lines(self, cell)
369 cell_verts_in_atom_coords = get_cell_vertex_points(cell, disp)
370 cell_vertices = self.to_image_plane_positions(
371 cell_verts_in_atom_coords)
372 T = update_line_order_for_atoms(L, T, D, self.atoms, self.radii)
373 # D are a positions in the image plane,
374 # not sure why it's setup like this
375 D = (self.to_image_plane_positions(D) + self.offset)[:, :2]
376 positions = np.concatenate((positions, L), axis=0)
377 else:
378 L = np.empty((0, 3))
379 T = None
380 D = None
381 cell_vertices = None
382 # just a rotations and scaling since offset is currently [0,0,0]
383 positions = self.to_image_plane_positions(positions)
384 self.positions = positions
385 # list of 2D cell points in the imageplane without the offset
386 self.D = D
387 # integers, probably z-order for lines?
388 self.T = T
389 self.cell_vertices = cell_vertices
391 # no displacement since it's a vector
392 cell_vec_im = self.scale * self.atoms.get_cell() @ self.rotation
393 self.cell = cell_vec_im
395 def to_image_plane_positions(self, positions):
396 """Converts atomic coordinates to image plane positions. The third
397 coordinate is distance above/below the image plane"""
398 im_positions = (positions @ self.rotation) * self.scale - self.offset
399 return im_positions
401 def to_atom_positions(self, im_positions):
402 """Converts image plane positions to atomic coordinates."""
403 positions = ((im_positions + self.offset) /
404 self.scale) @ self.rotation.T
405 return positions
407 def get_bbox_from_atoms(self, atoms, im_radii):
408 """Uses supplied atoms and radii to compute the bounding box of the
409 atoms in the image plane"""
410 im_positions = self.to_image_plane_positions(atoms.get_positions())
411 im_low = (im_positions - im_radii[:, None]).min(0)
412 im_high = (im_positions + im_radii[:, None]).max(0)
413 return np.array([im_low, im_high])
415 def get_bbox_from_cell(self, cell, disp=(0.0, 0.0, 0.0)):
416 """Uses supplied cell to compute the bounding box of the cell in the
417 image plane"""
418 displacement = np.array(disp)
419 cell_verts_in_atom_coords = get_cell_vertex_points(cell, displacement)
420 cell_vertices = self.to_image_plane_positions(cell_verts_in_atom_coords)
421 im_low = cell_vertices.min(0)
422 im_high = cell_vertices.max(0)
423 return np.array([im_low, im_high])
425 def get_image_plane_center(self):
426 return self.to_atom_positions(np.array([self.w / 2, self.h / 2, 0]))
428 def get_atom_direction(self, direction):
429 c0 = self.to_atom_positions([0, 0, 0]) # self.get_image_plane_center()
430 c1 = self.to_atom_positions(direction)
431 atom_direction = c1 - c0
432 return atom_direction / np.linalg.norm(atom_direction)
434 def get_camera_direction(self):
435 """Returns vector pointing away from camera toward atoms/cell in atomic
436 coordinates"""
437 return self.get_atom_direction([0, 0, -1])
439 def get_camera_up(self):
440 """Returns the image plane up direction in atomic coordinates"""
441 return self.get_atom_direction([0, 1, 0])
443 def get_camera_right(self):
444 """Returns the image plane right direction in atomic coordinates"""
445 return self.get_atom_direction([1, 0, 0])
448def cell_to_lines(writer, cell):
449 # XXX this needs to be updated for cell vectors that are zero.
450 # Cannot read the code though! (What are T and D? nn?)
451 nlines = 0
452 nsegments = []
453 for c in range(3):
454 d = np.sqrt((cell[c]**2).sum())
455 n = max(2, int(d / 0.3))
456 nsegments.append(n)
457 nlines += 4 * n
459 positions = np.empty((nlines, 3))
460 T = np.empty(nlines, int)
461 D = np.zeros((3, 3))
463 n1 = 0
464 for c in range(3):
465 n = nsegments[c]
466 dd = cell[c] / (4 * n - 2)
467 D[c] = dd
468 P = np.arange(1, 4 * n + 1, 4)[:, None] * dd
469 T[n1:] = c
470 for i, j in [(0, 0), (0, 1), (1, 0), (1, 1)]:
471 n2 = n1 + n
472 positions[n1:n2] = P + i * cell[c - 2] + j * cell[c - 1]
473 n1 = n2
475 return positions, T, D
478def make_patch_list(writer):
479 from matplotlib.patches import Circle, PathPatch, Wedge
480 from matplotlib.path import Path
482 indices = writer.positions[:, 2].argsort()
483 patch_list = []
484 for a in indices:
485 xy = writer.positions[a, :2]
486 if a < writer.natoms:
487 r = writer.d[a] / 2
488 if writer.frac_occ:
489 site_occ = writer.occs[str(writer.tags[a])]
490 # first an empty circle if a site is not fully occupied
491 if (np.sum([v for v in site_occ.values()])) < 1.0:
492 # fill with white
493 fill = '#ffffff'
494 patch = Circle(xy, r, facecolor=fill,
495 edgecolor='black')
496 patch_list.append(patch)
498 start = 0
499 # start with the dominant species
500 for sym, occ in sorted(site_occ.items(),
501 key=lambda x: x[1],
502 reverse=True):
503 if np.round(occ, decimals=4) == 1.0:
504 patch = Circle(xy, r, facecolor=writer.colors[a],
505 edgecolor='black')
506 patch_list.append(patch)
507 else:
508 # jmol colors for the moment
509 extent = 360. * occ
510 patch = Wedge(
511 xy, r, start, start + extent,
512 facecolor=default_colors[atomic_numbers[sym]],
513 edgecolor='black')
514 patch_list.append(patch)
515 start += extent
517 else:
518 if ((xy[1] + r > 0) and (xy[1] - r < writer.h) and
519 (xy[0] + r > 0) and (xy[0] - r < writer.w)):
520 patch = Circle(xy, r, facecolor=writer.colors[a],
521 edgecolor='black')
522 patch_list.append(patch)
523 else:
524 a -= writer.natoms
525 c = writer.T[a]
526 if c != -1:
527 hxy = writer.D[c]
528 patch = PathPatch(Path((xy + hxy, xy - hxy)))
529 patch_list.append(patch)
530 return patch_list
533class ImageChunk:
534 """Base Class for a file chunk which contains enough information to
535 reconstruct an atoms object."""
537 def build(self, **kwargs):
538 """Construct the atoms object from the stored information,
539 and return it"""
542class ImageIterator:
543 """Iterate over chunks, to return the corresponding Atoms objects.
544 Will only build the atoms objects which corresponds to the requested
545 indices when called.
546 Assumes ``ichunks`` is in iterator, which returns ``ImageChunk``
547 type objects. See extxyz.py:iread_xyz as an example.
548 """
550 def __init__(self, ichunks):
551 self.ichunks = ichunks
553 def __call__(self, fd: IO, index=None, **kwargs):
554 if isinstance(index, str):
555 index = string2index(index)
557 if index is None or index == ':':
558 index = slice(None, None, None)
560 if not isinstance(index, (slice, str)):
561 index = slice(index, (index + 1) or None)
563 for chunk in self._getslice(fd, index):
564 yield chunk.build(**kwargs)
566 def _getslice(self, fd: IO, indices: slice):
567 try:
568 iterator = islice(self.ichunks(fd),
569 indices.start, indices.stop,
570 indices.step)
571 except ValueError:
572 # Negative indices. Go through the whole thing to get the length,
573 # which allows us to evaluate the slice, and then read it again
574 if not hasattr(fd, 'seekable') or not fd.seekable():
575 raise ValueError('Negative indices only supported for '
576 'seekable streams')
578 startpos = fd.tell()
579 nchunks = 0
580 for _ in self.ichunks(fd):
581 nchunks += 1
582 fd.seek(startpos)
583 indices_tuple = indices.indices(nchunks)
584 iterator = islice(self.ichunks(fd), *indices_tuple)
585 return iterator
588def verify_cell_for_export(cell, check_orthorhombric=True):
589 """Function to verify if the cell size is defined and if the cell is
591 Parameters:
593 cell: cell object
594 cell to be checked.
596 check_orthorhombric: bool
597 If True, check if the cell is orthorhombric, raise an ``ValueError`` if
598 the cell is orthorhombric. If False, doesn't check if the cell is
599 orthorhombric.
601 Raise a ``ValueError`` if the cell if not suitable for export to mustem xtl
602 file or prismatic/computem xyz format:
603 - if cell is not orthorhombic (only when check_orthorhombric=True)
604 - if cell size is not defined
605 """
607 if check_orthorhombric and not cell.orthorhombic:
608 raise ValueError('To export to this format, the cell needs to be '
609 'orthorhombic.')
610 if cell.rank < 3:
611 raise ValueError('To export to this format, the cell size needs '
612 'to be set: current cell is {}.'.format(cell))
615def verify_dictionary(atoms, dictionary, dictionary_name):
616 """
617 Verify a dictionary have a key for each symbol present in the atoms object.
619 Parameters:
621 dictionary: dict
622 Dictionary to be checked.
625 dictionary_name: dict
626 Name of the dictionary to be displayed in the error message.
628 cell: cell object
629 cell to be checked.
632 Raise a ``ValueError`` if the key doesn't match the atoms present in the
633 cell.
634 """
635 # Check if we have enough key
636 for key in set(atoms.symbols):
637 if key not in dictionary:
638 raise ValueError('Missing the {} key in the `{}` dictionary.'
639 ''.format(key, dictionary_name))
642def segment_list(data, segment_size):
643 """Segments a list into sublists of a specified size."""
644 return [data[i:i + segment_size] for i in range(0, len(data), segment_size)]