groundTruth_press.py

import taichi as ti
import numpy as np
#ti.init(arch=ti.vulkan)

ti.init(arch=ti.cuda,default_fp=ti.f64)
#ti.init(arch=ti.cpu)

#pressure = 0
pressure = 3e6/2/2/4

n = 128
quad_size = 1.0 / n
dt = 4e-2 / n

gravity = ti.Vector([0, -9.8, 0]) #ti.Vector([0, 0, 0])
spring_Y = 1e4#3e4
dashpot_damping = 3e4#1e4
drag_damping = 3#1

ball_radius = 0.3
ball_center = ti.Vector.field(3, dtype=float, shape=(1, ))
ball_center[0] = [0, 0, 0]

x = ti.Vector.field(3, dtype=float, shape=(n, n))
v = ti.Vector.field(3, dtype=float, shape=(n, n))
pr = ti.Vector.field(3, dtype=float, shape=(n, n))
if_bc = ti.field(int, shape=(n, n)) # you can assign various values to it, for different kinds of bc.

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

bending_springs = False

@ti.kernel
def select_bc():
    for num in ti.grouped(x):
        #if num[0] == 0:
        if num[0] == 0 or num[0] == n-1 or num[1] == 0 or num[1] == n-1 : 
            if_bc[num] = 1
        else:
            if_bc[num] = 0

@ti.kernel
def initialize_mass_points():
    random_offset = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 0.1

    for i, j in x:
        x[i, j] = [
            i * quad_size - 0.5 + random_offset[0], 0.6,
            j * quad_size - 0.5 + random_offset[1]
        ]
        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 ( (j) // 4) % 3 == 0:
            colors[i * n + j] = (126/255, 47/255, 142/255)
        elif ( (j) // 4) % 3 == 1:
            colors[i * n + j] = (119/255, 172/255, 48/255)
        else:
            colors[i * n + j] = (162/255, 20/255, 47/255)

initialize_mesh_indices()

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

else:
    for i in range(-2, 3):
        for j in range(-2, 3):
            if (i, j) != (0, 0) and abs(i) + abs(j) <= 2:
                spring_offsets.append(ti.Vector([i, j]))

spring_neighbors.append( ti.Vector([-1, 0]) )
spring_neighbors.append( ti.Vector([ 0, 1]) )
spring_neighbors.append( ti.Vector([ 1, 0]) )
spring_neighbors.append( ti.Vector([ 0,-1]) )
#spring_neighbors.append( ti.Vector([-1, 0]) )

@ti.kernel
def substep():
    for i in ti.grouped(x):
        v[i] += gravity * dt

    for i in ti.grouped(x):
        force = ti.Vector([0.0, 0.0, 0.0])
        j2 = ti.Vector([0, 0]) 
        for spring_offset in ti.static(spring_offsets): 
            j = i + spring_offset
            if 0 <= j[0] < n and 0 <= j[1] < n: # you may use a bool list to store inner points if not a rectangle shape!
                x_ij = x[i] - x[j]
                v_ij = v[i] - v[j]
                d = x_ij.normalized()
                current_dist = x_ij.norm()
                original_dist = quad_size * float(i - j).norm()
                # Spring force
                force += -spring_Y * d * (current_dist / original_dist - 1)
                # Dashpot damping
                force += -v_ij.dot(d) * d * dashpot_damping * quad_size
        # Pressure
        #pr_=ti.Vector([0.0, 0.0, 0.0])
        pr[i] *=0
        #for pair_ in range(4):
        for j1 in ti.static(spring_neighbors):
            #print(jack[0])
            #j2 = ti.Vector([j1[1],-j1[0]])  #which is faster?
            j2[0]=j1[1]
            j2[1]=-j1[0]
            ij1=i+j1
            ij2=i+j2
            #print(j2)
            #j1 = spring_neighbors[pair_]
            #j2 = spring_neighbors[pair_+1]
            if 0 <= ij1[0] < n and 0 <= ij1[1] < n and 0 <= ij2[0] < n and 0 <= ij2[1] < n : #as above
                x_ij1= x[i]-x[ij1]
                x_ij2= x[i]-x[ij2]
                pr[i] += pressure*ti.math.cross(x_ij1, x_ij2)
                force += pressure*ti.math.cross(x_ij1, x_ij2)
        
        v[i] += force * dt
        #pr[i]=pr_

    for i in ti.grouped(x):
        v[i] *= ti.exp(-drag_damping * dt)
        #### Collison handling
        '''
        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
        '''
        ####  
        #### BC
        if if_bc[i]  == 1:
            v[i] = ti.Vector([0.0, 0.0, 0.0]) 
        ####
        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", 
                        (2000,2000),
                        #(1000, 1000), 
                        show_window = False,
                        vsync=True)
#window2 = ti.ui.Window("Taichi Cloth Simulation on GGUI 2", 
#                        (1000,1000),
#                        #(720, 720), 
#                        show_window = False,
#                        vsync=True)
canvas = window.get_canvas()
#canvas2 = window2.get_canvas()
canvas.set_background_color((1, 1, 1))
scene = ti.ui.Scene()
camera = ti.ui.Camera()

current_t = 0.0
initialize_mass_points()
select_bc()

trainDX = []
trainDV = []
trainDP = []
trainDX.append(x.to_numpy())
trainDV.append(v.to_numpy())
trainDP.append(pr.to_numpy())
#trainData = x.to_numpy().copy()

while window.running:
    if current_t > 1.5: #1.5 
        print("Done simulation!")
        trainDataX = np.stack(trainDX)
        print(trainDataX.dtype)
        print(trainDataX.shape)
        trainDataV = np.stack(trainDV)
        print(trainDataV.dtype)
        print(trainDataV.shape)
        trainDataP = np.stack(trainDP)
        print(trainDataP.dtype)
        print(trainDataP.shape)
        #np.savetxt("trainData.txt",trainData)
        np.savez("trainPressData",dataX=trainDataX,dataV=trainDataV,dataP=trainDataP)
        print("All done!")
        exit(0)
        # Reset
        initialize_mass_points()
        current_t = 0

    for i in range(10):
        substep()
        current_t += dt
        #np.stack( (trainData, x.to_numpy()) )
        trainDX.append(x.to_numpy())
        trainDV.append(v.to_numpy())
        trainDP.append(pr.to_numpy())
    update_vertices()
    print("time:  "+str(current_t))

    camera.position(0, 0.3, 2)
    camera.lookat(0, 0.4, 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)
    
    #canvas2.set_image(window.get_image_buffer_as_numpy())
#    window.show()
    #window2.show()