Is there a way to utilize the array indices within a vectorized numpy equation?
Specifically, I have this looping code that sets each value of a 2d array to the distance to some arbitrary center point.
img=np.ndarray((size[0],size[1]))
for x in range(size[0]):
for y in range(size[1]):
img[x,y]=math.sqrt((x-center[0])**2+(y-center[1])**2)
How might I vectorize that?
You can solve this easily using broadcasting:
import numpy as np
size = (64, 64)
center = (32, 32)
x = np.arange(size[0])
y = np.arange(size[1])
img = np.sqrt((x - center[0]) ** 2 + (y[:, None] - center[1]) ** 2)
Some help from Pandas would make this task relatively easy:
import itertools
import pandas as pd
import numpy as np
# get all of the xy pairs
xys = pd.DataFrame(list(itertools.product(range(size[0]), range(size[1]))))
# calculate distance
xys["distance"] = np.sqrt((xys[0] - center[0]) ** 2 + (xys[1] - center[1]) ** 2)
# transform to a 2d array
img = xys.set_index([0, 1])["distance"].unstack()
# if you want just the Numpy array, not a Pandas DataFrame
img.values
Yes, there is.
import numpy as np
size = (6, 4)
center = (3, 2)
img_xy = np.array([[(x, y) for x in range(size[0])] for y in range(size[1])])
img = np.sum((img_xy - center) ** 2, axis=2) ** 0.5
print('\nPlan1:\n', img)
img = np.linalg.norm(img_xy - center, axis=2)
print('\nPlan2:\n', img)
You will get this:
Plan1:
[[3.60555128 2.82842712 2.23606798 2. 2.23606798 2.82842712]
[3.16227766 2.23606798 1.41421356 1. 1.41421356 2.23606798]
[3. 2. 1. 0. 1. 2. ]
[3.16227766 2.23606798 1.41421356 1. 1.41421356 2.23606798]]
Plan2:
[[3.60555128 2.82842712 2.23606798 2. 2.23606798 2.82842712]
[3.16227766 2.23606798 1.41421356 1. 1.41421356 2.23606798]
[3. 2. 1. 0. 1. 2. ]
[3.16227766 2.23606798 1.41421356 1. 1.41421356 2.23606798]]
If you have any question, you could ask me.
