def get_k_bounds(data, spacing, packed=True):
"""
Return the bounds of wavenumber values for the specified grid.
"""
nx, ny, nz = transform.expanded_shape(data, packed=packed)
k0 = (2 * np.pi) / spacing
k_min = k0 / max(nx, ny, nz)
k_max = k0 * np.sqrt(3) / 2
return k_min, k_max
def create_ksq_grids(data, spacing, packed):
nx, ny, nz = transform.expanded_shape(data, packed=packed)
lambda0 = spacing / (2 * np.pi)
kx = np.fft.fftfreq(nx, lambda0)
ky = np.fft.fftfreq(ny, lambda0)
kz = np.fft.fftfreq(nz, lambda0)
if packed:
kz = kz[:nz//2 + 1]
# With the sparse option, the memory usage of these grids is
# O(nx+ny+nz) rather than O(nx*ny*nz).
return np.meshgrid(kx**2, ky**2, kz**2, sparse=True, indexing='ij')
def tabulate_sigmas(data, power, spacing, packed=True):
"""
Replace an array of log10(k) values with the corresponding sigmas.
Note that the scaling from P(k) to variance depends on convention for
normalizing the inverse FFT. Since we divide the inverse FFT by
nx * ny * nz, the appropriate scaling here is::
sigma**2 = (nx * ny * nz) * P(k) / (2 * Vbox)
= P(k) / (2 * spacing**3)
"""
validate_power(power)
nx, ny, nz = transform.expanded_shape(data, packed=packed)
N3 = nx * ny * nz
Vbox = N3 * spacing**3
power_k_min, power_k_max = np.min(power['k']), np.max(power['k'])
if power_k_min <= 0:
raise ValueError('Power uses min(k) <= 0: {0}.'.format(power_k_min))
data_k_min, data_k_max = get_k_bounds(data, spacing=spacing, packed=packed)
if power_k_min > data_k_min or power_k_max < data_k_max:
raise ValueError(
'Power k range [{0}:{1}] does not cover data k range [{2}:{3}].'
.format(power_k_min, power_k_max, data_k_min, data_k_max))
# Build an interpolater of sigma(|k|) that is linear in log10(|k|).
log10_k = np.log10(power['k'])
sigma = N3 * np.sqrt(power['Pk'] / (2 * Vbox))
interpolator = scipy.interpolate.interp1d(
log10_k, sigma, kind='linear', copy=False,
bounds_error=False, fill_value=0)
# Calculate interpolated values of sigma.
# We would ideally do this in place, but scipy.interpolate methods
# do not support this. A slower but memory efficient alternative would
# be to interpolate in batches. For now we just do the simplest thing.
data.real = interpolator(data.real)
return data
