Added some bispectrum computer

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)

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
 )
 

python/_cosmo_bispectrum.cpp

@@ -0,0 +1,194 @@
+#include <iostream>
+#include <boost/format.hpp>
+#include <boost/multi_array.hpp>
+#include <sys/types.h>
+#include <cmath>
+
+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<std::complex<double>, 3> a_delta(reinterpret_cast<std::complex<double>*>(delta_hat), boost::extents[Nx][Ny][kNz]); 
+    boost::multi_array_ref<size_t, 3> a_Nt(Ntriangles, boost::extents[Nk][Nk][Nk]); 
+    boost::multi_array_ref<std::complex<double>, 3> a_B(reinterpret_cast<std::complex<double>*>(B), boost::extents[Nk][Nk][Nk]); 
+    typedef std::complex<double> 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<std::complex<double>, 3> a_delta(reinterpret_cast<std::complex<double>*>(delta_hat), boost::extents[Nx][Ny][kNz]); 
+    boost::multi_array_ref<size_t, 1> a_Nc(Ncounts, boost::extents[Nk]); 
+    boost::multi_array_ref<double, 1> a_P(reinterpret_cast<double*>(P), boost::extents[Nk]); 
+    typedef std::complex<double> 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);
+    }
+}

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

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()

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)
+