def test_gpu_dedisp_dmt_crop():
file = os.path.join(_install_dir, 'data/28.fil')
cand = Candidate(fp=file,
dm=10,
tcand=2.0288800,
width=2,
label=-1,
snr=16.8128,
min_samp=256,
device=0)
cand.get_chunk()
cand = gpu_dedisp_and_dmt_crop(cand)
g_dmt = cand.dmt
g_dedisp = cand.dedispersed
assert cand.dedispersed.shape[0] == 256
assert cand.dmt.shape[1] == 256
cand.dedisperse()
cand.dmtime()
crop_start_sample_ft = cand.dedispersed.shape[0] // 2 - 256 // 2
crop_start_sample_dmt = cand.dmt.shape[1] // 2 - 256 // 2
c_dmt = crop(cand.dmt, crop_start_sample_dmt, 256, 1)
c_dedisp = crop(cand.dedispersed, crop_start_sample_ft, 256, 0)
assert np.isclose(np.sum(g_dmt - c_dmt), 0, atol=1)
assert np.isclose(np.sum(g_dedisp - c_dedisp), 0, atol=1)
def cand2h5(cand_val):
"""
TODO: Add option to use cand.resize for reshaping FT and DMT
Generates h5 file of candidate with resized frequency-time and DM-time arrays
:param cand_val: List of candidate parameters (filename, snr, width, dm, label, tcand(s))
:type cand_val: Candidate
:return: None
"""
filename, snr, width, dm, label, tcand, kill_mask_path, args, gpu_id = cand_val
if os.path.exists(str(kill_mask_path)):
logger.info(f'Using mask {kill_mask_path}')
kill_chans = np.loadtxt(kill_mask_path, dtype=np.int)
else:
logger.debug('No Kill Mask')
fname, ext = os.path.splitext(filename)
if ext == ".fits" or ext == ".sf":
files = glob.glob(fname[:-5] + '*fits')
elif ext == '.fil':
files = [filename]
else:
raise TypeError("Can only work with list of fits file or filterbanks")
logger.debug(f'Source file list: {files}')
cand = Candidate(files, snr=snr, width=width, dm=dm, label=label, tcand=tcand, device=gpu_id)
if os.path.exists(str(kill_mask_path)):
kill_mask = np.zeros(cand.nchans, dtype=np.bool)
kill_mask[kill_chans] = True
cand.kill_mask = kill_mask
cand.get_chunk()
if cand.isfil:
cand.fp.close()
logger.info('Got Chunk')
if gpu_id >= 0:
logger.debug(f"Using the GPU {gpu_id}")
try:
cand = gpu_dedisp_and_dmt_crop(cand, device=gpu_id)
except CudaAPIError:
logger.info("Ran into a CudaAPIError, using the CPU version for this candidate")
cand = cpu_dedisp_dmt(cand, args)
else:
cand = cpu_dedisp_dmt(cand, args)
cand.resize(key='ft', size=args.frequency_size, axis=1, anti_aliasing=True, mode='constant')
logger.info(f'Resized Frequency axis of FT to fsize: {cand.dedispersed.shape[1]}')
cand.dmt = normalise(cand.dmt)
cand.dedispersed = normalise(cand.dedispersed)
fout = cand.save_h5(out_dir=args.fout)
logger.debug(f"Filesize is {os.path.getsize(fout)}")
if not os.path.isfile(fout):
raise IOError(f"File with {cand.id} not written")
if os.path.getsize(fout) < 200 * 1024:
raise ValueError(f"File with {cand.id} has issues!")
logger.info(fout)
return None
def test_candidate_chunk_and_dedispersion():
fil_file = os.path.join(_install_dir, 'data/28.fil')
cand = Candidate(fp=fil_file, dm=475.28400, tcand=2.0288800, width=2, label=-1, snr=16.8128, min_samp=256, device=0)
cand.get_chunk()
assert np.isclose(np.mean(cand.data), 128, atol=1)
cand.dedisperse()
assert np.isclose(np.max(cand.dedispersed.T.sum(0)), 47527, atol=1)
assert np.isclose(np.max(cand.dedispersets()), 47527, atol=1)
cand.dmtime()
assert cand.dmt.shape[0] == 256
fnout = cand.save_h5()
assert os.path.isfile(fnout)
os.remove(fnout)
def test_gpu_dedisperse():
file = os.path.join(_install_dir, 'data/28.fil')
cand = Candidate(fp=file,
dm=475.28400,
tcand=2.0288800,
width=2,
label=-1,
snr=16.8128,
min_samp=256,
device=0)
cand.get_chunk()
cand.dedisperse(target='GPU')
g_dedisp = cand.dedispersed
cand.dedisperse(target='CPU')
c_dedisp = cand.dedispersed
assert np.isclose(np.mean(g_dedisp - c_dedisp), 0, atol=1)
assert np.isclose(np.max(cand.dedispersed.T.sum(0)), 47527, atol=1)
def test_gpu_dmt():
file = os.path.join(_install_dir, 'data/28.fil')
cand = Candidate(fp=file,
dm=10,
tcand=2.0288800,
width=2,
label=-1,
snr=16.8128,
min_samp=256,
device=0)
cand.get_chunk()
cand.dmtime(target='GPU')
g_dmt = cand.dmt
cand.dmtime(target='CPU')
c_dmt = cand.dmt
assert cand.dmt.shape[0] == 256
assert np.isclose(np.mean(g_dmt - c_dmt), 0, atol=1)
assert np.max(g_dmt - c_dmt) / np.max(g_dmt) < 0.05
def test_kill_mask():
fil_file = os.path.join(_install_dir, "data/28.fil")
km = np.zeros(336, dtype="bool")
km[[10, 12, 25, 100, 300]] = True
cand = Candidate(
fp=fil_file,
dm=475.28400,
tcand=2.0288800,
width=2,
label=-1,
snr=16.8128,
min_samp=256,
device=0,
flag_rfi=False,
kill_mask=km,
)
cand.get_chunk()
assert cand.data[:, cand.kill_mask].sum() == 0
assert cand.data[:, [10, 12, 300]].sum() == 0
assert cand.data[:, ~cand.kill_mask].sum() != 0
def test_rfi_mask():
fil_file = os.path.join(_install_dir, "data/28.fil")
cand = Candidate(
fp=fil_file,
dm=475.28400,
tcand=2.0288800,
width=2,
label=-1,
snr=16.8128,
min_samp=256,
device=0,
spectral_kurtosis_sigma=4,
savgol_frequency_window=15,
savgol_sigma=4,
flag_rfi=True,
)
cand.get_chunk()
assert cand.data[:, cand.rfi_mask].sum() == 0
assert cand.data[:, 172:177].sum() == 0
assert cand.data[:, ~cand.rfi_mask].sum() != 0
def cand2h5(cand_val):
"""
TODO: Add option to use cand.resize for reshaping FT and DMT
Generates h5 file of candidate with resized frequency-time and DM-time arrays
:param cand_val: List of candidate parameters (filename, snr, width, dm, label, tcand(s))
:type cand_val: Candidate
:return: None
"""
(
filename,
snr,
width,
dm,
label,
tcand,
kill_mask_path,
num_files,
args,
gpu_id,
) = cand_val
if os.path.exists(str(kill_mask_path)):
logger.info(f"Using mask {kill_mask_path}")
kill_chans = np.loadtxt(kill_mask_path, dtype=np.int)
else:
logger.debug("No Kill Mask")
fname, ext = os.path.splitext(filename)
if ext == ".fits" or ext == ".sf":
if num_files == 1:
files = [filename]
else:
files = glob.glob(fname[:-5] + "*fits")
if len(files) != num_files:
raise ValueError(
"Number of fits files found was not equal to num_files in cand csv."
)
elif ext == ".fil":
files = [filename]
else:
raise TypeError("Can only work with list of fits file or filterbanks")
logger.debug(f"Source file list: {files}")
cand = Candidate(
files,
snr=snr,
width=width,
dm=dm,
label=label,
tcand=tcand,
device=gpu_id,
spectral_kurtosis_sigma=args.spectral_kurtosis_sigma,
savgol_frequency_window=args.savgol_frequency_window,
savgol_sigma=args.savgol_sigma,
flag_rfi=args.flag_rfi,
)
if os.path.exists(str(kill_mask_path)):
kill_mask = np.zeros(cand.nchans, dtype=np.bool)
kill_mask[kill_chans] = True
cand.kill_mask = kill_mask
cand.get_chunk(for_preprocessing=True)
if cand.format == "fil":
cand.fp.close()
logger.info("Got Chunk")
if gpu_id >= 0:
logger.debug(f"Using the GPU {gpu_id}")
try:
cand = gpu_dedisp_and_dmt_crop(cand, device=gpu_id)
except CudaAPIError:
logger.info(
"Ran into a CudaAPIError, using the CPU version for this candidate"
)
cand = cpu_dedisp_dmt(cand, args)
else:
cand = cpu_dedisp_dmt(cand, args)
cand.resize(key="ft",
size=args.frequency_size,
axis=1,
anti_aliasing=True,
mode="constant")
logger.info(
f"Resized Frequency axis of FT to fsize: {cand.dedispersed.shape[1]}")
cand.dmt = normalise(cand.dmt)
cand.dedispersed = normalise(cand.dedispersed)
fout = cand.save_h5(out_dir=args.fout)
logger.debug(f"Filesize of {fout} is {os.path.getsize(fout)}")
if not os.path.isfile(fout):
raise IOError(f"File with {cand.id} not written")
if os.path.getsize(fout) < 100 * 1024:
raise ValueError(
f"File with id: {cand.id} has issues! Its size is too less.")
logger.info(fout)
del cand
return None