python - 对多行应用蒙版(语法糖？)

Question

我正在寻找一种优雅(或更优雅)的方式来在 numpy 中编写特定用例的代码。该用例是一个大型数据集(因此效率很重要)，包含 100 多个字段、超过 1,000 行代码以及多个代码段，其中我只想处理其中的一部分字段。只要我处理所有观察结果，这在普通 numpy 中就是干净且高效的:

wages = np.arange(40000,60000,2000)
cg    = np.arange(0,100000,10000)
ded   = np.repeat([6000,9000],5) 
exem  = np.repeat([2000,4000],5) 

agi  = wages + cg
tinc = agi - ded
tinc = tinc - exem

但是在许多代码小节中，我只想处理观察结果的子集，例如 30 行代码，这是我能想到的最好的方法:

agi  = wages + cg
mask = wages < 50001
tinc = agi

tinc[mask] = agi[mask] - ded[mask]
tinc[mask] = tinc[mask] - exem[mask]

这并不可怕，不要误会我的意思，但是将其乘以数百个变量和数百行代码。有没有办法在不诉诸 cython/numba 循环的情况下执行类似以下操作？

# fake code, just to convey the desired syntax
agi  = wages + cg
tinc = agi

mask( wages < 50001 ):    # i.e. inside a python loop, would be "if wages < 50001:"
   tinc = agi - ded
   tinc = tinc - exem

换句话说，我想定义代码的子部分，并指定完全相同的掩码应应用于代码部分中的每个数组，而无需为每个单个数组显式键入掩码。

(顺便说一句，我知道通过 pandas 可能有一些替代方法，但现在更愿意通过 numpy 探索我的最佳选择。稍后我可能会用 pandas 标签重新问这个问题。)

Answer 1

我不推荐这样做，但是......你可以使用一个非常神奇的上下文管理器来做到这一点。例如:

@contextlib.contextmanager
def masking(namespace, mask):
    # If you don't have a fixed set of maskable variables, make it
    # an instance/global/local variables, like `_names`, or just
    # [name for name, value in namespace.items() if isiinstance(value, np.ndarray)]
    names = 'tinc agi dec exem'.split()
    stash = {name: namespace[name] for name in names}
    for name in names:
        namespace[name] = namespace[name][mask]
    try:
        yield
    finally:
        for name in names:
            namespace[name] = stash[name]

现在你可以这样做:

with masking(globals(), wages < 50001):
    tinc = agi - dec
    tinc = tinc - exem

with masking(self.__dict__, self.wages < 50001):
    self.tinc = self.agi - self.dec
    self.tinc = self.tinc - self.exem

# etc.