numpy中复杂赋值逻辑的向量化



我在模拟中有一些复杂的赋值逻辑,我想优化性能。当前的逻辑是在各种numpy数组上作为一组嵌套的for循环实现的。我想对这个赋值逻辑进行矢量化,但不知道这是否可行

import numpy as np
from itertools import izip
def reverse_enumerate(l):
    return izip(xrange(len(l)-1, -1, -1), reversed(l))
materials = np.array([[1, 0, 1, 1],
                     [1, 1, 0, 0],
                     [0, 1, 1, 1],
                     [1, 0, 0, 1]])
vectors = np.array([[1, 1, 0, 0],
                   [0, 0, 1, 1]])
prices = np.array([10, 20, 30, 40])
demands = np.array([1, 1, 1, 1])
supply_by_vector = np.zeros(len(vectors)).astype(int)
#go through each material and assign it to the first vector that the material covers
for m_indx, material in enumerate(materials):
    #find the first vector where the material covers the SKU
    for v_indx, vector in enumerate(vectors):
        if (vector <= material).all():
            supply_by_vector[v_indx] = supply_by_vector[v_indx] + 1
            break
original_supply_by_vector = np.copy(supply_by_vector)
profit_by_vector = np.zeros(len(vectors))
remaining_ask_by_sku = np.copy(demands)
#calculate profit by assigning material from vectors to SKUs to satisfy demand
#go through vectors in reverse order (so lowest priority vectors are used up first)
profit = 0.0
for v_indx, vector in reverse_enumerate(vectors):
    for sku_indx, price in enumerate(prices):
        available = supply_by_vector[v_indx]
        if available == 0:
            continue
        ask = remaining_ask_by_sku[sku_indx]
        if ask <= 0:
            continue
        if vector[sku_indx]:
            assign = ask if available > ask else available
            remaining_ask_by_sku[sku_indx] = remaining_ask_by_sku[sku_indx] - assign
            supply_by_vector[v_indx] = supply_by_vector[v_indx] - assign
            profit_by_vector[v_indx] = profit_by_vector[v_indx] + assign*price
            profit = profit + assign * price
print 'total profit:', profit
print 'unfulfilled demand:', remaining_ask_by_sku
print 'original supply:', original_supply_by_vector
结果:

total profit: 80.0
unfulfilled demand: [0 1 0 0]
original supply: [1 2]

似乎在嵌套循环的第二部分/组中最内层嵌套循环中的迭代之间存在依赖关系,这对我来说似乎很难(如果不是不可能的话)向量化。所以,这篇文章基本上是一个部分解决方案,试图对第一组两个嵌套循环进行矢量化,它们是-

supply_by_vector = np.zeros(len(vectors)).astype(int)
for m_indx, material in enumerate(materials):
    #find the first vector where the material covers the SKU
    for v_indx, vector in enumerate(vectors):
        if (vector <= material).all():
            supply_by_vector[v_indx] = supply_by_vector[v_indx] + 1
            break

整个部分可以用一行向量化代码代替,像这样-

supply_by_vector = ((vectors[:,None] <= materials).all(2)).sum(1)

最新更新