Untitled

mail@pastecode.io avatar
unknown
plain_text
a year ago
26 kB
2
Indexable
Never
import pylab
import pygame
import random
import numpy as np
import pandas as pd
from copy import deepcopy
from numba import jit
import matplotlib as plt
plt.use("Agg")
import pygame
from pygame.locals import *
plt.rcParams.update({'figure.max_open_warning': 0})
import matplotlib.backends.backend_agg as agg

#%%

epoch = 366 # период который рассматриваем 366
limit_time_illness = 21 # T1 время болезни
limit_time_vaccine = 90 # время действия иммунитета от вакцины
limit_time_immun_con = 90 # T2 время действия иммунитета после болезни
probability_immun_abs = 0 # p1
probability_infect = 0.15 # p2
probability_lethal = 0.01 # p3
probability_virus_carrier = 0.1 # p4 вероятность появления больного
time_found_illness = 7

#%%
matrix_offsets = np.array(
    [[1,1,1],
     [1,1,1],
     [1,1,1]])
dy = matrix_offsets.shape[0] // 2
dx = matrix_offsets.shape[1] // 2
N = matrix_offsets.shape[0]
dict_of_colours = {0: "black",
                   1: "darkolivegreen3",
                   2: "lightcoral",
                   3: "lightskyblue",
                   4: "khaki1"}

#%%
dx, dy

#%%
data = {"вероятность заражения при контакте": [],
        "вероятность смерти во время болезни": [],
        "процент людей с иммунитетом": [],
        "число локальных контактов": [],
        "население": [],
        "число живых": [],
        "процент мёртвых": []}
df = pd.DataFrame(data=data)
df

#%%
info = pd.DataFrame(data=data)
info

#%%
def gameOver(surface):
    surface.fill(pygame.Color('black'))
    font = pygame.font.SysFont("comicsansms", 32)
    red = (255, 0, 0)
    green = (0, 255, 0)
    words = font.render("Вирус захватил мир!", 1, red, green)

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()

        surface.blit(words, (800,500))
        pygame.display.update()

#%%
def generate_left_border(states):
    rand_vector = np.random.sample(states.shape[2])
    states[1][0] = ((states[1][0] + rand_vector) <=
                    probability_virus_carrier) + 0

#%%
def generate_imun(states, add_abs_imun=0,
                   percent_vaccinated_people=0):
    rand_matrix = np.random.sample((states.shape[1]-1, states.shape[2]))
    states[2][1:states.shape[1]] = ((states[2][1:states.shape[1]] + rand_matrix) <=
                                    probability_immun_abs + add_abs_imun) + 0

    rand_matrix = np.random.sample((states.shape[1]-1, states.shape[2]))
    states[3][1:states.shape[1]] = ((states[3][1:states.shape[1]] + rand_matrix) <=
                                    percent_vaccinated_people) + 0

#%%
# TODO:
def create_info(number, states, add_abs_imun, local_limit):
    new_row = [probability_infect, probability_lethal, add_abs_imun, local_limit, states.shape[2] * states.shape[1],
               (states.shape[2] * states.shape[1] - np.sum(states[0])) / states.shape[2] / states.shape[1] * 100,
               np.sum(states[0]) / states.shape[2] / states.shape[1] * 100]
    df.loc[len(df)] = new_row
    return True


# %%
def visualisation(states, surface, WIDTH, HEIGHT, TILE):
    surface.fill(pygame.Color('white'))
    global green, black, red, yellow, blue

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()

    [pygame.draw.line(surface, pygame.Color('darkslategray'), (x, 0), (x, HEIGHT)) for x in range(0, WIDTH, TILE)]
    [pygame.draw.line(surface, pygame.Color('darkslategray'), (0, y), (WIDTH, y)) for y in range(0, HEIGHT, TILE)]

    green, black, red, yellow, blue = 0, 0, 0, 0, 0

    for x in range(0, states.shape[2]):
        for y in range(0, states.shape[1]):
            # жив - зелёный
            colour = dict_of_colours[1]
            # мёртв - чёрный
            if states[0][y][x] == 1:
                colour = dict_of_colours[0]
                black += 1
            # болен - красный
            elif states[1][y][x] == 1:
                colour = dict_of_colours[2]
                red += 1
            # приобретённый имун - жёлтый
            elif states[4][y][x] or states[3][y][x]:
                colour = dict_of_colours[4]
                yellow += 1
            # абс имун - синий
            elif states[2][y][x]:
                colour = dict_of_colours[3]
                blue += 1

            green = 2500 - yellow - red - blue - black

            pygame.draw.rect(surface, pygame.Color(colour), (x * TILE + 2, y * TILE + 2, TILE - 2, TILE - 2))

    # green = (WIDTH*HEIGHT)//10//TILE-yellow-red-blue-black
        return black, red, yellow, blue, green
        pygame.display.update() @jit(nopython=True)


def update_cell(states, times, local_limit=1):
    # если мёртв и абс имун - скип
    # если болен
    # return 0
    H = states.shape[1]
    W = states.shape[2]

    times[0] = times[0] + (times[0] != 0)

    # если умер
    rand_matrix = np.random.sample((H, W))
    died_people = (times[0] > 0) * (times[0] <= limit_time_illness) * (rand_matrix <= probability_lethal) * states[1]
    states[0][1:H] += died_people[1:H]

    # если переболел
    recovered_people = (times[0] > limit_time_illness) + 0
    states[4] += recovered_people
    times[2] += recovered_people
    times[0] = times[0] * (died_people == 0) * (recovered_people == 0)
    states[1] = states[1] * (died_people == 0) * (recovered_people == 0)

    # если переболел
    times[2] = times[2] + (times[2] != 0)
    with_imun = (times[2] <= limit_time_immun_con) * (times[2] > 0)
    states[4] = states[4] * with_imun
    times[2] = times[2] * with_imun

    # если вакцинирован
    times[1] = times[1] + (times[1] != 0)
    with_vaccine = (times[1] <= limit_time_vaccine)
    states[3] = states[3] * with_vaccine
    times[1] = times[1] * with_vaccine

    # если здоров
    infected_people = np.zeros((5, H + 2 * dy, W + 2 * dx))
    infected_people[:, dy:H + dy, dx:W + dx] = states

    width_time_0 = np.zeros((H + 2 * dy, W + 2 * dx))
    width_time_0[dy:H + dy, dx:W + dx] = times[0]
    number_infected_people = np.sum(states[1]) / (W * H - H) * local_limit * probability_infect
    for x in range(dx, W + dx):
        for y in range(dy + 1, H + dy):
            random_matrix = (np.random.sample((N, N)) <= (probability_infect + number_infected_people)) * matrix_offsets
            random_matrix = random_matrix * (width_time_0[y - dy:y + dy + 1, x-dx:x + dx + 1] <= time_found_illness)
            random_matrix = random_matrix * np.min((states[1:5, y - dy, x - dx] == 0))
            random_matrix = random_matrix * (infected_people[0][y-dy:y + dy + 1, x - dx:x + dx + 1] == 0)
            random_matrix = random_matrix * infected_people[1][y-dy:y + dy + 1, x - dx:x + dx + 1]
            random_matrix = random_matrix * (infected_people[3][y-dy:y + dy + 1, x - dx:x + dx + 1] == 0)
            random_matrix = random_matrix * (infected_people[2][y-dy:y + dy + 1, x - dx:x + dx + 1] == 0)
            random_matrix = random_matrix * (infected_people[4][y-dy:y + dy + 1, x - dx:x + dx + 1] == 0)

            states[1][y - dy][x - dx] += np.max(random_matrix)
            times[0][y - dy][x - dx] += np.max(random_matrix)







import matplotlib.pyplot as plt
def gist(red, black, yellow, blue):
    index = [0, 1, 2, 3, 4]
    values = [red, black, yellow, blue]

    # plt.ylim([0, 1500])
    # plt.xlim([0, 4])

    plt.bar(index[0], values[0], color='r')
    plt.bar(index[1], values[1], color='k')
    plt.bar(index[2], values[2], color='y')
    plt.bar(index[3], values[3], color='b')

    return plt.savefig('foo.png')


def aviabillity_of_leak(add_abs_imun=0.1, number=1, local_limit=0,
                        percent_vaccinated_people=0.01, WIDTH=2000, HEIGHT=2000,
                        show=False, info=False):
    TILE = 40
    if show:
        WIDTH = 500
        HEIGHT = 500
        TILE = 10

    pygame.init()
    RES = WIDTH, HEIGHT
    surface = pygame.display.set_mode((1000, 500))
    surf_left = pygame.Surface((200, 150))
    surf_left.fill((0, 0, 0))
    surface.blit(surf_left, (50, 25))
    clock = pygame.time.Clock()

    W, H = WIDTH // TILE, HEIGHT // TILE
    FPS = 100  # скорость протекания игры
    n = 5  # количество состояний
    # 0 - жив/ мёртв
    # 1 - болен
    # 2 - абсолютный иммунитет
    # 3 - вакцинирован
    # 4 - переболел
    states = np.zeros((n, H, W))

    # 0 - время болезни
    # 1 - время вакцины
    # 2 - время после болезни
    times = np.zeros((3, H, W))

    generate_imun(states, add_abs_imun, percent_vaccinated_people)

    iter_count = 1
    while iter_count <= 365:
        generate_left_border(states)

        if show:
            visualisation(states, surface, WIDTH, HEIGHT, TILE)

        green = 2500 - yellow - red - blue - black
        f1 = pygame.font.SysFont('arial', 24)
        text0 = f1.render('День № ' + str(iter_count), True, (128, 128, 128))
        surface.blit(text0, (510, 20))
        text1 = f1.render('Заболело ' + str(red), True, (255, 0, 0))
        surface.blit(text1, (510, 50))
        text2 = f1.render('Умерло ' + str(black), True, (0, 0, 0))
        surface.blit(text2, (510, 80))
        text3 = f1.render('Обладают временным иммунитетом ' + str(yellow),
                          True, (255, 255, 0))
        surface.blit(text3, (510, 110))
        text4 = f1.render('Обладают абсолютным иммунитетом ' + str(blue),
                          True, (0, 0, 255))
        surface.blit(text4, (510, 140))
        text5 = f1.render('Здоровы ' + str(green), True, (0, 180, 0))
        surface.blit(text5, (510, 170))

        gist(red, black, yellow, blue)
        png = pygame.image.load('foo.png')
        surface.blit(png, ((510, 230)))
        pygame.display.flip()
        clock.tick(FPS)

        update_cell(states, times, local_limit)

        if np.max(states[1][H - 1]) == 1:
            if info:
                create_info(number, states, add_abs_imun, 1)
            return False

        iter_count += 1

    if info:
        create_info(number, states, add_abs_imun, local_limit)

    return True


def calculate_critical_value(number=0, local_limit=1, eps=0.00001, info=True):
    start, current, end = 0, 0.5, 1
    i = 0
    while abs(start - end) >= eps:
        i += 1

        if not (aviabillity_of_leak(number=number, local_limit=local_limit,
                                    add_abs_imun=current, info=info)):
            current, start = (current + end) / 2, current
        else:
            current, end = (start + current) / 2, current

    return current


def calculate_average_critical_value(local_limit, repeats=1, eps=0.001):
    critical_value = sum([calculate_critical_value(number=i, local_limit=local_limit)
                          for i in range(repeats)]) / repeats
    return critical_value


# %%
def aviabillity_of_leak_repeats(add_abs_imun=0.05, local_limit=1, repeat=10):
    data = {"вероятность заражения при контакте": [],
            "вероятность смерти во время болезни": [],
            "процент людей с иммунитетои": [],
            "число локальных контактов": [],
            "население": [], "число живых": [],
            "процент мёртвых": []}
    global df
    df = pd.DataFrame(data=data)

    [aviabillity_of_leak(add_abs_imun=add_abs_imun, local_limit=local_limit,
                         info=True) for _ in range(repeat)]
    info.loc[len(info)] = (df.mean()).tolist()


aviabillity_of_leak(show=True, info=True)

import matplotlib.pyplot as plt


def gist(red, black, yellow, blue):
    index = [0, 1, 2, 3]
    values = [red, black, yellow, blue]
    # plt.ylim([0, 1500])
    # plt.xlim([0, 4])
    plt.bar(index[0], values[0], color='r')
    plt.bar(index[1], values[1], color='k')
    plt.bar(index[2], values[2], color='y')
    plt.bar(index[3], values[3], color='b')
    return plt.savefig('foo.png')



def aviabillity_of_leak(add_abs_imun=0.1, number=1, local_contact=0.5, #я туууууут нахуй измени меня на 0.5
                        percent_vaccinated_people=0.01, WIDTH=2000, HEIGHT=2000,
                        show=False, info=False):
    TILE = 40 # было 40
    if show:
        WIDTH = 500
        HEIGHT = 500 #### тут было по 500
        TILE = 10 # было 10

    pygame.init()
    RES = WIDTH, HEIGHT
    surface = pygame.display.set_mode((1000, 500))
    surf_left = pygame.Surface((200, 150))
    surf_left.fill((0, 0, 0))
    surface.blit(surf_left, (50, 25))
    clock = pygame.time.Clock()

    W, H = WIDTH // TILE, HEIGHT // TILE
    FPS = 100  # скорость протекания игры
    n = 5  # количество состояний
    # 0 - жив/ мёртв
    # 1 - болен
    # 2 - абсолютный иммунитет
    # 3 - вакцинирован
    # 4 - переболел
    states = np.zeros((n, H, W))

    # 0 - время болезни
    # 1 - время вакцины
    # 2 - время после болезни
    times = np.zeros((3, H, W))

    generate_imun(states, add_abs_imun, percent_vaccinated_people)

    iter_count = 1
    while iter_count <= 365:
        generate_left_border(states)

        if show:
            visualisation(states, surface, WIDTH, HEIGHT, TILE)

        f1 = pygame.font.SysFont('arial', 20)
        text0 = f1.render('День № ' + str(iter_count), True, (128, 128, 128))
        surface.blit(text0, (510, 20))
        text1 = f1.render('Заболело ' + str(red), True, (255, 0, 0))
        surface.blit(text1, (510, 40))
        text2 = f1.render('Умерло ' + str(black), True, (0, 0, 0))
        surface.blit(text2, (510, 60))
        text3 = f1.render('Обладают временным иммунитетом ' + str(yellow),
                          True, (255, 255, 0))
        surface.blit(text3, (510, 80))
        text4 = f1.render('Обладают абсолютным иммунитетом ' + str(blue),
                          True, (0, 0, 255))
        surface.blit(text4, (510, 100))
        text5 = f1.render('Здоровы ' + str(green), True, (0, 180, 0))
        surface.blit(text5, (510, 120))

        gist(red, black, yellow, blue)
        png = pygame.image.load('foo.png')
        surface.blit(png, ((510, 180)))
        pygame.display.flip()
        clock.tick(FPS)

        update_cell(states, times, local_contact)

        if np.max(states[1][H - 1]) == 1:
            if info:
                create_info(number, states, add_abs_imun, 0.5)
            return False

        iter_count += 1

    if info:
        create_info(number, states, add_abs_imun, local_contact)

    return True

# %%
def aviabillity_of_leak(add_abs_imun=0.1, number=0, local_contact=1,
                        percent_vaccinated_people=0.1, WIDTH=2000, HEIGHT=2000,
                        show=False, info=False):
    TILE = 40
    if show:
        WIDTH = 500
        HEIGHT = 500
        TILE = 10

    pygame.init()
    RES = WIDTH, HEIGHT
    surface = pygame.display.set_mode((1000, 500))
    surf_left = pygame.Surface((200, 150))
    surf_left.fill((0, 0, 0))
    surface.blit(surf_left, (50, 25))
    clock = pygame.time.Clock()

    W, H = WIDTH // TILE, HEIGHT // TILE
    FPS = 100  # скорость протекания игры
    n = 5  # количество состояний
    # 0 - жив/ мёртв
    # 1 - болен
    # 2 - абсолютный иммунитет
    # 3 - вакцинирован
    # 4 - переболел
    states = np.zeros((n, H, W))

    # 0 - время болезни
    # 1 - время вакцины
    # 2 - время после болезни
    times = np.zeros((3, H, W))

    generate_imun(states, add_abs_imun, percent_vaccinated_people)

    iter_count = 1
    while iter_count <= 100:
        generate_left_border(states)

        if show:
            visualisation(states, surface, WIDTH, HEIGHT, TILE)

        size = WIDTH * HEIGHT // 100

        f1 = pygame.font.SysFont('arial', 20)
        text0 = f1.render('День № ' + str(iter_count), True, (128, 128, 128))
        surface.blit(text0, (510, 20))
        text6 = f1.render('Население ' + str(size), True, (128, 128, 128))
        surface.blit(text6, (510, 40))
        text1 = f1.render('Заболело ' + str(red), True, (255, 0, 0))
        surface.blit(text1, (510, 60))
        text2 = f1.render('Умерло ' + str(black), True, (0, 0, 0))
        surface.blit(text2, (510, 80))
        text3 = f1.render('Обладают временным иммунитетом ' + str(yellow),
                          True, (255, 255, 0))
        surface.blit(text3, (510, 100))
        text4 = f1.render('Обладают абсолютным иммунитетом ' + str(blue),
                          True, (0, 0, 255))
        surface.blit(text4, (510, 120))
        text5 = f1.render('Здоровы ' + str(green), True, (0, 180, 0))
        surface.blit(text5, (510, 140))

        gist(red, black, yellow, blue)
        png = pygame.image.load('foo.png')
        surface.blit(png, ((510, 180)))
        pygame.display.flip()
        clock.tick(FPS)

        update_cell(states, times, local_contact)

        iter_count += 1

        if info:
            create_info(number, states, add_abs_imun, local_contact)

    return True

# %%
    aviabillity_of_leak(show=True, info=True)
    вакцина

    def aviabillity_of_leak(add_abs_imun=0.1, number=0,
                            local_contact=1,
                            percent_vaccinated_people=0.5,
                            WIDTH=2000, HEIGHT=2000, show=False,
                            info=False):
        TILE = 40
        if show:
            WIDTH = 500
        HEIGHT = 500
        TILE = 10
        pygame.init()
        RES = WIDTH, HEIGHT
        surface = pygame.display.set_mode((1000, 500))
        surf_left = pygame.Surface((200, 150))
        surf_left.fill((0, 0, 0))
        surface.blit(surf_left, (50, 25))
        clock = pygame.time.Clock()

        W, H = WIDTH // TILE, HEIGHT // TILE
        FPS = 100  # скорость протекания игры
        n = 5  # количество состояний
        # 0 - жив/ мёртв
        # 1 - болен
        # 2 - абсолютный иммунитет
        # 3 - вакцинирован
        # 4 - переболел
        states = np.zeros((n, H, W))
        # 0 - время болезни
        # 1 - время вакцины
        # 2 - время после болезни
        times = np.zeros((3, H, W))
        generate_imun(states, add_abs_imun,
                      percent_vaccinated_people)
        iter_count = 1
        while iter_count <= 100:
            generate_left_border(states)
            if show:
                visualisation(states, surface, WIDTH, HEIGHT, TILE)

            size = WIDTH * HEIGHT // 100

            f1 = pygame.font.SysFont('arial', 20)
            text0 = f1.render('День № ' + str(iter_count), True, (128, 128, 128))
            surface.blit(text0, (510, 20))
            text6 = f1.render('Население ' + str(size), True, (128, 128, 128))
            surface.blit(text6, (510, 40))
            text1 = f1.render('Заболело ' + str(red), True, (255, 0, 0))
            surface.blit(text1, (510, 60))
            text2 = f1.render('Умерло ' + str(black), True, (0, 0, 0))
            surface.blit(text2, (510, 80))
            text3 = f1.render('Обладают временным иммунитетом ' + str(yellow),
                              True, (255, 255, 0))
            surface.blit(text3, (510, 100))
            text4 = f1.render('Обладают абсолютным иммунитетом ' + str(blue),
                              True, (0, 0, 255))
            surface.blit(text4, (510, 120))
            text5 = f1.render('Здоровы ' + str(green), True, (0, 180, 0))
            surface.blit(text5, (510, 140))
            text7 = f1.render('Процент вакцинированных ' + str(percent_vaccinated_people * 100) + '%',
                              True, (128, 128, 128))
            surface.blit(text7, (510, 160))

            gist(red, black, yellow, blue)
            png = pygame.image.load('foo.png')
            surface.blit(png, ((510, 180)))
            pygame.display.flip()
            clock.tick(FPS)

            update_cell(states, times, local_contact)

            iter_count += 1

            if info:
                create_info(number, states, add_abs_imun, local_contact)

        return True

        # %%
        aviabillity_of_leak(show=True, info=True)
        локальные
        контакты

        def aviabillity_of_leak(add_abs_imun=0.1, number=0,
                                local_contact=0.8,
                                percent_vaccinated_people=0.45,
                                WIDTH=2000, HEIGHT=2000, show=False,
                                info=False):
            TILE = 40
            if show:
                WIDTH = 500
                HEIGHT = 500
                TILE = 10
                pygame.init()
                RES = WIDTH, HEIGHT
                surface = pygame.display.set_mode((1000, 500))
                surf_left = pygame.Surface((200, 150))
                surf_left.fill((0, 0, 0))
                surface.blit(surf_left, (50, 25))
                clock = pygame.time.Clock()
            W, H = WIDTH // TILE, HEIGHT // TILE
            FPS = 100  # скорость протекания игры
            n = 5  # количество состояний
            # 0 - жив/ мёртв
            # 1 - болен
            # 2 - абсолютный иммунитет
            # 3 - вакцинирован
            # 4 - переболел
            states = np.zeros((n, H, W))
            # 0 - время болезни
            # 1 - время вакцины
            # 2 - время после болезни
            times = np.zeros((3, H, W))
            generate_imun(states, add_abs_imun,
                          percent_vaccinated_people)
            iter_count = 1
            while iter_count <= 100:
                generate_left_border(states)
                if show:
                    visualisation(states, surface, WIDTH, HEIGHT, TILE)

                size = WIDTH * HEIGHT // 100

                f1 = pygame.font.SysFont('arial', 20)
                text0 = f1.render('День № ' + str(iter_count), True, (128, 128, 128))
                surface.blit(text0, (510, 20))
                text6 = f1.render('Население ' + str(size), True, (128, 128, 128))
                surface.blit(text6, (510, 40))
                text1 = f1.render('Заболело ' + str(red), True, (255, 0, 0))
                surface.blit(text1, (510, 60))
                text2 = f1.render('Умерло ' + str(black), True, (0, 0, 0))
                surface.blit(text2, (510, 80))
                text3 = f1.render('Обладают временным иммунитетом ' + str(yellow),
                                  True, (255, 255, 0))
                surface.blit(text3, (510, 100))
                text4 = f1.render('Обладают абсолютным иммунитетом ' + str(blue),
                                  True, (0, 0, 255))
                surface.blit(text4, (510, 120))
                text5 = f1.render('Здоровы ' + str(green), True, (0, 180, 0))
                surface.blit(text5, (510, 140))
                text7 = f1.render('Процент ограничений ' + str(local_contact * 100) + '%',
                                  True, (128, 128, 128))
                surface.blit(text7, (510, 160))

                gist(red, black, yellow, blue)
                png = pygame.image.load('foo.png')
                surface.blit(png, ((510, 180)))
                pygame.display.flip()
                clock.tick(FPS)

                update_cell(states, times, local_contact)

                iter_count += 1

                if info:
                    create_info(number, states, add_abs_imun, local_contact)

            return True

            # %%
            aviabillity_of_leak(show=True, info=True)