iso6.py

# Nicolas Seriot
# 2021-09-19
# https://gist.github.com/nst/070dda2056925eec07446c5f3f3c67df
# Thread: https://twitter.com/nst021/status/1437889678947110912
# Typical output: https://seriot.ch/visualization/iso6.gif

import cairo
import random
import numpy as np

import cProfile, pstats, io
from pstats import SortKey

import imageio
import os

MARGIN = 10
X_MAX, Y_MAX, Z_MAX = 24,24,16 # space size
DW, DH = 28, 14 # diamond size
DW2, DH2 = int(DW/2), int(DH/2)
TTL = 32

def draw_surface(c, points, color, vertices = []):
    
    c.save()

    c.set_source_rgb(*color)
    
    [c.line_to(*p) for p in points]

    c.fill()
    
    c.set_source_rgb(0,0,0)
    c.set_line_width(1)
    
    for pts in vertices:
        p1, p2 = pts
        c.move_to(*p1)
        c.line_to(*p2)
        c.stroke()
    
    c.restore()

def draw_grid(c):

    for x in range(X_MAX):
        for y in range(Y_MAX):        
            draw_cube(c, x, y, 0, flat=True)

def model_to_canvas(x, y):

    o_x = DW2 * Y_MAX
    o_y = DH2
    
    return (o_x + DW2*(x-y), o_y + DH2*(x+y))

def draw_cube(c, x, y, z,
              nx = False, nx_ = False, ny = False, ny_ = False, nz = False, nz_ = False,
              nxy_=False, nxz = False, nx_z_ = False, nyz = False, ny_z_ = False,
              flat=False):
    
    x, y = model_to_canvas(x, y)
    
    c.save()

    c.translate(x, y)

    """
      5
    2   6
      3      - DH
    1   7
      4
      | z_offset
    """
    
    z_offset = DH * z

    p1 = (0,   z_offset + DH2)
    p2 = (0,   z_offset + DH2 + DH)
    p3 = (DW2, z_offset + DH)
    p4 = (DW2, z_offset + 0)
    p5 = (DW2, z_offset + 2*DH)
    p6 = (DW,  z_offset + DH2 + DH)
    p7 = (DW,  z_offset + DH2)
    
    #COLOR_RIGHT = (0.5,0.5,0.5)
    #COLOR_TOP = (1,1,1)
    #COLOR_LEFT = (0,0,0)
    
    z_ratio = 0.4 + z*0.6/Z_MAX
    
    COLOR_TOP   = (1 * z_ratio,   0.5 * z_ratio, 0.5 * z_ratio)
    COLOR_LEFT  = (0.6 * z_ratio, 0,             0)
    COLOR_RIGHT = (1 * z_ratio,   0,             0)

    vertices = []
    
    if flat:
        vertices.append((p1,p3))
        vertices.append((p3,p7))
        vertices.append((p7,p4))
        vertices.append((p4,p1))
        draw_surface(c, [], (1,1,1), vertices)
    else:
    
        if not nx:#
            vertices.append((p5, p6))
            vertices.append((p6, p7))
        if not ny:
            vertices.append((p2, p5))
            vertices.append((p2, p1))
        if not nz_:
            vertices.append((p4, p7))
            vertices.append((p1, p4))
        if not nx_ and not ny_:
            vertices.append((p3, p4))
        if not nz and not ny_:
            vertices.append((p6, p3))
        if not nz and not nx_:
            vertices.append((p2, p3))
        if nxy_:
            vertices.append((p6, p7))
        if nx_z_:
            vertices.append((p1, p4))
        if nxz:
            vertices.append((p5, p6))
        if nyz:
            vertices.append((p2, p5))
        if ny_z_:
            vertices.append((p4, p7))
        
        draw_surface(c, [p2, p5, p6, p3], COLOR_TOP, vertices)
        draw_surface(c, [p1, p2, p3, p4], COLOR_LEFT, vertices)
        draw_surface(c, [p3, p6, p7, p4], COLOR_RIGHT, vertices)
    
    c.restore()

def visibility_matrix(m):

    X,Y,Z = m.shape

    # nothing is visible except the three visible faces of the space

    v = np.full(m.shape, False)
    
    for x in range(X):
        for y in range(Y):
            v[x][y][Z-1] = True

    for x in range(X):
        for z in range(Z):
            v[x][0][z] = True

    for y in range(Y):
        for z in range(Z):
            v[0][y][z] = True

    # iterating from user's standpoint
    
    for x in range(X):
        for y in range(Y):
            for z in range(Z)[::-1]:

                # if not visible
                # no need to update visibility of "back" cube
                # continue to the next cube
                if not v[x][y][z]:
                    continue

                # if m is empty
                # "back" cube becomes visible
                if not m[x][y][z]:
                    if x < (X-1) and y < (Y-1) and z > 0:
                        v[x+1][y+1][z-1] = True
        
    return v

def offset_is_valid(m,x,y,z,a,b,c):
    X,Y,Z = m.shape
    return 0 <= (x+a) < X and 0 <= (y+b) < Y and 0 <= (z+c) < Z

def draw_model(ctx, m, oids):
    
    X,Y,Z = m.shape
    
    v = visibility_matrix(m)
        
    for x in range(X)[::-1]:
        for y in range(Y)[::-1]:
            for z in range(Z_MAX):
                
                if not v[x][y][z]:
                    continue

                oid = oids[x][y][z]
                if oid == 0:
                    continue
                
                (a,b,c) = (1,0,0)
                nx = offset_is_valid(m,x,y,z,a,b,c) and oids[x+a,y+b,z+c] == oid
                
                (a,b,c) = (-1,0,0)
                nx_ = offset_is_valid(m,x,y,z,a,b,c) and oids[x+a,y+b,z+c] == oid

                (a,b,c) = (0,1,0)
                ny = offset_is_valid(m,x,y,z,a,b,c) and oids[x+a,y+b,z+c] == oid

                (a,b,c) = (0,-1,0)
                ny_ = offset_is_valid(m,x,y,z,a,b,c) and oids[x+a,y+b,z+c] == oid

                (a,b,c) = (0,0,1)
                nz = offset_is_valid(m,x,y,z,a,b,c) and oids[x+a,y+b,z+c] == oid
                
                (a,b,c) = (0,0,-1)
                nz_ = offset_is_valid(m,x,y,z,a,b,c) and oids[x+a,y+b,z+c] == oid

                nxy_  = m[x+1][y-1][z] if (x+1) < X  and (y-1) >= 0 else False

                nxz   = m[x+1][y][z+1] if (x+1) < X  and (z+1) < Z  else False
                nx_z_ = m[x-1][y][z-1] if (x-1) >= 0 and (z-1) >= 0 else False

                nyz   = m[x][y+1][z+1] if (y+1) < Y  and (z+1) < Z  else False
                ny_z_ = m[x][y-1][z-1] if (y-1) >= 0 and (z-1) >= 0 else False
                
                draw_cube(ctx,x,y,z,nx,nx_,ny,ny_,nz,nz_,nxy_,nxz,nx_z_,nyz,ny_z_)

def images_paths():
    # ensure we can read images when working from another dir
    cwd = os.getcwd()
    dir_path = os.path.dirname(os.path.realpath(__file__))
    os.chdir(dir_path)
    files = [os.path.sep.join([dir_path, s])
             for s in os.listdir(dir_path)
             if s.endswith('.png')]
    files.sort()
    os.chdir(cwd)
    return files

def draw_gif(filename):

    images = []
    
    for s in images_paths():
        images.append(imageio.imread(s))
    
    imageio.mimsave(filename, images, format='GIF', duration=0.1)

def draw_png(m, oids, filename):

    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32,
                                 DW + (X_MAX+Y_MAX)*DW2 + MARGIN*2,
                                 DH + (X_MAX+Y_MAX)*DH2 + Z_MAX*DH + MARGIN*2)
    c = cairo.Context(surface)
    
    cm = cairo.Matrix(yy=-1, y0=surface.get_height())
    c.transform(cm)
    
    c.translate(MARGIN, MARGIN)
    
    c.set_antialias(cairo.ANTIALIAS_NONE)
    
    #background
    c.set_source_rgb(1,1,1)
    c.paint()
    
    #pen
    c.set_source_rgb(0,0,0)
    c.set_line_width(1)
    
    #draw_grid(c)
    draw_model(c, m, oids)

    surface.write_to_png(filename)

def can_move_to(m, x, y, z):

    X,Y,Z = m.shape

    if not x in range(X):
        return False

    if not y in range(Y):
        return False

    if not z in range(Z):
        return False
    
    if m[x][y][z] > 0:
        return False

    return True    

def next_head(m, head, direction):

    x,y,z = head
    a,b,c = direction
    
    must_turn = can_move_to(m, x+a, y+b, z+c) == False

    keep_direction = random.uniform(0, 1) < 0.7

    if not must_turn and keep_direction:
        return (x+a, y+b, z+c), direction
        
    directions = [(-1,0,0), (1,0,0),
                  (0,-1,0), (0,1,0),
                  (0,0,-1), (0,0,1)]
               
    possible_directions = []

    for a,b,c in directions:
    
        x_ = x+a
        y_ = y+b
        z_ = z+c
        
        if x_ in range(X_MAX) \
            and y_ in range(Y_MAX) \
            and z_ in range(Z_MAX) \
            and m[x_][y_][z_] == 0:
            possible_directions.append((a,b,c))
        
    if len(possible_directions) == 0:
        print("** can't move")
        return head, direction
        
    a,b,c = random.choice(possible_directions)
    new_head = (x+a,y+b,z+c)
    
    return new_head, (a,b,c)

def main():

    #random.seed(0)

    m    = np.full((X_MAX, Y_MAX, Z_MAX), 0) # TTLs
    oids = np.full((X_MAX, Y_MAX, Z_MAX), 0) # object ids
    
    directions = [(-1,0,0), (1,0,0),
                  (0,-1,0), (0,1,0),
                  (0,0,-1), (0,0,1)]
    
    snakes = []
    for _ in range(40):
        d = random.choice(directions)
        x = random.randrange(X_MAX)
        y = random.randrange(Y_MAX)
        z = random.randrange(Z_MAX)
        snake = {"h":(x,y,z), "d":d}
        snakes.append(snake)
    
    for oid,snake in enumerate(snakes):
        x,y,z = snake["h"]
        m[x][y][z] = TTL # time to live
        oids[x][y][z] = oid+1 # 0 means no object
        
    draw_png(m, oids, "000.png")

#    pr = cProfile.Profile()
#    pr.enable()

    for i in range(1, 250):
    
        for oid, s in enumerate(snakes):
            h = s["h"]
            d = s["d"]
            
            h_, d_ = next_head(m, h, d)
            s["h"] = h_
            s["d"] = d_
            x,y,z = h_
            m[x][y][z] = TTL
            oids[x][y][z] = oid+1

        for index, ttl in np.ndenumerate(m):
            new_ttl = max(0, ttl-1)
            m[index] = new_ttl
            if new_ttl == 0:
                oids[index] = 0
            
        filename = f'{i:03}.png'
        draw_png(m, oids, filename)
        print(filename)

#    pr.disable()
#    s = io.StringIO()
#    sortby = SortKey.CUMULATIVE
#    ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
#    ps.print_stats()
#    print(s.getvalue())

    draw_gif("iso6.gif")

if __name__ == "__main__":
    
    main()