partialator: Add B-factor scaling and rejection

1 files changed, 148 insertions, 22 deletions

diff --git a/src/hrs-scaling.c b/src/hrs-scaling.c
index 7f15f54d..0225d625 100644
--- a/src/hrs-scaling.c
+++ b/src/hrs-scaling.c
@@ -3,11 +3,11 @@
  *
  * Intensity scaling using generalised HRS target function
  *
- * Copyright © 2012-2014 Deutsches Elektronen-Synchrotron DESY,
+ * Copyright © 2012-2015 Deutsches Elektronen-Synchrotron DESY,
  *                       a research centre of the Helmholtz Association.
  *
  * Authors:
- *   2010-2014 Thomas White <taw@physics.org>
+ *   2010-2015 Thomas White <taw@physics.org>
  *
  * This file is part of CrystFEL.
  *
@@ -37,6 +37,7 @@
 #include <gsl/gsl_vector.h>
 #include <gsl/gsl_linalg.h>
 #include <gsl/gsl_eigen.h>
+#include <gsl/gsl_fit.h>
 
 #include "image.h"
 #include "peaks.h"
@@ -45,6 +46,7 @@
 #include "cell.h"
 #include "utils.h"
 #include "reflist.h"
+#include "cell-utils.h"
 
 
 /* Minimum partiality of a reflection for it to be used for scaling */
@@ -63,6 +65,7 @@ struct scale_queue_args
 	Crystal **crystals;
 	int n_started;
 	PartialityModel pmodel;
+	UnitCell *cell;
 };
 
 
@@ -71,6 +74,7 @@ struct scale_worker_args
 	Crystal *crystal;
 	RefList *reference;
 	PartialityModel pmodel;
+	int crystal_number;
 };
 
 
@@ -83,6 +87,7 @@ static void *create_scale_job(void *vqargs)
 	wargs->reference = qargs->reference;
 	wargs->pmodel = qargs->pmodel;
 
+	wargs->crystal_number = qargs->n_started;
 	wargs->crystal = qargs->crystals[qargs->n_started++];
 
 	return wargs;
@@ -96,9 +101,35 @@ static void run_scale_job(void *vwargs, int cookie)
 	RefList *reference = wargs->reference;
 	Reflection *refl;
 	RefListIterator *iter;
-	double num = 0.0;
-	double den = 0.0;
-	double g;
+	int n = 0;
+	double *x;
+	double *y;
+	double *w;
+	int max_n = 256;
+	double c0, c1, cov00, cov01, cov11, chisq;
+	double G, B;
+	int r;
+	int debug = 0;
+	FILE *fh;
+
+	/* If this crystal's scaling was dodgy, it doesn't contribute to the
+	 * merged intensities */
+	if ( crystal_get_user_flag(cr) != 0 ) return;
+
+	x = malloc(max_n*sizeof(double));
+	y = malloc(max_n*sizeof(double));
+	w = malloc(max_n*sizeof(double));
+	if ( (x==NULL) || (y==NULL) || (w==NULL) ) {
+		ERROR("Failed to allocate memory for scaling.\n");
+		return;
+	}
+
+	if ( wargs->crystal_number == 0 ) {
+		char fn[256];
+		snprintf(fn, 255, "scale-debug.%i", wargs->crystal_number);
+		debug = 1;
+		fh = fopen(fn, "w");
+	}
 
 	for ( refl = first_refl(crystal_get_reflections(cr), &iter);
 	      refl != NULL;
@@ -107,30 +138,80 @@ static void run_scale_job(void *vwargs, int cookie)
 		signed int h, k, l;
 		double Ih, Ihl, corr;
 		Reflection *r;
+		double res;
+
+		get_indices(refl, &h, &k, &l);
+		res = resolution(crystal_get_cell(cr), h, k, l);
 
 		if ( (get_partiality(refl) < MIN_PART_SCALE)
-		  || (get_intensity(refl) < 3.0*get_esd_intensity(refl)) ) {
+		  || (get_intensity(refl) < 5.0*get_esd_intensity(refl)) ) {
 			continue;
 		}
 
 		/* Look up by asymmetric indices */
-		get_indices(refl, &h, &k, &l);
 		r = find_refl(reference, h, k, l);
-		if ( r == NULL ) continue;
+		if ( r == NULL ) {
+			continue;
+		}
 
 		Ih = get_intensity(r);
-
 		corr = get_partiality(refl) * get_lorentz(refl);
-
 		Ihl = get_intensity(refl) / corr;
 
-		num += Ih * Ihl;
-		den += Ih * Ih;
+		if ( Ihl <= 0.0 ) continue;
+		if ( Ih <= 0.0 ) continue;
+		if ( isnan(Ih) || isinf(Ihl) ) continue;
+		if ( isnan(Ih) || isinf(Ih) ) continue;
+
+		if ( n == max_n ) {
+			max_n *= 2;
+			x = realloc(x, max_n*sizeof(double));
+			y = realloc(y, max_n*sizeof(double));
+			w = realloc(w, max_n*sizeof(double));
+			if ( (x==NULL) || (y==NULL) || (w==NULL) ) {
+				ERROR("Failed to allocate memory for scaling.\n");
+				return;
+			}
+		}
+
+		x[n] = res*res;
+		y[n] = log(Ihl/Ih);
+		w[n] = 1.0;//log(get_esd_intensity(refl)*corr);
+		if ( debug ) {
+			fprintf(fh, "%i %10.2e %10.2f %10.2f %.2f %.2f %.2f\n",
+			        n, x[n], y[n], w[n], Ihl, Ih, res/1e9);
+		}
+		n++;
 
 	}
 
-	g = num / den;
-	crystal_set_osf(cr, g);  /* If it's NaN, it'll get rejected later */
+	if ( n < 2 ) {
+		crystal_set_user_flag(cr, 1);
+		return;
+	}
+
+	r = gsl_fit_wlinear(x, 1, w, 1, y, 1, n, &c0, &c1,
+	                     &cov00, &cov01, &cov11, &chisq);
+
+	G = 1.0/exp(c0);
+	B = -c1/2.0;
+	if ( debug ) {
+		printf("intercept = %e, gradient = %e\n", c0, c1);
+		printf("Scale factor = %f\n", G);
+		printf("Relative B = %.2f A^2\n", B * 1e20);
+		fclose(fh);
+	}
+
+	free(x);
+	free(y);
+	free(w);
+	if ( r || isnan(c0) ) {
+		crystal_set_user_flag(cr, 1);
+		return;
+	}
+
+	crystal_set_osf(cr, G);
+	crystal_set_Bfac(cr, B);
 }
 
 
@@ -174,6 +255,7 @@ struct merge_worker_args
 	RefList *full;
 	pthread_rwlock_t *full_lock;
 	PartialityModel pmodel;
+	int crystal_number;
 };
 
 
@@ -187,6 +269,7 @@ static void *create_merge_job(void *vqargs)
 	wargs->full_lock = &qargs->full_lock;
 	wargs->pmodel = qargs->pmodel;
 
+	wargs->crystal_number = qargs->n_started;
 	wargs->crystal = qargs->crystals[qargs->n_started++];
 
 	return wargs;
@@ -200,13 +283,23 @@ static void run_merge_job(void *vwargs, int cookie)
 	RefList *full = wargs->full;
 	Reflection *refl;
 	RefListIterator *iter;
-	double G;
+	double G, B;
+	FILE *fh;
+	int debug = 0;
 
 	/* If this crystal's scaling was dodgy, it doesn't contribute to the
 	 * merged intensities */
 	if ( crystal_get_user_flag(cr) != 0 ) return;
 
 	G = crystal_get_osf(cr);
+	B = crystal_get_Bfac(cr);
+
+	if ( wargs->crystal_number == 0 ) {
+		char fn[256];
+		snprintf(fn, 255, "merge-debug.%i", wargs->crystal_number);
+		debug = 1;
+		fh = fopen(fn, "w");
+	}
 
 	for ( refl = first_refl(crystal_get_reflections(cr), &iter);
 	      refl != NULL;
@@ -216,7 +309,7 @@ static void run_merge_job(void *vwargs, int cookie)
 		signed int h, k, l;
 		double num, den;
 		int red;
-		double Ihl, corr;
+		double Ihl, corr, res;
 
 		if ( get_partiality(refl) < MIN_PART_MERGE ) continue;
 
@@ -262,19 +355,27 @@ static void run_merge_job(void *vwargs, int cookie)
 
 		}
 
+		res = resolution(crystal_get_cell(cr), h, k, l);
 		corr = get_partiality(refl) * get_lorentz(refl);
 
 		Ihl = get_intensity(refl) / corr;
 
-		num += Ihl / G;
+		num += Ihl * G * exp(2.0*B*res);
 		den += 1.0;
 		red++;
 
+		if ( debug ) {
+			fprintf(fh, "%.2f %.2f %f %e\n", res/1e9,
+			                           Ihl*G*exp(2.0*B*res),
+						   G, B);
+		}
+
 		set_temp1(f, num);
 		set_temp2(f, den);
 		set_redundancy(f, red);
 		unlock_reflection(f);
 	}
+	if ( debug ) fclose(fh);
 }
 
 
@@ -449,10 +550,33 @@ static void reject_outliers(double *old_osfs, int n, Crystal **crystals)
 	int i;
 
 	for ( i=0; i<n; i++ ) {
+
 		double osf = crystal_get_osf(crystals[i]);
-		if ( isnan(osf) || (osf < 0.0) || (osf > 3.0) ) {
+		double Bfac = crystal_get_Bfac(crystals[i]);
+		int bad = 0;
+
+		if ( isnan(osf) || (osf < 0.0) || (osf > 10.0) ) bad = 1;
+		if ( isnan(Bfac) || (Bfac<-40e-20) || (Bfac>40e-20) ) bad = 1;
+
+		if ( bad ) {
 			crystal_set_user_flag(crystals[i], 1);
+			//STATUS("crystal %i bad: osf=%f B=%f A^2\n",
+			//       i, osf, Bfac*1e20);
+		} else {
+			//STATUS("OK %i: osf=%f B=%f A^2\n",
+			//       i, osf, Bfac*1e20);
 		}
+
+	}
+}
+
+
+static void reset_scaling_flag(Crystal *crystal)
+{
+	/* Every pattern gets another chance at being scaled, but stays flagged
+	 * if it was flagged for any other reason */
+	if ( crystal_get_user_flag(crystal) == 1 ) {
+		crystal_set_user_flag(crystal, 0);
 	}
 }
 
@@ -490,8 +614,9 @@ RefList *scale_intensities(Crystal **crystals, int n,
 	int done;
 
 	for ( i=0; i<n; i++ ) {
-		crystal_set_user_flag(crystals[i], 0);
+		reset_scaling_flag(crystals[i]);
 		crystal_set_osf(crystals[i], 1.0);
+		crystal_set_Bfac(crystals[i], 0.0);
 	}
 
 	if ( noscale ) {
@@ -518,26 +643,27 @@ RefList *scale_intensities(Crystal **crystals, int n,
 
 		for ( j=0; j<n; j++ ) {
 			old_osfs[j] = crystal_get_osf(crystals[j]);
-			crystal_set_user_flag(crystals[j], 0);
+			reset_scaling_flag(crystals[j]);
 		}
 
 		iterate_scale(crystals, n, full, n_threads, pmodel);
+		reject_outliers(old_osfs, n, crystals);
 
 		/* Normalise the scale factors */
 		for ( j=0; j<n; j++ ) {
 			double osf =  crystal_get_osf(crystals[j]);
-			if ( !isnan(osf) ) {
+			if ( crystal_get_user_flag(crystals[j]) == 0 ) {
 				total_sf += osf;
 				n_sf++;
 			}
 		}
 		norm_sf = total_sf / n_sf;
+		STATUS("norm = %f\n", norm_sf);
 		for ( j=0; j<n; j++ ) {
 			crystal_set_osf(crystals[j],
 			                crystal_get_osf(crystals[j])/norm_sf);
 		}
 
-		reject_outliers(old_osfs, n, crystals);
 		done = test_convergence(old_osfs, n, crystals);
 
 		/* Generate list for next iteration */