Symmetry Operations

def decode_symm_int(coded_integer):
    """
    Decode an integer encoded symmetry operation. 

    A space group operation is compressed using ternary numerical system for
    rotation and duodecimal system for translation. This is achieved because
    each element of rotation matrix can have only one of {-1,0,1}, and the
    translation can have one of {0,2,3,4,6,8,9,10} divided by 12.  Therefore
    3^9 * 12^3 = 34012224 different values can map space group operations. In
    principle, octal numerical system can be used for translation, but
    duodecimal system is more convenient.

    >>> encode_symm_str(*decode_symm_int(16484))
    '+x,+y,+z'

    Args:
        coded_integer (int): integer encoding a symmetry operation

    Returns:
        Tuple[np.ndarray, np.ndarray]: (3,3) rotation matrix, (3) translation vector
    """
    r = coded_integer % 19683  # 19683 = 3**9
    shift = 6561  # 6561 = 3**8
    rotation = np.empty((3, 3), dtype=np.float64)
    translation = np.empty(3, dtype=np.float64)
    for i in (0, 1, 2):
        for j in (0, 1, 2):
            # we need integer division here
            rotation[i, j] = (r % (shift * 3)) // shift - 1
            shift //= 3

    t = coded_integer // 19683
    shift = 144
    for i in (0, 1, 2):
        # we need integer division here by shift
        translation[i] = ((t % (shift * 12)) // shift) / 12
        shift //= 12
    return rotation, translation

def decode_symm_str(s):
    """
    Decode a symmetry operation represented in the string
    form e.g. '1/2 + x, y, -z -0.25' into a rotation matrix
    and translation vector.

    >>> encode_symm_str(*decode_symm_str("x,y,z"))
    '+x,+y,+z'
    >>> encode_symm_str(*decode_symm_str("1/2 - x,y-0.3333333,z"))
    '1/2-x,2/3+y,+z'

    Args:
        s (str): the encoded symmetry operation string

    Returns:
        Tuple[np.ndarray, np.ndarray]: a (3,3) rotation matrix and a (3) translation vector
    """
    rotation = np.zeros((3, 3), dtype=np.float64)
    translation = np.zeros((3,), dtype=np.float64)
    tokens = s.lower().replace(" ", "").split(",")
    for i, row in enumerate(tokens):
        fac = 1
        row = row.strip()
        symbols = re.findall(SYMM_STR_SYMBOL_REGEX, row)
        for symbol in symbols:
            if "x" in symbol:
                idx = 0
                fac = -1 if "-x" in symbol else 1
                rotation[i, idx] = fac
            elif "y" in symbol:
                idx = 1
                fac = -1 if "-y" in symbol else 1
                rotation[i, idx] = fac
            elif "z" in symbol:
                idx = 2
                fac = -1 if "-z" in symbol else 1
                rotation[i, idx] = fac
            else:
                if "/" in symbol:
                    numerator, denominator = symbol.split("/")
                    translation[i] = Fraction(
                        Fraction(numerator), Fraction(denominator)
                    )
                else:
                    translation[i] += float(Fraction(symbol))
    translation = translation % 1
    return rotation, translation

def encode_symm_int(rotation, translation):
    """
    Encode an integer encoded symmetry from a rotation matrix and translation
    vector.

    A space group operation is compressed using ternary numerical system for
    rotation and duodecimal system for translation. This is achieved because
    each element of rotation matrix can have only one of {-1,0,1}, and the
    translation can have one of {0,2,3,4,6,8,9,10} divided by 12.  Therefore
    3^9 * 12^3 = 34012224 different values can map space group operations. In
    principle, octal numerical system can be used for translation, but
    duodecimal system is more convenient.

    >>> encode_symm_int(((1, 0, 0), (0, 1, 0), (0, 0, 1)), (0, 0, 0))
    16484
    >>> encode_symm_int(((1, 0, 0), (0, 1, 0), (0, 1, 1)), (0, 0.5, 0))
    1433663

     Args:
        rotation (array_like): (3,3) matrix of -1, 0, or 1s encoding the rotation component
            of the symmetry operation
        translation (array_like): (3) vector of rational numbers encoding the translation component
            of the symmetry operation

    Returns:
        int: the encoded symmetry operation
    """

    r = 0
    shift = 1
    # encode rotation component
    rotation = np.round(np.array(rotation)).astype(int) + 1
    for i in (2, 1, 0):
        for j in (2, 1, 0):
            r += rotation[i, j] * shift
            shift *= 3
    t = 0
    shift = 1
    translation = np.round(np.array(translation) * 12).astype(int)
    for i in (2, 1, 0):
        t += translation[i] * shift
        shift *= 12
    return r + t * 19683

def encode_symm_str(rotation, translation):
    """
    Encode a rotation matrix (of -1, 0, 1s) and (rational) translation vector
    into string form e.g. 1/2-x,z-1/3,-y-1/6

    >>> encode_symm_str(((-1, 0, 0), (0, 0, 1), (0, 1, 0)), (0, 0.5, 1/3))
    '-x,1/2+z,1/3+y'
    >>> encode_symm_str(((1, 1, 1), (1, 0, 1), (0, 1, 0)), (0, 0.5, 1/3))
    '+x+y+z,1/2+x+z,1/3+y'

    Args:
        rotation (array_like): (3,3) matrix of -1, 0, or 1s encoding the rotation component
            of the symmetry operation
        translation (array_like): (3) vector of rational numbers encoding the translation component
            of the symmetry operation

    Returns:
        str: the encoded symmetry operation
    """
    symbols = "xyz"
    res = []
    for i in (0, 1, 2):
        t = Fraction(translation[i]).limit_denominator(12)
        v = ""
        if t != 0:
            v += str(t)
        for j in range(0, 3):
            c = rotation[i][j]
            if c != 0:
                s = "-" if c < 0 else "+"
                v += s + symbols[j]
        res.append(v)
    res = ",".join(res)
    return res

def expanded_symmetry_list(reduced_symops, lattice_type):
    """
    Create an expanded list of symmetry operations from the minimum
    specification given a certain lattice type.

    Args:
        reduced_symops (List[SymmetryOperation]): reduced list of symmetry operations
        lattice_type (int): integer encoded lattice type with SHELX conventions, i.e.
            ```
            1: P,
            2: I,
            3: rhombohedral obverse on hexagonal axes,
            4: F,
            5: A,
            6: B,
            7: C
            ```
    Returns:
        List[SymmetryOperation]: an expanded list of symmetry operations given lattice type
    """
    lattice_type_value = abs(lattice_type)
    translations = LATTICE_TYPE_TRANSLATIONS[lattice_type_value]

    identity = SymmetryOperation.identity()
    if identity not in reduced_symops:
        LOG.debug("Appending identity symop %s to reduced_symops")
        reduced_symops.append(identity)
    LOG.debug("Reduced symmetry list contains %d symops", len(reduced_symops))

    full_symops = []

    for symop in reduced_symops:
        full_symops.append(symop)
        for t in translations:
            full_symops.append(symop + t)

    if lattice_type > 0:
        full_symops += [x.inverted() for x in full_symops]

    LOG.debug("Expanded symmetry list contains %d symops", len(full_symops))
    return full_symops

def reduced_symmetry_list(full_symops, lattice_type):
    """
    Reduce an expanded list of symmetry operations to the minimum
    specification given a certain lattice type.

    Args:
        full_symops (List[SymmetryOperation]): list of symmetry operations
        lattice_type (int): integer encoded lattice type with SHELX conventions, i.e.
            ```
            1: P,
            2: I,
            3: rhombohedral obverse on hexagonal axes,
            4: F,
            5: A,
            6: B,
            7: C
            ```
    Returns:
        List[SymmetryOperation]: minimal list of symmetry operations given lattice type
    """
    lattice_type_value = abs(lattice_type)
    translations = LATTICE_TYPE_TRANSLATIONS[lattice_type_value]

    reduced_symops = [SymmetryOperation.identity()]
    symops_to_process = list(full_symops)

    inversion = lattice_type > 0

    while symops_to_process:
        next_symop = symops_to_process.pop(0)
        if next_symop in reduced_symops:
            continue
        if inversion and next_symop.inverted() in reduced_symops:
            continue
        for t in translations:
            x = next_symop + t
            if inversion and x.inverted() in reduced_symops:
                break
            if x in reduced_symops:
                break
        else:
            reduced_symops.append(next_symop)

    LOG.debug("Reduced symmetry list contains %d symops", len(reduced_symops))
    return reduced_symops