From b56da52a612eaa3c3142ce20ae7137d214a6e4df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20S=C3=A1nchez-Puga?= Date: Wed, 18 Mar 2026 11:23:57 +0100 Subject: [PATCH 1/6] Create sss.py --- src/scippneutron/period_slice/sss.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 src/scippneutron/period_slice/sss.py diff --git a/src/scippneutron/period_slice/sss.py b/src/scippneutron/period_slice/sss.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/src/scippneutron/period_slice/sss.py @@ -0,0 +1 @@ + From b81077b87938eb24b5500b9c9926f3f0819db0f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20S=C3=A1nchez-Puga?= Date: Wed, 18 Mar 2026 11:25:21 +0100 Subject: [PATCH 2/6] Add files via upload --- .../period_slice/Rheo_reduction_Scipp_LR.py | 336 ++++++++++++++++++ 1 file changed, 336 insertions(+) create mode 100644 src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py diff --git a/src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py b/src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py new file mode 100644 index 000000000..dee87b447 --- /dev/null +++ b/src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py @@ -0,0 +1,336 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import scipp as sc +import numpy as np +import scipy as sp +import matplotlib.pyplot as plt +from scipy.special import lambertw + +def cosine_wave(time, A, T, phi, Amean): + return A * sc.cos(sc.scalar(2 * np.pi, unit='rad') * time / T + phi) + Amean + +def sine_wave(time, **par): + A = par['A'] + T = par['T'] + Toff = par['Toff'] + Aoff = par['Aoff'] + + return A * sc.sin(sc.scalar(2 * np.pi, unit='rad') * (time - Toff) / T) + Aoff + +def ss_fft(signal): + + diff = (signal.coords['time'][1:] - signal.coords['time'][:-1]) + mean_diff = sc.mean(diff) + dt = mean_diff.values / 1e9 + + Fs = 1 / dt + L = len(signal) + + Y = np.fft.rfft(signal.values, n=L) + P = np.abs(Y) / L + P[1:-1] = 2 * P[1:-1] + ph = np.angle(Y) + + f = Fs * np.arange(0, L // 2 + 1) / L + + return P, ph, f + +def rheo_waveform_params(rheo_binned, method): + + if method == 'curve_fit': + rheo_for_fit = rheo_binned.copy() + # set time origin to 0 + rheo_for_fit.coords['time'] = rheo_for_fit.coords['time'] - rheo_for_fit.coords['time'][0] + # rheo_for_fit.coords['time'] = rheo_for_fit.coords['time'].to(unit='s') + + # Get a guess for period using DFT + P, ph, f = ss_fft(rheo_for_fit) + maxindex = np.argmax(P[1:]) + 1 + T_0 = sc.scalar(1e9 / f[maxindex], unit='ns') + A_0 = ((sc.max(rheo_for_fit) - sc.min(rheo_for_fit)) / sc.scalar(2, unit='')).data + Aoff_0 = sc.mean(rheo_for_fit).data + mask_first_p = rheo_for_fit.coords['time'] < T_0 + Toff_0 = rheo_for_fit.coords['time'][mask_first_p][(rheo_for_fit[mask_first_p] == sc.max(rheo_for_fit[mask_first_p])).data][0] - T_0 / 4 + + print({'A': A_0, 'T': T_0, 'Toff': Toff_0, 'Aoff': Aoff_0}) + par, _ = sc.curve_fit(['time'], sine_wave, rheo_for_fit, p0 = {'A': A_0, 'T': T_0, 'Toff': Toff_0, 'Aoff': Aoff_0}) + par['T'] = par['T'].to(unit='ns') + par['Toff'] = par['Toff'].to(unit='ns') + if par['A'].values < 0: + par['A'] = -par['A'] + par['Toff'] = par['Toff'] + par['T'] / 2 + + elif method == 'fft': + P, ph, f = ss_fft(rheo_binned) + maxindex = np.argmax(P[1:]) + 1 + + A = P[maxindex] + T = 1 / f[maxindex] + Aoff = sc.mean(rheo_binned).values + phi = ph[maxindex] + Toff = -T * (phi + np.pi / 2) / (2 * np.pi) + + par = sc.DataGroup( + A = sc.DataArray(data=sc.Variable(dims=[], values=A, unit='')), + T = sc.DataArray(data=sc.Variable(dims=[], values=T * 1e9, unit='ns')), + Toff = sc.DataArray(data=sc.Variable(dims=[], values=Toff * 1e9, unit='ns')), + Aoff = sc.DataArray(data=sc.Variable(dims=[], values=Aoff, unit='')) + ) + + else: + print('Method not supported') + par = sc.DataGroup( + A = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='')), + T = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='ns')), + Toff = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='ns')), + Aoff = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='')) + ) + + return par + +def period_edges(signal, plot_flag, value=None): + + if value == None: + signal_zeroed = signal - signal[0] + else: + signal_zeroed = signal - value + + if signal_zeroed.values[0] > signal_zeroed.values[1]: + # Falling edges + period_edges = np.where(np.diff(np.sign(signal_zeroed.values)) < 0)[0] + else: + # Rising edges + period_edges = np.where(np.diff(np.sign(signal_zeroed.values)) > 0)[0] + + if plot_flag: + plt.figure() + plt.plot(np.diff(period_edges), label='', linestyle='', marker='o', markerfacecolor='none') + plt.xlabel('Samples') + plt.ylabel('diff(Edges)') + plt.legend() + plt.show() + + return period_edges + +def ppT_autocorrelation(signal, plot_flag=None): + + dt = np.mean(np.diff(signal.coords['time'].values)) / 1e9 + signal = signal.values + N = len(signal) + # zero-padding to the next power of 2 + Npad = 2 ** int(np.ceil(np.log2(2 * N))) + signal_padded = np.zeros(Npad) + signal_padded[:N] = signal + X = np.fft.rfft(signal_padded) + P = X * np.conj(X) + R = np.fft.irfft(P) + R = np.real(R[:N]) + # R_full = np.correlate(signal, signal, mode='full') + # N = len(signal) + # R = R_full[N - 1:] + # R = R / R[0] + + locs, _ = sp.signal.find_peaks(R[1:]) + locs = locs + 1 + + if len(locs) == 0: + ppT = np.nan + print("Warning: Peak not found") + else: + ppT = locs[0] + + # Parabolic interpolation + if ppT > 0 and ppT < len(R) - 1: + y1, y2, y3 = R[ppT-1], R[ppT], R[ppT + 1] + delta = 0.5 * (y1 - y3) / (y1 - 2 * y2 + y3) + ppT = ppT + delta + + freq = 1 / ppT / dt + + if plot_flag: + R = R / R[0] + plt.figure() + plt.plot(R, label='Autocorrelation') + plt.plot(locs, R[locs], marker='o', linestyle='None', markeredgecolor='r', markerfacecolor='none', label='Peaks') + plt.xlabel('Samples') + plt.ylabel('Amplitude') + plt.legend() + plt.show() + + return freq, ppT + +def ppT_fft(signal, plot_flag=None): + + P, ph, f = ss_fft(signal) + max_index = np.argmax(P[1:]) + 1 + + freq = f[max_index] + dt = np.mean(np.diff(signal.coords['time'].values)) / 1e9 + ppT = 1 / freq / dt + + if plot_flag: + plt.figure() + plt.plot(f, P) + plt.xlabel('f (Hz)') + plt.ylabel('|P1(f)|') + plt.title('Single-Sided Amplitude Spectrum') + plt.grid(True) + plt.yscale('log') + plt.xscale('log') + # plt.xlim(f[0], f[-1]) + # plt.ylim(1e-10, 1e-2) + plt.show() + + return freq, ppT + +def ppT_curve_fit(signal): + + par = rheo_waveform_params(signal, "curve_fit") + + dt = np.mean(np.diff(signal.coords['time'].values)) / 1e9 + freq = (1e9 / par['T'].values) + ppT = 1 / freq / dt + + return freq, ppT + +def plot_rheo_periods(rheo, rheo_fitted, N_period, plot=None, periods=None): + """ + rheo: rheometer scipp data array + t_period: period duration + periods: tupla (p_start, p_end) + """ + time_rheo = rheo.coords['time'] + time_fft = rheo_fitted.coords['time'] + + if plot in ['plot', 'overplot']: + if plot == 'plot': + plt.figure() + if periods is not None: + idx_start = (periods[0] - 1) * N_period + idx_end = periods[1] * N_period + plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo[idx_start:idx_end].values, marker='.', linestyle='', color='k', markersize=3) + plt.plot(time_fft[idx_start:idx_end].values - time_fft[idx_start].values, rheo_fitted[idx_start:idx_end].values, marker='', linestyle='-', color='r', markersize=1) + else: + plt.plot(time_rheo.values, rheo.values) + plt.plot(time_fft.values, rheo_fitted.values) + plt.xlabel('Time [ns]') + plt.ylabel('Rheometer Deflection Angle') + # plt.legend() + plt.show() + if plot == 'overplot': + plt.figure() + if periods is not None: + for i in range(periods[0] - 1, periods[1]): + idx_start = i * N_period + idx_end = (i + 1) * N_period + plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo[idx_start:idx_end].values, marker='.', linestyle='', color='k', markersize=1) + plt.plot(time_fft[idx_start:idx_end].values - time_fft[idx_start].values, rheo_fitted[idx_start:idx_end].values, marker='', linestyle='-', color='r', markersize=1) + else: + for i in range(len(rheo) // N_period): + idx_start = i * N_period + idx_end = (i + 1) * N_period + plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo[idx_start:idx_end].values, marker='.', linestyle='', color='k', markersize=1) + plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo_fitted[idx_start:idx_end].values, marker='', linestyle='-', color='r', markersize=1) + plt.xlabel('Time [ns]') + plt.ylabel('Rheometer Deflection Angle') + # plt.legend() + plt.show() + +def time_at_sample(event_time_zero, event_time_offset, distance_to_moderator, distance_to_sample): + return ( + # Think on this + event_time_offset.to(unit=event_time_zero.unit) * distance_to_sample / distance_to_moderator + event_time_zero + # event_time_offset * distance_to_sample / distance_to_moderator + event_time_zero.to(dtype=event_time_offset.dtype, unit=event_time_offset.unit) + # event_time_offset * distance_to_sample / distance_to_moderator + event_time_zero.to(unit=event_time_offset.unit) + ) + +def period_slice(time_at_sample, T, Toff, N_slices): + first_edge = (Toff % T - T / 4 - T / N_slices / 2).to(unit=time_at_sample.unit) + # print(first_edge) + return ( + ((N_slices * ((time_at_sample - first_edge) / (T.to(unit=time_at_sample.unit)))) % N_slices).to(dtype='int32') + # ((N_slices * ((time_at_sample) / (T.to(unit=time_at_sample.unit)))) % N_slices).to(dtype='int32') + ) + +def period_slice_trigger(time_at_sample, period_times, N_slices): + + func = sc.DataArray(period_times[1:] - period_times[:-1], coords={'time_at_sample': period_times}) + lookUpTable = sc.lookup(func, 'time_at_sample') + period_duration = lookUpTable[time_at_sample] + + func = sc.DataArray(period_times[:-1], coords={'time_at_sample': period_times}) + lookUpTable = sc.lookup(func, 'time_at_sample') + period_start = lookUpTable[time_at_sample] + + return ( + ((N_slices * ((time_at_sample - period_start) / period_duration)) % N_slices).to(dtype='int32') + ) + + +def dead_time_correction(events, bad_events, proton_charge, tdead, TOFmin, TOFmax, TOFbin): + import matplotlib.pyplot as plt + + + # Grid on TOF placed at the edges of the bins + tof_bin_edges = sc.arange('event_time_offset', TOFmin - TOFbin / 2, TOFmax + TOFbin / 2, TOFbin, unit='us', dtype='float64') + # tof_bins = sc.arange('event_time_offset', TOFmin, TOFmax, TOFbin, unit='us', dtype='float64') + + # Histogram the event_time_offset + mask = (proton_charge != sc.scalar(0, unit='pC')).data + CT = events.hist(event_time_offset=tof_bin_edges, dim=events.dims) + # CT = events[mask.rename_dims(time = 'event_time_zero')].hist(event_time_offset=tof_bin_edges, dim=events.dims) + BCT = bad_events.hist(event_time_offset=tof_bin_edges, dim=events.dims) + # BCT = bad_events[mask.rename_dims(time = 'event_time_zero')].hist(event_time_offset=tof_bin_edges, dim=events.dims) + CT = CT + BCT + CT = CT.values + CT = CT / len(proton_charge[mask]) + + # if useNP == 0: # (paralyzable) + b = -lambertw(-CT * tdead / TOFbin) / (tdead / TOFbin) + dtc = b / CT + dtc = dtc.real + dtc = np.nan_to_num(dtc, nan=1) + # if useNP == 1: dtc = 1 / (1 - CT * tdead / tofbin) # (non-paralyzable) + + # plt.plot(dtc) + # plt.show() + + return sc.DataArray(sc.array(dims=['event_time_offset'], values=dtc), coords={'event_time_offset': tof_bin_edges}) + + #--------------------------------------------------------------------------- + +def gravity_correct(LAM, ThetaIn, dSamp, dSlit): + + dSamp=dSamp/1000 #dSamp is m from sample to incident slit + dSlit=dSlit/1000 #dSlit is m between slits + + #calculation from the ILL paper. this works for inclined beams. + g=9.8067 #m/s^2 + h=6.6260715e-34 #Js=kg m^2/s + mn=1.67492749804e-27 #kg + V=h/(mn*LAM*1e-10) + K=g/(2*V**2) + + #define the sample position as x=0, y=0. increasing x is towards moderator + xs=0 + ys=0 + + #positions of slits + x1=dSamp + x2=(dSamp+dSlit) + + #height of slits determined by incident theta, y=0 is the sample height + y1=x1*np.tan(ThetaIn*np.pi/180) + y2=x2*np.tan(ThetaIn*np.pi/180) + + #this is the location of the top of the parabola + x0=(y1-y2+K*(x1**2-x2**2))/(2*K*(x1-x2)) + y0=y2+K*(x2-x0)**2 + xs=x0-np.sqrt(y0/K) + ThetaSam=np.arctan(2*K*(x0-xs))*180/np.pi #angle is arctan(dy/dx) at sample + dTheta=ThetaSam-ThetaIn + + return dTheta + +#----------------------------------------------------------------------------- \ No newline at end of file From f561a2f56bbe31e0faaf2e09745caee8abae3dab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20S=C3=A1nchez-Puga?= Date: Wed, 18 Mar 2026 11:27:10 +0100 Subject: [PATCH 3/6] Remove empty sss.py file From 4a22ca193077950378cac7ac20b5a25d0dc5a362 Mon Sep 17 00:00:00 2001 From: psanchez0046 Date: Wed, 18 Mar 2026 11:29:29 +0100 Subject: [PATCH 4/6] Delete sss.py --- src/scippneutron/period_slice/sss.py | 1 - 1 file changed, 1 deletion(-) delete mode 100644 src/scippneutron/period_slice/sss.py diff --git a/src/scippneutron/period_slice/sss.py b/src/scippneutron/period_slice/sss.py deleted file mode 100644 index 8b1378917..000000000 --- a/src/scippneutron/period_slice/sss.py +++ /dev/null @@ -1 +0,0 @@ - From 142ff42df2d941a97e1f23b2442c473fd6afc174 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20S=C3=A1nchez-Puga?= Date: Wed, 18 Mar 2026 11:38:15 +0100 Subject: [PATCH 5/6] Rename Rheo_reduction_Scipp_LR.py to Rheo_reduction_Scipp.py --- .../{Rheo_reduction_Scipp_LR.py => Rheo_reduction_Scipp.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename src/scippneutron/period_slice/{Rheo_reduction_Scipp_LR.py => Rheo_reduction_Scipp.py} (99%) diff --git a/src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py b/src/scippneutron/period_slice/Rheo_reduction_Scipp.py similarity index 99% rename from src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py rename to src/scippneutron/period_slice/Rheo_reduction_Scipp.py index dee87b447..10f20346d 100644 --- a/src/scippneutron/period_slice/Rheo_reduction_Scipp_LR.py +++ b/src/scippneutron/period_slice/Rheo_reduction_Scipp.py @@ -333,4 +333,4 @@ def gravity_correct(LAM, ThetaIn, dSamp, dSlit): return dTheta -#----------------------------------------------------------------------------- \ No newline at end of file +#----------------------------------------------------------------------------- From 93ca552cf3592c50a4f7d873eb5ad7ae48cdb0cd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci-lite[bot]" <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com> Date: Thu, 19 Mar 2026 14:53:05 +0000 Subject: [PATCH 6/6] Apply automatic formatting --- .../period_slice/Rheo_reduction_Scipp.py | 267 ++++++++++++------ 1 file changed, 186 insertions(+), 81 deletions(-) diff --git a/src/scippneutron/period_slice/Rheo_reduction_Scipp.py b/src/scippneutron/period_slice/Rheo_reduction_Scipp.py index 10f20346d..3639e7609 100644 --- a/src/scippneutron/period_slice/Rheo_reduction_Scipp.py +++ b/src/scippneutron/period_slice/Rheo_reduction_Scipp.py @@ -1,15 +1,17 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -import scipp as sc +import matplotlib.pyplot as plt import numpy as np +import scipp as sc import scipy as sp -import matplotlib.pyplot as plt -from scipy.special import lambertw +from scipy.special import lambertw + def cosine_wave(time, A, T, phi, Amean): return A * sc.cos(sc.scalar(2 * np.pi, unit='rad') * time / T + phi) + Amean + def sine_wave(time, **par): A = par['A'] T = par['T'] @@ -18,13 +20,13 @@ def sine_wave(time, **par): return A * sc.sin(sc.scalar(2 * np.pi, unit='rad') * (time - Toff) / T) + Aoff -def ss_fft(signal): - diff = (signal.coords['time'][1:] - signal.coords['time'][:-1]) +def ss_fft(signal): + diff = signal.coords['time'][1:] - signal.coords['time'][:-1] mean_diff = sc.mean(diff) dt = mean_diff.values / 1e9 - Fs = 1 / dt + Fs = 1 / dt L = len(signal) Y = np.fft.rfft(signal.values, n=L) @@ -36,25 +38,39 @@ def ss_fft(signal): return P, ph, f -def rheo_waveform_params(rheo_binned, method): +def rheo_waveform_params(rheo_binned, method): if method == 'curve_fit': rheo_for_fit = rheo_binned.copy() # set time origin to 0 - rheo_for_fit.coords['time'] = rheo_for_fit.coords['time'] - rheo_for_fit.coords['time'][0] + rheo_for_fit.coords['time'] = ( + rheo_for_fit.coords['time'] - rheo_for_fit.coords['time'][0] + ) # rheo_for_fit.coords['time'] = rheo_for_fit.coords['time'].to(unit='s') # Get a guess for period using DFT P, ph, f = ss_fft(rheo_for_fit) maxindex = np.argmax(P[1:]) + 1 T_0 = sc.scalar(1e9 / f[maxindex], unit='ns') - A_0 = ((sc.max(rheo_for_fit) - sc.min(rheo_for_fit)) / sc.scalar(2, unit='')).data + A_0 = ( + (sc.max(rheo_for_fit) - sc.min(rheo_for_fit)) / sc.scalar(2, unit='') + ).data Aoff_0 = sc.mean(rheo_for_fit).data mask_first_p = rheo_for_fit.coords['time'] < T_0 - Toff_0 = rheo_for_fit.coords['time'][mask_first_p][(rheo_for_fit[mask_first_p] == sc.max(rheo_for_fit[mask_first_p])).data][0] - T_0 / 4 + Toff_0 = ( + rheo_for_fit.coords['time'][mask_first_p][ + (rheo_for_fit[mask_first_p] == sc.max(rheo_for_fit[mask_first_p])).data + ][0] + - T_0 / 4 + ) print({'A': A_0, 'T': T_0, 'Toff': Toff_0, 'Aoff': Aoff_0}) - par, _ = sc.curve_fit(['time'], sine_wave, rheo_for_fit, p0 = {'A': A_0, 'T': T_0, 'Toff': Toff_0, 'Aoff': Aoff_0}) + par, _ = sc.curve_fit( + ['time'], + sine_wave, + rheo_for_fit, + p0={'A': A_0, 'T': T_0, 'Toff': Toff_0, 'Aoff': Aoff_0}, + ) par['T'] = par['T'].to(unit='ns') par['Toff'] = par['Toff'].to(unit='ns') if par['A'].values < 0: @@ -64,7 +80,7 @@ def rheo_waveform_params(rheo_binned, method): elif method == 'fft': P, ph, f = ss_fft(rheo_binned) maxindex = np.argmax(P[1:]) + 1 - + A = P[maxindex] T = 1 / f[maxindex] Aoff = sc.mean(rheo_binned).values @@ -72,30 +88,30 @@ def rheo_waveform_params(rheo_binned, method): Toff = -T * (phi + np.pi / 2) / (2 * np.pi) par = sc.DataGroup( - A = sc.DataArray(data=sc.Variable(dims=[], values=A, unit='')), - T = sc.DataArray(data=sc.Variable(dims=[], values=T * 1e9, unit='ns')), - Toff = sc.DataArray(data=sc.Variable(dims=[], values=Toff * 1e9, unit='ns')), - Aoff = sc.DataArray(data=sc.Variable(dims=[], values=Aoff, unit='')) + A=sc.DataArray(data=sc.Variable(dims=[], values=A, unit='')), + T=sc.DataArray(data=sc.Variable(dims=[], values=T * 1e9, unit='ns')), + Toff=sc.DataArray(data=sc.Variable(dims=[], values=Toff * 1e9, unit='ns')), + Aoff=sc.DataArray(data=sc.Variable(dims=[], values=Aoff, unit='')), ) else: print('Method not supported') par = sc.DataGroup( - A = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='')), - T = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='ns')), - Toff = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='ns')), - Aoff = sc.DataArray(data=sc.Variable(dims=[], values=0, unit='')) + A=sc.DataArray(data=sc.Variable(dims=[], values=0, unit='')), + T=sc.DataArray(data=sc.Variable(dims=[], values=0, unit='ns')), + Toff=sc.DataArray(data=sc.Variable(dims=[], values=0, unit='ns')), + Aoff=sc.DataArray(data=sc.Variable(dims=[], values=0, unit='')), ) return par - -def period_edges(signal, plot_flag, value=None): + +def period_edges(signal, plot_flag, value=None): if value == None: signal_zeroed = signal - signal[0] - else: + else: signal_zeroed = signal - value - + if signal_zeroed.values[0] > signal_zeroed.values[1]: # Falling edges period_edges = np.where(np.diff(np.sign(signal_zeroed.values)) < 0)[0] @@ -105,7 +121,13 @@ def period_edges(signal, plot_flag, value=None): if plot_flag: plt.figure() - plt.plot(np.diff(period_edges), label='', linestyle='', marker='o', markerfacecolor='none') + plt.plot( + np.diff(period_edges), + label='', + linestyle='', + marker='o', + markerfacecolor='none', + ) plt.xlabel('Samples') plt.ylabel('diff(Edges)') plt.legend() @@ -113,8 +135,8 @@ def period_edges(signal, plot_flag, value=None): return period_edges -def ppT_autocorrelation(signal, plot_flag=None): +def ppT_autocorrelation(signal, plot_flag=None): dt = np.mean(np.diff(signal.coords['time'].values)) / 1e9 signal = signal.values N = len(signal) @@ -142,7 +164,7 @@ def ppT_autocorrelation(signal, plot_flag=None): # Parabolic interpolation if ppT > 0 and ppT < len(R) - 1: - y1, y2, y3 = R[ppT-1], R[ppT], R[ppT + 1] + y1, y2, y3 = R[ppT - 1], R[ppT], R[ppT + 1] delta = 0.5 * (y1 - y3) / (y1 - 2 * y2 + y3) ppT = ppT + delta @@ -152,7 +174,15 @@ def ppT_autocorrelation(signal, plot_flag=None): R = R / R[0] plt.figure() plt.plot(R, label='Autocorrelation') - plt.plot(locs, R[locs], marker='o', linestyle='None', markeredgecolor='r', markerfacecolor='none', label='Peaks') + plt.plot( + locs, + R[locs], + marker='o', + linestyle='None', + markeredgecolor='r', + markerfacecolor='none', + label='Peaks', + ) plt.xlabel('Samples') plt.ylabel('Amplitude') plt.legend() @@ -160,8 +190,8 @@ def ppT_autocorrelation(signal, plot_flag=None): return freq, ppT -def ppT_fft(signal, plot_flag=None): +def ppT_fft(signal, plot_flag=None): P, ph, f = ss_fft(signal) max_index = np.argmax(P[1:]) + 1 @@ -184,15 +214,16 @@ def ppT_fft(signal, plot_flag=None): return freq, ppT -def ppT_curve_fit(signal): +def ppT_curve_fit(signal): par = rheo_waveform_params(signal, "curve_fit") dt = np.mean(np.diff(signal.coords['time'].values)) / 1e9 - freq = (1e9 / par['T'].values) + freq = 1e9 / par['T'].values ppT = 1 / freq / dt - return freq, ppT + return freq, ppT + def plot_rheo_periods(rheo, rheo_fitted, N_period, plot=None, periods=None): """ @@ -201,16 +232,30 @@ def plot_rheo_periods(rheo, rheo_fitted, N_period, plot=None, periods=None): periods: tupla (p_start, p_end) """ time_rheo = rheo.coords['time'] - time_fft = rheo_fitted.coords['time'] + time_fft = rheo_fitted.coords['time'] - if plot in ['plot', 'overplot']: + if plot in ['plot', 'overplot']: if plot == 'plot': plt.figure() if periods is not None: idx_start = (periods[0] - 1) * N_period idx_end = periods[1] * N_period - plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo[idx_start:idx_end].values, marker='.', linestyle='', color='k', markersize=3) - plt.plot(time_fft[idx_start:idx_end].values - time_fft[idx_start].values, rheo_fitted[idx_start:idx_end].values, marker='', linestyle='-', color='r', markersize=1) + plt.plot( + time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, + rheo[idx_start:idx_end].values, + marker='.', + linestyle='', + color='k', + markersize=3, + ) + plt.plot( + time_fft[idx_start:idx_end].values - time_fft[idx_start].values, + rheo_fitted[idx_start:idx_end].values, + marker='', + linestyle='-', + color='r', + markersize=1, + ) else: plt.plot(time_rheo.values, rheo.values) plt.plot(time_fft.values, rheo_fitted.values) @@ -224,38 +269,84 @@ def plot_rheo_periods(rheo, rheo_fitted, N_period, plot=None, periods=None): for i in range(periods[0] - 1, periods[1]): idx_start = i * N_period idx_end = (i + 1) * N_period - plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo[idx_start:idx_end].values, marker='.', linestyle='', color='k', markersize=1) - plt.plot(time_fft[idx_start:idx_end].values - time_fft[idx_start].values, rheo_fitted[idx_start:idx_end].values, marker='', linestyle='-', color='r', markersize=1) + plt.plot( + time_rheo[idx_start:idx_end].values + - time_rheo[idx_start].values, + rheo[idx_start:idx_end].values, + marker='.', + linestyle='', + color='k', + markersize=1, + ) + plt.plot( + time_fft[idx_start:idx_end].values - time_fft[idx_start].values, + rheo_fitted[idx_start:idx_end].values, + marker='', + linestyle='-', + color='r', + markersize=1, + ) else: for i in range(len(rheo) // N_period): idx_start = i * N_period idx_end = (i + 1) * N_period - plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo[idx_start:idx_end].values, marker='.', linestyle='', color='k', markersize=1) - plt.plot(time_rheo[idx_start:idx_end].values - time_rheo[idx_start].values, rheo_fitted[idx_start:idx_end].values, marker='', linestyle='-', color='r', markersize=1) + plt.plot( + time_rheo[idx_start:idx_end].values + - time_rheo[idx_start].values, + rheo[idx_start:idx_end].values, + marker='.', + linestyle='', + color='k', + markersize=1, + ) + plt.plot( + time_rheo[idx_start:idx_end].values + - time_rheo[idx_start].values, + rheo_fitted[idx_start:idx_end].values, + marker='', + linestyle='-', + color='r', + markersize=1, + ) plt.xlabel('Time [ns]') plt.ylabel('Rheometer Deflection Angle') # plt.legend() plt.show() -def time_at_sample(event_time_zero, event_time_offset, distance_to_moderator, distance_to_sample): + +def time_at_sample( + event_time_zero, event_time_offset, distance_to_moderator, distance_to_sample +): return ( # Think on this - event_time_offset.to(unit=event_time_zero.unit) * distance_to_sample / distance_to_moderator + event_time_zero + event_time_offset.to(unit=event_time_zero.unit) + * distance_to_sample + / distance_to_moderator + + event_time_zero # event_time_offset * distance_to_sample / distance_to_moderator + event_time_zero.to(dtype=event_time_offset.dtype, unit=event_time_offset.unit) # event_time_offset * distance_to_sample / distance_to_moderator + event_time_zero.to(unit=event_time_offset.unit) ) + def period_slice(time_at_sample, T, Toff, N_slices): first_edge = (Toff % T - T / 4 - T / N_slices / 2).to(unit=time_at_sample.unit) # print(first_edge) return ( - ((N_slices * ((time_at_sample - first_edge) / (T.to(unit=time_at_sample.unit)))) % N_slices).to(dtype='int32') + ( + ( + N_slices + * ((time_at_sample - first_edge) / (T.to(unit=time_at_sample.unit))) + ) + % N_slices + ).to(dtype='int32') # ((N_slices * ((time_at_sample) / (T.to(unit=time_at_sample.unit)))) % N_slices).to(dtype='int32') ) + def period_slice_trigger(time_at_sample, period_times, N_slices): - - func = sc.DataArray(period_times[1:] - period_times[:-1], coords={'time_at_sample': period_times}) + func = sc.DataArray( + period_times[1:] - period_times[:-1], coords={'time_at_sample': period_times} + ) lookUpTable = sc.lookup(func, 'time_at_sample') period_duration = lookUpTable[time_at_sample] @@ -264,24 +355,32 @@ def period_slice_trigger(time_at_sample, period_times, N_slices): period_start = lookUpTable[time_at_sample] return ( - ((N_slices * ((time_at_sample - period_start) / period_duration)) % N_slices).to(dtype='int32') - ) + (N_slices * ((time_at_sample - period_start) / period_duration)) % N_slices + ).to(dtype='int32') -def dead_time_correction(events, bad_events, proton_charge, tdead, TOFmin, TOFmax, TOFbin): - import matplotlib.pyplot as plt - +def dead_time_correction( + events, bad_events, proton_charge, tdead, TOFmin, TOFmax, TOFbin +): + import matplotlib.pyplot as plt # Grid on TOF placed at the edges of the bins - tof_bin_edges = sc.arange('event_time_offset', TOFmin - TOFbin / 2, TOFmax + TOFbin / 2, TOFbin, unit='us', dtype='float64') + tof_bin_edges = sc.arange( + 'event_time_offset', + TOFmin - TOFbin / 2, + TOFmax + TOFbin / 2, + TOFbin, + unit='us', + dtype='float64', + ) # tof_bins = sc.arange('event_time_offset', TOFmin, TOFmax, TOFbin, unit='us', dtype='float64') # Histogram the event_time_offset mask = (proton_charge != sc.scalar(0, unit='pC')).data CT = events.hist(event_time_offset=tof_bin_edges, dim=events.dims) # CT = events[mask.rename_dims(time = 'event_time_zero')].hist(event_time_offset=tof_bin_edges, dim=events.dims) - BCT = bad_events.hist(event_time_offset=tof_bin_edges, dim=events.dims) - # BCT = bad_events[mask.rename_dims(time = 'event_time_zero')].hist(event_time_offset=tof_bin_edges, dim=events.dims) + BCT = bad_events.hist(event_time_offset=tof_bin_edges, dim=events.dims) + # BCT = bad_events[mask.rename_dims(time = 'event_time_zero')].hist(event_time_offset=tof_bin_edges, dim=events.dims) CT = CT + BCT CT = CT.values CT = CT / len(proton_charge[mask]) @@ -290,47 +389,53 @@ def dead_time_correction(events, bad_events, proton_charge, tdead, TOFmin, TOFma b = -lambertw(-CT * tdead / TOFbin) / (tdead / TOFbin) dtc = b / CT dtc = dtc.real - dtc = np.nan_to_num(dtc, nan=1) + dtc = np.nan_to_num(dtc, nan=1) # if useNP == 1: dtc = 1 / (1 - CT * tdead / tofbin) # (non-paralyzable) # plt.plot(dtc) # plt.show() - return sc.DataArray(sc.array(dims=['event_time_offset'], values=dtc), coords={'event_time_offset': tof_bin_edges}) + return sc.DataArray( + sc.array(dims=['event_time_offset'], values=dtc), + coords={'event_time_offset': tof_bin_edges}, + ) - #--------------------------------------------------------------------------- + # --------------------------------------------------------------------------- -def gravity_correct(LAM, ThetaIn, dSamp, dSlit): - dSamp=dSamp/1000 #dSamp is m from sample to incident slit - dSlit=dSlit/1000 #dSlit is m between slits +def gravity_correct(LAM, ThetaIn, dSamp, dSlit): + dSamp = dSamp / 1000 # dSamp is m from sample to incident slit + dSlit = dSlit / 1000 # dSlit is m between slits - #calculation from the ILL paper. this works for inclined beams. - g=9.8067 #m/s^2 - h=6.6260715e-34 #Js=kg m^2/s - mn=1.67492749804e-27 #kg - V=h/(mn*LAM*1e-10) - K=g/(2*V**2) + # calculation from the ILL paper. this works for inclined beams. + g = 9.8067 # m/s^2 + h = 6.6260715e-34 # Js=kg m^2/s + mn = 1.67492749804e-27 # kg + V = h / (mn * LAM * 1e-10) + K = g / (2 * V**2) - #define the sample position as x=0, y=0. increasing x is towards moderator - xs=0 - ys=0 + # define the sample position as x=0, y=0. increasing x is towards moderator + xs = 0 + ys = 0 - #positions of slits - x1=dSamp - x2=(dSamp+dSlit) + # positions of slits + x1 = dSamp + x2 = dSamp + dSlit - #height of slits determined by incident theta, y=0 is the sample height - y1=x1*np.tan(ThetaIn*np.pi/180) - y2=x2*np.tan(ThetaIn*np.pi/180) + # height of slits determined by incident theta, y=0 is the sample height + y1 = x1 * np.tan(ThetaIn * np.pi / 180) + y2 = x2 * np.tan(ThetaIn * np.pi / 180) - #this is the location of the top of the parabola - x0=(y1-y2+K*(x1**2-x2**2))/(2*K*(x1-x2)) - y0=y2+K*(x2-x0)**2 - xs=x0-np.sqrt(y0/K) - ThetaSam=np.arctan(2*K*(x0-xs))*180/np.pi #angle is arctan(dy/dx) at sample - dTheta=ThetaSam-ThetaIn + # this is the location of the top of the parabola + x0 = (y1 - y2 + K * (x1**2 - x2**2)) / (2 * K * (x1 - x2)) + y0 = y2 + K * (x2 - x0) ** 2 + xs = x0 - np.sqrt(y0 / K) + ThetaSam = ( + np.arctan(2 * K * (x0 - xs)) * 180 / np.pi + ) # angle is arctan(dy/dx) at sample + dTheta = ThetaSam - ThetaIn return dTheta -#----------------------------------------------------------------------------- + +# -----------------------------------------------------------------------------