Added difference output

Author: Oier Lauzirika Zarrabeitia

src/xmipp/applications/scripts/volume_consensus/volume_consensus.py

@@ -28,9 +28,12 @@
 **************************************************************************
 """
 import math
+import os
+import itertools
 import numpy as np
 import pywt
 import pywt.data
+
 from scipy.ndimage import zoom
 from xmipp_base import XmippScript
 import xmippLib
@@ -53,6 +56,25 @@ def run(self):
         outVolFn = self.getParam('-o')
         self.computeVolumeConsensus(inputFile, outVolFn)
 
+    def resize(self, image, dim):
+        imageFt = np.fft.rfftn(image)
+        resultFt = np.zeros(dim[:-1] + (dim[-1]//2+1,), dtype=imageFt.dtype)
+
+        copyExtent = np.minimum(image.shape, dim) // 2
+        srcCornerStart = image.shape-copyExtent
+        dstCornerStart = dim-copyExtent
+        for corners in itertools.product(range(2), repeat=len(dim)-1):
+            corners = np.array(corners + (0, ))
+            srcStart = np.where(corners, srcCornerStart, 0)
+            srcEnd = srcStart + copyExtent
+            dstStart = np.where(corners, dstCornerStart, 0)
+            dstEnd = dstStart + copyExtent
+            srcSlices = [slice(s, e) for s, e in zip(srcStart, srcEnd)]
+            dstSlices = [slice(s, e) for s, e in zip(dstStart, dstEnd)]
+            resultFt[tuple(dstSlices)] = imageFt[tuple(srcSlices)]
+            
+        return np.fft.irfftn(resultFt)
+        
     def computeVolumeConsensus(self, inputFile, outVolFn, wavelet='sym11'):
         outputWt = None
         outputMin = None
@@ -77,7 +99,8 @@ def computeVolumeConsensus(self, inputFile, outVolFn, wavelet='sym11'):
                     zdim2 = 2**(math.ceil(math.log2(zdimOrig)))  # Next power of 2
 
                 if xdimOrig!=xdim2 or ydimOrig!=ydim2 or zdimOrig!=zdim2:
-                    volume = zoom(volume, (xdim2/xdimOrig,ydim2/ydimOrig,zdim2/zdimOrig))
+                    #volume = zoom(volume, (xdim2/xdimOrig,ydim2/ydimOrig,zdim2/zdimOrig))
+                    volume = self.resize(volume, (zdim2, ydim2, xdim2))
 
                 nlevel = pywt.dwtn_max_level(volume.shape, wavelet=wavelet)
                 wt = pywt.wavedecn(
@@ -88,8 +111,14 @@ def computeVolumeConsensus(self, inputFile, outVolFn, wavelet='sym11'):
 
                 if outputWt == None:
                     outputWt = wt
-                    #outputMin = np.zeros_like(wt[1]['aaa'])
+                    outputMin = np.zeros_like(volume)
                 else:
+                    diff = np.abs(np.abs(wt[0]) - np.abs(outputWt[0]))
+                    diff = self.resize(diff, outputMin.shape)
+                    np.maximum(
+                        diff, outputMin,
+                        out=outputMin
+                    )
                     outputWt[0] = np.where(
                         np.abs(wt[0]) > np.abs(outputWt[0]),
                         wt[0], outputWt[0]
@@ -102,28 +131,30 @@ def computeVolumeConsensus(self, inputFile, outVolFn, wavelet='sym11'):
                             wtLevelDetail = wtLevel[detail]
                             outputWtLevelDetail = outputWtLevel[detail]
 
-                            outputWtLevelDetail[...] = np.where(
-                                np.abs(wtLevelDetail) > np.abs(outputWtLevelDetail),
-                                wtLevelDetail, outputWtLevelDetail
-                            )
-                            
-                            """
                             diff = np.abs(np.abs(wtLevelDetail) - np.abs(outputWtLevelDetail))
+                            diff = self.resize(diff, outputMin.shape)
                             np.maximum(
                                 diff, outputMin,
                                 out=outputMin
                             )
-                            """
 
+                            outputWtLevelDetail[...] = np.where(
+                                np.abs(wtLevelDetail) > np.abs(outputWtLevelDetail),
+                                wtLevelDetail, outputWtLevelDetail
+                            )
+                            
+
             f.close()
         consensus = pywt.waverecn(outputWt, wavelet)
         if xdimOrig!=xdim2 or ydimOrig!=ydim2 or zdimOrig!=zdim2:
-            consensus = zoom(consensus, (xdimOrig/xdim2,ydimOrig/ydim2,zdimOrig/zdim2))
+            consensus = self.resize(consensus, (zdimOrig, ydimOrig, xdimOrig))
         image.setData(consensus)
         image.write(outVolFn)
-        #image.setData(outputMin)
-        #outVolFn2 = splitext(outVolFn)[0] + '_diff.mrc'
-        #image.write(outVolFn2)
+        if xdimOrig!=xdim2 or ydimOrig!=ydim2 or zdimOrig!=zdim2:
+            outputMin = self.resize(outputMin, (zdimOrig, ydimOrig, xdimOrig))
+        image.setData(outputMin)
+        outVolFn2 = os.path.splitext(outVolFn)[0] + '_diff.mrc'
+        image.write(outVolFn2)
 
 
 if __name__=="__main__":