Skip to content

PCST

Exctraction of subgraph using Prize-Collecting Steiner Tree (PCST) algorithm.

PCSTPruning

Bases: NamedTuple

Prize-Collecting Steiner Tree (PCST) pruning algorithm implementation inspired by G-Retriever (He et al., 'G-Retriever: Retrieval-Augmented Generation for Textual Graph Understanding and Question Answering', NeurIPS 2024) paper. https://arxiv.org/abs/2402.07630 https://github.com/XiaoxinHe/G-Retriever/blob/main/src/dataset/utils/retrieval.py

Parameters:

Name Type Description Default
topk

The number of top nodes to consider.

 required
topk_e

The number of top edges to consider.

 required
cost_e

The cost of the edges.

 required
c_const

The constant value for the cost of the edges computation.

 required
root

The root node of the subgraph, -1 for unrooted.

 required
num_clusters

The number of clusters.

 required
pruning

The pruning strategy to use.

 required
verbosity_level

The verbosity level.

 required
Source code in aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
class PCSTPruning(NamedTuple):
    """
    Prize-Collecting Steiner Tree (PCST) pruning algorithm implementation inspired by G-Retriever
    (He et al., 'G-Retriever: Retrieval-Augmented Generation for Textual Graph Understanding and
    Question Answering', NeurIPS 2024) paper.
    https://arxiv.org/abs/2402.07630
    https://github.com/XiaoxinHe/G-Retriever/blob/main/src/dataset/utils/retrieval.py

    Args:
        topk: The number of top nodes to consider.
        topk_e: The number of top edges to consider.
        cost_e: The cost of the edges.
        c_const: The constant value for the cost of the edges computation.
        root: The root node of the subgraph, -1 for unrooted.
        num_clusters: The number of clusters.
        pruning: The pruning strategy to use.
        verbosity_level: The verbosity level.
    """
    topk: int = 3
    topk_e: int = 3
    cost_e: float = 0.5
    c_const: float = 0.01
    root: int = -1
    num_clusters: int = 1
    pruning: str = "gw"
    verbosity_level: int = 0

    def compute_prizes(self, graph: Data, query_emb: torch.Tensor) -> np.ndarray:
        """
        Compute the node prizes based on the cosine similarity between the query and nodes,
        as well as the edge prizes based on the cosine similarity between the query and edges.
        Note that the node and edge embeddings shall use the same embedding model and dimensions
        with the query.

        Args:
            graph: The knowledge graph in PyTorch Geometric Data format.
            query_emb: The query embedding in PyTorch Tensor format.

        Returns:
            The prizes of the nodes and edges.
        """
        # Compute prizes for nodes
        n_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.x)
        topk = min(self.topk, graph.num_nodes)
        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
        n_prizes = torch.zeros_like(n_prizes)
        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()

        # Compute prizes for edges
        # e_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.edge_attr)
        # topk_e = min(self.topk_e, e_prizes.unique().size(0))
        # topk_e_values, _ = torch.topk(e_prizes.unique(), topk_e, largest=True)
        # e_prizes[e_prizes < topk_e_values[-1]] = 0.0
        # last_topk_e_value = topk_e
        # for k in range(topk_e):
        #     indices = e_prizes == topk_e_values[k]
        #     value = min((topk_e - k) / sum(indices), last_topk_e_value)
        #     e_prizes[indices] = value
        #     last_topk_e_value = value * (1 - self.c_const)

        # Optimized version of the above code
        e_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.edge_attr)
        unique_prizes, inverse_indices = e_prizes.unique(return_inverse=True)
        topk_e = min(self.topk_e, unique_prizes.size(0))
        topk_e_values, _ = torch.topk(unique_prizes, topk_e, largest=True)
        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
        last_topk_e_value = topk_e
        for k in range(topk_e):
            indices = inverse_indices == (
                unique_prizes == topk_e_values[k]
            ).nonzero(as_tuple=True)[0]
            value = min((topk_e - k) / indices.sum().item(), last_topk_e_value)
            e_prizes[indices] = value
            last_topk_e_value = value * (1 - self.c_const)

        return {"nodes": n_prizes, "edges": e_prizes}

    def compute_subgraph_costs(
        self, graph: Data, prizes: dict
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Compute the costs in constructing the subgraph proposed by G-Retriever paper.

        Args:
            graph: The knowledge graph in PyTorch Geometric Data format.
            prizes: The prizes of the nodes and the edges.

        Returns:
            edges: The edges of the subgraph, consisting of edges and number of edges without
                virtual edges.
            prizes: The prizes of the subgraph.
            costs: The costs of the subgraph.
        """
        # Logic to reduce the cost of the edges such that at least one edge is selected
        updated_cost_e = min(
            self.cost_e,
            prizes["edges"].max().item() * (1 - self.c_const / 2),
        )

        # Initialize variables
        edges = []
        costs = []
        virtual = {
            "n_prizes": [],
            "edges": [],
            "costs": [],
        }
        mapping = {"nodes": {}, "edges": {}}

        # Compute the costs, edges, and virtual variables based on the prizes
        for i, (src, dst) in enumerate(graph.edge_index.T.numpy()):
            prize_e = prizes["edges"][i]
            if prize_e <= updated_cost_e:
                mapping["edges"][len(edges)] = i
                edges.append((src, dst))
                costs.append(updated_cost_e - prize_e)
            else:
                virtual_node_id = graph.num_nodes + len(virtual["n_prizes"])
                mapping["nodes"][virtual_node_id] = i
                virtual["edges"].append((src, virtual_node_id))
                virtual["edges"].append((virtual_node_id, dst))
                virtual["costs"].append(0)
                virtual["costs"].append(0)
                virtual["n_prizes"].append(prize_e - updated_cost_e)
        prizes = np.concatenate([prizes["nodes"], np.array(virtual["n_prizes"])])
        edges_dict = {}
        edges_dict["edges"] = edges
        edges_dict["num_prior_edges"] = len(edges)
        # Final computation of the costs and edges based on the virtual costs and virtual edges
        if len(virtual["costs"]) > 0:
            costs = np.array(costs + virtual["costs"])
            edges = np.array(edges + virtual["edges"])
            edges_dict["edges"] = edges

        return edges_dict, prizes, costs, mapping

    def get_subgraph_nodes_edges(
        self, graph: Data, vertices: np.ndarray, edges_dict: dict, mapping: dict,
    ) -> dict:
        """
        Get the selected nodes and edges of the subgraph based on the vertices and edges computed
        by the PCST algorithm.

        Args:
            graph: The knowledge graph in PyTorch Geometric Data format.
            vertices: The vertices of the subgraph computed by the PCST algorithm.
            edges_dict: The dictionary of edges of the subgraph computed by the PCST algorithm,
                and the number of prior edges (without virtual edges).
            mapping: The mapping dictionary of the nodes and edges.
            num_prior_edges: The number of edges before adding virtual edges.

        Returns:
            The selected nodes and edges of the extracted subgraph.
        """
        # Get edges information
        edges = edges_dict["edges"]
        num_prior_edges = edges_dict["num_prior_edges"]
        # Retrieve the selected nodes and edges based on the given vertices and edges
        subgraph_nodes = vertices[vertices < graph.num_nodes]
        subgraph_edges = [mapping["edges"][e] for e in edges if e < num_prior_edges]
        virtual_vertices = vertices[vertices >= graph.num_nodes]
        if len(virtual_vertices) > 0:
            virtual_vertices = vertices[vertices >= graph.num_nodes]
            virtual_edges = [mapping["nodes"][i] for i in virtual_vertices]
            subgraph_edges = np.array(subgraph_edges + virtual_edges)
        edge_index = graph.edge_index[:, subgraph_edges]
        subgraph_nodes = np.unique(
            np.concatenate(
                [subgraph_nodes, edge_index[0].numpy(), edge_index[1].numpy()]
            )
        )

        return {"nodes": subgraph_nodes, "edges": subgraph_edges}

    def extract_subgraph(self, graph: Data, query_emb: torch.Tensor) -> dict:
        """
        Perform the Prize-Collecting Steiner Tree (PCST) algorithm to extract the subgraph.

        Args:
            graph: The knowledge graph in PyTorch Geometric Data format.
            query_emb: The query embedding.

        Returns:
            The selected nodes and edges of the subgraph.
        """
        # Assert the topk and topk_e values for subgraph retrieval
        assert self.topk > 0, "topk must be greater than or equal to 0"
        assert self.topk_e > 0, "topk_e must be greater than or equal to 0"

        # Retrieve the top-k nodes and edges based on the query embedding
        prizes = self.compute_prizes(graph, query_emb)

        # Compute costs in constructing the subgraph
        edges_dict, prizes, costs, mapping = self.compute_subgraph_costs(
            graph, prizes
        )

        # Retrieve the subgraph using the PCST algorithm
        result_vertices, result_edges = pcst_fast.pcst_fast(
            edges_dict["edges"],
            prizes,
            costs,
            self.root,
            self.num_clusters,
            self.pruning,
            self.verbosity_level,
        )

        subgraph = self.get_subgraph_nodes_edges(
            graph,
            result_vertices,
            {"edges": result_edges, "num_prior_edges": edges_dict["num_prior_edges"]},
            mapping)

        return subgraph

compute_prizes(graph, query_emb)

Compute the node prizes based on the cosine similarity between the query and nodes, as well as the edge prizes based on the cosine similarity between the query and edges. Note that the node and edge embeddings shall use the same embedding model and dimensions with the query.

Parameters:

Name Type Description Default
graph Data

The knowledge graph in PyTorch Geometric Data format.

 required
query_emb Tensor

The query embedding in PyTorch Tensor format.

 required

Returns:

Type Description
ndarray

The prizes of the nodes and edges.
Source code in aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def compute_prizes(self, graph: Data, query_emb: torch.Tensor) -> np.ndarray:
    """
    Compute the node prizes based on the cosine similarity between the query and nodes,
    as well as the edge prizes based on the cosine similarity between the query and edges.
    Note that the node and edge embeddings shall use the same embedding model and dimensions
    with the query.

    Args:
        graph: The knowledge graph in PyTorch Geometric Data format.
        query_emb: The query embedding in PyTorch Tensor format.

    Returns:
        The prizes of the nodes and edges.
    """
    # Compute prizes for nodes
    n_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.x)
    topk = min(self.topk, graph.num_nodes)
    _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
    n_prizes = torch.zeros_like(n_prizes)
    n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()

    # Compute prizes for edges
    # e_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.edge_attr)
    # topk_e = min(self.topk_e, e_prizes.unique().size(0))
    # topk_e_values, _ = torch.topk(e_prizes.unique(), topk_e, largest=True)
    # e_prizes[e_prizes < topk_e_values[-1]] = 0.0
    # last_topk_e_value = topk_e
    # for k in range(topk_e):
    #     indices = e_prizes == topk_e_values[k]
    #     value = min((topk_e - k) / sum(indices), last_topk_e_value)
    #     e_prizes[indices] = value
    #     last_topk_e_value = value * (1 - self.c_const)

    # Optimized version of the above code
    e_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.edge_attr)
    unique_prizes, inverse_indices = e_prizes.unique(return_inverse=True)
    topk_e = min(self.topk_e, unique_prizes.size(0))
    topk_e_values, _ = torch.topk(unique_prizes, topk_e, largest=True)
    e_prizes[e_prizes < topk_e_values[-1]] = 0.0
    last_topk_e_value = topk_e
    for k in range(topk_e):
        indices = inverse_indices == (
            unique_prizes == topk_e_values[k]
        ).nonzero(as_tuple=True)[0]
        value = min((topk_e - k) / indices.sum().item(), last_topk_e_value)
        e_prizes[indices] = value
        last_topk_e_value = value * (1 - self.c_const)

    return {"nodes": n_prizes, "edges": e_prizes}

compute_subgraph_costs(graph, prizes)

Compute the costs in constructing the subgraph proposed by G-Retriever paper.

Parameters:

Name Type Description Default
graph Data

The knowledge graph in PyTorch Geometric Data format.

 required
prizes dict

The prizes of the nodes and the edges.

 required

Returns:

Name Type Description
edges ndarray

The edges of the subgraph, consisting of edges and number of edges without virtual edges.
prizes ndarray

The prizes of the subgraph.
costs ndarray

The costs of the subgraph.
Source code in aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def compute_subgraph_costs(
    self, graph: Data, prizes: dict
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Compute the costs in constructing the subgraph proposed by G-Retriever paper.

    Args:
        graph: The knowledge graph in PyTorch Geometric Data format.
        prizes: The prizes of the nodes and the edges.

    Returns:
        edges: The edges of the subgraph, consisting of edges and number of edges without
            virtual edges.
        prizes: The prizes of the subgraph.
        costs: The costs of the subgraph.
    """
    # Logic to reduce the cost of the edges such that at least one edge is selected
    updated_cost_e = min(
        self.cost_e,
        prizes["edges"].max().item() * (1 - self.c_const / 2),
    )

    # Initialize variables
    edges = []
    costs = []
    virtual = {
        "n_prizes": [],
        "edges": [],
        "costs": [],
    }
    mapping = {"nodes": {}, "edges": {}}

    # Compute the costs, edges, and virtual variables based on the prizes
    for i, (src, dst) in enumerate(graph.edge_index.T.numpy()):
        prize_e = prizes["edges"][i]
        if prize_e <= updated_cost_e:
            mapping["edges"][len(edges)] = i
            edges.append((src, dst))
            costs.append(updated_cost_e - prize_e)
        else:
            virtual_node_id = graph.num_nodes + len(virtual["n_prizes"])
            mapping["nodes"][virtual_node_id] = i
            virtual["edges"].append((src, virtual_node_id))
            virtual["edges"].append((virtual_node_id, dst))
            virtual["costs"].append(0)
            virtual["costs"].append(0)
            virtual["n_prizes"].append(prize_e - updated_cost_e)
    prizes = np.concatenate([prizes["nodes"], np.array(virtual["n_prizes"])])
    edges_dict = {}
    edges_dict["edges"] = edges
    edges_dict["num_prior_edges"] = len(edges)
    # Final computation of the costs and edges based on the virtual costs and virtual edges
    if len(virtual["costs"]) > 0:
        costs = np.array(costs + virtual["costs"])
        edges = np.array(edges + virtual["edges"])
        edges_dict["edges"] = edges

    return edges_dict, prizes, costs, mapping

extract_subgraph(graph, query_emb)

Perform the Prize-Collecting Steiner Tree (PCST) algorithm to extract the subgraph.

Parameters:

Name Type Description Default
graph Data

The knowledge graph in PyTorch Geometric Data format.

 required
query_emb Tensor

The query embedding.

 required

Returns:

Type Description
dict

The selected nodes and edges of the subgraph.
Source code in aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def extract_subgraph(self, graph: Data, query_emb: torch.Tensor) -> dict:
    """
    Perform the Prize-Collecting Steiner Tree (PCST) algorithm to extract the subgraph.

    Args:
        graph: The knowledge graph in PyTorch Geometric Data format.
        query_emb: The query embedding.

    Returns:
        The selected nodes and edges of the subgraph.
    """
    # Assert the topk and topk_e values for subgraph retrieval
    assert self.topk > 0, "topk must be greater than or equal to 0"
    assert self.topk_e > 0, "topk_e must be greater than or equal to 0"

    # Retrieve the top-k nodes and edges based on the query embedding
    prizes = self.compute_prizes(graph, query_emb)

    # Compute costs in constructing the subgraph
    edges_dict, prizes, costs, mapping = self.compute_subgraph_costs(
        graph, prizes
    )

    # Retrieve the subgraph using the PCST algorithm
    result_vertices, result_edges = pcst_fast.pcst_fast(
        edges_dict["edges"],
        prizes,
        costs,
        self.root,
        self.num_clusters,
        self.pruning,
        self.verbosity_level,
    )

    subgraph = self.get_subgraph_nodes_edges(
        graph,
        result_vertices,
        {"edges": result_edges, "num_prior_edges": edges_dict["num_prior_edges"]},
        mapping)

    return subgraph

get_subgraph_nodes_edges(graph, vertices, edges_dict, mapping)

Get the selected nodes and edges of the subgraph based on the vertices and edges computed by the PCST algorithm.

Parameters:

Name Type Description Default
graph Data

The knowledge graph in PyTorch Geometric Data format.

 required
vertices ndarray

The vertices of the subgraph computed by the PCST algorithm.

 required
edges_dict dict

The dictionary of edges of the subgraph computed by the PCST algorithm, and the number of prior edges (without virtual edges).

 required
mapping dict

The mapping dictionary of the nodes and edges.

 required
num_prior_edges

The number of edges before adding virtual edges.

 required

Returns:

Type Description
dict

The selected nodes and edges of the extracted subgraph.
Source code in aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def get_subgraph_nodes_edges(
    self, graph: Data, vertices: np.ndarray, edges_dict: dict, mapping: dict,
) -> dict:
    """
    Get the selected nodes and edges of the subgraph based on the vertices and edges computed
    by the PCST algorithm.

    Args:
        graph: The knowledge graph in PyTorch Geometric Data format.
        vertices: The vertices of the subgraph computed by the PCST algorithm.
        edges_dict: The dictionary of edges of the subgraph computed by the PCST algorithm,
            and the number of prior edges (without virtual edges).
        mapping: The mapping dictionary of the nodes and edges.
        num_prior_edges: The number of edges before adding virtual edges.

    Returns:
        The selected nodes and edges of the extracted subgraph.
    """
    # Get edges information
    edges = edges_dict["edges"]
    num_prior_edges = edges_dict["num_prior_edges"]
    # Retrieve the selected nodes and edges based on the given vertices and edges
    subgraph_nodes = vertices[vertices < graph.num_nodes]
    subgraph_edges = [mapping["edges"][e] for e in edges if e < num_prior_edges]
    virtual_vertices = vertices[vertices >= graph.num_nodes]
    if len(virtual_vertices) > 0:
        virtual_vertices = vertices[vertices >= graph.num_nodes]
        virtual_edges = [mapping["nodes"][i] for i in virtual_vertices]
        subgraph_edges = np.array(subgraph_edges + virtual_edges)
    edge_index = graph.edge_index[:, subgraph_edges]
    subgraph_nodes = np.unique(
        np.concatenate(
            [subgraph_nodes, edge_index[0].numpy(), edge_index[1].numpy()]
        )
    )

    return {"nodes": subgraph_nodes, "edges": subgraph_edges}