air.py

import pygame
import random


class Node:
    def __init__(self, x, y, air=0.0, wall=False):
        self.x = x
        self.y = y
        self.air = air
        self.wall = wall


# Init pygame
pygame.init()
size = width, height = 1024, 1024
screen = pygame.display.set_mode(size)
pygame.display.set_caption("Air Stuff")

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Init grid of nodes
grid_x, grid_y = 32, 32
grid = []
for gx in range(grid_x):
    grid.append([])
    for gy in range(grid_y):
        grid[gx].append(Node(gx, gy))


def get_neighbors(x, y):
    neighbors = []
    for nx in range(max(x - 1, 0), min(x + 1, grid_x - 1) + 1):
        for ny in range(max(y - 1, 0), min(y + 1, grid_y - 1) + 1):
            if (nx, ny) != (x, y):
                neighbors.append(grid[nx][ny])
    return neighbors


done = False
clock = pygame.time.Clock()
FRAME_RATE = 144
wall_paint = None
while not done:
    clock.tick(FRAME_RATE)
    frame_rate_ratio = clock.get_time() / (1000 / FRAME_RATE)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            done = True

    # Click to add / remove air
    left, mid, right = pygame.mouse.get_pressed()
    if left or mid or right:
        m_pos = pygame.mouse.get_pos()
        click_node = grid[int(m_pos[0] / (width / grid_x))][int(m_pos[1] // (height / grid_y))]
        if left:
            click_node.air = 5
        elif right:
            click_node.air = -5
        elif mid:
            if wall_paint is None:
                wall_paint = not click_node.wall
            click_node.wall = wall_paint
    else:
        wall_paint = None

    # Update air pressure
    for gx in range(grid_x):
        for gy in range(grid_y):
            node = grid[gx][gy]
            if node.wall:
                continue

            neighbors_low = [n for n in get_neighbors(gx, gy)
                             if not n.wall and n.air < node.air]
            random.shuffle(neighbors_low)
            if neighbors_low:
                delta_air = frame_rate_ratio * max(0.1, 0.1 * node.air) / len(neighbors_low)
                for neighbor in neighbors_low:
                    if neighbor.air + delta_air > node.air - delta_air:
                        neighbor.air = node.air = (neighbor.air + node.air) / 2
                    else:
                        neighbor.air += delta_air
                        node.air -= delta_air

    # Draw grid
    screen.fill(WHITE)
    for gx in range(grid_x):
        for gy in range(grid_y):
            node = grid[gx][gy]
            node_pos = (gx * (width / grid_x), gy * (height / grid_y))
            node_size = (width / grid_x, height / grid_y)

            node_color_val = max(0, min(255, int(255 * (1 - node.air))))
            node_color = (255, node_color_val, node_color_val)
            if node.air < 0:
                node_color_val = max(0, min(255, int(255 * (1 + node.air))))
                node_color = (node_color_val, node_color_val, 255)

            pygame.draw.rect(screen, BLACK if node.wall else node_color, [*node_pos, *node_size])
            # pygame.draw.rect(screen, BLACK, [*node_pos, *node_size], 1) # grid

    pygame.display.flip()

pygame.quit()