Moved subsampling to the loaders to be able to do online subsampling (only implemented for gadget and multidark so far, the other coming in...)

c_tools/mock/generateMock.cpp

@@ -1,3 +1,4 @@
+#include <gsl/gsl_rng.h>
 #include <boost/function.hpp>
 #include <boost/bind.hpp>
 #include <boost/format.hpp>
@@ -244,11 +245,28 @@ void selectBox(SimuData *simu, std::vector<long>& targets, generateMock_info& ar
 	  (simu->Pos[j][i] < ranges[j][1]);	
       }
       
-      if (acceptance && (drand48() <= subsample))
+      if (acceptance)
 	targets.push_back(i);
     }
 }
 
+class PreselectParticles: public SimulationPreprocessor
+{
+private:
+  gsl_rng *rng;
+  double subsample;
+public:
+  PreselectParticles(gsl_rng *r, double s)
+    : subsample(s), rng(r)
+  {
+  }
+
+  bool accept(const SingleParticle& p)
+  {
+    return gsl_rng_uniform(rng) < subsample;
+  }
+};
+
 void createBox(SimuData *simu, vector<long>& targets, vector<long>& snapshot_split, SimuData *& boxed, generateMock_info& args_info)
 {
   double *ranges = new double[6];
@@ -486,6 +504,7 @@ int main(int argc, char **argv)
   generateMock_conf_params args_params;
   SimuData *simu, *simuOut;
   SimulationLoader *loader;
+  gsl_rng *rng = gsl_rng_alloc(gsl_rng_default);
  
   generateMock_conf_init(&args_info);
   generateMock_conf_params_init(&args_params);
@@ -511,13 +530,16 @@ int main(int argc, char **argv)
   
   generateMock_conf_print_version();
 
+  gsl_rng_set(rng, args_info.subsample_seed_arg);
+  SimulationPreprocessor *preselector = new PreselectParticles(rng, args_info.subsample_arg);
+
   if (args_info.ramsesBase_given || args_info.ramsesId_given)
     {
       if (args_info.ramsesBase_given && args_info.ramsesId_given)   {
 	loader = ramsesLoader(args_info.ramsesBase_arg, 
 			      args_info.ramsesId_arg, 
 			      false,
-			      NEED_POSITION|NEED_VELOCITY|NEED_GADGET_ID);
+			      NEED_POSITION|NEED_VELOCITY|NEED_GADGET_ID, preselector);
       }
       else
 	{
@@ -527,15 +549,15 @@ int main(int argc, char **argv)
     }
   else if (args_info.gadget_given)
     {
-      loader = gadgetLoader(args_info.gadget_arg, 1/args_info.gadgetUnit_arg, NEED_POSITION|NEED_VELOCITY|NEED_GADGET_ID);
+      loader = gadgetLoader(args_info.gadget_arg, 1/args_info.gadgetUnit_arg, NEED_POSITION|NEED_VELOCITY|NEED_GADGET_ID, preselector);
     }
   else if (args_info.flash_given)
     {
-      loader = flashLoader(args_info.flash_arg, NEED_POSITION|NEED_VELOCITY|NEED_GADGET_ID); 
+      loader = flashLoader(args_info.flash_arg, NEED_POSITION|NEED_VELOCITY|NEED_GADGET_ID, preselector); 
   }
   else if (args_info.multidark_given)
     {
-      loader = multidarkLoader(args_info.multidark_arg);      
+      loader = multidarkLoader(args_info.multidark_arg, preselector);      
     }
   else
     {
@@ -598,7 +620,9 @@ int main(int argc, char **argv)
   saveBox(simuOut, args_info.outputParameter_arg);
   generateOutput(simuOut, args_info.axis_arg, 
                  args_info.output_arg);
+  delete preselector;
   
+  gsl_rng_free(rng);
  
   printf("Done!\n"); 
   return 0;

c_tools/mock/generateMock.ggo

@@ -33,3 +33,4 @@ option "subsample" - "Subsample the input simulation by the specified amount" do
 option "inputParameter" - "Input geometry (optional, warning!)" string optional
 option "gadgetUnit" - "Unit of length in gadget file in Mpc/h" double optional default="0.001"
 
+option "subsample_seed" - "Seed for random number generation to select the subsample" int optional default="190524"

c_tools/mock/loaders/basic_loader.cpp

@@ -0,0 +1,128 @@
+#include <cassert>
+#include <string>
+#include <iostream>
+#include <boost/format.hpp>
+#include <fstream>
+#include <CosmoTool/yorick.hpp>
+#include "simulation_loader.hpp"
+
+using namespace std;
+using namespace CosmoTool;
+using boost::format;
+using boost::str;
+
+class BasicGroupLoader: public SimulationLoader
+{
+private:
+  string storage;
+  SimuData *header;
+  int flags, numFiles;
+public:
+  BasicGroupLoader(const string& storage_path, SimuData *header, int flags, int numfiles)
+  {
+    this->header = header;
+    this->storage = storage_path;
+    this->flags = flags;
+    this->numFiles = numfiles;
+  }
+  
+  SimuData *getHeader() {
+    return header;
+  }
+  
+  int num_files() {
+    return numFiles;
+  }
+  
+  SimuData *loadFile(int id) {    
+    if (id < 0 || id >= numFiles)
+      return 0;
+
+    SimuData *simu = new SimuData;
+    uint32_t *dimlist, dimrank;
+    string fname;
+
+    simu->time = header->time;
+    simu->TotalNumPart = header->NumPart;
+    simu->Omega_M = header->Omega_M;
+    simu->Omega_Lambda = header->Omega_Lambda;
+    simu->NumPart = -1;
+
+    if (flags & NEED_POSITION)
+      {
+	loadArray(str(format("%s/x_%d.nc") % storage % id), 
+		  simu->Pos[0], dimlist, dimrank);
+	assert(dimrank == 1);
+	simu->NumPart = dimlist[0];
+	
+	loadArray(str(format("%s/y_%d.nc") % storage % id), 
+		  simu->Pos[1], dimlist, dimrank);
+	assert(dimrank == 1);
+	assert(simu->NumPart == dimlist[0]);
+	
+	loadArray(str(format("%s/z_%d.nc") % storage % id), 
+		  simu->Pos[2], dimlist, dimrank);
+	assert(dimrank == 1);
+	assert(simu->NumPart == dimlist[0]);
+      }
+
+    if (flags & NEED_VELOCITY)
+      {
+	loadArray(str(format("%s/vx_%d.nc") % storage % id), 
+		  simu->Vel[0], dimlist, dimrank);
+	assert(dimrank == 1);
+	if (simu->NumPart < 0)
+	  simu->NumPart = dimlist[0];
+
+	assert(simu->NumPart == dimlist[0]);
+	
+	loadArray(str(format("%s/vy_%d.nc") % storage % id), 
+		  simu->Vel[0], dimlist, dimrank);
+	assert(dimrank == 1);
+	assert(simu->NumPart == dimlist[0]);
+	
+	loadArray(str(format("%s/vz_%d.nc") % storage % id), 
+		  simu->Vel[2], dimlist, dimrank);
+	assert(dimrank == 1);
+	assert(simu->NumPart == dimlist[0]);
+      }
+
+    if (flags & NEED_GADGET_ID)
+      {
+	loadArray(str(format("%s/id_%d.nc") % storage % id), 
+		  simu->Id, dimlist, dimrank);
+	assert(dimrank == 1);
+	if (simu->NumPart < 0)
+	  simu->NumPart = dimlist[0];
+
+	assert(simu->NumPart == dimlist[0]);
+      }
+    
+    return simu;
+  }
+
+  ~BasicGroupLoader()
+  {
+    delete header;
+  }
+};
+
+SimulationLoader *basicGroupLoader(const std::string& simupath, 
+				   int flags)
+{
+  SimuData *header;
+  ifstream f;
+  string header_path = simupath + "/header.txt";
+  int numFiles;
+
+  header = new SimuData;
+  f.open(header_path.c_str());
+
+  f >> header->time
+    >> header->Omega_M
+    >> header->Omega_Lambda 
+    >> header->NumPart
+    >> numFiles;
+
+  return new BasicGroupLoader(simupath, header, flags, numFiles);
+}

c_tools/mock/loaders/flash_loader.cpp

@@ -64,7 +64,7 @@ public:
 };
 
 
-SimulationLoader *flashLoader(const std::string& snapshot, int flags)
+SimulationLoader *flashLoader(const std::string& snapshot, int flags, SimulationPreprocessor *p)
 {
   bool singleFile;
   int num_files;

c_tools/mock/loaders/gadget_loader.cpp

@@ -1,3 +1,4 @@
+#include <vector>
 #include <cassert>
 #include <string>
 #include <CosmoTool/loadGadget.hpp>
@@ -16,9 +17,10 @@ private:
   double unitMpc;
   SimuData *gadget_header;
   string snapshot_name;
+  SimulationPreprocessor *preproc;
 public:
-  GadgetLoader(const string& basename, SimuData *header, int flags, bool singleFile, int _num, double unit)
-    : snapshot_name(basename), load_flags(flags), onefile(singleFile), _num_files(_num), unitMpc(1/unit), gadget_header(header)
+  GadgetLoader(const string& basename, SimuData *header, int flags, bool singleFile, int _num, double unit, SimulationPreprocessor *p)
+    : snapshot_name(basename), load_flags(flags), onefile(singleFile), _num_files(_num), unitMpc(1/unit), gadget_header(header), preproc(p)
   {
   }
   
@@ -34,7 +36,7 @@ public:
   int num_files() {
     return _num_files;
   }
-  
+
   SimuData *loadFile(int id) {
     SimuData *d;
     
@@ -70,12 +72,38 @@ public:
 
     applyTransformations(d);
 
+    long numAccepted = 0;
+    bool *accepted = new bool[d->NumPart];
+    for (long i = 0; i < d->NumPart; i++)
+      {
+	SingleParticle p;
+
+	for (int k = 0; k < 3; k++)
+	  {
+	    p.Pos[k] = (d->Pos[k]) ? 0 : d->Pos[k][i];
+	    p.Vel[k] = (d->Vel[k]) ? 0 : d->Vel[k][i];
+	  }
+	p.ID = (d->Id) ? 0 : d->Id[i];
+	
+	accepted[i] = preproc->accept(p);
+	numAccepted += accepted[i];
+      }
+    
+    for (int k = 0; k < 3; k++)
+      {
+	filteredCopy(d->Pos[k], accepted, d->NumPart);
+	filteredCopy(d->Vel[k], accepted, d->NumPart);
+      }
+    filteredCopy(d->Id, accepted, d->NumPart);
+    filteredCopy(d->type, accepted, d->NumPart);
+    delete[] accepted;
+
     return d;
   }
 };
 
 
-SimulationLoader *gadgetLoader(const std::string& snapshot, double Mpc_unitLength, int flags)
+SimulationLoader *gadgetLoader(const std::string& snapshot, double Mpc_unitLength, int flags, SimulationPreprocessor *p)
 {
   bool singleFile = false;
   int num_files;
@@ -120,5 +148,5 @@ SimulationLoader *gadgetLoader(const std::string& snapshot, double Mpc_unitLengt
 	}
     }
     
-  return new GadgetLoader(snapshot, header, flags, singleFile, num_files, Mpc_unitLength);
+  return new GadgetLoader(snapshot, header, flags, singleFile, num_files, Mpc_unitLength, p);
 }

c_tools/mock/loaders/multidark_loader.cpp

@@ -14,9 +14,10 @@ class MultiDarkLoader: public SimulationLoader
 protected:
   SimuData *header;
   string darkname;
+  SimulationPreprocessor *preproc;
 public:
-  MultiDarkLoader(const std::string& name, SimuData *h)
-    : darkname(name), header(h)
+  MultiDarkLoader(const std::string& name, SimuData *h, SimulationPreprocessor *p)
+    : preproc(p), darkname(name), header(h)
   {
   }
 
@@ -48,40 +49,59 @@ public:
     simu->TotalNumPart = simu->NumPart;
     simu->Omega_Lambda = 1.0 - simu->Omega_M;
 
+    long estimated = (preproc == 0) ? simu->NumPart : preproc->getEstimatedPostprocessed(simu->NumPart);
+    long allocated = estimated; 
+
     for (int k = 0; k < 3; k++)
-      simu->Pos[k] = new float[simu->NumPart];
-    simu->Vel[2] = new float[simu->NumPart];
-    simu->Id = new long[simu->NumPart];
-    long *uniqueID = new long[simu->NumPart];
+      simu->Pos[k] = new float[allocated];
+    simu->Vel[2] = new float[allocated];
+    simu->Id = new long[allocated];
+    long *uniqueID = new long[allocated];
+    long *index = new long[allocated];
+
     double tempData;
 
     simu->new_attribute("uniqueID", uniqueID, delete_adaptor<long>);
+    simu->new_attribute("index", index, delete_adaptor<long>);
 
     cout << "loading multidark particles" << endl;
     long actualNumPart = 0;
+
     for (long i = 0; i < simu->NumPart; i++) {
+      SingleParticle p;
 
-      fp >> simu->Id[i] >> simu->Pos[0][i] >> simu->Pos[1][i]
-	 >> simu->Pos[2][i] >> simu->Vel[2][i] >> tempData >> tempData;
+      fp >> p.ID >> p.Pos[0] >> p.Pos[1]
+         >> p.Pos[2] >> p.Vel[2] >> tempData >> tempData;
 
-      uniqueID[i] = simu->Id[i];
-
-      if (simu->Id[i] == -99 && 
-          simu->Pos[0][i] == -99 && simu->Pos[1][i] == -99 && 
-          simu->Pos[2][i] == -99 && simu->Vel[2][i] == -99) {
+      if (p.ID == -99 && 
+          p.Pos[0] == -99 && p.Pos[1] == -99 && 
+          p.Pos[2] == -99 && p.Vel[2] == -99) {
         break;
-      } else {
-        actualNumPart++;
+
+      if (preproc != 0 && !preproc->accept(p))
+        continue;
+
+      copyParticleToSimu(p, simu, actualNumPart);
+      uniqueID[actualNumPart]= p.ID;
+      index[actualNumPart] = i;
+      actualNumPart++;
+      if (actualNumPart == allocated)
+        {
+          allocated += (estimated+9)/10;
+          reallocSimu(simu, allocated); 
+          reallocArray(uniqueID, allocated, actualNumPart);
+          reallocArray(index, allocated, actualNumPart);
+        }
       }
     }
-
+    applyTransformations(simu);
     simu->NumPart = actualNumPart;
     simu->TotalNumPart = actualNumPart;
     return simu;
   }
 };
 
-SimulationLoader *multidarkLoader(const string& multidarkname)
+SimulationLoader *multidarkLoader(const string& multidarkname, SimulationPreprocessor *p)
 {
   SimuData *header;
   int actualNumPart;
@@ -101,7 +121,7 @@ SimulationLoader *multidarkLoader(const string& multidarkname)
   header->TotalNumPart = header->NumPart;
   header->Omega_Lambda = 1.0 - header->Omega_M;
 
-  return new MultiDarkLoader(multidarkname, header);
+  return new MultiDarkLoader(multidarkname, header, p);
 }
 
   

c_tools/mock/loaders/ramses_loader.cpp

@@ -61,7 +61,7 @@ public:
   }
 };
 
-SimulationLoader *ramsesLoader(const std::string& snapshot, int baseid, bool double_precision, int flags)
+SimulationLoader *ramsesLoader(const std::string& snapshot, int baseid, bool double_precision, int flags, SimulationPreprocessor *p)
 {
   SimuData *d, *header;
   int num_files = 0;

c_tools/mock/loaders/simulation_loader.cpp

@@ -1,8 +1,34 @@
+#include <cmath>
 #include <CosmoTool/loadSimu.hpp>
 #include "simulation_loader.hpp"
 
+using std::min;
 using namespace CosmoTool;
 
+template<typename T> void reallocArray(T *& a, long newSize, long toCopy)
+{
+  T *b = new T[newSize];
+  if (a != 0)
+    {
+      memcpy(b, a, sizeof(T)*toCopy);
+      delete[] a;
+    }
+  a = b;
+}
+
+void SimulationLoader::reallocSimu(SimuData *s, long newNumPart)
+{
+  long to_copy = min(newNumPart, s->NumPart);
+  
+  for (int j = 0; j < 3; j++)
+    {
+      reallocArray(s->Pos[j], newNumPart, to_copy);
+      reallocArray(s->Vel[j], newNumPart, to_copy);
+    }
+  reallocArray(s->Id, newNumPart, to_copy);
+  reallocArray(s->type, newNumPart, to_copy);
+}
+
 void SimulationLoader::applyTransformations(SimuData *s)
 {
   float redshift_gravity = do_redshift ? 1.0 : 0.0;

c_tools/mock/loaders/simulation_loader.hpp

@@ -2,8 +2,17 @@
 #define _MOCK_SIMULATION_LOADER_HPP
 
 #include <string>
+#include <algorithm>
 #include <CosmoTool/loadSimu.hpp>
 
+struct SingleParticle
+{
+  float Pos[3];
+  float Vel[3];
+  long Index;
+  long ID;
+};
+
 class SimulationLoader
 {
 protected:
@@ -15,9 +24,37 @@ protected:
     do_redshift = false;
     redshift_axis = 2;
   }
+
+  template<typename T> void reallocArray(T *& a, long newSize, long toCopy)
+  {
+    T *b = new T[newSize];
+    if (a != 0)
+     {
+       std::copy(a, a+toCopy, b);
+       delete[] a;
+     }
+   a = b;
+  } 
+
+
+
+  void reallocSimu(CosmoTool::SimuData *s, long newNumPart);
  
   void applyTransformations(CosmoTool::SimuData *s);
 
+  void copyParticleToSimu(const SingleParticle& p, CosmoTool::SimuData *s, long index)
+  {
+    s->Pos[0][index] = p.Pos[0];
+    s->Pos[1][index] = p.Pos[1];
+    s->Pos[2][index] = p.Pos[2];
+    if (s->Vel[0])
+      s->Vel[0][index] = p.Vel[0];
+    if (s->Vel[1])
+      s->Vel[1][index] = p.Vel[1];
+    if (s->Vel[2])
+      s->Vel[2][index] = p.Vel[2];
+    s->Id[index] = p.ID;
+  }
 public:
   virtual ~SimulationLoader() {}
   
@@ -28,6 +65,38 @@ public:
   virtual int num_files() = 0;
   virtual CosmoTool::SimuData* loadFile(int id) = 0;
 
+  
+  template<typename T>
+  void filteredCopy(T *a, bool *accepted, long N)
+  {
+    long i = 0, j = 0;
+
+    if (a == 0)
+      return;
+
+    while (i < N)
+      {
+	if (!accepted[i])
+	  {
+	    i++;
+	    continue;
+	  }
+
+	a[j] = a[i];
+	j++;
+	i++;
+      }
+  }
+};
+
+class SimulationPreprocessor
+{
+public:
+  SimulationPreprocessor() {}
+  virtual ~SimulationPreprocessor() {}
+
+  virtual long getEstimatedPostprocessed(long numParticles) { return numParticles; };
+  virtual bool accept(const SingleParticle& p) { return true; }
 };
 
 template<typename T>
@@ -40,10 +109,10 @@ void delete_adaptor(void *ptr)
 
 
 // Unit length is the size of one Mpc in the simulation units
-SimulationLoader *gadgetLoader(const std::string& snapshot, double Mpc_unitLength, int flags);
-SimulationLoader *flashLoader(const std::string& snapshot, int flags);
-SimulationLoader *multidarkLoader(const std::string& snapshot);
-SimulationLoader *ramsesLoader(const std::string& snapshot, int baseid, bool double_precision, int flags);
+SimulationLoader *gadgetLoader(const std::string& snapshot, double Mpc_unitLength, int flags, SimulationPreprocessor *p);
+SimulationLoader *flashLoader(const std::string& snapshot, int flags, SimulationPreprocessor *p);
+SimulationLoader *multidarkLoader(const std::string& snapshot, SimulationPreprocessor *p);
+SimulationLoader *ramsesLoader(const std::string& snapshot, int baseid, bool double_precision, int flags, SimulationPreprocessor *p);
 
 
 #endif