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)