tree_graphs.py

import time
import numpy as np
from numba import njit
import networkx as nx
import multiprocessing


@njit
def generate_trees_old(last_tree: np.ndarray) -> np.ndarray:
    """Creates duplicates, do not use
    :param last_tree: ndarray
    :return: ndarray
    """

    num_trees = len(last_tree)
    num_nodes = len(last_tree[0])
    new_tree = np.zeros((num_nodes*num_trees, num_nodes+1), dtype=np.int64)

    count = 0
    for index in range(num_trees):
        for index2 in range(num_nodes):
            new_tree[count][0:num_nodes] = last_tree[index]
            new_tree[count, num_nodes] = index2
            count += 1

    return new_tree


def generate_trees(n: int) -> tuple[list[list[list[int]]], np.ndarray]:
    """Generates all nonisomorphic tree graphs with n nodes.

    :param n: Number of nodes in the trees
    :return: list of trees
    """
    trees: list[nx.Graph] = nx.generators.nonisomorphic_trees(n)
    new_trees = []
    for tree in trees:
        tree_list = list(list(int(node) for node in line.split()) for line in nx.generate_adjlist(tree))
        tree_list[0], tree_list[1] = tree_list[1], tree_list[0]
        new_list = list([] for _ in tree_list)
        for index, line in enumerate(tree_list):
            for index2 in range(1, len(line)):
                new_list[index].append(line[index2])
                new_list[line[index2]].append(line[0])

        new_trees.append(new_list)

    # if numpy_:
    trees_ = np.ones((len(new_trees), len(new_trees[0]), n), dtype=np.int64)
    # contains node connection information for each node in each tree,
    # last digit contains number of connections for the node
    trees_ = trees_*(-1)
    for index, tree in enumerate(new_trees):
        for index2, line in enumerate(tree):
            trees_[index][index2][0:len(line)] = np.array(line, dtype=np.int64)
            trees_[index][index2][-1] = len(line)

    return new_trees, trees_


def find_unfolding_3d(tree: list[list[int]]) -> list[list[int]]:
    """Finds all unfoldings of the given tree.

    :param tree: the tree to find unfoldings of
    :return: list of unfoldings
    """
    unfoldings = []
    for index in range(2, len(tree)):
        for index2 in range(2, len(tree)):
            for index3 in range(3, len(tree)):
                unfolding: list[int] = list(-1 for _ in tree)
                unfolding[0] = index
                unfolding[index] = 0

                if index2 not in unfolding:
                    for index_ in range(1, len(tree)):
                        if index_ not in unfolding and index_ != index2:
                            unfolding[index_] = index2
                            unfolding[index2] = index_
                            break

                if index3 not in unfolding:
                    for index__ in range(2, len(tree)):
                        if index__ not in unfolding and index__ != index3:
                            unfolding[index__] = index3
                            unfolding[index3] = index__
                            break

                if -1 not in unfolding:
                    if unfolding not in unfoldings:
                        for index4, edge in enumerate(unfolding):
                            if edge in tree[index4]:
                                break

                        else:  # no break
                            unfoldings.append(unfolding)

    return unfoldings


def find_unfolding_nd_old(tree: list[list[int]]) -> list[list[int]]:
    """Finds all unfoldings of the given tree.

    :param tree: the tree to find unfoldings of
    :return: list of unfoldings
    """
    unfoldings = []
    for index in range(2, len(tree)):
        for index2 in range(2, len(tree)):
            for index3 in range(3, len(tree)):
                unfolding: list[int] = list(-1 for _ in tree)
                unfolding[0] = index
                unfolding[index] = 0

                for index__, i in enumerate((index2, index3)):
                    if i not in unfolding:
                        for index_ in range(index__+1, len(tree)):
                            if index_ not in unfolding and index_ != i:
                                unfolding[index_] = i
                                unfolding[i] = index_
                                break

                if -1 not in unfolding:
                    if unfolding not in unfoldings:
                        for index4, edge in enumerate(unfolding):
                            if edge in tree[index4]:
                                break

                        else:
                            unfoldings.append(unfolding)

    return unfoldings


def find_unfolding_nd(tree: list[list[int]]) -> list[list[int]]:
    """Finds all unfoldings of the given tree.

    :param tree: the tree to find unfoldings of
    :return: list of unfoldings
    """
    def helper(count: int):
        """helper function for find_unfolding_nd"""

        if count < int(number_of_nodes/2 - 2):
            count += 1
            for indexx in range(count + 2, number_of_nodes):
                list_of_indexes[count] = indexx
                helper(count)

        else:
            # unfolding: list[int] = list(-1 for _ in tree)
            unfolding = []
            for _ in tree:
                unfolding.append(-1)
            unfolding[0] = index
            unfolding[index] = 0

            for index__, i in enumerate(list_of_indexes):
                if i not in unfolding:
                    for index_ in range(index__ + 1, len(tree)):
                        if index_ not in unfolding and index_ != i:
                            unfolding[index_] = i
                            unfolding[i] = index_
                            break

            if -1 not in unfolding:
                if unfolding not in unfoldings:
                    for index4, edge in enumerate(unfolding):
                        if edge in tree[index4]:
                            break

                    else:
                        unfoldings.append(unfolding)

    number_of_nodes = len(tree)
    count_ = -1
    unfoldings = []
    # list_of_indexes = list(0 for _ in range(int(number_of_nodes/2 - 1)))
    list_of_indexes = []
    for _ in range(int(number_of_nodes/2 - 1)):
        list_of_indexes.append(0)
    for index in range(2, number_of_nodes):
        helper(count_)

    return unfoldings


@njit
def helper_old(count: int, tree: np.ndarray, list_of_indexes: np.ndarray,
               unfoldings: np.ndarray, num_unfoldings: int, index: int):
    """helper function for find_unfolding_nd"""

    if count < int(len(tree)/2 - 2):
        count += 1
        for indexx in range(count + 2, len(tree)):
            list_of_indexes[count] = indexx
            num_unfoldings = helper(count, tree, list_of_indexes, unfoldings, num_unfoldings, index)

    else:
        unfolding = np.ones(len(tree), dtype=np.int8)
        unfolding *= -1
        unfolding[0] = index
        unfolding[index] = 0

        for index__, i in enumerate(list_of_indexes):
            if i not in unfolding:
                for index_ in range(index__ + 1, len(tree)):
                    if index_ not in unfolding and index_ != i:
                        unfolding[index_] = i
                        unfolding[i] = index_
                        break

        if -1 not in unfolding:
            for index4, edge in enumerate(unfolding):
                node = tree[index4, :tree[index4, -1]]
                if edge in node:
                    break
            else:  # no break
                for unfolding_ in unfoldings[:num_unfoldings]:
                    if np.all(unfolding_ == unfolding):
                        break

                else:  # no break
                    unfoldings[num_unfoldings] = unfolding
                    num_unfoldings += 1

    return num_unfoldings


@njit
def helper(count: int, tree: np.ndarray, list_of_indexes: np.ndarray,
           unfoldings: np.ndarray, num_unfoldings: int, index: int, num_unfoldings_old: int):
    """helper function for find_unfolding_nd"""

    if count < int(len(tree)/2 - 2):
        count += 1
        for indexx in range(count + 2, len(tree)):  # creates nested for loops
            list_of_indexes[count] = indexx
            num_unfoldings = helper(count, tree, list_of_indexes, unfoldings, num_unfoldings, index, num_unfoldings_old)

    else:
        if index not in tree[0, :tree[0, -1]]:  # check if node[index] is adjacent to node[0]
            unfolding = np.full(len(tree), -1, dtype=np.int8)
            unfolding[0] = index
            unfolding[index] = 0

            for index__, i in enumerate(list_of_indexes):
                if i not in unfolding:
                    for index_ in range(index__ + 1, len(tree)):
                        if index_ not in unfolding and index_ != i:
                            if i not in tree[index_, :tree[index_, -1]]:  # check if node[i] is adjacent to node[index_]
                                unfolding[index_] = i
                                unfolding[i] = index_
                                break

            if -1 not in unfolding:
                for unfolding_ in unfoldings[num_unfoldings_old:num_unfoldings]:  # check if unfolding is already found
                    if np.all(unfolding_ == unfolding):
                        break

                else:  # no break
                    unfoldings[num_unfoldings] = unfolding
                    num_unfoldings += 1

    return num_unfoldings


@njit
def find_unfolding_nd_njit(tree_: np.ndarray) -> np.ndarray:
    """Finds all unfoldings of the given tree.

    :param tree_: the tree to find unfoldings of
    :return: list of unfoldings
    """

    count_ = -1
    max_num_unfoldings = int(4.8*10**(len(tree_)/2-3)+3)  # optimized for memory usage
    unfoldings_ = np.empty((max_num_unfoldings, len(tree_)), dtype=np.int8)
    num_unfoldings_ = 0
    list_of_indexes_ = np.empty(int(len(tree_)/2 - 1), dtype=np.int8)
    for index in range(2, len(tree_)):
        num_unfoldings_old = num_unfoldings_
        num_unfoldings_ = helper(count_, tree_, list_of_indexes_, unfoldings_,
                                 num_unfoldings_, index, num_unfoldings_old)

    return unfoldings_[:num_unfoldings_]


def find_unfolding_3d2(tree: list[list[int]]) -> list[list[int]]:
    """Finds almost all unfoldings of the given tree.

    :param tree: The tree to find unfoldings of
    :return: list of unfoldings
    """
    unfoldings = []
    for index in range(2, len(tree)):
        for index2 in range(2, len(tree)):
            for index3 in range(3, len(tree)):
                unfolding: list[int] = list(-1 for _ in tree)
                unfolding[0] = index
                unfolding[index] = 0

                if index2 not in unfolding:
                    index_ = 1
                    if index_ not in unfolding and index_ != index2:
                        unfolding[index_] = index2
                        unfolding[index2] = index_

                if index3 not in unfolding:
                    index__ = 2
                    if index == 2:
                        index__ = 3
                    if index__ not in unfolding and index__ != index3:
                        unfolding[index__] = index3
                        unfolding[index3] = index__

                if -1 not in unfolding:
                    if unfolding not in unfoldings:
                        for index4, edge in enumerate(unfolding):
                            if edge in tree[index4]:
                                break

                        else:
                            unfoldings.append(unfolding)

    return unfoldings


def generate_unfold_order(tree: list[list[int]]) -> np.ndarray:
    """Defines an unfold-order for a given tree.

    :param tree: A tree-graph
    :return: An array of indexes
    """
    order = np.full(len(tree), -1, dtype=np.int8)
    order_index = 0
    start_index = 0
    for index, node in enumerate(tree):
        if len(node) == 1:
            order[order_index] = index
            order_index += 1

    stop_index = order_index
    while order_index < len(order):
        for index in range(start_index, stop_index):
            tree_index = order[index]
            node = tree[tree_index]
            for index2 in node:
                if index2 not in order:
                    order[order_index] = index2
                    order_index += 1
                    break

        start_index = stop_index
        stop_index = order_index

    return order[:-1]


def trees_and_unfoldings_nd(n: int) -> tuple[list[list[list[int]]], list[np.ndarray], np.ndarray, int]:
    """Generates all trees with n nodes and their respective unfoldings.

    :param n: The number of nodes for the trees
    :return: list of trees, list of unfoldings, number of unfoldings
    """
    trees_list, trees_np = generate_trees(n)
    unfoldings = []
    unfold_order = np.zeros((len(trees_list), len(trees_list[0])-1), dtype=int)
    num_unfoldings = 0
    print(f"Number of trees: {len(trees_list)}")
    for index, tree in enumerate(zip(trees_list, trees_np)):
        print(f"Tree number: {index}")
        unfold_order[index] = generate_unfold_order(tree[0])
        tree_unfold = find_unfolding_nd_njit(tree[1])
        unfoldings.append(tree_unfold)
        num_unfoldings += len(tree_unfold)

    return trees_list, unfoldings, unfold_order, num_unfoldings


def func(tree: tuple[list[list[int]], np.ndarray]) -> tuple[np.ndarray, np.ndarray]:
    """function for multi-processed unfold-finding.

    :param tree: Tree to unfold
    :return: unfold_order, tree_unfold
    """
    print("start tree")
    unfold_order = generate_unfold_order(tree[0])
    tree_unfold = find_unfolding_nd_njit(tree[1])
    print("end tree")
    return unfold_order, tree_unfold


def trees_and_unfoldings_nd_multiprocessing(n: int
                                            ) -> tuple[list[list[list[int]]], list[np.ndarray], np.ndarray, int]:
    """Generates all trees with n nodes and their respective unfoldings.

    :param n: The number of nodes for the trees
    :return: list of trees, list of unfoldings, unfold_order, number of unfoldings
    """
    trees_list, trees_np = generate_trees(n)
    unfoldings = []
    unfold_order = np.empty((len(trees_list), len(trees_list[0])-1), dtype=np.int8)
    num_unfoldings = 0
    print(f"Number of trees: {len(trees_list)}")

    with multiprocessing.Pool() as pool:
        for index, result in enumerate(pool.map(func, zip(trees_list, trees_np))):
            print(f"tree {index}")
            unfold_order_, tree_unfold = result
            unfold_order[index] = unfold_order_
            unfoldings.append(tree_unfold)
            num_unfoldings += len(tree_unfold)

    return trees_list, unfoldings, unfold_order, num_unfoldings


def main():
    """The main function"""
    while True:
        try:
            n = int(input("Number of nodes: "))
            break
        except ValueError:
            pass

    start = time.time()
    # trees_list, unfoldings, unfold_order, num_unfoldings = trees_and_unfoldings_nd_multiprocessing(n)
    trees_and_unfoldings_nd_multiprocessing(n)
    stop = time.time()
    print(f"Time to compute: {stop - start}")


if __name__ == "__main__":
    main()