Source code for schrodinger.application.matsci.nano.tools

"""
Tools for periodic systems (crystals, slabs) manipulation.

Copyright Schrodinger, LLC. All rights reserved."""

import numpy

from schrodinger.application.matsci.nano import xtal
from schrodinger.structutils import transform


def origin_shift_maestro(struct, atom_index, shift_to_origin=True):
    """
    Shift origin of the structure to origin or center.

    :param schrodinger.structure.Structure struct: Input structure

    :param int atom_index: Atom index to shift to

    :param bool shift_to_origin: If True, shift to origin, if False shift to
        middle of the cell ([0.5, 0.5, 0.5])
    """

    # This function is called by maestro, don't raise or return anything except
    # structure object no matter what.

    struct = struct.copy()

    if not (1 <= atom_index <= struct.atom_total):
        return struct

    if not xtal.sync_pbc2(struct):
        return struct

    vecs = xtal.get_vectors_from_chorus(struct)
    xyz = struct.atom[atom_index].xyz

    fracs = xtal.trans_cart_to_frac_from_vecs(xyz, *vecs)

    if shift_to_origin:
        origin = numpy.zeros(3)
    else:
        origin = numpy.array([0.5, 0.5, 0.5])

    oshift = origin - fracs

    oshift = xtal.trans_frac_to_cart_from_vecs(oshift, *vecs)
    transform.translate_structure(struct, *oshift)
    struct = xtal.move_atoms_into_cell(struct, preserve_bonding=True)

    return struct