├── .coveragerc ├── .github └── workflows │ ├── persubmit.yml │ └── release.yml ├── .gitignore ├── .readthedocs.yml ├── CHANGELOG.md ├── README.md ├── a.py ├── docs ├── Makefile ├── conf.py ├── index.rst └── requirements.txt ├── examples ├── blueorange.py ├── export_video.py ├── fluid_v1.py ├── fluid_v2.py ├── fractal.py ├── mass_spring.py ├── old_uv.py ├── paint_arrow.py ├── particles.py ├── rand.py ├── rand3d.py ├── sdf_rt1.py ├── sdf_rt2.py ├── simple_vortices.py ├── smoke.py ├── step_uv.py ├── taichi_logo.py ├── taichi_logo_aa.py ├── uv.py └── vortex_frog.py ├── jupyter_notebook_config.py ├── main.py ├── setup.py ├── taichi_glsl ├── __init__.py ├── classes.py ├── colormap.py ├── experimental_array.py ├── experimental_cfd.py ├── experimental_transform.py ├── gui.py ├── hack.py ├── lagrangian.py ├── mkimages.py ├── odop.py ├── painting.py ├── randgen.py ├── sampling.py ├── scalar.py ├── sphcore.py ├── vector.py └── version.py └── tests ├── test_array.py ├── test_odop.py ├── test_rand.py ├── test_scalar.py └── test_vector.py /.coveragerc: -------------------------------------------------------------------------------- 1 | [run] 2 | branch = True 3 | source = taichi_glsl 4 | -------------------------------------------------------------------------------- /.github/workflows/persubmit.yml: -------------------------------------------------------------------------------- 1 | name: Persubmit Checks 2 | on: 3 | push: 4 | branches: [master] 5 | pull_request: 6 | types: [opened, synchronize, reopened] 7 | 8 | 9 | jobs: 10 | run_tests: 11 | name: Run Tests 12 | strategy: 13 | matrix: 14 | taichi: 15 | - 0.7.18 16 | python: 17 | - 3.6 18 | - 3.7 19 | - 3.8 20 | runs-on: ubuntu-latest 21 | steps: 22 | - uses: actions/checkout@v2 23 | - uses: actions/setup-python@v2 24 | with: 25 | python-version: ${{ matrix.python }} 26 | 27 | - name: Install Requirements 28 | run: | 29 | pip install taichi==${{ matrix.taichi }} 30 | pip install pytest coverage 31 | 32 | - name: Functionallity Test 33 | run: | 34 | export PYTHONPATH=`pwd` 35 | coverage run -m pytest tests 36 | 37 | - name: Upload to Codecov 38 | if: ${{ matrix.python == '3.8' }} 39 | run: | 40 | coverage report 41 | bash <(curl -s https://codecov.io/bash) 42 | env: 43 | CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} 44 | 45 | check_format: 46 | name: Check Code Format 47 | runs-on: ubuntu-latest 48 | steps: 49 | - uses: actions/checkout@v2 50 | - uses: actions/setup-python@v2 51 | with: 52 | python-version: 3.8 53 | 54 | - name: Install YAPF 55 | run: | 56 | pip install yapf 57 | 58 | - name: Run YAPF 59 | run: | 60 | yapf -ir . 61 | 62 | - name: Check Git Diff 63 | run: | 64 | git diff --exit-code 65 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release Automation 2 | on: 3 | release: 4 | types: [created, published] 5 | 6 | jobs: 7 | test_pypi_publish: 8 | name: Run Tests 9 | runs-on: ubuntu-latest 10 | steps: 11 | - uses: actions/checkout@v2 12 | - uses: actions/setup-python@v2 13 | with: 14 | python-version: 3.8 15 | 16 | - name: Install Requirements 17 | run: | 18 | pip install taichi 19 | 20 | - name: Functionallity Test 21 | run: | 22 | export PYTHONPATH=`pwd` 23 | pytest tests -n4 24 | 25 | - name: Setup packages 26 | run: | 27 | python setup.py bdist_wheel 28 | 29 | - name: Publish package to TestPyPI 30 | uses: pypa/gh-action-pypi-publish@master 31 | with: 32 | user: __token__ 33 | password: ${{ secrets.TEST_PYPI_PASSWORD }} 34 | repository_url: https://test.pypi.org/legacy/ 35 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /dist/ 2 | /build/ 3 | /*.egg-info/ 4 | /.coverage 5 | /htmlcov 6 | /docs/_build 7 | /docs/api 8 | /.*_localrc 9 | /GNUmakefile 10 | /tags 11 | __pycache__ 12 | /.vscode 13 | /.idea -------------------------------------------------------------------------------- /.readthedocs.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | 3 | sphinx: 4 | configuration: docs/conf.py 5 | 6 | python: 7 | version: 3.8 8 | install: 9 | - requirements: docs/requirements.txt 10 | - method: setuptools 11 | path: . 12 | system_packages: true 13 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | **PyPI packages**: 2 | 3 | - Remove unexpected `p` and `s` directory in wheel file (sorry!) 4 | - Bump Taichi requirement to v0.6.32. 5 | 6 | **GUI system**: 7 | 8 | - Support choosing `ipython` or `matplotlib` as backend in `ts.Animation`, e.g.: 9 | ```py 10 | self.gui_backend = 'ipython' 11 | ``` 12 | 13 | This would allows you to display GUI result in Jupyter notebook, checkout the [config file](https://github.com/taichi-dev/taichi_glsl/blob/master/jupyter_notebook_config.py). 14 | 15 | **Field sampling**: 16 | 17 | - `dflSample(field, P, dfl)`: return a default value `dfl` when `P` out of range. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Taichi GLSL 2 | =========== 3 | 4 | Taichi GLSL is an extension library of the [Taichi Programming Language](https://github.com/taichi-dev/taichi), which provides a set of useful helper functions including but not limited to: 5 | 6 | 1. Handy scalar functions like ``clamp``, ``smoothstep``, ``mix``, ``round``. 7 | 2. GLSL-alike vector functions like ``normalize``, ``distance``, ``reflect``. 8 | 3. Well-behaved random generators including ``randUnit3D``, ``randNDRange``. 9 | 4. Handy vector and matrix initializer: ``vec`` and ``mat``. 10 | 5. Handy vector component shuffle accessor like ``v.xy``. 11 | 6. Handy field sampler including ``bilerp`` and ``sample``. 12 | 7. Useful physics helper functions like ``boundReflect``. 13 | 8. Shadertoy-alike inputed GUI base class ``Animation``. 14 | 15 | [[Clike me for documentation]](https://taichi-glsl.readthedocs.io) 16 | 17 | [![Build Status](https://img.shields.io/github/workflow/status/taichi-dev/taichi_glsl/Persubmit%20Checks)](https://github.com/taichi-dev/taichi_glsl/actions?query=workflow%3A%22Persubmit+Checks%22) 18 | [![Documentation Status](https://readthedocs.org/projects/taichi-glsl/badge?version=latest)](https://taichi-glsl.readthedocs.io/en/latest) 19 | [![Coverage Status](https://img.shields.io/codecov/c/github/taichi-dev/taichi_glsl)](https://codecov.io/gh/taichi-dev/taichi_glsl) 20 | [![Downloads](https://pepy.tech/badge/taichi-glsl/month)](https://pepy.tech/project/taichi-glsl/month) 21 | [![Latest Release](https://img.shields.io/github/v/release/taichi-dev/taichi_glsl)](https://github.com/taichi-dev/taichi_glsl/releases) 22 | 23 | 24 | Installation 25 | ------------ 26 | 27 | Install Taichi and Taichi GLSL with `pip`: 28 | 29 | ```bash 30 | # Python 3.6/3.7/3.8 (64-bit) 31 | pip install taichi taichi_glsl 32 | ``` 33 | 34 | 35 | How to play 36 | ----------- 37 | 38 | First, import Taichi and Taichi GLSL: 39 | ```py 40 | import taichi as ti 41 | import taichi_glsl as ts 42 | ``` 43 | 44 | --- 45 | 46 | Then, use `ts.xxx` helper functions in your Taichi kernels like this: 47 | ```py 48 | @ti.kernel 49 | def kern(): 50 | a = ts.vec(2.2, -3.3) # deduced to be vec2 51 | b = ts.normalize(a) # get normalized vector 52 | c = ts.clamp(a) # element-wise, clamp to range [0, 1] 53 | d = int(a) # cast to ivec2, vector of integers 54 | print(b, c, d) # [0.554700, -0.832050] [1.000000, 0.000000] [2, -3] 55 | ``` 56 | 57 | Hints 58 | ----- 59 | 60 | If you don't like the `ts.` prefix, import using: 61 | 62 | ```py 63 | from taichi_glsl import * 64 | 65 | @ti.kernel 66 | def kern(): 67 | a = vec(2.33, 6.66) 68 | b = normalize(a) 69 | ... 70 | ``` 71 | 72 | Note that this will import `taichi` as name `ti` as well. 73 | 74 | --- 75 | 76 | `vec2`, `vec3` and `vec4` are simply `vec` in Taichi GLSL: 77 | 78 | ```py 79 | v = vec(2.0, 3.0) # vec2 80 | v = vec(2.0, 3.0, 4.0) # vec3 81 | v = vec(2.0, 3.0, 4.0, 5.0) # vec4 82 | v = vec(2, 3) # ivec2 (since 2 is an integer) 83 | ``` 84 | 85 | Thanks to the python syntax of `vec(*args)`. 86 | 87 | Example 88 | ------- 89 | 90 | The following codes shows up an [Shadertoy](https://shadertoy.com/new)-style rainbow UV 91 | in the window: 92 | 93 | ```py 94 | import taichi as ti 95 | import taichi_glsl as ts 96 | 97 | ti.init() 98 | 99 | 100 | class MyAnimation(ts.Animation): 101 | def on_init(self): 102 | self.img = ti.Vector(3, ti.f32, (512, 512)) 103 | self.define_input() 104 | 105 | @ti.kernel 106 | def on_render(self): 107 | for I in ti.grouped(self.img): 108 | uv = I / self.iResolution 109 | self.img[I] = ti.cos(uv.xyx + self.iTime + 110 | ts.vec(0, 2, 4)) * 0.5 + 0.5 111 | 112 | 113 | MyAnimation().start() 114 | ``` 115 | 116 | Check out more examples in the `examples/` folder. 117 | 118 | 119 | Links 120 | ----- 121 | 122 | * [Taichi Main Repo](https://github.com/taichi-dev/taichi) 123 | * [Taichi GLSL Documentation](https://taichi-glsl.readthedocs.io) 124 | * [Taichi THREE Repo](https://github.com/taichi-dev/taichi_three) 125 | * [Taichi Documentation](https://taichi.readthedocs.io/en/stable) 126 | * [Taichi 中文文档](https://taichi.readthedocs.io/zh_CN/stable) 127 | * [Forum Thread for Taichi GLSL](https://forum.taichi.graphics/t/taichi-glsl-a-handy-extension-library-for-taichi/867) 128 | -------------------------------------------------------------------------------- /a.py: -------------------------------------------------------------------------------- 1 | import taichi_glsl as tl 2 | 3 | u = tl.vec(2, 3, 4, 5) 4 | 5 | print(u.y_xW) 6 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = . 9 | BUILDDIR = _build 10 | 11 | default: html 12 | @: 13 | 14 | # Put it first so that "make" without argument is like "make help". 15 | help: 16 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 17 | 18 | .PHONY: help Makefile 19 | 20 | # Catch-all target: route all unknown targets to Sphinx using the new 21 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 22 | %: Makefile 23 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 24 | 25 | ### Custom commands 26 | 27 | view: html 28 | firefox _build/html/index.html 29 | -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | # import os 14 | # import sys 15 | # sys.path.insert(0, os.path.abspath('.')) 16 | 17 | # -- Project information ----------------------------------------------------- 18 | 19 | project = 'taichi-glsl' 20 | copyright = '2020, 彭于斌' 21 | author = '彭于斌' 22 | 23 | # The full version, including alpha/beta/rc tags 24 | with open('../taichi_glsl/version.py') as f: 25 | env = dict() 26 | exec(f.read(), env) 27 | release = '.'.join(map(str, env['version'])) 28 | 29 | # -- General configuration --------------------------------------------------- 30 | 31 | # Add any Sphinx extension module names here, as strings. They can be 32 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 33 | # ones. 34 | extensions = [ 35 | 'sphinx_automodapi.automodapi', 36 | ] 37 | 38 | # Add any paths that contain templates here, relative to this directory. 39 | templates_path = ['_templates'] 40 | 41 | # List of patterns, relative to source directory, that match files and 42 | # directories to ignore when looking for source files. 43 | # This pattern also affects html_static_path and html_extra_path. 44 | exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] 45 | 46 | # -- Options for HTML output ------------------------------------------------- 47 | 48 | # The theme to use for HTML and HTML Help pages. See the documentation for 49 | # a list of builtin themes. 50 | # 51 | #html_theme = 'sphinx_catalystcloud_theme' 52 | html_theme = 'sphinx_rtd_theme' 53 | 54 | # Add any paths that contain custom static files (such as style sheets) here, 55 | # relative to this directory. They are copied after the builtin static files, 56 | # so a file named "default.css" will overwrite the builtin "default.css". 57 | #html_static_path = ['_static'] 58 | -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | Taichi GLSL API references 2 | ========================== 3 | 4 | Taichi GLSL is an extension library of the `Taichi Programming Language `_, which provides a set of useful helper functions including but not limited to: 5 | 6 | 1. Handy scalar functions like ``clamp``, ``smoothstep``, ``mix``, ``round``. 7 | 2. GLSL-alike vector functions like ``normalize``, ``distance``, ``reflect``. 8 | 3. Well-behaved random generators including ``randUnit3D``, ``randNDRange``. 9 | 4. Handy vector and matrix initializer: ``vec`` and ``mat``. 10 | 5. Handy vector component shuffle accessor like ``v.xy``. 11 | 6. Handy field sampler including ``bilerp`` and ``sample``. 12 | 7. Useful physics helper functions like ``boundReflect``. 13 | 8. Shadertoy-alike inputed GUI base class ``Animation``. 14 | 15 | Let me know if you encountered bugs or got a good idea by `opening an 16 | issue on GitHub `_. 17 | 18 | .. note:: 19 | 20 | Here's the documentation of **Taichi GLSL, the extension library**. 21 | For the documentation of **Taichi itself**, please click 22 | `here `_ | 23 | `这里 `_. 24 | 25 | Installation 26 | ------------ 27 | 28 | To install, make sure you have installed ``taichi`` first, then install 29 | ``taichi_glsl`` via ``pip``: 30 | 31 | .. code-block:: bash 32 | 33 | python3 -m pip install --user taichi_glsl 34 | 35 | Import me using: 36 | 37 | .. code-block:: python 38 | 39 | import taichi as ti 40 | import taichi_glsl as ts 41 | 42 | Or simply: 43 | 44 | .. code-block:: python 45 | 46 | from taichi_glsl import * 47 | 48 | Note that this will import ``taichi`` as name ``ti`` as well. 49 | 50 | 51 | Scalar math 52 | ----------- 53 | 54 | .. automodapi:: taichi_glsl.scalar 55 | 56 | :no-heading: 57 | 58 | :skip: acos, asin, ceil, cos, exp, floor, log, sin, sqrt, tan 59 | 60 | Vector math 61 | ----------- 62 | 63 | .. automodapi:: taichi_glsl.vector 64 | 65 | :no-heading: 66 | 67 | Field sampling 68 | -------------- 69 | 70 | .. automodapi:: taichi_glsl.sampling 71 | 72 | :no-heading: 73 | 74 | Taichi Classes 75 | -------------- 76 | 77 | .. automodapi:: taichi_glsl.classes 78 | 79 | :no-heading: 80 | :no-inheritance-diagram: 81 | 82 | GUI Base Class 83 | -------------- 84 | 85 | .. automodapi:: taichi_glsl.gui 86 | 87 | :no-heading: 88 | :no-inheritance-diagram: 89 | 90 | Random generator 91 | ---------------- 92 | 93 | .. automodapi:: taichi_glsl.randgen 94 | 95 | :no-heading: 96 | 97 | Particle simluation 98 | ------------------- 99 | 100 | .. automodapi:: taichi_glsl.lagrangian 101 | 102 | :no-heading: 103 | 104 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | sphinx==2.4.4 2 | sphinx-rtd-theme==0.4.3 3 | sphinx-catalystcloud-theme==0.0.1 4 | sphinx-automodapi==0.12 5 | taichi==0.6.17 6 | -------------------------------------------------------------------------------- /examples/blueorange.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as tl 3 | 4 | ti.init() 5 | 6 | 7 | class MyAnimation(tl.Animation): 8 | def on_init(self, n=320): 9 | self.n = n 10 | self.title = 'Julia Set' 11 | self.img = ti.Vector(3, ti.f32, (self.n * 2, self.n)) 12 | self.define_input() 13 | 14 | @ti.kernel 15 | def on_render(self): 16 | for p in ti.grouped(self.img): 17 | c = tl.Complex(self.iMouse * 2 - 1) 18 | z = tl.Complex(p / self.n - tl.vec(1, 0.5)) @ 2 19 | iterations = 0 20 | while z.mag2 < 4 and iterations < 50: 21 | z = z @ z + c 22 | iterations += 1 23 | self.img[p] = tl.blueorange(1 - iterations * 0.02) 24 | 25 | 26 | if __name__ == '__main__': 27 | animation = MyAnimation() 28 | animation.start() 29 | -------------------------------------------------------------------------------- /examples/export_video.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | 4 | ti.init() 5 | 6 | 7 | class MyAnimation(ts.Animation): 8 | def on_init(self): 9 | self.img = ti.Vector.field(3, ti.f32, (512, 512)) 10 | self.set_output_video('/tmp/video.gif') 11 | self.define_input() 12 | 13 | @ti.kernel 14 | def on_render(self): 15 | for I in ti.grouped(self.img): 16 | uv = I / self.iResolution 17 | self.img[I] = ti.cos(uv.xyx + self.iTime + 18 | ts.vec(0, 2, 4)) * 0.5 + 0.5 19 | 20 | 21 | if __name__ == '__main__': 22 | animation = MyAnimation() 23 | animation.start() 24 | -------------------------------------------------------------------------------- /examples/fluid_v1.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class MyAnimation(ts.Animation): 9 | def on_init(self): 10 | res = 512, 512 11 | self.color = ts.Maccormack.make(lambda: ti.var(ti.f32, res)) 12 | self.img = self.color.new 13 | self.dt = 0.04 14 | self.dx = 1 / res[0] 15 | self.define_input() 16 | 17 | @ti.func 18 | def velocity(self, P): 19 | p = P * self.dx 20 | return (2 * p - 1).Yx 21 | 22 | @ti.kernel 23 | def on_start(self): 24 | for I in ti.grouped(self.color): 25 | self.color.old[I] = ts.imageTaichi(I / self.iResolution, 0.9) 26 | 27 | @ti.kernel 28 | def on_advance(self): 29 | self.color.advance(self) 30 | self.color.update() 31 | 32 | 33 | if __name__ == '__main__': 34 | animation = MyAnimation() 35 | animation.start() -------------------------------------------------------------------------------- /examples/fluid_v2.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class MyAnimation(ts.Animation): 9 | def on_init(self): 10 | res = 512, 512 11 | self.color = ts.Maccormack.make(lambda: ti.var(ti.f32, res)) 12 | self.vel = ts.Maccormack.make(lambda: ti.Vector(2, ti.f32, res)) 13 | self.pre = ts.Pair.make(lambda: ti.var(ti.f32, res)) 14 | self.div = ti.var(ti.f32, res) 15 | self.img = self.color.new 16 | self.dt = 0.004 17 | self.dx = 1 / res[0] 18 | self.define_input() 19 | 20 | @ti.func 21 | def velocity(self, P): 22 | return ts.bilerp(self.vel.old, P) 23 | 24 | @ti.func 25 | def clearDiv(self): 26 | for I in ti.grouped(ti.ndrange(*self.pre.old.shape)): 27 | self.div[I] = ts.vgridDivergence(self.vel.new, I) 28 | for _ in ti.static(range(5)): 29 | for I in ti.grouped(ti.ndrange(*self.pre.old.shape)): 30 | pa = ts.vgridSumAround(self.pre.old, I) 31 | self.pre.new[I] = (pa - self.div[I] * self.dx**2) * 0.25 32 | self.pre.update() 33 | for I in ti.grouped(ti.ndrange(*self.pre.old.shape)): 34 | grad = ts.vgridGradient(self.div, I) / self.dx 35 | self.vel.new[I] -= grad * self.dt * 0.03 36 | 37 | @ti.kernel 38 | def on_start(self): 39 | for I in ti.grouped(self.color): 40 | self.color.old[I] = ts.imageChess(I / self.iResolution) 41 | for P in ti.grouped(self.color): 42 | p = P * self.dx 43 | self.vel.old[P] = (2 * p - 1).yx * ts.vec(-1, 1) 44 | 45 | @ti.kernel 46 | def on_advance(self): 47 | self.vel.advance(self) 48 | self.color.advance(self) 49 | self.vel.update() 50 | self.color.update() 51 | for i in ti.static(range(5)): 52 | self.clearDiv() 53 | 54 | 55 | if __name__ == '__main__': 56 | animation = MyAnimation() 57 | animation.start() 58 | -------------------------------------------------------------------------------- /examples/fractal.py: -------------------------------------------------------------------------------- 1 | import matplotlib.cm as cm 2 | import taichi as ti 3 | 4 | import taichi_glsl as tl 5 | 6 | ti.init() 7 | 8 | 9 | class MyAnimation(tl.Animation): 10 | def on_init(self, n=512): 11 | self.n = n 12 | self.title = 'Julia Set' 13 | self.img = ti.var(ti.f32, (self.n * 2, self.n)) 14 | self.colormap = cm.get_cmap('magma') 15 | self.define_input() 16 | 17 | @ti.kernel 18 | def on_render(self): 19 | for p in ti.grouped(self.img): 20 | c = tl.Complex(self.iMouse * 2 - 1) 21 | z = tl.Complex(p / self.n - tl.vec(1, 0.5)) @ 2 22 | iterations = 0 23 | while z.mag2 < 4 and iterations < 100: 24 | z = z @ z + c 25 | iterations += 1 26 | self.img[p] = 1 - iterations * 0.01 27 | 28 | 29 | if __name__ == '__main__': 30 | animation = MyAnimation() 31 | animation.start() 32 | -------------------------------------------------------------------------------- /examples/mass_spring.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as tl 3 | import taichi_three as t3 4 | 5 | import numpy as np 6 | import math 7 | 8 | ti.init(ti.gpu) 9 | 10 | ### Parameters 11 | 12 | dt, beta, steps = 5e-3, 0, 12 13 | #dt, beta, steps = 1e-2, 0.5, 5 14 | beta_dt = beta * dt 15 | alpha_dt = (1 - beta) * dt 16 | jacobi_steps = 15 17 | 18 | N = 128 19 | NN = N, N 20 | W = 1 21 | L = W / N 22 | gravity = 0.4 23 | stiff = 20 24 | damp = 2.6 25 | 26 | ### Generic helpers 27 | 28 | x = ti.Vector(3, ti.f32, NN) 29 | v = ti.Vector(3, ti.f32, NN) 30 | b = ti.Vector(3, ti.f32, NN) 31 | F = ti.Vector(3, ti.f32, NN) 32 | 33 | 34 | @ti.kernel 35 | def init(): 36 | for i in ti.grouped(x): 37 | x[i] = tl.vec((i + 0.5) * L - 0.5, 0.8).xzy 38 | 39 | 40 | links = [tl.vec(*_) for _ in [(-1, 0), (1, 0), (0, -1), (0, 1)]] 41 | 42 | 43 | @ti.func 44 | def accel(v: ti.template(), x: ti.template(), dt): 45 | for i in ti.grouped(x): 46 | acc = x[i] * 0 47 | for d in ti.static(links): 48 | disp = x[tl.clamp(i + d, 0, tl.vec(*NN) - 1)] - x[i] 49 | dis = disp.norm() 50 | acc += disp * (dis - L) / L**2 51 | v[i] += stiff * acc * dt 52 | v[i] *= ti.exp(-damp * dt) 53 | 54 | 55 | @ti.kernel 56 | def collide(x: ti.template(), v: ti.template()): 57 | for i in ti.grouped(x): 58 | v[i].y -= gravity * dt 59 | for i in ti.grouped(x): 60 | v[i] = ballBoundReflect(x[i], v[i], tl.vec(+0.0, +0.2, -0.0), 0.4) 61 | 62 | 63 | @ti.func 64 | def ballBoundReflect(pos, vel, center, radius): 65 | ret = vel 66 | above = tl.distance(pos, center) - radius 67 | if above <= 0: 68 | normal = tl.normalize(pos - center) 69 | NoV = tl.dot(vel, normal) - 6 * tl.smoothstep(above, 0, -0.1) 70 | if NoV < 0: 71 | ret -= NoV * normal 72 | return ret 73 | 74 | 75 | ### Explicit method 76 | 77 | 78 | @ti.kernel 79 | def explicit_accel(): 80 | accel(v, x, dt) 81 | 82 | 83 | @ti.kernel 84 | def explicit_update(): 85 | for i in ti.grouped(x): 86 | v[i] *= math.exp(dt * -damp) 87 | x[i] += dt * v[i] 88 | 89 | 90 | def explicit(): 91 | ''' 92 | v' = v + Mdt @ v 93 | x' = x + v'dt 94 | ''' 95 | explicit_accel() 96 | collide(x, v) 97 | explicit_update() 98 | 99 | 100 | ### Implicit method 101 | ''' 102 | v' = v + Mdt @ v' 103 | (I - Mdt) @ v' = v 104 | ''' 105 | ''' 106 | v' = v + Mdt @ [beta v' + alpha v] 107 | (I - beta Mdt) @ v' = (I + alpha Mdt) @ v 108 | ''' 109 | 110 | 111 | @ti.kernel 112 | def prepare(): 113 | if ti.static(beta != 1): 114 | for i in ti.grouped(x): 115 | x[i] += v[i] * alpha_dt 116 | accel(v, x, alpha_dt) 117 | for i in ti.grouped(x): 118 | b[i] = x[i] 119 | x[i] += v[i] * beta_dt 120 | 121 | 122 | @ti.kernel 123 | def jacobi(): 124 | for i in ti.grouped(x): 125 | b[i] = x[i] + F[i] * beta_dt**2 126 | F[i] = b[i] * 0 127 | accel(F, b, 1) 128 | 129 | 130 | @ti.kernel 131 | def update_pos(): 132 | for i in ti.grouped(x): 133 | x[i] = b[i] 134 | v[i] += F[i] * beta_dt 135 | 136 | 137 | def implicit(): 138 | prepare() 139 | for i in range(jacobi_steps): 140 | jacobi() 141 | collide(b, v) 142 | update_pos() 143 | 144 | 145 | ### Rendering GUI 146 | 147 | scene = t3.Scene() 148 | model = t3.Model() 149 | scene.add_model(model) 150 | 151 | faces = t3.Face.var(N**2 * 2) 152 | #lines = t3.Line.var(N**2 * 2) 153 | vertices = t3.Vertex.var(N**2) 154 | model.set_vertices(vertices) 155 | model.add_geometry(faces) 156 | #model.add_geometry(lines) 157 | 158 | 159 | @ti.kernel 160 | def init_display(): 161 | for i_ in ti.grouped(ti.ndrange(N - 1, N - 1)): 162 | i = i_ 163 | a = i.dot(tl.vec(N, 1)) 164 | i.x += 1 165 | b = i.dot(tl.vec(N, 1)) 166 | i.y += 1 167 | c = i.dot(tl.vec(N, 1)) 168 | i.x -= 1 169 | d = i.dot(tl.vec(N, 1)) 170 | i.y -= 1 171 | faces[a * 2 + 0].idx = tl.vec(a, c, b) 172 | faces[a * 2 + 1].idx = tl.vec(a, d, c) 173 | #faces[a * 4 + 2].idx = tl.vec(b, a, d) 174 | #faces[a * 4 + 3].idx = tl.vec(b, d, c) 175 | #lines[a * 2 + 0].idx = tl.vec(a, b) 176 | #lines[a * 2 + 1].idx = tl.vec(a, d) 177 | 178 | 179 | @ti.kernel 180 | def update_display(): 181 | for i in ti.grouped(x): 182 | vertices[i.dot(tl.vec(N, 1))].pos = x[i] 183 | 184 | 185 | init() 186 | init_display() 187 | scene.set_light_dir([0.4, -1.5, -1.8]) 188 | 189 | print('[Hint] mouse drag to orbit camera') 190 | with ti.GUI('Mass Spring') as gui: 191 | gui.frame = 0 192 | while gui.running and not gui.get_event(gui.ESCAPE): 193 | if not gui.is_pressed(gui.SPACE): 194 | for i in range(steps): 195 | if beta == 0: 196 | explicit() 197 | else: 198 | implicit() 199 | update_display() 200 | if gui.is_pressed(gui.LMB): 201 | scene.camera.from_mouse(gui) 202 | else: 203 | scene.camera.from_mouse([0.5, 0.4]) 204 | 205 | scene.render() 206 | gui.set_image(scene.img) 207 | #gui.show(f'/tmp/{gui.frame:06d}.png') 208 | gui.show() 209 | gui.frame += 1 210 | -------------------------------------------------------------------------------- /examples/old_uv.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | image = vec_array(3, float, 512, 512) 4 | 5 | 6 | @ti.kernel 7 | def paint(): 8 | for i, j in image: 9 | coor = view(image, i, j) 10 | image[i, j] = vec(coor.xy, 0.0) 11 | 12 | 13 | if __name__ == '__main__': 14 | paint() 15 | ti.imshow(image) 16 | -------------------------------------------------------------------------------- /examples/paint_arrow.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class MyAnimation(ts.Animation): 9 | def on_init(self): 10 | self.title = 'Test Arrow' 11 | self.img = ti.var(ti.f32, (512, 512)) 12 | self.auto_clean = True 13 | self.define_input() 14 | 15 | @ti.kernel 16 | def on_render(self): 17 | ts.paintArrow(self.img, ts.vec2(0.0), self.iMouse) 18 | 19 | 20 | if __name__ == '__main__': 21 | animation = MyAnimation() 22 | animation.start() 23 | -------------------------------------------------------------------------------- /examples/particles.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class MyAnimation(ts.Animation): 9 | def on_init(self): 10 | self.N = 8192 11 | self.dt = 0.01 12 | self.pos = ti.Vector(2, ti.f32, self.N) 13 | self.vel = ti.Vector(2, ti.f32, self.N) 14 | self.circles = self.pos # alias to make ts.Animation know 15 | self.attract_strength = ti.var(ti.f32, ()) 16 | self.attract_pos = ti.Vector(2, ti.f32, ()) 17 | self.resolution = (512, 512) 18 | self.title = 'Particles' 19 | self.define_input() 20 | 21 | @ti.kernel 22 | def on_start(self): 23 | for i in self.pos: 24 | self.pos[i] = ts.randND(2) 25 | self.vel[i] = ts.randSolid2D() 26 | 27 | @ti.kernel 28 | def on_advance(self): 29 | for i in self.pos: 30 | acc = ts.vec(0.0, -1.0) 31 | if any(self.iKeyDirection): # ASWD? 32 | acc = self.iKeyDirection 33 | if any(self.iMouseButton): 34 | dir = ts.normalize(self.iMouse - self.pos[i]) * 2 35 | if self.iMouseButton[0]: # LMB pressed? 36 | acc += dir 37 | if self.iMouseButton[1]: # RMB pressed? 38 | acc -= dir 39 | self.vel[i] += acc * self.dt 40 | for i in self.pos: 41 | self.vel[i] = ts.boundReflect(self.pos[i], self.vel[i], 0, 1, 0.8) 42 | for i in self.pos: 43 | self.pos[i] += self.vel[i] * self.dt 44 | 45 | 46 | if __name__ == '__main__': 47 | animation = MyAnimation() 48 | animation.start() 49 | -------------------------------------------------------------------------------- /examples/rand.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | res = 512 4 | 5 | img = array(float, res, res) 6 | 7 | 8 | @ti.kernel 9 | def render(): 10 | for i, j in img: 11 | img[i, j] *= 0.94 12 | for i in range(200): 13 | # Try comment and uncomment these: 14 | # p = randND(2) 15 | # p = 0.5 + 0.3 * randUnit2D() 16 | p = 0.5 + 0.3 * randSolid2D() 17 | img[int(p * res + 0.5)] = 1 18 | 19 | 20 | gui = ti.GUI('Random') 21 | if __name__ == '__main__': 22 | while not gui.get_event(ti.GUI.ESCAPE): 23 | render() 24 | gui.set_image(img) 25 | gui.show() 26 | -------------------------------------------------------------------------------- /examples/rand3d.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | res = 512 4 | 5 | img = vec_array(3, float, res, res) 6 | 7 | 8 | @ti.kernel 9 | def render(): 10 | for i, j in img: 11 | img[i, j] *= 0.94 12 | for i in range(200): 13 | p = 0.5 + 0.3 * randUnit3D() 14 | img[int(shuffle(p, 0, 1) * res + 0.5)][0] = 1 15 | img[int(shuffle(p, 1, 2) * res + 0.5)][1] = 1 16 | img[int(shuffle(p, 2, 0) * res + 0.5)][2] = 1 17 | 18 | 19 | gui = ti.GUI('Random') 20 | while not gui.get_event(ti.GUI.ESCAPE): 21 | render() 22 | gui.set_image(img) 23 | gui.show() 24 | -------------------------------------------------------------------------------- /examples/sdf_rt1.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | ti.init(ti.opengl) 4 | 5 | fov = math.tan(math.radians(26)) 6 | img = vec_array(3, float, 512, 512) 7 | mouse = vec_uniform(2, float) 8 | eps = 1e-5 9 | 10 | 11 | @ti.func 12 | def sdfSphere(p, c, r): 13 | return distance(p, c) - r 14 | 15 | 16 | @ti.func 17 | def mUnion(a, b): 18 | if a[0] > b[0]: 19 | a = b 20 | return a 21 | 22 | 23 | @ti.func 24 | def mIntersect(a, b): 25 | if a[0] < b[0]: 26 | a = b 27 | return a 28 | 29 | 30 | @ti.func 31 | def mScene(p): 32 | m = mouse[None] 33 | return mUnion( # center rad emi spec 34 | vec(sdfSphere(p, vec(m.x, m.y, 0.0), 0.2), 1.0, 0.0), 35 | vec(sdfSphere(p, vec(0.0, 0.0, 0.4), 0.5), 0.2, 1.0)) 36 | 37 | 38 | @ti.func 39 | def gradScene(p, sdf): 40 | return normalize( 41 | vec( 42 | mScene(p + vec(eps, 0, 0))[0], 43 | mScene(p + vec(0, eps, 0))[0], 44 | mScene(p + vec(0, 0, eps))[0]) - sdf) 45 | 46 | 47 | @ti.func 48 | def radiance(eye, dir): 49 | p = eye 50 | clr = vec3(0.0) 51 | depth = 0.0 52 | for i in range(50): 53 | res = mScene(p) 54 | sdf, emi, spec = res[0], res[1], res[2] 55 | if sdf < eps: 56 | clr += vec3(emi) 57 | if spec != 0: 58 | if rand() < spec: 59 | norm = gradScene(p, sdf) 60 | dir = reflect(dir, norm) 61 | p += 2 * eps * dir 62 | continue 63 | break 64 | depth += sdf 65 | p += sdf * dir 66 | return clr 67 | 68 | 69 | @ti.kernel 70 | def render(): 71 | eye = vec(0.0, 0.0, -1.8) 72 | for i, j in img: 73 | coor = fov * (view(img, i, j) * 2.0 - 1.0) 74 | dir = normalize(vec(coor, 1.0)) 75 | img[i, j] = radiance(eye, dir) 76 | 77 | 78 | mouse[None] = [-0.5, 0.5] 79 | gui = ti.GUI('SDF-RT1') 80 | while gui.running: 81 | gui.running = not gui.get_event(gui.ESCAPE) 82 | if gui.is_pressed(gui.LMB): 83 | x, y = gui.get_cursor_pos() 84 | mouse[None] = [x * 2 - 1, y * 2 - 1] 85 | render() 86 | gui.set_image(img) 87 | gui.show() 88 | -------------------------------------------------------------------------------- /examples/sdf_rt2.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | ti.init(ti.opengl) 4 | 5 | fov = math.tan(math.radians(26)) 6 | img = vec_array(3, float, 512, 512) 7 | mouse = vec_uniform(2, float) 8 | eps = 1e-5 9 | 10 | 11 | @ti.func 12 | def sdfSphere(p, c, r): 13 | return distance(p, c) - r 14 | 15 | 16 | @ti.func 17 | def mUnion(a, b): 18 | if a[0] > b[0]: 19 | a = b 20 | return a 21 | 22 | 23 | @ti.func 24 | def mIntersect(a, b): 25 | if a[0] < b[0]: 26 | a = b 27 | return a 28 | 29 | 30 | @ti.func 31 | def mScene(p): 32 | m = mouse[None] 33 | return mUnion( # center rad emi spec diff 34 | vec(sdfSphere(p, vec(m.x, m.y, 0.0), 0.2), 1.0, 0.0, 0.0), 35 | vec(sdfSphere(p, vec(0.0, 0.0, 0.4), 0.5), 0.0, 0.0, 1.0)) 36 | 37 | 38 | @ti.func 39 | def gradScene(p, sdf): 40 | return normalize( 41 | vec( 42 | mScene(p + vec(eps, 0, 0))[0], 43 | mScene(p + vec(0, eps, 0))[0], 44 | mScene(p + vec(0, 0, eps))[0]) - sdf) 45 | 46 | 47 | @ti.func 48 | def radiance(eye, dir): 49 | p = eye 50 | clr = vec3(0.0) 51 | depth = 0.0 52 | for i in range(60): 53 | res = mScene(p) 54 | sdf, emi, spec, diff = res 55 | if sdf < eps: 56 | clr += vec3(emi) 57 | if spec != 0: 58 | if rand() < spec: 59 | norm = gradScene(p, sdf) 60 | dir = reflect(dir, norm) 61 | p += 2 * eps * dir 62 | continue 63 | if diff != 0: 64 | if rand() < diff: 65 | norm = gradScene(p, sdf) 66 | dir = randUnit3D() 67 | NoD = dot(dir, norm) 68 | if NoD < 0: 69 | dir -= 2 * NoD * norm 70 | p += 2 * eps * dir 71 | continue 72 | break 73 | depth += sdf 74 | p += sdf * dir 75 | return clr 76 | 77 | 78 | @ti.kernel 79 | def render(): 80 | eye = vec(0.0, 0.0, -1.8) 81 | for i, j in img: 82 | coor = fov * (view(img, i, j) * 2.0 - 1.0) 83 | dir = normalize(vec(coor, 1.0)) 84 | for t in range(16): 85 | img[i, j] += radiance(eye, dir) 86 | 87 | 88 | mouse[None] = [-0.5, 0.5] 89 | gui = ti.GUI('SDF-RT2') 90 | count = 0 91 | while gui.running: 92 | gui.running = not gui.get_event(gui.ESCAPE) 93 | if gui.is_pressed(gui.LMB): 94 | x, y = gui.get_cursor_pos() 95 | mouse[None] = [x * 2 - 1, y * 2 - 1] 96 | img.fill(0) 97 | count = 0 98 | render() 99 | gui.set_image(np.sqrt(img.to_numpy() * (1 / 16 / (count + 1)))) 100 | gui.show() 101 | count += 1 102 | -------------------------------------------------------------------------------- /examples/simple_vortices.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class MyAnimation(ts.Animation): 9 | def on_init(self): 10 | res = 512, 512 11 | self.color = ts.Maccormack.make(lambda: ti.var(ti.f32, res)) 12 | self.vorts = ti.Vector(2, ti.f32, 4) 13 | self.circles = self.vorts 14 | self.img = self.color.new 15 | self.dt = 0.04 16 | self.dx = 1 / res[0] 17 | self.define_input() 18 | 19 | @ti.func 20 | def velocity(self, P): 21 | p = P * self.dx 22 | vel = ts.vec2(0.0) 23 | for v in range(self.vorts.shape[0]): 24 | dis = (p - self.vorts[v]) 25 | dis = ts.normalizePow(dis, -1, 0.001) 26 | dis = dis.yx * ts.vec(1, -1) 27 | vel += dis 28 | return vel * 0.01 29 | 30 | @ti.kernel 31 | def on_start(self): 32 | for I in ti.grouped(self.color): 33 | p = I / self.iResolution 34 | self.color.old[I] = ts.imageGrid(p) 35 | self.vorts[0] = ts.vec(0.25, 0.5) 36 | self.vorts[1] = ts.vec(0.75, 0.5) 37 | self.vorts[2] = ts.vec(0.5, 0.25) 38 | self.vorts[3] = ts.vec(0.5, 0.75) 39 | 40 | @ti.kernel 41 | def on_advance(self): 42 | self.color.advance(self) 43 | self.color.update() 44 | 45 | 46 | if __name__ == '__main__': 47 | animation = MyAnimation() 48 | animation.start() 49 | -------------------------------------------------------------------------------- /examples/smoke.py: -------------------------------------------------------------------------------- 1 | import matplotlib.cm as cm 2 | import taichi as ti 3 | 4 | import taichi_glsl as ts 5 | 6 | ti.init(arch=ti.gpu) 7 | 8 | cmap = cm.get_cmap('magma') 9 | 10 | N = 512 11 | dx = 1 / N 12 | dt = 0.001 13 | 14 | 15 | class Pair: 16 | def __init__(self, f): 17 | self.n = f() 18 | self.o = f() 19 | 20 | def swap(self): 21 | self.n, self.o = self.o, self.n 22 | 23 | 24 | dye = Pair(lambda: ts.array(float, N, N)) 25 | pre = Pair(lambda: ts.array(float, N, N)) 26 | vel = Pair(lambda: ti.Vector(2, ti.f32, (N, N))) 27 | div = ts.array(float, N, N) 28 | 29 | 30 | @ti.kernel 31 | def initdye(): 32 | for I in ti.grouped(dye.o): 33 | dye.o[I] = ts.imageChess(I / N) 34 | 35 | 36 | @ti.kernel 37 | def initrotv(): 38 | for I in ti.grouped(dye.o): 39 | vel.o[I] = (I - N / 2).yx * ts.D.zx 40 | 41 | 42 | @ti.func 43 | def backtrace(v: ti.template(), I, dt): 44 | midI = I - 0.5 * ts.bilerp(v, I) * dt 45 | finI = I - dt * ts.bilerp(v, midI) 46 | return finI 47 | 48 | 49 | @ti.kernel 50 | def advect(fn: ti.template(), f: ti.template(), v: ti.template()): 51 | for I in ti.grouped(f): 52 | btI = backtrace(v, I, dt) 53 | ftI = backtrace(v, btI, -dt) 54 | f_btI = ts.bilerp(f, btI) 55 | f_ftI = ts.bilerp(f, ftI) 56 | fn[I] = f_btI + 0.5 * (f_ftI - f[I]) 57 | 58 | 59 | @ti.kernel 60 | def compute_div(v: ti.template()): 61 | for I in ti.grouped(v): 62 | l = ts.sample(v, I + ts.D.zy).x 63 | r = ts.sample(v, I + ts.D.xy).x 64 | b = ts.sample(v, I + ts.D.yz).y 65 | t = ts.sample(v, I + ts.D.yx).y 66 | d = r - l + t - b 67 | div[I] = d * (0.5 / dx) 68 | 69 | 70 | @ti.kernel 71 | def jacobi(pn: ti.template(), p: ti.template()): 72 | for I in ti.grouped(p): 73 | l = ts.sample(p, I + ts.D.zy) 74 | r = ts.sample(p, I + ts.D.xy) 75 | b = ts.sample(p, I + ts.D.yz) 76 | t = ts.sample(p, I + ts.D.yx) 77 | sa = r + l + t + b 78 | pn[I] = (sa - dx**2 * div[I]) * 0.25 79 | 80 | 81 | @ti.kernel 82 | def gauss_seidel(pn: ti.template(), p: ti.template()): 83 | for I in ti.grouped(p): 84 | if I.sum() % 2 == 0: 85 | l = ts.sample(p, I + ts.D.zy) 86 | r = ts.sample(p, I + ts.D.xy) 87 | b = ts.sample(p, I + ts.D.yz) 88 | t = ts.sample(p, I + ts.D.yx) 89 | sa = r + l + t + b 90 | pn[I] = (sa - dx**2 * div[I]) * 0.25 91 | for I in ti.grouped(p): 92 | if I.sum() % 2 == 1: 93 | l = ts.sample(pn, I + ts.D.zy) 94 | r = ts.sample(pn, I + ts.D.xy) 95 | b = ts.sample(pn, I + ts.D.yz) 96 | t = ts.sample(pn, I + ts.D.yx) 97 | sa = r + l + t + b 98 | pn[I] = (sa - dx**2 * div[I]) * 0.25 99 | 100 | 101 | @ti.kernel 102 | def subgrad(v: ti.template(), p: ti.template()): 103 | for I in ti.grouped(v): 104 | l = ts.sample(p, I + ts.D.zy) 105 | r = ts.sample(p, I + ts.D.xy) 106 | b = ts.sample(p, I + ts.D.yz) 107 | t = ts.sample(p, I + ts.D.yx) 108 | g = ts.vec(r - l, t - b) 109 | v[I] = v[I] - g * (0.5 / dx) 110 | 111 | 112 | @ti.kernel 113 | def pump(v: ti.template(), d: ti.template(), a: ti.f32): 114 | pump_strength = ti.static(0.1) 115 | X, Y = ti.static(15, 15) 116 | for x, y in ti.ndrange((-X, X + 1), (-Y + 1, Y)): 117 | I = ts.vec(N // 2 + x, Y + y) 118 | s = ((Y - abs(y)) / Y * (X - abs(x)) / X)**2 119 | v[I] += ts.vecAngle(a + ts.pi / 2) * s * (pump_strength / dt) * 7.8 120 | d[I] += s * (dt / pump_strength) * 21.3 121 | 122 | 123 | if __name__ == '__main__': 124 | gui = ti.GUI('advect', N) 125 | while gui.running: 126 | for e in gui.get_events(gui.PRESS): 127 | if e.key == gui.ESCAPE: 128 | gui.running = False 129 | 130 | if not gui.is_pressed(gui.SPACE): 131 | a = 0 132 | if gui.is_pressed('a', gui.LEFT): 133 | a += 0.7 134 | if gui.is_pressed('d', gui.RIGHT): 135 | a -= 0.7 136 | pump(vel.o, dye.o, a) 137 | 138 | advect(dye.n, dye.o, vel.o) 139 | advect(vel.n, vel.o, vel.o) 140 | dye.swap() 141 | vel.swap() 142 | 143 | compute_div(vel.o) 144 | for _ in range(5): 145 | jacobi(pre.n, pre.o) 146 | pre.swap() 147 | 148 | for _ in range(20): 149 | gauss_seidel(pre.n, pre.o) 150 | pre.swap() 151 | 152 | subgrad(vel.o, pre.o) 153 | 154 | gui.set_image(cmap(dye.o.to_numpy())) 155 | gui.show() 156 | -------------------------------------------------------------------------------- /examples/step_uv.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | ti.init() 4 | 5 | 6 | class MyAnimation(Animation): 7 | def on_init(self): 8 | self.title = 'Step + UV' 9 | self.img = vec_array(3, float, 512, 512) 10 | self.define_input() 11 | 12 | @ti.kernel 13 | def on_render(self): 14 | t = 0.4 + 0.4 * cos(self.iTime) 15 | for i, j in self.img: 16 | # TODO(archibate): for coor in view(self.img): 17 | uv = view(self.img, i, j) 18 | self.img[i, j] = smoothstep(distance(uv, vec(0.5, 0.5)), t, 19 | t - 0.06) * vec(uv.xy, 0.0) 20 | 21 | 22 | if __name__ == '__main__': 23 | animation = MyAnimation() 24 | animation.start() 25 | -------------------------------------------------------------------------------- /examples/taichi_logo.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | 4 | ti.init() 5 | 6 | 7 | class TaichiLogo(ts.Animation): 8 | def on_init(self): 9 | self.title = 'Taichi Logo' 10 | self.img = ti.var(ti.f32, (512, 512)) 11 | 12 | @ti.kernel 13 | def on_render(self): 14 | for i, j in self.img: 15 | uv = ts.view(self.img, i, j) 16 | self.img[i, j] = ts.imageTaichi(uv) 17 | 18 | 19 | if __name__ == '__main__': 20 | animation = TaichiLogo() 21 | animation.start() 22 | -------------------------------------------------------------------------------- /examples/taichi_logo_aa.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class MyAnimation(ts.Animation): 9 | def on_init(self): 10 | self.title = 'Taichi Logo AA2x2' 11 | self.img = ti.var(ti.f32, (512, 512)) 12 | 13 | @ti.func 14 | def fetch(self, uv): 15 | return ts.imageTaichi(uv) 16 | 17 | @ti.kernel 18 | def on_render(self): 19 | for i, j in self.img: 20 | uv = ts.view(self.img, i, j) 21 | self.img[i, j] = ts.superSample2x2(self.fetch, uv, 0.5 / 512) 22 | 23 | 24 | if __name__ == '__main__': 25 | animation = MyAnimation() 26 | animation.start() 27 | -------------------------------------------------------------------------------- /examples/uv.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | ti.init() 4 | 5 | 6 | class MyAnimation(Animation): 7 | def on_init(self): 8 | self.title = 'UV' 9 | self.img = vec_array(3, float, 512, 512) 10 | 11 | @ti.kernel 12 | def on_render(self): 13 | for i, j in self.img: 14 | uv = view(self.img, i, j) 15 | self.img[i, j] = vec(uv.xy, 0.0) 16 | 17 | 18 | if __name__ == '__main__': 19 | animation = MyAnimation() 20 | animation.start() 21 | -------------------------------------------------------------------------------- /examples/vortex_frog.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | 3 | import taichi_glsl as ts 4 | 5 | ti.init() 6 | 7 | 8 | class VortexFrog(ts.Animation): 9 | def on_init(self): 10 | res = 512, 512 11 | self.color = ts.Maccormack.make(lambda: ti.var(ti.f32, res)) 12 | self.vorts = ti.Vector(2, ti.f32, 4) 13 | self.vorty = ti.var(ti.f32, 4) 14 | self.circles = self.vorts 15 | self.circle_radius = 3 16 | self.circle_color = 0x000000 17 | self.img = self.color.new 18 | self.dt = 0.04 19 | self.dx = 1 / res[0] 20 | self.define_input() 21 | 22 | @ti.func 23 | def velocity(self, P): 24 | p = P * self.dx 25 | vel = ts.vec2(0.0) 26 | for v in range(self.vorts.shape[0]): 27 | dis = (p - self.vorts[v]) 28 | dis = ts.normalizePow(dis, -1, 0.001) 29 | dis = dis.yx * ts.vec(-1, 1) * self.vorty[v] 30 | vel += dis 31 | return vel * 0.01 32 | 33 | @ti.kernel 34 | def on_start(self): 35 | for I in ti.grouped(self.color): 36 | p = I / self.iResolution 37 | self.color.old[I] = ts.imageChess(p) 38 | self.vorts[0] = ts.vec(0.25, 0.25) 39 | self.vorts[1] = ts.vec(0.25, 0.40) 40 | self.vorts[2] = ts.vec(0.25, 0.60) 41 | self.vorts[3] = ts.vec(0.25, 0.75) 42 | self.vorty[0] = -1.0 43 | self.vorty[1] = -1.0 44 | self.vorty[2] = +1.0 45 | self.vorty[3] = +1.0 46 | 47 | @ti.kernel 48 | def on_advance(self): 49 | for v in self.vorts: 50 | vel = self.velocity(self.vorts[v] / self.dx) 51 | self.vorts[v] += vel * self.dt 52 | self.color.advance(self) 53 | self.color.update() 54 | 55 | 56 | if __name__ == '__main__': 57 | animation = VortexFrog() 58 | animation.start() 59 | -------------------------------------------------------------------------------- /jupyter_notebook_config.py: -------------------------------------------------------------------------------- 1 | # Configuration file for jupyter-notebook. 2 | 3 | c.JupyterApp.answer_yes = True 4 | c.NotebookApp.allow_password_change = False 5 | c.NotebookApp.allow_remote_access = True 6 | c.NotebookApp.answer_yes = True 7 | c.NotebookApp.autoreload = True 8 | c.NotebookApp.base_url = '/' 9 | c.NotebookApp.default_url = '/tree' 10 | c.NotebookApp.ip = 'winner' 11 | c.NotebookApp.nbserver_extensions = {} 12 | c.NotebookApp.notebook_dir = 'F:\\Documents\\JupyterRoot' 13 | c.NotebookApp.open_browser = True 14 | c.NotebookApp.password = 'argon2:$argon2id$v=19$m=10240,t=10,p=8$mWVgG+OTUw2T2afUmHQY4w$1xCmsxdz+m4Z6CJjRiuf0g' 15 | c.NotebookApp.password_required = True 16 | c.NotebookApp.port = 80 17 | c.NotebookApp.port_retries = 1 18 | c.NotebookApp.quit_button = False 19 | c.NotebookApp.rate_limit_window = 3 20 | c.NotebookApp.token = '' 21 | c.ContentsManager.root_dir = '/' 22 | c.FileContentsManager.delete_to_trash = False 23 | 24 | import os 25 | 26 | os.environ['TI_GUI_BACKEND'] = 'ipython' -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import sys, os 3 | 4 | 5 | class Main: 6 | def __init__(self): 7 | self.parser = argparse.ArgumentParser( 8 | description='Taichi GLSL CLI Tools') 9 | self.parser.add_argument('command', default='test', nargs='?') 10 | 11 | def __call__(self, argv): 12 | args = self.parser.parse_args(argv[1:2]) 13 | getattr(self, f'cmd_{args.command}')(argv[2:]) 14 | 15 | def cmd_test(self, argv): 16 | pytest = 'pytest -nauto ' 17 | tests = 'tests' 18 | if len(argv): 19 | tests = ' '.join(f'tests/test_{a}.py' for a in argv) 20 | 21 | return os.system(pytest + tests) 22 | 23 | def cmd_coverage(self, argv): 24 | return os.system('coverage run -m pytest tests') 25 | 26 | def cmd_report(self, argv): 27 | return os.system('coverage report && coverage html') 28 | 29 | def cmd_format(self, argv): 30 | return os.system(f'yapf -ir .') 31 | 32 | def cmd_doc(self, argv): 33 | return os.system(f'make -C docs ' + ' '.join(argv)) 34 | 35 | def cmd_dist(self, argv): 36 | return os.system(f'{sys.executable} setup.py bdist_wheel') 37 | 38 | def cmd_install(self, argv): 39 | return os.system(f'{sys.executable} setup.py bdist_wheel && ' 40 | 'pip install --user --upgrade dist/*.whl') 41 | 42 | 43 | if __name__ == '__main__': 44 | cli = Main() 45 | exit(cli(sys.argv)) 46 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | project_name = 'taichi_glsl' 2 | from taichi_glsl.version import version, taichi_version 3 | 4 | version = '.'.join(map(str, version)) 5 | taichi_version = '.'.join(map(str, taichi_version)) 6 | description = 'A Taichi extension library providing a set of GLSL-alike helper functions' 7 | long_description = ''' 8 | Taichi GLSL 9 | =========== 10 | 11 | Taichi GLSL is an extension library of the `Taichi Programming Language `_, which provides a set of useful helper functions including but not limited to: 12 | 13 | 1. Handy scalar functions like ``clamp``, ``smoothstep``, ``mix``, ``round``. 14 | 2. GLSL-alike vector functions like ``normalize``, ``distance``, ``reflect``. 15 | 3. Well-behaved random generators including ``randUnit3D``, ``randNDRange``. 16 | 4. Handy vector and matrix initializer: ``vec`` and ``mat``. 17 | 5. Handy vector component shuffle accessor like ``v.xy``. 18 | 6. Handy field sampler including ``bilerp`` and ``sample``. 19 | 7. Useful physics helper functions like ``boundReflect``. 20 | 8. Shadertoy-alike inputed GUI base class ``Animation``. 21 | 22 | `[Clike me for documentation] `_ 23 | 24 | `[Clike me for GitHub repo] `_ 25 | ''' 26 | classifiers = [ 27 | 'Intended Audience :: Developers', 28 | 'Intended Audience :: Science/Research', 29 | 'Topic :: Multimedia :: Graphics', 30 | 'Topic :: Games/Entertainment :: Simulation', 31 | 'Operating System :: OS Independent', 32 | ] 33 | python_requires = '>=3.6' 34 | install_requires = [ 35 | 'taichi>=' + taichi_version, 36 | ] 37 | 38 | import setuptools 39 | 40 | setuptools.setup( 41 | name=project_name, 42 | version=version, 43 | author='彭于斌', 44 | author_email='1931127624@qq.com', 45 | description=description, 46 | long_description=long_description, 47 | classifiers=classifiers, 48 | python_requires=python_requires, 49 | install_requires=install_requires, 50 | packages=setuptools.find_packages(), 51 | ) 52 | -------------------------------------------------------------------------------- /taichi_glsl/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Manipulate Taichi with GLSL-alike helper functions. 3 | ''' 4 | 5 | import taichi as ti 6 | import numpy as np 7 | import math 8 | 9 | from .version import version as __version__ 10 | from .hack import * 11 | from .scalar import * 12 | from .vector import * 13 | from .odop import * 14 | from .randgen import * 15 | from .sampling import * 16 | from .lagrangian import * 17 | from .experimental_array import * 18 | from .experimental_transform import * 19 | from .experimental_cfd import * 20 | from .colormap import * 21 | from .painting import * 22 | from .mkimages import * 23 | from .sphcore import * 24 | from .classes import * 25 | from .gui import * 26 | -------------------------------------------------------------------------------- /taichi_glsl/classes.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as tl 3 | 4 | 5 | class Complex(tl.TaichiClass, ti.lang.common_ops.TaichiOperations): 6 | ''' 7 | Complex number support in Taichi GLSL 8 | 9 | :note: 10 | `complex * complex` is the **element-wise** multiplication. 11 | Use `complex @ complex` if you mean to use complex number multiplication. 12 | ''' 13 | def __init__(self, x, y=None): 14 | '''Construct a complex from a vector / two scalars''' 15 | if y is None: 16 | if isinstance(x, (ti.Matrix, Complex)): 17 | x, y = x.entries 18 | else: 19 | y = x * 0 20 | self.entries = [x, y] 21 | 22 | @property 23 | def x(self): 24 | '''Real part of the complex''' 25 | return self.entries[0] 26 | 27 | re = x 28 | 29 | @property 30 | def y(self): 31 | '''Imaginary part of the complex''' 32 | return self.entries[1] 33 | 34 | im = y 35 | 36 | @property 37 | def mag2(self): 38 | '''Squared magnitude of the complex''' 39 | return self.x**2 + self.y**2 40 | 41 | @property 42 | def mag(self): 43 | '''Magnitude of the complex''' 44 | return ti.sqrt(self.mag2) 45 | 46 | @property 47 | def ang(self): 48 | '''Phase angle of the complex''' 49 | return ti.atan2(self.y, self.x) 50 | 51 | @classmethod 52 | def _field(cls, *args, **kwargs): 53 | x = ti.field(*args, **kwargs) 54 | y = ti.field(*args, **kwargs) 55 | return x, y 56 | 57 | def element_wise_unary(a, foo): 58 | return Complex(foo(a.x), foo(a.y)) 59 | 60 | def element_wise_binary(a, foo, b): 61 | if not isinstance(b, Complex): 62 | b = Complex(b) 63 | return Complex(foo(a.x, b.x), foo(a.y, b.y)) 64 | 65 | def element_wise_writeback_binary(a, foo, b): 66 | if ti.is_taichi_class(b): 67 | b = b.variable() 68 | if not isinstance(b, Complex): 69 | b = Complex(b) 70 | return Complex(foo(a.x, b.x), foo(a.y, b.y)) 71 | 72 | def __ti_repr__(self): 73 | yield '(' 74 | yield self.x 75 | yield ' + ' 76 | yield self.y 77 | yield 'j' 78 | yield ')' 79 | 80 | # TODO: intergrate into __mul__? 81 | def __matmul__(a, b): 82 | '''Multiply two complex numbers''' 83 | # XXX: SSA violation on random? 84 | if not isinstance(b, Complex): 85 | b = Complex(b) 86 | return Complex(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x) 87 | 88 | __rmatmul__ = __matmul__ 89 | 90 | 91 | # TODO: add Affine and Subtensor 92 | -------------------------------------------------------------------------------- /taichi_glsl/colormap.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | 4 | 5 | @ti.lang.kernel_impl.pyfunc 6 | def grayscale(rgb): 7 | ''' 8 | Convert RGB value (vector) into grayscale (scalar). 9 | 10 | :parameter rgb: (3D vector) 11 | The RGB value, X compoment is the R value, can be eitier float or int. 12 | 13 | :return: 14 | The return value is calculated as `dot(vec(0.2989, 0.587, 0.114), rgb)`. 15 | ''' 16 | ret = rgb.x 17 | ret = ts.vec(0.2989, 0.587, 0.114).dot(rgb) 18 | return ret 19 | 20 | 21 | @ti.lang.kernel_impl.pyfunc 22 | def normalmap(n): 23 | ''' 24 | Convert XYZ normal vector into RGB values. 25 | 26 | :parameter n: (3D vector) 27 | The XYZ normal vector, should be normalized to get the desired output. 28 | 29 | :return: 30 | The return value is calculated as `n * 0.5 + 0.5`. 31 | ''' 32 | return n * 0.5 + 0.5 33 | 34 | 35 | @ti.lang.kernel_impl.pyfunc 36 | def blueorange(rgb): 37 | ''' 38 | Convert RGB value (vector) into blue-orange colormap (vector). 39 | 40 | :parameter rgb: (3D vector) 41 | The RGB value, X compoment is the R value, can be eitier float or int. 42 | 43 | :return: 44 | The return value is calculated as `sqrt(mix(vec(0, 0.01, 0.05), vec(1.19, 1.04, 0.98), color))`. 45 | ''' 46 | blue = ts.vec(0.00, 0.01, 0.05) 47 | orange = ts.vec(1.19, 1.04, 0.98) 48 | return ti.sqrt(ts.mix(blue, orange, rgb)) 49 | -------------------------------------------------------------------------------- /taichi_glsl/experimental_array.py: -------------------------------------------------------------------------------- 1 | ''' 2 | (Experimental) Function aliases on Taichi tensor definitions. 3 | ''' 4 | 5 | import taichi as ti 6 | from typing import List, NewType 7 | 8 | 9 | def dtype(dt): 10 | dt_map = { 11 | 'float32': ti.f32, 12 | 'float64': ti.f64, 13 | 'int8': ti.i8, 14 | 'int16': ti.i16, 15 | 'int32': ti.i32, 16 | 'int64': ti.i64, 17 | 'uint8': ti.u8, 18 | 'uint16': ti.u16, 19 | 'uint32': ti.u32, 20 | 'uint64': ti.u64, 21 | 'float': ti.f32, 22 | 'int': ti.i32, 23 | float: ti.f32, 24 | int: ti.i32, 25 | } 26 | if dt in dt_map.keys(): 27 | return dt_map[dt] 28 | else: 29 | return dt 30 | 31 | 32 | DataType = NewType('DataType', (str, type)) 33 | 34 | 35 | def array(dt: DataType, *shape: List[int]): 36 | return ti.field(dtype(dt), shape=shape) 37 | 38 | 39 | # TODO: use array(Vec[n, dt], *shape) instead: 40 | def vec_array(n: int, dt, *shape: List[int]): 41 | return ti.Vector.field(n, dtype(dt), shape=shape) 42 | 43 | 44 | def mat_array(n: int, m: int, dt: DataType, *shape: List[int]): 45 | return ti.Matrix.field(n, m, dtype(dt), shape=shape) 46 | 47 | 48 | def uniform(dt: DataType): 49 | return ti.field(dtype(dt), shape=()) 50 | 51 | 52 | def vec_uniform(n: int, dt: DataType): 53 | return ti.Vector.field(n, dtype(dt), shape=()) 54 | 55 | 56 | def mat_uniform(n: int, m: int, dt: DataType): 57 | return ti.Matrix.field(n, m, dtype(dt), shape=()) 58 | 59 | 60 | def tensor(dt: DataType, shape: (List[int], None) = None): 61 | return ti.field(dtype(dt), shape=shape) 62 | 63 | 64 | def vec_tensor(n: int, dt: DataType, shape: (List[int], None) = None): 65 | return ti.Vector.field(n, dtype(dt), shape=shape) 66 | 67 | 68 | def mat_tensor(n: int, m: int, dt: DataType, shape: (List[int], None) = None): 69 | return ti.Matrix.field(n, m, dtype(dt), shape=shape) 70 | -------------------------------------------------------------------------------- /taichi_glsl/experimental_cfd.py: -------------------------------------------------------------------------------- 1 | ''' 2 | (Experimental) Some ODOP tools that might be useful in fluid dynamics simulation. 3 | ''' 4 | 5 | import taichi as ti 6 | import taichi_glsl as ts 7 | 8 | from taichi_glsl.sampling import * 9 | 10 | 11 | @ti.func 12 | def vgridDivergence(field: ti.template(), I): 13 | return (sample(field, I + D.xy).x + sample(field, I + D.yx).y - 14 | sample(field, I + D.zy).x - sample(field, I + D.yz).y) 15 | 16 | 17 | @ti.func 18 | def vgridGradient(field: ti.template(), I): 19 | return ts.vec2( 20 | sample(field, I + D.xy) - sample(field, I + D.xz), 21 | sample(field, I + D.yx) - sample(field, I + D.zx)) 22 | 23 | 24 | @ti.func 25 | def vgridSumAround(field: ti.template(), I): 26 | return (sample(field, I + D.yx) + sample(field, I + D.zx) + 27 | sample(field, I + D.xy) + sample(field, I + D.xz)) 28 | 29 | 30 | class Pair(TaichiClass): 31 | @classmethod 32 | def make(cls, init): 33 | return cls(init(), init()) 34 | 35 | @property 36 | def old(self): 37 | return self.entries[0] 38 | 39 | @property 40 | def new(self): 41 | return self.entries[1] 42 | 43 | @ti.func 44 | def update(self): 45 | for I in ti.grouped(self.old): 46 | self.old[I] = self.new[I] 47 | 48 | 49 | class SemiLagrangianRK1(Pair): 50 | @ti.func 51 | def advance(self, world): 52 | for I in ti.grouped(self.old): 53 | btI = I - world.velocity(I) * (world.dt / world.dx) 54 | self.new[I] = bilerp(self.old, btI) 55 | 56 | 57 | class SemiLagrangianRK2(Pair): 58 | @ti.func 59 | def advance(self, world): 60 | for I in ti.grouped(self.old): 61 | spI = I - world.velocity(I) * (0.5 * world.dt / world.dx) 62 | btI = I - world.velocity(spI) * (world.dt / world.dx) 63 | self.new[I] = bilerp(self.old, btI) 64 | 65 | 66 | class Maccormack(Pair): 67 | def __init__(self, a, b, c, base=None): 68 | base = base or SemiLagrangianRK2 69 | super(Maccormack, self).__init__(a, b, c) 70 | self.forth, self.back = base(a, b), base(b, c) 71 | 72 | @classmethod 73 | def make(cls, init, base=None): 74 | a, b, c = init(), init(), init() 75 | return cls(a, b, c, base) 76 | 77 | @property 78 | def aux(self): 79 | return self.entries[2] 80 | 81 | @ti.func 82 | def advance(self, world): 83 | self.forth.advance(world) 84 | world.dt = ti.static(-world.dt) 85 | self.back.advance(world) 86 | world.dt = ti.static(-world.dt) 87 | 88 | for I in ti.grouped(self.old): 89 | self.new[I] += 0.5 * (self.old[I] - self.aux[I]) 90 | -------------------------------------------------------------------------------- /taichi_glsl/experimental_transform.py: -------------------------------------------------------------------------------- 1 | ''' 2 | (Experimental) OpenGL-alike view transformations. 3 | ''' 4 | 5 | import taichi as ti 6 | 7 | 8 | def _tuple_to_vector(x): 9 | if len(x) != 1 or not isinstance(x[0], Matrix): 10 | x = ti.Matrix(x) 11 | return x 12 | 13 | 14 | def view(image, *indices): 15 | indices = _tuple_to_vector(indices) 16 | return indices / image.shape() 17 | 18 | 19 | def tex(image, *coors): 20 | coors = _tuple_to_vector(coors) 21 | return indices / image.shape() 22 | -------------------------------------------------------------------------------- /taichi_glsl/gui.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Display images or animations using Taichi GUI (WIP) 3 | ''' 4 | 5 | import taichi as ti 6 | import taichi_glsl as ts 7 | import time 8 | import os 9 | 10 | 11 | class Animation(ts.DataOriented): 12 | ''' 13 | Handy Shadertoy-alike GUI base class. 14 | 15 | I'm able to: 16 | 17 | 1. Enable you to focus on computation, no need to hand-write a GUI event loop. 18 | 19 | 2. Easy-to-use image / video export wrapper like ``self.set_output_video(path)``. 20 | 21 | 3. Shadertoy-alike input variables including ``self.iMouse``, ``self.iKeyDirection``. 22 | 23 | 4. Callback style event processing system incuding ``self.on_click(x, y)``. 24 | 25 | 5. Easily port to Jupyter notebook by simply ``self.gui_backend = 'ipython'``. 26 | 27 | See `examples/export_video.py `_ for example: 28 | 29 | .. code-block:: python 30 | 31 | import taichi as ti 32 | import taichi_glsl as ts 33 | 34 | ti.init() 35 | 36 | 37 | class MyAnimation(ts.Animation): 38 | def on_init(self): 39 | self.img = ti.Vector.field(3, ti.f32, (512, 512)) 40 | self.set_output_video('/tmp/video.gif') 41 | self.define_input() 42 | 43 | @ti.kernel 44 | def on_render(self): 45 | for I in ti.grouped(self.img): 46 | uv = I / self.iResolution 47 | self.img[I] = ti.cos(uv.xyx + self.iTime + 48 | ts.vec(0, 2, 4)) * 0.5 + 0.5 49 | 50 | 51 | MyAnimation().start() 52 | 53 | And what's more, `examples/particles.py `_: 54 | 55 | .. code-block:: python 56 | 57 | import taichi as ti 58 | import taichi_glsl as ts 59 | 60 | ti.init() 61 | 62 | 63 | class MyAnimation(ts.Animation): 64 | def on_init(self): 65 | self.N = 8192 66 | self.dt = 0.01 67 | self.pos = ti.Vector.field(2, ti.f32, self.N) 68 | self.vel = ti.Vector.field(2, ti.f32, self.N) 69 | self.circles = self.pos # alias to make ts.Animation know 70 | self.attract_strength = ti.field(ti.f32, ()) 71 | self.attract_pos = ti.Vector.field(2, ti.f32, ()) 72 | self.resolution = (512, 512) 73 | self.title = 'Particles' 74 | self.define_input() 75 | 76 | @ti.kernel 77 | def on_start(self): 78 | for i in self.pos: 79 | self.pos[i] = ts.randND(2) 80 | self.vel[i] = ts.randSolid2D() 81 | 82 | @ti.kernel 83 | def on_advance(self): 84 | for i in self.pos: 85 | acc = ts.vec(0.0, -1.0) 86 | if any(self.iKeyDirection): # ASWD? 87 | acc = self.iKeyDirection 88 | if any(self.iMouseButton): 89 | dir = ts.normalize(self.iMouse - self.pos[i]) * 2 90 | if self.iMouseButton[0]: # LMB pressed? 91 | acc += dir 92 | if self.iMouseButton[1]: # RMB pressed? 93 | acc -= dir 94 | self.vel[i] += acc * self.dt 95 | for i in self.pos: 96 | self.vel[i] = ts.boundReflect(self.pos[i], self.vel[i], 0, 1, 0.8) 97 | for i in self.pos: 98 | self.pos[i] += self.vel[i] * self.dt 99 | 100 | 101 | MyAnimation().start() 102 | ''' 103 | def __init__(self, 104 | img=None, 105 | circles=None, 106 | title='Animation', 107 | res=(512, 512), 108 | **kwargs): 109 | self.title = title 110 | self.img = img 111 | self.circles = circles 112 | self.circle_color = 0xffffff 113 | self.circle_radius = 1 114 | self.background_color = 0x000000 115 | self.gui = None 116 | self.has_input = False 117 | self.auto_clean = False 118 | self.colormap = None 119 | self.screenshot_dir = None 120 | self.output_video = None 121 | self.gui_backend = os.environ.get('TI_GUI_BACKEND', 'native') 122 | self.start_time = time.time() 123 | self._resolution = res 124 | self.on_init(**kwargs) 125 | 126 | def set_output_video(self, path, framerate=24): 127 | ''' 128 | Export frames painted in GUI to a video. 129 | 130 | FIXME: Only work for ``self.img`` render, doesn't work for ``self.circles`` for now. 131 | Use ``self.screenshot_dir = '/tmp'``, then ``cd /tmp && ti video`` if you wish to. 132 | ''' 133 | output_dir = os.path.dirname(path) 134 | output_file = os.path.basename(path) 135 | try: 136 | output_ext = output_file.split(os.path.extsep)[-1] 137 | assert output_ext in ['gif', 'mp4'] 138 | except: 139 | output_ext = None 140 | self.video_manager = ti.VideoManager(output_dir, 141 | framerate=framerate, 142 | automatic_build=False) 143 | self.video_manager.taichi_glsl_output_ext = output_ext 144 | 145 | def _get_output_filename(suffix): 146 | if output_ext is not None: 147 | return path[:-(len(output_ext) + 1)] + suffix 148 | else: 149 | return path + suffix 150 | 151 | self.video_manager.get_output_filename = _get_output_filename 152 | 153 | def on_init(self): 154 | ''' 155 | Called when initializing ``Animation()``. 156 | 157 | Set up self.* properties for application usage here: 158 | 159 | +--------------------+--------------+-----------------+-------------------------------+ 160 | | Property | Type | Default | Description | 161 | +--------------------+--------------+-----------------+-------------------------------+ 162 | | ``img`` | ``np.array`` | ``None`` | Image to display. | 163 | +--------------------+--------------+-----------------+-------------------------------+ 164 | | ``circles`` | ``np.array`` | ``None`` | Circles to paint. | 165 | +--------------------+--------------+-----------------+-------------------------------+ 166 | | ``circle_radius`` | scalar | ``1`` | Radius of circles. | 167 | +--------------------+--------------+-----------------+-------------------------------+ 168 | | ``circle_color`` | RGB hex | ``0x000000`` | Color of circles. | 169 | +--------------------+--------------+-----------------+-------------------------------+ 170 | |``background_color``| RGB hex | ``0x000000`` | background color of window. | 171 | +--------------------+--------------+-----------------+-------------------------------+ 172 | | ``title`` | string | ``"Animation"`` | Title of the window. | 173 | +--------------------+--------------+-----------------+-------------------------------+ 174 | | ``screenshot_dir`` | string | ``None`` | Path to save screenshots. | 175 | +--------------------+--------------+-----------------+-------------------------------+ 176 | | ``resolution`` | tuple | ``img.shape`` | The size of window / screen. | 177 | +--------------------+--------------+-----------------+-------------------------------+ 178 | | ``auto_clean`` | boolean | ``False`` | Zero the image before render. | 179 | +--------------------+--------------+-----------------+-------------------------------+ 180 | | ``colormap`` | MPL CMap | ``None`` | ``matplotlib.cm`` color map. | 181 | +--------------------+--------------+-----------------+-------------------------------+ 182 | | ``gui_backend`` | string | ``native`` | native, matplotlib, or none | 183 | +--------------------+--------------+-----------------+-------------------------------+ 184 | ''' 185 | pass 186 | 187 | def on_advance(self): 188 | ''' 189 | Called to advance / step the physics scene. 190 | 191 | I.e. update ``self.circles`` if you're using it. 192 | ''' 193 | pass 194 | 195 | def on_pre_render(self): 196 | if self.auto_clean: 197 | self.img.fill(0) 198 | 199 | def on_render(self): 200 | ''' 201 | Called to render the displayed image. 202 | 203 | I.e. update ``self.img`` if it's used. 204 | ''' 205 | pass 206 | 207 | def on_show(self): 208 | pass 209 | 210 | def on_press(self, key): 211 | pass 212 | 213 | def on_pressing(self, key): 214 | pass 215 | 216 | def on_not_pressing(self): 217 | pass 218 | 219 | def on_release(self, key): 220 | pass 221 | 222 | def on_click(self, x, y, btn): 223 | pass 224 | 225 | def on_clicking(self, x, y, btn): 226 | pass 227 | 228 | def on_not_clicking(self, x, y): 229 | pass 230 | 231 | def on_unclick(self, x, y, btn): 232 | pass 233 | 234 | def on_hover(self, x, y): 235 | pass 236 | 237 | def on_drag(self, x, y): 238 | # TODO: add (dx, dy) for drag 239 | pass 240 | 241 | def on_wheel(self, x, y, dx, dy): 242 | pass 243 | 244 | def on_close(self): 245 | self.gui.running = False 246 | 247 | def on_escape(self): 248 | self.on_close() 249 | 250 | @property 251 | def _img(self): 252 | if self.colormap is not None: 253 | return self.colormap(self.img.to_numpy()) 254 | else: 255 | return self.img 256 | 257 | def on_pre_event(self): 258 | ''' 259 | Called per GUI main loop. 260 | ''' 261 | MOUSE = [ti.GUI.LMB, ti.GUI.MMB, ti.GUI.RMB] 262 | had_any = False 263 | for btn in MOUSE: 264 | if self.gui.is_pressed(btn): 265 | self.on_clicking(*self.mouse, btn) 266 | had_any = True 267 | if not had_any: 268 | self.on_not_clicking(*self.mouse) 269 | had_any = False 270 | for key in self.gui.key_pressed: 271 | if key not in MOUSE: 272 | self.on_pressing(key) 273 | had_any = True 274 | if not had_any: 275 | self.on_not_pressing() 276 | 277 | def on_start(self): 278 | ''' 279 | Called when GUI main loop started, i.e. ``Animation().start()``. 280 | ''' 281 | pass 282 | 283 | def on_post_render(self): 284 | pass 285 | 286 | def on_pre_exit(self): 287 | if hasattr(self, 'video_manager'): 288 | ext = self.video_manager.taichi_glsl_output_ext 289 | ti.info('Saving result to {}.{}', 290 | self.video_manager.get_output_filename(''), ext 291 | or 'gif and mp4') 292 | self.video_manager.make_video(gif=(not ext or ext == 'gif'), 293 | mp4=(not ext or ext == 'mp4')) 294 | 295 | def on_exit(self): 296 | pass 297 | 298 | def define_input(self): 299 | ''' 300 | Should be called if you wish to use ``self.iXXX`` as uniform scalars. 301 | 302 | If you are familiar with `Shadertoy `_, then this is for you :) 303 | ''' 304 | self._iTime = ti.field(ti.f32, ()) 305 | self._iFrame = ti.field(ti.i32, ()) 306 | self._iMouse = ti.Vector.field(2, ti.f32, ()) 307 | self._iMouseButton = ti.Vector.field(3, ti.i32, ()) 308 | self._iKeyDirection = ti.Vector.field(2, ti.f32, ()) 309 | self.has_input = True 310 | 311 | def on_update_input(self): 312 | if not self.has_input: 313 | return 314 | self._iTime[None] = self.time 315 | self._iFrame[None] = self.frame 316 | self._iMouse[None] = self.mouse 317 | ip = lambda *x: int(self.gui.is_pressed(*x)) 318 | lmb, mmb, rmb = ip(ti.GUI.LMB), ip(ti.GUI.MMB), ip(ti.GUI.RMB) 319 | self._iMouseButton[None] = [lmb, mmb, rmb] 320 | dx = ip('d', ti.GUI.RIGHT) - ip('a', ti.GUI.LEFT) 321 | dy = ip('w', ti.GUI.UP) - ip('s', ti.GUI.DOWN) 322 | self._iKeyDirection[None] = [dx, dy] 323 | 324 | @property 325 | def iTime(self): 326 | ''' 327 | (TS, float32, RO) Current time in seconds. 328 | ''' 329 | if not self.has_input: 330 | raise Exception( 331 | 'Add ``self.define_input()`` to ``on_init`` if you ' 332 | 'wish to use inputs') 333 | if ti.inside_kernel(): 334 | return ti.subscript(self._iTime, None) 335 | else: 336 | return self._iTime[None] 337 | 338 | @property 339 | def iFrame(self): 340 | ''' 341 | (TS, int32, RO) Current frame number start from 0. 342 | ''' 343 | if not self.has_input: 344 | raise Exception( 345 | 'Add ``self.define_input()`` to ``on_init`` if you ' 346 | 'wish to use inputs') 347 | if ti.inside_kernel(): 348 | return ti.subscript(self._iFrame, None) 349 | else: 350 | return self._iFrame[None] 351 | 352 | @property 353 | def iMouse(self): 354 | ''' 355 | (TS, 2D float32 vector, RO) Current mouse position from 0 to 1. 356 | ''' 357 | if not self.has_input: 358 | raise Exception( 359 | 'Add ``self.define_input()`` to ``on_init`` if you ' 360 | 'wish to use inputs') 361 | if ti.inside_kernel(): 362 | return self._iMouse.subscript(None) 363 | else: 364 | return self._iMouse[None] 365 | 366 | @property 367 | def iMouseButton(self): 368 | ''' 369 | (TS, 3D int32 vector, RO) Current mouse button status. 370 | 371 | ``self.iMouseButton[0]`` is ``1`` if LMB is pressed. 372 | ``self.iMouseButton[1]`` is ``1`` if MMB is pressed. 373 | ``self.iMouseButton[2]`` is ``1`` if RMB is pressed. 374 | Otherwise, ``0``. 375 | ''' 376 | if not self.has_input: 377 | raise Exception( 378 | 'Add ``self.define_input()`` to ``on_init`` if you ' 379 | 'wish to use inputs') 380 | if ti.inside_kernel(): 381 | return self._iMouseButton.subscript(None) 382 | else: 383 | return self._iMouseButton[None] 384 | 385 | @property 386 | def iKeyDirection(self): 387 | ''' 388 | (TS, 2D float32 vector, RO) Direction according to ASWD / arrow keys. 389 | 390 | If A or left arrow is pressed, then ``self.iKeyDirection`` is ``vec(-1.0, 0.0)``. 391 | If D or right arrow is pressed, then ``self.iKeyDirection`` is ``vec(1.0, 0.0)``. 392 | If W or up arrow is pressed, then ``self.iKeyDirection`` is ``vec(0.0, 1.0)``. 393 | If S or down arrow is pressed, then ``self.iKeyDirection`` is ``vec(0.0, -1.0)``. 394 | ''' 395 | if not self.has_input: 396 | raise Exception( 397 | 'Add ``self.define_input()`` to ``on_init`` if you ' 398 | 'wish to use inputs') 399 | if ti.inside_kernel(): 400 | return self._iKeyDirection.subscript(None) 401 | else: 402 | return self._iKeyDirection[None] 403 | 404 | @property 405 | def iResolution(self): 406 | ''' 407 | (TS, 2D int32 vector, RO) Window size / screen resolution. 408 | ''' 409 | return self.resolution 410 | 411 | def on_event(self, e): 412 | ''' 413 | Called when a event occurred, hook me if you want a raw event control. 414 | ''' 415 | MOUSE = [ti.GUI.LMB, ti.GUI.MMB, ti.GUI.RMB] 416 | if e.type == ti.GUI.PRESS: 417 | if e.key in MOUSE: 418 | self.on_click(*e.pos, e.key) 419 | else: 420 | if e.key == ti.GUI.ESCAPE: 421 | self.on_escape() 422 | if e.key == ti.GUI.EXIT: 423 | self.on_close() 424 | self.on_press(e.key) 425 | elif e.type == ti.GUI.RELEASE: 426 | if e.key in MOUSE: 427 | self.on_unclick(*e.pos, e.key) 428 | else: 429 | self.on_release(e.key) 430 | elif e.type == ti.GUI.MOTION: 431 | if e.type == ti.GUI.MOVE: 432 | had_any = False 433 | for btn in MOUSE: 434 | if self.gui.is_pressed(btn): 435 | self.on_drag(*e.pos, btn) 436 | had_any = True 437 | if not had_any: 438 | self.on_hover(*e.pos) 439 | elif e.type == ti.GUI.WHEEL: 440 | self.on_wheel(*e.pos, *e.delta) 441 | 442 | @property 443 | def mouse(self): 444 | ''' 445 | (PS, tuple of two float, RO) Get mouse position / cursor coordinate from 0 to 1. 446 | ''' 447 | return self.gui.get_cursor_pos() 448 | 449 | @property 450 | def resolution(self): 451 | ''' 452 | (PS, tuple of two int, RW) Get or set window size / screen resolution. 453 | ''' 454 | if self.img is not None: 455 | return self.img.shape()[0:2] 456 | else: 457 | return self._resolution 458 | 459 | @resolution.setter 460 | def resolution(self, value): 461 | self._resolution = value 462 | 463 | @property 464 | def time(self): 465 | ''' 466 | (PS, float32, RO) Get current time in seconds. 467 | ''' 468 | return time.time() - self.start_time 469 | 470 | @property 471 | def frame(self): 472 | ''' 473 | (PS, int, RO) Get current frame number start from 0. 474 | ''' 475 | return self.gui.frame 476 | 477 | def _show_mpl_animation(self): 478 | import numpy as np 479 | import matplotlib.pyplot as plt 480 | from matplotlib.animation import FuncAnimation 481 | 482 | fig = plt.figure(figsize=(4, 4)) 483 | fim = plt.figimage( 484 | np.zeros((*self.resolution, 3), dtype=np.float32) + 0.5) 485 | 486 | def update(frame): 487 | img = self._per_loop(do_get_img=True)[:, :, :3] 488 | fim.set_array(img) 489 | return fim, 490 | 491 | anim = FuncAnimation(fig, update, frames=50, interval=1) 492 | plt.show() 493 | 494 | def start(self): 495 | ''' 496 | Call me when GUI is ready to start shows up. 497 | 498 | A common usage for application can be: ``MyAnimation().start()``. 499 | ''' 500 | self.on_start() 501 | with ti.GUI(self.title, 502 | self.resolution, 503 | background_color=self.background_color, 504 | show_gui=(self.gui_backend == 'native')) as self.gui: 505 | if self.gui_backend == 'matplotlib': 506 | self._show_mpl_animation() 507 | else: 508 | while self.gui.running: 509 | self._per_loop() 510 | self.on_pre_exit() 511 | self.on_exit() 512 | self.gui = None 513 | 514 | def _per_loop(self, do_get_img=False): 515 | self.on_pre_event() 516 | for e in self.gui.get_events(): 517 | self.on_event(e) 518 | self.on_update_input() 519 | self.on_advance() 520 | self.on_pre_render() 521 | self.on_render() 522 | self.on_post_render() 523 | if self.img is not None: 524 | self.gui.set_image(self._img) 525 | if self.circles is not None: 526 | self.gui.circles(self.circles.to_numpy(), self.circle_color, 527 | self.circle_radius) 528 | if self.gui_backend == 'ipython': 529 | from IPython.display import clear_output 530 | clear_output(wait=True) 531 | img = self.gui.get_image()[:, :, :3] 532 | ti.imdisplay(img) 533 | if do_get_img: 534 | img = self.gui.get_image() 535 | self.on_show() 536 | if self.screenshot_dir is None: 537 | self.gui.show() 538 | else: 539 | self.gui.show(f'{self.screenshot_dir}/{self.frame:06d}.png') 540 | if hasattr(self, 'video_manager'): 541 | self.video_manager.write_frame(self.img.to_numpy()) 542 | ti.debug('Frame {} recorded', self.gui.frame) 543 | if do_get_img: 544 | return img 545 | else: 546 | return None 547 | -------------------------------------------------------------------------------- /taichi_glsl/hack.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Some hacks / hooks on Taichi to make Taichi GLSL work 3 | ''' 4 | 5 | import taichi as ti 6 | 7 | ## Get rid of the annoying deprecation warnings: 8 | __import__('warnings').filterwarnings('ignore') 9 | 10 | # Get rid of `maybe you want to use a.fill(b)?` limitation. 11 | _old_element_wise_binary = ti.Matrix._element_wise_binary 12 | 13 | 14 | def _new_element_wise_binary(self, foo, other): 15 | if foo.__name__ == 'assign': 16 | foo.__name__ = 'dummy_assign' 17 | return _old_element_wise_binary(self, foo, other) 18 | 19 | 20 | ti.Matrix._element_wise_binary = _new_element_wise_binary 21 | 22 | 23 | # Add ti.Matrix.product method: 24 | @ti.lang.kernel_impl.pyfunc 25 | def _vector_product(self: ti.template()): 26 | ret = self[0] 27 | for i in ti.static(range(1, self.n)): 28 | ret *= self[i] 29 | return ret 30 | 31 | 32 | ti.Matrix.product = _vector_product 33 | 34 | 35 | # Add ti.Matrix.{L,L2,N,T,det,inv,tr} property: 36 | def _vector_L(self): 37 | return self.norm() 38 | 39 | 40 | def _vector_L2(self): 41 | return self.norm_sqr() 42 | 43 | 44 | def _vector_N(self): 45 | return self.normalized() 46 | 47 | 48 | def _matrix_T(self): 49 | return self.transpose() 50 | 51 | 52 | def _matrix_det(self): 53 | return self.determinant() 54 | 55 | 56 | def _matrix_inv(self): 57 | return self.inverse() 58 | 59 | 60 | def _matrix_tr(self): 61 | return self.trace() 62 | 63 | 64 | ti.Matrix.L = property(_vector_L) 65 | ti.Matrix.L2 = property(_vector_L2) 66 | ti.Matrix.N = property(_vector_N) 67 | ti.Matrix.T = property(_matrix_T) 68 | ti.Matrix.det = property(_matrix_det) 69 | ti.Matrix.inv = property(_matrix_inv) 70 | ti.Matrix.tr = property(_matrix_tr) 71 | 72 | # Add ti.pi and ti.tau: 73 | import math 74 | 75 | pi = math.pi 76 | tau = math.tau 77 | ti.pi = pi 78 | ti.tau = tau 79 | -------------------------------------------------------------------------------- /taichi_glsl/lagrangian.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Some helper functions that might be useful in physics simulation. 3 | ''' 4 | 5 | import taichi as ti 6 | import taichi_glsl as tl 7 | 8 | 9 | @ti.func 10 | def boundReflect(pos, vel, pmin=0, pmax=1, gamma=1, gamma_perpendicular=1): 11 | ''' 12 | Reflect particle velocity from a rectangular boundary (if collides). 13 | 14 | `boundaryReflect` takes particle position, velocity and other parameters. 15 | Detect if the particle collides with the rect boundary given by ``pmin`` 16 | and ``pmax``, if collide, returns the velocity after bounced with boundary, 17 | otherwise return the original velocity without any change. 18 | 19 | :parameter pos: (Vector) 20 | The particle position. 21 | 22 | :parameter vel: (Vector) 23 | The particle velocity. 24 | 25 | :parameter pmin: (scalar or Vector) 26 | The position lower boundary. If vector, it's the bottom-left of rect. 27 | 28 | :parameter pmin: (scalar or Vector) 29 | The position upper boundary. If vector, it's the top-right of rect. 30 | ''' 31 | cond = pos < pmin and vel < 0 or pos > pmax and vel > 0 32 | for j in ti.static(range(pos.n)): 33 | if cond[j]: 34 | vel[j] *= -gamma 35 | for k in ti.static(range(pos.n)): 36 | if k != j: 37 | vel[k] *= gamma_perpendicular 38 | return vel 39 | 40 | 41 | @ti.func 42 | def ballBoundReflect(pos, vel, center, radius, anti_fall=0, anti_depth=0.1): 43 | ret = vel 44 | above = tl.distance(pos, center) - radius 45 | if above <= 0: 46 | normal = tl.normalize(pos - center) 47 | NoV = tl.dot(vel, normal) 48 | if ti.static(anti_fall): 49 | NoV -= anti_fall * tl.smoothstep(above, 0, -anti_depth) 50 | if NoV < 0: 51 | ret -= NoV * normal 52 | return ret 53 | 54 | 55 | @ti.func 56 | def momentumExchange(v1, v2, disp, m1=1, m2=1, gamma=1): 57 | ''' 58 | Exchange momentum (bounce) between two objects. 59 | 60 | `momentumExchange` should be invocated when a bounce occurred. 61 | It takes the velocity of two objects before bounce, and returns the 62 | velocity of two objects after bounce. 63 | 64 | This function is most useful in rigid-body simulation with collision. 65 | For example:: 66 | 67 | if distance(pos[i], pos[j]) < radius[i] + radius[j]: 68 | # Collision detected! Perform a momentum exchange: 69 | vel[i], vel[j] = momentumExchange( 70 | vel[i], vel[j], mass[i], mass[j], pos[i] - pos[j], 0.8) 71 | 72 | :parameter v1: (Vector) 73 | The velocity vector of the first object to bounce. 74 | Or, the velocity vector at the collision point in first object. 75 | 76 | :parameter v2: (Vector) 77 | The velocity vector of the second object to bounce. 78 | Or, the velocity vector at the collision point in second object. 79 | 80 | :parameter disp: (Vector) 81 | The displacement vector from between two object. 82 | Or, the normal vector of collision surface. 83 | Specifically, for balls or circles, `disp` is `pos1 - pos2`. 84 | 85 | :parameter m1: (scalar) 86 | The mass of the first object to bounce. 87 | 88 | :parameter m2: (scalar) 89 | The mass of the second object to bounce. 90 | 91 | :parameter gamma: (scalar) 92 | The decrease factor of bounce, in range [0, 1], determines how 93 | much energy is conserved after the bounce process. If 1, then 94 | no energy is loss; if 0, then the collided objects will stops 95 | immediately. 96 | 97 | :return: (tuple of Vector) 98 | The return value is a tuple of velocity of two objects after bounce. 99 | Specifically the first element is for the velocity of first object 100 | (previously to be `v1`), and same to the second element. 101 | 102 | :note: 103 | For usage example, check out this: 104 | https://github.com/taichi-dev/taichi_three/blob/master/examples/many_balls.py 105 | ''' 106 | vel1 = v1.dot(disp) 107 | vel2 = v2.dot(disp) 108 | 109 | sm1 = ti.sqrt(m1) 110 | sm2 = ti.sqrt(m2) 111 | itsm = 1 / ti.sqrt(m1 + m2) 112 | 113 | kero1 = vel1 * sm1 114 | kero2 = vel2 * sm2 115 | 116 | smd1 = sm2 * itsm 117 | smd2 = -sm1 * itsm 118 | 119 | kos = 2 * (kero1 * smd1 + kero2 * smd2) 120 | kero1 -= kos * smd1 121 | kero2 -= kos * smd2 122 | 123 | vel1 = kero1 / sm1 124 | vel2 = kero2 / sm2 125 | 126 | disp *= gamma 127 | 128 | v1 -= v1.dot(disp) * disp 129 | v2 -= v2.dot(disp) * disp 130 | 131 | v1 += vel1 * disp 132 | v2 += vel2 * disp 133 | 134 | return v1, v2 135 | -------------------------------------------------------------------------------- /taichi_glsl/mkimages.py: -------------------------------------------------------------------------------- 1 | ''' 2 | (Experimental) Some stuffs for fun! 3 | ''' 4 | 5 | import taichi as ti 6 | import taichi_glsl as ts 7 | 8 | 9 | @ti.func 10 | def _inside(p, c, r): 11 | return (p - c).norm_sqr() <= r * r 12 | 13 | 14 | @ti.func 15 | def imageTaichi(p, size=1): 16 | p = ti.Vector([0.5, 0.5]) + (p - ti.Vector([0.5, 0.5])) * 1.11 / size 17 | ret = -1 18 | if not _inside(p, ti.Vector([0.50, 0.50]), 0.55): 19 | if ret == -1: 20 | ret = 0 21 | if not _inside(p, ti.Vector([0.50, 0.50]), 0.50): 22 | if ret == -1: 23 | ret = 1 24 | if _inside(p, ti.Vector([0.50, 0.25]), 0.09): 25 | if ret == -1: 26 | ret = 1 27 | if _inside(p, ti.Vector([0.50, 0.75]), 0.09): 28 | if ret == -1: 29 | ret = 0 30 | if _inside(p, ti.Vector([0.50, 0.25]), 0.25): 31 | if ret == -1: 32 | ret = 0 33 | if _inside(p, ti.Vector([0.50, 0.75]), 0.25): 34 | if ret == -1: 35 | ret = 1 36 | if p[0] < 0.5: 37 | if ret == -1: 38 | ret = 1 39 | else: 40 | if ret == -1: 41 | ret = 0 42 | return 1 - ret 43 | 44 | 45 | @ti.func 46 | def imageChess(p, n=20): 47 | return (p // (1 / n)).sum() % 2 48 | 49 | 50 | @ti.func 51 | def imageGrid(p, n=20, tol=0.2): 52 | return ts.smoothstep(abs(ts.fract(p * n) - 0.5).min(), 0.0, tol) 53 | -------------------------------------------------------------------------------- /taichi_glsl/odop.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Enhanced Taichi Objective Data-Oriented Programming (ODOP) 3 | ''' 4 | 5 | import taichi as ti 6 | 7 | 8 | @ti.data_oriented 9 | class DataOriented: 10 | pass 11 | 12 | 13 | class TaichiClass: 14 | is_taichi_class = True 15 | 16 | def __init__(self, *entries): 17 | self.entries = entries 18 | 19 | @classmethod 20 | @ti.lang.util.python_scope 21 | def field(cls, *args, **kwargs): 22 | field_list = cls._field(*args, **kwargs) 23 | if not isinstance(field_list, (list, tuple)): 24 | field_list = [field_list] 25 | return cls(*field_list) 26 | 27 | @classmethod 28 | def _field(cls): 29 | raise NotImplementedError 30 | 31 | def _subscript(self, *indices): 32 | args = [ti.subscript(e, *indices) for e in self.entries] 33 | return self.__class__(*args) 34 | 35 | @ti.lang.util.taichi_scope 36 | def subscript(self, *indices): 37 | return self._subscript(*indices) 38 | 39 | def loop_range(self): 40 | return self.entries[0].loop_range() 41 | 42 | def get_field_members(self): 43 | if hasattr(e, 'get_field_members'): 44 | e = e.get_field_members() 45 | else: 46 | e = [e] 47 | ret += e 48 | 49 | @ti.lang.util.taichi_scope 50 | def variable(self): 51 | return self.__class__(*(ti.expr_init(e) for e in self.entries)) 52 | 53 | @property 54 | def snode(self): 55 | return self.loop_range().snode 56 | 57 | @property 58 | def shape(self): 59 | return self.snode().shape 60 | 61 | @ti.lang.util.taichi_scope 62 | def __ti_repr__(self): 63 | raise NotImplementedError 64 | 65 | def __repr__(self): 66 | ret = [] 67 | for e in self.__ti_repr__(): 68 | ret.append(e) 69 | return ''.join(map(str, ret)) 70 | -------------------------------------------------------------------------------- /taichi_glsl/painting.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as ts 3 | import math 4 | 5 | 6 | @ti.func 7 | def sdLine(u, v, p): 8 | pu = p - u 9 | vp = v - p 10 | vu = v - u 11 | puXvu = pu.cross(vu) 12 | puOvu = pu.dot(vu) 13 | vpOvu = vp.dot(vu) 14 | ret = 0.0 15 | if puOvu < 0: 16 | ret = ts.length(pu) 17 | elif vpOvu < 0: 18 | ret = ts.length(vp) 19 | else: 20 | ret = puXvu * ts.invLength(vu) 21 | return ret 22 | 23 | 24 | @ti.func 25 | def paintArrow(img: ti.template(), 26 | orig, 27 | dir, 28 | color=1, 29 | width=3, 30 | max_size=12, 31 | min_scale=0.4): 32 | res = ts.vec(*img.shape) 33 | I = orig * res 34 | D = dir * res 35 | J = I + D 36 | DL = ts.length(D) 37 | S = min(max_size, DL * min_scale) 38 | DS = D / (DL + 1e-4) * S 39 | SW = S + width 40 | D1 = ti.Matrix.rotation2d(+math.pi * 3 / 4) @ DS 41 | D2 = ti.Matrix.rotation2d(-math.pi * 3 / 4) @ DS 42 | bmin, bmax = ti.floor(max(0, 43 | min(I, J) - SW)), ti.ceil( 44 | min(res - 1, 45 | max(I, J) + SW)) 46 | for P in ti.grouped(ti.ndrange((bmin.x, bmax.x), (bmin.y, bmax.y))): 47 | c0 = ts.smoothstep(abs(sdLine(I, J, P)), width, width / 2) 48 | c1 = ts.smoothstep(abs(sdLine(J, J + D1, P)), width, width / 2) 49 | c2 = ts.smoothstep(abs(sdLine(J, J + D2, P)), width, width / 2) 50 | ti.atomic_max(img[P], max(c0, c1, c2) * color) 51 | -------------------------------------------------------------------------------- /taichi_glsl/randgen.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Pseudo-random number or noise generators. 3 | ''' 4 | 5 | import taichi as ti 6 | import math 7 | 8 | 9 | def rand(): 10 | ''' 11 | Generate a random floating number distributed evenly in range [0, 1]. 12 | 13 | :return: 14 | The return value a pseudo-random floating number in range [0, 1]. 15 | 16 | .. note:: 17 | 18 | Depending on Taichi backend design, the generated random number 19 | may have the **same seed** on start up. And Taichi doesn't provide 20 | any way to set a random seed yet, so does Taichi GLSL. 21 | ''' 22 | 23 | return ti.random() 24 | 25 | 26 | @ti.func 27 | def randInt(a, b): 28 | ''' 29 | Generate a random integer in range [a, b], including both end points. 30 | 31 | :parameter a: 32 | Specify the start point of range. 33 | :parameter b: 34 | Specify the end point of range. 35 | 36 | :return: 37 | The return value is calculated as `ti.random(ti.i32) % (b - a) + a`. 38 | 39 | :note: 40 | This means randInt could not fullfill its promise of "distributed 41 | evenly" when (b - a) is large (at the scalr of INT_MAX). 42 | ''' 43 | 44 | return abs(ti.random(ti.i32)) % (b - a) + a 45 | 46 | 47 | def randND(n): 48 | ''' 49 | Generate a n-D random vector in a n-D cube ([0, 1]). 50 | The return value is a n-D vector, each comonment is a random floating 51 | number in range [0, 1] generated independently. 52 | 53 | :parameter n: 54 | Specify the dimension of random vector to generate. 55 | 56 | :return: 57 | The return value is calculated as `vec(rand() for _ in range(n))`. 58 | ''' 59 | 60 | return ti.Vector([rand() for _ in range(n)]) 61 | 62 | 63 | @ti.func 64 | def randRange(a=0, b=1): 65 | ''' 66 | Generate random floating numbers in range [a, b]. 67 | If `a` and `b` are n-D vectors, then the return value will also be a 68 | n-D vector, whose i-th component is a random number distributed evenly 69 | in range [a[i], b[i]], each comonment is generated independently. 70 | 71 | :parameter a: 72 | Specify the start point of range. 73 | :parameter b: 74 | Specify the end point of range. 75 | 76 | :return: 77 | The return value is calculated as `a + rand() * (b - a)`. 78 | ''' 79 | return a + rand() * (b - a) 80 | 81 | 82 | @ti.func 83 | def randNDRange(a, b): 84 | ''' 85 | Generate a n-D random vector in a n-D cube. 86 | `a` and `b` should be n-D vectors. The return value is also a n-D 87 | vector, whose i-th component is a random number distributed evenly 88 | in range [a[i], b[i]], each componment is generated independently. 89 | 90 | :parameter a: 91 | Specify the staring point / coordinate of cube. 92 | :parameter b: 93 | Specify the ending point / coordinate of cube. 94 | 95 | :return: 96 | The return value is calculated as `a + randND((b - a).n) * (b - a)`. 97 | ''' 98 | c = b - a 99 | return a + randND(c.n) * c 100 | 101 | 102 | @ti.func 103 | def randUnit2D(): 104 | ''' 105 | Generate a 2-D random unit vector whose length is equal to 1.0. 106 | The return value is a 2-D vector, whose tip distributed evenly 107 | **on the border** of a unit circle. 108 | 109 | :return: 110 | The return value is computed as:: 111 | 112 | a = rand() * math.tau 113 | return vec(ti.cos(a), ti.sin(a)) 114 | ''' 115 | a = rand() * math.tau 116 | return ti.Vector([ti.cos(a), ti.sin(a)]) 117 | 118 | 119 | @ti.func 120 | def randSolid2D(): 121 | ''' 122 | Generate a 2-D random unit vector whose length is <= 1.0. 123 | The return value is a 2-D vector, whose tip distributed evenly 124 | **inside** a unit circle. 125 | 126 | :return: 127 | The return value is computed as:: 128 | 129 | a = rand() * math.tau 130 | r = sqrt(rand()) 131 | return vec(cos(a), sin(a)) * r 132 | ''' 133 | a = rand() * math.tau 134 | r = ti.sqrt(rand()) 135 | return ti.Vector([ti.cos(a), ti.sin(a)]) * r 136 | 137 | 138 | @ti.func 139 | def randUnit3D(): 140 | ''' 141 | Generate a 3-D random unit vector whose length is equal to 1.0. 142 | The return value is a 3-D vector, whose tip distributed evenly 143 | **on the surface** of a unit sphere. 144 | 145 | :return: 146 | The return value is computed as:: 147 | 148 | u = randUnit2D() 149 | s = rand() * 2 - 1 150 | c = sqrt(1 - s ** 2) 151 | return vec3(c * u, s) 152 | ''' 153 | from .vector import vec3 154 | u = randUnit2D() 155 | s = rand() * 2 - 1 156 | c = ti.sqrt(1 - s**2) 157 | return vec3(c * u, s) 158 | -------------------------------------------------------------------------------- /taichi_glsl/sampling.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Some helper functions in processing fields & images. 3 | ''' 4 | 5 | import taichi as ti 6 | import taichi_glsl as ts 7 | 8 | from .odop import TaichiClass 9 | 10 | D = ts.vec(1, 0, -1) 11 | 12 | 13 | @ti.func 14 | def sample(field: ti.template(), P): 15 | ''' 16 | Sampling a field with indices clampped into the field shape. 17 | 18 | :parameter field: (Tensor) 19 | Specify the field to sample. 20 | 21 | :parameter P: (Vector) 22 | Specify the index in field. 23 | 24 | :return: 25 | The return value is calcuated as:: 26 | 27 | P = clamp(P, 0, vec(*field.shape) - 1) 28 | return field[int(P)] 29 | ''' 30 | shape = ti.Vector(field.shape) 31 | P = ts.clamp(P, 0, shape - 1) 32 | return field[int(P)] 33 | 34 | 35 | @ti.func 36 | def dflSample(field: ti.template(), P, dfl): 37 | ''' 38 | Sampling a field, when indices out of the field shape, return the given default value. 39 | 40 | :parameter field: (Tensor) 41 | Specify the field to sample. 42 | 43 | :parameter P: (Vector) 44 | Specify the index in field. 45 | 46 | :parameter dfl: (with the same type of field) 47 | Specify the index in field. 48 | 49 | :return: 50 | The return value is calcuated as:: 51 | 52 | return field[int(P)] if 0 <= P < vec(*field.shape) else dfl 53 | ''' 54 | shape = ti.Vector(field.shape) 55 | inside = 0 <= P < shape 56 | return field[int(P)] if inside else dfl 57 | 58 | 59 | @ti.func 60 | def bilerp(field: ti.template(), P): 61 | ''' 62 | Bilinear sampling an 2D field with a real index. 63 | 64 | :parameter field: (2D Tensor) 65 | Specify the field to sample. 66 | 67 | :parameter P: (2D Vector of float) 68 | Specify the index in field. 69 | 70 | :note: 71 | If one of the element to be accessed is out of `field.shape`, then 72 | `bilerp` will automatically do a clamp for you, see :func:`sample`. 73 | 74 | :return: 75 | The return value is calcuated as:: 76 | 77 | I = int(P) 78 | x = fract(P) 79 | y = 1 - x 80 | return (sample(field, I + D.xx) * x.x * x.y + 81 | sample(field, I + D.xy) * x.x * y.y + 82 | sample(field, I + D.yy) * y.x * y.y + 83 | sample(field, I + D.yx) * y.x * x.y) 84 | 85 | .. where D = vec(1, 0, -1) 86 | ''' 87 | I = int(P) 88 | x = ts.fract(P) 89 | y = 1 - x 90 | return (sample(field, I + D.xx) * x.x * x.y + 91 | sample(field, I + D.xy) * x.x * y.y + 92 | sample(field, I + D.yy) * y.x * y.y + 93 | sample(field, I + D.yx) * y.x * x.y) 94 | 95 | 96 | @ti.func 97 | def trilerp(field: ti.template(), P): 98 | ''' 99 | Tilinear sampling an 3D field with a real index. 100 | 101 | :parameter field: (3D Tensor) 102 | Specify the field to sample. 103 | 104 | :parameter P: (3D Vector of float) 105 | Specify the index in field. 106 | 107 | :note: 108 | If one of the element to be accessed is out of `field.shape`, then 109 | `Tilerp` will automatically do a clamp for you, see :func:`sample`. 110 | 111 | Syntax ref : https://en.wikipedia.org/wiki/Trilinear_interpolation. 112 | 113 | :return: 114 | The return value is calcuated as:: 115 | 116 | I = int(P) 117 | w0 = ts.fract(P) 118 | w1 = 1.0 - w0 119 | c00 = ts.sample(field, I + ts.D.yyy) * w1.x + ts.sample(field, I + ts.D.xyy) * w0.x 120 | c01 = ts.sample(field, I + ts.D.yyx) * w1.x + ts.sample(field, I + ts.D.xyx) * w0.x 121 | c10 = ts.sample(field, I + ts.D.yxy) * w1.x + ts.sample(field, I + ts.D.xxy) * w0.x 122 | c11 = ts.sample(field, I + ts.D.yxx) * w1.x + ts.sample(field, I + ts.D.xxx) * w0.x 123 | 124 | c0 = c00 * w1.y + c10 * w0.y 125 | c1 = c01 * w1.y + c11 * w0.y 126 | 127 | return c0 * w1.z + c1 * w0.z 128 | 129 | .. where D = vec(1, 0, -1) 130 | ''' 131 | I = int(P) 132 | w0 = ts.fract(P) 133 | w1 = 1.0 - w0 134 | 135 | c00 = ts.sample(field, I + ts.D.yyy) * w1.x + ts.sample( 136 | field, I + ts.D.xyy) * w0.x 137 | c01 = ts.sample(field, I + ts.D.yyx) * w1.x + ts.sample( 138 | field, I + ts.D.xyx) * w0.x 139 | c10 = ts.sample(field, I + ts.D.yxy) * w1.x + ts.sample( 140 | field, I + ts.D.xxy) * w0.x 141 | c11 = ts.sample(field, I + ts.D.yxx) * w1.x + ts.sample( 142 | field, I + ts.D.xxx) * w0.x 143 | 144 | c0 = c00 * w1.y + c10 * w0.y 145 | c1 = c01 * w1.y + c11 * w0.y 146 | 147 | return c0 * w1.z + c1 * w0.z 148 | 149 | 150 | @ti.func 151 | def superSample2x2(fieldFunc: ti.template(), P, dx=1): 152 | dD = dx / 2 * D 153 | return (fieldFunc(P + dD.yy) + fieldFunc(P + dD.yz) + 154 | fieldFunc(P + dD.zz) + fieldFunc(P + dD.zy)) * 0.25 155 | -------------------------------------------------------------------------------- /taichi_glsl/scalar.py: -------------------------------------------------------------------------------- 1 | ''' 2 | GLSL-alike scalar arithmetic functions. 3 | ''' 4 | 5 | import taichi as ti 6 | from taichi import sin, cos, tan, asin, acos, floor, ceil, sqrt, exp, log 7 | 8 | 9 | @ti.lang.kernel_impl.pyfunc 10 | def clamp(x, xmin=0, xmax=1): 11 | ''' 12 | Constrain a value to lie between two further values. 13 | `clamp` returns the value of x constrained to the range xmin to xmax. 14 | 15 | :parameter x: 16 | Specify the value to constrain. 17 | :parameter xmin: 18 | Specify the lower end of the range into which to constrain x. 19 | :parameter xmax: 20 | Specify the upper end of the range into which to constrain x. 21 | 22 | :return: 23 | The returned value is computed as `min(xmax, max(xmin, x))`. 24 | ''' 25 | return min(xmax, max(xmin, x)) 26 | 27 | 28 | @ti.lang.kernel_impl.pyfunc 29 | def mix(x, y, a): 30 | ''' 31 | Linearly interpolate between two values. 32 | 33 | `mix` performs a linear interpolation between x and y using a to 34 | weight between them. 35 | 36 | :parameter x: 37 | Specify the start of the range in which to interpolate. 38 | :parameter y: 39 | Specify the end of the range in which to interpolate. 40 | :parameter a: 41 | Specify the value to use to interpolate between x and y. 42 | 43 | :return: 44 | The return value is computed as `x * (1 - a) + y * a`. 45 | ''' 46 | return x * (1 - a) + y * a 47 | 48 | 49 | @ti.lang.kernel_impl.pyfunc 50 | def sign(x, edge=0): 51 | ''' 52 | Extract the sign of the parameter. 53 | 54 | `sign` returns -1.0 if x is less than edge, 0.0 if x is 55 | equal to edge, and +1.0 if x is greater than edge. 56 | 57 | :parameter x: 58 | Specify the value from which to extract the sign. 59 | :parameter edge: 60 | Specify a custom location of the edge instead of 0. 61 | 62 | :note: 63 | `sign(x, edge)` is equivalent with `sign(x - edge)`. 64 | Currently Taichi use -1 to represent True. Consequently we add truth here. 65 | 66 | :return: 67 | The return value is computed as `(x >= edge) - (x <= edge)`, 68 | with type promoted. 69 | ''' 70 | ret = x + edge # type promotion 71 | ret = truth(x >= edge) - truth(x <= edge) 72 | return ret 73 | 74 | 75 | @ti.lang.kernel_impl.pyfunc 76 | def step(edge, x): 77 | ''' 78 | Generate a step function by comparing two values. 79 | 80 | `step` generates a step function by comparing x to edge. 81 | For element i of the return value, 0.0 is returned if x[i] < edge[i], and 1.0 is returned otherwise. 82 | 83 | :parameter edge: 84 | Specify the location of the edge of the step function. 85 | :parameter x: 86 | Specify the value to be used to generate the step function. 87 | 88 | :return: 89 | The return value is computed as `x >= edge`, with type 90 | promoted. 91 | ''' 92 | ret = x + edge # type promotion 93 | ret = (x >= edge) 94 | return ret 95 | 96 | 97 | @ti.lang.kernel_impl.pyfunc 98 | def atan(y, x=1): 99 | ''' 100 | Return the arc-tangent of the parameters 101 | 102 | `atan(y, x)` or `atan(y_over_x)`: 103 | 104 | `atan` returns the angle whose trigonometric arctangent is y / x or 105 | y_over_x, depending on which overload is invoked. 106 | 107 | In the first overload, the signs of y and x are used to determine 108 | the quadrant that the angle lies in. The values returned by atan in 109 | this case are in the range [−pi, pi]. Results are undefined if x 110 | is zero. 111 | 112 | For the second overload, atan returns the angle whose tangent 113 | is y_over_x. Values returned in this case are in the range 114 | [−pi/2, pi/2]. 115 | 116 | :parameter y: 117 | The numerator of the fraction whose arctangent to return. 118 | 119 | :parameter x: 120 | The denominator of the fraction whose arctangent to return. 121 | 122 | :return: 123 | The return value is `arctan(x / y)`. 124 | ''' 125 | return ti.atan2(y, x) 126 | 127 | 128 | @ti.lang.kernel_impl.pyfunc 129 | def fract(x): 130 | ''' 131 | Compute the fractional part of the argument. 132 | 133 | `fract` returns the fractional part of x. 134 | 135 | :parameter x: 136 | Specify the value to evaluate. 137 | 138 | :return: 139 | The return value is calculated as `x - floor(x)`. 140 | ''' 141 | return x - floor(float(x)) 142 | 143 | 144 | @ti.lang.kernel_impl.pyfunc 145 | def round(x): 146 | ''' 147 | Find the nearest integer less than or equal to the parameter. 148 | 149 | `round` returns a value equal to the nearest integer to x. 150 | The fraction 0.5 will round toward the larger integer. i.e. 151 | `round(0.5) = 1.0`. 152 | 153 | :parameter x: 154 | Specify the value to evaluate. 155 | 156 | :return: 157 | The return value is calculated as `floor(x + 0.5)`. 158 | ''' 159 | return floor(x + 0.5) 160 | 161 | 162 | @ti.lang.kernel_impl.pyfunc 163 | def smoothstep(x, a=0, b=1): 164 | ''' 165 | Perform Hermite interpolation between two values. 166 | 167 | `smoothstep` performs smooth Hermite interpolation between 0 and 1 168 | when a < x < b. This is useful in cases where a threshold function 169 | with a smooth transition is desired. 170 | 171 | Results are undefined if a >= b. 172 | 173 | :parameter a: 174 | Specifies the value of the lower edge of the Hermite function. 175 | :parameter b: 176 | Specifies the value of the upper edge of the Hermite function. 177 | :parameter x: 178 | Specifies the source value for interpolation. 179 | 180 | :return: 181 | The return value is is computed as:: 182 | 183 | t = clamp((x - a) / (b - a), 0, 1) 184 | return t * t * (3 - 2 * t) 185 | ''' 186 | t = clamp((x - a) / (b - a)) 187 | return t * t * (3 - 2 * t) 188 | 189 | 190 | @ti.lang.kernel_impl.pyfunc 191 | def inversesqrt(x): 192 | ''' 193 | Return the inverse of the square root of the parameter. 194 | 195 | `inversesqrt` returns the inverse of the square root of x; i.e. 196 | the value `1 / sqrt(x)`. The result is undefined if x <= 0. 197 | 198 | :parameter x: 199 | Specify the value of which to take the inverse of the square root. 200 | 201 | :return: 202 | The return value can be calculated as `1 / sqrt(x)` or `pow(x, -0.5)`. 203 | ''' 204 | return 1 / ti.sqrt(x) 205 | 206 | 207 | @ti.func 208 | def isnan(x): 209 | ''' 210 | Determine whether the parameter is a number. 211 | 212 | For each element element i of the result, `isnan` returns True 213 | if x[i] is posititve or negative floating point NaN (Not a Number) 214 | and False otherwise. 215 | 216 | :parameter x: 217 | Specifies the value to test for NaN. 218 | 219 | :return: 220 | The return value is computed as `not (x >= 0 or x <= 0)`. 221 | ''' 222 | return not (x >= 0 or x <= 0) 223 | 224 | 225 | @ti.func 226 | def isinf(x): 227 | ''' 228 | Determine whether the parameter is positive or negative infinity. 229 | 230 | For each element element i of the result, `isinf` returns True 231 | if x[i] is posititve or negative floating point infinity and 232 | False otherwise. 233 | 234 | :parameter x: 235 | Specifies the value to test for infinity. 236 | 237 | :return: 238 | The return value is computed as `2 * x == x and x != 0`. 239 | ''' 240 | return 2 * x == x and x != 0 241 | 242 | 243 | @ti.lang.kernel_impl.pyfunc 244 | def truth(cond): 245 | ''' 246 | return 1 if condition is not false, return 0 vice versa 247 | :param x: 248 | 249 | :return: 250 | The return value is computed as `return ti.select(cond, 1, 0)` 251 | note: Currently Taichi use -1 to represent True. 252 | This function serves as a helper function to those who want 1 to represent True 253 | 254 | ''' 255 | return ti.select(cond, 1, 0) 256 | -------------------------------------------------------------------------------- /taichi_glsl/sphcore.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Some often-used SPH smoother kernels. 3 | ''' 4 | 5 | import taichi as ti 6 | 7 | 8 | @ti.func 9 | def spiky(r, h): 10 | ''' 11 | (h - r)**3 12 | ''' 13 | return (h - r)**3 14 | 15 | 16 | @ti.func 17 | def poly6(r, h): 18 | ''' 19 | (h**2 - r**2)**3 20 | ''' 21 | return (h**2 - r**2)**3 22 | 23 | 24 | @ti.func 25 | def dspiky(r, h): 26 | ''' 27 | -3 * (h - r)**2 28 | ''' 29 | return -3 * (h - r)**2 30 | 31 | 32 | @ti.func 33 | def dpoly6(r, h): 34 | ''' 35 | -6 * r * (h**2 - r**2)**2 36 | ''' 37 | return -6 * r * (h**2 - r**2)**2 38 | -------------------------------------------------------------------------------- /taichi_glsl/vector.py: -------------------------------------------------------------------------------- 1 | ''' 2 | GLSL-alike linear algebra helper functions / aliases. 3 | ''' 4 | 5 | import taichi as ti 6 | 7 | 8 | def vecSimple(*xs): 9 | ''' 10 | Create a n-D vector whose components are specified by arguments. 11 | 12 | :parameter xs: 13 | Specify the scalar values to initialize the vector components. 14 | 15 | :return: 16 | The return value is a n-D vector, where `n = len(xs)`. 17 | And the i-th component is initialized to be `xs[i]`. 18 | 19 | :see also: 20 | :func:`vec`, :func:`vecFill`. 21 | ''' 22 | return ti.Vector(xs) 23 | 24 | 25 | def vec(*xs): 26 | ''' 27 | Create a vector by scalars or vectors in arguments (GLSL-alike). 28 | 29 | The return vector dimension depends on **the count of scalars in xs**. 30 | 31 | :parameter xs: 32 | Specify the scalar or vector values to initialize the 33 | vector components. If xs[i] is vector, it will be **flatten** 34 | into a series of scalars. 35 | 36 | :return: 37 | A n-D vector initialized as described above. 38 | 39 | :see also: 40 | :func:`vecND`, :func:`vecSimple`, :func:`vecFill`. 41 | ''' 42 | ys = [] 43 | for x in xs: 44 | if isinstance(x, ti.Matrix): 45 | ys.extend(x.entries) 46 | else: 47 | ys.append(x) 48 | return vecSimple(*ys) 49 | 50 | 51 | def vecND(n, *xs): 52 | ''' 53 | Create a n-D vector by scalars or vectors in arguments (GLSL-alike). 54 | 55 | The return vector dimension depends on **the specified argument n**. 56 | If the dimension mismatch the count of scalars in xs, an error 57 | will be raised. 58 | However, if only one scalar is specified, then `vecND(n, x)` is 59 | equivalent to `vecFill(n, x)`, see :func:`vecFill`. 60 | 61 | :parameter n: 62 | Specify the dimension of return vector. 63 | :parameter xs: 64 | Specify the scalar or vector values to initialize the 65 | vector components. If xs[i] is vector, it will be **flatten** 66 | into a series of scalars. 67 | 68 | :return: 69 | A n-D vector initialized as described above. 70 | 71 | :note: 72 | `vecND(n, x)` -> `vecFill(n, x)` 73 | `vecND(n, x, y, ...)` -> `vec(x, y, ...)` 74 | 75 | :see also: 76 | :func:`vec`, :func:`vec2`, :func:`vec3`. 77 | ''' 78 | # This's not inside of `@ti.func`, so it's safe to multi-return. 79 | # The `if` statement here is branched at compile-time. 80 | if len(xs) == 1 and not isinstance(xs[0], ti.Matrix): 81 | return vecFill(n, xs[0]) 82 | 83 | ys = [] 84 | for x in xs: 85 | if isinstance(x, ti.Matrix): 86 | ys.extend(x.entries) 87 | else: 88 | ys.append(x) 89 | 90 | # if len(ys) > n: 91 | # ys = ys[:n] 92 | # else: 93 | # ys.extend((n - len(ys)) * [0]) 94 | if len(ys) != n: 95 | raise ValueError(f'Cannot generate {n}-D vector from ' 96 | f'{len(ys)} scalars') 97 | 98 | return vecSimple(*ys) 99 | 100 | 101 | def vecFill(n, x): 102 | ''' 103 | Create a n-D vector whose all components are initialized by x. 104 | 105 | :parameter x: 106 | Specify the scalar value to fill the vector. 107 | 108 | :return: 109 | The return value is calculated as `vec(x for _ in range(n))`. 110 | 111 | :see also: 112 | :func:`vecND`, :func:`vec`, :func:`vec3`. 113 | ''' 114 | return ti.Vector([x for _ in range(n)]) 115 | 116 | 117 | def vec2(*xs): 118 | ''' 119 | An alias for `vecND(2, *xs)`. 120 | 121 | :see also: 122 | :func:`vecND`, :func:`vec`, :func:`vec3`. 123 | ''' 124 | return vecND(2, *xs) 125 | 126 | 127 | def vec3(*xs): 128 | ''' 129 | An alias for `vecND(3, *xs)`. 130 | 131 | :see also: 132 | :func:`vecND`, :func:`vec`, :func:`vec2`. 133 | ''' 134 | return vecND(3, *xs) 135 | 136 | 137 | def vec4(*xs): 138 | ''' 139 | An alias for `vecND(4, *xs)`. 140 | 141 | :see also: 142 | :func:`vecND`, :func:`vec`, :func:`vec3`. 143 | ''' 144 | return vecND(4, *xs) 145 | 146 | 147 | def vecAngle(a): 148 | ''' 149 | Return a 2D vector of specific phase angle. 150 | 151 | :parameter a: 152 | Specify the phase angle of vector. 153 | 154 | :return: 155 | The return value is computed as `vec(cos(a), sin(a))`. 156 | ''' 157 | return vecSimple(ti.cos(a), ti.sin(a)) 158 | 159 | 160 | def mat(*xs): 161 | ''' 162 | Matrix initializer (WIP). 163 | 164 | :parameter xs: 165 | A row-major list of list, which contains the elements of matrix. 166 | 167 | :return: 168 | A matrix, size depends on the input xs. 169 | 170 | For example:: 171 | mat([1, 2], [3, 4]) 172 | 173 | (TODO: WIP) 174 | ''' 175 | return ti.Matrix(xs) 176 | 177 | 178 | def dot(a, b): 179 | ''' 180 | Calculate the dot product of two vectors. 181 | 182 | `dot` returns the dot product of two vectors, i.e., 183 | `x[0] * y[0] + x[1] * y[1] + ...`. 184 | 185 | :parameter x: 186 | Specifies the first of two vectors. 187 | :parameter y: 188 | Specifies the second of two vectors. 189 | 190 | :return: 191 | The return value can be calculated as `summation(a * b)`. 192 | 193 | :see also: 194 | :func:`cross`, :func:`length`, :func:`outerProduct`, 195 | :func:`summation`. 196 | ''' 197 | return a.dot(b) 198 | 199 | 200 | def normalize(v): 201 | ''' 202 | Calculates the unit vector in the same direction as the original vector. 203 | 204 | `normalize` returns a vector with the same direction as its parameter, 205 | `v`, but with length 1. 206 | 207 | :parameter v: 208 | Specifies the vector to normalize. 209 | 210 | :return: 211 | The return value can be calculated as `v * invLength(v)`. 212 | 213 | :see also: 214 | :func:`dot`, :func:`length`, :func:`invLength`. 215 | ''' 216 | return v.normalized() 217 | 218 | 219 | @ti.func 220 | def normalizePow(v, n, eps=0): 221 | ''' 222 | Return a vector with same direction but with a `n`-powered length. 223 | 224 | This can be used in calculating gravitational force or speed around vortex. 225 | 226 | :parameter v: 227 | Specifies the vector to normalize. 228 | 229 | :parameter n: 230 | Specifies the power number to tweak the length. 231 | 232 | :return: 233 | The return value can be calculated as `normalize(v) * length(v) ** n`. 234 | ''' 235 | l2 = v.norm_sqr() + eps 236 | return v * (l2**((n - 1) / 2)) 237 | 238 | 239 | def summation(v): 240 | ''' 241 | Calculate the sum of all elements in a vector. 242 | 243 | `summation` returns the sum of vector elements, i.e., 244 | `v[0] + v[1] + ...`. 245 | 246 | :parameter v: 247 | Specifies the vector whose elements to sum up. 248 | 249 | :return: 250 | The return value is the sum of all elements in `v`. 251 | 252 | :see also: 253 | :func:`minimum`, :func:`maximum`. 254 | ''' 255 | return v.sum() 256 | 257 | 258 | def maximum(v): 259 | ''' 260 | Find the maximum value in all elements of a vector. 261 | 262 | `maximum` returns the maximum value in vector elements, i.e., 263 | `max(v[0], v[1] + ...)`. 264 | 265 | :parameter v: 266 | Specifies the vector where to find the maximum element. 267 | 268 | :return: 269 | The return value is the maximum value in vector elements. 270 | 271 | :see also: 272 | :func:`minimum`, :func:`summation`. 273 | ''' 274 | return v.max() 275 | 276 | 277 | def minimum(v): 278 | ''' 279 | Find the minimum value in all elements of a vector. 280 | 281 | `minimum` returns the minimum value in vector elements, i.e., 282 | `max(v[0], v[1] + ...)`. 283 | 284 | :parameter v: 285 | Specifies the vector where to find the minimum element. 286 | 287 | :return: 288 | The return value is the minimum value in vector elements. 289 | 290 | :see also: 291 | :func:`maximum`, :func:`summation`. 292 | ''' 293 | return v.min() 294 | 295 | 296 | def cross(a, b): 297 | ''' 298 | Calculate the cross product of two vectors. 299 | The argument can be 2D or 3D. If 3D, the return value is 3D; 300 | If 2D, the return value is a scalar. 301 | The computation is done in a right-handed coordinate system, i.e.:: 302 | 303 | cross(vec(1, 0, 0), vec(0, 1, 0)) == vec(0, 0, 1) 304 | cross(vec(1, 0), vec(0, 1)) == 1 305 | 306 | :parameter x: 307 | Specifies the first of two vectors, can be 2D or 3D. 308 | :parameter y: 309 | Specifies the second of two vectors, can be 2D or 3D. 310 | 311 | :return: 312 | The return value can be calculated as `summation(a * b)`. 313 | 314 | For example:: 315 | 316 | a = vec(1, 2, 3) 317 | b = vec(4, 5, 6) 318 | c = cross(a, b) 319 | # c = [2*6 - 5*3, 4*3 - 1*6, 1*5 - 4*2] = [-3, 6, -3] 320 | 321 | p = vec(1, 2) 322 | q = vec(4, 5) 323 | r = cross(a, b) 324 | # r = 1*5 - 4*2 = -3 325 | ''' 326 | return a.cross(b) 327 | 328 | 329 | def outerProduct(a, b): 330 | return a.outer_product(b) 331 | 332 | 333 | def sqrLength(x): 334 | ''' 335 | Calculate the square of the length of a vector. 336 | 337 | `sqrLength` returns the square of length of the vector, i.e., 338 | `x[0] ** 2 + x[1] * 2 + ...`. 339 | 340 | :parameter x: 341 | Specifies a vector of which to calculate the length. 342 | 343 | :return: 344 | The return value can be calculated as `dot(x, x)`. 345 | 346 | :see also: 347 | :func:`length`, :func:`dot`, :func:`invLength`. 348 | ''' 349 | return x.norm_sqr() 350 | 351 | 352 | def length(x): 353 | ''' 354 | Calculate the length of a vector. 355 | 356 | `length` returns the length or magnitude of the vector, i.e., 357 | `sqrt(x[0] ** 2 + x[1] * 2 + ...)`. 358 | 359 | :parameter x: 360 | Specifies a vector of which to calculate the length. 361 | 362 | :return: 363 | The return value can be calculated as `sqrt(dot(x, x))`. 364 | 365 | :see also: 366 | :func:`distance`, :func:`invLength`, :func:`dot`. 367 | ''' 368 | return x.norm() 369 | 370 | 371 | def invLength(x): 372 | ''' 373 | Calculate the inverse of length of a vector. 374 | 375 | `invLength` returns the inverse of the magnitude of the vector, i.e., 376 | `1 / sqrt(x[0] ** 2 + x[1] * 2 + ...)`. 377 | 378 | :parameter x: 379 | Specifies a vector of which to calculate the length. 380 | 381 | :return: 382 | The return value can be calculated as `inversesqrt(dot(x, x))`. 383 | 384 | :see also: 385 | :func:`length`, :func:`normalize`, :func:`inversesqrt`. 386 | ''' 387 | try: 388 | return x.norm_inv() 389 | except: 390 | return 1 / x.norm() 391 | 392 | 393 | @ti.func 394 | def distance(a, b): 395 | ''' 396 | Calculate the distance between two points. 397 | 398 | :parameter a: 399 | Specifies the first of two points. 400 | 401 | :parameter b: 402 | Specifies the second of two points. 403 | 404 | :return: 405 | The return value is calculated as `length(a - b)`. 406 | 407 | :see also: 408 | :func:`length`, :func:`normalize`, :func:`dot`. 409 | ''' 410 | return (a - b).norm() 411 | 412 | 413 | @ti.func 414 | def reflect(I, N): 415 | ''' 416 | Calculate the reflection direction for an incident vector. 417 | 418 | :parameter I: 419 | Specifies the incident vector. 420 | :parameter N: 421 | Specifies the normal vector. 422 | 423 | :return: 424 | For a given incident vector I and surface normal N, `reflect` 425 | returns the reflection direction calculated as 426 | `I - 2 * dot(N, I) * N`. 427 | 428 | :note: 429 | N should be normalized in order to achieve the desired result. 430 | 431 | :see also: 432 | :func:`refract`, :func:`normalize`, :func:`dot`. 433 | ''' 434 | return I - 2 * dot(N, I) * N 435 | 436 | 437 | @ti.func 438 | def refract(I, N, eta): 439 | ''' 440 | Calculate the refraction direction for an incident vector. 441 | 442 | For a given incident vector I, surface normal N and ratio of indices 443 | of refraction, eta, refract returns the refraction vector, R. 444 | 445 | :parameter I: 446 | Specifies the incident vector. 447 | :parameter N: 448 | Specifies the normal vector. 449 | :parameter eta: 450 | Specifies the ratio of indices of refraction. 451 | 452 | :return: 453 | The return value is calculated as:: 454 | 455 | k = 1 - eta * eta * (1 - dot(N, I) * dot(N, I)) 456 | R = I * 0 457 | if k >= 0: 458 | R = eta * I - (eta * dot(N, I) + sqrt(k)) * N 459 | return R 460 | 461 | :note: 462 | The input parameters I and N should be normalized in order to 463 | achieve the desired result. 464 | 465 | :see also: 466 | :func:`reflect`, :func:`normalize`, :func:`dot`. 467 | ''' 468 | NoI = dot(N, I) 469 | k = 1 - eta**2 * (1 - NoI**2) 470 | R = I * 0 471 | if k >= 0: 472 | R = eta * I - (eta * NoI + ti.sqrt(k)) * N 473 | 474 | 475 | def shuffle(a, *ks): 476 | ret = [] 477 | for k in ks: 478 | t = a.subscript(k) 479 | ret.append(ti.expr_init(t)) 480 | return ti.Vector(ret) 481 | 482 | 483 | def _vector_getattr(self, key): 484 | ret = [] 485 | stk = [] 486 | for k in key: 487 | sgn = 0 488 | i = 0 489 | if k != '_': 490 | sgn = 1 491 | i = 'xyzw'.find(k) 492 | if i == -1: 493 | sgn = -1 494 | i = 'XYZW'.find(k) 495 | if i == -1: 496 | break 497 | 498 | stk.append((i, sgn)) 499 | 500 | else: 501 | for i, sgn in stk: 502 | if ti.inside_kernel(): 503 | t = self.subscript(i) * sgn 504 | ret.append(ti.expr_init(t)) 505 | else: 506 | t = self[i] * sgn 507 | ret.append(t) 508 | 509 | return ti.Vector(ret) 510 | 511 | _taichi_skip_traceback = 1 512 | raise AttributeError(f"'Matrix' object has no attribute {key}") 513 | 514 | 515 | ti.Matrix.__getattr__ = _vector_getattr 516 | -------------------------------------------------------------------------------- /taichi_glsl/version.py: -------------------------------------------------------------------------------- 1 | version = (0, 0, 12) 2 | taichi_version = (0, 9, 0) 3 | 4 | print(f'[TaiGLSL] version {".".join(map(str, version))}') 5 | -------------------------------------------------------------------------------- /tests/test_array.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | import pytest 3 | 4 | dt_map = { 5 | 'float32': ti.f32, 6 | 'float64': ti.f64, 7 | 'int8': ti.i8, 8 | 'int16': ti.i16, 9 | 'int32': ti.i32, 10 | 'int64': ti.i64, 11 | 'uint8': ti.u8, 12 | 'uint16': ti.u16, 13 | 'uint32': ti.u32, 14 | 'uint64': ti.u64, 15 | 'float': ti.f32, 16 | 'int': ti.i32, 17 | float: ti.f32, 18 | int: ti.i32, 19 | } 20 | 21 | shapes = [(), (2, ), (3, 4)] 22 | 23 | 24 | @pytest.mark.parametrize('dt_name,dt', dt_map.items()) 25 | def test_dtype(dt_name, dt): 26 | assert dtype(dt_name) == dt 27 | 28 | 29 | @pytest.mark.parametrize('shape', shapes) 30 | @pytest.mark.parametrize('dt', set(dt_map.values())) 31 | @ti.host_arch_only 32 | def test_array(dt, shape): 33 | x = array(dt, *shape) 34 | 35 | @ti.kernel 36 | def func0(): 37 | x[None] = 2 38 | 39 | @ti.kernel 40 | def func1(): 41 | for i in x: 42 | x[i] = i + 1 43 | 44 | @ti.kernel 45 | def func2(): 46 | for i, j in x: 47 | x[i, j] = i + j + 1 48 | 49 | eval(f'func{len(shape)}')() 50 | 51 | 52 | @pytest.mark.parametrize('dt', [float, int]) 53 | @ti.host_arch_only 54 | def test_mat_array(dt): 55 | x = mat_array(2, 3, dt, 4, 3) 56 | 57 | @ti.kernel 58 | def func(): 59 | for i, j in x: 60 | x[i, j] = mat([2, 3, 4], [5, 6, 7]) 61 | 62 | func() 63 | 64 | 65 | @pytest.mark.parametrize('dt', [float, int]) 66 | @ti.host_arch_only 67 | def test_vec_array(dt): 68 | x = vec_array(3, dt, 2, 3) 69 | 70 | @ti.kernel 71 | def func(): 72 | for i, j in x: 73 | x[i, j] = vec(2, 3, 4) 74 | 75 | func() 76 | -------------------------------------------------------------------------------- /tests/test_odop.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import taichi_glsl as tl 3 | 4 | 5 | @ti.host_arch_only 6 | def test_subscript_0d(): 7 | c = tl.Complex.field(ti.f32, ()) 8 | 9 | @ti.kernel 10 | def func() -> ti.f32: 11 | return c[None].x + c[None].y 12 | 13 | # tl.TaichiClass.Proxy W.I.P. 14 | c.x[None] = 2 15 | c.y[None] = 11 16 | assert func() == 13 17 | 18 | 19 | @ti.host_arch_only 20 | def test_subscript_1d(): 21 | c = tl.Complex.field(ti.f32, 4) 22 | 23 | @ti.kernel 24 | def func(i: ti.i32) -> ti.f32: 25 | return c[i].x + c[i].y 26 | 27 | c.x[2] = 2 28 | c.x[1] = 11 29 | c.y[2] = 7 30 | assert func(2) == 9 31 | 32 | 33 | @ti.host_arch_only 34 | def test_loop_range_1d(): 35 | c = tl.Complex.field(ti.f32, 4) 36 | count = ti.field(ti.i32, ()) 37 | 38 | @ti.kernel 39 | def func() -> ti.f32: 40 | res = 0.0 41 | for i in c: 42 | count[None] += i 43 | res += c[i].x 44 | return res 45 | 46 | c.x[2] = 2 47 | c.x[1] = 11 48 | c.y[2] = 7 49 | c.y[2] = 8 50 | assert func() == 13 51 | assert count[None] == 6 52 | 53 | 54 | @ti.host_arch_only 55 | def test_loop_range_2d(): 56 | c = tl.Complex.field(ti.f32, (2, 4)) 57 | count = ti.field(ti.i32, ()) 58 | 59 | @ti.kernel 60 | def func() -> ti.f32: 61 | res = 0.0 62 | for i, j in c: 63 | count[None] += j 64 | res += c[i, j].x 65 | return res 66 | 67 | c.x[0, 2] = 2 68 | c.x[1, 1] = 11 69 | c.x[1, 2] = 8 70 | c.y[1, 0] = 7 71 | c.y[1, 2] = 8 72 | assert func() == 21 73 | assert count[None] == 12 74 | 75 | 76 | @ti.host_arch_only 77 | def test_loop_atomic_add(): 78 | c = tl.Complex.field(ti.f32, 4) 79 | r = tl.Complex.field(ti.f32, ()) 80 | 81 | @ti.kernel 82 | def func(): 83 | r[None] = 0.0 84 | for i in c: 85 | r[None] += c[i] 86 | 87 | c.x[2] = 2 88 | c.x[1] = 11 89 | c.y[2] = 7 90 | c.y[3] = 13 91 | func() 92 | assert r.x[None] == 13 93 | assert r.y[None] == 20 94 | 95 | 96 | @ti.host_arch_only 97 | def test_loop_matmul(): 98 | c = tl.Complex.field(ti.f32, 2) 99 | r = tl.Complex.field(ti.f32, ()) 100 | 101 | @ti.kernel 102 | def func(): 103 | r[None] = tl.Complex(1.0, 0.0) 104 | for i in ti.static(range(2)): 105 | print('r[None] =', r[None]) 106 | print('c[i] =', c[i]) 107 | print('b', i, r[None] @ c[i]) 108 | r[None] = r[None] @ c[i] 109 | print('a', i, r[None]) 110 | 111 | c.x[0] = 2 112 | c.y[0] = 7 113 | c.x[1] = 11 114 | c.y[1] = 13 115 | func() 116 | print('r', tl.Complex(r.x[None], r.y[None])) 117 | assert r.x[None] == 2 * 11 - 7 * 13 118 | assert r.y[None] == 2 * 13 + 7 * 11 119 | -------------------------------------------------------------------------------- /tests/test_rand.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | from pytest import approx 3 | 4 | 5 | def mean_approx_test(x, xmin, xmax, rel=1e-2): 6 | if not isinstance(x, np.ndarray): 7 | x = x.to_numpy() 8 | x = (x - xmin) / (xmax - xmin) 9 | for i in range(1, 4): 10 | assert (x**i).mean() == approx(1 / (i + 1), rel=rel) 11 | 12 | 13 | @ti.host_arch_only 14 | def test_rand(): 15 | n = 1024**2 16 | x = array(float, n) 17 | 18 | @ti.kernel 19 | def fill(): 20 | for i in x: 21 | x[i] = rand() 22 | 23 | fill() 24 | mean_approx_test(x, 0, 1) 25 | 26 | 27 | @ti.host_arch_only 28 | def test_rand_independent_product(): 29 | n = 1024**2 30 | x = array(float, n) 31 | 32 | @ti.kernel 33 | def fill(): 34 | for i in x: 35 | x[i] = rand() * rand() 36 | 37 | fill() 38 | assert x.to_numpy().mean() == approx(1 / 4, rel=1e-2) 39 | 40 | 41 | @ti.host_arch_only 42 | def test_rand_range(): 43 | n = 1024**2 44 | a, b = 0.6, 1.4 45 | x = array(ti.f32, n) 46 | 47 | @ti.kernel 48 | def fill(): 49 | for i in x: 50 | x[i] = randRange(a, b) 51 | 52 | fill() 53 | mean_approx_test(x, a, b) 54 | 55 | 56 | @ti.host_arch_only 57 | def test_rand_2d(): 58 | n = 8192 59 | x = vec_array(2, float, n) 60 | 61 | @ti.kernel 62 | def fill(): 63 | for i in x: 64 | x[i] = randND(2) 65 | 66 | fill() 67 | x = x.to_numpy() 68 | counters = [0 for _ in range(4)] 69 | for i in range(n): 70 | c = int(x[i, 0] < 0.5) * 2 + int(x[i, 1] < 0.5) 71 | counters[c] += 1 72 | 73 | for c in range(4): 74 | assert counters[c] / n == approx(1 / 4, rel=0.2) 75 | 76 | 77 | @ti.host_arch_only 78 | def test_rand_int(): 79 | n = 1024**2 80 | a, b = 768, 1131 81 | x = array(float, n) 82 | 83 | @ti.kernel 84 | def fill(): 85 | for i in x: 86 | x[i] = randInt(a, b) 87 | 88 | fill() 89 | mean_approx_test(x, a, b, rel=1e-1) 90 | 91 | 92 | @ti.host_arch_only 93 | def test_rand_unit_2d(): 94 | n = 1024 95 | x = vec_array(2, float, n) 96 | 97 | @ti.kernel 98 | def fill(): 99 | for i in x: 100 | x[i] = randUnit2D() 101 | 102 | fill() 103 | x = x.to_numpy() 104 | len = np.sum(x**2, axis=1) 105 | ang = np.arctan2(x[:, 1], x[:, 0]) 106 | assert len == approx(np.ones(n)) 107 | mean_approx_test(ang, -math.pi, math.pi, rel=4e-1) 108 | 109 | 110 | test_rand_int() 111 | -------------------------------------------------------------------------------- /tests/test_scalar.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | import pytest 3 | 4 | test_range = [ 5 | -1e3 + 0.1, 6 | -1e-4 + 1, 7 | -15.9, 8 | -1.501, 9 | -1.499, 10 | -1.0, 11 | -0.34, 12 | 0.5, 13 | 0.0, 14 | 0.33, 15 | 0.5, 16 | 1.0, 17 | 1.499, 18 | 1.501, 19 | 4.0, 20 | 9.99, 21 | 1e5, 22 | ] 23 | test_range_2 = [ 24 | (-1.0, 0.0), 25 | (-1.0, 0.5), 26 | (0.0, 0.5), 27 | (0.0, -1.0), 28 | (0.5, 1.0), 29 | (0.5, 1.5), 30 | (1e5, 0.0), 31 | (1e5, 0.5), 32 | (0.5, -1e5), 33 | (-3.5, 0.5), 34 | ] 35 | 36 | 37 | @pytest.mark.parametrize('a', test_range) 38 | def test_sign(a): 39 | @ti.kernel 40 | def calc(a: ti.f32) -> ti.f32: 41 | return sign(a) 42 | 43 | r = calc(a) 44 | if a > 0: 45 | assert r == 1 46 | elif a == 0: 47 | assert r == 0 48 | else: 49 | assert r == -1 50 | 51 | 52 | @pytest.mark.parametrize('a', test_range) 53 | def test_round(a): 54 | @ti.kernel 55 | def calc(a: ti.f32) -> ti.f32: 56 | return round(a) 57 | 58 | r = calc(a) 59 | # According to https://www.khronos.org/registry/OpenGL-Refpages/gl4/html/round.xhtml: 60 | # round returns a value equal to the nearest integer to x. The fraction 0.5 will round in a direction chosen by the implementation, presumably the direction that is fastest. This includes the possibility that round(x) returns the same value as roundEven(x) for all values of x. 61 | # So our implementation could: ts.round(0.5) = 1.0 62 | # However, np.round(0.5) = 0.0 63 | # So let's play a trick on it: 64 | assert r == pytest.approx(np.round(a + 1e-7)) 65 | 66 | 67 | @pytest.mark.parametrize('a', test_range) 68 | def test_floor(a): 69 | @ti.kernel 70 | def calc(a: ti.f32) -> ti.f32: 71 | return floor(a) 72 | 73 | r = calc(a) 74 | assert r == pytest.approx(np.floor(a), rel=1e-3) 75 | 76 | 77 | @pytest.mark.parametrize('a', test_range) 78 | def test_fract(a): 79 | @ti.kernel 80 | def calc(a: ti.f32) -> ti.f32: 81 | return fract(a) 82 | 83 | r = calc(a) 84 | assert r == pytest.approx(a - np.floor(a), rel=1e-3) 85 | 86 | 87 | @pytest.mark.parametrize('a', 88 | [4e-2, 0.3, 0.8, 1.0, 1.55, 2.0, 4.2, 141.0, 1e3]) 89 | def test_inversesqrt(a): 90 | @ti.kernel 91 | def calc(a: ti.f32) -> ti.f32: 92 | return inversesqrt(a) 93 | 94 | r = calc(a) 95 | assert r == pytest.approx(a**-0.5, rel=1e-3) 96 | 97 | 98 | @pytest.mark.parametrize('a', test_range) 99 | def test_atan(a): 100 | @ti.kernel 101 | def calc(a: ti.f32) -> ti.f32: 102 | return atan(a) 103 | 104 | r = calc(a) 105 | assert r == pytest.approx(np.arctan(a)) 106 | 107 | 108 | @pytest.mark.parametrize('a,b', test_range_2) 109 | def test_atan2(a, b): 110 | @ti.kernel 111 | def calc(a: ti.f32, b: ti.f32) -> ti.f32: 112 | return atan(a, b) 113 | 114 | r = calc(a, b) 115 | assert r == pytest.approx(np.arctan2(a, b)) 116 | 117 | 118 | @pytest.mark.parametrize('a,b', test_range_2) 119 | def test_step(a, b): 120 | @ti.kernel 121 | def calc(a: ti.f32, b: ti.f32) -> ti.f32: 122 | return step(a, b) 123 | 124 | r = calc(a, b) 125 | if a < b: 126 | assert r == 1 127 | else: 128 | assert r == 0 129 | 130 | 131 | @pytest.mark.parametrize('a,b', test_range_2) 132 | def test_sign2(a, b): 133 | @ti.kernel 134 | def calc(a: ti.f32, b: ti.f32) -> ti.f32: 135 | return sign(a, b) 136 | 137 | r = calc(a, b) 138 | if a > b: 139 | assert r == 1 140 | elif a == b: 141 | assert r == 0 142 | else: 143 | assert r == -1 144 | 145 | 146 | @pytest.mark.parametrize('a,b,c', [ 147 | (0, 0, 1), 148 | (1, 0, 0), 149 | (0, 1, 0), 150 | (1, 1, 0), 151 | ]) 152 | def test_isnan(a, b, c): 153 | @ti.kernel 154 | def calc(a: ti.f32, b: ti.f32) -> ti.i32: 155 | return isnan(a / b) 156 | 157 | r = calc(a, b) 158 | assert r == c 159 | 160 | 161 | @pytest.mark.parametrize('a,b,c', [ 162 | (0, 0, 0), 163 | (1, 0, 1), 164 | (0, 1, 0), 165 | (1, 1, 0), 166 | ]) 167 | def test_isinf(a, b, c): 168 | @ti.kernel 169 | def calc(a: ti.f32, b: ti.f32) -> ti.i32: 170 | return isinf(a / b) 171 | 172 | r = calc(a, b) 173 | assert r == c 174 | -------------------------------------------------------------------------------- /tests/test_vector.py: -------------------------------------------------------------------------------- 1 | from taichi_glsl import * 2 | 3 | 4 | @ti.host_arch_only 5 | def test_vec_fill(): 6 | p = vec_uniform(2, int) 7 | q = vec_uniform(3, int) 8 | r = vec_uniform(4, int) 9 | 10 | @ti.kernel 11 | def func(): 12 | p[None] = vec2(5) 13 | q[None] = vec3(5) 14 | r[None] = vec4(5) 15 | 16 | func() 17 | assert np.allclose(p.to_numpy(), np.ones(2) * 5) 18 | assert np.allclose(q.to_numpy(), np.ones(3) * 5) 19 | assert np.allclose(r.to_numpy(), np.ones(4) * 5) 20 | 21 | 22 | @ti.host_arch_only 23 | def test_vec_compose(): 24 | q = vec_uniform(3, int) 25 | r = vec_uniform(4, int) 26 | s = vec_uniform(3, int) 27 | 28 | @ti.kernel 29 | def func(): 30 | p1 = vec(2, 3) 31 | p2 = vec(5, 6) 32 | q[None] = vec(1, p1) 33 | r[None] = vec(p1, p2) 34 | s[None] = vec3(q[None]) 35 | 36 | func() 37 | assert np.allclose(q.to_numpy(), np.array([1, 2, 3])) 38 | assert np.allclose(r.to_numpy(), np.array([2, 3, 5, 6])) 39 | assert np.allclose(s.to_numpy(), np.array([1, 2, 3])) 40 | 41 | 42 | @ti.host_arch_only 43 | def test_vec_compose_nd(): 44 | q = vec_uniform(3, int) 45 | r = vec_uniform(4, int) 46 | 47 | @ti.kernel 48 | def func(): 49 | p1 = vec2(2, 3) 50 | p2 = vec2(5, 6) 51 | q[None] = vec3(1, p1) 52 | r[None] = vec4(p1, p2) 53 | 54 | func() 55 | assert np.allclose(q.to_numpy(), np.array([1, 2, 3])) 56 | assert np.allclose(r.to_numpy(), np.array([2, 3, 5, 6])) 57 | 58 | 59 | @ti.host_arch_only 60 | @ti.must_throw(ValueError) 61 | def test_vec_compose_mismatch_less(): 62 | q = vec_uniform(3, int) 63 | 64 | @ti.kernel 65 | def func(): 66 | p1 = vec2(2, 3) 67 | q[None] = vec2(1, p1) 68 | 69 | func() 70 | 71 | 72 | @ti.host_arch_only 73 | @ti.must_throw(ValueError) 74 | def test_vec_compose_mismatch_more(): 75 | q = vec_uniform(3, int) 76 | r = vec_uniform(4, int) 77 | 78 | @ti.kernel 79 | def func(): 80 | p1 = vec2(2, 3) 81 | q[None] = vec4(1, p1) 82 | 83 | func() 84 | --------------------------------------------------------------------------------