From c67b9a8b3858f504176e8563f54e8da7c1a4ae52 Mon Sep 17 00:00:00 2001
From: lhausamm <loic_hausammann@hotmail.com>
Date: Tue, 29 Aug 2017 17:13:59 +0200
Subject: [PATCH] Add index file writing.  - Currently too much
 copy-paste-modify. Will need to merge all the code when the code is ready  -
 need to remove a bug in the number of task created before testing this commit

---
 src/engine.c                     |   8 +-
 src/gravity/Default/gravity_io.h |  24 +++
 src/hydro/Gadget2/hydro_io.h     |  24 +++
 src/single_io.c                  | 252 ++++++++++++++++++++++++++++++-
 src/single_io.h                  |   3 +
 5 files changed, 307 insertions(+), 4 deletions(-)

diff --git a/src/engine.c b/src/engine.c
index ef477b2d48..961a2fae33 100644
--- a/src/engine.c
+++ b/src/engine.c
@@ -6453,8 +6453,12 @@ void engine_dump_snapshot(struct engine *e) {
                       MPI_INFO_NULL);
 #endif
 #else
-  write_output_single(e, e->snapshot_base_name, e->internal_units,
-                      e->snapshot_units);
+  if (e->policy & engine_policy_logger)
+    write_index_single(e, e->snapshotBaseName, e->internal_units,
+		       e->snapshotUnits);
+  else
+    write_output_single(e, e->snapshotBaseName, e->internal_units,
+			e->snapshotUnits);
 #endif
 #endif
 
diff --git a/src/gravity/Default/gravity_io.h b/src/gravity/Default/gravity_io.h
index 1ba3899e7e..d8fb1bb1fc 100644
--- a/src/gravity/Default/gravity_io.h
+++ b/src/gravity/Default/gravity_io.h
@@ -117,4 +117,28 @@ INLINE static void darkmatter_write_particles(const struct gpart* gparts,
                                  UNIT_CONV_NO_UNITS, gparts, id_or_neg_offset);
 }
 
+
+/**
+ * @brief Specifies which g-particle fields to write to a dataset
+ *
+ * @param gparts The g-particle array.
+ * @param list The list of i/o properties to write.
+ * @param num_fields The number of i/o fields to write.
+ * 
+ * In this version, we only want the ids and the offset.
+ */
+void darkmatter_write_index(struct gpart* gparts, struct io_props* list,
+			    int* num_fields) {
+
+  /* Say how much we want to read */
+  *num_fields = 2;
+
+  /* List what we want to read */
+  list[0] = io_make_output_field("ParticleIDs", ULONGLONG, 1,
+                                 UNIT_CONV_NO_UNITS, gparts, id_or_neg_offset);
+
+  list[1] = io_make_output_field("Offset", ULONGLONG, 1,
+                                 UNIT_CONV_NO_UNITS, gparts, last_offset);
+}
+
 #endif /* SWIFT_DEFAULT_GRAVITY_IO_H */
diff --git a/src/hydro/Gadget2/hydro_io.h b/src/hydro/Gadget2/hydro_io.h
index ec7d34f7ad..1b43306bb3 100644
--- a/src/hydro/Gadget2/hydro_io.h
+++ b/src/hydro/Gadget2/hydro_io.h
@@ -188,6 +188,30 @@ INLINE static void hydro_write_particles(const struct part* parts,
 #endif
 }
 
+
+/**
+ * @brief Specifies which particle fields to write to a dataset
+ *
+ * @param parts The particle array.
+ * @param list The list of i/o properties to write.
+ * @param num_fields The number of i/o fields to write.
+ * 
+ * In this version, we only want the ids and the offset.
+ */
+void hydro_write_index(struct part* parts, struct io_props* list,
+		       int* num_fields) {
+
+  *num_fields = 2;
+
+  /* List what we want to write */
+  list[0] = io_make_output_field("ParticleIDs", ULONGLONG, 1,
+                                 UNIT_CONV_NO_UNITS, parts, id);
+
+  list[1] = io_make_output_field("Offset", ULONGLONG, 1,
+                                 UNIT_CONV_NO_UNITS, parts, last_offset);
+}
+
+
 /**
  * @brief Writes the current model of SPH to the file
  * @param h_grpsph The HDF5 group in which to write
diff --git a/src/single_io.c b/src/single_io.c
index d0bc850251..dd072cf136 100644
--- a/src/single_io.c
+++ b/src/single_io.c
@@ -312,8 +312,9 @@ void writeArray(const struct engine* e, hid_t grp, char* fileName,
   if (h_err < 0) error("Error while writing data array '%s'.", props.name);
 
   /* Write XMF description for this data set */
-  xmf_write_line(xmfFile, fileName, partTypeGroupName, props.name, N,
-                 props.dimension, props.type);
+  if (xmfFile != NULL)
+    xmf_write_line(xmfFile, fileName, partTypeGroupName, props.name, N,
+		   props.dimension, props.type);
 
   /* Write unit conversion factors for this data set */
   char buffer[FIELD_BUFFER_SIZE];
@@ -685,6 +686,8 @@ void write_output_single(struct engine* e, const char* baseName,
 
   /* Write the relevant information */
   io_write_attribute(h_grp, "PeriodicBoundariesOn", INT, &periodic, 1);
+  int index = 0;
+  io_write_attribute(h_grp, "IsIndexFile", INT, &index, 1);
 
   /* Close runtime parameters */
   H5Gclose(h_grp);
@@ -994,4 +997,249 @@ void write_output_single(struct engine* e, const char* baseName,
   e->snapshot_output_count++;
 }
 
+
+/**
+ * @brief Writes an HDF5 index file
+ *
+ * @param e The engine containing all the system.
+ * @param baseName The common part of the snapshot file name.
+ * @param internal_units The #unit_system used internally
+ * @param snapshot_units The #unit_system used in the snapshots
+ *
+ * Creates an HDF5 output file and writes the offset and id of particles contained
+ * in the engine. If such a file already exists, it is erased and replaced
+ * by the new one.
+ *
+ * Calls #error() if an error occurs.
+ *
+ */
+void write_index_single(struct engine* e, const char* baseName,
+                         const struct unit_system* internal_units,
+                         const struct unit_system* snapshot_units) {
+
+  hid_t h_file = 0, h_grp = 0;
+  const size_t Ngas = e->s->nr_parts;
+  const size_t Nstars = e->s->nr_sparts;
+  const size_t Ntot = e->s->nr_gparts;
+  int periodic = e->s->periodic;
+  int numFiles = 1;
+  struct part* parts = e->s->parts;
+  struct gpart* gparts = e->s->gparts;
+  struct gpart* dmparts = NULL;
+  //struct spart* sparts = e->s->sparts;
+  static int outputCount = 0;
+
+  /* Number of unassociated gparts */
+  const size_t Ndm = Ntot > 0 ? Ntot - (Ngas + Nstars) : 0;
+
+  long long N_total[swift_type_count] = {Ngas, Ndm, 0, 0, Nstars, 0};
+
+  /* File name */
+  char fileName[FILENAME_BUFFER_SIZE];
+  snprintf(fileName, FILENAME_BUFFER_SIZE, "%s_%04i.hdf5", baseName,
+           outputCount);
+
+  /* Open file */
+  /* message("Opening file '%s'.", fileName); */
+  h_file = H5Fcreate(fileName, H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
+  if (h_file < 0) {
+    error("Error while opening file '%s'.", fileName);
+  }
+
+  /* Open header to write simulation properties */
+  /* message("Writing runtime parameters..."); */
+  h_grp =
+      H5Gcreate(h_file, "/RuntimePars", H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+  if (h_grp < 0) error("Error while creating runtime parameters group\n");
+
+  /* Write the relevant information */
+  io_write_attribute(h_grp, "PeriodicBoundariesOn", INT, &periodic, 1);
+  int index = 1;
+  io_write_attribute(h_grp, "IsIndexFile", INT, &index, 1);
+
+  /* Close runtime parameters */
+  H5Gclose(h_grp);
+
+  /* Open header to write simulation properties */
+  /* message("Writing file header..."); */
+  h_grp = H5Gcreate(h_file, "/Header", H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+  if (h_grp < 0) error("Error while creating file header\n");
+
+  /* Print the relevant information and print status */
+  io_write_attribute(h_grp, "BoxSize", DOUBLE, e->s->dim, 3);
+  double dblTime = e->time;
+  io_write_attribute(h_grp, "Time", DOUBLE, &dblTime, 1);
+  int dimension = (int)hydro_dimension;
+  io_write_attribute(h_grp, "Dimension", INT, &dimension, 1);
+
+  /* GADGET-2 legacy values */
+  /* Number of particles of each type */
+  unsigned int numParticles[swift_type_count] = {0};
+  unsigned int numParticlesHighWord[swift_type_count] = {0};
+  for (int ptype = 0; ptype < swift_type_count; ++ptype) {
+    numParticles[ptype] = (unsigned int)N_total[ptype];
+    numParticlesHighWord[ptype] = (unsigned int)(N_total[ptype] >> 32);
+  }
+  io_write_attribute(h_grp, "NumPart_ThisFile", LONGLONG, N_total,
+                     swift_type_count);
+  io_write_attribute(h_grp, "NumPart_Total", UINT, numParticles,
+                     swift_type_count);
+  io_write_attribute(h_grp, "NumPart_Total_HighWord", UINT,
+                     numParticlesHighWord, swift_type_count);
+  double MassTable[swift_type_count] = {0};
+  io_write_attribute(h_grp, "MassTable", DOUBLE, MassTable, swift_type_count);
+  unsigned int flagEntropy[swift_type_count] = {0};
+  flagEntropy[0] = writeEntropyFlag();
+  io_write_attribute(h_grp, "Flag_Entropy_ICs", UINT, flagEntropy,
+                     swift_type_count);
+  io_write_attribute(h_grp, "NumFilesPerSnapshot", INT, &numFiles, 1);
+
+  /* Close header */
+  H5Gclose(h_grp);
+
+  /* Print the code version */
+  io_write_code_description(h_file);
+
+  /* Print the SPH parameters */
+  if (e->policy & engine_policy_hydro) {
+    h_grp = H5Gcreate(h_file, "/HydroScheme", H5P_DEFAULT, H5P_DEFAULT,
+                      H5P_DEFAULT);
+    if (h_grp < 0) error("Error while creating SPH group");
+    hydro_props_print_snapshot(h_grp, e->hydro_properties);
+    writeSPHflavour(h_grp);
+    H5Gclose(h_grp);
+  }
+
+  /* Print the gravity parameters */
+  if (e->policy & engine_policy_self_gravity) {
+    h_grp = H5Gcreate(h_file, "/GravityScheme", H5P_DEFAULT, H5P_DEFAULT,
+                      H5P_DEFAULT);
+    if (h_grp < 0) error("Error while creating gravity group");
+    gravity_props_print_snapshot(h_grp, e->gravity_properties);
+    H5Gclose(h_grp);
+  }
+
+  /* Print the runtime parameters */
+  h_grp =
+      H5Gcreate(h_file, "/Parameters", H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+  if (h_grp < 0) error("Error while creating parameters group");
+  parser_write_params_to_hdf5(e->parameter_file, h_grp);
+  H5Gclose(h_grp);
+
+  /* Print the system of Units used in the spashot */
+  io_write_unit_system(h_file, snapshot_units, "Units");
+
+  /* Print the system of Units used internally */
+  io_write_unit_system(h_file, internal_units, "InternalCodeUnits");
+
+  /* Tell the user if a conversion will be needed */
+  if (e->verbose) {
+    if (units_are_equal(snapshot_units, internal_units)) {
+
+      message("Snapshot and internal units match. No conversion needed.");
+
+    } else {
+
+      message("Conversion needed from:");
+      message("(Snapshot) Unit system: U_M =      %e g.",
+              snapshot_units->UnitMass_in_cgs);
+      message("(Snapshot) Unit system: U_L =      %e cm.",
+              snapshot_units->UnitLength_in_cgs);
+      message("(Snapshot) Unit system: U_t =      %e s.",
+              snapshot_units->UnitTime_in_cgs);
+      message("(Snapshot) Unit system: U_I =      %e A.",
+              snapshot_units->UnitCurrent_in_cgs);
+      message("(Snapshot) Unit system: U_T =      %e K.",
+              snapshot_units->UnitTemperature_in_cgs);
+      message("to:");
+      message("(internal) Unit system: U_M = %e g.",
+              internal_units->UnitMass_in_cgs);
+      message("(internal) Unit system: U_L = %e cm.",
+              internal_units->UnitLength_in_cgs);
+      message("(internal) Unit system: U_t = %e s.",
+              internal_units->UnitTime_in_cgs);
+      message("(internal) Unit system: U_I = %e A.",
+              internal_units->UnitCurrent_in_cgs);
+      message("(internal) Unit system: U_T = %e K.",
+              internal_units->UnitTemperature_in_cgs);
+    }
+  }
+
+  /* Loop over all particle types */
+  for (int ptype = 0; ptype < swift_type_count; ptype++) {
+
+    /* Don't do anything if no particle of this kind */
+    if (numParticles[ptype] == 0) continue;
+
+    /* Open the particle group in the file */
+    char partTypeGroupName[PARTICLE_GROUP_BUFFER_SIZE];
+    snprintf(partTypeGroupName, PARTICLE_GROUP_BUFFER_SIZE, "/PartType%d",
+             ptype);
+    h_grp = H5Gcreate(h_file, partTypeGroupName, H5P_DEFAULT, H5P_DEFAULT,
+                      H5P_DEFAULT);
+    if (h_grp < 0) {
+      error("Error while creating particle group.\n");
+    }
+
+    int num_fields = 0;
+    struct io_props list[100];
+    size_t N = 0;
+
+    /* Write particle fields from the particle structure */
+    switch (ptype) {
+
+      case swift_type_gas:
+        N = Ngas;
+	hydro_write_index(parts, list, &num_fields);
+        break;
+
+      case swift_type_dark_matter:
+        /* Allocate temporary array */
+        if (posix_memalign((void*)&dmparts, gpart_align,
+                           Ndm * sizeof(struct gpart)) != 0)
+          error("Error while allocating temporart memory for DM particles");
+        bzero(dmparts, Ndm * sizeof(struct gpart));
+
+        /* Collect the DM particles from gpart */
+        io_collect_dm_gparts(gparts, Ntot, dmparts, Ndm);
+
+        /* Write DM particles */
+        N = Ndm;
+        darkmatter_write_index(dmparts, list, &num_fields);
+        break;
+
+      case swift_type_star:
+        N = Nstars;
+	error("TODO");
+        //star_write_index(sparts, list, &num_fields);
+        break;
+
+      default:
+        error("Particle Type %d not yet supported. Aborting", ptype);
+    }
+
+    /* Write everything */
+    for (int i = 0; i < num_fields; ++i)
+      writeArray(e, h_grp, fileName, NULL, partTypeGroupName, list[i], N,
+                 internal_units, snapshot_units);
+
+    /* Free temporary array */
+    if (dmparts) {
+      free(dmparts);
+      dmparts = NULL;
+    }
+
+    /* Close particle group */
+    H5Gclose(h_grp);
+
+  }
+
+  /* message("Done writing particles..."); */
+
+  /* Close file */
+  H5Fclose(h_file);
+
+  ++outputCount;
+}
+
 #endif /* HAVE_HDF5 */
diff --git a/src/single_io.h b/src/single_io.h
index 09d5d40e5e..b64ad17c29 100644
--- a/src/single_io.h
+++ b/src/single_io.h
@@ -42,6 +42,9 @@ void write_output_single(struct engine* e, const char* baseName,
                          const struct unit_system* internal_units,
                          const struct unit_system* snapshot_units);
 
+void write_index_single(struct engine* e, const char* baseName,
+			const struct unit_system* internal_units,
+			const struct unit_system* snapshot_units);
 #endif /* HAVE_HDF5 && !WITH_MPI */
 
 #endif /* SWIFT_SINGLE_IO_H */
-- 
GitLab