sgtbx_extra_groups.py

#!/usr/bin/env python

"""Quick and extremely dirty script for generating code for SpaceGroup
that are defined in cctbx, but not in mmLib.  It was used to generate
module sgtbxspacegroups.

This is a utility script that should not be included with code
distribution.

Not to be included with code distributions.
"""


import math
import re

import numpy
from cctbx import sgtbx

from diffpy.structure.spacegroups import SpaceGroup, SymOp, is_space_group_identifier, mmLibSpaceGroupList


def tuple_to_sg_array(tpl):
    if not _rtarrays:
        import diffpy.structure.SpaceGroups as sgmod

        for n in dir(sgmod):
            if not n.startswith("Rot_") and not n.startswith("Tr_"):
                continue
            a = getattr(sgmod, n)
            t = tuple(a.flatten())
            _rtarrays[t] = a
    if len(tpl) == 3:
        tpl = tuple([(x - math.floor(x)) for x in tpl])
        if tpl not in _rtarrays:
            _rtarrays[tpl] = numpy.array(tpl, dtype=float)
    return _rtarrays[tpl]


_rtarrays = {}


def mm_space_group_from_symbol(symbol):
    """Construct SpaceGroup instance from a string symbol using sgtbx
    data."""
    sginfo = sgtbx.space_group_info(symbol)
    symop_list = []
    symop_list = get_symop_list(sginfo.group())
    sgtype = sginfo.type()
    uhm = sgtype.lookup_symbol()
    sgsmbls = sgtbx.space_group_symbols(uhm)
    kw = {}
    kw["number"] = sgtype.number()
    kw["num_sym_equiv"] = len(symop_list)
    kw["num_primitive_sym_equiv"] = count_unique_rotations(symop_list)
    kw["short_name"] = sgsmbls.hermann_mauguin().replace(" ", "")
    pgt = sgsmbls.point_group_type()
    pgn = "PG" + re.sub(r"-(\d)", "\\1bar", pgt)
    kw["point_group_name"] = pgn
    kw["crystal_system"] = sgsmbls.crystal_system().upper()
    kw["pdb_name"] = sgsmbls.hermann_mauguin()
    kw["symop_list"] = symop_list
    mmsg = SpaceGroup(**kw)
    return mmsg


def adjust_mm_space_group_number(mmsg):
    sg0 = [x for x in mmLibSpaceGroupList if x.number == mmsg.number]
    if sg0 and cmp_space_groups(sg0[0], mmsg):
        return
    while mmsg.number in sgnumbers:
        mmsg.number += 1000
    sgnumbers.append(mmsg.number)


def get_symop_list(grp):
    symop_list = []
    for op in grp:
        r_sgtbx = op.r().as_double()
        t_sgtbx = op.t().as_double()
        R = tuple_to_sg_array(r_sgtbx)
        t = tuple_to_sg_array(t_sgtbx)
        symop_list.append(SymOp(R, t))
    return symop_list


def count_unique_rotations(symop_list):
    unique_rotations = set()
    for op in symop_list:
        tpl = tuple(op.R.flatten())
        unique_rotations.add(tpl)
    return len(unique_rotations)


def cmp_space_groups(sg0, sg1):
    if sg0 is sg1:
        return True
    s0 = hash_mm_space_group(sg0)
    s1 = hash_mm_space_group(sg1)
    return s0 == s1


def find_equivalent_mm_space_group(grp):
    if not _equivmmsg:
        for sgn in mmLibSpaceGroupList:
            ssgn = hash_mm_space_group(sgn)
            _equivmmsg.setdefault(ssgn, sgn)
    ssg = hash_sgtbx_group(grp)
    return _equivmmsg.get(ssg)


_equivmmsg = {}


def find_equivalent_sgtbx_space_group(sgmm):
    if not _equivsgtbx:
        for smbls in sgtbx.space_group_symbol_iterator():
            uhm = smbls.universal_hermann_mauguin()
            grp = sgtbx.space_group_info(uhm).group()
            hgrp = hash_sgtbx_group(grp)
            _equivsgtbx.setdefault(hgrp, grp)
    hgmm = hash_mm_space_group(sgmm)
    return _equivsgtbx.get(hgmm)


_equivsgtbx = {}


def hash_mm_space_group(sg):
    lines = [str(sg.number % 1000)] + sorted(map(str, sg.iter_symops()))
    s = "\n".join(lines)
    return s


def hash_sgtbx_group(grp):
    n = grp.type().number()
    lines = [str(n)] + sorted(map(str, get_symop_list(grp)))
    s = "\n".join(lines)
    return s


sgnumbers = [sg.number for sg in mmLibSpaceGroupList]

_SGsrc = """\
sg%(number)i = SpaceGroup(
    number = %(number)i,
    num_sym_equiv = %(num_sym_equiv)i,
    num_primitive_sym_equiv = %(num_sym_equiv)i,
    short_name = %(short_name)r,
    point_group_name = %(point_group_name)r,
    crystal_system = %(crystal_system)r,
    pdb_name = %(pdb_name)r,
    symop_list = [
        @SYMOPS@
    ]
)
"""


def sg_code(mmsg):
    src0 = _SGsrc % mmsg.__dict__
    src1 = src0.replace("@SYMOPS@", symops_code(mmsg))
    return src1


def symops_code(mmsg):
    lst = ["%8s%s," % ("", symop_code(op)) for op in mmsg.iter_symops()]
    src = "\n".join(lst).strip()
    return src


def symop_code(op):
    if not _rtnames:
        import diffpy.structure.SpaceGroups as sgmod

        for n in dir(sgmod):
            if not n.startswith("Rot_") and not n.startswith("Tr_"):
                continue
            a = getattr(sgmod, n)
            at = tuple(a.flatten())
            _rtnames[at] = "sgmod." + n
    nR = _rtnames[tuple(op.R.flatten())]
    nt = _rtnames[tuple(op.t)]
    src = "SymOp(%s, %s)" % (nR, nt)
    return src


_rtnames = {}


def main():
    duplicates = set()
    for smbls in sgtbx.space_group_symbol_iterator():
        uhm = smbls.universal_hermann_mauguin()
        grp = sgtbx.space_group_info(uhm).group()
        if find_equivalent_mm_space_group(grp):
            continue
        shn = smbls.hermann_mauguin().replace(" ", "")
        if is_space_group_identifier(shn):
            continue
        sg = mm_space_group_from_symbol(uhm)
        hsg = hash_mm_space_group(sg)
        if hsg in duplicates:
            continue
        adjust_mm_space_group_number(sg)
        duplicates.add(hsg)
        print(sg_code(sg))
    return


if __name__ == "__main__":
    main()