def compute_jacobian_realization(fr_module, num_samples, signal_dim, num_freq,
min_sep, fixed_freq, snr):
clean_signals, signal_frs, signal_num_freqs = gen_signal(
num_samples=num_samples,
signal_dim=signal_dim,
num_freq=num_freq,
min_sep=min_sep,
fixed_freq=fixed_freq)
noisy_signals = noise_torch(torch.as_tensor(clean_signals), snr,
'gaussian')
jacobian_realization = []
jacobian_realization_clean = []
for idx in range(clean_signals.shape[0]):
clean_signal, signal_fr, signal_num_freq = clean_signals[
idx], signal_frs[idx], signal_num_freqs[idx]
noisy_signal = noisy_signals[idx].cpu().numpy()
clean_signal_t = clean_signal[0] + clean_signal[1] * 1j
clean_signal_fft = np.fft.fft(clean_signal_t, n=1000)
clean_signal_fft = np.fft.fftshift(clean_signal_fft)
noisy_signal_t = noisy_signal[0] + 1j * noisy_signal[1]
noisy_signal_fft = np.fft.fft(noisy_signal_t, n=1000)
noisy_signal_fft = np.fft.fftshift(noisy_signal_fft)
noisy_signal = torch.as_tensor(noisy_signal).unsqueeze(dim=0)
jacobian, inputs, outputs = compute_jacobian_and_bias(
noisy_signal, fr_module)
clean_signal = torch.as_tensor(clean_signal).unsqueeze(dim=0)
jacobian_clean, inputs, outputs = compute_jacobian_and_bias(
clean_signal, fr_module)
jacobian_realization.append(jacobian)
jacobian_realization_clean.append(jacobian_clean)
return jacobian_realization, jacobian_realization_clean
def train_frequency_representation(args, fr_module, fr_optimizer, fr_criterion,
fr_scheduler, train_loader, val_loader,
xgrid, epoch, tb_writer):
"""
Train the frequency-representation module for one epoch
"""
epoch_start_time = time.time()
fr_module.train()
loss_train_fr = 0
for batch_idx, (clean_signal, target_fr, freq) in enumerate(train_loader):
if args.use_cuda:
clean_signal, target_fr = clean_signal.cuda(), target_fr.cuda()
snr = np.random.uniform(args.snrl, args.snrh + 1)
noisy_signal = noise_torch(clean_signal, snr, args.noise)
fr_optimizer.zero_grad()
output_fr = fr_module(noisy_signal)
loss_fr = fr_criterion(output_fr, target_fr)
loss_fr.backward()
fr_optimizer.step()
loss_train_fr += loss_fr.data.item()
fr_module.eval()
loss_val_fr, fnr_val = 0, 0
for batch_idx, (noisy_signal, _, target_fr, freq) in enumerate(val_loader):
if args.use_cuda:
noisy_signal, target_fr = noisy_signal.cuda(), target_fr.cuda()
with torch.no_grad():
output_fr = fr_module(noisy_signal)
loss_fr = fr_criterion(output_fr, target_fr)
loss_val_fr += loss_fr.data.item()
nfreq = (freq >= -0.5).sum(dim=1)
f_hat = fr.find_freq(output_fr.cpu().detach().numpy(), nfreq, xgrid)
fnr_val += fnr(f_hat, freq.cpu().numpy(), args.signal_dim)
loss_train_fr /= (args.n_training * (args.snrh - args.snrl + 1))
loss_val_fr /= (args.n_validation * (args.snrh - args.snrl + 1))
fnr_val *= 100 / (args.n_validation * (args.snrh - args.snrl + 1))
tb_writer.add_scalar('fr_l2_training', loss_train_fr, epoch)
tb_writer.add_scalar('fr_l2_validation', loss_val_fr, epoch)
tb_writer.add_scalar('fr_FNR', fnr_val, epoch)
fr_scheduler.step(loss_val_fr)
logger.info(
"Epochs: %d / %d, Time: %.1f, FR training L2 loss %.2f, FR validation L2 loss %.2f, FNR %.2f %%",
epoch, args.n_epochs_fr + args.n_epochs_fc,
time.time() - epoch_start_time, loss_train_fr, loss_val_fr, fnr_val)
def train_frequency_counting(args, fr_module, fc_module, fc_optimizer,
fc_criterion, fc_scheduler, train_loader,
val_loader, epoch, tb_writer):
"""
Train the frequency-counting module for one epoch
"""
epoch_start_time = time.time()
fr_module.eval()
fc_module.train()
loss_train_fc, acc_train_fc = 0, 0
for batch_idx, (clean_signal, target_fr, freq) in enumerate(train_loader):
if args.use_cuda:
clean_signal, target_fr, freq = clean_signal.cuda(
), target_fr.cuda(), freq.cuda()
snr = np.random.uniform(args.snrl, args.snrh + 1)
noisy_signal = noise_torch(clean_signal, snr, args.noise)
nfreq = (freq >= -0.5).sum(dim=1)
if args.use_cuda:
nfreq = nfreq.cuda()
if args.fc_module_type == 'classification':
nfreq = nfreq - 1
with torch.no_grad():
output_fr = fr_module(noisy_signal)
output_fr = output_fr.detach()
output_fc = fc_module(output_fr)
if args.fc_module_type == 'regression':
output_fc = output_fc.view(output_fc.size(0))
nfreq = nfreq.float()
loss_fc = fc_criterion(output_fc, nfreq)
if args.fc_module_type == 'classification':
estimate = output_fc.max(1)[1]
else:
estimate = torch.round(output_fc)
acc_train_fc += estimate.eq(nfreq).sum().cpu().item()
loss_train_fc += loss_fc.data.item()
fc_optimizer.zero_grad()
loss_fc.backward()
fc_optimizer.step()
loss_train_fc /= args.n_training
acc_train_fc *= 100 / args.n_training
fc_module.eval()
loss_val_fc = 0
acc_val_fc = 0
for batch_idx, (noisy_signal, _, target_fr, freq) in enumerate(val_loader):
if args.use_cuda:
noisy_signal, target_fr = noisy_signal.cuda(), target_fr.cuda()
nfreq = (freq >= -0.5).sum(dim=1)
if args.use_cuda:
nfreq = nfreq.cuda()
if args.fc_module_type == 'classification':
nfreq = nfreq - 1
with torch.no_grad():
output_fr = fr_module(noisy_signal)
output_fc = fc_module(output_fr)
if args.fc_module_type == 'regression':
output_fc = output_fc.view(output_fc.size(0))
nfreq = nfreq.float()
loss_fc = fc_criterion(output_fc, nfreq)
if args.fc_module_type == 'regression':
estimate = torch.round(output_fc)
elif args.fc_module_type == 'classification':
estimate = torch.argmax(output_fc, dim=1)
acc_val_fc += estimate.eq(nfreq).sum().item()
loss_val_fc += loss_fc.data.item()
loss_val_fc /= args.n_validation
acc_val_fc *= 100 / args.n_validation
fc_scheduler.step(acc_val_fc)
tb_writer.add_scalar('fc_loss_training', loss_train_fc,
epoch - args.n_epochs_fr)
tb_writer.add_scalar('fc_loss_validation', loss_val_fc,
epoch - args.n_epochs_fr)
tb_writer.add_scalar('fc_accuracy_training', acc_train_fc,
epoch - args.n_epochs_fr)
tb_writer.add_scalar('fc_accuracy_validation', acc_val_fc,
epoch - args.n_epochs_fr)
logger.info(
"Epochs: %d / %d, Time: %.1f, Training fc loss: %.2f, Vadidation fc loss: %.2f, "
"Training accuracy: %.2f %%, Validation accuracy: %.2f %%", epoch,
args.n_epochs_fr + args.n_epochs_fc,
time.time() - epoch_start_time, loss_train_fc, loss_val_fc,
acc_train_fc, acc_val_fc)
def generate_report_plots(num_samples=2,
signal_dim=50,
min_sep=1.,
snr=30,
fixed_freq=[0.4],
save=False):
if not os.path.exists('./plots'):
os.mkdir('./plots')
save_dir = './plots'
num_freq = len(fixed_freq)
for signal_idx in range(num_samples):
clean_signals, signal_frs, signal_num_freqs = gen_signal(
num_samples=num_samples,
signal_dim=signal_dim,
num_freq=num_freq,
min_sep=min_sep,
fixed_freq=fixed_freq)
noisy_signals = noise_torch(torch.as_tensor(clean_signals), snr,
'gaussian')
clean_signal, signal_fr, signal_num_freq = clean_signals[
signal_idx], signal_frs[signal_idx], signal_num_freqs[signal_idx]
noisy_signal = noisy_signals[signal_idx].cpu().numpy()
clean_signal_t = clean_signal[0] + clean_signal[1] * 1j
clean_signal_fft = np.fft.fft(clean_signal_t, n=1000)
clean_signal_fft = np.fft.fftshift(clean_signal_fft)
noisy_signal_t = noisy_signal[0] + 1j * noisy_signal[1]
noisy_signal_fft = np.fft.fft(noisy_signal_t, n=1000)
noisy_signal_fft = np.fft.fftshift(noisy_signal_fft)
if signal_idx == 0: # plot clean first
noisy_signal = torch.as_tensor(clean_signal).unsqueeze(dim=0)
else:
noisy_signal = torch.as_tensor(noisy_signal).unsqueeze(dim=0)
jacobian, inputs, outputs = compute_jacobian_and_bias(
noisy_signal, fr_module)
fft_filter = jacobian[0] - 1j * jacobian[1]
# plot 1
fig1, ax = plt.subplots(3, 1, figsize=(15, 6))
xgrid = np.linspace(-0.5, 0.5, fr_module.fr_size, endpoint=False)
ax[0].plot(xgrid, outputs[0])
ax[0].set_xticks(np.arange(-0.5, 0.5, 0.2))
ylim = ax[0].get_ylim()
for i in range(signal_fr.shape[0]):
ax[0].vlines(signal_fr[i],
ymin=ylim[0],
ymax=ylim[1],
label='target:{:4.2f}'.format(signal_fr[i]))
ax[0].legend()
ax[0].set_xlim(-0.5, 0.5)
ax[1].plot(xgrid, np.abs(clean_signal_fft), label='clean fft')
ax[1].plot(xgrid, np.abs(noisy_signal_fft), '--', label='noisy fft')
ax[1].set_xlim(-0.5, 0.5)
ylim = ax[1].get_ylim()
for i in range(signal_fr.shape[0]):
ax[1].vlines(signal_fr[i],
ymin=ylim[0],
ymax=ylim[1],
label='target:{:4.2f}'.format(signal_fr[i]))
ax[1].legend()
im = ax[2].imshow(np.abs(fft_filter))
ax[2].set_aspect(2.2)
ax[2].set_ylabel('jacobian', fontsize=13)
fig1.subplots_adjust(right=0.9)
cbar_ax = fig1.add_axes([0.91, 0.15, 0.01, 0.7])
fig1.colorbar(im, cax=cbar_ax)
plt.show()
# plot 2
fig2, ax = plt.subplots(3, 1, figsize=(15, 6))
xgrid = np.linspace(-0.5, 0.5, fr_module.fr_size, endpoint=False)
ax[0].plot(xgrid, outputs[0])
ax[0].set_xticks(np.arange(-0.5, 0.5, 0.2))
ylim = ax[0].get_ylim()
for i in range(signal_fr.shape[0]):
ax[0].vlines(signal_fr[i],
ymin=ylim[0],
ymax=ylim[1],
label='target:{:4.2f}'.format(signal_fr[i]))
ax[0].legend()
ax[1].plot(xgrid, np.abs(clean_signal_fft), label='clean fft')
ax[1].plot(xgrid, np.abs(noisy_signal_fft), '--', label='noisy fft')
ylim = ax[1].get_ylim()
for i in range(signal_fr.shape[0]):
ax[1].vlines(signal_fr[i],
ymin=ylim[0],
ymax=ylim[1],
label='target:{:4.2f}'.format(signal_fr[i]))
ax[1].legend()
fft_filter_norm = fft_filter * np.conjugate(fft_filter)
fft_filter_norm = fft_filter_norm.T.sum(axis=1).real
ax[2].plot(xgrid, fft_filter_norm)
ax[0].set_xlim(-0.5, 0.5)
ax[1].set_xlim(-0.5, 0.5)
ax[2].set_xlim(-0.5, 0.5)
fig2.subplots_adjust(right=0.9)
plt.show()
# plot. 3
indices = find_neariest_idx(signal_fr, xgrid)
target_filter = fft_filter.T[indices]
# time domain
fig3, ax = plt.subplots(target_filter.shape[0], 2)
for idx, filt in enumerate(target_filter):
if target_filter.shape[0] > 1:
ax[idx, 0].plot(filt.real)
ax[idx, 1].plot(filt.imag)
ax[idx, 0].set_title('Real Part of Freq:{:4.2f};'.format(
signal_fr[idx]))
ax[idx, 1].set_title('Complex Part of Freq:{:4.2f};'.format(
signal_fr[idx]))
else:
ax[0].plot(filt.real)
ax[1].plot(filt.imag)
ax[0].set_title('Real Part of Freq:{:4.2f};'.format(
signal_fr[idx]))
ax[1].set_title('Complex Part of Freq:{:4.2f};'.format(
signal_fr[idx]))
plt.tight_layout()
plt.show()
# fft domian of signals
fig4, ax = plt.subplots(target_filter.shape[0], 1, dpi=300)
for idx, filt in enumerate(target_filter):
filt_fft = np.fft.fft(filt, n=1000)
filt_fft = np.fft.fftshift(filt_fft)
magnitude = np.abs(filt_fft)
if target_filter.shape[0] > 1:
ax[idx].plot(xgrid, magnitude)
ax[idx].plot(signal_fr[idx], magnitude[indices[idx]], '*')
ax[idx].plot(-signal_fr[idx], magnitude[999 - indices[idx]],
'*')
ax[idx].set_title('Freq:{:4.2f}'.format(signal_fr[idx]))
else:
ax.plot(xgrid, magnitude)
ax.plot(signal_fr[idx], magnitude[indices[idx]], '*')
ax.plot(-signal_fr[idx], magnitude[999 - indices[idx]], '*')
ax.set_title('Freq:{:4.2f}'.format(signal_fr[idx]))
plt.tight_layout()
plt.show()
if save:
file_name = 'output_clean_{}.pdf' if signal_idx == 0 else 'output_{}.pdf'
pdf = matplotlib.backends.backend_pdf.PdfPages(
"./plots/" + file_name.format(signal_idx))
pdf.savefig(fig1)
pdf.savefig(fig2)
pdf.savefig(fig3)
pdf.savefig(fig4)
pdf.close()
.format(args.output_dir))
elif not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
with open(os.path.join(args.output_dir, 'data.args'), 'w') as f:
json.dump(args.__dict__, f, indent=2)
np.random.seed(args.numpy_seed)
torch.manual_seed(args.torch_seed)
s, f, nfreq = gen_signal(num_samples=args.n_test,
signal_dim=args.signal_dimension,
num_freq=args.max_freq,
min_sep=args.minimum_separation,
distance=args.distance,
amplitude=args.amplitude,
floor_amplitude=args.floor_amplitude,
variable_num_freq=True)
np.save(os.path.join(args.output_dir, 'infdB'), s)
np.save(os.path.join(args.output_dir, 'f'), f)
eval_snrs = [np.exp(np.log(10) * float(x) / 10) for x in args.dB]
for k, snr in enumerate(eval_snrs):
noisy_signals = noise.noise_torch(torch.tensor(s), snr,
'gaussian').cpu()
np.save(
os.path.join(args.output_dir, '{}dB'.format(float(args.dB[k]))),
noisy_signals)