将口罩应用于多行(语法糖？)

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]

这并不可怕，不要误解我的意思，但是把它乘以100个变量和数百行代码。有什么方法可以不使用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

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

(顺便说一句，我知道通过熊猫可能有其他的方法，但是现在我更愿意通过numpy来探索我最好的选择。我以后可能会用熊猫标签再问这个问题。)

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.

Answer 2

这样做的一种方法可能是使用Numpy蒙面阵列。

Numpy的文档将深入了解更多细节，我建议在开始之前查看它。它可能会增加更多的复杂性(您可能需要也可能不需要使用numpy.ma)，您必须注意代码中的所有操作都是正确的，但是如果您使用蒙面数组，您可以很容易地定义这样的上下文管理器：

@contextlib.contextmanager
def mask(clause, masked_arrays):
    for ma in masked_arrays:
        ma[~clause] = np.ma.masked
    yield
    for ma in masked_arrays:
        ma.mask = np.ma.nomask

注意，由于掩蔽操作必须就位，所以首先需要处理的数组是蒙面数组。然后你可以很容易地使用它，比如：

# Wrap arrays as masked arrays
mwages = np.ma.array(wages)
mcg = np.ma.array(cg)
... etc ...

with mask(mwages < 50001, [mwages, mcg, mded, ...]):
    # np.ma.compressed returns 1D array of non-masked items
    unmasked_tinc = np.ma.compressed(magi - mded - mexem)

    # or must change the masked one directly