interpolation and zoom for incoherent iterate stacks

fixed random guess creation?

proxtoolbox/experiments/orbitaltomography/degenerate_orbits.py

@@ -66,7 +66,7 @@ class DegenerateOrbital(PlanarMolecule):
         """
         self.u0 = np.array([self.data * np.exp(1j * 2 * np.pi * np.random.random_sample(self.data.shape))
                             for i in range(self.degeneracy)])
-        self.u0 = np.fft.fftn(self.u0, axes=(-1, -2))
+        self.u0 = shifted_fft(self.u0, axes=(-1, -2))
 
     def setupProxOperators(self):
         """
@@ -108,7 +108,7 @@ class DegenerateOrbital(PlanarMolecule):
         ax[1][-1].set_title("Local density of states")
         plt.show()
 
-    def show(self):
+    def show(self, **kwargs):
         """
         Create basic result plots of the phase retrieval procedure
         """
@@ -117,9 +117,14 @@ class DegenerateOrbital(PlanarMolecule):
         self.output['u1'] = self.algorithm.prox1.eval(self.algorithm.u)
         self.output['u2'] = self.algorithm.prox2.eval(self.algorithm.u)
 
+        figsize = kwargs.pop("figsize", (12,3))
         for i, operator in enumerate(self.algorithm.proxOperators):
             operator_name = self.proxOperators[i].__name__
-            f = self.plot_guess(operator.eval(self.algorithm.u), name="%s satisfied" % operator_name, show=False)
+            f = self.plot_guess(operator.eval(self.algorithm.u),
+                                name="%s satisfied" % operator_name,
+                                show=False,
+                                interpolate_and_zoom=self.interpolate_and_zoom,
+                                figsize=figsize)
             self.output['plots'].append(f)
 
         # f1 = self.plot_guess(self.output['u1'], name='Best approximation: physical constraint satisfied', show=False)
@@ -136,18 +141,20 @@ class DegenerateOrbital(PlanarMolecule):
     # def saveOutput(self, **kwargs):
     #     super(PlanarMolecule, self).saveOutput(**kwargs)
 
-    def plot_guess(self, u, name=None, show=True):
+    def plot_guess(self, u, name=None, show=True, interpolate_and_zoom=False, figsize=(12,3)):
         """"Given a list of fields, plot the individual fields and the combined intensity"""
-        fig, ax = plt.subplots(1, len(u) + 2, figsize=(12, 3), num=name)
+        prop = self.propagator(self)
+        u_hat = prop.eval(u)
+        fourier_intensity = np.sqrt(np.sum(abs(u_hat) ** 2, axis=0))
+        if interpolate_and_zoom:
+            u_show = self.interp_zoom_field(u)
+        fig, ax = plt.subplots(1, len(u) + 2, figsize=figsize, num=name)
         for ii in range(self.degeneracy):
-            im = ax[ii].imshow(complex_to_rgb(u[ii]))
+            im = ax[ii].imshow(complex_to_rgb(u_show[ii]))
             ax[ii].set_title("Degenerate orbit %d" % ii)
-        im = ax[-2].imshow(np.sum(abs(u) ** 2, axis=0))
+        im = ax[-2].imshow(np.sum(abs(u_show) ** 2, axis=0))
         ax[-2].set_title("Local density of states")
         plt.colorbar(im, ax=ax[-2])
-        prop = self.propagator(self)
-        u_hat = prop.eval(u)
-        fourier_intensity = np.sqrt(np.sum(abs(u_hat)**2, axis=0))
         im = ax[-1].imshow(fourier_intensity)
         ax[-1].set_title("Fourier domain intensity")
         plt.colorbar(im, ax=ax[-1])
@@ -156,9 +163,15 @@ class DegenerateOrbital(PlanarMolecule):
             plt.show()
         return fig
 
-    def interpolate_and_zoom(self, u, interpolation=2, zoom=0.5):
+    def interp_zoom_field(self, u, interpolation=2, zoom=0.5):
         """
         interpolate a field and zoom in to the center
         """
-        return u
+        nt, ny, nx = u.shape
+        zmy = int(ny * zoom) // 2
+        zmx = int(nx * zoom) // 2
+        zoomed = u[:, zmy:ny - zmy, zmx:nx - zmx]
+
+        interpolated = np.array([fourier_interpolate(u_i, factor=interpolation) for u_i in zoomed])
 
+        return interpolated