Source code for MDMC.MD.ase.conversions

"""This module enables conversion between MDMC Structure objects and ASE
Atom and Atoms objects.
from __future__ import annotations

from itertools import chain
from typing import Iterable, Union

import ase
# For some reason importing ase alone is not importing, therefore:
from import x3d
import numpy as np

from MDMC.MD.structures import Atom
from MDMC.MD.interactions import Bond

[docs]class ASEAtoms(ase.atoms.Atoms): """ A subclass of ``ase.atoms.Atoms`` with explicit bonds defined between atoms Attributes ---------- bonds : numpy.ndarray An ``array`` of ``tuple``, where each ``tuple`` is an atom pair, which are specified by the indexes (`int`) of each atom. Raises ------ ValueError If there are not the same number of ``ID`` as there are atoms """ def __init__(self, *args, **kwargs): bonds = kwargs.pop('bonds', None) IDs = kwargs.pop('IDs', None) super().__init__(*args, **kwargs) self.bonds = bonds if IDs and len(IDs) != len(self): raise ValueError('There must be an ID for every atom') self.IDs = IDs def __getitem__(self, i: Union[int, 'list[int]', slice]) -> Union[ase.atom.Atom, ASEAtoms]: """ Extends ``ase.atoms.Atoms._getitem__`` so that ``bonds`` are dealt with correctly. This means that any bond between two atoms in the subset defined by `i` are assigned to the new item returned. Parameters ---------- i : int, list of int, slice Returns an ``ase.atom.Atom`` (if `int`) or ``ASEAtoms`` (if `list` or slice). """ atoms = super().__getitem__(i) # If int then only single atom so bonds are irrelevant if isinstance(i, int): return atoms # Account for fact that ASE allows i to be a list of integers, as well # an int or slice if isinstance(i, slice): atoms.IDs = self.IDs[i] indexes = range(i.start or 0, i.stop, i.step or 1) else: atoms.IDs = [self.IDs[j] for j in i] indexes = i # Get all bonds where both atom indexes in a bond (which is a tuple of # two atom indexes) are in the of the subset determined by i bonds = [bond for bond in self.bonds if all(atom_index in indexes for atom_index in bond)] atoms.bonds = bonds return atoms # format (redefined builtin) is used as this is method signature in ASE base # class #pylint: disable=redefined-builtin
[docs] def write(self, filename: str, format: str = None, **kwargs: dict) -> None: """ Override ase.atoms.Atoms.write so that writing to X3D uses custom X3D class Parameters ---------- filename : str The name of the file to which the ``ASEAtoms`` object is written format : str The name of the format of the file **kwargs keyword arguments which are passed to the object which performs the writing """ if format and format.upper() in ['X3D', 'X3DOM']: X3D(self).write(filename, datatype=format.upper()) else: super().write(filename, format, **kwargs)
[docs]def convert_to_ase_atom(atom: Atom, index: int = None) -> ase.atom.Atom: """ Converts an MDMC ``Atom`` to an ``ase.atom.Atom`` Parameters ---------- atom : Atom An MDMC ``Atom`` object to be converted to an ``ase.atom.Atom`` object index : int, optional The ``index`` of the ``ase.atom.Atom`` object which is created. If this is not set, the MDMC ``Atom.ID`` is used. Returns ------- ase.atom.Atom An ``ASE.atom.Atom`` object which is equivalent to ``atom`` """ index = index if index else atom.ID return ase.atom.Atom(position=atom.position, index=index, mass=atom.mass, symbol=atom.element, charge=atom.charge)
[docs]def convert_from_ase_atom(ase_atom: ASEAtoms, atom_type: int = None, name: str = None, set_charge: bool = True, cutoff: float = None) -> 'Atom': """ Converts an ``ase.atom.Atom`` to an MDMC ``Atom``. As MDMC automatically generates atom ``ID``, ``ase_atom.index`` is not passed when initializing an ``Atom``. Parameters ---------- ase_atom : ASEAtom An ``ASEAtom`` object to be converted to an MDMC ``Atom`` object atom_type : int The atom_type of the MDMC ``Atom`` object. name : str, optional A name for the MDMC ``Atom``. The default is the element symbol. set_charge : bool, optional Whether the ``charge`` is set to the ``charge`` of the ``ase.atom.Atom, or left unset. All ``ase.atom.Atom`` objects have a ``charge``, which is set to 0. if it is uninitialized. As MDMC ``Atom`` objects can have ``charge=None`, in some cases it might be preferential to leave the ``charge`` unset. The default is to set the ``charge``. cutoff : float, optional The cutoff value for the atom's charge interaction. Must be set if set_charge is True. Returns ------- ``Atom`` An MDMC ``Atom`` object which is equivalent to ``ase_atom`` """ name = name if name else ase_atom.symbol kwargs = {'position': ase_atom.position, 'mass': ase_atom.mass, 'name': name} if set_charge: kwargs['charge'] = ase_atom.charge if cutoff is None: raise ValueError("If set_charge is True, a cutoff must also be set.") kwargs['cutoff'] = cutoff if atom_type: kwargs['atom_type'] = atom_type return Atom(ase_atom.symbol, **kwargs)
[docs]def get_ase_atoms(atoms: Iterable, cell: np.ndarray = None) -> ASEAtoms: """ Gets an ``ASEAtoms`` object equivalent to ``atoms``, including the bonding Parameters ---------- atoms : iterable An ``iterable`` of MDMC ``Atom`` objects to be converted to an ``ASEAtoms`` object cell : numpy.ndarray, optional A 3 element ``array`` specifying the unit cell of the ``ASEAtoms`` object. The default is `None`. Returns ------- ASEAtoms An ``ASEAtoms`` object which is equivalent to ``atoms`` """ # The ase.atoms.Atoms object unhelpfully overwrites the index attribute of # any ase.atom.Atom objects which belong to it (so Atoms[i].index == i, # regardless of the index of the Atom at that index). This means that the # atom IDs used for the bond atom pairs need to be converted. index_conv = {atom.ID: index for index, atom in enumerate(atoms)} bonds = set(chain.from_iterable([convert_bonds(atom.bonded_interactions, index_conv) for atom in atoms])) IDs = [atom.ID for atom in atoms] return ASEAtoms([convert_to_ase_atom(atom, index) for index, atom in enumerate(atoms)], cell=cell, bonds=bonds, IDs=IDs)
[docs]def convert_bond(bond: Bond, index_conv: dict = None) -> np.ndarray: """ Converts ``Bond`` objects into the form required by the ASE GUI Parameters ---------- bond : Bond The bond which will be converted. index_conv : dict A dictionary of ``MDMC_ID``: ``ASE_index`` pairs, where ``MDMC_ID`` is an `int` specifying an ``Atom.ID``, and ``ASE_index`` is the corresponding ``ase.atom.Atom.index``. The default is `None`, which means that the ``ID`` and ``index`` will be assumed to be identical. Returns ------- numpy.ndarray An ``array`` of 2 element `list` where each element is the `int` ``index`` of an atom between which the bond exists. """ indexing = (lambda x: index_conv[x.ID]) if index_conv else lambda x: x.ID # Ensure atom IDs are ordered in each atom pair return [tuple(sorted(map(indexing, atom_pair))) for atom_pair in bond.atoms]
[docs]def convert_bonds(bonds: 'list[Bond]', index_conv: dict = None) -> np.ndarray: """ Converts ``Bond`` objects into the form required by the ASE GUI Parameters ---------- bonds : list The `list` of ``Bond`` objects to be converted index_conv : dict A `dict` of ``MDMC_ID``: ``ASE_index`` pairs, where ``MDMC_ID`` is an `int` specifying an ``Atom.ID``, and ``ASE_index`` is the corresponding ``ase.atom.Atom.index``. The default is `None`, which means that the ``ID`` and ``index`` will be assumed to be identical. Returns ------- numpy.ndarray An ``array`` of 2 element `list` where each element is the `int` ``index`` of an atom between which the bond exists. """ # conditional because only bond objects are supported return list(chain.from_iterable([convert_bond(bond, index_conv) for bond in bonds if isinstance(bond, Bond)]))
[docs]class X3D(x3d.X3D): """ Class to write X3D or X3DOM (html) which can be read by web browers This is used for inline IPython (Jupyter) visualisation with X3DOM. It overrides the `write` method of the base class, in order to display bonds as defined within MDMC. Parameters ---------- atoms : ase.atoms.Atoms The ``ase.atoms.Atoms`` object to write """ def __init__(self, atoms: ase.atoms.Atoms): super().__init__(atoms) # 3000 atoms was picked because at this number of atoms the view is such # that a reduction in the sphere/cyinder resolution is insignificant. It # is also approximately half of the total number of atoms that can be # viewed on 2018 macbook pro (2.6GHz, 16GB DDR4, 4GB GDDR5) self.reduce_memory = len(self._atoms) > 3000
[docs] def write(self, fileobj: str, datatype: str = 'X3DOM'): """ Writes output to an 'X3DOM' (html) or 'X3D' file Parameters ---------- fileobj : str or file-like The file name to which the ``atoms`` are written datatype : str, optional The output format, which can be 'X3D' or 'X3DOM' (html). If `None`, format will be determined from the `filename` """ if datatype is None: if fileobj.endswith('.x3d'): datatype = 'X3D' elif fileobj.endswith('.html'): datatype = 'X3DOM' else: raise ValueError("filename must end in '.x3d' or '.html'.") # Write the header w = x3d.WriteToFile(fileobj) if datatype == 'X3DOM': w(0, '<html>') w(1, '<head>') w(2, '<title>MDMC atomic visualization</title>') w(2, '<link rel="stylesheet" type="text/css"') w(2, ' href="">') w(2, '</link>') w(2, '<script type="text/javascript"') w(2, ' src="">') w(2, '</script>') w(1, '</head>') w(1, '<body>') w(2, '<X3D width="800px" height="800px">') # Store next level of indentation and for scene scn_i = 3 elif datatype == 'X3D': w(0, '<?xml version="1.0" encoding="UTF-8"?>') w(0, '<!DOCTYPE X3D PUBLIC "ISO//Web3D//DTD X3D 3.2//EN" ' '"">') w(0, '<X3D profile="Interchange" version="3.2" ' 'xmlns:xsd="" ' 'xsd:noNamespaceSchemaLocation=' '"">') # Store next level of indentation and for scene scn_i = 1 w(scn_i, '<Scene>') w(scn_i, '<Viewpoint centerOfRotation="{0:.2f} {1:.2f} {2:.2f}"' ' position="{0:.2f} {1:.2f} {3:.2f}"></Viewpoint>'.format( *self.get_center_of_rotation(), self.get_viewpoint_z())) for atom in self._atoms: for indent, line in self.atom_lines(atom): w(4 + indent, line) bonds = getattr(self._atoms, 'bonds', []) if bonds: for bond in bonds: for indent, line in self.bond_lines(bond): w(4 + indent, line) if datatype == 'X3DOM': w(3, '</Scene>') w(2, '</X3D>') w(1, '</body>') w(0, '</html>') elif datatype == 'X3D': w(0, '</X3D>')
[docs] def atom_lines(self, atom: ase.atom.Atom) -> list: """ Generates a segment of X3D lines representing an atom Parameters ---------- atom : ase.atom.Atom The ``ase.atom.Atom`` for which the X3D html will be generated Returns ------- list A `list` of (indent, str) pairs for each line requied to describe the `atom` """ color = tuple([atom.number]) diffuse_color = 'diffuseColor="{0:.3f} {1:.3f} {2:.3f}"'.format(*color) specular_color = 'specularColor="0.5 0.5 0.5"'.format(*color) if self.reduce_memory: # Decrease the resolution of the spheres sphere_subdivision = ' subdivision=12,12' else: sphere_subdivision = '' lines = [(0, '<Transform translation="{0:.2f} {1:.2f} {2:.2f}">'.format( *atom.position))] lines += [(1, '<Shape>')] lines += [(2, '<Appearance>')] lines += [(3, '<Material {0} {1}>'.format(diffuse_color, specular_color))] lines += [(3, '</Material>')] lines += [(2, '</Appearance>')] # ASE covalent radii are too large if bonds are also going to be added lines += [(2, '<Sphere radius="{0:.2f}"{1}>'.format([atom.number] / 4., sphere_subdivision))] lines += [(2, '</Sphere>')] lines += [(1, '</Shape>')] lines += [(0, '</Transform>')] return lines
[docs] def bond_lines(self, bond: tuple) -> list: """ Generates a cylinder representing a bond Parameters ---------- bond : tuple A `tuple` containing the atom indexes for the bond for which the X3D html will be generated Returns ------- list A `list` of (indent, str) pairs for each line requied to describe the `bond` """ if self.reduce_memory: cylinder_subdivision = ' subdivision=16' else: cylinder_subdivision = '' positions = [self._atoms[index].position for index in bond] sub = (positions[1] - positions[0]) # Separation of the two atoms defines the height of the cylinder (bond) separation = np.linalg.norm(sub) # Move one end of the cylinder coincident with one atom origin = positions[0] + np.array([0., separation / 2., 0.]) # All cylinders (bonds) are oriented along y axis by default # Calculate axis-angle representation in order to set bond rotation cylinder = np.array([0., np.abs(separation), 0.]) def normalise(x): return x / np.linalg.norm(x) uvec1, uvec2 = normalise(cylinder), normalise(sub) axis = np.cross(uvec1, uvec2) angle = np.linalg.norm(np.arccos(, uvec2))) # Shift the transform center of the cylinder from the center (default) # to one end, so that it is rotated around one of the atoms lines = [(0, '<Transform center="0 {0:.4f} 0"' ' translation="{2:.4f} {3:.4f} {4:.4f}"' ' rotation="{5:.4f} {6:.4f} {7:.4f} {1:.4f}">'.format( -separation / 2., angle, *origin, *axis))] lines += [(1, '<Shape>')] lines += [(2, '<Appearance>')] lines += [(3, '<Material diffuseColor="0 0 0"' ' specularColor="0.5 0.5 0.5">')] lines += [(3, '</Material>')] lines += [(2, '</Appearance>')] lines += [(2, '<Cylinder height="{0:.4f}" radius="0.02"{1}>'.format( separation, cylinder_subdivision))] lines += [(2, '</Cylinder>')] lines += [(1, '</Shape>')] lines += [(0, '</Transform>')] return lines
[docs] def get_center_of_rotation(self) -> np.ndarray: """ Get the center of rotation for the viewpoint Returns ------- numpy.ndarray The center of `atoms.cell` if this has been set, or the center of the extents of `atoms.positions` """ atoms = self._atoms\ # As ASE always has a defined cell for atoms (it is just # Cell([0., 0., 0.]) if it has not been set), cannot simply test for # existence # Currently only consider orthorhombic cell cell = np.diagonal(atoms.cell) return (cell / 2. if not np.all(cell == np.array([0.] * 3)) else np.mean([np.max(atoms.positions, axis=0), np.min(atoms.positions, axis=0)], axis=0))
[docs] def get_viewpoint_z(self) -> float: """ Get the z position of the viewpoint which will display all of the `atoms` Returns ------- float The z position of the viewpoint """ # The viewpoint is always centered on the atoms (or the cell) in the xy # plane. It has been determined that an angle of ~0.30 radians between # the atom with the greatest extent (or the cell's greatest extent) will # be sufficient to display this atom (and therefore all atoms). VIEWPOINT_ANGLE = 0.30 # Currently only consider orthorhombic cell cell = np.diagonal(self._atoms.cell) if np.any(cell != np.array([0., 0., 0.])): extents = cell else: extents = np.max(self._atoms.positions, axis=0) # Calculate distance between extents (xy position) and z axis # (i.e. x == y == 0.) xydistance = np.linalg.norm(extents[:2] - self.get_center_of_rotation()[:2]) return xydistance / np.tan(VIEWPOINT_ANGLE) + extents[2]