Rumor vs Counterwave

from tys import probe, progress


def simulate(cfg: dict):
    """Simulate competing contagions on a random social graph."""

    import random
    import simpy

    env = simpy.Environment()
    rng = random.Random(cfg.get("seed", 12345))

    n = cfg["num_nodes"]               # number of users in the network
    degree = cfg["avg_degree"]         # average connections per user
    rumor_p = cfg["rumor_p"]           # chance a rumor holder forwards to each neighbour
    rumor_decay = cfg["rumor_decay"]   # daily chance a rumor fades
    correction_p = cfg["correction_p"] # chance a correction holder forwards
    correction_decay = cfg["correction_decay"]
    counter_delay = cfg["counter_delay"]   # days before fact-check seeding
    rumor_seed = cfg["rumor_seed"]         # initial rumor carriers
    counter_seed = cfg["counter_seed"]     # initial correction carriers
    sim_time = cfg["sim_time"]

    neighbors = {i: set() for i in range(n)}
    for i in range(n):
        while len(neighbors[i]) < degree:
            j = rng.randrange(n)
            if j != i and j not in neighbors[i]:
                neighbors[i].add(j)
                neighbors[j].add(i)

    rumor = set(rng.sample(range(n), rumor_seed))
    correction = set()

    initial_rumor = len(rumor)
    initial_correction = counter_seed
    rumor_spread = 0
    correction_spread = 0

    rumor_peak = len(rumor)
    rumor_peak_time = 0
    correction_half_time = None

    done = env.event()

    def dynamics():
        nonlocal rumor, correction, rumor_peak, rumor_peak_time, correction_half_time
        nonlocal rumor_spread, correction_spread

        for t in range(sim_time):
            new_rumor = set()
            new_correction = set()

            if t == counter_delay:
                seeds = rng.sample([x for x in range(n) if x not in correction], counter_seed)
                new_correction.update(seeds)

            for node in list(rumor):
                for neigh in neighbors[node]:
                    if neigh not in rumor and neigh not in correction:
                        if rng.random() < rumor_p:
                            new_rumor.add(neigh)
                            rumor_spread += 1
                if rng.random() < rumor_decay:
                    rumor.remove(node)

            for node in list(correction):
                for neigh in neighbors[node]:
                    if neigh not in correction:
                        if neigh in rumor:
                            rumor.discard(neigh)
                        if rng.random() < correction_p:
                            new_correction.add(neigh)
                            correction_spread += 1
                if rng.random() < correction_decay:
                    correction.remove(node)

            rumor.update(new_rumor - correction)
            correction.update(new_correction)

            probe("rumor", env.now, len(rumor))
            probe("correction", env.now, len(correction))

            if len(rumor) > rumor_peak:
                rumor_peak = len(rumor)
                rumor_peak_time = t
            if correction_half_time is None and len(correction) >= n / 2:
                correction_half_time = t

            progress(int(100 * t / sim_time))
            yield env.timeout(1)

        progress(100)
        final_coverage = len(correction) / n * 100
        time_gap = None
        if correction_half_time is not None:
            time_gap = correction_half_time - rumor_peak_time

        done.succeed({
            "peak_rumor": rumor_peak,
            "correction_coverage": final_coverage,
            "time_gap": time_gap,
            "R0_rumor": rumor_spread / initial_rumor if initial_rumor else 0,
            "R0_correction": correction_spread / initial_correction if initial_correction else 0,
        })

    env.process(dynamics())
    env.run(until=done)
    return done.value


def requirements():
    return {
        "builtin": ["micropip", "pyyaml"],
        "external": ["simpy==4.1.1"],
    }