SEIR python 易用代码

参考链接: https://zhuanlan.zhihu.com/p/388341199

fit 曲线 api

# 初始化模型
# 'seir', 'sis', 'sir'
method = 'seir'
problem = SEIR(delta=0.03, mu=0.06, lamda=0.3, number=1e6, i0=1e-3, e0=1e-3, method=method)

# fit 曲线
problem.fit(300)
problem.plot()

EGA 拟合给定数据 api

# 数据来自2023美赛C题
# 加载数据
x = data_loader()
# 拟合 x
problem.solver(x)

其中 solver 中 geatpy包的使用需要读者有一定基础。lb, ub是参数顺序delta, mu, lamda, number, i0, e0 的 下限 和 上限。其中 constraints 即为放缩倍数。

各个方法对比 api

# 各个方法对比
    problem.compare(len(x))

整体代码

import pandas as pd
import numpy as np
from scipy.integrate import odeint
import os
import matplotlib.pyplot as plt
import geatpy as ea
from copy import deepcopy

"""
参考代码: https://zhuanlan.zhihu.com/p/388341199
"""


class SEIR:
    delta, mu, lamda, number, i0, e0, y, method, method_ = [None] * 9

    # delta 发病率 mu 日治愈率 lamda 日接触率 number 总人数 i0 患病者比例的初始 e0潜伏者比例的初始
    def __init__(self, delta=0.03, mu=0.06, lamda=0.3, number=1e6, i0=1e-3, e0=1e-3, method='seir'):
        self.reload(delta, mu, lamda, number, i0, e0, method)

    def reload(self, delta, mu, lamda, number, i0, e0, method):
        self.delta, self.mu, self.lamda, self.number, self.i0, self.e0 = delta, mu, lamda, number, i0, e0
        self.y = None
        self.method_ = method
        try:
            method = ['seir', 'sis', 'sir'].index(method.lower())
        except:
            raise TypeError("Model choose incorrect, please choose ['seir', 'sis', 'sir']")
        self.method = [SEIR.dySEIR, SEIR.dySIS, SEIR.dySIR][method]

    # 输入天数
    def fit(self, tEnd):
        t = np.arange(0.0, tEnd, 1)  # (start,stop,step)
        s0 = 1 - self.i0  # 易感者比例的初值
        Y0 = (s0, self.e0, self.i0)  # 微分方程组的初值
        # odeint 数值解，求解微分方程初值问题
        ySEIR = odeint(self.method, Y0, t, args=(self.lamda, self.delta, self.mu))
        self.y = self.number * ySEIR
        return deepcopy(self.y)

    def compare(self, day):
        self.reload(delta=self.delta, mu=self.mu, lamda=self.lamda, number=self.number, i0=self.i0, e0=self.e0, method='seir')
        ySEIR = self.fit(day)
        self.reload(delta=self.delta, mu=self.mu, lamda=self.lamda, number=self.number, i0=self.i0, e0=self.e0, method='sir')
        ySIR = self.fit(day)
        self.reload(delta=self.delta, mu=self.mu, lamda=self.lamda, number=self.number, i0=self.i0, e0=self.e0, method='sis')
        ySIS = self.fit(day)
        self.reload(delta=self.delta, mu=0, lamda=self.lamda, number=self.number, i0=self.i0, e0=self.e0, method='sis')
        ySI = self.fit(day)

        plt.title("Comparison among SI, SIS, SIR and SEIR models")
        plt.xlabel('t-youcans')
        plt.plot(ySI[:, 2], 'cadetblue', label='i(t)-SI')
        plt.plot(ySIS[:, 2], 'steelblue', label='i(t)-SIS')
        plt.plot(ySIR[:, 2], 'cornflowerblue', label='i(t)-SIR')
        plt.plot(ySEIR[:, 2], '-', color='m', label='i(t)-SEIR')
        plt.legend(loc='right')
        plt.show()

    def plot(self, title=None):
        y = self.y
        if y is None:
            raise TypeError("y is None, please fit first")
        s, e, i = y[:, 0], y[:, 1], y[:, 2]
        if self.method is SEIR.dySEIR:
            plt.plot(s, '--', color='darkviolet', label='s')
            plt.plot(e, '-.', color='orchid', label='e')
            plt.plot(i, '-', color='m', label='i')
            plt.plot(self.number - s - e - i, ':', color='palevioletred', label='r')
        elif self.method is SEIR.dySIS:
            # plt.plot(s, '--', color='darkviolet', label='s')
            plt.plot(i, '-', color='m', label='i')
        else:
            plt.plot(s, '--', color='darkviolet', label='s')
            plt.plot(i, '-', color='m', label='i')
            plt.plot(self.number - s - i, ':', color='palevioletred', label='r')
        plt.xlabel('t')
        if title is not None:
            plt.title(title)
        plt.ylim(0)
        plt.legend(loc='upper right')
        plt.tight_layout()
        plt.show()

    @staticmethod
    def dySEIR(y, t, lamda, delta, mu):  # SEIR 模型，导数函数
        s, e, i = y
        ds_dt = -lamda * s * i  # ds/dt = -lamda*s*i
        de_dt = lamda * s * i - delta * e  # de/dt = lamda*s*i - delta*e
        di_dt = delta * e - mu * i  # di/dt = delta*e - mu*i
        return np.array([ds_dt, de_dt, di_dt])

    @staticmethod
    def dySIS(y, t, lamda, _, mu):  # SI/SIS 模型，导数函数
        _, _, i = y
        dy_dt = lamda * i * (1 - i) - mu * i  # di/dt = lamda*i*(1-i)-mu*i
        return [_, _, dy_dt]

    @staticmethod
    def dySIR(y, t, lamda, _, mu):  # SIR 模型，导数函数
        s, _, i = y
        ds_dt = -lamda * s * i  # ds/dt = -lamda*s*i
        di_dt = lamda * s * i - mu * i  # di/dt = lamda*s*i-mu*i
        return np.array([ds_dt, _, di_dt])

    def solver(self, x, ndind=50, maxgen=100):
        # delta, mu, lamda, number, i0, e0
        class seir(ea.Problem):
            def __init__(self, base):
                if "results" not in os.listdir():
                    os.makedirs("./results")
                name = 'SEIR'
                M = 1
                maxormins = [-1] * M
                Dim = 6
                varTypes = [1] * Dim
                lb = [10, 10, 10, 1e4, 1, 1]
                ub = [300, 100, 100, 1e7, 150, 150]
                lbin = [1] * Dim
                ubin = [1] * Dim
                self.maxans = float("-inf")
                self.maxanses = []
                self.averans = []
                self.base = base
                self.n = len(x)
                ea.Problem.__init__(self, name, M, maxormins, Dim, varTypes, lb, ub, lbin, ubin)

            # 目标函数即神经网络返回值
            def evalVars(self, Vars):
                ans = np.zeros(len(Vars), dtype=float).reshape(len(Vars), 1)
                for i in range(len(Vars)):
                    self.base.reload(*self.constraints(Vars[i]))
                    ans[i] = r2(self.base.fit(self.n)[:, 2], x)
                self.maxans = temp if (temp := ans.max()) > self.maxans else self.maxans
                self.maxanses.append(self.maxans)
                self.averans.append(ans.mean())
                return ans

            def constraints(self, var):
                return var[0] / 1000, var[1] / 1000, var[2] / 100, var[3], var[4] / 1000, var[5] / 1000, self.base.method_

        problem = seir(self)
        myAlgorithm = ea.soea_EGA_templet(problem, ea.Population(Encoding='RI', NIND=ndind), MAXGEN=maxgen, logTras=0)
        myAlgorithm.drawing = 1
        res = ea.optimize(myAlgorithm, seed=1, verbose=True, drawing=True, outputMsg=True, drawLog=True, saveFlag=True, dirName='results')
        res = problem.constraints(res["Vars"][0])
        print('parameter', res)
        self.reload(*res)
        self.fit(problem.n)
        self.plot()


def r2(y_test, y):
    return 1 - ((y_test - y) ** 2).sum() / ((y.mean() - y) ** 2).sum()


def data_loader():
    df = pd.read_excel("Problem_C_Data_Wordle_original.xlsx")
    x = np.array(df["Number of  reported results"])
    return x


if __name__ == "__main__":
    # 初始化模型
    method = 'seir'
    problem = SEIR(delta=0.03, mu=0.06, lamda=0.3, number=1e6, i0=1e-3, e0=1e-3, method=method)

    # fit 曲线
    problem.fit(100)
    problem.plot()

    # 加载数据
    x = data_loader()
    # 拟合 x
    problem.solver(x)

    # 各个方法对比
    problem.compare(len(x))