def test_ofamp():
"""
Testing function for `qetpy.ofamp`.
"""
signal, template, psd = create_example_data()
fs = 625e3
res1 = qp.ofamp(signal,
template,
psd,
fs,
lgcsigma=True,
nconstrain=100,
withdelay=True)
OF = qp.OptimumFilter(signal, template, psd, fs)
res2 = OF.ofamp_withdelay(nconstrain=100)
res_compare1 = res2 + (OF.energy_resolution(), )
res3 = qp.ofamp(signal, template, psd, fs, withdelay=False)
OF = qp.OptimumFilter(signal, template, psd, fs)
res4 = OF.ofamp_nodelay()
res_compare2 = (res4[0], 0.0, res4[1])
assert isclose(res1, res_compare1)
assert isclose(res3, res_compare2)
def test_estimate_g():
"""Testing function for `qetpy.utils.estimate_g`"""
p0 = 3e-12
tc = 40e-3
tbath = 0
p0_err = 1e-12
assert isclose(
estimate_g(p0, tc, tbath, p0_err=p0_err),
[3.75e-10, 1.25e-10],
)
sigp0 = 1e-12
sigtc = 1e-3
sigtbath = 1e-3
corr = 0.95
cov_test = np.array([
[sigp0**2, -corr * sigp0 * sigtc, -corr * sigp0 * sigtc],
[-corr * sigp0 * sigtc, sigtc**2, corr * sigtbath * sigtc],
[-corr * sigp0 * sigtc, corr * sigtbath * sigtc, sigtbath**2],
])
assert isclose(
estimate_g(p0, tc, tbath, cov=cov_test),
[3.75e-10, 1.339382436983552e-10],
)
def __init__(self, allocation, value=1.0):
if type(allocation) is LinearRamp:
if not isclose(sum(allocation.start), 1.0):
raise RuntimeError("Start allocations need to add up to 1.0")
if not isclose(sum(allocation.end), 1.0):
raise RuntimeError("End allocations need to add up to 1.0")
self.allocation = allocation.start
self.ramp = allocation
else:
if not isclose(sum(allocation), 1.0):
raise RuntimeError("Allocations need to add up to 1.0")
self.allocation = allocation
self.ramp = None
self.value = value
def __next__(self):
newAssets = self.interpolate()
if not isclose(sum(newAssets), 1.0):
raise RuntimeError("Allocations need to add up to 1.0")
self.count += 1
if self.count > self.expectedIterations:
raise StopIteration
return newAssets
def doTrinityTest(self, expectedRate, length, withdrawalRate, stocks,
bonds, ignoreInflation, endYear):
result = trinity(length, withdrawalRate * 1000000,
Portfolio(Assets(stocks, bonds), 1000000),
ignoreInflation, endYear)
# so its actually really f*****g hard to recreate these results perfectly.
# we are going to define an acceptable error rate that gets broader the further from 100% success we are
# the intuition is that if we are off by 10% when the expected success rate is only 10% we don't care
# but even errors of a few percent in the 90's matter
delta = .005 # this is the min delta given the trinity study's significant figures.
# so at an expected 1.0 success rate we allow no wiggle room
# and at 0.0 success rate we allow a full 10%
delta += .1 - expectedRate / 10
# we also scale by withdrawal rate. Do we really care if we are wrong at 12 withdrawal rate?
delta += withdrawalRate
# im pretty sure the stock data is wrong so we scale by that too
delta += stocks / 10
# again we don't really care about non 30 year simulations that much
delta += (30 - length) / 100
# Worst case delta...
# expected rate of 0.0; delta += 10%
# withdrawal rate of 12; delta += 12%
# 100% stocks; delta += 10%
# 15 year simulation; delta += 15%
# delta ~= 47%, which is a big gap.
# delta for a case we care about...
# expected rate of .96; delta += .4%
# withdrawal rate of 4%; delta += 4%
# 50/50 stocks/bonds; delta =+ 5%
# 30 year simulation; delta += 0%
# delta ~= 9.4%
# we also double the delta if trinity study is optimistic compared to us.
# and halve it if we are reporting a higher success rate.
actualSuccessRate = result.getSuccessRate()
if actualSuccessRate < expectedRate:
delta = delta * 2
if isclose(actualSuccessRate, expectedRate, abs_tol=delta):
return (True, actualSuccessRate)
else:
print(
"EXPECTED:\t{0:.2f}\tACTUAL:\t{1:.2f}\tRate tested:\t{2}\tYears\t{3}\tStocks:\t{4:.2f}\tBonds:\t{5:.2f}\tIgnoreInflations:\t{6}\tEndYear:\t{7}\tDelta:{8}"
.format(expectedRate, actualSuccessRate, withdrawalRate,
length, stocks, bonds, ignoreInflation, endYear,
delta))
self.assertAlmostEqual(actualSuccessRate,
expectedRate,
delta=delta)
def test_chi2_nopulse():
"""
Testing function for `qetpy.chi2_nopulse`.
"""
signal, template, psd = create_example_data()
fs = 625e3
res1 = qp.chi2_nopulse(signal, psd, fs)
OF = qp.OptimumFilter(signal, template, psd, fs)
res2 = OF.chi2_nopulse()
assert isclose(res1, res2)
def test_ofamp_pileup_stationary():
"""
Testing function for `qetpy.ofamp_pileup_stationary`.
"""
signal, template, psd = create_example_data()
fs = 625e3
res1 = qp.ofamp_pileup_stationary(signal, template, psd, fs)
OF = qp.OptimumFilter(signal, template, psd, fs)
res2 = OF.ofamp_pileup_stationary()
assert isclose(res1, res2)
def test_MuonTailFit():
"""
Testing function for `qetpy.MuonTailFit` class.
"""
signal, psd = create_example_muontail()
fs = 625e3
mtail = qp.MuonTailFit(psd, fs)
res = mtail.fitmuontail(signal, lgcfullrtn=False)
assert isclose(
res,
[4.5964397140760587e-07, 0.0201484585061281],
rtol=1e-8,
)
def test_ofamp_pileup():
"""
Testing function for `qetpy.ofamp_pileup`.
"""
signal, template, psd = create_example_data(lgcpileup=True)
fs = 625e3
res1 = qp.ofamp_pileup(signal, template, psd, fs)
OF = qp.OptimumFilter(signal, template, psd, fs)
res2 = OF.ofamp_withdelay()
res3 = OF.ofamp_pileup_iterative(res2[0], res2[1])
res_compare = res2[:-1] + res3
assert isclose(res1, res_compare)
def test_make_template():
"""Testing function for `qetpy.utils.make_template`."""
fs = 625e3
tau_r = 20e-6
tau_f = 80e-6
offset = -4
time = np.arange(1000) / fs
time_offset = (len(time) // 2) / fs + offset / fs
# calculate the pulse in an equivalent way, such that the result should
# be the same
pulse = np.heaviside(time - time_offset, 0)
pulse *= np.exp(-(time - time_offset) / tau_f, ) - np.exp(
-(time - time_offset) / tau_r, )
pulse /= pulse.max()
assert isclose(make_template(time, tau_r, tau_f, offset), pulse)
def getLongTermCorpBondsReturn(lyear, lmonth):
if (lyear, lmonth) not in longTermCorpBondsReturnCacne:
script = os.path.dirname(__file__)
path = os.path.join(script, "./data/longtermcorpbonds.txt")
with open(path, mode="r") as f:
for line in f:
if "YEAR" in line:
continue
split = line.split()
year = int(split[0])
yearRate = float(split[13]) / 10000
yearRateCalc = 1.0
for i in range(1, 13):
monthRate = float(split[i]) / 10000
yearRateCalc *= 1.0 + monthRate
longTermCorpBondsReturnCacne[(year, i)] = monthRate
if not isclose(yearRateCalc, yearRate + 1.0, rel_tol=.001):
raise RuntimeError
if not (lyear, lmonth) in longTermCorpBondsReturnCacne:
raise RuntimeError
return longTermCorpBondsReturnCacne[(lyear, lmonth)] + 1.0
def test_chi2lowfreq():
"""
Testing function for `qetpy.chi2lowfreq`.
"""
signal, template, psd = create_example_data()
fs = 625e3
res1 = qp.ofamp(signal, template, psd, fs)
chi2low = qp.chi2lowfreq(signal,
template,
res1[0],
res1[1],
psd,
fs,
fcutoff=10000)
OF = qp.OptimumFilter(signal, template, psd, fs)
res2 = OF.ofamp_withdelay()
chi2low_compare = OF.chi2_lowfreq(amp=res2[0], t0=res2[1], fcutoff=10000)
assert isclose(chi2low, chi2low_compare)
def test_OptimumFilter():
"""
Testing function for `qetpy.OptimumFilter` class.
"""
signal, template, psd = create_example_data()
fs = 625e3
OF = qp.OptimumFilter(signal, template, psd, fs)
res = OF.ofamp_nodelay()
assert isclose(res, (-1.589803642041125e-07, 2871569.457990007), rtol=1e-6)
res = OF.energy_resolution()
assert isclose(res, 2.3725914280425287e-09, rtol=1e-6)
res = OF.ofamp_withdelay()
assert isclose(res, (4.000884927004103e-06, 0.00016, 32474.45440205792),
rtol=1e-6)
res2 = OF.ofamp_nodelay(windowcenter=int(res[1] * fs))
assert isclose(res2, res[::2], rtol=1e-6)
res = OF.time_resolution(res[0])
assert isclose(res, 5.746611055379949e-09, rtol=1e-6)
res = OF.ofamp_withdelay(nconstrain=100)
assert isclose(res, (6.382904231454342e-07, 7.84e-05, 2803684.0424425197),
rtol=1e-6)
res = OF.ofamp_withdelay(nconstrain=100, lgcoutsidewindow=True)
assert isclose(res, (4.000884927004103e-06, 0.00016, 32474.45440205792),
rtol=1e-6)
res = OF.ofamp_withdelay(nconstrain=3, windowcenter=-5)
assert isclose(res, (-1.748136068514983e-07, -9.6e-06, 2870630.5945196496),
rtol=1e-6)
res = OF.chi2_lowfreq(amp=4.000884927004103e-06, t0=0.00016, fcutoff=10000)
assert isclose(res, 1052.9089578293142, rtol=1e-6)
res = OF.chi2_nopulse()
assert isclose(res, 2876059.4034037213, rtol=1e-6)
OF.update_signal(signal)
res = OF.ofamp_pileup_stationary()
assert isclose(res, (2.884298804131357e-09, 4.001001614298674e-06, 0.00016,
32472.978956471197),
rtol=1e-6)
signal, template, psd = create_example_data(lgcpileup=True)
OF.update_signal(signal)
res1 = OF.ofamp_withdelay()
res = OF.ofamp_pileup_iterative(res1[0], res1[1])
assert isclose(res, (4.000882414471985e-06, 0.00016, 32477.55571848713),
rtol=1e-6)
res = OF.ofamp_pileup_iterative(res1[0],
res1[1],
nconstrain=100,
lgcoutsidewindow=False)
assert isclose(res, (6.382879106117655e-07, 7.84e-05, 2803684.142039136),
rtol=1e-6)
res = OF.ofamp_pileup_iterative(res1[0],
res1[1],
nconstrain=100,
lgcoutsidewindow=True)
assert isclose(res, (4.000882414471985e-06, 0.00016, 32477.55571848713),
rtol=1e-6)
signal, template, psd = create_example_data(lgcbaseline=True)
OF.update_signal(signal)
res = OF.ofamp_baseline()
assert isclose(res, (4.000884927004102e-06, 0.00016, 32474.454402058076),
rtol=1e-6)
res = OF.ofamp_baseline(nconstrain=100)
assert isclose(res, (6.434754982839688e-07, 7.84e-05, 2806781.3450564747),
rtol=1e-6)
res = OF.ofamp_baseline(nconstrain=100, lgcoutsidewindow=True)
assert isclose(res, (4.000884927004102e-06, 0.00016, 32474.454402058076),
rtol=1e-6)
def test_OFnonlin():
"""
Testing function for `qetpy.OFnonlin` class.
"""
signal, template, psd = create_example_data()
fs = 625e3
signal = np.roll(
signal,
-100) # undo the roll in create_example_data to make test easier
nlin = qp.OFnonlin(psd, fs, template=template)
res1a = nlin.fit_falltimes(signal,
npolefit=1,
lgcfullrtn=False,
lgcplot=True,
taurise=20e-6,
scale_amplitude=True)
res1 = nlin.fit_falltimes(signal,
npolefit=1,
lgcfullrtn=False,
lgcplot=True,
taurise=20e-6,
scale_amplitude=False)
res2a = nlin.fit_falltimes(signal,
npolefit=2,
lgcfullrtn=False,
lgcplot=True,
scale_amplitude=True)
res2 = nlin.fit_falltimes(signal,
npolefit=2,
lgcfullrtn=False,
lgcplot=True,
scale_amplitude=False)
res3 = nlin.fit_falltimes(signal,
npolefit=3,
lgcfullrtn=False,
lgcplot=True)
res4 = nlin.fit_falltimes(signal,
npolefit=4,
lgcfullrtn=False,
lgcplot=True)
assert isclose(
res1a,
[4.008696926367952e-06, 6.577134966380607e-05, 2.600003126086262e-02])
assert isclose(
res1,
[9.690520626128428e-06, 6.577262665978902e-05, 2.600003114814408e-02],
rtol=1e-6)
assert isclose(res2a, [
4.010777893773002e-06, 1.952058681743050e-05, 6.667391354400327e-05,
2.600012092917421e-02
],
rtol=1e-6)
assert isclose(res2, [
9.501376001058713e-06, 1.962953013808533e-05, 6.638332141659392e-05,
2.600010755026570e-02
],
rtol=1e-6)
assert isclose(res3, [
9.308842323550344e-06, 1.332396374991919e-08, 1.930693061996180e-05,
6.697226655672301e-05, 1.502288275853276e-04, 2.600016234389370e-02
],
rtol=1e-6)
assert isclose(res4, [
9.491495350665769e-06, 3.023941433370170e-08, 5.523645346886680e-08,
1.976936973433418e-05, 6.566969025231684e-05, 9.213022382501221e-05,
2.246779922836221e-04, 2.600006569525561e-02
],
rtol=1e-6)
def test_ofnsmb_ttlfitting():
"""
Testing function for `qetpy.of_nsmb_setup and qetpy.of_nsmb`.
"""
fs = 625e3
ttlrate = 2e3
signal, template, psd = create_example_ttl_leakage_pulses(fs, ttlrate)
nbin = len(signal)
(backgroundtemplates,
backgroundtemplateshifts,
backgroundpolarityconstraint,
indwindow_nsmb) = qp.maketemplate_ttlfit_nsmb(template,
fs,
ttlrate,
lgcconstrainpolarity=True,
lgcpositivepolarity=False,
notch_window_size=1)
# add a second npdarray to the window list
# to test the window functionality
indwindow_nsmb.append(np.arange(100))
# concatenate signal and background template matrices and take FFT
sbtemplatef, sbtemplatet = qp.of_nsmb_ffttemplate(np.expand_dims(template, 1), backgroundtemplates)
(psddnu, phi, Pfs, P,
sbtemplatef, sbtemplatet, iB,
B, ns, nb, bitcomb, lfindex) = qp.of_nsmb_setup(template, backgroundtemplates, psd, fs)
sigpolarityconstraint = (-1)*np.ones(1)
s = signal
(amps_nsmb,t0_s_nsmb,
chi2_nsmb,chi2_nsmb_lf,
resid,amps_sig_nsmb_cwindow,
chi2_nsmb_cwindow,
t0_s_nsmb_cwindow,
amp_s_nsmb_int,
t0_s_nsmb_int,
chi2_nsmb_int,
amps_sig_nsmb_cwindow_int,
chi2_nsmb_cwindow_int,
t0_s_nsmb_cwindow_int) = qp.of_nsmb_con(s, phi, Pfs,
P, sbtemplatef.T, sbtemplatet,
psddnu.T, fs, indwindow_nsmb, ns, nb, bitcomb, lfindex,
background_templates_shifts=backgroundtemplateshifts,
bkgpolarityconstraint=backgroundpolarityconstraint,
sigpolarityconstraint=sigpolarityconstraint,
lgcplot=False, lgcsaveplots=False)
(ampsbonly_nsmb, chi2bonly_nsmb,
chi2bonly_nsmb_lf) = qp.of_mb(s, phi, sbtemplatef.T, sbtemplatet,
iB, B, psddnu.T, fs, ns, nb, lfindex,
background_templates_shifts=backgroundtemplateshifts,
bkgpolarityconstraint=backgroundpolarityconstraint,
sigpolarityconstraint=sigpolarityconstraint,
lgcplot=False, lgcsaveplots=False)
# check the signal amplitude and the first three
# background amplitudes
priorPulseAmp = -3.82861366e-08
priorB1Amp = -2.89749852e-08
priorB2Amp = -4.61737507e-09
priorB3Amp = -1.68752504e-08
priorbonlyB1 = -3.01203573e-08
priorbonlyB2 = -4.93831067e-09
priorampwindow = -4.79771072e-09
priorampwindow_int = -4.79831334e-09
savedVals = (priorPulseAmp, priorB1Amp,
priorB2Amp, priorB3Amp,
priorbonlyB1, priorbonlyB2,
priorampwindow, priorampwindow_int)
newVals = (amps_nsmb[0], amps_nsmb[1],
amps_nsmb[2], amps_nsmb[3],
ampsbonly_nsmb[0], ampsbonly_nsmb[1],
amps_sig_nsmb_cwindow[0,0], amps_sig_nsmb_cwindow_int[0,0])
rtol = 1e-6
assert isclose(newVals, savedVals, rtol=rtol)
def test_ofnsmb_muonsubtraction():
"""
Testing function for `qetpy.of_nsmb_setup and qetpy.of_nsmb`. This is a simple
test in that we only have two backgrounds, thus the matrices are only 3X3
"""
signal, template, psd = create_example_pulseplusmuontail(lgcbaseline=False)
fs = 625e3
nbin = len(signal)
# construct the background templates
backgroundtemplates, backgroundtemplatesshifts = qp.get_slope_dc_template_nsmb(nbin)
(psddnu,
phi,
Pfs,
P,
sbtemplatef,
sbtemplatet,
iB,
B,
ns,
nb,
bitcomb,
lfindex) = qp.of_nsmb_setup(template, backgroundtemplates, psd, fs)
iP = qp.of_nsmb_getiP(P)
# construct allowed window for signal template
indwindow = np.arange(len(template))
# make indwindow dimensions 1 X (time bins)
indwindow = indwindow[:, None].T
# find all indices within -lowInd and +highInd bins of background_template_shifts
lowInd = 1
highInd = 1
restrictInd = np.arange(-lowInd, highInd+1)
for ii in range(1, nb-1):
# start with the second index
restrictInd = np.concatenate((restrictInd,
np.arange(int(backgroundtemplatesshifts[ii] - lowInd),
int(backgroundtemplatesshifts[ii] + highInd + 1))))
# if values of restrictInd are negative
# wrap them around to the end of the window
lgcneg = restrictInd < 0
restrictInd[lgcneg] = len(template) + restrictInd[lgcneg]
# make restictInd 1 X (time bins)
restrictInd = restrictInd[:, None].T
# delete the restrictedInd from indwindow
indwindow = np.delete(indwindow, restrictInd)
# make indwindow dimensions 1 X (time bins)
indwindow = indwindow[:, None].T
indwindow_nsmb = [indwindow]
lgcplotnsmb = True
s = signal
(amps_nsmb, t0_s_nsmb, chi2_nsmb,
chi2_nsmb_lf, resid) = qp.of_nsmb(s,
phi,
sbtemplatef.T,
sbtemplatet,
iP,
psddnu.T,
fs,
indwindow_nsmb,
ns,
nb,
bitcomb,
lfindex,
lgcplot=lgcplotnsmb,
lgcsaveplots=False,
)
priorPulseAmp = -4.07338835e-08
priorMuonAmp = 1.13352442e-07
priorDC = -4.96896901e-07
savedVals = (priorPulseAmp, priorMuonAmp, priorDC)
rtol = 1e-6
assert isclose(amps_nsmb, savedVals, rtol=rtol)
