smoke_solver.py

# -----------------------------------------------------------------------------
# From Numpy to Python
# Copyright (2017) Nicolas P. Rougier - BSD license
# More information at https://github.com/rougier/numpy-book
# -----------------------------------------------------------------------------
"""
Real-Time Fluid Dynamics for Games by Jos Stam (2003).

Copyright (c) 2015 Alberto Santini - MIT License
Code adapted from Alberto Santini implementation available at:
https://github.com/albertosantini/python-fluid
"""
import numpy as np


def set_bnd(N, b, x):
    """We assume that the fluid is contained in a box with solid walls.

    No flow should exit the walls. This simply means that the horizontal
    component of the velocity should be zero on the vertical walls, while the
    vertical component of the velocity should be zero on the horizontal walls.
    For the density and other fields considered in the code we simply assume
    continuity. The following code implements these conditions.
    """

    if b == 1:
        x[0, 1:-1] = -x[1, 1:-1]
        x[-1, 1:-1] = -x[N, 1:-1]
    else:
        x[ 0, 1:-1] = x[1, 1:-1]
        x[-1, 1:-1] = x[N, 1:-1]
    if b == 2:
        x[1:-1,  0] = -x[1:-1, 1]
        x[1:-1, -1] = -x[1:-1, N]
    else:
        x[1:-1,  0] = x[1:-1, 1]
        x[1:-1, -1] = x[1:-1, N]

    x[ 0,  0] = 0.5 * (x[1,  0] + x[ 0, 1])
    x[ 0, -1] = 0.5 * (x[1, -1] + x[ 0, N])
    x[-1,  0] = 0.5 * (x[N,  0] + x[-1, 1])
    x[-1, -1] = 0.5 * (x[N, -1] + x[-1, N])


def lin_solve(N, b, x, x0, a, c):
    """lin_solve."""

    for k in range(20):
        x[1:-1, 1:-1] = (x0[1:-1, 1:-1] +
                         a * (x[:N, 1:-1] + x[2:, 1:-1] +
                              x[1:-1, :N] + x[1:-1, 2:])) / c
        set_bnd(N, b, x)


def add_source(N, x, s, dt):
    """Addition of forces: the density increases due to sources."""
    x += dt * s


def diffuse(N, b, x, x0, diff, dt):
    """Diffusion: the density diffuses at a certain rate.

    The basic idea behind our method is to find the densities which when
    diffused backward in time yield the densities we started with. The simplest
    iterative solver which works well in practice is Gauss-Seidel relaxation.
    """

    a = dt * diff * N * N
    lin_solve(N, b, x, x0, a, 1 + 4 * a)


def advect(N, b, d, d0, u, v, dt):
    """Advection: the density follows the velocity field.

    The basic idea behind the advection step. Instead of moving the cell
    centers forward in time through the velocity field, we look for the
    particles which end up exactly at the cell centers by tracing backwards in
    time from the cell centers.
    """

    dt0 = dt * N

    I, J = np.indices((N, N))
    I += 1
    J += 1
    X = I - dt0 * u[I, J]
    Y = J - dt0 * v[I, J]

    X = np.minimum(np.maximum(X, 0.5), N+0.5)
    I0 = X.astype(int)
    I1 = I0+1
    S1 = X - I0
    S0 = 1 - S1

    Y = np.minimum(np.maximum(Y, 0.5), N+0.5)
    J0 = Y.astype(int)
    J1 = J0 + 1
    T1 = Y - J0
    T0 = 1 - T1

    d[I, J] = (S0 * (T0 * d0[I0, J0] + T1 * d0[I0, J1])
             + S1 * (T0 * d0[I1, J0] + T1 * d0[I1, J1]))

    set_bnd(N, b, d)


def project(N, u, v, p, div):
    """ Projection """

    h = 1.0 / N
    div[1:-1, 1:-1] = (-0.5 * h *
                       (u[2:, 1:-1] - u[0:N, 1:-1] +
                        v[1:-1, 2:] - v[1:-1, 0:N]))
    p[1:-1, 1:-1] = 0
    set_bnd(N, 0, div)
    set_bnd(N, 0, p)
    lin_solve(N, 0, p, div, 1, 4)
    u[1:-1, 1:-1] -= 0.5 * (p[2:, 1:-1] - p[0:N, 1:-1]) / h
    v[1:-1, 1:-1] -= 0.5 * (p[1:-1, 2:] - p[1:-1, 0:N]) / h
    set_bnd(N, 1, u)
    set_bnd(N, 2, v)


def dens_step(N, x, x0, u, v, diff, dt):
    # Density step: advection, diffusion & addition of sources.
    add_source(N, x, x0, dt)
    x0, x = x, x0  # swap
    diffuse(N, 0, x, x0, diff, dt)
    x0, x = x, x0  # swap
    advect(N, 0, x, x0, u, v, dt)


def vel_step(N, u, v, u0, v0, visc, dt):
    # Velocity step: self-advection, viscous diffusion & addition of forces

    add_source(N, u, u0, dt)
    add_source(N, v, v0, dt)
    u0, u = u, u0  # swap
    
    diffuse(N, 1, u, u0, visc, dt)
    v0, v = v, v0  # swap
    
    diffuse(N, 2, v, v0, visc, dt)
    project(N, u, v, u0, v0)
    u0, u = u, u0  # swap
    v0, v = v, v0  # swap
    
    advect(N, 1, u, u0, u0, v0, dt)
    advect(N, 2, v, v0, u0, v0, dt)
    project(N, u, v, u0, v0)