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