File size: 3,092 Bytes
3de20f8 | 1 2 3 4 5 6 7 8 9 10 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 | from __future__ import annotations
import json
from pathlib import Path
import numpy as np
from sklearn.model_selection import train_test_split
PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
def build_graph(nodes: int = 600, seed: int = 2033) -> dict[str, np.ndarray]:
rng = np.random.default_rng(seed)
subnets = np.repeat(np.arange(6), nodes // 6)
adjacency = np.zeros((nodes, nodes), dtype=np.float32)
for left in range(nodes):
same_subnet = subnets == subnets[left]
probabilities = np.where(same_subnet, 0.045, 0.0025)
links = rng.random(nodes) < probabilities
links[: left + 1] = False
adjacency[left, links] = 1
adjacency = np.maximum(adjacency, adjacency.T)
compromised = np.zeros(nodes, dtype=bool)
seeds = rng.choice(nodes, size=14, replace=False)
compromised[seeds] = True
for _ in range(4):
exposure = adjacency @ compromised.astype(np.float32)
infection_probability = 1 - np.exp(-0.22 * exposure)
new_infections = (rng.random(nodes) < infection_probability) & ~compromised
compromised |= new_infections
base = rng.normal(0, 1, (nodes, 8)).astype(np.float32)
labels = compromised.astype(np.int64)
signal = labels[:, None].astype(np.float32)
features = base.copy()
features[:, 0:1] += signal * rng.normal(1.0, 0.5, (nodes, 1))
features[:, 1:2] += signal * rng.normal(0.8, 0.6, (nodes, 1))
features[:, 2:3] += signal * rng.normal(0.7, 0.6, (nodes, 1))
features[:, 3:4] += signal * rng.normal(0.5, 0.7, (nodes, 1))
features[:, 4] += subnets * 0.12
indices = np.arange(nodes)
train, remainder = train_test_split(
indices,
test_size=0.40,
stratify=labels,
random_state=seed,
)
validation, test = train_test_split(
remainder,
test_size=0.50,
stratify=labels[remainder],
random_state=seed,
)
return {
"features": features,
"adjacency": adjacency,
"labels": labels,
"subnets": subnets.astype(np.int64),
"train_indices": train,
"validation_indices": validation,
"test_indices": test,
"seed_nodes": seeds.astype(np.int64),
}
def main() -> None:
DATA_DIR.mkdir(parents=True, exist_ok=True)
graph = build_graph()
np.savez_compressed(DATA_DIR / "meshgraph.npz", **graph)
manifest = {
"nodes": len(graph["labels"]),
"edges": int(graph["adjacency"].sum() // 2),
"features": graph["features"].shape[1],
"subnets": len(np.unique(graph["subnets"])),
"compromised_rate": float(graph["labels"].mean()),
"train_nodes": len(graph["train_indices"]),
"validation_nodes": len(graph["validation_indices"]),
"test_nodes": len(graph["test_indices"]),
"path": "meshgraph.npz",
}
(DATA_DIR / "manifest.json").write_text(
json.dumps(manifest, indent=2),
encoding="utf-8",
)
print(json.dumps(manifest, indent=2))
if __name__ == "__main__":
main()
|