Python中的动态编程网格

Question

I want to minimize a cost function over a grid in Python. I have two variables x and y which can be computed as

x[i+1,j+1], y[i+1,j+1] = f(x[i,j], x[i+1,j], x[i,j+1], foo[i,j], bar[i,j])

In other words the grid point (i+1,j+1) depends on two kernels foo and bar , and it's neighboring nodes (i,j+1) (i+1,j) and (i,j). A toy example can be seen below

import numpy as np

N = 20
ivec = np.arange(N)
jvec = np.arange(N)

# Kernels
foo = np.sin(ivec[:,None] * jvec[None,:])
bar = np.cos(ivec[:,None] + jvec[None,:])

# We want to find the total cost for traversing over the matrix
d = np.zeros((N,N))

# And store the optimal path
indices = np.zeros((N,N), "int")

for i in range(N-1):
    for j in range(N-1):

        # Compute all posibilities for reaching current node
        dd = [
            d[i+1,j] + foo[i,j],
            d[i,j+1] + bar[i,j],
            d[i,j] + foo[i,j] * bar[i,j]
        ]

        # And find and store the minimim path
        indices[i+1,j+1] = np.argmin(dd)
        d[i+1,j+1] = dd[indices[i+1,j+1]]

print(d[-1,-1])

However this is a very inefficient solution. Especially as N might be arbitrarily large. So my question is: which is the most / more efficient way to compute this? Using iterators (I have tried np.nditer without great success), or using Numba, or are there some fancy trickery one can do in Numpy? I've started looking into ufuncs and ufunc.accumulate with Numpy but can't immediately see a solution.

Note that in foo , bar and dd will be more complicated than in the toy example.

Answer 1

As you mentioned, using numba could be the easiest way to make your code faster when N is large.

import numba
import numpy as np

@numba.jit(nopython=True)
def dp(foo, bar):
    N = foo.shape[0]

    # We want to find the total cost for traversing over the matrix
    d = np.zeros((N,N))

    # And store the optimal path
    indices = np.zeros((N,N), dtype=np.int64)

    for i in range(N-1):
        for j in range(N-1):

            # Compute all posibilities for reaching current node
            dd = np.array([
                d[i+1,j] + foo[i,j],
                d[i,j+1] + bar[i,j],
                d[i,j] + foo[i,j] * bar[i,j]
            ])

            # And find and store the minimim path
            indices[i+1,j+1] = np.argmin(dd)
            d[i+1,j+1] = dd[indices[i+1,j+1]]
    return d[-1,-1]