Pygame Strange Attractor.

main.py

"""
    main.py
    Author: Robert (Duality)
    Git: https://github.com/Duality4y

    notes:
        this is an implementation of the strange attractor in python.
        using pygame to draw it.
        details can be found here:
        https://en.wikipedia.org/wiki/Lorenz_system
    
        This program is hungry for cycles though, because it tries to do the
        Transformation as fast as possible.

        control:
        press i to increase the scale
        press d to decrease the scale
        press r to reset scale

        press c to clear the screen once
        press u to continuesly clear the screen.
    
    LICENSE:
        This Project
        Copyright (C) 2018

        This program is free software: you can redistribute it and/or modify
        it under the terms of the GNU General Public License as published by
        the Free Software Foundation, either version 3 of the License, or
        (at your option) any later version.

        This program is distributed in the hope that it will be useful,
        but WITHOUT ANY WARRANTY; without even the implied warranty of
        MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
        GNU General Public License for more details.

        You should have received a copy of the GNU General Public License
        along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import sys
import time
import random
from operator import add, sub, mul, truediv

import pygame

def draw_pixel(surface, color, point):
    # unpack 2 cordinates of the point (x, y) into the rect x, y
    pygame.draw.rect(surface, color, (*point[0:2:1], 1, 1), 1)

class Point(list):
    def __init__(self, *args, **kwargs):
        self.attrs = ["x", "y", "z"]
        if len(args):
            super(Point, self).__init__(args)
        else:
            self.append(kwargs.get("x", 0))
            self.append(kwargs.get("y", 0))
            self.append(kwargs.get("z", 0))

    @property
    def x(self):
        return self[0]

    @x.setter
    def x(self, value):
        self[0] = value

    @property
    def y(self):
        return self[1]

    @y.setter
    def y(self, value):
        self[1] = value

    @property
    def z(self):
        return self[2]

    @z.setter
    def z(self, value):
        self[2] = value

    def __add__(self, other):
        return Point([add(*value) for value in zip(self, other)])

    def __sub__(self, other):
        return Point([sub(*value) for value in zip(self, other)])

    def __mul__(self, other):
        return Point([mul(*value) for value in zip(self, other)])

    def __truediv__(self, other):
        return Point([truediv(*value) for value in zip(self, other)])

    def __repr__(self):
        return ("Point(%s)" % super(Point, self).__repr__())

if __name__ == "__main__":
    pygame.init()
    surface = pygame.display.set_mode((800, 600))

    # color definitions
    red = (0xff, 0, 0)
    green = (0, 0xff, 0)
    blue = (0, 0, 0xff)
    white = (0xff, 0xff, 0xff)
    black = (0, 0, 0)
    
    # transformation variables.
    p = Point(0.1, 0.42, 0.90)
    sigma = 10
    rho = 28.
    beta = 8. / 3.
    dt = 0.0001
    scale = 5

    # timeing variables.
    current = time.time()
    previous = 1.0
    interval = 1 / 30.

    # state and program variables
    should_clear = False
    clear_once = False
    bgcolor = black

    # main program loop
    # all the processing amd
    # all the screen updating happens here.
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYUP:
                if event.key == pygame.K_c:
                    clear_once = True
                elif event.key == pygame.K_u:
                    should_clear = not should_clear
                elif event.key == pygame.K_i:
                    scale *= 1.2
                    clear_once = True
                elif event.key == pygame.K_d:
                    scale /= 1.2
                    clear_once = True
                elif event.key == pygame.K_r:
                    scale = 5
                    clear_once = True

        # apply the transformation
        dx = (sigma * (p.y - p.x)) * dt
        dy = (p.x * (rho - p.z) - p.y) * dt
        dz = (p.x * p.y - beta * p.z) * dt
        p.x = p.x + dx
        p.y = p.y + dy
        p.z = p.z + dz

        # make a new point translated over to the center of the screen.
        # and also scale a bit.
        cp = Point(p.x * scale + 400, p.y * scale + 300, p.z * scale)

        # draw the pixel that is on the attractor.
        draw_pixel(surface, green, cp)

        # only update the screen 30 times a second.
        # such a performance hit to update every single itteration of the loop. 
        current = time.time()
        if (current - previous) >= interval:
            previous = current
            pygame.display.update()

            # only clear when it is allowed.
            # and also do it here so not to hava a to big performance hit.
            if should_clear:
                surface.fill(bgcolor)
            elif clear_once:
                clear_once = False
                surface.fill(bgcolor)