关于 `static for` 展开 `if` 的问题

当我想写类似下面的语句时,我可以用 static for

if obj.shape == 0:
    n = func0(p)
elif obj.shape == 1:
    n = func1(p)
else:
    ...
  • static for 版本:
for shape in ti.static(SHAPES):
    if obj.shape == shape:
        n = SHAPE_FUNC[shape](p)

但是它展开后,是没有 else 的,会造成很多次无效判断

有没有办法,解决这样的问题,而不用手写一大坨 if-else ?

PS: obj 是在并行时的一个变量,需要根据 obj.shape 的不同选择不同的函数执行

补充一下实例:

@ti.func
def normal(sdf_func: ti.template(), obj: SDFObject, p: vec3) -> vec3:
    pos, scale = calc_pos_scale(obj, p)
    n, h = vec3(0), 0.5773 * 0.005

    # from https://iquilezles.org/articles/normalsSDF/
    for i in ti.static(range(4)):
        e = 2.0*vec3((((i+3) >> 1) & 1), ((i >> 1) & 1), (i & 1))-1.0
        n += e*sdf_func(pos+e*h, scale)

    return normalize(n)
@ti.func
def calc_normal_1(obj: SDFObject, p: vec3) -> vec3:
    n = vec3(0)

    if obj.type == SHAPE.SPHERE:
        n = normal(SHAPE_FUNC[SHAPE.SPHERE], obj, p)
    elif obj.type == SHAPE.CYLINDER:
        n = normal(SHAPE_FUNC[SHAPE.CYLINDER], obj, p)
    elif obj.type == SHAPE.BOX:
        n = normal(SHAPE_FUNC[SHAPE.BOX], obj, p)
    elif obj.type == SHAPE.CYLINDER:
        n = normal(SHAPE_FUNC[SHAPE.CYLINDER], obj, p)

    # n = normal(SHAPE_FUNC[obj.type], obj, p)
    return n
@ti.func
def calc_normal_2(obj: SDFObject, p: vec3) -> vec3:
    n = vec3(0)
    for shape in ti.static(SHAPES):
        if obj.type == shape:
            n = normal(SHAPE_FUNC[shape], obj, p)
            # break
    return n
    ### 要是有优雅的解决方案就好了

目前还没有支持这样的feature的计划,可能得先用if else来处理

如果有个python运行时,taichi编译时的static switch就好了