initial import of wfm package

setup.py

@@ -0,0 +1,29 @@
+import setuptools
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+setuptools.setup(
+    name="dress",
+    version="0.0.1",
+    author="Neil Vaytet",
+    author_email="[email protected]",
+    description="Data Reduction for ESS",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/ess-dmsc-dram/dress",
+    packages=setuptools.find_packages("src"),
+    package_dir={"": "src"},
+    classifiers=[
+        "Programming Language :: Python :: 2",
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)",
+        "Operating System :: OS Independent",
+    ],
+    install_requires=[
+        "matplotlib",
+        "numpy",
+        "h5py",
+        "scipy"
+    ],
+)

src/dress/__init__.py

@@ -0,0 +1 @@
+from . import wfm

src/dress/wfm/__init__.py

@@ -0,0 +1,4 @@
+# from .wfm_to_tof import to_tof
+from .stitch import stitch
+from .wfmess import get_frames
+from . import v20

src/dress/wfm/chopper.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+
+def deg_to_rad(x):
+    return x * np.pi / 180.0
+
+
+class Chopper:
+
+    def __init__(self, frequency=0, openings=None, distance=0, phase=0,
+                 unit="rad", name=""):
+        # openings is list. First entry is start angle of the first cut-out
+        # second entry is end angle of first cut-out, etc.
+        self.frequency = frequency
+        self.openings = openings
+        self.omega = 2.0 * np.pi * frequency
+        self.distance = distance
+        self.phase = phase
+        if unit == "deg":
+            self.openings = deg_to_rad(self.openings)
+            self.phase = deg_to_rad(self.phase)
+        self.name = name

src/dress/wfm/frames_analytical.py

@@ -0,0 +1,127 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+
+
+
+def frames_analytical(instrument=None, plot=False):
+
+    info = instrument["info"]
+    choppers = instrument["choppers"]
+
+    # Seconds to microseconds
+    microseconds = 1.0e6
+
+    # Frame colors
+    # colors = ['b', 'k', 'g', 'r', 'cyan', 'magenta']
+
+    # Define and draw source pulse
+    x0 = 0.0
+    x1 = info["pulse_length"] * microseconds
+    y0 = 0.0
+    y1 = 0.0
+
+    # Make figure
+    if plot:
+        fig, ax = plt.subplots(1, 1, figsize=(9, 7))
+        ax.grid(True, color='lightgray', linestyle="dotted")
+        ax.set_axisbelow(True)
+
+        # Plot the chopper openings
+        for key, ch in choppers.items():
+            dist = [ch.distance, ch.distance]
+            for i in range(0, len(ch.openings), 2):
+                t1 = (ch.openings[i] + ch.phase) / ch.omega * microseconds
+                t2 = (ch.openings[i+1] + ch.phase) / ch.omega * microseconds
+                ax.plot([t1, t2], dist, color="C{}".format(i//2))
+            ax.text(t2 + (t2-t1), ch.distance, ch.name, ha="left", va="center")
+
+        psize = info["detector_position"] / 50.0
+        rect = Rectangle((x0, y0), x1, -psize, lw=1, fc='orange', ec='k', hatch="////", zorder=10)
+        ax.add_patch(rect)
+        ax.text(x0, -psize, "Source pulse (2.86 ms)", ha="left", va="top", fontsize=6)
+
+    # Now find frame boundaries and draw frames
+    frame_boundaries = []
+    frame_shifts = []
+
+    for i in range(info["nframes"]):
+
+        # Find the minimum and maximum slopes that are allowed through each frame
+        slope_min = 1.0e30
+        slope_max = -1.0e30
+        for key, ch in choppers.items():
+
+            # For now, ignore Wavelength band double chopper
+            if len(ch.openings) == info["nframes"] * 2:
+
+                xmin = (ch.openings[i*2] + ch.phase) / ch.omega * microseconds
+                xmax = (ch.openings[i*2+1] + ch.phase) / ch.omega * microseconds
+                slope1 = (ch.distance - y1) / (xmin - x1)
+                slope2 = (ch.distance - y0) / (xmax - x0)
+
+                if slope_min > slope1:
+                    x2 = xmin
+                    y2 = ch.distance
+                    slope_min = slope1
+                if slope_max < slope2:
+                    x3 = xmax
+                    y3 = ch.distance
+                    slope_max = slope2
+
+        # Compute line equation parameters y = a*x + b
+        a1 = (y3 - y0) / (x3 - x0)
+        a2 = (y2 - y1) / (x2 - x1)
+        b1 = y0 - a1 * x0
+        b2 = y1 - a2 * x1
+
+        y4 = info["detector_position"]
+        y5 = info["detector_position"]
+
+        # This is the frame boundaries
+        x5 = (y5 - b1)/a1
+        x4 = (y4 - b2)/a2
+        frame_boundaries.append([x4, x5])
+
+        # Compute frame shifts from fastest neutrons in frame
+        frame_shifts.append(-(info["wfm_choppers_midpoint"] - b2)/a2)
+
+        if plot:
+            col = "C{}".format(i)
+            ax.fill([x0, x1, x4, x5], [y0, y1, y4, y5], alpha=0.3, color=col)
+            ax.plot([x0, x5], [y0, y5], color=col, lw=1)
+            ax.plot([x1, x4], [y1, y4], color=col, lw=1)
+            ax.text(0.5*(x4+x5), info["detector_position"], "Frame {}".format(i+1), ha="center", va="top")
+
+
+    if plot:
+        # Plot detector location
+        ax.plot([0, np.amax(frame_boundaries)], [info["detector_position"], info["detector_position"]], lw=3, color='grey')
+        ax.text(0.0, info["detector_position"], "Detector", va="bottom", ha="left")
+        # Plot WFM choppers mid-point
+        ax.plot([0, np.amax(frame_boundaries)], [info["wfm_choppers_midpoint"], info["wfm_choppers_midpoint"]], lw=1, color='grey', ls="dashed")
+        ax.text(np.amax(frame_boundaries), info["wfm_choppers_midpoint"], "WFM chopper mid-point", va="bottom", ha="right")
+        # Save the figure
+        ax.set_xlabel("Time [microseconds]")
+        ax.set_ylabel("Distance [m]")
+        if isinstance(plot, str):
+            figname = plot
+        else:
+            figname = "frames_analytical.pdf"
+        fig.savefig(figname, bbox_inches="tight")
+
+    # if verbose:
+    #     # Print results
+    #     print("The frame boundaries are:", frame_boundaries)
+    #     print("The frame gaps are:", frame_gaps)
+    #     print("The frame shifts are:", frame_shifts)
+
+    frame_gaps = [0.5*(frame_boundaries[i][1]+frame_boundaries[i+1][0]) for i in range(len(frame_boundaries)-1)]
+
+    frames = {"left_edges": np.array([f[0] for f in frame_boundaries]),
+                    "right_edges": np.array([f[1] for f in frame_boundaries]),
+                    "gaps": np.array(frame_gaps),
+                    "shifts": np.array(frame_shifts)}
+
+
+    return frames

src/dress/wfm/frames_peakfinding.py

@@ -0,0 +1,190 @@
+"""
+  Find WFM frames using peak-finding.
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.signal import find_peaks
+from scipy.ndimage import gaussian_filter1d
+from matplotlib.patches import Rectangle
+
+
+def _tof_shifts(pscdata, psc_frequency=0):
+    cut_out_centre = np.reshape(pscdata,(len(pscdata)//2, 2)).mean(1)
+    cut_out_diffs = np.ediff1d(cut_out_centre)
+    return cut_out_diffs / (360.0*psc_frequency)
+
+
+def _get_frame_shifts(instrument):
+    """
+    This function generates a list of proper shifting parameters from psc data.
+
+    :param instrument: An instrument setup, containing chopper information.
+    :return: List of relative frame shifts in microseconds.
+    """
+
+    # Factor of 0.5 * 1.0e6  comes from taking mean and converting to microseconds
+    relative_shifts = (
+        _tof_shifts(instrument["choppers"]["WFM1"].openings,
+                    psc_frequency=instrument["choppers"]["WFM1"].frequency) +
+        _tof_shifts(instrument["choppers"]["WFM2"].openings,
+                    psc_frequency=instrument["choppers"]["WFM2"].frequency)
+        ) * 5.0e+05
+    return -relative_shifts
+
+
+def _make_frame_shifts(initial_shift, other_shifts=None):
+    """
+    This function constructs a list of proper shifting parameters for wavelength frames.
+    It expects the initial shift in microseconds, followed by an optional list/tuple of
+    shifts relative to the previous one. For example:
+
+        frame_shifts = make_frame_shifts(-6000, [-2000, -2000])
+
+    results in -6000,-8000,-10000.
+
+    The default other_shifts are the currently correct shifts for the V20 beamline at HZB.
+
+    :param initial_shift: Shift to apply to first wavelength frame in microseconds.
+    :param other_shifts: Iterable of shift increments with respect to the previous value
+                          (initial_shift for first value in list), also in microseconds.
+    :return: List of absolute frame shifts in microseconds.
+    """
+    frame_shift_increments = [initial_shift] + list(other_shifts)
+    frame_shifts = [sum(frame_shift_increments[:i + 1]) for i in
+                    range(len(frame_shift_increments))]
+
+    print("The frame_shifts are:", frame_shifts)
+
+    return frame_shifts
+
+
+def frames_peakfinding(data=None, instrument=None, initial_shift=-6630,
+                       plot=False, verbose=False, nbins=512,
+                       bg_threshold=0.05, peak_prominence=0.05,
+                       gsmooth=2, inter_frame_threshold=1.5,
+                       inter_frame_gaps=800.0):
+    """
+
+
+    """
+
+    frame_shifts = _make_frame_shifts(initial_shift, _get_frame_shifts(instrument))
+
+    if "events" in data:
+        # Find min and max in event time-of-arrival (TOA)
+        xmin = np.amin(data["events"])
+        xmax = np.amax(data["events"])
+        nevents = len(data["events"])
+        print("Read in {} events.".format(nevents))
+        print("The minimum and maximum tof values are:", xmin, xmax)
+        # Add padding
+        dx = (xmax - xmin) / float(nbins)
+        xmin -= dx
+        xmax += dx
+        frames_x = np.linspace(xmin, xmax, nbins + 1)
+        y, edges = np.histogram(data["events"], bins=frames_x)
+        x = 0.5 * (edges[:-1] + edges[1:])
+    else:
+        y = data["y"]
+        x = data["x"]
+
+    if plot:
+        # Histogram the events for plotting
+        fig, ax = plt.subplots()
+        ax.plot(x, y, color='k', label="Raw data", lw=3)
+
+    # Gaussian smooth the data
+    y = gaussian_filter1d(y, gsmooth)
+    if plot:
+        ax.plot(x, y, color="lightgray", label="Gaussian smoothed", lw=1)
+
+
+    # Minimum and maximum values in the histogram
+    ymin = np.amin(y)
+    ymax = np.amax(y)
+
+    # Determine background threshold by histogramming data
+    nx = len(x)
+    hist, edges = np.histogram(y, bins=50)
+    bg_raw = edges[np.argmax(hist) + 1]
+    bg_value = bg_raw + bg_threshold * (ymax - bg_raw)
+    hmin = np.amin(hist)
+    hmax = np.amax(hist)
+    # Find the leading and trailing edges; i.e. the leftmost and rightmost
+    # points that exceed the background value
+    i_start = 0
+    i_end = nx-1
+    for i in range(nx):
+        if y[i] > bg_value:
+            i_start = i
+            break
+    for i in range(nx-1,1,-1):
+        if y[i] > bg_value:
+            i_end = i
+            break
+
+
+    # Find valleys with scipy.signal.find_peaks by inverting the y array
+    peaks, params = find_peaks(-y, prominence=peak_prominence*(ymax-ymin),
+                               distance=nbins//20)
+
+    # Check for unwanted peaks:
+    # 1. If a peak is found inside the noise leading the signal, it will often
+    # have a zero left base
+    to_be_removed = []
+    for p in range(len(peaks)):
+        if params["left_bases"][p] < nbins//10:
+            print("Removed peak number {} at x,y = {},{} because of a bad left base".format(p,x[peaks[p]],y[peaks[p]]))
+            to_be_removed.append(p)
+    # 2. If there are peaks in the middle of a frame, then the y value is high
+    threshold = inter_frame_threshold * bg_value
+    for p in range(len(peaks)):
+        if y[peaks[p]] > threshold:
+            print("Removed peak number {} at x,y = {},{} because the y value exceeds the threshold {}".format(p, x[peaks[p]], y[peaks[p]], threshold))
+            to_be_removed.append(p)
+    # Actually remove the peaks
+    peaks = np.delete(peaks, to_be_removed)
+
+    frame_gaps = x[peaks]
+    print("The frame_gaps are:", frame_gaps)
+
+    # Compute boundaries
+    if not isinstance(inter_frame_gaps, list):
+        inter_frame_gaps = [inter_frame_gaps] * len(frame_gaps)
+    frame_boundaries = np.zeros([len(frame_gaps) + 1, 2])
+    frame_boundaries[0, 0] = x[i_start]
+    frame_boundaries[-1, 1] = x[i_end]
+    for i, g in enumerate(frame_gaps):
+        frame_boundaries[i, 1] = g - 0.5 * inter_frame_gaps[i]
+        frame_boundaries[i+1, 0] = g + 0.5 * inter_frame_gaps[i]
+
+    # Plot the diagnostics
+    if plot:
+        peaks = np.concatenate([[i_start], peaks, [i_end]])
+        ax.plot(x[peaks], y[peaks], 'o', color='r', zorder=10)
+        ax.plot([x[0], x[-1]], [threshold, threshold], ls="dashed", color="g")
+        xl = ax.get_xlim()
+        yl = ax.get_ylim()
+        for i in range(instrument["info"]["nframes"]):
+            col = "C{}".format(i)
+            ax.add_patch(Rectangle((frame_boundaries[i, 0], yl[0]),
+                frame_boundaries[i, 1]-frame_boundaries[i, 0], yl[1]-yl[0],
+                fc=col, alpha=0.4, ec="none"))
+            ax.axvline(x=frame_boundaries[i, 0], color=col)
+            ax.axvline(x=frame_boundaries[i, 1], color=col)
+        ax.add_patch(Rectangle((xl[0], yl[0]), xl[1]-xl[0], bg_value-yl[0],
+            fc='lightgray', ec='none', zorder=-10))
+        ax.set_xlim(xl)
+        ax.set_ylim(yl)
+        if isinstance(plot, str):
+            figname = plot
+        else:
+            figname = "frames_peakfinding.pdf"
+        fig.savefig(figname, bbox_inches="tight")
+
+    frames = {"left_edge": np.array([f[0] for f in frame_boundaries]),
+                    "right_edge": np.array([f[1] for f in frame_boundaries]),
+                    "gaps": np.array(frame_gaps),
+                    "shifts": np.array(frame_shifts)}
+
+    return frames

src/dress/wfm/stitch.py

@@ -0,0 +1,12 @@
+from .stitch_histogram import stitch_histogram
+from .stitch_events import stitch_events
+from .stitch_files import stitch_files
+
+
+def stitch(x=None, y=None, files=None, frames=None, **kwargs):
+    if y is not None:
+        return stitch_histogram(x=x, y=y, frames=frames, **kwargs)
+    if events is not None:
+        return stitch_events(events=x, frames=frames, **kwargs)
+    if files is not None:
+        return stitch_files(files=files, **kwargs)