Goal¶
- improve the count model
Tasks¶
- [ ] load the data
- [ ] compute the bottlenecks
- [ ] add a simple linear model on top (trained with sklearn)
TODO¶
- [ ] Update weights in the main BPNet model
- make sure you also un-do the inverse scaling
- [ ] re-do the scatterplots for the main figure
- [ ] re-do the importance scores
Required files¶
-
In [1]:
# Imports
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
from basepair.imports import *
hv.extension('bokeh')
In [2]:
paper_config()
In [3]:
from basepair.config import valid_chr, test_chr
In [4]:
# Common paths
model_dir = Path(f"{ddir}/processed/chipnexus/exp/models/oct-sox-nanog-klf/models/n_dil_layers=9/")
modisco_dir = model_dir / f"modisco/all/profile/"
output_dir = Path("/srv/www/kundaje/avsec/chipnexus/oct-sox-nanog-klf/models/n_dil_layers=9/modisco/all/profile")
Load the data¶
In [5]:
dataspec_file = model_dir / "dataspec.yaml"
ds = DataSpec.load(dataspec_file)
tasks = list(ds.task_specs)
In [6]:
from basepair.cli.imp_score import ImpScoreFile
In [22]:
imp_file = ImpScoreFile(model_dir / "grad.all.h5")
In [23]:
seq = imp_file.f.f['/inputs'][:]
Get bottleneck predictions¶
In [15]:
create_tf_session(0)
Out[15]:
In [16]:
bpnet = BPNetPredictor.from_mdir(model_dir)
In [17]:
bpnet.model.summary()
In [18]:
bottleneck_model = Model(bpnet.model.inputs, bpnet.model.get_layer("add_9").output)
In [19]:
bottleneck_predictions = bottleneck_model.predict(seq, batch_size=32, verbose=1)
Get target values¶
In [25]:
chromosomes = pd.Series(imp_file.f.f['/metadata/range/chr'][:])
In [26]:
# generate the data split
in_train = ~ chromosomes.isin(valid_chr + test_chr)
in_valid = chromosomes.isin(valid_chr)
in_test = chromosomes.isin(test_chr)
assert not np.any(in_train & in_valid)
assert not np.any(in_train & in_test)
assert not np.any(in_valid & in_test)
In [27]:
profiles = imp_file.get_profiles()
In [28]:
imp_file.close()
In [29]:
# Get the count matrix
total_counts = np.concatenate([profiles[t].sum(axis=1) for t in bpnet.tasks], axis=1)
columns_names = [f"{t}/{s}" for t in bpnet.tasks for s in ['pos', 'neg']]
total_counts = pd.DataFrame(total_counts.astype(int), columns = columns_names)
In [30]:
del profiles
In [31]:
total_counts.head()
Out[31]:
Save to file¶
In [32]:
from sklearn.preprocessing import StandardScaler
In [33]:
log_total_counts = np.log(1 + total_counts).values
scaler = StandardScaler()
In [34]:
train = (bottleneck_predictions[in_train], scaler.fit_transform(log_total_counts[in_train]))
valid = (bottleneck_predictions[in_valid], scaler.transform(log_total_counts[in_valid]))
test = (bottleneck_predictions[in_valid], scaler.transform(log_total_counts[in_test]))
In [35]:
train_averaged = train[0].mean(axis=1)
valid_averaged = valid[0].mean(axis=1)
Train the top model using sklearn¶
In [36]:
from basepair.plot.evaluate import regression_eval
Linear model¶
In [37]:
from sklearn.multioutput import MultiOutputRegressor
from sklearn.linear_model import LinearRegression
In [38]:
m = MultiOutputRegressor(LinearRegression())
In [39]:
m.fit(train_averaged, log_total_counts[in_train]) # Don't use the scaled version when training the linear model
y_pred = m.predict(valid_averaged)
In [92]:
fig, axes = plt.subplots(1, len(bpnet.tasks), figsize=get_figsize(frac=1.5, aspect=0.3), sharex=True, sharey=True)
for i, (a, ax) in enumerate(zip(bpnet.tasks, axes)):
s = slice(2*i, 2*(i+1))
regression_eval(log_total_counts[in_valid][:,s].mean(axis=1), y_pred[:,s].mean(axis=1), alpha=0.05, task=a, ax=ax);
plt.tight_layout()
Do the model surgery on bpnet¶
In [93]:
replace_weights = False
plt.figure(figsize=(4,4))
for task_i, task in enumerate(tasks):
W_lm = np.stack([m.estimators_[2*task_i + strand_i].coef_ for strand_i in range(2)], axis=-1)
b_lm = np.stack([m.estimators_[2*task_i + strand_i].intercept_ for strand_i in range(2)])
W, b = bpnet.model.get_layer(f'counts/{task}').get_weights()
if replace_weights:
bpnet.model.get_layer(f'counts/{task}').set_weights((W_lm, b_lm))
plt.scatter(W.ravel(), W_lm.ravel(), alpha=0.5, s=4);
plt.xlabel("Neural net")
plt.ylabel("Linear model");
In [94]:
# Save back the model
# bpnet.model.save(model_dir / 'model.h5')
Random forrest¶
In [102]:
from sklearn.ensemble import RandomForestRegressor
In [103]:
m = MultiOutputRegressor(RandomForestRegressor(n_estimators=20), n_jobs=8)
In [104]:
m.fit(train_averaged, train[1])
y_pred = m.predict(valid_averaged)
In [105]:
fig, axes = plt.subplots(1, len(bpnet.tasks), figsize=get_figsize(frac=1.5, aspect=0.3), sharex=True, sharey=True)
for i, (a, ax) in enumerate(zip(bpnet.tasks, axes)):
s = slice(2*i, 2*(i+1))
regression_eval(valid[1][:,s].mean(axis=1), y_pred[:,s].mean(axis=1), alpha=0.05, task=a, ax=ax);
plt.tight_layout()
Linear model on binned counts¶
In [108]:
from basepair.preproc import bin_counts
In [118]:
# bin and flatten
train_avg_pool = bin_counts(train[0], 50).reshape((len(train[0]), -1))
valid_avg_pool = bin_counts(valid[0], 50).reshape((len(valid[0]), -1))
In [119]:
m = MultiOutputRegressor(LinearRegression())
In [120]:
m.fit(train_avg_pool, train[1])
y_pred = m.predict(valid_avg_pool)
In [121]:
fig, axes = plt.subplots(1, len(bpnet.tasks), figsize=get_figsize(frac=1.5, aspect=0.3), sharex=True, sharey=True)
for i, (a, ax) in enumerate(zip(bpnet.tasks, axes)):
s = slice(2*i, 2*(i+1))
regression_eval(valid[1][:,s].mean(axis=1), y_pred[:,s].mean(axis=1), alpha=0.05, task=a, ax=ax);
plt.tight_layout()
MLP¶
In [158]:
from keras.models import Sequential
from keras.optimizers import Adam
from keras.callbacks import EarlyStopping, ModelCheckpoint
i = 0
In [159]:
!mkdir -p {model_dir}/count-models
In [244]:
m = Sequential([kl.BatchNormalization(input_shape = train_averaged.shape[1:]),
kl.Dense(64, activation='relu'),
kl.BatchNormalization(),
kl.Dense(len(tasks))])
m.compile(Adam(0.01), 'mse')
In [245]:
i+=1
ckp_file = str(model_dir / f"count-models/{i}.h5")
m.fit(train_averaged, train[1],
batch_size=512,
epochs=50,
callbacks=[EarlyStopping(patience=5, ),
ModelCheckpoint(ckp_file, save_best_only=True, )],
validation_data=(valid_averaged, valid[1]))
m = load_model(ckp_file)
y_pred = m.predict(valid_averaged)
In [246]:
fig, axes = plt.subplots(1, len(bpnet.tasks), figsize=get_figsize(frac=1.5, aspect=0.3), sharex=True, sharey=True)
for i, (a, ax) in enumerate(zip(bpnet.tasks, axes)):
s = slice(2*i, 2*(i+1))
regression_eval(valid[1][:,s].mean(axis=1), y_pred[:,s].mean(axis=1), alpha=0.05, task=a, ax=ax);
plt.tight_layout()
MLP on strided pooling¶
In [226]:
mp = Sequential([kl.AveragePooling1D(100, input_shape = train[0].shape[1:]),
kl.Flatten()])
train_maxpool = mp.predict(train[0])
valid_maxpool = mp.predict(valid[0])
In [238]:
m = Sequential([kl.BatchNormalization(input_shape = train_maxpool.shape[1:]),
kl.Dense(64, activation='relu'),
kl.Dropout(0.3),
kl.BatchNormalization(),
kl.Dense(len(tasks))])
m.compile(Adam(0.01), 'mse')
In [239]:
i+=1
ckp_file = str(model_dir / f"count-models/{i}.h5")
m.fit(train_maxpool, train[1],
batch_size=512,
epochs=50,
callbacks=[EarlyStopping(patience=5, ),
ModelCheckpoint(ckp_file, save_best_only=True, )],
validation_data=(valid_maxpool, valid[1]))
m = load_model(ckp_file)
y_pred = m.predict(valid_maxpool)
In [240]:
fig, axes = plt.subplots(1, len(bpnet.tasks), figsize=get_figsize(frac=1.5, aspect=0.3), sharex=True, sharey=True)
for i, (a, ax) in enumerate(zip(bpnet.tasks, axes)):
s = slice(2*i, 2*(i+1))
regression_eval(valid[1][:,s].mean(axis=1), y_pred[:,s].mean(axis=1), alpha=0.05, task=a, ax=ax);
plt.tight_layout()
