process_hkl,partialator: Allow arbitrary direction and degree of polarisation

6 files changed, 147 insertions, 32 deletions

diff --git a/doc/man/partialator.1 b/doc/man/partialator.1
index f0491616..63b7578e 100644
--- a/doc/man/partialator.1
+++ b/doc/man/partialator.1
@@ -93,9 +93,16 @@ Specify the partiality model.  See the list below for possible choices.
 Run \fIn\fR analyses in parallel.
 
 .PD 0
+.IP \fB--polarisation=\fItype\fR
+.PD
+Specify the polarisation of the incident radiation.  \fItype\fR can be \fBhoriz\fR, \fBvert\fR or \fBnone\fR to indicate 100% polarisation of the electric field in the horizontal plane, vertical plane or completely unpolarised radiation respectively.  Alternatively, \fItype\fR can be a direction followed by a percentage polarisation fraction.  For example, \fB45deg90\fR means that 90% of the radiation is polarised with its electric field in a direction 45 degrees from horizontal, and \fB10deg100\fR means that all the radiation is polarised at 10 degrees from horizontal.  The angle is specified clockwise from horizontal as viewed along the beam direction, i.e. as shown by \fBhdfsee\fR.  The beam is unpolarised when the fraction is 50% (equal parts of the radiation have their electric field in the specified plane).  If the polarisation fraction is 100%, it can be omitted.  For example \fB10deg\fR or \fBhoriz\fR.
+
+The default is \fB--polarisation=horiz\fR.
+
+.PD 0
 .IP \fB--no-polarisation\fR
 .PD
-Disable the polarisation correction.
+Synonym for \fB--polarisation=none\fR.
 
 .PD 0
 .IP \fB--max-adu=\fR\fIn\fR
diff --git a/doc/man/process_hkl.1 b/doc/man/process_hkl.1
index 658f42bf..d01cfc5d 100644
--- a/doc/man/process_hkl.1
+++ b/doc/man/process_hkl.1
@@ -84,9 +84,16 @@ Perform a second pass through the input, scaling each crystal's intensities to b
 Use \fBpartialator\fR if you need more advanced merging techniques.
 
 .PD 0
+.IP \fB--polarisation=\fItype\fR
+.PD
+Specify the polarisation of the incident radiation.  \fItype\fR can be \fBhoriz\fR, \fBvert\fR or \fBnone\fR to indicate 100% polarisation of the electric field in the horizontal plane, vertical plane or completely unpolarised radiation respectively.  Alternatively, \fItype\fR can be a direction followed by a percentage polarisation fraction.  For example, \fB45deg90\fR means that 90% of the radiation is polarised with its electric field in a direction 45 degrees from horizontal, and \fB10deg100\fR means that all the radiation is polarised at 10 degrees from horizontal.  The angle is specified clockwise from horizontal as viewed along the beam direction, i.e. as shown by \fBhdfsee\fR.  The beam is unpolarised when the fraction is 50% (equal parts of the radiation have their electric field in the specified plane).  If the polarisation fraction is 100%, it can be omitted.  For example \fB10deg\fR or \fBhoriz\fR.
+
+The default is \fB--polarisation=horiz\fR.
+
+.PD 0
 .IP \fB--no-polarisation\fR
 .PD
-Disable the polarisation correction.
+Synonym for \fB--polarisation=none\fR.
 
 .PD 0
 .IP \fB--min-measurements=\fR\fIn\fR
diff --git a/libcrystfel/src/geometry.c b/libcrystfel/src/geometry.c
index 4d25b89f..b99fdd04 100644
--- a/libcrystfel/src/geometry.c
+++ b/libcrystfel/src/geometry.c
@@ -36,6 +36,8 @@
 #include <fenv.h>
 #include <gsl/gsl_sf_erf.h>
 #include <gsl/gsl_linalg.h>
+#include <ctype.h>
+#include <math.h>
 
 #include "utils.h"
 #include "cell.h"
@@ -952,7 +954,72 @@ void update_predictions(Crystal *cryst)
 }
 
 
-void polarisation_correction(RefList *list, UnitCell *cell, struct image *image)
+struct polarisation parse_polarisation(const char *text)
+{
+	struct polarisation p;
+	char angle[14];
+	char deg[14];
+	char frac[14];
+	int i, n;
+
+	if ( strlen(text) > 13 ) goto err;
+
+	if ( strcmp(text, "none") == 0 ) {
+		p.fraction = 0.5;
+		p.angle = 0.0;
+		return p;
+	}
+
+	p.fraction = 1.0;
+
+	i = 0;
+	n = 0;
+	while ( isdigit(text[i]) ) {
+		angle[n++] = text[i++];
+	}
+	angle[n] = '\0';
+
+	n = 0;
+	while ( isalpha(text[i]) ) {
+		deg[n++] = text[i++];
+	}
+	deg[n] = '\0';
+
+	n = 0;
+	while ( isdigit(text[i]) ) {
+		frac[n++] = text[i++];
+	}
+	frac[n] = '\0';
+
+	if ( text[i] != '\0' ) goto err;
+
+	if ( frac[0] != '\0' ) {
+		int nfrac = atoi(frac);
+		if ( nfrac > 100 ) goto err;
+		if ( nfrac < 0 ) goto err;
+		p.fraction = nfrac / 100.0;
+	}
+
+	if ( strcmp(deg, "deg") == 0 ) {
+		p.angle = deg2rad(atoi(angle));
+	} else {
+		if ( angle[0] != '\0' ) goto err;
+		if ( strncmp(deg, "horiz", 5) == 0 ) p.angle = 0.0;
+		if ( strncmp(deg, "vert", 4) == 0 ) p.angle = M_PI_2;
+	}
+
+	return p;
+
+err:
+	ERROR("Invalid polarisation '%s'\n", text);
+	p.fraction = 0.5;
+	p.angle = 0.0;
+	return p;
+}
+
+
+void polarisation_correction(RefList *list, UnitCell *cell, double lambda,
+                             struct polarisation p)
 {
 	Reflection *refl;
 	RefListIterator *iter;
@@ -970,16 +1037,23 @@ void polarisation_correction(RefList *list, UnitCell *cell, struct image *image)
 	{
 		double pol;
 		double intensity;
-		double xl, yl;
+		double xl, yl, zl;
+		double phi, tt, kpred;
 		signed int h, k, l;
-		const double P = 1.0;  /* degree of polarisation */
+		const double P = p.fraction;  /* degree of polarisation */
 
 		get_indices(refl, &h, &k, &l);
 
-		xl = (h*asx + k*bsx + l*csx)*image->lambda;
-		yl = (h*asy + k*bsy + l*csy)*image->lambda;
+		xl = h*asx + k*bsx + l*csx;
+		yl = h*asy + k*bsy + l*csy;
+		zl = h*asz + k*bsz + l*csz;
+
+		kpred = get_kpred(refl);
+		tt = angle_between(0.0, 0.0, 1.0, xl, yl, zl+kpred);
+		phi = atan2(yl, xl) - p.angle;
 
-		pol = P*(1.0 - xl*xl) + (1.0-P)*(1.0 - yl*yl);
+		pol =         P*(1.0 - cos(phi)*cos(phi)*sin(tt)*sin(tt))
+		     +  (1.0-P)*(1.0 - sin(phi)*sin(phi)*sin(tt)*sin(tt));
 
 		intensity = get_intensity(refl);
 		set_intensity(refl, intensity / pol);
diff --git a/libcrystfel/src/geometry.h b/libcrystfel/src/geometry.h
index e6e862a5..4bd1624f 100644
--- a/libcrystfel/src/geometry.h
+++ b/libcrystfel/src/geometry.h
@@ -91,6 +91,17 @@ enum gparam {
 };
 
 
+/**
+ * This structure represents the polarisation of the incident radiation
+ */
+struct polarisation
+{
+	double fraction;   /**< Polarisation fraction (0 to 1) */
+	double angle;      /**< Angle of electron beam, radians, clockwise from
+	                    * horizontal when looking along beam */
+};
+
+
 extern RefList *predict_to_res(Crystal *cryst, double max_res);
 
 extern void calculate_partialities(Crystal *cryst, PartialityModel pmodel);
@@ -98,8 +109,9 @@ extern void calculate_partialities(Crystal *cryst, PartialityModel pmodel);
 extern double r_gradient(UnitCell *cell, int k, Reflection *refl,
                          struct image *image);
 extern void update_predictions(Crystal *cryst);
-extern void polarisation_correction(RefList *list, UnitCell *cell,
-                                    struct image *image);
+extern struct polarisation parse_polarisation(const char *text);
+extern void polarisation_correction(RefList *list, UnitCell *cell, double lambda,
+                                    struct polarisation p);
 
 extern double sphere_fraction(double rlow, double rhigh, double pr);
 extern double gaussian_fraction(double rlow, double rhigh, double pr);
diff --git a/src/partialator.c b/src/partialator.c
index 68c886c9..ef6d0f80 100644
--- a/src/partialator.c
+++ b/src/partialator.c
@@ -325,6 +325,7 @@ static void show_help(const char *s)
 "  -m, --model=<model>        Specify partiality model.\n"
 "      --min-measurements=<n> Minimum number of measurements to require.\n"
 "      --no-polarisation      Disable polarisation correction.\n"
+"      --polarisation=<p>     Specify type of polarisation correction.\n"
 "      --max-adu=<n>          Saturation value of detector.\n"
 "      --min-res=<n>          Merge only crystals which diffract above <n> A.\n"
 "      --push-res=<n>         Merge higher than apparent resolution cutoff.\n"
@@ -951,7 +952,7 @@ int main(int argc, char *argv[])
 	PartialityModel pmodel = PMODEL_XSPHERE;
 	int min_measurements = 2;
 	char *rval;
-	int polarisation = 1;
+	struct polarisation polarisation = {.fraction = 1.0, .angle= 0.0};
 	int start_after = 0;
 	int stop_after = 0;
 	double max_adu = +INFINITY;
@@ -1005,14 +1006,14 @@ int main(int argc, char *argv[])
 		{"force-radius",       1, NULL,               11},
 		{"min-res",            1, NULL,               12},
 		{"force-lambda",       1, NULL,               13},
+		{"polarisation",       1, NULL,               14},
+		{"polarization",       1, NULL,               14}, /* compat */
+		{"no-polarisation",    0, NULL,               15},
+		{"no-polarization",    0, NULL,               15}, /* compat */
 
 		{"no-scale",           0, &no_scale,           1},
 		{"no-Bscale",          0, &no_Bscale,          1},
 		{"no-pr",              0, &no_pr,              1},
-		{"no-polarisation",    0, &polarisation,       0},
-		{"no-polarization",    0, &polarisation,       0}, /* compat */
-		{"polarisation",       0, &polarisation,       1},
-		{"polarization",       0, &polarisation,       1}, /* compat */
 		{"no-free",            0, &no_free,            1},
 		{"output-every-cycle", 0, &output_everycycle,  1},
 		{"no-logs",            0, &no_logs,            1},
@@ -1184,6 +1185,13 @@ int main(int argc, char *argv[])
 			force_lambda *= 1e-10;
 			break;
 
+			case 14 :
+			polarisation = parse_polarisation(optarg);
+			break;
+
+			case 15 :
+			polarisation.fraction = 0.5;
+			break;
 
 			case 0 :
 			break;
@@ -1427,11 +1435,8 @@ int main(int argc, char *argv[])
 
 			cr_refl = apply_max_adu(cr_refl, max_adu);
 
-			if ( polarisation ) {
-				polarisation_correction(cr_refl,
-						        crystal_get_cell(cr),
-						        image);
-			}
+			polarisation_correction(cr_refl, crystal_get_cell(cr),
+			                        image->lambda, polarisation);
 
 			if ( !no_free ) select_free_reflections(cr_refl, rng);
 
diff --git a/src/process_hkl.c b/src/process_hkl.c
index 393dc3e8..5e17f31c 100644
--- a/src/process_hkl.c
+++ b/src/process_hkl.c
@@ -80,7 +80,8 @@ static void show_help(const char *s)
 "                             model.\n"
 "      --even-only           Merge even numbered crystals only\n"
 "      --odd-only            Merge odd numbered crystals only\n"
-"      --no-polarisation     Disable polarisation correction.\n"
+"      --no-polarisation      Disable polarisation correction.\n"
+"      --polarisation=<p>     Specify type of polarisation correction.\n"
 "      --min-measurements=<n> Require at least <n> measurements before a\n"
 "                             reflection appears in the output.  Default: 2\n"
 "      --min-snr=<n>         Require individual intensity measurements to\n"
@@ -258,7 +259,7 @@ static int merge_crystal(RefList *model, struct image *image, Crystal *cr,
                          RefList *reference, const SymOpList *sym,
                          double **hist_vals, signed int hist_h,
                          signed int hist_k, signed int hist_l, int *hist_n,
-                         int config_nopolar, double min_snr, double max_adu,
+                         struct polarisation p, double min_snr, double max_adu,
                          double push_res, double min_cc, int do_scale,
                          FILE *stat)
 {
@@ -270,9 +271,8 @@ static int merge_crystal(RefList *model, struct image *image, Crystal *cr,
 	new_refl = crystal_get_reflections(cr);
 
 	/* First, correct for polarisation */
-	if ( !config_nopolar ) {
-		polarisation_correction(new_refl, crystal_get_cell(cr), image);
-	}
+	polarisation_correction(new_refl, crystal_get_cell(cr),
+	                        image->lambda, p);
 
 	if ( reference != NULL ) {
 		if ( do_scale ) {
@@ -385,7 +385,7 @@ static int merge_all(Stream *st, RefList *model, RefList *reference,
                      const SymOpList *sym,
                      double **hist_vals, signed int hist_h,
                      signed int hist_k, signed int hist_l,
-                     int *hist_i, int config_nopolar, int min_measurements,
+                     int *hist_i, struct polarisation p, int min_measurements,
                      double min_snr, double max_adu,
                      int start_after, int stop_after, double min_res,
                      double push_res, double min_cc, int do_scale, int flag_even_odd, char *stat_output)
@@ -435,7 +435,7 @@ static int merge_all(Stream *st, RefList *model, RefList *reference,
 				r = merge_crystal(model, &image, cr, reference,
 				                  sym, hist_vals,
 						  hist_h, hist_k, hist_l,
-				                  hist_i, config_nopolar,
+				                  hist_i, p,
 						  min_snr, max_adu, push_res,
 						  min_cc, do_scale, stat);
 				if ( r == 0 ) n_crystals_used++;
@@ -506,7 +506,7 @@ int main(int argc, char *argv[])
 	int hist_i;
 	int space_for_hist = 0;
 	char *histo_params = NULL;
-	int config_nopolar = 0;
+	struct polarisation polarisation = {.fraction = 1.0, .angle= 0.0};
 	char *rval;
 	int min_measurements = 2;
 	int r;
@@ -530,8 +530,6 @@ int main(int argc, char *argv[])
 		{"scale",              0, &config_scale,       1},
 		{"even-only",          0, &config_evenonly,    1},
 		{"odd-only",           0, &config_oddonly,     1},
-		{"no-polarisation",    0, &config_nopolar,     1},
-		{"no-polarization",    0, &config_nopolar,     1},
 		{"symmetry",           1, NULL,               'y'},
 		{"histogram",          1, NULL,               'g'},
 		{"hist-parameters",    1, NULL,               'z'},
@@ -544,6 +542,10 @@ int main(int argc, char *argv[])
 		{"version",            0, NULL,                7},
 		{"min-cc",             1, NULL,                8},
 		{"stat",               1, NULL,                9},
+		{"polarisation",       1, NULL,               10},
+		{"polarization",       1, NULL,               10}, /* compat */
+		{"no-polarisation",    0, NULL,               11},
+		{"no-polarization",    0, NULL,               11}, /* compat */
 		{0, 0, NULL, 0}
 	};
 
@@ -671,6 +673,14 @@ int main(int argc, char *argv[])
 			twopass = 1;
 			break;
 
+			case 10 :
+			polarisation = parse_polarisation(optarg);
+			break;
+
+			case 11 :
+			polarisation.fraction = 0.5;
+			break;
+
 			case 0 :
 			break;
 
@@ -760,7 +770,7 @@ int main(int argc, char *argv[])
 
 	hist_i = 0;
 	r = merge_all(st, model, NULL, sym, &hist_vals, hist_h, hist_k, hist_l,
-	              &hist_i, config_nopolar, min_measurements, min_snr,
+	              &hist_i, polarisation, min_measurements, min_snr,
 	              max_adu, start_after, stop_after, min_res, push_res,
 	              min_cc, config_scale, flag_even_odd, stat_output);
 	fprintf(stderr, "\n");
@@ -795,7 +805,7 @@ int main(int argc, char *argv[])
 
 			r = merge_all(st, model, reference, sym, &hist_vals,
 			              hist_h, hist_k, hist_l, &hist_i,
-				      config_nopolar, min_measurements, min_snr,
+				      polarisation, min_measurements, min_snr,
 				      max_adu, start_after, stop_after, min_res,
 				      push_res, min_cc, config_scale,
 				      flag_even_odd, stat_output);