text: from __future__ import absolute_import
  from __future__ import print_function
  from builtins import range
  import autograd.numpy as np
  from autograd import value_and_grad
  
  from scipy.optimize import minimize
  
  import matplotlib.pyplot as plt
  import os
  
  rows, cols = 40, 60
  
  # Fluid simulation code based on
  # "Real-Time Fluid Dynamics for Games" by Jos Stam
  # http://www.intpowertechcorp.com/GDC03.pdf
  
  def occlude(f, occlusion):
      return f * (1 - occlusion)
  
  def project(vx, vy, occlusion):
      """Project the velocity field to be approximately mass-conserving,
         using a few iterations of Gauss-Seidel."""
      p = np.zeros(vx.shape)
      div = -0.5 * (np.roll(vx, -1, axis=1) - np.roll(vx, 1, axis=1)
                  + np.roll(vy, -1, axis=0) - np.roll(vy, 1, axis=0))
      div = make_continuous(div, occlusion)
  
      for k in range(50):
          p = (div + np.roll(p, 1, axis=1) + np.roll(p, -1, axis=1)
                   + np.roll(p, 1, axis=0) + np.roll(p, -1, axis=0))/4.0
          p = make_continuous(p, occlusion)
  
      vx = vx - 0.5*(np.roll(p, -1, axis=1) - np.roll(p, 1, axis=1))
      vy = vy - 0.5*(np.roll(p, -1, axis=0) - np.roll(p, 1, axis=0))
  
      vx = occlude(vx, occlusion)
      vy = occlude(vy, occlusion)
      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_xs, cell_ys = np.meshgrid(np.arange(cols), np.arange(rows))
      center_xs = (cell_xs - vx).ravel()
      center_ys = (cell_ys - vy).ravel()
  
      # Compute indices of source cells.
      left_ix = np.floor(center_ys).astype(int)
      top_ix  = np.floor(center_xs).astype(int)
      rw = center_ys - left_ix              # Relative weight of right-hand cells.
      bw = center_xs - 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 make_continuous(f, occlusion):
      non_occluded = 1 - occlusion
      num = np.roll(f,  1, axis=0) * np.roll(non_occluded,  1, axis=0)\
          + np.roll(f, -1, axis=0) * np.roll(non_occluded, -1, axis=0)\
          + np.roll(f,  1, axis=1) * np.roll(non_occluded,  1, axis=1)\
          + np.roll(f, -1, axis=1) * np.roll(non_occluded, -1, axis=1)
      den = np.roll(non_occluded,  1, axis=0)\
          + np.roll(non_occluded, -1, axis=0)\
          + np.roll(non_occluded,  1, axis=1)\
          + np.roll(non_occluded, -1, axis=1)
      return f * non_occluded + (1 - non_occluded) * num / ( den + 0.001)
  
  def sigmoid(x):
      return 0.5*(np.tanh(x) + 1.0)   # Output ranges from 0 to 1.
  
  def simulate(vx, vy, num_time_steps, occlusion, ax=None, render=False):
      occlusion = sigmoid(occlusion)
  
      # Disallow occlusion outside a certain area.
      mask = np.zeros((rows, cols))
      mask[10:30, 10:30] = 1.0
      occlusion = occlusion * mask
  
      # Initialize smoke bands.
      red_smoke = np.zeros((rows, cols))
      red_smoke[rows//4:rows//2] = 1
      blue_smoke = np.zeros((rows, cols))
      blue_smoke[rows//2:3*rows//4] = 1
  
      print("Running simulation...")
      vx, vy = project(vx, vy, occlusion)
      for t in range(num_time_steps):
          plot_matrix(ax, red_smoke, occlusion, blue_smoke, t, render)
          vx_updated = advect(vx, vx, vy)
          vy_updated = advect(vy, vx, vy)
          vx, vy = project(vx_updated, vy_updated, occlusion)
          red_smoke = advect(red_smoke, vx, vy)
          red_smoke = occlude(red_smoke, occlusion)
          blue_smoke = advect(blue_smoke, vx, vy)
          blue_smoke = occlude(blue_smoke, occlusion)
      plot_matrix(ax, red_smoke, occlusion, blue_smoke, num_time_steps, render)
      return vx, vy
  
  def plot_matrix(ax, r, g, b, t, render=False):
      if ax:
          plt.cla()
          ax.imshow(np.concatenate((r[...,np.newaxis], g[...,np.newaxis], b[...,np.newaxis]), axis=2))
          ax.set_xticks([])
          ax.set_yticks([])
          plt.draw()
          if render:
              plt.savefig('step{0:03d}.png'.format(t), bbox_inches='tight')
          plt.pause(0.001)
  
  
  if __name__ == '__main__':
  
      simulation_timesteps = 20
  
      print("Loading initial and target states...")
      init_vx = np.ones((rows, cols))
      init_vy = np.zeros((rows, cols))
  
      # Initialize the occlusion to be a block.
      init_occlusion = -np.ones((rows, cols))
      init_occlusion[15:25, 15:25] = 0.0
      init_occlusion = init_occlusion.ravel()
  
      def drag(vx): return np.mean(init_vx - vx)
      def lift(vy): return np.mean(vy - init_vy)
  
      def objective(params):
          cur_occlusion = np.reshape(params, (rows, cols))
          final_vx, final_vy = simulate(init_vx, init_vy, simulation_timesteps, cur_occlusion)
          return -lift(final_vy) / drag(final_vx)
  
      # 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(weights):
          cur_occlusion = np.reshape(weights, (rows, cols))
          simulate(init_vx, init_vy, simulation_timesteps, cur_occlusion, ax)
  
      print("Rendering initial flow...")
      callback(init_occlusion)
  
      print("Optimizing initial conditions...")
      result = minimize(objective_with_grad, init_occlusion, jac=True, method='CG',
                        options={'maxiter':50, 'disp':True}, callback=callback)
  
      print("Rendering optimized flow...")
      final_occlusion = np.reshape(result.x, (rows, cols))
      simulate(init_vx, init_vy, simulation_timesteps, final_occlusion, ax, render=True)
  
      print("Converting frames to an animated GIF...")   # Using imagemagick.
      os.system("convert -delay 5 -loop 0 step*.png "
                "-delay 250 step{0:03d}.png wing.gif".format(simulation_timesteps))
      os.system("rm step*.png")
  