├── requirements.txt ├── paint.py ├── LICENSE ├── flow_visualization.py ├── .gitignore ├── README.md ├── test ├── forward_fct.py └── diff_fct.py ├── 2dvof.py ├── diff_vof.py ├── diff_vof_replaced.py └── 3dvof.py /requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib==3.5.1 2 | numpy==1.22.3 3 | taichi==1.4.1 4 | -------------------------------------------------------------------------------- /paint.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | ti.init(arch=ti.cpu) 4 | 5 | buf_size = (400, 400) 6 | f = ti.field(dtype=ti.f32, shape=buf_size) 7 | 8 | gui = ti.GUI("Paint", buf_size) 9 | 10 | @ti.kernel 11 | def set_pixel(x:ti.f32, y:ti.f32): 12 | xcord = ti.i32(x * buf_size[0]) 13 | ycord = ti.i32(y * buf_size[1]) 14 | for i, j in ti.ndrange((xcord - 10, xcord + 10),(ycord-10, ycord+10) ): 15 | if i > 0 and j > 0: 16 | f[i, j] = 1.0 17 | 18 | while gui.running: 19 | gui.set_image(f) 20 | for e in gui.get_events(ti.GUI.PRESS): 21 | if e.key == ti.GUI.LMB: 22 | x, y = gui.get_cursor_pos() 23 | set_pixel(x, y) 24 | elif e.key == ti.GUI.ESCAPE: 25 | gui.running = False 26 | gui.show() -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Qian Bao 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 | -------------------------------------------------------------------------------- /flow_visualization.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import taichi as ti 3 | 4 | def plot_vector_field(vector_field, arrow_spacing, gui): 5 | ''' 6 | Plot the velocity vector field using gui.lines 7 | and gui.triangles. 8 | Might slow down the gui refresh rate because of the 9 | overhead of the nested for-loop. 10 | ''' 11 | V_np = vector_field.to_numpy() 12 | V_norm = np.linalg.norm(V_np, axis=-1) 13 | nx, ny = V_np.shape[0], V_np.shape[1] 14 | 15 | max_magnitude = np.max(V_norm) 16 | scale_factor = min(nx, ny) * 0.1 / (max_magnitude + 1e-16) 17 | arrowhead_size = 0.3 * arrow_spacing / min(nx, ny) 18 | 19 | for i in range(1, nx, arrow_spacing): 20 | for j in range(1, ny, arrow_spacing): 21 | x = i / nx 22 | y = j / ny 23 | u_arrow = V_np[i, j][0] * scale_factor / nx 24 | v_arrow = V_np[i, j][1] * scale_factor / ny 25 | begin = np.array(([x, y],)) 26 | end = np.array(([x + u_arrow, y + v_arrow],)) 27 | dx, dy = u_arrow, v_arrow 28 | arrow_dir = np.array([dx, dy]) / np.linalg.norm(np.array([dx, dy])) 29 | arrow_normal = np.array([-arrow_dir[1], arrow_dir[0]]) 30 | arrowhead_a = end - arrowhead_size * arrow_dir + 0.5 * arrowhead_size * arrow_normal 31 | arrowhead_b = end - arrowhead_size * arrow_dir - 0.5 * arrowhead_size * arrow_normal 32 | gui.triangles(a=end, b=arrowhead_a, c=arrowhead_b, color=0x000000) 33 | gui.lines(begin=begin, end=end, radius=1, color=0x000000) 34 | 35 | def plot_arrow_field(vector_field, arrow_spacing, gui): 36 | ''' 37 | Plot the velocity vector field using the built-in 38 | gui.arrows. 39 | ''' 40 | V_np = vector_field.to_numpy() 41 | V_norm = np.linalg.norm(V_np, axis=-1) 42 | nx, ny, ndim = V_np.shape 43 | 44 | max_magnitude = np.max(V_norm) 45 | scale_factor = min(nx, ny) * 0.1 / (max_magnitude + 1e-16) 46 | 47 | x = np.arange(0, 1, arrow_spacing / nx) 48 | y = np.arange(0, 1, arrow_spacing / ny) 49 | 50 | X, Y = np.meshgrid(x, y) 51 | begin = np.dstack((X, Y)).reshape(-1, 2, order='F') 52 | incre = (V_np[::arrow_spacing,::arrow_spacing] \ 53 | * np.array([scale_factor / nx, scale_factor / ny])) \ 54 | .reshape(-1, 2, order='C') 55 | gui.arrows(orig=begin, direction=incre, radius=1, color=0x000000) 56 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # Output dir 132 | output/ 133 | .DS_Store 134 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # taichi-2d-vof 2 | A single file VOF fluid solver implementation in Taichi. 3 | 4 | ![VOF blog](https://user-images.githubusercontent.com/2747993/225257322-8b55cf4e-16fa-4801-912d-8f1eb89a93c5.gif) 5 | 6 | Simply install Taichi and run! 7 | ```bash 8 | $ python3 -m pip install taichi 9 | $ python3 2dvof.py 10 | ``` 11 | 12 | In the GUI window, you can press **SPACE** to switch between visualization methods, or press **q** on your keyboard to terminate the simulation. 13 | 14 | ## Usage 15 | You can execute the script with the following arguments to control its behavior: 16 | ```bash 17 | $ python3 2dvof.py -ic 1 # ic stands for Initial Condition; Default value is 1 18 | $ python3 2dvof.py -s # Save png files to the ./output/ folder 19 | ``` 20 | Currently, there are three types of initial condition settings: 21 | 1. Dam break 22 | 2. Rising bubble 23 | 3. Droping liquid 24 | 25 | You can also tweak the `set_init_F()` kernel to implement your own. Be aware 26 | that the solver might become unstable under various settings. The material 27 | properties are carefully chosed to produce best results. 28 | 29 | During the operation, press **SPACE** to switch visualization method. Currently, 30 | the script can display 31 | 1. VOF field 32 | 2. U velocity 33 | 3. V velocity 34 | 4. Velocity norm 35 | 5. Velocity vectors (Implemented in `flow_visualization.py` module) 36 | 37 | ![display](https://user-images.githubusercontent.com/2747993/226554195-cd767de2-f386-46aa-8be7-00f1ed0c7f7a.png) 38 | 39 | You can follow the steps below if you wish to output a video file: 40 | ```bash 41 | $ python3 2dvof.py -ic 1 -s # Add -s argument to enable saving 42 | $ cd output 43 | $ ti video # Use Taichi's video utility to generate a mp4 file 44 | $ ti gif -i video.mp4 # Use Taichi's gif utility to generate a gif file 45 | ``` 46 | 47 | ## Extension to 3-Dimension 48 | The script in this repo can be easily extended to 3-dimension with minor changes in the implementation. 49 | An experimental dam-break demo is provided in the `3dvof.py` file. The following rendered images are produced by the [Taitopia](https://taitopia.design/) cloud-render. 50 | The VOF field data are exported to `.obj` file format using Paraview. 51 | 52 | Screenshot 2023-09-07 at 22 42 46 53 | 54 | The following 3D animation is created in Blender by rendering the `.obj` sequences generated in Paraview: 55 | 56 | ![ezgif com-crop](https://github.com/houkensjtu/taichi-2d-vof/assets/2747993/53f1f947-3e14-43cc-8127-c5c88d5c7c02) 57 | 58 | More details about the 3D implementation will be released in another repo, stay tuned! 59 | 60 | ## Implementation 61 | WIP. 62 | 63 | ## References 64 | 1. Volume of Fluid (VOF) Method for the Dynamics of Free Boundaries, C. W. Hirt and B. D. Nichols 65 | 2. Direct Numerical Simulations of Gas–Liquid Multiphase Flows, Grétar Tryggvason, Ruben Scardovelli, Stéphane Zaleski 66 | 3. Fully Multidimensional Flux-Corrected Transport Algorithms for Fluids, Steven T. Zalesak 67 | 4. Volume-Tracking Methods For Interfacial Flow Calculations, Murray Rudman 68 | 5. A Continuum Method for Modeling Surface Tension, J.U. Brackbill, D.B. Kothe, C. Zemach 69 | 70 | 71 | ## Acknowledgement 72 | The code in this repository is jointly developed by @houkensjtu and @zju-zhoucl. The original early version can be found 73 | in @zju-zhoucl's repo [here](https://github.com/zju-zhoucl/taichi_VOF). 74 | -------------------------------------------------------------------------------- /test/forward_fct.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | from math import pi as pi 5 | import os 6 | 7 | ti.init(arch=ti.cpu, default_fp=ti.f32) 8 | 9 | nx = 500 # Number of grid points in the x direction 10 | ny = 500 # Number of grid points in the y direction 11 | 12 | Lx = pi # The length of the domain 13 | Ly = pi # The width of the domain 14 | 15 | # Solution parameters 16 | dt = 1e-4 # Use smaller dt for higher density ratio 17 | imin = 1 18 | imax = imin + nx - 1 19 | jmin = 1 20 | jmax = jmin + ny - 1 21 | tmax = 1000 22 | 23 | F = ti.field(float, shape=(2 * tmax + 1, imax + 2, jmax + 2)) 24 | Ftd_x = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 25 | Ftd_y = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 26 | Ftarget = ti.field(float, shape=(imax + 2, jmax + 2)) 27 | Fd = ti.field(float, shape=(imax + 2, jmax + 2)) 28 | 29 | ax = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 30 | ay = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 31 | cx = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 32 | cy = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 33 | rp_x = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 34 | rm_x = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 35 | rp_y = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 36 | rm_y = ti.field(float, shape=(tmax, imax + 2, jmax + 2)) 37 | 38 | u = ti.field(float, shape=(imax + 2, jmax + 2)) 39 | v = ti.field(float, shape=(imax + 2, jmax + 2)) 40 | 41 | x = ti.field(float, shape=imax + 3) 42 | y = ti.field(float, shape=jmax + 3) 43 | x.from_numpy(np.hstack((0.0, np.linspace(0, Lx, nx + 1), Lx))) # [0, 0, ... 1, 1] 44 | y.from_numpy(np.hstack((0.0, np.linspace(0, Ly, ny + 1), Ly))) # [0, 0, ... 1, 1] 45 | 46 | xm = ti.field(float, shape=imax + 2) 47 | ym = ti.field(float, shape=jmax + 2) 48 | dx = x[imin + 2] - x[imin + 1] 49 | dy = y[jmin + 2] - y[jmin + 1] 50 | dxi = 1 / dx 51 | dyi = 1 / dy 52 | 53 | print(f'>>> VOF scheme testing') 54 | print(f'>>> Grid resolution: {nx} x {ny}, dt = {dt:4.2e}') 55 | 56 | @ti.kernel 57 | 58 | def grid_staggered(): # 11/3 Checked 59 | ''' 60 | Calculate the center position of cells. 61 | ''' 62 | for i in xm: # xm[0] = 0.0, xm[33] = 1.0 63 | xm[i] = 0.5 * (x[i] + x[i + 1]) 64 | for j in ym: 65 | ym[j] = 0.5 * (y[j] + y[j + 1]) 66 | 67 | 68 | @ti.func 69 | def find_area(i, j, cx, cy, r): 70 | a = 0.0 71 | xcoord_ct = (i - imin) * dx + dx / 2 72 | ycoord_ct = (j - jmin) * dy + dy / 2 73 | 74 | xcoord_lu = xcoord_ct - dx / 2 75 | ycoord_lu = ycoord_ct + dy / 2 76 | 77 | xcoord_ld = xcoord_ct - dx / 2 78 | ycoord_ld = ycoord_ct - dy / 2 79 | 80 | xcoord_ru = xcoord_ct + dx / 2 81 | ycoord_ru = ycoord_ct + dy / 2 82 | 83 | xcoord_rd = xcoord_ct + dx / 2 84 | ycoord_rd = ycoord_ct - dy / 2 85 | 86 | dist_ct = ti.sqrt((xcoord_ct - cx) ** 2 + (ycoord_ct - cy) ** 2) 87 | dist_lu = ti.sqrt((xcoord_lu - cx) ** 2 + (ycoord_lu - cy) ** 2) 88 | dist_ld = ti.sqrt((xcoord_ld - cx) ** 2 + (ycoord_ld - cy) ** 2) 89 | dist_ru = ti.sqrt((xcoord_ru - cx) ** 2 + (ycoord_ru - cy) ** 2) 90 | dist_rd = ti.sqrt((xcoord_rd - cx) ** 2 + (ycoord_rd - cy) ** 2) 91 | 92 | if dist_lu > r and dist_ld > r and dist_ru > r and dist_rd > r: 93 | a = 1.0 94 | elif dist_lu < r and dist_ld < r and dist_ru < r and dist_rd < r: 95 | a = 0.0 96 | else: 97 | a = 0.5 + 0.5 * (dist_ct - r) / (ti.sqrt(2.0) * dx) 98 | a = var(a, 0, 1) 99 | 100 | return a 101 | 102 | 103 | @ti.kernel 104 | def set_init_F(): 105 | # Sets the initial volume fraction 106 | for i, j in ti.ndrange(imax+2, jmax+2): 107 | F[0, i, j] = 1.0 108 | ''' 109 | # Dambreak 110 | # The initial volume fraction of the domain 111 | x1 = 0.0 112 | x2 = Lx / 2 113 | y1 = 0.0 114 | y2 = Ly / 3 115 | for i, j in F: # [0,33], [0,33] 116 | if (xm[i] >= x1) and (xm[i] <= x2) and (ym[j] >= y1) and (ym[j] <= y2): 117 | F[i, j] = 1.0 118 | Fn[i, j] = F[i, j] 119 | 120 | # Moving square 121 | for i, j in F: 122 | x = xm[i] 123 | y = ym[j] 124 | cx, cy = 0.05, 0.02 125 | l = 0.01 126 | if ( ti.abs(x - cx) < l) and ( ti.abs(y - cy) < l): 127 | F[i, j] = 0.0 128 | Fn[i, j] = 0.0 129 | ''' 130 | # Moving circle 131 | # for i, j in F: 132 | for i, j in ti.ndrange(imax+2, jmax+2): 133 | x = xm[i] 134 | y = ym[j] 135 | cx, cy = Lx / 2, Ly * 3 / 4 136 | r = Lx / 10 137 | F[0, i, j] = find_area(i, j, cx, cy, r) 138 | # Fn[i, j] = find_area(i, j, cx, cy, r) 139 | 140 | ''' 141 | # Slot disk 142 | for i, j in ti.ndrange(imax + 2, jmax + 2): 143 | x = xm[i] 144 | y = ym[j] 145 | cx, cy = Lx * 3. / 4, Ly * 3 / 4 146 | r = Lx / 10 147 | F[0, i, j] = find_area(i, j, cx, cy, r) 148 | 149 | # Ftarget[i, j] = find_area(i, j, cx, cy, r) 150 | sw = r / 6.0 151 | sh = r * 0.8 152 | 153 | if ti.abs(x - cx) < sw and ti.abs(y - cy + r / 4) < sh: 154 | F[0, i, j] = 1.0 155 | # Ftarget[i, j] = 1.0 156 | # Fn[i, j] = 1.0 157 | 158 | 159 | for i, j in ti.ndrange(imax + 2, jmax + 2): 160 | ix = i // 11 161 | jx = j // 11 162 | idx = ix + jx 163 | if idx % 2 != 0: 164 | F[0, i, j] = 1.0 165 | else: 166 | F[0, i, j] = 0.0 167 | ''' 168 | for i, j in ti.ndrange(imax + 2, jmax + 2): 169 | x = xm[i] 170 | y = ym[j] 171 | cx, cy = Lx / 2., Ly * 3. / 4 172 | r = Lx / 10 173 | Ftarget[i, j] = find_area(i, j, cx, cy, r) 174 | # F[0, i, j] = find_area(i, j, cx, cy, r) 175 | 176 | 177 | @ti.kernel 178 | def init_uv(): 179 | ''' 180 | # Simple translation 181 | for I in ti.grouped(u): 182 | u[I] = Lx / nx / dt 183 | v[I] = 0. 184 | 185 | # Zalesak's slot disk 186 | w = 3.0 187 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 188 | ux = xm[i] - dx / 2 189 | uy = ym[j] 190 | vx = xm[i] 191 | vy = ym[j] - dy / 2 192 | u[i, j] = - w * (uy - Ly / 2) 193 | v[i, j] = w * (vx - Lx / 2) 194 | ''' 195 | # Kother Rider test 196 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 197 | ux = xm[i] - dx / 2 198 | uy = ym[j] 199 | vx = xm[i] 200 | vy = ym[j] - dy / 2 201 | # u[i, j] = ti.cos(ux) * ti.sin(uy) 202 | # v[i, j] = - ti.sin(vx) * ti.cos(vy) 203 | u[i, j] = - ti.sin(ux) ** 2 * ti.sin(2 * uy) * (Lx*1.0/dt/tmax * 2) #(Lx/nx/dt) 204 | v[i, j] = ti.sin(vy) ** 2 * ti.sin(2 * vx) * (Lx*1.0/dt/tmax * 2) # (Lx/nx/dt) 205 | 206 | # Boundary conditions 207 | for i in ti.ndrange(imax + 2): 208 | # bottom: slip 209 | u[i, jmin - 1] = u[i, jmin] 210 | v[i, jmin] = v[i, jmin + 1] 211 | # top: open 212 | u[i, jmax + 1] = u[i, jmax] 213 | v[i, jmax + 1] = v[i, jmax] 214 | for j in ti.ndrange(jmax + 2): 215 | # left: slip 216 | u[imin, j] = u[imin + 1, j] 217 | v[imin - 1, j] = v[imin, j] 218 | # right: slip 219 | u[imax + 1, j] = u[imax, j] 220 | v[imax + 1, j] = v[imax, j] 221 | 222 | 223 | @ti.kernel 224 | def set_BC(t:ti.i32, target:ti.template()): 225 | for i in ti.ndrange(imax + 2): 226 | # bottom: slip 227 | target[t, i, jmin - 1] = F[t, i, jmin] 228 | # top: open 229 | target[t, i, jmax + 1] = F[t, i, jmax] 230 | for j in ti.ndrange(jmax + 2): 231 | # left: slip 232 | target[t, imin - 1, j] = F[t, imin, j] 233 | # right: slip 234 | target[t, imax + 1, j] = F[t, imax, j] 235 | 236 | 237 | @ti.func 238 | def var(a, b, c): 239 | # Find the median of a,b, and c 240 | center = a + b + c - ti.max(a, b, c) - ti.min(a, b, c) 241 | return center 242 | 243 | 244 | @ti.kernel 245 | def solve_VOF_upwind(t:ti.i32): 246 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 247 | fl = u[i, j] * dt * F[t, i - 1, j] if u[i, j] > 0 else u[i, j] * dt * F[t, i, j] 248 | fr = u[i + 1, j] * dt * F[t, i, j] if u[i + 1, j] > 0 else u[i + 1, j] * dt * F[t, i + 1, j] 249 | ft = v[i, j + 1] * dt * F[t, i, j] if v[i, j + 1] > 0 else v[i, j + 1] * dt * F[t, i, j + 1] 250 | fb = v[i, j] * dt * F[t, i, j - 1] if v[i, j] > 0 else v[i, j] * dt * F[t, i, j] 251 | F[t + 1, i, j] = F[t, i, j] + (fl - fr + fb - ft) * dy / (dx * dy) 252 | 253 | 254 | def solve_VOF_rudman(t, eps_value): 255 | # FCT Method described in Rudman's 1997 paper 256 | if t % 2 == 0: 257 | fct_y_sweep(t, 0, eps_value) 258 | set_BC(2 * t + 1, F) 259 | fct_x_sweep(t, 1, eps_value) 260 | set_BC(2 * t + 2, F) 261 | else: 262 | fct_x_sweep(t, 0, eps_value) 263 | set_BC(2 * t + 1, F) 264 | fct_y_sweep(t, 1, eps_value) 265 | set_BC(2 * t + 2, F) 266 | 267 | @ti.kernel 268 | def fct_x_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 269 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 270 | dv = dx * dy - dt * dy * (u[i + 1, j] - u[i, j]) 271 | fl_L = u[i, j] * dt * F[2 * t + offset, i - 1, j] if u[i, j] >= 0 else u[i, j] * dt * F[2 * t + offset, i, j] 272 | fr_L = u[i + 1, j] * dt * F[2 * t + offset, i, j] if u[i + 1, j] >= 0 else u[i + 1, j] * dt * F[2 * t + offset, i + 1, j] 273 | Ftd_x[t, i, j] = F[2 * t + offset, i, j] + (fl_L - fr_L) * dy / (dx * dy) * dx * dy / dv 274 | 275 | for i, j in ti.ndrange((imin, imax + 2), (jmin, jmax + 1)): 276 | fl_L = u[i, j] * dt * F[2 * t + offset, i - 1, j] if u[i, j] >= 0 else u[i, j] * dt * F[2 * t + offset, i, j] 277 | fl_H = u[i, j] * dt * F[2 * t + offset, i - 1, j] if u[i, j] <= 0 else u[i, j] * dt * F[2 * t + offset, i, j] 278 | ax[t, i, j] = fl_H - fl_L 279 | 280 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 281 | fmax = ti.max(Ftd_x[t, i, j], Ftd_x[t, i - 1, j], Ftd_x[t, i + 1, j]) 282 | fmin = ti.min(Ftd_x[t, i, j], Ftd_x[t, i - 1, j], Ftd_x[t, i + 1, j]) 283 | 284 | pp = ti.max(0, ax[t, i, j]) - ti.min(0, ax[t, i + 1, j]) 285 | qp = (fmax - Ftd_x[t, i, j]) * dx 286 | if pp > 0: 287 | rp_x[t, i, j] = ti.min(1, qp / (pp + eps)) 288 | else: 289 | rp_x[t, i, j] = 0.0 290 | 291 | pm = ti.max(0, ax[t, i + 1, j]) - ti.min(0, ax[t, i, j]) 292 | qm = (Ftd_x[t, i, j] - fmin) * dx 293 | if pm > 0: 294 | rm_x[t, i, j] = ti.min(1, qm / (pm + eps)) 295 | else: 296 | rm_x[t, i, j] = 0.0 297 | 298 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 299 | if ax[t, i + 1, j] >= 0: 300 | cx[t, i + 1, j] = ti.min(rp_x[t, i + 1, j], rm_x[t, i, j]) 301 | else: 302 | cx[t, i + 1, j] = ti.min(rp_x[t, i, j], rm_x[t, i + 1, j]) 303 | 304 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 305 | dv = dx * dy - dt * dy * (u[i + 1, j] - u[i, j]) 306 | F[2 * t + offset + 1, i, j] = Ftd_x[t, i, j] - ((ax[t, i + 1, j] * cx[t, i + 1, j] - \ 307 | ax[t, i, j] * cx[t, i, j]) / dy) * dx * dy / dv 308 | 309 | 310 | @ti.kernel 311 | def fct_y_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 312 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 313 | dv = dx * dy - dt * dx * (v[i, j + 1] - v[i, j]) 314 | ft_L = v[i, j + 1] * dt * F[2 * t + offset, i, j] if v[i, j + 1] >= 0 else v[i, j + 1] * dt * F[2 * t + offset, i, j + 1] 315 | fb_L = v[i, j] * dt * F[2 * t + offset, i, j - 1] if v[i, j] >= 0 else v[i, j] * dt * F[2 * t + offset, i, j] 316 | Ftd_y[t, i, j] = F[2 * t + offset, i, j] + (fb_L - ft_L) * dy / (dx * dy) * dx * dy / dv 317 | 318 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 2)): 319 | fb_L = v[i, j] * dt * F[2 * t + offset, i, j - 1] if v[i, j] >= 0 else v[i, j] * dt * F[2 * t + offset, i, j] 320 | fb_H = v[i, j] * dt * F[2 * t + offset, i, j - 1] if v[i, j] <= 0 else v[i, j] * dt * F[2 * t + offset, i, j] 321 | ay[t, i, j] = fb_H - fb_L 322 | 323 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 324 | fmax = ti.max(Ftd_y[t, i, j], Ftd_y[t, i, j - 1], Ftd_y[t, i, j + 1]) 325 | fmin = ti.min(Ftd_y[t, i, j], Ftd_y[t, i, j - 1], Ftd_y[t, i, j + 1]) 326 | 327 | # eps = 1e-4 328 | pp = ti.max(0, ay[t, i, j]) - ti.min(0, ay[t, i, j + 1]) 329 | qp = (fmax - Ftd_y[t, i, j]) * dx 330 | if pp > 0: 331 | rp_y[t, i, j] = ti.min(1, qp / (pp + eps)) 332 | else: 333 | rp_y[t, i, j] = 0.0 334 | 335 | pm = ti.max(0, ay[t, i, j + 1]) - ti.min(0, ay[t, i, j]) 336 | qm = (Ftd_y[t, i, j] - fmin) * dx 337 | if pm > 0: 338 | rm_y[t, i, j] = ti.min(1, qm / (pm + eps)) 339 | else: 340 | rm_y[t, i, j] = 0.0 341 | 342 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 343 | if ay[t, i, j + 1] >= 0: 344 | cy[t, i, j + 1] = ti.min(rp_y[t, i, j + 1], rm_y[t, i, j]) 345 | else: 346 | cy[t, i, j + 1] = ti.min(rp_y[t, i, j], rm_y[t, i, j + 1]) 347 | 348 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 349 | dv = dx * dy - dt * dx * (v[i, j + 1] - v[i, j]) 350 | F[2 * t + offset + 1, i, j] = Ftd_y[t, i, j] - ((ay[t, i, j + 1] * cy[t, i, j + 1] -\ 351 | ay[t, i, j] * cy[t, i, j]) / dy) * dx * dy / dv 352 | 353 | 354 | @ti.kernel 355 | def get_vof(t:ti.i32): 356 | # Copy the current field to display 357 | for i, j in ti.ndrange(imax + 2, jmax + 2): 358 | Fd[i, j] = F[t, i, j] 359 | 360 | 361 | def forward(eps_value): 362 | nstep = 5 363 | plot_contour = gui.contour 364 | for istep in range(tmax): 365 | # solve_VOF_upwind(istep) # Upwind scheme; F(2 * istep) -> F(2 * istep + 2) 366 | solve_VOF_rudman(istep, eps_value) 367 | # set_BC(istep) 368 | get_vof(2 * istep + 2) # Get F at (istep + 1) 369 | if (istep % nstep) == 0: # Output data every steps 370 | print(f'>>> Current step: {istep}') 371 | plot_contour(Fd) 372 | gui.show(f'./output/forward-step-{istep:05d}.png') 373 | 374 | 375 | # Start Main-loop 376 | grid_staggered() 377 | set_init_F() 378 | init_uv() 379 | set_BC(0, F) 380 | gui = ti.GUI('FCT test', res=(400,400)) 381 | os.makedirs('output', exist_ok=True) # Make dir for output 382 | forward(eps_value=1.0e-4) 383 | 384 | -------------------------------------------------------------------------------- /test/diff_fct.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | from math import pi as pi 5 | import os 6 | 7 | ti.init(arch=ti.cpu, default_fp=ti.f32) 8 | 9 | SAVE_FIG = True 10 | 11 | nx = 500 # Number of grid points in the x direction 12 | ny = 500 # Number of grid points in the y direction 13 | 14 | Lx = pi # The length of the domain 15 | Ly = pi # The width of the domain 16 | 17 | # Solution parameters 18 | dt = 1e-4 # Use smaller dt for higher density ratio 19 | imin = 1 20 | imax = imin + nx - 1 21 | jmin = 1 22 | jmax = jmin + ny - 1 23 | tmax = 200 24 | 25 | loss = ti.field(float, shape=(), needs_grad=True) 26 | learning_rate = ti.field(float, shape=()) 27 | learning_rate[None] = 0.1 28 | 29 | F = ti.field(float, shape=(2 * tmax + 1, imax + 2, jmax + 2), needs_grad=True) 30 | Ftd_x = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 31 | Ftd_y = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 32 | Ftarget = ti.field(float, shape=(imax + 2, jmax + 2)) 33 | Fd = ti.field(float, shape=(3*(imax + 2), jmax + 2)) 34 | 35 | ax = ti.field(float, shape=(tmax, imax + 2, jmax + 2),needs_grad=True) 36 | ay = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 37 | cx = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 38 | cy = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 39 | rp_x = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 40 | rm_x = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 41 | rp_y = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 42 | rm_y = ti.field(float, shape=(tmax, imax + 2, jmax + 2), needs_grad=True) 43 | 44 | u = ti.field(float, shape=(imax + 2, jmax + 2)) 45 | v = ti.field(float, shape=(imax + 2, jmax + 2)) 46 | 47 | x = ti.field(float, shape=imax + 3) 48 | y = ti.field(float, shape=jmax + 3) 49 | x.from_numpy(np.hstack((0.0, np.linspace(0, Lx, nx + 1), Lx))) # [0, 0, ... 1, 1] 50 | y.from_numpy(np.hstack((0.0, np.linspace(0, Ly, ny + 1), Ly))) # [0, 0, ... 1, 1] 51 | 52 | xm = ti.field(float, shape=imax + 2) 53 | ym = ti.field(float, shape=jmax + 2) 54 | dx = x[imin + 2] - x[imin + 1] 55 | dy = y[jmin + 2] - y[jmin + 1] 56 | dxi = 1 / dx 57 | dyi = 1 / dy 58 | 59 | print(f'>>> VOF scheme testing') 60 | print(f'>>> Grid resolution: {nx} x {ny}, dt = {dt:4.2e}') 61 | 62 | @ti.kernel 63 | def grid_staggered(): # 11/3 Checked 64 | ''' 65 | Calculate the center position of cells. 66 | ''' 67 | for i in xm: # xm[0] = 0.0, xm[33] = 1.0 68 | xm[i] = 0.5 * (x[i] + x[i + 1]) 69 | for j in ym: 70 | ym[j] = 0.5 * (y[j] + y[j + 1]) 71 | 72 | 73 | @ti.func 74 | def find_area(i, j, cx, cy, r): 75 | a = 0.0 76 | xcoord_ct = (i - imin) * dx + dx / 2 77 | ycoord_ct = (j - jmin) * dy + dy / 2 78 | 79 | xcoord_lu = xcoord_ct - dx / 2 80 | ycoord_lu = ycoord_ct + dy / 2 81 | 82 | xcoord_ld = xcoord_ct - dx / 2 83 | ycoord_ld = ycoord_ct - dy / 2 84 | 85 | xcoord_ru = xcoord_ct + dx / 2 86 | ycoord_ru = ycoord_ct + dy / 2 87 | 88 | xcoord_rd = xcoord_ct + dx / 2 89 | ycoord_rd = ycoord_ct - dy / 2 90 | 91 | dist_ct = ti.sqrt((xcoord_ct - cx) ** 2 + (ycoord_ct - cy) ** 2) 92 | dist_lu = ti.sqrt((xcoord_lu - cx) ** 2 + (ycoord_lu - cy) ** 2) 93 | dist_ld = ti.sqrt((xcoord_ld - cx) ** 2 + (ycoord_ld - cy) ** 2) 94 | dist_ru = ti.sqrt((xcoord_ru - cx) ** 2 + (ycoord_ru - cy) ** 2) 95 | dist_rd = ti.sqrt((xcoord_rd - cx) ** 2 + (ycoord_rd - cy) ** 2) 96 | 97 | if dist_lu > r and dist_ld > r and dist_ru > r and dist_rd > r: 98 | a = 1.0 99 | elif dist_lu < r and dist_ld < r and dist_ru < r and dist_rd < r: 100 | a = 0.0 101 | else: 102 | a = 0.5 + 0.5 * (dist_ct - r) / (ti.sqrt(2.0) * dx) 103 | a = var(a, 0, 1) 104 | 105 | return a 106 | 107 | 108 | @ti.kernel 109 | def set_init_F(): 110 | # Sets the initial volume fraction 111 | for i, j in ti.ndrange(imax+2, jmax+2): 112 | F[0, i, j] = 1.0 113 | ''' 114 | # Dambreak 115 | # The initial volume fraction of the domain 116 | x1 = 0.0 117 | x2 = Lx / 2 118 | y1 = 0.0 119 | y2 = Ly / 3 120 | for i, j in F: # [0,33], [0,33] 121 | if (xm[i] >= x1) and (xm[i] <= x2) and (ym[j] >= y1) and (ym[j] <= y2): 122 | F[i, j] = 1.0 123 | Fn[i, j] = F[i, j] 124 | 125 | # Moving square 126 | for i, j in F: 127 | x = xm[i] 128 | y = ym[j] 129 | cx, cy = 0.05, 0.02 130 | l = 0.01 131 | if ( ti.abs(x - cx) < l) and ( ti.abs(y - cy) < l): 132 | F[i, j] = 0.0 133 | Fn[i, j] = 0.0 134 | 135 | # Moving circle 136 | for i, j in F: 137 | x = xm[i] 138 | y = ym[j] 139 | cx, cy = Lx / 2, Ly * 3 / 4 140 | r = Lx / 10 141 | F[i, j] = find_area(i, j, cx, cy, r) 142 | Fn[i, j] = find_area(i, j, cx, cy, r) 143 | 144 | 145 | # Slot disk 146 | for i, j in ti.ndrange(imax + 2, jmax + 2): 147 | x = xm[i] 148 | y = ym[j] 149 | cx, cy = Lx * 3. / 4, Ly * 3 / 4 150 | r = Lx / 10 151 | F[0, i, j] = find_area(i, j, cx, cy, r) 152 | 153 | # Ftarget[i, j] = find_area(i, j, cx, cy, r) 154 | sw = r / 6.0 155 | sh = r * 0.8 156 | 157 | if ti.abs(x - cx) < sw and ti.abs(y - cy + r / 4) < sh: 158 | F[0, i, j] = 1.0 159 | # Ftarget[i, j] = 1.0 160 | # Fn[i, j] = 1.0 161 | 162 | 163 | for i, j in ti.ndrange(imax + 2, jmax + 2): 164 | ix = i // 11 165 | jx = j // 11 166 | idx = ix + jx 167 | if idx % 2 != 0: 168 | F[0, i, j] = 1.0 169 | else: 170 | F[0, i, j] = 0.0 171 | ''' 172 | for i, j in ti.ndrange(imax + 2, jmax + 2): 173 | x = xm[i] 174 | y = ym[j] 175 | cx, cy = Lx / 2., Ly * 3. / 4 176 | r = Lx / 10 177 | Ftarget[i, j] = find_area(i, j, cx, cy, r) 178 | # F[0, i, j] = find_area(i, j, cx, cy, r) 179 | 180 | 181 | @ti.kernel 182 | def init_uv(): 183 | ''' 184 | # Simple translation 185 | for I in ti.grouped(u): 186 | u[I] = Lx / nx / dt 187 | v[I] = 0. 188 | 189 | # Zalesak's slot disk 190 | w = 3.0 191 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 192 | ux = xm[i] - dx / 2 193 | uy = ym[j] 194 | vx = xm[i] 195 | vy = ym[j] - dy / 2 196 | u[i, j] = - w * (uy - Ly / 2) 197 | v[i, j] = w * (vx - Lx / 2) 198 | ''' 199 | # Kother Rider test 200 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 201 | ux = xm[i] - dx / 2 202 | uy = ym[j] 203 | vx = xm[i] 204 | vy = ym[j] - dy / 2 205 | # u[i, j] = ti.cos(ux) * ti.sin(uy) 206 | # v[i, j] = - ti.sin(vx) * ti.cos(vy) 207 | u[i, j] = - ti.sin(ux) ** 2 * ti.sin(2 * uy) * (Lx*1.0/dt/tmax/3) #(Lx/nx/dt) 208 | v[i, j] = ti.sin(vy) ** 2 * ti.sin(2 * vx) * (Lx*1.0/dt/tmax/3) # (Lx/nx/dt) 209 | 210 | # Boundary conditions 211 | for i in ti.ndrange(imax + 2): 212 | # bottom: slip 213 | u[i, jmin - 1] = u[i, jmin] 214 | v[i, jmin] = v[i, jmin + 1] 215 | # top: open 216 | u[i, jmax + 1] = u[i, jmax] 217 | v[i, jmax + 1] = v[i, jmax] 218 | for j in ti.ndrange(jmax + 2): 219 | # left: slip 220 | u[imin, j] = u[imin + 1, j] 221 | v[imin - 1, j] = v[imin, j] 222 | # right: slip 223 | u[imax + 1, j] = u[imax, j] 224 | v[imax + 1, j] = v[imax, j] 225 | 226 | 227 | @ti.ad.no_grad 228 | @ti.kernel 229 | def set_BC(t:ti.i32, target:ti.template()): 230 | for i in ti.ndrange(imax + 2): 231 | # bottom: slip 232 | target[t, i, jmin - 1] = F[t, i, jmin] 233 | # top: open 234 | target[t, i, jmax + 1] = F[t, i, jmax] 235 | for j in ti.ndrange(jmax + 2): 236 | # left: slip 237 | target[t, imin - 1, j] = F[t, imin, j] 238 | # right: slip 239 | target[t, imax + 1, j] = F[t, imax, j] 240 | 241 | 242 | @ti.func 243 | def var(a, b, c): 244 | # Find the median of a,b, and c 245 | center = a + b + c - ti.max(a, b, c) - ti.min(a, b, c) 246 | return center 247 | 248 | 249 | @ti.kernel 250 | def solve_VOF_upwind(t:ti.i32): 251 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 252 | fl = u[i, j] * dt * F[t, i - 1, j] if u[i, j] > 0 else u[i, j] * dt * F[t, i, j] 253 | fr = u[i + 1, j] * dt * F[t, i, j] if u[i + 1, j] > 0 else u[i + 1, j] * dt * F[t, i + 1, j] 254 | ft = v[i, j + 1] * dt * F[t, i, j] if v[i, j + 1] > 0 else v[i, j + 1] * dt * F[t, i, j + 1] 255 | fb = v[i, j] * dt * F[t, i, j - 1] if v[i, j] > 0 else v[i, j] * dt * F[t, i, j] 256 | F[t + 1, i, j] = F[t, i, j] + (fl - fr + fb - ft) * dy / (dx * dy) 257 | 258 | 259 | def solve_VOF_rudman(t, eps_value): 260 | # FCT Method described in Rudman's 1997 paper 261 | if t % 2 == 0: 262 | fct_y_sweep(t, 0, eps_value) 263 | set_BC(2 * t + 1, F) 264 | fct_x_sweep(t, 1, eps_value) 265 | set_BC(2 * t + 2, F) 266 | else: 267 | fct_x_sweep(t, 0, eps_value) 268 | set_BC(2 * t + 1, F) 269 | fct_y_sweep(t, 1, eps_value) 270 | set_BC(2 * t + 2, F) 271 | 272 | @ti.kernel 273 | def fct_x_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 274 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 275 | dv = dx * dy - dt * dy * (u[i + 1, j] - u[i, j]) 276 | fl_L = u[i, j] * dt * F[2 * t + offset, i - 1, j] if u[i, j] >= 0 else u[i, j] * dt * F[2 * t + offset, i, j] 277 | fr_L = u[i + 1, j] * dt * F[2 * t + offset, i, j] if u[i + 1, j] >= 0 else u[i + 1, j] * dt * F[2 * t + offset, i + 1, j] 278 | Ftd_x[t, i, j] = F[2 * t + offset, i, j] + (fl_L - fr_L) * dy / (dx * dy) * dx * dy / dv 279 | 280 | for i, j in ti.ndrange((imin, imax + 2), (jmin, jmax + 1)): 281 | fl_L = u[i, j] * dt * F[2 * t + offset, i - 1, j] if u[i, j] >= 0 else u[i, j] * dt * F[2 * t + offset, i, j] 282 | fl_H = u[i, j] * dt * F[2 * t + offset, i - 1, j] if u[i, j] <= 0 else u[i, j] * dt * F[2 * t + offset, i, j] 283 | ax[t, i, j] = fl_H - fl_L 284 | 285 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 286 | fmax = ti.max(Ftd_x[t, i, j], Ftd_x[t, i - 1, j], Ftd_x[t, i + 1, j]) 287 | fmin = ti.min(Ftd_x[t, i, j], Ftd_x[t, i - 1, j], Ftd_x[t, i + 1, j]) 288 | 289 | pp = ti.max(0, ax[t, i, j]) - ti.min(0, ax[t, i + 1, j]) 290 | qp = (fmax - Ftd_x[t, i, j]) * dx 291 | if pp > eps: 292 | rp_x[t, i, j] = ti.min(1, qp / pp) 293 | else: 294 | rp_x[t, i, j] = 0.0 295 | 296 | pm = ti.max(0, ax[t, i + 1, j]) - ti.min(0, ax[t, i, j]) 297 | qm = (Ftd_x[t, i, j] - fmin) * dx 298 | if pm > eps: 299 | rm_x[t, i, j] = ti.min(1, qm / pm) 300 | else: 301 | rm_x[t, i, j] = 0.0 302 | 303 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 304 | if ax[t, i + 1, j] >= 0: 305 | cx[t, i + 1, j] = ti.min(rp_x[t, i + 1, j], rm_x[t, i, j]) 306 | else: 307 | cx[t, i + 1, j] = ti.min(rp_x[t, i, j], rm_x[t, i + 1, j]) 308 | 309 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 310 | dv = dx * dy - dt * dy * (u[i + 1, j] - u[i, j]) 311 | F[2 * t + offset + 1, i, j] = Ftd_x[t, i, j] - ((ax[t, i + 1, j] * cx[t, i + 1, j] - \ 312 | ax[t, i, j] * cx[t, i, j]) / dy) * dx * dy / dv 313 | 314 | 315 | @ti.kernel 316 | def fct_y_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 317 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 318 | dv = dx * dy - dt * dx * (v[i, j + 1] - v[i, j]) 319 | ft_L = v[i, j + 1] * dt * F[2 * t + offset, i, j] if v[i, j + 1] >= 0 else v[i, j + 1] * dt * F[2 * t + offset, i, j + 1] 320 | fb_L = v[i, j] * dt * F[2 * t + offset, i, j - 1] if v[i, j] >= 0 else v[i, j] * dt * F[2 * t + offset, i, j] 321 | Ftd_y[t, i, j] = F[2 * t + offset, i, j] + (fb_L - ft_L) * dy / (dx * dy) * dx * dy / dv 322 | 323 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 2)): 324 | fb_L = v[i, j] * dt * F[2 * t + offset, i, j - 1] if v[i, j] >= 0 else v[i, j] * dt * F[2 * t + offset, i, j] 325 | fb_H = v[i, j] * dt * F[2 * t + offset, i, j - 1] if v[i, j] <= 0 else v[i, j] * dt * F[2 * t + offset, i, j] 326 | ay[t, i, j] = fb_H - fb_L 327 | 328 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 329 | fmax = ti.max(Ftd_y[t, i, j], Ftd_y[t, i, j - 1], Ftd_y[t, i, j + 1]) 330 | fmin = ti.min(Ftd_y[t, i, j], Ftd_y[t, i, j - 1], Ftd_y[t, i, j + 1]) 331 | 332 | # eps = 1e-4 333 | pp = ti.max(0, ay[t, i, j]) - ti.min(0, ay[t, i, j + 1]) 334 | qp = (fmax - Ftd_y[t, i, j]) * dx 335 | if pp > eps: 336 | rp_y[t, i, j] = ti.min(1, qp / pp) 337 | else: 338 | rp_y[t, i, j] = 0.0 339 | 340 | pm = ti.max(0, ay[t, i, j + 1]) - ti.min(0, ay[t, i, j]) 341 | qm = (Ftd_y[t, i, j] - fmin) * dx 342 | if pm > eps: 343 | rm_y[t, i, j] = ti.min(1, qm / pm) 344 | else: 345 | rm_y[t, i, j] = 0.0 346 | 347 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 348 | if ay[t, i, j + 1] >= 0: 349 | cy[t, i, j + 1] = ti.min(rp_y[t, i, j + 1], rm_y[t, i, j]) 350 | else: 351 | cy[t, i, j + 1] = ti.min(rp_y[t, i, j], rm_y[t, i, j + 1]) 352 | 353 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 354 | dv = dx * dy - dt * dx * (v[i, j + 1] - v[i, j]) 355 | F[2 * t + offset + 1, i, j] = Ftd_y[t, i, j] - ((ay[t, i, j + 1] * cy[t, i, j + 1] -\ 356 | ay[t, i, j] * cy[t, i, j]) / dy) * dx * dy / dv 357 | 358 | 359 | @ti.ad.no_grad 360 | @ti.kernel 361 | def get_vof(t:ti.i32): 362 | # Copy the current field to display 363 | for i, j in ti.ndrange(imax + 2, jmax + 2): 364 | Fd[i, j] = F[t, i, j] 365 | # Copy the target field to display 366 | for i, j in ti.ndrange(imax + 2, jmax + 2): 367 | Fd[i + imax + 2, j] = Ftarget[i, j] 368 | 369 | 370 | @ti.ad.no_grad 371 | @ti.kernel 372 | def get_vof_grad(): 373 | # Copy the target field to display 374 | for i, j in ti.ndrange(imax + 2, jmax + 2): 375 | Fd[i + 2 * imax + 4, j] = F.grad[0, i, j] 376 | 377 | 378 | @ti.kernel 379 | def compute_loss(): 380 | for i, j in ti.ndrange((imin, imax + 2), (jmin, jmax + 2)): 381 | loss[None] += ti.abs(Ftarget[i, j] - F[2 * tmax, i, j]) # maybe ti.abs is required 382 | 383 | 384 | @ti.kernel 385 | def apply_grad(): 386 | for i, j in ti.ndrange(imax + 2, jmax + 2): 387 | # if ti.abs(F.grad[0, i, j]) < 2.: 388 | F[0, i, j] -= learning_rate[None] * F.grad[0, i, j] 389 | F[0, i, j] = var(0, 1, F[0, i, j]) 390 | 391 | 392 | def forward(eps_value): 393 | nstep = 5 394 | plot_contour = ti.ad.no_grad(gui.contour) 395 | for istep in range(tmax): 396 | # solve_VOF_upwind(istep) # Upwind scheme; F(2 * istep) -> F(2 * istep + 2) 397 | solve_VOF_rudman(istep, eps_value) 398 | # set_BC(istep) 399 | get_vof(2 * istep + 2) # Get F at (istep + 1) 400 | if (istep % nstep) == 0: # Output data every steps 401 | print(f'>>> Current step: {istep}') 402 | plot_contour(Fd) 403 | gui.show(f'./output/opt-{opt:03d}-step-{istep:05d}.png') 404 | compute_loss() 405 | print(f'>>> >>> Current loss: {loss[None]}') 406 | 407 | 408 | # Start Main-loop 409 | grid_staggered() 410 | set_init_F() 411 | init_uv() 412 | set_BC(0, F) 413 | gui = ti.GUI('FCT test', res=(1200,400)) 414 | os.makedirs('output', exist_ok=True) # Make dir for output 415 | for opt in range(1000): 416 | with ti.ad.Tape(loss): 417 | forward(eps_value=5e-4) 418 | get_vof_grad() 419 | apply_grad() 420 | -------------------------------------------------------------------------------- /2dvof.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import matplotlib.cm as cm 5 | import argparse 6 | import os 7 | import flow_visualization as fv 8 | 9 | ti.init(arch=ti.cpu, default_fp=ti.f32) # Set default fp so that float=ti.f32 10 | 11 | parser = argparse.ArgumentParser() # Get the initial condition 12 | # 1 - Dam Break; 2 - Rising Bubble; 3 - Droping liquid 13 | parser.add_argument('-ic', type=int, choices=[1, 2, 3], default=1) 14 | parser.add_argument('-s', action='store_true') 15 | args = parser.parse_args() 16 | initial_condition = args.ic 17 | SAVE_FIG = args.s 18 | 19 | nx = 200 # Number of grid points in the x direction 20 | ny = 200 # Number of grid points in the y direction 21 | 22 | Lx = 0.1 # The length of the domain 23 | Ly = 0.1 # The width of the domain 24 | rho_l = 1000.0 25 | rho_g = 50.0 26 | nu_l = 1.0e-6 # kinematic viscosity, nu = mu / rho 27 | nu_g = 1.5e-5 28 | sigma = ti.field(dtype=float, shape=()) 29 | sigma[None] = 0.007 30 | gx = 0 31 | gy = -5 32 | 33 | dt = 4e-6 # Use smaller dt for higher density ratio 34 | eps = 1e-6 # Threshold used in vfconv and f post processings 35 | 36 | # Mesh information 37 | imin = 1 38 | imax = imin + nx - 1 39 | jmin = 1 40 | jmax = jmin + ny - 1 41 | x = ti.field(float, shape=imax + 3) 42 | y = ti.field(float, shape=jmax + 3) 43 | xnp = np.hstack((0.0, np.linspace(0, Lx, nx + 1), Lx)).astype(np.float32) # [0, 0, ... 1, 1] 44 | x.from_numpy(xnp) 45 | ynp = np.hstack((0.0, np.linspace(0, Ly, ny + 1), Ly)).astype(np.float32) # [0, 0, ... 1, 1] 46 | y.from_numpy(ynp) 47 | dx = x[imin + 2] - x[imin + 1] 48 | dy = y[jmin + 2] - y[jmin + 1] 49 | dxi = 1 / dx 50 | dyi = 1 / dy 51 | 52 | # Variables for VOF function 53 | F = ti.field(float, shape=(imax + 2, jmax + 2)) 54 | Ftd = ti.field(float, shape=(imax + 2, jmax + 2)) 55 | ax = ti.field(float, shape=(imax + 2, jmax + 2)) 56 | ay = ti.field(float, shape=(imax + 2, jmax + 2)) 57 | cx = ti.field(float, shape=(imax + 2, jmax + 2)) 58 | cy = ti.field(float, shape=(imax + 2, jmax + 2)) 59 | rp = ti.field(float, shape=(imax + 2, jmax + 2)) 60 | rm = ti.field(float, shape=(imax + 2, jmax + 2)) 61 | 62 | # Variables for N-S equation 63 | u = ti.field(float, shape=(imax + 2, jmax + 2)) 64 | v = ti.field(float, shape=(imax + 2, jmax + 2)) 65 | u_star = ti.field(float, shape=(imax + 2, jmax + 2)) 66 | v_star = ti.field(float, shape=(imax + 2, jmax + 2)) 67 | p = ti.field(float, shape=(imax + 2, jmax + 2)) 68 | pt = ti.field(float, shape=(imax + 2, jmax + 2)) 69 | Ap = ti.field(float, shape=(imax + 2, jmax + 2)) 70 | rhs = ti.field(float, shape=(imax + 2, jmax + 2)) 71 | rho = ti.field(float, shape=(imax + 2, jmax + 2)) 72 | nu = ti.field(float, shape=(imax + 2, jmax + 2)) 73 | V = ti.Vector.field(2, dtype=float, shape=(imax + 2, jmax + 2)) 74 | 75 | # Variables for interface reconstruction 76 | mx1 = ti.field(float, shape=(imax + 2, jmax + 2)) 77 | my1 = ti.field(float, shape=(imax + 2, jmax + 2)) 78 | mx2 = ti.field(float, shape=(imax + 2, jmax + 2)) 79 | my2 = ti.field(float, shape=(imax + 2, jmax + 2)) 80 | mx3 = ti.field(float, shape=(imax + 2, jmax + 2)) 81 | my3 = ti.field(float, shape=(imax + 2, jmax + 2)) 82 | mx4 = ti.field(float, shape=(imax + 2, jmax + 2)) 83 | my4 = ti.field(float, shape=(imax + 2, jmax + 2)) 84 | mxsum = ti.field(float, shape=(imax + 2, jmax + 2)) 85 | mysum = ti.field(float, shape=(imax + 2, jmax + 2)) 86 | mx = ti.field(float, shape=(imax+2, jmax+2)) 87 | my = ti.field(float, shape=(imax+2, jmax+2)) 88 | kappa = ti.field(float, shape=(imax + 2, jmax + 2)) # interface curvature 89 | magnitude = ti.field(float, shape=(imax+2, jmax+2)) 90 | 91 | # For visualization 92 | resolution = (nx * 2, ny * 2) 93 | rgb_buf = ti.field(dtype=float, shape=resolution) 94 | 95 | print(f'>>> A VOF solver written in Taichi; Press q to exit.') 96 | print(f'>>> Grid resolution: {nx} x {ny}, dt = {dt:4.2e}') 97 | print(f'>>> Density ratio: {rho_l / rho_g : 4.2f}, gravity : {gy : 4.2f}, sigma : {sigma[None] : 4.2f}') 98 | print(f'>>> Viscosity ratio: {nu_l / nu_g : 4.2f}') 99 | print(f'>>> Please wait a few seconds to let the kernels compile...') 100 | 101 | 102 | @ti.func 103 | def find_area(i, j, cx, cy, r): 104 | a = 0.0 105 | xcoord_ct = (i - imin) * dx + dx / 2 106 | ycoord_ct = (j - jmin) * dy + dy / 2 107 | 108 | xcoord_lu = xcoord_ct - dx / 2 109 | ycoord_lu = ycoord_ct + dy / 2 110 | 111 | xcoord_ld = xcoord_ct - dx / 2 112 | ycoord_ld = ycoord_ct - dy / 2 113 | 114 | xcoord_ru = xcoord_ct + dx / 2 115 | ycoord_ru = ycoord_ct + dy / 2 116 | 117 | xcoord_rd = xcoord_ct + dx / 2 118 | ycoord_rd = ycoord_ct - dy / 2 119 | 120 | dist_ct = ti.sqrt((xcoord_ct - cx) ** 2 + (ycoord_ct - cy) ** 2) 121 | dist_lu = ti.sqrt((xcoord_lu - cx) ** 2 + (ycoord_lu - cy) ** 2) 122 | dist_ld = ti.sqrt((xcoord_ld - cx) ** 2 + (ycoord_ld - cy) ** 2) 123 | dist_ru = ti.sqrt((xcoord_ru - cx) ** 2 + (ycoord_ru - cy) ** 2) 124 | dist_rd = ti.sqrt((xcoord_rd - cx) ** 2 + (ycoord_rd - cy) ** 2) 125 | 126 | if dist_lu > r and dist_ld > r and dist_ru > r and dist_rd > r: 127 | a = 1.0 128 | elif dist_lu < r and dist_ld < r and dist_ru < r and dist_rd < r: 129 | a = 0.0 130 | else: 131 | a = 0.5 + 0.5 * (dist_ct - r) / (ti.sqrt(2.0) * dx) 132 | a = var(a, 0, 1) 133 | 134 | return a 135 | 136 | 137 | @ti.kernel 138 | def set_init_F(ic:ti.i32): 139 | # Sets the initial volume fraction 140 | if ic == 1: # Dambreak 141 | x1 = 0.0 142 | x2 = Lx / 3 143 | y1 = 0.0 144 | y2 = Ly / 2 145 | for i, j in ti.ndrange(imax + 2, jmax + 2): 146 | if (x[i] >= x1) and (x[i] <= x2) and (y[j] >= y1) and (y[j] <= y2): 147 | F[i, j] = 1.0 148 | elif ic == 2: # Rising bubble 149 | for i, j in ti.ndrange(imax + 2, jmax + 2): 150 | r = Lx / 12 151 | cx, cy = Lx / 2, 2 * r 152 | F[i, j] = find_area(i, j, cx, cy, r) 153 | elif ic == 3: # Liquid drop 154 | for i, j in ti.ndrange(imax + 2, jmax + 2): 155 | r = Lx / 12 156 | cx, cy = Lx / 2, Ly - 3 * r 157 | F[i, j] = 1.0 - find_area(i, j, cx, cy, r) 158 | if y[j] < Ly * 0.37: 159 | F[i, j] = 1.0 160 | 161 | 162 | @ti.kernel 163 | def set_BC(): 164 | for i in ti.ndrange(imax + 2): 165 | # bottom: slip 166 | u[i, jmin - 1] = u[i, jmin] 167 | v[i, jmin] = 0 168 | F[i, jmin - 1] = F[i, jmin] 169 | p[i, jmin - 1] = p[i, jmin] 170 | rho[i, jmin - 1] = rho[i, jmin] 171 | # top: open 172 | u[i, jmax + 1] = u[i, jmax] 173 | v[i, jmax + 1] = 0 #v[i, jmax] 174 | F[i, jmax + 1] = F[i, jmax] 175 | p[i, jmax + 1] = p[i, jmax] 176 | rho[i, jmax + 1] = rho[i, jmax] 177 | for j in ti.ndrange(jmax + 2): 178 | # left: slip 179 | u[imin, j] = 0 180 | v[imin - 1, j] = v[imin, j] 181 | F[imin - 1, j] = F[imin, j] 182 | p[imin - 1, j] = p[imin, j] 183 | rho[imin - 1, j] = rho[imin, j] 184 | # right: slip 185 | u[imax + 1, j] = 0 186 | v[imax + 1, j] = v[imax, j] 187 | F[imax + 1, j] = F[imax, j] 188 | p[imax + 1, j] = p[imax, j] 189 | rho[imax + 1, j] = rho[imax, j] 190 | 191 | 192 | @ti.func 193 | def var(a, b, c): # Find the median of a,b, and c 194 | center = a + b + c - ti.max(a, b, c) - ti.min(a, b, c) 195 | return center 196 | 197 | 198 | @ti.kernel 199 | def cal_nu_rho(): 200 | for I in ti.grouped(rho): 201 | F = var(0.0, 1.0, F[I]) 202 | rho[I] = rho_g * (1 - F) + rho_l * F 203 | nu[I] = nu_l * F + nu_g * (1.0 - F) 204 | 205 | 206 | @ti.kernel 207 | def advect_upwind(): 208 | for i, j in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1)): 209 | v_here = 0.25 * (v[i - 1, j] + v[i - 1, j + 1] + v[i, j] + v[i, j + 1]) 210 | dudx = (u[i,j] - u[i-1,j]) * dxi if u[i,j] > 0 else (u[i+1,j]-u[i,j])*dxi 211 | dudy = (u[i,j] - u[i,j-1]) * dyi if v_here > 0 else (u[i,j+1]-u[i,j])*dyi 212 | kappa_ave = (kappa[i, j] + kappa[i - 1, j]) / 2.0 213 | fx_kappa = - sigma[None] * (F[i, j] - F[i - 1, j]) * kappa_ave / dx 214 | u_star[i, j] = ( 215 | u[i, j] + dt * 216 | (nu[i, j] * (u[i - 1, j] - 2 * u[i, j] + u[i + 1, j]) * dxi**2 217 | + nu[i, j] * (u[i, j - 1] - 2 * u[i, j] + u[i, j + 1]) * dyi**2 218 | - u[i, j] * dudx - v_here * dudy 219 | + gx + fx_kappa * 2 / (rho[i, j] + rho[i - 1, j])) 220 | ) 221 | for i, j in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1)): 222 | u_here = 0.25 * (u[i, j - 1] + u[i, j] + u[i + 1, j - 1] + u[i + 1, j]) 223 | dvdx = (v[i,j] - v[i-1,j]) * dxi if u_here > 0 else (v[i+1,j] - v[i,j]) * dxi 224 | dvdy = (v[i,j] - v[i,j-1]) * dyi if v[i,j] > 0 else (v[i,j+1] - v[i,j]) * dyi 225 | kappa_ave = (kappa[i, j] + kappa[i, j - 1]) / 2.0 226 | fy_kappa = - sigma[None] * (F[i, j] - F[i, j - 1]) * kappa_ave / dy 227 | v_star[i, j] = ( 228 | v[i, j] + dt * 229 | (nu[i, j] * (v[i - 1, j] - 2 * v[i, j] + v[i + 1, j]) * dxi**2 230 | + nu[i, j] * (v[i, j - 1] - 2 * v[i, j] + v[i, j + 1]) * dyi**2 231 | - u_here * dvdx - v[i, j] * dvdy 232 | + gy + fy_kappa * 2 / (rho[i, j] + rho[i, j - 1])) 233 | ) 234 | 235 | 236 | @ti.kernel 237 | def solve_p_jacobi(): 238 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 239 | rhs = rho[i, j] / dt * \ 240 | ((u_star[i + 1, j] - u_star[i, j]) * dxi + 241 | (v_star[i, j + 1] - v_star[i, j]) * dyi) 242 | ''' istep is compile time constant; so the den_corr actually has no effect 243 | # Calculate the term due to density gradient 244 | drhox1 = (rho[i + 1, j - 1] + rho[i + 1, j] + rho[i + 1, j + 1]) / 3 245 | drhox2 = (rho[i - 1, j - 1] + rho[i - 1, j] + rho[i - 1, j + 1]) / 3 246 | drhodx = (dt / drhox1 - dt / drhox2) / (2 * dx) 247 | drhoy1 = (rho[i - 1, j + 1] + rho[i, j + 1] + rho[i + 1, j + 1]) / 3 248 | drhoy2 = (rho[i - 1, j - 1] + rho[i, j - 1] + rho[i + 1, j - 1]) / 3 249 | drhody = (dt / drhoy1 - dt / drhoy2) / (2 * dy) 250 | dpdx = (p[i + 1, j] - p[i - 1, j]) / (2 * dx) 251 | dpdy = (p[i, j + 1] - p[i, j - 1]) / (2 * dy) 252 | den_corr = (drhodx * dpdx + drhody * dpdy) * rho[i, j] / dt 253 | if istep < 2: 254 | pass 255 | else: 256 | rhs -= den_corr 257 | ''' 258 | ae = dxi ** 2 if i != imax else 0.0 259 | aw = dxi ** 2 if i != imin else 0.0 260 | an = dyi ** 2 if j != jmax else 0.0 261 | a_s = dyi ** 2 if j != jmin else 0.0 262 | ap = - 1.0 * (ae + aw + an + a_s) 263 | pt[i, j] = (rhs - ae * p[i+1,j] - aw * p[i-1,j] - an * p[i,j+1] - a_s * p[i,j-1]) / ap 264 | 265 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 266 | p[i, j] = pt[i, j] 267 | 268 | 269 | @ti.kernel 270 | def update_uv(): 271 | for i, j in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1)): 272 | r = (rho[i, j] + rho[i-1, j]) * 0.5 273 | u[i, j] = u_star[i, j] - dt / r * (p[i, j] - p[i - 1, j]) * dxi 274 | if u[i, j] * dt > 0.25 * dx: 275 | print(f'U velocity courant number > 1, u[{i},{j}] = {u[i,j]}') 276 | for i, j in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1)): 277 | r = (rho[i, j] + rho[i, j-1]) * 0.5 278 | v[i, j] = v_star[i, j] - dt / r * (p[i, j] - p[i, j - 1]) * dyi 279 | if v[i, j] * dt > 0.25 * dy: 280 | print(f'V velocity courant number > 1, v[{i},{j}] = {v[i,j]}') 281 | 282 | 283 | @ti.kernel 284 | def get_normal_young(): 285 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 286 | # Points in between the outermost boundaries 287 | mx1[i, j] = -1 / (2 * dx) * (F[i + 1, j + 1] + F[i + 1, j] - F[i, j + 1] - F[i, j]) 288 | my1[i, j] = -1 / (2 * dy) * (F[i + 1, j + 1] - F[i + 1, j] + F[i, j + 1] - F[i, j]) 289 | mx2[i, j] = -1 / (2 * dx) * (F[i + 1, j] + F[i + 1, j - 1] - F[i, j] - F[i, j - 1]) 290 | my2[i, j] = -1 / (2 * dy) * (F[i + 1, j] - F[i + 1, j - 1] + F[i, j] - F[i, j - 1]) 291 | mx3[i, j] = -1 / (2 * dx) * (F[i, j] + F[i, j - 1] - F[i - 1, j] - F[i - 1, j - 1]) 292 | my3[i, j] = -1 / (2 * dy) * (F[i, j] - F[i, j - 1] + F[i - 1, j] - F[i - 1, j - 1]) 293 | mx4[i, j] = -1 / (2 * dx) * (F[i, j + 1] + F[i, j] - F[i - 1, j + 1] - F[i - 1, j]) 294 | my4[i, j] = -1 / (2 * dy) * (F[i, j + 1] - F[i, j] + F[i - 1, j + 1] - F[i - 1, j]) 295 | # Summing of mx and my components for normal vector 296 | mxsum[i, j] = (mx1[i, j] + mx2[i, j] + mx3[i, j] + mx4[i, j]) / 4 297 | mysum[i, j] = (my1[i, j] + my2[i, j] + my3[i, j] + my4[i, j]) / 4 298 | 299 | # Normalizing the normal vector into unit vectors 300 | if abs(mxsum[i, j]) < 1e-10 and abs(mysum[i, j])< 1e-10: 301 | mx[i, j] = mxsum[i, j] 302 | my[i, j] = mysum[i, j] 303 | else: 304 | magnitude[i, j] = ti.sqrt(mxsum[i, j] * mxsum[i, j] + mysum[i, j] * mysum[i, j]) 305 | mx[i, j] = mxsum[i, j] / magnitude[i, j] 306 | my[i, j] = mysum[i, j] / magnitude[i, j] 307 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 308 | kappa[i, j] = -(1 / dx / 2 * (mx[i + 1, j] - mx[i - 1, j]) + \ 309 | 1 / dy / 2 * (my[i, j + 1] - my[i, j - 1])) 310 | 311 | 312 | def solve_VOF_rudman(): 313 | if istep % 2 == 0: 314 | fct_y_sweep() 315 | fct_x_sweep() 316 | else: 317 | fct_x_sweep() 318 | fct_y_sweep() 319 | 320 | 321 | @ti.kernel 322 | def fct_x_sweep(): 323 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 324 | dv = dx * dy - dt * dy * (u[i + 1, j] - u[i, j]) 325 | fl_L = u[i, j] * dt * F[i - 1, j] if u[i, j] >= 0 else u[i, j] * dt * F[i, j] 326 | fr_L = u[i + 1, j] * dt * F[i, j] if u[i + 1, j] >= 0 else u[i + 1, j] * dt * F[i + 1, j] 327 | ft_L = 0 328 | fb_L = 0 329 | Ftd[i, j] = (F[i, j] + (fl_L - fr_L + fb_L - ft_L) * dy / (dx * dy)) * dx * dy / dv 330 | if Ftd[i, j] > 1. or Ftd[i, j] < 0: 331 | Ftd[i, j] = var(0, 1, Ftd[i, j]) 332 | 333 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 334 | fmax = ti.max(Ftd[i, j], Ftd[i - 1, j], Ftd[i + 1, j]) 335 | fmin = ti.min(Ftd[i, j], Ftd[i - 1, j], Ftd[i + 1, j]) 336 | 337 | fl_L = u[i, j] * dt * F[i - 1, j] if u[i, j] >= 0 else u[i, j] * dt * F[i, j] 338 | fr_L = u[i + 1, j] * dt * F[i, j] if u[i + 1, j] >= 0 else u[i + 1, j] * dt * F[i + 1, j] 339 | ft_L = 0 340 | fb_L = 0 341 | 342 | fl_H = u[i, j] * dt * F[i - 1, j] if u[i, j] <= 0 else u[i, j] * dt * F[i, j] 343 | fr_H = u[i + 1, j] * dt * F[i, j] if u[i + 1, j] <= 0 else u[i + 1, j] * dt * F[i + 1, j] 344 | ft_H = 0 345 | fb_H = 0 346 | 347 | ax[i + 1, j] = fr_H - fr_L 348 | ax[i, j] = fl_H - fl_L 349 | ay[i, j + 1] = 0 350 | ay[i, j] = 0 351 | 352 | pp = ti.max(0, ax[i, j]) - ti.min(0, ax[i + 1, j]) + ti.max(0, ay[i, j]) - ti.min(0, ay[i, j + 1]) 353 | qp = (fmax - Ftd[i, j]) * dx 354 | if pp > 0: 355 | rp[i, j] = ti.min(1, qp / pp) 356 | else: 357 | rp[i, j] = 0.0 358 | pm = ti.max(0, ax[i + 1, j]) - ti.min(0, ax[i, j]) + ti.max(0, ay[i, j + 1]) - ti.min(0, ay[i, j]) 359 | qm = (Ftd[i, j] - fmin) * dx 360 | if pm > 0: 361 | rm[i, j] = ti.min(1, qm / pm) 362 | else: 363 | rm[i, j] = 0.0 364 | 365 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 366 | if ax[i + 1, j] >= 0: 367 | cx[i + 1, j] = ti.min(rp[i + 1, j], rm[i, j]) 368 | else: 369 | cx[i + 1, j] = ti.min(rp[i, j], rm[i + 1, j]) 370 | 371 | if ay[i, j + 1] >= 0: 372 | cy[i, j + 1] = ti.min(rp[i, j + 1], rm[i, j]) 373 | else: 374 | cy[i, j + 1] = ti.min(rp[i, j], rm[i, j + 1]) 375 | 376 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 377 | dv = dx * dy - dt * dy * (u[i + 1, j] - u[i, j]) 378 | F[i, j] = Ftd[i, j] - ((ax[i + 1, j] * cx[i + 1, j] - \ 379 | ax[i, j] * cx[i, j] + \ 380 | ay[i, j + 1] * cy[i, j + 1] -\ 381 | ay[i, j] * cy[i, j]) / (dy)) * dx * dy / dv 382 | F[i, j] = var(0, 1, F[i, j]) 383 | 384 | 385 | @ti.kernel 386 | def fct_y_sweep(): 387 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 388 | dv = dx * dy - dt * dx * (v[i, j + 1] - v[i, j]) 389 | fl_L = 0 390 | fr_L = 0 391 | ft_L = v[i, j + 1] * dt * F[i, j] if v[i, j + 1] >= 0 else v[i, j + 1] * dt * F[i, j + 1] 392 | fb_L = v[i, j] * dt * F[i, j - 1] if v[i, j] >= 0 else v[i, j] * dt * F[i, j] 393 | Ftd[i, j] = (F[i, j] + (fl_L - fr_L + fb_L - ft_L) * dy / (dx * dy)) * dx * dy / dv 394 | if Ftd[i, j] > 1. or Ftd[i, j] < 0: 395 | Ftd[i, j] = var(0, 1, Ftd[i, j]) 396 | 397 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 398 | fmax = ti.max(Ftd[i, j], Ftd[i, j - 1], Ftd[i, j + 1]) 399 | fmin = ti.min(Ftd[i, j], Ftd[i, j - 1], Ftd[i, j + 1]) 400 | 401 | fl_L = 0 402 | fr_L = 0 403 | ft_L = v[i, j + 1] * dt * F[i, j] if v[i, j + 1] >= 0 else v[i, j + 1] * dt * F[i, j + 1] 404 | fb_L = v[i, j] * dt * F[i, j - 1] if v[i, j] >= 0 else v[i, j] * dt * F[i, j] 405 | 406 | fl_H = 0 407 | fr_H = 0 408 | ft_H = v[i, j + 1] * dt * F[i, j] if v[i, j + 1] <= 0 else v[i, j + 1] * dt * F[i, j + 1] 409 | fb_H = v[i, j] * dt * F[i, j - 1] if v[i, j] <= 0 else v[i, j] * dt * F[i, j] 410 | 411 | ax[i + 1, j] = 0 412 | ax[i, j] = 0 413 | ay[i, j + 1] = ft_H - ft_L 414 | ay[i, j] = fb_H - fb_L 415 | 416 | pp = ti.max(0, ax[i, j]) - ti.min(0, ax[i + 1, j]) + ti.max(0, ay[i, j]) - ti.min(0, ay[i, j + 1]) 417 | qp = (fmax - Ftd[i, j]) * dx 418 | if pp > 0: 419 | rp[i, j] = ti.min(1, qp / pp) 420 | else: 421 | rp[i, j] = 0.0 422 | pm = ti.max(0, ax[i + 1, j]) - ti.min(0, ax[i, j]) + ti.max(0, ay[i, j + 1]) - ti.min(0, ay[i, j]) 423 | qm = (Ftd[i, j] - fmin) * dx 424 | if pm > 0: 425 | rm[i, j] = ti.min(1, qm / pm) 426 | else: 427 | rm[i, j] = 0.0 428 | 429 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 430 | if ax[i + 1, j] >= 0: 431 | cx[i + 1, j] = ti.min(rp[i + 1, j], rm[i, j]) 432 | else: 433 | cx[i + 1, j] = ti.min(rp[i, j], rm[i + 1, j]) 434 | 435 | if ay[i, j + 1] >= 0: 436 | cy[i, j + 1] = ti.min(rp[i, j + 1], rm[i, j]) 437 | else: 438 | cy[i, j + 1] = ti.min(rp[i, j], rm[i, j + 1]) 439 | 440 | 441 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 442 | dv = dx * dy - dt * dx * (v[i, j + 1] - v[i, j]) 443 | F[i, j] = Ftd[i, j] - ((ax[i + 1, j] * cx[i + 1, j] - \ 444 | ax[i, j] * cx[i, j] + \ 445 | ay[i, j + 1] * cy[i, j + 1] -\ 446 | ay[i, j] * cy[i, j]) / (dy)) * dx * dy / dv 447 | 448 | F[i, j] = var(0, 1, F[i, j]) 449 | 450 | 451 | 452 | @ti.kernel 453 | def post_process_f(): 454 | for i, j in F: 455 | F[i, j] = var(F[i, j], 0, 1) 456 | 457 | 458 | @ti.kernel 459 | def get_vof_field(): 460 | r = resolution[0] // nx 461 | for I in ti.grouped(rgb_buf): 462 | rgb_buf[I] = F[I // r] 463 | 464 | 465 | @ti.kernel 466 | def get_u_field(): 467 | r = resolution[0] // nx 468 | max = Lx / 0.2 469 | for I in ti.grouped(rgb_buf): 470 | rgb_buf[I] = (u[I // r]) / max 471 | 472 | 473 | @ti.kernel 474 | def get_v_field(): 475 | r = resolution[0] // nx 476 | max = Ly / 0.2 477 | for I in ti.grouped(rgb_buf): 478 | rgb_buf[I] = (v[I // r]) / max 479 | 480 | 481 | @ti.kernel 482 | def get_vnorm_field(): 483 | r = resolution[0] // nx 484 | max = Ly / 0.2 485 | for I in ti.grouped(rgb_buf): 486 | rgb_buf[I] = ti.sqrt(u[I // r] ** 2 + v[I // r] ** 2) / max 487 | 488 | 489 | @ti.kernel 490 | def interp_velocity(): 491 | for i, j in ti.ndrange((1, imax + 2), (1, jmax + 1)): 492 | V[i, j] = ti.Vector([(u[i, j] + u[i + 1, j])/2, (v[i, j] + v[i, j + 1])/2]) 493 | 494 | 495 | # Start main script 496 | istep = 0 497 | nstep = 100 # Interval to update GUI 498 | set_init_F(initial_condition) 499 | 500 | os.makedirs('output', exist_ok=True) # Make dir for output 501 | os.makedirs('data', exist_ok=True) # Make dir for data save; only used for debugging 502 | gui = ti.GUI('VOF Solver', resolution, background_color=0xFFFFFF) 503 | vis_option = 0 # Tag for display 504 | 505 | while gui.running: 506 | istep += 1 507 | for e in gui.get_events(gui.RELEASE): 508 | if e.key == gui.SPACE: 509 | vis_option += 1 510 | elif e.key == 'q': 511 | gui.running = False 512 | 513 | cal_nu_rho() 514 | get_normal_young() 515 | 516 | # Advection 517 | advect_upwind() 518 | set_BC() 519 | 520 | # Pressure projection 521 | for _ in range(10): 522 | solve_p_jacobi() 523 | 524 | update_uv() 525 | set_BC() 526 | solve_VOF_rudman() 527 | post_process_f() 528 | set_BC() 529 | 530 | num_options = 5 531 | if (istep % nstep) == 0: # Output data every steps 532 | if vis_option % num_options == 0: # Display VOF distribution 533 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying VOF field.') 534 | get_vof_field() 535 | rgbnp = rgb_buf.to_numpy() 536 | gui.set_image(cm.Blues(rgbnp)) 537 | 538 | if vis_option % num_options == 1: # Display the u field 539 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying u velocity.') 540 | get_u_field() 541 | rgbnp = rgb_buf.to_numpy() 542 | gui.set_image(cm.coolwarm(rgbnp)) 543 | 544 | if vis_option % num_options == 2: # Display the v field 545 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying v velocity.') 546 | get_v_field() 547 | rgbnp = rgb_buf.to_numpy() 548 | gui.set_image(cm.coolwarm(rgbnp)) 549 | 550 | if vis_option % num_options == 3: # Display velocity norm 551 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity norm.') 552 | get_vnorm_field() 553 | rgbnp = rgb_buf.to_numpy() 554 | gui.set_image(cm.plasma(rgbnp)) 555 | 556 | if vis_option % num_options == 4: # Display velocity vectors 557 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity vectors.') 558 | interp_velocity() 559 | fv.plot_arrow_field(vector_field=V, arrow_spacing=4, gui=gui) 560 | 561 | gui.show() 562 | 563 | if SAVE_FIG: 564 | count = istep // nstep - 1 565 | Fnp = F.to_numpy() 566 | fx, fy = 5, Ly / Lx * 5 567 | plt.figure(figsize=(fx, fy)) 568 | plt.axis('off') 569 | plt.contourf(Fnp.T, cmap=plt.cm.Blues) 570 | plt.savefig(f'output/{count:06d}-f.png') 571 | plt.close() 572 | -------------------------------------------------------------------------------- /diff_vof.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import matplotlib.cm as cm 5 | import argparse 6 | import os 7 | import flow_visualization as fv 8 | 9 | ti.init(arch=ti.cpu, default_fp=ti.f32, debug=False, device_memory_fraction=0.9) # Set default fp so that float=ti.f32 10 | 11 | parser = argparse.ArgumentParser() # Get the initial condition 12 | # 1 - Dam Break; 2 - Rising Bubble; 3 - Droping liquid 13 | parser.add_argument('-ic', type=int, choices=[1, 2, 3], default=1) 14 | parser.add_argument('-s', action='store_true') 15 | args = parser.parse_args() 16 | initial_condition = args.ic 17 | SAVE_FIG = args.s 18 | 19 | nx = 80 # Number of grid points in the x direction 20 | ny = 80 # Number of grid points in the y direction 21 | 22 | Lx = 0.1 # The length of the domain 23 | Ly = 0.1 # The width of the domain 24 | rho_l = 1000.0 25 | rho_g = 50.0 26 | nu_l = 1.0e-6 # kinematic viscosity, nu = mu / rho 27 | nu_g = 1.5e-5 28 | sigma = ti.field(dtype=float, shape=()) 29 | sigma[None] = 0.007 30 | gx = 0 31 | gy = -1000 32 | 33 | dt = 4e-6 # Use smaller dt for higher density ratio 34 | eps = 1e-6 # Threshold used in vfconv and f post processings 35 | 36 | MAX_TIME_STEPS = 1000 37 | MAX_ITER = 10 38 | OPT_ITER = 100 39 | learning_rate = 0.02 40 | # Mesh information 41 | imin = 1 42 | imax = imin + nx - 1 43 | jmin = 1 44 | jmax = jmin + ny - 1 45 | x = ti.field(float, shape=imax + 3) 46 | y = ti.field(float, shape=jmax + 3) 47 | xnp = np.hstack((0.0, np.linspace(0, Lx, nx + 1), Lx)).astype(np.float32) # [0, 0, ... 1, 1] 48 | x.from_numpy(xnp) 49 | ynp = np.hstack((0.0, np.linspace(0, Ly, ny + 1), Ly)).astype(np.float32) # [0, 0, ... 1, 1] 50 | y.from_numpy(ynp) 51 | dx = x[imin + 2] - x[imin + 1] 52 | dy = y[jmin + 2] - y[jmin + 1] 53 | dxi = 1 / dx 54 | dyi = 1 / dy 55 | 56 | # Field shapes 57 | field_shape = (imax + 2, jmax + 2, MAX_TIME_STEPS) 58 | p_shape = (imax + 2, jmax + 2, MAX_TIME_STEPS * (MAX_ITER + 1)) 59 | 60 | # Variables for VOF function 61 | F = ti.field(float, shape=(imax + 2, jmax + 2, 2 * MAX_TIME_STEPS + 1), needs_grad=True) 62 | Ftd_x = ti.field(float, shape=field_shape, needs_grad=True) 63 | Ftd_y = ti.field(float, shape=field_shape, needs_grad=True) 64 | Ftarget = ti.field(float, shape=(field_shape[0], field_shape[1]), needs_grad=True) 65 | loss = ti.field(float, shape=(), needs_grad=True) 66 | 67 | ax = ti.field(float, shape=field_shape, needs_grad=True) 68 | ay = ti.field(float, shape=field_shape, needs_grad=True) 69 | cx = ti.field(float, shape=field_shape, needs_grad=True) 70 | cy = ti.field(float, shape=field_shape, needs_grad=True) 71 | rp_x = ti.field(float, shape=field_shape, needs_grad=True) 72 | rm_x = ti.field(float, shape=field_shape, needs_grad=True) 73 | rp_y = ti.field(float, shape=field_shape, needs_grad=True) 74 | rm_y = ti.field(float, shape=field_shape, needs_grad=True) 75 | 76 | # Variables for N-S equation 77 | u = ti.field(float, shape=field_shape, needs_grad=True) 78 | v = ti.field(float, shape=field_shape, needs_grad=True) 79 | u_star = ti.field(float, shape=field_shape, needs_grad=True) 80 | v_star = ti.field(float, shape=field_shape, needs_grad=True) 81 | 82 | # Pressure field shape should be different 83 | p = ti.field(float, shape=p_shape, needs_grad=True) 84 | rhs = ti.field(float, shape=field_shape, needs_grad=True) 85 | 86 | rho = ti.field(float, shape=field_shape, needs_grad=True) 87 | nu = ti.field(float, shape=field_shape, needs_grad=True) 88 | V = ti.Vector.field(2, dtype=float, shape=(field_shape[0], field_shape[1])) # For displaying velocity field 89 | 90 | # Variables for interface reconstruction 91 | mx1 = ti.field(float, shape=field_shape, needs_grad=True) 92 | my1 = ti.field(float, shape=field_shape, needs_grad=True) 93 | mx2 = ti.field(float, shape=field_shape, needs_grad=True) 94 | my2 = ti.field(float, shape=field_shape, needs_grad=True) 95 | mx3 = ti.field(float, shape=field_shape, needs_grad=True) 96 | my3 = ti.field(float, shape=field_shape, needs_grad=True) 97 | mx4 = ti.field(float, shape=field_shape, needs_grad=True) 98 | my4 = ti.field(float, shape=field_shape, needs_grad=True) 99 | mxsum = ti.field(float, shape=field_shape, needs_grad=True) 100 | mysum = ti.field(float, shape=field_shape, needs_grad=True) 101 | mx = ti.field(float, shape=field_shape, needs_grad=True) 102 | my = ti.field(float, shape=field_shape, needs_grad=True) 103 | kappa = ti.field(float, shape=field_shape, needs_grad=True) # interface curvature 104 | magnitude = ti.field(float, shape=field_shape, needs_grad=True) 105 | 106 | # For visualization 107 | resolution = (800, 400) 108 | rgb_buf = ti.field(dtype=float, shape=(2 * field_shape[0], field_shape[1])) 109 | 110 | print(f'>>> A VOF solver written in Taichi; Press q to exit.') 111 | print(f'>>> Grid resolution: {nx} x {ny}, dt = {dt:4.2e}') 112 | print(f'>>> Density ratio: {rho_l / rho_g : 4.2f}, gravity : {gy : 4.2f}, sigma : {sigma[None] : 6.3f}') 113 | print(f'>>> Viscosity ratio: {nu_l / nu_g : 4.2f}') 114 | print(f'>>> Please wait a few seconds to let the kernels compile...') 115 | 116 | 117 | @ti.func 118 | def find_area(i, j, cx, cy, r): 119 | a = 0.0 120 | xcoord_ct = (i - imin) * dx + dx / 2 121 | ycoord_ct = (j - jmin) * dy + dy / 2 122 | 123 | xcoord_lu = xcoord_ct - dx / 2 124 | ycoord_lu = ycoord_ct + dy / 2 125 | 126 | xcoord_ld = xcoord_ct - dx / 2 127 | ycoord_ld = ycoord_ct - dy / 2 128 | 129 | xcoord_ru = xcoord_ct + dx / 2 130 | ycoord_ru = ycoord_ct + dy / 2 131 | 132 | xcoord_rd = xcoord_ct + dx / 2 133 | ycoord_rd = ycoord_ct - dy / 2 134 | 135 | dist_ct = ti.sqrt((xcoord_ct - cx) ** 2 + (ycoord_ct - cy) ** 2) 136 | dist_lu = ti.sqrt((xcoord_lu - cx) ** 2 + (ycoord_lu - cy) ** 2) 137 | dist_ld = ti.sqrt((xcoord_ld - cx) ** 2 + (ycoord_ld - cy) ** 2) 138 | dist_ru = ti.sqrt((xcoord_ru - cx) ** 2 + (ycoord_ru - cy) ** 2) 139 | dist_rd = ti.sqrt((xcoord_rd - cx) ** 2 + (ycoord_rd - cy) ** 2) 140 | 141 | if dist_lu > r and dist_ld > r and dist_ru > r and dist_rd > r: 142 | a = 1.0 143 | elif dist_lu < r and dist_ld < r and dist_ru < r and dist_rd < r: 144 | a = 0.0 145 | else: 146 | a = 0.5 + 0.5 * (dist_ct - r) / (ti.sqrt(2.0) * dx) 147 | a = var(a, 0, 1) 148 | 149 | return a 150 | 151 | 152 | @ti.kernel 153 | def set_init_F(ic:ti.i32): 154 | # Sets the initial volume fraction 155 | if ic == 1: # Dambreak 156 | x1 = Lx / 3 * 1.0 157 | x2 = Lx / 3 * 2.0 158 | y1 = 0.0 159 | y2 = Ly / 2 160 | r = Ly / 4 161 | for i, j in ti.ndrange(imax + 2, jmax + 2): 162 | if (x[i] >= x1) and (x[i] <= x2) and (y[j] >= y1) and (y[j] <= y2): 163 | Ftarget[i, j] = 1.0 164 | 165 | elif ic == 2: # Rising bubble 166 | for i, j in ti.ndrange(imax + 2, jmax + 2): 167 | r = Lx / 12 168 | cx, cy = Lx / 2, Ly / 2 169 | Ftarget[i, j] = find_area(i, j, cx, cy, r) 170 | F[i, j, 0] = 1.0 171 | 172 | elif ic == 3: # Liquid drop 173 | for i, j in ti.ndrange(imax + 2, jmax + 2): 174 | r = Lx / 12 175 | cx, cy = Lx / 2, Ly / 2 176 | Ftarget[i, j] = 1.0 - find_area(i, j, cx, cy, r) 177 | 178 | 179 | @ti.kernel 180 | def set_pixel(x:ti.f32, y:ti.f32, f:ti.template()): 181 | xcord = ti.i32(x * imax) 182 | ycord = ti.i32(y * jmax) 183 | for i, j in ti.ndrange((xcord-2, xcord + 2),(ycord-2, ycord+2) ): 184 | if i >= 0 and j >= 0: 185 | f[i, j] = 1.0 186 | 187 | 188 | def set_init_by_paint(): 189 | gui = ti.GUI("Paint your initial", ) 190 | while gui.running: 191 | gui.contour(Ftarget) 192 | gui.get_event() 193 | if gui.is_pressed(ti.GUI.ESCAPE): 194 | gui.running = False 195 | if gui.is_pressed(ti.GUI.LMB): 196 | x, y = gui.get_cursor_pos() 197 | set_pixel(x, y, Ftarget) 198 | gui.show() 199 | 200 | 201 | @ti.kernel 202 | def set_BC(t:ti.i32): 203 | for i in ti.ndrange(imax + 2): 204 | # bottom: slip 205 | u[i, jmin - 1, t] = u[i, jmin, t] 206 | v[i, jmin, t] = 0 207 | F[i, jmin - 1, 2 * t] = F[i, jmin, 2 * t] 208 | p[i, jmin - 1, t * (MAX_ITER + 1)] = p[i, jmin, t * (MAX_ITER + 1)] 209 | rho[i, jmin - 1, t] = rho[i, jmin, t] 210 | # top: open 211 | u[i, jmax + 1, t] = u[i, jmax, t] 212 | v[i, jmax + 1, t] = 0 #v[i, jmax, t] 213 | F[i, jmax + 1, 2 * t] = F[i, jmax, 2 * t] 214 | p[i, jmax + 1, t * (MAX_ITER + 1)] = p[i, jmax, t * (MAX_ITER + 1)] 215 | rho[i, jmax + 1, t] = rho[i, jmax, t] 216 | for j in ti.ndrange(jmax + 2): 217 | # left: slip 218 | u[imin, j, t] = 0 219 | v[imin - 1, j, t] = v[imin, j, t] 220 | F[imin - 1, j, 2 * t] = F[imin, j, 2 * t] 221 | p[imin - 1, j, t * (MAX_ITER + 1)] = p[imin, j, t * (MAX_ITER + 1)] 222 | rho[imin - 1, j, t] = rho[imin, j, t] 223 | # right: slip 224 | u[imax + 1, j, t] = 0 225 | v[imax + 1, j, t] = v[imax, j, t] 226 | F[imax + 1, j, 2 * t] = F[imax, j, 2 * t] 227 | p[imax + 1, j, t * (MAX_ITER + 1)] = p[imax, j, t * (MAX_ITER + 1)] 228 | rho[imax + 1, j, t] = rho[imax, j, t] 229 | 230 | 231 | @ti.func 232 | def var(a, b, c): # Find the median of a,b, and c 233 | center = a + b + c - ti.max(a, b, c) - ti.min(a, b, c) 234 | return center 235 | 236 | 237 | @ti.kernel 238 | def cal_nu_rho(t:ti.i32): 239 | for i, j in ti.ndrange(field_shape[0], field_shape[1]): 240 | F = var(0.0, 1.0, F[i, j, 2 * t]) 241 | rho[i, j, t] = rho_g * (1 - F) + rho_l * F 242 | nu[i, j, t] = nu_l * F + nu_g * (1.0 - F) 243 | 244 | 245 | @ti.kernel 246 | def advect_upwind(t:ti.i32): 247 | for i, j in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1)): 248 | v_here = 0.25 * (v[i - 1, j, t] + v[i - 1, j + 1, t] + v[i, j, t] + v[i, j + 1, t]) 249 | dudx = (u[i, j, t] - u[i - 1, j, t]) * dxi if u[i, j, t] > 0 else (u[i + 1, j, t] - u[i, j, t]) * dxi 250 | dudy = (u[i, j, t] - u[i, j - 1, t]) * dyi if v_here > 0 else (u[i, j + 1, t] - u[i, j, t]) * dyi 251 | kappa_ave = (kappa[i, j, t] + kappa[i - 1, j, t]) / 2.0 252 | fx_kappa = - sigma[None] * (F[i, j, 2 * t] - F[i - 1, j, 2 * t]) * kappa_ave / dx # F(2*t) is F at t time step 253 | u_star[i, j, t] = ( 254 | u[i, j, t] + dt * 255 | (nu[i, j, t] * (u[i - 1, j, t] - 2 * u[i, j, t] + u[i + 1, j, t]) * dxi**2 256 | + nu[i, j, t] * (u[i, j - 1, t] - 2 * u[i, j, t] + u[i, j + 1, t]) * dyi**2 257 | - u[i, j, t] * dudx - v_here * dudy 258 | + gx + fx_kappa * 2 / (rho[i, j, t] + rho[i - 1, j, t])) 259 | ) 260 | for i, j in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1)): 261 | u_here = 0.25 * (u[i, j - 1, t] + u[i, j, t] + u[i + 1, j - 1, t] + u[i + 1, j, t]) 262 | dvdx = (v[i, j, t] - v[i - 1, j, t]) * dxi if u_here > 0 else (v[i + 1, j, t] - v[i, j, t]) * dxi 263 | dvdy = (v[i, j, t] - v[i, j - 1, t]) * dyi if v[i, j, t] > 0 else (v[i, j + 1, t] - v[i, j, t]) * dyi 264 | kappa_ave = (kappa[i, j, t] + kappa[i, j - 1, t]) / 2.0 265 | fy_kappa = - sigma[None] * (F[i, j, 2 * t] - F[i, j - 1, 2 * t]) * kappa_ave / dy 266 | v_star[i, j, t] = ( 267 | v[i, j, t] + dt * 268 | (nu[i, j, t] * (v[i - 1, j, t] - 2 * v[i, j, t] + v[i + 1, j, t]) * dxi**2 269 | + nu[i, j, t] * (v[i, j - 1, t] - 2 * v[i, j, t] + v[i, j + 1, t]) * dyi**2 270 | - u_here * dvdx - v[i, j, t] * dvdy 271 | + gy + fy_kappa * 2 / (rho[i, j, t] + rho[i, j - 1, t])) 272 | ) 273 | 274 | 275 | @ti.kernel 276 | def solve_p_jacobi(t:ti.i32, k:ti.i32): 277 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 278 | rhs = rho[i, j, t] / dt * \ 279 | ((u_star[i + 1, j, t] - u_star[i, j, t]) * dxi + 280 | (v_star[i, j + 1, t] - v_star[i, j, t]) * dyi) 281 | ae = dxi ** 2 if i != imax else 0.0 282 | aw = dxi ** 2 if i != imin else 0.0 283 | an = dyi ** 2 if j != jmax else 0.0 284 | a_s = dyi ** 2 if j != jmin else 0.0 285 | ap = - 1.0 * (ae + aw + an + a_s) 286 | p[i, j, t * (MAX_ITER + 1) + k + 1] = (rhs \ 287 | - ae * p[i+1,j, t * (MAX_ITER + 1) + k] \ 288 | - aw * p[i-1,j, t * (MAX_ITER + 1) + k] \ 289 | - an * p[i,j+1, t * (MAX_ITER + 1) + k] \ 290 | - a_s * p[i,j-1, t * (MAX_ITER + 1) +k] )\ 291 | / ap 292 | 293 | @ti.kernel 294 | def copy_p_field(t:ti.i32, k:ti.i32): 295 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 296 | p[i, j, (t + 1) * (MAX_ITER + 1)] = p[i, j, t * (MAX_ITER + 1) + MAX_ITER] 297 | 298 | 299 | @ti.kernel 300 | def update_uv(t:ti.i32): 301 | for i, j in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1)): 302 | r = (rho[i, j, t] + rho[i - 1, j, t]) * 0.5 303 | u[i, j, t + 1] = u_star[i, j, t] \ 304 | - dt / r * \ 305 | (p[i, j, t*(MAX_ITER+1)+MAX_ITER] - p[i - 1, j, t*(MAX_ITER+1)+MAX_ITER]) * dxi 306 | if u[i, j, t + 1] * dt > 0.25 * dx: 307 | print(f'U velocity courant number > 1, u[{i},{j},{t+1}] = {u[i, j, t+1]}') 308 | for i, j in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1)): 309 | r = (rho[i, j, t] + rho[i, j - 1, t]) * 0.5 310 | v[i, j, t + 1] = v_star[i, j, t] \ 311 | - dt / r \ 312 | * (p[i, j, t*(MAX_ITER+1)+MAX_ITER] - p[i, j - 1, t*(MAX_ITER+1)+MAX_ITER]) * dyi 313 | if v[i, j, t + 1] * dt > 0.25 * dy: 314 | print(f'V velocity courant number > 1, v[{i},{j},{t+1}] = {v[i,j,t+1]}') 315 | 316 | 317 | @ti.kernel 318 | def get_normal_young(t:ti.i32): 319 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 320 | # Points in between the outermost boundaries 321 | mx1[i, j, t] = -1 / (2 * dx) * (F[i + 1, j + 1, 2 * t] + F[i + 1, j, 2 * t] - F[i, j + 1, 2 * t] - F[i, j, 2 * t]) 322 | my1[i, j, t] = -1 / (2 * dy) * (F[i + 1, j + 1, 2 * t] - F[i + 1, j, 2 * t] + F[i, j + 1, 2 * t] - F[i, j, 2 * t]) 323 | mx2[i, j, t] = -1 / (2 * dx) * (F[i + 1, j, 2 * t] + F[i + 1, j - 1, 2 * t] - F[i, j, 2 * t] - F[i, j - 1, 2 * t]) 324 | my2[i, j, t] = -1 / (2 * dy) * (F[i + 1, j, 2 * t] - F[i + 1, j - 1, 2 * t] + F[i, j, 2 * t] - F[i, j - 1, 2 * t]) 325 | mx3[i, j, t] = -1 / (2 * dx) * (F[i, j, 2 * t] + F[i, j - 1, 2 * t] - F[i - 1, j, 2 * t] - F[i - 1, j - 1, 2 * t]) 326 | my3[i, j, t] = -1 / (2 * dy) * (F[i, j, 2 * t] - F[i, j - 1, 2 * t] + F[i - 1, j, 2 * t] - F[i - 1, j - 1, 2 * t]) 327 | mx4[i, j, t] = -1 / (2 * dx) * (F[i, j + 1, 2 * t] + F[i, j, 2 * t] - F[i - 1, j + 1, 2 * t] - F[i - 1, j, 2 * t]) 328 | my4[i, j, t] = -1 / (2 * dy) * (F[i, j + 1, 2 * t] - F[i, j, 2 * t] + F[i - 1, j + 1, 2 * t] - F[i - 1, j, 2 * t]) 329 | # Summing of mx and my components for normal vector 330 | mxsum[i, j, t] = (mx1[i, j, t] + mx2[i, j, t] + mx3[i, j, t] + mx4[i, j, t]) / 4 331 | mysum[i, j, t] = (my1[i, j, t] + my2[i, j, t] + my3[i, j, t] + my4[i, j, t]) / 4 332 | # Normalizing the normal vector into unit vectors 333 | if abs(mxsum[i, j, t]) < 1e-10 and abs(mysum[i, j, t])< 1e-10: 334 | mx[i, j, t] = mxsum[i, j, t] 335 | my[i, j, t] = mysum[i, j, t] 336 | else: 337 | magnitude[i, j, t] = ti.sqrt(mxsum[i, j, t] * mxsum[i, j, t] + mysum[i, j, t] * mysum[i, j, t]) 338 | mx[i, j, t] = mxsum[i, j, t] / magnitude[i, j, t] 339 | my[i, j, t] = mysum[i, j, t] / magnitude[i, j, t] 340 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 341 | kappa[i, j, t] = -(1 / dx / 2 * (mx[i + 1, j, t] - mx[i - 1, j, t]) + \ 342 | 1 / dy / 2 * (my[i, j + 1, t] - my[i, j - 1, t])) 343 | 344 | 345 | def solve_VOF_rudman(t): 346 | if t % 2 == 0: 347 | fct_y_sweep(t, 0, 1e-6) 348 | fct_x_sweep(t, 1, 1e-6) 349 | else: 350 | fct_x_sweep(t, 0, 1e-6) 351 | fct_y_sweep(t, 1, 1e-6) 352 | 353 | 354 | @ti.kernel 355 | def fct_x_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 356 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 357 | dv = dx * dy - dt * dy * (u[i + 1, j, t + 1] - u[i, j, t + 1]) 358 | fl_L = u[i, j, t + 1] * dt * F[i - 1, j, 2 * t + offset] if u[i, j, t + 1] >= 0 else u[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 359 | fr_L = u[i + 1, j, t + 1] * dt * F[i, j, 2 * t + offset] if u[i + 1, j, t + 1] >= 0 else u[i + 1, j, t + 1] * dt * F[i + 1, j, 2 * t + offset] 360 | Ftd_x[i, j, t] = F[i, j, 2 * t + offset] + (fl_L - fr_L) * dy / (dx * dy) * dx * dy / dv 361 | 362 | for i, j in ti.ndrange((imin, imax + 2), (jmin, jmax + 1)): 363 | fl_L = u[i, j, t + 1] * dt * F[i - 1, j, 2 * t + offset] if u[i, j, t + 1] >= 0 else u[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 364 | fl_H = u[i, j, t + 1] * dt * F[i - 1, j, 2 * t + offset] if u[i, j, t + 1] <= 0 else u[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 365 | ax[i, j, t] = fl_H - fl_L 366 | 367 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 368 | fmax = ti.max(Ftd_x[i, j, t], Ftd_x[i - 1, j, t], Ftd_x[i + 1, j, t]) 369 | fmin = ti.min(Ftd_x[i, j, t], Ftd_x[i - 1, j, t], Ftd_x[i + 1, j, t]) 370 | 371 | pp = ti.max(0, ax[i, j, t]) - ti.min(0, ax[i + 1, j, t]) 372 | qp = (fmax - Ftd_x[i, j, t]) * dx 373 | if pp > eps: 374 | rp_x[i, j, t] = ti.min(1, qp / pp) 375 | else: 376 | rp_x[i, j, t] = 0.0 377 | 378 | pm = ti.max(0, ax[i + 1, j, t]) - ti.min(0, ax[i, j, t]) 379 | qm = (Ftd_x[i, j, t] - fmin) * dx 380 | if pm > eps: 381 | rm_x[i, j, t] = ti.min(1, qm / pm) 382 | else: 383 | rm_x[i, j, t] = 0.0 384 | 385 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 386 | if ax[i + 1, j, t] >= 0: 387 | cx[i + 1, j, t] = ti.min(rp_x[i + 1, j, t], rm_x[i, j, t]) 388 | else: 389 | cx[i + 1, j, t] = ti.min(rp_x[i, j, t], rm_x[i + 1, j, t]) 390 | 391 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 392 | dv = dx * dy - dt * dy * (u[i + 1, j, t + 1] - u[i, j, t + 1]) 393 | F[i, j, 2 * t + offset + 1] = Ftd_x[i, j, t] - ((ax[i + 1, j, t] * cx[i + 1, j, t] - \ 394 | ax[i, j, t] * cx[i, j, t]) / dy) * dx * dy / dv 395 | 396 | 397 | @ti.kernel 398 | def fct_y_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 399 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 400 | dv = dx * dy - dt * dx * (v[i, j + 1, t + 1] - v[i, j, t + 1]) 401 | ft_L = v[i, j + 1, t + 1] * dt * F[i, j, 2 * t + offset] if v[i, j + 1, t + 1] >= 0 else v[i, j + 1, t + 1] * dt * F[i, j + 1, 2 * t + offset] 402 | fb_L = v[i, j, t + 1] * dt * F[i, j - 1, 2 * t + offset] if v[i, j, t + 1] >= 0 else v[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 403 | Ftd_y[i, j, t] = F[i, j, 2 * t + offset] + (fb_L - ft_L) * dy / (dx * dy) * dx * dy / dv 404 | 405 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 2)): 406 | fb_L = v[i, j, t + 1] * dt * F[i, j - 1, 2 * t + offset] if v[i, j, t + 1] >= 0 else v[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 407 | fb_H = v[i, j, t + 1] * dt * F[i, j - 1, 2 * t + offset] if v[i, j, t + 1] <= 0 else v[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 408 | ay[i, j, t] = fb_H - fb_L 409 | 410 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 411 | fmax = ti.max(Ftd_y[i, j, t], Ftd_y[i, j - 1, t], Ftd_y[i, j + 1, t]) 412 | fmin = ti.min(Ftd_y[i, j, t], Ftd_y[i, j - 1, t], Ftd_y[i, j + 1, t]) 413 | 414 | # eps = 1e-4 415 | pp = ti.max(0, ay[i, j, t]) - ti.min(0, ay[i, j + 1, t]) 416 | qp = (fmax - Ftd_y[i, j, t]) * dx 417 | if pp > eps: 418 | rp_y[i, j, t] = ti.min(1, qp / pp) 419 | else: 420 | rp_y[i, j, t] = 0.0 421 | 422 | pm = ti.max(0, ay[i, j + 1, t]) - ti.min(0, ay[i, j, t]) 423 | qm = (Ftd_y[i, j, t] - fmin) * dx 424 | if pm > eps: 425 | rm_y[i, j, t] = ti.min(1, qm / pm) 426 | else: 427 | rm_y[i, j, t] = 0.0 428 | 429 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 430 | if ay[i, j + 1, t] >= 0: 431 | cy[i, j + 1, t] = ti.min(rp_y[i, j + 1, t], rm_y[i, j, t]) 432 | else: 433 | cy[i, j + 1, t] = ti.min(rp_y[i, j, t], rm_y[i, j + 1, t]) 434 | 435 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 436 | dv = dx * dy - dt * dx * (v[i, j + 1, t + 1] - v[i, j, t + 1]) 437 | F[i, j, 2 * t + offset + 1] = Ftd_y[i, j, t] - ((ay[i, j + 1, t] * cy[i, j + 1, t] -\ 438 | ay[i, j, t] * cy[i, j, t]) / dy) * dx * dy / dv 439 | 440 | 441 | 442 | @ti.kernel 443 | def post_process_f(t:ti.i32): 444 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 445 | F[i, j, 2 * t + 2] = var(F[i, j, 2 * t + 2], 0, 1) 446 | 447 | 448 | @ti.ad.no_grad 449 | @ti.kernel 450 | def get_field_to_buf(src:ti.template(), t:ti.i32): 451 | for i, j in Ftarget: 452 | rgb_buf[i, j] = src[i, j, t] 453 | for i, j in Ftarget: 454 | rgb_buf[i + imax + 1, j] = Ftarget[i, j] 455 | 456 | 457 | @ti.ad.no_grad 458 | @ti.kernel 459 | def get_vnorm_field(t:ti.i32): 460 | for i, j in rgb_buf: 461 | rgb_buf[i, j] = ti.sqrt(u[i, j, t] ** 2 + v[i, j, t] ** 2) 462 | 463 | 464 | @ti.ad.no_grad 465 | @ti.kernel 466 | def interp_velocity(t:ti.i32): 467 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 468 | V[i, j] = ti.Vector([(u[i, j, t] + u[i + 1, j, t])/2, (v[i, j, t] + v[i, j + 1, t])/2]) 469 | 470 | 471 | @ti.kernel 472 | def compute_loss(): 473 | for i, j in ti.ndrange(imax + 2, jmax + 2): 474 | loss[None] += ti.abs(Ftarget[i, j] - F[i, j, 2 * MAX_TIME_STEPS - 2]) 475 | 476 | 477 | @ti.kernel 478 | def apply_grad(): 479 | for i, j in ti.ndrange((1, imax + 1), (1, jmax + 1)): 480 | if ti.abs(F.grad[i, j, 0]) < 5.: 481 | F[i, j, 0] -= learning_rate * F.grad[i, j, 0] 482 | F[i, j, 0] = var(0, 1, F[i, j, 0]) 483 | 484 | 485 | def forward(): 486 | vis_option = 0 # Tag for display 487 | for istep in range(MAX_TIME_STEPS - 1): 488 | for e in gui.get_events(gui.RELEASE): 489 | if e.key == gui.SPACE: 490 | vis_option += 1 491 | elif e.key == 'q': 492 | gui.running = False 493 | 494 | # Calculate initial F 495 | cal_nu_rho(istep) 496 | get_normal_young(istep) 497 | 498 | # Advection 499 | advect_upwind(istep) 500 | ''' 501 | Not necessary; boundary u,v = 0 already, and no update be made on those boundary 502 | For p solving, only those u,v=0 will be used. 503 | ''' 504 | # set_BC(istep) 505 | 506 | # Pressure projection 507 | for iter in range(MAX_ITER): 508 | solve_p_jacobi(istep, iter) 509 | copy_p_field(istep, iter) 510 | 511 | # Velocity correction 512 | update_uv(istep) 513 | ''' 514 | Not necessary. For VOF advection, only u, v = 0 will be used, which are untouched. 515 | And F on the boundary is set at the end of previous step's set_BC 516 | ''' 517 | # set_BC(istep + 1) 518 | 519 | # Advect the VOF function 520 | solve_VOF_rudman(istep) 521 | post_process_f(istep) # Post-processing violates GDAR, but necessary for stablize. 522 | set_BC(istep + 1) 523 | 524 | # Visualization 525 | num_options = 5 526 | plot_contour = ti.ad.no_grad(gui.contour) 527 | plot_vector = ti.ad.no_grad(gui.vector_field) 528 | if (istep % nstep) == 0: # Output data every steps 529 | if vis_option % num_options == 0: # Display VOF distribution 530 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying VOF field.') 531 | get_field_to_buf(F, 2 * (istep + 1)) 532 | plot_contour(rgb_buf) 533 | 534 | if vis_option % num_options == 1: # Display the u field 535 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying u velocity.') 536 | get_field_to_buf(u, istep) 537 | plot_contour(rgb_buf) 538 | 539 | if vis_option % num_options == 2: # Display the v field 540 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying v velocity.') 541 | get_field_to_buf(v, istep) 542 | plot_contour(rgb_buf) 543 | 544 | if vis_option % num_options == 3: # Display velocity norm 545 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity norm.') 546 | get_vnorm_field(istep) 547 | plot_contour(rgb_buf) 548 | 549 | if vis_option % num_options == 4: # Display velocity vectors 550 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity vectors.') 551 | interp_velocity(istep) 552 | plot_vector(V, arrow_spacing=2, color=0x000000) 553 | 554 | gui.show(f'./output/{opt:03d}-{istep:04d}.png') 555 | 556 | # Compute loss as the last step of forward() pass 557 | compute_loss() 558 | 559 | 560 | # Start main script 561 | istep = 0 562 | nstep = 20 # Interval to update GUI 563 | # set_init_F(initial_condition) # Set initial VOF by fixed shape 564 | set_init_by_paint() # Set initial VOF by user painting 565 | os.makedirs('output', exist_ok=True) # Make dir for output 566 | gui = ti.GUI('VOF Solver', resolution, background_color=0xFFFFFF) 567 | vis_option = 0 568 | 569 | for opt in range(OPT_ITER): 570 | print(f'>>> >>> Optimization cycle {opt}') 571 | with ti.ad.Tape(loss): 572 | forward() 573 | print(f'>>> >>> Current total loss is {loss[None]}') 574 | apply_grad() # Apply gradient should be outside the Tape() 575 | print(f'>>> >>> Gradient applied.') 576 | -------------------------------------------------------------------------------- /diff_vof_replaced.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import matplotlib.cm as cm 5 | import argparse 6 | import os 7 | import flow_visualization as fv 8 | 9 | ti.init(arch=ti.cpu, default_fp=ti.f32, debug=False, device_memory_fraction=0.9) # Set default fp so that float=ti.f32 10 | 11 | parser = argparse.ArgumentParser() # Get the initial condition 12 | # 1 - Dam Break; 2 - Rising Bubble; 3 - Droping liquid 13 | parser.add_argument('-ic', type=int, choices=[1, 2, 3], default=1) 14 | parser.add_argument('-s', action='store_true') 15 | args = parser.parse_args() 16 | initial_condition = args.ic 17 | SAVE_FIG = args.s 18 | 19 | nx = 80 # Number of grid points in the x direction 20 | ny = 80 # Number of grid points in the y direction 21 | 22 | Lx = 0.1 # The length of the domain 23 | Ly = 0.1 # The width of the domain 24 | rho_l = 1000.0 25 | rho_g = 50.0 26 | nu_l = 1.0e-6 # kinematic viscosity, nu = mu / rho 27 | nu_g = 1.5e-5 28 | sigma = ti.field(dtype=float, shape=()) 29 | sigma[None] = 0.007 30 | gx = 0 31 | gy = -1000 32 | 33 | dt = 4e-6 # Use smaller dt for higher density ratio 34 | eps = 1e-6 # Threshold used in vfconv and f post processings 35 | 36 | MAX_TIME_STEPS = 1000 37 | MAX_ITER = 20 38 | OPT_ITER = 100 39 | learning_rate = 0.02 40 | # Mesh information 41 | imin = 1 42 | imax = imin + nx - 1 43 | jmin = 1 44 | jmax = jmin + ny - 1 45 | x = ti.field(float, shape=imax + 3) 46 | y = ti.field(float, shape=jmax + 3) 47 | xnp = np.hstack((0.0, np.linspace(0, Lx, nx + 1), Lx)).astype(np.float32) # [0, 0, ... 1, 1] 48 | x.from_numpy(xnp) 49 | ynp = np.hstack((0.0, np.linspace(0, Ly, ny + 1), Ly)).astype(np.float32) # [0, 0, ... 1, 1] 50 | y.from_numpy(ynp) 51 | dx = x[imin + 2] - x[imin + 1] 52 | dy = y[jmin + 2] - y[jmin + 1] 53 | dxi = 1 / dx 54 | dyi = 1 / dy 55 | 56 | # Field shapes 57 | field_shape = (imax + 2, jmax + 2, MAX_TIME_STEPS) 58 | # p_shape = (imax + 2, jmax + 2, MAX_TIME_STEPS * (MAX_ITER + 1)) 59 | p_shape = (imax + 2, jmax + 2, MAX_TIME_STEPS) 60 | 61 | # Variables for VOF function 62 | F = ti.field(float, shape=(imax + 2, jmax + 2, 2 * MAX_TIME_STEPS + 1), needs_grad=True) 63 | Ftd_x = ti.field(float, shape=field_shape, needs_grad=True) 64 | Ftd_y = ti.field(float, shape=field_shape, needs_grad=True) 65 | Ftarget = ti.field(float, shape=(field_shape[0], field_shape[1]), needs_grad=True) 66 | loss = ti.field(float, shape=(), needs_grad=True) 67 | 68 | ax = ti.field(float, shape=field_shape, needs_grad=True) 69 | ay = ti.field(float, shape=field_shape, needs_grad=True) 70 | cx = ti.field(float, shape=field_shape, needs_grad=True) 71 | cy = ti.field(float, shape=field_shape, needs_grad=True) 72 | rp_x = ti.field(float, shape=field_shape, needs_grad=True) 73 | rm_x = ti.field(float, shape=field_shape, needs_grad=True) 74 | rp_y = ti.field(float, shape=field_shape, needs_grad=True) 75 | rm_y = ti.field(float, shape=field_shape, needs_grad=True) 76 | 77 | # Variables for N-S equation 78 | u = ti.field(float, shape=field_shape, needs_grad=True) 79 | v = ti.field(float, shape=field_shape, needs_grad=True) 80 | u_star = ti.field(float, shape=field_shape, needs_grad=True) 81 | v_star = ti.field(float, shape=field_shape, needs_grad=True) 82 | 83 | # Pressure field shape should be different 84 | p = ti.field(float, shape=p_shape, needs_grad=True) 85 | p_tmp = ti.field(float, shape=(p_shape[0], p_shape[1]), needs_grad=False) 86 | rhs = ti.field(float, shape=field_shape, needs_grad=True) 87 | rhs_tmp = ti.field(float, shape=(field_shape[0], field_shape[1]), needs_grad=False) 88 | rho = ti.field(float, shape=field_shape, needs_grad=True) 89 | nu = ti.field(float, shape=field_shape, needs_grad=True) 90 | V = ti.Vector.field(2, dtype=float, shape=(field_shape[0], field_shape[1])) # For displaying velocity field 91 | 92 | # Variables for interface reconstruction 93 | mx1 = ti.field(float, shape=field_shape, needs_grad=True) 94 | my1 = ti.field(float, shape=field_shape, needs_grad=True) 95 | mx2 = ti.field(float, shape=field_shape, needs_grad=True) 96 | my2 = ti.field(float, shape=field_shape, needs_grad=True) 97 | mx3 = ti.field(float, shape=field_shape, needs_grad=True) 98 | my3 = ti.field(float, shape=field_shape, needs_grad=True) 99 | mx4 = ti.field(float, shape=field_shape, needs_grad=True) 100 | my4 = ti.field(float, shape=field_shape, needs_grad=True) 101 | mxsum = ti.field(float, shape=field_shape, needs_grad=True) 102 | mysum = ti.field(float, shape=field_shape, needs_grad=True) 103 | mx = ti.field(float, shape=field_shape, needs_grad=True) 104 | my = ti.field(float, shape=field_shape, needs_grad=True) 105 | kappa = ti.field(float, shape=field_shape, needs_grad=True) # interface curvature 106 | magnitude = ti.field(float, shape=field_shape, needs_grad=True) 107 | 108 | # For visualization 109 | resolution = (800, 400) 110 | rgb_buf = ti.field(dtype=float, shape=(2 * field_shape[0], field_shape[1])) 111 | 112 | print(f'>>> A VOF solver written in Taichi; Press q to exit.') 113 | print(f'>>> Grid resolution: {nx} x {ny}, dt = {dt:4.2e}') 114 | print(f'>>> Density ratio: {rho_l / rho_g : 4.2f}, gravity : {gy : 4.2f}, sigma : {sigma[None] : 6.3f}') 115 | print(f'>>> Viscosity ratio: {nu_l / nu_g : 4.2f}') 116 | print(f'>>> Please wait a few seconds to let the kernels compile...') 117 | 118 | 119 | @ti.func 120 | def find_area(i, j, cx, cy, r): 121 | a = 0.0 122 | xcoord_ct = (i - imin) * dx + dx / 2 123 | ycoord_ct = (j - jmin) * dy + dy / 2 124 | 125 | xcoord_lu = xcoord_ct - dx / 2 126 | ycoord_lu = ycoord_ct + dy / 2 127 | 128 | xcoord_ld = xcoord_ct - dx / 2 129 | ycoord_ld = ycoord_ct - dy / 2 130 | 131 | xcoord_ru = xcoord_ct + dx / 2 132 | ycoord_ru = ycoord_ct + dy / 2 133 | 134 | xcoord_rd = xcoord_ct + dx / 2 135 | ycoord_rd = ycoord_ct - dy / 2 136 | 137 | dist_ct = ti.sqrt((xcoord_ct - cx) ** 2 + (ycoord_ct - cy) ** 2) 138 | dist_lu = ti.sqrt((xcoord_lu - cx) ** 2 + (ycoord_lu - cy) ** 2) 139 | dist_ld = ti.sqrt((xcoord_ld - cx) ** 2 + (ycoord_ld - cy) ** 2) 140 | dist_ru = ti.sqrt((xcoord_ru - cx) ** 2 + (ycoord_ru - cy) ** 2) 141 | dist_rd = ti.sqrt((xcoord_rd - cx) ** 2 + (ycoord_rd - cy) ** 2) 142 | 143 | if dist_lu > r and dist_ld > r and dist_ru > r and dist_rd > r: 144 | a = 1.0 145 | elif dist_lu < r and dist_ld < r and dist_ru < r and dist_rd < r: 146 | a = 0.0 147 | else: 148 | a = 0.5 + 0.5 * (dist_ct - r) / (ti.sqrt(2.0) * dx) 149 | a = var(a, 0, 1) 150 | 151 | return a 152 | 153 | 154 | @ti.kernel 155 | def set_init_F(ic:ti.i32): 156 | # Sets the initial volume fraction 157 | if ic == 1: # Dambreak 158 | x1 = Lx / 3 * 1.0 159 | x2 = Lx / 3 * 2.0 160 | y1 = 0.0 161 | y2 = Ly / 2 162 | r = Ly / 4 163 | for i, j in ti.ndrange(imax + 2, jmax + 2): 164 | if (x[i] >= x1) and (x[i] <= x2) and (y[j] >= y1) and (y[j] <= y2): 165 | Ftarget[i, j] = 1.0 166 | 167 | elif ic == 2: # Rising bubble 168 | for i, j in ti.ndrange(imax + 2, jmax + 2): 169 | r = Lx / 12 170 | cx, cy = Lx / 2, Ly / 2 171 | Ftarget[i, j] = find_area(i, j, cx, cy, r) 172 | F[i, j, 0] = 1.0 173 | 174 | elif ic == 3: # Liquid drop 175 | for i, j in ti.ndrange(imax + 2, jmax + 2): 176 | r = Lx / 12 177 | cx, cy = Lx / 2, Ly / 2 178 | Ftarget[i, j] = 1.0 - find_area(i, j, cx, cy, r) 179 | 180 | 181 | @ti.kernel 182 | def set_pixel(x:ti.f32, y:ti.f32, f:ti.template()): 183 | xcord = ti.i32(x * imax) 184 | ycord = ti.i32(y * jmax) 185 | for i, j in ti.ndrange((xcord-2, xcord + 2),(ycord-2, ycord+2) ): 186 | if i >= 0 and j >= 0: 187 | f[i, j] = 1.0 188 | 189 | 190 | def set_init_by_paint(): 191 | gui = ti.GUI("Paint your initial", ) 192 | while gui.running: 193 | gui.contour(Ftarget) 194 | gui.get_event() 195 | if gui.is_pressed(ti.GUI.ESCAPE): 196 | gui.running = False 197 | if gui.is_pressed(ti.GUI.LMB): 198 | x, y = gui.get_cursor_pos() 199 | set_pixel(x, y, Ftarget) 200 | gui.show() 201 | 202 | 203 | @ti.kernel 204 | def set_BC(t:ti.i32): 205 | for i in ti.ndrange(imax + 2): 206 | # bottom: slip 207 | u[i, jmin - 1, t] = u[i, jmin, t] 208 | v[i, jmin, t] = 0 209 | F[i, jmin - 1, 2 * t] = F[i, jmin, 2 * t] 210 | p[i, jmin - 1, t] = p[i, jmin, t] 211 | rho[i, jmin - 1, t] = rho[i, jmin, t] 212 | # top: open 213 | u[i, jmax + 1, t] = u[i, jmax, t] 214 | v[i, jmax + 1, t] = 0 #v[i, jmax, t] 215 | F[i, jmax + 1, 2 * t] = F[i, jmax, 2 * t] 216 | p[i, jmax + 1, t] = p[i, jmax, t] 217 | rho[i, jmax + 1, t] = rho[i, jmax, t] 218 | for j in ti.ndrange(jmax + 2): 219 | # left: slip 220 | u[imin, j, t] = 0 221 | v[imin - 1, j, t] = v[imin, j, t] 222 | F[imin - 1, j, 2 * t] = F[imin, j, 2 * t] 223 | p[imin - 1, j, t] = p[imin, j, t] 224 | rho[imin - 1, j, t] = rho[imin, j, t] 225 | # right: slip 226 | u[imax + 1, j, t] = 0 227 | v[imax + 1, j, t] = v[imax, j, t] 228 | F[imax + 1, j, 2 * t] = F[imax, j, 2 * t] 229 | p[imax + 1, j, t] = p[imax, j, t] 230 | rho[imax + 1, j, t] = rho[imax, j, t] 231 | 232 | 233 | @ti.func 234 | def var(a, b, c): # Find the median of a,b, and c 235 | center = a + b + c - ti.max(a, b, c) - ti.min(a, b, c) 236 | return center 237 | 238 | 239 | @ti.kernel 240 | def cal_nu_rho(t:ti.i32): 241 | for i, j in ti.ndrange(field_shape[0], field_shape[1]): 242 | F = var(0.0, 1.0, F[i, j, 2 * t]) 243 | rho[i, j, t] = rho_g * (1 - F) + rho_l * F 244 | nu[i, j, t] = nu_l * F + nu_g * (1.0 - F) 245 | 246 | 247 | @ti.kernel 248 | def advect_upwind(t:ti.i32): 249 | for i, j in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1)): 250 | v_here = 0.25 * (v[i - 1, j, t] + v[i - 1, j + 1, t] + v[i, j, t] + v[i, j + 1, t]) 251 | dudx = (u[i, j, t] - u[i - 1, j, t]) * dxi if u[i, j, t] > 0 else (u[i + 1, j, t] - u[i, j, t]) * dxi 252 | dudy = (u[i, j, t] - u[i, j - 1, t]) * dyi if v_here > 0 else (u[i, j + 1, t] - u[i, j, t]) * dyi 253 | kappa_ave = (kappa[i, j, t] + kappa[i - 1, j, t]) / 2.0 254 | fx_kappa = - sigma[None] * (F[i, j, 2 * t] - F[i - 1, j, 2 * t]) * kappa_ave / dx # F(2*t) is F at t time step 255 | u_star[i, j, t] = ( 256 | u[i, j, t] + dt * 257 | (nu[i, j, t] * (u[i - 1, j, t] - 2 * u[i, j, t] + u[i + 1, j, t]) * dxi**2 258 | + nu[i, j, t] * (u[i, j - 1, t] - 2 * u[i, j, t] + u[i, j + 1, t]) * dyi**2 259 | - u[i, j, t] * dudx - v_here * dudy 260 | + gx + fx_kappa * 2 / (rho[i, j, t] + rho[i - 1, j, t])) 261 | ) 262 | for i, j in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1)): 263 | u_here = 0.25 * (u[i, j - 1, t] + u[i, j, t] + u[i + 1, j - 1, t] + u[i + 1, j, t]) 264 | dvdx = (v[i, j, t] - v[i - 1, j, t]) * dxi if u_here > 0 else (v[i + 1, j, t] - v[i, j, t]) * dxi 265 | dvdy = (v[i, j, t] - v[i, j - 1, t]) * dyi if v[i, j, t] > 0 else (v[i, j + 1, t] - v[i, j, t]) * dyi 266 | kappa_ave = (kappa[i, j, t] + kappa[i, j - 1, t]) / 2.0 267 | fy_kappa = - sigma[None] * (F[i, j, 2 * t] - F[i, j - 1, 2 * t]) * kappa_ave / dy 268 | v_star[i, j, t] = ( 269 | v[i, j, t] + dt * 270 | (nu[i, j, t] * (v[i - 1, j, t] - 2 * v[i, j, t] + v[i + 1, j, t]) * dxi**2 271 | + nu[i, j, t] * (v[i, j - 1, t] - 2 * v[i, j, t] + v[i, j + 1, t]) * dyi**2 272 | - u_here * dvdx - v[i, j, t] * dvdy 273 | + gy + fy_kappa * 2 / (rho[i, j, t] + rho[i, j - 1, t])) 274 | ) 275 | 276 | 277 | @ti.kernel 278 | def cal_velocity_div(t:ti.i32): 279 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 280 | rhs[i, j, t] = rho[i, j, t] / dt * \ 281 | ((u_star[i + 1, j, t] - u_star[i, j, t]) * dxi + 282 | (v_star[i, j + 1, t] - v_star[i, j, t]) * dyi) 283 | 284 | 285 | @ti.kernel 286 | def solve_p_jacobi(t:ti.i32): 287 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 288 | ae = dxi ** 2 if i != imax else 0.0 289 | aw = dxi ** 2 if i != imin else 0.0 290 | an = dyi ** 2 if j != jmax else 0.0 291 | a_s = dyi ** 2 if j != jmin else 0.0 292 | ap = - 1.0 * (ae + aw + an + a_s) 293 | p[i, j, t + 1] = (rhs[i, j, t] \ 294 | - ae * p_tmp[i + 1, j] \ 295 | - aw * p_tmp[i - 1, j] \ 296 | - an * p_tmp[i, j + 1] \ 297 | - a_s * p_tmp[i, j - 1]\ 298 | ) / ap 299 | for i, j in p_tmp: 300 | p_tmp[i, j] = p[i, j, t + 1] 301 | 302 | 303 | @ti.kernel 304 | def solve_p_grad(t:ti.i32): 305 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 306 | ae = dxi ** 2 if i != imax else 0.0 307 | aw = dxi ** 2 if i != imin else 0.0 308 | an = dyi ** 2 if j != jmax else 0.0 309 | a_s = dyi ** 2 if j != jmin else 0.0 310 | ap = - 1.0 * (ae + aw + an + a_s) 311 | rhs_tmp[i, j] = (p.grad[i, j, t + 1] \ 312 | - ae * rhs.grad[i + 1, j, t] \ 313 | - aw * rhs.grad[i - 1, j, t] \ 314 | - an * rhs.grad[i, j + 1, t] \ 315 | - a_s * rhs.grad[i, j - 1, t]\ 316 | ) / ap 317 | for i, j in rhs_tmp: 318 | rhs.grad[i, j, t] = rhs_tmp[i, j] 319 | 320 | 321 | @ti.ad.grad_replaced 322 | def solve_p_iter(t): 323 | for _ in range(MAX_ITER): 324 | solve_p_jacobi(t) 325 | 326 | 327 | @ti.ad.grad_for(solve_p_iter) 328 | def solve_p_grad_iter(t): 329 | for _ in range(MAX_ITER): 330 | solve_p_grad(t) 331 | 332 | 333 | @ti.kernel 334 | def update_uv(t:ti.i32): 335 | for i, j in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1)): 336 | r = (rho[i, j, t] + rho[i - 1, j, t]) * 0.5 337 | u[i, j, t + 1] = u_star[i, j, t] \ 338 | - dt / r * \ 339 | (p[i, j, t + 1] - p[i - 1, j, t + 1]) * dxi 340 | if u[i, j, t + 1] * dt > 0.25 * dx: 341 | print(f'U velocity courant number > 1, u[{i},{j},{t+1}] = {u[i, j, t+1]}') 342 | for i, j in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1)): 343 | r = (rho[i, j, t] + rho[i, j - 1, t]) * 0.5 344 | v[i, j, t + 1] = v_star[i, j, t] \ 345 | - dt / r \ 346 | * (p[i, j, t + 1] - p[i, j - 1, t + 1]) * dyi 347 | if v[i, j, t + 1] * dt > 0.25 * dy: 348 | print(f'V velocity courant number > 1, v[{i},{j},{t+1}] = {v[i,j,t+1]}') 349 | 350 | 351 | @ti.kernel 352 | def get_normal_young(t:ti.i32): 353 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 354 | # Points in between the outermost boundaries 355 | mx1[i, j, t] = -1 / (2 * dx) * (F[i + 1, j + 1, 2 * t] + F[i + 1, j, 2 * t] - F[i, j + 1, 2 * t] - F[i, j, 2 * t]) 356 | my1[i, j, t] = -1 / (2 * dy) * (F[i + 1, j + 1, 2 * t] - F[i + 1, j, 2 * t] + F[i, j + 1, 2 * t] - F[i, j, 2 * t]) 357 | mx2[i, j, t] = -1 / (2 * dx) * (F[i + 1, j, 2 * t] + F[i + 1, j - 1, 2 * t] - F[i, j, 2 * t] - F[i, j - 1, 2 * t]) 358 | my2[i, j, t] = -1 / (2 * dy) * (F[i + 1, j, 2 * t] - F[i + 1, j - 1, 2 * t] + F[i, j, 2 * t] - F[i, j - 1, 2 * t]) 359 | mx3[i, j, t] = -1 / (2 * dx) * (F[i, j, 2 * t] + F[i, j - 1, 2 * t] - F[i - 1, j, 2 * t] - F[i - 1, j - 1, 2 * t]) 360 | my3[i, j, t] = -1 / (2 * dy) * (F[i, j, 2 * t] - F[i, j - 1, 2 * t] + F[i - 1, j, 2 * t] - F[i - 1, j - 1, 2 * t]) 361 | mx4[i, j, t] = -1 / (2 * dx) * (F[i, j + 1, 2 * t] + F[i, j, 2 * t] - F[i - 1, j + 1, 2 * t] - F[i - 1, j, 2 * t]) 362 | my4[i, j, t] = -1 / (2 * dy) * (F[i, j + 1, 2 * t] - F[i, j, 2 * t] + F[i - 1, j + 1, 2 * t] - F[i - 1, j, 2 * t]) 363 | # Summing of mx and my components for normal vector 364 | mxsum[i, j, t] = (mx1[i, j, t] + mx2[i, j, t] + mx3[i, j, t] + mx4[i, j, t]) / 4 365 | mysum[i, j, t] = (my1[i, j, t] + my2[i, j, t] + my3[i, j, t] + my4[i, j, t]) / 4 366 | # Normalizing the normal vector into unit vectors 367 | if abs(mxsum[i, j, t]) < 1e-10 and abs(mysum[i, j, t])< 1e-10: 368 | mx[i, j, t] = mxsum[i, j, t] 369 | my[i, j, t] = mysum[i, j, t] 370 | else: 371 | magnitude[i, j, t] = ti.sqrt(mxsum[i, j, t] * mxsum[i, j, t] + mysum[i, j, t] * mysum[i, j, t]) 372 | mx[i, j, t] = mxsum[i, j, t] / magnitude[i, j, t] 373 | my[i, j, t] = mysum[i, j, t] / magnitude[i, j, t] 374 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 375 | kappa[i, j, t] = -(1 / dx / 2 * (mx[i + 1, j, t] - mx[i - 1, j, t]) + \ 376 | 1 / dy / 2 * (my[i, j + 1, t] - my[i, j - 1, t])) 377 | 378 | 379 | def solve_VOF_rudman(t): 380 | if t % 2 == 0: 381 | fct_y_sweep(t, 0, 1e-6) 382 | fct_x_sweep(t, 1, 1e-6) 383 | else: 384 | fct_x_sweep(t, 0, 1e-6) 385 | fct_y_sweep(t, 1, 1e-6) 386 | 387 | 388 | @ti.kernel 389 | def fct_x_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 390 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 391 | dv = dx * dy - dt * dy * (u[i + 1, j, t + 1] - u[i, j, t + 1]) 392 | fl_L = u[i, j, t + 1] * dt * F[i - 1, j, 2 * t + offset] if u[i, j, t + 1] >= 0 else u[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 393 | fr_L = u[i + 1, j, t + 1] * dt * F[i, j, 2 * t + offset] if u[i + 1, j, t + 1] >= 0 else u[i + 1, j, t + 1] * dt * F[i + 1, j, 2 * t + offset] 394 | Ftd_x[i, j, t] = F[i, j, 2 * t + offset] + (fl_L - fr_L) * dy / (dx * dy) * dx * dy / dv 395 | 396 | for i, j in ti.ndrange((imin, imax + 2), (jmin, jmax + 1)): 397 | fl_L = u[i, j, t + 1] * dt * F[i - 1, j, 2 * t + offset] if u[i, j, t + 1] >= 0 else u[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 398 | fl_H = u[i, j, t + 1] * dt * F[i - 1, j, 2 * t + offset] if u[i, j, t + 1] <= 0 else u[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 399 | ax[i, j, t] = fl_H - fl_L 400 | 401 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 402 | fmax = ti.max(Ftd_x[i, j, t], Ftd_x[i - 1, j, t], Ftd_x[i + 1, j, t]) 403 | fmin = ti.min(Ftd_x[i, j, t], Ftd_x[i - 1, j, t], Ftd_x[i + 1, j, t]) 404 | 405 | pp = ti.max(0, ax[i, j, t]) - ti.min(0, ax[i + 1, j, t]) 406 | qp = (fmax - Ftd_x[i, j, t]) * dx 407 | if pp > eps: 408 | rp_x[i, j, t] = ti.min(1, qp / pp) 409 | else: 410 | rp_x[i, j, t] = 0.0 411 | 412 | pm = ti.max(0, ax[i + 1, j, t]) - ti.min(0, ax[i, j, t]) 413 | qm = (Ftd_x[i, j, t] - fmin) * dx 414 | if pm > eps: 415 | rm_x[i, j, t] = ti.min(1, qm / pm) 416 | else: 417 | rm_x[i, j, t] = 0.0 418 | 419 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 420 | if ax[i + 1, j, t] >= 0: 421 | cx[i + 1, j, t] = ti.min(rp_x[i + 1, j, t], rm_x[i, j, t]) 422 | else: 423 | cx[i + 1, j, t] = ti.min(rp_x[i, j, t], rm_x[i + 1, j, t]) 424 | 425 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 426 | dv = dx * dy - dt * dy * (u[i + 1, j, t + 1] - u[i, j, t + 1]) 427 | F[i, j, 2 * t + offset + 1] = Ftd_x[i, j, t] - ((ax[i + 1, j, t] * cx[i + 1, j, t] - \ 428 | ax[i, j, t] * cx[i, j, t]) / dy) * dx * dy / dv 429 | 430 | 431 | @ti.kernel 432 | def fct_y_sweep(t:ti.i32, offset:ti.i32, eps:ti.f32): 433 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 434 | dv = dx * dy - dt * dx * (v[i, j + 1, t + 1] - v[i, j, t + 1]) 435 | ft_L = v[i, j + 1, t + 1] * dt * F[i, j, 2 * t + offset] if v[i, j + 1, t + 1] >= 0 else v[i, j + 1, t + 1] * dt * F[i, j + 1, 2 * t + offset] 436 | fb_L = v[i, j, t + 1] * dt * F[i, j - 1, 2 * t + offset] if v[i, j, t + 1] >= 0 else v[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 437 | Ftd_y[i, j, t] = F[i, j, 2 * t + offset] + (fb_L - ft_L) * dy / (dx * dy) * dx * dy / dv 438 | 439 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 2)): 440 | fb_L = v[i, j, t + 1] * dt * F[i, j - 1, 2 * t + offset] if v[i, j, t + 1] >= 0 else v[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 441 | fb_H = v[i, j, t + 1] * dt * F[i, j - 1, 2 * t + offset] if v[i, j, t + 1] <= 0 else v[i, j, t + 1] * dt * F[i, j, 2 * t + offset] 442 | ay[i, j, t] = fb_H - fb_L 443 | 444 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 445 | fmax = ti.max(Ftd_y[i, j, t], Ftd_y[i, j - 1, t], Ftd_y[i, j + 1, t]) 446 | fmin = ti.min(Ftd_y[i, j, t], Ftd_y[i, j - 1, t], Ftd_y[i, j + 1, t]) 447 | 448 | # eps = 1e-4 449 | pp = ti.max(0, ay[i, j, t]) - ti.min(0, ay[i, j + 1, t]) 450 | qp = (fmax - Ftd_y[i, j, t]) * dx 451 | if pp > eps: 452 | rp_y[i, j, t] = ti.min(1, qp / pp) 453 | else: 454 | rp_y[i, j, t] = 0.0 455 | 456 | pm = ti.max(0, ay[i, j + 1, t]) - ti.min(0, ay[i, j, t]) 457 | qm = (Ftd_y[i, j, t] - fmin) * dx 458 | if pm > eps: 459 | rm_y[i, j, t] = ti.min(1, qm / pm) 460 | else: 461 | rm_y[i, j, t] = 0.0 462 | 463 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 464 | if ay[i, j + 1, t] >= 0: 465 | cy[i, j + 1, t] = ti.min(rp_y[i, j + 1, t], rm_y[i, j, t]) 466 | else: 467 | cy[i, j + 1, t] = ti.min(rp_y[i, j, t], rm_y[i, j + 1, t]) 468 | 469 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 470 | dv = dx * dy - dt * dx * (v[i, j + 1, t + 1] - v[i, j, t + 1]) 471 | F[i, j, 2 * t + offset + 1] = Ftd_y[i, j, t] - ((ay[i, j + 1, t] * cy[i, j + 1, t] -\ 472 | ay[i, j, t] * cy[i, j, t]) / dy) * dx * dy / dv 473 | 474 | 475 | 476 | @ti.kernel 477 | def post_process_f(t:ti.i32): 478 | for i, j in ti.ndrange((imin, imax+1), (jmin, jmax+1)): 479 | F[i, j, 2 * t + 2] = var(F[i, j, 2 * t + 2], 0, 1) 480 | 481 | 482 | @ti.ad.no_grad 483 | @ti.kernel 484 | def get_field_to_buf(src:ti.template(), t:ti.i32): 485 | for i, j in Ftarget: 486 | rgb_buf[i, j] = src[i, j, t] 487 | for i, j in Ftarget: 488 | rgb_buf[i + imax + 1, j] = Ftarget[i, j] 489 | 490 | 491 | @ti.ad.no_grad 492 | @ti.kernel 493 | def get_vnorm_field(t:ti.i32): 494 | for i, j in rgb_buf: 495 | rgb_buf[i, j] = ti.sqrt(u[i, j, t] ** 2 + v[i, j, t] ** 2) 496 | 497 | 498 | @ti.ad.no_grad 499 | @ti.kernel 500 | def interp_velocity(t:ti.i32): 501 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 502 | V[i, j] = ti.Vector([(u[i, j, t] + u[i + 1, j, t])/2, (v[i, j, t] + v[i, j + 1, t])/2]) 503 | 504 | 505 | @ti.kernel 506 | def compute_loss(): 507 | for i, j in ti.ndrange(imax + 2, jmax + 2): 508 | loss[None] += ti.abs(Ftarget[i, j] - F[i, j, 2 * MAX_TIME_STEPS - 2]) 509 | 510 | 511 | @ti.kernel 512 | def apply_grad(): 513 | for i, j in ti.ndrange((1, imax + 1), (1, jmax + 1)): 514 | if ti.abs(F.grad[i, j, 0]) < 5.: 515 | F[i, j, 0] -= learning_rate * F.grad[i, j, 0] 516 | F[i, j, 0] = var(0, 1, F[i, j, 0]) 517 | 518 | 519 | def forward(): 520 | vis_option = 0 # Tag for display 521 | for istep in range(MAX_TIME_STEPS - 1): 522 | for e in gui.get_events(gui.RELEASE): 523 | if e.key == gui.SPACE: 524 | vis_option += 1 525 | elif e.key == 'q': 526 | gui.running = False 527 | 528 | # Calculate initial F 529 | cal_nu_rho(istep) 530 | get_normal_young(istep) 531 | 532 | # Advection 533 | advect_upwind(istep) 534 | ''' 535 | Not necessary; boundary u,v = 0 already, and no update be made on those boundary 536 | For p solving, only those u,v=0 will be used. 537 | ''' 538 | # set_BC(istep) 539 | 540 | # Calculate the velocity divergence -> rhs 541 | cal_velocity_div(istep) 542 | # Pressure projection 543 | # for iter in range(MAX_ITER): 544 | # solve_p_jacobi(istep) 545 | solve_p_iter(istep) 546 | # copy_p_field(istep, iter) # Don't need copy in kernel replaced version 547 | 548 | # Velocity correction 549 | update_uv(istep) 550 | ''' 551 | Not necessary. For VOF advection, only u, v = 0 will be used, which are untouched. 552 | And F on the boundary is set at the end of previous step's set_BC 553 | ''' 554 | # set_BC(istep + 1) 555 | 556 | # Advect the VOF function 557 | solve_VOF_rudman(istep) 558 | post_process_f(istep) # Post-processing violates GDAR, but necessary for stablize. 559 | set_BC(istep + 1) 560 | 561 | # Visualization 562 | num_options = 5 563 | plot_contour = ti.ad.no_grad(gui.contour) 564 | plot_vector = ti.ad.no_grad(gui.vector_field) 565 | if (istep % nstep) == 0: # Output data every steps 566 | if vis_option % num_options == 0: # Display VOF distribution 567 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying VOF field.') 568 | get_field_to_buf(F, 2 * (istep + 1)) 569 | plot_contour(rgb_buf) 570 | 571 | if vis_option % num_options == 1: # Display the u field 572 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying u velocity.') 573 | get_field_to_buf(u, istep) 574 | plot_contour(rgb_buf) 575 | 576 | if vis_option % num_options == 2: # Display the v field 577 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying v velocity.') 578 | get_field_to_buf(v, istep) 579 | plot_contour(rgb_buf) 580 | 581 | if vis_option % num_options == 3: # Display velocity norm 582 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity norm.') 583 | get_vnorm_field(istep) 584 | plot_contour(rgb_buf) 585 | 586 | if vis_option % num_options == 4: # Display velocity vectors 587 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity vectors.') 588 | interp_velocity(istep) 589 | plot_vector(V, arrow_spacing=2, color=0x000000) 590 | 591 | gui.show(f'./output/{opt:03d}-{istep:04d}.png') 592 | 593 | # Compute loss as the last step of forward() pass 594 | compute_loss() 595 | 596 | 597 | # Start main script 598 | istep = 0 599 | nstep = 20 # Interval to update GUI 600 | set_init_F(initial_condition) # Set initial VOF by fixed shape 601 | # set_init_by_paint() # Set initial VOF by user painting 602 | os.makedirs('output', exist_ok=True) # Make dir for output 603 | gui = ti.GUI('VOF Solver', resolution, background_color=0xFFFFFF) 604 | vis_option = 0 605 | 606 | for opt in range(OPT_ITER): 607 | print(f'>>> >>> Optimization cycle {opt}') 608 | with ti.ad.Tape(loss): 609 | forward() 610 | print(f'>>> >>> Current total loss is {loss[None]}') 611 | apply_grad() # Apply gradient should be outside the Tape() 612 | print(f'>>> >>> Gradient applied.') 613 | -------------------------------------------------------------------------------- /3dvof.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import matplotlib.cm as cm 5 | import argparse 6 | import os 7 | import flow_visualization as fv 8 | from pyevtk.hl import gridToVTK 9 | 10 | ti.init(arch=ti.gpu, default_fp=ti.f32) # Set default fp so that float=ti.f32 11 | 12 | parser = argparse.ArgumentParser() # Get the initial condition 13 | # 1 - Dam Break; 2 - Rising Bubble; 3 - Droping liquid 14 | parser.add_argument('-ic', type=int, choices=[1, 2, 3], default=1) 15 | parser.add_argument('-s', action='store_true') 16 | args = parser.parse_args() 17 | initial_condition = args.ic 18 | SAVE_FIG = args.s 19 | 20 | nx = 200 # Number of grid points in the x direction 21 | ny = 200 # Number of grid points in the y direction 22 | nz = 200 # Number of grid points in the z direction 23 | 24 | Lx = 0.1 # The length of the domain 25 | Ly = 0.1 # The width of the domain 26 | Lz = 0.1 27 | rho_l = 1000.0 28 | rho_g = 50.0 29 | nu_l = 1.0e-6 # kinematic viscosity, nu = mu / rho 30 | nu_g = 1.5e-5 31 | sigma = ti.field(dtype=float, shape=()) 32 | sigma[None] = 0.007 33 | gx = 0 34 | gy = -5 35 | gz = 0 36 | 37 | dt = 4e-6 # Use smaller dt for higher density ratio 38 | eps = 1e-6 # Threshold used in vfconv and f post processings 39 | 40 | # Mesh information 41 | imin = 1 42 | imax = imin + nx - 1 43 | jmin = 1 44 | jmax = jmin + ny - 1 45 | kmin = 1 46 | kmax = kmin + nz - 1 47 | 48 | x = ti.field(float, shape=imax + 3) 49 | y = ti.field(float, shape=jmax + 3) 50 | z = ti.field(float, shape=kmax + 3) 51 | xnp = np.hstack((0.0, np.linspace(0, Lx, nx + 1), Lx)).astype(np.float32) # [0, 0, ... 1, 1] 52 | x.from_numpy(xnp) 53 | ynp = np.hstack((0.0, np.linspace(0, Ly, ny + 1), Ly)).astype(np.float32) # [0, 0, ... 1, 1] 54 | y.from_numpy(ynp) 55 | znp = np.hstack((0.0, np.linspace(0, Lz, nz + 1), Lz)).astype(np.float32) # [0, 0, ... 1, 1] 56 | z.from_numpy(znp) 57 | 58 | # For vtk file export 59 | xcor = np.linspace(0.0, 1.0, imax + 2).astype(np.float32) 60 | ycor = np.linspace(0.0, 1.0, jmax + 2).astype(np.float32) 61 | zcor = np.linspace(0.0, 1.0, kmax + 2).astype(np.float32) 62 | 63 | dx = x[imin + 2] - x[imin + 1] 64 | dy = y[jmin + 2] - y[jmin + 1] 65 | dz = z[jmin + 2] - z[jmin + 1] 66 | dxi = 1 / dx 67 | dyi = 1 / dy 68 | dzi = 1 / dz 69 | 70 | field_shape = (imax + 2, jmax + 2, kmax + 2) 71 | # Variables for VOF function 72 | F = ti.field(float, shape=field_shape) 73 | Ftd = ti.field(float, shape=field_shape) 74 | ax = ti.field(float, shape=field_shape) 75 | ay = ti.field(float, shape=field_shape) 76 | az = ti.field(float, shape=field_shape) 77 | cx = ti.field(float, shape=field_shape) 78 | cy = ti.field(float, shape=field_shape) 79 | cz = ti.field(float, shape=field_shape) 80 | 81 | rp = ti.field(float, shape=field_shape) 82 | rm = ti.field(float, shape=field_shape) 83 | 84 | # Variables for N-S equation 85 | u = ti.field(float, shape=field_shape) 86 | v = ti.field(float, shape=field_shape) 87 | w = ti.field(float, shape=field_shape) 88 | u_star = ti.field(float, shape=field_shape) 89 | v_star = ti.field(float, shape=field_shape) 90 | w_star = ti.field(float, shape=field_shape) 91 | p = ti.field(float, shape=field_shape) 92 | pt = ti.field(float, shape=field_shape) 93 | Ap = ti.field(float, shape=field_shape) 94 | rhs = ti.field(float, shape=field_shape) 95 | rho = ti.field(float, shape=field_shape) 96 | nu = ti.field(float, shape=field_shape) 97 | V = ti.Vector.field(3, dtype=float, shape=field_shape) 98 | 99 | # Variables for interface reconstruction 100 | mx1 = ti.field(float, shape=field_shape) 101 | my1 = ti.field(float, shape=field_shape) 102 | mx2 = ti.field(float, shape=field_shape) 103 | my2 = ti.field(float, shape=field_shape) 104 | mx3 = ti.field(float, shape=field_shape) 105 | my3 = ti.field(float, shape=field_shape) 106 | mx4 = ti.field(float, shape=field_shape) 107 | my4 = ti.field(float, shape=field_shape) 108 | mxsum = ti.field(float, shape=field_shape) 109 | mysum = ti.field(float, shape=field_shape) 110 | mx = ti.field(float, shape=field_shape) 111 | my = ti.field(float, shape=field_shape) 112 | kappa = ti.field(float, shape=field_shape) # interface curvature 113 | magnitude = ti.field(float, shape=field_shape) 114 | 115 | # For visualization 116 | resolution = (nx * 2, ny * 2) 117 | rgb_buf = ti.field(dtype=float, shape=resolution) 118 | 119 | print(f'>>> A 3D VOF solver written in Taichi; Press q to exit.') 120 | print(f'>>> Grid resolution: {nx} x {ny} x {nz}, dt = {dt:4.2e}') 121 | print(f'>>> Density ratio: {rho_l / rho_g : 4.2f}, gravity : {gy : 4.2f}, sigma : {sigma[None] : 4.2f}') 122 | print(f'>>> Viscosity ratio: {nu_l / nu_g : 4.2f}') 123 | print(f'>>> Please wait a few seconds to let the kernels compile...') 124 | 125 | 126 | @ti.kernel 127 | def set_init_F(ic:ti.i32): 128 | # Sets the initial volume fraction 129 | if ic == 1: # Dambreak 130 | x1 = 0.0 131 | x2 = Lx / 3 132 | y1 = 0.0 133 | y2 = Ly / 2 134 | z1 = 0.0 135 | z2 = Lz / 3 136 | for i, j, k in ti.ndrange(imax + 2, jmax + 2, kmax + 2): 137 | if (x[i] >= x1) and (x[i] <= x2) and (y[j] >= y1) and (y[j] <= y2) and (z[k] >= z1) and (z[k] <= z2): 138 | F[i, j, k] = 1.0 139 | 140 | 141 | @ti.kernel 142 | def set_BC(): 143 | # TODO: Observe the common pattern of the 3 dimensions; apply abstraction to write more concise code. 144 | for i, k in ti.ndrange(imax + 2, kmax + 2): 145 | # bottom: slip 146 | u[i, jmin - 1, k] = u[i, jmin, k] 147 | v[i, jmin, k] = 0 148 | w[i, jmin - 1, k] = w[i, jmin, k] 149 | F[i, jmin - 1, k] = F[i, jmin, k] 150 | p[i, jmin - 1, k] = p[i, jmin, k] 151 | rho[i, jmin - 1, k] = rho[i, jmin, k] 152 | # top: open 153 | u[i, jmax + 1, k] = u[i, jmax, k] 154 | v[i, jmax + 1, k] = 0 #v[i, jmax] 155 | w[i, jmax + 1, k] = w[i, jmax, k] 156 | F[i, jmax + 1, k] = F[i, jmax, k] 157 | p[i, jmax + 1, k] = p[i, jmax, k] 158 | rho[i, jmax + 1, k] = rho[i, jmax, k] 159 | 160 | for j, k in ti.ndrange(jmax + 2, kmax + 2): 161 | # left: slip 162 | u[imin, j, k] = 0 163 | v[imin - 1, j, k] = v[imin, j, k] 164 | w[imin - 1, j, k] = w[imin, j, k] 165 | F[imin - 1, j, k] = F[imin, j, k] 166 | p[imin - 1, j, k] = p[imin, j, k] 167 | rho[imin - 1, j, k] = rho[imin, j, k] 168 | # right: slip 169 | u[imax + 1, j, k] = 0 170 | v[imax + 1, j, k] = v[imax, j, k] 171 | w[imax + 1, j, k] = w[imax, j, k] 172 | F[imax + 1, j, k] = F[imax, j, k] 173 | p[imax + 1, j, k] = p[imax, j, k] 174 | rho[imax + 1, j, k] = rho[imax, j, k] 175 | 176 | for i, j in ti.ndrange(imax + 2, jmax + 2): 177 | # front: slip 178 | u[i, j, kmin - 1] = u[i, j, kmin] 179 | v[i, j, kmin - 1] = v[i, j, kmin] 180 | w[i, j, kmin] = 0 181 | F[i, j, kmin - 1] = F[i, j, kmin] 182 | p[i, j, kmin - 1] = p[i, j, kmin] 183 | rho[i, j, kmin - 1] = rho[i, j, kmin] 184 | # back: slip 185 | u[i, j, kmax + 1] = u[i, j, kmax] 186 | v[i, j, kmax + 1] = v[i, j, kmax] 187 | w[i, j, kmax + 1] = 0 188 | F[i, j, kmax + 1] = F[i, j, kmax] 189 | p[i, j, kmax + 1] = p[i, j, kmax] 190 | rho[i, j, kmax + 1] = rho[i, j, kmax] 191 | 192 | 193 | @ti.func 194 | def var(a, b, c): # Find the median of a,b, and c 195 | center = a + b + c - ti.max(a, b, c) - ti.min(a, b, c) 196 | return center 197 | 198 | 199 | @ti.kernel 200 | def cal_nu_rho(): 201 | for I in ti.grouped(rho): 202 | F = var(0.0, 1.0, F[I]) 203 | rho[I] = rho_g * (1 - F) + rho_l * F 204 | nu[I] = nu_l * F + nu_g * (1.0 - F) 205 | 206 | 207 | @ti.kernel 208 | def advect_upwind(): 209 | # Full 3D version of the momentum equation; upwind advection used. 210 | # 3D curvature is not correct; need to be corrected. 211 | for i, j, k in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 212 | v_here = 0.25 * (v[i - 1, j, k] + v[i - 1, j + 1, k] + v[i, j, k] + v[i, j + 1, k]) 213 | w_here = 0.25 * (w[i - 1, j, k] + w[i - 1, j, k + 1] + w[i, j, k] + w[i, j, k + 1] ) 214 | dudx = (u[i, j, k] - u[i - 1, j, k]) * dxi if u[i, j, k] > 0 else (u[i + 1, j, k] - u[i, j, k]) * dxi 215 | dudy = (u[i, j, k] - u[i, j - 1, k]) * dyi if v_here > 0 else (u[i, j + 1, k] - u[i, j, k]) * dyi 216 | dudz = (u[i, j, k] - u[i, j, k - 1]) * dzi if w_here > 0 else (u[i, j, k + 1] - u[i, j, k]) * dzi 217 | kappa_ave = (kappa[i, j, k] + kappa[i - 1, j, k]) / 2.0 218 | fx_kappa = - sigma[None] * (F[i, j, k] - F[i - 1, j, k]) * kappa_ave / dx 219 | u_star[i, j, k] = ( 220 | u[i, j, k] + dt * 221 | ( nu[i, j, k] * (u[i - 1, j, k] - 2 * u[i, j, k] + u[i + 1, j, k]) * dxi**2 222 | + nu[i, j, k] * (u[i, j - 1, k] - 2 * u[i, j, k] + u[i, j + 1, k]) * dyi**2 223 | + nu[i, j, k] * (u[i, j, k - 1] - 2 * u[i, j, k] + u[i, j, k + 1]) * dzi**2 224 | - u[i, j, k] * dudx - v_here * dudy - w_here * dudz 225 | + gx + fx_kappa * 2 / (rho[i, j, k] + rho[i - 1, j, k])) 226 | ) 227 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1), (kmin, kmax + 1)): 228 | u_here = 0.25 * (u[i, j - 1, k] + u[i, j, k] + u[i + 1, j - 1, k] + u[i + 1, j, k]) 229 | w_here = 0.25 * (w[i, j - 1, k + 1] + w[i, j - 1, k] + w[i, j, k] + w[i, j, k + 1]) 230 | dvdx = (v[i, j, k] - v[i - 1, j, k]) * dxi if u_here > 0 else (v[i + 1, j, k] - v[i, j, k]) * dxi 231 | dvdy = (v[i, j, k] - v[i, j - 1, k]) * dyi if v[i, j, k] > 0 else (v[i, j + 1, k] - v[i, j, k]) * dyi 232 | dvdz = (v[i, j, k] - v[i, j, k - 1]) * dzi if w_here > 0 else (v[i, j, k + 1] - v[i, j, k]) * dzi 233 | kappa_ave = (kappa[i, j, k] + kappa[i, j - 1, k]) / 2.0 234 | fy_kappa = - sigma[None] * (F[i, j, k] - F[i, j - 1, k]) * kappa_ave / dy 235 | v_star[i, j, k] = ( 236 | v[i, j, k] + dt * 237 | ( nu[i, j, k] * (v[i - 1, j, k] - 2 * v[i, j, k] + v[i + 1, j, k]) * dxi**2 238 | + nu[i, j, k] * (v[i, j - 1, k] - 2 * v[i, j, k] + v[i, j + 1, k]) * dyi**2 239 | + nu[i, j, k] * (v[i, j, k - 1] - 2 * v[i, j, k] + v[i, j, k + 1]) * dzi**2 240 | - u_here * dvdx - v[i, j, k] * dvdy - w_here * dvdz 241 | + gy + fy_kappa * 2 / (rho[i, j, k] + rho[i, j - 1, k])) 242 | ) 243 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin + 1, kmax + 1)): 244 | u_here = 0.25 * (u[i + 1, j, k - 1] + u[i, j, k - 1] + u[i + 1, j, k] + u[i, j, k]) 245 | v_here = 0.25 * (v[i, j + 1, k - 1] + v[i, j, k - 1] + v[i, j, k] + v[i, j + 1, k]) 246 | dwdx = (w[i, j, k] - w[i - 1, j, k]) * dxi if u_here > 0 else (w[i + 1, j, k] - w[i, j, k]) * dxi 247 | dwdy = (w[i, j, k] - w[i, j - 1, k]) * dyi if v_here > 0 else (w[i, j + 1, k] - w[i, j, k]) * dyi 248 | dwdz = (w[i, j, k] - w[i, j, k - 1]) * dzi if w[i, j, k] > 0 else (w[i, j, k + 1] - w[i, j, k]) * dzi 249 | kappa_ave = (kappa[i, j, k] + kappa[i, j, k - 1]) / 2.0 250 | fz_kappa = - sigma[None] * (F[i, j, k] - F[i, j, k - 1]) * kappa_ave / dz 251 | w_star[i, j, k] = ( 252 | w[i, j, k] + dt * 253 | ( nu[i, j, k] * (w[i - 1, j, k] - 2 * w[i, j, k] + w[i + 1, j, k]) * dxi**2 254 | + nu[i, j, k] * (w[i, j - 1, k] - 2 * w[i, j, k] + w[i, j + 1, k]) * dyi**2 255 | + nu[i, j, k] * (w[i, j, k - 1] - 2 * w[i, j, k] + w[i, j, k + 1]) * dzi**2 256 | - u_here * dwdx - v_here * dwdy - w[i, j, k] * dwdz 257 | + gz + fz_kappa * 2 / (rho[i, j, k] + rho[i, j, k - 1])) 258 | ) 259 | 260 | 261 | @ti.kernel 262 | def solve_p_jacobi(): 263 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 264 | rhs = rho[i, j, k] / dt * \ 265 | ((u_star[i + 1, j, k] - u_star[i, j, k]) * dxi + 266 | (v_star[i, j + 1, k] - v_star[i, j, k]) * dyi + 267 | (w_star[i, j, k + 1] - w_star[i, j, k]) * dzi) 268 | # No dencorr in the 3d version since dencorr was not used in 2d version either. 269 | ae = dxi ** 2 if i != imax else 0.0 270 | aw = dxi ** 2 if i != imin else 0.0 271 | an = dyi ** 2 if j != jmax else 0.0 272 | a_s = dyi ** 2 if j != jmin else 0.0 273 | af = dzi ** 2 if k != kmax else 0.0 274 | ab = dzi ** 2 if k != kmin else 0.0 275 | ap = - 1.0 * (ae + aw + an + a_s + ab + af) 276 | pt[i, j, k] = (rhs - ae * p[i + 1, j, k] \ 277 | - aw * p[i - 1, j, k] \ 278 | - an * p[i, j + 1, k] \ 279 | - a_s * p[i, j - 1, k]\ 280 | - af * p[i, j, k + 1] \ 281 | - ab * p[i, j, k - 1] ) / ap 282 | for i, j, k in ti.ndrange((imin, imax+1), (jmin, jmax+1), (kmin, kmax + 1)): 283 | p[i, j, k] = pt[i, j, k] 284 | 285 | 286 | @ti.kernel 287 | def update_uv(): 288 | for i, j, k in ti.ndrange((imin + 1, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 289 | r = (rho[i, j, k] + rho[i - 1, j, k]) * 0.5 290 | u[i, j, k] = u_star[i, j, k] - dt / r * (p[i, j, k] - p[i - 1, j, k]) * dxi 291 | if u[i, j, k] * dt > 0.25 * dx: 292 | print(f'U velocity courant number > 1, u[{i},{j},{k}] = {u[i, j, k]}') 293 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin + 1, jmax + 1), (kmin, kmax + 1)): 294 | r = (rho[i, j, k] + rho[i, j - 1, k]) * 0.5 295 | v[i, j, k] = v_star[i, j, k] - dt / r * (p[i, j, k] - p[i, j - 1, k]) * dyi 296 | if v[i, j, k] * dt > 0.25 * dy: 297 | print(f'V velocity courant number > 1, v[{i},{j},{k}] = {v[i, j, k]}') 298 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin + 1, kmax + 1)): 299 | r = (rho[i, j, k] + rho[i, j, k - 1]) * 0.5 300 | w[i, j, k] = w_star[i, j, k] - dt / r * (p[i, j, k] - p[i, j, k - 1]) * dzi 301 | if w[i, j, k] * dt > 0.25 * dx: 302 | print(f'W velocity courant number > 1, w[{i},{j},{k}] = {w[i, j, k]}') 303 | 304 | ''' 305 | @ti.kernel 306 | def get_normal_young(): 307 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 308 | # Points in between the outermost boundaries 309 | mx1[i, j] = -1 / (2 * dx) * (F[i + 1, j + 1] + F[i + 1, j] - F[i, j + 1] - F[i, j]) 310 | my1[i, j] = -1 / (2 * dy) * (F[i + 1, j + 1] - F[i + 1, j] + F[i, j + 1] - F[i, j]) 311 | mx2[i, j] = -1 / (2 * dx) * (F[i + 1, j] + F[i + 1, j - 1] - F[i, j] - F[i, j - 1]) 312 | my2[i, j] = -1 / (2 * dy) * (F[i + 1, j] - F[i + 1, j - 1] + F[i, j] - F[i, j - 1]) 313 | mx3[i, j] = -1 / (2 * dx) * (F[i, j] + F[i, j - 1] - F[i - 1, j] - F[i - 1, j - 1]) 314 | my3[i, j] = -1 / (2 * dy) * (F[i, j] - F[i, j - 1] + F[i - 1, j] - F[i - 1, j - 1]) 315 | mx4[i, j] = -1 / (2 * dx) * (F[i, j + 1] + F[i, j] - F[i - 1, j + 1] - F[i - 1, j]) 316 | my4[i, j] = -1 / (2 * dy) * (F[i, j + 1] - F[i, j] + F[i - 1, j + 1] - F[i - 1, j]) 317 | # Summing of mx and my components for normal vector 318 | mxsum[i, j] = (mx1[i, j] + mx2[i, j] + mx3[i, j] + mx4[i, j]) / 4 319 | mysum[i, j] = (my1[i, j] + my2[i, j] + my3[i, j] + my4[i, j]) / 4 320 | 321 | # Normalizing the normal vector into unit vectors 322 | if abs(mxsum[i, j]) < 1e-10 and abs(mysum[i, j])< 1e-10: 323 | mx[i, j] = mxsum[i, j] 324 | my[i, j] = mysum[i, j] 325 | else: 326 | magnitude[i, j] = ti.sqrt(mxsum[i, j] * mxsum[i, j] + mysum[i, j] * mysum[i, j]) 327 | mx[i, j] = mxsum[i, j] / magnitude[i, j] 328 | my[i, j] = mysum[i, j] / magnitude[i, j] 329 | for i, j in ti.ndrange((imin, imax + 1), (jmin, jmax + 1)): 330 | kappa[i, j] = -(1 / dx / 2 * (mx[i + 1, j] - mx[i - 1, j]) + \ 331 | 1 / dy / 2 * (my[i, j + 1] - my[i, j - 1])) 332 | ''' 333 | 334 | 335 | @ti.kernel 336 | def solve_VOF_upwind(): 337 | # Upwind advection of VOF; use Ftd as a buffer. 338 | for I in ti.grouped(Ftd): 339 | Ftd[I] = F[I] 340 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 341 | fl = u[i, j, k] * dt * Ftd[i - 1, j, k] if u[i, j, k] > 0 else u[i, j, k] * dt * Ftd[i, j, k] 342 | fr = u[i + 1, j, k] * dt * Ftd[i, j, k] if u[i + 1, j, k] > 0 else u[i + 1, j, k] * dt * Ftd[i + 1, j, k] 343 | fs = v[i, j, k] * dt * Ftd[i, j - 1, k] if v[i, j, k] > 0 else v[i, j, k] * dt * Ftd[i, j, k] 344 | fn = v[i, j + 1, k] * dt * Ftd[i, j, k] if v[i, j + 1, k] > 0 else v[i, j + 1, k] * dt * Ftd[i, j + 1, k] 345 | fb = w[i, j, k] * dt * Ftd[i, j, k - 1] if w[i, j, k] > 0 else w[i, j, k] * dt * Ftd[i, j, k] 346 | ff = w[i, j, k + 1] * dt * Ftd[i, j, k] if w[i, j, k + 1] > 0 else w[i, j, k + 1] * dt * Ftd[i, j, k + 1] 347 | F[i, j, k] += (fl - fr + fs - fn + fb - ff) * dx * dy / (dx * dy * dz) 348 | 349 | 350 | 351 | def solve_VOF_rudman(): 352 | if istep % 3 == 0: 353 | fct_x_sweep() 354 | fct_y_sweep() 355 | fct_z_sweep() 356 | elif istep % 3 == 1: 357 | fct_y_sweep() 358 | fct_z_sweep() 359 | fct_x_sweep() 360 | else: 361 | fct_z_sweep() 362 | fct_x_sweep() 363 | fct_y_sweep() 364 | 365 | 366 | @ti.kernel 367 | def fct_x_sweep(): 368 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 369 | dv = dx * dy * dz - dt * dy * dz * (u[i + 1, j, k] - u[i, j, k]) 370 | fl_L = u[i, j, k] * dt * F[i - 1, j, k] if u[i, j, k] >= 0 else u[i, j, k] * dt * F[i, j, k] 371 | fr_L = u[i + 1, j, k] * dt * F[i, j, k] if u[i + 1, j, k] >= 0 else u[i + 1, j, k] * dt * F[i + 1, j, k] 372 | ft_L = 0 373 | fb_L = 0 374 | Ftd[i, j, k] = (F[i, j, k] + (fl_L - fr_L) * dy * dz / (dx * dy * dz)) * dx * dy * dz / dv 375 | if Ftd[i, j, k] > 1. or Ftd[i, j, k] < 0: 376 | Ftd[i, j, k] = var(0, 1, Ftd[i, j, k]) 377 | 378 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 379 | fmax = ti.max(Ftd[i, j, k], Ftd[i - 1, j, k], Ftd[i + 1, j, k]) 380 | fmin = ti.min(Ftd[i, j, k], Ftd[i - 1, j, k], Ftd[i + 1, j, k]) 381 | 382 | fl_L = u[i, j, k] * dt * F[i - 1, j, k] if u[i, j, k] >= 0 else u[i, j, k] * dt * F[i, j, k] 383 | fr_L = u[i + 1, j, k] * dt * F[i, j, k] if u[i + 1, j, k] >= 0 else u[i + 1, j, k] * dt * F[i + 1, j, k] 384 | ft_L = 0 385 | fb_L = 0 386 | 387 | fl_H = u[i, j, k] * dt * F[i - 1, j, k] if u[i, j, k] <= 0 else u[i, j, k] * dt * F[i, j, k] 388 | fr_H = u[i + 1, j, k] * dt * F[i, j, k] if u[i + 1, j, k] <= 0 else u[i + 1, j, k] * dt * F[i + 1, j, k] 389 | ft_H = 0 390 | fb_H = 0 391 | 392 | ax[i + 1, j, k] = fr_H - fr_L 393 | ax[i, j, k] = fl_H - fl_L 394 | ay[i, j + 1, k] = 0 395 | ay[i, j, k] = 0 396 | 397 | pp = ti.max(0, ax[i, j, k]) - ti.min(0, ax[i + 1, j, k]) + ti.max(0, ay[i, j, k]) - ti.min(0, ay[i, j + 1, k]) 398 | qp = (fmax - Ftd[i, j, k]) * dx 399 | if pp > 0: 400 | rp[i, j, k] = ti.min(1, qp / pp) 401 | else: 402 | rp[i, j, k] = 0.0 403 | pm = ti.max(0, ax[i + 1, j, k]) - ti.min(0, ax[i, j, k]) + ti.max(0, ay[i, j + 1, k]) - ti.min(0, ay[i, j, k]) 404 | qm = (Ftd[i, j, k] - fmin) * dx 405 | if pm > 0: 406 | rm[i, j, k] = ti.min(1, qm / pm) 407 | else: 408 | rm[i, j, k] = 0.0 409 | 410 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 411 | if ax[i + 1, j, k] >= 0: 412 | cx[i + 1, j, k] = ti.min(rp[i + 1, j, k], rm[i, j, k]) 413 | else: 414 | cx[i + 1, j, k] = ti.min(rp[i, j, k], rm[i + 1, j, k]) 415 | 416 | if ay[i, j + 1, k] >= 0: 417 | cy[i, j + 1, k] = ti.min(rp[i, j + 1, k], rm[i, j, k]) 418 | else: 419 | cy[i, j + 1, k] = ti.min(rp[i, j, k], rm[i, j + 1, k]) 420 | 421 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 422 | dv = dx * dy * dz - dt * dy * dz * (u[i + 1, j, k] - u[i, j, k]) 423 | F[i, j, k] = Ftd[i, j, k] - ((ax[i + 1, j, k] * cx[i + 1, j, k] - \ 424 | ax[i, j, k] * cx[i, j, k] + \ 425 | ay[i, j + 1, k] * cy[i, j + 1, k] -\ 426 | ay[i, j, k] * cy[i, j, k]) / (dy)) * dx * dy * dz / dv 427 | F[i, j, k] = var(0, 1, F[i, j, k]) 428 | 429 | 430 | @ti.kernel 431 | def fct_y_sweep(): 432 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 433 | dv = dx * dy * dz - dt * dx * dz * (v[i, j + 1, k] - v[i, j, k]) 434 | fl_L = 0 435 | fr_L = 0 436 | ft_L = v[i, j + 1, k] * dt * F[i, j, k] if v[i, j + 1, k] >= 0 else v[i, j + 1, k] * dt * F[i, j + 1, k] 437 | fb_L = v[i, j, k] * dt * F[i, j - 1, k] if v[i, j, k] >= 0 else v[i, j, k] * dt * F[i, j, k] 438 | Ftd[i, j, k] = (F[i, j, k] + (fl_L - fr_L + fb_L - ft_L) * dy / (dx * dy)) * dx * dy * dz / dv 439 | if Ftd[i, j, k] > 1. or Ftd[i, j, k] < 0: 440 | Ftd[i, j, k] = var(0, 1, Ftd[i, j, k]) 441 | 442 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 443 | fmax = ti.max(Ftd[i, j, k], Ftd[i, j - 1, k], Ftd[i, j + 1, k]) 444 | fmin = ti.min(Ftd[i, j, k], Ftd[i, j - 1, k], Ftd[i, j + 1, k]) 445 | 446 | fl_L = 0 447 | fr_L = 0 448 | ft_L = v[i, j + 1, k] * dt * F[i, j, k] if v[i, j + 1, k] >= 0 else v[i, j + 1, k] * dt * F[i, j + 1, k] 449 | fb_L = v[i, j, k] * dt * F[i, j - 1, k] if v[i, j, k] >= 0 else v[i, j, k] * dt * F[i, j, k] 450 | 451 | fl_H = 0 452 | fr_H = 0 453 | ft_H = v[i, j + 1, k] * dt * F[i, j, k] if v[i, j + 1, k] <= 0 else v[i, j + 1, k] * dt * F[i, j + 1, k] 454 | fb_H = v[i, j, k] * dt * F[i, j - 1, k] if v[i, j, k] <= 0 else v[i, j, k] * dt * F[i, j, k] 455 | 456 | ax[i + 1, j, k] = 0 457 | ax[i, j, k] = 0 458 | ay[i, j + 1, k] = ft_H - ft_L 459 | ay[i, j, k] = fb_H - fb_L 460 | 461 | pp = ti.max(0, ax[i, j, k]) - ti.min(0, ax[i + 1, j, k]) + ti.max(0, ay[i, j, k]) - ti.min(0, ay[i, j + 1, k]) 462 | qp = (fmax - Ftd[i, j, k]) * dx 463 | if pp > 0: 464 | rp[i, j, k] = ti.min(1, qp / pp) 465 | else: 466 | rp[i, j, k] = 0.0 467 | pm = ti.max(0, ax[i + 1, j, k]) - ti.min(0, ax[i, j, k]) + ti.max(0, ay[i, j + 1, k]) - ti.min(0, ay[i, j, k]) 468 | qm = (Ftd[i, j, k] - fmin) * dx 469 | if pm > 0: 470 | rm[i, j, k] = ti.min(1, qm / pm) 471 | else: 472 | rm[i, j, k] = 0.0 473 | 474 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 475 | if ax[i + 1, j, k] >= 0: 476 | cx[i + 1, j, k] = ti.min(rp[i + 1, j, k], rm[i, j, k]) 477 | else: 478 | cx[i + 1, j, k] = ti.min(rp[i, j, k], rm[i + 1, j, k]) 479 | 480 | if ay[i, j + 1, k] >= 0: 481 | cy[i, j + 1, k] = ti.min(rp[i, j + 1, k], rm[i, j, k]) 482 | else: 483 | cy[i, j + 1, k] = ti.min(rp[i, j, k], rm[i, j + 1, k]) 484 | 485 | 486 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 487 | dv = dx * dy * dz - dt * dx * dz * (v[i, j + 1, k] - v[i, j, k]) 488 | F[i, j, k] = Ftd[i, j, k] - ((ax[i + 1, j, k] * cx[i + 1, j, k] - \ 489 | ax[i, j, k] * cx[i, j, k] + \ 490 | ay[i, j + 1, k] * cy[i, j + 1, k] -\ 491 | ay[i, j, k] * cy[i, j, k]) / (dy)) * dx * dy * dz / dv 492 | F[i, j, k] = var(0, 1, F[i, j, k]) 493 | 494 | 495 | @ti.kernel 496 | def fct_z_sweep(): 497 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 498 | dv = dx * dy * dz - dt * dx * dy * (w[i, j, k + 1] - w[i, j, k]) 499 | ff_L = w[i, j, k + 1] * dt * F[i, j, k] if w[i, j, k + 1] >= 0 else w[i, j, k + 1] * dt * F[i, j, k + 1] 500 | fb_L = w[i, j, k] * dt * F[i, j, k - 1] if w[i, j, k] >= 0 else w[i, j, k] * dt * F[i, j, k] 501 | Ftd[i, j, k] = (F[i, j, k] + (fb_L - ff_L) * dy * dx / (dx * dy * dz)) * dx * dy * dz / dv 502 | if Ftd[i, j, k] > 1. or Ftd[i, j, k] < 0: 503 | Ftd[i, j, k] = var(0, 1, Ftd[i, j, k]) 504 | 505 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 506 | fmax = ti.max(Ftd[i, j, k], Ftd[i, j, k - 1], Ftd[i, j, k + 1]) 507 | fmin = ti.min(Ftd[i, j, k], Ftd[i, j, k - 1], Ftd[i, j, k + 1]) 508 | 509 | ff_L = w[i, j, k + 1] * dt * F[i, j, k] if w[i, j, k + 1] >= 0 else w[i, j, k + 1] * dt * F[i, j, k + 1] 510 | fb_L = w[i, j, k] * dt * F[i, j, k - 1] if w[i, j, k] >= 0 else w[i, j, k] * dt * F[i, j, k] 511 | 512 | ff_H = w[i, j, k + 1] * dt * F[i, j, k] if w[i, j, k + 1] <= 0 else w[i, j, k + 1] * dt * F[i, j, k + 1] 513 | fb_H = w[i, j, k] * dt * F[i, j, k - 1] if w[i, j, k] <= 0 else w[i, j, k] * dt * F[i, j, k] 514 | 515 | az[i, j, k + 1] = ff_H - ff_L 516 | az[i, j, k] = fb_H - fb_L 517 | 518 | pp = ti.max(0, az[i, j, k]) - ti.min(0, az[i, j, k + 1]) 519 | qp = (fmax - Ftd[i, j, k]) * dz 520 | if pp > 0: 521 | rp[i, j, k] = ti.min(1, qp / pp) 522 | else: 523 | rp[i, j, k] = 0.0 524 | pm = ti.max(0, az[i, j, k + 1]) - ti.min(0, az[i, j, k]) 525 | qm = (Ftd[i, j, k] - fmin) * dz 526 | if pm > 0: 527 | rm[i, j, k] = ti.min(1, qm / pm) 528 | else: 529 | rm[i, j, k] = 0.0 530 | 531 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 532 | if az[i, j, k + 1] >= 0: 533 | cz[i, j, k + 1] = ti.min(rp[i, j, k + 1], rm[i, j, k]) 534 | else: 535 | cz[i, j, k + 1] = ti.min(rp[i, j, k], rm[i, j, k + 1]) 536 | 537 | for i, j, k in ti.ndrange((imin, imax + 1), (jmin, jmax + 1), (kmin, kmax + 1)): 538 | dv = dx * dy * dz - dt * dx * dy * (w[i, j, k + 1] - w[i, j, k]) 539 | F[i, j, k] = Ftd[i, j, k] - ((az[i, j, k + 1] * cz[i, j, k + 1] - \ 540 | az[i, j, k] * cz[i, j, k]) / (dz)) * dx * dy * dz / dv 541 | F[i, j, k] = var(0, 1, F[i, j, k]) 542 | 543 | 544 | @ti.kernel 545 | def post_process_f(): 546 | for i, j, k in F: 547 | F[i, j, k] = var(F[i, j, k], 0, 1) 548 | 549 | 550 | @ti.kernel 551 | def get_vof_field(): 552 | r = resolution[0] // nx 553 | for I in ti.grouped(rgb_buf): 554 | rgb_buf[I] = F[I // r] 555 | 556 | ''' 557 | @ti.kernel 558 | def get_u_field(): 559 | r = resolution[0] // nx 560 | max = Lx / 0.2 561 | for I in ti.grouped(rgb_buf): 562 | rgb_buf[I] = (u[I // r]) / max 563 | 564 | 565 | @ti.kernel 566 | def get_v_field(): 567 | r = resolution[0] // nx 568 | max = Ly / 0.2 569 | for I in ti.grouped(rgb_buf): 570 | rgb_buf[I] = (v[I // r]) / max 571 | 572 | 573 | @ti.kernel 574 | def get_vnorm_field(): 575 | r = resolution[0] // nx 576 | max = Ly / 0.2 577 | for I in ti.grouped(rgb_buf): 578 | rgb_buf[I] = ti.sqrt(u[I // r] ** 2 + v[I // r] ** 2) / max 579 | 580 | 581 | @ti.kernel 582 | def interp_velocity(): 583 | for i, j in ti.ndrange((1, imax + 2), (1, jmax + 1)): 584 | V[i, j] = ti.Vector([(u[i, j] + u[i + 1, j])/2, (v[i, j] + v[i, j + 1])/2]) 585 | ''' 586 | 587 | # Start main script 588 | istep = 0 589 | nstep = 100 # Interval to update GUI 590 | set_init_F(initial_condition) 591 | 592 | os.makedirs('output', exist_ok=True) # Make dir for output 593 | os.makedirs('data', exist_ok=True) # Make dir for data save; only used for debugging 594 | gui = ti.GUI('VOF Solver', resolution, background_color=0xFFFFFF) 595 | vis_option = 0 # Tag for display 596 | 597 | 598 | while gui.running: 599 | istep += 1 600 | for e in gui.get_events(gui.RELEASE): 601 | if e.key == gui.SPACE: 602 | vis_option += 1 603 | elif e.key == 'q': 604 | gui.running = False 605 | 606 | cal_nu_rho() 607 | # get_normal_young() 608 | 609 | # Advection 610 | advect_upwind() 611 | set_BC() 612 | 613 | # Pressure projection 614 | for _ in range(10): 615 | solve_p_jacobi() 616 | 617 | update_uv() 618 | set_BC() 619 | 620 | # solve_VOF_upwind() 621 | solve_VOF_rudman() 622 | post_process_f() 623 | set_BC() 624 | if (istep % nstep) == 0: 625 | print(f'>>> Exporting step-{istep:05d} result...') 626 | gridToVTK(f'./output/step-{istep:05d}', xcor, ycor, zcor, \ 627 | pointData={"VOF": np.ascontiguousarray(F.to_numpy())}) 628 | ''' 629 | num_options = 5 630 | if (istep % nstep) == 0: # Output data every steps 631 | if vis_option % num_options == 0: # Display VOF distribution 632 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying VOF field.') 633 | get_vof_field() 634 | rgbnp = rgb_buf.to_numpy() 635 | gui.set_image(cm.Blues(rgbnp)) 636 | 637 | if vis_option % num_options == 1: # Display the u field 638 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying u velocity.') 639 | get_u_field() 640 | rgbnp = rgb_buf.to_numpy() 641 | gui.set_image(cm.coolwarm(rgbnp)) 642 | 643 | if vis_option % num_options == 2: # Display the v field 644 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying v velocity.') 645 | get_v_field() 646 | rgbnp = rgb_buf.to_numpy() 647 | gui.set_image(cm.coolwarm(rgbnp)) 648 | 649 | if vis_option % num_options == 3: # Display velocity norm 650 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity norm.') 651 | get_vnorm_field() 652 | rgbnp = rgb_buf.to_numpy() 653 | gui.set_image(cm.plasma(rgbnp)) 654 | 655 | if vis_option % num_options == 4: # Display velocity vectors 656 | print(f'>>> Number of steps:{istep:<5d}, Time:{istep*dt:5.2e} sec. Displaying velocity vectors.') 657 | interp_velocity() 658 | fv.plot_arrow_field(vector_field=V, arrow_spacing=4, gui=gui) 659 | 660 | gui.show() 661 | 662 | if SAVE_FIG: 663 | count = istep // nstep - 1 664 | Fnp = F.to_numpy() 665 | fx, fy = 5, Ly / Lx * 5 666 | plt.figure(figsize=(fx, fy)) 667 | plt.axis('off') 668 | plt.contourf(Fnp.T, cmap=plt.cm.Blues) 669 | plt.savefig(f'output/{count:06d}-f.png') 670 | plt.close() 671 | ''' 672 | --------------------------------------------------------------------------------