tmp20022301_mdl.py

#! python3
# coding: utf-8

import math
import numpy as np
import matplotlib.pyplot as plt
from functools import reduce
from pprint import pprint as ppr

_gv = lambda v: lambda d: lambda k: d[k] if k in d else v
_rfi = lambda fi: (fi + 1) % 2

class c_model:

    def __init__(self, ps, init):
        gp = _gv(0)(ps)
        self.p = {
            'S2EbI': (gp('S2EbIf'), gp('S2EbIi')),
            'S2EbE': (gp('S2EbEf'), gp('S2EbEi')),
            'E2I': gp('E2I'),
            'I2R': (gp('I2Rf'), gp('I2Ri')),
            'I2E': (gp('I2Ef'), gp('I2Ei')),
            'I2D': (gp('I2Df'), gp('I2Di')),
            'E2R': gp('E2R'),
            'R2S': gp('R2S'),
            'Sfi': (gp('Sf2i'), gp('Si2f')),
            'Efi': (gp('Ef2i'), gp('Ei2f')),
            'Rfi': (0, gp('Ri2f')),
            'Ifi': (gp('If2i'), 0),
            'rc': (gp('rcf'), gp('rci')),
        }
        self.s = {
            'N': [init['N'], 0],
            'I': [init['I'], 0],
            'E': [init['E'], 0],
            'R': [0, 0],
            'S': [init['N'] - init['I'] - init['E'], 0],
            'D': 0,
            'r': [init['rf'], init['ri']],
            #'imx': init['imx'],
        }
        self.sk = [
            i for k in (
                (n + 'f', n + 'i') for n in list('SEIR')
            ) for i in k
        ] + ['D']

    def dS2E(self, fi):
        #_p = lambda p, n: p * n
        _p = lambda p, n: 1 - (1 - p) ** n
        rfi = _rfi(fi)
        rn = self.s['r'][0] * self.s['N'][0] + self.s['r'][1] * self.s['N'][1]
        ps = [
            self.s['r'][fi] * self.p['S2EbI'][fi] * self.s['I'][fi] / rn,
            self.s['r'][fi] * self.p['S2EbE'][fi] * self.s['E'][fi] / rn,
            self.s['r'][rfi] * self.p['S2EbI'][rfi] * self.s['I'][rfi] / rn,
            self.s['r'][rfi] * self.p['S2EbE'][rfi] * self.s['E'][rfi] / rn,
        ]
        #if ps[-1] > 1:
        #    ppr((ps, self.s['r'], self.p['S2EbE'][rfi], self.s['E'], self.s['N']))
        pa = reduce(lambda v, a: v + _p(a, self.s['r'][fi]), ps)
        return self.s['S'][fi] * pa

    def dE2I(self, fi):
        return self.p['E2I'] * self.s['E'][fi]

    def dI2Rf(self, fi):
        return self.p['I2R'][fi] * self.s['I'][fi]

    def dI2Ef(self, fi):
        return self.p['I2E'][fi] * self.s['I'][fi]

    def dI2D(self, fi):
        return self.p['I2D'][fi] * self.s['I'][fi]

    def dE2R(self, fi):
        return self.p['E2R'] * self.s['E'][fi]

    def dRf2Sf(self):
        return self.p['R2S'] * self.s['R'][0]

    def dSf2Si(self):
        return self.dI2I(0) * self.p['Sfi'][0] * (1 - self.p['S2EbI'][0])

    def dSi2Sf(self):
        return self.p['Sfi'][1] * self.s['S'][1]

    def dEf2Ei(self):
        return self.dI2I(0) * self.p['Efi'][0] * self.p['S2EbI'][0]

    def dEi2Ef(self):
        return self.p['Efi'][1] * self.s['E'][1]

    def dRi2Rf(self):
        return self.p['Rfi'][1] * self.s['R'][1]

    def dI2I(self, fi):
        return self.p['Ifi'][fi] * self.s['I'][fi]

    def _drf(self, ds):
        #if(self.p['rc'][0]):
        #    ppr(ds)
        #    print(self.s['r'][0], ds['Ii'], self.s['I'][1])
        return max( - self.s['r'][0] + self.s['r'][1],
            - self.p['rc'][0] * ds['Ii'] / self.s['I'][1] if abs(self.s['I'][1]) > abs(ds['Ii']) else 0)

    def drf(self, ds):
        if not hasattr(self, '_lst_dii'):
            self._lst_dii = ds['Ii']
        ddi = (ds['Ii'] - self._lst_dii) / abs(ds['Ii'] + 1 )
        di = ds['Ii'] / self.s['I'][1] if abs(self.s['I'][1]) > abs(ds['Ii']) else 0
        self._lst_dii = ds['Ii']
        rrng = self.s['r'][1], self.s['r'][1] * 30
        drmn, drmx = [i - self.s['r'][0] for i in rrng]
        return min(drmx, max( drmn,
            - self.p['rc'][0] * (math.copysign(di ** 2, di) + ddi)))
        

    def update(self, dbg = False):
        df = lambda s, d: 'd' + s + '2' + d
        fi = lambda k: 0 if k[1] == 'f' else 1
        srfi = lambda k: k if len(k) < 2 else k[0] + 'f' if k[1] == 'i' else k[0] + 'i'
        rd = {k: {k: 0 for k in self.sk} for k in self.sk}
        for ks in self.sk:
            for kd in self.sk:
                if ks == kd:
                    continue
                if hasattr(self, df(ks, kd)):
                    rd[ks][kd] = getattr(self, df(ks, kd))()
                if hasattr(self, df(ks[0], kd)):
                    rd[ks][kd] = getattr(self, df(ks[0], kd))(fi(ks))
                if hasattr(self, df(ks[0], kd[0])):
                    if ks[0] != kd[0] != 'D' and fi(ks) != fi(kd):
                        continue
                    rd[ks][kd] = getattr(self, df(ks[0], kd[0]))(fi(ks))
        if dbg and 0:
            ppr(rd)
        ds = {k: 0 for k in self.sk}
        for ks in self.sk:
            d = 0
            for kd in self.sk:
                if ks == kd:
                    continue
                dd = rd[kd][ks] - rd[ks][kd]
                d += dd
                ds[ks] += dd
            if ks == 'D':
                self.s[ks] += d
            else:
                sn, si = ks[0], fi(ks)
                self.s[sn][si] += d
                #try:
                #    d < 0
                #except:
                #    print(d, ks, kd)
                #    ppr(rd)
                if self.s[sn][si] < 0:
                    if dbg:
                        ppr(self.s)
                    rv = self.s[sn][si]
                    self.s[sn][si] = 0
                    self.s[sn][_rfi(si)] += rv
                    self.s['N'][si] -= rv
                    self.s['N'][_rfi(si)] += rv
                    ds[ks] -= rv
                    ds[srfi(ks)] += rv
                    if dbg and not (sn in 'SE' and si == 0):
                        print(ks, ds)
                        ppr(self.s)
                        ppr(rd)
                        assert(0)
                    assert(sn in 'SE' and si == 0)
                self.s['N'][fi(ks)] += d
        self.s['r'][0] += self.drf(ds);

    def st_arr(self):
        rs = []
        for k in 'NSEIRD':
            s = self.s[k]
            if isinstance(s, list):
                rs.extend(s)
                rs.append(sum(s))
            else:
                rs.append(s)
        rs.append(self.s['r'][0])
        return rs

    def k_arr(self):
        rs = []
        for k in 'NSEIRD':
            s = self.s[k]
            if isinstance(s, list):
                rs.append(k + 'f')
                rs.append(k + 'i')
            rs.append(k)
        rs.append('rf')
        return rs

def make_format(current, other):
    # current and other are axes
    def format_coord(x, y):
        # x, y are data coordinates
        # convert to display coords
        display_coord = current.transData.transform((x,y))
        inv = other.transData.inverted()
        # convert back to data coords with respect to ax
        ax_coord = inv.transform(display_coord)
        coords = [ax_coord, (x, y)]
        return ('Left: {:<40}    Right: {:<}'
                .format(*['({:.3f}, {:.3f})'.format(x, y) for x,y in coords]))
    return format_coord

def calc_model(n, *args, flt = None):
    mdl = c_model(*args)
    xs = np.arange(n)
    yss = [mdl.st_arr()]
    for i in range(n - 1):
        mdl.update()#i == 35)
        yss.append(mdl.st_arr())
    yss = list(zip(*yss))
    ax = plt.subplot()
    ax.cla()
    lgd = []
    for k, ys in zip(mdl.k_arr(), yss):
        if flt and not flt(k):
            continue
        if k == 'rf':
            _l, = ax.plot([0], [0], '.')
            _c = _l.get_color()
            _l.remove()
            ax2 = ax.twinx()
            #ax2.set_zorder(-100)
            ax2.format_coord = make_format(ax2, ax)
            ax2.cla()
            ax2.plot(xs, ys, '-', color = _c)
            ax2.legend('rf', loc='lower right')
            #ax2.set_ylim(ymin=0)
        else:
            ax.plot(xs, ys, '-')
            lgd.append(k)
    ax.legend(lgd, loc='upper right')
    plt.show()
    return mdl

if __name__ == '__main__':

    params_standard = {
        'S2EbIf': 0.3,
        'S2EbEf': 0.2,
        'S2EbIi': 0.03,
        'S2EbEi': 0.02,
        'E2I': 1/15,
        'E2R': 0.07,
        'I2Rf': 0.1,
        'I2Ri': 0.15,
        'I2Df': 0.08,
        'I2Di': 0.025,
    }
    
    init_standard = {
        'N': 1e7,
        'I': 100,
        'E': 0,
        'rf': 2,
        'ri': 0.5,
    }

    params_isolate = {
        **params_standard,
        'Sf2i': 5,
        'Ef2i': 5,
        'If2i': 0.7,
        'Si2f': 1/15,
        'Ei2f': 1/15,
        'Ri2f': 1/15,
    }

    params_recur = {
        **params_isolate,
        'I2Ei': 0.05,
        'R2S': 0.01,
    }

    init_social = {
        **init_standard,
        'rf': 8,
    }

    params_social = lambda p: {
        **p,
        'rcf': 1.5,
    }

    flt = None
    #flt = lambda k: k in ['Ni', 'If', 'Ii', 'D']
    #flt = lambda k: k[-1] == 'i'
    #flt = lambda k: k != 'N' and len(k) < 2 or k == 'Ni'
    flt = lambda k: not k in 'NSR' and len(k) < 2 or k == 'Ni'
    fltr = lambda k: not k in 'NSR' and len(k) < 2 or k == 'Ni' or k == 'rf'
    #flt = lambda k: not k[0] in 'NS'
    #flt = lambda k: k in 'IRD'
    #flt = lambda k: k == 'rf'
    #flt = lambda k: k == 'Ii'
    #flt = lambda k: k == 'Ri'
    #flt = lambda k: k[0] in 'S'
    fltn = lambda *lb: lambda k: k in lb

    tm = 150
    mdl = calc_model(tm, params_standard, init_standard, flt=fltn('E', 'I', 'D'))
    mdl = calc_model(tm, params_isolate, init_standard, flt=fltn('Ni', 'E', 'I', 'D'))
    #mdl = calc_model(tm, params_recur, init_standard, flt=fltn('Ni', 'E', 'I', 'D'))

    tm = 350
    #mdl = calc_model(tm, params_standard, init_social, flt=fltn('Ni', 'E', 'I', 'D'))
    
    #mdl = calc_model(tm, params_isolate, init_social, flt=fltn('Ni', 'E', 'I', 'D'))
    mdl = calc_model(tm, params_social(params_isolate), init_social, flt=fltn('Ni', 'E', 'I', 'D', 'rf'))
    #mdl = calc_model(tm, params_social(params_isolate), init_social, flt=fltn('E', 'I', 'Ii'))
    #mdl = calc_model(tm, params_social(params_isolate), init_social, flt=fltn('Ii', 'rf'))

    #mdl = calc_model(tm, params_recur, init_social, flt=fltn('Ni', 'E', 'I', 'D'))
    mdl = calc_model(tm, params_social(params_recur), init_social, flt=fltn('Ni', 'E', 'I', 'D', 'rf'))