cosmotool/python/_cosmotool.pyx

from libcpp cimport bool
import numpy as np
cimport numpy as np
from cpython cimport PyObject, Py_INCREF

np.import_array()


cdef extern from "loadSimu.hpp" namespace "CosmoTool":

  cdef cppclass SimuData:
    np.float_t BoxSize
    np.float_t time
    np.float_t Hubble

    np.float_t Omega_M
    np.float_t Omega_Lambda
    np.int64_t TotalNumPart
    np.int64_t NumPart
    np.int64_t *Id
    float *Pos[3]
    float *Vel[3]
    int *type

  cdef const int NEED_GADGET_ID 
  cdef const int NEED_POSITION
  cdef const int NEED_VELOCITY
  cdef const int NEED_TYPE

cdef extern from "loadGadget.hpp" namespace "CosmoTool":

  SimuData *loadGadgetMulti(const char *fname, int id, int flags) except +

cdef extern from "loadRamses.hpp" namespace "CosmoTool":
  SimuData *loadRamsesSimu(const char *basename, int id, int cpuid, bool dp, int flags) except +

cdef class Simulation:

  cdef list positions
  cdef list velocities
  cdef object identifiers

  cdef SimuData *data

  property BoxSize:
    def __get__(Simulation self):
        return self.data.BoxSize
        
  property time:
    def __get__(Simulation self):
        return self.data.time

  property Hubble:
    def __get__(Simulation self):
        return self.data.Hubble

  property Omega_M:
    def __get__(Simulation self):
        return self.data.Omega_M
        
  property positions:
    def __get__(Simulation self):
        return self.positions
        
  property velocities:
    def __get__(Simulation self):
        return self.velocities
        
  property numParticles:
    def __get__(Simulation self):
        return self.data.NumPart

  def __cinit__(Simulation self):
    self.data = <SimuData *>0
  
  def __dealloc__(Simulation self):
    if self.data != <SimuData *>0:
      print("Clearing simulation data")
      del self.data


cdef class ArrayWrapper:
    cdef void* data_ptr
    cdef int size
    cdef int type_array
    
    cdef set_data(self, int size, int type_array, void* data_ptr):
        """ Set the data of the array

This cannot be done in the constructor as it must recieve C-level
arguments.

Parameters:
-----------
size: int
Length of the array.
data_ptr: void*
Pointer to the data

    """
        self.data_ptr = data_ptr
        self.size = size
        self.type_array = type_array
     
    def __array__(self):
        """ Here we use the __array__ method, that is called when numpy
tries to get an array from the object."""
        cdef np.npy_intp shape[1]
        
        shape[0] = <np.npy_intp> self.size
        # Create a 1D array, of length 'size'
        ndarray = np.PyArray_SimpleNewFromData(1, shape, self.type_array, self.data_ptr)
        return ndarray
     
    def __dealloc__(self):
        """ Frees the array. This is called by Python when all the
references to the object are gone. """
        pass
 
cdef object wrap_array(void *p, np.uint64_t s, int typ):
    cdef np.ndarray ndarray
    cdef ArrayWrapper wrapper

    wrapper = ArrayWrapper()
    wrapper.set_data(s, typ, p)
    ndarray = np.array(wrapper, copy=False)
    ndarray.base = <PyObject*> wrapper
    Py_INCREF(wrapper)
    
    return ndarray


cdef object wrap_float_array(float *p, np.uint64_t s):
    return wrap_array(<void *>p, s, np.NPY_FLOAT)

cdef object wrap_int64_array(np.int64_t* p, np.uint64_t s):
    return wrap_array(<void *>p, s, np.NPY_INT64)

cdef object wrap_simudata(SimuData *data, int flags):
  cdef Simulation simu

  simu = Simulation()
  simu.data = data
  if flags & NEED_POSITION:
     simu.positions = [wrap_float_array(data.Pos[i], data.NumPart) for i in xrange(3)]
  if flags & NEED_VELOCITY:
     simu.velocities = [wrap_float_array(data.Vel[i], data.NumPart) for i in xrange(3)]
  if flags & NEED_GADGET_ID:
     simu.identifiers = wrap_int64_array(data.Id, data.NumPart)
  return simu

def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False):

  cdef int flags
  cdef SimuData *data
  cdef Simulation simu

  flags = 0
  if loadPosition:
     flags |= NEED_POSITION
  if loadVelocity:
     flags |= NEED_VELOCITY
  if loadId:
     flags |= NEED_GADGET_ID

  data = loadGadgetMulti(filename, snapshot_id, flags)
  if data == <SimuData*>0:
    return None

  return wrap_simudata(data, flags)

def loadRamses(str basepath, int snapshot_id, int cpu_id, bool doublePrecision = False, bool loadPosition = True, bool loadVelocity = False):
  """ loadRamses(basepath, snapshot_id, cpu_id, doublePrecision=False, loadPosition=True, loadVelocity=False)
  Loads the indicated snapshot based on the cpu id, snapshot id and basepath. It is important to specify the correct precision in doublePrecision.
"""
  cdef int flags
  cdef SimuData *data
  cdef Simulation simu

  flags = 0
  if loadPosition:
     flags |= NEED_POSITION
  if loadVelocity:
     flags |= NEED_VELOCITY

  data = loadRamsesSimu(basepath, snapshot_id, cpu_id, doublePrecision, flags)
  if data == <SimuData*>0:
    return None

  return wrap_simudata(data, flags)