python-arcade-pymunk实现数据模拟渲染

from arcade import Window, close_window, run, start_render, set_background_color, color
from arcade import SpriteList, SpriteCircle, draw_text, draw_lines
from pymunk import Body, Space, Segment, Circle, Arbiter
from random import uniform, randint, choice
import matplotlib.pyplot as plt
import timeit
import time

WIDTH = 1000
HEIGHT = 700
TITLE = "数据模拟渲染"
FPS = 60
RADIUS = 9
RECOVERY_TIME = 100

INIT_XY = [[20, 180, 439, 700], [215, 385, 439, 700], [415, 585, 439, 700],
           [615, 785, 439, 700], [815, 995, 439, 700], [20, 180, 10, 250],
           [215, 385, 10, 250], [415, 585, 10, 250], [615, 785, 10, 250],
           [815, 995, 10, 250]]

wall_XY = [[(200, 450), (200, 700)], [(200, 0), (200, 250)],
           [(400, 450), (400, 700)], [(400, 0), (400, 250)],
           [(600, 450), (600, 700)], [(600, 0), (600, 250)],
           [(800, 450), (800, 700)], [(800, 0), (800, 250)],
           [(0, 456), (100, 456)], [(205, 456), (305, 456)],
           [(405, 456), (505, 456)], [(605, 456), (705, 456)],
           [(805, 456), (905, 456)], [(105, 245), (195, 245)],
           [(305, 245), (405, 245)], [(505, 245), (605, 245)],
           [(705, 245), (805, 245)], [(905, 245), (1005, 245)]]

exterior_Wall_XY = [[(0, 0), (0, HEIGHT)], [(0, 0), (WIDTH, 0)],
                    [(WIDTH, 0), (WIDTH, HEIGHT)],
                    [(0, HEIGHT), (WIDTH, HEIGHT)]]

door_XY = [[(100, 456), (195, 456)], [(305, 456), (395, 456)],
           [(505, 456), (595, 456)], [(705, 456), (795, 456)],
           [(905, 456), (999, 456)], [(0, 250), (105, 247)],
           [(205, 247), (305, 247)], [(405, 247), (505, 247)],
           [(605, 247), (705, 247)], [(805, 247), (905, 247)]]

COLOR = [(35, 76, 69), (16, 150, 216), (125, 110, 13), (97, 193, 209),
         (208, 208, 244), (242, 134, 152), (244, 180, 181), (65, 234, 188),
         (0, 58, 86), (81, 117, 76)]

space = Space()
infected_count = []


class MySegment(Segment):
    def __init__(self, body, xy, index=99999999):
        super().__init__(body, xy[0], xy[1], 5)
        self.type = 1
        self.class_room = index


class ExteriorWall:
    def __init__(self):
        for exterior_wall_xy in exterior_Wall_XY:
            self.body = Body(body_type=Body.STATIC)
            self.shape = MySegment(self.body, exterior_wall_xy)
            self.shape.elasticity = 1
            space.add(self.shape)


class Wall:
    def __init__(self):
        self.wall_xy = []
        for wall_xy in wall_XY:
            self.body = Body(body_type=Body.STATIC)
            self.shape = MySegment(self.body, wall_xy)
            self.shape.elasticity = 1
            space.add(self.shape)
            self.wall_xy.append(wall_xy[0])
            self.wall_xy.append(wall_xy[1])

    def draw(self):
        draw_lines(self.wall_xy, color.WHITE, 10)


class Door:
    def __init__(self):
        for index, door_xy in enumerate(door_XY):
            self.body = Body(body_type=Body.STATIC)
            self.shape = MySegment(self.body, door_xy, index + 1)
            self.shape.elasticity = 1
            space.add(self.shape)


class MyCircle(Circle):
    def __init__(self, body, index):
        super().__init__(body, RADIUS)
        self.type = 2
        self.class_room = index


class Student(SpriteCircle):
    def __init__(self, index, circle_color, xy):
        super().__init__(RADIUS, circle_color)
        self.center_x = uniform(xy[0], xy[1])
        self.center_y = uniform(xy[2], xy[3])
        self.body = Body()
        self.body.position = self.center_x, self.center_y
        self.body.velocity = uniform(-300, 300), uniform(-300, 300)
        self.shape = MyCircle(self.body, index)
        self.shape.density = 1
        self.shape.elasticity = 0.99999
        space.add(self.body, self.shape)

        self.infected_time = 0
        self.infected = False
        self.recovered = False

    def update(self):
        if self.infected:
            self.color = color.RED
        elif self.recovered:
            self.color = color.BLUE
        else:
            self.color = color.WHITE
        self.center_x, self.center_y = self.body.position
        self.pass_time()

    def pass_time(self):
        if self.infected:
            self.infected_time += 1
        if self.infected_time >= RECOVERY_TIME:
            self.infected = False
            self.recovered = True
            self.shape.collision_type = 99999999999999

    def infect(self, space=0, arbiter=0, data=0):
        self.infected = True
        self.shape.collision_type = 1
        self.body.velocity = 0, 0
        self.shape.density = 10


class Main(Window):
    def __init__(self, WIDTH, HEIGHT, TITLE):
        super().__init__(WIDTH, HEIGHT, TITLE)
        self.processing_time = 0
        self.draw_time = 0
        self.frame_count = 0
        self.fps_start_timer = None
        self.fps = None

        self.wall = Wall()
        self.exterior_wall = ExteriorWall()
        self.door = Door()
        self.student_sprite_list = SpriteList()
        self.setup()
        for index, student in enumerate(self.student_sprite_list):
            student.shape.collision_type = index + 2
            handler = space.add_collision_handler(index + 2, 1)
            handler.separate = student.infect

        choice(self.student_sprite_list).infect()
        handler = space.add_default_collision_handler()
        handler.begin = self.call_begin

    def call_begin(self, arbiter: Arbiter, space, data):
        if arbiter:
            shapes = arbiter.shapes
            if len(shapes) == 2:
                if shapes[0].class_room == shapes[
                        1].class_room and shapes[0].type != shapes[1].type:
                    return False
        return True

    def setup(self):
        for index in range(10):
            students = [
                Student(index + 1, color.WHITE, INIT_XY[index])
                for _ in range(randint(50, 81))
            ]
            self.student_sprite_list.extend(students)

    def on_show(self):
        set_background_color((50, 50, 50))

    def on_update(self, delta_time):
        start_time = timeit.default_timer()
        space.step(1.0 / FPS)
        self.student_sprite_list.update()
        self.processing_time = timeit.default_timer() - start_time

    def on_draw(self):
        draw_start_time = timeit.default_timer()
        if self.frame_count % 60 == 0:
            if self.fps_start_timer is not None:
                total_time = timeit.default_timer() - self.fps_start_timer
                self.fps = 60 / total_time
            self.fps_start_timer = timeit.default_timer()
        self.frame_count += 1
        start_render()
        set_background_color((50, 50, 50))
        self.wall.draw()
        self.student_sprite_list.draw()
        infected_count_thisframe = 0
        for student in self.student_sprite_list:
            student.pass_time()
            if student.infected:
                infected_count_thisframe += 1
        if len(infected_count) <= 1000:
            infected_count.append(infected_count_thisframe)

        if self.fps is not None:
            draw_text(f'FPS: {self.fps:.0f}', 1, HEIGHT - 20, color.WHITE, 16)
        else:
            draw_text('FPS: 0.0', 1, HEIGHT - 20, color.WHITE, 16)
        draw_text(f'Drawing time: {self.draw_time:.3f}', 1, HEIGHT - 50,
                  color.WHITE, 16)
        dt = time.strftime("%Y-%m-%d %H:%M:%S",
                           time.localtime(int(time.time())))
        draw_text(dt, 1, HEIGHT - 80, color.WHITE, 16)
        self.draw_time = timeit.default_timer() - draw_start_time


def main():
    Main(WIDTH, HEIGHT, TITLE)
    run()
    close_window()
    plt.plot(range(0, len(infected_count), 1), infected_count)
    plt.show()


if __name__ == "__main__":
    main()