# ----------------------------------------------------------------------------- # 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)