Split scaling and post-refinement apart (again)

Scaling needs the log residual, PR needs the linear one

1 files changed, 369 insertions, 91 deletions

diff --git a/src/post-refinement.c b/src/post-refinement.c
index 16fc1791..a35bc655 100644
--- a/src/post-refinement.c
+++ b/src/post-refinement.c
@@ -301,12 +301,6 @@ static void apply_shift(Crystal *cr, int k, double shift)
 	double t;
 	struct image *image = crystal_get_image(cr);
 
-	if ( isnan(shift) ) {
-		ERROR("Refusing NaN shift for parameter %i\n", k);
-		ERROR("Image serial %i\n", image->serial);
-		return;
-	}
-
 	switch ( k ) {
 
 		case GPARAM_DIV :
@@ -354,10 +348,9 @@ static void apply_shift(Crystal *cr, int k, double shift)
 }
 
 
-/* Perform one cycle of post refinement on 'image' against 'full' */
-static double pr_iterate(Crystal *cr, const RefList *full,
-                         PartialityModel pmodel, int no_scale, int no_pr,
-                         int *n_filtered, int verbose)
+/* Perform one cycle of scaling of 'cr' against 'full' */
+static double scale_iterate(Crystal *cr, const RefList *full,
+                            PartialityModel pmodel)
 {
 	gsl_matrix *M;
 	gsl_vector *v;
@@ -372,40 +365,17 @@ static double pr_iterate(Crystal *cr, const RefList *full,
 	enum gparam rv[32];
 	double G, B;
 
-	*n_filtered = 0;
-
-	/* If partiality model is anything other than "unity", refine all the
-	 * geometrical parameters */
-	if ( (pmodel != PMODEL_UNITY) && !no_pr ) {
-		rv[num_params++] = GPARAM_ASX;
-		rv[num_params++] = GPARAM_ASY;
-		rv[num_params++] = GPARAM_ASZ;
-		rv[num_params++] = GPARAM_BSX;
-		rv[num_params++] = GPARAM_BSY;
-		rv[num_params++] = GPARAM_BSZ;
-		rv[num_params++] = GPARAM_CSX;
-		rv[num_params++] = GPARAM_CSY;
-		rv[num_params++] = GPARAM_CSZ;
-		//rv[num_params++] = GPARAM_R;
-	}
-
-	/* If we are scaling, refine scale factors (duh) */
-	if ( !no_scale ) {
-		rv[num_params++] = GPARAM_OSF;
-		rv[num_params++] = GPARAM_BFAC;
-	}
-
-	if ( num_params == 0 ) return 0.0;
-
-	reflections = crystal_get_reflections(cr);
+	rv[num_params++] = GPARAM_OSF;
+	rv[num_params++] = GPARAM_BFAC;
 
 	M = gsl_matrix_calloc(num_params, num_params);
 	v = gsl_vector_calloc(num_params);
 
+	reflections = crystal_get_reflections(cr);
 	G = crystal_get_osf(cr);
 	B = crystal_get_Bfac(cr);
 
-	/* Construct the equations, one per reflection in this image */
+	/* Scaling terms */
 	for ( refl = first_refl(reflections, &iter);
 	      refl != NULL;
 	      refl = next_refl(refl, iter) )
@@ -434,12 +404,14 @@ static double pr_iterate(Crystal *cr, const RefList *full,
 		/* Actual measurement of this reflection from this pattern */
 		I_partial = get_intensity(refl);
 		esd = get_esd_intensity(refl);
+		p = get_partiality(refl);
 
-		if ( (get_partiality(refl) < MIN_PART_REFINE)
-		  || (get_redundancy(match) < 2)
-		  || (I_full <= 0) || (I_partial < 3*esd) ) continue;
+		/* Scale only using strong reflections */
+		if ( I_partial <= 3.0*esd ) continue; /* Also because of log */
+		if ( get_redundancy(match) < 2 ) continue;
+		if ( I_full <= 0 ) continue;  /* Because log */
+		if ( p <= 0.0 ) continue; /* Because of log */
 
-		p = get_partiality(refl);
 		L = get_lorentz(refl);
 		s = resolution(crystal_get_cell(cr), ha, ka, la);
 
@@ -449,12 +421,6 @@ static double pr_iterate(Crystal *cr, const RefList *full,
 		/* Calculate all gradients for this reflection */
 		for ( k=0; k<num_params; k++ ) {
 			gradients[k] = gradient(cr, rv[k], refl, pmodel);
-
-			/* Gradients affecting the "p" part of the residual
-			 * need to account for it being ln(p) */
-			if ( (rv[k] != GPARAM_OSF) && (rv[k] != GPARAM_BFAC) ) {
-				gradients[k] /= p;
-			}
 		}
 
 		for ( k=0; k<num_params; k++ ) {
@@ -487,9 +453,147 @@ static double pr_iterate(Crystal *cr, const RefList *full,
 
 		nref++;
 	}
+
+	if ( nref < num_params ) {
+		crystal_set_user_flag(cr, PRFLAG_FEWREFL);
+		gsl_matrix_free(M);
+		gsl_vector_free(v);
+		return 0.0;
+	}
+
+	max_shift = 0.0;
+	shifts = solve_svd(v, M, NULL, 0);
+	if ( shifts != NULL ) {
+
+		for ( param=0; param<num_params; param++ ) {
+			double shift = gsl_vector_get(shifts, param);
+			apply_shift(cr, rv[param], shift);
+			if ( fabs(shift) > max_shift ) max_shift = fabs(shift);
+		}
+
+	} else {
+		crystal_set_user_flag(cr, PRFLAG_SOLVEFAIL);
+	}
+
+	gsl_matrix_free(M);
+	gsl_vector_free(v);
+	gsl_vector_free(shifts);
+
+	return max_shift;
+}
+
+
+/* Perform one cycle of post refinement on 'image' against 'full' */
+static double pr_iterate(Crystal *cr, const RefList *full,
+                         PartialityModel pmodel,
+                         int *n_filtered, int verbose)
+{
+	gsl_matrix *M;
+	gsl_vector *v;
+	gsl_vector *shifts;
+	int param;
+	Reflection *refl;
+	RefListIterator *iter;
+	RefList *reflections;
+	double max_shift;
+	int nref = 0;
+	int num_params = 0;
+	enum gparam rv[32];
+	double G, B;
+
+	if ( n_filtered != NULL ) *n_filtered = 0;
+
+	rv[num_params++] = GPARAM_ASX;
+	rv[num_params++] = GPARAM_ASY;
+	rv[num_params++] = GPARAM_ASZ;
+	rv[num_params++] = GPARAM_BSX;
+	rv[num_params++] = GPARAM_BSY;
+	rv[num_params++] = GPARAM_BSZ;
+	rv[num_params++] = GPARAM_CSX;
+	rv[num_params++] = GPARAM_CSY;
+	rv[num_params++] = GPARAM_CSZ;
+//	rv[num_params++] = GPARAM_R;
+//	rv[num_params++] = GPARAM_DIV;
+
+	M = gsl_matrix_calloc(num_params, num_params);
+	v = gsl_vector_calloc(num_params);
+
+	reflections = crystal_get_reflections(cr);
+	G = crystal_get_osf(cr);
+	B = crystal_get_Bfac(cr);
+
+	/* Post-refinement terms */
+	for ( refl = first_refl(reflections, &iter);
+	      refl != NULL;
+	      refl = next_refl(refl, iter) )
+	{
+		signed int ha, ka, la;
+		double I_full, delta_I, esd;
+		double w;
+		double I_partial;
+		int k;
+		double p, L, s;
+		double fx;
+		Reflection *match;
+		double gradients[num_params];
+
+		if ( get_flag(refl) ) continue;
+
+		get_indices(refl, &ha, &ka, &la);
+		match = find_refl(full, ha, ka, la);
+		if ( match == NULL ) continue;
+		I_full = get_intensity(match);
+
+		if ( get_redundancy(match) < 2 ) continue;
+
+		p = get_partiality(refl);
+		L = get_lorentz(refl);
+		I_partial = get_intensity(refl);
+		esd = get_esd_intensity(refl);
+		s = resolution(crystal_get_cell(cr), ha, ka, la);
+
+		/* Calculate the weight for this reflection */
+		w = 1.0;
+
+		/* Calculate all gradients for this reflection */
+		for ( k=0; k<num_params; k++ ) {
+			gradients[k] = gradient(cr, rv[k], refl, pmodel);
+			gradients[k] *= exp(G)*exp(-B*s*s)*I_full/L;
+		}
+
+		for ( k=0; k<num_params; k++ ) {
+
+			int g;
+			double v_c, v_curr;
+
+			for ( g=0; g<num_params; g++ ) {
+
+				double M_c, M_curr;
+
+				/* Matrix is symmetric */
+				if ( g > k ) continue;
+
+				M_c = w * gradients[g] * gradients[k];
+
+				M_curr = gsl_matrix_get(M, k, g);
+				gsl_matrix_set(M, k, g, M_curr + M_c);
+				gsl_matrix_set(M, g, k, M_curr + M_c);
+
+			}
+
+			fx = exp(G)*p*exp(-B*s*s)*I_full/L;
+			delta_I = I_partial - fx;
+			v_c = w * delta_I * gradients[k];
+			v_curr = gsl_vector_get(v, k);
+			gsl_vector_set(v, k, v_curr + v_c);
+
+		}
+
+		nref++;
+	}
 	if ( verbose ) {
-		STATUS("The original equation:\n");
-		show_matrix_eqn(M, v);
+		//STATUS("The original equation:\n");
+		//show_matrix_eqn(M, v);
 		STATUS("%i reflections went into the equations.\n", nref);
 	}
 
@@ -501,13 +605,20 @@ static double pr_iterate(Crystal *cr, const RefList *full,
 	}
 
 	max_shift = 0.0;
-	shifts = solve_svd(v, M, n_filtered, verbose);
+	shifts = solve_svd(v, M, n_filtered, 0);
 	if ( shifts != NULL ) {
 
 		for ( param=0; param<num_params; param++ ) {
 			double shift = gsl_vector_get(shifts, param);
 			if ( verbose ) STATUS("Shift %i: %e\n", param, shift);
-			apply_shift(cr, rv[param], shift);
+			if ( isnan(shift) ) {
+				//ERROR("NaN shift parameter %i (image ser %i), "
+				//       "%i reflections.\n", rv[param],
+				//       crystal_get_image(cr)->serial,
+				//       nref);
+			} else {
+				apply_shift(cr, rv[param], shift);
+			}
 			if ( fabs(shift) > max_shift ) max_shift = fabs(shift);
 		}
 
@@ -523,13 +634,62 @@ static double pr_iterate(Crystal *cr, const RefList *full,
 }
 
 
-static double residual(Crystal *cr, const RefList *full, int verbose, int free)
+static double log_residual(Crystal *cr, const RefList *full, int free)
+{
+	double dev = 0.0;
+	double G, B;
+	Reflection *refl;
+	RefListIterator *iter;
+
+	G = crystal_get_osf(cr);
+	B = crystal_get_Bfac(cr);
+
+	for ( refl = first_refl(crystal_get_reflections(cr), &iter);
+	      refl != NULL;
+	      refl = next_refl(refl, iter) )
+	{
+		double p, L, s, w;
+		signed int h, k, l;
+		Reflection *match;
+		double esd, I_full, I_partial;
+		double fx, dc;
+
+		if ( free && !get_flag(refl) ) continue;
+
+		get_indices(refl, &h, &k, &l);
+		match = find_refl(full, h, k, l);
+		if ( match == NULL ) continue;
+
+		p = get_partiality(refl);
+		L = get_lorentz(refl);
+		I_partial = get_intensity(refl);
+		I_full = get_intensity(match);
+		esd = get_esd_intensity(refl);
+		s = resolution(crystal_get_cell(cr), h, k, l);
+
+		if ( I_partial <= 3.0*esd ) continue; /* Also because of log */
+		if ( get_redundancy(match) < 2 ) continue;
+		if ( I_full <= 0 ) continue;  /* Because log */
+		if ( p <= 0.0 ) continue; /* Because of log */
+
+		fx = G + log(p) - log(L) - B*s*s + log(I_full);
+		dc = log(I_partial) - fx;
+		w = 1.0;
+		dev += w*dc*dc;
+	}
+
+	return dev;
+}
+
+static double residual(Crystal *cr, const RefList *full, int verbose, int free,
+                       int *pn_used)
 {
 	double dev = 0.0;
 	double G, B;
 	Reflection *refl;
 	RefListIterator *iter;
 	FILE *fh = NULL;
+	int n_used = 0;
 
 	if ( verbose ) {
 		fh = fopen("residual.dat", "w");
@@ -553,100 +713,218 @@ static double residual(Crystal *cr, const RefList *full, int verbose, int free)
 		get_indices(refl, &h, &k, &l);
 		match = find_refl(full, h, k, l);
 		if ( match == NULL ) continue;
+		I_full = get_intensity(match);
+
+		if ( get_redundancy(match) < 2 ) continue;
 
 		p = get_partiality(refl);
 		L = get_lorentz(refl);
 		I_partial = get_intensity(refl);
-		I_full = get_intensity(match);
 		esd = get_esd_intensity(refl);
 		s = resolution(crystal_get_cell(cr), h, k, l);
 
-		if ( (get_partiality(refl) < MIN_PART_REFINE)
-		  || (get_redundancy(match) < 2)
-		  || (I_full <= 0) || (I_partial < 3*esd) ) continue;
-
-		fx = G + log(p) - log(L) - B*s*s + log(I_full);
-		dc = log(I_partial) - fx;
+		fx = exp(G)*p*exp(-B*s*s)*I_full/L;
+		dc = I_partial - fx;
 		w = 1.0;
 		dev += w*dc*dc;
+		n_used++;
 
 		if ( fh != NULL ) {
-			fprintf(fh, "%4i %4i %4i %e %.2f %e %f %f\n",
-			        h, k, l, s, G, B, I_partial, I_full);
+			fprintf(fh, "%4i %4i %4i %e %.2f %e %f %f %f\n",
+			        h, k, l, s, G, B, I_partial, p, I_full);
 		}
 	}
 
 	if ( fh != NULL ) fclose(fh);
+	if ( pn_used != NULL ) *pn_used = n_used;
 
 	return dev;
 }
 
 
-static struct prdata pr_refine(Crystal *cr, const RefList *full,
-                               PartialityModel pmodel, int no_scale, int no_pr)
+static void write_residual_graph(Crystal *cr, const RefList *full)
 {
 	int i;
-	int verbose = 0;
-	struct prdata prdata;
-	int done = 0;
-	double old_dev;
+	double asx, asy, asz;
+	double bsx, bsy, bsz;
+	double csx, csy, csz;
+	double a;
+	double step;
+	double orig_asx;
+	FILE *fh;
+	UnitCell *cell;
+
+	cell = crystal_get_cell(cr);
+
+	fh = fopen("residual-plot.dat", "w");
+
+	cell_get_reciprocal(cell, &asx, &asy, &asz,
+	                          &bsx, &bsy, &bsz,
+	                          &csx, &csy, &csz);
+	a = modulus(asx, asy, asz);
+	orig_asx = asx;
+
+	step = a/100.0;
+
+	for ( i=-50; i<=50; i++ ) {
+
+		double res;
+		int n;
+		asx = orig_asx + (i*step);
+		cell_set_reciprocal(cell, asx, asy, asz,
+		                          bsx, bsy, bsz,
+		                          csx, csy, csz);
+		update_partialities(cr, PMODEL_SCSPHERE);
+		res = residual(cr, full, 0, 0, &n);
+		fprintf(fh, "%i %e %e %i\n", i, asx, res, n);
+	}
+
+	cell_set_reciprocal(cell, orig_asx, asy, asz,
+	                          bsx, bsy, bsz,
+	                          csx, csy, csz);
+	update_partialities(cr, PMODEL_SCSPHERE);
+	fclose(fh);
+}
 
-	prdata.refined = 0;
-	prdata.n_filtered = 0;
 
-	/* Don't refine crystal if scaling was bad */
-	if ( crystal_get_user_flag(cr) != 0 ) return prdata;
+static double do_scale_refine(Crystal *cr, const RefList *full,
+                              PartialityModel pmodel, int verbose)
+{
+	int i, done;
+	double old_dev;
 
-	old_dev = residual(cr, full, 0, 0);
-	prdata.initial_free_residual = residual(cr, full, 0, 1);
-	prdata.initial_residual = old_dev;
+	old_dev = log_residual(cr, full, 0);
 
 	if ( verbose ) {
-		STATUS("\n");  /* Deal with progress bar */
 		STATUS("Initial G=%.2f, B=%e\n",
 		       crystal_get_osf(cr), crystal_get_Bfac(cr));
-		STATUS("Initial  dev =  %10.5e, free dev = %10.5e\n",
-		       old_dev, residual(cr, full, 0, 1));
+		STATUS("Scaling initial  dev =  %10.5e, free dev = %10.5e\n",
+		       old_dev, log_residual(cr, full, 0));
 	}
 
 	i = 0;
+	done = 0;
 	do {
 
+		double dev;
+
+		scale_iterate(cr, full, pmodel);
+
+		dev = log_residual(cr, full, 0);
+		if ( fabs(dev - old_dev) < dev*0.01 ) done = 1;
+
+		if ( verbose ) {
+			STATUS("Scaling %2i: dev = %10.5e, free dev = %10.5e\n",
+			       i+1, dev, log_residual(cr, full, 0));
+		}
+
+		i++;
+		old_dev = dev;
+
+	} while ( i < MAX_CYCLES && !done );
+
+	if ( verbose ) {
+		STATUS("Final G=%.2f, B=%e\n", crystal_get_osf(cr),
+		       crystal_get_Bfac(cr));
+	}
+
+	return old_dev;
+}
+
+
+static double do_pr_refine(Crystal *cr, const RefList *full,
+                           PartialityModel pmodel, int verbose)
+{
+	int i, done;
+	double old_dev;
+	UnitCell *cell = crystal_get_cell(cr);
+
+	old_dev = residual(cr, full, 0, 0, NULL);
+
+	if ( verbose ) {
 		double asx, asy, asz;
 		double bsx, bsy, bsz;
 		double csx, csy, csz;
-		double dev;
+		cell_get_reciprocal(cell, &asx, &asy, &asz,
+		                          &bsx, &bsy, &bsz,
+		                          &csx, &csy, &csz);
+		STATUS("Initial asx = %e\n", asx);
+		STATUS("PR initial  dev =  %10.5e, free dev = %10.5e\n",
+		       old_dev, residual(cr, full, 0, 1, NULL));
+	}
+
+	i = 0;
+	done = 0;
+	do {
 
-		cell_get_reciprocal(crystal_get_cell(cr), &asx, &asy, &asz,
-			               &bsx, &bsy, &bsz, &csx, &csy, &csz);
+		double dev;
 
-		pr_iterate(cr, full, pmodel, no_scale, no_pr,
-		           &prdata.n_filtered, verbose);
+		pr_iterate(cr, full, pmodel, NULL, verbose);
 
 		update_partialities(cr, pmodel);
 
-		dev = residual(cr, full, 0, 0);
+		dev = residual(cr, full, 0, 0, NULL);
 		if ( fabs(dev - old_dev) < dev*0.01 ) done = 1;
 
 		if ( verbose ) {
-			STATUS("Iter %2i: dev = %10.5e, free dev = %10.5e\n",
-			       i+1, dev, residual(cr, full, 0, 1));
+			STATUS("PR iter %2i: dev = %10.5e, free dev = %10.5e\n",
+			       i+1, dev, residual(cr, full, 0, 1, NULL));
 		}
 
 		i++;
 		old_dev = dev;
 
-	} while ( i < MAX_CYCLES && !done );
+	} while ( i < 30 && !done );
+
+	if ( verbose ) {
+		double asx, asy, asz;
+		double bsx, bsy, bsz;
+		double csx, csy, csz;
+		cell_get_reciprocal(cell, &asx, &asy, &asz,
+		                          &bsx, &bsy, &bsz,
+		                          &csx, &csy, &csz);
+		STATUS("Final asx = %e\n", asx);
+	}
+
+	return old_dev;
+}
+
+
+static struct prdata pr_refine(Crystal *cr, const RefList *full,
+                               PartialityModel pmodel, int no_scale, int no_pr)
+{
+	int verbose = 0;
+	struct prdata prdata;
+	double final_dev;
+
+	prdata.refined = 0;
+	prdata.n_filtered = 0;
+
+	prdata.initial_residual = residual(cr, full, 0, 0, NULL);
+	prdata.initial_free_residual = residual(cr, full, 0, 1, NULL);
+	final_dev = prdata.initial_residual;
+
+	if ( !no_scale ) {
+		final_dev = do_scale_refine(cr, full, pmodel, verbose);
+	}
+
+	if ( verbose ) {
+		write_residual_graph(cr, full);
+	}
+
+	if ( !no_pr ) {
+		final_dev = do_pr_refine(cr, full, pmodel, verbose);
+	}
 
 	if ( crystal_get_user_flag(cr) == 0 ) {
 		prdata.refined = 1;
 	}
-	prdata.final_free_residual = residual(cr, full, 0, 1);
-	prdata.final_residual = old_dev;
+	prdata.final_free_residual = residual(cr, full, 0, 1, NULL);
+	prdata.final_residual = final_dev;
 
-	if ( verbose ) {
-		STATUS("Final G=%.2f, B=%e\n", crystal_get_osf(cr),
-		       crystal_get_Bfac(cr));
+	if ( isnan(final_dev) ) {
+		STATUS("nan residual!  Image serial %i\n",
+		       crystal_get_image(cr)->serial);
 	}
 
 	return prdata;