def apply_otfcorrections(st, segment, data, raw=False):
""" Phase shift data to single phase center
raw defines whether to correct data with all chans (metadata.freq_orig)
or selected set (st.freq).
"""
from rfpipe import util
# shift phasecenters to first phasecenter in segment
if len(st.otfcorrections[segment]) > 1:
# get reference phase center
pc0 = st.get_pc(segment)
ra0, dec0 = st.get_radec(pc=pc0)
logger.info("Correcting {0} phasecenters to pc {1}".format(
len(st.otfcorrections[segment]) - 1, pc0))
for pc, ints, ra_deg, dec_deg in st.otfcorrections[segment]:
if pc != pc0:
# using dl,dm
# l0 = np.radians((ra_deg-ra0)*np.cos(np.radians(dec0)))
# m0 = np.radians(dec_deg-dec0)
# uvw = util.get_uvw_segment(st, segment, pc=pc, raw=raw)
# util.phase_shift(data, uvw=uvw, dl=-l0, dm=-m0, ints=ints)
# using dw
u0, v0, w0 = util.get_uvw_segment(st,
segment,
pc_radec=pc,
pc_mjd=pc,
raw=raw)
u1, v1, w1 = util.get_uvw_segment(st,
segment,
pc_radec=pc0,
pc_mjd=pc,
raw=raw)
dw = w1 - w0
util.phase_shift(data, dw=dw, ints=ints)
def pipeline_canddata(st, candloc, data_dmdt=None, cpuonly=False, sig_ts=None,
kalman_coeffs=None, **kwargs):
""" Generate image and phased visibility data for candloc.
Phases to peak pixel in image of candidate.
Can optionally pass in corrected data, if available.
cpuonly argument not being used at present.
"""
from rfpipe import candidates
segment, candint, dmind, dtind, beamnum = candloc
dt = st.dtarr[dtind]
dm = st.dmarr[dmind]
uvw = util.get_uvw_segment(st, segment)
wisdom = rfpipe.search.set_wisdom(st.npixx, st.npixy)
if data_dmdt is None:
data_dmdt = pipeline_datacorrect(st, candloc)
if 'snrk' in st.features:
if data_dmdt.shape[0] > 1:
spec_std = data_dmdt.real.mean(axis=3).mean(axis=1).std(axis=0)
else:
spec_std = data_dmdt[0].real.mean(axis=2).std(axis=0)
if sig_ts is None or kalman_coeffs is None:
sig_ts, kalman_coeffs = kalman_prepare_coeffs(spec_std)
else:
spec_std, sig_ts, kalman_coeffs = None, None, None
# fftmode = 'fftw' if cpuonly else st.fftmode # can't remember why i did this!
image = rfpipe.search.grid_image(data_dmdt, uvw, st.npixx, st.npixy, st.uvres,
'fftw', st.prefs.nthread, wisdom=wisdom,
integrations=[candint])[0]
# TODO: allow dl,dm as args and reproduce detection for other SNRs
dl, dm = st.pixtolm(np.where(image == image.max()))
util.phase_shift(data_dmdt, uvw, dl, dm)
dataph = data_dmdt[max(0, candint-st.prefs.timewindow//2):candint+st.prefs.timewindow//2].mean(axis=1)
util.phase_shift(data_dmdt, uvw, -dl, -dm)
spec = data_dmdt.real.mean(axis=3).mean(axis=1)[candloc[1]]
if 'snrk' in st.features:
significance_kalman = -kalman_significance(spec, spec_std,
sig_ts=sig_ts,
coeffs=kalman_coeffs)
snrk = (2*significance_kalman)**0.5
logger.info("Calculated snrk of {0} after detection. Adding it to CandData.".format(snrk))
kwargs['snrk'] = snrk
canddata = candidates.CandData(state=st, loc=tuple(candloc), image=image,
data=dataph, **kwargs)
# output is as from searching functions
return canddata
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 pipeline_imdata(st, candloc, data_dmdt=None, cpuonly=False, **kwargs):
""" Generate image and phased visibility data for candloc.
Phases to peak pixel in image of candidate.
Can optionally pass in corrected data, if available.
"""
segment, candint, dmind, dtind, beamnum = candloc
dt = st.dtarr[dtind]
dm = st.dmarr[dmind]
uvw = util.get_uvw_segment(st, segment)
wisdom = rfpipe.search.set_wisdom(st.npixx, st.npixy)
if data_dmdt is None:
data_dmdt = pipeline_datacorrect(st, candloc)
fftmode = 'fftw' if cpuonly else st.fftmode
image = rfpipe.search.grid_image(data_dmdt,
uvw,
st.npixx,
st.npixy,
st.uvres,
fftmode,
st.prefs.nthread,
wisdom=wisdom,
integrations=[candint])[0]
# TODO: allow dl,dm as args and reproduce detection for other SNRs
dl, dm = st.pixtolm(np.where(image == image.max()))
util.phase_shift(data_dmdt, uvw, dl, dm)
dataph = data_dmdt[max(0, candint - st.prefs.timewindow // 2):candint +
st.prefs.timewindow // 2].mean(axis=1)
util.phase_shift(data_dmdt, uvw, -dl, -dm)
canddata = candidates.CandData(state=st,
loc=tuple(candloc),
image=image,
data=dataph,
**kwargs)
# output is as from searching functions
return canddata
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 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 search_thresh_armk(st, data, uvw, integrations=None, spec_std=None,
sig_ts=[], coeffs=[]):
"""
"""
if integrations is None:
integrations = list(range(len(data)))
elif isinstance(integrations, int):
integrations = [integrations]
if spec_std is None:
spec_std = data.real.mean(axis=1).mean(axis=2).std(axis=0)
if not len(sig_ts):
sig_ts = [x*np.median(spec_std) for x in [0.3, 0.1, 0.03, 0.01]]
if not len(coeffs):
sig_ts, coeffs = kalman_prepare_coeffs(spec_std, sig_ts)
n_max_cands = 10 # TODO set with function of sigma_arms
u, v, w = uvw
ch0 = 0
u0 = u[:, ch0]
v0 = v[:, ch0]
w0 = w[:, ch0]
order = ['N', 'E', 'W']
T012 = maparms(st=st, u0=u0, v0=v0, order=order)
arm0, arm1, arm2 = image_arms(st, data.take(integrations, axis=0), uvw,
order=order)
# TODO: This is not returning bright simulated transients. Why?
candinds, armlocs, snrarms = thresh_arms(arm0, arm1, arm2, T012,
st.prefs.sigma_arm,
st.prefs.sigma_arms,
n_max_cands)
# kalman filter integrated for now
T01U = maparms(st=st, u0=u0, v0=v0, order=[order[0], order[1]],
e2=(1., 0.))
T01V = maparms(st=st, u0=u0, v0=v0, order=[order[0], order[1]],
e2=(0., 1.))
T12U = maparms(st=st, u0=u0, v0=v0, order=[order[1], order[2]],
e2=(1., 0.))
T12V = maparms(st=st, u0=u0, v0=v0, order=[order[1], order[2]],
e2=(0., 1.))
T20U = maparms(st=st, u0=u0, v0=v0, order=[order[2], order[0]],
e2=(1., 0.))
T20V = maparms(st=st, u0=u0, v0=v0, order=[order[2], order[0]],
e2=(0., 1.))
npix = max(st.npixx_full, st.npixy_full)
kpeaks = []
for i in range(len(candinds)):
kpeak = ()
snrlast = 0. # initialize snr to find max per i
for j in range(n_max_cands):
if snrarms[i, j] > 0.:
spec = data.take([integrations[candinds[i, j]]], axis=0).copy()
armloc0, armloc1, armloc2 = armlocs[i, j]
# find x,y loc from common loc inferred from each arm pair
peakx01 = projectarms(armloc0-npix//2, armloc1-npix//2, T01U,
st.npixx_full)
peaky01 = projectarms(armloc0-npix//2, armloc1-npix//2, T01V,
st.npixy_full)
peakx12 = projectarms(armloc1-npix//2, armloc2-npix//2, T12U,
st.npixx_full)
peaky12 = projectarms(armloc1-npix//2, armloc2-npix//2, T12V,
st.npixy_full)
peakx20 = projectarms(armloc2-npix//2, armloc0-npix//2, T20U,
st.npixx_full)
peaky20 = projectarms(armloc2-npix//2, armloc0-npix//2, T20V,
st.npixy_full)
peakx = np.sort([peakx01, peakx12, peakx20])[1]
peaky = np.sort([peaky01, peaky12, peaky20])[1]
l, m = st.calclm(st.npixx_full, st.npixy_full, st.uvres, peakx,
peaky)
util.phase_shift(spec, uvw, l, m)
spec = spec[0].real.mean(axis=2).mean(axis=0)
significance_kalman = kalman_significance(spec, spec_std,
sig_ts=sig_ts,
coeffs=coeffs)
snrk = (2*significance_kalman)**0.5
snrtot = (snrk**2 + snrarms[i, j]**2)**0.5
if (snrtot > (st.prefs.sigma_kalman**2 + st.prefs.sigma_arms**2)**0.5) and (snrtot > snrlast):
kpeak = (integrations[candinds[i, j]], snrarms[i, j],
snrk, (armloc0, armloc1, armloc2), (peakx, peaky),
(l, m))
snrlast = snrtot
if len(kpeak):
kpeaks.append(kpeak)
return kpeaks
def dedisperse_search_cuda(st, segment, data, devicenum=None):
""" Run dedispersion, resample for all dm and dt.
Grid and image on GPU.
rfgpu is built from separate repo.
Uses state to define integrations to image based on segment, dm, and dt.
devicenum can force the gpu to use, but can be inferred via distributed.
"""
assert st.dtarr[0] == 1, "st.dtarr[0] assumed to be 1"
assert all([st.dtarr[dtind]*2 == st.dtarr[dtind+1]
for dtind in range(len(st.dtarr)-1)]), ("dtarr must increase "
"by factors of 2")
if not np.any(data):
logger.info("Data is all zeros. Skipping search.")
return candidates.CandCollection(prefs=st.prefs,
metadata=st.metadata)
if devicenum is None:
# assume first gpu, but try to infer from worker name
devicenum = 0
try:
from distributed import get_worker
name = get_worker().name
devicenum = int(name.split('g')[1])
logger.debug("Using name {0} to set GPU devicenum to {1}"
.format(name, devicenum))
except IndexError:
logger.warn("Could not parse worker name {0}. Using default GPU devicenum {1}"
.format(name, devicenum))
except ValueError:
logger.warn("No worker found. Using default GPU devicenum {0}"
.format(devicenum))
except ImportError:
logger.warn("distributed not available. Using default GPU devicenum {0}"
.format(devicenum))
rfgpu.cudaSetDevice(devicenum)
beamnum = 0
uvw = util.get_uvw_segment(st, segment)
upix = st.npixx
vpix = st.npixy//2 + 1
grid = rfgpu.Grid(st.nbl, st.nchan, st.readints, upix, vpix)
image = rfgpu.Image(st.npixx, st.npixy)
image.add_stat('rms')
image.add_stat('pix')
# Data buffers on GPU
vis_raw = rfgpu.GPUArrayComplex((st.nbl, st.nchan, st.readints))
vis_grid = rfgpu.GPUArrayComplex((upix, vpix))
img_grid = rfgpu.GPUArrayReal((st.npixx, st.npixy))
# Convert uv from lambda to us
u, v, w = uvw
u_us = 1e6*u[:, 0]/(1e9*st.freq[0])
v_us = 1e6*v[:, 0]/(1e9*st.freq[0])
# Q: set input units to be uv (lambda), freq in GHz?
grid.set_uv(u_us, v_us) # u, v in us
grid.set_freq(st.freq*1e3) # freq in MHz
grid.set_cell(st.uvres) # uv cell size in wavelengths (== 1/FoV(radians))
# Compute gridding transform
grid.compute()
# move Stokes I data in (assumes dual pol data)
vis_raw.data[:] = np.rollaxis(data.mean(axis=3), 0, 3)
vis_raw.h2d() # Send it to GPU memory
grid.conjugate(vis_raw)
# some prep if kalman filter is to be applied
if st.prefs.searchtype in ['imagek']:
# TODO: check that this is ok if pointing at bright source
spec_std = data.real.mean(axis=1).mean(axis=2).std(axis=0)
sig_ts, kalman_coeffs = kalman_prepare_coeffs(spec_std)
if not np.all(sig_ts):
logger.info("sig_ts all zeros. Skipping search.")
return candidates.CandCollection(prefs=st.prefs,
metadata=st.metadata)
# place to hold intermediate result lists
canddict = {}
canddict['candloc'] = []
for feat in st.features:
canddict[feat] = []
for dtind in range(len(st.dtarr)):
if dtind > 0:
grid.downsample(vis_raw)
for dmind in range(len(st.dmarr)):
delay = util.calc_delay(st.freq, st.freq.max(), st.dmarr[dmind],
st.inttime)
grid.set_shift(delay >> dtind) # dispersion shift per chan in samples
integrations = st.get_search_ints(segment, dmind, dtind)
if len(integrations) == 0:
continue
minint = min(integrations)
maxint = max(integrations)
logger.info('Imaging {0} ints ({1}-{2}) in seg {3} at DM/dt {4:.1f}/{5}'
' with image {6}x{7} (uvres {8}) with gpu {9}'
.format(len(integrations), minint, maxint, segment,
st.dmarr[dmind], st.dtarr[dtind], st.npixx,
st.npixy, st.uvres, devicenum))
for i in integrations:
# grid and FFT
grid.operate(vis_raw, vis_grid, i)
image.operate(vis_grid, img_grid)
# calc snr
stats = image.stats(img_grid)
if stats['rms'] != 0.:
snr1 = stats['max']/stats['rms']
else:
snr1 = 0.
logger.warn("rfgpu rms is 0 in int {0}. Skipping.".format(i))
# threshold image
if snr1 > st.prefs.sigma_image1:
candloc = (segment, i, dmind, dtind, beamnum)
xpeak = stats['xpeak']
ypeak = stats['ypeak']
l1, m1 = st.pixtolm((xpeak+st.npixx//2, ypeak+st.npixy//2))
if st.prefs.searchtype == 'image':
logger.info("Got one! SNR1 {0:.1f} candidate at {1} and (l, m) = ({2},{3})"
.format(snr1, candloc, l1, m1))
canddict['candloc'].append(candloc)
canddict['l1'].append(l1)
canddict['m1'].append(m1)
canddict['snr1'].append(snr1)
canddict['immax1'].append(stats['max'])
elif st.prefs.searchtype == 'imagek':
# TODO: implement phasing on GPU
data_corr = dedisperseresample(data, delay,
st.dtarr[dtind],
parallel=st.prefs.nthread > 1,
resamplefirst=st.fftmode=='cuda')
spec = data_corr.take([i], axis=0)
util.phase_shift(spec, uvw, l1, m1)
spec = spec[0].real.mean(axis=2).mean(axis=0)
# TODO: this significance can be biased low if averaging in long baselines that are not phased well
# TODO: spec should be calculated from baselines used to measure l,m?
significance_kalman = kalman_significance(spec,
spec_std,
sig_ts=sig_ts,
coeffs=kalman_coeffs)
snrk = (2*significance_kalman)**0.5
snrtot = (snrk**2 + snr1**2)**0.5
if snrtot > (st.prefs.sigma_kalman**2 + st.prefs.sigma_image1**2)**0.5:
logger.info("Got one! SNR1 {0:.1f} and SNRk {1:.1f} candidate at {2} and (l,m) = ({3},{4})"
.format(snr1, snrk, candloc, l1, m1))
canddict['candloc'].append(candloc)
canddict['l1'].append(l1)
canddict['m1'].append(m1)
canddict['snr1'].append(snr1)
canddict['immax1'].append(stats['max'])
canddict['snrk'].append(snrk)
elif st.prefs.searchtype == 'armkimage':
raise NotImplementedError
elif st.prefs.searchtype == 'armk':
raise NotImplementedError
else:
logger.warn("searchtype {0} not recognized"
.format(st.prefs.searchtype))
cc = candidates.make_candcollection(st, **canddict)
logger.info("First pass found {0} candidates in seg {1}."
.format(len(cc), segment))
if st.prefs.clustercands is not None:
cc = candidates.cluster_candidates(cc)
if st.prefs.savecands or st.prefs.saveplots:
# triggers optional plotting and saving
cc = reproduce_candcollection(cc, data)
candidates.save_cands(st, candcollection=cc)
return cc
def dedisperse_search_fftw(st, segment, data, wisdom=None):
""" Fuse the dediserpse, resample, search, threshold functions.
Returns list of CandData objects that define candidates with
candloc, image, and phased visibility data.
Integrations can define subset of all available in data to search.
Default will take integrations not searched in neighboring segments.
** only supports threshold > image max (no min)
** dmind, dtind, beamnum assumed to represent current state of data
"""
if not np.any(data):
logger.info("Data is all zeros. Skipping search.")
return candidates.CandCollection(prefs=st.prefs,
metadata=st.metadata)
# some prep if kalman filter is to be applied
if st.prefs.searchtype in ['imagek', 'armk', 'armkimage']:
# TODO: check that this is ok if pointing at bright source
spec_std = data.real.mean(axis=1).mean(axis=2).std(axis=0)
sig_ts, kalman_coeffs = kalman_prepare_coeffs(spec_std)
beamnum = 0
uvw = util.get_uvw_segment(st, segment)
# place to hold intermediate result lists
canddict = {}
canddict['candloc'] = []
for feat in st.features:
canddict[feat] = []
for dtind in range(len(st.dtarr)):
for dmind in range(len(st.dmarr)):
# set search integrations
integrations = st.get_search_ints(segment, dmind, dtind)
if len(integrations) == 0:
continue
minint = min(integrations)
maxint = max(integrations)
logger.info('{0} search of {1} ints ({2}-{3}) in seg {4} at DM/dt '
'{5:.1f}/{6} with image {7}x{8} (uvres {9}) with fftw'
.format(st.prefs.searchtype, len(integrations), minint,
maxint, segment, st.dmarr[dmind],
st.dtarr[dtind], st.npixx,
st.npixy, st.uvres))
# correct data
delay = util.calc_delay(st.freq, st.freq.max(), st.dmarr[dmind],
st.inttime)
data_corr = dedisperseresample(data, delay, st.dtarr[dtind],
parallel=st.prefs.nthread > 1,
resamplefirst=st.fftmode=='cuda')
# run search
if st.prefs.searchtype in ['image', 'imagek']:
images = grid_image(data_corr, uvw, st.npixx, st.npixy, st.uvres,
'fftw', st.prefs.nthread, wisdom=wisdom,
integrations=integrations)
for i, image in enumerate(images):
immax1 = image.max()
snr1 = immax1/image.std()
if snr1 > st.prefs.sigma_image1:
candloc = (segment, integrations[i], dmind, dtind, beamnum)
l1, m1 = st.pixtolm(np.where(image == immax1))
# if set, use sigma_kalman as second stage filter
if st.prefs.searchtype == 'imagek':
spec = data_corr.take([integrations[i]], axis=0)
util.phase_shift(spec, uvw, l1, m1)
spec = spec[0].real.mean(axis=2).mean(axis=0)
# TODO: this significance can be biased low if averaging in long baselines that are not phased well
# TODO: spec should be calculated from baselines used to measure l,m?
significance_kalman = kalman_significance(spec,
spec_std,
sig_ts=sig_ts,
coeffs=kalman_coeffs)
snrk = (2*significance_kalman)**0.5
snrtot = (snrk**2 + snr1**2)**0.5
if snrtot > (st.prefs.sigma_kalman**2 + st.prefs.sigma_image1**2)**0.5:
logger.info("Got one! SNR1 {0:.1f} and SNRk {1:.1f} candidate at {2} and (l,m) = ({3},{4})"
.format(snr1, snrk, candloc, l1, m1))
canddict['candloc'].append(candloc)
canddict['l1'].append(l1)
canddict['m1'].append(m1)
canddict['snr1'].append(snr1)
canddict['immax1'].append(immax1)
canddict['snrk'].append(snrk)
elif st.prefs.searchtype == 'image':
logger.info("Got one! SNR1 {0:.1f} candidate at {1} and (l, m) = ({2},{3})"
.format(snr1, candloc, l1, m1))
canddict['candloc'].append(candloc)
canddict['l1'].append(l1)
canddict['m1'].append(m1)
canddict['snr1'].append(snr1)
canddict['immax1'].append(immax1)
elif st.prefs.searchtype in ['armkimage', 'armk']:
armk_candidates = search_thresh_armk(st, data_corr, uvw,
integrations=integrations,
spec_std=spec_std,
sig_ts=sig_ts,
coeffs=kalman_coeffs)
for candind, snrarms, snrk, armloc, peakxy, lm in armk_candidates:
candloc = (segment, candind, dmind, dtind, beamnum)
# if set, use sigma_kalman as second stage filter
if st.prefs.searchtype == 'armkimage':
image = grid_image(data_corr, uvw, st.npixx_full,
st.npixy_full, st.uvres, 'fftw',
st.prefs.nthread,
wisdom=wisdom, integrations=candind)
peakx, peaky = np.where(image[0] == image[0].max())
l1, m1 = st.calclm(st.npixx_full, st.npixy_full,
st.uvres, peakx[0], peaky[0])
immax1 = image.max()
snr1 = immax1/image.std()
if snr1 > st.prefs.sigma_image1:
logger.info("Got one! SNRarms {0:.1f} and SNRk "
"{1:.1f} and SNR1 {2:.1f} candidate at"
" {3} and (l,m) = ({4},{5})"
.format(snrarms, snrk, snr1,
candloc, l1, m1))
canddict['candloc'].append(candloc)
canddict['l1'].append(l1)
canddict['m1'].append(m1)
canddict['snrarms'].append(snrarms)
canddict['snrk'].append(snrk)
canddict['snr1'].append(snr1)
canddict['immax1'].append(immax1)
elif st.prefs.searchtype == 'armk':
l1, m1 = lm
logger.info("Got one! SNRarms {0:.1f} and SNRk {1:.1f} "
"candidate at {2} and (l,m) = ({3},{4})"
.format(snrarms, snrk, candloc, l1, m1))
canddict['candloc'].append(candloc)
canddict['l1'].append(l1)
canddict['m1'].append(m1)
canddict['snrarms'].append(snrarms)
canddict['snrk'].append(snrk)
else:
raise NotImplemented("only searchtype=image, imagek, armk, armkimage implemented")
cc = candidates.make_candcollection(st, **canddict)
logger.info("First pass found {0} candidates in seg {1}."
.format(len(cc), segment))
if st.prefs.clustercands is not None:
cc = candidates.cluster_candidates(cc)
if st.prefs.savecands or st.prefs.saveplots:
# triggers optional plotting and saving
cc = reproduce_candcollection(cc, data)
candidates.save_cands(st, candcollection=cc)
return cc
def pipeline_canddata(st,
candloc,
data_dmdt=None,
spec_std=None,
cpuonly=False,
sig_ts=[],
kalman_coeffs=[],
**kwargs):
""" Generate image and phased visibility data for candloc.
Phases to peak pixel in image of candidate.
Can optionally pass in corrected data, if available.
cpuonly argument not being used at present.
"""
from rfpipe import candidates, util
import rfpipe.search
segment, candint, dmind, dtind, beamnum = candloc
dt = st.dtarr[dtind]
dm = st.dmarr[dmind]
pc = st.get_pc(segment)
uvw = util.get_uvw_segment(st, segment, pc_mjd=pc, pc_radec=pc)
wisdom = rfpipe.search.set_wisdom(st.npixx, st.npixy)
if data_dmdt is None:
data_dmdt = pipeline_datacorrect(st, candloc)
if ('snrk' in st.features and 'snrk' not in kwargs and
(spec_std is None or not len(sig_ts) or not len(kalman_coeffs))):
spec_std, sig_ts, kalman_coeffs = util.kalman_prep(data_dmdt)
# fftmode = 'fftw' if cpuonly else st.fftmode # can't remember why i did this!
image = rfpipe.search.grid_image(data_dmdt,
uvw,
st.npixx,
st.npixy,
st.uvres,
'fftw',
st.prefs.nthread,
wisdom=wisdom,
integrations=[candint])[0]
# TODO: allow dl,dm as args and reproduce detection for other SNRs
dl, dm = st.pixtolm(np.where(image == image.max()))
util.phase_shift(data_dmdt, uvw=uvw, dl=dl, dm=dm)
dataph = data_dmdt[max(0, candint - st.prefs.timewindow // 2):candint +
st.prefs.timewindow // 2].mean(axis=1)
util.phase_shift(data_dmdt, uvw=uvw, dl=-dl, dm=-dm)
# TODO: This probably needs to be masked to avoid averaging zeros in
spec = data_dmdt.real.mean(axis=3).mean(axis=1)[candloc[1]]
if 'snrk' in st.features and 'snrk' not in kwargs:
if np.count_nonzero(spec) / len(spec) > 1 - st.prefs.max_zerofrac:
significance_kalman = -kalman_significance(
spec, spec_std, sig_ts=sig_ts, coeffs=kalman_coeffs)
snrk = (2 * significance_kalman)**0.5
else:
logger.warning("snrk set to 0, since {0}/{1} are zeroed".format(
len(spec) - np.count_nonzero(spec), len(spec)))
snrk = 0.
logger.info(
"Calculated snrk of {0} after detection. Adding it to CandData.".
format(snrk))
kwargs['snrk'] = snrk
canddata = candidates.CandData(state=st,
loc=tuple(candloc),
image=image,
data=dataph,
**kwargs)
# output is as from searching functions
return canddata
def pipeline_canddata(st,
candloc,
data_dmdt=None,
cpuonly=False,
sig_ts=None,
kalman_coeffs=None,
**kwargs):
""" Generate image and phased visibility data for candloc.
Phases to peak pixel in image of candidate.
Can optionally pass in corrected data, if available.
cpuonly argument not being used at present.
"""
from rfpipe import candidates, util
import rfpipe.search
segment, candint, dmind, dtind, beamnum = candloc
dt = st.dtarr[dtind]
dm = st.dmarr[dmind]
uvw = util.get_uvw_segment(st, segment)
wisdom = rfpipe.search.set_wisdom(st.npixx, st.npixy)
if data_dmdt is None:
data_dmdt = pipeline_datacorrect(st, candloc)
if 'snrk' in st.features:
if data_dmdt.shape[0] > 1:
spec_std = data_dmdt.real.mean(axis=3).mean(axis=1).std(axis=0)
else:
spec_std = data_dmdt[0].real.mean(axis=2).std(axis=0)
if not np.any(spec_std):
logger.warning("spectrum std all zeros. Not estimating coeffs.")
kalman_coeffs = []
else:
sig_ts, kalman_coeffs = kalman_prepare_coeffs(spec_std)
if not np.all(np.nan_to_num(sig_ts)):
kalman_coeffs = []
else:
spec_std, sig_ts, kalman_coeffs = None, None, None
# fftmode = 'fftw' if cpuonly else st.fftmode # can't remember why i did this!
image = rfpipe.search.grid_image(data_dmdt,
uvw,
st.npixx,
st.npixy,
st.uvres,
'fftw',
st.prefs.nthread,
wisdom=wisdom,
integrations=[candint])[0]
# TODO: allow dl,dm as args and reproduce detection for other SNRs
dl, dm = st.pixtolm(np.where(image == image.max()))
util.phase_shift(data_dmdt, uvw, dl, dm)
dataph = data_dmdt[max(0, candint - st.prefs.timewindow // 2):candint +
st.prefs.timewindow // 2].mean(axis=1)
util.phase_shift(data_dmdt, uvw, -dl, -dm)
spec = data_dmdt.real.mean(axis=3).mean(axis=1)[candloc[1]]
if 'snrk' in st.features:
significance_kalman = -kalman_significance(
spec, spec_std, sig_ts=sig_ts, coeffs=kalman_coeffs)
snrk = (2 * significance_kalman)**0.5
logger.info(
"Calculated snrk of {0} after detection. Adding it to CandData.".
format(snrk))
kwargs['snrk'] = snrk
canddata = candidates.CandData(state=st,
loc=tuple(candloc),
image=image,
data=dataph,
**kwargs)
# output is as from searching functions
return canddata