From 8e2f2f44f46c18f7bd621a2ef9a3d0aa813d76d9 Mon Sep 17 00:00:00 2001
From: Thomas White <taw@physics.org>
Date: Mon, 20 Jan 2014 17:20:10 +0100
Subject: pattern_sim: Overhaul and add SASE spectrum simulation

---
 data/diffraction.cl         |  83 ++++++++++---
 libcrystfel/src/hdf5-file.c |  61 +++++++++-
 libcrystfel/src/image.h     |  15 +++
 src/diffraction-gpu.c       | 142 ++++++++++------------
 src/diffraction.c           | 281 ++++++++++++++++++++++++++++++++++++++------
 src/diffraction.h           |  12 +-
 src/pattern_sim.c           |  63 +++++++++-
 src/pattern_sim.h           |   2 +-
 tests/gpu_sim_check.c       |  11 +-
 9 files changed, 524 insertions(+), 146 deletions(-)

diff --git a/data/diffraction.cl b/data/diffraction.cl
index 75933927..78cf1cf4 100644
--- a/data/diffraction.cl
+++ b/data/diffraction.cl
@@ -3,9 +3,27 @@
  *
  * GPU calculation kernel for truncated lattice diffraction
  *
- * (c) 2006-2010 Thomas White <taw@physics.org>
+ * Copyright © 2012-2014 Deutsches Elektronen-Synchrotron DESY,
+ *                       a research centre of the Helmholtz Association.
  *
- * Part of CrystFEL - crystallography with a FEL
+ * Authors:
+ *   2009-2014 Thomas White <taw@physics.org>
+ *   2013      Alexandra Tolstikova
+ *
+ * This file is part of CrystFEL.
+ *
+ * CrystFEL is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * CrystFEL is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with CrystFEL.  If not, see <http://www.gnu.org/licenses/>.
  *
  */
 
@@ -13,14 +31,13 @@
 /* Maxmimum index to hold values up to (can be increased if necessary)
  * WARNING: Altering this value constitutes an ABI change, and means you must
  * update src/pattern_sim.h then recompile and reinstall everything. */
-#define INDMAX 120
+#define INDMAX 130
 #define IDIM (INDMAX*2 +1)
 
 #ifndef M_PI
 #define M_PI ((float)(3.14159265))
 #endif
 
-
 const sampler_t sampler_a = CLK_NORMALIZED_COORDS_TRUE
                              | CLK_ADDRESS_REPEAT
                              | CLK_FILTER_LINEAR;
@@ -122,7 +139,7 @@ float molecule_factor(global float *intensities, global float *flags,
 	signed int i;
 
 	#ifdef FLAT_INTENSITIES
-	return 1.0e5;
+	return 100.0;
 	#else
 
 	hf = cell.s0*q.x + cell.s1*q.y + cell.s2*q.z;  /* h */
@@ -154,6 +171,15 @@ float molecule_factor(global float *intensities, global float *flags,
 	val += lookup_flagged_intensity(intensities, flags, h, k, -l);
 	#endif /* PGMMM */
 
+	#ifdef PG321H
+	val += lookup_flagged_intensity(intensities, flags, h, k, l);
+	val += lookup_flagged_intensity(intensities, flags, i, h, l);
+	val += lookup_flagged_intensity(intensities, flags, k, h, -l);
+	val += lookup_flagged_intensity(intensities, flags, k, i, l);
+	val += lookup_flagged_intensity(intensities, flags, i, k, -l);
+	val += lookup_flagged_intensity(intensities, flags, h, i, -l);
+	#endif /* PG321H */
+
 	#ifdef PG6
 	val += lookup_flagged_intensity(intensities, flags,  h,  k,  l);
 	val += lookup_flagged_intensity(intensities, flags,  i,  h,  l);
@@ -248,33 +274,36 @@ float molecule_factor(global float *intensities, global float *flags,
 	val += lookup_flagged_intensity(intensities, flags, -l,  h,  k);
 	#endif /* PGM3 */
 
+	/* FIXME: Add the remaining point groups */
+
 	return val;
 	#endif /* FLAT_INTENSITIIES */
 }
 
 
-kernel void diffraction(global float *diff, float k,
+kernel void diffraction(global float *diff, float k, float weight,
                         int w, float corner_x, float corner_y,
+                        float fsx, float fsy, float ssx, float ssy,
                         float res, float clen, float16 cell,
-                        global float *intensities,
+                        global float *intensities, global float *flags,
                         read_only image2d_t func_a,
                         read_only image2d_t func_b,
                         read_only image2d_t func_c,
-                        global float *flags,
-                        float fsx, float fsy, float ssx, float ssy,
-                        float xo, float yo)
+                        local float *tmp)
 {
-	float tt;
 	float fs, ss;
 	float f_lattice, I_lattice;
 	float I_molecule;
 	float4 q;
-	float intensity;
-	int idx;
+	const int ls0 = get_local_size(0);
+	const int ls1 = get_local_size(1);
+	const int li0 = get_local_id(0);
+	const int li1 = get_local_id(1);
+	const int ls = ls0 * ls1;
 
 	/* Calculate fractional coordinates in fs/ss */
-	fs = convert_float(get_global_id(0)) + xo;
-	ss = convert_float(get_global_id(1)) + yo;
+	fs = convert_float(get_global_id(0)) / convert_float(ls0);
+	ss = convert_float(get_global_id(1)) / convert_float(ls1);
 
 	/* Get the scattering vector */
 	q = get_q(fs, ss, res, clen, k,
@@ -285,7 +314,25 @@ kernel void diffraction(global float *diff, float k,
 	I_molecule = molecule_factor(intensities, flags, cell, q);
 	I_lattice = pow(f_lattice, 2.0f);
 
-	/* Write the value to memory */
-	idx = convert_int_rtz(fs) + w*convert_int_rtz(ss);
-	diff[idx] = I_molecule * I_lattice;
+	tmp[li0 + ls0*li1] = I_molecule * I_lattice;
+
+	barrier(CLK_LOCAL_MEM_FENCE);
+
+	/* First thread in group sums the samples */
+	if ( li0 + li1 == 0 ) {
+
+		int i;
+		float sum = 0.0;
+		float val;
+		int idx;
+
+		idx = convert_int_rtz(fs) + w*convert_int_rtz(ss);
+
+		for ( i=0; i<ls; i++ ) sum += tmp[i];
+
+		val = weight * sum / convert_float(ls);
+		diff[idx] = val;
+
+	}
+
 }
diff --git a/libcrystfel/src/hdf5-file.c b/libcrystfel/src/hdf5-file.c
index 95ceb9c1..ad9495b5 100644
--- a/libcrystfel/src/hdf5-file.c
+++ b/libcrystfel/src/hdf5-file.c
@@ -295,6 +295,8 @@ int hdf5_write_image(const char *filename, struct image *image)
 	herr_t r;
 	hsize_t size[2];
 	double lambda, eV;
+	double *arr;
+	int i;
 
 	fh = H5Fcreate(filename, H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
 	if ( fh < 0 ) {
@@ -362,11 +364,7 @@ int hdf5_write_image(const char *filename, struct image *image)
 	eV = ph_lambda_to_eV(image->lambda);
 	r = H5Dwrite(dh, H5T_NATIVE_DOUBLE, H5S_ALL,
 	             H5S_ALL, H5P_DEFAULT, &eV);
-	if ( r < 0 ) {
-		H5Dclose(dh);
-		H5Fclose(fh);
-		return 1;
-	}
+
 	H5Dclose(dh);
 
 	dh = H5Dcreate2(fh, "/LCLS/photon_wavelength_A", H5T_NATIVE_DOUBLE, sh,
@@ -383,6 +381,59 @@ int hdf5_write_image(const char *filename, struct image *image)
 		H5Fclose(fh);
 		return 1;
 	}
+
+	H5Dclose(dh);
+
+	arr = malloc(image->spectrum_size*sizeof(double));
+	if ( arr == NULL ) {
+		H5Fclose(fh);
+		return 1;
+	}
+	for ( i=0; i<image->spectrum_size; i++ ) {
+		arr[i] = 1.0e10/image->spectrum[i].k;
+	}
+
+	size[0] = image->spectrum_size;
+	sh = H5Screate_simple(1, size, NULL);
+
+	dh = H5Dcreate2(gh, "spectrum_wavelengths_A", H5T_NATIVE_DOUBLE, sh,
+	                H5P_DEFAULT, H5S_ALL, H5P_DEFAULT);
+	if ( dh < 0 ) {
+		H5Fclose(fh);
+		return 1;
+	}
+	r = H5Dwrite(dh, H5T_NATIVE_DOUBLE, H5S_ALL,
+	             H5S_ALL, H5P_DEFAULT, arr);
+	H5Dclose(dh);
+
+	for ( i=0; i<image->spectrum_size; i++ ) {
+		arr[i] = image->spectrum[i].weight;
+	}
+	dh = H5Dcreate2(gh, "spectrum_weights", H5T_NATIVE_DOUBLE, sh,
+	                H5P_DEFAULT, H5S_ALL, H5P_DEFAULT);
+	if ( dh < 0 ) {
+		H5Fclose(fh);
+		return 1;
+	}
+	r = H5Dwrite(dh, H5T_NATIVE_DOUBLE, H5S_ALL,
+	             H5S_ALL, H5P_DEFAULT, arr);
+
+	H5Dclose(dh);
+	free(arr);
+
+	size[0] = 1;
+	sh = H5Screate_simple(1, size, NULL);
+
+	dh = H5Dcreate2(gh, "number_of_samples", H5T_NATIVE_INT, sh,
+	                H5P_DEFAULT, H5S_ALL, H5P_DEFAULT);
+	if ( dh < 0 ) {
+		H5Fclose(fh);
+		return 1;
+	}
+
+	r = H5Dwrite(dh, H5T_NATIVE_INT, H5S_ALL,
+	             H5S_ALL, H5P_DEFAULT, &image->nsamples);
+
 	H5Dclose(dh);
 
 	H5Gclose(gh);
diff --git a/libcrystfel/src/image.h b/libcrystfel/src/image.h
index a1ac75a3..0c189cad 100644
--- a/libcrystfel/src/image.h
+++ b/libcrystfel/src/image.h
@@ -44,6 +44,10 @@
 #include "crystal.h"
 #include "index.h"
 
+typedef enum {
+	SPECTRUM_TOPHAT,
+	SPECTRUM_SASE
+} SpectrumType;
 
 /* Structure describing a feature in an image */
 struct imagefeature {
@@ -67,6 +71,13 @@ struct imagefeature {
 /* An opaque type representing a list of image features */
 typedef struct _imagefeaturelist ImageFeatureList;
 
+/* Structure describing a wavelength sample from a spectrum */
+struct sample
+{
+	double k;
+	double weight;
+};
+
 
 /**
  * image:
@@ -144,6 +155,10 @@ struct image {
 	int                     id;   /* ID number of the thread
 	                               * handling this image */
 
+	struct sample           *spectrum;
+	int                     nsamples; /* Number of wavelengths */
+	int                     spectrum_size;  /* Size of "spectrum" */
+
 	/* Per-shot radiation values */
 	double                  lambda;        /* Wavelength in m */
 	double                  div;           /* Divergence in radians */
diff --git a/src/diffraction-gpu.c b/src/diffraction-gpu.c
index f41a45ec..ff4f2414 100644
--- a/src/diffraction-gpu.c
+++ b/src/diffraction-gpu.c
@@ -3,11 +3,13 @@
  *
  * Calculate diffraction patterns by Fourier methods (GPU version)
  *
- * Copyright © 2012 Deutsches Elektronen-Synchrotron DESY,
- *                  a research centre of the Helmholtz Association.
+ * Copyright © 2012-2014 Deutsches Elektronen-Synchrotron DESY,
+ *                       a research centre of the Helmholtz Association.
  *
  * Authors:
- *   2010-2012 Thomas White <taw@physics.org>
+ *   2009-2014 Thomas White <taw@physics.org>
+ *   2013      Alexandra Tolstikova
+ *   2013-2014 Chun Hong Yoon <chun.hong.yoon@desy.de>
  *
  * This file is part of CrystFEL.
  *
@@ -168,18 +170,20 @@ static int set_arg_mem(struct gpu_context *gctx, int idx, cl_mem val)
 
 
 static void do_panels(struct gpu_context *gctx, struct image *image,
-                      int *n_inf, int *n_neg, int *n_nan, double k, double frac,
-                      double xo, double yo)
+                      double k, double weight,
+                      int *n_inf, int *n_neg, int *n_nan)
 {
 	int i;
+	const int sampling = 4;  /* This, squared, number of samples / pixel */
 
 	if ( set_arg_float(gctx, 1, k) ) return;
+	if ( set_arg_float(gctx, 2, weight) ) return;
 
 	/* Iterate over panels */
 	for ( i=0; i<image->det->n_panels; i++ ) {
 
 		size_t dims[2];
-		size_t ldims[2] = {1, 1};
+		size_t ldims[2];
 		struct panel *p;
 		cl_mem diff;
 		size_t diff_size;
@@ -203,20 +207,29 @@ static void do_panels(struct gpu_context *gctx, struct image *image,
 		}
 
 		if ( set_arg_mem(gctx, 0, diff) ) return;
-		if ( set_arg_int(gctx, 2, pan_width) ) return;
-		if ( set_arg_float(gctx, 3, p->cnx) ) return;
-		if ( set_arg_float(gctx, 4, p->cny) ) return;
-		if ( set_arg_float(gctx, 5, p->res) ) return;
-		if ( set_arg_float(gctx, 6, p->clen) ) return;
-		if ( set_arg_float(gctx, 13, p->fsx) ) return;
-		if ( set_arg_float(gctx, 14, p->fsy) ) return;
-		if ( set_arg_float(gctx, 15, p->ssx) ) return;
-		if ( set_arg_float(gctx, 16, p->ssy) ) return;
-		if ( set_arg_float(gctx, 17, xo) ) return;
-		if ( set_arg_float(gctx, 18, yo) ) return;
-
-		dims[0] = pan_width;
-		dims[1] = pan_height;
+
+		if ( set_arg_int(gctx, 3, pan_width) ) return;
+		if ( set_arg_float(gctx, 4, p->cnx) ) return;
+		if ( set_arg_float(gctx, 5, p->cny) ) return;
+		if ( set_arg_float(gctx, 6, p->fsx) ) return;
+		if ( set_arg_float(gctx, 7, p->fsy) ) return;
+		if ( set_arg_float(gctx, 8, p->ssx) ) return;
+		if ( set_arg_float(gctx, 9, p->ssy) ) return;
+		if ( set_arg_float(gctx, 10, p->res) ) return;
+		if ( set_arg_float(gctx, 11, p->clen) ) return;
+
+		dims[0] = pan_width * sampling;
+		dims[1] = pan_height * sampling;
+
+		ldims[0] = sampling;
+		ldims[1] = sampling;
+
+		err = clSetKernelArg(gctx->kern, 18,
+		                     sampling*sampling*sizeof(cl_float), NULL);
+		if ( err != CL_SUCCESS ) {
+			ERROR("Couldn't set local memory: %s\n", clError(err));
+			return;
+		}
 
 		err = clEnqueueNDRangeKernel(gctx->cq, gctx->kern, 2, NULL,
 		                             dims, ldims, 0, NULL, NULL);
@@ -250,7 +263,7 @@ static void do_panels(struct gpu_context *gctx, struct image *image,
 
 			tfs = p->min_fs + fs;
 			tss = p->min_ss + ss;
-			image->data[tfs + image->width*tss] += frac*val;
+			image->data[tfs + image->width*tss] += val;
 
 		}
 		}
@@ -276,78 +289,54 @@ void get_diffraction_gpu(struct gpu_context *gctx, struct image *image,
 	int n_neg = 0;
 	int n_nan = 0;
 	int fs, ss;
-	const int nkstep = 10;
-	const int nxs = 4;
-	const int nys = 4;
-	int kstep;
-	double klow, khigh, kinc, w;
-	double frac;
-	int xs, ys;
-	int n, ntot;
+	int i;
 
 	if ( gctx == NULL ) {
 		ERROR("GPU setup failed.\n");
 		return;
 	}
 
-	cell_get_cartesian(ucell, &ax, &ay, &az, &bx, &by, &bz, &cx, &cy, &cz);
-	cell.s[0] = ax;  cell.s[1] = ay;  cell.s[2] = az;
-	cell.s[3] = bx;  cell.s[4] = by;  cell.s[5] = bz;
-	cell.s[6] = cx;  cell.s[7] = cy;  cell.s[8] = cz;
-
 	/* Ensure all required LUTs are available */
 	check_sinc_lut(gctx, na);
 	check_sinc_lut(gctx, nb);
 	check_sinc_lut(gctx, nc);
 
-	if ( set_arg_mem(gctx, 8, gctx->intensities) ) return;
-	if ( set_arg_mem(gctx, 9, gctx->sinc_luts[na-1]) ) return;
-	if ( set_arg_mem(gctx, 10, gctx->sinc_luts[nb-1]) ) return;
-	if ( set_arg_mem(gctx, 11, gctx->sinc_luts[nc-1]) ) return;
-	if ( set_arg_mem(gctx, 12, gctx->flags) ) return;
-
 	/* Unit cell */
-	err = clSetKernelArg(gctx->kern, 7, sizeof(cl_float16), &cell);
+	cell_get_cartesian(ucell, &ax, &ay, &az, &bx, &by, &bz, &cx, &cy, &cz);
+	cell.s[0] = ax;  cell.s[1] = ay;  cell.s[2] = az;
+	cell.s[3] = bx;  cell.s[4] = by;  cell.s[5] = bz;
+	cell.s[6] = cx;  cell.s[7] = cy;  cell.s[8] = cz;
+
+	err = clSetKernelArg(gctx->kern, 12, sizeof(cl_float16), &cell);
 	if ( err != CL_SUCCESS ) {
 		ERROR("Couldn't set unit cell: %s\n", clError(err));
 		return;
 	}
 
+	if ( set_arg_mem(gctx, 13, gctx->intensities) ) return;
+	if ( set_arg_mem(gctx, 14, gctx->flags) ) return;
+	if ( set_arg_mem(gctx, 15, gctx->sinc_luts[na-1]) ) return;
+	if ( set_arg_mem(gctx, 16, gctx->sinc_luts[nb-1]) ) return;
+	if ( set_arg_mem(gctx, 17, gctx->sinc_luts[nc-1]) ) return;
+
 	/* Allocate memory for the result */
 	image->data = calloc(image->width * image->height, sizeof(float));
 
-	w = image->lambda * image->bw;
-	klow  = 1.0 / (image->lambda + (w/2.0));  /* Smallest k */
-	khigh = 1.0 / (image->lambda - (w/2.0));  /* Largest k */
-	kinc = (khigh-klow) / (nkstep+1);
-
-	ntot = nkstep * nxs * nys;
-	frac = 1.0 / ntot;
-
-	n = 0;
-	for ( xs=0; xs<nxs; xs++ ) {
-	for ( ys=0; ys<nys; ys++ ) {
-
-		double xo, yo;
-
-		xo = (1.0/nxs) * xs;
-		yo = (1.0/nys) * ys;
+	double tot = 0.0;
+	for ( i=0; i<image->nsamples; i++ ) {
 
-		for ( kstep=0; kstep<nkstep; kstep++ ) {
+		printf("%.1f eV, weight = %.5f\n",
+		       ph_lambda_to_eV(1.0/image->spectrum[i].k),
+		       image->spectrum[i].weight);
 
-			double k;
+		do_panels(gctx, image, image->spectrum[i].k,
+		          image->spectrum[i].weight,
+		          &n_inf, &n_neg, &n_nan);
 
-			k = klow + kstep*kinc;
+		tot += image->spectrum[i].weight;
 
-			do_panels(gctx, image, &n_inf, &n_neg, &n_nan, k, frac,
-			          xo, yo);
-
-			n++;
-			progress_bar(n, ntot, "Simulating");
-
-		}
-	}
 	}
+	printf("total weight = %f\n", tot);
 
 	if ( n_neg + n_inf + n_nan ) {
 		ERROR("WARNING: The GPU calculation produced %i negative"
@@ -360,16 +349,9 @@ void get_diffraction_gpu(struct gpu_context *gctx, struct image *image,
 	for ( fs=0; fs<image->width; fs++ ) {
 	for ( ss=0; ss<image->height; ss++ ) {
 
-		double twotheta, k;
-		int idx;
-
-		/* Calculate k this time round */
-		k = 1.0/image->lambda;
-
-		get_q(image, fs, ss, &twotheta, k);
-
-		idx = fs + image->width*ss;
-		image->twotheta[idx] = twotheta;
+		double twotheta;
+		get_q(image, fs, ss, &twotheta, 1.0/image->lambda);
+		image->twotheta[fs + image->width*ss] = twotheta;
 
 	}
 	}
@@ -437,7 +419,7 @@ struct gpu_context *setup_gpu(int no_sfac,
 		}
 	} else {
 		for ( i=0; i<IDIM*IDIM*IDIM; i++ ) {
-			intensities_ptr[i] = 1e5;
+			intensities_ptr[i] = 100.0;  /* Does nothing */
 		}
 		strncat(cflags, "-DFLAT_INTENSITIES ", 511-strlen(cflags));
 	}
@@ -468,6 +450,8 @@ struct gpu_context *setup_gpu(int no_sfac,
 			strncat(cflags, "-DPG23 ", 511-strlen(cflags));
 		} else if ( strcmp(sym, "m-3") == 0 ) {
 			strncat(cflags, "-DPGM3 ", 511-strlen(cflags));
+		} else if ( strcmp(sym, "321_H") == 0 ) {
+			strncat(cflags, "-DPG321H ", 511-strlen(cflags));
 		} else {
 			ERROR("Sorry!  Point group '%s' is not currently"
 			      " supported on the GPU."
diff --git a/src/diffraction.c b/src/diffraction.c
index 2f764987..826aac77 100644
--- a/src/diffraction.c
+++ b/src/diffraction.c
@@ -3,11 +3,13 @@
  *
  * Calculate diffraction patterns by Fourier methods
  *
- * Copyright © 2012 Deutsches Elektronen-Synchrotron DESY,
- *                  a research centre of the Helmholtz Association.
+ * Copyright © 2012-2014 Deutsches Elektronen-Synchrotron DESY,
+ *                       a research centre of the Helmholtz Association.
  *
  * Authors:
- *   2009-2012 Thomas White <taw@physics.org>
+ *   2009-2014 Thomas White <taw@physics.org>
+ *   2013-2014 Chun Hong Yoon <chun.hong.yoon@desy.de>
+ *   2013      Alexandra Tolstikova
  *
  * This file is part of CrystFEL.
  *
@@ -325,7 +327,7 @@ static double molecule_factor(const double *intensities, const double *phases,
 	ld = q.u * cx + q.v * cy + q.w * cz;
 
 	/* No flags -> flat intensity distribution */
-	if ( flags == NULL ) return 1.0e5;
+	if ( flags == NULL ) return 100.0;
 
 	switch ( m ) {
 
@@ -358,18 +360,250 @@ static double molecule_factor(const double *intensities, const double *phases,
 }
 
 
+
+static void diffraction_at_k(struct image *image, const double *intensities,
+                             const double *phases, const unsigned char *flags,
+                             UnitCell *cell, GradientMethod m,
+                             const SymOpList *sym, double k,
+                             double ax, double ay, double az,
+                             double bx, double by, double bz,
+                             double cx, double cy, double cz,
+                             double *lut_a, double *lut_b, double *lut_c,
+                             double weight)
+{
+	unsigned int fs, ss;
+	const int nxs = 4;
+	const int nys = 4;
+
+	weight /= nxs*nys;
+
+	for ( fs=0; fs<image->width; fs++ ) {
+	for ( ss=0; ss<image->height; ss++ ) {
+
+		int idx;
+		double f_lattice, I_lattice;
+		double I_molecule;
+		struct rvec q;
+		double twotheta;
+		int xs, ys;
+		float xo, yo;
+
+		for ( xs=0; xs<nxs; xs++ ) {
+		for ( ys=0; ys<nys; ys++ ) {
+
+			xo = (1.0/nxs) * xs;
+			yo = (1.0/nys) * ys;
+
+			q = get_q(image, fs+xo, ss+yo, &twotheta, k);
+
+			f_lattice = lattice_factor(q, ax, ay, az,
+					           bx, by, bz,
+					           cx, cy, cz,
+					           lut_a, lut_b, lut_c);
+
+			I_molecule = molecule_factor(intensities,
+					             phases, flags, q,
+					             ax, ay, az,
+					             bx, by, bz,
+					             cx, cy, cz,
+					             m, sym);
+
+			I_lattice = pow(f_lattice, 2.0);
+
+			idx = fs + image->width*ss;
+			image->data[idx] += I_lattice * I_molecule * weight;
+			image->twotheta[idx] = twotheta;
+
+		}
+		}
+	}
+	progress_bar(fs, image->width-1, "Calculating diffraction");
+	}
+}
+
+
+static int compare_samples(const void *a, const void *b)
+{
+	struct sample *sample1 = (struct sample *)a;
+	struct sample *sample2 = (struct sample *)b;
+	if ( sample1->weight < sample2->weight ) {
+		return 1;
+	}
+	return -1;
+}
+
+
+static struct sample *get_gaussian_spectrum(double eV_cen, double eV_step,
+                                            double sigma, int spec_size)
+{
+	struct sample *spectrum;
+	int i;
+	double eV;
+
+	spectrum = malloc(spec_size * sizeof(struct sample));
+	if ( spectrum == NULL ) return NULL;
+
+	eV = eV_cen - (spec_size/2)*eV_step;
+	for ( i=0; i<spec_size; i++ ) {
+
+		spectrum[i].k = 1.0/ph_eV_to_lambda(eV);
+		spectrum[i].weight = exp(-(pow(eV - eV_cen, 2.0)
+		                         / (2.0*sigma*sigma)));
+		eV += eV_step;
+	}
+
+	return spectrum;
+}
+
+
+static int add_sase_noise(struct sample *spectrum, int nsteps)
+{
+	struct sample *noise;
+	int i, j;
+	double *gaussianNoise;
+	int shiftLim = 5;
+	double noise_mean = 0.0;
+	double noise_sigma = 1.0;
+
+	if ( shiftLim > nsteps ) shiftLim = nsteps;
+
+	noise = malloc(nsteps * sizeof(struct sample));
+	if ( noise == NULL ) return 1;
+
+	gaussianNoise = malloc(3 * nsteps * sizeof(double));
+	if ( gaussianNoise == NULL ) {
+		free(noise);
+		return 1;
+	}
+
+	/* Generate Gaussian noise of length of spectrum
+	 * (replicate on both ends for circshift below) */
+	for ( i=0; i<nsteps; i++) {
+
+		noise[i].weight = 0.0;
+
+		/* Gaussian noise with mean = 0, std = 1 */
+		gaussianNoise[i] = gaussian_noise(noise_mean, noise_sigma);
+		gaussianNoise[i+nsteps] = gaussianNoise[i];
+		gaussianNoise[i+2*nsteps] = gaussianNoise[i];
+	}
+
+	/* Sum Gaussian noise by circshifting by +/- shiftLim */
+	for ( i=nsteps; i<2*nsteps; i++ ) {
+		for ( j=-shiftLim; j<=shiftLim; j++ ) {
+			noise[i-nsteps].weight += gaussianNoise[i+j];
+		}
+	}
+
+	/* Normalize the number of circshift sum */
+	for ( i=0; i<nsteps; i++) {
+		noise[i].weight = noise[i].weight/(2*shiftLim+1);
+	}
+
+	/* Noise amplitude should have a 2 x Gaussian distribution */
+	for ( i=0; i<nsteps; i++ ) {
+		noise[i].weight = 2.0 * spectrum[i].weight * noise[i].weight;
+	}
+
+	/* Add noise to spectrum */
+	for ( i=0; i<nsteps; i++ ) {
+
+		spectrum[i].weight += noise[i].weight;
+
+		/* The final spectrum can not be negative */
+		if ( spectrum[i].weight < 0.0 ) spectrum[i].weight = 0.0;
+
+	}
+
+	return 0;
+}
+
+
+struct sample *generate_tophat(struct image *image)
+{
+	struct sample *spectrum;
+	int i;
+	double k, k_step;
+
+	double halfwidth = image->bw * image->lambda / 2.0;  /* m */
+	double mink = 1.0/(image->lambda + halfwidth);
+	double maxk = 1.0/(image->lambda - halfwidth);
+
+	spectrum = malloc(image->nsamples * sizeof(struct sample));
+	if ( spectrum == NULL ) return NULL;
+
+	k = mink;
+	k_step = (maxk-mink)/(image->nsamples-1);
+	for ( i=0; i<image->nsamples; i++ ) {
+		spectrum[i].k = k;
+		spectrum[i].weight = 1.0/(double)image->nsamples;
+		k += k_step;
+	}
+
+	image->spectrum_size = image->nsamples;
+
+	return spectrum;
+}
+
+
+struct sample *generate_SASE(struct image *image)
+{
+	struct sample *spectrum;
+	int i;
+	const int spec_size = 1024;
+	double eV_cen;  /* Central photon energy for this spectrum */
+	const double jitter_sigma_eV = 8.0;
+
+	/* Central wavelength jitters with Gaussian distribution */
+	eV_cen = gaussian_noise(ph_lambda_to_eV(image->lambda),
+	                        jitter_sigma_eV);
+
+	/* Convert FWHM to standard deviation.  Note that bandwidth is taken to
+	 * be "delta E over E" (E = photon energy), not the bandwidth in terms
+	 * of wavelength, but the difference should be very small */
+	double sigma = (image->bw*eV_cen) / (2.0*sqrt(2.0*log(2.0)));
+
+	/* The spectrum will be calculated to a resolution which spreads six
+	 * sigmas of the original (no SASE noise) Gaussian pulse over spec_size
+	 * points */
+	double eV_step = 6.0*sigma/(spec_size-1);
+
+	spectrum = get_gaussian_spectrum(eV_cen, eV_step, sigma, spec_size);
+
+	/* Add SASE-type noise to Gaussian spectrum */
+	add_sase_noise(spectrum, spec_size);
+
+	/* Normalise intensity (before taking restricted number of samples) */
+	double total_weight = 0.0;
+	for ( i=0; i<spec_size; i++ ) {
+		total_weight += spectrum[i].weight;
+	}
+	for ( i=0; i<spec_size; i++ ) {
+		spectrum[i].weight /= total_weight;
+	}
+
+	/* Sort samples in spectrum by weight.  Diffraction calculation will
+	 * take the requested number, starting from the brightest */
+	qsort(spectrum, spec_size, sizeof(struct sample), compare_samples);
+
+	image->spectrum_size = spec_size;
+
+	return spectrum;
+}
+
+
 void get_diffraction(struct image *image, int na, int nb, int nc,
                      const double *intensities, const double *phases,
                      const unsigned char *flags, UnitCell *cell,
                      GradientMethod m, const SymOpList *sym)
 {
-	unsigned int fs, ss;
 	double ax, ay, az;
 	double bx, by, bz;
 	double cx, cy, cz;
 	double *lut_a;
 	double *lut_b;
 	double *lut_c;
+	int i;
 
 	cell_get_cartesian(cell, &ax, &ay, &az, &bx, &by, &bz, &cx, &cy, &cz);
 
@@ -383,39 +617,18 @@ void get_diffraction(struct image *image, int na, int nb, int nc,
 	lut_b = get_sinc_lut(nb);
 	lut_c = get_sinc_lut(nc);
 
-	for ( fs=0; fs<image->width; fs++ ) {
-	for ( ss=0; ss<image->height; ss++ ) {
+	for ( i=0; i<image->nsamples; i++ ) {
 
-		int idx;
-		double k;
-		double f_lattice, I_lattice;
-		double I_molecule;
-		struct rvec q;
-		double twotheta;
+		printf("%.1f eV, weight = %.5f\n",
+		       ph_lambda_to_eV(1.0/image->spectrum[i].k),
+		       image->spectrum[i].weight);
 
-		/* Calculate k this time round */
-		k = 1.0/image->lambda;
+		diffraction_at_k(image, intensities, phases,
+		                flags, cell, m, sym, image->spectrum[i].k,
+		                ax, ay, az, bx, by, bz, cx, cy, cz,
+		                lut_a, lut_b, lut_c, image->spectrum[i].weight);
 
-		q = get_q(image, fs, ss, &twotheta, k);
 
-		f_lattice = lattice_factor(q, ax, ay, az,
-		                              bx, by, bz,
-		                              cx, cy, cz,
-		                              lut_a, lut_b, lut_c);
-
-		I_molecule = molecule_factor(intensities,
-		                             phases, flags, q,
-		                             ax,ay,az,bx,by,bz,cx,cy,cz,
-		                             m, sym);
-
-		I_lattice = pow(f_lattice, 2.0);
-
-		idx = fs + image->width*ss;
-		image->data[idx] = I_lattice * I_molecule;
-		image->twotheta[idx] = twotheta;
-
-	}
-	progress_bar(fs, image->width-1, "Calculating diffraction");
 	}
 
 	free(lut_a);
diff --git a/src/diffraction.h b/src/diffraction.h
index b7e34a0e..a903346e 100644
--- a/src/diffraction.h
+++ b/src/diffraction.h
@@ -3,11 +3,13 @@
  *
  * Calculate diffraction patterns by Fourier methods
  *
- * Copyright © 2012 Deutsches Elektronen-Synchrotron DESY,
- *                  a research centre of the Helmholtz Association.
+ * Copyright © 2012-2014 Deutsches Elektronen-Synchrotron DESY,
+ *                       a research centre of the Helmholtz Association.
  *
  * Authors:
- *   2009-2012 Thomas White <taw@physics.org>
+ *   2009-2014 Thomas White <taw@physics.org>
+ *   2013-2014 Chun Hong Yoon <chun.hong.yoon@desy.de>
+ *   2013      Alexandra Tolstikova
  *
  * This file is part of CrystFEL.
  *
@@ -48,4 +50,8 @@ extern void get_diffraction(struct image *image, int na, int nb, int nc,
                             const unsigned char *flags, UnitCell *cell,
                             GradientMethod m, const SymOpList *sym);
 
+extern struct sample *generate_tophat(struct image *image);
+
+extern struct sample *generate_SASE(struct image *image);
+
 #endif	/* DIFFRACTION_H */
diff --git a/src/pattern_sim.c b/src/pattern_sim.c
index 606a173c..81a94c37 100644
--- a/src/pattern_sim.c
+++ b/src/pattern_sim.c
@@ -110,6 +110,15 @@ static void show_help(const char *s)
 "     --max-size=<s>        Use <s> as the maximum crystal size in nm.\n"
 "                            --min-size is also required.\n"
 "     --no-noise            Do not calculate Poisson noise.\n"
+" -s, --sample-spectrum=<N> Use N samples from spectrum. Default 3.\n"
+" -x, --spectrum=<type>     Use <type> for the calculation of spectrum.\n"
+"                             Choose from:\n"
+"                             tophat :      Tophat spectrum. Bandwidth is\n"
+"                                           taken from beam parameters.\n"
+"                             SASE      :   SASE spectrum. Random SASE pulse \n"
+"                                           is generated from a model.\n"
+"                                           Bandwidth is taken from beam \n"
+"                                           parameters.\n"
 );
 }
 
@@ -253,7 +262,9 @@ int main(int argc, char *argv[])
 	char *outfile = NULL;
 	char *geometry = NULL;
 	char *beamfile = NULL;
+	char *spectrum_str = NULL;
 	GradientMethod grad;
+	SpectrumType spectrum_type;
 	int ndone = 0;    /* Number of simulations done (images or not) */
 	int number = 1;   /* Number used for filename of image */
 	int n_images = 1; /* Generate one image by default */
@@ -266,6 +277,7 @@ int main(int argc, char *argv[])
 	double max_size = 0.0;
 	char *sym_str = NULL;
 	SymOpList *sym;
+	int nsamples = 3;
 
 	/* Long options */
 	const struct option longopts[] = {
@@ -283,6 +295,9 @@ int main(int argc, char *argv[])
 		{"output",             1, NULL,               'o'},
 		{"geometry",           1, NULL,               'g'},
 		{"beam",               1, NULL,               'b'},
+		{"sample-spectrum",    1, NULL,               's'},
+		{"type-spectrum",      1, NULL,               'x'},
+		{"spectrum",           1, NULL,               'x'},
 		{"really-random",      0, &config_random,      1},
 		{"gpu-dev",            1, NULL,                2},
 		{"min-size",           1, NULL,                3},
@@ -291,7 +306,7 @@ int main(int argc, char *argv[])
 	};
 
 	/* Short options */
-	while ((c = getopt_long(argc, argv, "hrn:i:t:p:o:g:b:y:",
+	while ((c = getopt_long(argc, argv, "hrn:i:t:p:o:g:b:y:s:x:",
 	                        longopts, NULL)) != -1) {
 
 		switch (c) {
@@ -344,6 +359,18 @@ int main(int argc, char *argv[])
 			sym_str = strdup(optarg);
 			break;
 
+			case 's' :
+			nsamples = strtol(optarg, &rval, 10);
+			if ( *rval != '\0' ) {
+				ERROR("Invalid number of spectrum samples.\n");
+				return 1;
+			}
+			break;
+
+			case 'x' :
+			spectrum_str = strdup(optarg);
+			break;
+
 			case 2 :
 			gpu_dev = atoi(optarg);
 			break;
@@ -443,6 +470,20 @@ int main(int argc, char *argv[])
 		return 1;
 	}
 
+	if ( spectrum_str == NULL ) {
+		STATUS("You didn't specify a spectrum type, so"
+		       " I'm using a 'tophat' spectrum.\n");
+		spectrum_type = SPECTRUM_TOPHAT;
+	} else if ( strcmp(spectrum_str, "tophat") == 0) {
+		spectrum_type = SPECTRUM_TOPHAT;
+	} else if ( strcmp(spectrum_str, "SASE") == 0) {
+		spectrum_type = SPECTRUM_SASE;
+	} else {
+		ERROR("Unrecognised spectrum type '%s'\n", spectrum_str);
+		return 1;
+	}
+	free(spectrum_str);
+
 	if ( intfile == NULL ) {
 
 		/* Gentle reminder */
@@ -504,12 +545,13 @@ int main(int argc, char *argv[])
 		ERROR("Photon energy must be specified, not taken from the"
 		      " HDF5 file.  Please alter %s accordingly.\n", beamfile)
 		return 1;
-	} else {
-		double wl = ph_en_to_lambda(eV_to_J(image.beam->photon_energy));
-		image.lambda = wl;
 	}
+
+	double wl = ph_en_to_lambda(eV_to_J(image.beam->photon_energy));
+	image.lambda = wl;
 	image.bw = image.beam->bandwidth;
 	image.div = image.beam->divergence;
+	image.nsamples = nsamples;
 	free(beamfile);
 
 	/* Load unit cell */
@@ -571,6 +613,18 @@ int main(int argc, char *argv[])
 
 		cell = cell_rotate(input_cell, orientation);
 
+		switch ( spectrum_type ) {
+
+			case SPECTRUM_TOPHAT :
+			image.spectrum = generate_tophat(&image);
+			break;
+
+			case SPECTRUM_SASE :
+			image.spectrum = generate_SASE(&image);
+			break;
+
+		}
+
 		/* Ensure no residual information */
 		image.data = NULL;
 		image.twotheta = NULL;
@@ -638,6 +692,7 @@ int main(int argc, char *argv[])
 		/* Clean up */
 		free(image.data);
 		free(image.twotheta);
+
 		cell_free(cell);
 
 skip:
diff --git a/src/pattern_sim.h b/src/pattern_sim.h
index 61bd87d0..4269ea90 100644
--- a/src/pattern_sim.h
+++ b/src/pattern_sim.h
@@ -37,7 +37,7 @@
 /* Maxmimum index to hold values up to (can be increased if necessary)
  * WARNING: Altering this value constitutes an ABI change, and means you must
  * update data/diffraction.cl then recompile and reinstall everything. */
-#define INDMAX 120
+#define INDMAX 130
 
 /* Array size */
 #define IDIM (INDMAX*2 +1)
diff --git a/tests/gpu_sim_check.c b/tests/gpu_sim_check.c
index 5c45ea02..c183a2a2 100644
--- a/tests/gpu_sim_check.c
+++ b/tests/gpu_sim_check.c
@@ -142,14 +142,21 @@ int main(int argc, char *argv[])
 	beam = calloc(1, sizeof(struct beam_params));
 	beam->fluence = 1.0e15;  /* Does nothing */
 	beam->beam_radius = 1.0e-6;
-	beam->photon_energy = 9000.0;
-	beam->bandwidth = 0.1 / 100.0;
+	beam->photon_energy = 6000.0;
+	beam->bandwidth = 1.0 / 100.0;
 	beam->divergence = 0.0;
 	cpu_image.beam = beam;
 	gpu_image.beam = beam;
 
 	cpu_image.lambda = ph_en_to_lambda(eV_to_J(beam->photon_energy));
 	gpu_image.lambda = ph_en_to_lambda(eV_to_J(beam->photon_energy));
+	cpu_image.bw = beam->bandwidth;
+	gpu_image.bw = beam->bandwidth;
+
+	cpu_image.nsamples = 10;
+	gpu_image.nsamples = 10;
+	cpu_image.spectrum = generate_tophat(&cpu_image);
+	gpu_image.spectrum = generate_tophat(&gpu_image);
 
 	start = get_hires_seconds();
 	get_diffraction_gpu(gctx, &gpu_image, 8, 8, 8, cell);
-- 
cgit v1.2.3