diff --git a/CMakeLists.txt b/CMakeLists.txt index d8b4c96..c15f7d7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -53,6 +53,7 @@ ENDIF(YORICK_SUPPORT) find_program(CYTHON cython) find_library(ZLIB z) find_library(DL_LIBRARY dl) +find_library(MATH_LIBRARY m) set(NETCDF_FIND_REQUIRED ${YORICK_SUPPORT}) set(GSL_FIND_REQUIRED TRUE) diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index f054f0b..49e0da1 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -36,6 +36,7 @@ add_library(_cosmo_power MODULE ${CMAKE_CURRENT_BINARY_DIR}/_cosmo_power.cpp) add_library(_cosmo_cic MODULE ${CMAKE_CURRENT_BINARY_DIR}/_cosmo_cic.cpp) add_library(_fast_interp MODULE ${CMAKE_CURRENT_BINARY_DIR}/_fast_interp.cpp) add_library(_project MODULE ${CMAKE_CURRENT_BINARY_DIR}/_project.cpp) +add_library(_cosmo_bispectrum MODULE _cosmo_bispectrum.cpp) SET(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Bsymbolic-functions") @@ -44,6 +45,7 @@ target_link_libraries(_cosmo_power ${CosmoTool_local} ${PYTHON_LIBRARIES} ${GSL_ target_link_libraries(_cosmo_cic ${CosmoTool_local} ${PYTHON_LIBRARIES} ${GSL_LIBRARIES}) target_link_libraries(_project ${PYTHON_LIBRARIES}) target_link_libraries(_fast_interp ${CosmoTool_local} ${PYTHON_LIBRARIES}) +target_link_libraries(_cosmo_bispectrum ${MATH_LIBRARY}) # Discover where to put packages if (NOT PYTHON_SITE_PACKAGES) @@ -69,7 +71,7 @@ if (WIN32 AND NOT CYGWIN) SET_TARGET_PROPERTIES(_cosmotool PROPERTIES SUFFIX ".pyd") endif (WIN32 AND NOT CYGWIN) -INSTALL(TARGETS _cosmotool _project _cosmo_power _cosmo_cic _fast_interp +INSTALL(TARGETS _cosmotool _project _cosmo_power _cosmo_cic _fast_interp _cosmo_bispectrum LIBRARY DESTINATION ${PYTHON_SITE_PACKAGES}/cosmotool ) diff --git a/python/_cosmo_bispectrum.cpp b/python/_cosmo_bispectrum.cpp new file mode 100644 index 0000000..72a14a4 --- /dev/null +++ b/python/_cosmo_bispectrum.cpp @@ -0,0 +1,194 @@ +#include +#include +#include +#include +#include + +using std::cout; +using std::endl; +using boost::format; + + +#if defined _WIN32 || defined __CYGWIN__ + #ifdef BUILDING_DLL + #ifdef __GNUC__ + #define DLL_PUBLIC __attribute__ ((dllexport)) + #else + #define DLL_PUBLIC __declspec(dllexport) // Note: actually gcc seems to also supports this syntax. + #endif + #else + #ifdef __GNUC__ + #define DLL_PUBLIC __attribute__ ((dllimport)) + #else + #define DLL_PUBLIC __declspec(dllimport) // Note: actually gcc seems to also supports this syntax. + #endif + #endif + #define DLL_LOCAL +#else + #if __GNUC__ >= 4 + #define DLL_PUBLIC __attribute__ ((visibility ("default"))) + #define DLL_LOCAL __attribute__ ((visibility ("hidden"))) + #else + #define DLL_PUBLIC + #define DLL_LOCAL + #endif +#endif + + + +struct ModeSet +{ + size_t N1, N2, N3; + + struct TriangleIterator + { + ssize_t i1, i2, i3; + size_t N1, N2, N3; + + + TriangleIterator& operator++() { + i3++; + if (i3==N3) { i3 = 0; i2++; } + if (i2==N2) { i2 = 0; i1++; } + return *this; + } + + bool operator!=(const TriangleIterator& t) const { + return i1!=t.i1 || i2!=t.i2 || i3 != t.i3; + } + + bool in_box() const { + ssize_t hN1 = N1/2, hN2 = N2/2, hN3 = N3/2; + + return (i1 >= -hN1) && (i1 < hN1) && (i2 >= -hN2) && (i2 < hN2) && (i3 >= -hN3) && (i3 < hN3); + } + + TriangleIterator operator+(const TriangleIterator& other_t) const { + TriangleIterator t = *this; + + t.i1 = (t.i1+other_t.i1); + t.i2 = (t.i2+other_t.i2); + t.i3 = (t.i3+other_t.i3); + return t; + } + + TriangleIterator real() const { + TriangleIterator t = *this; + if (t.i1 >= N1/2) + t.i1 -= N1; + if (t.i2 >= N2/2) + t.i2 -= N2; + return t; + } + + double norm() const { + double r1 = i1, r2 = i3, r3 = i3; + return std::sqrt(r1*r1 + r2*r2 + r3*r3); + } + + TriangleIterator& operator*() { return *this; } + }; + + ModeSet(size_t N1_, size_t N2_, size_t N3_) + : N1(N1_), N2(N2_), N3(N3_) {} + TriangleIterator begin() const { + TriangleIterator t; + t.i1 = t.i2 = t.i3 = 0; + t.N1 = N1; + t.N2 = N2; + t.N3 = N3; + return t; + } + + TriangleIterator end() const { + TriangleIterator t; + t.i2 = t.i3 = 0; + t.i1 = N1; + t.N1 = N1; + t.N2 = N2; + t.N3 = N3; + return t; + } + +}; + +extern "C" DLL_PUBLIC +void CosmoTool_compute_bispectrum( + double *delta_hat, size_t Nx, size_t Ny, size_t Nz, + size_t *Ntriangles, + double* B, double delta_k, size_t Nk ) +{ + // First remap to multi_array for easy access + size_t kNz = Nz/2+1; + boost::multi_array_ref, 3> a_delta(reinterpret_cast*>(delta_hat), boost::extents[Nx][Ny][kNz]); + boost::multi_array_ref a_Nt(Ntriangles, boost::extents[Nk][Nk][Nk]); + boost::multi_array_ref, 3> a_B(reinterpret_cast*>(B), boost::extents[Nk][Nk][Nk]); + typedef std::complex CType; + + // First loop over m1 + for (auto m1 : ModeSet(Nx, Ny, kNz)) { + CType& v1 = a_delta[m1.i1][m1.i2][m1.i3]; + // Second mode m2 + for (auto m2 : ModeSet(Nx, Ny, kNz)) { + // Now derive m3 + CType& v2 = a_delta[m2.i1][m2.i2][m2.i3]; + auto m3 = (m1.real()+m2.real()); + + if (!m3.in_box()) + continue; + + size_t q1 = std::floor(m1.norm()/delta_k); + size_t q2 = std::floor(m2.norm()/delta_k); + size_t q3 = std::floor(m3.norm()/delta_k); + + + if (q1 > Nk || q2 > Nk || q3 > Nk) + continue; + + CType prod = v1*v2; + // We use hermitic symmetry to get -m3, it is just the mode in m3 but conjugated. + if (m3.i3 > 0) { + if (m3.i1 < 0) m3.i1 += m3.N1; + if (m3.i2 < 0) m3.i2 += m3.N2; + prod *= std::conj(a_delta[m3.i1][m3.i2][m3.i3]); + } else { + m3.i1 = -m3.i1; + m3.i2 = -m3.i2; + if (m3.i1 < 0) m3.i1 += m3.N1; + if (m3.i2 < 0) m3.i2 += m3.N2; + prod *= a_delta[m3.i1][m3.i2][m3.i3]; + } + +// a_Nt[q1][q2][q3] ++; +// a_B[q1][q2][q3] += prod; + } + } +} + + +extern "C" DLL_PUBLIC +void CosmoTool_compute_powerspectrum( + double *delta_hat, size_t Nx, size_t Ny, size_t Nz, + size_t *Ncounts, + double* P, double delta_k, size_t Nk ) +{ + // First remap to multi_array for easy access + size_t kNz = Nz/2+1; + boost::multi_array_ref, 3> a_delta(reinterpret_cast*>(delta_hat), boost::extents[Nx][Ny][kNz]); + boost::multi_array_ref a_Nc(Ncounts, boost::extents[Nk]); + boost::multi_array_ref a_P(reinterpret_cast(P), boost::extents[Nk]); + typedef std::complex CType; + + // First loop over m1 + for (auto m1 : ModeSet(Nx, Ny, kNz)) { + CType& v1 = a_delta[m1.i1][m1.i2][m1.i3]; + + size_t q1 = std::floor(m1.norm()/delta_k); + + if (q1 > Nk) + continue; + + a_Nc[q1] ++; + a_P[q1] += std::norm(v1); + } +} diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py index a73b7c5..7e76e2c 100644 --- a/python/cosmotool/__init__.py +++ b/python/cosmotool/__init__.py @@ -14,6 +14,7 @@ except: from .simu import loadRamsesAll, simpleWriteGadget, SimulationBare from .timing import time_block, timeit, timeit_quiet +from .bispectrum import bispectrum, powerspectrum try: from .fftw import CubeFT diff --git a/python/cosmotool/bispectrum.py b/python/cosmotool/bispectrum.py new file mode 100644 index 0000000..729796b --- /dev/null +++ b/python/cosmotool/bispectrum.py @@ -0,0 +1,121 @@ +import numpy as np + +try: + import cffi + import os + + _ffi = cffi.FFI() + _ffi.cdef(""" + + void CosmoTool_compute_bispectrum( + double *delta_hat, size_t Nx, size_t Ny, size_t Nz, + size_t *Ntriangles, + double* B, double delta_k, size_t Nk ) ; + """); + + _pathlib = os.path.dirname(os.path.abspath(__file__)) + _lib = _ffi.dlopen(os.path.join(_pathlib,"_cosmo_bispectrum.so")) +except Exception as e: + print(repr(e)) + raise RuntimeError("Failed to initialize _cosmo_bispectrum module") + + +def bispectrum(delta, delta_k, Nk, fourier=True): + """bispectrum(delta, fourier=True) + + Args: + * delta: a 3d density field, can be Fourier modes if fourier set to True + + + Return: + * A 3d array of the binned bispectrum +""" + + if len(delta.shape) != 3: + raise ValueError("Invalid shape for delta") + try: + delta_k = float(delta_k) + Nk = int(Nk) + except: + raise ValueError() + + if not fourier: + delta = np.fft.rfftn(delta) + N1,N2,N3 = delta.shape + delta_hat_buf = np.empty((N1*N2*N3*2),dtype=np.double) + delta_hat_buf[::2] = delta.real.ravel() + delta_hat_buf[1::2] = delta.imag.ravel() + + size_size = _ffi.sizeof("size_t") + if size_size == 4: + triangle_buf = np.zeros((Nk,Nk,Nk),dtype=np.int32) + elif size_size == 8: + triangle_buf = np.zeros((Nk,Nk,Nk),dtype=np.int64) + else: + raise RuntimeError("Internal error, do not know how to map size_t") + + B_buf = np.zeros((Nk*Nk*Nk*4), dtype=np.double) + + _lib.CosmoTool_compute_bispectrum( \ + _ffi.cast("double *", delta_hat_buf.ctypes.data), \ + N1, N2, N3, \ + _ffi.cast("size_t *", triangle_buf.ctypes.data), \ + _ffi.cast("double *", B_buf.ctypes.data), \ + delta_k, \ + Nk) + B_buf = B_buf.reshape((Nk,Nk,Nk,4)) + return triangle_buf, B_buf[...,0]+1j*B_buf[...,1] + +def powerspectrum(delta, delta_k, Nk, fourier=True): + """powerspectrum(delta, fourier=True) + + Args: + * delta: a 3d density field, can be Fourier modes if fourier set to True + + + Return: + * A 3d array of the binned bispectrum +""" + + if len(delta.shape) != 3: + raise ValueError("Invalid shape for delta") + try: + delta_k = float(delta_k) + Nk = int(Nk) + except: + raise ValueError() + + if not fourier: + delta = np.fft.rfftn(delta) + N1,N2,N3 = delta.shape + delta_hat_buf = np.empty((N1*N2*N3*2),dtype=np.double) + delta_hat_buf[::2] = delta.real.ravel() + delta_hat_buf[1::2] = delta.imag.ravel() + + size_size = _ffi.sizeof("size_t") + if size_size == 4: + count_buf = np.zeros((Nk,),dtype=np.int32) + elif size_size == 8: + count_buf = np.zeros((Nk,),dtype=np.int64) + else: + raise RuntimeError("Internal error, do not know how to map size_t") + + B_buf = np.zeros((Nk,), dtype=np.double) + + _lib.CosmoTool_compute_bispectrum( \ + _ffi.cast("double *", delta_hat_buf.ctypes.data), \ + N1, N2, N3, \ + _ffi.cast("size_t *", count_buf.ctypes.data), \ + _ffi.cast("double *", B_buf.ctypes.data), \ + delta_k, \ + Nk) + return count_buf, B_buf[...] + + +if __name__=="__main__": + delta=np.zeros((16,16,16)) + delta[0,0,0]=1 + delta[3,2,1]=1 + b = powerspectrum(delta, 1, 16, fourier=False) + a = bispectrum(delta, 1, 16, fourier=False) + print a[0].max() diff --git a/python_sample/test_bispectrum.py b/python_sample/test_bispectrum.py new file mode 100644 index 0000000..f1f247b --- /dev/null +++ b/python_sample/test_bispectrum.py @@ -0,0 +1,10 @@ +import numpy as np +import cosmotool as ct + +f=0.01 +d=np.random.normal(size=(16,16,16)) +rho = d + f *(d*d - np.average(d*d)) + +B = ct.bispectrum(rho, 1, 16, fourier=False) +P = ct.powerspectrum(rho, 1, 16, fourier=False) +