File size: 2,221 Bytes
90ca518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """数组化 + Numba 的「单层」八叉树编码,等价替换
MatrixOctree(max_depth=target_depth+1).build + OctreeFaceTokenBuilder(mc=8).encode_level_tokens(target_depth)
前提(per-row 栈已满足):max_depth = target_depth+1 -> 只有一次「叶子(td+1)->节点(td)」聚合,
不需要多层 neighbor 传播;且 max_children=8 时 reduce_node_children_to_max 是 no-op(活跃 octant ≤ 8)。
去掉 OctreeNode 对象 / dict-of-tuple / set-of-tuple:
* 叶子坐标量化 -> 边表(numpy) -> 节点坐标(leaf>>1)+ octant 局部索引(位运算, 向量化)
* 按节点聚合 intra 8×8 / 活跃掩码 F / 节点-节点邻接(向量化 scatter / np.unique)
* 三角形(3-clique)枚举走 Numba njit(CSR 两指针求交), 这是原 get_level_info 的热点
* inter 8×8 按 (节点,邻居) 分组填充
输出 (raw_t, out_t) 与原实现「逐 slot 解码后网格一致」(顺序可不同;decode 与顺序无关)。
"""
import numpy as np
from numba import njit
from utils.data_utils import discretize
_OCT_W = np.array([1, 2, 4], dtype=np.int64) # octant 局部索引 = x + 2y + 4z
@njit(cache=False)
def _triangles_csr(indptr, indices):
"""前向邻接 CSR(indices[indptr[u]:indptr[u+1]] = 与 u 相连且 >u 的节点, 升序)。
枚举 u<v<w 三角形, 返回 [[u,v,w],...]。两指针求 adj[u]∩adj[v]。"""
out_u = []
out_v = []
out_w = []
n = indptr.shape[0] - 1
for u in range(n):
su, eu = indptr[u], indptr[u + 1]
for pa in range(su, eu):
v = indices[pa]
sv, ev = indptr[v], indptr[v + 1]
# 交集: adj[u] 中 >v 的部分(即 pa+1..eu, 因为升序且后面都 >v) 与 adj[v](sv..ev)
i = pa + 1
j = sv
while i < eu and j < ev:
wi = indices[i]
wj = indices[j]
if wi == wj:
out_u.append(u); out_v.append(v); out_w.append(wi)
i += 1; j += 1
elif wi < wj:
i += 1
else:
j += 1
return (np.array(out_u, dtype=np.int64),
np.array(out_v, dtype=np.int64),
np.array(out_w, dtype=np.int64))
|