Repository: microsoft/graphrag
File: graphrag/index/operations/build_noun_graph/build_noun_graph.py (Python)

Code before:
        text_units_df.groupby("text_unit_id").agg({"title": list}).reset_index()
    )
    text_units_df["edges"] = text_units_df["title"].apply(
        lambda x: list(combinations(x, 2))
    )
    edge_df = text_units_df.explode("edges").loc[:, ["edges", "text_unit_id"]]

    edge_df["source"] = edge_df["edges"].apply(
        lambda x: x[0] if isinstance(x, tuple) else None
    )
    edge_df["target"] = edge_df["edges"].apply(
        lambda x: x[1] if isinstance(x, tuple) else None
    )
    edge_df = edge_df[(edge_df.source.notna()) & (edge_df.target.notna())]
    edge_df = edge_df.drop(columns=["edges"])

    # make sure source is always smaller than target
    edge_df["source"], edge_df["target"] = zip(
        *edge_df.apply(
            lambda x: (x["source"], x["target"])
            if x["source"] < x["target"]
            else (x["target"], x["source"]),
            axis=1,
        ),
        strict=False,
    )

    # group by source and target, count the number of text units
    grouped_edge_df = (
        edge_df.groupby(["source", "target"]).agg({"text_unit_id": list}).reset_index()
    )
    grouped_edge_df = grouped_edge_df.rename(columns={"text_unit_id": "text_unit_ids"})
    grouped_edge_df["weight"] = grouped_edge_df["text_unit_ids"].apply(len)

    grouped_edge_df = grouped_edge_df.loc[
        :, ["source", "target", "weight", "text_unit_ids"]
    ]

    if normalize_edge_weights:
        # use PMI weight instead of raw weight
        grouped_edge_df = _calculate_pmi_edge_weights(nodes_df, grouped_edge_df)

    return grouped_edge_df


def _calculate_pmi_edge_weights(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    node_name_col="title",
    node_freq_col="frequency",
    edge_weight_col="weight",

Review comment (suggestion):
This code can be replaced by this, to remove the applies used (also tested locally on my branch)

```python
edge_df[["source", "target"]] = edge_df["edges"].to_list()
edge_df["min_source"] = edge_df[["source", "target"]].min(axis=1)
edge_df["max_target"] = edge_df[["source", "target"]].max(axis=1)
edge_df = edge_df.drop(columns=["source", "target"]).rename(columns={"min_source": "source", "max_target": "target"})

edge_df = edge_df[(edge_df.source.notna()) & (edge_df.target.notna())]
edge_df = edge_df.drop(columns=["edges"])
```



Code after:
        text_units_df.groupby("text_unit_id")
        .agg({"title": lambda x: list(x) if len(x) > 1 else np.nan})
        .reset_index()
    )
    text_units_df = text_units_df.dropna()
    titles = text_units_df["title"].tolist()
    all_edges: list[list[tuple[str, str]]] = [list(combinations(t, 2)) for t in titles]

    text_units_df = text_units_df.assign(edges=all_edges)  # type: ignore
    edge_df = text_units_df.explode("edges")[["edges", "text_unit_id"]]

    edge_df[["source", "target"]] = edge_df.loc[:, "edges"].to_list()
    edge_df["min_source"] = edge_df[["source", "target"]].min(axis=1)
    edge_df["max_target"] = edge_df[["source", "target"]].max(axis=1)
    edge_df = edge_df.drop(columns=["source", "target"]).rename(
        columns={"min_source": "source", "max_target": "target"}  # type: ignore
    )

    edge_df = edge_df[(edge_df.source.notna()) & (edge_df.target.notna())]
    edge_df = edge_df.drop(columns=["edges"])
    # group by source and target, count the number of text units
    grouped_edge_df = (
        edge_df.groupby(["source", "target"]).agg({"text_unit_id": list}).reset_index()
    )
    grouped_edge_df = grouped_edge_df.rename(columns={"text_unit_id": "text_unit_ids"})
    grouped_edge_df["weight"] = grouped_edge_df["text_unit_ids"].apply(len)
    grouped_edge_df = grouped_edge_df.loc[
        :, ["source", "target", "weight", "text_unit_ids"]
    ]
    if normalize_edge_weights:
        # use PMI weight instead of raw weight
        grouped_edge_df = _calculate_pmi_edge_weights(nodes_df, grouped_edge_df)

    return grouped_edge_df


def _calculate_pmi_edge_weights(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    node_name_col="title",
    node_freq_col="frequency",
    edge_weight_col="weight",
    edge_source_col="source",
    edge_target_col="target",
) -> pd.DataFrame:
    """
    Calculate pointwise mutual information (PMI) edge weights.

    pmi(x,y) = log2(p(x,y) / (p(x)p(y)))
    p(x,y) = edge_weight(x,y) / total_edge_weights
    p(x) = freq_occurrence(x) / total_freq_occurrences