Rustでネットワーク上のウイルス感染シミュレーションを作った

# ネットワーク上のウイルス感染について

古典的な感染拡大のモデルとして[SIRモデル](https://ja.wikipedia.org/wiki/SIR%E3%83%A2%E3%83%87%E3%83%AB)、[SEIRモデル](https://ja.wikipedia.org/wiki/SEIR%E3%83%A2%E3%83%87%E3%83%AB)などがある。
実際には人同士のつながりは[複雑ネットワーク](https://ja.wikipedia.org/wiki/%E8%A4%87%E9%9B%91%E3%83%8D%E3%83%83%E3%83%88%E3%83%AF%E3%83%BC%E3%82%AF)になっている。wikipediaから引用すると以下のような性質が挙げられる。

> 現実世界に存在するネットワークは多様であり、巨大で複雑な構造を有しているが、一定の共通する性質を見出すことができる。それらの性質は「スケールフリー性」（次数分布のべき乗則）、「スモールワールド性」、「クラスター性」と呼ばれている。「スケールフリー性」（次数分布のべき乗則）とは、例えば、一部の人は非常にたくさんの知人を持っているが、大多数の人々の知人の数は少ないという性質である。「スモールワールド性」とは、例えば、「世間は狭い」と言われるように、一見赤の他人に見えても、実際は中間に少数の人を介するだけでつながっているという性質である。「クラスター性」とは、例えば、「自分と知人Aさんがいるときに、自分もAさんもどちらも知っている共通の知人Bさんのような人が1人もいない」という状況はまずありえないという性質である。

ウイルス感染も実際には人同士のつながりに応じて感染拡大の仕方が変わる。
"ネットワーク 感染症モデル"でググると解析や分析結果がたくさん出てくる。

本記事ではRustでシミュレーションを作ったので、そのコードを紹介する。


# 実装

```rust
extern crate yaml_rust;

use rand::distributions::Uniform;
use rand::Rng;
use regex::Regex;
use std::cell::RefCell;
use std::env;
use std::fs::{self, File};
use std::io::BufRead;
use std::io::BufReader;
use std::rc::Rc;
use yaml_rust::YamlLoader;

#[derive(Debug, Clone)]
struct Parameters {
    num_agents: u32,
    initial_infectious_rate: f64,
    simulation_time: u32,
    infection_rate: f64,
    recover_rate: f64,
    connection_file: String,
}

impl Parameters {
    fn from_file(filename: String) -> Self {
        let yaml_str = fs::read_to_string(&filename).unwrap();
        let docs = YamlLoader::load_from_str(&yaml_str.to_string()).unwrap();
        let doc = &docs[0]; // get first doc
        Parameters {
            num_agents: doc["num_agents"].as_i64().unwrap() as u32,
            initial_infectious_rate: doc["initial_infectious_rate"].as_f64().unwrap(),
            simulation_time: doc["simulation_time"].as_i64().unwrap() as u32,
            infection_rate: doc["infection_rate"].as_f64().unwrap(),
            recover_rate: doc["recover_rate"].as_f64().unwrap(),
            connection_file: doc["connection_file"].as_str().unwrap().to_string(),
        }
    }
}

struct Settings {
    connections: Vec<(usize, usize)>,
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum State {
    Susceptible,
    Infected,
    Recovered,
}

struct Agent {
    current_state: State,
    past_state: State,
    infection_rate: f64,
    recover_rate: f64,
    neighbors: Vec<Rc<RefCell<Agent>>>,
}

impl Agent {
    fn new(infection_rate: f64, recover_rate: f64) -> Self {
        Agent {
            current_state: State::Susceptible,
            past_state: State::Susceptible,
            infection_rate: infection_rate,
            recover_rate: recover_rate,
            neighbors: Vec::new(),
        }
    }

    fn add_neighbor(&mut self, neighbors: Rc<RefCell<Agent>>) {
        self.neighbors.push(neighbors);
    }

    fn update<T: rand::Rng>(&mut self, rng: &mut T) {
        match self.current_state {
            State::Susceptible => {
                let num_infectious_neighbors = self
                    .neighbors
                    .iter()
                    .filter(|x| x.borrow().current_state == State::Infected)
                    .count();
                for _ in 0..num_infectious_neighbors {
                    let p: f64 = rng.gen();
                    if p < self.infection_rate {
                        self.current_state = State::Infected;
                        break;
                    }
                }
            }
            State::Infected => {
                let p: f64 = rng.gen();
                if p < self.recover_rate {
                    self.current_state = State::Recovered;
                }
            }
            State::Recovered => {}
        }
    }

    fn initialize_state(&mut self, new_state: State) {
        self.current_state = new_state;
        self.past_state = self.current_state;
    }

    fn save_state(&mut self) {
        self.past_state = self.current_state;
    }
}

struct AgentNetwork {
    agents: Vec<Rc<RefCell<Agent>>>,
    changed_agents: Vec<Rc<RefCell<Agent>>>,
}

impl AgentNetwork {
    fn new(
        num_agents: u32,
        infection_rate: f64,
        recover_rate: f64,
        connections: Vec<(usize, usize)>,
    ) -> Self {
        let mut ag_state = AgentNetwork {
            agents: Vec::with_capacity(num_agents as usize),
            changed_agents: Vec::new(),
        };
        for _ in 0..num_agents {
            ag_state.agents.push(Rc::new(RefCell::new(Agent::new(
                infection_rate,
                recover_rate,
            ))));
        }
        ag_state.set_neighbors(connections.to_vec());
        ag_state
    }

    fn set_neighbors(&mut self, connections: Vec<(usize, usize)>) {
        for (s, d) in connections {
            self.agents[s]
                .borrow_mut()
                .add_neighbor(Rc::clone(&self.agents[d]));
            self.agents[d]
                .borrow_mut()
                .add_neighbor(Rc::clone(&self.agents[s]));
        }
    }

    fn update<T: rand::Rng>(&mut self, rng: &mut T) {
        self.agents.iter().for_each(|x| x.borrow_mut().update(rng));
    }

    fn save_state(&mut self) {
        self.changed_agents
            .iter_mut()
            .for_each(|x| x.borrow_mut().save_state());
        self.changed_agents.clear();
    }

    fn count_state(&self, state: State) -> usize {
        self.agents
            .iter()
            .filter(|x| x.borrow().current_state == state)
            .count()
    }

    fn rate(&self, state: State) -> f64 {
        self.count_state(state) as f64 / self.agents.len() as f64
    }
}

struct Simulation {
    num_agents: u32,
    initial_infectious_rate: f64,
    simulation_time: u32,
    agent_network: AgentNetwork,
    rng: rand::rngs::ThreadRng,
}

impl Simulation {
    fn new(param: &Parameters, settings: Settings) -> Self {
        Simulation {
            num_agents: param.num_agents,
            initial_infectious_rate: param.initial_infectious_rate,
            simulation_time: param.simulation_time,
            agent_network: AgentNetwork::new(
                param.num_agents,
                param.infection_rate,
                param.recover_rate,
                settings.connections,
            ),
            rng: rand::thread_rng(),
        }
    }

    fn initialize(&mut self) {
        let num_initial_infectious =
            ((self.num_agents as f64) * self.initial_infectious_rate) as usize;
        let range = Uniform::from(0..self.num_agents);
        self.rng
            .sample_iter(range)
            .take(num_initial_infectious)
            .for_each(|i| {
                self.agent_network.agents[i as usize]
                    .borrow_mut()
                    .initialize_state(State::Infected);
            });
    }

    fn update(&mut self) {
        self.agent_network.update(&mut self.rng);
        self.agent_network.save_state();
    }

    fn rate(&self, state: State) -> f64 {
        self.agent_network.rate(state)
    }

    fn run(&mut self) -> Vec<(f64, f64, f64)> {
        let mut result = Vec::with_capacity(self.simulation_time as usize);

        self.initialize();
        for _ in 0..self.simulation_time {
            self.update();
            result.push((
                self.rate(State::Susceptible),
                self.rate(State::Infected),
                self.rate(State::Recovered),
            ));
        }

        result
    }
}

struct Utils {}

impl Utils {
    fn load_connections<T: BufRead>(mut buf: T) -> Vec<(usize, usize)> {
        let re = Regex::new(r"\((\d+), (\d+)\)").unwrap();
        let mut line = String::new();
        let mut connections = Vec::new();
        while buf.read_line(&mut line).unwrap() > 0 {
            // each line can include multiple (s,d) pairs
            for cap in re.captures_iter(&line) {
                // entire match is stored in cap[0]
                let s: usize = cap[1].parse().unwrap();
                let d: usize = cap[2].parse().unwrap();
                connections.push((s, d));
            }
            line.clear();
        }
        connections
    }

    // fn load_connections_stdin() -> Vec<(usize, usize)> {
    //     load_connections(io::stdin().lock())
    // }

    fn load_connections_file(filename: &String) -> Vec<(usize, usize)> {
        let f = File::open(filename).unwrap();
        Utils::load_connections(BufReader::new(f))
    }

    fn parse_arg() -> String {
        let args: Vec<String> = env::args().collect();
        if args.len() < 2 {
            println!("Usage {:?} <parameter_file>", args[0]);
            panic!("Invalid args");
        }
        args[1].clone()
    }

    fn show(values: &Vec<(f64, f64, f64)>) {
        for v in values {
            println!("{}, {}, {}", v.0, v.1, v.2);
        }
    }
}

fn main() {
    let file = Utils::parse_arg();
    let param = Parameters::from_file(file);
    let settings = Settings {
        connections: Utils::load_connections_file(&param.connection_file),
    };
    let mut sim = Simulation::new(&param, settings);
    let result = sim.run();
    Utils::show(&result);
}
```

# 仕様など

**パラメータ**
yamlで指定する。
num_agents: シミュレーション対象の人数
initial_infectious_rate: 最初に感染している人の割合
simulation_time: シミュレーションステップ数
infection_rate: 各ステップで非感染者が感染者一人あたりから感染する確率
recover_rate: 感染者が各ステップごとに回復する確率
connection_file: ネットワークの構成を書いたファイル（辺を列挙する）

サンプル

```yaml
num_agents: 100000
initial_infectious_rate: 0.0001
simulation_time: 500
infection_rate: 0.03
recover_rate: 0.1
connection_file: "path/to/connection/file"
```

**コネクションファイル**
ネットワークの辺を列挙したファイル
頂点は数字で表し、辺の両端の頂点を`s,t`として`(s,t)`で辺を表す。
実装では正規表現で構文解析している。

頂点数3の完全グラフでは以下のようになる

```
(0, 1), (1, 2), (1, 2)
```

**結果**
標準出力に直接書き出す。

# pythonでネットワーク生成＋結果のグラフ表示

pythonだとnetworkxというライブラリで複雑ネットワークが簡単に作れる。
ネットワーク生成とグラフを保存するpythonのコードは次のようになる。

```python
#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import networkx as nx
import yaml
import subprocess

def write_file(g, fname):
    with open(fname, "w") as f:
        f.write(str(g.edges))

def write_yaml_file(data, fname):
    with open(fname, "w") as f:
        f.write(yaml.dump(data))

def set_params(param_file, con_file, g, data):
    write_yaml_file(data, param_file)
    write_file(g, con_file)

def make_data(param_file, con_file):
    n = 100000

    data = {
        "num_agents": n,
        "initial_infectious_rate": 0.0001,
        "simulation_time": 500,
        "infection_rate": 0.03,
        "recover_rate": 0.1,
        "connection_file": con_file,
    }

    g = nx.watts_strogatz_graph(n, k=10, p=0.3)
    #g = nx.barabasi_albert_graph(n, m=10)

    set_params(param_file, con_file, g, data)

def run(param_file, result_file):
    cmd = "cargo run --release {} > {}".format(param_file, result_file)
    subprocess.run(cmd, shell=True)

def plot(result_file, img_file, max_rate=1.0, max_time=None):
    with open(result_file, "r") as f:
        vals = np.array([[float(_.strip()) for _ in line.split(",")] for line in f.readlines()])
    if max_time is None:
        max_time = len(vals)
    s = vals[:max_time,0]
    i = vals[:max_time,1]
    r = vals[:max_time,2]
    plt.plot(s, label="s")
    plt.plot(i, label="i")
    plt.plot(r, label="r")
    plt.ylim(0, max_rate)
    plt.savefig(img_file)

def main():
    con_file = "data/connection" 
    param_file = "data/param.yaml"
    result_file = "result/rslt.dat"
    img_file = "result/img.png"

    make_data(param_file, con_file)
    run(param_file, result_file)
    plot(result_file, img_file)

if __name__ == "__main__":
    main()
```

# 結果

結果をいくつか示す。

ネットワークのモデル: [WSモデル](https://ja.wikipedia.org/wiki/%E8%A4%87%E9%9B%91%E3%83%8D%E3%83%83%E3%83%88%E3%83%AF%E3%83%BC%E3%82%AF#%E3%83%AF%E3%83%83%E3%83%84%E3%83%BB%E3%82%B9%E3%83%88%E3%83%AD%E3%82%AC%E3%83%83%E3%83%84%E3%83%A2%E3%83%87%E3%83%AB)

- WSモデルのパラメータ:
  - 平均次数k: 10
  - 辺の入れ替え確率p: 0.3
- シミュレーション対象の人数: 100000人
- 最初に感染している人の割合: 0.0001 (10人)
- 感染確率: 0.03
- 回復確率: 0.1

![img.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/156722/035add9f-bb49-7573-a313-b93181598d52.png)

- WSモデルのパラメータ:
  - 平均次数k: 10
  - 辺の入れ替え確率p: 0.3
- シミュレーション対象の人数: 100000人
- 最初に感染している人の割合: 0.0001 (10人)
- 感染確率: **0.02**
- 回復確率: 0.1

※太字が一つ目から変更したパラメータ

![img.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/156722/09b46c9d-fd51-11c4-bde7-2d25ba80403b.png)

- WSモデルのパラメータ:
  - 平均次数k: **20**
  - 辺の入れ替え確率p: 0.3
- シミュレーション対象の人数: 100000人
- 最初に感染している人の割合: 0.0001 (10人)
- 感染確率: **0.02**
- 回復確率: 0.1

※太字が一つ目から変更したパラメータ

![img.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/156722/42565451-bb30-edf2-ea14-e25e34e7ca0e.png)

- WSモデルのパラメータ:
  - 平均次数k: **5**
  - 辺の入れ替え確率p: 0.3
- シミュレーション対象の人数: **1000000**人
- 最初に感染している人の割合: 0.0001 (**100**人)
- 感染確率: **0.1**
- 回復確率: 0.1

※太字が一つ目から変更したパラメータ

![img.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/156722/9a469366-c417-2268-2ee0-1ae3408f3bc3.png)

- WSモデルのパラメータ:
  - 平均次数k: **5**
  - 辺の入れ替え確率p: 0.3
- シミュレーション対象の人数: **1000000**人
- 最初に感染している人の割合: 0.0001 (**100**人)
- 感染確率: **0.05**
- 回復確率: 0.1

※太字が一つ目から変更したパラメータ

![img.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/156722/f4a37dcd-2850-0167-afea-58b1bb6089a7.png)


平均次数や感染確率(一人あたりから感染する確率)によって大きく結果が変わるが、人同士の接触が多いと感染者がいなくなるまでに、大部分の人が一度は感染することになる。