Spaces:
Running
Running
File size: 5,753 Bytes
ba246bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | import pygame
import sys
from pygame import math
import warnings
edge_color_normal = pygame.Color(30, 30, 30) # Dark Grey
vertex_color_normal = pygame.Color(255, 0, 0) # Red
edge_color_select = pygame.Color(0, 255, 0) # Green
vertex_color_select = pygame.Color(128, 128, 128) # Grey
positive_irregularity_color = pygame.Color(255, 0, 255)
negative_irregularity_color = pygame.Color(0, 0, 255)
class Vertex:
def __init__(self, idx, x, y, value=0):
self.idx = idx
self.x = x
self.y = y
self.selected = False
self.color = vertex_color_normal
self.obj = None
self.value = round(value, 0)
def switch_selection(self):
if self.selected:
self.selected = False
self.color = vertex_color_normal
else:
self.selected = True
self.color = vertex_color_select
def collide_point(self, x, y):
return self.obj.collide_point(x, y)
def draw(self, window, font, w_data):
x = w_data.center.x + w_data.stretch * (self.x - w_data.scene_center.x)
y = w_data.center.y - w_data.stretch * (self.y - w_data.scene_center.y)
if self.value != 0 :
if self.value > 0:
self.obj = pygame.draw.circle(window, positive_irregularity_color,
(x, y),
15, 15)
text = font.render(str(self.value), True, (255, 255, 255))
self.blit = window.blit(text, text.get_rect(center=(x, y)))
else:
self.obj = pygame.draw.circle(window, negative_irregularity_color,
(x, y),
15, 15)
text = font.render(str(self.value), True, (255, 255, 255))
self.blit = window.blit(text, text.get_rect(center=(x, y)))
class Edge:
def __init__(self, v1, v2):
self.start = v1
self.end = v2
self.selected = False
self.color = edge_color_normal
self.obj = None
# static method to facilitate unit test
@staticmethod
def is_pt_on_segment(x1, y1, x2, y2, pt, tolerance) -> bool:
nv = pygame.math.Vector2(y1 - y2, x2 - x1)
lp = pygame.math.Vector2(x1, y1)
p = pygame.math.Vector2(pt)
xy = pygame.math.Vector2(x2 - x1, y2 - y1)
# distance from the straight line represented by the edge
distance_ok = abs(nv.normalize().dot(p - lp)) < tolerance
# on the segment ?
segment_ok = (0 <= xy.normalize().dot(p - lp)) and (xy.normalize().dot(p - lp) <= xy.length())
if distance_ok and segment_ok:
return True
return False
def switch_selection(self):
if self.selected:
self.selected = False
self.color = edge_color_normal
else:
self.selected = True
self.color = edge_color_select
def distance_point(self, pt, w_data):
x1 = w_data.center.x + w_data.stretch * (self.start.x - w_data.scene_center.x)
y1 = w_data.center.y - w_data.stretch * (self.start.y - w_data.scene_center.y)
x2 = w_data.center.x + w_data.stretch * (self.end.x - w_data.scene_center.x)
y2 = w_data.center.y - w_data.stretch * (self.end.y - w_data.scene_center.y)
return Edge.is_pt_on_segment(x1, y1, x2, y2, pt, w_data.edge_picking_pixel_tolerance)
def collide_point(self, x, y, w_data):
return self.distance_point(math.Vector2(x, y), w_data)
def draw(self, window, w_data):
self.obj = pygame.draw.line(window, self.color,
[w_data.center.x + w_data.stretch * (self.start.x - w_data.scene_center.x),
w_data.center.y - w_data.stretch * (self.start.y - w_data.scene_center.y)],
[w_data.center.x + w_data.stretch * (self.end.x - w_data.scene_center.x),
w_data.center.y - w_data.stretch * (self.end.y - w_data.scene_center.y)],
w_data.edge_thickness)
class Graph:
def __init__(self, vertices=[], edges=[], scores=[]):
self.clear()
self.update(vertices, edges, scores)
def clear(self):
self.vertices = []
self.edges = []
def update(self, vertices, edges, scores):
for n in vertices:
nodes_scores = scores[0]
n_value = nodes_scores[n[0]]
self.create_vertex(n[0], n[1], n[2], n_value)
for e in edges:
self.create_edge(e[0], e[1])
def create_vertex(self, id: int, x: int, y: int, n_value) -> int:
v = Vertex(id, x, y, n_value)
self.add_vertex(v)
return len(self.vertices) - 1
def create_edge(self, i1: int, i2: int) -> int:
n1, n2 = None, None
for v in self.vertices:
if v.idx == i1:
n1 = v
elif v.idx == i2:
n2 = v
if n1 is None or n2 is None:
warnings.warn("try to create an edge between nodes not found")
self.add_edge(Edge(n1, n2))
return len(self.edges) - 1
def add_vertex(self, n: Vertex) -> None:
self.vertices.append(n)
def add_edge(self, e):
self.edges.append(e)
def bounding_box(self):
x_min = sys.float_info.max
y_min = sys.float_info.max
x_max = sys.float_info.min
y_max = sys.float_info.min
for v in self.vertices:
x_min = min(v.x, x_min)
y_min = min(v.y, y_min)
x_max = max(v.x, x_max)
y_max = max(v.y, y_max)
return x_min, y_min, x_max, y_max
|