Add the possibility to have different partiality models

6 files changed, 64 insertions, 25 deletions

diff --git a/libcrystfel/src/geometry.c b/libcrystfel/src/geometry.c
index 621c9484..450ebfaf 100644
--- a/libcrystfel/src/geometry.c
+++ b/libcrystfel/src/geometry.c
@@ -346,7 +346,7 @@ RefList *select_intersections(struct image *image, Crystal *cryst)
 
 
 /* Calculate partialities and apply them to the image's reflections */
-void update_partialities(Crystal *cryst)
+void update_partialities(Crystal *cryst, PartialityModel pmodel)
 {
 	Reflection *refl;
 	RefListIterator *iter;
diff --git a/libcrystfel/src/geometry.h b/libcrystfel/src/geometry.h
index aecdc28a..8b726ed7 100644
--- a/libcrystfel/src/geometry.h
+++ b/libcrystfel/src/geometry.h
@@ -43,10 +43,26 @@
 extern "C" {
 #endif
 
+/**
+ * PartialityModel:
+ * @PMODEL_SPHERE : Intersection of sphere with excited volume of reciprocal
+ *   space.
+ * @PMODEL_UNITY : Set all all partialities and Lorentz factors to 1.
+ *
+ * A %PartialityModelModel describes a geometrical model which can be used to
+ * calculate spot partialities and Lorentz correction factors.
+ **/
+typedef enum {
+
+	PMODEL_SPHERE,
+	PMODEL_UNITY,
+
+} PartialityModel;
+
 extern RefList *find_intersections(struct image *image, Crystal *cryst);
 extern RefList *select_intersections(struct image *image, Crystal *cryst);
 
-extern void update_partialities(Crystal *cryst);
+extern void update_partialities(Crystal *cryst, PartialityModel pmodel);
 
 #ifdef __cplusplus
 }
diff --git a/src/partialator.c b/src/partialator.c
index a4af3c18..079e078e 100644
--- a/src/partialator.c
+++ b/src/partialator.c
@@ -88,6 +88,7 @@ struct refine_args
 {
 	RefList *full;
 	Crystal *crystal;
+	PartialityModel pmodel;
 };
 
 
@@ -106,7 +107,7 @@ static void refine_image(void *task, int id)
 	struct refine_args *pargs = task;
 	Crystal *cr = pargs->crystal;
 
-	pr_refine(cr, pargs->full);
+	pr_refine(cr, pargs->full, pargs->pmodel);
 }
 
 
@@ -314,6 +315,8 @@ int main(int argc, char *argv[])
 	int noscale = 0;
 	Stream *st;
 	Crystal **crystals;
+	char *pmodel_str = NULL;
+	PartialityModel pmodel = PMODEL_SPHERE;
 
 	/* Long options */
 	const struct option longopts[] = {
@@ -326,6 +329,7 @@ int main(int argc, char *argv[])
 		{"iterations",         1, NULL,               'n'},
 		{"no-scale",           0, &noscale,            1},
 		{"reference",          1, NULL,               'r'},
+		{"partiality",         1, NULL,               'm'},
 		{0, 0, NULL, 0}
 	};
 
@@ -336,7 +340,7 @@ int main(int argc, char *argv[])
 	}
 
 	/* Short options */
-	while ((c = getopt_long(argc, argv, "hi:o:g:b:y:n:r:j:",
+	while ((c = getopt_long(argc, argv, "hi:o:g:b:y:n:r:j:m:",
 	                        longopts, NULL)) != -1)
 	{
 
@@ -370,6 +374,10 @@ int main(int argc, char *argv[])
 			n_iter = atoi(optarg);
 			break;
 
+			case 'm' :
+			pmodel_str = strdup(optarg);
+			break;
+
 			case 'b' :
 			beam = get_beam_parameters(optarg);
 			if ( beam == NULL ) {
@@ -434,6 +442,17 @@ int main(int argc, char *argv[])
 		return 1;
 	}
 
+	if ( pmodel_str != NULL ) {
+		if ( strcmp(pmodel_str, "sphere") == 0 ) {
+			pmodel = PMODEL_SPHERE;
+		} else if ( strcmp(pmodel_str, "unity") == 0 ) {
+			pmodel = PMODEL_UNITY;
+		} else {
+			ERROR("Unknown partiality model '%s'.\n", pmodel_str);
+			return 1;
+		}
+	}
+
 	if ( reference_file != NULL ) {
 
 		RefList *list;
@@ -544,7 +563,7 @@ int main(int argc, char *argv[])
 			crystal_set_image(cryst, &images[i]);
 
 			/* Now it's safe to do the following */
-			update_partialities(cryst);
+			update_partialities(cryst, pmodel);
 			as = crystal_get_reflections(cryst);
 			nobs += select_scalable_reflections(as, reference);
 
diff --git a/src/post-refinement.c b/src/post-refinement.c
index 1439b148..f6f4caa5 100644
--- a/src/post-refinement.c
+++ b/src/post-refinement.c
@@ -88,7 +88,7 @@ static double partiality_rgradient(double r, double profile_radius)
 
 
 /* Return the gradient of parameter 'k' given the current status of 'image'. */
-double gradient(Crystal *cr, int k, Reflection *refl)
+double gradient(Crystal *cr, int k, Reflection *refl, PartialityModel pmodel)
 {
 	double ds, azix, aziy;
 	double ttlow, tthigh, tt;
@@ -364,7 +364,8 @@ static gsl_vector *solve_svd(gsl_vector *v, gsl_matrix *M)
 
 
 /* Perform one cycle of post refinement on 'image' against 'full' */
-static double pr_iterate(Crystal *cr, const RefList *full)
+static double pr_iterate(Crystal *cr, const RefList *full,
+                         PartialityModel pmodel)
 {
 	gsl_matrix *M;
 	gsl_vector *v;
@@ -420,7 +421,7 @@ static double pr_iterate(Crystal *cr, const RefList *full)
 		/* Calculate all gradients for this reflection */
 		for ( k=0; k<NUM_PARAMS; k++ ) {
 			double gr;
-			gr = gradient(cr, k, refl);
+			gr = gradient(cr, k, refl, pmodel);
 			gradients[k] = gr;
 		}
 
@@ -532,7 +533,7 @@ static double guide_dev(Crystal *cr, const RefList *full)
 }
 
 
-void pr_refine(Crystal *cr, const RefList *full)
+void pr_refine(Crystal *cr, const RefList *full, PartialityModel pmodel)
 {
 	double max_shift, dev;
 	int i;
@@ -557,9 +558,9 @@ void pr_refine(Crystal *cr, const RefList *full)
 		cell_get_reciprocal(crystal_get_cell(cr), &asx, &asy, &asz,
 			               &bsx, &bsy, &bsz, &csx, &csy, &csz);
 
-		max_shift = pr_iterate(cr, full);
+		max_shift = pr_iterate(cr, full, pmodel);
 
-		update_partialities(cr);
+		update_partialities(cr, pmodel);
 
 		if ( verbose ) {
 			dev = guide_dev(cr, full);
diff --git a/src/post-refinement.h b/src/post-refinement.h
index fe171882..7b822938 100644
--- a/src/post-refinement.h
+++ b/src/post-refinement.h
@@ -40,6 +40,7 @@
 #include "image.h"
 #include "utils.h"
 #include "crystal.h"
+#include "geometry.h"
 
 
 /* Refineable parameters */
@@ -59,10 +60,11 @@ enum {
 };
 
 
-extern void pr_refine(Crystal *cr, const RefList *full);
+extern void pr_refine(Crystal *cr, const RefList *full, PartialityModel pmodel);
 
 /* Exported so it can be poked by tests/pr_gradient_check */
-extern double gradient(Crystal *cr, int k, Reflection *refl);
+extern double gradient(Crystal *cr, int k, Reflection *refl,
+                       PartialityModel pmodel);
 
 
 #endif	/* POST_REFINEMENT_H */
diff --git a/tests/pr_gradient_check.c b/tests/pr_gradient_check.c
index 91b51536..c7eb0129 100644
--- a/tests/pr_gradient_check.c
+++ b/tests/pr_gradient_check.c
@@ -219,7 +219,7 @@ static void calc_either_side(Crystal *cr, double incr_val,
 
 
 static int test_gradients(Crystal *cr, double incr_val, int refine,
-                          const char *str)
+                          const char *str, PartialityModel pmodel)
 {
 	Reflection *refl;
 	RefListIterator *iter;
@@ -285,7 +285,7 @@ static int test_gradients(Crystal *cr, double incr_val, int refine,
 			grad2 = (vals[2][i] - vals[1][i]) / incr_val;
 			grad = (grad1 + grad2) / 2.0;
 
-			cgrad = gradient(cr, refine, refl);
+			cgrad = gradient(cr, refine, refl, pmodel);
 
 			get_partial(refl, &r1, &r2, &p, &cl, &ch);
 
@@ -346,6 +346,7 @@ int main(int argc, char *argv[])
 	struct quaternion orientation;
 	int i;
 	int val;
+	const PartialityModel pmodel = PMODEL_SPHERE;
 
 	image.width = 1024;
 	image.height = 1024;
@@ -388,28 +389,28 @@ int main(int argc, char *argv[])
 			            &bz, &cx, &cy, &cz);
 
 		incr_val = incr_frac * image.div;
-		val += test_gradients(cr, incr_val, REF_DIV, "div");
+		val += test_gradients(cr, incr_val, REF_DIV, "div", pmodel);
 
 		incr_val = incr_frac * ax;
-		val += test_gradients(cr, incr_val, REF_ASX, "ax*");
+		val += test_gradients(cr, incr_val, REF_ASX, "ax*", pmodel);
 		incr_val = incr_frac * ay;
-		val += test_gradients(cr, incr_val, REF_ASY, "ay*");
+		val += test_gradients(cr, incr_val, REF_ASY, "ay*", pmodel);
 		incr_val = incr_frac * az;
-		val += test_gradients(cr, incr_val, REF_ASZ, "az*");
+		val += test_gradients(cr, incr_val, REF_ASZ, "az*", pmodel);
 
 		incr_val = incr_frac * bx;
-		val += test_gradients(cr, incr_val, REF_BSX, "bx*");
+		val += test_gradients(cr, incr_val, REF_BSX, "bx*", pmodel);
 		incr_val = incr_frac * by;
-		val += test_gradients(cr, incr_val, REF_BSY, "by*");
+		val += test_gradients(cr, incr_val, REF_BSY, "by*", pmodel);
 		incr_val = incr_frac * bz;
-		val += test_gradients(cr, incr_val, REF_BSZ, "bz*");
+		val += test_gradients(cr, incr_val, REF_BSZ, "bz*", pmodel);
 
 		incr_val = incr_frac * cx;
-		val += test_gradients(cr, incr_val, REF_CSX, "cx*");
+		val += test_gradients(cr, incr_val, REF_CSX, "cx*", pmodel);
 		incr_val = incr_frac * cy;
-		val += test_gradients(cr, incr_val, REF_CSY, "cy*");
+		val += test_gradients(cr, incr_val, REF_CSY, "cy*", pmodel);
 		incr_val = incr_frac * cz;
-		val += test_gradients(cr, incr_val, REF_CSZ, "cz*");
+		val += test_gradients(cr, incr_val, REF_CSZ, "cz*", pmodel);
 
 	}