Add the original positions to the fancy_ptycho position plotting

src/cdtools/models/fancy_ptycho.py

@@ -798,6 +798,28 @@ class FancyPtycho(CDIModel):
             **kwargs),
 
         
+    def plot_translations_and_originals(self, fig, dataset):
+        """Only used to make a plot for the plot list."""
+        p.plot_translations(
+            dataset.translations,
+            fig=fig,
+            units=self.units,
+            label='original translations',
+            color='#CCCCCC',
+            marker='o',
+        )
+        p.plot_translations(
+            self.corrected_translations(dataset),
+            fig=fig,
+            units=self.units,
+            clear_fig=False,
+            label='refined translations',
+            color='k',
+            marker='.'
+        )
+        plt.legend()
+        
+        
     plot_list = [
         ('',
          lambda self, fig, dataset: self.plot_wavefront_variation(
@@ -895,7 +917,7 @@ class FancyPtycho(CDIModel):
          lambda self: self.exponentiate_obj),
 
         ('Corrected Translations',
-         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
+         lambda self, fig, dataset: self.plot_translations_and_originals(fig, dataset)),
         ('Background',
          lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)),
         ('Quantum Efficiency Mask',

src/cdtools/tools/plotting/plotting.py

@@ -522,7 +522,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
                       units=units, show_cbar=False, **kwargs)
 
 
-def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
+def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, clear_fig=True, label=None, color=None, marker='.', **kwargs):
     """Plots a set of probe translations in a nicely formatted way
 
     Parameters
@@ -537,6 +537,14 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
         Whether to plot lines indicating the path taken
     invert_xaxis : bool
         Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
+    clear_fig : bool
+        Default is True. Whether to clear the figure before plotting.
+    label : str
+        Default is None. A label to give the plotted markers for a legend.
+    color : str
+        Default is None. The color to plot the markers in. By default, will follow the matplotlib color cycle.
+    color : str
+        Default is '.'. The marker style to plot with.
     \\**kwargs
         All other args are passed to fig.add_subplot(111, \\**kwargs)
 
@@ -554,18 +562,24 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
         ax = fig.add_subplot(111, **kwargs)
     else:
         plt.figure(fig.number)
-        plt.gcf().clear()
+        if clear_fig:
+            plt.gcf().clear()
 
     if isinstance(translations, t.Tensor):
         translations = translations.detach().cpu().numpy()
 
     translations = translations * factor
-    plt.plot(translations[:,0], translations[:,1],'k.')
+
+    linestyle = '-' if lines else 'None'
+    linewidth = 1 if lines else 0
+    plt.plot(translations[:,0], translations[:,1],
+             marker=marker, linestyle=linestyle,
+             label=label, color=color, 
+             linewidth=linewidth)
+    
     if invert_xaxis:
         plt.gca().invert_xaxis()
         
-    if lines:
-        plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
     plt.xlabel('X (' + units + ')')
     plt.ylabel('Y (' + units + ')')