Revert refinement step if too many reflections are lost

4 files changed, 91 insertions, 17 deletions

diff --git a/libcrystfel/src/geometry.c b/libcrystfel/src/geometry.c
index 744c13ea..c6ae1f63 100644
--- a/libcrystfel/src/geometry.c
+++ b/libcrystfel/src/geometry.c
@@ -380,7 +380,8 @@ static void set_unity_partialities(Crystal *cryst)
 
 
 /* Calculate partialities and apply them to the image's reflections */
-void update_partialities(Crystal *cryst, PartialityModel pmodel)
+void update_partialities_2(Crystal *cryst, PartialityModel pmodel,
+                           int *n_gained, int *n_lost)
 {
 	Reflection *refl;
 	RefListIterator *iter;
@@ -420,26 +421,45 @@ void update_partialities(Crystal *cryst, PartialityModel pmodel)
 		vals = check_reflection(image, cryst, h, k, l, xl, yl, zl);
 
 		if ( vals == NULL ) {
-			set_redundancy(refl, 0);
-			continue;
-		}
-		set_redundancy(refl, 1);
 
-		/* Transfer partiality stuff */
-		get_partial(vals, &r1, &r2, &p, &clamp1, &clamp2);
-		set_partial(refl, r1, r2, p, clamp1, clamp2);
-		L = get_lorentz(vals);
-		set_lorentz(refl, L);
+			if ( get_redundancy(refl) != 0 ) {
+				(*n_lost)++;
+				set_redundancy(refl, 0);
+			}
+
+		} else {
+
+			if ( get_redundancy(refl) == 0 ) {
+				(*n_gained)++;
+				set_redundancy(refl, 1);
+			}
 
-		/* Transfer detector location */
-		get_detector_pos(vals, &x, &y);
-		set_detector_pos(refl, 0.0, x, y);
+			/* Transfer partiality stuff */
+			get_partial(vals, &r1, &r2, &p, &clamp1, &clamp2);
+			set_partial(refl, r1, r2, p, clamp1, clamp2);
+			L = get_lorentz(vals);
+			set_lorentz(refl, L);
+
+			/* Transfer detector location */
+			get_detector_pos(vals, &x, &y);
+			set_detector_pos(refl, 0.0, x, y);
+
+			reflection_free(vals);
+
+		}
 
-		reflection_free(vals);
 	}
 }
 
 
+void update_partialities(Crystal *cryst, PartialityModel pmodel)
+{
+	int n_gained = 0;
+	int n_lost = 0;
+	update_partialities_2(cryst, pmodel, &n_gained, &n_lost);
+}
+
+
 void polarisation_correction(RefList *list, UnitCell *cell, struct image *image)
 {
 	Reflection *refl;
diff --git a/libcrystfel/src/geometry.h b/libcrystfel/src/geometry.h
index a6708b2d..bceb97e0 100644
--- a/libcrystfel/src/geometry.h
+++ b/libcrystfel/src/geometry.h
@@ -63,6 +63,8 @@ extern RefList *find_intersections(struct image *image, Crystal *cryst);
 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);
 extern void polarisation_correction(RefList *list, UnitCell *cell,
                                     struct image *image);
 
diff --git a/src/post-refinement.c b/src/post-refinement.c
index cbb4e5b5..1293b249 100644
--- a/src/post-refinement.c
+++ b/src/post-refinement.c
@@ -567,10 +567,47 @@ static double guide_dev(Crystal *cr, const RefList *full)
 }
 
 
+static Crystal *backup_crystal(Crystal *cr)
+{
+	Crystal *b;
+
+	b = crystal_new();
+	crystal_set_cell(b, cell_new_from_cell(crystal_get_cell(cr)));
+
+	return b;
+}
+
+
+static void revert_crystal(Crystal *cr, Crystal *backup)
+{
+	double asx, asy, asz;
+	double bsx, bsy, bsz;
+	double csx, csy, csz;
+
+	cell_get_reciprocal(crystal_get_cell(backup),
+	                    &asx, &asy, &asz,
+	                    &bsx, &bsy, &bsz,
+	                    &csx, &csy, &csz);
+
+	cell_set_reciprocal(crystal_get_cell(cr),
+	                    asx, asy, asz,
+	                    bsx, bsy, bsz,
+	                    csx, csy, csz);
+}
+
+
+static void free_backup_crystal(Crystal *cr)
+{
+	cell_free(crystal_get_cell(cr));
+	crystal_free(cr);
+}
+
+
 void pr_refine(Crystal *cr, const RefList *full, PartialityModel pmodel)
 {
 	double max_shift, dev;
 	int i;
+	Crystal *backup;
 	const int verbose = 0;
 
 	if ( verbose ) {
@@ -580,6 +617,8 @@ void pr_refine(Crystal *cr, const RefList *full, PartialityModel pmodel)
 		       dev);
 	}
 
+	backup = backup_crystal(cr);
+
 	i = 0;
 	crystal_set_user_flag(cr, 0);
 	do {
@@ -588,21 +627,33 @@ void pr_refine(Crystal *cr, const RefList *full, PartialityModel pmodel)
 		double bsx, bsy, bsz;
 		double csx, csy, csz;
 		double dev;
+		int n_total;
+		int n_gained = 0;
+		int n_lost = 0;
 
+		n_total = num_reflections(crystal_get_reflections(cr));
 		cell_get_reciprocal(crystal_get_cell(cr), &asx, &asy, &asz,
 			               &bsx, &bsy, &bsz, &csx, &csy, &csz);
 
 		max_shift = pr_iterate(cr, full, pmodel);
 
-		update_partialities(cr, pmodel);
+		update_partialities_2(cr, pmodel, &n_gained, &n_lost);
 
 		if ( verbose ) {
 			dev = guide_dev(cr, full);
 			STATUS("PR Iteration %2i: max shift = %10.2f"
-			       " dev = %10.5e\n", i+1, max_shift, dev);
+			       " dev = %10.5e, %i gained, %i lost, %i total\n",
+			       i+1, max_shift, dev, n_gained, n_lost, n_total);
+		}
+
+		if ( 10*n_lost > n_total ) {
+			revert_crystal(cr, backup);
+			crystal_set_user_flag(cr, 1);
 		}
 
 		i++;
 
 	} while ( (max_shift > 50.0) && (i < MAX_CYCLES) );
+
+	free_backup_crystal(backup);
 }
diff --git a/src/post-refinement.h b/src/post-refinement.h
index f565a880..5e0c184f 100644
--- a/src/post-refinement.h
+++ b/src/post-refinement.h
@@ -43,7 +43,8 @@
 #include "geometry.h"
 
 
-/* Refineable parameters */
+/* Refineable parameters.
+ * Don't forget to add new things to backup_crystal() et al.! */
 enum {
 	REF_ASX,
 	REF_ASY,