content: from __future__ import absolute_import
  from __future__ import print_function
  import autograd.numpy as np
  from autograd import value_and_grad
  
  from scipy.optimize import minimize
  from scipy.misc import imread
  
  import matplotlib
  import matplotlib.pyplot as plt
  import os
  from builtins import range
  
  # Fluid simulation code based on
  # "Real-Time Fluid Dynamics for Games" by Jos Stam
  # http://www.intpowertechcorp.com/GDC03.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('step{0:03d}.png'.format(t), mat)
      plt.pause(0.001)
  
  
  if __name__ == '__main__':
  
      simulation_timesteps = 100
  
      print("Loading initial and target states...")
      init_smoke = imread('init_smoke.png')[:,:,0]
      #target = imread('peace.png')[::2,::2,3]
      target = imread('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")
  