def mainImage(iMouse: ti.Vector, iTime: ti.f32, i: ti.i32,
j: ti.i32) -> ti.Vector:
fragCoord = ti.Vector([i, j])
p = -1.0 + 2.0 * fragCoord / iResolution
m = -1.0 + 2.0 * iMouse # iMouse 已经归一化
a1 = ti.atan2(p[1] - m[1], p[0] - m[0])
a2 = ti.atan2(p[1] + m[1], p[0] + m[0])
r1 = ti.sqrt((p - m).dot(p - m))
r2 = ti.sqrt((p + m).dot(p + m))
uv = ti.Vector(
[0.2 * iTime + (r1 - r2) * 0.25,
ti.asin(ti.sin(a1 - a2)) / 3.1416])
w = ti.exp(-15.0 * r1 * r1) + ti.exp(-15.0 * r2 * r2)
w += 0.25 * smoothstep(0.93, 1.0, ti.sin(128.0 * uv[0]))
w += 0.25 * smoothstep(0.93, 1.0, ti.sin(128.0 * uv[1]))
# 可以使用纹理
# vec3 col = texture( iChannel0, 0.125*uv ).zyx;
col = ti.Vector([0.0, 0.0, 0.0])
fragColor = col + w
return fragColor
def rotateY(angle):
return ti.Matrix([
[ ti.cos(angle), 0, ti.sin(angle), 0],
[ 0, 1, 0, 0],
[-ti.sin(angle), 0, ti.cos(angle), 0],
[ 0, 0, 0, 1],
])
def rotateX(angle):
return ti.Matrix([
[1, 0, 0, 0],
[0, ti.cos(angle), -ti.sin(angle), 0],
[0, ti.sin(angle), ti.cos(angle), 0],
[0, 0, 0, 1],
])
def render(t: ti.f32):
# camera & camera tx
ro = ti.Vector([ti.sin(t * 0.16), 0.0, ti.cos(t * 0.1)])
ta = ro + ti.Vector([ti.sin(t * 0.15), ti.sin(t * 0.18), ti.cos(t * 0.24)])
cw = (ta - ro).normalized()
cp = ti.Vector([0.0, 1.0, 0.0])
cu = cp.cross(cw).normalized()
for i, j in pixels:
# coords-trans
coords = -1.0 + 2.0 * ti.Vector([i / m, j / n])
coords[0] *= (m / n)
rd = (coords[0] * cu + coords[1] * cp + 2.0 * cw).normalized()
v = ti.Vector([0.0, 0.0, 0.0])
for k1 in range(50):
s1 = (k1 + 1) * 0.1
p = ro + rd * s1
for k2 in range(8):
s2 = 0.1 + 0.12 * k2
p = abs(p) / (p + ti.sin(t * 0.1) * 0.21).dot(p) - 0.85
a = p.norm() * 0.12
v += ti.Vector([a * s2, a * s2 ** 2, a * s2 ** 3])
pixels[i, j] = v * 0.01
def glCircle(center, dir1, dir2, N: ti.template()):
for n in range(N):
a = n * ts.math.tau / N
b = a - ts.math.tau / N
p = center + dir1 * ti.cos(a) + dir2 * ti.sin(a)
q = center + dir1 * ti.cos(b) + dir2 * ti.sin(b)
glTri(p, center, q)
def getPixel(coord, timef32):
uv = -1.0 + 2.0 * ti.Vector([coord[0] / WIDTH, coord[1] / HEIGHT])
uv[0] *= (WIDTH / HEIGHT)
# ray
ang = ti.Vector(
[ti.sin(timef32 * 3.0) * 0.1,
ti.sin(timef32) * 0.2 + 0.3, timef32])
ori = ti.Vector([0.0, 3.5, timef32 * 5.0])
dir = ti.Vector([uv[0], uv[1], -2.0]).normalized()
dir[2] += uv.norm() * 0.14
dir_r = dir.normalized()
me = fromEuler(ang)
dir = ti.Vector([
dir_r.dot(ti.Vector([me[0, 0], me[0, 1], me[0, 2]])),
dir_r.dot(ti.Vector([me[1, 0], me[1, 1], me[1, 2]])),
dir_r.dot(ti.Vector([me[2, 0], me[2, 1], me[2, 2]]))
])
# tracing
p = ti.Vector([0.0, 0.0, 0.0])
p = heightMapTracing(ori, dir, p, timef32)
dist = p - ori
n = getNormal(p, dist.dot(dist) * EPSILON_NRM, timef32)
light = ti.Vector([0.0, 1.0, 0.8]).normalized()
# color
return mix(getSkyColor(dir), getSeaColor(p, n, light, dir, dist),
pow(smoothstep(0.0, -0.02, dir[1]), 0.2))
def build_b(self):
self.b = np.zeros(shape=(self.N, self.N))
for i, j in np.ndindex(self.N, self.N):
xl = i / self.N
yl = j / self.N
self.b[i, j] = ti.sin(2.0 * np.pi * xl) * ti.sin(2.0 * np.pi * yl)
self.b = self.b.flatten(order='C')
def CheckSdfCapsule(pos, vel):
cap_rot = capsule_rotation[None][0]
capsule_rotmat = ti.Matrix([
[ti.cos(cap_rot), -ti.sin(cap_rot)],
[ti.sin(cap_rot), ti.cos(cap_rot)],
])
local_pos = WorldSpaceToMaterialSpace(pos, capsule_translation[None],
capsule_rotmat)
phi = SdfCapsule(local_pos, capsule_radius, capsule_half_length)
inside = False
dotnv = 0.
diff_vel = ti.Vector.zero(float, 2)
n = ti.Vector.zero(float, 2)
if phi < 0.0:
n = capsule_rotmat @ SdfNormalCapsule(local_pos, capsule_radius,
capsule_half_length)
solid_vel = ti.Vector([
capsule_trans_vel[None][0] - capsule_angular_vel *
(pos[1] - capsule_translation[None][1]),
capsule_trans_vel[None][1] + capsule_angular_vel *
(pos[0] - capsule_translation[None][0]),
])
diff_vel = solid_vel - vel
dotnv = n.dot(diff_vel)
if dotnv > 0.0:
inside = True
return inside, dotnv, diff_vel, n
def noise(p):
i = ti.floor(p)
a = i.dot(vec(1.0, 57.0, 21.0)) + vec(0.0, 57.0, 21.0, 78.0)
f = ti.cos((p - i) * math.acos(-1.0)) * (-0.5) + 0.5
a = mix(ti.sin(ti.cos(a) * a), ti.sin(ti.cos(1.0 + a) * (1.0 + a)), f[0])
a[0] = mix(a[0], a[1], f[1])
a[1] = mix(a[2], a[3], f[1])
return mix(a[0], a[1], f[2])
def test_inner_loops_local_variable():
for i in arr:
for j in range(3):
s = 0.0
t = 0.0
for k in range(3):
s += ti.sin(x[None]) + 1.0
t += ti.sin(x[None])
loss[None] += s + t
def glCircle(center, dir1, dir2, N: ti.template()):
for n in range(N):
a = n * ts.math.tau / N
b = a - ts.math.tau / N
p = center + dir1 * ti.cos(a) + dir2 * ti.sin(a)
q = center + dir1 * ti.cos(b) + dir2 * ti.sin(b)
# make the face double-sided
glTri(q, center, p)
glTri(p, center, q)
def get_transform_matrix(a):
A = ti.Matrix([[1., 0., -center[None][0]], [0., 1., -center[None][1]],
[0., 0., 1.]])
theta = a / 180. * pi
A2 = ti.Matrix([[ti.cos(theta), -ti.sin(theta), 0.],
[ti.sin(theta), ti.cos(theta), 0.], [0., 0., 1.]])
A2 = A2 @ A
A = ti.Matrix([[1., 0., center[None][0]], [0., 1., center[None][1]],
[0., 0., 1.]])
transform[None] = A @ A2
def glCylinder(center, polar, dir1, dir2, N: ti.template()):
for n in range(N):
a = n * ts.math.tau / N
b = a + ts.math.tau / N
p = center + dir1 * ti.cos(a) + dir2 * ti.sin(a)
q = center + dir1 * ti.cos(b) + dir2 * ti.sin(b)
glQuad(p, q, q + polar, p + polar)
glCircle(center, dir2, dir1, N)
glCircle(center + polar, dir1, dir2, N)
def test_more_inner_loops_local_variable():
for i in arr:
for j in range(2):
s = 0.0
for k in range(3):
u = 0.0
s += ti.sin(x[None]) + 1.0
for l in range(2):
u += ti.sin(x[None])
loss[None] += u
loss[None] += s
def init():
for i, j in ti.ndrange((N_ext, N_tot - N_ext), (N_ext, N_tot - N_ext)):
# xl, yl, zl = [0,1]
xl = (i - N_ext) * 2.0 / N_tot
yl = (j - N_ext) * 2.0 / N_tot
# r[0] = b - Ax, where x = 0; therefore r[0] = b
r[i, j] = ti.sin(2.0 * np.pi * xl) * ti.sin(2.0 * np.pi * yl)
b[i, j] = ti.sin(2.0 * np.pi * xl) * ti.sin(2.0 * np.pi * yl)
Ap[i, j] = 0.0
Ax[i, j] = 0.0
p[i, j] = 0.0
x[i, j] = 0.0
def init(self):
for i, j in ti.ndrange((self.N_ext, self.N_tot - self.N_ext),
(self.N_ext, self.N_tot - self.N_ext)):
# xl, yl, zl = [0,1]
xl = (i - self.N_ext) * 2.0 / self.N_tot
yl = (j - self.N_ext) * 2.0 / self.N_tot
# r[0] = b - Ax, where x = 0; therefore r[0] = b
self.r[i, j] = ti.sin(2.0 * np.pi * xl) * ti.sin(2.0 * np.pi * yl)
self.b[i, j] = ti.sin(2.0 * np.pi * xl) * ti.sin(2.0 * np.pi * yl)
self.Ap[i, j] = 0.0
self.Ax[i, j] = 0.0
self.p[i, j] = 0.0
def init():
for i, j, k in ti.ndrange((N_ext, N_tot - N_ext), (N_ext, N_tot - N_ext),
(N_ext, N_tot - N_ext)):
xl = (i - N_ext) * 2.0 / N_tot
yl = (j - N_ext) * 2.0 / N_tot
zl = (k - N_ext) * 2.0 / N_tot
r[0][i, j, k] = ti.sin(2.0 * np.pi * xl) * ti.sin(
2.0 * np.pi * yl) * ti.sin(2.0 * np.pi * zl)
z[0][i, j, k] = 0.0
Ap[i, j, k] = 0.0
p[i, j, k] = 0.0
x[i, j, k] = 0.0
def init(self):
for i, j, k in ti.ndrange((self.N_ext, self.N_tot - self.N_ext),
(self.N_ext, self.N_tot - self.N_ext),
(self.N_ext, self.N_tot - self.N_ext)):
xl = (i - self.N_ext) * 2.0 / self.N_tot
yl = (j - self.N_ext) * 2.0 / self.N_tot
zl = (k - self.N_ext) * 2.0 / self.N_tot
self.r[0][i, j, k] = ti.sin(2.0 * np.pi * xl) * ti.sin(
2.0 * np.pi * yl) * ti.sin(2.0 * np.pi * zl)
self.z[0][i, j, k] = 0.0
self.Ap[i, j, k] = 0.0
self.p[i, j, k] = 0.0
self.x[i, j, k] = 0.0
def test_stacked_inner_loops_local_variable():
for i in arr:
loss[None] += ti.sin(x[None])
for j in range(3):
s = 0.0
for k in range(3):
s += ti.sin(x[None]) + 1.0
loss[None] += s
for j in range(3):
s = 0.0
for k in range(3):
s += ti.sin(x[None]) + 1.0
loss[None] += s
def fromEuler(ang):
a1 = ti.Vector([ti.sin(ang[0]), ti.cos(ang[0])])
a2 = ti.Vector([ti.sin(ang[1]), ti.cos(ang[1])])
a3 = ti.Vector([ti.sin(ang[2]), ti.cos(ang[2])])
m = ti.Matrix([[
a1[1] * a3[1] + a1[0] * a2[0] * a3[0],
a1[1] * a2[0] * a3[0] + a3[1] * a1[0], -1.0 * a2[1] * a3[0]
], [-1.0 * a2[1] * a1[0], a1[1] * a2[1], a2[0]],
[
a3[1] * a1[0] * a2[0] + a1[1] * a3[0],
a1[0] * a3[0] - a1[1] * a3[1] * a2[0], a2[1] * a3[1]
]])
return m
def ggx_sample(n, wo, roughness):
u = ti.Vector([1.0, 0.0, 0.0])
if abs(n[1]) < 1.0 - eps:
u = n.cross(ti.Vector([0.0, 1.0, 0.0])).normalized()
v = n.cross(u)
r0 = ti.random()
r1 = ti.random()
a = roughness ** 2.0
a2 = a * a
theta = ti.acos(ti.sqrt((1.0 - r0) / ((a2-1.0)*r0 + 1.0)))
phi = 2.0 * math.pi * r1
wm = ti.cos(theta) * n + ti.sin(theta) * ti.cos(phi) * \
u + ti.sin(theta) * ti.sin(phi) * v
wi = reflect(-wo, wm)
return wi
def computeDstToRefTransform(dst_mean_x, dst_mean_y):
"""
Compute the rotation matrix and translation vectors to move dst p.c. into ref p.c
Beased on https://lucidar.me/en/mathematics/singular-value-decomposition-of-a-2x2-matrix/
"""
a = A[0, 0]
b = A[1, 0]
c = A[0, 1]
d = A[1, 1]
teta = 0.5 * ti.atan2(2 * a * c + 2 * b * d, a * a + b * b - c * c - d * d)
U[0, 0] = ti.cos(teta)
U[1, 0] = -ti.sin(teta)
U[0, 1] = ti.sin(teta)
U[1, 1] = ti.cos(teta)
# We don't need the Sigma matrix for ICP
# s1 = a*a + b*b + c*c + d*d
# s2 = ti.sqrt((a*a + b*b - c*c - d*d) * (a*a + b*b - c*c - d*d) + 4 * (a*c + b*d) * (a*c + b*d))
# sigma1 = ti.sqrt((s1 + s2) * 0.5)
# sigma2 = ti.sqrt((s1 - s2) * 0.5)
# S[0,0] = sigma1
# S[1,0] = 0
# S[0,1] = 0
# S[1,1] = sigma2
phi = 0.5 * ti.atan2(2 * a * b + 2 * c * d, a * a - b * b + c * c - d * d)
s11 = (a * ti.cos(teta) + c * ti.sin(teta)) * ti.cos(phi) + (
b * ti.cos(teta) + d * ti.sin(teta)) * ti.sin(phi)
if s11 != 0:
s11 = s11 / ti.abs(s11)
s22 = (a * ti.sin(teta) - c * ti.cos(teta)) * ti.sin(phi) + (
-b * ti.sin(teta) + d * ti.cos(teta)) * ti.cos(phi)
if s22 != 0:
s22 = s22 / ti.abs(s22)
V[0, 0] = s11 * ti.cos(phi)
V[1, 0] = -s22 * ti.sin(phi)
V[0, 1] = s11 * ti.sin(phi)
V[1, 1] = s22 * ti.cos(phi)
Rot[0, 0] = V[0, 0] * U[0, 0] + V[1, 0] * U[1, 0]
Rot[1, 0] = V[0, 0] * U[0, 1] + V[1, 0] * U[1, 1]
Rot[0, 1] = V[0, 1] * U[0, 0] + V[1, 1] * U[1, 0]
Rot[1, 1] = V[0, 1] * U[0, 1] + V[1, 1] * U[1, 1]
Trans[0] = ref_mean[0] - Rot[0, 0] * dst_mean_x - Rot[1, 0] * dst_mean_y
Trans[1] = ref_mean[1] - Rot[0, 1] * dst_mean_x - Rot[1, 1] * dst_mean_y
def nn1(t: ti.i32):
for i in range(n_hidden):
actuation = 0.0
for j in ti.static(range(n_sin_waves)):
actuation += weights1[i, j] * ti.sin(spring_omega * t * dt +
2 * math.pi / n_sin_waves * j)
for j in ti.static(range(n_objects)):
offset = x[t, j] - x[t, head_id]
# use a smaller weight since there are too many of them
actuation += weights1[i, j * 6 + n_sin_waves] * offset[0] * 0.05
actuation += weights1[i,
j * 6 + n_sin_waves + 1] * offset[1] * 0.05
actuation += weights1[i, j * 6 + n_sin_waves + 2] * v[t,
i][0] * 0.05
actuation += weights1[i, j * 6 + n_sin_waves + 3] * v[t,
i][1] * 0.05
actuation += weights1[i, j * 6 + n_sin_waves +
4] * rotation[t, i] * 0.05
actuation += weights1[i, j * 6 + n_sin_waves + 5] * omega[t,
i] * 0.05
actuation += weights1[i, n_objects * 6 + n_sin_waves] * goal[None][0]
actuation += weights1[i,
n_objects * 6 + n_sin_waves + 1] * goal[None][1]
actuation += bias1[i]
actuation = ti.tanh(actuation)
hidden[t, i] = actuation
def step(phase: ti.i32):
# move
for i in position:
pos, ang = position[i], heading[i]
l = sense(phase, pos, ang - SENSE_ANGLE)
c = sense(phase, pos, ang)
r = sense(phase, pos, ang + SENSE_ANGLE)
if l < c < r:
ang += MOVE_ANGLE
elif l > c > r:
ang -= MOVE_ANGLE
elif c < l and c < r:
ang += MOVE_ANGLE * (2 * (ti.random() < 0.5) - 1)
pos += ti.Vector([ti.cos(ang), ti.sin(ang)]) * MOVE_STEP
position[i], heading[i] = pos, ang
# deposit
for i in position:
ipos = position[i].cast(int) % GRID_SIZE
grid[phase, ipos] += 1.0
# diffuse
for i, j in ti.ndrange(GRID_SIZE, GRID_SIZE):
a = 0.0
for di in ti.static(range(-1, 2)):
for dj in ti.static(range(-1, 2)):
a += grid[phase, (i + di) % GRID_SIZE, (j + dj) % GRID_SIZE]
a *= EVAPORATION / 9.0
grid[1 - phase, i, j] = a
def activate(t: ti.f32):
for i, j in ti.ndrange(n, n):
p = ti.Vector([i, j]) / n
p = ti.Matrix.rotation2d(ti.sin(t)) @ (p - 0.5) + 0.5
if ti.taichi_logo(p) == 0:
x[i, j] = 1
def init_mesh():
mesh.nrm[0] = [0, 0, 1]
mesh.pos[N] = [0, 0, 0]
for i in range(N):
a = i / N * t3.tau
mesh.pos[i] = [ti.sin(a), ti.cos(a), 0]
mesh.faces[i] = [[i, 0, 0], [(i + 1) % N, 0, 0], [N, 0, 0]]
def rotate3d(p, axis, ang):
"""Rotates a vector in 3d space, given an axis and angle of rotation.
The vector `axis` should be a unit vector.
See "https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula"
Args:
p (:class:`~taichi.math.vec3`): The 3d vector to rotate.
axis (:class:`~taichi.math.vec3`): Axis of rotation.
ang (float): Angle of rotation, in radians.
Example::
>>> from taichi.math import *
>>> @ti.kernel
>>> def test():
>>> v = vec3(1, 0, 0)
>>> axis = normalize(vec3(1, 1, 1))
>>> print(rotate3d(v, axis, radians(30)))
[0.910684, 0.333333, -0.244017]
Returns:
:class:`~taichi.math.vec3`: The vector after rotation.
"""
ca, sa = ti.cos(ang), ti.sin(ang)
return mix(dot(p, axis) * axis, p, ca) + cross(axis, p) * sa
def paint(t:ti.f32) :
for i, j in pixels:
p = ti.Vector([2*i-Width, 2*j-Height])/min(Width,Height)
#background-color
bcol = ti.Vector([1.0,0.8,0.7-0.07*p[1]]) * (1.0-0.25*p.norm())
#animate
tt = mod(t, 1.5)/1.5
ss = pow(tt, 0.2)*0.5+0.5
ss = 1.0 + ss*0.5*ti.sin(tt*6.2831*3.0+p[1]*0.5)*ti.exp(-4.0*tt)
p *= ti.Vector([0.5, 1.5]) + ss*ti.Vector([0.5, -0.5])
#shape
p[1] -= 0.25
a = ti.atan2(p[0], p[1]) / 3.141593
r = p.norm()
h = abs(a)
d = (13.0*h - 22.0*h*h + 10.0*h*h*h)/(6.0-5.0*h)
#color
s = 0.75 + 0.75*p[0]
s *= 1.0 - 0.4*r
s = 0.3 + 0.7*s
s *= 0.5 + 0.5*pow(1.0-clamp(r/d, 0.0, 1.0), 1.0)
hcol = ti.Vector([1.0, 0.5*r, 0.3])*s
pixels[i,j] = mix(bcol , hcol ,smoothstep(-0.01, 0.01, d-r))
def rot3(axis, ang):
"""Returns the matrix representation of a 3d rotation,
given the axis and angle of rotation.
Args:
axis (:class:`~taichi.math.vec3`): Axis of rotation.
ang (float): Angle of rotation in radians.
Returns:
:class:`~taichi.math.mat3`: 3x3 rotation matrix.
Example::
>>> from taichi.math import *
>>> @ti.kernel
>>> def test():
>>> M = rot3(normalize(vec3(1, 1, 1)), radians(30))
[[0.732051, -0.366025, 0.633975],
[0.633975, 0.732051, -0.366025],
[-0.366025, 0.633975, 0.732051]]
"""
ca, sa = ti.cos(ang), ti.sin(ang)
x, y, z = axis
I = eye(3)
K = mat3([[0, -z, y], [z, 0, -x], [-y, x, 0]])
return I + sa * K + (1.0 - ca) * K @ K
def UniformSampleSphere(self, u1, u2):
z = 1.0 - 2.0 * u1
r = ti.sqrt(ts.clamp(1.0 - z * z, 0.0, 1.0))
phi = 2.0 * 3.1415926 * u2
x = r * ti.cos(phi)
y = r * ti.sin(phi)
return ti.Vector([x, y, z])
