Index: sasmodels/data.py
===================================================================
--- sasmodels/data.py	(revision 3b4243d6f7b05b3af937e85ebba3c60b8f564534)
+++ sasmodels/data.py	(revision 3b4243d6f7b05b3af937e85ebba3c60b8f564534)
@@ -0,0 +1,436 @@
+"""
+SAS data representations.
+
+Plotting functions for data sets:
+
+    :func:`plot_data` plots the data file.
+
+    :func:`plot_theory` plots a calculated result from the model.
+
+Wrappers for the sasview data loader and data manipulations:
+
+    :func:`load_data` loads a sasview data file.
+
+    :func:`set_beam_stop` masks the beam stop from the data.
+
+    :func:`set_half` selects the right or left half of the data, which can
+    be useful for shear measurements which have not been properly corrected
+    for path length and reflections.
+
+    :func:`set_top` cuts the top part off the data.
+
+
+Empty data sets for evaluating models without data:
+
+    :func:`empty_data1D` creates an empty dataset, which is useful for plotting
+    a theory function before the data is measured.
+
+    :func:`empty_data2D` creates an empty 2D dataset.
+
+Note that the empty datasets use a minimal representation of the SasView
+objects so that models can be run without SasView on the path.  You could
+also use these for your own data loader.
+
+"""
+import traceback
+
+import numpy as np
+
+def load_data(filename):
+    """
+    Load data using a sasview loader.
+    """
+    from sas.dataloader.loader import Loader
+    loader = Loader()
+    data = loader.load(filename)
+    if data is None:
+        raise IOError("Data %r could not be loaded" % filename)
+    return data
+
+
+def set_beam_stop(data, radius, outer=None):
+    """
+    Add a beam stop of the given *radius*.  If *outer*, make an annulus.
+    """
+    from sas.dataloader.manipulations import Ringcut
+    if hasattr(data, 'qx_data'):
+        data.mask = Ringcut(0, radius)(data)
+        if outer is not None:
+            data.mask += Ringcut(outer, np.inf)(data)
+    else:
+        data.mask = (data.x < radius)
+        if outer is not None:
+            data.mask |= (data.x >= outer)
+
+
+def set_half(data, half):
+    """
+    Select half of the data, either "right" or "left".
+    """
+    from sas.dataloader.manipulations import Boxcut
+    if half == 'right':
+        data.mask += \
+            Boxcut(x_min=-np.inf, x_max=0.0, y_min=-np.inf, y_max=np.inf)(data)
+    if half == 'left':
+        data.mask += \
+            Boxcut(x_min=0.0, x_max=np.inf, y_min=-np.inf, y_max=np.inf)(data)
+
+
+def set_top(data, cutoff):
+    """
+    Chop the top off the data, above *cutoff*.
+    """
+    from sas.dataloader.manipulations import Boxcut
+    data.mask += \
+        Boxcut(x_min=-np.inf, x_max=np.inf, y_min=-np.inf, y_max=cutoff)(data)
+
+
+class Data1D(object):
+    def __init__(self, x=None, y=None, dx=None, dy=None):
+        self.x, self.y, self.dx, self.dy = x, y, dx, dy
+        self.dxl = None
+
+    def xaxis(self, label, unit):
+        """
+        set the x axis label and unit
+        """
+        self._xaxis = label
+        self._xunit = unit
+
+    def yaxis(self, label, unit):
+        """
+        set the y axis label and unit
+        """
+        self._yaxis = label
+        self._yunit = unit
+
+
+
+class Data2D(object):
+    def __init__(self):
+        self.detector = []
+        self.source = Source()
+
+    def xaxis(self, label, unit):
+        """
+        set the x axis label and unit
+        """
+        self._xaxis = label
+        self._xunit = unit
+
+    def yaxis(self, label, unit):
+        """
+        set the y axis label and unit
+        """
+        self._yaxis = label
+        self._yunit = unit
+
+    def zaxis(self, label, unit):
+        """
+        set the y axis label and unit
+        """
+        self._zaxis = label
+        self._zunit = unit
+
+
+class Vector(object):
+    def __init__(self, x=None, y=None, z=None):
+        self.x, self.y, self.z = x, y, z
+
+class Detector(object):
+    def __init__(self):
+        self.pixel_size = Vector()
+
+class Source(object):
+    pass
+
+
+def empty_data1D(q, resolution=0.05):
+    """
+    Create empty 1D data using the given *q* as the x value.
+
+    *resolution* dq/q defaults to 5%.
+    """
+
+    #Iq = 100 * np.ones_like(q)
+    #dIq = np.sqrt(Iq)
+    Iq, dIq = None, None
+    data = Data1D(q, Iq, dx=resolution * q, dy=dIq)
+    data.filename = "fake data"
+    data.qmin, data.qmax = q.min(), q.max()
+    data.mask = np.zeros(len(q), dtype='bool')
+    return data
+
+
+def empty_data2D(qx, qy=None, resolution=0.05):
+    """
+    Create empty 2D data using the given mesh.
+
+    If *qy* is missing, create a square mesh with *qy=qx*.
+
+    *resolution* dq/q defaults to 5%.
+    """
+    if qy is None:
+        qy = qx
+    Qx, Qy = np.meshgrid(qx, qy)
+    Qx, Qy = Qx.flatten(), Qy.flatten()
+    Iq = 100 * np.ones_like(Qx)
+    dIq = np.sqrt(Iq)
+    mask = np.ones(len(Iq), dtype='bool')
+
+    data = Data2D()
+    data.filename = "fake data"
+    data.qx_data = Qx
+    data.qy_data = Qy
+    data.data = Iq
+    data.err_data = dIq
+    data.mask = mask
+    data.qmin = 1e-16
+    data.qmax = np.inf
+
+    # 5% dQ/Q resolution
+    if resolution != 0:
+        # https://www.ncnr.nist.gov/staff/hammouda/distance_learning/chapter_15.pdf
+        # Should have an additional constant which depends on distances and
+        # radii of the aperture, pixel dimensions and wavelength spread
+        # Instead, assume radial dQ/Q is constant, and perpendicular matches
+        # radial (which instead it should be inverse).
+        Q = np.sqrt(Qx**2 + Qy**2)
+        data.dqx_data = resolution * Q
+        data.dqy_data = resolution * Q
+
+    detector = Detector()
+    detector.pixel_size.x = 5 # mm
+    detector.pixel_size.y = 5 # mm
+    detector.distance = 4 # m
+    data.detector.append(detector)
+    data.xbins = qx
+    data.ybins = qy
+    data.source.wavelength = 5 # angstroms
+    data.source.wavelength_unit = "A"
+    data.Q_unit = "1/A"
+    data.I_unit = "1/cm"
+    data.q_data = np.sqrt(Qx ** 2 + Qy ** 2)
+    data.xaxis("Q_x", "A^{-1}")
+    data.yaxis("Q_y", "A^{-1}")
+    data.zaxis("Intensity", r"\text{cm}^{-1}")
+    return data
+
+
+def plot_data(data, view='log'):
+    """
+    Plot data loaded by the sasview loader.
+    """
+    # Note: kind of weird using the plot result functions to plot just the
+    # data, but they already handle the masking and graph markup already, so
+    # do not repeat.
+    if hasattr(data, 'lam'):
+        _plot_result_sesans(data, None, None, plot_data=True)
+    elif hasattr(data, 'qx_data'):
+        _plot_result2D(data, None, None, view, plot_data=True)
+    else:
+        _plot_result1D(data, None, None, view, plot_data=True)
+
+
+def plot_theory(data, theory, resid=None, view='log', plot_data=True):
+    if hasattr(data, 'lam'):
+        _plot_result_sesans(data, theory, resid, plot_data=True)
+    elif hasattr(data, 'qx_data'):
+        _plot_result2D(data, theory, resid, view, plot_data)
+    else:
+        _plot_result1D(data, theory, resid, view, plot_data)
+
+
+def protect(fn):
+    def wrapper(*args, **kw):
+        try:
+            return fn(*args, **kw)
+        except:
+            traceback.print_exc()
+            pass
+
+    return wrapper
+
+
+@protect
+def _plot_result1D(data, theory, resid, view, plot_data):
+    """
+    Plot the data and residuals for 1D data.
+    """
+    import matplotlib.pyplot as plt
+    from numpy.ma import masked_array, masked
+
+    plot_theory = theory is not None
+    plot_resid = resid is not None
+
+    if data.y is None:
+        plot_data = False
+    scale = data.x**4 if view == 'q4' else 1.0
+
+    if plot_data or plot_theory:
+        if plot_resid:
+            plt.subplot(121)
+
+        #print vmin, vmax
+        positive = False
+        if plot_data:
+            mdata = masked_array(data.y, data.mask)
+            mdata[~np.isfinite(mdata)] = masked
+            if view is 'log':
+                mdata[mdata <= 0] = masked
+            plt.errorbar(data.x, scale*mdata, yerr=data.dy, fmt='.')
+            positive = positive or (mdata>0).any()
+
+        if plot_theory:
+            mtheory = masked_array(theory, data.mask)
+            if view is 'log':
+                mtheory[mtheory<= 0] = masked
+            plt.plot(data.x, scale*mtheory, '-', hold=True)
+            positive = positive or (mtheory>0).any()
+
+        plt.xscale(view)
+        plt.yscale('linear' if view == 'q4' or not positive else view)
+        plt.xlabel('Q')
+        plt.ylabel('I(Q)')
+
+    if plot_resid:
+        if plot_data or plot_theory:
+            plt.subplot(122)
+
+        mresid = masked_array(resid, data.mask)
+        plt.plot(data.x, mresid, '-')
+        plt.ylabel('residuals')
+        plt.xlabel('Q')
+        plt.xscale(view)
+
+
+@protect
+def _plot_result_sesans(data, theory, resid, plot_data):
+    import matplotlib.pyplot as plt
+    if data.y is None:
+        plot_data = False
+    plot_theory = theory is not None
+    plot_resid = resid is not None
+
+    if plot_data or plot_theory:
+        if plot_resid:
+            plt.subplot(121)
+        if plot_data:
+            plt.errorbar(data.x, data.y, yerr=data.dy)
+        if theory is not None:
+            plt.plot(data.x, theory, '-', hold=True)
+        plt.xlabel('spin echo length (nm)')
+        plt.ylabel('polarization (P/P0)')
+
+    if resid is not None:
+        if plot_data or plot_theory:
+            plt.subplot(122)
+
+        plt.plot(data.x, resid, 'x')
+        plt.xlabel('spin echo length (nm)')
+        plt.ylabel('residuals (P/P0)')
+
+
+@protect
+def _plot_result2D(data, theory, resid, view, plot_data):
+    """
+    Plot the data and residuals for 2D data.
+    """
+    import matplotlib.pyplot as plt
+    if data.data is None:
+        plot_data = False
+    plot_theory = theory is not None
+    plot_resid = resid is not None
+
+    # Put theory and data on a common colormap scale
+    vmin, vmax = np.inf, -np.inf
+    if plot_data:
+        target = data.data[~data.mask]
+        datamin = target[target>0].min() if view == 'log' else target.min()
+        datamax = target.max()
+        vmin = min(vmin, datamin)
+        vmax = max(vmax, datamax)
+    if plot_theory:
+        theorymin = theory[theory>0].min() if view == 'log' else theory.min()
+        theorymax = theory.max()
+        vmin = min(vmin, theorymin)
+        vmax = max(vmax, theorymax)
+
+    if plot_data:
+        if plot_theory and plot_resid:
+            plt.subplot(131)
+        elif plot_theory or plot_resid:
+            plt.subplot(121)
+        _plot_2d_signal(data, target, view=view, vmin=vmin, vmax=vmax)
+        plt.title('data')
+        plt.colorbar()
+
+    if plot_theory:
+        if plot_data and plot_resid:
+            plt.subplot(132)
+        elif plot_data:
+            plt.subplot(122)
+        elif plot_resid:
+            plt.subplot(121)
+        _plot_2d_signal(data, theory, view=view, vmin=vmin, vmax=vmax)
+        plt.title('theory')
+        plt.colorbar()
+
+    #if plot_data or plot_theory:
+    #    plt.colorbar()
+
+    if plot_resid:
+        if plot_data and plot_theory:
+            plt.subplot(133)
+        elif plot_data or plot_theory:
+            plt.subplot(122)
+        _plot_2d_signal(data, resid, view='linear')
+        plt.colorbar()
+        plt.title('residuals')
+
+
+@protect
+def _plot_2d_signal(data, signal, vmin=None, vmax=None, view='log'):
+    """
+    Plot the target value for the data.  This could be the data itself,
+    the theory calculation, or the residuals.
+
+    *scale* can be 'log' for log scale data, or 'linear'.
+    """
+    import matplotlib.pyplot as plt
+    from numpy.ma import masked_array
+
+    image = np.zeros_like(data.qx_data)
+    image[~data.mask] = signal
+    valid = np.isfinite(image)
+    if view == 'log':
+        valid[valid] = (image[valid] > 0)
+        image[valid] = np.log10(image[valid])
+    elif view == 'q4':
+        image[valid] *= (data.qx_data[valid]**2+data.qy_data[valid]**2)**2
+    image[~valid | data.mask] = 0
+    #plottable = Iq
+    plottable = masked_array(image, ~valid | data.mask)
+    xmin, xmax = min(data.qx_data), max(data.qx_data)
+    ymin, ymax = min(data.qy_data), max(data.qy_data)
+    # TODO: fix vmin, vmax so it is shared for theory/resid
+    vmin = vmax = None
+    try:
+        if vmin is None: vmin = image[valid & ~data.mask].min()
+        if vmax is None: vmax = image[valid & ~data.mask].max()
+    except:
+        vmin, vmax = 0, 1
+    plt.imshow(plottable.reshape(128, 128),
+               interpolation='nearest', aspect=1, origin='upper',
+               extent=[xmin, xmax, ymin, ymax], vmin=vmin, vmax=vmax)
+
+
+def demo():
+    data = load_data('DEC07086.DAT')
+    set_beam_stop(data, 0.004)
+    plot_data(data)
+    import matplotlib.pyplot as plt; plt.show()
+
+
+if __name__ == "__main__":
+    demo()
