test_basic.py

"""
AgentTorch 기반 루머 확산 최소 테스트
SIR 모델 변형: Unaware → Spreader → Stifler (DK 모델)
- Unaware(U): 루머를 모름
- Spreader(S): 루머를 믿고 퍼뜨림
- Stifler(R): 루머를 접했으나 확산 중단 (사실확인 등)
"""

import torch
import torch.nn as nn
import numpy as np
import random

# ── 시뮬레이션 파라미터 ──────────────────────────────────────
N = 10_000          # 에이전트 수
STEPS = 50          # 시뮬레이션 스텝 수
BETA = 0.3          # 루머 전파율 (U → S)
GAMMA = 0.1         # 루머 소멸율 (S → R)
SEED_RATIO = 0.01   # 초기 Spreader 비율 (1%)

DEVICE = "cpu"
torch.manual_seed(42)

# ── 상태 정의 ───────────────────────────────────────────────
# 0 = Unaware, 1 = Spreader, 2 = Stifler
state = torch.zeros(N, dtype=torch.long)

# 초기 Spreader 설정
n_seed = int(N * SEED_RATIO)
seed_idx = torch.randperm(N)[:n_seed]
state[seed_idx] = 1

# ── 네트워크: 랜덤 그래프 (각 에이전트 평균 10 연결) ──────────
k = 10  # 평균 연결 수
edges = set()
for i in range(N):
    neighbors = random.sample(range(N), k)
    for j in neighbors:
        if i != j:
            edges.add((min(i,j), max(i,j)))

# adjacency: 각 노드의 이웃 리스트 (고정 크기 텐서로)
adj = [[] for _ in range(N)]
for i, j in edges:
    adj[i].append(j)
    adj[j].append(i)

# ── 시뮬레이션 루프 ─────────────────────────────────────────
history = {"U": [], "S": [], "R": []}

print(f"에이전트 수: {N:,}")
print(f"초기 Spreader: {n_seed} ({SEED_RATIO*100:.1f}%)")
print(f"전파율(β): {BETA}, 소멸율(γ): {GAMMA}")
print(f"{'Step':>4} | {'Unaware':>8} | {'Spreader':>8} | {'Stifler':>8}")
print("-" * 40)

for step in range(STEPS):
    new_state = state.clone()

    spreaders = (state == 1).nonzero(as_tuple=True)[0].tolist()

    for s in spreaders:
        # Spreader → Stifler 전환
        if random.random() < GAMMA:
            new_state[s] = 2
            continue
        # 이웃에게 루머 전파
        for neighbor in adj[s]:
            if state[neighbor] == 0 and random.random() < BETA:
                new_state[neighbor] = 1

    state = new_state

    n_u = (state == 0).sum().item()
    n_s = (state == 1).sum().item()
    n_r = (state == 2).sum().item()

    history["U"].append(n_u)
    history["S"].append(n_s)
    history["R"].append(n_r)

    if step % 10 == 0 or step == STEPS - 1:
        print(f"{step:>4} | {n_u:>8,} | {n_s:>8,} | {n_r:>8,}")

    if n_s == 0:
        print(f"\n[Step {step}] Spreader 소멸 → 시뮬레이션 종료")
        break

print("\n── 최종 결과 ──────────────────────────────────────────")
print(f"루머 노출(Spreader+Stifler): {N - history['U'][-1]:,} ({(N - history['U'][-1])/N*100:.1f}%)")
print(f"최대 동시 Spreader: {max(history['S']):,}")
print(f"최종 Unaware: {history['U'][-1]:,} ({history['U'][-1]/N*100:.1f}%)")

# ── 결과 시각화 ─────────────────────────────────────────────
try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    steps_ran = len(history["U"])
    x = range(steps_ran)

    plt.figure(figsize=(10, 5))
    plt.plot(x, history["U"], label="Unaware (모름)", color="steelblue")
    plt.plot(x, history["S"], label="Spreader (전파 중)", color="crimson")
    plt.plot(x, history["R"], label="Stifler (전파 중단)", color="seagreen")
    plt.xlabel("시뮬레이션 스텝")
    plt.ylabel("에이전트 수")
    plt.title(f"루머 확산 시뮬레이션 (N={N:,}, β={BETA}, γ={GAMMA})")
    plt.legend()
    plt.tight_layout()
    plt.savefig("result.png", dpi=150)
    print("\n결과 그래프 저장: result.png")
except Exception as e:
    print(f"\n시각화 생략: {e}")