请问numpy如何简化以下代码？

Question

这绝对是我写过的最蠢的代码……

这是一个3d模型的数据结构:

me.vertices: [v0, v1, ...]
me.edges: [e0, e1, ...]
e0.vertices: [i0, i1]
me.vertice[i0] == v0
v0.co: [x, y, z]

大致流程:

• 遍历所有边
• 获取边的两个顶点座标

想要的结果:

[
    [0, 0, 0], // 边 1 的第一个顶点座标
    [1, 1, 1], // 边 1 的第二个顶点座标
    [1, 1, 1], // 边 2 的第一个顶点座标
    [2, 2, 2], // 边 2 的第二个顶点座标
    ...
]

len_edges = len(me.edges)
verts = me.vertices
len_verts = len_edges * 2
vs = np.zeros((len_verts * 3, ), dtype=np.float32, )

for i in range(len_edges):
    i0, i1 = me.edges[i].vertices
    v0 = verts[i0].co
    v1 = verts[i1].co
    s = i * 6
    vs[s] = v0[0]
    vs[s + 1] = v0[1]
    vs[s + 2] = v0[2]
    vs[s + 3] = v1[0]
    vs[s + 4] = v1[1]
    vs[s + 5] = v1[2]

vs.shape = (-1, 3, )

以上代码功能正常，但太蠢了，目前api有一个语法糖:

collection.foreach_get(attr, some_seq)

# Python equivalent
for i in range(len(seq)):
    some_seq[i] = getattr(collection[i], attr)

这么用就可以:

vs = np.zeros((len_verts * 3, ), dtype=np.float32, )
me.vertices.foreach_get('co', vs)
vs.shape = (-1, 3, )

但我想不到该怎么改……

感谢关注！

Answer 1

import numpy as np

len_edges = len(me.edges)
verts = np.array([v.co for v in me.vertices], dtype=np.float32)
edges_indices = np.array([e.vertices for e in me.edges], dtype=int)

vs = np.zeros((len_edges * 2, 3), dtype=np.float32)

vs[::2] = verts[edges_indices[:, 0]]
vs[1::2] = verts[edges_indices[:, 1]]