├── requirements.txt ├── data └── direct_vs_cg.gif ├── LICENSE ├── .gitignore ├── README.md └── implicit_mass_spring.py /requirements.txt: -------------------------------------------------------------------------------- 1 | taichi 2 | -------------------------------------------------------------------------------- /data/direct_vs_cg.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/taichiCourse01/taichi_course_final_project/HEAD/data/direct_vs_cg.gif -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 TaichiCourse 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 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Pycharm 2 | .idea 3 | 4 | # Byte-compiled / optimized / DLL files 5 | __pycache__/ 6 | *.py[cod] 7 | *$py.class 8 | 9 | # C extensions 10 | *.so 11 | 12 | # Distribution / packaging 13 | .Python 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | wheels/ 26 | share/python-wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | MANIFEST 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .nox/ 46 | .coverage 47 | .coverage.* 48 | .cache 49 | nosetests.xml 50 | coverage.xml 51 | *.cover 52 | *.py,cover 53 | .hypothesis/ 54 | .pytest_cache/ 55 | cover/ 56 | 57 | # Translations 58 | *.mo 59 | *.pot 60 | 61 | # Django stuff: 62 | *.log 63 | local_settings.py 64 | db.sqlite3 65 | db.sqlite3-journal 66 | 67 | # Flask stuff: 68 | instance/ 69 | .webassets-cache 70 | 71 | # Scrapy stuff: 72 | .scrapy 73 | 74 | # Sphinx documentation 75 | docs/_build/ 76 | 77 | # PyBuilder 78 | .pybuilder/ 79 | target/ 80 | 81 | # Jupyter Notebook 82 | .ipynb_checkpoints 83 | 84 | # IPython 85 | profile_default/ 86 | ipython_config.py 87 | 88 | # pyenv 89 | # For a library or package, you might want to ignore these files since the code is 90 | # intended to run in multiple environments; otherwise, check them in: 91 | # .python-version 92 | 93 | # pipenv 94 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 95 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 96 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 97 | # install all needed dependencies. 98 | #Pipfile.lock 99 | 100 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 101 | __pypackages__/ 102 | 103 | # Celery stuff 104 | celerybeat-schedule 105 | celerybeat.pid 106 | 107 | # SageMath parsed files 108 | *.sage.py 109 | 110 | # Environments 111 | .env 112 | .venv 113 | env/ 114 | venv/ 115 | ENV/ 116 | env.bak/ 117 | venv.bak/ 118 | 119 | # Spyder project settings 120 | .spyderproject 121 | .spyproject 122 | 123 | # Rope project settings 124 | .ropeproject 125 | 126 | # mkdocs documentation 127 | /site 128 | 129 | # mypy 130 | .mypy_cache/ 131 | .dmypy.json 132 | dmypy.json 133 | 134 | # Pyre type checker 135 | .pyre/ 136 | 137 | # pytype static type analyzer 138 | .pytype/ 139 | 140 | # Cython debug symbols 141 | cython_debug/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 太极图形课S1-大作业 2 | 3 | ## 作业来源 4 | > 请你介绍大作业的灵感来源，可以是自己原创的想法，也可以是复现一篇论文。如果有参考论文、文章或者参考代码，请**务必**给出来源。 5 | 6 | ## 运行方式 7 | > 展示项目运行方式，让更多的人在自己的机器上运行你的代码查看运行效果。 8 | > 为了复现运行效果，请给出你的运行环境。 9 | > 如果项目有依赖，请给出项目的依赖。 10 | 11 | #### 运行环境： 12 | > 请列出`Taichi`程序运行环境，可以在命令行中输入命令：`$ ti`，然后复制输出的第一行。 13 | 14 | #### 运行： 15 | > 请列出哪些代码可以运行以及如何运行，如果有命令行参数可以列出参数以及参数对应的含义。 16 | 17 | ## 效果展示 18 | > 这里可以展示这份作业运行起来后的可视化效果，可以让其他人更直观感受到你的工作。可以是**图片**（比如渲染效果）、**动图**或**视频**。 19 | 20 | ## 整体结构 21 | > 脉络清晰的结构能完整展示你的设计思想，以及实现方式，方便读者快速代入。Python的代码，可以使用命令：`yapf -i xx.py`来格式化。可以在repo的目录中包含如下内容： 22 | ``` 23 | -LICENSE 24 | -|data 25 | -README.MD 26 | -xx.py 27 | ``` 28 | 29 | ## 实现细节： 30 | > 请给出代码的具体实现流程和细节，让感兴趣的人能够更深入的了解你的代码。 31 | 32 | # 示例 33 | 下面给出一个模版供大家参考，在你提交大作业的时候可以删除示例部分。按上面的要求填写自己大作业的内容。 34 | 35 | ## 作业来源 36 | 布料仿真是计算机图形学中一个重要的仿真场景，在电影和游戏中都有很广泛的应用。一个很经典的布料仿真方法是 [*Large Steps in Cloth Simulation.*](https://www.cs.cmu.edu/~baraff/papers/sig98.pdf) David Baraff, Andrew Witkin. Siggraph 1998. 这也是皮克斯 (Pixar) 动画仿真器 **Fizt** 使用的布料仿真方法。在SCA 2020中有有一篇[论文](https://www.tkim.graphics/FEMBW/)描述了该论文中布料仿真模型对应的FEM形式。 37 | 38 | 本次大作业就是对Baraff-Witkin布料仿真的实现。在实现过程中，可以了解能量、力和刚度矩阵的计算。除此之外，还可以了解Taichi语言中稀疏矩阵 (Sparse Matrix) 的使用，共轭梯度法 (Conjugate Gradient, CG) 的实现。 39 | 40 | ## 运行方式 41 | #### 运行环境： 42 | `[Taichi] version 0.8.7, llvm 10.0.0, commit 1c3c705a, osx, python 3.8.8` 43 | 44 | #### 运行： 45 | 在运行 `implicit_mass_spring_system.py`时，可以通过命令行参数 `-cg` 来控制是否使用 **CG** solver。 46 | 47 | - 使用 Direct solver: 48 | `python implicit_mass_spring_system.py` 49 | 50 | - 使用 CG solver: 51 | `python implicit_mass_spring_system.py -cg` 52 | 53 | 54 | ## 效果展示 55 | 左侧：Direct solver；右侧： Conjugate Gradient (CG) solver。 56 | ![mass spring demo](./data/direct_vs_cg.gif) 57 | 58 | ## 整体结构 59 | ``` 60 | -LICENSE 61 | -|data 62 | -README.MD 63 | -implicit_mass_spring_system.py 64 | ``` 65 | 66 | ## 实现细节： 67 | `implicit_mass_spring_system.py`是项目源代码，其中包含一个Cloth类和一个`main`函数。具体公式推导，可以看: [here](https://github.com/taichiCourse01/taichiCourse01/blob/main/material/09_implicit_integration.pdf) (P45-46)。 68 | 69 | ### 整体流程 70 | 1. 布料的初始化 71 | 2. 创建一个GUI来显示布料 72 | 3. 根据命令行选项更新布料，默认使用 direct solver 来更新布料。可以命令行参数 `-cg` 来使用 CG solver。 73 | 4. 在GUI中显示布料 74 | 75 | ### 布料类 76 | 1. [初始化](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L13) 77 | - 设置field，确定数据结构 78 | - 初始化field，设置位置，速度，质量和每个弹簧两个端点索引 79 | 80 | 2. 布料更新 Direct Solver: [`update_direct(h)`](https://github.com/FantasyVR/taichi_course_final_project/blob/8f79e0026237e75ec3abe7d09b39be0a2fadc994/implicit_mass_spring.py#L156) 81 | 82 | **计算力**: [`compute_force()`](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L98) 83 | - 重力 84 | - 弹簧力 85 | - 约束力: 固定布料的两个顶点，使用一个刚度极大的弹簧来模拟。 86 | 87 | **计算力的导数**:[`compute_force_Jacobians`](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L120) 88 | - 弹簧力的导数 89 | - 约束力的导数 90 | 91 | **组装刚度矩阵**: [`assemble_K`](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L138) 92 | - 遍历每个弹簧，组建度矩阵 93 | 94 | **组装系统矩阵**：[here](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L162) 95 | - 系统矩阵 96 | 97 | **计算系统矩阵右侧b**: [here](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L169) 98 | - 99 | 100 | **求解线性系统**: [here](https://github.com/FantasyVR/taichi_course_final_project/blob/304a32dfa686862adcb54f737ed6970e21fe8d5b/implicit_mass_spring.py#L171) 101 | - 102 | 103 | **更新速度和位置**: [`directUpdatePosVel`](https://github.com/FantasyVR/taichi_course_final_project/blob/8f79e0026237e75ec3abe7d09b39be0a2fadc994/implicit_mass_spring.py#L151) 104 | - 105 | 106 | 3. 布料更新 CG solver: [`update_cg(h)`](https://github.com/FantasyVR/taichi_course_final_project/blob/8f79e0026237e75ec3abe7d09b39be0a2fadc994/implicit_mass_spring.py#L249) 107 | 108 | 在使用CG进行布料更新的时候，同样需要计算力和计算力的导数。但我们使用CG的时候，不需要显式的构建出刚度矩阵，只需要计算矩阵-向量乘就可以。具体算法请参考：[here](https://github.com/taichiCourse01/taichiCourse01/blob/main/material/09_implicit_integration.pdf) (P105-110)。 109 | 110 | -------------------------------------------------------------------------------- /implicit_mass_spring.py: -------------------------------------------------------------------------------- 1 | # https://www.cs.cmu.edu/~baraff/papers/sig98.pdf 2 | import argparse 3 | 4 | import numpy as np 5 | 6 | import taichi as ti 7 | 8 | import time 9 | 10 | 11 | @ti.data_oriented 12 | class Cloth: 13 | def __init__(self, N): 14 | self.N = N 15 | self.NF = 2 * N**2 # number of faces 16 | self.NV = (N + 1)**2 # number of vertices 17 | self.NE = 2 * N * (N + 1) + 2 * N * N # numbser of edges 18 | self.initPos = ti.Vector.field(2, ti.f32, self.NV) 19 | self.pos = ti.Vector.field(2, ti.f32, self.NV) 20 | self.vel = ti.Vector.field(2, ti.f32, self.NV) 21 | self.force = ti.Vector.field(2, ti.f32, self.NV) 22 | self.mass = ti.field(ti.f32, self.NV) 23 | self.spring = ti.Vector.field(2, ti.i32, self.NE) 24 | self.rest_len = ti.field(ti.f32, self.NE) 25 | self.ks = 1000.0 # spring stiffness 26 | self.kf = 1.0e5 # Attachment point stiffness 27 | self.Jx = ti.Matrix.field(2, 2, ti.f32, self.NE) # Force Jacobian 28 | self.Jf = ti.Matrix.field(2, 2, ti.f32, 2) # Attachment Jacobian 29 | 30 | self.init_pos() 31 | self.init_edges() 32 | 33 | # For sparse matrix solver， PPT: P45 34 | max_num_triplets = 10000 35 | self.MBuilder = ti.linalg.SparseMatrixBuilder(2 * self.NV, 2 * self.NV, 36 | max_num_triplets) 37 | self.init_mass_sp(self.MBuilder) 38 | self.M = self.MBuilder.build() 39 | self.KBuilder = ti.linalg.SparseMatrixBuilder(2 * self.NV, 2 * self.NV, 40 | max_num_triplets) 41 | 42 | # For conjugate gradient method, PPT: P106 43 | self.x = ti.Vector.field(2, ti.f32, self.NV) 44 | self.Ax = ti.Vector.field(2, ti.f32, self.NV) 45 | self.b = ti.Vector.field(2, ti.f32, self.NV) 46 | self.r = ti.Vector.field(2, ti.f32, self.NV) 47 | self.d = ti.Vector.field(2, ti.f32, self.NV) 48 | self.Ad = ti.Vector.field(2, ti.f32, self.NV) 49 | 50 | @ti.kernel 51 | def init_pos(self): 52 | for i, j in ti.ndrange(self.N + 1, self.N + 1): 53 | k = i * (self.N + 1) + j 54 | self.initPos[k] = ti.Vector([i, j]) / self.N * 0.5 + ti.Vector( 55 | [0.25, 0.25]) 56 | self.pos[k] = self.initPos[k] 57 | self.vel[k] = ti.Vector([0, 0]) 58 | self.mass[k] = 1.0 59 | 60 | @ti.kernel 61 | def init_edges(self): 62 | pos, spring, N, rest_len = ti.static(self.pos, self.spring, self.N, 63 | self.rest_len) 64 | for i, j in ti.ndrange(N + 1, N): 65 | idx, idx1 = i * N + j, i * (N + 1) + j 66 | spring[idx] = ti.Vector([idx1, idx1 + 1]) 67 | start = N * (N + 1) 68 | for i, j in ti.ndrange(N, N + 1): 69 | idx, idx1, idx2 = start + i + j * N, i * (N + 1) + j, i * ( 70 | N + 1) + j + N + 1 71 | spring[idx] = ti.Vector([idx1, idx2]) 72 | start = 2 * N * (N + 1) 73 | for i, j in ti.ndrange(N, N): 74 | idx, idx1, idx2 = start + i * N + j, i * (N + 1) + j, (i + 1) * ( 75 | N + 1) + j + 1 76 | spring[idx] = ti.Vector([idx1, idx2]) 77 | start = 2 * N * (N + 1) + N * N 78 | for i, j in ti.ndrange(N, N): 79 | idx, idx1, idx2 = start + i * N + j, i * (N + 1) + j + 1, ( 80 | i + 1) * (N + 1) + j 81 | spring[idx] = ti.Vector([idx1, idx2]) 82 | for i in range(self.NE): 83 | idx1, idx2 = spring[i] 84 | rest_len[i] = (pos[idx1] - pos[idx2]).norm() 85 | 86 | @ti.kernel 87 | def init_mass_sp(self, M: ti.linalg.sparse_matrix_builder()): 88 | for i in range(self.NV): 89 | M[2 * i + 0, 2 * i + 0] += self.mass[i] 90 | M[2 * i + 1, 2 * i + 1] += self.mass[i] 91 | 92 | @ti.func 93 | def clear_force(self): 94 | for i in self.force: 95 | self.force[i] = ti.Vector([0.0, 0.0]) 96 | 97 | @ti.kernel 98 | def compute_force(self): 99 | self.clear_force() 100 | gravity = ti.Vector([0.0, -2.0]) 101 | for i in self.force: 102 | self.force[i] += gravity * self.mass[i] 103 | 104 | for i in self.spring: 105 | idx1, idx2 = self.spring[i][0], self.spring[i][1] 106 | pos1, pos2 = self.pos[idx1], self.pos[idx2] 107 | dis = pos1 - pos2 108 | # Hook's law 109 | force = self.ks * (dis.norm() - 110 | self.rest_len[i]) * dis.normalized() 111 | self.force[idx1] -= force 112 | self.force[idx2] += force 113 | # Attachment constraint force 114 | self.force[self.N] += self.kf * (self.initPos[self.N] - 115 | self.pos[self.N]) 116 | self.force[self.NV - 1] += self.kf * (self.initPos[self.NV - 1] - 117 | self.pos[self.NV - 1]) 118 | 119 | @ti.kernel 120 | def compute_force_Jacobians(self): 121 | for i in self.spring: 122 | idx1, idx2 = self.spring[i][0], self.spring[i][1] 123 | pos1, pos2 = self.pos[idx1], self.pos[idx2] 124 | dx = pos1 - pos2 125 | I = ti.Matrix([[1.0, 0.0], [0.0, 1.0]]) 126 | dxtdx = ti.Matrix([[dx[0] * dx[0], dx[0] * dx[1]], 127 | [dx[1] * dx[0], dx[1] * dx[1]]]) 128 | l = dx.norm() 129 | if l != 0.0: 130 | l = 1.0 / l 131 | self.Jx[i] = (I - self.rest_len[i] * l * 132 | (I - dxtdx * l**2)) * self.ks 133 | # Attachment constraint force Jacobian 134 | self.Jf[0] = ti.Matrix([[-self.kf, 0], [0, -self.kf]]) 135 | self.Jf[1] = ti.Matrix([[-self.kf, 0], [0, -self.kf]]) 136 | 137 | @ti.kernel 138 | def assemble_K(self, K: ti.linalg.sparse_matrix_builder()): 139 | for i in self.spring: 140 | idx1, idx2 = self.spring[i][0], self.spring[i][1] 141 | for m, n in ti.static(ti.ndrange(2, 2)): 142 | K[2 * idx1 + m, 2 * idx1 + n] -= self.Jx[i][m, n] 143 | K[2 * idx1 + m, 2 * idx2 + n] += self.Jx[i][m, n] 144 | K[2 * idx2 + m, 2 * idx1 + n] += self.Jx[i][m, n] 145 | K[2 * idx2 + m, 2 * idx2 + n] -= self.Jx[i][m, n] 146 | for m, n in ti.static(ti.ndrange(2, 2)): 147 | K[2 * self.N + m, 2 * self.N + n] += self.Jf[0][m, n] 148 | K[2 * (self.NV - 1) + m, 2 * (self.NV - 1) + n] += self.Jf[1][m, n] 149 | 150 | @ti.kernel 151 | def directUpdatePosVel(self, h: ti.f32, v_next: ti.ext_arr()): 152 | for i in self.pos: 153 | self.vel[i] = ti.Vector([v_next[2 * i], v_next[2 * i + 1]]) 154 | self.pos[i] += h * self.vel[i] 155 | 156 | def update_direct(self, h): 157 | self.compute_force() 158 | self.compute_force_Jacobians() 159 | # Assemble global system 160 | self.assemble_K(self.KBuilder) 161 | K = self.KBuilder.build() 162 | A = self.M - h**2 * K 163 | solver = ti.linalg.SparseSolver(solver_type="LLT") 164 | solver.analyze_pattern(A) 165 | solver.factorize(A) 166 | 167 | vel = self.vel.to_numpy().reshape(2 * self.NV) 168 | force = self.force.to_numpy().reshape(2 * self.NV) 169 | b = h * force + self.M @ vel 170 | 171 | v_next = solver.solve(b) 172 | # flag = solver.info() 173 | # print("solver flag: ", flag) 174 | self.directUpdatePosVel(h, v_next) 175 | 176 | @ti.kernel 177 | def cgUpdatePosVel(self, h: ti.f32): 178 | for i in self.pos: 179 | self.vel[i] = self.x[i] 180 | self.pos[i] += h * self.vel[i] 181 | 182 | @ti.kernel 183 | def compute_RHS(self, h: ti.f32): 184 | #rhs = b = h * force + M @ v 185 | for i in range(self.NV): 186 | self.b[i] = h * self.force[i] + self.mass[i] * self.vel[i] 187 | 188 | @ti.func 189 | def dot(self, v1, v2): 190 | result = 0.0 191 | for i in range(self.NV): 192 | result += v1[i][0] * v2[i][0] 193 | result += v1[i][1] * v2[i][1] 194 | return result 195 | 196 | @ti.func 197 | def A_mult_x(self, h, dst, src): 198 | coeff = -h**2 199 | for i in range(self.NV): 200 | dst[i] = self.mass[i] * src[i] 201 | for i in range(self.NE): 202 | idx1, idx2 = self.spring[i][0], self.spring[i][1] 203 | temp = self.Jx[i] @ (src[idx1] - src[idx2]) 204 | dst[idx1] -= coeff * temp 205 | dst[idx2] += coeff * temp 206 | # Attachment constraint 207 | Attachment1, Attachment2 = self.N, self.NV - 1 208 | dst[Attachment1] -= coeff * self.kf * src[Attachment1] 209 | dst[Attachment2] -= coeff * self.kf * src[Attachment2] 210 | 211 | # conjugate gradient solving 212 | # https://www.cs.cmu.edu/~quake-papers/painless-conjugate-gradient.pdf 213 | 214 | @ti.kernel 215 | def before_ite(self) -> ti.f32: 216 | for i in range(self.NV): 217 | self.x[i] = ti.Vector([0.0, 0.0]) 218 | self.A_mult_x(h, self.Ax, self.x) # Ax = A @ x 219 | for i in range(self.NV): # r = b - A @ x 220 | self.r[i] = self.b[i] - self.Ax[i] 221 | for i in range(self.NV): # d = r 222 | self.d[i] = self.r[i] 223 | delta_new = self.dot(self.r, self.r) 224 | return delta_new 225 | 226 | @ti.kernel 227 | def run_iteration(self, delta_new: ti.f32) -> ti.f32: 228 | self.A_mult_x(h, self.Ad, self.d) # Ad = A @ d 229 | alpha = delta_new / self.dot(self.d, 230 | self.Ad) # alpha = (r^T * r) / dot(d, Ad) 231 | for i in range(self.NV): 232 | self.x[i] += alpha * self.d[i] # x^{i+1} = x^{i} + alpha * d 233 | self.r[i] -= alpha * self.Ad[i] # r^{i+1} = r^{i} + alpha * Ad 234 | delta_old = delta_new 235 | delta_new = self.dot(self.r, self.r) 236 | beta = delta_new / delta_old 237 | for i in range(self.NV): 238 | self.d[i] = self.r[i] + beta * self.d[ 239 | i] #p^{i+1} = r^{i+1} + beta * p^{i} 240 | return delta_new 241 | 242 | def cg(self, h: ti.f32): 243 | delta_new = self.before_ite() 244 | ite, iteMax = 0, 2 * self.NV 245 | while ite < iteMax and delta_new > 1.0e-6: 246 | delta_new = self.run_iteration(delta_new) 247 | ite += 1 248 | 249 | def update_cg(self, h): 250 | self.compute_force() 251 | self.compute_force_Jacobians() 252 | self.compute_RHS(h) 253 | self.cg(h) 254 | self.cgUpdatePosVel(h) 255 | 256 | def display(self, gui, radius=5, color=0xffffff): 257 | springs, pos = self.spring.to_numpy(), self.pos.to_numpy() 258 | line_Begin = np.zeros(shape=(springs.shape[0], 2)) 259 | line_End = np.zeros(shape=(springs.shape[0], 2)) 260 | for i in range(springs.shape[0]): 261 | idx1, idx2 = springs[i][0], springs[i][1] 262 | line_Begin[i], line_End[i] = pos[idx1], pos[idx2] 263 | gui.lines(line_Begin, line_End, radius=2, color=0x0000ff) 264 | gui.circles(self.pos.to_numpy(), radius, color) 265 | 266 | 267 | if __name__ == "__main__": 268 | ti.init(arch=ti.cpu) 269 | cloth = Cloth(N=5) 270 | parser = argparse.ArgumentParser() 271 | parser.add_argument('-cg', 272 | '--use_cg', 273 | action='store_true', 274 | help='Solve Ax=b with conjugate gradient method (CG).') 275 | args, unknowns = parser.parse_known_args() 276 | use_cg = args.use_cg 277 | 278 | gui = ti.GUI('Implicit Mass Spring System', res=(500, 500)) 279 | pause = False 280 | h, max_step = 0.01, 3 281 | while gui.running: 282 | for e in gui.get_events(): 283 | if e.key == gui.ESCAPE: 284 | gui.running = False 285 | elif e.key == gui.SPACE: 286 | pause = not pause 287 | if not pause: 288 | for i in range(max_step): 289 | if use_cg: 290 | cloth.update_cg(h) 291 | else: 292 | cloth.update_direct(h) 293 | cloth.display(gui) 294 | gui.show() --------------------------------------------------------------------------------