Init commit

.gitignore

+__pycache__/
+*.py[cod]
+*$py.class
+explore*
+*data/
+*backup/
+*vdb/
+build/
+*.npy
+*.csv
+.vscode
+.venv*
+*.dll
\ No newline at end of file

LICENCE

+MIT License
+
+Copyright © 2021, ETH Zurich, Jingwei Tang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

+# Shrunk-Domain
+
+Official PyTorch implementation for [Honey, I Shrunk the Domain: Frequency-aware Force Field Reduction for Efficient Fluids Optimization](https://cgl.ethz.ch/publications/papers/paperTan21a.php).
+
+### Installation
+``` bash
+# create a conda env
+conda create -n shrunk-domain
+conda activate shrunk-domain
+# install pytorch
+conda install pytorch==1.7.1 torchvision==0.8.2 cudatoolkit=11.0 -c pytorch
+# install other requirements
+pip install -r requirements.txt
+```
+
+### Simulation
+The sequences for the benchmark tests can be generated through simulation by running the provided simulation scripts. For example, to generate 2D uniform-wind simulation:
+```bash
+bash scripts/simulate_2d_uniform.sh
+```
+
+### Optimization
+Optimizations can be used for keyframe matching. Before running the optimization scripts, we need to put the keyframe data in `.npz` format into `simulations/data` directory. For example, the above 2D 128x128 uniform-wind simulation will create `.npz` files under `simulations/data/benchmark/2d_128_uniform`. A corresponding `opt.txt` file is also needed to specify simulation configurations as guidance to the simulations used in the optimization process.
+
+After the data and config files are prepared, we can start to run the optimization. The provided scripts give examples to run different optimization methods: baseline+TV `bash scripts/optimize_2d_baseline.sh`, progressive upsampling `bash scripts/optimize_2d_prog_force.sh` and ours `bash scripts/optimize_2d_prog_freq.sh`.
+
+After the optimizaiton is done, we can run a simulation by loading the saved optimized parameter to get to the target frame:
+
+``` bash
+# need to specify EXPERIMENT_NAME in this script first
+bash scripts/optimize_simulate.sh
+```
+
+### Structure of the code
+- `lib/dfluids` implements differentiable operations needed to run smoke simulations.
+- `lib/PyTorch_LBFGS.py` is adapted from [PyTorch-LBFGS](https://github.com/hjmshi/PyTorch-LBFGS), the optimization processes use this implementation of LBFGS optimizer.
+- `simulations` defines argparse interfaces for users to define simulation config, and uses operations from `dfluids` to run smoke simulations.
+- `optimizations/models` defines models to wrap optimization parameters and how the parameters are used in the simulation.
+- `optimizations/optimizations` specifies optimization processes, e.g. procedures to compute objective function, to optimize the parameters, etc.
+- `forces` contains scripts to generate force fields used for the benchmark tests.

forces/generate_multi_vortex_128.py

+import os
+import sys
+import torch
+import numpy as np
+
+sys.path.append(os.getcwd())
+from lib.dfluids.solver import FluidSolver
+from lib.dfluids.grid import GridType
+import lib.dfluids.utils as utils
+
+# 2d big curl force
+res = 128
+solver = FluidSolver(grid_size=(res, res))
+ext = solver.create_extforces()
+mac_op = solver.create_grid_operator(GridType.MAC)
+real_op = solver.create_grid_operator(GridType.REAL)
+node_op = solver.create_grid_operator(GridType.NODE)
+flags, phi_obs = ext.init_flags(phi=True, boundary_width=0, wall="xXyY")
+
+curl = 0
+
+# large scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, 0.0]),
+    sig=torch.Tensor([[0.1, 0], [0, 0.1]]),
+    amp=torch.Tensor([-0.02]),
+)
+curl = curl + node_op.get_curl(phi)
+
+# middle scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.4, -0.4]),
+    sig=torch.Tensor([[0.04, 0], [0, 0.04]]),
+    amp=torch.Tensor([0.015]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.35, 0.1]),
+    sig=torch.Tensor([[0.01, 0], [0, 0.01]]),
+    amp=torch.Tensor([0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+# small scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.2, 0.2]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.006]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.25, 0.4]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([-0.006]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.5, 0.35]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.006]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.7, 0.2]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.006]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.7, 0.0]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([-0.006]),
+)
+curl = curl + node_op.get_curl(phi)
+
+utils.mkdir("forces/data/2d_128_multi_vortex/")
+mac_op.save_img(curl, "forces/data/2d_128_multi_vortex/lic.png", colormap="LIC")
+mac_op.save_img(curl, "forces/data/2d_128_multi_vortex/arrow.png", colormap="ARROW")
+np.savez("forces/data/2d_128_multi_vortex/000.npz", force=mac_op.grid2array(curl))

forces/generate_multi_vortex_256.py

+import os
+import sys
+import torch
+import numpy as np
+
+sys.path.append(os.getcwd())
+from lib.dfluids.solver import FluidSolver
+from lib.dfluids.grid import GridType
+import lib.dfluids.utils as utils
+
+# 2d big curl force
+res = 256
+solver = FluidSolver(grid_size=(res, res))
+ext = solver.create_extforces()
+mac_op = solver.create_grid_operator(GridType.MAC)
+real_op = solver.create_grid_operator(GridType.REAL)
+node_op = solver.create_grid_operator(GridType.NODE)
+flags, phi_obs = ext.init_flags(phi=True, boundary_width=0, wall="xXyY")
+
+curl = 0
+
+# large scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, 0.0]),
+    sig=torch.Tensor([[0.1, 0], [0, 0.1]]),
+    amp=torch.Tensor([-0.03]),
+)
+curl = curl + node_op.get_curl(phi)
+
+# middle scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.4, -0.4]),
+    sig=torch.Tensor([[0.04, 0], [0, 0.04]]),
+    amp=torch.Tensor([0.03]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.35, 0.1]),
+    sig=torch.Tensor([[0.01, 0], [0, 0.01]]),
+    amp=torch.Tensor([0.02]),
+)
+curl = curl + node_op.get_curl(phi)
+
+# small scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.2, 0.2]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.25, 0.4]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([-0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.5, 0.35]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.7, 0.2]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.7, 0.0]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([-0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+
+utils.mkdir("forces/data/2d_256_multi_vortex/")
+mac_op.save_img(curl, "forces/data/2d_256_multi_vortex/lic.png", colormap="LIC")
+mac_op.save_img(curl, "forces/data/2d_256_multi_vortex/arrow.png", colormap="ARROW")
+np.savez("forces/data/2d_256_multi_vortex/000.npz", force=mac_op.grid2array(curl))

forces/generate_multi_vortex_64.py

+import os
+import sys
+import torch
+import numpy as np
+
+sys.path.append(os.getcwd())
+from lib.dfluids.solver import FluidSolver
+from lib.dfluids.grid import GridType
+import lib.dfluids.utils as utils
+
+# 2d big curl force
+res = 64
+solver = FluidSolver(grid_size=(res, res))
+ext = solver.create_extforces()
+mac_op = solver.create_grid_operator(GridType.MAC)
+real_op = solver.create_grid_operator(GridType.REAL)
+node_op = solver.create_grid_operator(GridType.NODE)
+flags, phi_obs = ext.init_flags(phi=True, boundary_width=0, wall="xXyY")
+
+curl = 0
+
+# large scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, 0.0]),
+    sig=torch.Tensor([[0.1, 0], [0, 0.1]]),
+    amp=torch.Tensor([-0.015]),
+)
+curl = curl + node_op.get_curl(phi)
+
+# middle scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.4, -0.4]),
+    sig=torch.Tensor([[0.04, 0], [0, 0.04]]),
+    amp=torch.Tensor([0.01]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.35, 0.1]),
+    sig=torch.Tensor([[0.01, 0], [0, 0.01]]),
+    amp=torch.Tensor([0.005]),
+)
+curl = curl + node_op.get_curl(phi)
+
+# small scale
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.2, 0.2]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.002]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.25, 0.4]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([-0.002]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.5, 0.35]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.002]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.7, 0.2]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([0.002]),
+)
+curl = curl + node_op.get_curl(phi)
+
+phi = node_op.get_gaussian(
+    mu=torch.Tensor([-0.7, 0.0]),
+    sig=torch.Tensor([[0.001, 0], [0, 0.001]]),
+    amp=torch.Tensor([-0.002]),
+)
+curl = curl + node_op.get_curl(phi)
+
+
+utils.mkdir("forces/data/2d_64_multi_vortex/")
+mac_op.save_img(curl, "forces/data/2d_64_multi_vortex/lic.png", colormap="LIC")
+mac_op.save_img(curl, "forces/data/2d_64_multi_vortex/arrow.png", colormap="ARROW")
+np.savez("forces/data/2d_64_multi_vortex/000.npz", force=mac_op.grid2array(curl))

forces/generate_s_shape_128.py

+import os
+import sys
+import torch
+import numpy as np
+
+sys.path.append(os.getcwd())
+from lib.dfluids.solver import FluidSolver
+from lib.dfluids.grid import GridType
+import lib.dfluids.utils as utils
+
+# 2d big curl force
+res = 128
+solver = FluidSolver(grid_size=(res, res))
+ext = solver.create_extforces()
+mac_op = solver.create_grid_operator(GridType.MAC)
+real_op = solver.create_grid_operator(GridType.REAL)
+node_op = solver.create_grid_operator(GridType.NODE)
+flags, phi_obs = ext.init_flags(phi=True, boundary_width=0, wall="xXyY")
+
+phi1 = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, -0.5]),
+    sig=torch.Tensor([[40, 0], [0, 1]]),
+    amp=torch.Tensor([-0.08]),
+)
+
+phi2 = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, 0.5]),
+    sig=torch.Tensor([[40, 0], [0, 1]]),
+    amp=torch.Tensor([0.08]),
+)
+
+curl1 = node_op.get_curl(phi1)
+curl2 = node_op.get_curl(phi2)
+
+curl = curl1
+curl[:, int(res / 2) :, :] = curl2[:, int(res / 2) :, :]
+
+
+utils.mkdir("forces/data/2d_128_s_shape/")
+mac_op.save_img(curl, "forces/data/2d_128_s_shape/lic.png", colormap="LIC")
+mac_op.save_img(curl, "forces/data/2d_128_s_shape/arrow.png", colormap="ARROW")
+np.savez("forces/data/2d_128_s_shape/000.npz", force=mac_op.grid2array(curl))

forces/generate_s_shape_256.py

+import os
+import sys
+import torch
+import numpy as np
+
+sys.path.append(os.getcwd())
+from lib.dfluids.solver import FluidSolver
+from lib.dfluids.grid import GridType
+import lib.dfluids.utils as utils
+
+
+res = 256
+solver = FluidSolver(grid_size=(res, res))
+ext = solver.create_extforces()
+mac_op = solver.create_grid_operator(GridType.MAC)
+real_op = solver.create_grid_operator(GridType.REAL)
+node_op = solver.create_grid_operator(GridType.NODE)
+flags, phi_obs = ext.init_flags(phi=True, boundary_width=0, wall="xXyY")
+
+phi1 = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, -0.5]),
+    sig=torch.Tensor([[40, 0], [0, 1]]),
+    amp=torch.Tensor([-0.15]),
+)
+
+phi2 = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, 0.5]),
+    sig=torch.Tensor([[40, 0], [0, 1]]),
+    amp=torch.Tensor([0.15]),
+)
+
+curl1 = node_op.get_curl(phi1)
+curl2 = node_op.get_curl(phi2)
+
+curl = curl1
+curl[:, int(res / 2) :, :] = curl2[:, int(res / 2) :, :]
+
+
+utils.mkdir("forces/data/2d_256_s_shape/")
+mac_op.save_img(curl, "forces/data/2d_256_s_shape/lic.png", colormap="LIC")
+mac_op.save_img(curl, "forces/data/2d_256_s_shape/arrow.png", colormap="ARROW")
+np.savez("forces/data/2d_256_s_shape/000.npz", force=mac_op.grid2array(curl))

forces/generate_s_shape_64.py

+import os
+import sys
+import torch
+import numpy as np
+
+sys.path.append(os.getcwd())
+from lib.dfluids.solver import FluidSolver
+from lib.dfluids.grid import GridType
+import lib.dfluids.utils as utils
+
+# 2d big curl force
+res = 64
+solver = FluidSolver(grid_size=(res, res))
+ext = solver.create_extforces()
+mac_op = solver.create_grid_operator(GridType.MAC)
+real_op = solver.create_grid_operator(GridType.REAL)
+node_op = solver.create_grid_operator(GridType.NODE)
+flags, phi_obs = ext.init_flags(phi=True, boundary_width=0, wall="xXyY")
+
+phi1 = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, -0.5]),
+    sig=torch.Tensor([[40, 0], [0, 1]]),
+    amp=torch.Tensor([-0.04]),
+)
+
+phi2 = node_op.get_gaussian(
+    mu=torch.Tensor([0.0, 0.5]),
+    sig=torch.Tensor([[40, 0], [0, 1]]),
+    amp=torch.Tensor([0.04]),
+)
+
+curl1 = node_op.get_curl(phi1)
+curl2 = node_op.get_curl(phi2)
+
+curl = curl1
+curl[:, int(res / 2) :, :] = curl2[:, int(res / 2) :, :]
+
+utils.mkdir("forces/data/2d_64_s_shape/")
+mac_op.save_img(curl, "forces/data/2d_64_s_shape/lic.png", colormap="LIC")
+mac_op.save_img(curl, "forces/data/2d_64_s_shape/arrow.png", colormap="ARROW")
+np.savez("forces/data/2d_64_s_shape/000.npz", force=mac_op.grid2array(curl))

lib/PyTorch_LBFGS.py

+# LBFGS code from https://github.com/hjmshi/PyTorch-LBFGS
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from functools import reduce
+from copy import deepcopy
+from torch.optim import Optimizer
+
+#%% Helper Functions for L-BFGS
+
+def is_legal(v):
+    """
+    Checks that tensor is not NaN or Inf.
+
+    Inputs:
+        v (tensor): tensor to be checked
+
+    """
+    legal = not torch.isnan(v).any() and not torch.isinf(v)
+
+    return legal
+
+def polyinterp(points, x_min_bound=None, x_max_bound=None, plot=False):
+    """
+    Gives the minimizer and minimum of the interpolating polynomial over given points
+    based on function and derivative information. Defaults to bisection if no critical
+    points are valid.
+
+    Based on polyinterp.m Matlab function in minFunc by Mark Schmidt with some slight
+    modifications.
+
+    Implemented by: Hao-Jun Michael Shi and Dheevatsa Mudigere
+    Last edited 12/6/18.
+
+    Inputs:
+        points (nparray): two-dimensional array with each point of form [x f g]
+        x_min_bound (float): minimum value that brackets minimum (default: minimum of points)
+        x_max_bound (float): maximum value that brackets minimum (default: maximum of points)
+        plot (bool): plot interpolating polynomial
+
+    Outputs:
+        x_sol (float): minimizer of interpolating polynomial
+        F_min (float): minimum of interpolating polynomial
+
+    Note:
+      . Set f or g to np.nan if they are unknown
+
+    """
+    no_points = points.shape[0]
+    order = np.sum(1 - np.isnan(points[:,1:3]).astype('int')) - 1
+
+    x_min = np.min(points[:, 0])
+    x_max = np.max(points[:, 0])
+
+    # compute bounds of interpolation area
+    if(x_min_bound is None):
+        x_min_bound = x_min
+    if(x_max_bound is None):
+        x_max_bound = x_max
+
+    # explicit formula for quadratic interpolation
+    if no_points == 2 and order == 2 and plot is False:
+        # Solution to quadratic interpolation is given by:
+        # a = -(f1 - f2 - g1(x1 - x2))/(x1 - x2)^2
+        # x_min = x1 - g1/(2a)
+        # if x1 = 0, then is given by:
+        # x_min = - (g1*x2^2)/(2(f2 - f1 - g1*x2))
+
+        if(points[0, 0] == 0):
+            x_sol = -points[0, 2]*points[1, 0]**2/(2*(points[1, 1] - points[0, 1] - points[0, 2]*points[1, 0]))
+        else:
+            a = -(points[0, 1] - points[1, 1] - points[0, 2]*(points[0, 0] - points[1, 0]))/(points[0, 0] - points[1, 0])**2
+            x_sol = points[0, 0] - points[0, 2]/(2*a)
+
+        x_sol = np.minimum(np.maximum(x_min_bound, x_sol), x_max_bound)
+
+    # explicit formula for cubic interpolation
+    elif no_points == 2 and order == 3 and plot is False:
+        # Solution to cubic interpolation is given by:
+        # d1 = g1 + g2 - 3((f1 - f2)/(x1 - x2))
+        # d2 = sqrt(d1^2 - g1*g2)
+        # x_min = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2))
+        d1 = points[0, 2] + points[1, 2] - 3*((points[0, 1] - points[1, 1])/(points[0, 0] - points[1, 0]))
+        d2 = np.sqrt(d1**2 - points[0, 2]*points[1, 2])
+        if np.isreal(d2):
+            x_sol = points[1, 0] - (points[1, 0] - points[0, 0])*((points[1, 2] + d2 - d1)/(points[1, 2] - points[0, 2] + 2*d2))
+            x_sol = np.minimum(np.maximum(x_min_bound, x_sol), x_max_bound)
+        else:
+            x_sol = (x_max_bound + x_min_bound)/2
+
+    # solve linear system
+    else:
+        # define linear constraints
+        A = np.zeros((0, order+1))
+        b = np.zeros((0, 1))
+
+        # add linear constraints on function values
+        for i in range(no_points):
+            if not np.isnan(points[i, 1]):
+                constraint = np.zeros((1, order+1))
+                for j in range(order, -1, -1):
+                    constraint[0, order - j] = points[i, 0]**j
+                A = np.append(A, constraint, 0)
+                b = np.append(b, points[i, 1])
+
+        # add linear constraints on gradient values
+        for i in range(no_points):
+            if not np.isnan(points[i, 2]):
+                constraint = np.zeros((1, order+1))
+                for j in range(order):
+                    constraint[0, j] = (order-j)*points[i,0]**(order-j-1)
+                A = np.append(A, constraint, 0)
+                b = np.append(b, points[i, 2])
+
+        # check if system is solvable
+        if(A.shape[0] != A.shape[1] or np.linalg.matrix_rank(A) != A.shape[0]):
+            x_sol = (x_min_bound + x_max_bound)/2
+            f_min = np.Inf
+        else:
+            # solve linear system for interpolating polynomial
+            coeff = np.linalg.solve(A, b)
+
+            # compute critical points
+            dcoeff = np.zeros(order)
+            for i in range(len(coeff) - 1):
+                dcoeff[i] = coeff[i]*(order-i)
+
+            crit_pts = np.array([x_min_bound, x_max_b