indexamajig: Add --copy-hdf5-field option for copying something to the HDF5 file

8 files changed, 111 insertions, 5 deletions

diff --git a/contrib/alter_stream.c b/contrib/alter_stream.c
index 2c5d207b..5487c257 100644
--- a/contrib/alter_stream.c
+++ b/contrib/alter_stream.c
@@ -201,7 +201,7 @@ int main(int argc, char *argv[])
 			image.reflections = find_intersections(&image,
 			                               image.indexed_cell);
 
-			write_chunk(ofh, &image, STREAM_INTEGRATED);
+			write_chunk(ofh, &image, NULL, STREAM_INTEGRATED);
 
 		}
 
diff --git a/src/hdf5-file.c b/src/hdf5-file.c
index b43297ec..6215cc5c 100644
--- a/src/hdf5-file.c
+++ b/src/hdf5-file.c
@@ -18,6 +18,7 @@
 #include <stdio.h>
 #include <stdint.h>
 #include <hdf5.h>
+#include <assert.h>
 
 #include "image.h"
 #include "hdf5-file.h"
@@ -538,6 +539,74 @@ double get_value(struct hdfile *f, const char *name)
 }
 
 
+struct copy_hdf5_field
+{
+	char **fields;
+	int n_fields;
+	int max_fields;
+};
+
+
+struct copy_hdf5_field *new_copy_hdf5_field_list()
+{
+	struct copy_hdf5_field *n;
+
+	n = calloc(1, sizeof(struct copy_hdf5_field));
+	if ( n == NULL ) return NULL;
+
+	n->max_fields = 32;
+	n->fields = malloc(n->max_fields*sizeof(char *));
+	if ( n->fields == NULL ) {
+		free(n);
+		return NULL;
+	}
+
+	return n;
+}
+
+
+void add_copy_hdf5_field(struct copy_hdf5_field *copyme,
+                         const char *name)
+{
+	assert(copyme->n_fields < copyme->max_fields);  /* FIXME */
+
+	copyme->fields[copyme->n_fields] = strdup(name);
+	if ( copyme->fields[copyme->n_fields] == NULL ) {
+		ERROR("Failed to add field for copying '%s'\n", name);
+		return;
+	}
+
+	copyme->n_fields++;
+}
+
+
+void copy_hdf5_fields(struct hdfile *f, const struct copy_hdf5_field *copyme,
+                      FILE *fh)
+{
+	int i;
+
+	if ( copyme == NULL ) return;
+
+	for ( i=0; i<copyme->n_fields; i++ ) {
+
+		char *val;
+		char *field;
+
+		field = copyme->fields[i];
+		val = hdfile_get_string_value(f, field);
+
+		if ( field[0] == '/' ) {
+			fprintf(fh, "hdf5%s = %s\n", field, val);
+		} else {
+			fprintf(fh, "hdf5/%s = %s\n", field, val);
+		}
+
+		free(val);
+
+	}
+}
+
+
 char *hdfile_get_string_value(struct hdfile *f, const char *name)
 {
 	hid_t dh;
diff --git a/src/hdf5-file.h b/src/hdf5-file.h
index c0431dc6..811bfb1f 100644
--- a/src/hdf5-file.h
+++ b/src/hdf5-file.h
@@ -23,6 +23,8 @@
 
 struct hdfile;
 
+struct copy_hdf5_field;
+
 
 extern int hdf5_write(const char *filename, const void *data,
                       int width, int height, int type);
@@ -42,4 +44,11 @@ extern char *hdfile_get_string_value(struct hdfile *f, const char *name);
 extern int get_peaks(struct image *image, struct hdfile *f, const char *p);
 extern double get_value(struct hdfile *f, const char *name);
 
+extern struct copy_hdf5_field *new_copy_hdf5_field_list(void);
+extern void copy_hdf5_fields(struct hdfile *f,
+                             const struct copy_hdf5_field *copyme, FILE *fh);
+extern void add_copy_hdf5_field(struct copy_hdf5_field *copyme,
+                                const char *name);
+
+
 #endif	/* HDF5_H */
diff --git a/src/image.h b/src/image.h
index f60af673..4d4ac2d7 100644
--- a/src/image.h
+++ b/src/image.h
@@ -70,6 +70,7 @@ typedef struct _imagefeaturelist ImageFeatureList;
  *    struct detector         *det;
  *    struct beam_params      *beam;
  *    char                    *filename;
+ *    const struct copy_hdf5_field *copyme;
  *
  *    int                     id;
  *
@@ -112,6 +113,9 @@ typedef struct _imagefeaturelist ImageFeatureList;
  * after cell reduction or matching has been performed.  The job of the cell
  * reduction is to convert the list of candidate cells into a single indexed
  * cell, or <function>NULL</function> on failure.
+ *
+ * <structfield>copyme</structfield> represents a list of HDF5 fields to copy
+ * to the output stream.
  **/
 struct image;
 
@@ -128,6 +132,7 @@ struct image {
 	struct detector         *det;
 	struct beam_params      *beam;  /* The nominal beam parameters */
 	char                    *filename;
+	const struct copy_hdf5_field *copyme;
 
 	int                     id;   /* ID number of the thread
 	                               * handling this image */
diff --git a/src/indexamajig.c b/src/indexamajig.c
index d27ff4d4..79bf69e6 100644
--- a/src/indexamajig.c
+++ b/src/indexamajig.c
@@ -80,6 +80,7 @@ struct static_index_args
 	/* Output stream */
 	pthread_mutex_t *output_mutex;  /* Protects the output stream */
 	FILE *ofh;
+	const struct copy_hdf5_field *copyme;
 };
 
 
@@ -185,6 +186,10 @@ static void show_help(const char *s)
 "                           Example: /data/data0.\n"
 "                           Default: The first one found.\n"
 "\n"
+"\nFor time-resolved stuff, you might want to use:\n\n"
+"     --copy-hdf5-field <f>  Copy the value of field <f> into the stream. You\n"
+"                             can use this option as many times as you need.\n"
+"\n"
 "\nOptions for greater performance or verbosity:\n\n"
 "     --verbose            Be verbose about indexing.\n"
 " -j <n>                   Run <n> analyses in parallel.  Default 1.\n"
@@ -232,6 +237,7 @@ static void process_image(void *pp, int cookie)
 	image.id = cookie;
 	image.filename = filename;
 	image.det = copy_geom(pargs->static_args.det);
+	image.copyme = pargs->static_args.copyme;
 
 	pargs->indexable = 0;
 
@@ -338,7 +344,7 @@ static void process_image(void *pp, int cookie)
 	}
 
 	pthread_mutex_lock(pargs->static_args.output_mutex);
-	write_chunk(pargs->static_args.ofh, &image,
+	write_chunk(pargs->static_args.ofh, &image, hdfile,
 	            pargs->static_args.stream_flags);
 	pthread_mutex_unlock(pargs->static_args.output_mutex);
 
@@ -528,6 +534,13 @@ int main(int argc, char *argv[])
 	int cpu_offset = 0;
 	char *endptr;
 	char *hdf5_peak_path = NULL;
+	struct copy_hdf5_field *copyme;
+
+	copyme = new_copy_hdf5_field_list();
+	if ( copyme == NULL ) {
+		ERROR("Couldn't allocate HDF5 field list.\n");
+		return 1;
+	}
 
 	/* Long options */
 	const struct option longopts[] = {
@@ -562,6 +575,7 @@ int main(int argc, char *argv[])
 		{"bg-sub",             0, &config_bgsub,       1}, /* Compat */
 		{"no-bg-sub",          0, &config_bgsub,       0},
 		{"hdf5-peaks",         1, NULL,                9},
+		{"copy-hdf5-field",    1, NULL,               10},
 		{0, 0, NULL, 0}
 	};
 
@@ -676,6 +690,10 @@ int main(int argc, char *argv[])
 			hdf5_peak_path = strdup(optarg);
 			break;
 
+		case 10 :
+			add_copy_hdf5_field(copyme, optarg);
+			break;
+
 		case 0 :
 			break;
 
@@ -891,6 +909,7 @@ int main(int argc, char *argv[])
 	qargs.static_args.element = element;
 	qargs.static_args.stream_flags = stream_flags;
 	qargs.static_args.hdf5_peak_path = hdf5_peak_path;
+	qargs.static_args.copyme = copyme;
 
 	qargs.fh = fh;
 	qargs.prefix = prefix;
diff --git a/src/partial_sim.c b/src/partial_sim.c
index 99b0e0cf..b3e8561f 100644
--- a/src/partial_sim.c
+++ b/src/partial_sim.c
@@ -218,7 +218,7 @@ static void finalise_job(void *vqargs, void *vwargs)
 	struct worker_args *wargs = vwargs;
 	struct queue_args *qargs = vqargs;
 
-	write_chunk(qargs->stream, &wargs->image, STREAM_INTEGRATED);
+	write_chunk(qargs->stream, &wargs->image, NULL, STREAM_INTEGRATED);
 
 	reflist_free(wargs->image.reflections);
 	cell_free(wargs->image.indexed_cell);
diff --git a/src/stream.c b/src/stream.c
index af92fc14..a7cdc2d9 100644
--- a/src/stream.c
+++ b/src/stream.c
@@ -196,7 +196,7 @@ static void write_peaks(struct image *image, FILE *ofh)
 }
 
 
-void write_chunk(FILE *ofh, struct image *i, int f)
+void write_chunk(FILE *ofh, struct image *i, struct hdfile *hdfile, int f)
 {
 	double asx, asy, asz;
 	double bsx, bsy, bsz;
@@ -252,6 +252,8 @@ void write_chunk(FILE *ofh, struct image *i, int f)
 
 	}
 
+	copy_hdf5_fields(hdfile, i->copyme, ofh);
+
 	if ( (f & STREAM_PEAKS)
 	  || ((f & STREAM_PEAKS_IF_INDEXED) && (i->indexed_cell != NULL))
 	  || ((f & STREAM_PEAKS_IF_NOT_INDEXED) && (i->indexed_cell == NULL)) )
diff --git a/src/stream.h b/src/stream.h
index 3c77cfef..ba218fb9 100644
--- a/src/stream.h
+++ b/src/stream.h
@@ -18,6 +18,7 @@
 
 
 struct image;
+struct hdfile;
 
 /* Possible options dictating what goes into the output stream */
 enum
@@ -34,7 +35,8 @@ extern int count_patterns(FILE *fh);
 
 extern void write_stream_header(FILE *ofh, int argc, char *argv[]);
 
-extern void write_chunk(FILE *ofh, struct image *image, int flags);
+extern void write_chunk(FILE *ofh, struct image *image, struct hdfile *hdfile,
+                        int flags);
 
 extern int parse_stream_flags(const char *a);