From 5e55b263a1806a596ce2f5301bdfbb641c7f888f Mon Sep 17 00:00:00 2001
From: Guilhem Lavaux <lavaux@iap.fr>
Date: Thu, 20 Nov 2014 14:45:00 +0100
Subject: [PATCH] Added parallel gadget loader to saturate I/O

---
 python/_cosmotool.pyx  | 51 +++++++++++++++++++++++++++++++++++++++++-
 python/safe_gadget.hpp | 29 ++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 python/safe_gadget.hpp

diff --git a/python/_cosmotool.pyx b/python/_cosmotool.pyx
index 380f146..83590b2 100644
--- a/python/_cosmotool.pyx
+++ b/python/_cosmotool.pyx
@@ -1,5 +1,7 @@
 from libcpp cimport bool
 from libcpp cimport string as cppstring
+from libcpp.vector cimport vector as cppvector
+from cython.parallel cimport prange
 import numpy as np
 cimport numpy as np
 from cpython cimport PyObject, Py_INCREF
@@ -34,9 +36,13 @@ cdef extern from "loadSimu.hpp" namespace "CosmoTool":
 
 cdef extern from "loadGadget.hpp" namespace "CosmoTool":
 
-  SimuData *loadGadgetMulti(const char *fname, int id, int flags) except +
+  SimuData *loadGadgetMulti(const char *fname, int id, int flags) nogil except +
   void cxx_writeGadget "CosmoTool::writeGadget" (const char * s, SimuData *data) except +
 
+cdef extern from "safe_gadget.hpp":
+  SimuData *loadGadgetMulti_safe(cppstring.string s, int flags) nogil
+  SimuData **alloc_simudata(int num) nogil
+  void del_simudata(SimuData **d) nogil
 
 cdef extern from "loadRamses.hpp" namespace "CosmoTool":
   SimuData *loadRamsesSimu(const char *basename, int id, int cpuid, bool dp, int flags) except +
@@ -295,6 +301,49 @@ It loads a gadget-1 snapshot and return a PySimulationBase object. If cpu_id is
 
   return _PySimulationAdaptor(wrap_simudata(data, flags))
 
+def loadParallelGadget(list filename_list, bool loadPosition = True, bool loadVelocity = False, bool loadId = False, bool loadType = False, bool loadMass=False):
+  cdef int flags, i, num_files
+  cdef list out_arrays
+  cdef SimuData ** data
+  cdef SimuData * local_data
+  cdef Simulation simu
+  cdef cppvector[cppstring.string] filenames
+
+  flags = 0
+  if loadPosition:
+     flags |= NEED_POSITION
+  if loadVelocity:
+     flags |= NEED_VELOCITY
+  if loadId:
+     flags |= NEED_GADGET_ID
+  if loadType:
+     flags |= NEED_TYPE
+  if loadMass:
+     flags |= NEED_MASS
+
+  num_files = len(filename_list)
+  filenames.resize(num_files)
+  data = alloc_simudata(num_files)
+  for i,l in enumerate(filename_list):
+    filenames[i] = l
+  
+  with nogil:
+    for i in prange(num_files):
+      local_data = loadGadgetMulti_safe(filenames[i], flags)
+      data[i] = local_data
+#      data[i] = loadGadgetMulti(filenames[i].c_str(), -1, flags)
+  
+  out_arrays = None    
+  for i in xrange(num_files):
+    if data[i] == <SimuData*>0:        
+      out_arrays.append(None)
+    else:
+      out_arrays.append(_PySimulationAdaptor(wrap_simudata(data[i], flags)))
+
+  del_simudata(data)
+
+  return out_arrays
+
 def writeGadget(str filename, object simulation):
 
   cdef SimuData simdata
diff --git a/python/safe_gadget.hpp b/python/safe_gadget.hpp
new file mode 100644
index 0000000..3f36e4c
--- /dev/null
+++ b/python/safe_gadget.hpp
@@ -0,0 +1,29 @@
+#include "config.hpp"
+#include "loadGadget.hpp"
+#include <string>
+
+static inline
+CosmoTool::SimuData *loadGadgetMulti_safe(const std::string& fname, int flags)
+{
+  try
+    {
+      return CosmoTool::loadGadgetMulti(fname.c_str(), -1, flags);
+    }
+  catch (const CosmoTool::Exception& e)
+    {
+      return 0;
+    }
+}
+
+
+static inline
+CosmoTool::SimuData **alloc_simudata(int n)
+{
+  return new CosmoTool::SimuData *[n];
+}
+
+static inline
+void del_simudata(CosmoTool::SimuData **s)
+{
+  delete[] s;
+}