fluidsim.py

import os

import matplotlib
import matplotlib.pyplot as plt
from scipy.misc import imread
from scipy.optimize import minimize

import autograd.numpy as np
from autograd import value_and_grad

# Fluid simulation code based on
# "Real-Time Fluid Dynamics for Games" by Jos Stam
# https://www.josstam.com/_files/ugd/cf1fd6_9989229efbd34a26ba5ccd913721a2ac.pdf


def project(vx, vy):
    """Project the velocity field to be approximately mass-conserving,
    using a few iterations of Gauss-Seidel."""
    p = np.zeros(vx.shape)
    h = 1.0 / vx.shape[0]
    div = (
        -0.5
        * h
        * (
            np.roll(vx, -1, axis=0)
            - np.roll(vx, 1, axis=0)
            + np.roll(vy, -1, axis=1)
            - np.roll(vy, 1, axis=1)
        )
    )

    for k in range(10):
        p = (
            div
            + np.roll(p, 1, axis=0)
            + np.roll(p, -1, axis=0)
            + np.roll(p, 1, axis=1)
            + np.roll(p, -1, axis=1)
        ) / 4.0

    vx -= 0.5 * (np.roll(p, -1, axis=0) - np.roll(p, 1, axis=0)) / h
    vy -= 0.5 * (np.roll(p, -1, axis=1) - np.roll(p, 1, axis=1)) / h
    return vx, vy


def advect(f, vx, vy):
    """Move field f according to x and y velocities (u and v)
    using an implicit Euler integrator."""
    rows, cols = f.shape
    cell_ys, cell_xs = np.meshgrid(np.arange(rows), np.arange(cols))
    center_xs = (cell_xs - vx).ravel()
    center_ys = (cell_ys - vy).ravel()

    # Compute indices of source cells.
    left_ix = np.floor(center_xs).astype(int)
    top_ix = np.floor(center_ys).astype(int)
    rw = center_xs - left_ix  # Relative weight of right-hand cells.
    bw = center_ys - top_ix  # Relative weight of bottom cells.
    left_ix = np.mod(left_ix, rows)  # Wrap around edges of simulation.
    right_ix = np.mod(left_ix + 1, rows)
    top_ix = np.mod(top_ix, cols)
    bot_ix = np.mod(top_ix + 1, cols)

    # A linearly-weighted sum of the 4 surrounding cells.
    flat_f = (1 - rw) * ((1 - bw) * f[left_ix, top_ix] + bw * f[left_ix, bot_ix]) + rw * (
        (1 - bw) * f[right_ix, top_ix] + bw * f[right_ix, bot_ix]
    )
    return np.reshape(flat_f, (rows, cols))


def simulate(vx, vy, smoke, num_time_steps, ax=None, render=False):
    print("Running simulation...")
    for t in range(num_time_steps):
        if ax:
            plot_matrix(ax, smoke, t, render)
        vx_updated = advect(vx, vx, vy)
        vy_updated = advect(vy, vx, vy)
        vx, vy = project(vx_updated, vy_updated)
        smoke = advect(smoke, vx, vy)
    if ax:
        plot_matrix(ax, smoke, num_time_steps, render)
    return smoke


def plot_matrix(ax, mat, t, render=False):
    plt.cla()
    ax.matshow(mat)
    ax.set_xticks([])
    ax.set_yticks([])
    plt.draw()
    if render:
        matplotlib.image.imsave(f"step{t:03d}.png", mat)
    plt.pause(0.001)


if __name__ == "__main__":
    simulation_timesteps = 100
    basepath = os.path.dirname(__file__)

    print("Loading initial and target states...")
    init_smoke = imread(os.path.join(basepath, "init_smoke.png"))[:, :, 0]
    # target = imread('peace.png')[::2,::2,3]
    target = imread(os.path.join(basepath, "skull.png"))[::2, ::2]
    rows, cols = target.shape

    init_dx_and_dy = np.zeros((2, rows, cols)).ravel()

    def distance_from_target_image(smoke):
        return np.mean((target - smoke) ** 2)

    def convert_param_vector_to_matrices(params):
        vx = np.reshape(params[: (rows * cols)], (rows, cols))
        vy = np.reshape(params[(rows * cols) :], (rows, cols))
        return vx, vy

    def objective(params):
        init_vx, init_vy = convert_param_vector_to_matrices(params)
        final_smoke = simulate(init_vx, init_vy, init_smoke, simulation_timesteps)
        return distance_from_target_image(final_smoke)

    # Specify gradient of objective function using autograd.
    objective_with_grad = value_and_grad(objective)

    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, frameon=False)

    def callback(params):
        init_vx, init_vy = convert_param_vector_to_matrices(params)
        simulate(init_vx, init_vy, init_smoke, simulation_timesteps, ax)

    print("Optimizing initial conditions...")
    result = minimize(
        objective_with_grad,
        init_dx_and_dy,
        jac=True,
        method="CG",
        options={"maxiter": 25, "disp": True},
        callback=callback,
    )

    print("Rendering optimized flow...")
    init_vx, init_vy = convert_param_vector_to_matrices(result.x)
    simulate(init_vx, init_vy, init_smoke, simulation_timesteps, ax, render=True)

    print("Converting frames to an animated GIF...")
    os.system(
        "convert -delay 5 -loop 0 step*.png" " -delay 250 step100.png surprise.gif"
    )  # Using imagemagick.
    os.system("rm step*.png")