Skip to content

Commit

Permalink
yapf
Browse files Browse the repository at this point in the history
  • Loading branch information
nvaytet committed Jun 3, 2020
1 parent e3ee1e2 commit 609c3a8
Show file tree
Hide file tree
Showing 11 changed files with 210 additions and 122 deletions.
7 changes: 1 addition & 6 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,5 @@
"License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)",
"Operating System :: OS Independent",
],
install_requires=[
"matplotlib",
"numpy",
"h5py",
"scipy"
],
install_requires=["matplotlib", "numpy", "h5py", "scipy"],
)
1 change: 0 additions & 1 deletion src/dress/wfm/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# from .wfm_to_tof import to_tof
from .stitch import stitch
from .wfm import get_frames
from . import v20
10 changes: 7 additions & 3 deletions src/dress/wfm/chopper.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,13 @@ def deg_to_rad(x):


class Chopper:

def __init__(self, frequency=0, openings=None, distance=0, phase=0,
unit="rad", name=""):
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
Expand Down
87 changes: 63 additions & 24 deletions src/dress/wfm/frames_analytical.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from matplotlib.patches import Rectangle



def frames_analytical(instrument=None, plot=False):

info = instrument["info"]
Expand Down Expand Up @@ -32,14 +31,30 @@ def frames_analytical(instrument=None, plot=False):
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")
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)
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)
ax.text(x0,
-psize,
"Source pulse (2.86 ms)",
ha="left",
va="top",
fontsize=6)

# Now find frame boundaries and draw frames
frame_boundaries = []
Expand All @@ -55,8 +70,10 @@ def frames_analytical(instrument=None, plot=False):
# 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
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)

Expand All @@ -79,28 +96,46 @@ def frames_analytical(instrument=None, plot=False):
y5 = info["detector_position"]

# This is the frame boundaries
x5 = (y5 - b1)/a1
x4 = (y4 - b2)/a2
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)
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")

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")
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")
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]")
Expand All @@ -116,12 +151,16 @@ def frames_analytical(instrument=None, plot=False):
# 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)}

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
92 changes: 57 additions & 35 deletions src/dress/wfm/frames_peakfinding.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@


def _tof_shifts(pscdata, psc_frequency=0):
cut_out_centre = np.reshape(pscdata,(len(pscdata)//2, 2)).mean(1)
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)
return cut_out_diffs / (360.0 * psc_frequency)


def _get_frame_shifts(instrument):
Expand All @@ -23,12 +23,11 @@ def _get_frame_shifts(instrument):
"""

# 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
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


Expand All @@ -50,26 +49,35 @@ def _make_frame_shifts(initial_shift, other_shifts=None):
: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))]
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,
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))

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"])
Expand Down Expand Up @@ -98,7 +106,6 @@ def frames_peakfinding(data=None, instrument=None, initial_shift=-6630,
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)
Expand All @@ -113,34 +120,38 @@ def frames_peakfinding(data=None, instrument=None, initial_shift=-6630,
# 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
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):
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)
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]]))
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))
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)
Expand All @@ -156,7 +167,7 @@ def frames_peakfinding(data=None, instrument=None, initial_shift=-6630,
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]
frame_boundaries[i + 1, 0] = g + 0.5 * inter_frame_gaps[i]

# Plot the diagnostics
if plot:
Expand All @@ -167,13 +178,22 @@ def frames_peakfinding(data=None, instrument=None, initial_shift=-6630,
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.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.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):
Expand All @@ -182,9 +202,11 @@ def frames_peakfinding(data=None, instrument=None, initial_shift=-6630,
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)}
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
Loading

0 comments on commit 609c3a8

Please sign in to comment.