Improvements to the power spectrum and zcv module

[boryanah]

The power spectrum code can be optimized and made to run faster. While for now we do not want to enable MPI on it, it is worth implementing multithreading/parallelization improvements. Several quick things include (but we don't need to limit ourselves to these):

• Swapping scipy.fft with pyfftw which is multithreaded.
• Painting the TSC/CIC grid can be done using multithreading straightforwardly.
• For more robust performance (i.e. avoiding numerical roundoff errors), we can do the power spectrum binning in integers of the wavemodes rather than using floats.

On the side of the ZCV module:

• Currently the zenbu calls are very slow, so there is room to improve those.
• Also, it would be great to rewrite the parts that use Joe DeRose's code and add comments to them (i.e. tools_jdr.py)

Finally, it would be helpful to add notebooks with examples of how to use the power spectrum and the ZCV module.

[lgarrison]

On the comparison with nbodykit, in test_power.py the first 11 bins seem to give nearly machine-precision with nbodykit. After that, the differences are a fraction of a percent. This may be where our binning of modes and nbodykit's begin to differ. Possibly changing to integer mode handling will make this more robust, but if nbodykit's not doing that, then it may not help the comparison.

But the agreement of the large-scale modes is very encouraging. It may not test the interlacing and compensation as well as the small-scale modes, though.

[boryanah]

Summary of discussion today:

• Implement pyfftw: can either use it in the "estimate" mode which is easier to use but slower or "wisdom mode." Can see example [here].(https://github.com/abacusorg/abacus/blob/development/Analysis/PowerSpectrum/PowerSpectrum/PowerSpectrum.py)
• Lehman will try to speed up TSC/CIC (maybe with sorting?)
• We can shave off a bunch of memory asks if we do not pre-load mu_box and k_box but compute them on the fly.
• We should make this object-oriented (lots of tedious argument carrying happening right now)
• We should work on the documentation.
• Future: we could make this into its own package (could be faster and have less memory requirements than nbodykit). pypower could import functions directly from abacusutils or abacuspower.

Minimum example for Lehman:

from pathlib import Path

import numpy as np
import pytest

from abacusnbody.hod.power_spectrum import calc_power, get_k_mu_box_edges, calc_pk3d, get_field_fft, get_W_compensated

np.random.seed(300)

# specifications of the power spectrum computation
paste = "TSC"
compensated = interlaced = True
nmesh = 256
nbins_mu = 4
logk = False
k_hMpc_max = 1.
nbins_k = 100

# create synthetic particles
Lbox = 1000.
N = int(1.e7)
x = np.random.rand(N)*Lbox
y = np.random.rand(N)*Lbox
z = np.random.rand(N)*Lbox
tracer_pos = np.vstack((x, y, z)).T
poles = []

# this is if you just want to get the power spectrum

# compute power
k_binc, mu_binc, Pkmu, Nkmu, binned_poles, Npoles = calc_power(x, y, z, nbins_k, nbins_mu, k_hMpc_max, logk,
                                                                    Lbox, paste, nmesh, compensated, interlaced)
print(Pkmu)

# this is what it does under the hood (broken down into the different processes; obviously would be better as object-oriented)

# obtain the field in fourier space
w = None
W = get_W_compensated(Lbox, nmesh, paste, interlaced)
tracer_field_fft = get_field_fft(tracer_pos, Lbox, nmesh, paste, w, W, compensated, interlaced)

# compute power from the fourier field
n_perp = n_los = nmesh # cubic box
k_box, mu_box, k_bin_edges, mu_bin_edges = get_k_mu_box_edges(Lbox, n_perp, n_los, nbins_k, nbins_mu, k_hMpc_max, logk) # for mode counting
pk3d, N3d, binned_poles, Npoles = calc_pk3d(tracer_field_fft, Lbox, k_box, mu_box, k_bin_edges, mu_bin_edges, logk, poles=poles)
print(pk3d)

[boryanah]

A quick comment would be that perhaps we should also enable the computing of the power spectrum in a non-periodic box (as would be the case in a CutSky mock or on the light cone or using data).