# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# %matplotlib inline

import matplotlib
import numpy as np
import matplotlib.pyplot as plt


# +
def calc_can_flow(obstacles):
    flow = np.zeros_like(obstacles)

    flow[-1, :] = 1

    for x in range(obstacles.shape[0]-2, -1, -1):
        flow[x] = np.logical_and(flow[x+1], np.logical_not(obstacles[x]))
        
        # FLOOD FILL
        seen = np.zeros(obstacles.shape[1])
        opens = []
        for i in range(obstacles.shape[1]):
            if flow[x, i]:
                opens.append((i))
        while opens:
            i = opens.pop()
            if i > 0 and not seen[i-1] and not obstacles[x, i-1]:
                flow[x, i-1] = 1
                opens.append(i-1)
                seen[i-1] = 1

            if i < obstacles.shape[1]-1 and not seen[i+1] and not obstacles[x, i+1]:
                flow[x, i+1] = 1
                opens.append(i+1)
                seen[i+1] = 1
                
    return flow

def flow_particle(x, y, v, grid, can_flow, dy=0):
    y2 = y
    while not can_flow[x+1, y2]:
        if can_flow[x, y2+dy]:
            y2 += dy
        else:
            return
    
    if y == y2:
        return
    
    v /= abs(y2 - y)

    while y != y2:
        y += dy
        grid[x, y] += v

def calculte_flow(can_flow):
    flow = np.zeros_like(can_flow, dtype='f')
    nx = can_flow.shape[0]
    ny = can_flow.shape[1]
    flow[0] = 1
    
    for x in range(1, nx):
        xm1 = flow[x-1].copy()

        for y in range(ny):
            if not can_flow[x, y] and flow[x-1, y]:
                if y > 0:
                    flow_particle(x-1, y, xm1[y]/2, flow, can_flow, -1)
                if y < ny - 1:
                    flow_particle(x-1, y, xm1[y]/2, flow, can_flow, 1)

        for y in range(ny):
            if can_flow[x, y]:
                above = flow[x-1, y-1] if y > 0 and not obstacles[x,y-1] else flow[x-1, y]
                below = flow[x-1, y+1] if y < ny-1 and not obstacles[x,y+1] else flow[x-1, y]
                flow[x, y] = (.5*flow[x-1, y] + .25*above + .25*below)

    return flow



# +
shape = (20, 10)
obstacles = np.zeros(shape, dtype='uint8')
flow = np.zeros(shape+(2,), dtype='uint8')

obstacles[10:12, 2:8] = 1
obstacles[8:10, 2] = 1
obstacles[8:10, 7] = 1
obstacles[15, 1:5] = 1

#  obstacles[10:12, 4:6] = 0

# -

plt.imshow(obstacles.T)

can_flow = calc_can_flow(obstacles)
plt.imshow(can_flow.T)

flow = calculte_flow(can_flow)
_ = plt.imshow(flow.T)

np.set_printoptions(precision=2)
print(flow)



StarcoderdataJupyterScriptsDedupFiltered