<s> from __future__ import absolute_import <EOL> from __future__ import print_function <EOL> import autograd . numpy as np <EOL> from autograd import value_and_grad <EOL> from scipy . optimize import minimize <EOL> from scipy . misc import imread <EOL> import matplotlib <EOL> import matplotlib . pyplot as plt <EOL> import os <EOL> from builtins import range <EOL> def project ( vx , vy ) : <EOL> """<STR_LIT>""" <EOL> p = np . zeros ( vx . shape ) <EOL> h = <NUM_LIT:1.0> / vx . shape [ <NUM_LIT:0> ] <EOL> div = - <NUM_LIT:0.5> * h * ( np . roll ( vx , - <NUM_LIT:1> , axis = <NUM_LIT:0> ) - np . roll ( vx , <NUM_LIT:1> , axis = <NUM_LIT:0> ) <EOL> + np . roll ( vy , - <NUM_LIT:1> , axis = <NUM_LIT:1> ) - np . roll ( vy , <NUM_LIT:1> , axis = <NUM_LIT:1> ) ) <EOL> for k in range ( <NUM_LIT:10> ) : <EOL> p = ( div + np . roll ( p , <NUM_LIT:1> , axis = <NUM_LIT:0> ) + np . roll ( p , - <NUM_LIT:1> , axis = <NUM_LIT:0> ) <EOL> + np . roll ( p , <NUM_LIT:1> , axis = <NUM_LIT:1> ) + np . roll ( p , - <NUM_LIT:1> , axis = <NUM_LIT:1> ) ) / <NUM_LIT> <EOL> vx -= <NUM_LIT:0.5> * ( np . roll ( p , - <NUM_LIT:1> , axis = <NUM_LIT:0> ) - np . roll ( p , <NUM_LIT:1> , axis = <NUM_LIT:0> ) ) / h <EOL> vy -= <NUM_LIT:0.5> * ( np . roll ( p , - <NUM_LIT:1> , axis = <NUM_LIT:1> ) - np . roll ( p , <NUM_LIT:1> , axis = <NUM_LIT:1> ) ) / h <EOL> return vx , vy <EOL> def advect ( f , vx , vy ) : <EOL> """<STR_LIT>""" <EOL> rows , cols = f . shape <EOL> cell_ys , cell_xs = np . meshgrid ( np . arange ( rows ) , np . arange ( cols ) ) <EOL> center_xs = ( cell_xs - vx ) . ravel ( ) <EOL> center_ys = ( cell_ys - vy ) . ravel ( ) <EOL> left_ix = np . floor ( center_xs ) . astype ( int ) <EOL> top_ix = np . floor ( center_ys ) . astype ( int ) <EOL> rw = center_xs - left_ix <EOL> bw = center_ys - top_ix <EOL> left_ix = np . mod ( left_ix , rows ) <EOL> right_ix = np . mod ( left_ix + <NUM_LIT:1> , rows ) <EOL> top_ix = np . mod ( top_ix , cols ) <EOL> bot_ix = np . mod ( top_ix + <NUM_LIT:1> , cols ) <EOL> flat_f = ( <NUM_LIT:1> - rw ) * ( ( <NUM_LIT:1> - bw ) * f [ left_ix , top_ix ] + bw * f [ left_ix , bot_ix ] ) + rw * ( ( <NUM_LIT:1> - bw ) * f [ right_ix , top_ix ] + bw * f [ right_ix , bot_ix ] ) <EOL> return np . reshape ( flat_f , ( rows , cols ) ) <EOL> def simulate ( vx , vy , smoke , num_time_steps , ax = None , render = False ) : <EOL> print ( "<STR_LIT>" ) <EOL> for t in range ( num_time_steps ) : <EOL> if ax : plot_matrix ( ax , smoke , t , render ) <EOL> vx_updated = advect ( vx , vx , vy ) <EOL> vy_updated = advect ( vy , vx , vy ) <EOL> vx , vy = project ( vx_updated , vy_updated ) <EOL> smoke = advect ( smoke , vx , vy ) <EOL> if ax : plot_matrix ( ax , smoke , num_time_steps , render ) <EOL> return smoke <EOL> def plot_matrix ( ax , mat , t , render = False ) : <EOL> plt . cla ( ) <EOL> ax . matshow ( mat ) <EOL> ax . set_xticks ( [ ] ) <EOL> ax . set_yticks ( [ ] ) <EOL> plt . draw ( ) <EOL> if render : <EOL> matplotlib . image . imsave ( '<STR_LIT>' . format ( t ) , mat ) <EOL> plt . pause ( <NUM_LIT> ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> simulation_timesteps = <NUM_LIT:100> <EOL> print ( "<STR_LIT>" ) <EOL> init_smoke = imread ( '<STR_LIT>' ) [ : , : , <NUM_LIT:0> ] <EOL> target = imread ( '<STR_LIT>' ) [ : : <NUM_LIT:2> , : : <NUM_LIT:2> ] <EOL> rows , cols = target . shape <EOL> init_dx_and_dy = np . zeros ( ( <NUM_LIT:2> , rows , cols ) ) . ravel ( ) <EOL> def distance_from_target_image ( smoke ) : <EOL> return np . mean ( ( target - smoke ) ** <NUM_LIT:2> ) <EOL> def convert_param_vector_to_matrices ( params ) : <EOL> vx = np . reshape ( params [ : ( rows * cols ) ] , ( rows , cols ) ) <EOL> vy = np . reshape ( params [ ( rows * cols ) : ] , ( rows , cols ) ) <EOL> return vx , vy <EOL> def objective ( params ) : <EOL> init_vx , init_vy = convert_param_vector_to_matrices ( params ) <EOL> final_smoke = simulate ( init_vx , init_vy , init_smoke , simulation_timesteps ) <EOL> return distance_from_target_image ( final_smoke ) <EOL> objective_with_grad = value_and_grad ( objective ) <EOL> fig = plt . figure ( figsize = ( <NUM_LIT:8> , <NUM_LIT:8> ) ) <EOL> ax = fig . add_subplot ( <NUM_LIT> , frameon = False ) <EOL> def callback ( params ) : <EOL> init_vx , init_vy = convert_param_vector_to_matrices ( params ) <EOL> simulate ( init_vx , init_vy , init_smoke , simulation_timesteps , ax ) <EOL> print ( "<STR_LIT>" ) <EOL> result = minimize ( objective_with_grad , init_dx_and_dy , jac = True , method = '<STR_LIT>' , <EOL> options = { '<STR_LIT>' : <NUM_LIT> , '<STR_LIT>' : True } , callback = callback ) <EOL> print ( "<STR_LIT>" ) <EOL> init_vx , init_vy = convert_param_vector_to_matrices ( result . x ) <EOL> simulate ( init_vx , init_vy , init_smoke , simulation_timesteps , ax , render = True ) <EOL> print ( "<STR_LIT>" ) <EOL> os . system ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> os . system ( "<STR_LIT>" ) </s>
