├── .gitignore ├── LICENSE ├── README.md ├── flip_extension.py ├── fluid_simulator.py ├── initializer_2d.py ├── level_set.py ├── main_2d.py ├── mgpcg.py ├── power_pic.py ├── pressure_project.py ├── requirements.txt ├── utils.py ├── visualizer_2d.py └── volume_control.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | pip-wheel-metadata/ 25 | share/python-wheels/ 26 | *.egg-info/ 27 | .installed.cfg 28 | *.egg 29 | MANIFEST 30 | 31 | # PyInstaller 32 | # Usually these files are written by a python script from a template 33 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 34 | *.manifest 35 | *.spec 36 | 37 | # Installer logs 38 | pip-log.txt 39 | pip-delete-this-directory.txt 40 | 41 | # Unit test / coverage reports 42 | htmlcov/ 43 | .tox/ 44 | .nox/ 45 | .coverage 46 | .coverage.* 47 | .cache 48 | nosetests.xml 49 | coverage.xml 50 | *.cover 51 | *.py,cover 52 | .hypothesis/ 53 | .pytest_cache/ 54 | 55 | # Translations 56 | *.mo 57 | *.pot 58 | 59 | # Django stuff: 60 | *.log 61 | local_settings.py 62 | db.sqlite3 63 | db.sqlite3-journal 64 | 65 | # Flask stuff: 66 | instance/ 67 | .webassets-cache 68 | 69 | # Scrapy stuff: 70 | .scrapy 71 | 72 | # Sphinx documentation 73 | docs/_build/ 74 | 75 | # PyBuilder 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | .python-version 87 | 88 | # pipenv 89 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 90 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 91 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 92 | # install all needed dependencies. 93 | #Pipfile.lock 94 | 95 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 96 | __pypackages__/ 97 | 98 | # Celery stuff 99 | celerybeat-schedule 100 | celerybeat.pid 101 | 102 | # SageMath parsed files 103 | *.sage.py 104 | 105 | # Environments 106 | .env 107 | .venv 108 | env/ 109 | venv/ 110 | ENV/ 111 | env.bak/ 112 | venv.bak/ 113 | 114 | # Spyder project settings 115 | .spyderproject 116 | .spyproject 117 | 118 | # Rope project settings 119 | .ropeproject 120 | 121 | # mkdocs documentation 122 | /site 123 | 124 | # mypy 125 | .mypy_cache/ 126 | .dmypy.json 127 | dmypy.json 128 | 129 | # Pyre type checker 130 | .pyre/ 131 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Yu Chang 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 | # Power-PIC 2 | Taichi Implementation of "The Power Particle-in-Cell Method" 3 | 4 | Running the Demo 5 | ---------------- 6 | 7 | ```bash 8 | python3 main_2d.py 9 | ``` 10 | 11 | parameters: 12 | 13 | ```bash 14 | usage: main_2d.py [-h] [--flip] [--save] [--res RES] [--t_res T_RES] [--show SHOW] 15 | 16 | optional arguments: 17 | -h, --help show this help message and exit 18 | --flip 19 | --save 20 | --res RES 21 | --t_res T_RES 22 | --show SHOW 23 | ``` 24 | 25 | References 26 | ---------------- 27 | 28 | [1] Ziyin Qu, Minchen Li, Fernando De Goes, and Chenfanfu Jiang. 2022. The power particle-in-cell method. ACM Trans. Graph. 41, 4, Article 118 (July 2022), 13 pages. https://doi.org/10.1145/3528223.3530066 -------------------------------------------------------------------------------- /flip_extension.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | from fluid_simulator import FluidSimulator 4 | from utils import * 5 | import utils 6 | 7 | import time 8 | import numpy as np 9 | 10 | @ti.data_oriented 11 | class FLIPSimulator(FluidSimulator): 12 | def __init__(self, 13 | dim = 2, 14 | res = (128, 128), 15 | dt = 1.25e-2, 16 | substeps = 1, 17 | dx = 1.0, 18 | rho = 1000.0, 19 | gravity = [0, -9.8], 20 | p0 = 1e-3, 21 | real = float, 22 | free_surface = True): 23 | super().__init__(dim, res, dt, substeps, dx, rho, gravity, p0, real) 24 | self.free_surface=free_surface 25 | self.total_fluid = ti.field(ti.i32, shape=()) 26 | self.p_v = ti.Vector.field(dim, dtype=real, shape=self.max_particles) # velocities 27 | self.color_p = ti.Vector.field(3, real, shape=self.max_particles) # random color for visualization 28 | 29 | @ti.kernel 30 | def p2g(self): 31 | for p in range(self.total_mk[None]): 32 | for k in ti.static(range(self.dim)): 33 | utils.splat(self.velocity[k], self.velocity_backup[k], self.p_v[p][k], self.p_x[p] / self.dx - 0.5 * (1 - ti.Vector.unit(self.dim, k))) 34 | 35 | for k in ti.static(range(self.dim)): 36 | for I in ti.grouped(self.velocity_backup[k]): # reuse velocity_backup as weight 37 | if self.velocity_backup[k][I] > 0: 38 | self.velocity[k][I] /= self.velocity_backup[k][I] 39 | 40 | @ti.func 41 | def vel_old_interp(self, pos): 42 | v = ti.Vector.zero(self.real, self.dim) 43 | for k in ti.static(range(self.dim)): 44 | v[k] = utils.sample(self.velocity_backup[k], pos / self.dx - 0.5 * (1 - ti.Vector.unit(self.dim, k))) 45 | return v 46 | 47 | @ti.kernel 48 | def g2p(self, dt : ti.f32): 49 | for p in range(self.total_mk[None]): 50 | old_v = self.p_v[p] 51 | self.p_v[p] = self.vel_interp(self.p_x[p]) + 0.99 * (old_v - self.vel_old_interp(self.p_x[p])) 52 | mispos = self.p_x[p] + self.vel_interp(self.p_x[p]) * (0.5 * dt) 53 | self.p_x[p] += self.vel_interp(mispos) * dt 54 | 55 | @ti.kernel 56 | def apply_markers_p(self): 57 | for I in ti.grouped(self.cell_type): 58 | if self.cell_type[I] != utils.SOLID: 59 | self.cell_type[I] = utils.AIR 60 | 61 | for p in range(self.total_mk[None]): 62 | I = int(self.p_x[p] / self.dx - 0.5) 63 | for offset in ti.grouped(ti.ndrange(*((-1, 2), ) * self.dim)): 64 | if self.cell_type[I+offset] != utils.SOLID and self.is_valid(I+offset): 65 | self.cell_type[I+offset] = utils.FLUID 66 | 67 | def substep(self, dt): 68 | if self.free_surface: 69 | # self.level_set.build_from_markers(self.p_x, self.total_mk[None]) 70 | self.apply_markers_p() 71 | 72 | for k in range(self.dim): 73 | self.velocity[k].fill(0) 74 | self.velocity_backup[k].fill(0) 75 | self.p2g() 76 | for k in range(self.dim): 77 | self.velocity_backup[k].copy_from(self.velocity[k]) 78 | 79 | self.extrap_velocity() 80 | self.enforce_boundary() 81 | 82 | self.add_gravity(dt) 83 | self.enforce_boundary() 84 | self.solve_pressure(dt, self.strategy) 85 | 86 | if self.verbose: 87 | prs = np.max(self.pressure.to_numpy()) 88 | print(f'\033[36mMax pressure: {prs}\033[0m') 89 | 90 | self.apply_pressure(dt) 91 | self.g2p(dt) 92 | # self.advect_markers(dt) 93 | self.total_t += self.dt 94 | 95 | @ti.kernel 96 | def init_markers(self): 97 | self.total_mk[None] = 0 98 | for I in ti.grouped(self.cell_type): 99 | if self.cell_type[I] == utils.FLUID: 100 | self.total_fluid[None] += 1 101 | for offset in ti.grouped(ti.ndrange(*((0, self.p_per_axis), ) * self.dim)): 102 | num = ti.atomic_add(self.total_mk[None], 1) 103 | self.p_x[num] = (I + (offset + [ti.random() for _ in ti.static(range(self.dim))]) / self.p_per_axis) * self.dx 104 | self.p_v[num] = self.vel_interp(self.p_x[num]) 105 | self.color_p[num] = ti.Vector([ti.random(), ti.random(), ti.random()]) 106 | 107 | -------------------------------------------------------------------------------- /fluid_simulator.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import utils 4 | from utils import * 5 | from mgpcg import MGPCGPoissonSolver 6 | from pressure_project import PressureProjectStrategy 7 | from level_set import FastSweepingLevelSet 8 | from volume_control import PressureProjectWithVolumeControlStrategy 9 | 10 | from functools import reduce 11 | import time 12 | import numpy as np 13 | 14 | ti.init(arch=ti.cuda, kernel_profiler=False, device_memory_GB=4.0) 15 | 16 | ADVECT_REDISTANCE = 0 17 | MARKERS = 1 18 | 19 | FAST_SWEEPING_METHOD = 0 20 | FAST_MARCHING_METHOD = 1 21 | 22 | @ti.data_oriented 23 | class FluidSimulator: 24 | def __init__(self, 25 | dim = 2, 26 | res = (128, 128), 27 | dt = 1.25e-2, 28 | substeps = 1, 29 | dx = 1.0, 30 | rho = 1000.0, 31 | gravity = [0, -9.8], 32 | p0 = 1e-3, 33 | real = float): 34 | 35 | # ADVECT_REDISTANCE: advect the level-set with Semi-Lagrangian, then redistance it (Standard) 36 | # MARKERS: advect markers with Semi-Lagrangian, then build the level-set from markers 37 | self.solver_type = ADVECT_REDISTANCE 38 | 39 | self.dim = dim 40 | self.real = real 41 | self.res = res 42 | self.dx = dx 43 | self.dt = dt 44 | 45 | self.total_t = 0.0 # total simulation time 46 | 47 | self.p0 = p0 # the standard atmospheric pressure 48 | self.rho = rho # density 49 | self.gravity = gravity # body force 50 | self.substeps = substeps 51 | 52 | # cell_type 53 | self.cell_type = ti.field(dtype=ti.i32) 54 | 55 | self.velocity = [ti.field(dtype=real) for _ in range(self.dim)] # MAC grid 56 | self.velocity_backup = [ti.field(dtype=real) for _ in range(self.dim)] # backup / use as weight in apic update 57 | self.pressure = ti.field(dtype=real) 58 | 59 | # extrap utils 60 | self.valid = ti.field(dtype=ti.i32) 61 | self.valid_temp = ti.field(dtype=ti.i32) 62 | 63 | # marker/apic particles 64 | self.total_mk = ti.field(dtype=ti.i32, shape = ()) # total number of particles/markers 65 | self.p_x = ti.Vector.field(dim, dtype=real) # positions 66 | 67 | self.indices = ti.ijk if self.dim == 3 else ti.ij 68 | self.p_per_axis = 2 69 | self.ppc = self.p_per_axis ** dim 70 | self.max_particles = reduce(lambda x, y : x * y, res) * (4 ** dim) 71 | ti.root.dense(ti.i, self.max_particles).place(self.p_x) 72 | 73 | ti.root.dense(self.indices, res).place(self.cell_type, self.pressure) 74 | ti.root.dense(self.indices, [res[_] + 1 for _ in range(self.dim)]).place(self.valid, self.valid_temp) 75 | for d in range(self.dim): 76 | ti.root.dense(self.indices, [res[_] + (d == _) for _ in range(self.dim)]).place(self.velocity[d], self.velocity_backup[d]) 77 | 78 | # Level-Set 79 | self.level_set = FastSweepingLevelSet(self.dim, 80 | self.res, 81 | self.dx, 82 | self.real) 83 | 84 | # MGPCG 85 | self.n_mg_levels = 4 86 | self.pre_and_post_smoothing = 2 87 | self.bottom_smoothing = 10 88 | self.iterations = 50 89 | self.verbose = False 90 | self.poisson_solver = MGPCGPoissonSolver(self.dim, 91 | self.res, 92 | self.n_mg_levels, 93 | self.pre_and_post_smoothing, 94 | self.bottom_smoothing, 95 | self.real) 96 | 97 | # Pressure Solve 98 | self.ghost_fluid_method = False # Gibou et al. [GFCK02] 99 | self.volume_control = False 100 | if self.volume_control: 101 | self.strategy = PressureProjectWithVolumeControlStrategy(self.dim, 102 | self.velocity, 103 | self.ghost_fluid_method, 104 | self.level_set.phi, 105 | self.p0, 106 | self.level_set, 107 | self.dt) # [Losasso et al. 2008] 108 | else: 109 | self.strategy = PressureProjectStrategy(self.dim, 110 | self.velocity, 111 | self.ghost_fluid_method, 112 | self.level_set.phi, 113 | self.p0) 114 | 115 | @ti.func 116 | def is_valid(self, I): 117 | return all(I >= 0) and all(I < self.res) 118 | 119 | @ti.func 120 | def is_fluid(self, I): 121 | return self.is_valid(I) and self.cell_type[I] == utils.FLUID 122 | 123 | @ti.func 124 | def is_solid(self, I): 125 | return not self.is_valid(I) or self.cell_type[I] == utils.SOLID 126 | 127 | @ti.func 128 | def is_air(self, I): 129 | return self.is_valid(I) and self.cell_type[I] == utils.AIR 130 | 131 | @ti.func 132 | def vel_interp(self, pos): 133 | v = ti.Vector.zero(self.real, self.dim) 134 | for k in ti.static(range(self.dim)): 135 | v[k] = utils.sample(self.velocity[k], pos / self.dx - 0.5 * (1 - ti.Vector.unit(self.dim, k))) 136 | return v 137 | 138 | @ti.kernel 139 | def advect_markers(self, dt : ti.f32): 140 | for p in range(self.total_mk[None]): 141 | midpos = self.p_x[p] + self.vel_interp(self.p_x[p]) * (0.5 * dt) 142 | self.p_x[p] += self.vel_interp(midpos) * dt 143 | 144 | @ti.kernel 145 | def apply_markers(self): 146 | for I in ti.grouped(self.cell_type): 147 | if self.cell_type[I] != utils.SOLID: 148 | self.cell_type[I] = utils.AIR 149 | 150 | for I in ti.grouped(self.cell_type): 151 | if self.cell_type[I] != utils.SOLID and self.level_set.phi[I] <= 0: 152 | self.cell_type[I] = utils.FLUID 153 | 154 | @ti.kernel 155 | def add_gravity(self, dt : ti.f32): 156 | for k in ti.static(range(self.dim)): 157 | if ti.static(self.gravity[k] != 0): 158 | g = self.gravity[k] 159 | for I in ti.grouped(self.velocity[k]): 160 | self.velocity[k][I] += g * dt 161 | 162 | @ti.kernel 163 | def enforce_boundary(self): 164 | for I in ti.grouped(self.cell_type): 165 | if self.cell_type[I] == utils.SOLID: 166 | for k in ti.static(range(self.dim)): 167 | self.velocity[k][I] = 0 168 | self.velocity[k][I + ti.Vector.unit(self.dim, k)] = 0 169 | 170 | def solve_pressure(self, dt, strategy): 171 | strategy.scale_A = dt / (self.rho * self.dx * self.dx) 172 | strategy.scale_b = 1 / self.dx 173 | 174 | start1 = time.perf_counter() 175 | self.poisson_solver.reinitialize(self.cell_type, strategy) 176 | end1 = time.perf_counter() 177 | 178 | start2 = time.perf_counter() 179 | self.poisson_solver.solve(self.iterations, self.verbose) 180 | end2 = time.perf_counter() 181 | 182 | print(f'\033[33minit cost {end1 - start1}s, solve cost {end2 - start2}s\033[0m') 183 | self.pressure.copy_from(self.poisson_solver.x) 184 | 185 | @ti.kernel 186 | def apply_pressure(self, dt : ti.f32): 187 | scale = dt / (self.rho * self.dx) 188 | 189 | for k in ti.static(range(self.dim)): 190 | for I in ti.grouped(self.cell_type): 191 | I_1 = I - ti.Vector.unit(self.dim, k) 192 | if self.is_fluid(I_1) or self.is_fluid(I): 193 | if self.is_solid(I_1) or self.is_solid(I): self.velocity[k][I] = 0 194 | # FLuid-Air 195 | elif self.is_air(I): 196 | if ti.static(self.ghost_fluid_method): 197 | c = (self.level_set.phi[I_1] - self.level_set.phi[I]) / self.level_set.phi[I_1] 198 | self.velocity[k][I] -= scale * (self.p0 - self.pressure[I_1]) * min(c, 1e3) # # limit the coefficient 199 | else: self.velocity[k][I] -= scale * (self.p0 - self.pressure[I_1]) 200 | # Air-Fluid 201 | elif self.is_air(I_1): 202 | if ti.static(self.ghost_fluid_method): 203 | c = (self.level_set.phi[I] - self.level_set.phi[I_1]) / self.level_set.phi[I] 204 | self.velocity[k][I] -= scale * (self.pressure[I] - self.p0) * min(c, 1e3) 205 | else: self.velocity[k][I] -= scale * (self.pressure[I] - self.p0) 206 | # Fluid-Fluid 207 | else: self.velocity[k][I] -= scale * (self.pressure[I] - self.pressure[I_1]) 208 | 209 | @ti.func 210 | def advect(self, I, dst, src, offset, dt): 211 | pos = (I + offset) * self.dx 212 | midpos = pos - self.vel_interp(pos) * (0.5 * dt) 213 | p0 = pos - self.vel_interp(midpos) * dt 214 | dst[I] = utils.sample(src, p0 / self.dx - offset) 215 | 216 | @ti.kernel 217 | def advect_quantity(self, dt : ti.f32): 218 | if ti.static(self.solver_type == ADVECT_REDISTANCE): 219 | for I in ti.grouped(self.level_set.phi): 220 | self.advect(I, self.level_set.phi_temp, self.level_set.phi, 0.5, dt) 221 | 222 | for k in ti.static(range(self.dim)): 223 | offset = 0.5 * (1 - ti.Vector.unit(self.dim, k)) 224 | for I in ti.grouped(self.velocity_backup[k]): 225 | self.advect(I, self.velocity_backup[k], self.velocity[k], offset, dt) 226 | 227 | def update_quantity(self): 228 | if ti.static(self.solver_type == ADVECT_REDISTANCE): 229 | self.level_set.phi.copy_from(self.level_set.phi_temp) 230 | for k in range(self.dim): 231 | self.velocity[k].copy_from(self.velocity_backup[k]) 232 | 233 | @ti.kernel 234 | def mark_valid(self, k : ti.template()): 235 | for I in ti.grouped(self.velocity[k]): 236 | # NOTE that the the air-liquid interface is valid 237 | I_1 = I - ti.Vector.unit(self.dim, k) 238 | if self.is_fluid(I_1) or self.is_fluid(I): 239 | self.valid[I] = 1 240 | else: 241 | self.valid[I] = 0 242 | 243 | @ti.kernel 244 | def diffuse_quantity(self, dst : ti.template(), src : ti.template(), valid_dst : ti.template(), valid : ti.template()): 245 | for I in ti.grouped(dst): 246 | if valid[I] == 0: 247 | tot = ti.cast(0, self.real) 248 | cnt = 0 249 | for offset in ti.static(ti.grouped(ti.ndrange(*((-1, 2), ) * self.dim))): 250 | if valid[I + offset] == 1: 251 | tot += src[I + offset] 252 | cnt += 1 253 | if cnt > 0: 254 | dst[I] = tot / ti.cast(cnt, self.real) 255 | valid_dst[I] = 1 256 | 257 | def extrap_velocity(self): 258 | for k in range(self.dim): 259 | self.mark_valid(k) 260 | for i in range(10): 261 | self.velocity_backup[k].copy_from(self.velocity[k]) 262 | self.valid_temp.copy_from(self.valid) 263 | self.diffuse_quantity(self.velocity[k], self.velocity_backup[k], self.valid, self.valid_temp) 264 | 265 | def begin_substep(self, dt): 266 | self.advect_markers(dt) 267 | self.advect_quantity(dt) 268 | self.update_quantity() 269 | 270 | if self.solver_type == MARKERS: 271 | self.level_set.build_from_markers(self.p_x, self.total_mk) 272 | else: 273 | self.level_set.redistance() 274 | 275 | self.apply_markers() 276 | self.enforce_boundary() 277 | 278 | if self.verbose: 279 | mks = max(np.max(self.velocity[0].to_numpy()), np.max(self.velocity[1].to_numpy())) 280 | print(f'\033[36mMax advect velocity: {mks}\033[0m') 281 | 282 | def end_substep(self, dt): 283 | self.extrap_velocity() 284 | self.enforce_boundary() 285 | 286 | self.total_t += self.dt 287 | 288 | def substep(self, dt): 289 | self.begin_substep(dt) 290 | 291 | self.add_gravity(dt) 292 | self.enforce_boundary() 293 | 294 | self.extrap_velocity() 295 | self.enforce_boundary() 296 | 297 | self.solve_pressure(dt, self.strategy) 298 | if self.verbose: 299 | prs = np.max(self.pressure.to_numpy()) 300 | print(f'\033[36mMax pressure: {prs}\033[0m') 301 | self.apply_pressure(dt) 302 | self.extrap_velocity() 303 | self.enforce_boundary() 304 | 305 | self.end_substep(dt) 306 | 307 | def run(self, max_steps, visualizer, verbose = True): 308 | self.verbose = verbose 309 | step = 0 310 | 311 | while step < max_steps or max_steps == -1: 312 | print(f'Current progress: ({step} / {max_steps})') 313 | for substep in range(self.substeps): 314 | self.substep(self.dt) 315 | visualizer.visualize(self) 316 | step += 1 317 | 318 | visualizer.end() 319 | 320 | @ti.kernel 321 | def init_boundary(self): 322 | for I in ti.grouped(self.cell_type): 323 | if any(I == 0) or any(I + 1 == self.res): 324 | self.cell_type[I] = utils.SOLID 325 | 326 | @ti.kernel 327 | def init_markers(self): 328 | self.total_mk[None] = 0 329 | for I in ti.grouped(self.cell_type): 330 | if self.cell_type[I] == utils.FLUID: 331 | for offset in ti.static(ti.grouped(ti.ndrange(*((0, self.p_per_axis), ) * self.dim))): 332 | num = ti.atomic_add(self.total_mk[None], 1) 333 | self.p_x[num] = (I + (offset + [ti.random() for _ in ti.static(range(self.dim))]) / self.p_per_axis) * self.dx 334 | 335 | @ti.kernel 336 | def reinitialize(self): 337 | for I in ti.grouped(ti.ndrange(* [self.res[_] for _ in range(self.dim)])): 338 | self.cell_type[I] = 0 339 | self.pressure[I] = 0 340 | for k in ti.static(range(self.dim)): 341 | I_1 = I + ti.Vector.unit(self.dim, k) 342 | self.velocity[k][I] = 0 343 | self.velocity[k][I_1] = 0 344 | self.velocity_backup[k][I] = 0 345 | self.velocity_backup[k][I_1] = 0 346 | 347 | def initialize(self, initializer): 348 | self.reinitialize() 349 | 350 | self.cell_type.fill(utils.AIR) 351 | initializer.init_scene(self) 352 | 353 | self.init_boundary() 354 | self.init_markers() 355 | -------------------------------------------------------------------------------- /initializer_2d.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import utils 3 | from fluid_simulator import * 4 | 5 | @ti.data_oriented 6 | class Initializer2D: 7 | def __init__(self, res, x0, y0, x1, y1, blocks): 8 | self.res = res 9 | self.x0 = int(res * x0) 10 | self.y0 = int(res * y0) 11 | self.x1 = int(res * x1) 12 | self.y1 = int(res * y1) 13 | self.blocks = blocks 14 | 15 | @ti.kernel 16 | def init_kernel(self, cell_type : ti.template()): 17 | for i, j in cell_type: 18 | in_block = False 19 | for k in ti.static(range(len(self.blocks))): 20 | if i >= self.blocks[k][0] * self.res and i <= self.blocks[k][2] * self.res and \ 21 | j >= self.blocks[k][1] * self.res and j <= self.blocks[k][3] * self.res: 22 | in_block = True 23 | if in_block: 24 | cell_type[i, j] = utils.SOLID 25 | elif i >= self.x0 and i <= self.x1 and j >= self.y0 and j <= self.y1: 26 | cell_type[i, j] = utils.FLUID 27 | 28 | def init_scene(self, simulator): 29 | simulator.level_set.initialize_with_aabb((self.x0 * simulator.dx, self.y0 * simulator.dx), (self.x1 * simulator.dx, self.y1 * simulator.dx)) 30 | self.init_kernel(simulator.cell_type) 31 | 32 | @ti.data_oriented 33 | class SphereInitializer2D: 34 | def __init__(self, res, x0, y0, r, free_surface=True): 35 | self.res = res 36 | self.x0 = int(res * x0) 37 | self.y0 = int(res * y0) 38 | self.r = int(res * r) 39 | self.free_surface = free_surface 40 | 41 | @ti.kernel 42 | def init_kernel(self, cell_type : ti.template()): 43 | for i, j in cell_type: 44 | if (i - self.x0) ** 2 + (j - self.y0) ** 2 <= self.r ** 2: 45 | cell_type[i, j] = utils.FLUID 46 | 47 | @ti.kernel 48 | def init_kernel1(self, cell_type : ti.template(), u0 : ti.template()): 49 | for i, j in cell_type: 50 | cell_type[i, j] = utils.FLUID 51 | if (i - self.x0) ** 2 + (j - self.y0) ** 2 <= self.r ** 2: 52 | u0[i, j] = -10.0 53 | u0[i, j+1] = -10.0 54 | 55 | def init_scene(self, simulator): 56 | simulator.level_set.initialize_with_sphere((self.x0 * simulator.dx, self.y0 * simulator.dx), self.r * simulator.dx) 57 | if self.free_surface: 58 | self.init_kernel(simulator.cell_type) 59 | else: 60 | self.init_kernel1(simulator.cell_type, simulator.velocity[1]) 61 | -------------------------------------------------------------------------------- /level_set.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import utils 4 | from utils import * 5 | 6 | from functools import reduce 7 | 8 | @ti.data_oriented 9 | class LevelSet: 10 | def __init__(self, dim, res, dx, real): 11 | self.dim = dim 12 | self.res = res 13 | self.dx = dx 14 | self.real = real 15 | 16 | self.valid = ti.field(dtype=ti.i32, shape=res) # indices to the closest points / reuse as visit sign 17 | self.phi = ti.field(dtype=real, shape=res) 18 | self.phi_temp = ti.field(dtype=real, shape=res) 19 | 20 | self.eps = 1 * self.dx # the "thickness" of the interface, O(∆x) and is smaller than the local feature size 21 | 22 | @ti.func 23 | def theta(self, phi): # smoothed step Heaviside function 24 | theta = ti.cast(0, self.real) 25 | if phi <= -self.eps: theta = 0 26 | elif phi >= self.eps: theta = 1 27 | else: theta = 1/2 + phi/(2*self.eps) + 1/(2*ti.math.pi) * ti.sin(ti.math.pi*phi/self.eps) 28 | return theta 29 | 30 | @ti.func 31 | def delta(self, phi): # smoothed regular Dirac delta function 32 | delta = ti.cast(0, self.real) 33 | if phi <= -self.eps or phi >= self.eps: delta = 0 34 | else: delta = (1 + ti.cos(ti.math.pi*phi/self.eps)) / (2*self.eps) 35 | return delta 36 | 37 | 38 | @ti.func 39 | def distance_of_aabb(self, x, x0, x1): 40 | phi = ti.cast(0, self.real) 41 | if all(x > x0) and all(x < x1): # (inside) 42 | phi = (ti.max(x0 - x, x - x1)).max() 43 | else: # (outside) 44 | # Find the closest point (p,q,r) on the surface of the box 45 | p = ti.zero(x) 46 | for k in ti.static(range(self.dim)): 47 | if x[k] < x0[k]: p[k] = x0[k] 48 | elif x[k] > x1[k]: p[k] = x1[k] 49 | else: p[k] = x[k] 50 | phi = (x - p).norm() 51 | 52 | return phi 53 | 54 | @ti.kernel 55 | def initialize_with_aabb(self, x0 : ti.template(), x1 : ti.template()): 56 | for I in ti.grouped(self.phi): 57 | self.phi[I] = self.distance_of_aabb((I + 0.5) * self.dx, x0, x1) 58 | 59 | @ti.kernel 60 | def initialize_with_sphere(self, x0 : ti.template(), r : ti.template()): 61 | for I in ti.grouped(self.phi): 62 | self.phi[I] = ((I + 0.5) * self.dx - x0).norm() - r 63 | 64 | @ti.kernel 65 | def target_surface(self): 66 | for I in ti.grouped(self.phi): 67 | sign_change = False 68 | est = ti.cast(1e20, self.real) 69 | for k in ti.static(range(self.dim)): 70 | for s in ti.static((-1, 1)): 71 | offset = ti.Vector.unit(self.dim, k) * s 72 | I1 = I + offset 73 | if I1[k] >= 0 and I1[k] < self.res[k] and \ 74 | ti.math.sign(self.phi[I]) != ti.math.sign(self.phi[I1]): 75 | theta = self.phi[I] / (self.phi[I] - self.phi[I1]) 76 | est0 = ti.math.sign(self.phi[I]) * theta * self.dx 77 | est = est0 if ti.abs(est0) < ti.abs(est) else est 78 | sign_change = True 79 | 80 | if sign_change: 81 | self.phi_temp[I] = est 82 | self.valid[I] = 0 83 | else: 84 | self.phi_temp[I] = ti.cast(1e20, self.real) # an upper bound for all possible distances 85 | 86 | @ti.func 87 | def update_from_neighbor(self, I): 88 | # solve the Eikonal equation 89 | nb = ti.Vector.zero(self.real, self.dim) 90 | for k in ti.static(range(self.dim)): 91 | o = ti.Vector.unit(self.dim, k) 92 | if I[k] == 0 or (I[k] < self.res[k] - 1 and ti.abs(self.phi_temp[I + o]) < ti.abs(self.phi_temp[I - o])): nb[k] = ti.abs(self.phi_temp[I + o]) 93 | else: nb[k] = ti.abs(self.phi_temp[I - o]) 94 | 95 | # sort 96 | for i in ti.static(range(self.dim-1)): 97 | for j in ti.static(range(self.dim-1-i)): 98 | if nb[j] > nb[j + 1]: nb[j], nb[j + 1] = nb[j + 1], nb[j] 99 | 100 | # (Try just the closest neighbor) 101 | d = nb[0] + self.dx 102 | if d > nb[1]: 103 | # (Try the two closest neighbors) 104 | d = (1/2) * (nb[0] + nb[1] + ti.sqrt(2 * (self.dx ** 2) - (nb[1] - nb[0]) ** 2)) 105 | if ti.static(self.dim == 3): 106 | if d > nb[2]: 107 | # (Use all three neighbors) 108 | d = (1/3) * (nb[0] + nb[1] + nb[2] + ti.sqrt(ti.max(0, (nb[0] + nb[1] + nb[2]) ** 2 - 3 * (nb[0] ** 2 + nb[1] ** 2 + nb[2] ** 2 - self.dx ** 2)))) 109 | 110 | return d 111 | 112 | @ti.kernel 113 | def distance_to_markers(self, markers : ti.template(), total_mk : ti.template()): 114 | # (Initialize the arrays near the input geometry) 115 | for p in range(total_mk): 116 | I = (markers[p] / self.dx).cast(int) 117 | d = (markers[p] - (I + 0.5) * self.dx).norm() 118 | if all(I >= 0 and I < self.res) and d < self.phi[I]: 119 | self.phi[I] = d 120 | self.valid[I] = p 121 | 122 | @ti.kernel 123 | def target_minus(self): 124 | for I in ti.grouped(self.phi): 125 | self.phi[I] -= (0.99 * self.dx) # the particle radius r (typically just a little less than the grid cell size dx) 126 | 127 | for I in ti.grouped(self.phi): 128 | sign_change = False 129 | for k in ti.static(range(self.dim)): 130 | for s in ti.static((-1, 1)): 131 | offset = ti.Vector.unit(self.dim, k) * s 132 | I1 = I + offset 133 | if I1[k] >= 0 and I1[k] < self.res[k] and \ 134 | ti.math.sign(self.phi[I]) != ti.math.sign(self.phi[I1]): 135 | sign_change = True 136 | 137 | if sign_change and self.phi[I] <= 0: 138 | self.valid[I] = 0 139 | self.phi_temp[I] = self.phi[I] 140 | elif self.phi[I] <= 0: 141 | self.phi_temp[I] = ti.cast(-1, self.real) 142 | else: 143 | self.phi_temp[I] = self.phi[I] 144 | self.valid[I] = 0 145 | 146 | @ti.kernel 147 | def smoothing(self, phi : ti.template(), phi_temp : ti.template()): 148 | for I in ti.grouped(phi_temp): 149 | phi_avg = ti.cast(0, self.real) 150 | tot = ti.cast(0, int) 151 | for k in ti.static(range(self.dim)): 152 | for s in ti.static((-1, 1)): 153 | offset = ti.Vector.unit(self.dim, k) * s 154 | I1 = I + offset 155 | if I1[k] >= 0 and I1[k] < self.res[k]: 156 | phi_avg += phi_temp[I1] 157 | tot += 1 158 | 159 | phi_avg /= tot 160 | phi[I] = phi_avg if phi_avg < phi_temp[I] else phi_temp[I] 161 | 162 | 163 | # H. Zhao. A fast sweeping method for Eikonal equations. Math. Comp., 74:603–627, 2005. 164 | @ti.data_oriented 165 | class FastSweepingLevelSet(LevelSet): 166 | def __init__(self, dim, res, dx, real): 167 | super().__init__(dim, res, dx, real) 168 | 169 | self.repeat_times = 10 170 | 171 | @ti.func 172 | def propagate_update(self, I, s): 173 | if self.valid[I] == -1: 174 | d = self.update_from_neighbor(I) 175 | if ti.abs(d) < ti.abs(self.phi_temp[I]): self.phi_temp[I] = d * ti.math.sign(self.phi[I]) 176 | return s 177 | 178 | @ti.func 179 | def markers_propagate_update(self, markers, lI, o, s): 180 | I, offset = ti.Vector(lI), ti.Vector(o) 181 | if all(I + offset >= 0) and all(I + offset < self.res): 182 | d = (markers[self.valid[I + offset]] - (I + 0.5) * self.dx).norm() 183 | if d < self.phi[I]: 184 | self.phi[I] = d 185 | self.valid[I] = self.valid[I + o] 186 | return s 187 | 188 | @ti.kernel 189 | def propagate(self): 190 | if ti.static(self.dim == 2): 191 | for i in range(self.res[0]): 192 | j = 0 193 | while j < self.res[1]: j += self.propagate_update([i, j], 1) 194 | 195 | for i in range(self.res[0]): 196 | j = self.res[1] - 1 197 | while j >= 0: j += self.propagate_update([i, j], -1) 198 | 199 | for j in range(self.res[1]): 200 | i = 0 201 | while i < self.res[1]: i += self.propagate_update([i, j], 1) 202 | 203 | for j in range(self.res[1]): 204 | i = self.res[1] - 1 205 | while i >= 0: i += self.propagate_update([i, j], -1) 206 | 207 | if ti.static(self.dim == 3): 208 | for i, j in ti.ndrange(self.res[0], self.res[1]): 209 | k = 0 210 | while k < self.res[2]: k += self.propagate_update([i, j, k], 1) 211 | 212 | for i, j in ti.ndrange(self.res[0], self.res[1]): 213 | k = self.res[2] - 1 214 | while k >= 0: k += self.propagate_update([i, j, k], -1) 215 | 216 | for i, k in ti.ndrange(self.res[0], self.res[2]): 217 | j = 0 218 | while j < self.res[1]: j += self.propagate_update([i, j, k], 1) 219 | 220 | for i, k in ti.ndrange(self.res[0], self.res[2]): 221 | j = self.res[1] - 1 222 | while j >= 0: j += self.propagate_update([i, j, k], -1) 223 | 224 | for j, k in ti.ndrange(self.res[1], self.res[2]): 225 | i = 0 226 | while i < self.res[1]: i += self.propagate_update([i, j, k], 1) 227 | 228 | for j, k in ti.ndrange(self.res[1], self.res[2]): 229 | i = self.res[0] - 1 230 | while i >= 0: i += self.propagate_update([i, j, k], -1) 231 | 232 | def redistance(self): 233 | self.valid.fill(-1) 234 | self.target_surface() 235 | self.propagate() 236 | self.phi.copy_from(self.phi_temp) 237 | 238 | @ti.kernel 239 | def markers_propagate(self, markers : ti.template(), total_mk : ti.template()): 240 | if ti.static(self.dim == 2): 241 | for i in range(self.res[0]): 242 | j = 0 243 | while j < self.res[1]: j += self.markers_propagate_update(markers, [i, j], [0, 1], 1) 244 | 245 | for i in range(self.res[0]): 246 | j = self.res[1] - 1 247 | while j >= 0: j += self.markers_propagate_update(markers, [i, j], [0, -1], -1) 248 | 249 | for j in range(self.res[1]): 250 | i = 0 251 | while i < self.res[1]: i += self.markers_propagate_update(markers, [i, j], [1, 0], 1) 252 | 253 | for j in range(self.res[1]): 254 | i = self.res[1] - 1 255 | while i >= 0: i += self.markers_propagate_update(markers, [i, j], [-1, 0], -1) 256 | 257 | if ti.static(self.dim == 3): 258 | for i, j in ti.ndrange(self.res[0], self.res[1]): 259 | k = 0 260 | while k < self.res[2]: k += self.markers_propagate_update(markers, [i, j, k], [0, 0, 1], 1) 261 | 262 | for i, j in ti.ndrange(self.res[0], self.res[1]): 263 | k = self.res[2] - 1 264 | while k >= 0: k += self.markers_propagate_update(markers, [i, j, k], [0, 0, -1], -1) 265 | 266 | for i, k in ti.ndrange(self.res[0], self.res[2]): 267 | j = 0 268 | while j < self.res[1]: j += self.markers_propagate_update(markers, [i, j, k], [0, 1, 0], 1) 269 | 270 | for i, k in ti.ndrange(self.res[0], self.res[2]): 271 | j = self.res[1] - 1 272 | while j >= 0: j += self.markers_propagate_update(markers, [i, j, k], [0, -1, 0], -1) 273 | 274 | for j, k in ti.ndrange(self.res[1], self.res[2]): 275 | i = 0 276 | while i < self.res[1]: i += self.markers_propagate_update(markers, [i, j, k], [1, 0, 0], 1) 277 | 278 | for j, k in ti.ndrange(self.res[1], self.res[2]): 279 | i = self.res[0] - 1 280 | while i >= 0: i += self.markers_propagate_update(markers, [i, j, k], [-1, 0, 0], -1) 281 | 282 | def build_from_markers(self, markers, total_mk): 283 | self.phi.fill(1e20) 284 | self.valid.fill(-1) 285 | self.distance_to_markers(markers, total_mk) 286 | for i in range(self.repeat_times): 287 | self.markers_propagate(markers, total_mk) 288 | self.valid.fill(-1) 289 | self.target_minus() 290 | for i in range(self.repeat_times): 291 | self.propagate() 292 | self.smoothing(self.phi, self.phi_temp) 293 | self.smoothing(self.phi_temp, self.phi) 294 | self.smoothing(self.phi, self.phi_temp) 295 | 296 | -------------------------------------------------------------------------------- /main_2d.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import taichi as ti 3 | from fluid_simulator import * 4 | from flip_extension import * 5 | from power_pic import * 6 | from initializer_2d import * 7 | from visualizer_2d import * 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--flip', action="store_true") 11 | parser.add_argument('--save', action="store_true") 12 | parser.add_argument('--res', type=int, default=64) 13 | parser.add_argument('--t_res', type=int, default=256) 14 | parser.add_argument('--show', type=str, default="p") 15 | args = parser.parse_args() 16 | 17 | if __name__ == '__main__': 18 | fs = False # not support free-surface yet 19 | initializer = SphereInitializer2D(args.res, 0.5, 0.5, 0.2, free_surface=fs) 20 | if args.flip: 21 | solver = FLIPSimulator(dim=2, res=(args.res, args.res), dt=1e-2, substeps=1, dx=1 / args.res, p0=0, 22 | gravity = [0.0, -9.8] if fs else [0.0, 0.0], 23 | free_surface=fs) 24 | else: 25 | solver = PowerPICSimulator(dim=2, res=(args.res, args.res), t_res=(args.t_res, args.t_res), 26 | dt=1e-2, substeps=1, dx=1 / args.res, p0=0, 27 | gravity = [0.0, -9.8] if fs else [0.0, 0.0], 28 | free_surface=fs) 29 | 30 | 31 | visualizer = GUIVisualizer2D(args.res if args.flip else args.t_res, 512, args.show, export= "results" if args.save else "") 32 | solver.initialize(initializer) 33 | solver.run(24 * 60 if args.save else -1, visualizer) 34 | -------------------------------------------------------------------------------- /mgpcg.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import utils 3 | 4 | @ti.data_oriented 5 | class MGPCGPoissonSolver: 6 | def __init__(self, dim, res, n_mg_levels = 4, pre_and_post_smoothing = 2, bottom_smoothing = 50, real = float): 7 | 8 | self.FLUID = utils.FLUID 9 | self.SOLID = utils.SOLID 10 | self.AIR = utils.AIR 11 | 12 | # grid parameters 13 | self.dim = dim 14 | self.res = res 15 | self.n_mg_levels = n_mg_levels 16 | self.pre_and_post_smoothing = pre_and_post_smoothing 17 | self.bottom_smoothing = bottom_smoothing 18 | self.real = real 19 | 20 | # rhs of linear system 21 | self.b = ti.field(dtype=real, shape=res) # Ax=b 22 | 23 | self.r = [ti.field(dtype=real, shape=[res[_] // 2**l for _ in range(dim)]) 24 | for l in range(self.n_mg_levels)] # residual 25 | self.z = [ti.field(dtype=real, shape=[res[_] // 2**l for _ in range(dim)]) 26 | for l in range(self.n_mg_levels)] # M^-1 self.r 27 | 28 | # grid type 29 | self.grid_type = [ti.field(dtype=ti.i32, shape=[res[_] // 2**l for _ in range(dim)]) 30 | for l in range(self.n_mg_levels)] 31 | 32 | # lhs of linear system and its corresponding form in coarse grids 33 | self.Adiag = [ti.field(dtype=real, shape=[res[_] // 2**l for _ in range(dim)]) 34 | for l in range(self.n_mg_levels)] # A(i,j,k)(i,j,k) 35 | self.Ax = [ti.Vector.field(dim, dtype=real, shape=[res[_] // 2**l for _ in range(dim)]) 36 | for l in range(self.n_mg_levels)] # Ax=A(i,j,k)(i+1,j,k), Ay=A(i,j,k)(i,j+1,k), Az=A(i,j,k)(i,j,k+1) 37 | 38 | 39 | self.x = ti.field(dtype=real, shape=res) # solution 40 | self.p = ti.field(dtype=real, shape=res) # conjugate gradient 41 | self.Ap = ti.field(dtype=real, shape=res) # matrix-vector product 42 | self.sum = ti.field(dtype=real, shape=()) # storage for reductions 43 | self.alpha = ti.field(dtype=real, shape=()) # step size 44 | self.beta = ti.field(dtype=real, shape=()) # step size 45 | 46 | @ti.kernel 47 | def init_gridtype(self, grid0 : ti.template(), grid : ti.template()): 48 | for I in ti.grouped(grid): 49 | I2 = I * 2 50 | 51 | tot_fluid = 0 52 | tot_air = 0 53 | for offset in ti.static(ti.grouped(ti.ndrange(*((0, 2), ) * self.dim))): 54 | attr = int(grid0[I2 + offset]) 55 | if attr == self.AIR: tot_air += 1 56 | elif attr == self.FLUID: tot_fluid += 1 57 | 58 | if tot_air > 0: grid[I] = self.AIR 59 | elif tot_fluid > 0: grid[I] = self.FLUID 60 | else: grid[I] = self.SOLID 61 | 62 | 63 | @ti.kernel 64 | def initialize(self): 65 | for I in ti.grouped(ti.ndrange(* [self.res[_] for _ in range(self.dim)])): 66 | self.r[0][I] = 0 67 | self.z[0][I] = 0 68 | self.Ap[I] = 0 69 | self.p[I] = 0 70 | self.x[I] = 0 71 | self.b[I] = 0 72 | 73 | for l in ti.static(range(self.n_mg_levels)): 74 | for I in ti.grouped(ti.ndrange(* [self.res[_] // (2**l) for _ in range(self.dim)])): 75 | self.grid_type[l][I] = 0 76 | self.Adiag[l][I] = 0 77 | self.Ax[l][I] = ti.zero(self.Ax[l][I]) 78 | 79 | def reinitialize(self, cell_type, strategy): 80 | self.initialize() 81 | self.grid_type[0].copy_from(cell_type) 82 | strategy.build_b(self) 83 | strategy.build_A(self, 0) 84 | 85 | for l in range(1, self.n_mg_levels): 86 | self.init_gridtype(self.grid_type[l - 1], self.grid_type[l]) 87 | strategy.build_A(self, l) 88 | 89 | def full_reinitialize(self, strategy): 90 | self.initialize() 91 | strategy.build_b(self) 92 | for l in range(self.n_mg_levels): 93 | self.grid_type[l].fill(utils.FLUID) 94 | strategy.build_A(self, l) 95 | 96 | @ti.func 97 | def neighbor_sum(self, Ax, x, I): 98 | ret = ti.cast(0.0, self.real) 99 | for i in ti.static(range(self.dim)): 100 | offset = ti.Vector.unit(self.dim, i) 101 | ret += Ax[I - offset][i] * x[I - offset] + Ax[I][i] * x[I + offset] 102 | return ret 103 | 104 | @ti.kernel 105 | def smooth(self, l: ti.template(), phase: ti.template()): 106 | # phase = red/black Gauss-Seidel phase 107 | for I in ti.grouped(self.r[l]): 108 | if (I.sum()) & 1 == phase and self.grid_type[l][I] == self.FLUID: 109 | self.z[l][I] = (self.r[l][I] - self.neighbor_sum(self.Ax[l], self.z[l], I)) / self.Adiag[l][I] 110 | 111 | @ti.kernel 112 | def restrict(self, l: ti.template()): 113 | for I in ti.grouped(self.r[l]): 114 | if self.grid_type[l][I] == self.FLUID: 115 | Az = self.Adiag[l][I] * self.z[l][I] 116 | Az += self.neighbor_sum(self.Ax[l], self.z[l], I) 117 | res = self.r[l][I] - Az 118 | self.r[l + 1][I // 2] += res 119 | 120 | @ti.kernel 121 | def prolongate(self, l: ti.template()): 122 | for I in ti.grouped(self.z[l]): 123 | self.z[l][I] += self.z[l + 1][I // 2] 124 | 125 | def v_cycle(self): 126 | self.z[0].fill(0.0) 127 | for l in range(self.n_mg_levels - 1): 128 | for i in range(self.pre_and_post_smoothing): 129 | self.smooth(l, 0) 130 | self.smooth(l, 1) 131 | 132 | self.r[l + 1].fill(0.0) 133 | self.z[l + 1].fill(0.0) 134 | self.restrict(l) 135 | 136 | # solve Az = r on the coarse grid 137 | for i in range(self.bottom_smoothing // 2): 138 | self.smooth(self.n_mg_levels - 1, 0) 139 | self.smooth(self.n_mg_levels - 1, 1) 140 | for i in range(self.bottom_smoothing // 2): 141 | self.smooth(self.n_mg_levels - 1, 1) 142 | self.smooth(self.n_mg_levels - 1, 0) 143 | 144 | for l in reversed(range(self.n_mg_levels - 1)): 145 | self.prolongate(l) 146 | for i in range(self.pre_and_post_smoothing): 147 | self.smooth(l, 1) 148 | self.smooth(l, 0) 149 | 150 | def solve(self, 151 | max_iters=-1, 152 | verbose=False, 153 | rel_tol=1e-12, 154 | abs_tol=1e-14, 155 | eps=1e-12): 156 | 157 | self.r[0].copy_from(self.b) 158 | self.reduce(self.r[0], self.r[0]) 159 | initial_rTr = self.sum[None] 160 | 161 | if verbose: 162 | print(f"init rtr = {initial_rTr}") 163 | 164 | tol = max(abs_tol, initial_rTr * rel_tol) 165 | 166 | # self.r = b - Ax = b since self.x = 0 167 | # self.p = self.r = self.r + 0 self.p 168 | self.v_cycle() 169 | self.update_p() 170 | 171 | self.reduce(self.z[0], self.r[0]) 172 | old_zTr = self.sum[None] 173 | 174 | # Conjugate gradients 175 | iter = 0 176 | while max_iters == -1 or iter < max_iters: 177 | # self.alpha = rTr / pTAp 178 | self.compute_Ap() 179 | self.reduce(self.p, self.Ap) 180 | pAp = self.sum[None] 181 | self.alpha[None] = old_zTr / (pAp + eps) 182 | 183 | # self.x = self.x + self.alpha self.p 184 | # self.r = self.r - self.alpha self.Ap 185 | self.update_xr() 186 | 187 | # check for convergence 188 | self.reduce(self.r[0], self.r[0]) 189 | rTr = self.sum[None] 190 | 191 | if verbose: 192 | print(f'iter {iter}, |residual|_2={ti.sqrt(rTr)}') 193 | 194 | if rTr < tol: 195 | break 196 | 197 | # self.z = M^-1 self.r 198 | self.v_cycle() 199 | 200 | # self.beta = new_rTr / old_rTr 201 | self.reduce(self.z[0], self.r[0]) 202 | new_zTr = self.sum[None] 203 | self.beta[None] = new_zTr / (old_zTr + eps) 204 | 205 | # self.p = self.z + self.beta self.p 206 | self.update_p() 207 | old_zTr = new_zTr 208 | 209 | iter += 1 210 | 211 | @ti.kernel 212 | def reduce(self, p: ti.template(), q: ti.template()): 213 | self.sum[None] = 0 214 | for I in ti.grouped(p): 215 | if self.grid_type[0][I] == self.FLUID: 216 | self.sum[None] += p[I] * q[I] 217 | 218 | @ti.kernel 219 | def compute_Ap(self): 220 | for I in ti.grouped(self.Ap): 221 | if self.grid_type[0][I] == self.FLUID: 222 | r = self.Adiag[0][I] * self.p[I] 223 | r += self.neighbor_sum(self.Ax[0], self.p, I) 224 | self.Ap[I] = r 225 | 226 | @ti.kernel 227 | def update_xr(self): 228 | alpha = self.alpha[None] 229 | for I in ti.grouped(self.p): 230 | if self.grid_type[0][I] == self.FLUID: 231 | self.x[I] += alpha * self.p[I] 232 | self.r[0][I] -= alpha * self.Ap[I] 233 | 234 | @ti.kernel 235 | def update_p(self): 236 | for I in ti.grouped(self.p): 237 | if self.grid_type[0][I] == self.FLUID: 238 | self.p[I] = self.z[0][I] + self.beta[None] * self.p[I] 239 | -------------------------------------------------------------------------------- /power_pic.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | from flip_extension import FLIPSimulator 4 | from utils import * 5 | import utils 6 | import numpy as np 7 | 8 | @ti.data_oriented 9 | class PowerPICSimulator(FLIPSimulator): 10 | def __init__(self, 11 | dim = 2, 12 | res = (128, 128), 13 | t_res = (256, 256), 14 | dt = 1.25e-2, 15 | substeps = 1, 16 | dx = 1.0, 17 | rho = 1000.0, 18 | gravity = [0, -9.8], 19 | p0 = 1e-3, 20 | real = float, 21 | free_surface = True): 22 | super().__init__(dim, res, dt, substeps, dx, rho, gravity, p0, real, free_surface) 23 | self.s_p = ti.field(real, shape=self.max_particles) 24 | self.sum_p = ti.field(real, shape=self.max_particles) 25 | self.Z_p = ti.field(real, shape=self.max_particles) # log-domain Z=max{a[i]+b[i]} 26 | self.c_p = ti.Vector.field(dim, real, shape=self.max_particles) 27 | self.u_p = ti.Vector.field(dim, real, shape=self.max_particles) 28 | 29 | self.t_res = t_res 30 | self.t_dx = 1 / self.t_res[0] 31 | self.s_j = ti.field(real) 32 | self.sum_j = ti.field(real) 33 | self.Z_j = ti.field(real) # log-domain Z=max{a[i]+b[i]} 34 | ti.root.dense(self.indices, self.t_res).place(self.s_j, self.sum_j, self.Z_j) 35 | 36 | self.entropy_eps = 2.0 * (self.t_dx ** 2) 37 | self.sinkhorn_delta = 0.1 38 | self.V_j = self.t_dx ** self.dim 39 | 40 | self.R_2 = 6 # 3 * sqrt(2) + EPS 41 | self.R = self.R_2 * 2 42 | 43 | @ti.func 44 | def K_pj(self, x_p, x_j): 45 | return -(x_p - x_j).norm_sqr() / self.entropy_eps 46 | 47 | @ti.func 48 | def check_Tg(self, I): 49 | return all(I >= 0) and all(I < self.t_res) 50 | 51 | @ti.func 52 | def T(self, p, I): 53 | pos = (I + 0.5) * self.t_dx 54 | return ti.math.exp(self.K_pj(self.p_x[p], pos) + self.s_p[p] + self.s_j[I]) 55 | 56 | @ti.func 57 | def get_base(self, I): 58 | p, I_i = I[0], I[1:] 59 | base = (self.p_x[p] / self.t_dx).cast(ti.i32) 60 | I_x = base + I_i - self.R_2 61 | pos = (I_x + 0.5) * self.t_dx 62 | return p, I_x, pos 63 | 64 | def init_sinkhorn_algo(self): 65 | self.s_p.fill(0.0) 66 | self.s_j.fill(0.0) 67 | 68 | @ti.kernel 69 | def calc_max_sum_j(self) -> ti.f32: 70 | for I in ti.grouped(self.sum_j): self.sum_j[I] = 0.0 71 | for I in ti.grouped(ti.ndrange(self.total_mk[None], *(self.R, ) * self.dim)): 72 | p, I_x, pos = self.get_base(I) 73 | if self.check_Tg(I_x): 74 | self.sum_j[I_x] += self.T(p, I_x) 75 | 76 | ret = 0.0 77 | for I in ti.grouped(self.sum_j): 78 | if self.check_Tg(I): 79 | val = ti.abs(self.sum_j[I] / self.V_j - 1) 80 | ti.atomic_max(ret, val) 81 | return ret 82 | 83 | @ti.kernel 84 | def sinkhorn_algo(self): 85 | V_p = 1.0 / self.total_mk[None] 86 | log_Vp = ti.math.log(V_p) 87 | log_Vj = ti.math.log(self.V_j) 88 | stencil = ti.static((self.total_mk[None], ) + (self.R, ) * self.dim) 89 | 90 | for I in ti.grouped(self.sum_j): 91 | self.sum_j[I] = 0.0 92 | self.Z_j[I] = -1e10 93 | 94 | # solve in log-domain 95 | for I in ti.grouped(ti.ndrange(*stencil)): 96 | p, I_x, pos = self.get_base(I) 97 | if self.check_Tg(I_x): 98 | ti.atomic_max(self.Z_j[I_x], self.K_pj(self.p_x[p], pos) + self.s_p[p]) 99 | for I in ti.grouped(ti.ndrange(*stencil)): 100 | p, I_x, pos = self.get_base(I) 101 | if self.check_Tg(I_x): 102 | self.sum_j[I_x] += ti.math.exp(self.K_pj(self.p_x[p], pos) + self.s_p[p] - self.Z_j[I_x]) 103 | 104 | for I in ti.grouped(self.s_j): 105 | self.s_j[I] = log_Vj - (self.Z_j[I] + ti.math.log(self.sum_j[I])) 106 | 107 | 108 | for p in self.sum_p: 109 | self.sum_p[p] = 0.0 110 | self.Z_p[p] = -1e10 111 | for I in ti.grouped(ti.ndrange(*stencil)): 112 | p, I_x, pos = self.get_base(I) 113 | if self.check_Tg(I_x): 114 | ti.atomic_max(self.Z_p[p], self.K_pj(self.p_x[p], pos) + self.s_j[I_x]) 115 | for I in ti.grouped(ti.ndrange(*stencil)): 116 | p, I_x, pos = self.get_base(I) 117 | if self.check_Tg(I_x): 118 | self.sum_p[p] += ti.math.exp(self.K_pj(self.p_x[p], pos) + self.s_j[I_x] - self.Z_p[p]) 119 | 120 | for p in self.s_p: 121 | self.s_p[p] = log_Vp - (self.Z_p[p] + ti.math.log(self.sum_p[p])) 122 | 123 | @ti.kernel 124 | def p2g(self): 125 | V_p = 1.0 / self.total_mk[None] 126 | for I in ti.grouped(ti.ndrange(self.total_mk[None], *(self.R, ) * self.dim)): 127 | p, I_x, pos = self.get_base(I) 128 | if self.check_Tg(I_x): 129 | for k in ti.static(range(self.dim)): 130 | utils.splat_w(self.velocity[k], self.velocity_backup[k], self.p_v[p][k], pos / self.dx - 0.5 * (1 - ti.Vector.unit(self.dim, k)), self.T(p, I_x) / V_p) 131 | 132 | for k in ti.static(range(self.dim)): 133 | for I in ti.grouped(self.velocity_backup[k]): # reuse velocity_backup as weight 134 | if self.velocity_backup[k][I] > 0: 135 | self.velocity[k][I] /= self.velocity_backup[k][I] 136 | 137 | @ti.kernel 138 | def g2p(self, dt : ti.f32): 139 | V_p = 1.0 / self.total_mk[None] 140 | for p in range(self.total_mk[None]): 141 | self.c_p[p].fill(0.0) 142 | self.u_p[p].fill(0.0) 143 | for I in ti.grouped(ti.ndrange(self.total_mk[None], *(self.R, ) * self.dim)): 144 | p, I_x, pos = self.get_base(I) 145 | if self.check_Tg(I_x): 146 | T = self.T(p, I_x) 147 | self.p_v[p] += T * (self.vel_interp(pos) - self.vel_old_interp(pos)) / V_p 148 | self.c_p[p] += pos * T / V_p 149 | self.u_p[p] += T * self.vel_interp(pos) / V_p 150 | 151 | for p in range(self.total_mk[None]): 152 | self.p_x[p] = self.c_p[p] + dt * self.u_p[p] 153 | 154 | def substep(self, dt): 155 | self.init_sinkhorn_algo() 156 | while self.calc_max_sum_j() > self.sinkhorn_delta: 157 | self.sinkhorn_algo() 158 | 159 | for k in range(self.dim): 160 | self.velocity[k].fill(0) 161 | self.velocity_backup[k].fill(0) 162 | self.p2g() 163 | for k in range(self.dim): 164 | self.velocity_backup[k].copy_from(self.velocity[k]) 165 | 166 | # self.extrap_velocity() 167 | # self.enforce_boundary() 168 | 169 | self.add_gravity(dt) 170 | self.enforce_boundary() 171 | self.solve_pressure(dt, self.strategy) 172 | 173 | if self.verbose: 174 | prs = np.max(self.pressure.to_numpy()) 175 | print(f'\033[36mMax pressure: {prs}\033[0m') 176 | 177 | self.apply_pressure(dt) 178 | self.g2p(dt) 179 | self.total_t += self.dt -------------------------------------------------------------------------------- /pressure_project.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | from mgpcg import MGPCGPoissonSolver 3 | import utils 4 | 5 | @ti.data_oriented 6 | class PressureProjectStrategy: 7 | def __init__(self, dim, velocity, ghost_fluid_method, phi, p0): 8 | self.dim = dim 9 | 10 | self.velocity = velocity 11 | self.ghost_fluid_method = ghost_fluid_method 12 | self.phi = phi 13 | self.p0 = p0 # the standard atmospheric pressure 14 | 15 | @ti.kernel 16 | def build_b_kernel(self, 17 | cell_type : ti.template(), 18 | b : ti.template()): 19 | for I in ti.grouped(cell_type): 20 | if cell_type[I] == utils.FLUID: 21 | for k in ti.static(range(self.dim)): 22 | offset = ti.Vector.unit(self.dim, k) 23 | b[I] += (self.velocity[k][I] - self.velocity[k][I + offset]) 24 | b[I] *= self.scale_b 25 | 26 | for I in ti.grouped(cell_type): 27 | if cell_type[I] == utils.FLUID: 28 | for k in ti.static(range(self.dim)): 29 | for s in ti.static((-1, 1)): 30 | offset = ti.Vector.unit(self.dim, k) * s 31 | if cell_type[I + offset] == utils.SOLID: 32 | if s < 0: b[I] -= self.scale_b * (self.velocity[k][I] - 0) 33 | else: b[I] += self.scale_b * (self.velocity[k][I + offset] - 0) 34 | elif cell_type[I + offset] == utils.AIR: 35 | if ti.static(self.ghost_fluid_method): 36 | c = (self.phi[I] - self.phi[I + offset]) / self.phi[I] 37 | b[I] += self.scale_A * min(c, 1e3) * self.p0 38 | else: 39 | b[I] += self.scale_A * self.p0 40 | 41 | def build_b(self, solver : MGPCGPoissonSolver): 42 | self.build_b_kernel(solver.grid_type[0], 43 | solver.b) 44 | 45 | @ti.kernel 46 | def build_A_kernel(self, 47 | level : ti.template(), 48 | grid_type : ti.template(), 49 | Adiag : ti.template(), 50 | Ax : ti.template()): 51 | for I in ti.grouped(grid_type): 52 | if grid_type[I] == utils.FLUID: 53 | for k in ti.static(range(self.dim)): 54 | for s in ti.static((-1, 1)): 55 | offset = ti.Vector.unit(self.dim, k) * s 56 | if grid_type[I + offset] == utils.FLUID: 57 | Adiag[I] += self.scale_A 58 | if ti.static(s > 0): Ax[I][k] = -self.scale_A 59 | elif grid_type[I + offset] == utils.AIR: 60 | if ti.static(self.ghost_fluid_method and level == 0): 61 | c = (self.phi[I] - self.phi[I + offset]) / self.phi[I] 62 | Adiag[I] += self.scale_A * min(c, 1e3) 63 | else: 64 | Adiag[I] += self.scale_A 65 | 66 | def build_A(self, solver : MGPCGPoissonSolver, level): 67 | self.build_A_kernel(level, 68 | solver.grid_type[level], 69 | solver.Adiag[level], 70 | solver.Ax[level]) 71 | 72 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | taichi 2 | numpy 3 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | FLUID = 0 4 | AIR = 1 5 | SOLID = 2 6 | 7 | @ti.pyfunc 8 | def clamp(x, a, b): 9 | return max(a, min(b, x)) 10 | 11 | @ti.func 12 | def sample(data, pos): 13 | tot = data.shape 14 | # static unfold for efficiency 15 | if ti.static(len(data.shape) == 2): 16 | i, j = ti.math.clamp(int(pos[0]), 0, tot[0] - 1), ti.math.clamp(int(pos[1]), 0, tot[1] - 1) 17 | ip, jp = ti.math.clamp(i + 1, 0, tot[0] - 1), ti.math.clamp(j + 1, 0, tot[1] - 1) 18 | s, t = ti.math.clamp(pos[0] - i, 0.0, 1.0), ti.math.clamp(pos[1] - j, 0.0, 1.0) 19 | return \ 20 | (data[i, j] * (1 - s) + data[ip, j] * s) * (1 - t) + \ 21 | (data[i, jp] * (1 - s) + data[ip, jp] * s) * t 22 | 23 | else: 24 | i, j, k = ti.math.clamp(int(pos[0]), 0, tot[0] - 1), ti.math.clamp(int(pos[1]), 0, tot[1] - 1), ti.math.clamp(int(pos[2]), 0, tot[2] - 1) 25 | ip, jp, kp = ti.math.clamp(i + 1, 0, tot[0] - 1), ti.math.clamp(j + 1, 0, tot[1] - 1), ti.math.clamp(k + 1, 0, tot[2] - 1) 26 | s, t, u = ti.math.clamp(pos[0] - i, 0.0, 1.0), ti.math.clamp(pos[1] - j, 0.0, 1.0), ti.math.clamp(pos[2] - k, 0.0, 1.0) 27 | return \ 28 | ((data[i, j, k] * (1 - s) + data[ip, j, k] * s) * (1 - t) + \ 29 | (data[i, jp, k] * (1 - s) + data[ip, jp, k] * s) * t) * (1 - u) + \ 30 | ((data[i, j, kp] * (1 - s) + data[ip, j, kp] * s) * (1 - t) + \ 31 | (data[i, jp, kp] * (1 - s) + data[ip, jp, kp] * s) * t) * u 32 | 33 | @ti.func 34 | def splat_w(data, weights, v, pos, T0): 35 | tot = data.shape 36 | dim = ti.static(len(tot)) 37 | 38 | I0 = ti.Vector.zero(int, dim) 39 | I1 = ti.Vector.zero(int, dim) 40 | w = ti.zero(pos) 41 | 42 | for k in ti.static(range(dim)): 43 | I0[k] = ti.math.clamp(int(pos[k]), 0, tot[k] - 1) 44 | I1[k] = ti.math.clamp(I0[k] + 1, 0, tot[k] - 1) 45 | w[k] = ti.math.clamp(pos[k] - I0[k], 0.0, 1.0) 46 | 47 | for u in ti.static(ti.grouped(ti.ndrange(*((0, 2), ) * dim))): 48 | dpos = ti.zero(pos) 49 | I = ti.Vector.zero(int, dim) 50 | W = 1.0 51 | for k in ti.static(range(dim)): 52 | dpos[k] = (pos[k] - I0[k]) if u[k] == 0 else (pos[k] - I1[k]) 53 | I[k] = I0[k] if u[k] == 0 else I1[k] 54 | W *= (1 - w[k]) if u[k] == 0 else w[k] 55 | data[I] += v * W * T0 56 | weights[I] += W * T0 57 | 58 | @ti.func 59 | def splat(data, weights, v, pos): 60 | splat_w(data, weights, v, pos, 1.0) -------------------------------------------------------------------------------- /visualizer_2d.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import utils 3 | from fluid_simulator import * 4 | import numpy as np 5 | 6 | @ti.data_oriented 7 | class Visualizer2D: 8 | def __init__(self, grid_res, res): 9 | self.grid_res = grid_res 10 | self.res = res 11 | self.tmp = ti.Vector.field(3, dtype=ti.f32, shape=(self.grid_res, self.grid_res)) 12 | self.tmp_w = ti.field(dtype=ti.f32, shape=(self.grid_res, self.grid_res)) 13 | self.color_buffer = ti.Vector.field(3, dtype=ti.f32, shape=(self.res, self.res)) 14 | 15 | @ti.func 16 | def ij_to_xy(self, i, j): 17 | return int((i + 0.5) / self.res * self.grid_res), \ 18 | int((j + 0.5) / self.res * self.grid_res) 19 | 20 | @ti.kernel 21 | def fill_power(self, sim : ti.template()): 22 | for i, j in self.tmp: 23 | self.tmp[i, j].fill(0.0) 24 | self.tmp_w[i, j] = 0.0 25 | for I in ti.grouped(ti.ndrange(sim.total_mk[None], *(sim.R, ) * sim.dim)): 26 | p, I_x, pos = self.get_base(I) 27 | if self.check_Tg(I_x): 28 | T = sim.T(p, I_x) 29 | self.tmp[I_x] += T * sim.color_p[p] 30 | self.tmp_w[I_x] += T 31 | 32 | V_p = (1.0 / sim.total_mk[None]) 33 | for i, j in self.color_buffer: 34 | x, y = self.ij_to_xy(i, j) 35 | self.color_buffer[i, j] = self.tmp[x, y] / self.tmp_w[x, y] 36 | 37 | def visualize_factory(self, simulator): 38 | if self.mode == 'power': 39 | self.fill_power(simulator) 40 | 41 | def visualize(self, simulator): 42 | assert 0, 'Please use GUIVisualizer2D' 43 | 44 | def end(self): 45 | pass 46 | 47 | @ti.data_oriented 48 | class GUIVisualizer2D(Visualizer2D): 49 | def __init__(self, grid_res, res, mode, title = 'demo', export=""): 50 | super().__init__(grid_res, res) 51 | self.mode = mode 52 | self.window = ti.ui.Window(title, (res, res), vsync=True) 53 | self.canvas = self.window.get_canvas() 54 | self.frame = 0 55 | self.export = False 56 | if export != "": 57 | self.export = True 58 | self.video_manager = ti.tools.VideoManager(export) 59 | 60 | def visualize(self, simulator): 61 | self.canvas.set_background_color(color=(0.0, 0.0, 0.0)) 62 | if self.mode == 'p': 63 | self.canvas.circles(simulator.p_x, 0.002, per_vertex_color=simulator.color_p) 64 | else: 65 | self.visualize_factory(simulator) 66 | self.canvas.set_image(self.color_buffer) 67 | 68 | if self.export: 69 | img = self.window.get_image_buffer_as_numpy() 70 | self.video_manager.write_frame(img) 71 | 72 | self.window.show() 73 | self.frame += 1 74 | 75 | def end(self): 76 | if self.export: 77 | self.video_manager.make_video(gif=True, mp4=True) -------------------------------------------------------------------------------- /volume_control.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | from mgpcg import MGPCGPoissonSolver 3 | from pressure_project import PressureProjectStrategy 4 | import utils 5 | 6 | @ti.data_oriented 7 | class PressureProjectWithVolumeControlStrategy(PressureProjectStrategy): 8 | def __init__(self, dim, velocity, ghost_fluid_method, phi, p0, level_set, dt): 9 | super().__init__(dim, velocity, ghost_fluid_method, phi, p0) 10 | self.level_set = level_set 11 | self.dt = dt 12 | 13 | self.step = 0 14 | self.vol_0 = 0.0 # the desired value 15 | # proportional-integral (PI) controller 16 | self.y = 0.0 17 | self.n_p = 25 18 | self.k_p = 2.3 / (self.n_p * self.dt) # proportional gain 19 | self.zeta = 2 # suppress the noise coming from the volume computation 20 | self.k_i = (self.k_p / (2 * self.zeta)) ** 2 # PI gain 21 | 22 | @ti.kernel 23 | def calc_volume(self) -> ti.f32: 24 | vol = ti.cast(0, ti.f32) 25 | for I in ti.grouped(self.phi): 26 | cell_vol = ti.cast(0, ti.f32) 27 | for offset in ti.static(ti.grouped(ti.ndrange(*((0, 2), ) * self.dim))): 28 | cell_vol += self.level_set.theta(-utils.sample(self.phi, I + offset - 0.5)) 29 | vol += cell_vol / (2 ** self.dim) 30 | 31 | return vol 32 | 33 | @ti.kernel 34 | def add_c(self, cell_type : ti.template(), b : ti.template(), c : ti.f32): 35 | for I in ti.grouped(b): 36 | if cell_type[I] == utils.FLUID: 37 | b[I] += c 38 | 39 | def build_b(self, solver : MGPCGPoissonSolver): 40 | self.step += 1 41 | super().build_b(solver) 42 | vol = self.calc_volume() 43 | if self.step == 1: 44 | self.vol_0 = vol 45 | return 46 | 47 | x = (vol - self.vol_0) / self.vol_0 # the normalized difference between the current and desired volume 48 | self.y += x * self.dt # drift error can be removed by integrating the volume error over time 49 | c = (-self.k_p * x - self.k_i * self.y) / (x + 1) # the required divergence 50 | print(f'\033[31mvolume = {vol}, c = {c}\033[0m') 51 | self.add_c(solver.grid_type[0], solver.b, c) 52 | --------------------------------------------------------------------------------