File size: 8,237 Bytes
e048d40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 |
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
import os
from sklearn.metrics import r2_score
import matplotlib.colors as mcolors
from fuson_plm.utils.visualizing import set_font
global default_cmap_dict
default_cmap_dict = {
'Asphericity': '#785EF0',
'End-to-End Distance (Re)': '#DC267F',
'Radius of Gyration (Rg)': '#FE6100',
'Scaling Exponent': '#FFB000'
}
# Method for lengthening the model name
def lengthen_model_name(model_name, model_epoch):
if 'esm' in model_name:
return model_name
return f'{model_name}_e{model_epoch}'
def plot_train_val_test_values_hist(train_values_list, val_values_list, test_values_list, dataset_name="Data", color="black", save_path=None, ax=None):
"""
Plot Histogram to show the ranges of values
"""
set_font()
if ax is None:
fig, ax = plt.subplots(1, 1, figsize=(6, 4), dpi=300)
total_seqs = len(train_values_list)+len(val_values_list)+len(test_values_list)
ax.hist(train_values_list, color=color, alpha=0.7,label=f"train (n={len(train_values_list)})")
if not(test_values_list is None):
ax.hist(test_values_list, color='black',alpha=0.7,label=f"test (n={len(test_values_list)})")
if not(val_values_list is None):
ax.hist(val_values_list, color='grey',alpha=0.7,label=f"val (n={len(val_values_list)})")
ax.grid(True)
ax.set_axisbelow(True)
ax.set_title(f'{dataset_name} Distribution (n={total_seqs})')
ax.set_xlabel(dataset_name)
ax.legend()
plt.tight_layout()
if save_path is not None:
plt.savefig(save_path)
def plot_values_hist(values_list, dataset_name="Data", color="black", save_path=None, ax=None):
"""
Plot Histogram to show the ranges of values
"""
set_font()
if ax is None:
fig, ax = plt.subplots(1, 1, figsize=(6, 4), dpi=300)
ax.hist(values_list, color=color)
ax.grid(True)
ax.set_axisbelow(True)
ax.set_title(f'{dataset_name} Distribution')
ax.set_xlabel(dataset_name)
plt.tight_layout()
if save_path is not None:
plt.savefig(save_path)
def plot_all_values_hist_grid(values_dict, cmap_dict=default_cmap_dict, save_path="processed_data/value_histograms.png"):
"""
Args:
values_dict: dictionary where keys are dataset names and values are value lists
cmap_dict: dictioanry where keys are dataset names (same as in values dict) and values are value lists
"""
fig, axes = plt.subplots(2, 2, figsize=(12, 8), dpi=300)
axes = axes.flatten()
for i, (dataset_name, values_list) in enumerate(values_dict.items()):
ax = axes[i]
plot_values_hist(values_list, dataset_name=dataset_name, color=cmap_dict[dataset_name], ax=ax)
fig.set_tight_layout(True)
fig.savefig(save_path)
def plot_all_train_val_test_values_hist_grid(values_dict, cmap_dict=default_cmap_dict, save_path="processed_data/value_histograms.png"):
"""
Args:
values_dict: dictionary where keys are dataset names and values are another dict: {'train': train_values_list, 'test': test_values_list}
cmap_dict: dictioanry where keys are dataset names (same as in values dict) and values are value lists
"""
fig, axes = plt.subplots(2, 2, figsize=(12, 8), dpi=300)
axes = axes.flatten()
for i, (dataset_name, train_val_test_dict) in enumerate(values_dict.items()):
ax = axes[i]
train_values_list = train_val_test_dict['train']
test_values_list, val_values_list = None, None
if 'test' in train_val_test_dict:
test_values_list = train_val_test_dict['test']
if 'val' in train_val_test_dict:
val_values_list = train_val_test_dict['val']
plot_train_val_test_values_hist(train_values_list, val_values_list, test_values_list, dataset_name=dataset_name, color=cmap_dict[dataset_name], ax=ax)
fig.set_tight_layout(True)
fig.savefig(save_path)
#only need to change labels at bottom depending on what embeddings+dimension is being looked at
def plot_r2(model_type, idr_property, test_preds, save_path):
set_font()
# prepare ylabels from idr_property
ylabel_dict = {'asph': 'Asphericity',
'scaled_re': 'End-to-End Radius, $R_e$',
'scaled_rg': 'Radius of Gyration, $R_g$',
'scaling_exp': 'Polymer Scaling Exponent'}
y_unitlabel_dict = {'asph': 'Asphericity',
'scaled_re': '$R_e$ (Å)',
'scaled_rg': '$R_g$ (Å)',
'scaling_exp': 'Exponent'
}
y_label = ylabel_dict[idr_property]
y_unitlabel = y_unitlabel_dict[idr_property]
# get true values and predictions
true_values = test_preds['true_values'].tolist()
predictions = test_preds['predictions'].tolist()
# save this source data, including the IDs of the sequences
test_df = pd.read_csv(f"splits/{idr_property}/test_df.csv")
processed_data = pd.read_csv("processed_data/all_albatross_seqs_and_properties.csv")
seq_id_dict = dict(zip(processed_data['Sequence'],processed_data['IDs']))
test_df['IDs'] = test_df['Sequence'].map(seq_id_dict)
test_df_with_preds = test_preds[['true_values','predictions']]
test_df_with_preds['IDs'] = test_df['IDs']
print("number of sequences with no ID: ", len(test_df_with_preds.loc[test_df_with_preds['IDs'].isna()]))
test_df_with_preds.to_csv(save_path.replace(".png","_source_data.csv"),index=False)
r2 = r2_score(true_values, predictions)
# Plotting
plt.figure(figsize=(10, 8))
plt.scatter(true_values, predictions, alpha=0.5, label='Predictions')
plt.plot([min(true_values), max(true_values)], [min(true_values), max(true_values)], 'r--', label='Ideal Fit')
plt.text(0.65, 0.35, f"$R^2$ = {r2:.2f}", transform=plt.gca().transAxes, fontsize=44)
# Adjusting font sizes and setting font properties
plt.xlabel(f'True {y_unitlabel}',size=44)
plt.ylabel(f'Predicted {y_unitlabel}',size=44)
plt.title(f"{y_label}",size=50) #: {model_type}\n($R^2$={r2:.2f})",size=44)
# Create legend and set font properties
legend = plt.legend(fontsize=32)
for text in legend.get_texts():
text.set_fontsize(32)
# Adjust marker size in the legend
for handle in legend.legendHandles:
handle._sizes = [100]
# Enable grid
plt.grid(True)
# Adjusting tick labels font size
plt.xticks(fontsize=36)
plt.yticks(fontsize=36)
# Setting font properties for tick labels (another way to adjust them individually)
for label in plt.gca().get_xticklabels():
label.set_fontsize(32)
for label in plt.gca().get_yticklabels():
label.set_fontsize(32)
plt.tight_layout()
plt.savefig(save_path, dpi=300, bbox_inches='tight')
plt.show()
def plot_all_r2(output_dir, idr_properties):
for idr_property in idr_properties:
# make the R^2 Plots for the BEST one
best_results = pd.read_csv(f"{output_dir}/{idr_property}_best_test_r2.csv")
model_type_to_path_dict = dict(zip(best_results['model_type'],best_results['path_to_model']))
for model_type, path_to_model in model_type_to_path_dict.items():
model_preds_folder = path_to_model.split('/best-checkpoint.ckpt')[0]
test_preds = pd.read_csv(f"{model_preds_folder}/{idr_property}_test_predictions.csv")
# make paths for R^2 plots
if not os.path.exists(f"{output_dir}/r2_plots"):
os.makedirs(f"{output_dir}/r2_plots")
os.makedirs(f"{output_dir}/r2_plots/{idr_property}", exist_ok=True)
model_type_dict = {
'fuson_plm': 'FusOn-pLM',
'esm2_t33_650M_UR50D': 'ESM-2'
}
r2_save_path = f"{output_dir}/r2_plots/{idr_property}/{model_type}_{idr_property}_R2.png"
plot_r2(model_type_dict[model_type], idr_property, test_preds, r2_save_path)
def main():
plot_all_r2("results/final", ["asph","scaled_re","scaled_rg","scaling_exp"])
if __name__ == '__main__':
main() |