Stop PR when mean change in partiality is small

4 files changed, 28 insertions, 10 deletions

diff --git a/libcrystfel/src/geometry.c b/libcrystfel/src/geometry.c
index c6ae1f63..9ca6262c 100644
--- a/libcrystfel/src/geometry.c
+++ b/libcrystfel/src/geometry.c
@@ -381,7 +381,7 @@ static void set_unity_partialities(Crystal *cryst)
 
 /* Calculate partialities and apply them to the image's reflections */
 void update_partialities_2(Crystal *cryst, PartialityModel pmodel,
-                           int *n_gained, int *n_lost)
+                           int *n_gained, int *n_lost, double *mean_p_change)
 {
 	Reflection *refl;
 	RefListIterator *iter;
@@ -389,6 +389,8 @@ void update_partialities_2(Crystal *cryst, PartialityModel pmodel,
 	double bsx, bsy, bsz;
 	double csx, csy, csz;
 	struct image *image = crystal_get_image(cryst);
+	double total_p_change = 0.0;
+	int n = 0;
 
 	if ( pmodel == PMODEL_UNITY ) {
 		/* It isn't strictly necessary to set the partialities to 1,
@@ -410,8 +412,10 @@ void update_partialities_2(Crystal *cryst, PartialityModel pmodel,
 		double xl, yl, zl;
 		signed int h, k, l;
 		int clamp1, clamp2;
+		double old_p;
 
 		get_symmetric_indices(refl, &h, &k, &l);
+		old_p = get_partiality(refl);
 
 		/* Get the coordinates of the reciprocal lattice point */
 		xl = h*asx + k*bsx + l*csx;
@@ -446,17 +450,25 @@ void update_partialities_2(Crystal *cryst, PartialityModel pmodel,
 
 			reflection_free(vals);
 
+			total_p_change += fabs(p - old_p);
+			n++;
+
 		}
 
 	}
+
+	*mean_p_change = total_p_change / n;
 }
 
 
+/* Wrapper to maintain API compatibility */
 void update_partialities(Crystal *cryst, PartialityModel pmodel)
 {
 	int n_gained = 0;
 	int n_lost = 0;
-	update_partialities_2(cryst, pmodel, &n_gained, &n_lost);
+	double mean_p_change = 0.0;
+	update_partialities_2(cryst, pmodel, &n_gained, &n_lost,
+	                      &mean_p_change);
 }
 
 
diff --git a/libcrystfel/src/geometry.h b/libcrystfel/src/geometry.h
index bceb97e0..de0259a7 100644
--- a/libcrystfel/src/geometry.h
+++ b/libcrystfel/src/geometry.h
@@ -64,7 +64,8 @@ extern RefList *select_intersections(struct image *image, Crystal *cryst);
 
 extern void update_partialities(Crystal *cryst, PartialityModel pmodel);
 extern void update_partialities_2(Crystal *cryst, PartialityModel pmodel,
-                                  int *n_gained, int *n_lost);
+                                  int *n_gained, int *n_lost,
+                                  double *mean_p_change);
 extern void polarisation_correction(RefList *list, UnitCell *cell,
                                     struct image *image);
 
diff --git a/src/partialator.c b/src/partialator.c
index ad6d9a49..97d2b689 100644
--- a/src/partialator.c
+++ b/src/partialator.c
@@ -607,13 +607,15 @@ int main(int argc, char *argv[])
 			RefList *as;
 			int n_gained = 0;
 			int n_lost = 0;
+			double mean_p_change = 0.0;
 
 			cryst = images[i].crystals[j];
 			crystal_set_image(cryst, &images[i]);
 
 			/* Now it's safe to do the following */
 			update_partialities_2(cryst, pmodel,
-			                      &n_gained, &n_lost);
+			                      &n_gained, &n_lost,
+			                      &mean_p_change);
 			assert(n_gained == 0);  /* That'd just be silly */
 			as = crystal_get_reflections(cryst);
 			nobs += select_scalable_reflections(as, reference);
diff --git a/src/post-refinement.c b/src/post-refinement.c
index 05e1829a..43cc002b 100644
--- a/src/post-refinement.c
+++ b/src/post-refinement.c
@@ -661,11 +661,12 @@ static void free_backup_crystal(Crystal *cr)
 struct prdata pr_refine(Crystal *cr, const RefList *full,
                         PartialityModel pmodel)
 {
-	double max_shift, dev;
+	double dev;
 	int i;
 	Crystal *backup;
 	const int verbose = 0;
 	struct prdata prdata;
+	double mean_p_change = 0.0;
 
 	if ( verbose ) {
 		dev = guide_dev(cr, full);
@@ -692,15 +693,17 @@ struct prdata pr_refine(Crystal *cr, const RefList *full,
 			               &bsx, &bsy, &bsz, &csx, &csy, &csz);
 
 		prdata.n_filtered = 0;
-		max_shift = pr_iterate(cr, full, pmodel, &prdata);
+		pr_iterate(cr, full, pmodel, &prdata);
 
-		update_partialities_2(cr, pmodel, &n_gained, &n_lost);
+		update_partialities_2(cr, pmodel, &n_gained, &n_lost,
+		                      &mean_p_change);
 
 		if ( verbose ) {
 			dev = guide_dev(cr, full);
-			STATUS("PR Iteration %2i: max shift = %10.2f"
+			STATUS("PR Iteration %2i: mean p change = %10.2f"
 			       " dev = %10.5e, %i gained, %i lost, %i total\n",
-			       i+1, max_shift, dev, n_gained, n_lost, n_total);
+			       i+1, mean_p_change, dev,
+			       n_gained, n_lost, n_total);
 		}
 
 		if ( 3*n_lost > n_total ) {
@@ -710,7 +713,7 @@ struct prdata pr_refine(Crystal *cr, const RefList *full,
 
 		i++;
 
-	} while ( (max_shift > 50.0) && (i < MAX_CYCLES) );
+	} while ( (mean_p_change > 0.01) && (i < MAX_CYCLES) );
 
 	free_backup_crystal(backup);