Merge pull request #331 from jonwright/master

Fix #330

ImageD11/sinograms/sinogram.py

@@ -109,10 +109,17 @@ def prepare_peaks_from_4d(self, cf_4d, gord, inds, grain_label, hkltol=0.25):
 
         self.cf_for_sino = flt
 
-    def build_sinogram(self):
+    def build_sinogram(self, columns=('omega',)):
         """
         Computes sinogram for this grain using all peaks in self.cf_for_sino
+
+        columns = list of columns to produce sinograms of intensity*value
+                for example: "ds", "eta", "omega", etc
         """
+        for a in columns:
+            assert a in self.cf_for_sino.titles
+        assert 'omega' in columns
+
         NY = len(self.ds.ybincens)  # number of y translations
         iy = np.round(
             (self.cf_for_sino.dty - self.ds.ybincens[0]) / (self.ds.ybincens[1] - self.ds.ybincens[0])).astype(
@@ -135,26 +142,37 @@ def build_sinogram(self):
         pkrow = np.cumsum(pkmsk.ravel()).reshape(pkmsk.shape) - 1  #
         # counting where we hit an HKL position with a found peak
         # e.g (-10, -9, -10) didn't get hit, but the next one did, so increment
-
+        #
+        # found peak indexes into the pkmsk array
+        pkhkle = np.arange( np.prod( pkmsk.shape ), dtype=int )[ pkmsk.flat == 1 ]
+        # hkl indices (transpose for 3,N versus N,3)
+        pkindices = np.array( np.unravel_index(pkhkle, pkmsk.shape) )
+        pkindices[:3] += hklmin[:,np.newaxis]
         npks = pkmsk.sum()
         destRow = pkrow[dh[0], dh[1], dh[2], de]
         sino = np.zeros((npks, NY), 'f')
         hits = np.zeros((npks, NY), 'f')
-        angs = np.zeros((npks, NY), 'f')
         adr = destRow * NY + iy
         # Just accumulate
         sig = self.cf_for_sino.sum_intensity
         ImageD11.cImageD11.put_incr64(sino, adr, sig)
         ImageD11.cImageD11.put_incr64(hits, adr, np.ones(len(de), dtype='f'))
-        ImageD11.cImageD11.put_incr64(angs, adr, self.cf_for_sino.omega)
-
-        sinoangles = angs.sum(axis=1) / hits.sum(axis=1)
+        # intensity weighted sums
+        angs = {}
+        for name in columns:
+            angs[name] = np.zeros((npks, NY), 'f')
+            ImageD11.cImageD11.put_incr64(angs[name], adr, self.cf_for_sino[name] * sig )
+        sinoangles = angs['omega'] .sum(axis=1) / sino.sum(axis=1)
         # Normalise:
-        self.sino = (sino.T / sino.max(axis=1)).T
+        self.proj_scale = sino.max(axis=1)
+        self.sino = sino / self.proj_scale[:, np.newaxis]
         # Sort (cosmetic):
         order = np.lexsort((np.arange(npks), sinoangles))
         self.sinoangles = sinoangles[order]
         self.ssino = self.sino[order].T
+        if len(columns)>1:
+            self.angle_wt_sinos = { name : angs[name][order].T for name in columns }
+        self.hkle = pkindices[:, order]  # dims are [ (h,k,l,sign(eta)) , nprojections ]
 
     def update_lab_position_from_peaks(self, cf_4d, grain_label):
         """Updates translation of self.grain using peaks in assigned 4D colfile.
@@ -215,32 +233,56 @@ def update_recon_parameters(self, pad=None, shift=None, mask=None, niter=None, y
         if y0 is not None:
             self.recon_y0 = y0
 
-    def recon(self, method="iradon", workers=1, **extra_args):
-        """Performs reconstruction given reconstruction method"""
+    def recon(self,
+              method="iradon",
+              workers=1,
+              projections = None,
+              **extra_args):
+        """Performs reconstruction given reconstruction method
+
+        method = "iradon", "mlem", "astra"
+        workers = threads to use - passed to the method
+
+        projections = which projections to use for reconstruction using a subset of data (twins)
+                iradon( self.sino[ :, projections ], theta = self.sinoangles[ projections ], ... )
+                effectively a sinogram mask - perhaps this should move to iradon code.
+
+        extra_args = passed to the projection function
+
+           iradon -> ImageD11.sinograms.roi_iradon.run_iradon
+           { pad=20, shift=0, workers=1, mask=None,
+                      apply_halfmask=False, mask_central_zingers=False, central_mask_radius=25 }
+           mlem : ImageD11.sinograms.roi_iradon.run_mlem
+              { mask=None, pad=20, shift=0, workers=1, niter=20, apply_halfmask=False,
+             mask_central_zingers=False, central_mask_radius=25 }
+           astra : run_astra in this file
+               { shift=0, pad=0, mask=None, niter=100, astra_method='SIRT_CUDA', workers=None }
+        """
         if method not in ["iradon", "mlem", "astra"]:
             raise ValueError("Unsupported method!")
         if self.ssino is None or self.sinoangles is None:
             raise ValueError("Sorted sino or sinoangles are missing/have not been computed, unable to reconstruct.")
-
-        sino = self.ssino
-        angles = self.sinoangles
-
+        if projections is None:
+            sino = self.ssino
+            angles = self.sinoangles
+        else:
+            sino = self.ssino[:, projections]
+            angles = self.sinoangles[ projections ]
+            print("Using subset",sino.shape,"from", self.ssino.shape)
         if method == "iradon":
             recon_function = ImageD11.sinograms.roi_iradon.run_iradon
         elif method == "mlem":
             # MLEM has niter as an extra argument
+            recon_function = ImageD11.sinograms.roi_iradon.run_mlem
             # Overwrite the default argument if self.recon_niter is set
+            if self.recon_niter is not None and "niter" not in extra_args:
+                extra_args[ "niter" ] = self.recon_niter
             # TODO: We should think about default reconstruction arguments in more detail
             # At the moment, we could pass None for pad or shift etc to recon_function if they are not manually set, which seems dangerous
             # Do we check for None at the start of this function?
             # Or do we initialise them to sensible default values inside __init__?
-            if self.recon_niter is not None:
-                recon_function = partial(ImageD11.sinograms.roi_iradon.run_mlem, niter=self.recon_niter)
-            else:
-                recon_function = ImageD11.sinograms.roi_iradon.run_mlem
         elif method == "astra":
             recon_function = run_astra
-
         recon = recon_function(sino=sino,
                                angles=angles,
                                pad=self.recon_pad,
@@ -250,6 +292,7 @@ def recon(self, method="iradon", workers=1, **extra_args):
                                **extra_args)
 
         self.recons[method] = recon
+        return recon
 
     def get_shape_mask(self, method="iradon", cutoff=None):
         """Gets a boolean mask representing the grain shape from the grain reconstruction.