To receive notifications about scheduled maintenance, please subscribe to the mailing-list gitlab-operations@sympa.ethz.ch. You can subscribe to the mailing-list at https://sympa.ethz.ch

Commit 14694a27 authored by Jingwei Tang's avatar Jingwei Tang
Browse files

Init commit

parents
__pycache__/
*.py[cod]
*$py.class
explore*
*data/
*backup/
*vdb/
build/
*.npy
*.csv
.vscode
.venv*
*.dll
\ No newline at end of file
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.
# 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.
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))
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))
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))
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))
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))
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))
This diff is collapsed.
import torch
import torch.nn.functional as F
from .grid import GridType
from . import utils
class Advection(object):
def __init__(self, solver):
self.solver = solver
self.device = self.solver.get_device()
self.mac_op = self.solver.create_grid_operator(GridType.MAC)
self.real_op = self.solver.create_grid_operator(GridType.REAL)
self.node_op = self.solver.create_grid_operator(GridType.NODE)
self.ext = self.solver.create_extforces()
def advect_semi_lagrange(
self, flags, phi_obs, vel, rho, dt=None, order=1, clamp_mode=2, rk_order=3
):
if dt is None:
dt = self.solver.get_dt()
if order == 1:
rho_adv = self._advect_semi_lagrange(vel, rho, dt, rk_order)
elif order == 2:
rho_fwd = self._advect_semi_lagrange(vel, rho, dt, rk_order)
rho_bwd = self._advect_semi_lagrange(vel, rho_fwd, -dt, rk_order)
rho_corr = rho_fwd + (rho - rho_bwd) * 0.5
rho_adv = self._maccormack_clamp(
vel, rho_corr, rho_fwd, rho, dt, clamp_mode, rk_order
)
else:
raise ValueError("Advection order must be 1 or 2.")
rho_adv = self.ext.set_wall_bcs(phi_obs, rho_adv)
return rho_adv
def _advect_semi_lagrange(self, vel, rho, dt, rk_order):
if rho.size() == self.solver.get_torch_size_real():
return self._advect_real_semi_lagrange(vel, rho, dt, rk_order)
elif rho.size() == self.solver.get_torch_size_mac():
return self._advect_mac_semi_lagrange(vel, rho, dt, rk_order)
# WARNING: does not distinguish 3D MAC and 3D NODE
elif rho.size() == self.solver.get_torch_size_node():
return self._advect_node_semi_lagrange(vel, rho, dt, rk_order)
else:
raise ValueError("rho has a wrong size: {}".format(rho.size()))
def _advect_real_semi_lagrange(self, vel, rho, dt, rk_order):
return self._advect_semi_lagrange_single_channel(
vel=vel, rho=rho, dt=dt, rk_order=rk_order, sample_pos="center"
)
def _advect_node_semi_lagrange(self, vel, rho, dt, rk_order):
return self._advect_semi_lagrange_single_channel(
vel=vel, rho=rho, dt=dt, rk_order=rk_order, sample_pos="node"
)
def _advect_mac_semi_lagrange_inaccurate(self, vel, rho, dt, rk_order):
rho_adv_centered = self._advect_semi_lagrange_single_channel(
vel=vel,
rho=self.mac_op.get_centered(rho),
dt=dt,
rk_order=rk_order,
sample_pos="center",
)
rho_adv = self.mac_op.get_mac(rho_adv_centered)
return rho_adv
def _advect_mac_semi_lagrange(self, vel, rho, dt, rk_order):
# advect in x direction
rho_adv_x = self._advect_semi_lagrange_single_channel(
vel=vel, rho=rho[0:1, ...], dt=dt, rk_order=rk_order, sample_pos="mac_x"
)
# advect in y direction
rho_adv_y = self._advect_semi_lagrange_single_channel(
vel=vel, rho=rho[1:2, ...], dt=dt, rk_order=rk_order, sample_pos="mac_y"
)
# concatenate back advected values into the same tensor
rho_adv = torch.cat((rho_adv_x, rho_adv_y), dim=0)
# advect in z direction in 3d case
if self.solver.is_3d():
rho_adv_z = self._advect_semi_lagrange_single_channel(
vel=vel, rho=rho[2:3, ...], dt=dt, rk_order=rk_order, sample_pos="mac_z"
)
rho_adv = torch.cat((rho_adv, rho_adv_z), dim=0)
return rho_adv
def _advect_semi_lagrange_single_channel(self, vel, rho, dt, rk_order, sample_pos):
# create mesh grid
mgrid = utils.create_mesh_grid(
rho.size(), device=self.device, dtype=self.solver.dtype
)
# use runge kutta or euler method to create the back trace field
back_trace = self._runge_kutta(vel, mgrid, dt, rk_order, sample_pos)
# use warped grid sample function from PyTorch to interpolate the grid
# and bring the interporlated values to the current position
rho_adv = utils.grid_sample(rho, back_trace)
return rho_adv
def _euler(self, vel, mgrid, dt, sample_pos):
if sample_pos == "node":
raise NotImplementedError(
"Euler step not implemented " "for node data advection."
)
get_pos_fcn = self._map_sample_pos(sample_pos, rk=False)
return mgrid - get_pos_fcn(vel) * dt
def _runge_kutta(self, vel, mgrid, dt, rk_order, sample_pos):
grid_sample_fcn = self._map_sample_pos(sample_pos, rk=True)
if rk_order == 1:
# return mgrid - grid_sample_fcn(vel, mgrid)*dt
return self._euler(vel, mgrid, dt, sample_pos)
elif rk_order == 2:
# find the mid point position using half of time step
# back_trace_mid = mgrid - grid_sample_fcn(vel, mgrid)*dt*0.5
back_trace_mid = self._euler(vel, mgrid, 0.5 * dt, sample_pos)
# find the mid point velocity
return mgrid - grid_sample_fcn(vel, back_trace_mid) * dt
elif rk_order == 3:
# k1 = grid_sample_fcn(vel, mgrid)
k1 = (mgrid - self._euler(vel, mgrid, dt, sample_pos)) / dt
k2 = grid_sample_fcn(vel, mgrid - 0.5 * dt * k1)