polya.py

from functools import lru_cache
from itertools import product

from scipy.special import binom


def polya_counting(permutation_group, num_color):
    cnt = 0
    for perm in permutation_group:
        type_of_perm = get_type_of_permutation(perm)
        cnt += num_color ** sum(type_of_perm)

    assert(cnt % len(permutation_group) == 0)
    cnt //= len(permutation_group)
    return cnt


@lru_cache(maxsize=None)
def get_inventory_coefficient(type_of_perm: tuple, didx, num_elements_of_each_color: tuple):
    # termination
    if didx == 0:
        if all([e == 0 for e in num_elements_of_each_color]):
            return 1
        else:
            return 0

    ret = 0

    for nec in product(*[range(0, e + 1, didx) for e in num_elements_of_each_color]):
        list_k = [e // didx for e in nec]
        if sum(list_k) != type_of_perm[didx - 1]:
            continue
        complement = tuple([e1 - e2 for e1, e2 in zip(num_elements_of_each_color, nec)])
        ret += get_multinomial_coefficient(list_k) \
            * get_inventory_coefficient(type_of_perm, didx - 1, complement)

    return ret


# ref: https://stackoverflow.com/questions/46374185/does-python-have-a-function-which-computes-multinomial-coefficients
def get_multinomial_coefficient(params):
    if len(params) == 1:
        return 1
    if params[-1] == 0:
        return get_multinomial_coefficient(params[:-1])
    return int(binom(sum(params), params[-1])) * get_multinomial_coefficient(params[:-1])


def polya_fixed_degrees_counting(permutation_group, num_color, num_elements_of_each_color):
    num_elements = len(permutation_group[0])
    coeffs = 0

    for perm in permutation_group:
        type_of_perm = tuple(get_type_of_permutation(perm))
        # import pdb; pdb.set_trace()
        coeffs_perm = get_inventory_coefficient(type_of_perm, num_elements,
                                                tuple(num_elements_of_each_color))
        coeffs += coeffs_perm

    assert(coeffs % len(permutation_group) == 0)
    coeffs //= len(permutation_group)
    return coeffs


def get_type_of_permutation(permutation):
    num_elements = len(permutation)
    type_of_perm = [0 for _ in range(num_elements)]

    flags = [False for _ in range(num_elements)]
    for i in range(num_elements):
        if flags[i]:
            continue
        flags[i] = True
        pos = permutation[i]
        cnt = 1
        while pos != i:
            flags[pos] = True
            pos = permutation[pos]
            cnt += 1

        type_of_perm[cnt - 1] += 1

    return type_of_perm


if __name__ == '__main__':
    Dihedral4 = [
        [0, 1, 2, 3],  # E
        [3, 0, 1, 2],  # C4
        [2, 3, 0, 1],
        [1, 2, 3, 0],
        [1, 0, 3, 2],  # mirror
        [2, 1, 0, 3],
        [3, 2, 1, 0],
        [0, 3, 2, 1]
    ]

    actual = polya_counting(Dihedral4, num_color=2)
    assert(actual == 6)

    actual = polya_fixed_degrees_counting(Dihedral4, num_color=2,
                                          num_elements_of_each_color=[3, 1])
    assert(actual == 1)

    actual = polya_fixed_degrees_counting(Dihedral4, num_color=2,
                                          num_elements_of_each_color=[2, 2])
    assert(actual == 2)