Feuerwerk!

generators/fireworks.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+
+import sys
+import argparse
+import time
+import math
+import colorsys
+import blup.frame
+import blup.output
+import blup.animation
+import blup.writebml
+import numpy as np
+import random
+
+WIDTH = 22
+HEIGHT = 16
+DEPTH = 256
+
+G = np.array([0, 9.81])
+
+class Fireball:
+    def __init__(self, startpos, startspeed, startcolor, lifetime, t0):
+        self.startpos, self.startspeed, self.startcolor, self.lifetime = \
+            startpos, startspeed, startcolor, lifetime
+        self.t0 = t0
+        self.pos = startpos
+        self.color = startcolor
+        self.dead = False
+
+    def update_pos(self, t):
+        t = t - self.t0
+        if t > self.lifetime:
+            self.dead = True
+            return
+        self.pos = 0.5 * G * t**2 + self.startspeed * t + self.startpos
+        self.color = self.startcolor * (1 - t/self.lifetime)
+        if random.randint(1, 10) > 8:
+            self.color = np.array([0, 0, 0])
+
+    def draw(self, frame):
+        if not self.dead:
+            frame.setPixel(int(round(self.pos[0])), int(round(self.pos[1])),
+                           tuple(map(lambda x: int(round(x)), self.color)))
+
+class Rocket:
+    def __init__(self, startpos, startspeed, explode_height, t0, color):
+        self.startpos, self.startspeed, self.explode_height, self.t0 = \
+            startpos, startspeed, explode_height, t0
+        self.color = color
+        self.pos = None
+        self.exploded = False
+
+    def update(self, t):
+        t = t - self.t0
+        if t < 0 or self.exploded:
+            return
+        self.pos = 0.5 * G * t**2 + self.startspeed * t + self.startpos
+        if self.pos[1] <= self.explode_height:
+            self.exploded = True
+            self.color = np.array([0, 0, 0])
+
+    def draw(self, frame):
+        if self.pos is None or self.exploded:
+            return
+        frame.setPixel(int(round(self.pos[0])), int(round(self.pos[1])),
+                       tuple(map(lambda x: int(round(x)), self.color)))
+
+        
+
+
+def getRandomColor(maxval):
+    return list(map(lambda x: int(round(x*maxval)), colorsys.hsv_to_rgb(random.random(), 1, 1)))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Generate plasma animations')
+    parser.add_argument('-d', '--delay', type=int, default=50)
+    parser.add_argument('-n', '--num-rockets', type=int, default=3)
+    parser.add_argument('output_file')
+    args = parser.parse_args()
+
+    dim = blup.frame.FrameDimension(WIDTH, HEIGHT, DEPTH, 3)
+    anim = blup.animation.Animation(dim)
+    anim.tags['description'] = ' '.join(sys.argv)
+
+    rockets = []
+    for i in range(args.num_rockets):
+        rockets.append(
+            Rocket(
+                np.array([random.randint(0, dim.width-1), dim.height-1]),
+                np.array([random.randint(-8, 8), -20 + random.randint(-2, 2)]),
+                4 + random.randint(-2, 2),
+                random.randint(0, 40) / 10,
+                np.array([dim.depth-1] * 3)
+            )
+        )
+
+    fbs = []
+    t = 0
+    while True:
+        frame = blup.animation.AnimationFrame(dim, args.delay)
+
+        for fb in fbs:
+            fb.update_pos(t)
+            try:
+                fb.draw(frame)
+            except ValueError:
+                continue
+        fbs = list(filter(lambda f: not f.dead, fbs))
+
+        for r in rockets:
+            r.update(t)
+            try:
+                r.draw(frame)
+            except ValueError:
+                continue
+            if r.exploded:
+                color = np.array(getRandomColor(dim.depth-1))
+                for i in range(random.randint(15, 25)):
+                    v = np.array([
+                        random.choice([-1, 1]) * random.randint(0, 25),
+                        random.choice([-1, 1]) * random.randint(0, 25),
+                    ])
+                    fbs.append(Fireball(r.pos, v, color, 1.5 + random.randint(-1, 1), t))
+        rockets = list(filter(lambda r: not r.exploded, rockets))
+
+        anim.addFrame(frame)
+
+        if len(rockets) == 0 and len(fbs) == 0:
+            break
+        t += args.delay / 1000
+
+    blup.writebml.writeBml(anim, args.output_file)