<jupyter_start><jupyter_text>Get a prediction trajectory, and make a movie!<jupyter_code>gaussian_trajs_concat = np.load('../bayes_implicit_solvent/rjmc_experiments/elaborate_tree_rjmc_march31_run_n_compounds=315_n_iter=100000_gaussian_ll_small_proposals.npz')['within_model_trajs']
student_t_trajs_concat = np.load('../bayes_implicit_solvent/rjmc_experiments/elaborate_tree_rjmc_march31_run_n_compounds=315_n_iter=100000_student-t_ll_small_proposals.npz')['within_model_trajs']

len(set([tuple(t) for t in gaussian_trajs_concat]))
len(gaussian_trajs_concat) /4000
gaussian_trajs = [gaussian_trajs_concat[26*i:26*(i+1)] for i in range(len(gaussian_trees))]
student_t_trajs = [student_t_trajs_concat[26*i:26*(i+1)] for i in range(len(student_t_trees))]
gaussian_thetas = [traj[-1] for traj in gaussian_trajs]
student_t_thetas = [traj[-1] for traj in student_t_trajs]
set([len(t) for t in gaussian_trajs]), set([len(t) for t in student_t_trajs])
plt.plot([len(t) for t in gaussian_trajs_concat])
N_s_continuous = np.array([int(len(t)/2) for t in gaussian_trajs_concat])
len(N_s_continuous)
26 * (len(gaussian_trees))
plt.plot([len(t) for t in gaussian_trajs])
len(N_s_continuous)
len(N_s_continuous)
okay_inds = []
problem_inds = []
for i in tqdm(range(len(gaussian_trees))):
    n_continuous = N_s_continuous[i * 26]
    n_discrete = len(gaussian_trees[i].ordered_nodes)
    if n_continuous == n_discrete:
        okay_inds.append(i)
    else:
        problem_inds.append(i)
len(problem_inds)
preds = get_predictions(
    student_t_thetas[-1],
    t_typings[tuple(student_t_trees[-1].ordered_nodes)])
def rmse(x, y):
    return np.sqrt(np.mean((x - y)**2))
rmse(unreduce(preds), expt_means)
gaussian_trajs[0][0]
gaussian_trees[0]
gaussian_prediction_traj = []
for i in tqdm(range(len(gaussian_trees))[::5]):
    gaussian_prediction_traj.append(get_predictions(
        gaussian_thetas[i],
        gaussian_typings[tuple(gaussian_trees[i].ordered_nodes)]))
diag = np.linspace(-15,5)
plt.scatter(expt_means, gaussian_prediction_traj[-1])
plt.plot(diag, diag, c='grey')
unreduced_gaussian_prediction_traj = unreduce(np.array(gaussian_prediction_traj))
diag = np.linspace(-15,5)
plt.scatter(expt_means, unreduced_gaussian_prediction_traj[-1])
plt.plot(diag, diag, c='grey')
len(set([tuple(p) for p in unreduced_gaussian_prediction_traj]))
len(set([tuple(t.get_radii()) for t in gaussian_trees]))
np.save('march30_gaussian_ll_rjmc_prediction_traj.npy', unreduced_gaussian_prediction_traj)
plt.plot([rmse(p, expt_means) for p in gaussian_prediction_traj])
plt.plot([rmse(p, expt_means) for p in unreduced_gaussian_prediction_traj])
plt.plot([rmse(p, expt_means) for p in unreduced_gaussian_prediction_traj[10:]])
rmse(unreduced_gaussian_prediction_traj[-1], expt_means)
rmse(gaussian_prediction_traj[-1], expt_means)
train_inds = np.array([  1,   2,   3,   5,   6,   8,  10,  13,  16,  17,  18,  20,  21,
        22,  23,  24,  28,  32,  33,  36,  37,  39,  41,  43,  44,  46,
        48,  49,  51,  52,  54,  56,  57,  62,  64,  66,  68,  70,  73,
        79,  82,  83,  84,  85,  92,  94,  97, 100, 101, 104, 105, 107,
       109, 110, 111, 115, 117, 118, 119, 120, 121, 122, 123, 126, 129,
       131, 132, 138, 139, 142, 143, 144, 145, 146, 151, 152, 154, 156,
       157, 159, 162, 163, 164, 166, 172, 173, 175, 176, 177, 183, 184,
       187, 188, 190, 195, 196, 199, 200, 201, 204, 205, 208, 214, 215,
       216, 223, 226, 227, 229, 230, 232, 233, 238, 239, 240, 243, 244,
       245, 246, 254, 255, 262, 263, 264, 266, 267, 270, 271, 273, 274,
       275, 277, 278, 279, 289, 290, 291, 292, 295, 296, 298, 299, 301,
       302, 303, 304, 305, 306, 307, 308, 311, 312, 319, 320, 321, 322,
       324, 326, 328, 329, 333, 336, 337, 340, 341, 342, 345, 346, 347,
       351, 353, 354, 357, 359, 360, 363, 366, 367, 369, 371, 374, 376,
       379, 380, 382, 386, 391, 392, 393, 394, 395, 396, 397, 399, 400,
       401, 402, 404, 408, 411, 412, 413, 414, 415, 417, 419, 420, 421,
       428, 429, 430, 431, 432, 433, 434, 435, 436, 440, 441, 442, 443,
       445, 447, 448, 452, 453, 455, 457, 458, 459, 461, 462, 463, 464,
       466, 468, 469, 470, 471, 472, 474, 482, 483, 484, 485, 486, 487,
       490, 491, 494, 495, 496, 499, 500, 502, 503, 509, 511, 516, 517,
       519, 523, 524, 527, 528, 529, 530, 531, 533, 534, 538, 539, 541,
       542, 543, 545, 546, 547, 548, 553, 561, 562, 565, 566, 568, 570,
       573, 574, 577, 579, 580, 581, 583, 584, 585, 589, 590, 593, 594,
       595, 596, 597, 603, 605, 606, 609, 615, 616, 617, 618, 619, 623,
       624, 625, 628])
test_inds = np.array([  0,   4,   7,   9,  11,  12,  14,  15,  19,  25,  26,  27,  29,
        30,  31,  34,  35,  38,  40,  42,  45,  47,  50,  53,  55,  58,
        59,  60,  61,  63,  65,  67,  69,  71,  72,  74,  75,  76,  77,
        78,  80,  81,  86,  87,  88,  89,  90,  91,  93,  95,  96,  98,
        99, 102, 103, 106, 108, 112, 113, 114, 116, 124, 125, 127, 128,
       130, 133, 134, 135, 136, 137, 140, 141, 147, 148, 149, 150, 153,
       155, 158, 160, 161, 165, 167, 168, 169, 170, 171, 174, 178, 179,
       180, 181, 182, 185, 186, 189, 191, 192, 193, 194, 197, 198, 202,
       203, 206, 207, 209, 210, 211, 212, 213, 217, 218, 219, 220, 221,
       222, 224, 225, 228, 231, 234, 235, 236, 237, 241, 242, 247, 248,
       249, 250, 251, 252, 253, 256, 257, 258, 259, 260, 261, 265, 268,
       269, 272, 276, 280, 281, 282, 283, 284, 285, 286, 287, 288, 293,
       294, 297, 300, 309, 310, 313, 314, 315, 316, 317, 318, 323, 325,
       327, 330, 331, 332, 334, 335, 338, 339, 343, 344, 348, 349, 350,
       352, 355, 356, 358, 361, 362, 364, 365, 368, 370, 372, 373, 375,
       377, 378, 381, 383, 384, 385, 387, 388, 389, 390, 398, 403, 405,
       406, 407, 409, 410, 416, 418, 422, 423, 424, 425, 426, 427, 437,
       438, 439, 444, 446, 449, 450, 451, 454, 456, 460, 465, 467, 473,
       475, 476, 477, 478, 479, 480, 481, 488, 489, 492, 493, 497, 498,
       501, 504, 505, 506, 507, 508, 510, 512, 513, 514, 515, 518, 520,
       521, 522, 525, 526, 532, 535, 536, 537, 540, 544, 549, 550, 551,
       552, 554, 555, 556, 557, 558, 559, 560, 563, 564, 567, 569, 571,
       572, 575, 576, 578, 582, 586, 587, 588, 591, 592, 598, 599, 600,
       601, 602, 604, 607, 608, 610, 611, 612, 613, 614, 620, 621, 622,
       626, 627, 629, 630])
def train_rmse(preds):
    return rmse(preds[train_inds], expt_means[train_inds])
def test_rmse(preds):
    return rmse(preds[test_inds], expt_means[test_inds])
gaussian_train_rmse_traj = np.array(list(map(train_rmse, unreduced_gaussian_prediction_traj)))
gaussian_test_rmse_traj = np.array(list(map(test_rmse, unreduced_gaussian_prediction_traj)))

plt.plot(gaussian_train_rmse_traj)
plt.plot(gaussian_test_rmse_traj)
plt.plot(gaussian_train_rmse_traj[5:])
plt.plot(gaussian_test_rmse_traj[5:])
gaussian_trees[-1].get_radii()
len(set(gaussian_test_rmse_traj))<jupyter_output><empty_output><jupyter_text>Make a movie of the atom-typing scheme!<jupyter_code>list(zip(discrete_trees[0], np.bincount(np.hstack(list(gaussian_typings.values())[0]))))
student_t_trees[-1]
def expt_unc_contained(preds, i, desired_coverage=0.95):
    alpha = 100 * ((1 - desired_coverage) / 2)
    upper, lower = norm.cdf(np.percentile(preds, q=[100 - alpha, alpha]), loc=expt_means[i], scale=expt_uncs[i])
    return upper - lower
from scipy.stats import norm
from bayes_implicit_solvent.utils import remove_top_right_spines
desired_coverages = np.linspace(0,1)


actual_coverages = np.vstack([np.array([expt_unc_contained(unreduced_gaussian_prediction_traj[:,i], i, desired_coverage=p) for i in range(len(mols))]) for p in desired_coverages])



ax = plt.subplot(1,1,1)
remove_top_right_spines(ax)
plt.plot(desired_coverages, desired_coverages, c='grey', linestyle='--')

plt.plot(desired_coverages, np.median(actual_coverages, 1), label='gaussian')
plt.fill_between(desired_coverages, np.median(actual_coverages, 1), alpha=0.25)
np.trapz(np.mean(actual_coverages, 1), desired_coverages)
student_t_prediction_traj = []
for i in tqdm(range(len(student_t_trees))[::10]):
    student_t_prediction_traj.append(get_predictions(
        student_t_thetas[i],
        t_typings[tuple(student_t_trees[i].ordered_nodes)]))
student_t_prediction_traj = np.array(student_t_prediction_traj)
unreduced_student_t_prediction_traj = unreduce(np.array(student_t_prediction_traj))
t_actual_coverages = np.vstack([np.array([expt_unc_contained(unreduced_student_t_prediction_traj[:,i], i, desired_coverage=p) for i in range(len(mols))]) for p in desired_coverages])
ax = plt.subplot(1,1,1)
remove_top_right_spines(ax)
plt.plot(desired_coverages, desired_coverages, c='grey', linestyle='--')

#for a in actual_coverages.T[10:]:
#    plt.plot(desired_coverages, a, c='lightblue', alpha=0.5)
#plt.plot(desired_coverages, np.mean(actual_coverages[:,10:], 1), label='gaussian')
#plt.fill_between(desired_coverages, np.mean(actual_coverages[:,10:], 1), alpha=0.25)

plt.plot(desired_coverages, np.mean(t_actual_coverages, 1), label='student-t')
plt.fill_between(desired_coverages, np.mean(t_actual_coverages, 1), alpha=0.25)

plt.legend()
plt.xlabel('desired coverage probability')
plt.ylabel('actual coverage probability')
plt.xlim(0,1)
plt.ylim(0,1)
ax = plt.subplot(1,1,1)
remove_top_right_spines(ax)
plt.plot(desired_coverages, desired_coverages, c='grey', linestyle='--')

plt.plot(desired_coverages, np.mean(actual_coverages, 1), label='gaussian')
plt.fill_between(desired_coverages, np.mean(actual_coverages, 1), alpha=0.25)

plt.plot(desired_coverages, np.mean(t_actual_coverages, 1), label='student-t')
plt.fill_between(desired_coverages, np.mean(t_actual_coverages, 1), alpha=0.25)

plt.legend()
plt.xlabel('desired coverage probability')
plt.ylabel('actual coverage probability')
plt.xlim(0,1)
plt.ylim(0,1)
np.trapz(np.mean(t_actual_coverages, 1), desired_coverages)
gaussian_train_rmse_traj = np.array(list(map(train_rmse, unreduced_gaussian_prediction_traj)))
gaussian_test_rmse_traj = np.array(list(map(test_rmse, unreduced_gaussian_prediction_traj)))

t_train_rmse_traj = np.array(list(map(train_rmse, unreduced_student_t_prediction_traj)))
t_test_rmse_traj = np.array(list(map(test_rmse, unreduced_student_t_prediction_traj)))


plt.plot(gaussian_train_rmse_traj)
plt.plot(gaussian_test_rmse_traj)
plt.figure()
plt.plot(t_train_rmse_traj)
plt.plot(t_test_rmse_traj)
plt.plot(t_train_rmse_traj[1:])
#plt.plot(t_test_rmse_traj[1:])
t_test_rmse_traj
t_test_rmse_traj[-1]
train_rmse(np.mean(unreduced_student_t_prediction_traj[10:], 0))
train_rmse(np.mean(unreduced_gaussian_prediction_traj[10:], 0))
test_rmse(np.mean(unreduced_student_t_prediction_traj[10:], 0))
np.min(t_test_rmse_traj)
test_rmse(np.mean(unreduced_gaussian_prediction_traj[10:], 0))
# oh, wait, each one of the predictions has an uncertainty I've neglected!
@jit
def compute_exp_uncertainty(w_F):
    x = np.exp(w_F - np.max(w_F))
    Ex = np.mean(x)
    dx = np.std(x) / np.sqrt(len(x))
    return dx / Ex
@jit
def predict_solvation_free_energy_jax_w_unc(theta, distance_matrices, charges, element_ind_array):
    N = int(len(theta) / 2)
    radii_, scaling_factors_ = theta[:N], theta[N:]

    radii = radii_[element_ind_array]
    scaling_factors = scaling_factors_[element_ind_array]

    @jit
    def compute_component(distance_matrix):
        return compute_OBC_energy_vectorized(distance_matrix, radii, scaling_factors, charges)

    W_F = vmap(compute_component)(distance_matrices)

    w_F = W_F * kj_mol_to_kT
    unc = compute_exp_uncertainty(w_F)
    return one_sided_exp(w_F), unc
import jax.numpy as np
def get_predictions_and_uncs(theta, typings):
    
    ## TypeError: 'FilledConstant' object does not support item assignment
    #preds = np.zeros(len(typings))
    #uncs = np.zeros(len(typings))
    #for i in range(len(charges)):
    #    preds[i], uncs[i] = predict_solvation_free_energy_jax_w_unc(theta, distance_matrices[i], charges[i], typings[i])
    
    # OKAY, fine a list comprehension...
    preds_and_uncs = [predict_solvation_free_energy_jax_w_unc(theta, distance_matrices[i], charges[i], typings[i]) for i in range(len(typings))]
    preds = np.array([p for (p,u) in preds_and_uncs])
    uncs = np.array([u for (p,u) in preds_and_uncs])
    return preds, uncs
student_t_prediction_traj = []
student_t_unc_traj = []
for i in tqdm(range(len(student_t_trees))[::5]):
    preds, uncs = get_predictions_and_uncs(
        student_t_thetas[i],
        t_typings[tuple(student_t_trees[i].ordered_nodes)])
    student_t_prediction_traj.append(preds)
    student_t_unc_traj.append(uncs)
student_t_prediction_traj = np.array(student_t_prediction_traj)
student_t_unc_traj = np.array(student_t_unc_traj)
student_t_unc_traj.flatten().max()
plt.hist(student_t_unc_traj.flatten(), bins=50);
preds, uncs = student_t_prediction_traj[:,0], student_t_unc_traj[:,0]
x_ = np.linspace(-15,5,1000)
plt.hist(preds, density=True);
def uncertainty_density(x, preds, uncs):
    return sum([norm.pdf(x, loc=mu, scale=sigma) for (mu, sigma) in zip(preds, uncs)]) / len(preds)
y_ = uncertainty_density(x_, preds, uncs)
plt.plot(x_, y_)
plt.xlim(-12,-4)
plt.hist(preds, density=True);
plt.plot(x_, norm.pdf(x_, expt_means[0], expt_uncs[0]))
from numpy import trapz
trapz(y_, x_)
plt.plot(x_, norm.pdf(x_, preds[0], uncs[0]))
plt.plot(x_, norm.pdf(x_, expt_means[0], expt_uncs[0]))
i = train_inds[10]
preds, uncs = student_t_prediction_traj[:,i], student_t_unc_traj[:,i]

plt.plot(x_, norm.pdf(x_, expt_means[i], expt_uncs[i]))
plt.plot(x_, norm.pdf(x_, preds[0], uncs[0]))
plt.xlim(-8,-1)
plt.plot(x_, norm.pdf(x_, expt_means[i], expt_uncs[i]))
y_ = uncertainty_density(x_, preds, uncs)
plt.plot(x_, y_)
plt.xlim(-8,-1)
diag = np.linspace(-15,5)
plt.scatter(expt_means, unreduced_student_t_prediction_traj[-1])
plt.plot(diag, diag, c='grey')
diag = np.linspace(-15,5)
plt.scatter(expt_means[train_inds], unreduced_student_t_prediction_traj[-1][train_inds])
plt.plot(diag, diag, c='grey')
diag = np.linspace(-15,5)
plt.scatter(expt_means[test_inds], unreduced_student_t_prediction_traj[-1][test_inds])
plt.plot(diag, diag, c='grey')
averaged_prediction = np.mean(unreduced_student_t_prediction_traj, 0)
median_prediction = np.median(unreduced_student_t_prediction_traj, 0)
diag = np.linspace(-15,5)
plt.scatter(averaged_prediction[test_inds], expt_means[test_inds])
plt.plot(diag, diag, c='grey')
diag = np.linspace(-15,5)
plt.scatter(median_prediction[test_inds], expt_means[test_inds])
plt.plot(diag, diag, c='grey')
averaged_prediction = np.mean(unreduced_student_t_prediction_traj, 0)
median_prediction = np.median(unreduced_student_t_prediction_traj, 0)

rmse(median_prediction[test_inds], expt_means[test_inds]), rmse(averaged_prediction[test_inds], expt_means[test_inds])
averaged_prediction = np.mean(unreduced_gaussian_prediction_traj, 0)
median_prediction = np.median(unreduced_gaussian_prediction_traj, 0)

rmse(median_prediction[test_inds], expt_means[test_inds]), rmse(averaged_prediction[test_inds], expt_means[test_inds])
burn_in = 10
running_median_prediction = np.array([np.median(unreduced_student_t_prediction_traj[burn_in:t], 0) for t in range(burn_in, len(unreduced_student_t_prediction_traj))])


t_train_rmse_traj_med = np.array(list(map(train_rmse, running_median_prediction)))
t_test_rmse_traj_med = np.array(list(map(test_rmse, running_median_prediction)))
plt.title('instantaneous predictions')
plt.plot(t_train_rmse_traj[burn_in:], label='train')
plt.plot(t_test_rmse_traj[burn_in:], label='test')
plt.legend()
plt.title('running median predictions')
plt.plot(t_train_rmse_traj_med, label='train')
plt.plot(t_test_rmse_traj_med, label='test')
plt.legend()
Deltas = np.zeros((len(test_inds), len(test_inds)))

for i in range(len(test_inds)):
    for j in range(len(test_inds)):
        Deltas[i,j] = expt_means[test_inds[i]] - expt_means[test_inds[j]]
flat_Deltas = []

for i in range(len(test_inds)):
    for j in range(i+1, len(test_inds)):
        flat_Deltas.append(expt_means[test_inds[i]] - expt_means[test_inds[j]])
flat_Deltas = np.array(flat_Deltas)
preds = running_median_prediction[-1]
predicted_Deltas = np.zeros((len(test_inds), len(test_inds)))

for i in range(len(test_inds)):
    for j in range(len(test_inds)):
        predicted_Deltas[i,j] = preds[test_inds[i]] - preds[test_inds[j]]
flat_predicted_Deltas = []

for i in range(len(test_inds)):
    for j in range(i+1, len(test_inds)):
        flat_predicted_Deltas.append(preds[test_inds[i]] - preds[test_inds[j]])
flat_predicted_Deltas = np.array(flat_predicted_Deltas)
np.mean(np.sign(flat_Deltas) == np.sign(flat_predicted_Deltas))
disagreements = np.sign(flat_Deltas) != np.sign(flat_predicted_Deltas)
np.mean(disagreements)
plt.scatter(flat_Deltas[disagreements], flat_predicted_Deltas[disagreements], s=1)
plt.hexbin(flat_Deltas[disagreements], flat_predicted_Deltas[disagreements], bins='log', cmap='Blues')<jupyter_output><empty_output><jupyter_text>Instead of comparing the means here, should be doing a statistical test or something!  Need to compare with predictions from original OBC2 model!  Should move these analyses into a submodule, maybe performance_metrics or similar...<jupyter_code>np.mean((Deltas - predicted_Deltas)**2)
plt.scatter(Deltas.flatten(), predicted_Deltas.flatten(), s=1)
plt.hexbin(flat_Deltas, predicted_Deltas, bins='log', cmap='Blues')
np.argmax(Deltas - predicted_Deltas)
mbondi2_preds, mbondi2_uncs = unreduce(np.load('mbondi2_preds.npy')).T
rmse(mbondi2_preds, expt_means)
from bayes_implicit_solvent.gb_models.obc2_parameters import mbondi2_model
mbondi2_types = mbondi2_model.apply_to_molecule_list(oemols)
mbondi2_radii = mbondi2_model.get_radii()
mbondi2_scales = mbondi2_model.get_scale_factors()
mbondi2_theta = np.hstack([mbondi2_radii, mbondi2_scales])
reduced_mbondi2_preds, reduced_mbondi2_uncs = get_predictions_and_uncs(mbondi2_theta, mbondi2_types)
mbondi2_preds = unreduce(reduced_mbondi2_preds)
mbondi2_uncs = unreduce(reduced_mbondi2_uncs)
rmse(mbondi2_preds, expt_means)
OBC2_Deltas = np.zeros((len(test_inds), len(test_inds)))

for i in range(len(test_inds)):
    for j in range(len(test_inds)):
        OBC2_Deltas[i,j] = mbondi2_preds[test_inds[i]] - mbondi2_preds[test_inds[j]]
flat_OBC2_Deltas = []

for i in range(len(test_inds)):
    for j in range(i+1, len(test_inds)):
        flat_OBC2_Deltas.append(mbondi2_preds[test_inds[i]] - mbondi2_preds[test_inds[j]])
flat_OBC2_Deltas = np.array(flat_OBC2_Deltas)
plt.hexbin(flat_Deltas, flat_OBC2_Deltas, bins='log', cmap='Blues')
np.median((flat_Deltas - flat_OBC2_Deltas)**2)
np.median((flat_Deltas - flat_predicted_Deltas)**2)
disagreements = np.sign(flat_Deltas) != np.sign(flat_OBC2_Deltas)
1 - np.mean(disagreements)
plt.hexbin(flat_OBC2_Deltas[disagreements],flat_Deltas[disagreements],  bins='log', cmap='Blues')
plt.xlabel('OBC2 prediction')
plt.ylabel('experimental value')
np.max(flat_OBC2_Deltas[disagreements]-flat_Deltas[disagreements])
np.min(flat_OBC2_Deltas[disagreements]-flat_Deltas[disagreements])
plt.hist(flat_OBC2_Deltas[disagreements], bins=50);
plt.hist(flat_Deltas[disagreements], bins=50);
# for each one of these pairs, can I do a substructure analysis to see how similar they are?
# then I could pick some extreme examples on both ends, and some moderate examples:
# * a very similar-looking pair with a very big difference in hydration free energy
# * a very different-looking pair with a very small difference in hydration free energy
# * less extreme examples...
from openeye import oechem


mcss.SetMCSFunc(oechem.OEMCSMaxAtoms())
def size_of_mcs(m1, m2):
    atomexpr = oechem.OEExprOpts_DefaultAtoms
    bondexpr = oechem.OEExprOpts_DefaultBonds
    mcss = oechem.OEMCSSearch(m1, atomexpr, bondexpr, oechem.OEMCSType_Approximate)
    mcss.SetMCSFunc(oechem.OEMCSMaxAtoms())

    match_sizes = [0]
    for count, match in enumerate(mcss.Match(m2, True)):
        m = oechem.OEGraphMol()
        oechem.OESubsetMol(m, match, True)
        match_sizes.append(m.NumAtoms())
    return max(match_sizes)
test_mols = [oemols[i] for i in test_inds]
mcs_sizes = np.zeros((len(test_mols), len(test_mols)))
for i in tqdm(range(len(test_mols))):
    for j in range(len(test_mols)):
        mcs_sizes[i,j] = size_of_mcs(test_mols[i], test_mols[j])
oechem.OEMCSType_Approximate
sizes = np.array([m.NumAtoms() for m in test_mols])
pair_min_sizes = np.minimum(np.outer(sizes, np.ones(len(sizes))).T, np.outer(sizes, np.ones(len(sizes))))

pair_max_sizes = np.maximum(np.outer(sizes, np.ones(len(sizes))).T, np.outer(sizes, np.ones(len(sizes))))
plt.imshow(mcs_sizes / pair_max_sizes)
plt.colorbar()
affinities = mcs_sizes / pair_max_sizes
from sklearn.cluster.bicluster import SpectralBiclustering
bic = SpectralBiclustering(n_clusters=6)
bic.fit(affinities)
inds = np.argsort(bic.column_labels_)
plt.imshow(affinities[inds][:,inds])
plt.colorbar()
plt.scatter(affinities.flatten(), np.abs(Deltas).flatten(), s=1)
plt.hexbin(affinities.flatten(), np.abs(Deltas).flatten(),bins='log', cmap='Blues')
plt.hexbin(affinities.flatten(), Deltas.flatten(),bins='log', cmap='Blues')
plt.xlabel('|MCS(m1,m2)| / max(|m1|, |m2|)')
plt.ylabel(r'$\Delta G_{hyd}(m_1) - \Delta G_{hyd}(m_2)$')
# find pareto front of MCS size and absolute value of Delta

X = np.vstack([affinities.flatten(), np.abs(Deltas.flatten())]).T

pareto_front = [0]

for i in range(len(X)):
    pareto_optimal = True
    
    for j in pareto_front:
        if (X[j] > X[i]).all():
            pareto_optimal = False
            break
    
    if pareto_optimal:
        pareto_front.append(i)
len(pareto_front)
plt.scatter(*X[pareto_front].T)
np.max(np.abs(Deltas.flatten())[affinities.flatten() > 0.9])
x_affinities = sorted(list(set(affinities.flatten())))[:-1]
y_Deltas = [np.max(np.abs(Deltas.flatten())[affinities.flatten() > thresh]) for thresh in x_affinities]
plt.scatter(x_affinities, np.abs(y_Deltas))

# maybe instead of "affinities" say "atoms not in common"...
atoms_not_in_common = pair_max_sizes - mcs_sizes
np.min(atoms_not_in_common), np.max(atoms_not_in_common)
x_atoms_not_in_common = sorted(list(set(atoms_not_in_common.flatten())))[1:]
y_Deltas = [np.max(np.abs(Deltas.flatten())[atoms_not_in_common.flatten() < thresh]) for thresh in x_atoms_not_in_common]
plt.scatter(x_atoms_not_in_common, np.abs(y_Deltas))
plt.xlabel('atoms not in common')

y_Deltas = [np.max(np.abs(OBC2_Deltas.flatten())[atoms_not_in_common.flatten() < thresh]) for thresh in x_atoms_not_in_common]
plt.scatter(x_atoms_not_in_common, np.abs(y_Deltas))

y_Deltas = [np.max(np.abs(predicted_Deltas.flatten())[atoms_not_in_common.flatten() < thresh]) for thresh in x_atoms_not_in_common]
plt.scatter(x_atoms_not_in_common, np.abs(y_Deltas))
x_atoms_not_in_common = sorted(list(set(atoms_not_in_common.flatten())))[1:]
y_Deltas = [np.max(np.abs(Deltas.flatten())[atoms_not_in_common.flatten() < thresh]) for thresh in x_atoms_not_in_common]
plt.scatter(x_atoms_not_in_common, np.abs(y_Deltas))
plt.xlabel('atoms not in common')
plt.xlim(0,6.5)
y_Deltas
i_s, j_s = np.where(Deltas == 7.18)
expt_means[test_inds][i_s] - expt_means[test_inds][j_s]
atoms_not_in_common[i_s, j_s]
z = np.argmin(atoms_not_in_common[i_s, j_s])
i,j = i_s[z], j_s[z]
smiles_i, smiles_j = mols[test_inds[i]].smiles, mols[test_inds[j]].smiles
smiles_i, smiles_j
print('ClC1')
test_inds[i], test_inds[j]
from bayes_implicit_solvent.freesolv import smiles_to_cid
smiles_to_cid[smiles_i], smiles_to_cid[smiles_j]
x_atoms_not_in_common = sorted(list(set(atoms_not_in_common.flatten())))[:-1]
y_Deltas = [np.min(np.abs(Deltas.flatten())[atoms_not_in_common.flatten() > thresh]) for thresh in x_atoms_not_in_common]
plt.scatter(x_atoms_not_in_common, np.abs(y_Deltas))
plt.xlabel('atoms not in common')

x_atoms_not_in_common
X[0]
np.vstack([affinities.flatten(), np.abs(Deltas.flatten())]).T.shape
plt.hexbin(affinities.flatten(), OBC2_Deltas.flatten(),bins='log', cmap='Blues')
plt.hexbin(mcs_sizes.flatten(), Deltas.flatten(),bins='log', cmap='Blues')
plt.xlabel('|MCS(m1,m2)|')
plt.ylabel(r'$\Delta G_{hyd}(m_1) - \Delta G_{hyd}(m_2)$')<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup