├── .gitignore ├── figures ├── inflow.gif ├── karman.gif ├── taylorgreen.gif └── vortexsheet.gif ├── utils ├── __init__.py ├── math_utils.py ├── args_utils.py └── vis_utils.py ├── requirements.txt ├── configs ├── inflow.json ├── karman.json ├── taylorgreen.json └── vortexsheet.json ├── examples ├── __init__.py ├── taylorgreen.py ├── vortexsheet.py ├── inflow.py └── karman.py ├── LICENSE ├── README.md ├── run.py └── fluid.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | outputs 3 | -------------------------------------------------------------------------------- /figures/inflow.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ChrisWu1997/StableFluids-python/HEAD/figures/inflow.gif -------------------------------------------------------------------------------- /figures/karman.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ChrisWu1997/StableFluids-python/HEAD/figures/karman.gif -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .math_utils import * 2 | from .vis_utils import * 3 | from .args_utils import * -------------------------------------------------------------------------------- /figures/taylorgreen.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ChrisWu1997/StableFluids-python/HEAD/figures/taylorgreen.gif -------------------------------------------------------------------------------- /figures/vortexsheet.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ChrisWu1997/StableFluids-python/HEAD/figures/vortexsheet.gif -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib==3.6.0 2 | numpy==1.23.3 3 | Pillow==9.2.0 4 | pip==22.2.2 5 | scipy==1.9.1 6 | tqdm==4.64.1 7 | -------------------------------------------------------------------------------- /configs/inflow.json: -------------------------------------------------------------------------------- 1 | { 2 | "tag": "inflow", 3 | "example": "inflow", 4 | "N": 256, 5 | "dt": 1, 6 | "T": 200, 7 | "diff": 0, 8 | "visc": 0, 9 | "draw": "mix", 10 | "fps": 30 11 | } -------------------------------------------------------------------------------- /configs/karman.json: -------------------------------------------------------------------------------- 1 | { 2 | "tag": "karman", 3 | "example": "karman", 4 | "N": 512, 5 | "dt": 0.1, 6 | "T": 600, 7 | "diff": 0, 8 | "visc": 0, 9 | "draw": "curl", 10 | "fps": 40 11 | } -------------------------------------------------------------------------------- /configs/taylorgreen.json: -------------------------------------------------------------------------------- 1 | { 2 | "tag": "taylorgreen", 3 | "example": "taylorgreen", 4 | "N": 32, 5 | "dt": 0.05, 6 | "T": 100, 7 | "diff": 0, 8 | "visc": 0, 9 | "draw": "velocity", 10 | "fps": 20 11 | } -------------------------------------------------------------------------------- /configs/vortexsheet.json: -------------------------------------------------------------------------------- 1 | { 2 | "tag": "vortexsheet", 3 | "example": "vortexsheet", 4 | "N": 512, 5 | "dt": 0.025, 6 | "T": 200, 7 | "diff": 0, 8 | "visc": 0, 9 | "draw": "density", 10 | "fps": 20 11 | } -------------------------------------------------------------------------------- /examples/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | 3 | 4 | def get_example_setup(name): 5 | try: 6 | module = importlib.import_module("examples." + name) 7 | setup = getattr(module, "setup") 8 | return setup 9 | except Exception as e: 10 | print(e) 11 | raise RuntimeError(f"Cannot find example setup. Example name: {name}.") 12 | -------------------------------------------------------------------------------- /examples/taylorgreen.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def taylorgreen_velocity(coords: np.ndarray): 5 | A = 1 6 | a = 1 7 | B = -1 8 | b = 1 9 | x = coords[..., 0] 10 | y = coords[..., 1] 11 | u = A * np.sin(a * x) * np.cos(b * y) 12 | v = B * np.cos(a * x) * np.sin(b * y) 13 | vel = np.stack([u, v], axis=-1) 14 | return vel 15 | 16 | 17 | setup = { 18 | "domain": [[0, 2 * np.pi], [0, 2 * np.pi]], 19 | "vsource": taylorgreen_velocity, 20 | "dsource": None, 21 | "src_duration": 1, 22 | "boundary_func": None 23 | } 24 | -------------------------------------------------------------------------------- /examples/vortexsheet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def vortexsheet_velocity(coords: np.ndarray, rigidR=0.5, rate=0.1): 5 | w = 1 * 1.0 / rate 6 | R = np.linalg.norm(coords, 2, axis=-1) 7 | mask = R < rigidR 8 | weight = 1 9 | u = w * coords[..., 1] * weight 10 | v = -w * coords[..., 0] * weight 11 | u[~mask] = 0 12 | v[~mask] = 0 13 | vel = np.stack([u, v], axis=-1) 14 | return vel 15 | 16 | 17 | def vortexsheet_density(coords: np.ndarray, rigidR=0.5): 18 | R = np.linalg.norm(coords, 2, axis=-1) 19 | den = np.zeros(coords.shape[:-1]) 20 | den[R < rigidR] = 1.0 21 | return den 22 | 23 | 24 | setup = { 25 | "domain": [[-1, 1], [-1, 1]], 26 | "vsource": vortexsheet_velocity, 27 | "dsource": vortexsheet_density, 28 | "src_duration": 1, 29 | "boundary_func": None 30 | } 31 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Rundi Wu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # StableFluids-python 2 | 3 | A pure-python implementation of Stable Fluids with stagger grids. 4 | 5 | ## Demo examples 6 | 7 | | Taylor–Green vortex | Vortexsheet | 8 | | :------------------------: |:-------------:| 9 | | ```python run.py -c configs/taylorgreen.json``` | ```python run.py -c configs/vortexsheet.json``` | 10 | | velocity | density | 11 | | | | 12 | | **Kármán vortex street** | **Inflow** | 13 | | ```python run.py -c configs/karman.json``` | ```python run.py -c configs/inflow.json``` | 14 | | curl(vorticity) | R: curl, G: 1, B: density | 15 | | | | 16 | 17 | To setup new examples, simply add a cooresponding python file under `examples` folder. 18 | 19 | ## Reference 20 | - [Joe Stam's Stable Fluid paper](https://pages.cs.wisc.edu/~chaol/data/cs777/stam-stable_fluids.pdf) 21 | - [Course material from Columbia Computer Animation course](http://www.cs.columbia.edu/~cxz/teaching/w4167_f21/#) 22 | - [GregTJ's git repo](https://github.com/GregTJ/stable-fluids) -------------------------------------------------------------------------------- /utils/math_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy import sparse 3 | 4 | 5 | def build_laplacian_matrix(M: int, N: int): 6 | """build laplacian matrix with neumann boundary condition, 7 | i.e., gradient at boundary equals zero. 8 | This eliminates one degree of freedom. 9 | 10 | Args: 11 | M (int): number of rows 12 | N (int): number of cols 13 | 14 | Returns: 15 | np.array: laplacian matrix 16 | """ 17 | main_diag = np.full(M * N, -4) 18 | main_diag[[0, N - 1, -N, -1]] = -2 19 | main_diag[[*range(1, N - 1), *range(-N + 1, -1), 20 | *range(N, (M - 2) * N + 1, N), *range(2 * N - 1, (M - 1) * N, N)]] = -3 21 | side_diag = np.ones(M * N - 1) 22 | side_diag[[*range(N - 1, M * N - 1, N)]] = 0 23 | data = [np.ones(M * N - N), side_diag, main_diag, side_diag, np.ones(M * N - N)] 24 | offsets = [-N, -1, 0, 1, N] 25 | mat = sparse.diags(data, offsets) 26 | return mat 27 | 28 | 29 | def compute_curl(velocity_field: np.ndarray, h: float): 30 | """compute curl(vorticity) of a 2D velocity field 31 | 32 | Args: 33 | velocity_field (np.ndarray): velocity field of shape (M, N, 2) 34 | h (float): grid size 35 | """ 36 | dvy_dx = np.gradient(velocity_field[..., 1], h)[0] 37 | dvx_dy = np.gradient(velocity_field[..., 0], h)[1] 38 | curl = dvy_dx - dvx_dy 39 | return curl 40 | -------------------------------------------------------------------------------- /examples/inflow.py: -------------------------------------------------------------------------------- 1 | # from https://github.com/GregTJ/stable-fluids/blob/main/example.py 2 | import numpy as np 3 | 4 | 5 | def create_circular_points(point_dist: float, point_count: int): 6 | points = np.linspace(-np.pi, np.pi, point_count, endpoint=False) 7 | points = tuple(np.array((np.cos(p), np.sin(p))) for p in points) 8 | normals = tuple(-p for p in points) 9 | points = tuple(point_dist * p for p in points) 10 | return points, normals 11 | 12 | 13 | def inflow_velocity(coords: np.ndarray, vel: float=0.01, point_radius: float=0.05, point_dist: float=0.8, point_count: int=4): 14 | points, normals = create_circular_points(point_dist, point_count) 15 | 16 | inflow_velocity = np.zeros_like(coords) 17 | for p, n in zip(points, normals): 18 | mask = np.linalg.norm(coords - p[None, None], 2, axis=-1) <= point_radius 19 | inflow_velocity[mask] += n[None] * vel 20 | return inflow_velocity 21 | 22 | 23 | def inflow_density(coords: np.ndarray, point_radius: float=0.05, point_dist: float=0.8, point_count: int=4): 24 | points, normals = create_circular_points(point_dist, point_count) 25 | 26 | inflow_density = np.zeros(coords.shape[:-1]) 27 | for p in points: 28 | mask = np.linalg.norm(coords - p[None, None], 2, axis=-1) <= point_radius 29 | inflow_density[mask] = 1 30 | return inflow_density 31 | 32 | 33 | setup = { 34 | "domain": [[-1, 1], [-1, 1]], 35 | "vsource": inflow_velocity, 36 | "dsource": inflow_density, 37 | "src_duration": 60, 38 | "boundary_func": None 39 | } 40 | -------------------------------------------------------------------------------- /run.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | from tqdm import tqdm 4 | from fluid import StableFluids 5 | from examples import get_example_setup 6 | from utils import parse_args, frames2gif 7 | 8 | # read config 9 | cfg = parse_args() 10 | output_dir = f"outputs/{cfg['tag']}" 11 | os.makedirs(output_dir, exist_ok=True) 12 | draw_dir = os.path.join(output_dir, cfg["draw"]) 13 | os.makedirs(draw_dir, exist_ok=True) 14 | path = os.path.join(output_dir, "saved_config.json") 15 | with open(path, "w") as fp: 16 | json.dump(cfg, fp, indent=2) 17 | 18 | # get example setup 19 | setup = get_example_setup(cfg["example"]) 20 | 21 | # init 22 | fluid = StableFluids(cfg["N"], cfg["dt"], setup["domain"], cfg["visc"], cfg["diff"], setup["boundary_func"]) 23 | 24 | # set initial condition 25 | fluid.add_source("velocity", setup["vsource"]) 26 | fluid.add_source("density", setup["dsource"]) 27 | 28 | # run simulation and draw frames 29 | fluid.draw(cfg["draw"], os.path.join(draw_dir, f"{0:04d}.png")) 30 | pbar = tqdm(range(1, cfg["T"] + 1), desc="Simulate and draw") 31 | for i in pbar: 32 | runtime = fluid.step() 33 | 34 | drawtime = fluid.draw(cfg["draw"], os.path.join(draw_dir, f"{i:04d}.png")) 35 | 36 | if i < setup["src_duration"]: 37 | fluid.add_source("velocity", setup["vsource"]) 38 | fluid.add_source("density", setup["dsource"]) 39 | 40 | pbar.set_postfix({"runtime": round(runtime, 6), "drawtime": round(drawtime, 6)}) 41 | 42 | # frames to gif animation 43 | save_path = os.path.join(os.path.dirname(draw_dir), f"anim_{cfg['draw']}.gif") 44 | frames2gif(draw_dir, save_path, cfg["fps"]) 45 | -------------------------------------------------------------------------------- /examples/karman.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | KARMAN_VEL = 0.5 5 | 6 | def karman_velocity(coords: np.ndarray, init_vel: float=KARMAN_VEL): 7 | vel = np.zeros_like(coords) 8 | vel[..., 1] = init_vel # constant horizontal velocity 9 | return vel 10 | 11 | 12 | def sphere_obstacle(coords: np.ndarray): 13 | center = np.array([np.pi / 2, np.pi / 4]) 14 | radius = np.pi / 15 15 | center = center.reshape(*[1]*len(coords.shape[:-1]), 2) 16 | sign_dist = np.linalg.norm(coords - center, 2, axis=-1) - radius 17 | return sign_dist 18 | 19 | 20 | def karman_boundary(d_grid, u_grid, v_grid, h): 21 | """d, u, v from stagger grid. h: grid size""" 22 | grid_indices = np.indices(d_grid.shape) 23 | 24 | def _transform_coords(coords, offset): 25 | """transform coords in grid space to original domain""" 26 | minxy = np.array([0, 0])[None, None] 27 | offset = np.array(offset)[None, None] 28 | coords = coords.transpose((1, 2, 0)).astype(float) + offset 29 | coords = coords * h + minxy 30 | return coords 31 | 32 | # solid sphere obstacle 33 | u_coords = _transform_coords(grid_indices[:, :, :-1], [-0.5, 0]) 34 | v_coords = _transform_coords(grid_indices[:, :-1], [0, -0.5]) 35 | mask_u = sphere_obstacle(u_coords) < 0 36 | mask_v = sphere_obstacle(v_coords) < 0 37 | u_grid[mask_u] = 0 38 | v_grid[mask_v] = 0 39 | 40 | # domain boundary: open right side 41 | u_grid[0, :] = 0 42 | u_grid[-1, :] = 0 43 | u_grid[:, 0] = KARMAN_VEL 44 | 45 | v_grid[0, :] = 0 46 | v_grid[-1, :] = 0 47 | v_grid[:, 0] = 0 48 | v_grid[:, -1] = 0 49 | 50 | d_grid[0, :] = 0 51 | d_grid[-1, :] = 0 52 | d_grid[:, 0] = 0 53 | d_grid[:, -1] = 0 54 | 55 | 56 | setup = { 57 | "domain": [[0, np.pi], [0, 2 * np.pi]], 58 | "vsource": karman_velocity, 59 | "dsource": None, 60 | "src_duration": 1, 61 | "boundary_func": karman_boundary 62 | } 63 | -------------------------------------------------------------------------------- /utils/args_utils.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | 4 | 5 | def parse_args() -> dict: 6 | """parse command-line arguments. 7 | 8 | Returns: 9 | config (dict): config dictionary with all parameters 10 | """ 11 | 12 | # load config json file (should contain all necessary parameters) 13 | conf_parser = argparse.ArgumentParser(description=__doc__, 14 | formatter_class=argparse.RawDescriptionHelpFormatter, add_help=False) 15 | conf_parser.add_argument("-c", "--conf_file", default="configs/taylorgreen.json", 16 | help="Specify config file", metavar="FILE") 17 | args, remaining_argv = conf_parser.parse_known_args() 18 | 19 | with open(args.conf_file, 'r') as f: 20 | loaded_params = json.load(f) 21 | config = loaded_params 22 | 23 | # additional command-line arguments to modify config 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--tag', type=str, default=None, help='name tag') 26 | parser.add_argument('-e', '--example', type=str, default=None, 27 | help='example setup (module name under examples folder)') 28 | parser.add_argument('--N', type=int, default=None, help='grid size') 29 | parser.add_argument('--dt', type=float, default=None, help='time step size') 30 | parser.add_argument('--T', type=int, default=None, help='total time steps') 31 | parser.add_argument('--diff', type=float, default=None, help='diffusion coefficent') 32 | parser.add_argument('--visc', type=float, default=None, help='viscosity coefficent') 33 | parser.add_argument('--draw', type=str, default=None, choices=['velocity', 'density', 'curl', 'mix']) 34 | parser.add_argument('--fps', type=int, default=None, help="fps for saved animation") 35 | parser.add_argument('--save_grids', type=int, default=None, help="save grids") 36 | args = parser.parse_args(remaining_argv) 37 | 38 | for k, v in args.__dict__.items(): 39 | if v is not None: 40 | config[k] = v 41 | 42 | # set command-line arguments as attributes 43 | print("----Configuration-----") 44 | for k, v in config.items(): 45 | print("{0:20}".format(k), v) 46 | 47 | return config 48 | -------------------------------------------------------------------------------- /utils/vis_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | import matplotlib.cm as cm 4 | from PIL import Image 5 | import os 6 | from scipy.special import erf 7 | 8 | 9 | def draw_velocity(arr: np.ndarray, save_path=None, figsize=(4, 4)): 10 | """draw 2D velocity field""" 11 | if arr.shape[0] > 32: 12 | s = arr.shape[0] // 32 + 1 13 | arr = arr[::s, ::s] 14 | fig, ax = plt.subplots(figsize=figsize) 15 | indices = np.indices(arr.shape[:-1]) 16 | ax.quiver(indices[0], indices[1], arr[..., 0], arr[..., 1], 17 | scale=arr.shape[0], scale_units='width') 18 | fig.tight_layout() 19 | if save_path is not None: 20 | save_figure(fig, save_path, axis_off=True) 21 | return fig 22 | 23 | 24 | def draw_curl(curl: np.ndarray, save_path=None): 25 | """draw 2D curl(vorticity) field""" 26 | curl = (erf(curl) + 1) / 2 # map range to 0~1 27 | img = cm.bwr(curl) 28 | img = Image.fromarray((img * 255).astype('uint8')) 29 | if save_path is not None: 30 | img.save(save_path) 31 | return img 32 | 33 | 34 | def draw_density(density: np.ndarray, save_path=None): 35 | """draw 2D density field""" 36 | density = erf(np.clip(density, 0, None) * 2) 37 | img = cm.cividis(density) 38 | img = Image.fromarray((img * 255).astype('uint8')) 39 | if save_path is not None: 40 | img.save(save_path) 41 | return img 42 | 43 | 44 | def draw_mix(curl: np.ndarray, density: np.ndarray, save_path=None): 45 | """R: curl, G: 1, B: density""" 46 | curl = (erf(curl) + 1) / 2 47 | img = np.dstack((curl, np.ones_like(curl), density)) 48 | img = (np.clip(img, 0, 1) * 255).astype('uint8') 49 | img = Image.fromarray(img, mode='HSV').convert('RGB') 50 | if save_path is not None: 51 | img.save(save_path) 52 | return img 53 | 54 | 55 | def save_figure(fig, save_path, close=True, axis_off=False): 56 | if axis_off: 57 | # plt.axis('off') 58 | plt.xticks([]) 59 | plt.yticks([]) 60 | plt.savefig(save_path, bbox_inches='tight') 61 | if close: 62 | plt.close(fig) 63 | 64 | 65 | def frames2gif(src_dir, save_path, fps=30): 66 | print("Convert frames to gif...") 67 | filenames = sorted([x for x in os.listdir(src_dir) if x.endswith('.png')]) 68 | img_list = [Image.open(os.path.join(src_dir, name)) for name in filenames] 69 | img = img_list[0] 70 | img.save(fp=save_path, append_images=img_list[1:], 71 | save_all=True, duration=1 / fps * 1000, loop=0) 72 | print("Done.") 73 | -------------------------------------------------------------------------------- /fluid.py: -------------------------------------------------------------------------------- 1 | import time 2 | import numpy as np 3 | from functools import partial 4 | from scipy.ndimage import map_coordinates 5 | from scipy.sparse.linalg import factorized, spsolve 6 | from utils import build_laplacian_matrix, draw_velocity, compute_curl, draw_density, draw_curl, draw_mix 7 | 8 | 9 | class StableFluids(object): 10 | def __init__(self, N: int, dt: float, domain: list=[[0, 1], [0, 1]], visc: float=0, diff: float=0, boundary_func=None): 11 | 12 | """Stable Fluids solver with stagger grid discretization. 13 | Density(dye) and pressure values are stored at the center of grids. 14 | Horizontal (u) and vertical (v) velocity values are stored at edges. 15 | The velocity along (i, j) indexing directions are (v, u). 16 | A layer of boundary is warpped outside. 17 | TODO: support different pressure solvers 18 | TODO: support external force 19 | TODO: support arbitrary rho. now assume rho=1 everywhere. 20 | 21 | ---v-----v--- 22 | | | | 23 | u d u d u 24 | | | | 25 | ---v-----v--- 26 | | | | 27 | u d u d u 28 | | | | 29 | ---v-----v--- 30 | 31 | Args: 32 | N (int): grid resolution along the longest dimension. 33 | dt (float): timestep size 34 | domain (list, optional): 2D domain ([[x_min, x_max], [y_min, y_max]]). 35 | Defaults to [[0, 1], [0, 1]]. 36 | visc (float, optional): viscosity coefficient. Defaults to 0. 37 | diff (float, optional): diffusion coefficient. Defaults to 0. 38 | boundary_func (function): function to set boundary condition, 39 | func(d_grid, u_grid, v_grid) -> None. Defaults to None, using solid boundary. 40 | """ 41 | len_x = (domain[0][1] - domain[0][0]) 42 | len_y = (domain[1][1] - domain[1][0]) 43 | self.h = max(len_x, len_y) / N 44 | self.M = int(len_x / self.h) 45 | self.N = int(len_y / self.h) 46 | 47 | self.dt = dt 48 | self.visc = visc 49 | self.diff = diff 50 | self.domain = domain 51 | self.timestep = 0 52 | 53 | self._d_grid = np.zeros((self.M + 2, self.N + 2)) 54 | self._u_grid = np.zeros((self.M + 2, self.N + 1)) 55 | self._v_grid = np.zeros((self.M + 1, self.N + 2)) 56 | 57 | # grid coordinates 58 | self._grid_indices = np.indices(self._d_grid.shape) 59 | 60 | # interpolation function 61 | self._interp_func = partial(map_coordinates, 62 | order=1, prefilter=False, mode='constant', cval=0) 63 | 64 | # boundary condition function 65 | if boundary_func is None: # assume solid boundary by default 66 | self.boundary_func = self._set_solid_boundary 67 | else: 68 | self.boundary_func = boundary_func 69 | 70 | # linear system solver 71 | print("Build pre-factorized linear system solver. Could take a while.") 72 | self.lap_mat = build_laplacian_matrix(self.M, self.N) 73 | self.pressure_solver = factorized(self.lap_mat) 74 | 75 | if self.diff > 0: 76 | self.diffD_solver = factorized(np.identity(self.M * self.N) - 77 | diff * dt / self.h / self.h * build_laplacian_matrix(self.M, self.N)) 78 | if self.visc > 0: 79 | self.diffU_solver = factorized(np.identity(self.M * (self.N + 1)) - 80 | visc * dt / self.h / self.h * build_laplacian_matrix(self.M, self.N + 1)) 81 | self.diffV_solver = factorized(np.identity((self.M + 1) * self.N) - 82 | visc * dt / self.h / self.h * build_laplacian_matrix(self.M + 1, self.N)) 83 | 84 | @property 85 | def grid_density(self): 86 | """density values at grid centers""" 87 | return self._d_grid[1:-1, 1:-1] 88 | 89 | @property 90 | def grid_velocity(self): 91 | """velocity values at grid centers""" 92 | u = (self._u_grid[1:-1, 1:] + self._u_grid[1:-1, :-1]) / 2 93 | v = (self._v_grid[1:, 1:-1] + self._v_grid[:-1, 1:-1]) / 2 94 | vel = np.stack([v, u], axis=-1) 95 | return vel 96 | 97 | @property 98 | def grid_curl(self): 99 | """curl(vorticity) values at grid centers""" 100 | curl = compute_curl(self.grid_velocity, self.h) 101 | return curl 102 | 103 | def _interpolateD(self, i_arr, j_arr): 104 | """interpolate d_grid with global indices""" 105 | i_arr = np.clip(i_arr, 1, self.M) 106 | j_arr = np.clip(j_arr, 1, self.N) 107 | return self._interp_func(self._d_grid, np.stack([i_arr, j_arr])) 108 | 109 | def _interpolateU(self, i_arr, j_arr): 110 | """interpolate u_grid with global indices""" 111 | i_arr = np.clip(i_arr, 1, self.M) 112 | j_arr = np.clip(j_arr - 0.5, 0, self.N) 113 | return self._interp_func(self._u_grid, np.stack([i_arr, j_arr])) 114 | 115 | def _interpolateV(self, i_arr, j_arr): 116 | """interpolate v_grid with global indices""" 117 | i_arr = np.clip(i_arr - 0.5, 0, self.M) 118 | j_arr = np.clip(j_arr, 1, self.N) 119 | return self._interp_func(self._v_grid, np.stack([i_arr, j_arr])) 120 | 121 | def _transform_coords(self, coords, offset): 122 | """transform coords in grid space to original domain""" 123 | minxy = np.array([self.domain[0][0], self.domain[1][0]])[None, None] 124 | offset = np.array(offset)[None, None] 125 | coords = coords.transpose((1, 2, 0)).astype(float) + offset 126 | coords = coords * self.h + minxy 127 | return coords 128 | 129 | def add_source(self, attr: str, source_func): 130 | """add source to density(d) or velocity field(u, v) 131 | 132 | Args: 133 | attr (str): "velocity" or "density" 134 | source_func (function): attr(x) = source_func(x) 135 | 136 | Raises: 137 | ValueError: _description_ 138 | """ 139 | if source_func is None: 140 | return 141 | 142 | if attr == "velocity": 143 | u_indices = self._transform_coords(self._grid_indices[:, :, :-1], [-0.5, 0]) 144 | self._u_grid += source_func(u_indices)[..., 1] 145 | 146 | v_indices = self._transform_coords(self._grid_indices[:, :-1], [0, -0.5]) 147 | self._v_grid += source_func(v_indices)[..., 0] 148 | elif attr == "density": 149 | d_indices = self._transform_coords(self._grid_indices, [-0.5, -0.5]) 150 | self._d_grid += source_func(d_indices) 151 | else: 152 | raise ValueError(f"attr must be velocity or density, but got {attr}.") 153 | self.boundary_func(self._d_grid, self._u_grid, self._v_grid, self.h) 154 | 155 | def step(self): 156 | """Integrates the system forward in time by dt.""" 157 | since = time.time() 158 | 159 | self._velocity_step() 160 | self._density_step() 161 | self.timestep += 1 162 | 163 | timecost = time.time() - since 164 | return timecost 165 | 166 | def _density_step(self): 167 | """update density field by one timestep""" 168 | # diffusion 169 | if self.diff > 0: 170 | self._diffuseD() 171 | 172 | # advection 173 | self._advectD() 174 | 175 | def _velocity_step(self): 176 | """update density field by one timestep""" 177 | # external force 178 | pass 179 | 180 | # advection 181 | self._advectVel() 182 | 183 | # diffusion 184 | if self.visc > 0: 185 | self._diffuseU() 186 | self._diffuseV() 187 | 188 | # projection 189 | self._project() 190 | 191 | def _diffuseD(self): 192 | """diffusion step for d ([1, M], [1, N]) using implicit method""" 193 | self._d_grid[1:-1, 1:-1] = self.diffD_solver(self._d_grid[1:-1, 1:-1].flatten()).reshape(self.M, self.N) 194 | 195 | def _diffuseU(self): 196 | """diffusion step for u ([1, M], [0, N]) using implicit method""" 197 | self._u_grid[1:-1, :] = self.diffU_solver(self._u_grid[1:-1].flatten()).reshape(self.M, self.N + 1) 198 | 199 | def _diffuseV(self): 200 | """diffusion step for v ([0, M], [1, N]) using implicit method""" 201 | self._v_grid[:, 1:-1] = self.diffV_solver(self._v_grid[:, 1:-1].flatten()).reshape(self.M + 1, self.N) 202 | 203 | def _advectD(self): 204 | """advect density for ([1, M], [1, N])""" 205 | ij_arr = self._grid_indices[:, 1:-1, 1:-1].astype(float) 206 | i_back = ij_arr[0] - self.dt / self.h * self._interpolateV(ij_arr[0], ij_arr[1]) 207 | j_back = ij_arr[1] - self.dt / self.h * self._interpolateU(ij_arr[0], ij_arr[1]) 208 | self._d_grid[1:-1, 1:-1] = self._interpolateD(i_back, j_back) 209 | 210 | def _advectVel(self): 211 | """addvect velocity field""" 212 | new_u_grid = self._advectU() 213 | new_v_grid = self._advectV() 214 | self._u_grid[1:-1, :] = new_u_grid 215 | self._v_grid[:, 1:-1] = new_v_grid 216 | 217 | def _advectU(self): 218 | """advect horizontal velocity (u) for ([1, M], [0, N])""" 219 | ij_arr = self._grid_indices[:, 1:-1, :-1].astype(float) 220 | ij_arr[1] += 0.5 221 | i_back = ij_arr[0] - self.dt / self.h * self._interpolateV(ij_arr[0], ij_arr[1]) 222 | j_back = ij_arr[1] - self.dt / self.h * self._interpolateU(ij_arr[0], ij_arr[1]) 223 | new_u_grid = self._interpolateU(i_back, j_back) 224 | return new_u_grid 225 | 226 | def _advectV(self): 227 | """advect vertical velocity (v) for ([0, M], [1, N])""" 228 | ij_arr = self._grid_indices[:, :-1, 1:-1].astype(float) 229 | ij_arr[0] += 0.5 230 | i_back = ij_arr[0] - self.dt / self.h * self._interpolateV(ij_arr[0], ij_arr[1]) 231 | j_back = ij_arr[1] - self.dt / self.h * self._interpolateU(ij_arr[0], ij_arr[1]) 232 | new_v_grid = self._interpolateV(i_back, j_back) 233 | return new_v_grid 234 | 235 | def _solve_pressure(self): 236 | """solve pressure field for laplacian(pressure) = divergence(velocity)""" 237 | # compute divergence of velocity field 238 | h = self.h 239 | div = (self._u_grid[1:-1, 1:] - self._u_grid[1:-1, :-1] + 240 | self._v_grid[1:, 1:-1] - self._v_grid[:-1, 1:-1]) / h 241 | 242 | # solve poisson equation as a linear system 243 | rhs = div.flatten() * h * h 244 | p_grid = self.pressure_solver(rhs).reshape(self.M, self.N) 245 | return p_grid 246 | 247 | def _project(self): 248 | """projection step to enforce divergence free (incompressible flow)""" 249 | # set solid boundary condition 250 | self.boundary_func(self._d_grid, self._u_grid, self._v_grid, self.h) 251 | 252 | p_grid = self._solve_pressure() 253 | 254 | # apply gradient of pressure to correct velocity field 255 | # ([1, M], [1, N-1]) for u, ([1, M-1], [1, N]) for v. FIXME: this range holds only for solid boundary. 256 | self._u_grid[1:-1, 1:-1] -= (p_grid[:, 1:] - p_grid[:, :-1]) / self.h 257 | self._v_grid[1:-1, 1:-1] -= (p_grid[1:, :] - p_grid[:-1, :]) / self.h 258 | 259 | self.boundary_func(self._d_grid, self._u_grid, self._v_grid, self.h) 260 | 261 | def _set_solid_boundary(self, d_grid, u_grid, v_grid, h): 262 | """set solid (zero Dirichlet) boundary condition""" 263 | for grid in [d_grid, u_grid, v_grid]: 264 | grid[0, :] = 0 265 | grid[-1, :] = 0 266 | grid[:, 0] = 0 267 | grid[:, -1] = 0 268 | 269 | def draw(self, attr: str, save_path: str): 270 | """draw a frame""" 271 | since = time.time() 272 | 273 | if attr == "velocity": 274 | draw_velocity(self.grid_velocity, save_path) 275 | elif attr == "curl": 276 | draw_curl(self.grid_curl, save_path) 277 | elif attr == "density": 278 | draw_density(self.grid_density, save_path) 279 | elif attr == "mix": 280 | draw_mix(self.grid_curl, self.grid_density, save_path) 281 | else: 282 | raise NotImplementedError 283 | 284 | timecost = time.time() - since 285 | return timecost 286 | --------------------------------------------------------------------------------