Source code for garnett.ciffilewriter

# Copyright (c) 2020 The Regents of the University of Michigan
# All rights reserved.
# This software is licensed under the BSD 3-Clause License.
"""cif-file writer for the Glotzer Group, University of Michigan.

Authors: Julia Dshemuchadse, Carl Simon Adorf

.. code::

    writer = CifFileWriter()

    # write to screen:
    write.write(trajectory)

    # write to file:
    with open('a_ciffile.pos', 'w') as ciffile:
        writer.write(trajectory, ciffile)
"""

import io
import sys
import logging
import math
import datetime
from collections import defaultdict

import numpy as np

logger = logging.getLogger(__name__)


def _determine_unitcell(box):
    lengths = np.sqrt(np.sum(np.array(box.get_box_matrix())**2, axis=0))  # a, b, c
    gamma = math.degrees(
        np.arccos(
            box.xy / math.sqrt(1 + box.xy ** 2)))
    beta = math.degrees(
        np.arccos(
            box.xz / math.sqrt(1 + box.xz ** 2 + box.yz ** 2)))
    alpha = math.degrees(
        np.arccos(
            (box.xy * box.xz + box.yz) /
            (math.sqrt(1 + box.xy ** 2)
                * math.sqrt(1 + box.xz ** 2 + box.yz ** 2))))
    angles = np.array([alpha, beta, gamma])
    return lengths, angles


[docs]class CifFileWriter(object): """cif-file writer for the Glotzer Group, University of Michigan. Authors: Julia Dshemuchadse, Carl Simon Adorf .. code:: writer = CifFileWriter() # write to screen: write.write(trajectory) # write to file: with open('a_ciffile.pos', 'w') as ciffile: writer.write(trajectory, ciffile) """ def _write_frame(self, frame, file, data, occupancies, fractional, raw): from . import __version__ if occupancies is None: occupancies = np.ones(frame.position.shape[0]) def _write(msg='', end='\n'): file.write(msg + end) unitcell_lengths, unitcell_angles = _determine_unitcell(frame.box) # write title _write("data_" + data) # _write("data_" + os.path.splitext(ciffilename)[0]) _write("_audit_creation_method " "'generated by garnett {}'".format(__version__)) _write("_audit_creation_date {}".format(str(datetime.date.today()))) # write unit cell parameters _write("_cell_length_a {}".format(unitcell_lengths[0])) _write("_cell_length_b {}".format(unitcell_lengths[1])) _write("_cell_length_c {}".format(unitcell_lengths[2])) _write("_cell_angle_alpha {}".format(unitcell_angles[0])) _write("_cell_angle_beta {}".format(unitcell_angles[1])) _write("_cell_angle_gamma {}".format(unitcell_angles[2])) _write() # write symmetry - P1 _write("_symmetry_space_group_name_H-M " + "'P 1'") _write("_symmetry_Int_Tables_number " + str(1)) _write() # write header for particle positions _write("loop_") _write("_atom_site_label") _write("_atom_site_type_symbol") _write("_atom_site_occupancy") _write("_atom_site_fract_x") _write("_atom_site_fract_y") _write("_atom_site_fract_z") # write header particle positions if raw: fractions = frame.cif_coordinates.copy() else: if fractional: fractions = frame.position.copy() else: invbox = np.linalg.inv(frame.box.get_box_matrix()) fractions = np.dot(invbox, frame.position.T).T fractions += 0.5 type_counter = defaultdict(int) n_digits = len(str(len(frame.position))) particle_str = "{ptype}{pnum:0" + str(n_digits) + "d} {ptype} {occ:3.2f} {position}" for i, (position, typeid, occupancy) in enumerate(zip(fractions, frame.typeid, occupancies)): particle_type = frame.types[typeid] _write(particle_str.format( pnum=type_counter[particle_type], ptype=particle_type, occ=occupancy, position=' '.join(map(str, position)))) type_counter[particle_type] += 1
[docs] def write(self, trajectory, file=sys.stdout, data='simulation', occupancy=None, fractional=False, raw=False): """Serialize a trajectory into cif-format and write it to file. :param trajectory: The trajectory to serialize :type trajectory: :class:`~garnett.trajectory.Trajectory` :param file: The file to write the trajectory to. :type file: A file-like object. :param data: Identifier which be will written to the file, signifying the origin of the data. :type data: str :param occupancy: The default occupancy of individual particles. :type occupancy: int :param fractional: Whether or not the input coordinates are fractional :type fractional: bool :param raw: Whether or not to write the raw CIF coordinates (with no transformations) :type raw: bool """ for i, frame in enumerate(trajectory): self._write_frame( frame=frame, file=file, data='{}_frame_{}'.format(data, i), occupancies=occupancy, fractional=fractional, raw=raw) logger.debug("Wrote frame {}.".format(i + 1)) logger.info("Wrote {} frames.".format(i + 1))
[docs] def dump(self, trajectory, data='simulation', occupancy=None, fractional=False, raw=False): """Serialize trajectory into cif-format. :param trajectory: The trajectory to serialize. :type trajectory: :class:`~garnett.trajectory.Trajectory` :param data: Identifier which be will written to the file, signifying the origin of the data. :type data: str :param occupancy: The default occupancy of individual particles. :type occupancy: numpy.array :rtype: str """ f = io.StringIO() self.write(trajectory, f, occupancy=occupancy, fractional=fractional, raw=raw) return f.getvalue()