Fluid simulation from 1D cellular automata

eca_fluid_sim.py

#!/usr/bin/env python

# Idea: YankeeMinstrel
# Python code: torstengrust
# Animated GIF: Nicolas Seriot, 2017-07-01
# https://www.reddit.com/r/cellular_automata/comments/6jhdfw/i_used_1dimensional_cellular_automata_to_make_a/

from PIL import Image, ImageDraw
import os

def image(terrain, water):
    assert(len(terrain) == len(water))

    NB_CELLS = len(terrain)
    CELL_WIDTH = 48
    W, H = (NB_CELLS * CELL_WIDTH, 480)

    H_FACTOR = 10
    
    img = Image.new("RGB", (W, H), "white")
    draw = ImageDraw.Draw(img)

    for i in range(NB_CELLS):
        t_x0 = i * CELL_WIDTH
        t_y0 = 0
        t_x1 = t_x0 + CELL_WIDTH
        t_y1 = terrain[i] * H_FACTOR
        
        w_x0 = t_x0
        w_y0 = t_y1
        w_x1 = t_x1
        w_y1 = w_y0 + water[i] * H_FACTOR

        draw.rectangle([(t_x0, t_y0), (t_x1, t_y1)], fill="brown", outline="white")
        draw.rectangle([(w_x0, w_y0), (w_x1, w_y1)], fill="blue", outline="white")    
    
    return img.transpose(Image.FLIP_TOP_BOTTOM)

images = []

##

ground = [50, 1, 5, 8,12, 9, 5, 3, 2, 4, 8,12,15,20,50]  # elevation
water  = [ 0,20,20, 0, 0, 0, 0, 0, 0, 0,20,25,20,20, 0]  # water volume
width = len(ground)

energy = [0]*width

for iter in range(24*5):  # number of iterations to run
    
    dwater  = [0]*width
    denergy = [0]*width
    
    for x in range(1,width-1):
        if ground[x] + water[x] - energy[x] > ground[x-1] + water[x-1] + energy[x-1]:
            flow = min(water[x], ground[x] + water[x] - energy[x] - ground[x-1] - water[x-1] - energy[x-1]) / 4.0
            dwater[x-1]  += flow
            dwater[x]    += -flow
            denergy[x-1] += -energy[x-1] / 2 - flow
    
        if ground[x] + water[x] + energy[x] > ground[x+1] + water[x+1] - energy[x+1]:
            flow = min(water[x], ground[x] + water[x] + energy[x] - ground[x+1] - water[x+1] + energy[x+1]) / 4.0
            dwater[x+1]  += flow
            dwater[x]    += -flow
            denergy[x+1] += -energy[x+1] / 2 + flow
    
    water  = [sum(w) for w in zip(water, dwater)]
    energy = [sum(e) for e in zip(energy, denergy)]

    img = image(ground, water)
    
    images.append(img)

count = 0
for i in images:
    count += 1
    i.save("img" + str(count).zfill(4) + ".png")

def play(format):
   os.system("ffmpeg -i img%04d.png movie." + format)
   os.system("ffplay movie." + format)

play("gif")
#os.system("open -a Safari %s" % filename)