def data_prep(st, data): """ Applies calibration, flags, and subtracts time mean for data. """ # ** need to make this portable or at least reimplement in rfpipe if np.any(data): if st.metadata.datasource != 'sim': if os.path.exists(st.gainfile): data = calibration.apply_telcal(st, data) else: logger.warn( 'Telcal file not found. No calibration being applied.') else: logger.info('Not applying telcal solutions for simulated data') # ** dataflag points to rtpipe for now. # changes memory in place, so need to force writability data = util.dataflag(st, np.require(data, requirements='W')) if st.prefs.timesub == 'mean': logger.info('Subtracting mean visibility in time.') data = util.meantsub(data) else: logger.info('No visibility subtraction done.') if st.prefs.savenoise: logger.warn("Saving of noise properties not implemented yet.") return data
def prepare_data(sdmfile, gainfile, delta_l, delta_m, segment=0, dm=0, dt=1, spws=None): """ Applies Calibration, flagging, dedispersion and other data preparation steps from rfpipe. Then phaseshifts the data to the location of the candidate. """ st = state.State(sdmfile=sdmfile, sdmscan=1, inprefs={ 'gainfile': gainfile, 'workdir': '.', 'maxdm': 0, 'flaglist': [] }, showsummary=False) if spws: st.prefs.spw = spws data = source.read_segment(st, segment) takepol = [st.metadata.pols_orig.index(pol) for pol in st.pols] takebls = [ st.metadata.blarr_orig.tolist().index(list(bl)) for bl in st.blarr ] datap = np.require(data, requirements='W').take(takepol, axis=3).take( st.chans, axis=2).take(takebls, axis=1) datap = source.prep_standard(st, segment, datap) datap = calibration.apply_telcal(st, datap) datap = flagging.flag_data(st, datap) delay = calc_delay(st.freq, st.freq.max(), dm, st.inttime) data_dmdt = dedisperseresample(datap, delay, dt) print(f'shape of data_dmdt is {data_dmdt.shape}') uvw = get_uvw_segment(st, segment) phase_shift(data_dmdt, uvw=uvw, dl=delta_l, dm=delta_m) dataret = data_dmdt return dataret, st
def data_prep(st, segment, data, flagversion="latest"): """ Applies calibration, flags, and subtracts time mean for data. flagversion can be "latest" or "rtpipe". Optionally prepares data with antenna flags, fixing out of order data, calibration, downsampling, etc.. """ if not np.any(data): return data # take pols of interest takepol = [st.metadata.pols_orig.index(pol) for pol in st.pols] logger.debug('Selecting pols {0} and chans {1}'.format(st.pols, st.chans)) # TODO: check on reusing 'data' to save memory datap = np.nan_to_num(data.take(takepol, axis=3).take(st.chans, axis=2), copy=True) datap = prep_standard(st, segment, datap) if not np.any(datap): logger.info("All data zeros after prep_standard") return datap if st.gainfile is not None: datap = calibration.apply_telcal(st, datap) if not np.any(datap): logger.info("All data zeros after apply_telcal") return datap else: logger.info("No gainfile found, so not applying calibration.") # support backwards compatibility for reproducible flagging if flagversion == "latest": datap = flagging.flag_data(st, datap) elif flagversion == "rtpipe": datap = flagging.flag_data_rtpipe(st, datap) if st.prefs.timesub == 'mean': logger.info('Subtracting mean visibility in time.') datap = util.meantsub(datap, parallel=st.prefs.nthread > 1) else: logger.info('No visibility subtraction done.') if st.prefs.savenoise: save_noise(st, segment, datap) return datap
def make_transient_params(st, ntr=1, segment=None, dmind=None, dtind=None, i=None, amp=None, lm=None, snr=None, data=None): """ Given a state, create ntr randomized detectable transients. Returns list of ntr tuples of parameters. If data provided, it is used to inject transient at fixed SNR. selects random value from dmarr and dtarr. Mock transient will have either l or m equal to 0. Option exists to overload random selection with fixed segment, dmind, etc. """ segment0 = segment dmind0 = dmind dtind0 = dtind i0 = i amp0 = amp lm0 = lm snr0 = snr mocks = [] for tr in range(ntr): if segment is None: segment = random.choice(range(st.nsegment)) if dmind is not None: dm = st.dmarr[dmind] # dmind = random.choice(range(len(st.dmarr))) else: dm = np.random.uniform(min(st.dmarr), max(st.dmarr)) # pc /cc dmarr = np.array(calc_dmarr(st)) if dm > np.max(dmarr): logging.warning( "Dm of injected transient is greater than the max DM searched." ) dmind = len(dmarr) - 1 else: dmind = np.argmax(dmarr > dm) if dtind is not None: dt = st.inttime * min(st.dtarr[dtind], 2) # dt>2 not yet supported else: #dtind = random.choice(range(len(st.dtarr))) dt = st.inttime * np.random.uniform(min(st.dtarr), max( st.dtarr)) # s #like an alias for "dt" if dt < st.inttime: dtind = 0 else: dtind = int(np.log2(dt / st.inttime)) if dtind >= len(st.dtarr): dtind = len(st.dtarr) - 1 logging.warning( "Width of transient is greater than max dt searched.") # TODO: add support for arb dm/dt # dm = random.uniform(min(st.dmarr), max(st.dmarr)) # dt = random.uniform(min(st.dtarr), max(st.dtarr)) if i is None: i = random.choice(st.get_search_ints(segment, dmind, dtind)) if amp is None: if data is None: amp = random.uniform(0.1, 0.5) logger.info("Setting mock amp to {0}".format(amp)) else: if snr is None: snr = random.uniform(10, 50) logger.info("Setting mock snr to {0}".format(snr)) # TODO: support flagged data in size calc and injection datap = calibration.apply_telcal(st, data) sig = madtostd(datap[i].real) / np.sqrt( datap[i].size * st.dtarr[dtind]) amp = snr * sig logger.info("Setting mock amp as {0}*{1}={2}".format( snr, sig, amp)) if lm is None: # flip a coin to set either l or m if random.choice([0, 1]): l = math.radians( random.uniform(-st.fieldsize_deg / 2, st.fieldsize_deg / 2)) m = 0. else: l = 0. m = math.radians( random.uniform(-st.fieldsize_deg / 2, st.fieldsize_deg / 2)) else: assert len(lm) == 2, "lm must be 2-tuple" l, m = lm mocks.append((segment, i, dm, dt, amp, l, m)) (segment, dmind, dtind, i, amp, lm, snr) = (segment0, dmind0, dtind0, i0, amp0, lm0, snr0) return mocks
def make_transient_params(st, ntr=1, segment=None, dmind=None, dtind=None, i=None, amp=None, lm=None, snr=None, data=None): """ Given a state, create ntr randomized detectable transients. Returns list of ntr tuples of parameters. If data provided, it is used to inject transient at fixed apparent SNR. selects random value from dmarr and dtarr. Mock transient will have either l or m equal to 0. Option exists to overload random selection with fixed segment, dmind, etc. """ segment0 = segment dmind0 = dmind dtind0 = dtind i0 = i amp0 = amp lm0 = lm snr0 = snr mocks = [] for tr in range(ntr): if segment is None: segment = random.choice(range(st.nsegment)) if dmind is not None: dm = st.dmarr[dmind] # dmind = random.choice(range(len(st.dmarr))) else: dm = np.random.uniform(min(st.dmarr), max(st.dmarr)) # pc /cc dmarr = np.array(calc_dmarr(st)) if dm > np.max(dmarr): logging.warning( "Dm of injected transient is greater than the max DM searched." ) dmind = len(dmarr) - 1 else: dmind = np.argmax(dmarr > dm) if dtind is not None: dt = st.inttime * st.dtarr[dtind] else: #max_width = 0.04/st.inttime #dt = st.inttime*np.random.uniform(0, max_width) # s #like an alias for "dt" dt = np.random.uniform(0.001, 0.04) if dt < st.inttime: dtind = 0 else: boxcar_widths = np.array(st.dtarr) * st.inttime if dt > np.max(boxcar_widths): logging.warning( "Width of transient is greater than max dt searched.") dtind = np.argmin(np.abs(boxcar_widths - dt)) # else: # dtind = int(np.round(np.log2(dt/st.inttime))) # if dtind >= len(st.dtarr): # dtind = len(st.dtarr) - 1 # logging.warning("Width of transient is greater than max dt searched.") # TODO: add support for arb dm/dt # dm = random.uniform(min(st.dmarr), max(st.dmarr)) # dt = random.uniform(min(st.dtarr), max(st.dtarr)) if i is None: ints = np.array(st.get_search_ints(segment, dmind, dtind)) * st.dtarr[dtind] i = np.random.randint(min(ints), max(ints)) #i = np.random.randint(0,st.readints) #i = random.choice(st.get_search_ints(segment, dmind, dtind))*st.dtarr[dtind] if amp is None: if data is None: amp = random.uniform(0.1, 0.5) logger.info("Setting mock amp to {0}".format(amp)) else: if snr is None: snr = random.uniform(10, 50) # TODO: support flagged data in size calc and injection if data.shape != st.datashape: logger.info( "Looks like raw data passed in. Selecting and calibrating." ) takepol = [ st.metadata.pols_orig.index(pol) for pol in st.pols ] data = calibration.apply_telcal( st, data.take(takepol, axis=3).take(st.chans, axis=2)) # noise = madtostd(data[i].real)/np.sqrt(data[i].size*st.dtarr[dtind]) noise = madtostd(data[i].real) / np.sqrt( data[i].size) #*st.dtarr[dtind]) #width_factor = (dt//st.inttime)/np.sqrt(st.dtarr[dtind]) width_factor = (dt / st.inttime) / np.sqrt(st.dtarr[dtind]) amp = snr * noise * width_factor #*(st.inttime/dt) logger.info("Setting mock amp as {0}*{1}*{2}={3}".format( snr, noise, width_factor, amp)) #else: # if data is None: # logger.info("Setting mock amp to {0}".format(amp)) # else: # if data.shape != st.datashape: # logger.info("Looks like raw data passed in. Selecting and calibrating.") # takepol = [st.metadata.pols_orig.index(pol) for pol in st.pols] # data = calibration.apply_telcal(st, data.take(takepol, axis=3).take(st.chans, axis=2)) # bkgrnd = np.mean(data[i].real) # amp = amp0 + bkgrnd # logger.info("Setting mock amp to {0}+{1} = {2}".format(amp0, bkgrnd, amp)) # noise = madtostd(data[i].real)/np.sqrt(data[i].size) # snr_est = amp/(noise*np.sqrt(st.dtarr[dtind])) # logger.info("Estimated SNR is {0}".format(snr_est)) if lm is None: # flip a coin to set either l or m if random.choice([0, 1]): l = math.radians( random.uniform(-st.fieldsize_deg / 2, st.fieldsize_deg / 2)) m = 0. else: l = 0. m = math.radians( random.uniform(-st.fieldsize_deg / 2, st.fieldsize_deg / 2)) else: if lm == -1: l = math.radians( random.uniform(-st.fieldsize_deg / 2, st.fieldsize_deg / 2)) m = math.radians( random.uniform(-st.fieldsize_deg / 2, st.fieldsize_deg / 2)) else: assert len(lm) == 2, "lm must be 2-tuple or -1" l, m = lm mocks.append((segment, i, dm, dt, amp, l, m)) (segment, dmind, dtind, i, amp, lm, snr) = (segment0, dmind0, dtind0, i0, amp0, lm0, snr0) return mocks
def data_prep(st, segment, data, flagversion="latest", returnsoltime=False): """ Applies calibration, flags, and subtracts time mean for data. flagversion can be "latest" or "rtpipe". Optionally prepares data with antenna flags, fixing out of order data, calibration, downsampling, OTF rephasing... """ from rfpipe import calibration, flagging, util if not np.any(data): return data # take pols of interest takepol = [st.metadata.pols_orig.index(pol) for pol in st.pols] logger.debug('Selecting pols {0} and chans {1}'.format(st.pols, st.chans)) # TODO: check on reusing 'data' to save memory datap = np.nan_to_num( np.require(data, requirements='W').take(takepol, axis=3).take(st.chans, axis=2)) datap = prep_standard(st, segment, datap) if not np.any(datap): logger.info("All data zeros after prep_standard") return datap if st.gainfile is not None: logger.info("Applying calibration with {0}".format(st.gainfile)) ret = calibration.apply_telcal(st, datap, savesols=st.prefs.savesols, returnsoltime=returnsoltime) if returnsoltime: datap, soltime = ret else: datap = ret if not np.any(datap): logger.info("All data zeros after apply_telcal") return datap else: logger.info("No gainfile found, so not applying calibration.") # support backwards compatibility for reproducible flagging logger.info("Flagging with version: {0}".format(flagversion)) if flagversion == "latest": datap = flagging.flag_data(st, datap) elif flagversion == "rtpipe": datap = flagging.flag_data_rtpipe(st, datap) zerofrac = 1 - np.count_nonzero(datap) / datap.size if zerofrac > 0.8: logger.warning( 'Flagged {0:.1f}% of data. Zeroing all if greater than 80%.'. format(zerofrac * 100)) return np.array([]) if st.prefs.timesub == 'mean': logger.info('Subtracting mean visibility in time.') datap = util.meantsub(datap, parallel=st.prefs.nthread > 1) else: logger.info('No visibility subtraction done.') if (st.prefs.apply_chweights or st.prefs.apply_blweights) and st.readints > 3: if st.prefs.apply_chweights: # TODO: find better estimator. Currently loses sensitivity to FRB 121102 bursts. chvar = np.std(np.abs(datap).mean(axis=1), axis=0) chvar_norm = np.mean(1 / chvar**2, axis=0) if st.prefs.apply_blweights: blvar = np.std(np.abs(datap).mean(axis=2), axis=0) blvar_norm = np.mean(1 / blvar**2, axis=0) if st.prefs.apply_chweights: logger.info('Reweighting data by channel variances') datap = (datap / chvar[None, None, :, :]) / chvar_norm[None, None, None, :] if st.prefs.apply_blweights: logger.info('Reweighting data by baseline variances') datap = (datap / blvar[None, :, None, :]) / blvar_norm[None, None, None, :] if st.prefs.savenoise: save_noise(st, segment, datap) if returnsoltime: return datap, soltime else: return datap
def prep_standard(st, segment, data): """ Common first data prep stages, incl online flags, resampling, and mock transients. """ from rfpipe import calibration, flagging, util if not np.any(data): return data # read and apply flags for given ant/time range. 0=bad, 1=good if st.prefs.applyonlineflags and st.metadata.datasource in ['vys', 'sdm']: flags = flagging.getonlineflags(st, segment) data = np.where(flags[None, :, None, None], data, 0j) else: logger.info('Not applying online flags.') if not np.any(data): return data if st.prefs.simulated_transient is not None or st.otfcorrections is not None: uvw = util.get_uvw_segment(st, segment) # optionally integrate (downsample) if ((st.prefs.read_tdownsample > 1) or (st.prefs.read_fdownsample > 1)): data2 = np.zeros(st.datashape, dtype='complex64') if st.prefs.read_tdownsample > 1: logger.info('Downsampling in time by {0}'.format( st.prefs.read_tdownsample)) for i in range(st.datashape[0]): data2[i] = data[i * st.prefs.read_tdownsample:(i + 1) * st.prefs.read_tdownsample].mean(axis=0) if st.prefs.read_fdownsample > 1: logger.info('Downsampling in frequency by {0}'.format( st.prefs.read_fdownsample)) for i in range(st.datashape[2]): data2[:, :, i, :] = data[:, :, i * st.prefs.read_fdownsample:(i + 1) * st.prefs.read_fdownsample].mean(axis=2) data = data2 # optionally add transients if st.prefs.simulated_transient is not None: # for an int type, overload prefs.simulated_transient random mocks if isinstance(st.prefs.simulated_transient, int): logger.info( "Filling simulated_transient with {0} random transients". format(st.prefs.simulated_transient)) st.prefs.simulated_transient = util.make_transient_params( st, segment=segment, ntr=st.prefs.simulated_transient, data=data) assert isinstance(st.prefs.simulated_transient, list), "Simulated transient must be list of tuples." for params in st.prefs.simulated_transient: assert len(params) == 7 or len(params) == 8, ( "Transient requires 7 or 8 parameters: " "(segment, i0/int, dm/pc/cm3, dt/s, " "amp/sys, dl/rad, dm/rad) and optionally " "ampslope/sys") if segment == params[0]: if len(params) == 7: (mock_segment, i0, dm, dt, amp, l, m) = params ampslope = 0 logger.info( "Adding transient to segment {0} at int {1}, " "DM {2}, dt {3} with amp {4} and l,m={5},{6}".format( mock_segment, i0, dm, dt, amp, l, m)) elif len(params) == 8: (mock_segment, i0, dm, dt, amp, l, m, ampslope) = params logger.info("Adding transient to segment {0} at int {1}, " " DM {2}, dt {3} with amp {4}-{5} and " "l,m={6},{7}".format(mock_segment, i0, dm, dt, amp, amp + ampslope, l, m)) try: model = np.require(np.broadcast_to( util.make_transient_data( st, amp, i0, dm, dt, ampslope=ampslope).transpose()[:, None, :, None], st.datashape), requirements='W') except IndexError: logger.warning( "IndexError while adding transient. Skipping...") continue if st.gainfile is not None: model = calibration.apply_telcal(st, model, sign=-1) util.phase_shift(model, uvw, -l, -m) data += model if st.otfcorrections is not None: # shift phasecenters to first phasecenter in segment if len(st.otfcorrections[segment]) > 1: ints, ra0, dec0 = st.otfcorrections[segment][ 0] # new phase center for segment logger.info( "Correcting {0} phasecenters to first at RA,Dec = {1},{2}". format(len(st.otfcorrections[segment]) - 1, ra0, dec0)) for ints, ra_deg, dec_deg in st.otfcorrections[segment][1:]: l0 = np.radians(ra_deg - ra0) m0 = np.radians(dec_deg - dec0) util.phase_shift(data, uvw, l0, m0, ints=ints) return data
def data_prep(st, segment, data, flagversion="latest", returnsoltime=False): """ Applies calibration, flags, and subtracts time mean for data. flagversion can be "latest" or "rtpipe". Optionally prepares data with antenna flags, fixing out of order data, calibration, downsampling, OTF rephasing... """ from rfpipe import calibration, flagging if not np.any(data): return data # take pols of interest takepol = [st.metadata.pols_orig.index(pol) for pol in st.pols] logger.debug('Selecting pols {0} and chans {1}'.format(st.pols, st.chans)) # TODO: check on reusing 'data' to save memory datap = np.nan_to_num(np.require(data, requirements='W').take(takepol, axis=3).take(st.chans, axis=2)) datap = prep_standard(st, segment, datap) if not np.any(datap): logger.info("All data zeros after prep_standard") return datap if st.gainfile is not None: logger.info("Applying calibration with {0}".format(st.gainfile)) ret = calibration.apply_telcal(st, datap, savesols=st.prefs.savesols, returnsoltime=returnsoltime) if returnsoltime: datap, soltime = ret else: datap = ret if not np.any(datap): logger.info("All data zeros after apply_telcal") return datap else: logger.info("No gainfile found, so not applying calibration.") # support backwards compatibility for reproducible flagging logger.info("Flagging with version: {0}".format(flagversion)) if flagversion == "latest": datap = flagging.flag_data(st, datap) elif flagversion == "rtpipe": datap = flagging.flag_data_rtpipe(st, datap) zerofrac = 1-np.count_nonzero(datap)/datap.size if zerofrac > 0.8: logger.warning('Flagged {0:.1f}% of data. Zeroing all if greater than 80%.'.format(zerofrac*100)) return np.array([]) if st.prefs.timesub == 'mean': logger.info('Subtracting mean visibility in time.') datap = util.meantsub(datap, parallel=st.prefs.nthread > 1) else: logger.info('No visibility subtraction done.') if (st.prefs.apply_chweights or st.prefs.apply_blweights) and st.readints > 3: if st.prefs.apply_chweights: # TODO: find better estimator. Currently loses sensitivity to FRB 121102 bursts. chvar = np.std(np.abs(datap).mean(axis=1), axis=0) chvar_norm = np.mean(1/chvar**2, axis=0) if st.prefs.apply_blweights: blvar = np.std(np.abs(datap).mean(axis=2), axis=0) blvar_norm = np.mean(1/blvar**2, axis=0) if st.prefs.apply_chweights: logger.info('Reweighting data by channel variances') datap = (datap/chvar[None, None, :, :])/chvar_norm[None, None, None, :] if st.prefs.apply_blweights: logger.info('Reweighting data by baseline variances') datap = (datap/blvar[None, :, None, :])/blvar_norm[None, None, None, :] if st.prefs.savenoise: save_noise(st, segment, datap) if returnsoltime: return datap, soltime else: return datap
def prep_standard(st, segment, data): """ Common first data prep stages, incl online flags, resampling, and mock transients. """ from rfpipe import calibration, flagging if not np.any(data): return data # read and apply flags for given ant/time range. 0=bad, 1=good if st.prefs.applyonlineflags and st.metadata.datasource in ['vys', 'sdm']: flags = flagging.getonlineflags(st, segment) data = np.where(flags[None, :, None, None], data, 0j) else: logger.info('Not applying online flags.') if not np.any(data): return data if st.prefs.simulated_transient is not None or st.otfcorrections is not None: uvw = util.get_uvw_segment(st, segment) # optionally integrate (downsample) if ((st.prefs.read_tdownsample > 1) or (st.prefs.read_fdownsample > 1)): data2 = np.zeros(st.datashape, dtype='complex64') if st.prefs.read_tdownsample > 1: logger.info('Downsampling in time by {0}' .format(st.prefs.read_tdownsample)) for i in range(st.datashape[0]): data2[i] = data[ i*st.prefs.read_tdownsample:(i+1)*st.prefs.read_tdownsample].mean(axis=0) if st.prefs.read_fdownsample > 1: logger.info('Downsampling in frequency by {0}' .format(st.prefs.read_fdownsample)) for i in range(st.datashape[2]): data2[:, :, i, :] = data[:, :, i*st.prefs.read_fdownsample:(i+1)*st.prefs.read_fdownsample].mean(axis=2) data = data2 # optionally add transients if st.prefs.simulated_transient is not None: # for an int type, overload prefs.simulated_transient random mocks if isinstance(st.prefs.simulated_transient, int): logger.info("Filling simulated_transient with {0} random transients" .format(st.prefs.simulated_transient)) st.prefs.simulated_transient = util.make_transient_params(st, segment=segment, ntr=st.prefs.simulated_transient, data=data) assert isinstance(st.prefs.simulated_transient, list), "Simulated transient must be list of tuples." for params in st.prefs.simulated_transient: assert len(params) == 7 or len(params) == 8, ("Transient requires 7 or 8 parameters: " "(segment, i0/int, dm/pc/cm3, dt/s, " "amp/sys, dl/rad, dm/rad) and optionally " "ampslope/sys") if segment == params[0]: if len(params) == 7: (mock_segment, i0, dm, dt, amp, l, m) = params ampslope = 0 logger.info("Adding transient to segment {0} at int {1}, " "DM {2}, dt {3} with amp {4} and l,m={5},{6}" .format(mock_segment, i0, dm, dt, amp, l, m)) elif len(params) == 8: (mock_segment, i0, dm, dt, amp, l, m, ampslope) = params logger.info("Adding transient to segment {0} at int {1}, " " DM {2}, dt {3} with amp {4}-{5} and " "l,m={6},{7}" .format(mock_segment, i0, dm, dt, amp, amp+ampslope, l, m)) try: model = np.require(np.broadcast_to(util.make_transient_data(st, amp, i0, dm, dt, ampslope=ampslope) .transpose()[:, None, :, None], st.datashape), requirements='W') except IndexError: logger.warning("IndexError while adding transient. Skipping...") continue if st.gainfile is not None: model = calibration.apply_telcal(st, model, sign=-1) util.phase_shift(model, uvw, -l, -m) data += model if st.otfcorrections is not None: # shift phasecenters to first phasecenter in segment if len(st.otfcorrections[segment]) > 1: ints, ra0, dec0 = st.otfcorrections[segment][0] # new phase center for segment logger.info("Correcting {0} phasecenters to first at RA,Dec = {1},{2}" .format(len(st.otfcorrections[segment])-1, ra0, dec0)) for ints, ra_deg, dec_deg in st.otfcorrections[segment][1:]: l0 = np.radians(ra_deg-ra0) m0 = np.radians(dec_deg-dec0) util.phase_shift(data, uvw, l0, m0, ints=ints) return data
def make_transient_params(st, ntr=1, segment=None, dmind=None, dtind=None, i=None, amp=None, lm=None, snr=None, data=None): """ Given a state, create ntr randomized detectable transients. Returns list of ntr tuples of parameters. If data provided, it is used to inject transient at fixed SNR. selects random value from dmarr and dtarr. Mock transient will have either l or m equal to 0. Option exists to overload random selection with fixed segment, dmind, etc. """ segment0 = segment dmind0 = dmind dtind0 = dtind i0 = i amp0 = amp lm0 = lm snr0 = snr mocks = [] for tr in range(ntr): if segment is None: segment = random.choice(range(st.nsegment)) if dmind is not None: dm = st.dmarr[dmind] # dmind = random.choice(range(len(st.dmarr))) else: dm = np.random.uniform(min(st.dmarr), max(st.dmarr)) # pc /cc dmarr = np.array(calc_dmarr(st)) if dm > np.max(dmarr): logging.warning("Dm of injected transient is greater than the max DM searched.") dmind = len(dmarr) - 1 else: dmind = np.argmax(dmarr>dm) if dtind is not None: dt = st.inttime*min(st.dtarr[dtind], 2) # dt>2 not yet supported else: #dtind = random.choice(range(len(st.dtarr))) dt = st.inttime*np.random.uniform(min(st.dtarr), max(st.dtarr)) # s #like an alias for "dt" if dt < st.inttime: dtind = 0 else: dtind = int(np.log2(dt/st.inttime)) if dtind >= len(st.dtarr): dtind = len(st.dtarr) - 1 logging.warning("Width of transient is greater than max dt searched.") # TODO: add support for arb dm/dt # dm = random.uniform(min(st.dmarr), max(st.dmarr)) # dt = random.uniform(min(st.dtarr), max(st.dtarr)) if i is None: i = random.choice(st.get_search_ints(segment, dmind, dtind)) if amp is None: if data is None: amp = random.uniform(0.1, 0.5) logger.info("Setting mock amp to {0}".format(amp)) else: if snr is None: snr = random.uniform(10, 50) logger.info("Setting mock snr to {0}".format(snr)) # TODO: support flagged data in size calc and injection datap = calibration.apply_telcal(st, data) sig = madtostd(datap[i].real)/np.sqrt(datap[i].size*st.dtarr[dtind]) amp = snr*sig logger.info("Setting mock amp as {0}*{1}={2}".format(snr, sig, amp)) if lm is None: # flip a coin to set either l or m if random.choice([0, 1]): l = math.radians(random.uniform(-st.fieldsize_deg/2, st.fieldsize_deg/2)) m = 0. else: l = 0. m = math.radians(random.uniform(-st.fieldsize_deg/2, st.fieldsize_deg/2)) else: assert len(lm) == 2, "lm must be 2-tuple" l, m = lm mocks.append((segment, i, dm, dt, amp, l, m)) (segment, dmind, dtind, i, amp, lm, snr) = (segment0, dmind0, dtind0, i0, amp0, lm0, snr0) return mocks