python - 不同维度数组的 Numpy getter/setter

Question

如果我想:

如果 x 是一维的，返回 x[I]

如果 x 是二维的，返回 x[i,:]

如果 x 是三维的，返回 x[i, j, :]

...

x 是 numpy 数组

class Foo:
   def __init__(self, x):
       self.x = x  # an n-d array

   def pop_q(q, tuple):
       q = zeros(tuple)

   def set_q(q, i, j=None, qq):
       if q.ndim == 1:
           q[i] = qq
       else if q.ndim == 2:
           q[i, :] = qq
       else if q.ndim == 3:
           q[i, j, :] = qq

   def get_q(q, i, j=None):
       if q.ndim == 1:
           return q[i]
       else if q.ndim == 2:
           return q[i, :]
       else if q.ndim == 3:
           return q[i, j, :]

但是有更好的写法吗？ (可能使用 *args 到任意维数组)

我之所以这样封装是为了保证模块化

Answer 1

我认为这应该可行:

def get_q(self, i, j=None):
    slicing = (i, j, slice(None))[:self.x.ndim]
    return self.x[slicing]

记住slice(None)对应:

所以我创建了一个对应于 (i, j, :) 的元组，但是我只保留了这个元组的第一个 self.x.ndim 元素

如果您想要一个适用于任意数量维度的解决方案:

def get_q(self, *indices):
    slicing = (indices + (slice(None),) * self.x.ndim)[:self.x.ndim]
    return self.x[slicing]

Setter 应该以完全相同的方式工作。