From 05cacf92d9de7ecd13694a9af0edfbe04cef6638 Mon Sep 17 00:00:00 2001
From: Matthieu Schaller <schaller@strw.leidenuniv.nl>
Date: Tue, 23 Oct 2018 16:58:53 +0200
Subject: [PATCH] When dumping a snapshot, fish out the particles that have not
 been inhibited.

---
 src/common_io.c | 100 +++++++++++++++++++++++++++++------
 src/common_io.h |  14 ++++-
 src/single_io.c | 138 ++++++++++++++++++++++++++++++++++++------------
 3 files changed, 198 insertions(+), 54 deletions(-)

diff --git a/src/common_io.c b/src/common_io.c
index c381fdc755..65dd4cfec1 100644
--- a/src/common_io.c
+++ b/src/common_io.c
@@ -787,36 +787,102 @@ void io_duplicate_stars_gparts(struct threadpool* tp,
 }
 
 /**
- * @brief Copy every DM #gpart into the dmparts array.
+ * @brief Copy every non-inhibited #part into the parts_written array.
+ *
+ * @param parts The array of #part containing all particles.
+ * @param xparts The array of #xpart containing all particles.
+ * @param parts_written The array of #part to fill with particles we want to write.
+ * @param xparts_written The array of #xpart  to fill with particles we want to write.
+ * @param Nparts The total number of #part.
+ * @param Nparts_written The total number of #part to write.
+ */
+void io_collect_parts_to_write(const struct part* restrict parts,
+			       const struct xpart* restrict xparts,
+			       struct part* restrict parts_written,
+			       struct xpart* restrict xparts_written,
+			       const size_t Nparts, const size_t Nparts_written) {
+
+  size_t count = 0;
+
+  /* Loop over all parts */
+  for (size_t i = 0; i < Nparts; ++i) {
+
+    /* And collect the ones that have not been removed */
+    if(parts[i].time_bin != time_bin_inhibited) {
+
+      parts_written[count] = parts[i];
+      xparts_written[count] = xparts[i];
+      count++;
+    }
+  }
+
+  /* Check that everything is fine */
+  if (count != Nparts_written)
+    error("Collected the wrong number of particles (%zu vs. %zu expected)",
+          count, Nparts_written);
+}
+
+/**
+ * @brief Copy every non-inhibited #spart into the sparts_written array.
+ *
+ * @param sparts The array of #spart containing all particles.
+ * @param sparts_written The array of #spart to fill with particles we want to write.
+ * @param Nsparts The total number of #part.
+ * @param Nsparts_written The total number of #part to write.
+ */
+void io_collect_sparts_to_write(const struct spart* restrict sparts,
+				struct spart* restrict sparts_written,
+				const size_t Nsparts, const size_t Nsparts_written) {
+
+  size_t count = 0;
+
+  /* Loop over all parts */
+  for (size_t i = 0; i < Nsparts; ++i) {
+
+    /* And collect the ones that have not been removed */
+    if(sparts[i].time_bin != time_bin_inhibited) {
+
+      sparts_written[count] = sparts[i];
+      count++;
+    }
+  }
+
+  /* Check that everything is fine */
+  if (count != Nsparts_written)
+    error("Collected the wrong number of s-particles (%zu vs. %zu expected)",
+          count, Nsparts_written);
+}
+
+/**
+ * @brief Copy every non-inhibited DM #gpart into the gparts_written array.
  *
  * @param gparts The array of #gpart containing all particles.
- * @param Ntot The number of #gpart.
- * @param dmparts The array of #gpart containg DM particles to be filled.
- * @param Ndm The number of DM particles.
+ * @param gparts_written The array of #gpart to fill with particles we want to write.
+ * @param Ngparts The total number of #part.
+ * @param Ngparts_written The total number of #part to write.
  */
-void io_collect_dm_gparts(const struct gpart* const gparts, size_t Ntot,
-                          struct gpart* const dmparts, size_t Ndm) {
+void io_collect_gparts_to_write(const struct gpart* restrict gparts,
+				struct gpart* restrict gparts_written,
+				const size_t Ngparts, const size_t Ngparts_written) {
 
   size_t count = 0;
 
-  /* Loop over all gparts */
-  for (size_t i = 0; i < Ntot; ++i) {
+  /* Loop over all parts */
+  for (size_t i = 0; i < Ngparts; ++i) {
 
-    /* message("i=%zd count=%zd id=%lld part=%p", i, count, gparts[i].id,
-     * gparts[i].part); */
+    /* And collect the ones that have not been removed */
+    if((gparts[i].time_bin != time_bin_inhibited) &&
+       (gparts[i].type == swift_type_dark_matter)) {
 
-    /* And collect the DM ones that have not been removed */
-    if (gparts[i].type == swift_type_dark_matter &&
-        gparts[i].time_bin != time_bin_inhibited) {
-      dmparts[count] = gparts[i];
+      gparts_written[count] = gparts[i];
       count++;
     }
   }
 
   /* Check that everything is fine */
-  if (count != Ndm)
-    error("Collected the wrong number of dm particles (%zu vs. %zu expected)",
-          count, Ndm);
+  if (count != Ngparts_written)
+    error("Collected the wrong number of s-particles (%zu vs. %zu expected)",
+          count, Ngparts_written);
 }
 
 /**
diff --git a/src/common_io.h b/src/common_io.h
index 9ab49109c3..a79acdb16c 100644
--- a/src/common_io.h
+++ b/src/common_io.h
@@ -35,6 +35,7 @@
 struct part;
 struct gpart;
 struct spart;
+struct xpart;
 struct io_props;
 struct engine;
 struct threadpool;
@@ -91,8 +92,17 @@ void io_copy_temp_buffer(void* temp, const struct engine* e,
 size_t io_sizeof_type(enum IO_DATA_TYPE type);
 int io_is_double_precision(enum IO_DATA_TYPE type);
 
-void io_collect_dm_gparts(const struct gpart* const gparts, size_t Ntot,
-                          struct gpart* const dmparts, size_t Ndm);
+void io_collect_parts_to_write(const struct part* restrict parts,
+			       const struct xpart* restrict xparts,
+			       struct part* restrict parts_written,
+			       struct xpart* restrict xparts_written,
+			       const size_t Nparts, const size_t Nparts_written);
+void io_collect_sparts_to_write(const struct spart* restrict sparts,
+				struct spart* restrict sparts_written,
+				const size_t Nsparts, const size_t Nsparts_written);
+void io_collect_gparts_to_write(const struct gpart* restrict gparts,
+				struct gpart* restrict gparts_written,
+				const size_t Ngparts, const size_t Ngparts_written);
 void io_prepare_dm_gparts(struct threadpool* tp, struct gpart* const gparts,
                           size_t Ndm);
 void io_duplicate_hydro_gparts(struct threadpool* tp, struct part* const parts,
diff --git a/src/single_io.c b/src/single_io.c
index 63c3d7473a..4a8e9ed25e 100644
--- a/src/single_io.c
+++ b/src/single_io.c
@@ -603,9 +603,6 @@ void write_output_single(struct engine* e, const char* baseName,
                          const struct unit_system* snapshot_units) {
 
   hid_t h_file = 0, h_grp = 0;
-  const size_t Ngas = e->total_nr_parts - e->nr_inhibited_parts;
-  const size_t Nstars = e->total_nr_sparts - e->nr_inhibited_sparts;
-  const size_t Ntot = e->total_nr_gparts - e->nr_inhibited_gparts;
   int periodic = e->s->periodic;
   int numFiles = 1;
   const struct part* parts = e->s->parts;
@@ -614,12 +611,24 @@ void write_output_single(struct engine* e, const char* baseName,
   const struct spart* sparts = e->s->sparts;
   const struct cooling_function_data* cooling = e->cooling_func;
   struct swift_params* params = e->parameter_file;
-
-  /* Number of unassociated gparts */
-  const size_t Ndm = Ntot > 0 ? Ntot - (Ngas + Nstars) : 0;
-
+  
+  /* Number of particles currently in the arrays */
+  const size_t Ntot = e->s->nr_gparts;
+  const size_t Ngas = e->s->nr_parts;
+  const size_t Nstars = e->s->nr_sparts;
+  //const size_t Nbaryons = Ngas + Nstars;
+  //const size_t Ndm = Ntot > 0 ? Ntot - Nbaryons : 0;
+
+  /* Number of particles that we will write */
+  const size_t Ntot_written = e->s->nr_gparts - e->s->nr_inhibited_sparts;  
+  const size_t Ngas_written = e->s->nr_parts - e->s->nr_inhibited_parts;
+  const size_t Nstars_written = e->s->nr_sparts - e->s->nr_inhibited_gparts;
+  const size_t Nbaryons_written = Ngas_written + Nstars_written;
+  const size_t Ndm_written = Ntot_written > 0 ? Ntot_written - Nbaryons_written : 0;
+
+  /* Format things in a Gadget-friendly array */
   long long N_total[swift_type_count] = {
-      (long long)Ngas, (long long)Ndm, 0, 0, (long long)Nstars, 0};
+    (long long)Ngas_written, (long long)Ndm_written, 0, 0, (long long)Nstars_written, 0};
 
   /* File name */
   char fileName[FILENAME_BUFFER_SIZE];
@@ -827,38 +836,97 @@ void write_output_single(struct engine* e, const char* baseName,
     struct io_props list[100];
     size_t N = 0;
 
-    struct gpart* dmparts = NULL;
+    struct part* parts_written = NULL;
+    struct xpart* xparts_written = NULL;
+    struct gpart* gparts_written = NULL;
+    struct spart* sparts_written = NULL;
 
     /* Write particle fields from the particle structure */
     switch (ptype) {
 
       case swift_type_gas:
-        N = Ngas;
-        hydro_write_particles(parts, xparts, list, &num_fields);
-        num_fields += chemistry_write_particles(parts, list + num_fields);
-        num_fields +=
-            cooling_write_particles(xparts, list + num_fields, cooling);
+	{
+	  if (Ngas == Ngas_written) {
+
+	    /* No inhibted particles: easy case */
+	    N = Ngas;
+	    hydro_write_particles(parts, xparts, list, &num_fields);
+	    num_fields += chemistry_write_particles(parts, list + num_fields);
+	    num_fields += cooling_write_particles(xparts, list + num_fields, cooling);
+	  }
+	  else {
+
+	    /* Ok, we need to fish out the particles we want */
+	    N = Ngas_written;
+
+	    /* Allocate temporary arrays */
+	    if (posix_memalign((void**)&parts_written, part_align, Ngas_written * sizeof(struct part)) != 0)
+	      error("Error while allocating temporart memory for parts");
+	    if (posix_memalign((void**)&xparts_written, xpart_align, Ngas_written * sizeof(struct xpart)) != 0)
+	      error("Error while allocating temporart memory for xparts");
+
+	    /* Collect the particles we want to write */
+	    io_collect_parts_to_write(parts, xparts, parts_written, xparts_written, Ngas, Ngas_written);
+
+	    /* Select the fields to write */
+	    hydro_write_particles(parts_written, xparts_written, list, &num_fields);
+	    num_fields += chemistry_write_particles(parts_written, list + num_fields);
+	    num_fields += cooling_write_particles(xparts_written, list + num_fields, cooling);
+	  }	  
+	}
         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 non-inhibited DM particles from gpart */
-        io_collect_dm_gparts(gparts, Ntot, dmparts, Ndm);
-
-        /* Write DM particles */
-        N = Ndm;
-        darkmatter_write_particles(dmparts, list, &num_fields);
+	{
+	  if(Ntot == Ndm_written) {
+
+	    /* This is a DM-only run without inhibited particles */
+	    N = Ntot;
+	    darkmatter_write_particles(gparts, list, &num_fields);
+	  }
+	  else {
+
+	    /* Ok, we need to fish out the particles we want */
+	    N = Ndm_written;
+	  
+	    /* Allocate temporary array */
+	    if (posix_memalign((void**)&gparts_written, gpart_align, Ndm_written * sizeof(struct gpart)) != 0)
+	      error("Error while allocating temporart memory for gparts");
+	    
+	    /* Collect the non-inhibited DM particles from gpart */
+	    io_collect_gparts_to_write(gparts, gparts_written, Ntot, Ndm_written);
+	    
+	    /* Write DM particles */
+	    darkmatter_write_particles(gparts_written, list, &num_fields);
+	  }
+	}
         break;
 
       case swift_type_stars:
-        N = Nstars;
-        stars_write_particles(sparts, list, &num_fields);
-        break;
+	{
+	  if(Nstars == Nstars_written) {
+
+	    /* No inhibted particles: easy case */
+	    N = Nstars;
+	    stars_write_particles(sparts, list, &num_fields);
+	  }
+	  else {
+
+	    /* Ok, we need to fish out the particles we want */
+	    N = Nstars_written;
+
+	    /* Allocate temporary arrays */
+	    if (posix_memalign((void**)&sparts_written, spart_align, Nstars_written * sizeof(struct spart)) != 0)
+	      error("Error while allocating temporart memory for sparts");
+
+	    /* Collect the particles we want to write */
+	    io_collect_sparts_to_write(sparts, sparts_written, Nstars, Nstars_written);
+
+	    /* Select the fields to write */
+	    stars_write_particles(sparts, list, &num_fields);
+	  }
+	}
+	break;
 
       default:
         error("Particle Type %d not yet supported. Aborting", ptype);
@@ -878,12 +946,12 @@ void write_output_single(struct engine* e, const char* baseName,
                    internal_units, snapshot_units);
     }
 
-    /* Free temporary array */
-    if (dmparts) {
-      free(dmparts);
-      dmparts = NULL;
-    }
-
+    /* Free temporary arrays */
+    if(parts_written) free(parts_written);
+    if(xparts_written) free(xparts_written);
+    if(gparts_written) free(gparts_written);
+    if(sparts_written) free(sparts_written);
+	 
     /* Close particle group */
     H5Gclose(h_grp);
 
-- 
GitLab