This script calculates points in reciprocal space for hexagonal 2D lattices, then uses the cartesian product from itertools to add each vector from one lattice to all of the vectors of the other in the line
np.array([a+b for a, b in list(itertools.product(p1.T, p2.T))])
It's slow right now because as written its instantiating millions of tiny numpy arrays.
I'm aware of:
- Cartesian product of x and y array points into single array of 2D points
- Using numpy to build an array of all combinations of two arrays
- cartesian products in numPy
and I suspect there's some way to do this, possibly using np.meshgrid
or np.mgrid
that's faster, uses less memory and looks cleaner, but I can not figure out how.
I will use the output in an optimization loop matching these positions to measured positions, so it needs to be callable several hundred times in a row, so reusing large array spaces rather than instantiating and garbage collecting them might have some advantages.
import numpy as np
import matplotlib.pyplot as plt
import itertools
def rotatem(xy, rot):
r3o2, twopi, to_degs, to_rads = np.sqrt(3)/2., 2*np.pi, 180/np.pi, np.pi/180
c, s = [f(to_rads*rot) for f in (np.cos, np.sin)]
x, y = xy
xr = c*x - s*y
yr = c*y + s*x
return np.vstack((xr, yr))
def get_points(a=1.0, nmax=5, rot=0):
r3o2, twopi, to_degs, to_rads = np.sqrt(3)/2., 2*np.pi, 180/np.pi, np.pi/180
g = twopi / (r3o2 * a)
i = np.arange(-nmax, nmax+1)
I, J = [thing.flatten() for thing in np.meshgrid(i, i)]
keep = np.abs(I + J) <= nmax
I, J = [thing[keep] for thing in (I, J)]
xy = np.vstack((I+0.5*J, r3o2*J))
return g * rotatem(xy, rot=rot)
r3o2, twopi, to_degs, to_rads = np.sqrt(3)/2., 2*np.pi, 180/np.pi, np.pi/180
a1, a2, rot = 1.0, 2**0.2, 22
p1 = get_points(a=a1, nmax=20)
p2 = get_points(a=a2, nmax=20, rot=rot)
p3 = get_points(a=a2, nmax=20, rot=-rot)
d12 = np.array([a+b for a, b in list(itertools.product(p1.T, p2.T))])
d13 = np.array([a+b for a, b in list(itertools.product(p1.T, p3.T))])
d12, d13 = [d[((d**2).sum(axis=1)<4.)] for d in (d12, d13)]
if True:
plt.figure()
for d in (d12, d13):
plt.plot(*d.T, 'o', ms=2)
plt.gca().set_aspect('equal')
plt.show()
np.array(a+b for a, b in itertools.product(p1.T, p2.T))
. The same point could be made about the creation ofd12
andd13
in your larger code example. \$\endgroup\$d12 = np.array(a+b for a, b in list(itertools.product(p1.T, p2.T)))
is accepted, the next stepd12 = np.array(d12[((d12**2).sum(axis=1)<4.)])
throws an exception "unsupported operand type(s) for ** or pow(): 'generator' and 'int" \$\endgroup\$