def make_square_stop_origins(side_length: float,
num_rays: int,
x_axis: int = -2,
y_axis: int = -1):
ox = mathx.reshape_vec(np.linspace(-0.5, 0.5, num_rays),
x_axis) * side_length
oy = mathx.reshape_vec(np.linspace(-0.5, 0.5, num_rays),
y_axis) * side_length
return ox, oy
def transform_to_arb(self, E, R, k, axis=None, deriv=0):
"""
Args:
E (array): input vector, with r running along axis
R (scalar): input aperture
k (array): transverse k to which to transform
axis (int): axis along which r runs in E.
deriv (int): derivative order to calculate
Returns:
"""
if axis == None:
axis = last_axis(E)
k = np.asarray(k)
E = expand_dims(E, axis)
# Now E.shape[axis]=1 and r runs along axis-1.
if axis > -k.ndim:
# If k spans axis, then to keep all dimensions aligned need to shift leading axes of k too. New axis in k
# should be new r axis.
k = expand_dims(k, axis - 1)
K = self.conj_R(R)
j = mathx.reshape_vec(self.j, axis - 1)
J1sqd = mathx.reshape_vec(self.J1sqd, axis - 1)
def calc_Et(E, k):
T = 2 * jvp(0, k * j / K, deriv) * ((j / K)**deriv / J1sqd) / K**2
Et = (T * E).sum(axis - 1)
return (Et, )
# shape returns 32 bit int, need 64 bit to avoid overflow
working_size = np.prod(
np.array(mathx.broadcasted_shape(E.shape, k.shape),
dtype=np.int64))
# TODO better available memory test
size_threshold = 1e7
if working_size > size_threshold and k.ndim > 0:
chunk_axis = -k.ndim + np.argmax(k.shape)
num_chunks = working_size / size_threshold
chunk_length = int(math.ceil(k.shape[chunk_axis] / num_chunks))
str = 'QDHT working size %g exceeds %g. Making %d chunks of length %d along axis %d ...'
logger.info(3, str, working_size, size_threshold, num_chunks,
chunk_length, chunk_axis)
# BUG! If E runs along chunk_axis too, then this fails.
# Et=mathx.eval_array_fun_chunked(calc_Et_from_k,k,chunk_axis,chunk_length)
# Et=mathx.iterate_broadcast_op(calc_Et,(E,k),chunk_axis,chunk_length)
Et = mathx.eval_iterated(calc_Et, (E, k),
iter_dims=(chunk_axis, ),
keep_iter_dims=True,
iter_chunk_size=chunk_length,
print_progress=logger.level < 3)[0]
logger.info(3, '... QDHT done')
return Et
else:
return calc_Et(E, k)[0]
def make_square_field_vectors(side_length: float,
num_spots: int,
x_axis: int = -2,
y_axis: int = -1):
# Generate square grid of plane wave vectors along -1 and -2 axes.
thetax = mathx.reshape_vec(np.linspace(-0.5, 0.5, num_spots),
x_axis) * side_length
thetay = mathx.reshape_vec(np.linspace(-0.5, 0.5, num_spots),
y_axis) * side_length
vx = np.sin(thetax)
vy = np.sin(thetay)
vz = (1 - vx**2 - vy**2)**0.5
return (thetax, thetay), (vx, vy, vz)
def pad(v, b, a, axis=None):
if axis is None:
axis = get_axis(v)
l = np.shape(v)[axis]
vp = mathx.reshape_vec(np.arange(-b, l + a), axis) * delta(v, axis) + zero(
v, axis)
return vp
def interp1d_to_uniform(x, y, axis=None):
"""Resample array to uniformly sampled axis.
Has some limitations due to use of scipy interp1d.
Args:
x (vector): independent variable
y (array): dependent variable, must broadcast with x
axis (int): axis along which to resample
Returns:
xu: uniformly spaced independent variable
yu: dependent resampled at xu
"""
x = np.asarray(x)
y = np.asarray(y)
if axis is None:
axis = mathx.vector_dim(x)
num = x.shape[axis]
mn = x.min(axis, keepdims=True)
mx = x.max(axis, keepdims=True)
# Limitation of scipy interp1d
x = x.squeeze()
mn = mn.squeeze()
mx = mx.squeeze()
assert x.ndim == 1
xu = np.arange(num) / (num - 1) * (mx - mn) + mn
yu = scipy.interpolate.interp1d(x.squeeze(),
y,
axis=axis,
bounds_error=False)(xu)
return mathx.reshape_vec(xu, axis), yu
def scale_fac(self, R=1, dim=-1):
"""Parseval's theorem scaling vector.
The scaling vector s_n such that
sum_{n=0}^{N-1} abs(f(r_n))**2 s_n
equals the norm-squared of the signal f(r_n), where r_n is given by QDHT.points.
"""
return 4 * pi * R**2 / reshape_vec(self.j_Np1**2 * self.J1sqd, dim)
def array(**kwargs):
"""
Args:
axis: dimension
zero: first element
last: last element
delta: spacing
N: number of samples
Returns:
numpy.ndarray
"""
axis = kwargs.pop('axis', -1)
if 'vector' in kwargs:
a = mathx.reshape_vec(kwargs.pop('vector'), axis)
elif 'array' in kwargs:
a = kwargs.pop('array')
else:
if all(k in kwargs for k in ('zero', 'N', 'delta')):
zero = np.array(kwargs.pop('zero'))
N = kwargs.pop('N')
delta = np.array(kwargs.pop('delta'))
elif all(k in kwargs for k in ('zero', 'last', 'delta')):
zero = np.array(kwargs.pop('zero'))
last = np.array(kwargs.pop('last'))
delta = kwargs.pop('delta')
N = set(round((last - zero) / delta).flatten().tolist())
if len(N) != 1:
raise ValueError('N must be same for all sub-vectors')
N = N.pop()
elif all(k in kwargs for k in ('zero', 'N', 'last')):
zero = np.array(kwargs.pop('zero', 0))
last = np.array(kwargs.pop('last'))
N = kwargs.pop('N')
delta = (last - zero) / (N - 1)
else:
raise ValueError('Unknown argument combination ' +
str(list(kwargs.keys())))
a = zero + mathx.reshape_vec(np.arange(N), axis) * delta
if len(kwargs) > 0:
raise ValueError('Unused arguments ' + str(list(kwargs.keys())))
return a
def pad_sampled(v, f, b, a, axis=None, value=0):
f = np.asarray(f)
if axis is None:
axis = get_axis(v)
shape_a = list(f.shape)
shape_a[axis] = a
shape_b = list(f.shape)
shape_b[axis] = b
fp = np.concatenate(
(value * np.ones(shape_b), f, value * np.ones(shape_a)), axis)
vp = mathx.reshape_vec(np.arange(-b, f.shape[axis] + a), axis) * delta(
v, axis) + zero(v, axis)
return vp, fp
def conj_axis(t, offset=0, offset_type='middle', axis=None):
if axis is None:
axis = mx.last_axis(t)
Dt = np.diff(np.take(t, [0, 1], axis), axis=axis)
N = t.shape[axis]
Dw = 2 * pi / (N * Dt)
if offset_type == 'middle':
n_o = N / 2
elif offset_type == 'middle_index':
n_o = int(N / 2)
elif offset_type == 'start':
n_o = 0
elif offset_type == 'end':
n_o = N - 1
else:
raise ValueError('Unknown offset_type ', offset_type)
n = mx.reshape_vec(np.arange(N), axis) - n_o
return offset + n * Dw
def test_arb_trans():
##
ht = QDHT(64)
R = 5
r = ht.points(R)
k = ht.conj_points(R)
Er = exp(-r**2 / 2)
Ek = ht.transform(Er, R)
Eka = ht.transform_to_arb(Er, R, k)
assert np.allclose(Ek, Eka)
Eka = ht.transform_to_arb(Er, R, mathx.reshape_vec(k, -3))
assert np.allclose(Ek, Eka.squeeze()) and Eka.shape == (64, 1, 1)
R = ht.j_Np1**0.5
r = ht.points(R)
k = ht.conj_points(R)
Er = exp(-r**2 / 2)
Erp = -r * exp(-r**2 / 2)
plt = pg.plot(r, Erp, pen=pg.mkPen('b', width=10))
plt.plot(k, ht.transform_to_arb(Er, R, k, deriv=1), pen='r')
def test_arb_trans():
##
ht=QDHT(64)
R=5
r=ht.points(R)
k=ht.conj_points(R)
Er=np.exp(-r**2/2)
Ek=ht.transform(Er,R)
Eka=ht.transform_to_arb(Er,R,k)
assert np.allclose(Ek,Eka)
Eka=ht.transform_to_arb(Er,R,mathx.reshape_vec(k,-3))
assert np.allclose(Ek,Eka.squeeze()) and Eka.shape==(64,1,1)
Era=ht.inv_transform_to_arb(Ek,R,r)
assert np.allclose(Er,Era)
R=ht.j_Np1**0.5 # for same r & k axes
r=ht.points(R)
k=ht.conj_points(R)
Er=np.exp(-r**2/2)
Erp=-r*np.exp(-r**2/2)
Ekp=ht.transform_to_arb(Er,R,k,deriv=1)
assert np.allclose(Erp,Ekp)
def __init__(self, sign=-1, axis=None, **kwargs):
"""N is number of samples, x_0 is first x sample, Dx is x spacing, i.e. x_n=x_0+n*Dx,
X is N*Dx. x_m is the middle = x_0+(N-1)*Dx/2. All of these are defined
similarly for k. Can specify any combination of them with sufficient information
to define the sampling axes. If unspecified, x_0 and k_0 default to 0."""
for key in kwargs.keys():
if key not in ('N', 'x', 'x_0', 'Dx', 'x_m', 'x_middle_ind', 'X',
'k', 'k_0', 'Dk', 'k_m', 'k_middle_ind', 'K'):
raise ValueError('Unknown argument %s' % key)
self.sign = sign
if axis is not None:
if axis >= 0:
raise TypeError(
'axis must be negative (because numpy broadcasts starting from trailing dimensions)'
)
self.axis = axis
elif 'x' in kwargs:
self.axis = mx.last_axis(kwargs['x'])
elif 'k' in kwargs:
self.axis = mx.last_axis(kwargs['k'])
else:
self.axis = -1
# Determine N
if 'N' in kwargs:
self.N = kwargs['N']
elif 'x' in kwargs:
self.N = kwargs['x'].shape[self.axis]
elif 'k' in kwargs:
self.N = kwargs['k'].shape[self.axis]
elif 'Dx' in kwargs and 'X' in kwargs:
self.N = math.ceil(kwargs['X'] / kwargs['Dx'])
elif 'Dk' in kwargs and 'K' in kwargs:
self.N = math.ceil(kwargs['K'] / kwargs['Dk'])
elif 'Dx' in kwargs and 'Dk' in kwargs:
self.N = math.ceil(2 * pi / (kwargs['Dx'] * kwargs['Dk']))
elif 'X' in kwargs and 'K' in kwargs:
self.N = math.ceil(kwargs['X'] * kwargs['K'] / (2 * pi))
else:
raise ValueError('N undetermined by arguments')
self.N = np.unique(self.N)
if len(self.N) > 1:
raise ValueError('Only one N allowed')
self.N = self.N[0]
# Determine Dx
Dx = None
if 'x' in kwargs:
Dx = np.diff(np.take(kwargs['x'], [0, 1], self.axis),
axis=self.axis)
elif 'Dx' in kwargs:
Dx = kwargs['Dx']
elif 'X' in kwargs:
Dx = kwargs['X'] / self.N
if Dx is not None:
self.Dx = np.array(Dx)
self.Dk = 2 * pi / (self.N * self.Dx)
else:
if 'k' in kwargs:
Dk = np.diff(np.take(kwargs['k'], [0, 1], axis=self.axis),
axis=self.axis)
elif 'Dk' in kwargs:
Dk = kwargs['Dk']
elif 'K' in kwargs:
Dk = kwargs['K'] / self.N
else:
raise ValueError('Dx and Dk undetermined by arguments')
self.Dk = np.asarray(Dk)
self.Dx = 2 * pi / (self.N * self.Dk)
# Determine x_0
if 'x' in kwargs:
x_0 = np.take(kwargs['x'], [0], axis=self.axis)
elif 'x_0' in kwargs:
x_0 = kwargs['x_0']
elif 'x_m' in kwargs:
x_0 = kwargs['x_m'] - (self.N - 1) * self.Dx / 2
elif 'x_middle_ind' in kwargs:
x_0 = kwargs['x_middle_ind'] - round(self.N / 2) * self.Dx
else:
x_0 = 0
self.x_0 = np.asarray(x_0)
# Determine k_0
if 'k' in kwargs:
k_0 = np.take(kwargs['k'], [0], axis=self.axis)
elif 'k_0' in kwargs:
k_0 = kwargs['k_0']
elif 'k_m' in kwargs:
k_0 = kwargs['k_m'] - (self.N - 1) * self.Dk / 2
elif 'k_middle_ind' in kwargs:
k_0 = kwargs['k_middle_ind'] - round(self.N / 2) * self.Dk
else:
k_0 = 0
self.k_0 = np.asarray(k_0)
# Setup arrays
n = mx.reshape_vec(np.arange(self.N), self.axis)
self.x = self.x_0 + self.Dx * n
self.k = self.k_0 + self.Dk * n
self.X = self.Dx * self.N
self.K = self.Dk * self.N
self.x_fac = np.exp(1j * self.sign * self.k_0 * self.Dx * n)
self.k_fac = np.exp(1j * self.sign * self.x_0 * self.Dk * n)
self.fac = np.exp(1j * self.sign * self.x_0 * self.k_0) * sqrt(
self.N / (2 * pi))
def polar_grid(Ntheta, Nphi, dim_theta=-2, dim_phi=-1):
theta = np.pi*(mathx.reshape_vec(np.arange(Ntheta), dim_theta) + 1)/(Ntheta + 1)
phi = np.pi*2*mathx.reshape_vec(np.arange(Nphi), dim_phi)/Nphi
return theta, phi
def points(self, R=1, dim=-1):
"""Compute r sampling points."""
return reshape_vec(self.j, dim) / self.j_Np1 * R