basis.py 文件源码

def car2sph(sh, cart):
    """
    Turns symbolic solid harmonic functions into a dictionary of
    arrays containing cartesian to spherical transformation matrices.

    Args
        sh (OrderedDict): the result of solid_harmonics(l_tot)
        cart (dict): dictionary of l, cartesian l, m, n ordering
    """
    keys, conv = {}, {}
    prevL, mscnt = 0, 0
    for L in cart:
        for idx, (l, m, n) in enumerate(cart[L]):
            key = ''
            if l: key += 'x'
            if l > 1: key += str(l)
            if m: key += 'y'
            if m > 1: key += str(m)
            if n: key += 'z'
            if n > 1: key += str(n)
            keys[key] = idx
    # TODO: six compatibility
    for key, sym in sh.items():
        L = key[0]
        mscnt = mscnt if prevL == L else 0
        conv.setdefault(L, np.zeros((cart_lml_count[L],
                                     spher_lml_count[L]),
                                     dtype=np.float64))
        if isinstance(sym, (Mul, Add)):
            string = (str(sym.expand())
                      .replace(' + ', ' ')
                      .replace(' - ', ' -'))
            for chnk in string.split():
                pre, exp = chnk.split('*', 1)
                if L == 1: conv[L] = np.array(cart[L])
                else: conv[L][keys[exp.replace('*', '')], mscnt] = float(pre)
        prevL = L
        mscnt += 1
    conv[0] = np.array([[1]])
    return conv