documentation + plotInputData improvement

proxtoolbox/experiments/orbitaltomography/planar_molecule.py

@@ -7,16 +7,17 @@ from proxtoolbox.experiments.orbitaltomography.orbitExperiment import OrbitalTom
 from proxtoolbox.utils.visualization.complex_field_visualization import complex_to_rgb
 from proxtoolbox.utils.orbitaltomog import bin_array, shifted_fft, shifted_ifft, fourier_interpolate, roll_to_pos
 
-#for downloading data
+# for downloading data
 import proxtoolbox.utils.GetData as GetData
 from proxtoolbox.utils.GetData import datadir
 
+
 class PlanarMolecule(OrbitalTomographyExperiment):
     @staticmethod
     def getDefaultParameters():
         # TODO: optimize parameters and proxoperators to get good & consistent phase retrieval using the demo
         defaultParams = {
-            'experiment_name': 'noisy 2D ARPES', # '2D ARPES', #
+            'experiment_name': 'noisy 2D ARPES',  # '2D ARPES', #
             'data_filename': datadir/'OrbitalTomog'/'coronene_homo1.tif',
             'from_intensity_data': False,
             'object': 'real',
@@ -48,6 +49,16 @@ class PlanarMolecule(OrbitalTomographyExperiment):
         return defaultParams
 
     def __init__(self, **kwargs):
+        """
+        Orbital tomography for planar molecules, yielding 2D reconstructions from single-photon-energy ARPES
+        measurements.
+
+        Kwargs can include:
+            - data_filename: full path to data file
+            - support_filename: full path to file indicating the object support
+            - sparsity_parameter: integer number of pixels of the object
+            - and more...
+        """
         super(PlanarMolecule, self).__init__(**kwargs)
 
         # do here any data member initialization
@@ -140,16 +151,17 @@ class PlanarMolecule(OrbitalTomographyExperiment):
 
     def plotInputData(self):
         """Quick plotting routine to show the data, initial guess and the sparsity support"""
-        fig, ax = plt.subplots(1, 3, figsize=(12, 3.5))
+        fig, ax = plt.subplots(1, 3, figsize=(8, 3))
         im = ax[0].imshow(self.data)
-        plt.colorbar(im, ax=ax[0])
-        plt.title("Photoelectron spectrum")
+        plt.colorbar(im, ax=ax[0], shrink=0.6)
+        ax[0].set_title("Photoelectron spectrum")
         ax[1].imshow(complex_to_rgb(self.u0))
-        plt.title("Initial guess")
+        ax[1].set_title("Initial guess")
         if self.sparsity_support is not None:
             im = ax[2].imshow(self.sparsity_support, cmap='gray')
             # plt.colorbar(im, ax=ax[2])
-            plt.title("Sparsity support")
+            ax[2].set_title("Sparsity support")
+        plt.tight_layout()
         plt.show()
 
     def show(self):
@@ -206,12 +218,12 @@ class PlanarMolecule(OrbitalTomographyExperiment):
         ax[0].set_title("Measured data")
         prop = self.propagator(self)
         guess = self.algorithm.prox2.eval(self.algorithm.u)
-        guess = roll_to_pos(guess, pos=tuple([s//2 for s in guess.shape]), move_maximum=True)
+        guess = roll_to_pos(guess, pos=tuple([s // 2 for s in guess.shape]), move_maximum=True)
         guess = roll_to_pos(guess, pos=tuple([s // 2 for s in guess.shape]))
         u_hat = prop.eval(self.algorithm.prox1.eval(guess))
         ax[1].imshow(abs(u_hat))
         ax[1].set_title("Predicted measurement intensity")
-        ax[2].imshow(complex_to_rgb(u_hat))
+        ax[2].imshow(complex_to_rgb(u_hat, background='light'))
         ax[2].set_title("Predicted phase (by color)")
         h.tight_layout()
         plt.show()