Add two colour spectrum

4 files changed, 75 insertions, 4 deletions

diff --git a/libcrystfel/src/image.h b/libcrystfel/src/image.h
index 0c189cad..6b49ad91 100644
--- a/libcrystfel/src/image.h
+++ b/libcrystfel/src/image.h
@@ -46,7 +46,8 @@
 
 typedef enum {
 	SPECTRUM_TOPHAT,
-	SPECTRUM_SASE
+	SPECTRUM_SASE,
+	SPECTRUM_TWOCOLOUR
 } SpectrumType;
 
 /* Structure describing a feature in an image */
diff --git a/src/diffraction.c b/src/diffraction.c
index 92e18521..82e5c132 100644
--- a/src/diffraction.c
+++ b/src/diffraction.c
@@ -434,7 +434,7 @@ static int compare_samples(const void *a, const void *b)
 
 
 static struct sample *get_gaussian_spectrum(double eV_cen, double eV_step,
-                                            double sigma, int spec_size)
+                                            double sigma, int spec_size, double eV_start)
 {
 	struct sample *spectrum;
 	int i;
@@ -443,7 +443,12 @@ static struct sample *get_gaussian_spectrum(double eV_cen, double eV_step,
 	spectrum = malloc(spec_size * sizeof(struct sample));
 	if ( spectrum == NULL ) return NULL;
 
-	eV = eV_cen - (spec_size/2)*eV_step;
+	if (eV_start == 0) { /* eV starts at 3 sigma below the mean*/
+		eV = eV_cen - (spec_size/2)*eV_step;
+	} else {
+		eV = eV_start;
+	}
+
 	for ( i=0; i<spec_size; i++ ) {
 
 		spectrum[i].k = 1.0/ph_eV_to_lambda(eV);
@@ -568,7 +573,7 @@ struct sample *generate_SASE(struct image *image, gsl_rng *rng)
 	 * points */
 	double eV_step = 6.0*sigma/(spec_size-1);
 
-	spectrum = get_gaussian_spectrum(eV_cen, eV_step, sigma, spec_size);
+	spectrum = get_gaussian_spectrum(eV_cen, eV_step, sigma, spec_size,0);
 
 	/* Add SASE-type noise to Gaussian spectrum */
 	add_sase_noise(spectrum, spec_size, rng);
@@ -591,6 +596,60 @@ struct sample *generate_SASE(struct image *image, gsl_rng *rng)
 	return spectrum;
 }
 
+struct sample *generate_twocolour(struct image *image)
+{
+	struct sample *spectrum;
+	struct sample *spectrum1;
+	struct sample *spectrum2;
+        int i;
+
+        double halfwidth = image->bw/2; /* eV */
+
+	const int spec_size = 1024;
+        double eV_cen1;  /* Central photon energy for first colour */
+        double eV_cen2;  /* Central photon energy for second colour */
+	eV_cen1 = ph_lambda_to_eV(image->lambda) - halfwidth;
+	eV_cen2 = ph_lambda_to_eV(image->lambda) + halfwidth;
+
+        /* Hard-code sigma to be 1/5 of bandwidth */
+        double sigma = image->bw/5; /* eV */
+
+        /* 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_start = eV_cen1 - 3*sigma;
+        double eV_end = eV_cen2 + 3*sigma;
+        double eV_step = (eV_end - eV_start)/(spec_size-1);
+
+	spectrum1 = get_gaussian_spectrum(eV_cen1, eV_step, sigma, spec_size, eV_start);
+	spectrum2 = get_gaussian_spectrum(eV_cen2, eV_step, sigma, spec_size, eV_start);
+
+	spectrum = malloc(spec_size * sizeof(struct sample));
+        if ( spectrum == NULL ) return NULL;
+
+        for ( i=0; i<spec_size; i++ ) {
+                spectrum[i].weight = spectrum1[i].weight + spectrum2[i].weight;
+                spectrum[i].k = spectrum1[i].k;
+        }
+
+        /* 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,
diff --git a/src/diffraction.h b/src/diffraction.h
index 5b67c198..e397bd3a 100644
--- a/src/diffraction.h
+++ b/src/diffraction.h
@@ -57,4 +57,6 @@ extern struct sample *generate_tophat(struct image *image);
 
 extern struct sample *generate_SASE(struct image *image, gsl_rng *rng);
 
+extern struct sample *generate_twocolour(struct image *image);
+
 #endif	/* DIFFRACTION_H */
diff --git a/src/pattern_sim.c b/src/pattern_sim.c
index 08fc75dd..f23be972 100644
--- a/src/pattern_sim.c
+++ b/src/pattern_sim.c
@@ -451,6 +451,11 @@ int main(int argc, char *argv[])
 		spectrum_type = SPECTRUM_TOPHAT;
 	} else if ( strcasecmp(spectrum_str, "sase") == 0) {
 		spectrum_type = SPECTRUM_SASE;
+	} else if ( strcasecmp(spectrum_str, "twocolour") == 0 ||
+                    strcasecmp(spectrum_str, "twocolor") == 0 ||
+                    strcasecmp(spectrum_str, "twocolours") == 0 ||
+                    strcasecmp(spectrum_str, "twocolors") == 0) {
+		spectrum_type = SPECTRUM_TWOCOLOUR;
 	} else {
 		ERROR("Unrecognised spectrum type '%s'\n", spectrum_str);
 		return 1;
@@ -606,6 +611,10 @@ int main(int argc, char *argv[])
 			image.spectrum = generate_SASE(&image, rng);
 			break;
 
+                        case SPECTRUM_TWOCOLOUR :
+                        image.spectrum = generate_twocolour(&image);
+                        break;
+
 		}
 
 		/* Ensure no residual information */