我有一个形状为[sampling_size* sampling_size, 2]的2D网格。我使用它在Tensorflow中生成3D表面，如下所示：

def cube(G):

    res = []

    for (X, Y) in G:
        if X >= -1 and X < 1 and Y >= -1 and Y < 1:
            res.append(1.)
        else:
            res.append(0.)

    return tf.convert_to_tensor(res)
    
Z_cube = cube(grid)

cube_2d = tf.reshape(Z_cube, [sampling_size, sampling_size])
plot_surface(X, Y, cube_2d)

下面是另一个例子：

def prism(G):

    res = []

    for (X, Y) in G:
        if X >= -1 and X < 1 and Y >= -1 and Y < 1:
            res.append(X + 1.)
        else:
            res.append(0.)

    return tf.convert_to_tensor(res)

Z_prism = prism(grid)

prism_2d = tf.reshape(Z_prism, [sampling_size, sampling_size])

我的问题是：由于这使用循环，所以这种方法效率不高，生成一个立方体需要10秒。
我想知道是否有人知道一种更有效的矢量化方法来生成这些表面。
编辑：我使用以下代码生成网格

sampling_size = 100
limit = math.pi

def generate_grid(_from, _to, _step):

    range_ = tf.range(_from, _to, _step, dtype=float)

    x, y = tf.meshgrid(range_, range_)

    _x = tf.reshape(x, (-1,1))
    _y = tf.reshape(y, (-1,1))

    return tf.squeeze(tf.stack([_x, _y], axis=-1)), x, y

grid, X, Y = generate_grid(-limit, limit, 2*limit / sampling_size)

而对于标图：

import matplotlib.pyplot as plt
from matplotlib import cm

def plot_surface(X, Y, Z, a = 30, b = 15):
    fig = plt.figure()
    ax = plt.axes(projection='3d')

    ax.plot_surface(X, Y, Z, rstride=3, cstride=3, linewidth=1, antialiased=True,
                    cmap=cm.viridis)
    ax.view_init(a, b)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.show()