python - 越来越多的雪花:优化几何变换

我已经基于here编写了一个NumPy版本的Reiter算法，用于结晶雪花（Reiter Chaos，Solitons and Fractals 23（2005）1111-1119，有兴趣的用户可以使用pure-Python code submitted to ActiveState Code）。

毫不奇怪，实际的细胞自动机生长步骤现在要快得多，但是我无法在代码上进行很多改进，从而实现了几何变换，以将雪花从阵列中的表示“拉直”到正方形图像。那里的NumPy专家是否可以向我展示如何向量化或以其他方式提高其性能？

工作代码为：

import numpy as np
import pylab

# Size of CA field and centre position; maximum number of iterations
nx, ny = 250, 250
cx, cy = nx // 2, ny // 2
maxit = 200

# Snowflake parameters: see C. A. Reiter, Chaos, Solitons and Fractals 23,
# (2005) 1111–1119
alpha, beta, gamma = 1, 0.6, 0.01
# Initialize the CA field with the background water level
ca = beta * np.ones((ny, nx))

# dx, dy displacements from a cell to get to its hexagonal neighbours as
# represented in the ca array
dcell = [(-1,-1), (0,-1), (-1,0), (1,0), (0,1), (1,1)]

# Seed the snowflake at the centre of the field
ca[cx, cy] = 1

for i in range(maxit):
    print('{}/{}'.format(i+1,maxit))
    # Which cells are "receptive"?
    recep = (ca >= 1)
    for i,j in dcell:
        recep |= (np.roll(np.roll(ca, i, 0), j, 1) >= 1)

    # Contrinbutions from receptive and non-receptive cells
    ca_recep = recep * (ca + gamma)
    ca_nonrecep = (~recep) * ca

    # Form the weighted sum of non-receptive cells to the next iteration
    wsum = ca_nonrecep * (1 - alpha / 2)
    wnsum = sum(np.roll(np.roll(ca_nonrecep, i, 0), j, 1) for i,j in dcell)
    wsum += wnsum * alpha / 12
    ca = ca_recep + wsum

# Geometric transformation from CA field to image:
###### This bit is quite slow and could be optimized ######
# cos(-45deg), sin(-45deg), y-scaling factor
sn, cs, sc = np.sin(-np.pi/4), np.cos(-np.pi/4), np.sqrt(3)
# Image height and width; centre position in image
height, width = 2 * ny, 2 * nx
imcy, imcx = height / 2, width / 2
# Image array (0=water, 1=ice)
im = np.zeros((height, width))
for ky in range(height):
    for kx in range(width):
        # Stretch and rotate CA field onto image
        tx, ty = kx - imcx, (ky - imcy) * sc
        tx, ty = tx * cs - ty * sn + imcx, tx * sn + ty * cs + imcy
        # Don't stray outside image bounds
        if 0 <= tx <= width and 0 <= ty <= height:
            c = ca[int((cy-1)*ty/ny), int((cx-1)*tx/nx)]
            if c >= 1:
                im[ky, kx] = 1

pylab.imshow(im, interpolation="nearest")
pylab.show()

我尝试了以下方法：

tx, ty = np.arange(-width/2, width/2), np.arange(-height/2, height/2) * sc
tx, ty = tx * cs - ty * sn + imcx, tx*sn + ty * cs + imcy
im = (ca[((cy-1)*ty/ny).astype(int), ((cx-1)*tx/nx).astype(int)] >= 1)
pylab.imshow(im, interpolation="nearest")
pylab.show()

但只要得到IndexErrors，因为我在ca范围之外建立索引。

最佳答案

令人讨厌的东西，例如六角形网格，但是下面的代码应该照顾好它。

# Indices for the square mesh
ky, kx = np.indices([height, width])

# Map square mesh to hexagonal mesh
tx0 = kx - imcx
ty0 = (ky - imcy) * sc
tx = tx0*cs - ty0*sn + imcx
ty = tx0*sn + ty0*cs + imcy

# Boolean mask to stay within bounds
b = (0 <= tx) & (tx <= width) & (0 <= ty) & (ty <= height)

# Cast to integers and copy values
ci = ((cy-1)*ty[b]/ny).astype(int)
cj = ((cx-1)*tx[b]/nx).astype(int)
im = np.zeros([height, width])
im[b] = ca[ci, cj]

最后一行等效于可能更明显的内容：

im[ky[b],kx[b]] = ca[ci, cj]