在@ti.kernel里的for循环建立不同size的matrix

hello 大家好,
我想在for 循环里面解一个线性返程A x= b, 但是A的列数会变化,size是nX2,x是2X1矩阵 b是长度和A一样的矩阵,数值为1, 我尝试了ti.Matrix.zero, 但是里面的n不能是一个variable。不知道大家有没有什么建议。谢谢。
具体来说我想在SPH 模拟里面实现一个free surface detection 的算法,我用numpy 很快就能做出来,但是想复现到这个SPH 里面。
我用numpy做的实现:

import matplotlib.pyplot as plt 
import numpy as np 
import matplotlib.path as mplPath

#set up grid
x = np.linspace(-10, 10, 60)
y = np.linspace(-10, 10, 60)
x, y = np.meshgrid(x, y)
x, y = x.reshape((x.size)), y.reshape((y.size))
#set up circle points
circle = np.zeros((len(x), 2))
circle[:, 0], circle[:, 1] = x, y
for i in range(len(circle)):
    r = np.sqrt(circle[i, 0]**2 + circle[i, 1]**2)
    if r > 5 or r < 2:
        circle[i, 0], circle[i, 1] = 0, 0
    # if circle[i, 0] > 5 or circle[i, 0] < -5 or circle[i, 1] > 5 or circle[i, 1] < -5:
    #     circle[i, 0], circle[i, 1] = 0, 0
    # elif circle[i, 0] < 2 and circle[i, 0] > -2 and circle[i, 1] < 2 and circle[i, 1] > -2:
    #     circle[i, 0], circle[i, 1] = 0, 0
circle = circle[~np.all(circle == 0, axis=1)]

#calculate normal
spacing = np.min(np.sqrt((circle[0, 0] - circle[1, 0])**2 + (circle[0, 1] - circle[1, 1])**2)) #points distance
sr = 3 * spacing # support radius
scalar = 1.5       #unit normal scalar
theta = 45       #scan area angle
angle_criterion = np.cos(theta*np.pi/180)
normal_points = np.zeros((circle.shape))

free_surface_state = np.ones((circle.shape[0],1))
for i in range(circle.shape[0]):
    vectors = circle - circle[i, :]                                                                                 #calculate the vectors from one point to rest
    vector_distance = np.sqrt(vectors[:, 0]**2 + vectors[:, 1]**2).reshape((vectors.shape[0], 1))                   #calculate the distance from one point to rest
    vectors = vectors[np.all(vector_distance <= sr, axis=1)]                                                        #only keep neighbour distance
    d = np.ones((vectors.shape[0], 1))                                                                              #compute value one matrix for same directions
    normal = np.linalg.inv(vectors.T@vectors)@vectors.T@d                                                           #using normal equation to calculta normal
    unit_normal = - normal/(np.linalg.norm(normal + 1e-10) * scalar+ 1e-10)                                         #normalised normal and make it face outwards
    if np.linalg.norm(normal) == 0:                                                                            #points normal balanced by neighbour points
        free_surface_state[i] = 0
    else:
        normal_points[i, 0], normal_points[i, 1] =  circle[i, :][0] + unit_normal[0], circle[i, :][1] + unit_normal[1]  #normal point for plot use
        vectors_state = np.zeros((len(vectors), 1))                                                                     #setup each vectors' angles to normal direction
        for j in range(len(vectors)):
            vectors_state[j] = np.dot(vectors[j, :], unit_normal) / (np.linalg.norm(vectors[j, :]+1e-10) * np.linalg.norm(unit_normal+1e-10))
        vectors_state = vectors_state[np.all(vectors_state > angle_criterion, axis=1)]
        if len(vectors_state) > 0:
            free_surface_state[i] = 0

#compute surface points
surface_points = np.zeros((circle.shape))
for i in range(circle.shape[0]):
    if free_surface_state[i] == 0:
        surface_points[i] = np.array([0, 0])
    else:
        surface_points[i] = circle[i]
surface_points = surface_points[~np.all(surface_points == 0, axis=1)]

plt.figure(figsize=(10, 10), dpi=80)
plt.scatter(circle[:, 0], circle[:, 1], s=10, facecolors='none', edgecolors='blue')
plt.plot([circle[:, 0], normal_points[:, 0]], [circle[:, 1], normal_points[:, 1]])
# plt.xlim([-10, 10])
# plt.ylim([-10, 10])

plt.figure(figsize=(10, 10), dpi=80)
plt.scatter(circle[:, 0], circle[:, 1], s=10, facecolors='none', edgecolors='blue')
plt.scatter(surface_points[:, 0], surface_points[:, 1], s=10, facecolors='red', edgecolors='red')
# plt.xlim([-10, 10])
# plt.ylim([-10, 10])


plt.show()

我想把这个算法移植到SPH 的solver里所以我会遍历每个particle和每个particle的邻居比如:

for p_i in range(self.ps.particle_num[None]):
   x_i = self.ps.x[p_i]
   vectors = ti.Matrix.zero(self.ps.particle_neighbors_num[p_i], 2) ????
   for j in range(self.ps.particle_neighbors_num[p_i]):
       p_j = self.ps.particle_neighbors[p_i, j]
       x_j = self.ps.x[p_j]
       vectors.append(x_i - x_j) ????

这里vectors长度会不一样,因为邻居个数不一样,然后在第二个for循环需要把算出来的vector append到vectors里去。

添加这两个为matrix支持动态下标的函数:

def mat_read(A: ti.template(), i, j):
  ret = A[0,0]
  for ii in ti.static(range(A.m)):  # 这里m和n可能写反了,如果有问题可以试着调换一下
    for jj in ti.static(range(A.n)):
      if i == ii and j == jj:
        ret = A[ii,jj]
  return ret

def mat_write(A: ti.template(), i, j, val):
  for ii in ti.static(range(A.m)):
    for jj in ti.static(range(A.n)):
      if i == ii and j == jj:
        A[ii,jj] = val

然后

for p_i in range(self.ps.particle_num[None]):
   x_i = self.ps.x[p_i]
   vectors = ti.Matrix.zero(8, 2)  # 8 是一个足够大的数
   vector_len = 0
   for j in range(self.ps.particle_neighbors_num[p_i]):
       p_j = self.ps.particle_neighbors[p_i, j]
       x_j = self.ps.x[p_j]
       val_to_append = x_i - x_j
       mat_write(vectors,vector_len,0, val_to_append[0])
       mat_write(vectors,vector_len,1, val_to_append[1])
       vector_len += 1

感谢你的回答,

这确实是一个很好的思路。但是不知道为什么会有这样的bug:

def mat_write(A, i, j, val):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Function receives 4 argument(s) and 5 provided.

其次,把8改成一个比较大的数之后,数组会有很多【0, 0】 的列, 怎么在taichi里面把这些列去除呢?