Added (optional for the moment) SDF support in simulation_loaders.

c_tools/mock/loaders/sdf_loader.cpp

@@ -0,0 +1,232 @@
+#include <iostream>
+#include <boost/format.hpp>
+#include <vector>
+#include <cassert>
+#include <string>
+#include "sdfloader_internal.hpp"
+#include "simulation_loader.hpp"
+#include <libsdf/mpmy.h>
+#include <libsdf/SDF.h>
+#include <libsdf/error.h>
+#include <libsdf/cosmo.h>
+
+using boost::format;
+
+using namespace std;
+using namespace CosmoTool;
+
+class SDFLoader: public SimulationLoader
+{
+private:
+  int load_flags;
+  bool onefile;
+  int _num_files;
+  double unitMpc;
+  SimuData *sdf_header;
+  SDF *sdfp;
+  SimulationPreprocessor *preproc;
+  int num_splitting;
+public:
+  SDFLoader(SDF *sdf_file, SimuData *header, int flags, int num_splitting,
+            SimulationPreprocessor *p)
+    : sdfp(sdf_file), load_flags(flags),
+      sdf_header(header), preproc(p)
+  {
+    this->num_splitting = num_splitting;
+  }
+  
+  ~SDFLoader()
+  {
+    delete sdf_header;
+  }
+  
+  SimuData *getHeader() {
+    return sdf_header;
+  }
+  
+  int num_files() {
+    return num_splitting;
+  }
+
+
+  int64_t getStart(int id)
+  {
+    return sdf_header->TotalNumPart * int64_t(id) / num_splitting;
+  }
+
+  int64_t getNumberInSplit(int id)
+  {
+    return getStart(id+1)-getStart(id);
+  }
+
+  void rescaleParticles(SimuData *d,
+                        double rescale_position,
+                        double rescale_velocity)
+  {
+    float shift = 0.5*d->BoxSize;
+    rescale_position /= d->time;
+
+    if (d->Pos[0] != 0)
+      {
+        for (int k = 0; k < 3; k++)
+          {
+            for (int64_t i = 0; i < d->NumPart; i++)
+              {
+                d->Pos[k][i] = (d->Pos[k][i]*rescale_position + shift);
+              }
+          }
+      }
+    if (d->Vel[0] != 0)
+      {
+        for (int k = 0; k < 3; k++)
+          {
+            for (int64_t i = 0; i < d->NumPart; i++)
+              {
+                d->Vel[k][i] *= rescale_velocity;
+              }
+          }
+      }
+  }
+
+
+  SimuData *loadFile(int id) {
+    SimuData *d;
+    int64_t starts[7];
+    int nobjs[7];
+    int strides[7];
+    void *addrs[7];
+    const char *names[7];
+    int ret;
+
+    if (id >= num_splitting)
+      return 0;
+
+    d = new SimuData;
+    *d = *sdf_header;
+    d->NumPart = getNumberInSplit(id);
+
+    int64_t numPartToLoad = getNumberInSplit(id);
+    int64_t start = getStart(id);
+    int k = 0;
+#define SETUP_READ(name, vec) \
+    names[k] = name, starts[k] = start, nobjs[k] = numPartToLoad, strides[k] = sizeof(vec[0]), addrs[k] = vec, k++;
+
+    if (load_flags & NEED_POSITION)
+      {
+        const char *name_template[3] = { "x", "y", "z" };
+        for (int q = 0; q < 3; q++)
+          {
+            d->Pos[q] = new float[numPartToLoad];
+            SETUP_READ(name_template[q], d->Pos[q]);
+          }
+      }
+    if (load_flags & NEED_VELOCITY)
+      {
+        const char *name_template[3] = { "x", "y", "z" };
+        for (int q = 0; q < 3; q++)
+          {
+            d->Vel[q] = new float[numPartToLoad];
+            SETUP_READ(name_template[q], d->Vel[q]);
+          }
+      }
+    if (load_flags & NEED_GADGET_ID)
+      {
+        d->Id = new long[numPartToLoad];
+        SETUP_READ("ident", d->Id);
+      }
+#undef SETUP_READ
+
+    ret = SDFseekrdvecsarr(sdfp, k, (char**)names, starts, nobjs, addrs, strides);
+    if (ret != 0)
+      {
+        cerr << format("Splitting %d/%d, SDF read failure: %s") % id % num_splitting % SDFerrstring << endl;
+        abort();
+      }
+
+    if (load_flags & (NEED_POSITION | NEED_VELOCITY))
+      rescaleParticles(d, 0.001*d->Hubble, one_kpc/one_Gyr);
+
+    applyTransformations(d);
+    basicPreprocessing(d, preproc);
+
+    return d;
+  }
+};
+
+SimulationLoader *sdfLoader(const std::string& snapshot, int flags, 
+                            int num_splitting,
+                            SimulationPreprocessor *p)
+{
+  SDF *sdfp;
+  int fileversion;
+  SimuData *hdr;
+  int64_t gnobj;
+
+  sdfp = SDFopen(0, snapshot.c_str());
+  if (sdfp == 0)
+    {
+      return 0;
+    }
+  
+  SDFgetintOrDefault(sdfp, "version", &fileversion, 1);
+  if (fileversion == 1)
+    {
+      SDFgetfloatOrDie(sdfp, "Omega0", &hdr->Omega_M);
+      SDFgetfloatOrDie(sdfp, "Lambda_prime", &hdr->Omega_Lambda);
+      SDFgetfloatOrDie(sdfp, "H0", &hdr->Hubble);
+    }
+  else
+    {
+      float Or, Of;
+      SDFgetfloatOrDie(sdfp, "Omega0_m", &hdr->Omega_M);
+      SDFgetfloatOrDie(sdfp, "Omega0_lambda", &hdr->Omega_Lambda);
+      SDFgetfloatOrDie(sdfp, "Omega0_r", &Or);
+      SDFgetfloatOrDie(sdfp, "Omega0_fld", &Of);
+
+      hdr->Omega_M += Or;
+      hdr->Omega_Lambda += Of;
+      SDFgetfloatOrDie(sdfp, "H0", &hdr->Hubble);
+      
+    }
+  double h0;
+  h0 = hdr->Hubble*10.0*(one_kpc/one_Gyr);
+
+  if (SDFhasname("R0", sdfp))
+    {
+      double R0;
+
+      SDFgetdoubleOrDie(sdfp, "R0", &R0);
+      hdr->BoxSize = 2.0*0.001*R0*h0;
+    }
+
+  if (SDFhasname("Rx", sdfp))
+    {
+      double R0;
+
+      SDFgetdoubleOrDie(sdfp, "Rx", &R0);
+      hdr->BoxSize = 2.0*0.001*R0*h0;
+    }
+
+  if (SDFhasname("redshift", sdfp))
+    {
+      double redshift;
+
+      SDFgetdoubleOrDie(sdfp, "redshift", &redshift);
+      hdr->time = 1/(1+redshift);
+    }
+   
+  if (SDFgetint64(sdfp, "npart", &gnobj))
+    {
+      gnobj = SDFnrecs("x", sdfp);
+      cerr << format("[Warning] No 'npart' found in SDF file '%s', guessing npart=%ld from SDFnrecs") % snapshot % gnobj << endl;
+    }
+  hdr->NumPart = hdr->TotalNumPart = gnobj;
+
+  return new SDFLoader(sdfp, hdr, flags, num_splitting, p);
+}
+
+
+void sdfLoaderInit(int& argc, char **& argv)
+{
+  MPMY_Init(&argc, &argv);
+}