Implement named structure in numpy and implement them for GEAR and SPHENIX

logger/chemistry/GEAR/logger_chemistry.h

@@ -147,7 +147,16 @@ chemistry_logger_interpolate_field_spart(
 __attribute__((always_inline)) INLINE static void
 chemistry_logger_generate_python_part(struct logger_python_field *fields) {
   fields[chemistry_logger_field_part_all] =
-      logger_loader_python_field(2 * GEAR_CHEMISTRY_ELEMENT_COUNT, NPY_DOUBLE);
+      logger_loader_python_field(/* Dimension */ 2, CUSTOM_NPY_TYPE);
+
+  logger_loader_python_field_add_subfield(
+      &fields[chemistry_logger_field_part_all], "SmoothedMetalMassFractions",
+      /* offset */ 0, GEAR_CHEMISTRY_ELEMENT_COUNT, NPY_DOUBLE);
+
+  logger_loader_python_field_add_subfield(
+      &fields[chemistry_logger_field_part_all], "MetalMassFractions",
+      GEAR_CHEMISTRY_ELEMENT_COUNT * sizeof(double),
+      GEAR_CHEMISTRY_ELEMENT_COUNT, NPY_DOUBLE);
 }
 
 /**

logger/hydro/SPHENIX/logger_hydro.c

@@ -34,8 +34,6 @@ const int hydro_logger_field_size[hydro_logger_field_count] = {
     member_size(struct part, u),        // InternalEnergies
     member_size(struct part, id),       // IDs
     member_size(struct part, rho),      // density
-    sizeof(float),                      // Entropy
-    sizeof(float),                      // Pressure
-    3 * sizeof(float),                  // Viscosity / diffusion
-    2 * sizeof(float),                  // Velocity divergence + deriv
+    7 * sizeof(float),                  // entropy + pressure +
+    // viscosity + diffusion + laplacian u + Velocity divergence + deriv
 };

logger/hydro/SPHENIX/logger_hydro.h

@@ -104,7 +104,6 @@ hydro_logger_interpolate_field(const double t_before,
                                  /* dimension= */ 3);
       break;
     case hydro_logger_field_accelerations:
-    case hydro_logger_field_viscosity_diffusion:
       interpolate_linear_float_ND(t_before, before, t_after, after, output, t,
                                   /* dimension= */ 3);
       break;
@@ -113,8 +112,6 @@ hydro_logger_interpolate_field(const double t_before,
     case hydro_logger_field_smoothing_lengths:
     case hydro_logger_field_internal_energies:
     case hydro_logger_field_densities:
-    case hydro_logger_field_entropies:
-    case hydro_logger_field_pressures:
       interpolate_linear_float(t_before, before, t_after, after, output, t);
       break;
       /* Check the ids */
@@ -122,19 +119,27 @@ hydro_logger_interpolate_field(const double t_before,
       interpolate_ids(t_before, before, t_after, after, output, t);
       break;
 
-    case hydro_logger_field_velocity_divergences: {
+    case hydro_logger_field_secondary_fields: {
       /* Get some variables */
       const float wa = (t - t_before) / (t_after - t_before);
       const float wb = 1. - wa;
       float *x = (float *)output;
-      const float *div_bef = (float *)before->field;
-      const float *div_aft = (float *)after->field;
-
-      /* Use cubic hermite spline. */
-      x[0] = interpolate_cubic_hermite_spline(
-          t_before, div_bef[0], div_bef[1], t_after, div_aft[0], div_aft[1], t);
-      /* Use the linear interpolation */
-      x[1] = wa * div_aft[1] + wb * div_bef[1];
+      const float *bef = (float *)before->field;
+      const float *aft = (float *)after->field;
+
+      /* Entropy + pressure + Viscosity + Diffusion + Laplacian*/
+      const int n_linear = 5;
+      for (int i = 0; i < n_linear; i++) {
+        x[i] = wa * aft[i] + wb * bef[i];
+      }
+
+      /* Div v */
+      x[n_linear] = interpolate_cubic_hermite_spline(
+          t_before, bef[n_linear], bef[n_linear + 1], t_after, aft[n_linear],
+          aft[n_linear + 1], t);
+
+      /* d Div v / dt */
+      x[n_linear + 1] = wa * aft[n_linear + 1] + wb * bef[n_linear + 1];
       break;
     }
 
@@ -163,14 +168,39 @@ __attribute__((always_inline)) INLINE static void hydro_logger_generate_python(
       logger_loader_python_field(/* Dimension */ 1, NPY_LONGLONG);
   fields[hydro_logger_field_densities] =
       logger_loader_python_field(/* Dimension */ 1, NPY_FLOAT32);
-  fields[hydro_logger_field_entropies] =
-      logger_loader_python_field(/* Dimension */ 1, NPY_FLOAT32);
-  fields[hydro_logger_field_pressures] =
-      logger_loader_python_field(/* Dimension */ 1, NPY_FLOAT32);
-  fields[hydro_logger_field_viscosity_diffusion] =
-      logger_loader_python_field(/* Dimension */ 3, NPY_FLOAT32);
-  fields[hydro_logger_field_velocity_divergences] =
-      logger_loader_python_field(/* Dimension */ 2, NPY_FLOAT32);
+
+  /* Now separate the secondary fields */
+  fields[hydro_logger_field_secondary_fields] =
+      logger_loader_python_field(/* Dimension */ 7, CUSTOM_NPY_TYPE);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields], "Entropies", /* offset */ 0,
+      /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields], "Pressures",
+      /* offset */ sizeof(float), /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields], "ViscosityParameters",
+      /* offset */ 2 * sizeof(float), /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields], "DiffusionParameters",
+      /* offset */ 3 * sizeof(float), /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields], "LaplacianInternalEnergies",
+      /* offset */ 4 * sizeof(float), /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields], "VelocityDivergences",
+      /* offset */ 5 * sizeof(float), /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[hydro_logger_field_secondary_fields],
+      "VelocityDivergenceTimeDifferentials", /* offset */ 6 * sizeof(float),
+      /* Dimension */ 1, NPY_FLOAT32);
 }
 
 #endif  // HAVE_PYTHON

logger/logger_python_tools.h

@@ -28,14 +28,39 @@
 /* Include numpy */
 #include <numpy/arrayobject.h>
 
+#define CUSTOM_NPY_TYPE -1
+
+/* Structure that allows the user to easily define
+   subfield in a numpy array */
+struct logger_python_subfield {
+  /* Name of the subfield */
+  char name[STRING_SIZE];
+
+  /* Offset in the bytes */
+  size_t offset;
+
+  /* Numpy typenum of the subfield (e.g. NPY_FLOAT32, NPY_INT32). */
+  int typenum;
+
+  /* Number of dimensions for the subfield */
+  int dimension;
+};
+
 /* Structure that allows the user to easily define the
    fields in a numpy array. */
 struct logger_python_field {
   /* Dimension of the field (e.g. 1 for density and 3 for coordinates). */
   int dimension;
 
-  /* Numpy typenum of the array (e.g. NPY_FLOAT32, NPY_INT32). */
+  /* Numpy typenum of the array (e.g. NPY_FLOAT32, NPY_INT32 or
+   * CUSTOM_NPY_TYPE). */
   int typenum;
+
+  /* Number of subfields registred */
+  int subfields_registred;
+
+  /* List of subtypes for custom fields */
+  struct logger_python_subfield *subfields;
 };
 
 /**
@@ -53,10 +78,120 @@ logger_loader_python_field(int dimension, int numpy_type) {
 
   ret.typenum = numpy_type;
   ret.dimension = dimension;
+  ret.subfields_registred = 0;
+  ret.subfields = NULL;
+
+  /* We are done with basic types */
+  if (numpy_type != CUSTOM_NPY_TYPE) return ret;
+
+  /* Allocate the memory of the subfields if needed */
+  ret.subfields = (struct logger_python_subfield *)malloc(
+      dimension * sizeof(struct logger_python_subfield));
+
+  if (ret.subfields == NULL) {
+    error_python("Failed to allocate memory for a subfield");
+  }
 
   return ret;
 }
 
+/**
+ * @brief Define a subtype for a numpy array.
+ *
+ * @param field The #logger_python_field where to add the subfield.
+ * @param dimension The number of dimension for the field.
+ * @param type The numpy data type (e.g. NPY_FLOAT32, NPY_DOUBLE, NPY_LONGLONG,
+ * ...)
+ * @param name The name of the subfield.
+ */
+__attribute__((always_inline)) INLINE static void
+logger_loader_python_field_add_subfield(struct logger_python_field *field,
+                                        const char *name, size_t offset,
+                                        int dimension, int numpy_type) {
+
+  /* Ensure that we still have enough place */
+  if (field->subfields_registred == field->dimension) {
+    error_python(
+        "Trying to register too many subfields"
+        " (currently registering %s)",
+        name);
+  }
+
+  /* Get the current subfield and register it */
+  struct logger_python_subfield *subfield =
+      &field->subfields[field->subfields_registred];
+  field->subfields_registred += 1;
+
+  /* Set the variables for the subfield */
+  subfield->dimension = dimension;
+  subfield->typenum = numpy_type;
+  subfield->offset = offset;
+  strcpy(subfield->name, name);
+}
+
+/**
+ * @brief Free the memory allocated for the field.
+ *
+ * @param field The #logger_python_field to clean.
+ */
+__attribute__((always_inline)) INLINE static void
+logger_loader_python_field_free(struct logger_python_field *field) {
+  if (field->typenum == CUSTOM_NPY_TYPE) {
+    free(field->subfields);
+    field->subfields = NULL;
+  }
+}
+
+/**
+ * @brief Generate a python string for the format provided.
+ *
+ * @param format (output) The format string.
+ * @param typenum The C typenum (e.g. NPY_FLOAT32, NPY_DOUBLE, ...).
+ * @param dimension The number of dimensions.
+ */
+__attribute__((always_inline)) INLINE static PyObject *
+logger_python_tools_get_format_string(int typenum, int dimension) {
+
+  const char *type_format = NULL;
+  switch (typenum) {
+    case NPY_FLOAT32:
+      type_format = "f4";
+      break;
+    case NPY_FLOAT64:
+      type_format = "f8";
+      break;
+
+    case NPY_LONGLONG:
+      switch (sizeof(long long)) {
+        case 8:
+          type_format = "i8";
+          break;
+        case 4:
+          type_format = "i4";
+          break;
+
+        default:
+          error_python("This size of long long (%li) is not yet supported",
+                       sizeof(long long));
+      }
+      break;
+
+    default:
+      error_python("Format not implemented");
+  }
+
+  /* Now construct the complete string */
+  char format[STRING_SIZE];
+  if (dimension == 1) {
+    strcpy(format, type_format);
+  } else {
+    sprintf(format, "%i%s", dimension, type_format);
+  }
+
+  /* Convert it into python */
+  return PyUnicode_FromString(format);
+}
+
 #endif  // HAVE_PYTHON
 
 #endif  // SWIFT_LOGGER_PYTHON_TOOLS_H

logger/logger_python_wrapper.c

@@ -230,16 +230,74 @@ logger_loader_create_output(void **output, const int *field_indices,
       error_python("Failed to find the required field");
     }
     PyObject *array = NULL;
+    /* Simple case where we only have normal type */
+    if (current_field->typenum != CUSTOM_NPY_TYPE) {
       if (current_field->dimension > 1) {
         npy_intp dims[2] = {n_tot, current_field->dimension};
-      array =
-          PyArray_SimpleNewFromData(2, dims, current_field->typenum, output[i]);
+        array = PyArray_SimpleNewFromData(2, dims, current_field->typenum,
+                                          output[i]);
       } else {
         npy_intp dims = n_tot;
         array = PyArray_SimpleNewFromData(1, &dims, current_field->typenum,
                                           output[i]);
       }
+    }
+    /* Now the complex types */
+    else {
+      /* Ensure that the field is properly defined */
+      if (current_field->subfields_registred != current_field->dimension) {
+        error_python(
+            "It seems that you forgot to register a field (found %i and "
+            "expecting %i)",
+            current_field->subfields_registred, current_field->dimension);
+      }
+
+      /* Creates the dtype */
+      PyObject *names = PyList_New(current_field->dimension);
+      PyObject *formats = PyList_New(current_field->dimension);
+      PyObject *offsets = PyList_New(current_field->dimension);
+
+      /* Gather the data into the lists */
+      for (int k = 0; k < current_field->dimension; k++) {
+        struct logger_python_subfield *subfield = &current_field->subfields[k];
+
+        /* Transform the information into python */
+        PyObject *name = PyUnicode_FromString(subfield->name);
+        PyObject *offset = PyLong_FromSize_t(subfield->offset);
+        PyObject *format = logger_python_tools_get_format_string(
+            subfield->typenum, subfield->dimension);
+
+        /* Add everything into the lists */
+        PyList_SetItem(names, k, name);
+        PyList_SetItem(formats, k, format);
+        PyList_SetItem(offsets, k, offset);
+      }
+
+      /* Now create the dtype dictionary */
+      PyObject *dtype = PyDict_New();
+      PyDict_SetItemString(dtype, "names", names);
+      PyDict_SetItemString(dtype, "formats", formats);
+      PyDict_SetItemString(dtype, "offsets", offsets);
+
+      /* Cleanup a bit */
+      Py_DECREF(names);
+      Py_DECREF(formats);
+      Py_DECREF(offsets);
+
+      /* Now get the pyarray_descr */
+      PyArray_Descr *descr = NULL;
+      PyArray_DescrConverter(dtype, &descr);
+
+      /* Create the list of flags */
+
+      /* Create the array */
+      npy_intp dims = n_tot;
+      array = PyArray_NewFromDescr(&PyArray_Type, descr, /* nd */ 1, &dims,
+                                   /* strides */ NULL, output[i],
+                                   NPY_ARRAY_CARRAY, NULL);
+    }
 
+    logger_loader_python_field_free(current_field);
     PyList_SetItem(list, i, array);
   }
 

logger/star_formation/GEAR/logger_star_formation.h

@@ -95,9 +95,20 @@ star_formation_logger_interpolate_field(
 #ifdef HAVE_PYTHON
 __attribute__((always_inline)) INLINE static void
 star_formation_logger_generate_python(struct logger_python_field *fields) {
-  // TODO split fields
   fields[star_formation_logger_field_all] =
       logger_loader_python_field(/* Dimension */ 2, NPY_FLOAT);
+
+  logger_loader_python_field_add_subfield(
+      &fields[star_formation_logger_field_all], "BirthDensities",
+      /* offset */ 0, /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[star_formation_logger_field_all], "BirthMasses",
+      /* offset */ sizeof(float), /* Dimension */ 1, NPY_FLOAT32);
+
+  logger_loader_python_field_add_subfield(
+      &fields[star_formation_logger_field_all], "ProgenitorIDs",
+      /* offset */ 2 * sizeof(float), /* Dimension */ 1, NPY_LONGLONG);
 }
 
 #endif  // HAVE_PYTHON

src/hydro/SPHENIX/hydro_logger.c

@@ -31,8 +31,7 @@
 const char *hydro_logger_field_names[hydro_logger_field_count] = {
     "Coordinates", "Velocities",       "Accelerations",
     "Masses",      "SmoothingLengths", "InternalEnergies",
-    "ParticleIDs", "Densities",          "Entropies",
-    "Pressures",   "ViscosityDiffusion", "VelocityDivergences",
+    "ParticleIDs", "Densities",        "SecondaryFields",
 };
 
 #endif  // WITH_LOGGER

src/hydro/SPHENIX/hydro_logger.h

@@ -38,10 +38,12 @@ enum hydro_logger_fields {
   hydro_logger_field_internal_energies,
   hydro_logger_field_particle_ids,
   hydro_logger_field_densities,
-  hydro_logger_field_entropies,
-  hydro_logger_field_pressures,
-  hydro_logger_field_viscosity_diffusion,
-  hydro_logger_field_velocity_divergences,
+  /* Due to the low number of masks available,
+   * we group the fields together: pressure, entopries,
+   * viscosity and diffusion parameters, laplacian of the energy
+   * and the velocity divergence along its derivative.
+   */
+  hydro_logger_field_secondary_fields,
   hydro_logger_field_count,
 };
 
@@ -90,19 +92,10 @@ INLINE static int hydro_logger_writer_populate_mask_data(
   mask_data[hydro_logger_field_densities] = logger_create_mask_entry(
       hydro_logger_field_names[hydro_logger_field_densities], sizeof(float));
 
-  mask_data[hydro_logger_field_entropies] = logger_create_mask_entry(
-      hydro_logger_field_names[hydro_logger_field_entropies], sizeof(float));
-
-  mask_data[hydro_logger_field_pressures] = logger_create_mask_entry(
-      hydro_logger_field_names[hydro_logger_field_pressures], sizeof(float));
-
-  mask_data[hydro_logger_field_viscosity_diffusion] = logger_create_mask_entry(
-      hydro_logger_field_names[hydro_logger_field_viscosity_diffusion],
-      3 * sizeof(float));
-
-  mask_data[hydro_logger_field_velocity_divergences] = logger_create_mask_entry(
-      hydro_logger_field_names[hydro_logger_field_velocity_divergences],
-      2 * sizeof(float));
+  const size_t size_secondary = 7 * sizeof(float);
+  mask_data[hydro_logger_field_secondary_fields] = logger_create_mask_entry(
+      hydro_logger_field_names[hydro_logger_field_secondary_fields],
+      size_secondary);
 
   return hydro_logger_field_count;
 }
@@ -160,21 +153,9 @@ INLINE static void hydro_logger_compute_size_and_mask(
   *mask |= logger_add_field_to_mask(masks[hydro_logger_field_densities],
                                     buffer_size);
 
-  /* Add the entropies. */
-  *mask |= logger_add_field_to_mask(masks[hydro_logger_field_entropies],
+  /* Add the secondary fields. */
+  *mask |= logger_add_field_to_mask(masks[hydro_logger_field_secondary_fields],
                                     buffer_size);
-
-  /* Add the pressures. */
-  *mask |= logger_add_field_to_mask(masks[hydro_logger_field_pressures],
-                                    buffer_size);
-
-  /* Add the viscosity / diffusion. */
-  *mask |= logger_add_field_to_mask(
-      masks[hydro_logger_field_viscosity_diffusion], buffer_size);
-
-  /* Add the velocity divergences + derivative. */
-  *mask |= logger_add_field_to_mask(
-      masks[hydro_logger_field_velocity_divergences], buffer_size);
 }
 
 /**
@@ -262,40 +243,21 @@ INLINE static char *hydro_logger_write_particle(
     buff += sizeof(float);
   }
 
-  /* Write the entropy. */
-  if (logger_should_write_field(mask_data[hydro_logger_field_entropies],
+  /* Write the secondary fields. */
+  if (logger_should_write_field(mask_data[hydro_logger_field_secondary_fields],
                                 mask)) {
-    const float entropy = hydro_get_comoving_entropy(p, xp);
-    memcpy(buff, &entropy, sizeof(float));
-    buff += sizeof(float);
-  }
-
-  /* Write the pressures. */
-  if (logger_should_write_field(mask_data[hydro_logger_field_pressures],
-                                mask)) {
-    const float pressure = hydro_get_comoving_pressure(p);
-    memcpy(buff, &pressure, sizeof(float));
-    buff += sizeof(float);
-  }
-
-  /* Write the viscosity / diffusion. */
-  if (logger_should_write_field(
-          mask_data[hydro_logger_field_viscosity_diffusion], mask)) {
-    const float coef[3] = {p->viscosity.alpha * p->force.balsara,
-                           p->diffusion.alpha, p->diffusion.laplace_u};
-    memcpy(buff, coef, sizeof(coef));
-    buff += sizeof(coef);
-  }
-
-  /* Write the velocity divergence. */
-  if (logger_should_write_field(
-          mask_data[hydro_logger_field_velocity_divergences], mask)) {
-    const float div[2] = {
+    const float secondary[7] = {
+        hydro_get_comoving_entropy(p, xp),
+        hydro_get_comoving_pressure(p),
+        p->viscosity.alpha * p->force.balsara,
+        p->diffusion.alpha,
+        p->diffusion.laplace_u,
         p->viscosity.div_v,
         p->viscosity.div_v_dt,
     };
-    memcpy(buff, div, sizeof(div));
-    buff += sizeof(div);
+
+    memcpy(buff, &secondary, sizeof(secondary));
+    buff += sizeof(secondary);
   }
 
   return buff;