wmaee.core.visualize

  1# coding: utf-8
  2# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
  3# Distributed under the terms of "New BSD License", see the LICENSE file.
  4
  5import warnings
  6import numpy as np
  7from ase.io import Trajectory
  8from matplotlib.colors import rgb2hex
  9from scipy.interpolate import interp1d
 10from wmaee.core.requirements import requires
 11
 12__author__ = "Joerg Neugebauer, Sudarsan Surendralal"
 13__copyright__ = (
 14    "Copyright 2021, Max-Planck-Institut für Eisenforschung GmbH - "
 15    "Computational Materials Design (CM) Department"
 16)
 17__version__ = "1.0"
 18__maintainer__ = "Sudarsan Surendralal"
 19__email__ = "surendralal@mpie.de"
 20__status__ = "production"
 21__date__ = "Sep 1, 2017"
 22
 23
 24@requires("nglview")
 25def plot3d(
 26        structure,
 27        mode="NGLview",
 28        show_cell=True,
 29        show_axes=True,
 30        camera="orthographic",
 31        spacefill=True,
 32        particle_size=1.0,
 33        select_atoms=None,
 34        background="white",
 35        color_scheme=None,
 36        colors=None,
 37        scalar_field=None,
 38        scalar_start=None,
 39        scalar_end=None,
 40        scalar_cmap=None,
 41        vector_field=None,
 42        vector_color=None,
 43        magnetic_moments=False,
 44        view_plane=np.array([0, 0, 1]),
 45        distance_from_camera=1.0,
 46        opacity=1.0,
 47):
 48    """
 49    Plot3d relies on NGLView or plotly to visualize atomic structures. Here, we construct a string in the "protein database"
 50
 51    The final widget is returned. If it is assigned to a variable, the visualization is suppressed until that
 52    variable is evaluated, and in the meantime more NGL operations can be applied to it to modify the visualization.
 53
 54    Args:
 55        mode (str): `NGLView`, `plotly` or `ase`
 56        show_cell (bool): Whether or not to show the frame. (Default is True.)
 57        show_axes (bool): Whether or not to show xyz axes. (Default is True.)
 58        camera (str): 'perspective' or 'orthographic'. (Default is 'perspective'.)
 59        spacefill (bool): Whether to use a space-filling or ball-and-stick representation. (Default is True, use
 60            space-filling atoms.)
 61        particle_size (float): Size of the particles. (Default is 1.)
 62        select_atoms (numpy.ndarray): Indices of atoms to show, either as integers or a boolean array mask.
 63            (Default is None, show all atoms.)
 64        background (str): Background color. (Default is 'white'.)
 65        color_scheme (str): NGLView color scheme to use. (Default is None, color by element.)
 66        colors (numpy.ndarray): A per-atom array of HTML color names or hex color codes to use for atomic colors.
 67            (Default is None, use coloring scheme.)
 68        scalar_field (numpy.ndarray): Color each atom according to the array value (Default is None, use coloring
 69            scheme.)
 70        scalar_start (float): The scalar value to be mapped onto the low end of the color map (lower values are
 71            clipped). (Default is None, use the minimum value in `scalar_field`.)
 72        scalar_end (float): The scalar value to be mapped onto the high end of the color map (higher values are
 73            clipped). (Default is None, use the maximum value in `scalar_field`.)
 74        scalar_cmap (matplotlib.cm): The colormap to use. (Default is None, giving a blue-red divergent map.)
 75        vector_field (numpy.ndarray): Add vectors (3 values) originating at each atom. (Default is None, no
 76            vectors.)
 77        vector_color (numpy.ndarray): Colors for the vectors (only available with vector_field). (Default is None,
 78            vectors are colored by their direction.)
 79        magnetic_moments (bool): Plot magnetic moments as 'scalar_field' or 'vector_field'.
 80        view_plane (numpy.ndarray): A Nx3-array (N = 1,2,3); the first 3d-component of the array specifies
 81            which plane of the system to view (for example, [1, 0, 0], [1, 1, 0] or the [1, 1, 1] planes), the
 82            second 3d-component (if specified, otherwise [1, 0, 0]) gives the horizontal direction, and the third
 83            component (if specified) is the vertical component, which is ignored and calculated internally. The
 84            orthonormality of the orientation is internally ensured, and therefore is not required in the function
 85            call. (Default is np.array([0, 0, 1]), which is view normal to the x-y plane.)
 86        distance_from_camera (float): Distance of the camera from the structure. Higher = farther away.
 87            (Default is 14, which also seems to be the NGLView default value.)
 88
 89        Possible NGLView color schemes:
 90          " ", "picking", "random", "uniform", "atomindex", "residueindex",
 91          "chainindex", "modelindex", "sstruc", "element", "resname", "bfactor",
 92          "hydrophobicity", "value", "volume", "occupancy"
 93
 94    Returns:
 95        (nglview.NGLWidget): The NGLView widget itself, which can be operated on further or viewed as-is.
 96
 97    Warnings:
 98        * Many features only work with space-filling atoms (e.g. coloring by a scalar field).
 99        * The colour interpretation of some hex codes is weird, e.g. 'green'.
100    """
101    if mode == "NGLview":
102        return _plot3d(
103            structure=structure,
104            show_cell=show_cell,
105            show_axes=show_axes,
106            camera=camera,
107            spacefill=spacefill,
108            particle_size=particle_size,
109            select_atoms=select_atoms,
110            background=background,
111            color_scheme=color_scheme,
112            colors=colors,
113            scalar_field=scalar_field,
114            scalar_start=scalar_start,
115            scalar_end=scalar_end,
116            scalar_cmap=scalar_cmap,
117            vector_field=vector_field,
118            vector_color=vector_color,
119            magnetic_moments=magnetic_moments,
120            view_plane=view_plane,
121            distance_from_camera=distance_from_camera,
122        )
123    elif mode == "ase":
124        return _plot3d_ase(
125            structure=structure,
126            show_cell=show_cell,
127            show_axes=show_axes,
128            camera=camera,
129            spacefill=spacefill,
130            particle_size=particle_size,
131            background=background,
132            color_scheme=color_scheme,
133        )
134    else:
135        raise ValueError("plot method not recognized")
136
137
138def _plot3d(
139        structure,
140        show_cell=True,
141        show_axes=True,
142        camera="orthographic",
143        spacefill=True,
144        particle_size=1.0,
145        select_atoms=None,
146        background="white",
147        color_scheme=None,
148        colors=None,
149        scalar_field=None,
150        scalar_start=None,
151        scalar_end=None,
152        scalar_cmap=None,
153        vector_field=None,
154        vector_color=None,
155        magnetic_moments=False,
156        view_plane=np.array([0, 0, 1]),
157        distance_from_camera=1.0,
158):
159    """
160    Plot3d relies on NGLView to visualize atomic structures. Here, we construct a string in the "protein database"
161    ("pdb") format, then turn it into an NGLView "structure". PDB is a white-space sensitive format, so the
162    string snippets are carefully formatted.
163
164    The final widget is returned. If it is assigned to a variable, the visualization is suppressed until that
165    variable is evaluated, and in the meantime more NGL operations can be applied to it to modify the visualization.
166
167    Args:
168        show_cell (bool): Whether or not to show the frame. (Default is True.)
169        show_axes (bool): Whether or not to show xyz axes. (Default is True.)
170        camera (str): 'perspective' or 'orthographic'. (Default is 'perspective'.)
171        spacefill (bool): Whether to use a space-filling or ball-and-stick representation. (Default is True, use
172            space-filling atoms.)
173        particle_size (float): Size of the particles. (Default is 1.)
174        select_atoms (numpy.ndarray): Indices of atoms to show, either as integers or a boolean array mask.
175            (Default is None, show all atoms.)
176        background (str): Background color. (Default is 'white'.)
177        color_scheme (str): NGLView color scheme to use. (Default is None, color by element.)
178        colors (numpy.ndarray): A per-atom array of HTML color names or hex color codes to use for atomic colors.
179            (Default is None, use coloring scheme.)
180        scalar_field (numpy.ndarray): Color each atom according to the array value (Default is None, use coloring
181            scheme.)
182        scalar_start (float): The scalar value to be mapped onto the low end of the color map (lower values are
183            clipped). (Default is None, use the minimum value in `scalar_field`.)
184        scalar_end (float): The scalar value to be mapped onto the high end of the color map (higher values are
185            clipped). (Default is None, use the maximum value in `scalar_field`.)
186        scalar_cmap (matplotlib.cm): The colormap to use. (Default is None, giving a blue-red divergent map.)
187        vector_field (numpy.ndarray): Add vectors (3 values) originating at each atom. (Default is None, no
188            vectors.)
189        vector_color (numpy.ndarray): Colors for the vectors (only available with vector_field). (Default is None,
190            vectors are colored by their direction.)
191        magnetic_moments (bool): Plot magnetic moments as 'scalar_field' or 'vector_field'.
192        view_plane (numpy.ndarray): A Nx3-array (N = 1,2,3); the first 3d-component of the array specifies
193            which plane of the system to view (for example, [1, 0, 0], [1, 1, 0] or the [1, 1, 1] planes), the
194            second 3d-component (if specified, otherwise [1, 0, 0]) gives the horizontal direction, and the third
195            component (if specified) is the vertical component, which is ignored and calculated internally. The
196            orthonormality of the orientation is internally ensured, and therefore is not required in the function
197            call. (Default is np.array([0, 0, 1]), which is view normal to the x-y plane.)
198        distance_from_camera (float): Distance of the camera from the structure. Higher = farther away.
199            (Default is 14, which also seems to be the NGLView default value.)
200
201        Possible NGLView color schemes:
202          " ", "picking", "random", "uniform", "atomindex", "residueindex",
203          "chainindex", "modelindex", "sstruc", "element", "resname", "bfactor",
204          "hydrophobicity", "value", "volume", "occupancy"
205
206    Returns:
207        (nglview.NGLWidget): The NGLView widget itself, which can be operated on further or viewed as-is.
208
209    Warnings:
210        * Many features only work with space-filling atoms (e.g. coloring by a scalar field).
211        * The colour interpretation of some hex codes is weird, e.g. 'green'.
212    """
213    try:  # If the graphical packages are not available, the GUI will not work.
214        import nglview
215    except ImportError:
216        raise ImportError(
217            "The package nglview needs to be installed for the plot3d() function!"
218        )
219
220    if (
221            magnetic_moments is True
222            and np.sum(np.abs(structure.get_initial_magnetic_moments())) > 0
223    ):
224        if len(structure.get_initial_magnetic_moments().shape) == 1:
225            scalar_field = structure.get_initial_magnetic_moments()
226        else:
227            vector_field = structure.get_initial_magnetic_moments()
228
229    elements = structure.get_chemical_symbols()
230    atomic_numbers = structure.numbers
231    positions = structure.positions
232
233    # If `select_atoms` was given, visualize only a subset of the `parent_basis`
234    if select_atoms is not None:
235        select_atoms = np.array(select_atoms, dtype=int)
236        elements = elements[select_atoms]
237        atomic_numbers = atomic_numbers[select_atoms]
238        positions = positions[select_atoms]
239        if colors is not None:
240            colors = np.array(colors)
241            colors = colors[select_atoms]
242        if scalar_field is not None:
243            scalar_field = np.array(scalar_field)
244            scalar_field = scalar_field[select_atoms]
245        if vector_field is not None:
246            vector_field = np.array(vector_field)
247            vector_field = vector_field[select_atoms]
248        if vector_color is not None:
249            vector_color = np.array(vector_color)
250            vector_color = vector_color[select_atoms]
251
252    # Write the nglview protein-database-formatted string
253    struct = nglview.TextStructure(
254        _ngl_write_structure(elements, positions, structure.cell)
255    )
256
257    # Parse the string into the displayable widget
258    view = nglview.NGLWidget(struct)
259
260    if spacefill:
261        # Color by scheme
262        if color_scheme is not None:
263            if colors is not None:
264                warnings.warn("`color_scheme` is overriding `colors`")
265            if scalar_field is not None:
266                warnings.warn("`color_scheme` is overriding `scalar_field`")
267            view = _add_colorscheme_spacefill(
268                view, elements, atomic_numbers, particle_size, color_scheme
269            )
270        # Color by per-atom colors
271        elif colors is not None:
272            if scalar_field is not None:
273                warnings.warn("`colors` is overriding `scalar_field`")
274            view = _add_custom_color_spacefill(
275                view, atomic_numbers, particle_size, colors
276            )
277        # Color by per-atom scalars
278        elif scalar_field is not None:  # Color by per-atom scalars
279            colors = _scalars_to_hex_colors(
280                scalar_field, scalar_start, scalar_end, scalar_cmap
281            )
282            view = _add_custom_color_spacefill(
283                view, atomic_numbers, particle_size, colors
284            )
285        # Color by element
286        else:
287            view = _add_colorscheme_spacefill(
288                view, elements, atomic_numbers, particle_size
289            )
290        view.remove_ball_and_stick()
291    else:
292        view.add_ball_and_stick()
293
294    if show_cell:
295        if structure.cell is not None:
296            if all(np.max(structure.cell, axis=0) > 1e-2):
297                view.add_unitcell()
298
299    if vector_color is None and vector_field is not None:
300        vector_color = (
301                0.5
302                * np.array(vector_field)
303                / np.linalg.norm(vector_field, axis=-1)[:, np.newaxis]
304                + 0.5
305        )
306    elif (
307            vector_field is not None and vector_field is not None
308    ):  # WARNING: There must be a bug here...
309        try:
310            if vector_color.shape != np.ones((len(structure), 3)).shape:
311                vector_color = np.outer(
312                    np.ones(len(structure)),
313                    vector_color / np.linalg.norm(vector_color),
314                )
315        except AttributeError:
316            vector_color = np.ones((len(structure), 3)) * vector_color
317
318    if vector_field is not None:
319        for arr, pos, col in zip(vector_field, positions, vector_color):
320            view.shape.add_arrow(list(pos), list(pos + arr), list(col), 0.2)
321
322    if show_axes:  # Add axes
323        axes_origin = -np.ones(3)
324        arrow_radius = 0.1
325        text_size = 1
326        text_color = [0, 0, 0]
327        arrow_names = ["x", "y", "z"]
328
329        for n in [0, 1, 2]:
330            start = list(axes_origin)
331            shift = np.zeros(3)
332            shift[n] = 1
333            end = list(start + shift)
334            color = list(shift)
335            # We cast as list to avoid JSON warnings
336            view.shape.add_arrow(start, end, color, arrow_radius)
337            view.shape.add_text(end, text_color, text_size, arrow_names[n])
338
339    if camera != "perspective" and camera != "orthographic":
340        warnings.warn(
341            "Only perspective or orthographic is (likely to be) permitted for camera"
342        )
343
344    view.camera = camera
345    view.background = background
346
347    orientation = _get_flattened_orientation(
348        view_plane=view_plane, distance_from_camera=distance_from_camera * 14
349    )
350    view.control.orient(orientation)
351
352    return view
353
354
355def _plot3d_ase(
356        structure,
357        spacefill=True,
358        show_cell=True,
359        camera="perspective",
360        particle_size=0.5,
361        background="white",
362        color_scheme="element",
363        show_axes=True,
364):
365    """
366    Possible color schemes:
367      " ", "picking", "random", "uniform", "atomindex", "residueindex",
368      "chainindex", "modelindex", "sstruc", "element", "resname", "bfactor",
369      "hydrophobicity", "value", "volume", "occupancy"
370    Returns:
371    """
372    try:  # If the graphical packages are not available, the GUI will not work.
373        import nglview
374    except ImportError:
375        raise ImportError(
376            "The package nglview needs to be installed for the plot3d() function!"
377        )
378    # Always visualize the parent basis
379    view = nglview.show_ase(structure)
380    if spacefill:
381        view.add_spacefill(
382            radius_type="vdw", color_scheme=color_scheme, radius=particle_size
383        )
384        # view.add_spacefill(radius=1.0)
385        view.remove_ball_and_stick()
386    else:
387        view.add_ball_and_stick()
388    if show_cell:
389        if structure.cell is not None:
390            if all(np.max(structure.cell, axis=0) > 1e-2):
391                view.add_unitcell()
392    if show_axes:
393        view.shape.add_arrow([-2, -2, -2], [2, -2, -2], [1, 0, 0], 0.5)
394        view.shape.add_arrow([-2, -2, -2], [-2, 2, -2], [0, 1, 0], 0.5)
395        view.shape.add_arrow([-2, -2, -2], [-2, -2, 2], [0, 0, 1], 0.5)
396    if camera != "perspective" and camera != "orthographic":
397        print("Only perspective or orthographic is permitted")
398        return None
399    view.camera = camera
400    view.background = background
401    return view
402
403
404def _ngl_write_cell(a1, a2, a3, f1=90, f2=90, f3=90):
405    """
406    Writes a PDB-formatted line to represent the simulation cell.
407
408    Args:
409        a1, a2, a3 (float): Lengths of the cell vectors.
410        f1, f2, f3 (float): Angles between the cell vectors (which angles exactly?) (in degrees).
411
412    Returns:
413        (str): The line defining the cell in PDB format.
414    """
415    return "CRYST1 {:8.3f} {:8.3f} {:8.3f} {:6.2f} {:6.2f} {:6.2f} P 1\n".format(
416        a1, a2, a3, f1, f2, f3
417    )
418
419
420def _ngl_write_atom(
421        num,
422        species,
423        x,
424        y,
425        z,
426        group=None,
427        num2=None,
428        occupancy=1.0,
429        temperature_factor=0.0,
430):
431    """
432    Writes a PDB-formatted line to represent an atom.
433
434    Args:
435        num (int): Atomic index.
436        species (str): Elemental species.
437        x, y, z (float): Cartesian coordinates of the atom.
438        group (str): A...group name? (Default is None, repeat elemental species.)
439        num2 (int): An "alternate" index. (Don't ask me...) (Default is None, repeat first number.)
440        occupancy (float): PDB occupancy parameter. (Default is 1.)
441        temperature_factor (float): PDB temperature factor parameter. (Default is 0.
442
443    Returns:
444        (str): The line defining an atom in PDB format
445
446    Warnings:
447        * The [PDB docs](https://www.cgl.ucsf.edu/chimera/docs/UsersGuide/tutorials/pdbintro.html) indicate that
448            the xyz coordinates might need to be in some sort of orthogonal basis. If you have weird behaviour,
449            this might be a good place to investigate.
450    """
451    if group is None:
452        group = species
453    if num2 is None:
454        num2 = num
455    return "ATOM {:>6} {:>4} {:>4} {:>5} {:10.3f} {:7.3f} {:7.3f} {:5.2f} {:5.2f} {:>11} \n".format(
456        num, species, group, num2, x, y, z, occupancy, temperature_factor, species
457    )
458
459
460def _ngl_write_structure(elements, positions, cell):
461    """
462    Turns structure information into a NGLView-readable protein-database-formatted string.
463
464    Args:
465        elements (numpy.ndarray/list): Element symbol for each atom.
466        positions (numpy.ndarray/list): Vector of Cartesian atom positions.
467        cell (numpy.ndarray/list): Simulation cell Bravais matrix.
468
469    Returns:
470        (str): The PDB-formatted representation of the structure.
471    """
472    from ase.geometry import cell_to_cellpar, cellpar_to_cell
473
474    if cell is None or any(np.max(cell, axis=0) < 1e-2):
475        # Define a dummy cell if it doesn't exist (eg. for clusters)
476        max_pos = np.max(positions, axis=0) - np.min(positions, axis=0)
477        max_pos[np.abs(max_pos) < 1e-2] = 10
478        cell = np.eye(3) * max_pos
479    cellpar = cell_to_cellpar(cell)
480    exportedcell = cellpar_to_cell(cellpar)
481    rotation = np.linalg.solve(cell, exportedcell)
482
483    pdb_str = _ngl_write_cell(*cellpar)
484    pdb_str += "MODEL     1\n"
485
486    if rotation is not None:
487        positions = np.array(positions).dot(rotation)
488
489    for i, p in enumerate(positions):
490        pdb_str += _ngl_write_atom(i, elements[i], *p)
491
492    pdb_str += "ENDMDL \n"
493    return pdb_str
494
495
496def _atomic_number_to_radius(atomic_number, shift=0.2, slope=0.1, scale=1.0):
497    """
498    Give the atomic radius for plotting, which scales like the root of the atomic number.
499
500    Args:
501        atomic_number (int/float): The atomic number.
502        shift (float): A constant addition to the radius. (Default is 0.2.)
503        slope (float): A multiplier for the root of the atomic number. (Default is 0.1)
504        scale (float): How much to rescale the whole thing by.
505
506    Returns:
507        (float): The radius. (Not physical, just for visualization!)
508    """
509    return (shift + slope * np.sqrt(atomic_number)) * scale
510
511
512def _add_colorscheme_spacefill(
513        view, elements, atomic_numbers, particle_size, scheme="element"
514):
515    """
516    Set NGLView spacefill parameters according to a color-scheme.
517
518    Args:
519        view (NGLWidget): The widget to work on.
520        elements (numpy.ndarray/list): Elemental symbols.
521        atomic_numbers (numpy.ndarray/list): Integer atomic numbers for determining atomic size.
522        particle_size (float): A scale factor for the atomic size.
523        scheme (str): The scheme to use. (Default is "element".)
524
525        Possible NGLView color schemes:
526          " ", "picking", "random", "uniform", "atomindex", "residueindex",
527          "chainindex", "modelindex", "sstruc", "element", "resname", "bfactor",
528          "hydrophobicity", "value", "volume", "occupancy"
529
530    Returns:
531        (nglview.NGLWidget): The modified widget.
532    """
533    for elem, num in set(list(zip(elements, atomic_numbers))):
534        view.add_spacefill(
535            selection="#" + elem,
536            radius_type="vdw",
537            radius=_atomic_number_to_radius(num, scale=particle_size),
538            color_scheme=scheme,
539        )
540    return view
541
542
543def _add_custom_color_spacefill(view, atomic_numbers, particle_size, colors):
544    """
545    Set NGLView spacefill parameters according to per-atom colors.
546
547    Args:
548        view (NGLWidget): The widget to work on.
549        atomic_numbers (numpy.ndarray/list): Integer atomic numbers for determining atomic size.
550        particle_size (float): A scale factor for the atomic size.
551        colors (numpy.ndarray/list): A per-atom list of HTML or hex color codes.
552
553    Returns:
554        (nglview.NGLWidget): The modified widget.
555    """
556    for n, num in enumerate(atomic_numbers):
557        view.add_spacefill(
558            selection=[n],
559            radius_type="vdw",
560            radius=_atomic_number_to_radius(num, scale=particle_size),
561            color=colors[n],
562        )
563    return view
564
565
566def _scalars_to_hex_colors(scalar_field, start=None, end=None, cmap=None):
567    """
568    Convert scalar values to hex codes using a colormap.
569
570    Args:
571        scalar_field (numpy.ndarray/list): Scalars to convert.
572        start (float): Scalar value to map to the bottom of the colormap (values below are clipped). (Default is
573            None, use the minimal scalar value.)
574        end (float): Scalar value to map to the top of the colormap (values above are clipped).  (Default is
575            None, use the maximal scalar value.)
576        cmap (matplotlib.cm): The colormap to use. (Default is None, which gives a blue-red divergent map.)
577
578    Returns:
579        (list): The corresponding hex codes for each scalar value passed in.
580    """
581    if start is None:
582        start = np.amin(scalar_field)
583    if end is None:
584        end = np.amax(scalar_field)
585    interp = interp1d([start, end], [0, 1])
586    remapped_field = interp(np.clip(scalar_field, start, end))  # Map field onto [0,1]
587
588    if cmap is None:
589        try:
590            from seaborn import diverging_palette
591        except ImportError:
592            print(
593                "The package seaborn needs to be installed for the plot3d() function!"
594            )
595        cmap = diverging_palette(245, 15, as_cmap=True)  # A nice blue-red palette
596
597    return [
598        rgb2hex(cmap(scalar)[:3]) for scalar in remapped_field
599    ]  # The slice gets RGB but leaves alpha
600
601
602def _get_orientation(view_plane):
603    """
604    A helper method to plot3d, which generates a rotation matrix from the input `view_plane`, and returns a
605    flattened list of len = 16. This flattened list becomes the input argument to `view.control.orient`.
606
607    Args:
608        view_plane (numpy.ndarray/list): A Nx3-array/list (N = 1,2,3); the first 3d-component of the array
609            specifies which plane of the system to view (for example, [1, 0, 0], [1, 1, 0] or the [1, 1, 1] planes),
610            the second 3d-component (if specified, otherwise [1, 0, 0]) gives the horizontal direction, and the
611            third component (if specified) is the vertical component, which is ignored and calculated internally.
612            The orthonormality of the orientation is internally ensured, and therefore is not required in the
613            function call.
614
615    Returns:
616        (list): orientation tensor
617    """
618    if len(np.array(view_plane).flatten()) % 3 != 0:
619        raise ValueError(
620            "The shape of view plane should be (N, 3), where N = 1, 2 or 3. Refer docs for more info."
621        )
622    view_plane = np.array(view_plane).reshape(-1, 3)
623    rotation_matrix = np.roll(np.eye(3), -1, axis=0)
624    rotation_matrix[: len(view_plane)] = view_plane
625    rotation_matrix /= np.linalg.norm(rotation_matrix, axis=-1)[:, np.newaxis]
626    rotation_matrix[1] -= (
627            np.dot(rotation_matrix[0], rotation_matrix[1]) * rotation_matrix[0]
628    )  # Gran-Schmidt
629    rotation_matrix[2] = np.cross(
630        rotation_matrix[0], rotation_matrix[1]
631    )  # Specify third axis
632    if np.isclose(np.linalg.det(rotation_matrix), 0):
633        return np.eye(
634            3
635        )  # view_plane = [0,0,1] is the default view of NGLview, so we do not modify it
636    return np.roll(
637        rotation_matrix / np.linalg.norm(rotation_matrix, axis=-1)[:, np.newaxis],
638        2,
639        axis=0,
640    ).T
641
642
643def _get_flattened_orientation(view_plane, distance_from_camera):
644    """
645    A helper method to plot3d, which generates a rotation matrix from the input `view_plane`, and returns a
646    flattened list of len = 16. This flattened list becomes the input argument to `view.contol.orient`.
647
648    Args:
649        view_plane (numpy.ndarray/list): A Nx3-array/list (N = 1,2,3); the first 3d-component of the array
650            specifies which plane of the system to view (for example, [1, 0, 0], [1, 1, 0] or the [1, 1, 1] planes),
651            the second 3d-component (if specified, otherwise [1, 0, 0]) gives the horizontal direction, and the
652            third component (if specified) is the vertical component, which is ignored and calculated internally.
653            The orthonormality of the orientation is internally ensured, and therefore is not required in the
654            function call.
655        distance_from_camera (float): Distance of the camera from the structure. Higher = farther away.
656
657    Returns:
658        (list): Flattened list of len = 16, which is the input argument to `view.contol.orient`
659    """
660    if distance_from_camera <= 0:
661        raise ValueError("´distance_from_camera´ must be a positive float!")
662    flattened_orientation = np.eye(4)
663    flattened_orientation[:3, :3] = _get_orientation(view_plane)
664
665    return (distance_from_camera * flattened_orientation).ravel().tolist()
666
667
668@requires("nglview")
669def animate_trajectory(
670        trajectory: Trajectory,
671        spacefill: bool = True,
672        show_cell: bool = True,
673        center_of_mass: bool = False,
674        particle_size: float = 0.5,
675        camera: str = "orthographic",
676):
677    """
678    Animate a series of atomic structures.
679
680    :param trajectory: the trajectory to animate
681    :type trajectory: ase.io.Trajectory
682    :param spacefill: If True, then atoms are visualized in spacefill stype
683    :type spacefill: bool
684    :param show_cell: True if the cell boundaries of the structure is to be shown
685    :type show_cell: bool
686    :param particle_size: Scaling factor for the spheres representing the atoms. (The radius is determined by
687        the atomic number)
688    :type particle_size: float
689    :param center_of_mass: False (default) if the specified positions are w.r.t. the origin
690    :type center_of_mass: bool
691    :param camera: camera perspective, choose from "orthographic" or "perspective" (default is "orthographic")
692    :type camera: str
693    :return: nglview IPython widget
694    :rtype: nglview.NGLWidget
695    """
696
697    import nglview
698    animation = nglview.show_asetraj(trajectory)
699    if spacefill:
700        animation.add_spacefill(radius_type="vdw", scale=0.5, radius=particle_size)
701        animation.remove_ball_and_stick()
702    else:
703        animation.add_ball_and_stick()
704    if show_cell:
705        animation.add_unitcell()
706    animation.camera = camera
707    return animation