Нанесите маску на несколько строк (синтаксический сахар?)

Question

Я ищу элегантный (или более элегантный) способ кодирования конкретного варианта использования в numpy. Вариант использования - это большой набор данных (поэтому имеет значение эффективность) со 100 + полями, более 1000 строк кода и несколькими разделами кода, где я хотел бы обрабатывать только подмножество полей. Пока я обработка всех наблюдений, это экологически чистый и эффективный в простом 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's masked arrays.

Документация Numpy приводится в более подробной информации, и я бы рекомендовал просмотреть ее перед запуском. Вероятно, это добавит еще немного сложности (вам может понадобиться или не нужно использовать function under 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