molecular-python3

1
from typing import Tuple, Callable, List
2

3
import numpy as np
4
import pygame
5

6

7
class Screen:
8
    draw_circles_vectorized: Callable
9
    draw_lines_vectorized: Callable
10
    _screen: pygame.Surface
11
    _background_color: Tuple[int, int, int]
12
    _offset: List[int]
13
    _scale: float
14

15
    def __init__(self, screen: pygame.Surface, background_color: Tuple[int, int, int] = (0, 0, 0)):
16
        self._screen = screen
17
        self._offset = [0, 0]
18
        self._scale = 1
19
        self._background_color = background_color
20
        self.draw_circles_vectorized = np.vectorize(self.draw_circle)
21
        self.draw_lines_vectorized = np.vectorize(self.draw_line)
22

23
    def move_offset(self, dx: int, dy: int):
24
        self._offset[0] += dx
25
        self._offset[1] += dy
26

27
    def move_scale(self, dv):
28
        self._scale *= dv
29

30
    def clear(self) -> None:
31
        self._screen.fill(self._background_color)
32

33
    def draw_circle(self, center_x: float, center_y: float, radius: int, color_r: int, color_g: int, color_b: int) -> None:
34
        pygame.draw.circle(
35
            self._screen,
36
            (color_r, color_g, color_b),
37
            ((center_x + self._offset[0])*self._scale, (center_y + self._offset[1])*self._scale),
38
            radius*self._scale
39
        )
40

41
    def draw_line(self, start_x: float, start_y: float, end_x: float, end_y: float, color_r: int, color_g: int, color_b: int) -> None:
42
        pygame.draw.line(
43
            self._screen,
44
            (color_r, color_g, color_b),
45
            ((start_x + self._offset[0])*self._scale, (start_y + self._offset[1])*self._scale),
46
            ((end_x + self._offset[0])*self._scale, (end_y + self._offset[1])*self._scale),
47
        )
48

49
    @staticmethod
50
    def update() -> None:
51
        pygame.display.flip()
52