From 3eec9c0cec31775ec3d5c2ea9a04e7af47b5def5 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Sun, 19 Oct 2025 12:45:19 +0200
Subject: [PATCH 1/2] Add the original positions to the fancy_ptycho position
 plotting

---
 src/cdtools/models/fancy_ptycho.py     | 24 +++++++++++++++++++++++-
 src/cdtools/tools/plotting/plotting.py | 24 +++++++++++++++++++-----
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 6220299a..fa4757f6 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -798,6 +798,28 @@ def get_probes(idx):
             **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 @@ def get_probes(idx):
          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',
diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py
index abd8f841..f2357f78 100644
--- a/src/cdtools/tools/plotting/plotting.py
+++ b/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 + ')')
 

From 86e85e18530ea71b71fee6a2da696fffb6e95a7c Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Sun, 19 Oct 2025 13:08:16 +0200
Subject: [PATCH 2/2] Fix a bug where the x-axis was double flipped

---
 src/cdtools/tools/plotting/plotting.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py
index f2357f78..0f570ea1 100644
--- a/src/cdtools/tools/plotting/plotting.py
+++ b/src/cdtools/tools/plotting/plotting.py
@@ -578,7 +578,11 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
              linewidth=linewidth)
     
     if invert_xaxis:
-        plt.gca().invert_xaxis()
+        ax = plt.gca()
+        x_min, x_max = ax.get_xlim()
+        # Protect against flipping twice if plotting on top of existing graph
+        if x_min <= x_max:
+            ax.invert_xaxis()
         
     plt.xlabel('X (' + units + ')')
     plt.ylabel('Y (' + units + ')')