How to generate a numpy Manhattan distance array with torus geometry fast?

Viewed 300

I have a N x M field and I'd like to create a matrix with the Manhattan distances to a particular location. Additionally, the field wraps around on all ends (torus).

Can I generate this in Numpy fast (without doing a slow flood fill loop)?

For example the function would return

f(x=2, y=3, size_x=8, size_y=8) ->

array([[5, 4, 3, 2, 3, 4, 5, 6],
       [4, 3, 2, 1, 2, 3, 4, 5],
       [3, 2, 1, 0, 1, 2, 3, 4],
       [4, 3, 2, 1, 2, 3, 4, 5],
       [5, 4, 3, 2, 3, 4, 5, 6],
       [6, 5, 4, 3, 4, 5, 6, 7],
       [7, 6, 5, 4, 5, 6, 7, 8],
       [6, 5, 4, 3, 4, 5, 6, 7]])

Here position (2,3) has distance zero and all other positions the Manhattan/taxi-cab distance to (2,3) while being able to wrap around at the edges.

2 Answers

You can create templates for shifts in x and y axes with numpy.roll and then use numpy.add.outer to add them axes-wise.

import numpy as np

def f(x, y, size_x, size_y):
    """
    >>> f(x=2, y=3, size_x=8, size_y=8)
    array([[5, 4, 3, 2, 3, 4, 5, 6],
           [4, 3, 2, 1, 2, 3, 4, 5],
           [3, 2, 1, 0, 1, 2, 3, 4],
           [4, 3, 2, 1, 2, 3, 4, 5],
           [5, 4, 3, 2, 3, 4, 5, 6],
           [6, 5, 4, 3, 4, 5, 6, 7],
           [7, 6, 5, 4, 5, 6, 7, 8],
           [6, 5, 4, 3, 4, 5, 6, 7]])
    >>> f(x=1, y=1, size_x=3, size_y=3)
    array([[2, 1, 2],
           [1, 0, 1],
           [2, 1, 2]])
    """
    a, b = divmod(size_x, 2)
    x_template = np.r_[:a+b, a:0:-1] # [0 1 2 1] for size_x == 4 and [0 1 2 2 1] for size_x == 5
    x_template = np.roll(x_template, x) # for x == 2, size_x == 8: [2 1 0 1 2 3 4 3]
    a, b = divmod(size_y, 2)
    y_template = np.r_[:a+b, a:0:-1]
    y_template = np.roll(y_template, y)
    return np.add.outer(x_template, y_template)

There might be something faster but this should work:

def f(x=2, y=3, X=8, Y=8):
    return np.array([[min(abs(y-i), abs(Y+y-i))+min(abs(x-j), abs(X+x-j)) for i in range(Y)] for j in range(X)])
Related