def est_age(input_pkl, age_column, analyses, output_csv):
"""Estimate age using cgAgeR"""
import pandas as pd
import rpy2.robjects as robjects
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri, numpy2ri
pandas2ri.activate()
os.makedirs(output_csv[:output_csv.rfind('/')], exist_ok=True)
methyl_array = MethylationArray.from_pickle(input_pkl)
if age_column:
age_column = pandas2ri.py2ri(methyl_array.pheno[age_column])
else:
age_column = robjects.r('NULL')
run_analyses = dict(epitoc=False, horvath=False, hannum=False)
for analysis in analyses:
run_analyses[analysis] = True
importr('cgageR')
returned_ages = robjects.r(
"""function (beta, hannum, horvath, epitoc, age) {
return(getAgeR(beta, epitoc=epitoc, horvath=horvath, hannum=hannum, chrage=age))
}""")(methyl_array.beta.T, run_analyses['hannum'],
run_analyses['horvath'], run_analyses['epitoc'],
age_column)
result_dfs = []
return_data = lambda data_str: robjects.r(
"""function (results) results{}""".format(data_str))(returned_ages)
if 'hannum' in analyses:
result_dfs.append(
pandas2ri.ri2py(
return_data('$HannumClock.output$Hannum.Clock.Est')))
if 'epitoc' in analyses:
result_dfs.append(
pandas2ri.ri2py(return_data('$EpiTOC.output$EpiTOC.Est')))
if 'horvath' in analyses:
result_dfs.append(
pandas2ri.ri2py(return_data('$HorvathClock.output$Horvath.Est')))
df = pd.concat(result_dfs, axis=1)
df.index = methyl_array.pheno.index
df.to_csv(output_csv)
print(df)
def generate_embed(input_pkl, output_generate_pkl, output_embed_pkl, cuda,
input_vae_pkl, stratify_column, n_workers, batch_size):
import copy
from methylnet.models import AutoEncoder, TybaltTitusVAE
from methylnet.datasets import get_methylation_dataset
import torch
from torch.utils.data import DataLoader
os.makedirs(os.path.dirname(output_generate_pkl), exist_ok=True)
os.makedirs(os.path.dirname(output_embed_pkl), exist_ok=True)
methyl_array = MethylationArray.from_pickle(
input_pkl
) # generate results pickle to run through classification/regression report
if cuda:
model = torch.load(input_vae_pkl)
else:
model = torch.load(input_vae_pkl, map_location='cpu')
test_methyl_dataset = get_methylation_dataset(copy.deepcopy(methyl_array),
stratify_column)
test_methyl_dataloader = DataLoader(dataset=test_methyl_dataset,
num_workers=n_workers,
batch_size=min(
batch_size,
len(test_methyl_dataset)),
shuffle=False)
auto_encoder = AutoEncoder(autoencoder_model=model,
n_epochs=0,
loss_fn=None,
optimizer=None,
cuda=cuda,
kl_warm_up=None,
beta=None,
scheduler_opts={})
Z, _, _ = auto_encoder.transform(test_methyl_dataloader)
X_hat = auto_encoder.generate(test_methyl_dataloader)
methyl_array.beta.iloc[:, :] = X_hat
methyl_array.write_pickle(output_generate_pkl)
methyl_array.beta = pd.DataFrame(Z, index=methyl_array.beta.index)
methyl_array.write_pickle(output_embed_pkl)
def model_capsnet_(
train_methyl_array='train_val_test_sets/train_methyl_array.pkl',
val_methyl_array='train_val_test_sets/val_methyl_array.pkl',
interest_col='disease',
n_epochs=10,
n_bins=0,
bin_len=1000000,
min_capsule_len=300,
primary_caps_out_len=45,
caps_out_len=45,
hidden_topology='30,80,50',
gamma=1e-2,
decoder_topology='100,300',
learning_rate=1e-2,
routing_iterations=3,
overlap=0.,
custom_loss='none',
gamma2=1e-2,
job=0,
capsule_choice=['genomic_binned'],
custom_capsule_file='',
test_methyl_array='',
predict=False,
batch_size=16,
limited_capsule_names_file='',
gsea_superset='',
tissue='',
number_sets=25,
use_set=False,
gene_context=False,
select_subtypes=[],
fit_spw=False,
l1_l2='',
custom_capsule_file2='',
min_capsules=5):
capsule_choice = list(capsule_choice)
#custom_capsule_file=list(custom_capsule_file)
hlt_list = filter(None, hidden_topology.split(','))
if hlt_list:
hidden_topology = list(map(int, hlt_list))
else:
hidden_topology = []
hlt_list = filter(None, decoder_topology.split(','))
if hlt_list:
decoder_topology = list(map(int, hlt_list))
else:
decoder_topology = []
hidden_caps_layers = []
include_last = False
ma = MethylationArray.from_pickle(train_methyl_array)
ma_v = MethylationArray.from_pickle(val_methyl_array)
if test_methyl_array and predict:
ma_t = MethylationArray.from_pickle(test_methyl_array)
try:
ma.remove_na_samples(interest_col)
ma_v.remove_na_samples(interest_col)
if test_methyl_array and predict:
ma_t.remove_na_samples(interest_col)
except:
pass
if select_subtypes:
print(ma.pheno[interest_col].unique())
ma.pheno = ma.pheno.loc[ma.pheno[interest_col].isin(select_subtypes)]
ma.beta = ma.beta.loc[ma.pheno.index]
ma_v.pheno = ma_v.pheno.loc[ma_v.pheno[interest_col].isin(
select_subtypes)]
ma_v.beta = ma_v.beta.loc[ma_v.pheno.index]
print(ma.pheno[interest_col].unique())
if test_methyl_array and predict:
ma_t.pheno = ma_t.pheno.loc[ma_t.pheno[interest_col].isin(
select_subtypes)]
ma_t.beta = ma_t.beta.loc[ma_t.pheno.index]
if custom_capsule_file2 and os.path.exists(custom_capsule_file2):
capsules_dict = torch.load(custom_capsule_file2)
final_modules, modulecpgs, module_names = capsules_dict[
'final_modules'], capsules_dict['modulecpgs'], capsules_dict[
'module_names']
if min_capsule_len > 1:
include_capsules = [
len(x) > min_capsule_len for x in final_modules
]
final_modules = [
final_modules[i] for i in range(len(final_modules))
if include_capsules[i]
]
module_names = [
module_names[i] for i in range(len(module_names))
if include_capsules[i]
]
modulecpgs = (reduce(np.union1d, final_modules)).tolist()
else:
final_modules, modulecpgs, module_names = build_capsules(
capsule_choice, overlap, bin_len, ma, include_last,
min_capsule_len, custom_capsule_file, gsea_superset, tissue,
gene_context, use_set, number_sets, limited_capsule_names_file)
if custom_capsule_file2:
torch.save(
dict(final_modules=final_modules,
modulecpgs=modulecpgs,
module_names=module_names), custom_capsule_file2)
assert len(
final_modules) >= min_capsules, "Below the number of allowed capsules."
if fit_spw:
modulecpgs = list(reduce(lambda x, y: np.hstack((x, y)),
final_modules))
if not include_last: # ERROR HAPPENS HERE!
ma.beta = ma.beta.loc[:, modulecpgs]
ma_v.beta = ma_v.beta.loc[:, modulecpgs]
if test_methyl_array and predict:
ma_t.beta = ma_t.beta.loc[:, modulecpgs]
# https://github.com/higgsfield/Capsule-Network-Tutorial/blob/master/Capsule%20Network.ipynb
original_interest_col = interest_col
if n_bins:
new_interest_col = interest_col + '_binned'
ma.pheno.loc[:,
new_interest_col], bins = pd.cut(ma.pheno[interest_col],
bins=n_bins,
retbins=True)
ma_v.pheno.loc[:,
new_interest_col], _ = pd.cut(ma_v.pheno[interest_col],
bins=bins,
retbins=True)
if test_methyl_array and predict:
ma_t.pheno.loc[:, new_interest_col], _ = pd.cut(
ma_t.pheno[interest_col], bins=bins, retbins=True)
interest_col = new_interest_col
datasets = dict()
datasets['train'] = MethylationDataset(
ma,
interest_col,
modules=final_modules,
module_names=module_names,
original_interest_col=original_interest_col,
run_spw=fit_spw)
print(datasets['train'].X.isnull().sum().sum())
datasets['val'] = MethylationDataset(
ma_v,
interest_col,
modules=final_modules,
module_names=module_names,
original_interest_col=original_interest_col,
run_spw=fit_spw)
if test_methyl_array and predict:
datasets['test'] = MethylationDataset(
ma_t,
interest_col,
modules=final_modules,
module_names=module_names,
original_interest_col=original_interest_col,
run_spw=fit_spw)
dataloaders = dict()
dataloaders['train'] = DataLoader(datasets['train'],
batch_size=batch_size,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True)
dataloaders['val'] = DataLoader(datasets['val'],
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True,
drop_last=False)
n_primary = len(final_modules)
if test_methyl_array and predict:
dataloaders['test'] = DataLoader(datasets['test'],
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True,
drop_last=False)
n_inputs = list(map(len, final_modules))
n_out_caps = len(datasets['train'].y_unique)
if not fit_spw:
print("Not fitting MethylSPWNet")
primary_caps = PrimaryCaps(modules=final_modules,
hidden_topology=hidden_topology,
n_output=primary_caps_out_len)
hidden_caps = []
output_caps = CapsLayer(n_out_caps,
n_primary,
primary_caps_out_len,
caps_out_len,
routing_iterations=routing_iterations)
decoder = Decoder(n_out_caps * caps_out_len, len(list(ma.beta)),
decoder_topology)
model = CapsNet(primary_caps,
hidden_caps,
output_caps,
decoder,
gamma=gamma)
if test_methyl_array and predict:
model.load_state_dict(torch.load('capsnet_model.pkl'))
else:
print("Fitting MethylSPWNet")
module_lens = [len(x) for x in final_modules]
model = MethylSPWNet(module_lens,
hidden_topology,
dropout_p=0.2,
n_output=n_out_caps)
if test_methyl_array and predict:
model.load_state_dict(torch.load('spwnet_model.pkl'))
if torch.cuda.is_available():
model = model.cuda()
# extract all c_ij for all layers across all batches, or just last batch
if l1_l2 and fit_spw:
l1, l2 = list(map(float, l1_l2.split(',')))
elif fit_spw:
l1, l2 = 0., 0.
trainer = Trainer(model=model,
validation_dataloader=dataloaders['val'],
n_epochs=n_epochs,
lr=learning_rate,
n_primary=n_primary,
custom_loss=custom_loss,
gamma2=gamma2,
spw_mode=fit_spw,
l1=l1 if fit_spw else 0.,
l2=l2 if fit_spw else 0.)
if not predict:
try:
#assert 1==2
trainer.fit(dataloader=dataloaders['train'])
val_loss = min(trainer.val_losses)
torch.save(
trainer.model.state_dict(),
'capsnet_model.pkl' if not fit_spw else 'spwnet_model.pkl')
if fit_spw:
torch.save(
dict(final_modules=final_modules,
modulecpgs=modulecpgs,
module_names=module_names), 'spwnet_capsules.pkl')
torch.save(
dict(module_names=module_names,
module_lens=module_lens,
dropout_p=0.2,
hidden_topology=hidden_topology,
n_output=n_out_caps), 'spwnet_config.pkl')
except Exception as e:
print(e)
val_loss = -2
with sqlite3.connect('jobs.db', check_same_thread=False) as conn:
pd.DataFrame([job, val_loss],
index=['job', 'val_loss'],
columns=[0]).T.to_sql('val_loss',
conn,
if_exists='append')
else:
if test_methyl_array:
trainer.weights = 1.
Y = trainer.predict(dataloaders['test'])
pickle.dump(Y, open('predictions.pkl', 'wb'))
val_loss = -1
#print(val_loss)
# print([min(trainer.val_losses),n_epochs,
# n_bins,
# bin_len,
# min_capsule_len,
# primary_caps_out_len,
# caps_out_len,
# hidden_topology,
# gamma,
# decoder_topology,
# learning_rate,
# routing_iterations])
return val_loss
def bin_column(test_pkl, col, n_bins, output_test_pkl):
"""Convert continuous phenotype column into categorical by binning."""
os.makedirs(output_test_pkl[:output_test_pkl.rfind('/')], exist_ok=True)
test_methyl_array = MethylationArray.from_pickle(test_pkl)
new_col_name = test_methyl_array.bin_column(col, n_bins)
test_methyl_array.write_pickle(output_test_pkl)
def print_shape(input_pkl):
"""Print dimensions of beta matrix."""
print(MethylationArray.from_pickle(input_pkl).beta.shape)
def make_new_predictions(test_pkl, model_pickle, batch_size, n_workers,
interest_cols, categorical, cuda, categorical_encoder,
output_dir):
"""Run prediction model again to further assess outcome. Only evaluate prediction model."""
os.makedirs(output_dir, exist_ok=True)
test_methyl_array = MethylationArray.from_pickle(
test_pkl
) # generate results pickle to run through classification/regression report
if cuda:
model = torch.load(model_pickle)
model.vae.cuda_on = True
else:
model = torch.load(model_pickle, map_location='cpu')
model.vae.cuda_on = False
if not categorical:
test_methyl_array.remove_na_samples(
interest_cols if len(interest_cols) > 1 else interest_cols[0])
if os.path.exists(categorical_encoder):
categorical_encoder = pickle.load(open(categorical_encoder, 'rb'))
else:
categorical_encoder = None
test_methyl_dataset = get_methylation_dataset(
test_methyl_array,
interest_cols,
categorical=categorical,
predict=True,
categorical_encoder=categorical_encoder)
test_methyl_dataloader = DataLoader(dataset=test_methyl_dataset,
num_workers=n_workers,
batch_size=min(
batch_size,
len(test_methyl_dataset)),
shuffle=False)
vae_mlp = MLPFinetuneVAE(mlp_model=model,
categorical=categorical,
cuda=cuda)
Y_pred, Y_true, latent_projection, _ = vae_mlp.predict(
test_methyl_dataloader)
results = dict(test={})
results['test']['y_pred'], results['test']['y_true'] = copy.deepcopy(
Y_pred), copy.deepcopy(Y_true)
if categorical:
Y_true = Y_true.argmax(axis=1)[:, np.newaxis]
Y_pred = Y_pred.argmax(axis=1)[:, np.newaxis]
test_methyl_array = test_methyl_dataset.to_methyl_array()
Y_pred = pd.DataFrame(
Y_pred.flatten() if (np.array(Y_pred.shape) == 1).any() else Y_pred,
index=test_methyl_array.beta.index,
columns=(['y_pred'] if categorical else interest_cols))
Y_true = pd.DataFrame(
Y_true.flatten() if (np.array(Y_true.shape) == 1).any() else Y_true,
index=test_methyl_array.beta.index,
columns=(['y_true'] if categorical else interest_cols))
results_df = pd.concat([
Y_pred, Y_true
], axis=1) if categorical else pd.concat([
Y_pred.rename(columns={name: name + '_pred'
for name in list(Y_pred)}),
Y_true.rename(columns={name: name + '_true'
for name in list(Y_pred)})
],
axis=1) # FIXME
latent_projection = pd.DataFrame(latent_projection,
index=test_methyl_array.beta.index)
test_methyl_array.beta = latent_projection
output_file = join(output_dir, 'results.csv')
results_file = join(output_dir, 'results.p')
output_file_latent = join(output_dir, 'latent.csv')
output_pkl = join(output_dir, 'vae_mlp_methyl_arr.pkl')
test_methyl_array.write_pickle(output_pkl)
pickle.dump(results, open(results_file, 'wb'))
latent_projection.to_csv(output_file_latent)
results_df.to_csv(output_file)
def train_predict(train_pkl,
test_pkl,
input_vae_pkl,
output_dir,
cuda,
interest_cols,
categorical,
disease_only,
hidden_layer_topology,
learning_rate_vae,
learning_rate_mlp,
weight_decay,
dropout_p,
n_epochs,
scheduler='null',
decay=0.5,
t_max=10,
eta_min=1e-6,
t_mult=2,
batch_size=50,
val_pkl='val_methyl_array.pkl',
n_workers=8,
add_validation_set=False,
loss_reduction='sum',
add_softmax=False):
os.makedirs(output_dir, exist_ok=True)
output_file = join(output_dir, 'results.csv')
training_curve_file = join(output_dir, 'training_val_curve.p')
results_file = join(output_dir, 'results.p')
output_file_latent = join(output_dir, 'latent.csv')
output_model = join(output_dir, 'output_model.p')
output_pkl = join(output_dir, 'vae_mlp_methyl_arr.pkl')
output_onehot_encoder = join(output_dir, 'one_hot_encoder.p')
#input_dict = pickle.load(open(input_pkl,'rb'))
if cuda:
vae_model = torch.load(input_vae_pkl)
vae_model.cuda_on = True
else:
vae_model = torch.load(input_vae_pkl, map_location='cpu')
vae_model.cuda_on = False
train_methyl_array, val_methyl_array, test_methyl_array = MethylationArray.from_pickle(
train_pkl
), MethylationArray.from_pickle(val_pkl), MethylationArray.from_pickle(
test_pkl
) #methyl_array.split_train_test(train_p=train_percent, stratified=(True if categorical else False), disease_only=disease_only, key=interest_cols[0], subtype_delimiter=',')
if not categorical:
train_methyl_array.remove_na_samples(interest_cols)
val_methyl_array.remove_na_samples(interest_cols)
test_methyl_array.remove_na_samples(interest_cols)
print(train_methyl_array.beta.shape)
print(val_methyl_array.beta.shape)
print(test_methyl_array.beta.shape)
if len(interest_cols) == 1 and disease_only and interest_cols[0].endswith(
'_only') == False:
print(interest_cols)
interest_cols[0] += '_only'
print(train_methyl_array.pheno[interest_cols[0]].unique())
print(test_methyl_array.pheno[interest_cols[0]].unique())
train_methyl_dataset = get_methylation_dataset(
train_methyl_array,
interest_cols,
categorical=categorical,
predict=True) # train, test split? Add val set?
#print(list(train_methyl_dataset.encoder.get_feature_names()))
val_methyl_dataset = get_methylation_dataset(
val_methyl_array,
interest_cols,
categorical=categorical,
predict=True,
categorical_encoder=train_methyl_dataset.encoder)
test_methyl_dataset = get_methylation_dataset(
test_methyl_array,
interest_cols,
categorical=categorical,
predict=True,
categorical_encoder=train_methyl_dataset.encoder)
if not batch_size:
batch_size = len(train_methyl_dataset)
train_batch_size = min(batch_size, len(train_methyl_dataset))
val_batch_size = min(batch_size, len(val_methyl_dataset))
train_methyl_dataloader = DataLoader(dataset=train_methyl_dataset,
num_workers=n_workers,
batch_size=train_batch_size,
shuffle=True)
val_methyl_dataloader = DataLoader(dataset=val_methyl_dataset,
num_workers=n_workers,
batch_size=val_batch_size,
shuffle=True) # False
test_methyl_dataloader = DataLoader(dataset=test_methyl_dataset,
num_workers=n_workers,
batch_size=min(
batch_size,
len(test_methyl_dataset)),
shuffle=False)
scaling_factors = dict(
val=float(len(val_methyl_dataset)) /
((len(val_methyl_dataset) // val_batch_size) * val_batch_size),
train_batch_size=train_batch_size,
val_batch_size=val_batch_size)
model = VAE_MLP(vae_model=vae_model,
categorical=categorical,
hidden_layer_topology=hidden_layer_topology,
n_output=train_methyl_dataset.outcome_col.shape[1],
dropout_p=dropout_p,
add_softmax=add_softmax)
class_weights = []
if categorical:
out_weight = Counter(
np.argmax(train_methyl_dataset.outcome_col, axis=1))
#total_samples=sum(out_weight.values())
for k in sorted(list(out_weight.keys())):
class_weights.append(1. / float(out_weight[k])) # total_samples
class_weights = np.array(class_weights)
class_weights = (class_weights / class_weights.sum()).tolist()
print(class_weights)
if class_weights:
class_weights = torch.FloatTensor(class_weights)
if cuda:
class_weights = class_weights.cuda()
else:
class_weights = None
optimizer_vae = torch.optim.Adam(model.vae.parameters(),
lr=learning_rate_vae,
weight_decay=weight_decay)
optimizer_mlp = torch.optim.Adam(model.mlp.parameters(),
lr=learning_rate_mlp,
weight_decay=weight_decay)
loss_fn = CrossEntropyLoss(
reduction=loss_reduction,
weight=class_weights) if categorical else MSELoss(
reduction=loss_reduction) # 'sum'
scheduler_opts = dict(scheduler=scheduler,
lr_scheduler_decay=decay,
T_max=t_max,
eta_min=eta_min,
T_mult=t_mult)
vae_mlp = MLPFinetuneVAE(mlp_model=model,
n_epochs=n_epochs,
categorical=categorical,
loss_fn=loss_fn,
optimizer_vae=optimizer_vae,
optimizer_mlp=optimizer_mlp,
cuda=cuda,
scheduler_opts=scheduler_opts)
if add_validation_set:
vae_mlp.add_validation_set(val_methyl_dataloader)
vae_mlp = vae_mlp.fit(train_methyl_dataloader)
if 'encoder' in dir(train_methyl_dataset):
pickle.dump(train_methyl_dataset.encoder,
open(output_onehot_encoder, 'wb'))
results = dict(test={}, train={}, val={})
results['train']['y_pred'], results['train'][
'y_true'], _, _ = vae_mlp.predict(train_methyl_dataloader)
results['val']['y_pred'], results['val']['y_true'], _, _ = vae_mlp.predict(
val_methyl_dataloader)
del train_methyl_dataloader, train_methyl_dataset
"""methyl_dataset=get_methylation_dataset(methyl_array,interest_cols,predict=True)
methyl_dataset_loader = DataLoader(
dataset=methyl_dataset,
num_workers=9,
batch_size=1,
shuffle=False)"""
Y_pred, Y_true, latent_projection, _ = vae_mlp.predict(
test_methyl_dataloader
) # FIXME change to include predictions for all classes for AUC
results['test']['y_pred'], results['test']['y_true'] = copy.deepcopy(
Y_pred), copy.deepcopy(Y_true)
if categorical:
Y_true = Y_true.argmax(axis=1)[:, np.newaxis]
Y_pred = Y_pred.argmax(axis=1)[:, np.newaxis]
test_methyl_array = test_methyl_dataset.to_methyl_array()
"""if categorical:
Y_true=test_methyl_dataset.encoder.inverse_transform(Y_true)[:,np.newaxis]
Y_pred=test_methyl_dataset.encoder.inverse_transform(Y_pred)[:,np.newaxis]"""
#sample_names = np.array(list(test_methyl_array.beta.index)) # FIXME
#outcomes = np.array([outcome[0] for outcome in outcomes]) # FIXME
Y_pred = pd.DataFrame(
Y_pred.flatten() if (np.array(Y_pred.shape) == 1).any() else Y_pred,
index=test_methyl_array.beta.index,
columns=(['y_pred'] if categorical else
interest_cols)) #dict(zip(sample_names,outcomes))
Y_true = pd.DataFrame(
Y_true.flatten() if (np.array(Y_true.shape) == 1).any() else Y_true,
index=test_methyl_array.beta.index,
columns=(['y_true'] if categorical else interest_cols))
results_df = pd.concat([
Y_pred, Y_true
], axis=1) if categorical else pd.concat([
Y_pred.rename(columns={name: name + '_pred'
for name in list(Y_pred)}),
Y_true.rename(columns={name: name + '_true'
for name in list(Y_pred)})
],
axis=1) # FIXME
latent_projection = pd.DataFrame(latent_projection,
index=test_methyl_array.beta.index)
test_methyl_array.beta = latent_projection
test_methyl_array.write_pickle(output_pkl)
pickle.dump(results, open(results_file, 'wb'))
pickle.dump(vae_mlp.training_plot_data, open(training_curve_file, 'wb'))
latent_projection.to_csv(output_file_latent)
torch.save(vae_mlp.model, output_model)
results_df.to_csv(
output_file) #pickle.dump(outcome_dict, open(outcome_dict_file,'wb'))
return latent_projection, Y_pred, Y_true, vae_mlp, scaling_factors
def embed_vae(train_pkl,
output_dir,
cuda,
n_latent,
lr,
weight_decay,
n_epochs,
hidden_layer_encoder_topology,
kl_warm_up=0,
beta=1.,
scheduler='null',
decay=0.5,
t_max=10,
eta_min=1e-6,
t_mult=2,
bce_loss=False,
batch_size=50,
val_pkl='val_methyl_array.pkl',
n_workers=9,
convolutional=False,
height_kernel_sizes=[],
width_kernel_sizes=[],
add_validation_set=False,
loss_reduction='sum',
stratify_column='disease'):
from methylnet.models import AutoEncoder, TybaltTitusVAE
from methylnet.datasets import get_methylation_dataset
import torch
from torch.utils.data import DataLoader
from torch.nn import MSELoss, BCELoss
os.makedirs(output_dir, exist_ok=True)
output_file = join(output_dir, 'output_latent.csv')
output_model = join(output_dir, 'output_model.p')
training_curve_file = join(output_dir, 'training_val_curve.p')
outcome_dict_file = join(output_dir, 'output_outcomes.p')
output_pkl = join(output_dir, 'vae_methyl_arr.pkl')
#input_dict = pickle.load(open(input_pkl,'rb'))
#methyl_array=MethylationArray(*extract_pheno_beta_df_from_pickle_dict(input_dict))
#print(methyl_array.beta)
train_methyl_array, val_methyl_array = MethylationArray.from_pickle(
train_pkl
), MethylationArray.from_pickle(
val_pkl
) #methyl_array.split_train_test(train_p=train_percent, stratified=True, disease_only=True, key='disease', subtype_delimiter=',')
train_methyl_dataset = get_methylation_dataset(
train_methyl_array, stratify_column) # train, test split? Add val set?
val_methyl_dataset = get_methylation_dataset(val_methyl_array,
stratify_column)
if not batch_size:
batch_size = len(methyl_dataset)
train_batch_size = min(batch_size, len(train_methyl_dataset))
val_batch_size = min(batch_size, len(val_methyl_dataset))
train_methyl_dataloader = DataLoader(
dataset=train_methyl_dataset,
num_workers=n_workers, #n_workers
batch_size=train_batch_size,
shuffle=True,
pin_memory=False)
val_methyl_dataloader = DataLoader(dataset=val_methyl_dataset,
num_workers=n_workers,
batch_size=val_batch_size,
shuffle=True,
pin_memory=False)
scaling_factors = dict(
train=float(len(train_methyl_dataset)) /
((len(train_methyl_dataset) // train_batch_size) * train_batch_size),
val=float(len(val_methyl_dataset)) /
((len(val_methyl_dataset) // val_batch_size) * val_batch_size),
train_batch_size=train_batch_size,
val_batch_size=val_batch_size)
print('SCALE', len(train_methyl_dataset), len(val_methyl_dataset),
train_batch_size, val_batch_size, scaling_factors)
n_input = train_methyl_array.return_shape()[1]
if not convolutional:
model = TybaltTitusVAE(
n_input=n_input,
n_latent=n_latent,
hidden_layer_encoder_topology=hidden_layer_encoder_topology,
cuda=cuda)
else:
model = CVAE(n_latent=n_latent,
in_shape=methyl_dataset.new_shape,
kernel_heights=height_kernel_sizes,
kernel_widths=width_kernel_sizes,
n_pre_latent=n_latent * 2) # change soon
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
loss_fn = BCELoss(reduction=loss_reduction) if bce_loss else MSELoss(
reduction=loss_reduction) # 'sum'
scheduler_opts = dict(scheduler=scheduler,
lr_scheduler_decay=decay,
T_max=t_max,
eta_min=eta_min,
T_mult=t_mult)
auto_encoder = AutoEncoder(autoencoder_model=model,
n_epochs=n_epochs,
loss_fn=loss_fn,
optimizer=optimizer,
cuda=cuda,
kl_warm_up=kl_warm_up,
beta=beta,
scheduler_opts=scheduler_opts)
if add_validation_set:
auto_encoder.add_validation_set(val_methyl_dataloader)
auto_encoder = auto_encoder.fit(train_methyl_dataloader)
train_methyl_array = train_methyl_dataset.to_methyl_array()
val_methyl_array = val_methyl_dataset.to_methyl_array()
del val_methyl_dataloader, train_methyl_dataloader, val_methyl_dataset, train_methyl_dataset
methyl_dataset = get_methylation_dataset(
MethylationArrays([train_methyl_array, val_methyl_array]).combine(),
stratify_column)
methyl_dataset_loader = DataLoader(dataset=methyl_dataset,
num_workers=n_workers,
batch_size=1,
shuffle=False)
latent_projection, _, _ = auto_encoder.transform(methyl_dataset_loader)
#print(latent_projection.shape)
methyl_array = methyl_dataset.to_methyl_array()
#sample_names = np.array([sample_name[0] for sample_name in sample_names]) # FIXME
#outcomes = np.array([outcome[0] for outcome in outcomes]) # FIXME
#outcome_dict=dict(zip(sample_names,outcomes))
#print(methyl_array.beta)
latent_projection = pd.DataFrame(latent_projection,
index=methyl_array.beta.index)
methyl_array.beta = latent_projection
methyl_array.write_pickle(output_pkl)
latent_projection.to_csv(output_file)
pickle.dump(auto_encoder.training_plot_data, open(training_curve_file,
'wb'))
torch.save(auto_encoder.model, output_model)
#pickle.dump(outcome_dict, open(outcome_dict_file,'wb'))
return latent_projection, None, scaling_factors, n_input, auto_encoder
def main():
p = argparse.ArgumentParser()
p.add_argument('--interest_col', type=str)
p.add_argument('--n_bins', type=int)
args = p.parse_args()
bin_len = 1000000
min_capsule_len = 350
interest_col = args.interest_col
n_bins = args.n_bins
primary_caps_out_len = 40
caps_out_len = 20
n_epochs = 500
hidden_topology = [30, 80, 50]
gamma = 1e-2
decoder_top = [100, 300]
lr = 1e-3
routing_iterations = 3
if not os.path.exists('hg19.{}.bed'.format(bin_len)):
BedTool('hg19.genome').makewindows(g='hg19.genome', w=bin_len).saveas(
'hg19.{}.bed'.format(bin_len)) #.to_dataframe().shape
ma = MethylationArray.from_pickle(
'train_val_test_sets/train_methyl_array.pkl')
ma_v = MethylationArray.from_pickle(
'train_val_test_sets/val_methyl_array.pkl')
include_last = False
@pysnooper.snoop('get_mod.log')
def get_final_modules(ma=ma,
a='450kannotations.bed',
b='lola_vignette_data/activeDHS_universe.bed',
include_last=False,
min_capsule_len=2000):
allcpgs = ma.beta.columns.values
df = BedTool(a).to_dataframe()
df.iloc[:, 0] = df.iloc[:, 0].astype(str).map(
lambda x: 'chr' + x.split('.')[0])
df = df.set_index('name').loc[list(
ma.beta)].reset_index().iloc[:, [1, 2, 3, 0]]
df_bed = pd.read_table(b, header=None)
df_bed['features'] = np.arange(df_bed.shape[0])
df_bed = df_bed.iloc[:, [0, 1, 2, -1]]
b = BedTool.from_dataframe(df)
a = BedTool.from_dataframe(
df_bed) #('lola_vignette_data/activeDHS_universe.bed')
c = a.intersect(b, wa=True, wb=True).sort()
d = c.groupby(g=[1, 2, 3, 4], c=(8, 8), o=('count', 'distinct'))
df2 = d.to_dataframe()
df3 = df2.loc[df2.iloc[:, -2] > min_capsule_len]
modules = [cpgs.split(',') for cpgs in df3.iloc[:, -1].values]
modulecpgs = np.array(
list(set(list(reduce(lambda x, y: x + y, modules)))))
if include_last:
missing_cpgs = np.setdiff1d(allcpgs, modulecpgs).tolist()
final_modules = modules + ([missing_cpgs] if include_last else [])
module_names = (df3.iloc[:, 0] + '_' + df3.iloc[:, 1].astype(str) +
'_' + df3.iloc[:, 2].astype(str)).tolist()
return final_modules, modulecpgs, module_names
final_modules, modulecpgs, module_names = get_final_modules(
b='hg19.{}.bed'.format(bin_len),
include_last=include_last,
min_capsule_len=min_capsule_len)
print('LEN_MODULES', len(final_modules))
if not include_last:
ma.beta = ma.beta.loc[:, modulecpgs]
ma_v.beta = ma_v.beta.loc[:, modulecpgs]
# https://github.com/higgsfield/Capsule-Network-Tutorial/blob/master/Capsule%20Network.ipynb
def softmax(input_tensor, dim=1):
# transpose input
transposed_input = input_tensor.transpose(dim,
len(input_tensor.size()) - 1)
# calculate softmax
softmaxed_output = F.softmax(transposed_input.contiguous().view(
-1, transposed_input.size(-1)),
dim=-1)
# un-transpose result
return softmaxed_output.view(*transposed_input.size()).transpose(
dim,
len(input_tensor.size()) - 1)
class MLP(
nn.Module
): # add latent space extraction, and spits out csv line of SQL as text for UMAP
def __init__(self,
n_input,
hidden_topology,
dropout_p,
n_outputs=1,
binary=False,
softmax=False):
super(MLP, self).__init__()
self.hidden_topology = hidden_topology
self.topology = [n_input] + hidden_topology + [n_outputs]
layers = [
nn.Linear(self.topology[i], self.topology[i + 1])
for i in range(len(self.topology) - 2)
]
for layer in layers:
torch.nn.init.xavier_uniform_(layer.weight)
self.layers = [
nn.Sequential(layer, nn.ReLU(), nn.Dropout(p=dropout_p))
for layer in layers
]
self.output_layer = nn.Linear(self.topology[-2], self.topology[-1])
torch.nn.init.xavier_uniform_(self.output_layer.weight)
if binary:
output_transform = nn.Sigmoid()
elif softmax:
output_transform = nn.Softmax()
else:
output_transform = nn.Dropout(p=0.)
self.layers.append(
nn.Sequential(self.output_layer, output_transform))
self.mlp = nn.Sequential(*self.layers)
def forward(self, x):
#print(x.shape)
return self.mlp(x)
class MethylationDataset(Dataset):
def __init__(self,
methyl_arr,
outcome_col,
binarizer=None,
modules=[]):
if binarizer == None:
binarizer = LabelBinarizer()
binarizer.fit(methyl_arr.pheno[outcome_col].astype(str).values)
self.y = binarizer.transform(
methyl_arr.pheno[outcome_col].astype(str).values)
self.y_unique = np.unique(np.argmax(self.y, 1))
self.binarizer = binarizer
if not modules:
modules = [list(methyl_arr.beta)]
self.modules = modules
self.X = methyl_arr.beta
self.length = methyl_arr.beta.shape[0]
def __len__(self):
return self.length
def __getitem__(self, i):
return tuple([torch.FloatTensor(self.X.iloc[i].values)] + [
torch.FloatTensor(self.X.iloc[i].loc[module].values)
for module in self.modules
] + [torch.FloatTensor(self.y[i])])
class PrimaryCaps(nn.Module):
def __init__(self, modules, hidden_topology, n_output):
super(PrimaryCaps, self).__init__()
self.capsules = nn.ModuleList([
MLP(len(module), hidden_topology, 0., n_outputs=n_output)
for module in modules
])
def forward(self, x):
#print(self.capsules)
u = [self.capsules[i](x[i]) for i in range(len(self.capsules))]
u = torch.stack(u, dim=1)
#print(u.size())
return self.squash(u)
def squash(self, x):
squared_norm = (x**2).sum(-1, keepdim=True)
#print('prim_norm',squared_norm.size())
output_tensor = squared_norm * x / (
(1. + squared_norm) * torch.sqrt(squared_norm))
#print('z_init',output_tensor.size())
return output_tensor
def get_weights(self):
return list(
self.capsules[0].parameters()
)[0].data #self.state_dict()#[self.capsules[i].state_dict() for i in range(len(self.capsules))]
class CapsLayer(nn.Module):
def __init__(self,
n_capsules,
n_routes,
n_input,
n_output,
routing_iterations=3):
super(CapsLayer, self).__init__()
self.n_capsules = n_capsules
self.num_routes = n_routes
self.W = nn.Parameter(
torch.randn(1, n_routes, n_capsules, n_output, n_input))
self.routing_iterations = routing_iterations
self.c_ij = None
def forward(self, x):
batch_size = x.size(0)
x = torch.stack([x] * self.n_capsules, dim=2).unsqueeze(4)
W = torch.cat([self.W] * batch_size, dim=0)
#print('affine',W.size(),x.size())
u_hat = torch.matmul(W, x)
#print('affine_trans',u_hat.size())
b_ij = Variable(torch.zeros(1, self.num_routes, self.n_capsules,
1))
if torch.cuda.is_available():
b_ij = b_ij.cuda()
for iteration in range(self.routing_iterations):
self.c_ij = softmax(b_ij)
#print(c_ij)
c_ij = torch.cat([self.c_ij] * batch_size, dim=0).unsqueeze(4)
#print('coeff',c_ij.size())#[0,:,0,:])#.size())
s_j = (c_ij * u_hat).sum(dim=1, keepdim=True)
v_j = self.squash(s_j)
#print('z',v_j.size())
if iteration < self.routing_iterations - 1:
a_ij = torch.matmul(
u_hat.transpose(3, 4),
torch.cat([v_j] * self.num_routes, dim=1))
b_ij = b_ij + a_ij.squeeze(4).mean(dim=0, keepdim=True)
return v_j.squeeze(1)
def return_routing_coef(self):
return self.c_ij
def squash(self, x):
#print(x.size())
squared_norm = (x**2).sum(-1, keepdim=True)
#print('norm',squared_norm.size())
output_tensor = squared_norm * x / (
(1. + squared_norm) * torch.sqrt(squared_norm))
return output_tensor
class Decoder(nn.Module):
def __init__(self, n_input, n_output, hidden_topology):
super(Decoder, self).__init__()
self.decoder = MLP(n_input,
hidden_topology,
0.,
n_outputs=n_output,
binary=True)
def forward(self, x):
return self.decoder(x)
class CapsNet(nn.Module):
def __init__(self,
primary_caps,
caps_hidden_layers,
caps_output_layer,
decoder,
lr_balance=0.5,
gamma=0.005):
super(CapsNet, self).__init__()
self.primary_caps = primary_caps
self.caps_hidden_layers = caps_hidden_layers
self.caps_output_layer = caps_output_layer
self.decoder = decoder
self.recon_loss_fn = nn.BCELoss()
self.lr_balance = lr_balance
self.gamma = gamma
def forward(self, x_orig, modules_input):
x = self.primary_caps(modules_input)
primary_caps_out = x #.view(x.size(0),x.size(1)*x.size(2))
#print(x.size())
for layer in self.caps_hidden_layers:
x = layer(x)
y_pred = self.caps_output_layer(x) #.squeeze(-1)
#print(y_pred.shape)
classes = torch.sqrt((y_pred**2).sum(2))
classes = F.softmax(classes)
max_length_indices = classes.argmax(dim=1)
masked = torch.sparse.torch.eye(self.caps_output_layer.n_capsules)
if torch.cuda.is_available():
masked = masked.cuda()
masked = masked.index_select(
dim=0, index=max_length_indices.squeeze(1).data)
embedding = (y_pred * masked[:, :, None, None]).view(
y_pred.size(0), -1)
#print(y_pred.size())
x_hat = self.decoder(embedding) #.reshape(y_pred.size(0),-1))
return x_orig, x_hat, y_pred, embedding, primary_caps_out
def recon_loss(self, x_orig, x_hat):
return self.recon_loss_fn(x_hat, x_orig)
def margin_loss(self, x, labels):
batch_size = x.size(0)
v_c = torch.sqrt((x**2).sum(dim=2, keepdim=True))
#print(v_c)
left = (F.relu(0.9 - v_c)**2).view(batch_size, -1)
right = (F.relu(v_c - 0.1)**2).view(batch_size, -1)
#print(left)
#print(right)
#print(labels)
loss = labels * left + self.lr_balance * (1.0 - labels) * right
#print(loss.shape)
loss = loss.sum(dim=1).mean()
return loss
def calculate_loss(self, x_orig, x_hat, y_pred, y_true):
margin_loss = self.margin_loss(y_pred, y_true)
recon_loss = self.gamma * self.recon_loss(x_orig, x_hat)
loss = margin_loss + recon_loss
return loss, margin_loss, recon_loss
if n_bins:
ma.pheno.loc[:, interest_col], bins = pd.cut(ma.pheno[interest_col],
bins=n_bins,
retbins=True)
ma_v.pheno.loc[:, interest_col], bins = pd.cut(
ma_v.pheno[interest_col],
bins=bins,
retbins=True,
)
dataset = MethylationDataset(ma, interest_col, modules=final_modules)
dataset_v = MethylationDataset(ma_v, interest_col, modules=final_modules)
dataloader = DataLoader(dataset,
batch_size=16,
shuffle=True,
num_workers=8,
drop_last=True)
dataloader_v = DataLoader(dataset_v,
batch_size=16,
shuffle=False,
num_workers=8,
drop_last=False)
n_inputs = list(map(len, final_modules))
n_primary = len(final_modules)
primary_caps = PrimaryCaps(modules=final_modules,
hidden_topology=hidden_topology,
n_output=primary_caps_out_len)
hidden_caps = []
n_out_caps = len(dataset.y_unique)
output_caps = CapsLayer(n_out_caps,
n_primary,
primary_caps_out_len,
caps_out_len,
routing_iterations=routing_iterations)
decoder = Decoder(n_out_caps * caps_out_len, len(list(ma.beta)),
decoder_top)
capsnet = CapsNet(primary_caps,
hidden_caps,
output_caps,
decoder,
gamma=gamma)
if torch.cuda.is_available():
capsnet = capsnet.cuda()
for d in ['figures/embeddings' + x for x in ['', '2', '3']]:
os.makedirs(d, exist_ok=True)
os.makedirs('results/routing_weights', exist_ok=True)
# extract all c_ij for all layers across all batches, or just last batch
optimizer = Adam(capsnet.parameters(), lr)
scheduler = CosineAnnealingLR(optimizer,
T_max=10,
eta_min=0,
last_epoch=-1)
for epoch in range(n_epochs):
print(epoch)
capsnet.train(True)
running_loss = 0.
Y = {'true': [], 'pred': []}
for i, batch in enumerate(dataloader):
x_orig = batch[0]
#print(x_orig)
y_true = batch[-1]
module_x = batch[1:-1]
if torch.cuda.is_available():
x_orig = x_orig.cuda()
y_true = y_true.cuda()
module_x = [mod.cuda() for mod in module_x]
x_orig, x_hat, y_pred, embedding, primary_caps_out = capsnet(
x_orig, module_x)
loss, margin_loss, recon_loss = capsnet.calculate_loss(
x_orig, x_hat, y_pred, y_true)
Y['true'].extend(y_true.argmax(1).detach().cpu().numpy().tolist())
Y['pred'].extend(
F.softmax(torch.sqrt(
(y_pred**2
).sum(2))).argmax(1).detach().cpu().numpy().tolist())
train_loss = margin_loss.item() #print(loss)
running_loss += train_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
#print(capsnet.primary_caps.get_weights())
running_loss /= (i + 1)
print('Epoch {}: Train Loss {}, Train R2: {}, Train MAE: {}'.format(
epoch, running_loss, r2_score(Y['true'], Y['pred']),
mean_absolute_error(Y['true'], Y['pred'])))
print(classification_report(Y['true'], Y['pred']))
scheduler.step()
capsnet.train(False)
running_loss = np.zeros((3, )).astype(float)
Y = {
'true': [],
'pred': [],
'embeddings': [],
'embeddings2': [],
'embeddings3': [],
'routing_weights': []
}
with torch.no_grad():
for i, batch in enumerate(dataloader_v):
x_orig = batch[0]
y_true = batch[-1]
module_x = batch[1:-1]
if torch.cuda.is_available():
x_orig = x_orig.cuda()
y_true = y_true.cuda()
module_x = [mod.cuda() for mod in module_x]
x_orig, x_hat, y_pred, embedding, primary_caps_out = capsnet(
x_orig, module_x)
#print(primary_caps_out.size())
routing_coefs = capsnet.caps_output_layer.return_routing_coef(
).detach().cpu().numpy()
if not i:
Y['routing_weights'] = pd.DataFrame(
routing_coefs[0, ..., 0].T,
index=dataset.binarizer.classes_,
columns=module_names)
else:
Y['routing_weights'] += pd.DataFrame(
routing_coefs[0, ..., 0].T,
index=dataset.binarizer.classes_,
columns=module_names)
Y['embeddings3'].append(
torch.cat(
[primary_caps_out[i] for i in range(x_orig.size(0))],
dim=0).detach().cpu().numpy())
primary_caps_out = primary_caps_out.view(
primary_caps_out.size(0),
primary_caps_out.size(1) * primary_caps_out.size(2))
Y['embeddings'].append(embedding.detach().cpu().numpy())
Y['embeddings2'].append(
primary_caps_out.detach().cpu().numpy())
loss, margin_loss, recon_loss = capsnet.calculate_loss(
x_orig, x_hat, y_pred, y_true)
val_loss = margin_loss.item() #print(loss)
running_loss = running_loss + np.array(
[loss.item(), margin_loss,
recon_loss.item()])
Y['true'].extend(
y_true.argmax(1).detach().cpu().numpy().tolist())
Y['pred'].extend(
(y_pred**2
).sum(2).argmax(1).detach().cpu().numpy().tolist())
running_loss /= (i + 1)
Y['routing_weights'].iloc[:, :] = Y['routing_weights'].values / (
i + 1)
Y['pred'] = np.array(Y['pred']).astype(str)
Y['true'] = np.array(Y['true']).astype(str)
#np.save('results/routing_weights/routing_weights.{}.npy'.format(epoch),Y['routing_weights'])
pickle.dump(
Y['routing_weights'],
open('results/routing_weights/routing_weights.{}.p'.format(epoch),
'wb'))
Y['embeddings'] = pd.DataFrame(PCA(n_components=2).fit_transform(
np.vstack(Y['embeddings'])),
columns=['x', 'y'])
Y['embeddings2'] = pd.DataFrame(PCA(n_components=2).fit_transform(
np.vstack(Y['embeddings2'])),
columns=['x', 'y'])
#print(list(map(lambda x: x.shape,Y['embeddings3'])))
Y['embeddings3'] = pd.DataFrame(PCA(n_components=2).fit_transform(
np.vstack(Y['embeddings3'])),
columns=['x', 'y']) #'z'
Y['embeddings']['color'] = Y['true']
Y['embeddings2']['color'] = Y['true']
Y['embeddings3']['color'] = module_names * ma_v.beta.shape[
0] #Y['true']
Y['embeddings3']['name'] = list(
reduce(lambda x, y: x + y, [[i] * n_primary for i in Y['true']]))
fig = px.scatter(Y['embeddings3'],
x="x",
y="y",
color="color",
symbol='name') #, text='name')
py.plot(fig,
filename='figures/embeddings3/embeddings3.{}.pos.html'.format(
epoch),
auto_open=False)
#Y['embeddings3']['color']=list(reduce(lambda x,y:x+y,[[i]*n_primary for i in Y['true']]))
fig = px.scatter(Y['embeddings3'], x="x", y="y",
color="name") #, text='color')
py.plot(fig,
filename='figures/embeddings3/embeddings3.{}.true.html'.format(
epoch),
auto_open=False)
fig = px.scatter(Y['embeddings'], x="x", y="y", color="color")
py.plot(fig,
filename='figures/embeddings/embeddings.{}.true.html'.format(
epoch),
auto_open=False)
fig = px.scatter(Y['embeddings2'], x="x", y="y", color="color")
py.plot(fig,
filename='figures/embeddings2/embeddings2.{}.true.html'.format(
epoch),
auto_open=False)
Y['embeddings'].loc[:, 'color'] = Y['pred']
Y['embeddings2'].loc[:, 'color'] = Y['pred']
fig = px.scatter(Y['embeddings'], x="x", y="y", color="color")
py.plot(fig,
filename='figures/embeddings/embeddings.{}.pred.html'.format(
epoch),
auto_open=False)
fig = px.scatter(Y['embeddings2'], x="x", y="y", color="color")
py.plot(fig,
filename='figures/embeddings2/embeddings2.{}.pred.html'.format(
epoch),
auto_open=False)
print(
'Epoch {}: Val Loss {}, Margin Loss {}, Recon Loss {}, Val R2: {}, Val MAE: {}'
.format(
epoch, running_loss[0], running_loss[1], running_loss[2],
r2_score(Y['true'].astype(int), Y['pred'].astype(int)),
mean_absolute_error(Y['true'].astype(int),
Y['pred'].astype(int))))
print(classification_report(Y['true'], Y['pred']))
def return_spw_importances_(train_methyl_array,
val_methyl_array,
interest_col,
select_subtypes,
capsules_pickle,
include_last,
n_bins,
spw_config,
model_state_dict_pkl,
batch_size,
by_subtype=False):
ma = MethylationArray.from_pickle(train_methyl_array)
ma_v = MethylationArray.from_pickle(val_methyl_array)
try:
ma.remove_na_samples(interest_col)
ma_v.remove_na_samples(interest_col)
except:
pass
if select_subtypes:
ma.pheno = ma.pheno.loc[ma.pheno[interest_col].isin(select_subtypes)]
ma.beta = ma.beta.loc[ma.pheno.index]
ma_v.pheno = ma_v.pheno.loc[ma_v.pheno[interest_col].isin(
select_subtypes)]
ma_v.beta = ma_v.beta.loc[ma_v.pheno.index]
capsules_dict = torch.load(capsules_pickle)
final_modules, modulecpgs, module_names = capsules_dict[
'final_modules'], capsules_dict['modulecpgs'], capsules_dict[
'module_names']
if not include_last:
ma.beta = ma.beta.loc[:, modulecpgs]
ma_v.beta = ma_v.beta.loc[:, modulecpgs]
original_interest_col = interest_col
if n_bins:
new_interest_col = interest_col + '_binned'
ma.pheno.loc[:,
new_interest_col], bins = pd.cut(ma.pheno[interest_col],
bins=n_bins,
retbins=True)
ma_v.pheno.loc[:,
new_interest_col], _ = pd.cut(ma_v.pheno[interest_col],
bins=bins,
retbins=True)
interest_col = new_interest_col
datasets = dict()
datasets['train'] = MethylationDataset(
ma,
interest_col,
modules=final_modules,
module_names=module_names,
original_interest_col=original_interest_col,
run_spw=True)
datasets['val'] = MethylationDataset(
ma_v,
interest_col,
modules=final_modules,
module_names=module_names,
original_interest_col=original_interest_col,
run_spw=True)
y_val = datasets['val'].y_label
y_val_uniq = np.unique(y_val)
dataloaders = dict()
dataloaders['train'] = DataLoader(datasets['train'],
batch_size=batch_size,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True)
dataloaders['val'] = DataLoader(datasets['val'],
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True,
drop_last=False)
n_primary = len(final_modules)
spw_config = torch.load(spw_config)
spw_config.pop('module_names')
model = MethylSPWNet(**spw_config)
model.load_state_dict(torch.load(model_state_dict_pkl))
if torch.cuda.is_available():
model = model.cuda()
model.eval()
pathway_extractor = model.pathways
#extract_pathways = lambda modules_x:torch.cat([pathway_extractor[i](module_x) for i,module_x in enumerate(modules_x)],dim=1)
tensor_data = dict(train=dict(X=[], y=[]), val=dict(X=[], y=[]))
for k in tensor_data:
for i, (batch) in enumerate(dataloaders[k]):
x = batch[0]
y_true = batch[-1].argmax(1) #[-2]
modules_x = batch[1:-1] #2]
if torch.cuda.is_available():
x = x.cuda()
modules_x = modules_x[0].cuda(
) #[module.cuda() for module in modules_x]
tensor_data[k]['X'].append(
pathway_extractor(x, modules_x).detach().cpu()
) #extract_pathways(modules_x).detach().cpu())
tensor_data[k]['y'].append(y_true.flatten().view(-1, 1))
tensor_data[k]['X'] = torch.cat(tensor_data[k]['X'], dim=0)
tensor_data[k]['y'] = torch.cat(tensor_data[k]['y'], dim=0)
print(tensor_data[k]['X'].size(), tensor_data[k]['y'].size())
tensor_data[k] = TensorDataset(tensor_data[k]['X'],
tensor_data[k]['y'])
dataloaders[k] = DataLoader(tensor_data[k],
batch_size=32,
sampler=ImbalancedDatasetSampler(
tensor_data[k]))
model = model.output_net
to_cuda = lambda x: x.cuda() if torch.cuda.is_available() else x
y = np.unique(tensor_data['train'].tensors[1].numpy().flatten())
gs = GradientShap(model)
X_train = torch.cat(
[next(iter(dataloaders['train']))[0] for i in range(2)], dim=0)
if torch.cuda.is_available():
X_train = X_train.cuda()
#val_loader=iter(dataloaders['val'])
def return_importances(dataloaders, X_train):
attributions = []
for i in range(20):
batch = next(iter(dataloaders['val']))
X_test = to_cuda(batch[0])
y_test = to_cuda(batch[1].flatten())
attributions.append(
torch.abs(
gs.attribute(
X_test,
stdevs=0.03,
n_samples=200,
baselines=X_train,
target=y_test,
return_convergence_delta=False))) #torch.tensor(y_i)
attributions = torch.sum(torch.cat(attributions, dim=0), dim=0)
importances = pd.DataFrame(
pd.Series(attributions.detach().cpu().numpy(),
index=module_names).sort_values(ascending=False),
columns=['importances'])
return importances
if by_subtype:
importances = []
for k in y_val_uniq:
idx = np.where(y_val == k)[0]
if len(idx) > 2:
val_dataset = Subset(tensor_data['val'], idx)
n_concat = int(np.ceil(64. / len(idx)))
if n_concat > 1:
val_dataset = ConcatDataset([val_dataset] * n_concat)
#sampler=SubsetRandomSampler(idx)
dataloaders['val'] = DataLoader(val_dataset,
batch_size=32,
shuffle=True)
df = return_importances(dataloaders, X_train)
df['subtype'] = k
importances.append(df)
importances = pd.concat(importances)
else:
importances = return_importances(dataloaders, X_train)
return importances
def fit_group_lasso(
train_methyl_array='train_val_test_sets/train_methyl_array.pkl',
val_methyl_array='train_val_test_sets/val_methyl_array.pkl',
test_methyl_array='train_val_test_sets/test_methyl_array.pkl',
l1_vals=np.hstack((np.arange(0.01, 1.1,
0.1), np.array([10., 20., 50.,
100.]))).tolist(),
outcome_col='disease_only',
min_capsule_len=5,
capsule_choice=['gene'],
n_jobs=0,
n_epochs=10,
output_results_pickle='group_lasso_model.pkl',
output_file='group_lasso_importances.csv',
batch_size=1280,
lr=0.0001,
predict=False):
# if torch.cuda.is_available():
# torch.set_default_tensor_type('torch.cuda.FloatTensor')
datasets = dict(train=train_methyl_array,
val=val_methyl_array,
test=test_methyl_array)
# LogisticRegression = lambda ne, lr: net = NeuralNetClassifier(LogisticRegressionModel,max_epochs=ne,lr=lr,iterator_train__shuffle=True, callbacks=[EpochScoring(LASSO)])
X = dict()
Y = dict()
le = LabelEncoder()
for k in ['train', 'val', 'test']:
datasets[k] = MethylationArray.from_pickle(datasets[k])
capsules, cpgs, names, cpg_arr = return_final_capsules(
datasets['train'],
capsule_choice,
min_capsule_len,
None,
None,
0,
'',
'',
return_original_capsule_assignments=True)
cpgs = np.unique(cpgs)
cpg2idx = dict(zip(cpgs, np.arange(len(cpgs))))
cpg_arr.loc[:, 'cpg'] = cpg_arr.loc[:, 'cpg'].map(cpg2idx)
capsule_sizes = [len(capsule) for capsule in capsules]
for k in ['train', 'val', 'test']:
X[k] = datasets[k].beta.loc[:,
cpgs] #cudf.from_pandas(datasets[k].beta)#
X[k].loc[:, :] = beta2M(X[k].loc[:, :].values)
Y[k] = le.fit_transform(
datasets[k].pheno[outcome_col]) if k == 'train' else le.transform(
datasets[k].pheno[outcome_col]
) #cudf.Series(, dtype = np.float32 )
n_classes = len(np.unique(Y['train']))
n_cpgs = len(cpgs)
class_weights = torch.tensor(
compute_class_weight('balanced', np.unique(Y['train']),
Y['train'])).float()
if torch.cuda.is_available():
class_weights = class_weights.cuda()
dataloaders = {}
for k in ['train', 'val', 'test']:
X[k] = torch.tensor(X[k].values).float()
Y[k] = torch.tensor(Y[k]).long()
dataloaders[k] = DataLoader(TensorDataset(X[k], Y[k]),
batch_size=min(batch_size, X[k].shape[0]),
shuffle=(k == 'train'),
num_workers=n_jobs)
def get_res(logreg_model,
dataloader=dataloaders['test'],
return_pred=False):
y_res = {'pred': [], 'true': []}
for i, (x, y) in enumerate(dataloader):
if torch.cuda.is_available():
x = x.cuda()
y = y.cuda()
y_res['true'].extend(y.detach().cpu().numpy().flatten().tolist())
y_res['pred'].extend(
logreg_model(x).argmax(
1).detach().cpu().numpy().flatten().tolist())
y_true = le.inverse_transform(np.array(y_res['true']))
y_pred = le.inverse_transform(np.array(y_res['pred']))
print(classification_report(y_true, y_pred))
if return_pred:
return (y_true, y_pred)
return f1_score(y_true, y_pred, average='macro')
# @torchsnooper.snoop()
def train_model(l1, return_model=False):
logreg_model = GroupLasso(
n_cpgs, n_classes, names, cpg_arr, capsule_sizes,
l1) #nn.Module(nn.Linear())#nn.Sequential(,nn.LogSoftmax())
if torch.cuda.is_available():
logreg_model = logreg_model.cuda()
logreg_model.weights = logreg_model.weights.cuda()
# logreg_model.groups=logreg_model.groups.cuda()
# logreg_model.sqrt_group_sizes=logreg_model.sqrt_group_sizes.cuda()
optimizer = optim.Adam(logreg_model.parameters(), lr=lr)
# scheduler=
criterion = nn.CrossEntropyLoss(weight=class_weights)
for epoch in range(n_epochs):
running_loss = {'train': [], 'val': []}
for phase in ['train', 'val']:
logreg_model.train(phase == 'train')
for i, (x, y) in enumerate(dataloaders[phase]):
if torch.cuda.is_available():
x = x.cuda()
y = y.cuda()
optimizer.zero_grad()
#print(y.shape,logreg_model(x).shape)
y_pred = logreg_model(x)
loss = criterion(y_pred, y)
loss = loss + logreg_model.penalize()
#group_lasso()#group_lasso(logreg_model[0].weight)
if phase == 'train':
loss.backward()
optimizer.step()
# optimizer.zero_grad()
# penalty=logreg_model.penalize()
# if phase=='train':
# penalty.backward(retain_graph=True)
# optimizer.step()
# loss=loss+penalty
item_loss = loss.item()
print("Epoch {}[Batch {}] - {} Loss {} ".format(
epoch, i, phase, item_loss),
flush=True)
running_loss[phase].append(item_loss)
running_loss[phase] = np.mean(running_loss[phase])
print("Epoch {} - Train Loss {} , Val Loss {}".format(
epoch, running_loss['train'], running_loss['val']),
flush=True)
if epoch and running_loss['val'] <= min_val_loss:
min_val_loss = running_loss['val']
best_model_weights = copy.deepcopy(logreg_model.state_dict())
elif not epoch:
min_val_loss = running_loss['val']
logreg_model.load_state_dict(best_model_weights)
if not return_model:
return l1, get_res(logreg_model, dataloader=dataloaders['val'])
else:
return get_res(logreg_model,
dataloader=dataloaders['val'],
return_pred=False), logreg_model
# pool=ProcessPool(nodes=8)
# l1_f1=np.array(pool.map(train_model, l1_vals))
if len(l1_vals) > 1:
l1_f1 = np.array(
dask.compute(*[dask.delayed(train_model)(l1) for l1 in l1_vals],
scheduler='threading')
) #np.array([train_model(l1) for l1 in l1_vals])
l1 = l1_f1[np.argmax(l1_f1[:, 1]), 0]
else:
l1_f1 = None
l1 = l1_vals[0]
# group_lasso=GroupLasso(len(cpgs),len(np.unique(Y['train'])),names,cpg_arr,capsule_sizes,l1)
f1, logreg_model = train_model(l1, return_model=True)
logreg_model.train(False)
y_true, y_pred = get_res(logreg_model, return_pred=True)
torch.save(
dict(model=logreg_model, l1=l1_f1, y_true=y_true, y_pred=y_pred,
f1=f1), output_results_pickle)
weights = logreg_model.return_weights()
pd.DataFrame(dict(zip(names, weights)),
index=['importances']).T.to_csv(output_file)
def fit_logreg(train_methyl_array='train_val_test_sets/train_methyl_array.pkl',
val_methyl_array='train_val_test_sets/val_methyl_array.pkl',
test_methyl_array='train_val_test_sets/test_methyl_array.pkl',
l1_vals=np.hstack(
(np.arange(0.01, 1.1, 0.1), np.array([10., 20., 50.,
100.]))),
outcome_col='disease_only',
min_capsule_len=5,
capsule_choice=['gene'],
n_jobs=20):
datasets = dict(train=train_methyl_array,
val=val_methyl_array,
test=test_methyl_array)
# LogisticRegression = lambda ne, lr: net = NeuralNetClassifier(LogisticRegressionModel,max_epochs=ne,lr=lr,iterator_train__shuffle=True, callbacks=[EpochScoring(LASSO)])
X = dict()
Y = dict()
le = LabelEncoder()
for k in ['train', 'val', 'test']:
datasets[k] = MethylationArray.from_pickle(datasets[k])
X[k] = datasets[k].beta #cudf.from_pandas(datasets[k].beta)#
X[k].loc[:, :] = beta2M(X[k].loc[:, :].values)
Y[k] = le.fit_transform(
datasets[k].pheno[outcome_col]) if k == 'train' else le.transform(
datasets[k].pheno[outcome_col]
) #cudf.Series(, dtype = np.float32 )
capsules, _, names = return_final_capsules(datasets['train'],
capsule_choice, min_capsule_len,
None, None, 0, '', '')
# make_pipeline(CapsuleSelection(capsule,name), LogisticRegression(penalty='l1', C=1./l1,class_weight='balanced'))
# capsules=capsules[:2]#[capsule for capsule in capsules]
# names=names[:2]#[name for name in names]
build_stacking_model = lambda l1: ParallelStackingClassifier(
n_jobs=n_jobs,
meta_classifier=LogisticRegression(penalty='l1',
n_jobs=n_jobs,
C=1. / l1,
class_weight='balanced',
solver='saga'),
use_clones=False,
classifiers=[
make_pipeline(
CapsuleSelection(capsule, name),
LogisticRegression(penalty='l1',
C=1. / l1,
class_weight='balanced',
solver='saga'))
for capsule, name in zip(capsules, names) if len(capsule)
])
def get_score(l1, capsules):
print('Fitting l1: {}'.format(l1))
score = f1_score(build_stacking_model(l1).fit(
X['train'], Y['train'], capsules=capsules).predict(X['val']),
Y['val'],
average='macro')
return l1, score
scores = [get_score(l1, capsules) for l1 in l1_vals]
# pool=ProcessPool(nodes=8)
# scores=pool.map(lambda l1: get_score(l1,capsules), l1_vals)
# for l1 in l1_vals:
# scores.append(get_score(l1,capsules))#scores.append(dask.delayed(get_score)(l1))
# scores.append((l1,f1_score(reg.predict(X['val']).to_pandas().values.flatten().astype(int),Y['val'].to_pandas().values.flatten().astype(int),average='macro')))
scores = np.array(scores) #dask.compute(*scores,scheduler='processes')
np.save('l1_scores.npy', scores)
l1 = scores[np.argmin(scores[:, 1]), 0]
reg = build_stacking_model(
l1
) #LogisticRegression(penalty='l1', C=1./l1)#LogisticRegression(ne,lr)#
reg.fit(X['train'], Y['train'], capsules=capsules)
print(
classification_report(le.inverse_transform(Y['test']),
le.inverse_transform(reg.predict(X['test']))))
# print(classification_report(le.inverse_transform(Y['test'].to_pandas().values.flatten().astype(int)),le.inverse_transform(reg.predict(X['test']).to_pandas().values.flatten().astype(int))))
pickle.dump(dict(model=reg, features=names), open('stacked_model.pkl',
'wb'))