<jupyter_start><jupyter_text>load results<jupyter_code>dirs = [
    "db5_saliency_warmstart_seed=1"
]
dics, results, models = load_results(dirs)<jupyter_output><empty_output><jupyter_text>prediction accuracy, compression, and computation time<jupyter_code>r = {
    'Standard Wavelet (DB5)': [],
    'AWD (Ours)': [],
}

best_lam = []
best_wt = []

# params for feature transformer
sgn = "pos"
m = 6

index = 0
res = results[index]
mos = models[index]

# load data
(train_loader, test_loader) = get_dataloader(p.data_path,
                                             batch_size=p.batch_size,
                                             is_continuous=True)

# cross validation
reg_score = []

for i, wt in enumerate(list(dics[index]['wt'].values())):
    wt = wt.to('cpu')
    (X, y), (X_test, y_test) = max_transformer(wt,
                                               train_loader,
                                               test_loader,
                                               sgn=sgn,
                                               m=m)

    clf = LinearRegression()
    scores = cross_val_score(clf, X, y, cv=5, scoring='r2')
    reg_score.append(scores.mean())
reg_score = np.array(reg_score)

# select best wavelet
idx1, idx2 = list(dics[index]['wt'].keys())[np.argmax(reg_score).flatten()[0]]
idx = dics[index]['index'][(idx1, idx2)]
lamL1wave = dics[index]['lamL1wave'][(idx1, idx2)]
lamL1attr = dics[index]['lamL1attr'][(idx1, idx2)]
best_lam.append((lamL1wave, lamL1attr))
wt = dics[index]['wt'][(idx1, idx2)]
best_wt.append(wt)

# load pre-trained model
model = load_pretrained_model(p.model_path, device=device)

p.batch_size = 3000
(train_loader, test_loader) = get_dataloader(p.data_path,
                                             batch_size=p.batch_size,
                                             is_continuous=True)

(X, y), (X_test, y_test) = max_transformer(wt,
                                           train_loader,
                                           test_loader,
                                           sgn=sgn,
                                           m=m)

clf = LinearRegression().fit(X, y)
preds = clf.predict(X_test)
r['AWD (Ours)'].append(metrics.r2_score(y_test, preds))

wt_o = adaptive_wavelets.DWT1d(wave='db5', mode='zero', J=4, init_factor=1, noise_factor=0.0)
(X, y), (X_test, y_test) = max_transformer(wt_o,
                                           train_loader,
                                           test_loader,
                                           sgn=sgn,
                                           m=m)

clf_o = LinearRegression().fit(X, y)
preds = clf_o.predict(X_test)
r['Standard Wavelet (DB5)'].append(metrics.r2_score(y_test, preds))

print(r)
(train_loader, test_loader) = get_dataloader(p.data_path,
                                             batch_size=p.batch_size,
                                             is_continuous=True)

(X, y), (X_test, y_test) = max_transformer(wt,
                                           train_loader,
                                           test_loader,
                                           sgn=sgn,
                                           m=m)

clf = LinearRegression().fit(X, y)
preds = clf.predict(X_test)
print(metrics.r2_score(y_test, preds))<jupyter_output>0.2625306384308951<jupyter_text>Plotting top coefs<jupyter_code>wt = best_wt[0]
wt = wt.to('cpu')
J = wt.J

(Xs, y) = get_tensors(train_loader)
data_t = wt(Xs)
data_to = wt_o(Xs)

X = []
for j in range(J + 1):
    d = data_t[j].detach().squeeze().numpy()
    X.append(max_fun(d, sgn=sgn, m=m))
X = np.hstack(X)
y = y.detach().squeeze().numpy()

idx = 149
def helper(wc: tuple, scale, nums):
    """ 
    Threshold attributions retaining those with top absolute attributions 
    """
    wc = list(wc)
    batch_size = wc[0].shape[0]
    J = len(wc)

    for j in range(J):
        if j != scale:
            wc[j] = torch.zeros_like(wc[j])
        else:
            b = wc[j].reshape(batch_size, -1)
    sort_indexes = torch.argsort(b, dim=1, descending=True)

    m = torch.zeros_like(sort_indexes)
    for i in range(batch_size):
        for num in nums:
            m[i][sort_indexes[i, num]] = 1

    wc[scale] = torch.mul(m.reshape(wc[scale].shape), wc[scale])
    return wc
contr = X[idx] * clf.coef_
idx_pos = np.argwhere(contr > 0).squeeze()
idx_scale = []
idx_contr = []
for j in range(J + 1):
    r = [i - 6 * j for i in idx_pos if 6 * j <= i < 6 * (j + 1)]
    c = [contr[i] for i in idx_pos if 6 * j <= i < 6 * (j + 1)]
    idx_scale.append(r)
    idx_contr.append(c)
fig = plt.figure(figsize=(8, 5), dpi=200)
spec = gridspec.GridSpec(ncols=2, nrows=2, figure=fig)
colors = ['red', 'C0']

x = Xs[idx, 0].squeeze().detach().numpy()

ax_00 = fig.add_subplot(spec[0, 0])
dt = helper(data_t, 0, [0])
x_r = wt.inverse(dt)[idx, 0].squeeze().detach().numpy()
ax_00.plot(x, lw=1, color=colors[0], label="Clathrin amplitude")
ax_00.plot(x_r, lw=1, color=colors[1], alpha=0.5, label="Max feature with top 1 contribution")
ax_00.set_yticklabels([])
ax_00.legend(fontsize=5, loc='upper left', fancybox=True)

ax_01 = fig.add_subplot(spec[0, 1])
dt = helper(data_t, 0, [1])
x_r = wt.inverse(dt)[idx, 0].squeeze().detach().numpy()
ax_01.plot(x, lw=1, color=colors[0], label="Clathrin amplitude")
ax_01.plot(x_r, lw=1, color=colors[1], alpha=0.5, label="Max feature with top 2 contribution")
ax_01.set_yticklabels([])
ax_01.legend(fontsize=5, loc='upper left', fancybox=True)

ax_10 = fig.add_subplot(spec[1, 0])
dt = helper(data_t, 1, [5])
x_r = wt.inverse(dt)[idx, 0].squeeze().detach().numpy()
ax_10.plot(x, lw=1, color=colors[0], label="Clathrin amplitude")
ax_10.plot(x_r, lw=1, color=colors[1], alpha=0.5, label="Max feature with top 3 contribution")
ax_10.set_yticklabels([])
ax_10.legend(fontsize=5, loc='upper left', fancybox=True)

ax_11 = fig.add_subplot(spec[1, 1])
dt = helper(data_t, 3, [2])
x_r = wt.inverse(dt)[idx, 0].squeeze().detach().numpy()
ax_11.plot(x, lw=1, color=colors[0], label="Clathrin amplitude")
ax_11.plot(x_r, lw=1, color=colors[1], alpha=0.5, label="Max feature with top 4 contribution")
ax_11.set_yticklabels([])
ax_11.legend(fontsize=5, loc='upper left', fancybox=True)

# fig.savefig(
#     "figures/bio_viz_important_features.pdf", bbox_inches="tight", pad_inches=0.0
# )<jupyter_output><empty_output><jupyter_text>remade fig<jupyter_code>plt.rcParams['axes.spines.right'] = False
plt.rcParams['axes.spines.top'] = False

fig = plt.figure(figsize=(6.5, 8), dpi=200)

R, C = 5, 1
plt.subplot(R, C, 1)
x = Xs[idx, 0].squeeze().detach().numpy()
ax = plt.plot(x, lw=2, color=colors[0], label="Clathrin amplitude")
plt.title('A', fontweight='bold', loc='left')
plt.ylim(-0.5, 2)
plt.grid(alpha=0.4)
plt.gca().axes.xaxis.set_ticklabels([])
plt.ylabel('Fitted clathrin\namplitude')

N = 3
prefixes = ['Most', '2nd-most', '3rd-most']
lets = 'BCD'
for i, (l, n) in enumerate([(0, 0), (0, 1), (1, 5), (3, 2)][:N]):
    plt.subplot(R, C, 2 + i)
    plt.title(lets[i], fontweight='bold', loc='left')
    dt = helper(data_t, l, [n])
    x_r = wt.inverse(dt)[idx, 0].squeeze().detach().numpy()
    plt.plot(x, lw=2, color=colors[0], label="Clathrin amplitude", alpha=0.16)
    plt.plot(x_r, lw=2)
    plt.ylim(-0.5, 2)
    plt.grid(alpha=0.4)
    if i < N - 1:
        plt.gca().axes.xaxis.set_ticklabels([])
    plt.ylabel(prefixes[i] + '\nimportant\nfeature')
plt.xlabel('Time (s)')

# plt.tight_layout()
# fig.savefig(
#     "figures/bio_viz_important_features_merged.pdf", bbox_inches="tight", pad_inches=0.3
# )<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16369636963696368, 'fraction_numerical': 0.029174917491749174, 'mean_word_length': 2.68661800486618, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 45, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '118471', 'license': 'MIT', 'n_tokens_mistral': 2991, 'n_tokens_neox': 2832, 'n_words': 528, 'path': 'notebooks/biology/05_regression_plot.ipynb', 'repo_name': 'Yu-Group/adaptive-wavelets', 'size': '7575'}