Mock up wave equation example

Light strip outputs aren't hooked up yet
This commit is contained in:
Aaron Fenyes 2019-09-28 20:56:25 +02:00
parent e517eb1880
commit 36077fd872
1 changed files with 65 additions and 11 deletions

View File

@ -1,4 +1,4 @@
from math import pi, cos, sin
from math import pi, cos, sin, exp
import pygame, colorsys
@ -11,7 +11,7 @@ vertices = [(viewsize*(1 + shrink*x)/2, viewsize*(1 + shrink*y)/2) for x, y in
[(0, 0)] + [(cos(-n*pi/5), sin(-n*pi/5)) for n in range(10)] + [(0, 0)]
]
# edges as vertex pairs
# edges as vertex pairs. FireStar will rest on vertices 6 (A), 5 (B), 7 (C)
edges = [
(0, 1), (0, 3), (0, 5), (0, 7), (0, 9),
(2, 8), (2, 6), (6, 10), (10, 4), (4, 8),
@ -41,36 +41,90 @@ cyc_edge_adj = [
(14, 11), (11, 23), (23, 20), (20, 17), (17, 14)
]
# === interaction parameters ===
frame_rate = 24
inductance = 0.001
drain = False
drain_rate = 0.1
# === phase space ===
litness = 30*[0]
charge = 12*[0]
current = 30*[0]
# === main loop ===
if __name__ == '__main__':
# set up display and clock
pygame.display.init()
screen = pygame.display.set_mode((viewsize, viewsize))
pygame.init()
screen = pygame.display.set_mode((viewsize, viewsize + 100))
pygame.display.set_caption('Preview')
font = pygame.font.Font(None, 25)
clock = pygame.time.Clock()
# set up background
background = pygame.Surface(screen.get_size()).convert()
background.fill((32, 32, 32))
to_light = 0
while True:
# handle events
for event in pygame.event.get():
if event.type == pygame.QUIT:
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_j:
charge[0] += 10
charge[11] -= 10
elif event.key == pygame.K_k:
charge[0] += 100
charge[11] -= 100
elif event.key == pygame.K_SPACE:
drain = True
pygame.draw.aaline(screen, (255, 255, 255), (0, 0), (viewsize, viewsize))
elif event.type == pygame.KEYUP:
if event.key == pygame.K_SPACE:
drain = False
elif event.type == pygame.QUIT:
quit()
litness[to_light] = 1.
# print data
energy = 0.5*(sum(q*q for q in charge) + inductance*sum(i*i for i in current))
screen.blit(background, (0, 0))
text = font.render('energy', True, (255, 255, 255))
screen.blit(text, (20, viewsize))
text = font.render(str(energy), True, (255, 255, 255))
screen.blit(text, (120, viewsize))
text = font.render('charge[0]', True, (255, 255, 255))
screen.blit(text, (20, 30 + viewsize))
text = font.render(str(charge[0]), True, (255, 255, 255))
screen.blit(text, (120, 30 + viewsize))
# show state
for e in range(30):
litness = 1 - exp(-2*current[e]*current[e])
color = colorsys.hsv_to_rgb(0.45 + litness*0.25, 1 - 0.5*litness*litness, litness)
##color = colorsys.hsv_to_rgb(litness*0.167, 1 - 0.5*litness*litness, litness)
pygame.draw.aaline(
screen,
tuple(255*c for c in colorsys.hsv_to_rgb(e/30., litness[e], 0.2 + 0.8*litness[e])),
tuple(255*c for c in color),
vertices[edges[e][0]],
vertices[edges[e][1]],
)
litness[e] *= 0.9
# evolve state. use verlet integration, first updating the currents and then
# using the new currents to update the charges.
for e in range(30):
u, v = edges[e]
##current[e] *= 1 - resistance
current[e] += inductance*(charge[u] - charge[v])
for e in range(30):
u, v = edges[e]
charge[u] -= current[e]
charge[v] += current[e]
if drain:
for v in range(12):
charge[v] *= 1 - drain_rate
# step
to_light = (to_light + 1) % 30
pygame.display.flip()
clock.tick(24)
clock.tick(frame_rate)