python - 高效的布尔磁盘包装掩码

Question

我需要制作一个六边形包装圆盘的面具。下面的代码完成了这项工作，但我觉得它效率不高。我也在学习 python，所以我很想获得一些专家建议，了解如何提高计算效率。

r = 0.01
X, Y = np.mgrid[0:1:1000j, 0:1:1000j]
mask = np.full(X.shape, False)
px, py = np.mgrid[r : 1 : 2 * r * np.sqrt(3), r : 1 + r + np.finfo(float).eps: 2 * r]
px = np.vstack((px, px +  r * np.sqrt(3))) 
py = np.vstack((py, py - r)) 

fig, ax = plt.subplots(figsize= (12, 12), dpi=50)
img = ax.imshow(mask * 1, cmap = 'gray', vmin = 0, vmax = 1, extent = [0, 1, 0, 1])

for i, _ in np.ndenumerate(px):                #is this loop dumb and inefficient?
    C = (X - px[i]) ** 2 + (Y - py[i]) ** 2
    mask = mask | (C < r ** 2) 
    img.set_data(mask * 1)
ax.set_aspect(1)

特别是，有没有办法矢量化 for 循环？

谢谢

标签： pythonnumpyvectorization

Answer 1

创建图案的单个拼贴，然后根据需要水平和垂直重复它可能是有效的：

创建图块：

import numpy as np
import matplotlib.pyplot as plt

r = 0.01
sq3 = np.sqrt(3)
samples = 1000

X, Y = np.mgrid[1:(1 + 2 * sq3):int(sq3 * samples) * 1j, 0:2:samples * 1j]
XY = np.c_[X.flatten(), Y.flatten()]
# coordinates of centers of disks; suffices to take disks of radius 1 here
p = np.array([[1, 1], [(1 + sq3), 0], [(1 + sq3), 2], [(1 + 2 * sq3), 1]])
# compute the distance from each point of XY to each disk center
dist = (XY**2).sum(axis=1).reshape(-1, 1) + (p**2).sum(axis=1) - 2 * (XY @ p.T)
# mask points outside the disks
tile = (np.min(dist, axis=1) < 1).reshape(X.shape)

fig, ax = plt.subplots(figsize=(5, 5))
ax.set_aspect(1)
plt.imshow(tile, extent=[0, 2 * r, r, (1 + 2 * sq3) * r]);

它给：

重复瓷砖：

# the number of times to repeat the tile in the horizontal and vertical directions
h, w = 20, 30
# donwsample the tile as needed and replicate
sample_rate = 10
mask = np.tile(tile[::sample_rate, ::sample_rate], (h, w))
fig, ax = plt.subplots(figsize=(8, 8))
ax.set_aspect(1)
ax.imshow(mask, extent=[0, 2 * r * w, 0, 2 * sq3 * r * h]);

这给出了：