Monitoring

PtyLab.Monitor

Monitor

AbstractMonitor

Bases: object

This monitor implements all the basic features that you have to override for a custom monitor.

Alternatively, you can instantiate this class to create a monitor that does not do anything and therefore, will not take time to run.

Source code in PtyLab/Monitor/Monitor.py

class AbstractMonitor(object):
    """
    This monitor implements all the basic features that you have to override for a custom monitor.

    Alternatively, you can instantiate this class to create a monitor that does not do anything and therefore,
    will not take time to run.

    """

    def initializeMonitors(self):
        """
        This code is run after __init__, use it to build a.k.a. GUI elements.

        :return: Nothing
        """

        pass

    def update_focusing_metric(self, TV_value, AOI_image, metric_name, allmerits=None):
        """
        Show the total variation of the object estimate inside the area of interest.
        :param TV_value: Value of TV
        :return:
        """
        pass

    def writeEngineName(self, name):
        """
        Save the engine name for this particular iteration.

        :param name:  Name of the engine that is used for this iteration.

        :return:
        """
        pass

    def visualize_probe_engine(self, estimate):
        """
        For nonlinear imaging only. Return the electric field of the probe.
        """
        pass

    def updatePlot(
        self,
        object_estimate: np.ndarray,
        probe_estimate: np.ndarray,
        zo=None,
        encoder_positions=None,
    ):
        """
        Update the visualisation of both probe and and object estimate.

        Please note, that only the part that is defined by probeZoom and objectZoom is uploaded.

        #TODO: Change that, upload the entire image and make the monitor decide which parts to show.

        :param object_estimate: Array with the object estimate.
        :param probe_estimate:  Array with the probe estimate.
        :return:
        """
        pass

    def getOverlap(self, *args, **kwargs):
        """
        Get the overlap between subsequent probes.

        Todo(dbs660): Move this.

        :param args:
        :param kwargs:
        :return:
        """
        # TODO: remove this from here, it should go in some utils
        pass

    def update_positions(self, *args, **kwargs):
        """Update the position information."""
        pass

    def update_encoder(self, corrected_positions, original_positions, *args, **kwargs):
        """Update the image of the encoder positions."""
        pass

    def update_overlap(self, overlap_area, linear_overlap):
        pass

    def updateObjectProbeErrorMonitor(
        self,
        error: float,
        object_estimate: np.ndarray,
        probe_estimate: np.ndarray,
        zo=None,
        purity_probe=None,
        purity_object=None,
        encoder_positions=None,
    ):
        """
        Update the Object and Probe error monitor, and any associated metrics.

        :param error: Value of the loss for the current iteration.
        :param object_estimate: Current object estimate
        :param probe_estimate: Current probe estimate
        :param zo: Sample-detector distance
        :param purity_probe: purity of the probe
        :param purity_object: purity of the object
        :return:
        """

        pass

    def updateBeamWidth(self, beamwidth_x, beamwidth_y):
        """
        Update the beam width in X and Y

        Parameters
        ----------
        beamwidth_x: beamwidth x in meter
        beamwidth_y: beamwidth y in meter

        Returns
        -------

        """
        pass

    def updateDiffractionDataMonitor(self, Iestimated, Imeasured):
        """
        Update the diffraction data estimate at the current iteration.

        :param Iestimated: Estimated intensity at the current position
        :param Imeasured:  Measured intensity at the current position
        :return:
        """

        pass

initializeMonitors()

This code is run after init, use it to build a.k.a. GUI elements.

:return: Nothing

Source code in PtyLab/Monitor/Monitor.py

def initializeMonitors(self):
    """
    This code is run after __init__, use it to build a.k.a. GUI elements.

    :return: Nothing
    """

    pass

update_focusing_metric(TV_value, AOI_image, metric_name, allmerits=None)

Show the total variation of the object estimate inside the area of interest. :param TV_value: Value of TV :return:

Source code in PtyLab/Monitor/Monitor.py

def update_focusing_metric(self, TV_value, AOI_image, metric_name, allmerits=None):
    """
    Show the total variation of the object estimate inside the area of interest.
    :param TV_value: Value of TV
    :return:
    """
    pass

writeEngineName(name)

Save the engine name for this particular iteration.

:param name: Name of the engine that is used for this iteration.

:return:

Source code in PtyLab/Monitor/Monitor.py

def writeEngineName(self, name):
    """
    Save the engine name for this particular iteration.

    :param name:  Name of the engine that is used for this iteration.

    :return:
    """
    pass

visualize_probe_engine(estimate)

For nonlinear imaging only. Return the electric field of the probe.

Source code in PtyLab/Monitor/Monitor.py

def visualize_probe_engine(self, estimate):
    """
    For nonlinear imaging only. Return the electric field of the probe.
    """
    pass

updatePlot(object_estimate, probe_estimate, zo=None, encoder_positions=None)

Update the visualisation of both probe and and object estimate.

Please note, that only the part that is defined by probeZoom and objectZoom is uploaded.

TODO: Change that, upload the entire image and make the monitor decide which parts to show.

:param object_estimate: Array with the object estimate. :param probe_estimate: Array with the probe estimate. :return:

Source code in PtyLab/Monitor/Monitor.py

def updatePlot(
    self,
    object_estimate: np.ndarray,
    probe_estimate: np.ndarray,
    zo=None,
    encoder_positions=None,
):
    """
    Update the visualisation of both probe and and object estimate.

    Please note, that only the part that is defined by probeZoom and objectZoom is uploaded.

    #TODO: Change that, upload the entire image and make the monitor decide which parts to show.

    :param object_estimate: Array with the object estimate.
    :param probe_estimate:  Array with the probe estimate.
    :return:
    """
    pass

getOverlap(*args, **kwargs)

Get the overlap between subsequent probes.

Todo(dbs660): Move this.

:param args: :param kwargs: :return:

Source code in PtyLab/Monitor/Monitor.py

def getOverlap(self, *args, **kwargs):
    """
    Get the overlap between subsequent probes.

    Todo(dbs660): Move this.

    :param args:
    :param kwargs:
    :return:
    """
    # TODO: remove this from here, it should go in some utils
    pass

update_positions(*args, **kwargs)

Update the position information.

Source code in PtyLab/Monitor/Monitor.py

def update_positions(self, *args, **kwargs):
    """Update the position information."""
    pass

update_encoder(corrected_positions, original_positions, *args, **kwargs)

Update the image of the encoder positions.

Source code in PtyLab/Monitor/Monitor.py

def update_encoder(self, corrected_positions, original_positions, *args, **kwargs):
    """Update the image of the encoder positions."""
    pass

updateObjectProbeErrorMonitor(error, object_estimate, probe_estimate, zo=None, purity_probe=None, purity_object=None, encoder_positions=None)

Update the Object and Probe error monitor, and any associated metrics.

:param error: Value of the loss for the current iteration. :param object_estimate: Current object estimate :param probe_estimate: Current probe estimate :param zo: Sample-detector distance :param purity_probe: purity of the probe :param purity_object: purity of the object :return:

Source code in PtyLab/Monitor/Monitor.py

def updateObjectProbeErrorMonitor(
    self,
    error: float,
    object_estimate: np.ndarray,
    probe_estimate: np.ndarray,
    zo=None,
    purity_probe=None,
    purity_object=None,
    encoder_positions=None,
):
    """
    Update the Object and Probe error monitor, and any associated metrics.

    :param error: Value of the loss for the current iteration.
    :param object_estimate: Current object estimate
    :param probe_estimate: Current probe estimate
    :param zo: Sample-detector distance
    :param purity_probe: purity of the probe
    :param purity_object: purity of the object
    :return:
    """

    pass

updateBeamWidth(beamwidth_x, beamwidth_y)

Update the beam width in X and Y

Parameters

beamwidth_x: beamwidth x in meter beamwidth_y: beamwidth y in meter

Returns

Source code in PtyLab/Monitor/Monitor.py

def updateBeamWidth(self, beamwidth_x, beamwidth_y):
    """
    Update the beam width in X and Y

    Parameters
    ----------
    beamwidth_x: beamwidth x in meter
    beamwidth_y: beamwidth y in meter

    Returns
    -------

    """
    pass

updateDiffractionDataMonitor(Iestimated, Imeasured)

Update the diffraction data estimate at the current iteration.

:param Iestimated: Estimated intensity at the current position :param Imeasured: Measured intensity at the current position :return:

Source code in PtyLab/Monitor/Monitor.py

def updateDiffractionDataMonitor(self, Iestimated, Imeasured):
    """
    Update the diffraction data estimate at the current iteration.

    :param Iestimated: Estimated intensity at the current position
    :param Imeasured:  Measured intensity at the current position
    :return:
    """

    pass

Monitor

Bases: AbstractMonitor

Monitor contains two submonitors: ObjectProbeErrorPlot (object,probe,error) and DiffractionDataPlot (diffraction intensity estimate and measurement)

Source code in PtyLab/Monitor/Monitor.py

class Monitor(AbstractMonitor):
    """
    Monitor contains two submonitors: ObjectProbeErrorPlot (object,probe,error) and DiffractionDataPlot (diffraction
    intensity estimate and measurement)
    """

    def __init__(self):
        # settings for visualization
        self._figureUpdateFrequency = 1
        self.objectPlot = "complex"
        self._verboseLevel = "low"
        self.objectZoom = 1
        self.probeZoom = 1
        self.objectPlotContrast = 1
        self.probePlotContrast = 1
        self.reconstruction = None
        self.cmapDiffraction = setColorMap()
        self.defaultMonitor = None
        self.screenshot_directory = None
        self.diffractionDataMonitor = None

    @property
    def figureUpdateFrequency(self):
        return self._figureUpdateFrequency

    @figureUpdateFrequency.setter
    def figureUpdateFrequency(self, value):
        self._figureUpdateFrequency = value
        if is_inline() and self.figureUpdateFrequency < 5:
            warnings.simplefilter("always", UserWarning)
            warnings.warn(
                "For faster update of the reconstruction plot, set `monitor.figureUpdateFrequency = 5` or higher."
            )

    @property
    def verboseLevel(self):
        return self._verboseLevel

    @verboseLevel.setter
    def verboseLevel(self, value):
        self._verboseLevel = value
        if is_inline() and self._verboseLevel == "high":
            warnings.simplefilter("always", UserWarning)
            warnings.warn(
                "For diffraction data plot, preferably use an interactive matplotlib backend or"
                ' set `monitor.verboseLevel = "low"`. '
            )

    def initializeMonitors(self):
        """
        Create the figure and axes etc.
        :return:
        """
        # only initialize if it hasn't been done
        if self.defaultMonitor is None:
            self.defaultMonitor = ObjectProbeErrorPlot()

        if self.verboseLevel == "high":
            self.diffractionDataMonitor = DiffractionDataPlot()

    @property
    def objectPixelSize(self):
        """Pixel size of the object estimate that is plotted."""
        if self.reconstruction.data.operationMode == "FPM":
            return self.reconstruction.dxo_fpm
        return self.reconstruction.dxo

    @property
    def probePixelSize(self):
        """Axis step of the probe panel."""
        if self.reconstruction.data.operationMode == "FPM":
            return self.reconstruction.dfp
        return self.reconstruction.dxp

    @property
    def probeAxisUnit(self):
        """Unit of the probe panel axes: a length for CPM, a spatial frequency for FPM."""
        if self.reconstruction.data.operationMode == "FPM":
            return "1/um"
        return "mm"

    @property
    def probeLabel(self):
        """Title of the probe panel. FPM estimates a pupil instead of a probe."""
        if self.reconstruction.data.operationMode == "FPM":
            return "Pupil estimate"
        return "Probe estimate"

    def updateObjectProbeErrorMonitor(
        self,
        error,
        object_estimate,
        probe_estimate,
        zo=None,
        purity_probe=None,
        purity_object=None,
        encoder_positions=None,
    ):
        """
        update the probe object plots
        :param object_estimate:
        :return:
        """
        self.defaultMonitor.updateError(error)  # self.reconstruction.error)
        # print(f"Object plot: {self.objectPlot}")
        self.defaultMonitor.updateObject(
            object_estimate,
            self.reconstruction,
            objectPlot=self.objectPlot,
            pixelSize=self.objectPixelSize,
            axisUnit="mm",
            amplitudeScalingFactor=self.objectPlotContrast,
        )
        self.defaultMonitor.updateProbe(
            probe_estimate,
            self.reconstruction,
            pixelSize=self.probePixelSize,
            axisUnit=self.probeAxisUnit,
            label=self.probeLabel,
            amplitudeScalingFactor=self.probePlotContrast,
        )
        self.defaultMonitor.update_z(zo)
        self.defaultMonitor.drawNow()

        if self.screenshot_directory is not None:
            self.defaultMonitor.figure.savefig(
                Path(self.screenshot_directory) / f"frame_{len(error)}.png"
            )

    def describe_parameters(self, *args, **kwargs):
        pass

    def updateDiffractionDataMonitor(self, Iestimated, Imeasured):
        """
        update the diffraction plots
        """

        self.diffractionDataMonitor.update_view(
            Iestimated, Imeasured, cmap=self.cmapDiffraction
        )
        # self.diffractionDataMonitor.updateIestimated(Iestimated, cmap=self.cmapDiffraction)
        # self.diffractionDataMonitor.updateImeasured(Imeasured, cmap=self.cmapDiffraction)
        self.diffractionDataMonitor.drawNow()

objectPixelSize property

Pixel size of the object estimate that is plotted.

probePixelSize property

Axis step of the probe panel.

probeAxisUnit property

Unit of the probe panel axes: a length for CPM, a spatial frequency for FPM.

probeLabel property

Title of the probe panel. FPM estimates a pupil instead of a probe.

initializeMonitors()

Create the figure and axes etc. :return:

Source code in PtyLab/Monitor/Monitor.py

def initializeMonitors(self):
    """
    Create the figure and axes etc.
    :return:
    """
    # only initialize if it hasn't been done
    if self.defaultMonitor is None:
        self.defaultMonitor = ObjectProbeErrorPlot()

    if self.verboseLevel == "high":
        self.diffractionDataMonitor = DiffractionDataPlot()

updateObjectProbeErrorMonitor(error, object_estimate, probe_estimate, zo=None, purity_probe=None, purity_object=None, encoder_positions=None)

update the probe object plots :param object_estimate: :return:

Source code in PtyLab/Monitor/Monitor.py

def updateObjectProbeErrorMonitor(
    self,
    error,
    object_estimate,
    probe_estimate,
    zo=None,
    purity_probe=None,
    purity_object=None,
    encoder_positions=None,
):
    """
    update the probe object plots
    :param object_estimate:
    :return:
    """
    self.defaultMonitor.updateError(error)  # self.reconstruction.error)
    # print(f"Object plot: {self.objectPlot}")
    self.defaultMonitor.updateObject(
        object_estimate,
        self.reconstruction,
        objectPlot=self.objectPlot,
        pixelSize=self.objectPixelSize,
        axisUnit="mm",
        amplitudeScalingFactor=self.objectPlotContrast,
    )
    self.defaultMonitor.updateProbe(
        probe_estimate,
        self.reconstruction,
        pixelSize=self.probePixelSize,
        axisUnit=self.probeAxisUnit,
        label=self.probeLabel,
        amplitudeScalingFactor=self.probePlotContrast,
    )
    self.defaultMonitor.update_z(zo)
    self.defaultMonitor.drawNow()

    if self.screenshot_directory is not None:
        self.defaultMonitor.figure.savefig(
            Path(self.screenshot_directory) / f"frame_{len(error)}.png"
        )

updateDiffractionDataMonitor(Iestimated, Imeasured)

update the diffraction plots

Source code in PtyLab/Monitor/Monitor.py

def updateDiffractionDataMonitor(self, Iestimated, Imeasured):
    """
    update the diffraction plots
    """

    self.diffractionDataMonitor.update_view(
        Iestimated, Imeasured, cmap=self.cmapDiffraction
    )
    # self.diffractionDataMonitor.updateIestimated(Iestimated, cmap=self.cmapDiffraction)
    # self.diffractionDataMonitor.updateImeasured(Imeasured, cmap=self.cmapDiffraction)
    self.diffractionDataMonitor.drawNow()

DummyMonitor

Bases: object

Monitor without any visualisation so it won't consume any time

Source code in PtyLab/Monitor/Monitor.py

class DummyMonitor(object):
    """Monitor without any visualisation so it won't consume any time"""

    objectZoom = 1
    probeZoom = 1
    # remains from mPIE
    figureUpdateFrequency = 1000000
    verboseLevel = "low"

    def update_encoder(self, *args, **kwargs):
        pass

    def updateBeamWidth(self, *args, **kwargs):
        pass



    def updatePlot(self, object_estimate, probe_estimate):
        pass

    def getOverlap(self, ind1, ind2, probePixelsize):
        pass

    def initializeVisualisation(self):
        pass

    def initializeMonitors(self):
        pass

    def updateObjectProbeErrorMonitor(self, *args, **kwargs):
        pass

    def updateDiffractionDataMonitor(self, *args, **kwargs):
        pass

    def writeEngineName(self, *args, **kwargs):
        pass

NapariMonitor

Bases: DummyMonitor

Source code in PtyLab/Monitor/Monitor.py

class NapariMonitor(DummyMonitor):

    def initializeVisualisation(self):
        # currently a hacky way for this class, these napari implementations must
        # later be moved an optional sub-package.
        try:
            import napari

            self.viewer = napari.Viewer()
            self.viewer.show()
        except ImportError:
            msg = "Install napari to access this `NapariMonitor` implementation"
            raise ImportError(msg)

        self.viewer.add_image(name="object estimate", data=np.random.rand(100, 100))
        self.viewer.add_image(name="probe estimate", data=np.random.rand(100, 100))

        self.Iestimated = self.viewer.add_image(
            name="I estimated", data=np.random.rand(100, 100)
        )
        self.Imeasured = self.viewer.add_image(
            name="I measured", data=np.random.rand(100, 100)
        )
        self.rawdatamonitor = self.viewer.add_image(
            name="raw data (unset)", data=np.random.rand(100, 100), visible=False
        )

    def initializeMonitors(self):
        self.initializeVisualisation()

    def add_ptychogram(self, experimentalData):
        self.rawdatamonitor.name = "raw data"
        self.rawdatamonitor.data = experimentalData.ptychogram
        self.rawdatamonitor.visible = True

    def update_probe_image(self, new_probe):
        RGB_probe = complex2rgb(new_probe)
        self.viewer.layers["probe_estimate"].data = RGB_probe

    def update_object_image(self, object_estimate):
        RGB_object = complex2rgb(object_estimate)
        self.viewer.layers["object estimate"].data = RGB_object

    def updatePlot(self, object_estimate, probe_estimate):
        self.update_probe_image(probe_estimate)
        self.update_object_image(object_estimate)

    def updateDiffractionDataMonitor(self, Iestimated, Imeasured):
        self.Iestimated.data = Iestimated
        self.Imeasured.data = Imeasured

    def update_positions(self, *args, **kwargs):
        """
        Update the positions. USeful for position correction.
        :param args:
        :param kwargs:
        :return:
        """
        pass

update_positions(*args, **kwargs)

Update the positions. USeful for position correction. :param args: :param kwargs: :return:

Source code in PtyLab/Monitor/Monitor.py

def update_positions(self, *args, **kwargs):
    """
    Update the positions. USeful for position correction.
    :param args:
    :param kwargs:
    :return:
    """
    pass

Plots

ObjectProbeErrorPlot

Bases: object

Source code in PtyLab/Monitor/Plots.py

class ObjectProbeErrorPlot(object):
    def __init__(self, figNum=1):
        """Create a monitor.

        In principle, to use this method all you have to do is initialize the monitor and then call

        updateObject, updateProbe, updateErrorMetric and drawnow to ensure that something is drawn immediately.

        For example usage, see test_matplot_monitor.py.

        """
        self.figNum = figNum
        self._createFigure()

        # Get a reference to the figure canvas
        self.canvas = self.figure.canvas
        self.display_id = None

    def update_z(self, *args, **kwargs):
        """Update the sample-detector distance. Does nothing at the moment."""
        pass

    def _createFigure(self) -> None:
        """
        Create the figure.
        :return:
        """

        plt.ion()
        self.figure, axes = plt.subplot_mosaic(
            """Ape""",
            num=self.figNum,
            figsize=(10, 3),
            empty_sentinel=" ",
            constrained_layout=False,
        )

        self.ax_object = axes["A"]
        self.ax_probe = axes["p"]
        # self.ax_probe_ff = axes["P"]
        self.ax_error_metric = axes["e"]
        # self.ax_probe_ff.set_title("FF probe")
        self.ax_probe.set_title("Probe")
        # self.ax_object = axes[0][0]
        # self.ax_probe = axes[0][1]
        # self.ax_error_metric = axes[0][2]
        # self.ax_object.set_title
        self.txt_purityProbe = self.ax_probe.set_title("Probe estimate")
        self.txt_purityObject = self.ax_object.set_title("Object estimate")
        self.ax_error_metric.set_title("Error metric")
        self.ax_error_metric.grid(True)
        self.ax_error_metric.grid(
            animated=True, which="minor", color="#999999", linestyle="-", alpha=0.2
        )
        self.ax_error_metric.set_xlabel("iterations")
        self.ax_error_metric.set_ylabel("error")
        self.ax_error_metric.set_xscale("log")
        self.ax_error_metric.set_yscale("log")
        self.ax_error_metric.axis("image")
        self.figure.tight_layout()
        self.firstrun = True

    def updateObject(
        self,
        object_estimate,
        optimizable,
        objectPlot,
        amplitudeScalingFactor=1,
        **kwargs,
    ):
        OE = modeTile(object_estimate, normalize=True)
        if objectPlot == "complex":
            OE = complex2rgb(OE, amplitudeScalingFactor=amplitudeScalingFactor)

        elif objectPlot == "abs":
            # original
            # OE = OE / abs(OE).max()
            # better
            AOE = abs(OE)
            OE = OE / (AOE.mean() + np.std(AOE))
            # OE = OE / abs(OE.max())
            OE = abs(OE)
        elif objectPlot == "angle":
            OE = np.angle(OE)

        if self.firstrun:
            if objectPlot == "complex":
                self.im_object = complexPlot(OE, ax=self.ax_object, **kwargs)
            else:
                self.im_object = self.ax_object.imshow(OE, interpolation=None)
                divider = make_axes_locatable(self.ax_object)
                cax = divider.append_axes("right", size="5%", pad=0.1)
                self.objectCbar = plt.colorbar(
                    self.im_object, ax=self.ax_object, cax=cax
                )
        else:
            self.im_object.set_data(OE)
            if optimizable.nosm > 1:
                self.txt_purityObject.set_text(
                    "Object estimate\nPurity: %i" % (100 * optimizable.purityObject)
                    + "%"
                )

        self.im_object.autoscale()

    def updateProbe(
        self,
        probe_estimate,
        optimizable,
        amplitudeScalingFactor=1,
        label="Probe estimate",
        **kwargs,
    ):

        # from PtyLab.Operators.Operators import fft2c
        #
        # probe_estimate_ff = fft2c(probe_estimate)
        # PE_ff = complex2rgb(modeTile(probe_estimate_ff, normalize=True))

        PE = complex2rgb(
            modeTile(probe_estimate, normalize=True),
            amplitudeScalingFactor=amplitudeScalingFactor,
        )

        if self.firstrun:
            self.im_probe = complexPlot(PE, ax=self.ax_probe, **kwargs)
            self.txt_purityProbe = self.ax_probe.set_title(label)
            # self.im_probe_ff = complexPlot(PE_ff, self.ax_probe_ff, **kwargs)
        else:
            self.im_probe.set_data(PE)
            # self.im_probe_ff.set_data(PE_ff)
            if (
                optimizable.npsm > 1
                and optimizable.purityProbe == optimizable.purityProbe
            ):
                self.txt_purityProbe.set_text(
                    "%s\nPurity: %.2f" % (label, 100 * optimizable.purityProbe) + "%"
                )
        self.im_probe.autoscale()

    def updateError(self, error_estimate: np.ndarray) -> None:
        """
        Update the error estimate plot.
        :param error_estimate:
        :return:
        """

        if self.firstrun:
            self.error_metric_plot = self.ax_error_metric.plot(
                error_estimate, "o-", mfc="none"
            )[0]
        else:
            if len(error_estimate) > 1 and error_estimate[-1] == error_estimate[-1]:
                self.error_metric_plot.set_data(
                    np.arange(len(error_estimate)) + 1, error_estimate
                )
                self.ax_error_metric.set_xlim(1, len(error_estimate))
                self.ax_error_metric.set_ylim(
                    np.min(error_estimate), np.max(error_estimate)
                )
                data_aspect = np.log(
                    np.max(error_estimate) / np.min(error_estimate)
                ) / np.log(len(error_estimate))
                self.ax_error_metric.set_aspect(1 / data_aspect)
                self.ax_error_metric.set_title(
                    f"Error metric (it {len(error_estimate)})"
                )

    def drawNowScript(self):
        """
        Forces the image to be drawn
        :return:
        """
        if self.firstrun:
            self.figure.show()
            self.firstrun = False

        # Reopen the figure if the window is closed
        if not plt.fignum_exists(self.figNum):
            self.figure.show()

        self.canvas.draw_idle()
        self.canvas.flush_events()

    def drawNowIpython(self):
        if self.firstrun:
            self.display_id = display(self.figure, display_id=True)
            self.firstrun = False
        else:
            clear_output(wait=True)
            self.display_id = display(
                self.figure, display_id=self.display_id.display_id
            )
        self.canvas.draw_idle()
        self.canvas.flush_events()

    def drawNow(self):
        if is_inline():
            self.drawNowIpython()
        else:
            self.drawNowScript()

__init__(figNum=1)

Create a monitor.

In principle, to use this method all you have to do is initialize the monitor and then call

updateObject, updateProbe, updateErrorMetric and drawnow to ensure that something is drawn immediately.

For example usage, see test_matplot_monitor.py.

Source code in PtyLab/Monitor/Plots.py

def __init__(self, figNum=1):
    """Create a monitor.

    In principle, to use this method all you have to do is initialize the monitor and then call

    updateObject, updateProbe, updateErrorMetric and drawnow to ensure that something is drawn immediately.

    For example usage, see test_matplot_monitor.py.

    """
    self.figNum = figNum
    self._createFigure()

    # Get a reference to the figure canvas
    self.canvas = self.figure.canvas
    self.display_id = None

update_z(*args, **kwargs)

Update the sample-detector distance. Does nothing at the moment.

Source code in PtyLab/Monitor/Plots.py

def update_z(self, *args, **kwargs):
    """Update the sample-detector distance. Does nothing at the moment."""
    pass

updateError(error_estimate)

Update the error estimate plot. :param error_estimate: :return:

Source code in PtyLab/Monitor/Plots.py

def updateError(self, error_estimate: np.ndarray) -> None:
    """
    Update the error estimate plot.
    :param error_estimate:
    :return:
    """

    if self.firstrun:
        self.error_metric_plot = self.ax_error_metric.plot(
            error_estimate, "o-", mfc="none"
        )[0]
    else:
        if len(error_estimate) > 1 and error_estimate[-1] == error_estimate[-1]:
            self.error_metric_plot.set_data(
                np.arange(len(error_estimate)) + 1, error_estimate
            )
            self.ax_error_metric.set_xlim(1, len(error_estimate))
            self.ax_error_metric.set_ylim(
                np.min(error_estimate), np.max(error_estimate)
            )
            data_aspect = np.log(
                np.max(error_estimate) / np.min(error_estimate)
            ) / np.log(len(error_estimate))
            self.ax_error_metric.set_aspect(1 / data_aspect)
            self.ax_error_metric.set_title(
                f"Error metric (it {len(error_estimate)})"
            )

drawNowScript()

Forces the image to be drawn :return:

Source code in PtyLab/Monitor/Plots.py

def drawNowScript(self):
    """
    Forces the image to be drawn
    :return:
    """
    if self.firstrun:
        self.figure.show()
        self.firstrun = False

    # Reopen the figure if the window is closed
    if not plt.fignum_exists(self.figNum):
        self.figure.show()

    self.canvas.draw_idle()
    self.canvas.flush_events()

DiffractionDataPlot

Bases: object

Source code in PtyLab/Monitor/Plots.py

class DiffractionDataPlot(object):
    def __init__(self, figNum=2):
        """Create a monitor.

        In principle, to use this method all you have to do is initialize the monitor and then call

        updateImeasured, updateIestimated and drawnow to ensure that something is drawn immediately.

        For example usage, see test_matplot_monitor.py.

        """
        self.figNum = figNum
        self._createFigure()

        # Get a reference to the figure canvas
        self.canvas = self.figure.canvas
        self.display_id = None  # Added attribute

    def _createFigure(self) -> None:
        """
        Create the figure.
        :return:
        """

        # add an axis for the object
        plt.ion()
        self.figure, axes = plt.subplots(
            1, 2, num=self.figNum, squeeze=False, clear=True, figsize=(8, 3)
        )
        self.ax_Iestimated = axes[0][0]
        self.ax_Imeasured = axes[0][1]
        self.ax_Iestimated.set_title("Estimated intensity")
        self.ax_Imeasured.set_title("Measured intensity")
        self.figure.tight_layout()
        self.firstrun = True

    def updateIestimated(self, Iestimate, cmap="gray", **kwargs):
        # move it to CPU if it's on the GPU
        Iestimate = gpuUtils.asNumpyArray(Iestimate)

        if self.firstrun:

            self.im_Iestimated: AxesImage = self.ax_Iestimated.imshow(
                np.log10(np.squeeze(Iestimate + 1)), cmap=cmap, interpolation=None
            )

            divider = make_axes_locatable(self.ax_Iestimated)
            cax = divider.append_axes("right", size="5%", pad=0.1)
            self.IestimatedCbar = plt.colorbar(
                self.im_Iestimated, ax=self.ax_Iestimated, cax=cax
            )
            # scale it according to I measured

        else:
            self.im_Iestimated.set_data(np.log10(np.squeeze(Iestimate + 1)))
        # self.im_Iestimated.autoscale()
        # self.im_Iestimated.set_

    def updateImeasured(self, Imeasured, cmap="gray", **kwargs):
        Imeasured = gpuUtils.asNumpyArray(Imeasured)
        if self.firstrun:
            self.im_Imeasured: AxesImage = self.ax_Imeasured.imshow(
                np.log10(np.squeeze(Imeasured + 1)), cmap=cmap, interpolation=None
            )

            divider = make_axes_locatable(self.ax_Imeasured)
            cax = divider.append_axes("right", size="5%", pad=0.1)
            self.ImeasuredCbar = plt.colorbar(
                self.im_Imeasured, ax=self.ax_Imeasured, cax=cax
            )

        else:
            self.im_Imeasured.set_data(np.log10(np.squeeze(Imeasured + 1)))
        self.im_Imeasured.autoscale()

    def drawNowScript(self):
        """
        Forces the image to be drawn
        :return:
        """
        if self.firstrun:
            self.figure.show()
            self.firstrun = False

        # Reopen the figure if the window is closed
        if not plt.fignum_exists(self.figNum):
            self.figure.show()

        self.canvas.draw_idle()
        self.canvas.flush_events()

    def drawNowIpython(self):
        if self.firstrun:
            self.display_id = display(self.figure, display_id=True)
            self.firstrun = False
        else:
            clear_output(wait=True)
            self.display_id = display(
                self.figure, display_id=self.display_id.display_id
            )
        self.canvas.draw_idle()
        self.canvas.flush_events()

    def drawNow(self):
        if is_inline():
            self.drawNowIpython()
        else:
            self.drawNowScript()

    def update_view(self, Iestimated, Imeasured, cmap):
        """Update the I measured and I estimated and make sure that the colormaps have the same limits"""
        self.updateImeasured(Imeasured, cmap=cmap)
        self.updateIestimated(Iestimated, cmap=cmap)
        self._equalize_contrast()

    def _equalize_contrast(self):
        """Adopt the contrast limits from the measured data and apply them to the predicted"""
        clims = self.im_Imeasured.get_clim()
        self.im_Iestimated.set_clim(*clims)

__init__(figNum=2)

Create a monitor.

In principle, to use this method all you have to do is initialize the monitor and then call

updateImeasured, updateIestimated and drawnow to ensure that something is drawn immediately.

For example usage, see test_matplot_monitor.py.

Source code in PtyLab/Monitor/Plots.py

def __init__(self, figNum=2):
    """Create a monitor.

    In principle, to use this method all you have to do is initialize the monitor and then call

    updateImeasured, updateIestimated and drawnow to ensure that something is drawn immediately.

    For example usage, see test_matplot_monitor.py.

    """
    self.figNum = figNum
    self._createFigure()

    # Get a reference to the figure canvas
    self.canvas = self.figure.canvas
    self.display_id = None  # Added attribute

drawNowScript()

Forces the image to be drawn :return:

Source code in PtyLab/Monitor/Plots.py

def drawNowScript(self):
    """
    Forces the image to be drawn
    :return:
    """
    if self.firstrun:
        self.figure.show()
        self.firstrun = False

    # Reopen the figure if the window is closed
    if not plt.fignum_exists(self.figNum):
        self.figure.show()

    self.canvas.draw_idle()
    self.canvas.flush_events()

update_view(Iestimated, Imeasured, cmap)

Update the I measured and I estimated and make sure that the colormaps have the same limits

Source code in PtyLab/Monitor/Plots.py

def update_view(self, Iestimated, Imeasured, cmap):
    """Update the I measured and I estimated and make sure that the colormaps have the same limits"""
    self.updateImeasured(Imeasured, cmap=cmap)
    self.updateIestimated(Iestimated, cmap=cmap)
    self._equalize_contrast()

is_inline()

Default IPython (jupyter notebook) backend

Source code in PtyLab/Monitor/Plots.py

def is_inline():
    """Default IPython (jupyter notebook) backend"""
    return True if "inline" in mpl.get_backend().lower() else False

TensorboardMonitor

TensorboardMonitor

Bases: AbstractMonitor

Source code in PtyLab/Monitor/TensorboardMonitor.py

class TensorboardMonitor(AbstractMonitor):

    # maximum number of probe state mixtures that we want to show
    max_npsm = 10
    # maximum number of object state mixtures we want to show
    max_nosm = 10
    # remove any phase slant from the object
    center_angle_object = False
    # Turn on to show the FFT of the object. Usually not useful.
    show_farfield_object = False

    # the probe is shown as an inset in the object. This specifies how much to downsample it
    probe_downsampling = 2

    # downsample all images by the same amount to make it run a bit faster
    downsample_everything = 1

    def __init__(self, logdir="./logs_tensorboard", name=None):
        super(AbstractMonitor).__init__()
        # if true, all phases are centered in such a way that the average phase in the center of any RGB plot is zero.
        self.center_phases = True
        if name is None:
            starttime = time.strftime("%H%M")
            name = f"Start {starttime}"
            print("Name", name)
        path = Path(logdir) / name
        self.writer: tfs.SummaryWriter = tfs.create_file_writer(
            logdir=str(path), name=name
        )
        self.i = 0

    # These are the codes that have to be implemented

    def updatePlot(
        self,
        object_estimate,
        probe_estimate,
        highres=True,
        zo=None,
        encoder_positions=None,
    ):
        self.i += 1

        if not highres:
            Np = probe_estimate.shape[-1]
            No = object_estimate.shape[-1]
            xmax, ymax = np.clip(encoder_positions.max(axis=0) + 2 * Np // 3, 0, No)
            xmin, ymin = np.clip(encoder_positions.min(axis=0) + Np // 3, 0, No)

            probe_estimate = probe_estimate[..., ::4, ::4]
            object_estimate = object_estimate[..., ymin:ymax, xmin:xmax]
        probe_estimate_rgb = self._update_probe_estimate(
            probe_estimate, highres=highres
        )
        self._update_object_estimate(
            object_estimate, probe_estimate_rgb, highres=highres, zo=zo
        )

    def visualize_probe_engine(self, engine):
        RGB_image = complex2rgb_vectorized(engine.get_fundamental(), center_phase=self.center_phases)
        self.__safe_upload_image("original probe", np.squeeze(RGB_image), self.i)
        pass

    def updateObjectProbeErrorMonitor(
        self,
        error,
        object_estimate,
        probe_estimate,
        zo,
        purity_object=None,
        purity_probe=None,
        encoder_positions=None,
        *args,
        **kwargs,
    ):
        object_estimate = asNumpyArray(object_estimate)
        probe_estimate = asNumpyArray(probe_estimate)
        # The input can be either an empty array, an array with length 1 or a list of all the errors so far.
        # In the last case, we only want the last value.
        # or just a number. This part should account of all of them.
        if encoder_positions is None:
            raise ValueError("Please submit encoder positions or the code won't work.")
        self._update_error_estimate(error)

        self.updatePlot(
            object_estimate,
            probe_estimate,
            highres=self.i % 5 == 0,
            encoder_positions=encoder_positions,
            zo=zo,
        )

        self.update_z(zo)
        self.update_purities(asNumpyArray(purity_probe), asNumpyArray(purity_object))

    def updateDiffractionDataMonitor(self, Iestimated, Imeasured):

        Itotal = np.hstack([Iestimated, Imeasured])
        Itotal = Itotal[None, ..., None]
        Itotal = Itotal / Itotal.max() * 255
        Itotal = Itotal.astype(np.uint8)
        self.__safe_upload_image(
            "Estimated and measured intensity example", Itotal, self.i, 1
        )

    def writeEngineName(self, name, *args, **kwargs):
        with self.writer.as_default():
            tfs.text(
                "reconstructor type",
                name,
                step=self.i,
                description="reconstruction name",
            )

    def update_focusing_metric(self, TV_value, AOI_image, metric_name, allmerits=None):
        if TV_value is not None:
            self.__safe_upload_scalar(
                f"Autofocus {metric_name}",
                TV_value,
                self.i,
                "Total Variation of the object",
            )
        if AOI_image is not None:
            # print(AOI_image)

            self.__smart_upload_image_couldbecomplex(
                f"Autofocus {metric_name} AOI",
                AOI_image,
                self.i,
                1,
                "AOI used by autofocus",
                center_phase=True,
            )
        if allmerits is not None:
            import matplotlib.pyplot as plt
            allmerits, new_z = allmerits
            fig, ax = plt.subplot_mosaic('A')
            ax = ax['A']
            ax.plot(allmerits[0], allmerits[1], '-ro')
            ax.vlines(new_z, *ax.get_ylim())
            ax.set_title(f'{metric_name}')
            buf = io.BytesIO()
            fig.savefig(buf, format='png', dpi=70)
            plt.close(fig)
            buf.seek(0)
            with self.writer.as_default():
                img = image.decode_png(buf.getvalue(), channels=4)
                img = np.expand_dims(img, 0)
                tfs.image('Autofocus dz score', img, self.i)

    def updateBeamWidth(self, beamwidth_y, beamwidth_x):
        self.__safe_upload_scalar('beamwidth/x_um', beamwidth_x*1e6, step=self.i)
        self.__safe_upload_scalar('beamwidth/y_um', beamwidth_y*1e6, step=self.i)

    def update_overlap(self, overlap_area, linear_overlap):
        self.__safe_upload_scalar('overlap/area', overlap_area, step=self.i)
        self.__safe_upload_scalar('overlap/linear', linear_overlap, step=self.i)


    def update_encoder(
        self,
        corrected_positions: np.ndarray,
        original_positions: np.ndarray,
        scaling: float = 1.0,
            beamwidth=None
    ) -> None:
        """
        Update the stage position images.
        :param corrected_positions:
        :param original_positions:
        :param scaling:
        :return:
        """
        # convert the positions to mm
        corrected_positions = corrected_positions
        original_positions = original_positions
        # set them to mean 0
        corrected_positions = corrected_positions - corrected_positions.mean(
            axis=0, keepdims=True
        )
        original_positions = original_positions - original_positions.mean(
            axis=0, keepdims=True
        )



        matplotlib.use("Agg")  # no images output
        import matplotlib.pyplot as plt

        # make a fov that makes sense
        scale_0 = 1.1
        if scaling > scale_0:
            scale_0 = scaling

        position_range = np.min(original_positions.flatten()), np.max(
            original_positions.flatten()
        )
        diff = np.diff(position_range)
        mean = np.mean(position_range)
        position_range = mean - scale_0 * diff / 2, mean + scale_0 * diff / 2
        fig, axes = plt.subplot_mosaic(
            """
        ONS""",
            constrained_layout=True,
            figsize=(15, 8),
        )

        meandiff = np.mean(
            abs(1e6 * corrected_positions - 1e6 * original_positions)
        )

        self.__safe_upload_scalar(
            "mean position displacement in micron",
            meandiff,
            self.i,
            "mean absolute displacement",
        )

        axes["O"].set_title("Original positions")
        axes["N"].set_title("Updated positions")
        axes["S"].set_title(f"Diff. Mean: {meandiff} $\mu$m")

        # plot the original one everywhere

        for name, ax in axes.items():
            ax: plt.Axes = ax
            ax.scatter(
                original_positions[:, 0],
                original_positions[:, 1],
                color="C0",
                marker=".",
                label="original",
            )
            ax.set_xlabel("X")
            ax.set_ylabel("Y")
            ax.set_aspect(1)
            ax.xaxis.set_major_formatter(EngFormatter(unit="m"))
            ax.yaxis.set_major_formatter(EngFormatter(unit="m"))
            ax.set_xlim(ax.set_ylim(position_range))

        # plot the new one in the middle image
        axes["N"].scatter(
            corrected_positions[:, 0],
            corrected_positions[:, 1],
            color="C1",
            marker="x",
            label="new",
        )
        # scaled version on the right (should only show displacement, not magnification)
        diff = corrected_positions - original_positions
        axes["S"].quiver(
            original_positions[:, 0],
            original_positions[:, 1],
            diff[:, 0],
            diff[:, 1],
            angles="xy",
            units="xy",
            scale=1,
        )

        buf = io.BytesIO()
        fig.savefig(buf, format="png", dpi=70)
        plt.close(fig)
        buf.seek(0)
        with self.writer.as_default():
            img = image.decode_png(buf.getvalue(), channels=4)
            img = np.expand_dims(img, 0)
            tfs.image("position correction", img, self.i)

    # internal use for tensorboardMonitor
    def _update_object_estimate(
        self,
        object_estimate,
        probe_estimate_rgb,
        highres=True,
        zo=None,
    ):
        """
        Update the object estimate. This ensures that within the web interface the object and the probe estimate are available.


        :param object_estimate:
        :param probe_estimate_rgb:
        :param highres:
        :param z: Object-detector distance. Not required.
        :return:
        """
        if self.center_angle_object:
            object_estimate, shift1 = center_angle(object_estimate)
            object_estimate, shift2 = center_angle(object_estimate)
            print("Angle shifts: ", shift1, shift2)

        # convert the object estimate to colour
        object_estimate_rgb = complex2rgb_vectorized(object_estimate, center_phase=self.center_phases
                                                     )

        # ensure it's 4 d as that's what is needed by tensorflow
        if object_estimate_rgb.ndim == 3:
            object_estimate_rgb = object_estimate_rgb[None]

        # patch in an image of the first probe in the top right
        object_estimate_rgb = self.__pad_probe_in_object_estimate(
            probe_estimate_rgb, object_estimate_rgb
        )

        # Add an object, but if it's the low-resolution version add it to the (low_res) version.
        tag = "object estimate"
        if not highres:
            tag += "(low res)"
        if self.show_farfield_object:
            I_ff = asNumpyArray(
                abs(
                    fft2c(
                        object_estimate
                        - object_estimate.mean(axis=(-2, -1), keepdims=True)
                    )
                )
                ** 2
            ) ** (0.2)
            I_ff = I_ff / I_ff.max() * 255
            I_ff = np.clip(I_ff, 0, 255).astype(np.uint8)
            self.__safe_upload_image("I ff estimate", I_ff, self.i, self.max_nosm)

        self.__safe_upload_image(tag, object_estimate_rgb, self.i, self.max_npsm)

        if highres:
            I_object = abs(object_estimate**2)

            std_obj = I_object.std()
            mean_obj = I_object.mean()
            min_int = 0  # mean_obj - 2 * std_obj
            # max_int = np.min((mean_obj + 2 * std_obj, I_object.max()))
            N = I_object.shape[-1]
            roi = slice(N // 2 - N // 5, N // 2 + N // 5)
            # max_int = I_object[...,roi,roi].max()
            from scipy import ndimage

            max_int = ndimage.gaussian_filter(I_object, 3)[..., roi, roi].max()

            if min_int == max_int:
                max_int += 1

            I_object = (I_object - min_int) / (
                max_int - min_int
            )  # +1 to ensure that this always works
            logI = np.log(255 * I_object.astype(float) + 1)
            logI -= logI.min()
            logI /= logI.max() / 255
            I_object = np.clip(I_object * 255, 0, 255).astype(np.uint8)
            self.__safe_upload_image(
                "I object estimate", I_object, self.i, self.max_nosm
            )

            self.__safe_upload_image(
                "I object log estimate", logI.astype(np.uint8), self.i, self.max_nosm
            )
            self.save_intensities = False
            if self.save_intensities:
                p = Path("./intensities")
                p.mkdir(exist_ok=True)

                import matplotlib.pyplot as plt

                plt.clf()
                im = plt.imshow(I_object)
                plt.colorbar(im)
                if zo is not None:
                    plt.title(f"z = {zo*1e3:.3f} mm")
                plt.savefig(f"intensities/{self.i}.png")
                plt.savefig(f"intensities/AAA.png")

    def _update_probe_estimate(self, probe_estimate, highres=True):
        # first, convert it to images
        # while probe_estimate.ndim <= 3:
        #     probe_estimate = probe_estimate[None]
        probe_estimate_rgb = complex2rgb_vectorized(probe_estimate, center_phase=self.center_phases)
        # ensure it's 4 d as that's what is needed by tensorflow
        tag = "probe estimate"
        if not highres:
            tag += "(low res)"
        self.__safe_upload_image(tag, probe_estimate_rgb, self.i, self.max_npsm)

        if highres:
            ff_probe = fft2c(probe_estimate)
            ff_probe = complex2rgb_vectorized(ff_probe, center_phase=self.center_phases)
            self.__safe_upload_image("FF " + tag, ff_probe, self.i, self.max_npsm)

        # make a probe COM estimate
        from scipy import ndimage
        P = probe_estimate
        while P.ndim > 2:
            P = P[0]
        cy, cx = ndimage.center_of_mass(abs(P**2))
        N = probe_estimate.shape[-1]
        self.__safe_upload_scalar('com/cy', cy-N//2, self.i)
        self.__safe_upload_scalar('com/cx', cx-N//2, self.i)
        return probe_estimate_rgb



    def __smart_upload_image_couldbecomplex(
        self, name, data, step, max_outputs=3, description=None, center_phase=False,
    ):
        """
        Safely upload an image that could be complex. If it is, cast it to colour before uploading.

        """
        data = asNumpyArray(data)


        if np.iscomplexobj(data):
            if center_phase:
                phexp = data.sum((-2,-1), keepdims=True)
                phexp = phexp.conj() / (abs(phexp) + 1e-9)
            else:
                phexp = 1
            print("Got complex datatype")
            data = complex2rgb_vectorized(data*phexp)
        else:
            print("Got real datatype")
            # auto scale
            data = data / data.max() * 255
            data = data.astype(np.uint8)
            # convert to black-white

        self.__safe_upload_image(name, data, step, max_outputs, description)

    def __safe_upload_image(self, name, data, step, max_outputs=3, description=None):
        data = asNumpyArray(data)
        if data.shape[-1] not in [1, 3]:
            data = data[..., None]
        while data.ndim < 4:
            data = data[None]
        with self.writer.as_default():
            tfs.image(
                name,
                data[
                    ..., :: self.downsample_everything, :: self.downsample_everything, :
                ],
                step,
                max_outputs=max_outputs,
                description=description,
            )

    def __safe_upload_scalar(self, name, data, step, description=None):
        if isinstance(data, list):
            if data == []:
                return  # initialization, not required for tensorboard, ignore it
        data = asNumpyArray(data)
        try:
            # only take the last one in case of a list
            data = np.array(data).ravel()[-1]
        except:
            data = float(data)

        with self.writer.as_default():
            tfs.scalar(name, data, step, description)

    def update_z(self, z):
        self.__safe_upload_scalar(
            "zo (mm)", 1e3 * z, step=self.i, description="Propagation distance"
        )

    def _update_error_estimate(self, error):
        self.__safe_upload_scalar(
            "error metric", error, self.i, "Error metric (single image)"
        )

    def _update_probe_purity(self, probe_purity):
        if probe_purity is None:
            return
        self.__safe_upload_scalar("probe purity", probe_purity, self.i, "probe purity")

    def _update_object_purity(self, object_purity):

        if object_purity is None:
            return
        self.__safe_upload_scalar(
            "object purity", object_purity, self.i, "object purity"
        )

    def update_purities(self, probe_purity, object_purity):
        self._update_object_purity(object_purity)
        self._update_probe_purity(probe_purity)

    def describe_parameters(self, params: Params):
        text = "\n".join(["%s: %s" % (k, d) for (k, d) in params.__dict__.items()])
        with self.writer.as_default():
            tfs.text(
                "summary parameters", text, step=self.i, description="initial settings"
            )

    def __pad_probe_in_object_estimate(self, probe_estimate_rgb, object_estimate_rgb):
        probe_estimate_rgb_ss = probe_estimate_rgb[
            ..., :: self.probe_downsampling, :: self.probe_downsampling, :
        ]

        self.__safe_upload_scalar(
            "probe downsampling in inset",
            self.probe_downsampling,
            self.i,
            "probe downsampling in the inset images",
        )
        Ny, Nx, _ = probe_estimate_rgb_ss.shape[-3:]
        if object_estimate_rgb.shape[-2] < Nx:
            raise RuntimeError(
                "The downsampled probe size would be larger than the downsampled object size."
                ""
                "Try setting monitor.probe_downsampling higher or monitor.object_downsampling lower."
            )
        # add a red marker around the edge
        probe_estimate_rgb_ss[..., -1, -1] = 255
        probe_estimate_rgb_ss[..., -1, :, -1] = 255

        # last channel is for color
        # print(Ny, Nx, object_estimate_rgb.shape, probe_estimate_rgb_ss.shape)
        object_estimate_rgb[..., :Ny, :Nx, :] = probe_estimate_rgb_ss[0]

        return object_estimate_rgb

update_encoder(corrected_positions, original_positions, scaling=1.0, beamwidth=None)

Update the stage position images. :param corrected_positions: :param original_positions: :param scaling: :return:

Source code in PtyLab/Monitor/TensorboardMonitor.py

def update_encoder(
    self,
    corrected_positions: np.ndarray,
    original_positions: np.ndarray,
    scaling: float = 1.0,
        beamwidth=None
) -> None:
    """
    Update the stage position images.
    :param corrected_positions:
    :param original_positions:
    :param scaling:
    :return:
    """
    # convert the positions to mm
    corrected_positions = corrected_positions
    original_positions = original_positions
    # set them to mean 0
    corrected_positions = corrected_positions - corrected_positions.mean(
        axis=0, keepdims=True
    )
    original_positions = original_positions - original_positions.mean(
        axis=0, keepdims=True
    )



    matplotlib.use("Agg")  # no images output
    import matplotlib.pyplot as plt

    # make a fov that makes sense
    scale_0 = 1.1
    if scaling > scale_0:
        scale_0 = scaling

    position_range = np.min(original_positions.flatten()), np.max(
        original_positions.flatten()
    )
    diff = np.diff(position_range)
    mean = np.mean(position_range)
    position_range = mean - scale_0 * diff / 2, mean + scale_0 * diff / 2
    fig, axes = plt.subplot_mosaic(
        """
    ONS""",
        constrained_layout=True,
        figsize=(15, 8),
    )

    meandiff = np.mean(
        abs(1e6 * corrected_positions - 1e6 * original_positions)
    )

    self.__safe_upload_scalar(
        "mean position displacement in micron",
        meandiff,
        self.i,
        "mean absolute displacement",
    )

    axes["O"].set_title("Original positions")
    axes["N"].set_title("Updated positions")
    axes["S"].set_title(f"Diff. Mean: {meandiff} $\mu$m")

    # plot the original one everywhere

    for name, ax in axes.items():
        ax: plt.Axes = ax
        ax.scatter(
            original_positions[:, 0],
            original_positions[:, 1],
            color="C0",
            marker=".",
            label="original",
        )
        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.set_aspect(1)
        ax.xaxis.set_major_formatter(EngFormatter(unit="m"))
        ax.yaxis.set_major_formatter(EngFormatter(unit="m"))
        ax.set_xlim(ax.set_ylim(position_range))

    # plot the new one in the middle image
    axes["N"].scatter(
        corrected_positions[:, 0],
        corrected_positions[:, 1],
        color="C1",
        marker="x",
        label="new",
    )
    # scaled version on the right (should only show displacement, not magnification)
    diff = corrected_positions - original_positions
    axes["S"].quiver(
        original_positions[:, 0],
        original_positions[:, 1],
        diff[:, 0],
        diff[:, 1],
        angles="xy",
        units="xy",
        scale=1,
    )

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=70)
    plt.close(fig)
    buf.seek(0)
    with self.writer.as_default():
        img = image.decode_png(buf.getvalue(), channels=4)
        img = np.expand_dims(img, 0)
        tfs.image("position correction", img, self.i)

__smart_upload_image_couldbecomplex(name, data, step, max_outputs=3, description=None, center_phase=False)

Safely upload an image that could be complex. If it is, cast it to colour before uploading.

Source code in PtyLab/Monitor/TensorboardMonitor.py

def __smart_upload_image_couldbecomplex(
    self, name, data, step, max_outputs=3, description=None, center_phase=False,
):
    """
    Safely upload an image that could be complex. If it is, cast it to colour before uploading.

    """
    data = asNumpyArray(data)


    if np.iscomplexobj(data):
        if center_phase:
            phexp = data.sum((-2,-1), keepdims=True)
            phexp = phexp.conj() / (abs(phexp) + 1e-9)
        else:
            phexp = 1
        print("Got complex datatype")
        data = complex2rgb_vectorized(data*phexp)
    else:
        print("Got real datatype")
        # auto scale
        data = data / data.max() * 255
        data = data.astype(np.uint8)
        # convert to black-white

    self.__safe_upload_image(name, data, step, max_outputs, description)