Week 5 - PDEs¶

Learning about Tai Chi¶

Neil said we should

this is literally just the example code from the Tai Chi getting started tutorials: https://docs.taichi-lang.org/docs/hello_world

Saved here for reference! -- Launching the GUI from jupyter seems to kill the kernel upon exit fyi.

Julia set example¶

In [ ]:
import taichi as ti
import taichi.math as tm

ti.init(arch=ti.gpu) #tells ti what to run on, if no GPU found, falls back on CPU.

n = 320
pixels = ti.field(dtype=float, shape=(n * 2, n)) #field of shape 'shape' full of type dtype=float

@ti.func #must be called from within a kernel - DOES NOT REQUIRE type hinted args and returns
def complex_sqr(z):  # complex square of a 2D vector
    return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])

@ti.kernel #can be called from anywhere, kicks taichi into JIT. REQUIRES type hinted arguments and returns
def paint(t: float):
    for i, j in pixels:  # Parallelized over all pixels
        c = tm.vec2(-0.8, tm.cos(t) * 0.2)
        z = tm.vec2(i / n - 1, j / n - 0.5) * 2
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] = 1 - iterations * 0.02

gui = ti.GUI("Julia Set", res=(n * 2, n)) #sets title and resolution

for i in range(1000000):
    paint(i * 0.03)
    gui.set_image(pixels)
    gui.show()

primes example¶

In [ ]:
## add
"""Count the prime numbers in the range [1, n]
"""
import taichi as ti
import time

ti.init(arch=ti.gpu)

startTime = time.time()

# Checks if a positive integer is a prime number
@ti.func
def is_prime(n: int):
    result = True
    # Traverses the range between 2 and sqrt(n)
    # - Returns False if n can be divided by one of them;
    # - otherwise, returns True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

# Traverses the range between 2 and n
# Counts the primes according to the return of is_prime()
@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
           count += 1

    return count

print(count_primes(1000000))

executionTime = (time.time() - startTime)
print('Execution time in seconds: ' + str(executionTime))

lcs example¶

In [ ]:
import taichi as ti
import numpy as np

ti.init(arch=ti.cpu)

N = 15000
a_numpy = np.random.randint(0, 100, N, dtype=np.int32)
b_numpy = np.random.randint(0, 100, N, dtype=np.int32)

f = ti.field(dtype=ti.i32, shape=(N + 1, N + 1)) #NxN field

@ti.kernel
def compute_lcs(a: ti.types.ndarray(), b: ti.types.ndarray()) -> ti.i32:
    len_a, len_b = a.shape[0], b.shape[0]
    ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
    for i in range(1, len_a + 1):
        for j in range(1, len_b + 1):
            f[i, j] = ti.max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
                          ti.max(f[i - 1, j], f[i, j - 1]))

    return f[len_a, len_b]

print(compute_lcs(a_numpy, b_numpy))

cloth example¶

In [2]:
## so far, running this from Jupyter has killed the kernel when the window closes lmao so I guess maybe we won't do the visualizations from here. 
import taichi as ti

ti.init(arch=ti.cpu)

n = 128
dt = 4e-2 / n
substeps = int(1 / 60 // dt)

# x is an n x n field consisting of 3D floating-point vectors
# representing the mass points' positions
x = ti.Vector.field(3, dtype=float, shape=(n, n))
# v is an n x n field consisting of 3D floating-point vectors
# representing the mass points' velocities
v = ti.Vector.field(3, dtype=float, shape=(n, n))

# The n x n grid is normalized
# The distance between two x- or z-axis adjacent points
# is 1.0 / n
quad_size = 1.0 / n


num_triangles = (n - 1) * (n - 1) * 2
indices = ti.field(int, shape=num_triangles * 3)
vertices = ti.Vector.field(3, dtype=float, shape=n * n)
colors = ti.Vector.field(3, dtype=float, shape=n * n)

bending_springs = False

ball_radius = 0.3
# Use a 1D field for storing the position of the ball center
# The only element in the field is a 3-dimentional floating-point vector
ball_center = ti.Vector.field(3, dtype=float, shape=(1, ))
# Place the ball center at the original point
ball_center[0] = [0, 0, 0]

# substep() works out the *accumulative* effects
# of gravity, internal force, damping, and collision
# on the mass-spring system

# Gravity is a force applied in the negative direction of the y axis,
# and so is set to [0, -9.8, 0]
gravity = ti.Vector([0, -9.8, 0])

# Elastic coefficient of springs
spring_Y = 3e4
# Damping coefficient caused by
# the relative movement of the two mass points
# The assumption here is:
# A mass point can have at most 12 'influential` points
dashpot_damping = 1e4

# The cloth is modeled as a mass-spring grid. Assume that:
# a mass point, whose relative index is [0, 0],
# can be affected by at most 12 surrounding points
#
# spring_offsets is such a list storing
# the relative indices of these 'influential' points

# Damping coefficient of springs
drag_damping = 1


# The @ti.kernel decorator instructs Taichi to
# automatically parallelize all top-level for loops
# inside initialize_mass_points()
@ti.kernel
def initialize_mass_points():
    # A random offset to apply to each mass point
    random_offset = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 0.1

    # Field x stores the mass points' positions
    for i, j in x:
        # The piece of cloth is 0.6 (y-axis) above the original point
        #
        # By taking away 0.5 from each mass point's [x,z] coordinate value
        # you move the cloth right above the original point
        x[i, j] = [
            i * quad_size - 0.5 + random_offset[0], 0.6,
            j * quad_size - 0.5 + random_offset[1]
        ]
        # The initial velocity of each mass point is set to 0
        v[i, j] = [0, 0, 0]

@ti.kernel
def initialize_mesh_indices():
    for i, j in ti.ndrange(n - 1, n - 1):
        quad_id = (i * (n - 1)) + j
        # 1st triangle of the square
        indices[quad_id * 6 + 0] = i * n + j
        indices[quad_id * 6 + 1] = (i + 1) * n + j
        indices[quad_id * 6 + 2] = i * n + (j + 1)
        # 2nd triangle of the square
        indices[quad_id * 6 + 3] = (i + 1) * n + j + 1
        indices[quad_id * 6 + 4] = i * n + (j + 1)
        indices[quad_id * 6 + 5] = (i + 1) * n + j

    for i, j in ti.ndrange(n, n):
        if (i // 4 + j // 4) % 2 == 0:
            colors[i * n + j] = (0.22, 0.72, 0.52)
        else:
            colors[i * n + j] = (1, 0.334, 0.52)

initialize_mesh_indices()

spring_offsets = []
for i in range(-1, 2):
    for j in range(-1, 2):
        if (i, j) != (0, 0):
            spring_offsets.append(ti.Vector([i, j]))


@ti.kernel
def substep():
    # The for loop iterates over all elements of the field v
    for i in ti.grouped(x):
        v[i] += gravity * dt

    # Traverses the field x as a 1D array
    #
    # The `i` here refers to the *absolute* index
    # of an element in the field x
    #
    # Note that `i` is a 2-dimentional vector here
    for i in ti.grouped(x):
        # Initial force exerted to a specific mass point
        force = ti.Vector([0.0, 0.0, 0.0])
        # Traverse the surrounding mass points
        for spring_offset in ti.static(spring_offsets):
            # j is the *absolute* index of an 'influential' point
            # Note that j is a 2-dimensional vector here
            j = i + spring_offset
            # If the 'influential` point is in the n x n grid,
            # then work out the internal force that it exerts
            # on the current mass point
            if 0 <= j[0] < n and 0 <= j[1] < n:
                # The relative displacement of the two points
                # The internal force is related to it
                x_ij = x[i] - x[j]
                # The relative movement of the two points
                v_ij = v[i] - v[j]
                # d is a normalized vector (its norm is 1)
                d = x_ij.normalized()
                current_dist = x_ij.norm()
                original_dist = quad_size * float(i - j).norm()
                # Internal force of the spring
                force += -spring_Y * d * (current_dist / original_dist - 1)
                # Continues to apply the damping force
                # from the relative movement
                # of the two points
                force += -v_ij.dot(d) * d * dashpot_damping * quad_size

        # Continues to add the velocity caused by the internal forces
        # to the current velocity
        v[i] += force * dt
    
    # Traverse the elements in field v
    for i in ti.grouped(x):
        v[i] *= ti.exp(-drag_damping * dt)
    
    # Traverse the elements in field v
    for i in ti.grouped(x):

        offset_to_center = x[i] - ball_center[0]
        if offset_to_center.norm() <= ball_radius:
            # Velocity projection
            normal = offset_to_center.normalized()
            v[i] -= min(v[i].dot(normal), 0) * normal
        # After working out the accumulative v[i],
        # work out the positions of each mass point
        x[i] += dt * v[i]

@ti.kernel
def update_vertices():
    for i, j in ti.ndrange(n, n):
        vertices[i * n + j] = x[i, j]

window = ti.ui.Window("Taichi Cloth Simulation on GGUI", (1024, 1024),
                      vsync=True)
canvas = window.get_canvas()
canvas.set_background_color((1, 1, 1))
scene = ti.ui.Scene()
camera = ti.ui.Camera()

current_t = 0.0
initialize_mass_points()

while window.running:
    window.get_event()
    if window.is_pressed('q'):
        window.running = False
    if current_t > 1.5:
        # Reset
        initialize_mass_points()
        current_t = 0

    for i in range(substeps):
        substep()
        current_t += dt
    update_vertices()

    camera.position(0.0, 0.0, 3)
    camera.lookat(0.0, 0.0, 0)
    scene.set_camera(camera)

    scene.point_light(pos=(0, 1, 2), color=(1, 1, 1))
    scene.ambient_light((0.5, 0.5, 0.5))
    scene.mesh(vertices,
               indices=indices,
               per_vertex_color=colors,
               two_sided=True)

    # Draw a smaller ball to avoid visual penetration
    scene.particles(ball_center, radius=ball_radius * 0.95, color=(0.5, 0.42, 0.8))
    canvas.scene(scene)
    window.show()
Canceled future for execute_request message before replies were done
The Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details.

PROBLEM SET¶

the code is copied here for reference/readability. b/c the taichi gui kills the kernel I would recommend against running these in jupyter. The code is also saved in the py/ directory.

In [ ]:
import numpy as np
import sounddevice as sd
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt
import time

import taichi as ti
ti.init(arch=ti.cpu)

dx = 0.01 #space step
dt = 0.001 #time step

xmin = 0
xmax = 2*np.pi 

N = int(xmax/dx)

x_range = np.linspace(xmin,xmax,N)
substeps = int(1 / 60 // dt)

v = 1 #speed

gamma = 0.001
gamma0 = -0.05

u0 = ti.field(float, shape=(N,))#np.zeros((N))

# IMPULSE
i_index = int(N/2) #where the impulse is

#u0[i_index] = 0.5

# #not impulse
# i_index = int(N/2) #where the wave center is
# h = 0.5

# for i in range((N/2-5),(N/2+5)):
#     u[i] = 



#compute initial values
u_prev = u0
u_prevprev = ti.field(float, shape=(N,))#np.zeros(N)
for i in range(1,(N-2)):
    u_prevprev[i] = u_prev[i] + 1/2*((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])

u_next = ti.field(float, shape=(N,))#np.zeros(N)
# print(u_prev)
# print(u_prevprev)

#@ti.kernel
#def nextStep(u_prev: ti.types.ndarray(),u_prevprev:ti.types.ndarray()):
# def nextStep(u_prev,u_prevprev):
#     u_next = np.zeros(N)#ti.field(float, shape=(N))#
#     #ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
#     for i in range(1,(N-2)):
#         u_next[i] = -u_prevprev[i] + 2*u_prev[i] +((v**2)*(dx**2)/(dt**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
#     #boundary conditions
#     u_next[0] = 0
#     u_next[N-1] = 0
#     #update priors
#     return u_next

@ti.kernel
def nextStep():
    ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
    for i in range(1,(N-2)):
        #u_next[i] = -1*u_prevprev[i] + 2*u_prev[i] +((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
        #u_next[i] = -1*u_prevprev[i] + 2*u_prev[i] +((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1] + gamma0*((u_prev[i]-u_prevprev[i])/dt))
        u_next[i] = -1*u_prevprev[i] + 2*u_prev[i] +((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1] + gamma*((u_prev[i+1]-2*u_prev[i]+u_prev[i-1] - (u_prevprev[i+1]-2*u_prevprev[i]+u_prevprev[i-1]))/(dt)))
    #boundary conditions
    u_next[0] = 0
    u_next[N-1] = 0


@ti.kernel
def updatePrev():
    for i in range(N-1):
        u_prevprev[i] = u_prev[i]
        u_prev[i] = u_next[i]

# def update(u0,xmax,xmin,dx,dt,N):
#     u_next = np.zeros((N))
#     u_prev = u0
#     u_prevprev = np.zeros(N)
#     for i in range(1,(N-2)):
#         u_prevprev[i] = u_prev[i] + 1/2*((v**2)*(dx**2)/(dt**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
#     while True:
#         for i in range(1,(N-2)):
#             u_next[i] = -u_prevprev[i] + 2*u_prev[i] +((v**2)*(dx**2)/(dt**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
#         u_next[0] = 0
#         u_next[N-1] = 0
#         u_prevprev = u_prev
#         u_prev = u_next
#         yield u_next
# my_iter = update(u0,xmax,xmin,dx,dt,N)


x_range_resize = np.linspace(0,1,N)
offset = ti.field(float,shape=(N,))
offset.fill(0.5)

gui = ti.GUI('click to add impulse', res=(640, 400))
while gui.running:
    if gui.get_event(ti.GUI.PRESS,ti.GUI.RMB):
        mouse_x, mouse_y = gui.get_cursor_pos()
        u_prev[int(mouse_x*N)] = mouse_y*1.5 - 0.5
        u_prevprev[int(mouse_x*N)] = mouse_y*1.5 - 0.5
    for i in range(substeps):
        nextStep()
        updatePrev()
    pos = np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
    
    gui.lines(begin=pos[:(N-2)],end=pos[1:(N-1)],radius=2,color=0x068587)
    gui.show()

#sd.play(pos[:,1], 44100)

tune = pos[:,1] #just the y values.
tunef = fft(tune)
xf = fftfreq(N,dx)
plt.plot(xf,tunef)
plt.title('FFT of final wave')
plt.show()

tunemax = np.abs(tune[0:N//2].copy())
ind = []
for i in range(5):
     ind.append(np.argmax(tunemax))
     tunemax[np.argmax(tunemax)] = 0

print(xf[ind])

#temp_x = np.linspace()
f_s = 44100 #sampling rate
s = 3 #duration in seconds

def resynth(f):
    f_s = 44100 #sampling rate
    s = 3 #duration in seconds
    samples = np.zeros(round(f_s*s))
    W = 2*np.pi*f/f_s #frequency

    for n in range(round(f_s*s)):
            samples[n] = np.sin(W*n)
    return samples

samples = np.zeros(round(f_s*s))
for i in range(5):
    samples = np.add(samples, 10*resynth(30*xf[ind[i]]))


# f = 200 #frequency of signal
# f_s = 44100 #sampling rate
# s = 5 #duration in seconds
# samples = np.zeros(round(f_s*s))
# W = 2*np.pi*f/f_s #frequency

# for n in range(round(f_s*s)):
# 		samples[n] = np.sin(W*n)
		
plt.plot(samples)
plt.title('reconstituted wave, close window to play sound')
plt.show()

sd.play(samples,f_s)

time.sleep(3.0)

sd.stop()
In [ ]:
import numpy as np
import taichi as ti
import sounddevice as sd
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt

ti.init(arch=ti.cpu)

n = 200
dt = 0.1 / n
substeps = int(1 / 60 // dt)


x = ti.Vector.field(2, dtype=float, shape=(n,))

v = ti.Vector.field(2, dtype=float, shape=(n,))

quad_size = 1.0 / n

spring_Y = 10
dashpot_damping = 10
drag_damping = 1

@ti.kernel
def initializeXV():
    for i in x:
        x[i] = [i * quad_size,  0.5]
        v[i] = [0,0]

@ti.kernel
def substep():
    #ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
    # for i in range(1,(n-2)):
    for i in x:
        if i!=0 and i!=(n-1):
            force = ti.Vector([0.0,0.0])
            l0 = ti.Vector([quad_size,0.0])
            fdiff = (x[i+1] - x[i])
            vfdiff = (v[i+1]-v[i])
            force += -spring_Y *(fdiff-l0)
            force += -vfdiff * dashpot_damping 
            bdiff = -x[i-1] + x[i]
            vbdiff = -v[i-1]+v[i]
            force += -spring_Y *(bdiff-l0)
            force += -vbdiff * dashpot_damping 
            #print(force)
            v[i] += force * dt
        elif i==0:
            force = ti.Vector([0.0,0.0])
            l0 = ti.Vector([quad_size,0.0])
            fdiff = (x[i+1] - x[i])
            vfdiff = (v[i+1]-v[i])
            force += -spring_Y *(fdiff-l0)
            force += -vfdiff * dashpot_damping 
        else:
            force = ti.Vector([0.0,0.0])
            l0 = ti.Vector([quad_size,0.0])
            bdiff = -x[i-1] + x[i]
            vbdiff = -v[i-1]+v[i]
            force += -spring_Y *(bdiff-l0)
            force += -vbdiff * dashpot_damping 
    
    for i in x:
        v[i] *= ti.exp(-drag_damping * dt)
    for i in x:
        x[i] += dt * v[i]
        if i ==0:
            x[i] == [0,0.5]
            v[i] == [0,0]
        elif i == (n-1):
            x[i] == [0,0.5]
            v[i] == [0,0]

            

gui = ti.GUI('click to add impulse, with damping', res=(640, 400))

initializeXV()

while gui.running:
    if gui.get_event(ti.GUI.PRESS,ti.GUI.RMB):
        mouse_x, mouse_y = gui.get_cursor_pos()
        #print(x[int(mouse_x*n)])
        x[int(mouse_x*n)] = [x[int(mouse_x*n)][0],mouse_y]

    for i in range(substeps):
        substep()
    pos = x.to_numpy()#np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
    #print(pos)
    gui.lines(begin=pos[:(n-2)],end=pos[1:(n-1)],radius=2,color=0x068587)
    gui.show()

tune = pos[:,1] #just the y values.
tunef = fft(tune)
xf = fftfreq(n,quad_size)[:n//2]
plt.plot(xf,2.0/n * np.abs(tunef[0:n//2]))
plt.show()
In [ ]:
import numpy as np
import taichi as ti

ti.init(arch=ti.cpu)

N=500
D = 1
dt = 0.1

# substeps = int(1 / 60 // dt)
#print(substeps)

L = 1

#print(D*dt/((L/N)**2))
C = dt
u = ti.Vector.field(2, dtype=float, shape=(N,))
uprev = ti.Vector.field(2, dtype=float, shape=(N,))

@ti.kernel
def init():
    for i in u:
        if i==int(N/2):
            u[i] = [i*L/N,1]
            uprev[i] = [i*L/N,1]
        else:
            u[i] = [i*L/N,0]
            uprev[i] = [i*L/N,0]

@ti.kernel
def nextStep():
    for i in u:
        if i == 0:
            u[i][1] = 0
        elif i == N-1:
            u[i][1] = 0
        else:
            u[i][1] = uprev[i][1] +C*(uprev[i+1][1]-2*uprev[i][1]+uprev[i-1][1])
        uprev[i] = u[i]
        #print(u[i])
        


gui = ti.GUI('title', res=(640, 400))

init()
pos = u.to_numpy()#np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
#print(pos)
gui.lines(begin=pos[:(N-2)],end=pos[1:(N-1)],radius=2,color=0x068587)
gui.show()

while gui.running:
    nextStep()
    pos = u.to_numpy()#np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
    #print(pos)
    gui.lines(begin=pos[:(N-2)],end=pos[1:(N-1)],radius=2,color=0x068587)
    gui.show()
In [ ]:
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import taichi as ti

ti.init(arch=ti.cpu)


N = 200
dx = 1
dt = 0.1
D = 1
C = D*dt/(2*dx**2)

# u = np.random.random((N,N))#np.zeros((N,2))
# u_next = u.copy()

#NxN field of size 1 vectors
u = ti.Vector.field(1, dtype=float, shape=(N,N))
uprev = ti.Vector.field(1, dtype=float, shape=(N,N))

@ti.kernel
def init():
    for i,j in u:
        if i == 0 or j == 0 or i == (N-1) or j == (N-1):
            u[i,j] = [0]
        else:
            u[i,j] = [ti.random()]
        uprev[i,j] = u[i,j]


init()
# u_half = u.copy()#np.zeros((N,2))
# u_full = np.zeros((N,N))

# def getStep():
#     for i in range(1,N-2):
#         for j in range(1,N-2):
#             u_half[i,j] = u[i,j] + C*(u_half[i+1,j]-2*u_half[i,j]+u_half[i-1,j]+u[i+1,j]-2*u[i,j]+u[i-1,j])
#     for i in range(1,N-2):
#         for j in range(1,N-2):
#             u_full[i,j] = u_half[i,j] + C*(u_half[i+1,j]-2*u_half[i,j]+u_half[i-1,j]+u_full[i+1,j]-2*u[i,j]+u[i-1,j])
#     return u_half, u_full


#print(u.to_numpy)
un = u.to_numpy()
#print(np.shape(un))
print(un[:,:,0])

def update_frame(n):
    # u_half,u_full = getStep()
    # u = u_full
    un = u.to_numpy()
    return un


fig, axes = plt.subplots()
sns.set_theme(style="white")
cmap = sns.cubehelix_palette(start=0, light=1, as_cmap=True)

# sns.kdeplot(data=un[2],
#         x=un[0], y=un[1],
#         cmap=cmap, fill=True,
#         clip=(-2, 2), cut=10,
#         thresh=0, levels=15,
#         ax=axes,
#     )

sns.heatmap(un[:,:,0],
        cmap=cmap,
        ax=axes,
    )
axes.set_axis_off()

ani = animation.FuncAnimation(
    fig, update_frame, interval=20)

plt.show()