EnjunDu/GraphMaster
089
1import json
2import numpy as np
3import networkx as nx
4from collections import Counter, defaultdict
5import random
6import scipy.sparse as sp
7from scipy.sparse.linalg import eigsh
8import sys
9import os
10
11try:
12 import community as community_louvain
13except ImportError:
14 print("Warning: python-louvain package not found. Installing...")
15 import subprocess
16 subprocess.check_call([sys.executable, "-m", "pip", "install", "python-louvain"])
17 import community as community_louvain
18
19def load_graph_from_json(json_file):
20 """Load graph from a JSON file with nodes."""
21 nodes = []
22
23 try:
24 # First try to parse as a single JSON array or object
25 with open(json_file, 'r', encoding='utf-8') as f:
26 content = f.read().strip()
27 try:
28 data = json.loads(content)
29 if isinstance(data, list):
30 nodes = data
31 else:
32 nodes = [data]
33 except json.JSONDecodeError:
34 # Reset and try parsing line by line
35 nodes = []
36 with open(json_file, 'r') as f:
37 for line in f:
38 line = line.strip()
39 if line: # Skip empty lines
40 try:
41 node_data = json.loads(line)
42 nodes.append(node_data)
43 except json.JSONDecodeError:
44 continue
45 except Exception as e:
46 print(f"Error loading graph: {e}")
47 return []
48
49 return nodes
50
51def build_networkx_graph(nodes):
52 """Build a NetworkX graph from the loaded node data."""
53 G = nx.Graph()
54
55 # Add nodes with attributes
56 for node in nodes:
57 G.add_node(
58 node['node_id'],
59 label=node['label'],
60 text=node['text'],
61 mask=node['mask']
62 )
63
64 # Add edges
65 for node in nodes:
66 node_id = node['node_id']
67 for neighbor_id in node['neighbors']:
68 if G.has_node(neighbor_id): # Only add edge if both nodes exist
69 G.add_edge(node_id, neighbor_id)
70
71 return G
72
73def analyze_graph_properties(G):
74 """Analyze the properties of the graph as specified in the requirements."""
75 properties = {}
76
77 # Mask distribution (Train/Validation/Test)
78 masks = [G.nodes[n]['mask'] for n in G.nodes]
79 mask_distribution = Counter(masks)
80 properties['mask_distribution'] = {k: v/len(G.nodes) for k, v in mask_distribution.items()}
81
82 # Label distribution
83 labels = [G.nodes[n]['label'] for n in G.nodes]
84 label_distribution = Counter(labels)
85 properties['label_distribution'] = {k: v/len(G.nodes) for k, v in label_distribution.items()}
86
87 # Graph density
88 properties['density'] = nx.density(G)
89
90 # Degree distribution
91 degrees = [d for n, d in G.degree()]
92 degree_counts = Counter(degrees)
93 properties['degree_distribution'] = {k: v/len(G.nodes) for k, v in degree_counts.items()}
94
95 # Community structure (using Louvain algorithm)
96 try:
97 communities = community_louvain.best_partition(G)
98 community_counts = Counter(communities.values())
99 properties['community_distribution'] = {k: v/len(G.nodes) for k, v in community_counts.items()}
100 except:
101 properties['community_distribution'] = {}
102
103 # Spectral characteristics
104 if len(G) > 1:
105 try:
106 laplacian = nx.normalized_laplacian_matrix(G)
107 if sp.issparse(laplacian) and laplacian.shape[0] > 1:
108 try:
109 k = min(5, laplacian.shape[0]-1)
110 if k > 0:
111 eigenvalues = eigsh(laplacian, k=k, which='SM', return_eigenvectors=False)
112 properties['spectral_eigenvalues'] = sorted(eigenvalues.tolist())
113 else:
114 properties['spectral_eigenvalues'] = []
115 except:
116 properties['spectral_eigenvalues'] = []
117 else:
118 properties['spectral_eigenvalues'] = []
119 except:
120 properties['spectral_eigenvalues'] = []
121 else:
122 properties['spectral_eigenvalues'] = []
123
124 # Connectivity characteristics
125 properties['connected_components'] = nx.number_connected_components(G)
126 largest_cc = max(nx.connected_components(G), key=len)
127 properties['largest_cc_ratio'] = len(largest_cc) / len(G.nodes)
128
129 return properties
130
131def sample_graph_preserving_properties(G, percentage, original_properties):
132 """Sample a percentage of nodes while preserving graph properties."""
133 num_nodes = len(G.nodes)
134 num_nodes_to_sample = max(1, int(num_nodes * percentage / 100))
135
136 # If the graph is too small, just return it
137 if num_nodes <= num_nodes_to_sample:
138 return G, {n: n for n in G.nodes}
139
140 # 1. Preserve label and mask distribution (top priority per requirements)
141 mask_label_groups = defaultdict(list)
142 for node in G.nodes:
143 mask = G.nodes[node]['mask']
144 label = G.nodes[node]['label']
145 mask_label_groups[(mask, label)].append(node)
146
147 # Calculate how many nodes to sample from each mask-label group
148 group_counts = {}
149 for (mask, label), nodes in mask_label_groups.items():
150 mask_ratio = original_properties['mask_distribution'].get(mask, 0)
151 label_ratio = original_properties['label_distribution'].get(label, 0)
152
153 # Calculate joint probability
154 joint_ratio = mask_ratio * label_ratio / sum(
155 original_properties['mask_distribution'].get(m, 0) *
156 original_properties['label_distribution'].get(l, 0)
157 for m in original_properties['mask_distribution']
158 for l in original_properties['label_distribution']
159 )
160
161 target_count = int(num_nodes_to_sample * joint_ratio)
162 # Ensure at least one node from non-empty groups
163 group_counts[(mask, label)] = max(1, target_count) if nodes else 0
164
165 # Adjust to match the exact sample size
166 total_count = sum(group_counts.values())
167 if total_count != num_nodes_to_sample:
168 diff = num_nodes_to_sample - total_count
169 groups = list(group_counts.keys())
170
171 if diff > 0:
172 # Add nodes to groups proportionally to their size
173 group_sizes = [len(mask_label_groups[g]) for g in groups]
174 group_probs = [s/sum(group_sizes) for s in group_sizes]
175
176 for _ in range(diff):
177 group = random.choices(groups, weights=group_probs)[0]
178 if len(mask_label_groups[group]) > group_counts[group]:
179 group_counts[group] += 1
180 else:
181 # Remove nodes from groups with excess
182 groups_with_excess = [(g, c) for g, c in group_counts.items()
183 if c > 1 and c > len(mask_label_groups[g]) * 0.2]
184 groups_with_excess.sort(key=lambda x: x[1], reverse=True)
185
186 for i in range(min(-diff, len(groups_with_excess))):
187 group_counts[groups_with_excess[i][0]] -= 1
188
189 # 2. Sample nodes from each group, prioritizing connectivity and community structure
190 sampled_nodes = []
191
192 # First try to get community structure
193 try:
194 communities = community_louvain.best_partition(G)
195 except:
196 communities = {node: 0 for node in G.nodes} # Fallback if community detection fails
197
198 # Sample from each mask-label group
199 for (mask, label), count in group_counts.items():
200 candidates = mask_label_groups[(mask, label)]
201
202 if len(candidates) <= count:
203 # Take all nodes in this group
204 sampled_nodes.extend(candidates)
205 else:
206 # Score nodes based on degree and community representation
207 node_scores = {}
208 for node in candidates:
209 # Higher score for higher degree nodes (connectivity)
210 degree_score = G.degree(node) / max(1, max(d for n, d in G.degree()))
211
212 # Higher score for nodes in underrepresented communities
213 comm = communities.get(node, 0)
214 comm_sampled = sum(1 for n in sampled_nodes if communities.get(n, -1) == comm)
215 comm_total = sum(1 for n in G.nodes if communities.get(n, -1) == comm)
216 comm_score = 1 - (comm_sampled / max(1, comm_total))
217
218 # Combined score (prioritize connectivity slightly more)
219 node_scores[node] = 0.6 * degree_score + 0.4 * comm_score
220
221 # Sort candidates by score and select the top ones
222 sorted_candidates = sorted(candidates, key=lambda n: node_scores.get(n, 0), reverse=True)
223 sampled_nodes.extend(sorted_candidates[:count])
224
225 # 3. Create the sampled subgraph
226 sampled_G = G.subgraph(sampled_nodes).copy()
227
228 # 4. Improve connectivity if needed
229 if nx.number_connected_components(sampled_G) > original_properties['connected_components']:
230 # Try to improve connectivity by swapping nodes
231 non_sampled = [n for n in G.nodes if n not in sampled_nodes]
232
233 # Calculate betweenness centrality for non-sampled nodes
234 betweenness = {}
235 for node in non_sampled:
236 # Count how many different components this node would connect
237 neighbors = list(G.neighbors(node))
238 sampled_neighbors = [n for n in neighbors if n in sampled_nodes]
239
240 if not sampled_neighbors:
241 continue
242
243 components_connected = set()
244 for n in sampled_neighbors:
245 for comp_idx, comp in enumerate(nx.connected_components(sampled_G)):
246 if n in comp:
247 components_connected.add(comp_idx)
248 break
249
250 betweenness[node] = len(components_connected)
251
252 # Sort non-sampled nodes by how many components they would connect
253 connector_nodes = [(n, b) for n, b in betweenness.items() if b > 1]
254 connector_nodes.sort(key=lambda x: x[1], reverse=True)
255
256 # Try to improve connectivity by swapping nodes
257 for connector, _ in connector_nodes:
258 # Find a node to swap out (prefer low degree nodes from well-represented groups)
259 mask = G.nodes[connector]['mask']
260 label = G.nodes[connector]['label']
261
262 # Find nodes with the same mask and label
263 same_group = [n for n in sampled_nodes
264 if G.nodes[n]['mask'] == mask and G.nodes[n]['label'] == label]
265
266 if not same_group:
267 continue
268
269 # Sort by degree (ascending)
270 same_group.sort(key=lambda n: sampled_G.degree(n))
271
272 # Swap the node with lowest degree
273 to_remove = same_group[0]
274 sampled_nodes.remove(to_remove)
275 sampled_nodes.append(connector)
276
277 # Update the sampled subgraph
278 sampled_G = G.subgraph(sampled_nodes).copy()
279
280 # Stop if we've reached the desired connectivity
281 if nx.number_connected_components(sampled_G) <= original_properties['connected_components']:
282 break
283
284 # 5. Relabel nodes to have consecutive IDs starting from 0
285 node_mapping = {old_id: new_id for new_id, old_id in enumerate(sorted(sampled_nodes))}
286 relabeled_G = nx.relabel_nodes(sampled_G, node_mapping)
287
288 # Return the sampled graph and the inverse mapping (new_id -> original_id)
289 inverse_mapping = {new_id: old_id for old_id, new_id in node_mapping.items()}
290 return relabeled_G, inverse_mapping
291
292def graph_to_json_format(G):
293 """Convert a NetworkX graph to the required JSON format."""
294 result = []
295
296 for node_id in sorted(G.nodes):
297 node_data = {
298 "node_id": int(node_id),
299 "label": G.nodes[node_id]['label'],
300 "text": G.nodes[node_id]['text'],
301 "neighbors": sorted([int(n) for n in G.neighbors(node_id)]),
302 "mask": G.nodes[node_id]['mask']
303 }
304
305 result.append(node_data)
306
307 return result
308
309def sample_text_attribute_graph(input_file, output_file, percentage):
310 """Main function to sample a text attribute graph and preserve its properties."""
311 # Load the graph data
312 print(f"Loading graph from {input_file}...")
313 nodes = load_graph_from_json(input_file)
314
315 if not nodes:
316 print("Failed to load nodes from the input file.")
317 return None, None, None
318
319 print(f"Loaded {len(nodes)} nodes.")
320
321 # Build the NetworkX graph
322 print("Building graph...")
323 G = build_networkx_graph(nodes)
324 print(f"Built graph with {len(G.nodes)} nodes and {len(G.edges)} edges.")
325
326 # Analyze the original graph properties
327 print("Analyzing original graph properties...")
328 original_properties = analyze_graph_properties(G)
329
330 # Sample the graph
331 print(f"Sampling {percentage}% of the nodes...")
332 sampled_G, inverse_mapping = sample_graph_preserving_properties(G, percentage, original_properties)
333 print(f"Sampled graph has {len(sampled_G.nodes)} nodes and {len(sampled_G.edges)} edges.")
334
335 # Convert the sampled graph to JSON format
336 print("Converting sampled graph to JSON format...")
337 sampled_data = graph_to_json_format(sampled_G)
338
339 # Save the sampled graph
340 print(f"Saving sampled graph to {output_file}...")
341 with open(output_file, 'w') as f:
342 json.dump(sampled_data, f, indent=2)
343
344 # Analyze the sampled graph properties
345 print("Analyzing sampled graph properties...")
346 sampled_properties = analyze_graph_properties(sampled_G)
347
348 # Print comparison of original and sampled properties
349 print("\nComparison of Graph Properties:")
350 print(f"{'Property':<25} {'Original':<15} {'Sampled':<15}")
351 print("-" * 55)
352 print(f"{'Number of nodes':<25} {len(G.nodes):<15} {len(sampled_G.nodes):<15}")
353 print(f"{'Number of edges':<25} {len(G.edges):<15} {len(sampled_G.edges):<15}")
354 print(f"{'Density':<25} {original_properties['density']:.4f}{'':>10} {sampled_properties['density']:.4f}{'':>10}")
355
356 print("\nMask Distribution:")
357 print(f"{'Mask':<10} {'Original %':<15} {'Sampled %':<15}")
358 print("-" * 40)
359 for mask in sorted(set(original_properties['mask_distribution'].keys()) | set(sampled_properties['mask_distribution'].keys())):
360 orig_pct = original_properties['mask_distribution'].get(mask, 0) * 100
361 sampled_pct = sampled_properties['mask_distribution'].get(mask, 0) * 100
362 print(f"{mask:<10} {orig_pct:.2f}%{'':>9} {sampled_pct:.2f}%{'':>9}")
363
364 print("\nLabel Distribution:")
365 print(f"{'Label':<10} {'Original %':<15} {'Sampled %':<15}")
366 print("-" * 40)
367 for label in sorted(set(original_properties['label_distribution'].keys()) | set(sampled_properties['label_distribution'].keys())):
368 orig_pct = original_properties['label_distribution'].get(label, 0) * 100
369 sampled_pct = sampled_properties['label_distribution'].get(label, 0) * 100
370 print(f"{label:<10} {orig_pct:.2f}%{'':>9} {sampled_pct:.2f}%{'':>9}")
371
372 print("\nConnectivity:")
373 print(f"Connected components: {original_properties['connected_components']} (original) vs {sampled_properties['connected_components']} (sampled)")
374
375 return sampled_G, original_properties, sampled_properties
376
377def main():
378 """Command-line interface."""
379 if len(sys.argv) != 4:
380 print("Usage: python sample_graph.py input_file output_file percentage")
381 sys.exit(1)
382
383 input_file = sys.argv[1]
384 output_file = sys.argv[2]
385 try:
386 percentage = float(sys.argv[3])
387 if percentage <= 0 or percentage > 100:
388 raise ValueError("Percentage must be between 0 and 100")
389 except ValueError:
390 print("Error: Percentage must be a number between 0 and 100")
391 sys.exit(1)
392
393 sample_text_attribute_graph(input_file, output_file, percentage)
394
395if __name__ == "__main__":
396 main()