题目描述

1import networkx as nx
2import matplotlib.pyplot as plt
3
4# 创建一个包含500个节点、平均度为4、重连概率为0.5的合成图
5G = nx.watts_strogatz_graph(n=500, k=4, p=0.5, seed=42)
6
7# 使用matplotlib可视化图
8nx.draw(G, with_labels=False, node_size=50, width=0.5)

1import networkx as nx
2import numpy as np
3import random
4import torch
5import matplotlib.pyplot as plt
6
7G = nx.connected_watts_strogatz_graph(500, 4, 0.5) # 合成图
8
9# 为节点添加节点属性
10for node in G.nodes():
11    G.nodes[node]['age'] = np.random.randint(18, 65)
12    G.nodes[node]['sex'] = np.random.choice(['male', 'female'])
13    G.nodes[node]['location'] = np.random.choice(['urban', 'rural'])
14    G.nodes[node]['occupation'] = np.random.choice(['student', 'employee', 'self-employed', 'unemployed'])
15    G.nodes[node]['tested'] = np.random.choice([0, 1], p=[0.8, 0.2])
16    G.nodes[node]['symptoms'] = np.random.choice([0, 1], p=[0.5, 0.5])
17    G.nodes[node]['vaccinated'] = np.random.choice([0, 1], p=[0.6, 0.4])
18    G.nodes[node]['mobility'] = random.randint(1,10)
19
20# 添加接触数据
21for (u,v) in G.edges():
22    G[u][v]['contact'] = random.randint(1,10)
23
24# 使用某些逻辑添加标签
25for node in G.nodes():
26    mobility = G.nodes[node]['mobility']
27    tested = G.nodes[node]['tested']
28    vaccinated = G.nodes[node]['vaccinated']
29    symptoms = G.nodes[node]['symptoms']
30
31    if mobility > 7 and (tested == 1 or (vaccinated == 1 and random.random() > 0.5) or symptoms == 1 and random.random() < 0.5):
32        G.nodes[node]['label'] = 'infected'
33    else:
34        G.nodes[node]['label'] = 'not infected'
35
36# 使用颜色作为标签绘制图
37plt.figure(figsize=(10, 6),dpi=300)
38color_map = {'infected': 'red', 'not infected': 'green'}
39nx.draw(G, with_labels=False, node_size=50, width=0.5, node_color=[color_map[G.nodes[node]['label']] for node in G.nodes()])