不同维度数组的 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 的工作方式应该完全相同。