pixelserver2/apps/pendlum.py
2024-10-28 18:10:03 +01:00

82 lines
1.9 KiB
Python
Executable File

#!/usr/bin/env python3
import numpy as np
import sys
import os
import time
import colorsys
Nx = int(sys.argv[1])
Ny = int(sys.argv[2])
N = 100
iterations = 2
try:
N = int(sys.argv[3])
except:
pass
def rk4(f, x, dt):
k1 = f(x)
k2 = f(x+k1*dt/2)
k3 = f(x+k2*dt/2)
k4 = f(x+k3*dt)
return x+dt*(k1+2*k2+2*k3+k4)/6.0
dt = 0.08
m = 1
l = 1
g = 1
def sq(x): return x*x
def rhs(x):
phi1 = x[0]
p1 = x[1]
phi2 = x[2]
p2 = x[3]
dphi1 = 6/(m*l*l)*(2*p1-3*np.cos(phi1-phi2)*p2)/(16-9*sq(np.cos(phi1-phi2)))
dphi2 = 6/(m*l*l)*(8*p2-3*np.cos(phi1-phi2)*p1)/(16-9*sq(np.cos(phi1-phi2)))
dp1 = -0.5*m*l*l*(dphi1*dphi2*np.sin(phi1-phi2)+3*g/l*np.sin(phi1))
dp2 = -0.5*m*l*l*(-dphi1*dphi2*np.sin(phi1-phi2)+g/l*np.sin(phi2))
return np.array([dphi1, dp1, dphi2, dp2])
s = np.array([np.pi+0.1, 0, np.pi-0.1, 0])
angles = np.linspace(np.pi, np.pi/3, N)+np.random.uniform(-np.pi/6/N, np.pi/6/N, N)
angles = np.random.uniform(-np.pi/6, np.pi/3, N)+np.pi
states = [np.array([np.pi, 0, a, 0]) for a in angles]
def clamp(x, min_, max_):
return max(min_, min(max_, x))
buffer = bytearray(b"\x00"*Nx*Ny*3)
def setpixel(x, y, r, g, b):
xi = int(Nx*(x+1.2)/2.4)
yi = int(Ny*(y+1.2)/2.4)
if xi < 0 or xi >= Nx or yi < 0 or yi >= Ny:
return
idx = xi+Nx*yi
buffer[3*idx+0] = r
buffer[3*idx+1] = g
buffer[3*idx+2] = b
timestep = 0
timefac = 0.03
while True:
timestep += 1
for s, i in zip(states, range(len(states))):
phi1 = s[0]
phi2 = s[2]
x1 = np.sin(phi1)*l
y1 = np.cos(phi1)*l
x2 = np.sin(phi2)*l+x1
y2 = np.cos(phi2)*l+y1
h = 0.2*i/N+0*timestep*timefac/iterations
r, g, b = colorsys.hsv_to_rgb(h%1, 1, 1)
setpixel(x2, y2, int(r*255), int(g*255), int(b*255))
states[i] = rk4(rhs, s, dt/iterations)
if timestep > 10*iterations and timestep % iterations == 0:
#time.sleep(0.01)
os.write(1, buffer)