From cb4454792eecaf613164838aead6b26631d7a05a Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 14:38:43 -0400
Subject: [PATCH 1/9] work on probe

---
 refl1d/experiment.py          |   4 +-
 refl1d/ncnrdata.py            |   7 +-
 refl1d/plotting/__init__.py   |   0
 refl1d/plotting/plot_probe.py | 357 +++++++++++++++++++++++++++++
 refl1d/probe.py               | 410 ----------------------------------
 5 files changed, 364 insertions(+), 414 deletions(-)
 create mode 100644 refl1d/plotting/__init__.py
 create mode 100644 refl1d/plotting/plot_probe.py

diff --git a/refl1d/experiment.py b/refl1d/experiment.py
index 93d109a2..8e35bafc 100644
--- a/refl1d/experiment.py
+++ b/refl1d/experiment.py
@@ -31,6 +31,7 @@
 from .reflectivity import BASE_GUIDE_ANGLE as DEFAULT_THETA_M
 from . import model
 from .probe import Probe, NeutronProbe, PolarizedNeutronProbe
+from .plotting import plot_probe
 
 # print("Using pure python reflectivity calculator")
 # from .abeles import refl as reflamp
@@ -167,7 +168,8 @@ def nllf(self):
     def plot_reflectivity(self, show_resolution=False, view=None, plot_shift=None):
         n = self.interpolation
         QR = self.reflectivity(interpolation=n)
-        self.probe.plot(
+        plot_probe.plot(
+            self.probe,
             theory=QR,
             substrate=self._substrate,
             surface=self._surface,
diff --git a/refl1d/ncnrdata.py b/refl1d/ncnrdata.py
index 9d18d456..c8141189 100644
--- a/refl1d/ncnrdata.py
+++ b/refl1d/ncnrdata.py
@@ -16,23 +16,24 @@
 Example loading data:
 
     >>> from refl1d.names import *
+    >>> from refl1d import plotting
     >>> datafile = sample_data('chale207.refl')
     >>> instrument = NCNR.ANDR(Tlo=0.5, slits_at_Tlo=0.2, slits_below=0.1)
     >>> probe = instrument.load(datafile)
-    >>> probe.plot(view='log')
+    >>> plotting.plot_probe.plot(probe, view='log')
 
 Magnetic data has multiple cross sections and often has fixed slits:
 
     >>> datafile = sample_data('lha03_255G.refl')
     >>> instrument = NCNR.NG1(slits_at_Tlo=1)
     >>> probe = instrument.load_magnetic(datafile)
-    >>> probe.plot(view='SA', substrate=silicon) # Spin asymmetry view
+    >>> plotting.plot_probe.plot(probe, view='SA', substrate=silicon) # Spin asymmetry view
 
 For simulation, you need a probe and a sample:
 
     >>> instrument = NCNR.ANDR(Tlo=0.5, slits_at_Tlo=0.2, slits_below=0.1)
     >>> probe = instrument.probe(T=np.linspace(0, 5, 51))
-    >>> probe.plot_resolution()
+    >>> plotting.plot_probe.plot_resolution(probe)
     >>> sample = silicon(0, 10) | gold(100, 10) | air
     >>> M = Experiment(probe=probe, sample=sample)
     >>> M.simulate_data() # Optional
diff --git a/refl1d/plotting/__init__.py b/refl1d/plotting/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
new file mode 100644
index 00000000..aa569219
--- /dev/null
+++ b/refl1d/plotting/plot_probe.py
@@ -0,0 +1,357 @@
+import sys
+from refl1d.probe import Probe, PolarizedNeutronProbe, ProbeSet, spin_asymmetry
+from bumps.plotutil import coordinated_colors, auto_shift
+
+import matplotlib.pyplot as plt
+
+
+def plot(probe, view=None, theory=None, **kwargs):
+    if type(probe) is PolarizedNeutronProbe:
+        polarized_neutron_plot(probe, view, **kwargs)
+    elif type(probe) is ProbeSet:
+        for p, th in probe.parts(theory):
+            plot(p, theory=th, **kwargs)
+    else:
+        base_plot(probe, view, **kwargs)
+
+
+def _xs_plot(probe, plotter, theory=None, **kwargs):
+    """
+    Plot available cross sections
+    """
+
+    thismodule = sys.modules[__name__]
+    fn = getattr(thismodule, plotter)
+
+    if theory is None:
+        theory = (None, None, None, None)
+    for x_data, x_th, suffix in zip(probe.xs, theory, ("$^{--}$", "$^{-+}$", "$^{+-}$", "$^{++}$")):
+        if x_data is not None:
+            fn(theory=x_th, suffix=suffix, **kwargs)
+
+
+def _probeset_plot(probeset, plotter, theory=None, **kw):
+    thismodule = sys.modules[__name__]
+    fn = getattr(thismodule, plotter)
+
+    for p, th in probeset.parts(theory):
+        fn(theory=th, **kw)
+
+
+def polarized_neutron_plot(probe, view=None, **kwargs):
+    """
+    Plot theory against data.
+
+    Need substrate/surface for Fresnel-normalized reflectivity
+    """
+    view = view if view is not None else probe.view
+
+    if view is None:
+        view = Probe.view  # Default to Probe.view
+
+    if view == "linear":
+        _xs_plot(probe, "plot_linear", **kwargs)
+    elif view == "log":
+        _xs_plot(probe, "plot_log", **kwargs)
+    elif view == "fresnel":
+        _xs_plot(probe, "plot_fresnel", **kwargs)
+    elif view == "logfresnel":
+        _xs_plot(probe, "plot_logfresnel", **kwargs)
+    elif view == "q4":
+        _xs_plot(probe, "plot_Q4", **kwargs)
+    elif view.startswith("resid"):
+        _xs_plot(probe, "plot_residuals", **kwargs)
+    elif view == "SA":
+        plot_SA(probe, **kwargs)
+    elif view == "resolution":
+        _xs_plot(probe, "plot_resolution", **kwargs)
+    else:
+        raise TypeError("incorrect reflectivity view '%s'" % view)
+
+
+def probeset_plot(probeset, view=None, **kwargs):
+    """
+    Plot a ProbeSet.
+    """
+    view = view if view is not None else probe.view
+
+    if view is None:
+        view = ProbeSet.view  # Default to Probe.view
+
+    if view == "linear":
+        _probeset_plot(probeset, "plot_linear", **kwargs)
+    elif view == "log":
+        _probeset_plot(probeset, "plot_log", **kwargs)
+    elif view == "fresnel":
+        _probeset_plot(probeset, "plot_fresnel", **kwargs)
+    elif view == "logfresnel":
+        _probeset_plot(probeset, "plot_logfresnel", **kwargs)
+    elif view == "q4":
+        _probeset_plot(probeset, "plot_Q4", **kwargs)
+    elif view.startswith("resid"):
+        _probeset_plot(probeset, "plot_residuals", **kwargs)
+    elif view == "SA":
+        _probeset_plot(probeset, **kwargs)
+    elif view == "resolution":
+        for p in probeset.probes:
+            p.plot_resolution(**kwargs)
+    else:
+        raise TypeError("incorrect reflectivity view '%s'" % view)
+
+
+def base_plot(probe, view=None, **kwargs):
+    """
+    Plot theory against data.
+
+    Need substrate/surface for Fresnel-normalized reflectivity
+    """
+    view = view if view is not None else probe.view
+
+    if view == "linear":
+        plot_linear(probe, **kwargs)
+    elif view == "log":
+        plot_log(probe, **kwargs)
+    elif view == "fresnel":
+        plot_fresnel(probe, **kwargs)
+    elif view == "logfresnel":
+        plot_logfresnel(probe, **kwargs)
+    elif view == "q4":
+        plot_Q4(probe, **kwargs)
+    elif view == "resolution":
+        plot_resolution(probe, **kwargs)
+    elif view.startswith("resid"):
+        plot_residuals(probe, **kwargs)
+    elif view == "fft":
+        plot_fft(probe, **kwargs)
+    elif view == "SA":  # SA does not plot since it does not exist
+        pass
+    else:
+        raise TypeError("incorrect reflectivity view '%s'" % view)
+
+
+def plot_resolution(probe, suffix="", label=None, **kwargs):
+    plt.plot(probe.Q, probe.dQ, label=probe.label(prefix=label, suffix=suffix))
+    plt.xlabel(r"Q ($\AA^{-1}$)")
+    plt.ylabel(r"Q resolution ($1-\sigma \AA^{-1}$)")
+    plt.title("Measurement resolution")
+
+
+def plot_linear(probe, **kwargs):
+    """
+    Plot the data associated with probe.
+    """
+
+    probe._plot_pair(ylabel="Reflectivity", **kwargs)
+    plt.yscale("linear")
+
+
+def plot_log(probe, **kwargs):
+    """
+    Plot the data associated with probe.
+    """
+
+    _plot_pair(probe, ylabel="Reflectivity", **kwargs)
+    plt.yscale("log")
+
+
+def plot_logfresnel(probe, *args, **kw):
+    """
+    Plot the log Fresnel-normalized reflectivity associated with the probe.
+    """
+
+    probe.plot_fresnel(*args, **kw)
+    plt.yscale("log")
+
+
+def plot_fresnel(probe, substrate=None, surface=None, **kwargs):
+    r"""
+    Plot the Fresnel-normalized reflectivity associated with the probe.
+
+    Note that the Fresnel reflectivity has the intensity and background
+    applied before normalizing so that hydrogenated samples display more
+    cleanly.  The formula to reproduce the graph is:
+
+    .. math::
+
+            R' = R / (F(Q) I + B)
+            \Delta R' = \Delta R / (F(Q) I + B)
+
+    where $I$ is the intensity and $B$ is the background.
+    """
+
+    if substrate is None and surface is None:
+        raise TypeError("Fresnel-normalized reflectivity needs substrate or surface")
+    F = probe.fresnel(substrate=substrate, surface=surface)
+
+    # print("substrate", substrate, "surface", surface)
+    def scale(Q, dQ, R, dR, interpolation=0):
+        Q, fresnel = probe.apply_beam(probe.calc_Q, F(probe.calc_Q), interpolation=interpolation)
+        return Q, dQ, R / fresnel, dR / fresnel
+
+    if substrate is None:
+        name = "air:%s" % surface.name
+    elif surface is None or isinstance(surface, Vacuum):
+        name = substrate.name
+    else:
+        name = "%s:%s" % (substrate.name, surface.name)
+    _plot_pair(probe, scale=scale, ylabel="R/(R(%s)" % name, **kwargs)
+
+
+def plot_Q4(probe, **kwargs):
+    r"""
+    Plot the Q**4 reflectivity associated with the probe.
+
+    Note that Q**4 reflectivity has the intensity and background applied
+    so that hydrogenated samples display more cleanly.  The formula
+    to reproduce the graph is:
+
+    .. math::
+
+            R' = R / ( (100*Q)^{-4} I + B)
+            \Delta R' = \Delta R / ( (100*Q)^{-4} I + B )
+
+    where $B$ is the background.
+    """
+
+    def scale(Q, dQ, R, dR, interpolation=0):
+        # Q4 = np.maximum(1e-8*Q**-4, probe.background.value)
+        Q4 = 1e-8 * Q**-4 * probe.intensity.value + probe.background.value
+        return Q, dQ, R / Q4, dR / Q4
+
+    # Q4[Q4==0] = 1
+    _plot_pair(probe, scale=scale, ylabel="R (100 Q)^4", **kwargs)
+
+
+def _plot_pair(
+    probe,
+    theory=None,
+    scale=lambda Q, dQ, R, dR, interpolation=0: (Q, dQ, R, dR),
+    ylabel="",
+    suffix="",
+    label=None,
+    plot_shift=None,
+    **kwargs,
+):
+    c = coordinated_colors()
+    plot_shift = plot_shift if plot_shift is not None else Probe.plot_shift
+    trans = auto_shift(plot_shift)
+    if hasattr(probe, "R") and probe.R is not None:
+        Q, dQ, R, dR = scale(probe.Q, probe.dQ, probe.R, probe.dR)
+        if not probe.show_resolution:
+            dQ = None
+        plt.errorbar(
+            Q,
+            R,
+            yerr=dR,
+            xerr=dQ,
+            capsize=0,
+            fmt=".",
+            color=c["light"],
+            transform=trans,
+            label=probe.label(prefix=label, gloss="data", suffix=suffix),
+        )
+    if theory is not None:
+        # TODO: completely restructure interpolation handling
+        # Interpolation is used to show the theory curve between the
+        # data points.  The _calc_Q points used to predict theory at
+        # the measured data are used for the interpolated Q points, with
+        # the resolution function centered on each interpolated value.
+        # The result is that when interpolation != 0, there are more
+        # theory points than data points, and we will need to accomodate
+        # this when computing normalization curves for Fresnel and Q^4
+        # reflectivity.
+        # Issues with current implementation:
+        # * If the resolution is too tight, there won't be sufficient
+        #   support from _calc_Q to compute theory at Q interpolated.
+        # * dQ for the interpolated points uses linear interpolation
+        #   of dQ between neighbours.  If measurements with tight and
+        #   loose resolution are interleaved the result will look very
+        #   strange.
+        # * There are too many assumptions about interpolation shared
+        #   between Experiment and Probe objects.  In particular, the
+        #   Fresnel object needs to be computed at the same interpolated
+        #   points as the theory function.
+        # * If there are large gaps in the data the interpolation will
+        #   not fill them in correctly.  Perhaps we should set _Q_plot
+        #   and _calc_Q_plot independently from the data?
+        # * We sometimes show the theory curve without resolution
+        #   applied; this has not been tested with interpolation
+        interpolation = kwargs.get("interpolation", 0)
+        Q, R = theory
+        Q, dQ, R, dR = scale(Q, 0, R, 0, interpolation=interpolation)
+        plt.plot(
+            Q,
+            R,
+            "-",
+            color=c["dark"],
+            transform=trans,
+            label=probe.label(prefix=label, gloss="theory", suffix=suffix),
+        )
+        # from numpy.fft import fft
+        # x, y = Q[1::2], abs(fft(R)[:(len(R)-1)//2])
+        # y = y * (R.max()/y.max())
+        # plt.plot(x, y, '-')
+    plt.xlabel("Q (inv Angstroms)")
+    plt.ylabel(ylabel)
+    # plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
+    h = plt.legend(fancybox=True, numpoints=1)
+    h.get_frame().set_alpha(0.5)
+
+
+def plot_residuals(probe, theory=None, suffix="", label=None, plot_shift=None, **kwargs):
+    plot_shift = plot_shift if plot_shift is not None else Probe.residuals_shift
+    trans = auto_shift(plot_shift)
+    if theory is not None and probe.R is not None:
+        c = coordinated_colors()
+        Q, R = theory
+        # In case theory curve is evaluated at more/different points...
+        R = np.interp(probe.Q, Q, R)
+        residual = (R - probe.R) / probe.dR
+        plt.plot(
+            probe.Q, residual, ".", color=c["light"], transform=trans, label=probe.label(prefix=label, suffix=suffix)
+        )
+    plt.axhline(1, color="black", ls="--", lw=1)
+    plt.axhline(0, color="black", lw=1)
+    plt.axhline(-1, color="black", ls="--", lw=1)
+    plt.xlabel("Q (inv A)")
+    plt.ylabel("(theory-data)/error")
+    plt.legend(numpoints=1)
+
+
+def plot_SA(probe, theory=None, label=None, plot_shift=None, **kwargs):
+    if probe.pp is None or probe.mm is None:
+        raise TypeError("cannot form spin asymmetry plot without ++ and --")
+
+    plot_shift = plot_shift if plot_shift is not None else Probe.plot_shift
+    trans = auto_shift(plot_shift)
+    pp, mm = probe.pp, probe.mm
+    c = coordinated_colors()
+    if hasattr(pp, "R") and hasattr(mm, "R") and pp.R is not None and mm.R is not None:
+        Q, SA, dSA = spin_asymmetry(pp.Q, pp.R, pp.dR, mm.Q, mm.R, mm.dR)
+        if dSA is not None:
+            res = probe.show_resolution if probe.show_resolution is not None else Probe.show_resolution
+            dQ = pp.dQ if res else None
+            plt.errorbar(
+                Q,
+                SA,
+                yerr=dSA,
+                xerr=dQ,
+                fmt=".",
+                capsize=0,
+                label=pp.label(prefix=label, gloss="data"),
+                transform=trans,
+                color=c["light"],
+            )
+        else:
+            plt.plot(Q, SA, ".", label=pp.label(prefix=label, gloss="data"), transform=trans, color=c["light"])
+        # Set limits based on max theoretical SA, which is in (-1.0, 1.0)
+        # If the error bars are bigger than that, you usually don't care.
+        ylim_low, ylim_high = plt.ylim()
+        plt.ylim(max(ylim_low, -2.5), min(ylim_high, 2.5))
+    if theory is not None:
+        mm, mp, pm, pp = theory
+        Q, SA, _ = spin_asymmetry(pp[0], pp[1], None, mm[0], mm[1], None)
+        plt.plot(Q, SA, label=probe.pp.label(prefix=label, gloss="theory"), transform=trans, color=c["dark"])
+    plt.xlabel(r"Q (\AA^{-1})")
+    plt.ylabel(r"spin asymmetry $(R^{++} -\, R^{--}) / (R^{++} +\, R^{--})$")
+    plt.legend(numpoints=1)
diff --git a/refl1d/probe.py b/refl1d/probe.py
index 4eda9ae7..652ed2c5 100644
--- a/refl1d/probe.py
+++ b/refl1d/probe.py
@@ -12,8 +12,6 @@
 
 """
 
-from __future__ import with_statement, division, print_function
-
 # TOF stitching introduces a lot of complexity.
 # Theta offset:
 #   Q1 = 4 pi sin(T1 + offset)/L1
@@ -55,7 +53,6 @@
 
 from periodictable import nsf, xsf
 from bumps.parameter import Parameter, to_dict, Constant
-from bumps.plotutil import coordinated_colors, auto_shift
 from bumps.data import parse_multi, strip_quotes
 from bumps.util import USE_PYDANTIC
 
@@ -347,261 +344,6 @@ def save(self, filename, theory, substrate=None, surface=None):
             fid.write(asbytes(header))
             np.savetxt(fid, A.T, fmt="%20.15g")
 
-    def plot(self, view=None, **kwargs):
-        """
-        Plot theory against data.
-
-        Need substrate/surface for Fresnel-normalized reflectivity
-        """
-
-        view = view if view is not None else self.view
-
-        if view == "linear":
-            self.plot_linear(**kwargs)
-        elif view == "log":
-            self.plot_log(**kwargs)
-        elif view == "fresnel":
-            self.plot_fresnel(**kwargs)
-        elif view == "logfresnel":
-            self.plot_logfresnel(**kwargs)
-        elif view == "q4":
-            self.plot_Q4(**kwargs)
-        elif view == "resolution":
-            self.plot_resolution(**kwargs)
-        elif view.startswith("resid"):
-            self.plot_residuals(**kwargs)
-        elif view == "fft":
-            self.plot_fft(**kwargs)
-        elif view == "SA":  # SA does not plot since it does not exist
-            pass
-        else:
-            raise TypeError("incorrect reflectivity view '%s'" % view)
-
-    def plot_resolution(self, suffix="", label=None, **kwargs):
-        import matplotlib.pyplot as plt
-
-        plt.plot(self.Q, self.dQ, label=self.label(prefix=label, suffix=suffix))
-        plt.xlabel(r"Q ($\AA^{-1}$)")
-        plt.ylabel(r"Q resolution ($1-\sigma \AA^{-1}$)")
-        plt.title("Measurement resolution")
-
-    def plot_linear(self, **kwargs):
-        """
-        Plot the data associated with probe.
-        """
-        import matplotlib.pyplot as plt
-
-        self._plot_pair(ylabel="Reflectivity", **kwargs)
-        plt.yscale("linear")
-
-    def plot_log(self, **kwargs):
-        """
-        Plot the data associated with probe.
-        """
-        import matplotlib.pyplot as plt
-
-        self._plot_pair(ylabel="Reflectivity", **kwargs)
-        plt.yscale("log")
-
-    def plot_logfresnel(self, *args, **kw):
-        """
-        Plot the log Fresnel-normalized reflectivity associated with the probe.
-        """
-        import matplotlib.pyplot as plt
-
-        self.plot_fresnel(*args, **kw)
-        plt.yscale("log")
-
-    def plot_fresnel(self, substrate=None, surface=None, **kwargs):
-        r"""
-        Plot the Fresnel-normalized reflectivity associated with the probe.
-
-        Note that the Fresnel reflectivity has the intensity and background
-        applied before normalizing so that hydrogenated samples display more
-        cleanly.  The formula to reproduce the graph is:
-
-        .. math::
-
-             R' = R / (F(Q) I + B)
-             \Delta R' = \Delta R / (F(Q) I + B)
-
-        where $I$ is the intensity and $B$ is the background.
-        """
-        if substrate is None and surface is None:
-            raise TypeError("Fresnel-normalized reflectivity needs substrate or surface")
-        F = self.fresnel(substrate=substrate, surface=surface)
-
-        # print("substrate", substrate, "surface", surface)
-        def scale(Q, dQ, R, dR, interpolation=0):
-            Q, fresnel = self.apply_beam(self.calc_Q, F(self.calc_Q), interpolation=interpolation)
-            return Q, dQ, R / fresnel, dR / fresnel
-
-        if substrate is None:
-            name = "air:%s" % surface.name
-        elif surface is None or isinstance(surface, Vacuum):
-            name = substrate.name
-        else:
-            name = "%s:%s" % (substrate.name, surface.name)
-        self._plot_pair(scale=scale, ylabel="R/(R(%s)" % name, **kwargs)
-
-    def plot_Q4(self, **kwargs):
-        r"""
-        Plot the Q**4 reflectivity associated with the probe.
-
-        Note that Q**4 reflectivity has the intensity and background applied
-        so that hydrogenated samples display more cleanly.  The formula
-        to reproduce the graph is:
-
-        .. math::
-
-             R' = R / ( (100*Q)^{-4} I + B)
-             \Delta R' = \Delta R / ( (100*Q)^{-4} I + B )
-
-        where $B$ is the background.
-        """
-
-        def scale(Q, dQ, R, dR, interpolation=0):
-            # Q4 = np.maximum(1e-8*Q**-4, self.background.value)
-            Q4 = 1e-8 * Q**-4 * self.intensity.value + self.background.value
-            return Q, dQ, R / Q4, dR / Q4
-
-        # Q4[Q4==0] = 1
-        self._plot_pair(scale=scale, ylabel="R (100 Q)^4", **kwargs)
-
-    def _plot_pair(
-        self,
-        theory=None,
-        scale=lambda Q, dQ, R, dR, interpolation=0: (Q, dQ, R, dR),
-        ylabel="",
-        suffix="",
-        label=None,
-        plot_shift=None,
-        **kwargs,
-    ):
-        import matplotlib.pyplot as plt
-
-        c = coordinated_colors()
-        plot_shift = plot_shift if plot_shift is not None else Probe.plot_shift
-        trans = auto_shift(plot_shift)
-        if hasattr(self, "R") and self.R is not None:
-            Q, dQ, R, dR = scale(self.Q, self.dQ, self.R, self.dR)
-            if not self.show_resolution:
-                dQ = None
-            plt.errorbar(
-                Q,
-                R,
-                yerr=dR,
-                xerr=dQ,
-                capsize=0,
-                fmt=".",
-                color=c["light"],
-                transform=trans,
-                label=self.label(prefix=label, gloss="data", suffix=suffix),
-            )
-        if theory is not None:
-            # TODO: completely restructure interpolation handling
-            # Interpolation is used to show the theory curve between the
-            # data points.  The _calc_Q points used to predict theory at
-            # the measured data are used for the interpolated Q points, with
-            # the resolution function centered on each interpolated value.
-            # The result is that when interpolation != 0, there are more
-            # theory points than data points, and we will need to accomodate
-            # this when computing normalization curves for Fresnel and Q^4
-            # reflectivity.
-            # Issues with current implementation:
-            # * If the resolution is too tight, there won't be sufficient
-            #   support from _calc_Q to compute theory at Q interpolated.
-            # * dQ for the interpolated points uses linear interpolation
-            #   of dQ between neighbours.  If measurements with tight and
-            #   loose resolution are interleaved the result will look very
-            #   strange.
-            # * There are too many assumptions about interpolation shared
-            #   between Experiment and Probe objects.  In particular, the
-            #   Fresnel object needs to be computed at the same interpolated
-            #   points as the theory function.
-            # * If there are large gaps in the data the interpolation will
-            #   not fill them in correctly.  Perhaps we should set _Q_plot
-            #   and _calc_Q_plot independently from the data?
-            # * We sometimes show the theory curve without resolution
-            #   applied; this has not been tested with interpolation
-            interpolation = kwargs.get("interpolation", 0)
-            Q, R = theory
-            Q, dQ, R, dR = scale(Q, 0, R, 0, interpolation=interpolation)
-            plt.plot(
-                Q,
-                R,
-                "-",
-                color=c["dark"],
-                transform=trans,
-                label=self.label(prefix=label, gloss="theory", suffix=suffix),
-            )
-            # from numpy.fft import fft
-            # x, y = Q[1::2], abs(fft(R)[:(len(R)-1)//2])
-            # y = y * (R.max()/y.max())
-            # plt.plot(x, y, '-')
-        plt.xlabel("Q (inv Angstroms)")
-        plt.ylabel(ylabel)
-        # plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
-        h = plt.legend(fancybox=True, numpoints=1)
-        h.get_frame().set_alpha(0.5)
-
-    def plot_residuals(self, theory=None, suffix="", label=None, plot_shift=None, **kwargs):
-        import matplotlib.pyplot as plt
-
-        plot_shift = plot_shift if plot_shift is not None else Probe.residuals_shift
-        trans = auto_shift(plot_shift)
-        if theory is not None and self.R is not None:
-            c = coordinated_colors()
-            Q, R = theory
-            # In case theory curve is evaluated at more/different points...
-            R = np.interp(self.Q, Q, R)
-            residual = (R - self.R) / self.dR
-            plt.plot(
-                self.Q, residual, ".", color=c["light"], transform=trans, label=self.label(prefix=label, suffix=suffix)
-            )
-        plt.axhline(1, color="black", ls="--", lw=1)
-        plt.axhline(0, color="black", lw=1)
-        plt.axhline(-1, color="black", ls="--", lw=1)
-        plt.xlabel("Q (inv A)")
-        plt.ylabel("(theory-data)/error")
-        plt.legend(numpoints=1)
-
-    def plot_fft(self, theory=None, suffix="", label=None, substrate=None, surface=None, **kwargs):
-        """
-        FFT analysis of reflectivity signal.
-        """
-        raise NotImplementedError
-        import matplotlib.pyplot as plt
-
-        c = coordinated_colors()
-        if substrate is None and surface is None:
-            raise TypeError("FFT reflectivity needs substrate or surface")
-        F = self.fresnel(substrate=substrate, surface=surface)
-        # Qc = sqrt(16*pi*substrate)
-        Qc = 0
-        Qmax = max(self.Q)
-        T = np.linspace(Qc, Qmax, len(self.Q))
-        z = np.linspace(0, 2 * pi / Qmax, len(self.Q) // 2)
-        if hasattr(self, "R"):
-            signal = np.interp(T, self.Q, self.R / F(self.Q))
-            A = abs(numpy.fft.fft(signal - np.average(signal)))
-            A = A[: len(A) // 2]
-            plt.plot(z, A, ".-", color=c["light"], label=self.label(prefix=label, gloss="data", suffix=suffix))
-        if theory is not None:
-            Q, R = theory
-            signal = np.interp(T, Q, R / F(Q))
-            A = abs(numpy.fft.fft(signal - np.average(signal)))
-            A = A[: len(A) // 2]
-            plt.plot(z, A, "-", color=c["dark"], label=self.label(prefix=label, gloss="theory", suffix=suffix))
-        plt.xlabel("w (A)")
-        if substrate is None:
-            name = "air:%s" % surface.name
-        elif surface is None or isinstance(surface, Vacuum):
-            name = substrate.name
-        else:
-            name = "%s:%s" % (substrate.name, surface.name)
-        plt.ylabel("|FFT(R/R(%s))|" % name)
-
     def label(self, prefix=None, gloss="", suffix=""):
         if not prefix:
             prefix = self.name
@@ -1208,56 +950,6 @@ def save(self, filename, theory, substrate=None, surface=None):
 
     save.__doc__ = Probe.save.__doc__
 
-    def plot(self, theory=None, **kw):
-        import matplotlib.pyplot as plt
-
-        for p, th in self.parts(theory):
-            p.plot(theory=th, **kw)
-
-    plot.__doc__ = Probe.plot.__doc__
-
-    def plot_resolution(self, **kw):
-        for p in self.probes:
-            p.plot_resolution(**kw)
-
-    plot_resolution.__doc__ = Probe.plot_resolution.__doc__
-
-    def plot_linear(self, theory=None, **kw):
-        for p, th in self.parts(theory):
-            p.plot_linear(theory=th, **kw)
-
-    plot_linear.__doc__ = Probe.plot_linear.__doc__
-
-    def plot_log(self, theory=None, **kw):
-        for p, th in self.parts(theory):
-            p.plot_log(theory=th, **kw)
-
-    plot_log.__doc__ = Probe.plot_log.__doc__
-
-    def plot_fresnel(self, theory=None, **kw):
-        for p, th in self.parts(theory):
-            p.plot_fresnel(theory=th, **kw)
-
-    plot_fresnel.__doc__ = Probe.plot_fresnel.__doc__
-
-    def plot_logfresnel(self, theory=None, **kw):
-        for p, th in self.parts(theory):
-            p.plot_logfresnel(theory=th, **kw)
-
-    plot_logfresnel.__doc__ = Probe.plot_logfresnel.__doc__
-
-    def plot_Q4(self, theory=None, **kw):
-        for p, th in self.parts(theory):
-            p.plot_Q4(theory=th, **kw)
-
-    plot_Q4.__doc__ = Probe.plot_Q4.__doc__
-
-    def plot_residuals(self, theory=None, **kw):
-        for p, th in self.parts(theory):
-            p.plot_residuals(theory=th, **kw)
-
-    plot_residuals.__doc__ = Probe.plot_residuals.__doc__
-
     def parts(self, theory):
         if theory is None:
             for p in self.probes:
@@ -2060,108 +1752,6 @@ def save(self, filename, theory, substrate=None, surface=None):
 
     save.__doc__ = Probe.save.__doc__
 
-    def plot(self, view=None, **kwargs):
-        """
-        Plot theory against data.
-
-        Need substrate/surface for Fresnel-normalized reflectivity
-        """
-        view = view if view is not None else self.view
-
-        if view is None:
-            view = Probe.view  # Default to Probe.view
-
-        if view == "linear":
-            self.plot_linear(**kwargs)
-        elif view == "log":
-            self.plot_log(**kwargs)
-        elif view == "fresnel":
-            self.plot_fresnel(**kwargs)
-        elif view == "logfresnel":
-            self.plot_logfresnel(**kwargs)
-        elif view == "q4":
-            self.plot_Q4(**kwargs)
-        elif view.startswith("resid"):
-            self.plot_residuals(**kwargs)
-        elif view == "SA":
-            self.plot_SA(**kwargs)
-        elif view == "resolution":
-            self.plot_resolution(**kwargs)
-        else:
-            raise TypeError("incorrect reflectivity view '%s'" % view)
-
-    def plot_resolution(self, **kwargs):
-        self._xs_plot("plot_resolution", **kwargs)
-
-    def plot_linear(self, **kwargs):
-        self._xs_plot("plot_linear", **kwargs)
-
-    def plot_log(self, **kwargs):
-        self._xs_plot("plot_log", **kwargs)
-
-    def plot_fresnel(self, **kwargs):
-        self._xs_plot("plot_fresnel", **kwargs)
-
-    def plot_logfresnel(self, **kwargs):
-        self._xs_plot("plot_logfresnel", **kwargs)
-
-    def plot_Q4(self, **kwargs):
-        self._xs_plot("plot_Q4", **kwargs)
-
-    def plot_residuals(self, **kwargs):
-        self._xs_plot("plot_residuals", **kwargs)
-
-    def plot_SA(self, theory=None, label=None, plot_shift=None, **kwargs):
-        import matplotlib.pyplot as plt
-
-        if self.pp is None or self.mm is None:
-            raise TypeError("cannot form spin asymmetry plot without ++ and --")
-
-        plot_shift = plot_shift if plot_shift is not None else Probe.plot_shift
-        trans = auto_shift(plot_shift)
-        pp, mm = self.pp, self.mm
-        c = coordinated_colors()
-        if hasattr(pp, "R") and hasattr(mm, "R") and pp.R is not None and mm.R is not None:
-            Q, SA, dSA = spin_asymmetry(pp.Q, pp.R, pp.dR, mm.Q, mm.R, mm.dR)
-            if dSA is not None:
-                res = self.show_resolution if self.show_resolution is not None else Probe.show_resolution
-                dQ = pp.dQ if res else None
-                plt.errorbar(
-                    Q,
-                    SA,
-                    yerr=dSA,
-                    xerr=dQ,
-                    fmt=".",
-                    capsize=0,
-                    label=pp.label(prefix=label, gloss="data"),
-                    transform=trans,
-                    color=c["light"],
-                )
-            else:
-                plt.plot(Q, SA, ".", label=pp.label(prefix=label, gloss="data"), transform=trans, color=c["light"])
-            # Set limits based on max theoretical SA, which is in (-1.0, 1.0)
-            # If the error bars are bigger than that, you usually don't care.
-            ylim_low, ylim_high = plt.ylim()
-            plt.ylim(max(ylim_low, -2.5), min(ylim_high, 2.5))
-        if theory is not None:
-            mm, mp, pm, pp = theory
-            Q, SA, _ = spin_asymmetry(pp[0], pp[1], None, mm[0], mm[1], None)
-            plt.plot(Q, SA, label=self.pp.label(prefix=label, gloss="theory"), transform=trans, color=c["dark"])
-        plt.xlabel(r"Q (\AA^{-1})")
-        plt.ylabel(r"spin asymmetry $(R^{++} -\, R^{--}) / (R^{++} +\, R^{--})$")
-        plt.legend(numpoints=1)
-
-    def _xs_plot(self, plotter, theory=None, **kwargs):
-        import matplotlib.pyplot as plt
-
-        # Plot available cross sections
-        if theory is None:
-            theory = (None, None, None, None)
-        for x_data, x_th, suffix in zip(self.xs, theory, ("$^{--}$", "$^{-+}$", "$^{+-}$", "$^{++}$")):
-            if x_data is not None:
-                fn = getattr(x_data, plotter)
-                fn(theory=x_th, suffix=suffix, **kwargs)
-
 
 def spin_asymmetry(Qp, Rp, dRp, Qm, Rm, dRm):
     r"""

From 73ac7af2d5b42a9c079dd31fd78be3bea4fcc33c Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 14:52:53 -0400
Subject: [PATCH 2/9] Add docs

---
 refl1d/plotting/plot_probe.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index aa569219..ec8ad705 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -1,3 +1,11 @@
+"""
+Plotting functions for probes.
+
+Design note:
+These functions were originally part of the probe classes, and were moved out
+for future refactor to have a common plotting interface with webview.
+"""
+
 import sys
 from refl1d.probe import Probe, PolarizedNeutronProbe, ProbeSet, spin_asymmetry
 from bumps.plotutil import coordinated_colors, auto_shift
@@ -6,6 +14,9 @@
 
 
 def plot(probe, view=None, theory=None, **kwargs):
+    """
+    General plotting function for probes.
+    """
     if type(probe) is PolarizedNeutronProbe:
         polarized_neutron_plot(probe, view, **kwargs)
     elif type(probe) is ProbeSet:
@@ -31,6 +42,9 @@ def _xs_plot(probe, plotter, theory=None, **kwargs):
 
 
 def _probeset_plot(probeset, plotter, theory=None, **kw):
+    """
+    Plot a ProbeSet
+    """
     thismodule = sys.modules[__name__]
     fn = getattr(thismodule, plotter)
 

From 929a560b383aa0322f00b0114955e6dbd5b4a4d5 Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 15:05:23 -0400
Subject: [PATCH 3/9] Clean up

---
 refl1d/plotting/plot_probe.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index ec8ad705..36251323 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -87,8 +87,6 @@ def probeset_plot(probeset, view=None, **kwargs):
     """
     Plot a ProbeSet.
     """
-    view = view if view is not None else probe.view
-
     if view is None:
         view = ProbeSet.view  # Default to Probe.view
 

From 14e0b3814c0cd81b016f390e15d36cfa6ca9ca89 Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 16:21:25 -0400
Subject: [PATCH 4/9] plotter class

---
 refl1d/experiment.py          |   5 +-
 refl1d/plotting/plot_probe.py | 171 ++++++++++++++--------------------
 2 files changed, 72 insertions(+), 104 deletions(-)

diff --git a/refl1d/experiment.py b/refl1d/experiment.py
index 8e35bafc..bf0f595b 100644
--- a/refl1d/experiment.py
+++ b/refl1d/experiment.py
@@ -168,7 +168,7 @@ def nllf(self):
     def plot_reflectivity(self, show_resolution=False, view=None, plot_shift=None):
         n = self.interpolation
         QR = self.reflectivity(interpolation=n)
-        plot_probe.plot(
+        plotter = plot_probe.ProbePlotter(
             self.probe,
             theory=QR,
             substrate=self._substrate,
@@ -177,6 +177,9 @@ def plot_reflectivity(self, show_resolution=False, view=None, plot_shift=None):
             plot_shift=plot_shift,
             interpolation=n,
         )
+        plotter.plot()
+        plotter.y_log()
+        plotter.x_log()
 
         if show_resolution:
             import matplotlib.pyplot as plt
diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index 36251323..eb975b98 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -23,7 +23,7 @@ def plot(probe, view=None, theory=None, **kwargs):
         for p, th in probe.parts(theory):
             plot(p, theory=th, **kwargs)
     else:
-        base_plot(probe, view, **kwargs)
+        base_plot(probe, view, theory=theory, **kwargs)
 
 
 def _xs_plot(probe, plotter, theory=None, **kwargs):
@@ -111,6 +111,62 @@ def probeset_plot(probeset, view=None, **kwargs):
         raise TypeError("incorrect reflectivity view '%s'" % view)
 
 
+class ProbePlotter:
+    def __init__(self, probe, **kwargs):
+        self.probe = probe
+        self.kwargs = kwargs
+
+    def plot(self):
+        _plot_pair(self.probe, ylabel="Reflectivity", **self.kwargs)
+        plt.yscale("log")
+        plt.xscale("log")
+
+    def y_linear(self):
+        plt.yscale("linear")
+
+    def y_log(self):
+        plt.yscale("log")
+
+    def x_linear(self):
+        plt.xscale("linear")
+
+    def x_log(self):
+        plt.xscale("log")
+
+    def plot_fresnel(self, substrate=None, surface=None):
+        r"""
+        Plot the Fresnel-normalized reflectivity associated with the probe.
+
+        Note that the Fresnel reflectivity has the intensity and background
+        applied before normalizing so that hydrogenated samples display more
+        cleanly.  The formula to reproduce the graph is:
+
+        .. math::
+
+                R' = R / (F(Q) I + B)
+                \Delta R' = \Delta R / (F(Q) I + B)
+
+        where $I$ is the intensity and $B$ is the background.
+        """
+
+        if substrate is None and surface is None:
+            raise TypeError("Fresnel-normalized reflectivity needs substrate or surface")
+        F = self.probe.fresnel(substrate=substrate, surface=surface)
+
+        # print("substrate", substrate, "surface", surface)
+        def scale(Q, dQ, R, dR, interpolation=0):
+            Q, fresnel = self.probe.apply_beam(self.probe.calc_Q, F(self.probe.calc_Q), interpolation=interpolation)
+            return Q, dQ, R / fresnel, dR / fresnel
+
+        if substrate is None:
+            name = "air:%s" % surface.name
+        elif surface is None or isinstance(surface, Vacuum):
+            name = substrate.name
+        else:
+            name = "%s:%s" % (substrate.name, surface.name)
+        _plot_pair(self.probe, scale=scale, ylabel="R/(R(%s)" % name)
+
+
 def base_plot(probe, view=None, **kwargs):
     """
     Plot theory against data.
@@ -118,27 +174,27 @@ def base_plot(probe, view=None, **kwargs):
     Need substrate/surface for Fresnel-normalized reflectivity
     """
     view = view if view is not None else probe.view
+    if view in ["linear", "log"]:
+        plotter = ProbePlotter(probe, **kwargs)
 
-    if view == "linear":
-        plot_linear(probe, **kwargs)
-    elif view == "log":
-        plot_log(probe, **kwargs)
+    if "fresnel" in view:
+        plotter.plot_fresnel()
     elif view == "fresnel":
-        plot_fresnel(probe, **kwargs)
-    elif view == "logfresnel":
-        plot_logfresnel(probe, **kwargs)
+        plotter.plot_fresnel()
+        plotter.y_linear()
     elif view == "q4":
         plot_Q4(probe, **kwargs)
     elif view == "resolution":
         plot_resolution(probe, **kwargs)
     elif view.startswith("resid"):
         plot_residuals(probe, **kwargs)
-    elif view == "fft":
-        plot_fft(probe, **kwargs)
     elif view == "SA":  # SA does not plot since it does not exist
         pass
     else:
-        raise TypeError("incorrect reflectivity view '%s'" % view)
+        plotter.plot()
+
+    if "log" in view:
+        plotter.y_log()
 
 
 def plot_resolution(probe, suffix="", label=None, **kwargs):
@@ -148,67 +204,6 @@ def plot_resolution(probe, suffix="", label=None, **kwargs):
     plt.title("Measurement resolution")
 
 
-def plot_linear(probe, **kwargs):
-    """
-    Plot the data associated with probe.
-    """
-
-    probe._plot_pair(ylabel="Reflectivity", **kwargs)
-    plt.yscale("linear")
-
-
-def plot_log(probe, **kwargs):
-    """
-    Plot the data associated with probe.
-    """
-
-    _plot_pair(probe, ylabel="Reflectivity", **kwargs)
-    plt.yscale("log")
-
-
-def plot_logfresnel(probe, *args, **kw):
-    """
-    Plot the log Fresnel-normalized reflectivity associated with the probe.
-    """
-
-    probe.plot_fresnel(*args, **kw)
-    plt.yscale("log")
-
-
-def plot_fresnel(probe, substrate=None, surface=None, **kwargs):
-    r"""
-    Plot the Fresnel-normalized reflectivity associated with the probe.
-
-    Note that the Fresnel reflectivity has the intensity and background
-    applied before normalizing so that hydrogenated samples display more
-    cleanly.  The formula to reproduce the graph is:
-
-    .. math::
-
-            R' = R / (F(Q) I + B)
-            \Delta R' = \Delta R / (F(Q) I + B)
-
-    where $I$ is the intensity and $B$ is the background.
-    """
-
-    if substrate is None and surface is None:
-        raise TypeError("Fresnel-normalized reflectivity needs substrate or surface")
-    F = probe.fresnel(substrate=substrate, surface=surface)
-
-    # print("substrate", substrate, "surface", surface)
-    def scale(Q, dQ, R, dR, interpolation=0):
-        Q, fresnel = probe.apply_beam(probe.calc_Q, F(probe.calc_Q), interpolation=interpolation)
-        return Q, dQ, R / fresnel, dR / fresnel
-
-    if substrate is None:
-        name = "air:%s" % surface.name
-    elif surface is None or isinstance(surface, Vacuum):
-        name = substrate.name
-    else:
-        name = "%s:%s" % (substrate.name, surface.name)
-    _plot_pair(probe, scale=scale, ylabel="R/(R(%s)" % name, **kwargs)
-
-
 def plot_Q4(probe, **kwargs):
     r"""
     Plot the Q**4 reflectivity associated with the probe.
@@ -263,31 +258,6 @@ def _plot_pair(
             label=probe.label(prefix=label, gloss="data", suffix=suffix),
         )
     if theory is not None:
-        # TODO: completely restructure interpolation handling
-        # Interpolation is used to show the theory curve between the
-        # data points.  The _calc_Q points used to predict theory at
-        # the measured data are used for the interpolated Q points, with
-        # the resolution function centered on each interpolated value.
-        # The result is that when interpolation != 0, there are more
-        # theory points than data points, and we will need to accomodate
-        # this when computing normalization curves for Fresnel and Q^4
-        # reflectivity.
-        # Issues with current implementation:
-        # * If the resolution is too tight, there won't be sufficient
-        #   support from _calc_Q to compute theory at Q interpolated.
-        # * dQ for the interpolated points uses linear interpolation
-        #   of dQ between neighbours.  If measurements with tight and
-        #   loose resolution are interleaved the result will look very
-        #   strange.
-        # * There are too many assumptions about interpolation shared
-        #   between Experiment and Probe objects.  In particular, the
-        #   Fresnel object needs to be computed at the same interpolated
-        #   points as the theory function.
-        # * If there are large gaps in the data the interpolation will
-        #   not fill them in correctly.  Perhaps we should set _Q_plot
-        #   and _calc_Q_plot independently from the data?
-        # * We sometimes show the theory curve without resolution
-        #   applied; this has not been tested with interpolation
         interpolation = kwargs.get("interpolation", 0)
         Q, R = theory
         Q, dQ, R, dR = scale(Q, 0, R, 0, interpolation=interpolation)
@@ -299,13 +269,8 @@ def _plot_pair(
             transform=trans,
             label=probe.label(prefix=label, gloss="theory", suffix=suffix),
         )
-        # from numpy.fft import fft
-        # x, y = Q[1::2], abs(fft(R)[:(len(R)-1)//2])
-        # y = y * (R.max()/y.max())
-        # plt.plot(x, y, '-')
-    plt.xlabel("Q (inv Angstroms)")
+    plt.xlabel("Q (1/Angstroms)")
     plt.ylabel(ylabel)
-    # plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
     h = plt.legend(fancybox=True, numpoints=1)
     h.get_frame().set_alpha(0.5)
 

From d6f718b08317dce7af47e8870d051d14d9fd39ce Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 16:23:51 -0400
Subject: [PATCH 5/9] add Q4

---
 refl1d/plotting/plot_probe.py | 51 +++++++++++++++++------------------
 1 file changed, 25 insertions(+), 26 deletions(-)

diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index eb975b98..90d10959 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -164,7 +164,31 @@ def scale(Q, dQ, R, dR, interpolation=0):
             name = substrate.name
         else:
             name = "%s:%s" % (substrate.name, surface.name)
-        _plot_pair(self.probe, scale=scale, ylabel="R/(R(%s)" % name)
+        _plot_pair(self.probe, scale=scale, ylabel="R/(R(%s)" % name, **self.kwargs)
+
+    def plot_Q4(self):
+        r"""
+        Plot the Q**4 reflectivity associated with the probe.
+
+        Note that Q**4 reflectivity has the intensity and background applied
+        so that hydrogenated samples display more cleanly.  The formula
+        to reproduce the graph is:
+
+        .. math::
+
+                R' = R / ( (100*Q)^{-4} I + B)
+                \Delta R' = \Delta R / ( (100*Q)^{-4} I + B )
+
+        where $B$ is the background.
+        """
+
+        def scale(Q, dQ, R, dR, interpolation=0):
+            # Q4 = np.maximum(1e-8*Q**-4, probe.background.value)
+            Q4 = 1e-8 * Q**-4 * self.probe.intensity.value + self.probe.background.value
+            return Q, dQ, R / Q4, dR / Q4
+
+        # Q4[Q4==0] = 1
+        _plot_pair(self.probe, scale=scale, ylabel="R (100 Q)^4", **self.kwargs)
 
 
 def base_plot(probe, view=None, **kwargs):
@@ -204,31 +228,6 @@ def plot_resolution(probe, suffix="", label=None, **kwargs):
     plt.title("Measurement resolution")
 
 
-def plot_Q4(probe, **kwargs):
-    r"""
-    Plot the Q**4 reflectivity associated with the probe.
-
-    Note that Q**4 reflectivity has the intensity and background applied
-    so that hydrogenated samples display more cleanly.  The formula
-    to reproduce the graph is:
-
-    .. math::
-
-            R' = R / ( (100*Q)^{-4} I + B)
-            \Delta R' = \Delta R / ( (100*Q)^{-4} I + B )
-
-    where $B$ is the background.
-    """
-
-    def scale(Q, dQ, R, dR, interpolation=0):
-        # Q4 = np.maximum(1e-8*Q**-4, probe.background.value)
-        Q4 = 1e-8 * Q**-4 * probe.intensity.value + probe.background.value
-        return Q, dQ, R / Q4, dR / Q4
-
-    # Q4[Q4==0] = 1
-    _plot_pair(probe, scale=scale, ylabel="R (100 Q)^4", **kwargs)
-
-
 def _plot_pair(
     probe,
     theory=None,

From 5beaa29b86c3bdef369ee2bead78b053ae5ce7bf Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 16:39:13 -0400
Subject: [PATCH 6/9] Much better

---
 refl1d/experiment.py          |  5 +--
 refl1d/plotting/plot_probe.py | 82 +++++++++++++++++++++--------------
 2 files changed, 50 insertions(+), 37 deletions(-)

diff --git a/refl1d/experiment.py b/refl1d/experiment.py
index bf0f595b..8e35bafc 100644
--- a/refl1d/experiment.py
+++ b/refl1d/experiment.py
@@ -168,7 +168,7 @@ def nllf(self):
     def plot_reflectivity(self, show_resolution=False, view=None, plot_shift=None):
         n = self.interpolation
         QR = self.reflectivity(interpolation=n)
-        plotter = plot_probe.ProbePlotter(
+        plot_probe.plot(
             self.probe,
             theory=QR,
             substrate=self._substrate,
@@ -177,9 +177,6 @@ def plot_reflectivity(self, show_resolution=False, view=None, plot_shift=None):
             plot_shift=plot_shift,
             interpolation=n,
         )
-        plotter.plot()
-        plotter.y_log()
-        plotter.x_log()
 
         if show_resolution:
             import matplotlib.pyplot as plt
diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index 90d10959..30f352ae 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -7,6 +7,8 @@
 """
 
 import sys
+import numpy as np
+
 from refl1d.probe import Probe, PolarizedNeutronProbe, ProbeSet, spin_asymmetry
 from bumps.plotutil import coordinated_colors, auto_shift
 
@@ -133,7 +135,7 @@ def x_linear(self):
     def x_log(self):
         plt.xscale("log")
 
-    def plot_fresnel(self, substrate=None, surface=None):
+    def plot_fresnel(self):
         r"""
         Plot the Fresnel-normalized reflectivity associated with the probe.
 
@@ -148,6 +150,8 @@ def plot_fresnel(self, substrate=None, surface=None):
 
         where $I$ is the intensity and $B$ is the background.
         """
+        substrate = self.kwargs.get("substrate", None)
+        surface = self.kwargs.get("surface", None)
 
         if substrate is None and surface is None:
             raise TypeError("Fresnel-normalized reflectivity needs substrate or surface")
@@ -190,6 +194,45 @@ def scale(Q, dQ, R, dR, interpolation=0):
         # Q4[Q4==0] = 1
         _plot_pair(self.probe, scale=scale, ylabel="R (100 Q)^4", **self.kwargs)
 
+    def plot_resolution(self):
+        suffix = self.kwargs.get("suffix", "")
+        label = self.kwargs.get("label", None)
+
+        plt.plot(self.probe.Q, self.probe.dQ, label=self.probe.label(prefix=label, suffix=suffix))
+        plt.xlabel(r"Q ($\AA^{-1}$)")
+        plt.ylabel(r"Q resolution ($1-\sigma \AA^{-1}$)")
+        plt.title("Measurement resolution")
+
+    def plot_residuals(self):
+        suffix = self.kwargs.get("suffix", "")
+        label = self.kwargs.get("label", None)
+        plot_shift = self.kwargs.get("plot_shift", None)
+        theory = self.kwargs.get("theory", None)
+
+        plot_shift = plot_shift if plot_shift is not None else Probe.residuals_shift
+        trans = auto_shift(plot_shift)
+
+        if theory is not None and self.probe.R is not None:
+            c = coordinated_colors()
+            Q, R = theory
+            # In case theory curve is evaluated at more/different points...
+            R = np.interp(self.probe.Q, Q, R)
+            residual = (R - self.probe.R) / self.probe.dR
+            plt.plot(
+                self.probe.Q,
+                residual,
+                ".",
+                color=c["light"],
+                transform=trans,
+                label=self.probe.label(prefix=label, suffix=suffix),
+            )
+        plt.axhline(1, color="black", ls="--", lw=1)
+        plt.axhline(0, color="black", lw=1)
+        plt.axhline(-1, color="black", ls="--", lw=1)
+        plt.xlabel("Q (inv A)")
+        plt.ylabel("(theory-data)/error")
+        plt.legend(numpoints=1)
+
 
 def base_plot(probe, view=None, **kwargs):
     """
@@ -198,8 +241,8 @@ def base_plot(probe, view=None, **kwargs):
     Need substrate/surface for Fresnel-normalized reflectivity
     """
     view = view if view is not None else probe.view
-    if view in ["linear", "log"]:
-        plotter = ProbePlotter(probe, **kwargs)
+
+    plotter = ProbePlotter(probe, **kwargs)
 
     if "fresnel" in view:
         plotter.plot_fresnel()
@@ -207,11 +250,11 @@ def base_plot(probe, view=None, **kwargs):
         plotter.plot_fresnel()
         plotter.y_linear()
     elif view == "q4":
-        plot_Q4(probe, **kwargs)
+        plotter.plot_Q4()
     elif view == "resolution":
-        plot_resolution(probe, **kwargs)
+        plotter.plot_resolution()
     elif view.startswith("resid"):
-        plot_residuals(probe, **kwargs)
+        plotter.plot_residuals()
     elif view == "SA":  # SA does not plot since it does not exist
         pass
     else:
@@ -221,13 +264,6 @@ def base_plot(probe, view=None, **kwargs):
         plotter.y_log()
 
 
-def plot_resolution(probe, suffix="", label=None, **kwargs):
-    plt.plot(probe.Q, probe.dQ, label=probe.label(prefix=label, suffix=suffix))
-    plt.xlabel(r"Q ($\AA^{-1}$)")
-    plt.ylabel(r"Q resolution ($1-\sigma \AA^{-1}$)")
-    plt.title("Measurement resolution")
-
-
 def _plot_pair(
     probe,
     theory=None,
@@ -274,26 +310,6 @@ def _plot_pair(
     h.get_frame().set_alpha(0.5)
 
 
-def plot_residuals(probe, theory=None, suffix="", label=None, plot_shift=None, **kwargs):
-    plot_shift = plot_shift if plot_shift is not None else Probe.residuals_shift
-    trans = auto_shift(plot_shift)
-    if theory is not None and probe.R is not None:
-        c = coordinated_colors()
-        Q, R = theory
-        # In case theory curve is evaluated at more/different points...
-        R = np.interp(probe.Q, Q, R)
-        residual = (R - probe.R) / probe.dR
-        plt.plot(
-            probe.Q, residual, ".", color=c["light"], transform=trans, label=probe.label(prefix=label, suffix=suffix)
-        )
-    plt.axhline(1, color="black", ls="--", lw=1)
-    plt.axhline(0, color="black", lw=1)
-    plt.axhline(-1, color="black", ls="--", lw=1)
-    plt.xlabel("Q (inv A)")
-    plt.ylabel("(theory-data)/error")
-    plt.legend(numpoints=1)
-
-
 def plot_SA(probe, theory=None, label=None, plot_shift=None, **kwargs):
     if probe.pp is None or probe.mm is None:
         raise TypeError("cannot form spin asymmetry plot without ++ and --")

From 714e6a1dc30df4d907bf36b3abb1519898e1bde5 Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Thu, 12 Sep 2024 17:29:34 -0400
Subject: [PATCH 7/9] mag class

---
 refl1d/ncnrdata.py            |   2 +-
 refl1d/plotting/plot_probe.py | 144 +++++++++++++++++++---------------
 2 files changed, 80 insertions(+), 66 deletions(-)

diff --git a/refl1d/ncnrdata.py b/refl1d/ncnrdata.py
index c8141189..bf3c91c3 100644
--- a/refl1d/ncnrdata.py
+++ b/refl1d/ncnrdata.py
@@ -33,7 +33,7 @@
 
     >>> instrument = NCNR.ANDR(Tlo=0.5, slits_at_Tlo=0.2, slits_below=0.1)
     >>> probe = instrument.probe(T=np.linspace(0, 5, 51))
-    >>> plotting.plot_probe.plot_resolution(probe)
+    >>> plotting.plot_probe.plot(probe, view='resolution')
     >>> sample = silicon(0, 10) | gold(100, 10) | air
     >>> M = Experiment(probe=probe, sample=sample)
     >>> M.simulate_data() # Optional
diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index 30f352ae..d78d6002 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -20,12 +20,32 @@ def plot(probe, view=None, theory=None, **kwargs):
     General plotting function for probes.
     """
     if type(probe) is PolarizedNeutronProbe:
-        polarized_neutron_plot(probe, view, **kwargs)
+        plotter = MagneticProbePlotter(probe, theory=theory, **kwargs)
     elif type(probe) is ProbeSet:
         for p, th in probe.parts(theory):
             plot(p, theory=th, **kwargs)
     else:
-        base_plot(probe, view, theory=theory, **kwargs)
+        plotter = ProbePlotter(probe, theory=theory, **kwargs)
+
+    view = view if view is not None else probe.view
+
+    # Determine how to plot
+    if view == "fresnel":
+        plotter.plot_fresnel()
+        plotter.y_linear()
+    elif view == "q4":
+        plotter.plot_Q4()
+    elif view == "resolution":
+        plotter.plot_resolution()
+    elif view.startswith("resid"):
+        plotter.plot_residuals()
+    elif view == "SA":
+        plotter.plot_SA()
+    else:
+        plotter.plot()
+
+    if "log" in view:
+        plotter.y_log()
 
 
 def _xs_plot(probe, plotter, theory=None, **kwargs):
@@ -203,6 +223,12 @@ def plot_resolution(self):
         plt.ylabel(r"Q resolution ($1-\sigma \AA^{-1}$)")
         plt.title("Measurement resolution")
 
+    def plot_SA(self):
+        """
+        This plot is unavailable for this probe.
+        """
+        pass
+
     def plot_residuals(self):
         suffix = self.kwargs.get("suffix", "")
         label = self.kwargs.get("label", None)
@@ -234,34 +260,61 @@ def plot_residuals(self):
         plt.legend(numpoints=1)
 
 
-def base_plot(probe, view=None, **kwargs):
-    """
-    Plot theory against data.
+class MagneticProbePlotter(ProbePlotter):
+    def __init__(self, probe, **kwargs):
+        super().__init__(probe, **kwargs)
 
-    Need substrate/surface for Fresnel-normalized reflectivity
-    """
-    view = view if view is not None else probe.view
+    def plot(self):
+        theory = self.kwargs.get("theory", (None, None, None, None))
+        self.kwargs["theory"] = theory
 
-    plotter = ProbePlotter(probe, **kwargs)
+        for x_data, x_th, suffix in zip(self.probe.xs, theory, ("$^{--}$", "$^{-+}$", "$^{+-}$", "$^{++}$")):
+            if x_data is not None:
+                self.kwargs["suffix"] = suffix
+                self.kwargs["theory"] = x_th
+                super().plot()
 
-    if "fresnel" in view:
-        plotter.plot_fresnel()
-    elif view == "fresnel":
-        plotter.plot_fresnel()
-        plotter.y_linear()
-    elif view == "q4":
-        plotter.plot_Q4()
-    elif view == "resolution":
-        plotter.plot_resolution()
-    elif view.startswith("resid"):
-        plotter.plot_residuals()
-    elif view == "SA":  # SA does not plot since it does not exist
-        pass
-    else:
-        plotter.plot()
+    def plot_SA(self):
+        theory = self.kwargs.get("theory", (None, None, None, None))
+        label = self.kwargs.get("label", None)
+        plot_shift = self.kwargs.get("plot_shift", None)
 
-    if "log" in view:
-        plotter.y_log()
+        if self.probe.pp is None or self.probe.mm is None:
+            raise TypeError("cannot form spin asymmetry plot without ++ and --")
+
+        plot_shift = plot_shift if plot_shift is not None else Probe.plot_shift
+        trans = auto_shift(plot_shift)
+        pp, mm = self.probe.pp, self.probe.mm
+        c = coordinated_colors()
+        if hasattr(pp, "R") and hasattr(mm, "R") and pp.R is not None and mm.R is not None:
+            Q, SA, dSA = spin_asymmetry(pp.Q, pp.R, pp.dR, mm.Q, mm.R, mm.dR)
+            if dSA is not None:
+                res = self.probe.show_resolution if self.probe.show_resolution is not None else Probe.show_resolution
+                dQ = pp.dQ if res else None
+                plt.errorbar(
+                    Q,
+                    SA,
+                    yerr=dSA,
+                    xerr=dQ,
+                    fmt=".",
+                    capsize=0,
+                    label=pp.label(prefix=label, gloss="data"),
+                    transform=trans,
+                    color=c["light"],
+                )
+            else:
+                plt.plot(Q, SA, ".", label=pp.label(prefix=label, gloss="data"), transform=trans, color=c["light"])
+            # Set limits based on max theoretical SA, which is in (-1.0, 1.0)
+            # If the error bars are bigger than that, you usually don't care.
+            ylim_low, ylim_high = plt.ylim()
+            plt.ylim(max(ylim_low, -2.5), min(ylim_high, 2.5))
+        if theory is not None:
+            mm, mp, pm, pp = theory
+            Q, SA, _ = spin_asymmetry(pp[0], pp[1], None, mm[0], mm[1], None)
+            plt.plot(Q, SA, label=self.probe.pp.label(prefix=label, gloss="theory"), transform=trans, color=c["dark"])
+        plt.xlabel(r"Q (\AA^{-1})")
+        plt.ylabel(r"spin asymmetry $(R^{++} -\, R^{--}) / (R^{++} +\, R^{--})$")
+        plt.legend(numpoints=1)
 
 
 def _plot_pair(
@@ -308,42 +361,3 @@ def _plot_pair(
     plt.ylabel(ylabel)
     h = plt.legend(fancybox=True, numpoints=1)
     h.get_frame().set_alpha(0.5)
-
-
-def plot_SA(probe, theory=None, label=None, plot_shift=None, **kwargs):
-    if probe.pp is None or probe.mm is None:
-        raise TypeError("cannot form spin asymmetry plot without ++ and --")
-
-    plot_shift = plot_shift if plot_shift is not None else Probe.plot_shift
-    trans = auto_shift(plot_shift)
-    pp, mm = probe.pp, probe.mm
-    c = coordinated_colors()
-    if hasattr(pp, "R") and hasattr(mm, "R") and pp.R is not None and mm.R is not None:
-        Q, SA, dSA = spin_asymmetry(pp.Q, pp.R, pp.dR, mm.Q, mm.R, mm.dR)
-        if dSA is not None:
-            res = probe.show_resolution if probe.show_resolution is not None else Probe.show_resolution
-            dQ = pp.dQ if res else None
-            plt.errorbar(
-                Q,
-                SA,
-                yerr=dSA,
-                xerr=dQ,
-                fmt=".",
-                capsize=0,
-                label=pp.label(prefix=label, gloss="data"),
-                transform=trans,
-                color=c["light"],
-            )
-        else:
-            plt.plot(Q, SA, ".", label=pp.label(prefix=label, gloss="data"), transform=trans, color=c["light"])
-        # Set limits based on max theoretical SA, which is in (-1.0, 1.0)
-        # If the error bars are bigger than that, you usually don't care.
-        ylim_low, ylim_high = plt.ylim()
-        plt.ylim(max(ylim_low, -2.5), min(ylim_high, 2.5))
-    if theory is not None:
-        mm, mp, pm, pp = theory
-        Q, SA, _ = spin_asymmetry(pp[0], pp[1], None, mm[0], mm[1], None)
-        plt.plot(Q, SA, label=probe.pp.label(prefix=label, gloss="theory"), transform=trans, color=c["dark"])
-    plt.xlabel(r"Q (\AA^{-1})")
-    plt.ylabel(r"spin asymmetry $(R^{++} -\, R^{--}) / (R^{++} +\, R^{--})$")
-    plt.legend(numpoints=1)

From 2b3a530906a51f7a8bf28f5fa4d35070dd34947c Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Fri, 13 Sep 2024 13:58:41 -0400
Subject: [PATCH 8/9] Refactor plotters

---
 refl1d/plotting/plot_probe.py | 117 +++++++++-------------------------
 1 file changed, 30 insertions(+), 87 deletions(-)

diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index d78d6002..71885d2c 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -37,103 +37,22 @@ def plot(probe, view=None, theory=None, **kwargs):
         plotter.plot_Q4()
     elif view == "resolution":
         plotter.plot_resolution()
-    elif view.startswith("resid"):
+    elif view is not None and view.startswith("resid"):
         plotter.plot_residuals()
     elif view == "SA":
         plotter.plot_SA()
     else:
         plotter.plot()
 
-    if "log" in view:
+    if view is not None and "log" in view:
         plotter.y_log()
 
 
-def _xs_plot(probe, plotter, theory=None, **kwargs):
-    """
-    Plot available cross sections
-    """
-
-    thismodule = sys.modules[__name__]
-    fn = getattr(thismodule, plotter)
-
-    if theory is None:
-        theory = (None, None, None, None)
-    for x_data, x_th, suffix in zip(probe.xs, theory, ("$^{--}$", "$^{-+}$", "$^{+-}$", "$^{++}$")):
-        if x_data is not None:
-            fn(theory=x_th, suffix=suffix, **kwargs)
-
-
-def _probeset_plot(probeset, plotter, theory=None, **kw):
-    """
-    Plot a ProbeSet
+class ProbePlotter:
     """
-    thismodule = sys.modules[__name__]
-    fn = getattr(thismodule, plotter)
-
-    for p, th in probeset.parts(theory):
-        fn(theory=th, **kw)
-
-
-def polarized_neutron_plot(probe, view=None, **kwargs):
+    Plotter for non-magnetic probes
     """
-    Plot theory against data.
 
-    Need substrate/surface for Fresnel-normalized reflectivity
-    """
-    view = view if view is not None else probe.view
-
-    if view is None:
-        view = Probe.view  # Default to Probe.view
-
-    if view == "linear":
-        _xs_plot(probe, "plot_linear", **kwargs)
-    elif view == "log":
-        _xs_plot(probe, "plot_log", **kwargs)
-    elif view == "fresnel":
-        _xs_plot(probe, "plot_fresnel", **kwargs)
-    elif view == "logfresnel":
-        _xs_plot(probe, "plot_logfresnel", **kwargs)
-    elif view == "q4":
-        _xs_plot(probe, "plot_Q4", **kwargs)
-    elif view.startswith("resid"):
-        _xs_plot(probe, "plot_residuals", **kwargs)
-    elif view == "SA":
-        plot_SA(probe, **kwargs)
-    elif view == "resolution":
-        _xs_plot(probe, "plot_resolution", **kwargs)
-    else:
-        raise TypeError("incorrect reflectivity view '%s'" % view)
-
-
-def probeset_plot(probeset, view=None, **kwargs):
-    """
-    Plot a ProbeSet.
-    """
-    if view is None:
-        view = ProbeSet.view  # Default to Probe.view
-
-    if view == "linear":
-        _probeset_plot(probeset, "plot_linear", **kwargs)
-    elif view == "log":
-        _probeset_plot(probeset, "plot_log", **kwargs)
-    elif view == "fresnel":
-        _probeset_plot(probeset, "plot_fresnel", **kwargs)
-    elif view == "logfresnel":
-        _probeset_plot(probeset, "plot_logfresnel", **kwargs)
-    elif view == "q4":
-        _probeset_plot(probeset, "plot_Q4", **kwargs)
-    elif view.startswith("resid"):
-        _probeset_plot(probeset, "plot_residuals", **kwargs)
-    elif view == "SA":
-        _probeset_plot(probeset, **kwargs)
-    elif view == "resolution":
-        for p in probeset.probes:
-            p.plot_resolution(**kwargs)
-    else:
-        raise TypeError("incorrect reflectivity view '%s'" % view)
-
-
-class ProbePlotter:
     def __init__(self, probe, **kwargs):
         self.probe = probe
         self.kwargs = kwargs
@@ -261,10 +180,17 @@ def plot_residuals(self):
 
 
 class MagneticProbePlotter(ProbePlotter):
+    """
+    Plotter for magnetic probes
+    """
+
     def __init__(self, probe, **kwargs):
         super().__init__(probe, **kwargs)
 
-    def plot(self):
+    def magnetic(self, plot_function):
+        """
+        Loop over the cross sections and plot them
+        """
         theory = self.kwargs.get("theory", (None, None, None, None))
         self.kwargs["theory"] = theory
 
@@ -272,7 +198,24 @@ def plot(self):
             if x_data is not None:
                 self.kwargs["suffix"] = suffix
                 self.kwargs["theory"] = x_th
-                super().plot()
+                plotter = ProbePlotter(x_data, **self.kwargs)
+                fn = getattr(plotter, plot_function)
+                fn()
+
+    def plot(self):
+        self.magnetic("plot")
+
+    def plot_fresnel(self):
+        self.magnetic("plot_fresnel")
+
+    def plot_Q4(self):
+        self.magnetic("plot_Q4")
+
+    def plot_resolution(self):
+        self.magnetic("plot_resolution")
+
+    def plot_residuals(self):
+        self.magnetic("plot_residuals")
 
     def plot_SA(self):
         theory = self.kwargs.get("theory", (None, None, None, None))

From 22a000de5496ef472fa6aeec069a77c591597910 Mon Sep 17 00:00:00 2001
From: Mathieu Doucet <matd10@yahoo.com>
Date: Fri, 13 Sep 2024 14:05:12 -0400
Subject: [PATCH 9/9] Fix logic

---
 refl1d/plotting/plot_probe.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/refl1d/plotting/plot_probe.py b/refl1d/plotting/plot_probe.py
index 71885d2c..fab851ad 100644
--- a/refl1d/plotting/plot_probe.py
+++ b/refl1d/plotting/plot_probe.py
@@ -19,16 +19,18 @@ def plot(probe, view=None, theory=None, **kwargs):
     """
     General plotting function for probes.
     """
+    view = view if view is not None else probe.view
+
     if type(probe) is PolarizedNeutronProbe:
         plotter = MagneticProbePlotter(probe, theory=theory, **kwargs)
     elif type(probe) is ProbeSet:
+        # If we have a ProbeSet, plot each probe in the set and exit
         for p, th in probe.parts(theory):
             plot(p, theory=th, **kwargs)
+        return
     else:
         plotter = ProbePlotter(probe, theory=theory, **kwargs)
 
-    view = view if view is not None else probe.view
-
     # Determine how to plot
     if view == "fresnel":
         plotter.plot_fresnel()