def _build_graph()

in identity-resolution/notebooks/identity-graph/nepytune/usecase/similar_audience.py [0:0]


def _build_graph(average_buyer_categories, similar_audience):
    avg_buyer = "averageBuyer"

    graph = nx.Graph()
    graph.add_node(avg_buyer, label=avg_buyer, **average_buyer_categories)

    for avg_iab in average_buyer_categories.keys():
        graph.add_node(avg_iab, label="IAB", category=avg_iab)
        graph.add_edge(avg_buyer, avg_iab, label="interestedIn")

    for user in similar_audience:
        pid, cats, tids = user["pid"], user["iabs"], user["tids"]

        user_categories = dict(sorted(cats.items(), key=lambda x: x[1])[:3])
        comparison = {k: cats.get(k, 0) for k in average_buyer_categories.keys()}
        user_categories.update(comparison)

        user_comparisons = False
        for ucategory, value in user_categories.items():
            graph.add_node(ucategory, label="IAB", category=ucategory)
            label = "interestedIn"
            if value:
                if ucategory in average_buyer_categories:
                    if user_categories[ucategory] >= average_buyer_categories[ucategory]:
                        user_comparisons = True
                    else:
                        label = "interestedInButNotSufficient"
                graph.add_edge(pid, ucategory, label=label)

        opacity = 1 if user_comparisons else 0.5
        for tid in tids:
            graph.add_edge(pid, tid, label="hasIdentity")
            graph.add_node(tid, label="transientId", uid=tid, opacity=opacity)

        graph.add_node(
            pid, label="persistentId", pid=pid,
            opacity=opacity, **cats
        )

    return graph