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

"""
Utility functions for nanostructures.

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

# Contributor: Thomas F. Hughes

from past.utils import old_div

import numpy

from schrodinger.infra import mm
from schrodinger.structutils import transform

_version = '$Revision 0.0 $'


def get_atomic_element_symbols():
    """
    Return a set of symbols of atomic elements in the periodic table.

    :rtype: set of strs
    :return: symbols, set of atomic elements
    """

    max_atomic_number = mm.MMELEMENTS_MAX

    symbols = set()
    for atomicnumber in range(1, max_atomic_number + 1):
        atomicsymbol = mm.mmelement_get_symbol_by_atomic_number(atomicnumber)
        symbols.add(atomicsymbol)

    return symbols


def get_rotated_vector(invector, angle, length=None, axis=None):
    """
    Return a rotated vector of a given length.

    :type invector: numpy.array
    :param invector: input vector

    :type angle: float
    :param angle: rotation angle in radians

    :type length: float
    :param length: length of output vector

    :type axis: numpy.array
    :param axis: rotation axis

    :rtype: numpy.array
    :return: outvector, output vector
    """

    if not axis:
        axis = transform.Z_AXIS

    inveclen = numpy.linalg.norm(invector)
    inunitvec = transform.get_normalized_vector(invector)
    rotmatrix = transform.get_rotation_matrix(axis, angle)
    outunitvec = transform.transform_atom_coordinates(inunitvec, rotmatrix)
    if length:
        outvector = outunitvec * length
    else:
        outvector = outunitvec * inveclen

    return outvector


def get_inversion_index(index):
    """
    Return the atom index obtained by applying an inversion operation
    to the given index in a given hexagon.

    :type index: int
    :param index: atom index of an atom in a honeycomb cell

    :rtype: invindex, int
    :return: invindex, atom index after inversion
    """

    numatoms = 6

    invindex = index + old_div(numatoms, 2)
    if invindex > numatoms:
        invindex = invindex - numatoms

    return invindex