New Year 2021 PyWay demo

tree2021rt.py

# inspired by https://www.shadertoy.com/view/4st3DH

import array
import os
from math import *
import colorsys

W, H = 640, 480
OUT_FILE = 'test_2.ppm'

DOTS_TREE = 190
DOTS_STAR = 20
ANGLE = DOTS_TREE * 1.8
INTENSE_STAR = 1.0 / 2000.0
INTENSE_TREE = 1.0 / 5000.0
SCALE = 0.2
SCR_X, SCR_Y = 1.0, H / W
STAR_X, STAR_Y = SCR_X / 2.0, SCR_Y * 0.11


def clamp01(x):
    return min(1.0, max(0.0, x))


def hsv_to_rgb(h, s, v):
    r, g, b = colorsys.hsv_to_rgb(clamp01(h), clamp01(s), clamp01(v))
    return clamp01(r), clamp01(g), clamp01(b)


def add_tupples(t1, t2):
    return tuple(x1 + x2 for x1, x2 in zip(t1, t2))


def fract(x):
    return x - floor(x)


def distance(v1, v2):
    return max(1e-5, sqrt(sum((a1 - a2) ** 2 for a1, a2 in zip(v1, v2))))


def fragment_shader(x, y, t=0.5):
    u, v = x / W, y / W

    col = 0.0, 0.0, 0.0

    for i in range(DOTS_STAR):
        arg = 2.0 * pi * i / DOTS_STAR
        cos_arg = cos(arg)
        sin_arg = sin(arg)
        pos = (
            cos(t * 0.2) / 20.0 * cos_arg * SCALE + STAR_X,
            0.15 * sin_arg * SCALE + STAR_Y
        )
        d = distance((u, v), pos)
        col = add_tupples(col, hsv_to_rgb(
            i / DOTS_STAR,
            d / INTENSE_STAR,
            INTENSE_STAR / d
        ))

        pos = (
            0.12 * cos(arg + t * 0.2) * SCALE + STAR_X,
            0.08 * sin_arg * SCALE + STAR_Y
        )
        d = distance((u, v), pos)
        col = add_tupples(col, hsv_to_rgb(
            1.0 - i / DOTS_STAR,
            d / INTENSE_STAR,
            INTENSE_STAR / d
        ))

        pos = (
            0.12 * cos(arg + t * 0.2) * SCALE + STAR_X,
            -0.08 * sin_arg * SCALE + STAR_Y
        )
        d = distance((u, v), pos)
        col = add_tupples(col, hsv_to_rgb(
            i / DOTS_STAR,
            d / INTENSE_STAR,
            INTENSE_STAR / d
        ))

    for i in range(DOTS_TREE):
        pos = (
            SCR_X / 2.0 + sin(i / 2.0 - t * 0.2) / (3.0 / (i + 1.0) * ANGLE),
            SCR_Y * (i / DOTS_TREE + 0.21) * 0.8
        )

        d = distance((u, v), pos)
        col = add_tupples(col, hsv_to_rgb(
            1.5 * i / DOTS_TREE + fract(t / 4.0),
            d / INTENSE_STAR,
            INTENSE_STAR / d
        ))
    return col


def new_image(w, h):
    return array.array('B', [0, 0, 0] * w * h)


def write_ppm(filename, w, h, data):
    ppm_header = f'P6 {w} {h} 255\n'
    with open(filename, 'wb') as f:
        f.write(bytearray(ppm_header, 'ascii'))
        data.tofile(f)


if __name__ == '__main__':
    img = new_image(W, H)

    index = 0
    for y in range(H):
        for x in range(W):
            r, g, b = fragment_shader(x, y)
            img[index] = max(0, min(255, int(r * 255)))
            img[index + 1] = max(0, min(255, int(g * 255)))
            img[index + 2] = max(0, min(255, int(b * 255)))
            index += 3
        print(f'{y}/{H}')

    write_ppm(OUT_FILE, W, H, img)
    os.system(f'open "{OUT_FILE}"')  # MacOs