coupled phaser

.gitignore

@@ -5,9 +5,8 @@ phasingResult-2.mat
 *.html
 .ipynb_checkpoints/*
 .*.*.swp
+*.ipynb
 deprecated/*
 test_scripts/
 stracthpad*
 Changelog.txt
-singleScrewDislocation.mat
-guidedResult*.mat

Core.py

@@ -36,9 +36,8 @@ def UpdateSupport( self, support ):
         return
 
 # Writer function to manually reset image
-    def resetImage( self, cImg, fSup, reset_error=True ):
-        self._cImage = fftshift( cImg )
-        self._support = fftshift( fSup )
+    def ImageRestart( self, cImg, reset_error=True ):
+        self._cImage = cImg
         if reset_error:
             self._error = []
         return
@@ -131,17 +130,16 @@ def Retrieve( self ):
         return
 
 # Generates a package for the GPU module to read and generate tensors.
-    def generateGPUPackage( self, pcc=False, pcc_params=None ):
-        if pcc_params==None:
-            pcc_params =  np.array( [ 1., 1., 1., 0., 0., 0. ] )
+    def generateGPUPackage( self, pcc=False ):
         mydict = { 
             'array_shape':self._support.shape,
             'modulus':self._modulus, 
             'support':self._support, 
             'beta':self._beta, 
             'cImage':self._cImage,
-            'pcc':pcc, 
-            'pcc_params':pcc_params
+            'pcc':self._pcc,
+            'free_vox_mask':self._free_vox_mask
+
         }
         return mydict
 
@@ -154,7 +152,7 @@ def _initializeSupport( self, sigma=0.575 ):
         self._support = np.zeros( self._arraySize )
         self._support[ np.where( labeled==support_label ) ] = 1.
         self._support = fftshift( self._support )
-#        self.BinaryErosion( 1 )
+        self.BinaryErosion( 1 )
         return
 
 

GPUModule.py

@@ -12,6 +12,8 @@
 import numpy as np
 import tensorflow as tf
 
+import time
+
 import PostProcessing as post
 
 # Class 'Solver' inherits methods from the mixins defined in the following modules.
@@ -33,34 +35,35 @@ class Solver(
 
     def __init__( self, gpack ):   
 
-        self.manageGPUMemory()
+#         self.manageGPUMemory()
 
         # see Phaser.py for definition of gpack
         self.ImportCore( gpack )
         self.generateAlgoDict()
         if gpack[ 'pcc' ]==True: 
+            #self._cImage.assign( self._cImage * tf.cast( tf.reduce_sum( tf.abs( self._modulus )**2 ) / tf.reduce_sum( tf.abs( self._cImage )**2 ), dtype=tf.complex64 ) )
             self._cImage.assign( 
                 self._cImage * tf.sqrt( 
-                    tf.reduce_sum( 
-                        tf.cast( tf.abs( self._modulus ), dtype=tf.complex64 ) 
-                    ) / tf.reduce_sum( 
-                        tf.cast( tf.abs( tf.signal.fft3d( self._cImage ) )**2, dtype=tf.complex64 ) 
-                    )
+                    tf.reduce_sum( tf.cast( tf.abs( self._modulus ), dtype=tf.complex64 ) ) / tf.reduce_sum( tf.cast( tf.abs( tf.signal.fft3d( self._cImage ) )**2, dtype=tf.complex64 ) )
                 )
             )
             self._pccSolver = PCSolver( np.absolute( self._modulus )**2, gpack )
-            self._kernel_f = self._pccSolver.getBlurKernel()
+            self._kernel_f = self._pccSolver.getBlurKernel() 
             self._ModProject = self._ModProjectPC
             self._algodict[ 'PCC' ] = self.PCC
+
         return
 
 
     def manageGPUMemory( self ):
         physical_devices = tf.config.experimental.list_physical_devices( 'GPU' )
+        print(physical_devices)
         assert len(physical_devices) > 0, 'GPU(s) not found. '
-        for this_device in physical_devices:
-            tf.config.experimental.set_memory_growth( this_device, True )
-        return
+        self.__config__ = tf.config.experimental.set_memory_growth(
+            physical_devices[0], 
+            True
+        )
+
 
 
 

GPUModule_Core.py

@@ -13,9 +13,9 @@
 import tensorflow as tf
 import numpy as np
 import functools as ftools
-
+import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm
 import PostProcessing as post
-from scipy.spatial.transform import Rotation
 
 try:
     from pyfftw.interfaces.numpy_fft import fftshift
@@ -26,6 +26,17 @@ class Mixin:
 
     def ImportCore( self, varDict ):
         self._modulus = tf.constant( varDict[ 'modulus' ], dtype=tf.complex64 )
+
+        if varDict[ 'free_vox_mask' ] is None:
+            self.mask = np.ones_like(self._modulus.numpy())
+            self.unmask = self.mask
+
+        else:
+            self.mask = varDict[ 'free_vox_mask' ]
+            self.unmask = self.mask == 0
+
+        self.mask = tf.Variable(self.mask,dtype=tf.float32)
+        self.unmask = tf.Variable(self.unmask,dtype='bool')
         self._support = tf.Variable( varDict[ 'support' ], dtype=tf.complex64 )
         self._support_comp = tf.Variable( 1. - varDict[ 'support' ], dtype=tf.complex64 )
         self._beta = tf.constant( varDict[ 'beta' ], dtype=tf.complex64 )
@@ -36,6 +47,10 @@ def ImportCore( self, varDict ):
         self._cImage_fft_mod = tf.Variable( tf.abs( tf.signal.fft3d( self._cImage ) ) )
         self.BinaryErosion = self.__GPUErosion__
         self._error = []
+        msk = np.where(self._modulus.numpy() > 3,1,0)
+        self.nonzero_mask = tf.Variable(msk,dtype=tf.float32)
+#         self.fourier_sup = (msk*self.mask.numpy()).sum()
+
         self._UpdateError()
 
 
@@ -48,32 +63,16 @@ def ImportCore( self, varDict ):
 
         return
 
-    def resetImage( self, cImg, fSup, reset_error=True ):
-        self._cImage = tf.Variable( fftshift( cImg ), dtype=tf.complex64 )
-        self._support = tf.Variable( fftshift( fSup ), dtype=tf.complex64 )
-        if reset_error:
-            self._error = []
-        return
-
-    def resetParameterList( self, arr ):
-        self._pccSolver._resetParameterList( arr )
-        return
+    def _UpdateError( self ):
+        self._error.append(tf.reduce_sum((self._cImage_fft_mod - tf.abs( self._modulus ))**2).numpy())
+#         self._error.append( tf.reduce_mean((self._cImage_fft_mod - tf.abs( self._modulus ))**2 ,self.nonzero_mask*self.mask) /tf.boolean_mask(tf.abs(self._modulus)**2,self.nonzero_mask*self.mask)).numpy())
 
-        return
-
-    def Modulus( self ):
-        return np.absolute( tf.signal.fftshift( tf.signal.fft3d( self._cImage ) ).numpy() )
 
-    def _UpdateError( self ):
-        self._error.append( 
-            tf.reduce_sum(
-                ( self._cImage_fft_mod - tf.abs( self._modulus ) )**2
-            ).numpy()
-        )
         return
 
     def _UpdateMod( self ):
         self._cImage_fft_mod.assign( tf.abs( tf.signal.fft3d( self._cImage ) ) )
+
         return
 
     def UpdateSupport( self, support ):
@@ -86,11 +85,15 @@ def _CacheImage( self ):
         return
 
     def _UpdateHIOStep( self ):
-        self._cImage.assign( 
-            ( self._support * self._cImage ) +\
-            self._support_comp * ( self._cachedImage - self._beta * self._cImage )
-        )
+        self._cImage.assign(tf.signal.ifft3d( self._modulus * tf.exp( 1.j * tf.cast( 
+                    tf.math.angle( tf.signal.fft3d( self._cImage ) ), 
+                    dtype=tf.complex64 ) 
+            ) ))
+#         self._cImage.assign( tf.where(self.unmask,
+#             ( self._support * self._cImage ) +\
+#             self._support_comp * ( self._cachedImage - self._beta * self._cImage ),self._cImage))
         return
+
 
 # GPU-specific shrinkwrap routine
     def Shrinkwrap( self, sigma, thresh ):
@@ -104,15 +107,20 @@ def Shrinkwrap( self, sigma, thresh ):
         return
 
     def _ModProject( self ):
-        self._cImage.assign( 
-            tf.signal.ifft3d( self._modulus * tf.exp( 1.j * tf.cast( 
+        self._cImage.assign(tf.signal.ifft3d( self._modulus * tf.exp( 1.j * tf.cast( 
                     tf.math.angle( tf.signal.fft3d( self._cImage ) ), 
                     dtype=tf.complex64 
                 ) 
-            ) )
-        )
+            ) ))
+#         self._cImage.assign( tf.where(self.unmask,
+#             tf.signal.ifft3d( self._modulus * tf.exp( 1.j * tf.cast( 
+#                     tf.math.angle( tf.signal.fft3d( self._cImage ) ), 
+#                     dtype=tf.complex64 
+#                 ) 
+#             ) ),self._cImage))
         return
 
+
     def _SupProject( self ):
         self._cImage.assign( self._cImage * self._support )
         return
@@ -124,29 +132,14 @@ def _SupReflect( self ):
         return
 
     def Retrieve( self ):
+
+
         self.finalImage, self.finalSupport = post.centerObject( 
             self._cImage.numpy(), np.absolute( self._support.numpy() )
         )
-        if hasattr( self, '_pccSolver' ):
-            self.pccParameters = self._pccSolver.trainable_variables[0].numpy()
+        self._cImage = tf.Variable(self.finalImage,dtype=tf.complex64)
+        self._support = tf.Variable(self.finalSupport,dtype=tf.complex64)
+
         return
 
-    def getCovarianceMatrix( self ):
-        try:
-            ln = self.pccParameters.size
-        except NameError: # pccParameters does not exist yet
-            self.pccParameters = self._pccSolver.trainable_variables[0].numpy()
-
-        evalues = np.diag( self.pccParameters[:3] )**2
-        ang = self.pccParameters[3] # rotation angle in radians
-        th, ph = tuple( self.pccParameters[4:] )
-        ax = np.array( 
-            [ 
-                np.sin( th ) * np.cos( ph ), 
-                np.sin( th ) * np.sin( ph ), 
-                np.cos( th )
-            ]
-        )
-        evectors = Rotation.from_rotvec( ang*ax ).as_matrix()
-        C = evectors @ evalues @ evectors.T
-        return C
+

GaussPCC.py

@@ -29,6 +29,7 @@ class Mixin: # inherited by Phaser module
     def _ModProjectPC( self ):
         self._cImage_f = tf.signal.fft3d( self._cImage )
         self._patt = tf.signal.fft3d( tf.cast( tf.abs( self._cImage_f )**2, dtype=tf.complex64 ) )
+        #self._patt.assign( self._patt * tf.reduce_sum( self._modulus ) / tf.reduce_sum( self._patt )  )
         self._pcoh_est = tf.sqrt( tf.cast( tf.abs( tf.signal.ifft3d( self._patt * self._kernel_f ) ), dtype=tf.complex64 ) )
         self._cImage.assign( 
             tf.signal.ifft3d( 
@@ -38,22 +39,25 @@ def _ModProjectPC( self ):
         return
 
     def PCC( self, n_iterations, show_progress=False ):
+
+#         print(self._pccSolver.trainable_variables[0].numpy())
         self._pccSolver._setCoherentEstimate( ( tf.abs( tf.signal.fft3d( self._cImage ) )**2 ).numpy() )
         self._pccSolver.Deblur( iterations=n_iterations, show_progress=show_progress )
         self._pccSolver._setupAuxiliary()
         self._pccSolver._updateBlurKernel()
         self._kernel_f = self._pccSolver.getBlurKernel()
+
         return
 
 class PCSolver( tf.Module ): 
 
     def __init__( self, measured_intensity, gpack ):
         self._shape = gpack[ 'array_shape' ]        
         self._modulus_measured = tf.constant( np.sqrt( measured_intensity ), dtype=tf.float32 )
-        pts = self._setupDomain( gpack=gpack )
+        pts, parm_list = self._setupDomain( gpack=gpack )
         self._setupConstants( pts )
         self._setCoherentEstimate( np.absolute( fftn( gpack[ 'cImage' ] ) )**2 )
-        self._setupVariables()
+        self._setupVariables( parm_list )
         self._setupAuxiliary()
         self._updateBlurKernel()
         self._setupOptimizer( learning_rate=0.01, momentum=0.99 )
@@ -74,18 +78,12 @@ def getBlurKernel( self ):
     def _setupDomain( self, gpack ):
         x, y, z = tuple( fftshift( this ) for this in np.meshgrid( *[ np.arange( -n//2., n//2. ) for n in gpack[ 'support' ].shape ] ) )
         pts = np.concatenate( tuple( this.reshape( 1, -1 ) for this in [ x, y, z ] ), axis=0 )
-        if isinstance( gpack[ 'pcc_params' ], type( None ) ):
-            self.parm_list = 0.5, 0.5, 0.5, 0., 0., 0.
+        if 'initial_guess' not in gpack.keys():
+            #l1p, l2p, l3p, psip, thetap, phip = 2., 2., 2., 0., 0., 0.
+            parm_list = 0.5, 0.5, 0.5, 0., 0., 0.
         else:
-            self.parm_list = tuple( gpack[ 'pcc_params' ] )
-        return pts
-
-    def _resetParameterList( self, arr ):
-        self.parm_list = tuple( arr )
-        self._setupVariables()
-        self._setupAuxiliary()
-        self._updateBlurKernel()
-        return
+            parm_list = tuple( vardict[ 'initial_guess' ] )
+        return pts, parm_list
 
     def _setupConstants( self, pts ):
         self._q = tf.constant( pts, dtype=tf.float32 )
@@ -113,8 +111,11 @@ def _setupAuxiliary( self ):
         return
 
 
-    def _setupVariables( self ):
-        self._vars = tf.Variable( np.array( self.parm_list ), dtype=tf.float32 )
+    def _setupVariables( self, parm_list ):
+        #self._l1, self._l2, self._l3, self._psi, self._theta, self._phi = tuple( 
+        #    tf.Variable( this, dtype=tf.float32 ) for this in parm_list
+        #)
+        self._vars = tf.Variable( np.array( parm_list ), dtype=tf.float32 )
         return
 
     #@tf.function       # don't do this, it messes with eager execution
@@ -143,6 +144,8 @@ def Deblur( self, iterations, show_progress=False ):
             allIterations = tqdm( list( range( iterations ) ), desc='PCC' )
         else:
             allIterations = list( range( iterations ) )
+
+
 
         for n in allIterations: 
             with tf.GradientTape( persistent=True ) as tape:
@@ -151,6 +154,8 @@ def Deblur( self, iterations, show_progress=False ):
                 objfun = self.Objective()
 
             gradient = tape.gradient( objfun, self.trainable_variables )
+            #print( objfun )
+            #print( gradient )
             self._optimizer.apply_gradients( zip( gradient, self.trainable_variables ) )
 
         return