CFD Python (close to realism)

CFD.py

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

n = 50
density = 0
u = 0
v = 0
accuracy = 10
time_step = 0.5
diffusion_amount = 0.1
advection_amount = 10
im = 0
cb = 0
frametime = 0.01
pressure_global = 0
quiv = 0

def main():
    setup()
    display()
    while (True):
        #print(internal_sum(density))
        evolve_density()
        evolve_velocity()
        update()

def update():

    display = density

    global quiv
    im.set_data(display.T)
    vmin = min(-0.1, display.min())
    vmax = max(0.1, display.max())
    im.set_norm(colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax))
    quiv.remove()
    quiv = plt.quiver(u.T, -1 * v.T)
    plt.show()
    plt.pause(frametime)

def internal_sum(array):
    sum = 0
    for i in range(1, n+1):
        for j in range(1, n+1):
            sum += array[i,j]

    return sum

def display():
    global cb
    global im
    global quiv
    im = plt.imshow(density.T, interpolation='none', cmap='ocean')
    cb = plt.colorbar()
    quiv = plt.quiver(u.T,-1*v.T)

    plt.show()
    plt.pause(frametime)


def setup():
    global density, u, v, im
    plt.ion()
    density = set_bounds(0, density_init())
    u = set_bounds(1, -1*np.zeros((n + 2, n + 2)))
    v = set_bounds(2, -1*np.zeros((n + 2, n + 2)))


def density_init():
    density = np.zeros((n+2, n+2))
    #density[10,10] = 50
    #density[10, 40] = 50
    #density[40, 10] = 50
    #density[40, 40] = 50
    #density[44,46] = 50
    #density[46, 44] = 50
    return density


def evolve_density():
    global density
    density = diffusion(density, 0)
    density = advect(density, 0)
    #density[10, 10] = 50
    #density[10, 40] = 50
    #density[40, 10] = 50
    #density[40, 40] = 50
    #density[44, 46] = 50
    #density[46, 44] = 50
    #density[26, 25] = 50
    density[10,10] = 50
def evolve_velocity():
    global u, v

    u[9, 10] = 5

    u = diffusion(u, 1)
    v = diffusion(v, 2)
    conserve_mass()

    u = advect(u, 1)
    v = advect(v, 2)
    conserve_mass()

def diffusion(density, boundary_type):
    # factor = speed * dt
    # d_old = d_new - factor * (surrounding - 4 * d_new)
    # ...
    # d_new = (x_old + factor * (surrounding)) / (1 + 4 * factor)

    factor = time_step * diffusion_amount

    for k in range(accuracy):
        for x in range(1, n + 1):
            for y in range(1, n + 1):
                density[x, y] = (density[x, y] + factor * (density[x + 1, y] + density[x - 1, y] + density[x, y + 1] +
                                                           density[x, y - 1])) / (1 + 4 * factor)
        density = set_bounds(boundary_type, density)

    return density


def advect(array, boundary_type):
    factor = time_step * advection_amount
    new_array = np.zeros((n + 2, n + 2))
    for x in range(1, n + 1):
        for y in range(1, n + 1):
            from_x = clamp(x - factor * u[x, y], 0.5, n + 0.5)
            from_y = clamp(y - factor * v[x, y], 0.5, n + 0.5)
            x_floor = int(from_x)
            y_floor = int(from_y)
            ratio_x = from_x - x_floor
            ratio_y = from_y - y_floor
            #value_interpolate = linear_interpolate(

                #linear_interpolate(array[x_floor, y_floor], array[x_floor + 1, y_floor], ratio_x),
                #linear_interpolate(array[x_floor, y_floor + 1], array[x_floor + 1, y_floor + 1], ratio_x),
                #ratio_y)

            x_ceil = x_floor + 1
            y_ceil = y_floor + 1

            s1 = ratio_x
            s0 = 1 - s1
            t1 = ratio_y
            t0 = 1 - t1

            value_interpolate = s0*(t0*array[x_floor, y_floor]+t1*array[x_floor, y_ceil]) + s1*(t0*array[x_ceil, y_floor]+t1*array[x_ceil, y_ceil])


            new_array[x, y] = value_interpolate

    return set_bounds(boundary_type, new_array)


def linear_interpolate(a, b, ratio):
    return a * (1 - ratio) + b * ratio


def clamp(num, min_value, max_value):
    return max(min(num, max_value), min_value)


def divergence(u, v):
    h = 1 / n
    div = np.zeros((n + 2, n + 2))
    for x in range(1, n + 1):
        for y in range(1, n + 1):
            div[x, y] = -0.5 * h * ((u[x+1, y] - u[x-1, y]) + v[x, y+1] - v[x, y-1])
            # zeros can be left for boundaries; not used in any calculations (linear solve)
    return set_bounds(0, div)


def pressure_from_divergence(divergence):
    # difference in pressure = divergence
    # reverse this, solve for pressure from the equation
    # (sum p[surrounding]) - 4*(p[xy]) = divergence
    # p[xy] = (sum(p[surrounding]) - divergence)/4

    global pressure_global

    pressure = np.zeros((n + 2, n + 2))

    for k in range(accuracy):
        for x in range(1, n + 1):
            for y in range(1, n + 1):
                pressure[x, y] = (divergence[x,y] + pressure[x-1, y] + pressure[x+1, y] + pressure[x, y+1] + pressure[x, y-1])/4

        pressure = set_bounds(0, pressure)  # sets pressure boundaries to surrounding equal each repetition

    pressure_global = pressure
    return pressure


def gradient(pressure):
    # gradient of a vector field returns the direction of the steepest ascent, returns [u,v] of this vector
    # (pressure-induced-velocity)

    u = np.zeros((n + 2, n + 2))
    v = np.zeros((n + 2, n + 2))

    for x in range(1, n + 1):
        for y in range(1, n + 1):
            u[x, y] = 0.5*(pressure[x + 1, y] - pressure[x - 1, y]) / (1 / n)
            v[x, y] = 0.5*(pressure[x, y + 1] - pressure[x, y - 1]) / (1 / n)

    #u = set_bounds(1, u)
    #v = set_bounds(2, v)

    return [u, v]


def set_bounds(type, a):
    for i in range(1, n + 1):
        a[0, i] = (-1 * a[1, i]) if type == 1 else (a[1, i])
        a[n + 1, i] = (-1 * a[n, i]) if type == 1 else (a[n, i])
        a[i, 0] = (-1 * a[i, 1]) if type == 2 else (a[i, 1])
        a[i, n + 1] = (-1 * a[i, n]) if type == 2 else (a[i, n])

    a[0, 0] = a[1, 1] if type == 0 else 0
    a[0, n + 1] = a[1, n] if type == 0 else 0
    a[n + 1, 0] = a[n, 1] if type == 0 else 0
    a[n + 1, n + 1] = a[n, n] if type == 0 else 0

    return a


def conserve_mass():
    global u, v
    divergence_array = divergence(u, v)
    pressure = pressure_from_divergence(divergence_array)
    induced_velocity = gradient(pressure)
    compensate_x = induced_velocity[0]
    compensate_y = induced_velocity[1]
    u = u - compensate_x
    v = v - compensate_y
    u = set_bounds(1, u)
    v = set_bounds(2, v)

main()