Monitor modernization + few cleanups and fixes

.gitignore

@@ -5,8 +5,6 @@
 *.vscode
 *.venv
 logs_tensorboard
-*.hdf5
-*.lock
 .python-version
 .DS_Store
 .idea
@@ -18,4 +16,9 @@ site/
 .cache/mkdocs-jupyter/
 
 # misc
-notes/
+notes/
+
+# files
+*.hdf5
+*.png
+*.lock

PtyLab/Engines/OPR.py

@@ -12,7 +12,7 @@
 import logging
 import sys
 
-import tqdm
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.Engines.BaseEngine import BaseEngine
 from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
@@ -26,7 +26,6 @@
 
 
 class OPR(BaseEngine):
-
     def __init__(
         self,
         reconstruction: Reconstruction,
@@ -79,9 +78,7 @@ def reconstruct(self):
             )
 
         # actual reconstruction ePIE_engine
-        self.pbar = tqdm.trange(
-            self.numIterations, desc="OPR", file=sys.stdout, leave=True
-        )
+        self.pbar = trange(self.numIterations, desc="OPR", file=sys.stdout, leave=True)
         for loop in self.pbar:
             self.it = loop
             # set position order
@@ -205,8 +202,13 @@ def _orthogonalizeIncoherentModes_batched(self):
         # stays bounded regardless of frame count; the batched call is already
         # wide enough at a few dozen frames to hide launch overhead.
         elements_per_frame = nModes * n * n
-        chunk = int(max(1, self._orthogonalization_chunk_bytes //
-                        (elements_per_frame * stack.dtype.itemsize)))
+        chunk = int(
+            max(
+                1,
+                self._orthogonalization_chunk_bytes
+                // (elements_per_frame * stack.dtype.itemsize),
+            )
+        )
 
         flat = stack[0, 0, :, 0, :, :, :].reshape(nModes, n * n, nFrames)
         for start in range(0, nFrames, chunk):

PtyLab/Engines/aPIE.py

@@ -1,9 +1,9 @@
 from typing import Any
 
 import numpy as np
-import tqdm
 from matplotlib import pyplot as plt
 from scipy.interpolate import interp2d
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.utils.visualisation import hsvplot
 
@@ -83,9 +83,7 @@ def doReconstruction(self):
             self.thetaSearchRadiusMax, self.thetaSearchRadiusMin, self.numIterations
         )
 
-        self.pbar = tqdm.trange(
-            self.numIterations, desc="aPIE", file=sys.stdout, leave=True
-        )
+        self.pbar = trange(self.numIterations, desc="aPIE", file=sys.stdout, leave=True)
         for loop in self.pbar:
             # save theta search history
             self.reconstruction.thetaHistory = np.append(

PtyLab/Engines/e3PIE.py

@@ -1,6 +1,6 @@
 import numpy as np
-import tqdm
 from matplotlib import pyplot as plt
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 try:
     import cupy as cp
@@ -25,7 +25,6 @@
 
 
 class e3PIE(BaseEngine):
-
     def __init__(
         self,
         reconstruction: Reconstruction,
@@ -76,7 +75,7 @@ def reconstruct(self):
         # get module
         xp = getArrayModule(self.reconstruction.object)
 
-        self.pbar = tqdm.trange(
+        self.pbar = trange(
             self.numIterations, desc="e3PIE", file=sys.stdout, leave=True
         )
 
@@ -123,7 +122,6 @@ def reconstruct(self):
                 ]
                 # update object slice
                 for loopTemp in range(self.reconstruction.nslice - 1):
-
                     sliceLoop = self.reconstruction.nslice - 1 - loopTemp
 
                     # temp_delta = self.reconstruction.esw[..., sliceLoop, sy, sx]
@@ -137,7 +135,7 @@ def reconstruct(self):
                         )
                     )
                     # eswTemp update (here probe incident on last slice)
-                    beth = 0.9 # todo, why need beth, not betaProbe, changable?
+                    beth = 0.9  # todo, why need beth, not betaProbe, changable?
                     self.reconstruction.probe[:, :, :, sliceLoop, ...] = (
                         self.probeUpdate(
                             objectPatch[:, :, :, sliceLoop, ...],

PtyLab/Engines/ePIE.py

@@ -2,8 +2,8 @@
 import sys
 
 import numpy as np
-import tqdm
 from matplotlib import pyplot as plt
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.Engines.BaseEngine import BaseEngine
 from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
@@ -66,6 +66,7 @@ class ePIE(BaseEngine):
         imaging," Ultramicroscopy 109, 1256-1262 (2009).
         https://doi.org/10.1016/j.ultramic.2009.05.012
     """
+
     def __init__(
         self,
         reconstruction: Reconstruction,
@@ -169,8 +170,9 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None):
         probe estimate is updated sequentially across all scan positions.
 
         After all scan positions in an iteration have been processed,
-        `getErrorMetrics()` evaluates the reconstruction error and
-        `applyConstraints()` applies the enabled reconstruction constraints.
+        `getErrorMetrics()` evaluates the reconstruction error,
+        `applyConstraints()` applies the enabled reconstruction constraints and
+        `showReconstruction()` updates the monitor.
 
         The method yields after every scan-position update, allowing custom code
         to be interleaved with the reconstruction.
@@ -190,9 +192,7 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None):
         self._prepareReconstruction()
 
         # actual reconstruction ePIE_engine
-        self.pbar = tqdm.trange(
-            self.numIterations, desc="ePIE", file=sys.stdout, leave=True
-        )
+        self.pbar = trange(self.numIterations, desc="ePIE", file=sys.stdout, leave=True)
         for loop in self.pbar:
             # set position order
             self.setPositionOrder()
@@ -234,10 +234,10 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None):
                     object_patch = self.objectPatchUpdate(objectPatch, DELTA)
 
                 self.reconstruction.object[..., sy, sx] = object_patch
-                
-                #self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
+
+                # self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                 #    objectPatch, DELTA
-                #)
+                # )
 
                 # probe update
                 self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)
@@ -255,8 +255,8 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None):
             # apply Constraints
             self.applyConstraints(loop)
 
-            # show reconstruction
-            # self.showReconstruction(loop)
+            # update the monitor, as every other engine does after the constraints
+            self.showReconstruction(loop)
 
         if self.params.gpuFlag:
             self.logger.info("switch to cpu")

PtyLab/Engines/mPIE.py

@@ -13,7 +13,7 @@
 import sys
 import warnings
 
-import tqdm
+from tqdm.auto import tqdm, trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.Engines.BaseEngine import BaseEngine
 from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
@@ -98,6 +98,7 @@ class mPIE(BaseEngine):
             Baseline ePIE reconstruction without rPIE regularization or
             momentum acceleration.
     """
+
     def __init__(
         self,
         reconstruction: Reconstruction,
@@ -127,7 +128,7 @@ def __init__(
             monitor (Monitor):
                 Monitor used for reconstruction visualization and progress
                 reporting.
-    """
+        """
         super().__init__(reconstruction, experimentalData, params, monitor)
         self.logger = logging.getLogger("mPIE")
         self.logger.info("Successfully created mPIE engine")
@@ -200,7 +201,9 @@ def initializeReconstructionParams(self):
 
         self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)
 
-    def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each_iteration=None):
+    def reconstruct(
+        self, experimentalData=None, reconstruction=None, vis_after_each_iteration=None
+    ):
         r"""
         Run the mPIE reconstruction to completion.
 
@@ -273,13 +276,11 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each
         self.reconstruction.objectBuffer = self.reconstruction.object.copy()
         self.reconstruction.probeBuffer = self.reconstruction.probe.copy()
         # actual reconstruction MPIE_engine
-        self.pbar = tqdm.trange(
-            self.numIterations, desc="mPIE", file=sys.stdout, leave=True
-        )
+        self.pbar = trange(self.numIterations, desc="mPIE", file=sys.stdout, leave=True)
         for loop in self.pbar:
             # set position order
             self.setPositionOrder()
-            self.pbar_pos = tqdm.tqdm(
+            self.pbar_pos = tqdm(
                 self.positionIndices, leave=False, desc="ptychogram", file=sys.stdout
             )
             for positionLoop, positionIndex in enumerate(self.pbar_pos):
@@ -329,9 +330,7 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each
                 # self.reconstruction.push_probe_update(self.reconstruction.probe, positionIndex, self.experimentalData.ptychogram.shape[0])
 
                 if self.params.positionCorrectionSwitch:
-                    self.positionCorrection(
-                        objectPatch, positionIndex, sy, sx
-                    )
+                    self.positionCorrection(objectPatch, positionIndex, sy, sx)
                     # self.pbar_pos.write(f'Corr: {shifter[0]*1e6:.2f} um x {shifter[1]*1e6:.2f} um')
 
                 # momentum updates
@@ -610,4 +609,4 @@ def stepM(self):
 
     @stepM.setter
     def stepM(self, value):
-        self.frictionM = value
+        self.frictionM = value

PtyLab/Engines/mqNewton.py

@@ -12,7 +12,7 @@
 import logging
 import sys
 
-import tqdm
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.Engines.BaseEngine import BaseEngine
 from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
@@ -85,7 +85,7 @@ def reconstruct(self, experimentalData: ExperimentalData = None):
         self._prepareReconstruction()
         self.initializeAdaptiveMomentum()
 
-        self.pbar = tqdm.trange(
+        self.pbar = trange(
             self.numIterations, desc="mqNewton", file=sys.stdout, leave=True
         )
         for loop in self.pbar:

PtyLab/Engines/multiPIE.py

@@ -12,7 +12,7 @@
 import logging
 import sys
 
-import tqdm
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.Engines.BaseEngine import BaseEngine
 from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
@@ -68,7 +68,7 @@ def initializeReconstructionParams(self):
     def reconstruct(self):
         self._prepareReconstruction()
 
-        self.pbar = tqdm.trange(
+        self.pbar = trange(
             self.numIterations, desc="multiPIE", file=sys.stdout, leave=True
         )
 

PtyLab/Engines/qNewton.py

@@ -12,7 +12,7 @@
 import logging
 import sys
 
-import tqdm
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 from PtyLab.Engines.BaseEngine import BaseEngine
 from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
@@ -57,7 +57,7 @@ def reconstruct(self, experimentalData: ExperimentalData = None):
             self.experimentalData = experimentalData
         self._prepareReconstruction()
 
-        self.pbar = tqdm.trange(
+        self.pbar = trange(
             self.numIterations, desc="qNewton", file=sys.stdout, leave=True
         )
         for loop in self.pbar:

PtyLab/Engines/zPIE.py

@@ -1,6 +1,6 @@
 import numpy as np
-import tqdm
 from matplotlib import pyplot as plt
+from tqdm.auto import trange  # widget bar in Jupyter, text bar in a terminal
 
 try:
     import cupy as cp
@@ -112,7 +112,7 @@ def reconstruct(self, experimentalData=None, reconstruction=None):
         X, Y = xp.meshgrid(xp.arange(-n // 2, n // 2), xp.arange(-n // 2, n // 2))
         w = xp.exp(-((xp.sqrt(X**2 + Y**2) / self.reconstruction.Np) ** 4))
 
-        self.pbar = tqdm.trange(
+        self.pbar = trange(
             self.numIterations, desc="zPIE", file=sys.stdout, leave=True
         )  # in order to change description to the tqdm progress bar
         for loop in self.pbar: