Merge pull request #253 from jasondet/refactor

Merge main into refactor

pygama/dsp/_processors/Wiener_filter.py

@@ -0,0 +1,112 @@
+import numpy as np
+from pygama import lh5
+from numba import guvectorize
+from pygama.dsp.errors import DSPFatal
+
+
+def Wiener_filter(file_name_array):
+    """
+    Apply a Wiener filter to the waveform.  Note that this convolution is performed in the frequency domain
+    
+    Initialization Parameters
+    -------------------------
+    file_name_array : string
+            Array with path to an lh5 file containing the time domain version of the superpulse in one column and noise waveform in another,
+            the superpulse group must be titled 'spms/processed/superpulse' and the noise waveform must be called 'spms/processed/noise_wf'
+
+    Parameters
+    ----------
+    fft_w_in : array-like
+           The fourier transform input waveform
+    fft_w_out: array-like
+           The filtered waveform, in the frequency domain     
+           
+    Processing Chain Example
+    ------------------------
+    "wf_wiener": {
+        "function": "Wiener_filter",
+        "module": "pygama.dsp.processors",
+        "args": ["wf_bl_fft", "wf_wiener(2000,f)"],
+        "unit": "dB",
+        "prereqs": ["wf_bl_fft"],
+        "init_args": ["/path/to/file/wiener.lh5"]
+    }
+    """
+
+    sto = lh5.Store()
+
+    # Check that the file is valid and the data is in the correct format
+
+    try: file_name_array[0]
+    except: raise DSPFatal('init_args must be an array with the filename')
+
+    file_name = file_name_array[0]
+
+    try: f = sto.gimme_file(file_name, 'r')
+    except: raise DSPFatal('File must be a valid lh5 file') 
+
+    if 'spms/processed/superpulse' not in f:
+        raise DSPFatal('lh5 file must have \'spms/processed/superpulse\' as a group')
+
+    if 'spms/processed/noise_wf' not in f:
+        raise DSPFatal('lh5 file must have \'spms/processed/noise_wf\' as a group')
+
+    # Read in the data 
+
+    superpulse, _ = sto.read_object('spms/processed/superpulse', file_name)
+    superpulse = superpulse.nda
+
+    noise_wf, _ = sto.read_object('spms/processed/noise_wf', file_name)
+    noise_wf = noise_wf.nda
+
+    # Now check that the data are valid
+
+    if len(superpulse) <= 0:
+        raise DSPFatal('The length of the filter must be positive')
+
+    if len(superpulse) != len(noise_wf): 
+        raise DSPFatal('The length of the superpulse must be equal to the length of the noise waveform')
+
+    if np.argmax(superpulse) <= 0 or np.argmax(superpulse) > len(superpulse):
+        raise DSPFatal('The index of the maximum of the superpulse must occur within the waveform')
+
+    # Transform these to the frequency domain to eventually create the Wiener filter 
+
+    fft_superpulse = np.fft.fft(superpulse)
+    fft_noise_wf = np.fft.fft(noise_wf)
+
+    # Create the point spread function for the detector's response 
+
+    def PSF(superpulse, fft_superpulse):  
+
+        delta = np.zeros_like(superpulse)
+        arg_max = np.argmax(superpulse)
+        delta[arg_max] = np.amax(superpulse)
+
+        return fft_superpulse/np.fft.fft(delta)
+
+    # Now create the Wiener filter in the frequency domain 
+
+    fft_PSF = PSF(superpulse, fft_superpulse)
+    PSD_noise_wf = fft_noise_wf*np.conj(fft_noise_wf)
+    PSD_superpulse = fft_superpulse*np.conj(fft_superpulse)
+
+    w_filter = (np.conj(fft_PSF))/((fft_PSF*np.conj(fft_PSF))+(PSD_noise_wf/PSD_superpulse)) 
+
+    # Create a factory function that performs the convolution with the Wiener filter, the output is still in the frequency domain 
+
+    @guvectorize(["void(complex64[:], complex64[:])",
+                  "void(complex128[:], complex128[:])"],
+                 "(n)->(n)", forceobj=True)
+    def Wiener_out(fft_w_in, fft_w_out):
+        fft_w_out[:] = np.nan
+
+        if np.isnan(fft_w_in).any():
+            return
+
+        if len(w_filter) != len(fft_w_in):
+            raise DSPFatal('The filter is not the same length of the input waveform')
+
+        fft_w_out[:] = fft_w_in * w_filter
+
+    return Wiener_out

pygama/dsp/_processors/pulse_injector.py

@@ -0,0 +1,45 @@
+import numpy as np
+from numba import guvectorize
+from math import log,exp
+
+@guvectorize(["void(float32[:], float32, float32, float32, float32, float32[:])",
+              "void(float64[:], float64, float64,float64,float64,float64[:])"],
+              "(n),(),(),(), ()->(n)", nopython=True, cache=True)
+
+def inject_sig_pulse(wf_in,t0,rt,A,decay, wf_out):
+    """
+    Inject sigmoid pulse into existing waveform to simulate pileup
+    """
+
+
+    w_out[:] = np.nan
+
+    if np.isnan(w_in).any() or np.isnan(rt) or np.isnan(t0) or np.isnan(A) or np.isnan(decay):
+        return
+
+    rise = 4*log(99)/rt
+
+    wf_out[:] = wf_in[:]
+    for T in range(len(wf_out)):
+        wf_out[T] = wf_out[T]+ (A/(1+exp((-rise)*(T-(t0+rt/2))))*exp(-(1/decay)*(T-t0)))
+
+@guvectorize(["void(float32[:], float32, float32, float32, float32, float32[:])",
+              "void(float64[:], float64, float64,float64,float64,float64[:])"],
+              "(n),(),(),(), ()->(n)", nopython=True, cache=True)
+
+def inject_exp_pulse(wf_in,t0,rt,A,decay, wf_out):
+    """
+    Inject exponential pulse into existing waveform to simulate pileup
+    """
+
+    w_out[:] = np.nan
+
+    if np.isnan(w_in).any() or np.isnan(rt) or np.isnan(t0) or np.isnan(A) or np.isnan(decay):
+        return
+
+    wf_out[:] = wf_in[:]
+    for T in range(len(wf_out)):
+        if (t0<= T)& (T<= t0+rt):
+            wf_out[T] += ((A*exp((T-t0-rt)/(rise))*exp(-(1/decay)*(T-t0))))
+        elif (T>t0+rt):
+            wf_out[T] += (A*exp(-(1/decay)*(T-t0)))

pygama/dsp/processors.py

@@ -80,3 +80,5 @@
 from ._processors.get_multi_local_extrema import get_multi_local_extrema
 from ._processors.multi_t_filter import multi_t_filter, remove_duplicates
 from ._processors.multi_a_filter import multi_a_filter
+from ._processors.Wiener_filter import Wiener_filter
+from ._processors.pulse_injector import inject_sig_pulse,inject_exp_pulse

pygama/math/peak_fitting.py

@@ -333,7 +333,7 @@ def gauss_mode_width_max(hist, bins, var=None, mode_guess=None, n_bins=5,
         return None, None
     if pars[1] < 0: pars[1] = -pars[1]
     if inflate_errors:
-        chi2, dof = goodness_of_fit(hist, bins, var, gauss_basic, pars)
+        chi2, dof = goodness_of_fit(hist, bins, var, gauss_amp, pars)
         if chi2 > dof: cov *= chi2/dof
     return pars, cov
 
@@ -543,7 +543,7 @@ def unnorm_step_pdf(x,  mu, sigma, hstep):
     return step_f
 
 @nb.njit(**kwd)
-def step(x,  mu, sigma, hstep, lower_range=np.inf , upper_range=np.inf):
+def step_pdf(x,  mu, sigma, hstep, lower_range=np.inf , upper_range=np.inf):
 
     """
     Normalised step function w/args mu, sigma, hstep
@@ -584,7 +584,7 @@ def gauss_step_pdf(x,  n_sig, mu, sigma, n_bkg, hstep, lower_range=np.inf , uppe
     """
 
     try:
-        bkg= step(x, mu, sigma, hstep, lower_range, upper_range)
+        bkg= step_pdf(x, mu, sigma, hstep, lower_range, upper_range)
     except ZeroDivisionError:
         bkg = np.zeros_like(x, dtype=np.float64)
     if np.any(bkg<0):
@@ -629,7 +629,7 @@ def gauss_step_cdf(x,  n_sig, mu, sigma,n_bkg, hstep, lower_range=np.inf , upper
         return (1/(n_sig+n_bkg))*n_sig*gauss_cdf(x, mu, sigma), (1/(n_sig+n_bkg))*(n_bkg*bkg)
 
 @nb.njit(**kwd)
-def gauss_tail(x, mu, sigma, tau):
+def gauss_tail_pdf(x, mu, sigma, tau):
 
     """
     A gaussian tail function template
@@ -667,7 +667,7 @@ def gauss_tail_integral(x,mu,sigma,tau):
 
     abstau = np.abs(tau)
     part1 = (tau/(2*abstau)) * nb_erf((tau*(x-mu) )/(np.sqrt(2)*sigma*abstau))
-    part2 =    tau * gauss_tail(x,mu,sigma,tau)
+    part2 =    tau * gauss_tail_pdf(x,mu,sigma,tau)
     return part1+part2
 
 @nb.njit(**kwd)
@@ -677,7 +677,7 @@ def gauss_tail_norm(x,mu,sigma,tau, lower_range=np.inf , upper_range=np.inf):
     Normalised gauss tail. Note: this is only needed when the fitting range does not include the whole tail
     """
 
-    tail = gauss_tail(x,mu,sigma,tau)
+    tail = gauss_tail_pdf(x,mu,sigma,tau)
     if lower_range ==np.inf and upper_range ==np.inf:
         integral = gauss_tail_integral(np.array([np.nanmin(x), np.nanmax(x)]), mu, sigma, tau)
     else:
@@ -708,9 +708,15 @@ def gauss_with_tail_pdf(x, mu, sigma,  htail,tau, components=False):
     Pdf for gaussian with tail 
     """
 
+    if htail < 0 or htail > 1:
+        if components ==False:
+            return np.full_like(x, np.nan, dtype='float64')
+        else:
+            return np.full_like(x, np.nan, dtype='float64'), np.full_like(x, np.nan, dtype='float64')
+
     peak = gauss_norm(x,mu,sigma)
     try: 
-        tail = gauss_tail(x, mu, sigma, tau)
+        tail = gauss_tail_pdf(x, mu, sigma, tau)
     except ZeroDivisionError:
         tail = np.zeros_like(x, dtype=np.float64)
     if components ==False:
@@ -724,6 +730,12 @@ def gauss_with_tail_cdf(x, mu, sigma, htail,  tau, components=False):
     Cdf for gaussian with tail 
     """
 
+    if htail < 0 or htail > 1:
+        if components ==False:
+            return np.full_like(x, np.nan, dtype='float64')
+        else:
+            return np.full_like(x, np.nan, dtype='float64'), np.full_like(x, np.nan, dtype='float64')
+
     peak = gauss_cdf(x,mu,sigma)
     try: 
         tail = gauss_tail_cdf(x, mu, sigma, tau)
@@ -742,7 +754,7 @@ def radford_pdf(x, n_sig, mu, sigma, htail, tau, n_bkg, hstep,
     """
 
     try:
-        bkg= step(x, mu, sigma, hstep, lower_range, upper_range)
+        bkg= step_pdf(x, mu, sigma, hstep, lower_range, upper_range)
     except ZeroDivisionError:
         bkg = np.zeros_like(x, dtype=np.float64)
     if np.any(bkg<0):
@@ -798,6 +810,10 @@ def radford_fwhm(sigma, htail, tau,  cov = None):
     # optimize this to find max value
     def neg_radford_peak_bgfree(E, sigma, htail, tau):
         return -gauss_with_tail_pdf(np.array([E]), 0, sigma, htail, tau)[0]
+
+    if htail<0 or htail>1:
+        print("htail outside allowed limits of 0 and 1")
+        raise ValueError
 
     res = minimize_scalar( neg_radford_peak_bgfree,
                            args=(sigma, htail, tau),
@@ -853,19 +869,19 @@ def radford_peakshape_derivative(E, pars, step_norm):
     gaus = gauss_norm(E, mu, sigma)
     y = (E-mu)/sigma
     ret = -(1-htail)*(y/sigma)*gaus
-    ret -= htail/tau*(-gauss_tail(np.array([E,E-1]), mu, sigma, tau)[0]+gaus)
+    ret -= htail/tau*(-gauss_tail_pdf(np.array([E,E-1]), mu, sigma, tau)[0]+gaus)
 
     return n_sig*ret - n_bkg*hstep*gaus/step_norm #need norm factor for bkg
 
 def radford_parameter_gradient(E, pars, step_norm):
     n_sig, mu, sigma, htail, tau, n_bkg, hstep = pars 
 
     gaus = gauss_norm(np.array([E, E-1]), mu, sigma)[0] 
-    tailL = gauss_tail(np.array([E, E-1]), mu, sigma, tau)[0] 
+    tailL = gauss_tail_pdf(np.array([E, E-1]), mu, sigma, tau)[0] 
     if n_bkg ==0:
         step_f = 0
     else:
-        step_f = pgf.unnorm_step_pdf(np.array([E, E-1]), mu, sigma, hstep)[0] /step_norm
+        step_f = unnorm_step_pdf(np.array([E, E-1]), mu, sigma, hstep)[0] /step_norm
 
     #some unitless numbers that show up a bunch
     y = (E-mu)/sigma