<jupyter_start><jupyter_text>Evaluate the Performance of MPNN modelsGet all of the models, regardless how we trained them and evaluate their performance<jupyter_code>%matplotlib inline
from matplotlib import pyplot as plt
from datetime import datetime
from sklearn import metrics
from tqdm import tqdm
from glob import glob
import pandas as pd
import numpy as np
import json
import os<jupyter_output><empty_output><jupyter_text>Find the Models and Summarize ThemThere are `best_model.h5` files in subdirectories that contain data on their configuration.<jupyter_code>models = glob(os.path.join('**', 'test_predictions.csv'), recursive=True)
print(f'Found {len(models)} models')
def generate_summary(path):
    """Generate the summary of a model, given path to its output
    
    Args:
        path (str): Path ot the trained weights
    Returns:
        (dict) Model information
    """
    
    # Store the directory first
    dir_name = os.path.dirname(path)
    output = {'path': dir_name}
    
    # Get the host and run parameters
    for f in ['host_info.json', 'run_params.json']:
        with open(os.path.join(dir_name, f)) as fp:
            output.update(json.load(fp))
            
    # Compute the number of nodes
    output['n_nodes'] = output['total_ranks'] // output['ranks_per_node'] \
        if 'total_ranks' in output else 1
    
    # Convert the start time to a datetime
    output['start_time'] = datetime.fromisoformat(output['start_time'])
            
    # Get the log infomration
    log_file = os.path.join(dir_name, 'log.csv')
    log = pd.read_csv(log_file)
    output['completed_epochs'] = len(log)
    output['val_loss'] = log['val_loss'].min()
    output['loss'] = log['loss'].min()
    output['epoch_time'] = np.percentile(log['epoch_time'], 50)
    output['total_train_time'] = log['epoch_time'].sum()
    output['total_node_hours'] = output['total_train_time'] * output['n_nodes']
    
    # Compute performance on hold-out set
    results = pd.read_csv(os.path.join(output['path'], 'test_predictions.csv'))
    for m in ['r2_score', 'mean_squared_error', 'mean_absolute_error', 'median_absolute_error']:
        v = getattr(metrics, m)(results['y_true'], results['y_pred'])
        output[m] = v
    
    return output
model_info = pd.DataFrame([generate_summary(m) for m in models])
print(f'Found {len(model_info)} models')<jupyter_output>Found 35 models<jupyter_text>Print out Best PerformerWe are going to pick the one that has the best performance on the test set  Coarse NetworkSee how we did on the "node per water" network<jupyter_code>model = model_info.query('network_choice=="coarse"').sort_values('mean_absolute_error').iloc[0]
print(f'Model being evaluated: {model["path"]}')
model[['path', 'network_choice', 'activation', 'message_steps', 'dropout', 'features', 'batch_size']]
model[['loss', 'val_loss', 'mean_squared_error']]<jupyter_output><empty_output><jupyter_text>Plot the logs<jupyter_code>log = pd.read_csv(os.path.join(model['path'], 'log.csv'))
fig, ax = plt.subplots(figsize=(3.5, 2.5))

ax.semilogy(log['epoch'], log['loss'], label='Train')
ax.semilogy(log['epoch'], log['val_loss'], label='Validation')

ax.legend()
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')<jupyter_output><empty_output><jupyter_text>*Finding*: Huge variance in validation loss is indicative of overfitting Plot the performance on the test set<jupyter_code>results = pd.read_csv(os.path.join(model['path'], 'test_predictions.csv'))
for m in ['r2_score', 'mean_squared_error', 'mean_absolute_error']:
    v = getattr(metrics, m)(results['y_true'], results['y_pred'])
    print(f'{m}: {v: .2f}')<jupyter_output>r2_score:  1.00
mean_squared_error:  6.13
mean_absolute_error:  1.82<jupyter_text>Plot the true vs predicted<jupyter_code>fig, ax = plt.subplots()

ax.scatter(results['y_true'], results['y_pred'], s=0.5, alpha=0.2)

ax.plot(ax.get_xlim(), ax.get_ylim(), 'k--')
ax.set_xlabel('$E$, True')
ax.set_ylabel('$E$, ML')

fig.set_size_inches(3.5, 3.5)<jupyter_output><empty_output><jupyter_text>Plot only the largest cluster size<jupyter_code>subset = results.query(f'n_waters == {results["n_waters"].max()}')
print(f'Scores for the {len(subset)} largest molecules with {results["n_waters"].max()} waters')
for m in ['r2_score', 'mean_squared_error', 'mean_absolute_error']:
    v = getattr(metrics, m)(subset['y_true'], subset['y_pred'])
    print(f'{m}: {v: .2f}')
fig, ax = plt.subplots()

errors = subset['y_pred'] - subset['y_true']
bins = np.linspace(-10, 10, 256)
ax.hist(errors, bins=bins, density=False)

ax.set_xlabel('Error (kcal/mol)')
ax.set_ylabel('Frequency')
fig.set_size_inches(3.5, 2)
fig, ax = plt.subplots(figsize=(3.5, 3.5))

ax.scatter(subset['y_true'], subset['y_pred'], s=0.5, alpha=0.1)

ax.set_ylim(-340, -305)
ax.set_xlim(ax.get_ylim())
ax.set_ylim(ax.get_xlim())
ax.plot(ax.get_xlim(), ax.get_xlim(), 'k--')
ax.set_xlabel('$E$ (kcal/mol), True')
ax.set_ylabel('$E$ (kcal/mol), ML')

fig.tight_layout()<jupyter_output><empty_output><jupyter_text>Atomic NetworkSee how we did for the "node per atom" network<jupyter_code>model = model_info.query('network_choice=="atomic"').sort_values('mean_absolute_error').iloc[0]
print(f'Model being evaluated: {model["path"]}')
model[['path', 'network_choice', 'activation', 'message_steps', 'dropout', 'features', 'batch_size']]
model[['loss', 'val_loss', 'mean_squared_error']]<jupyter_output><empty_output><jupyter_text>Plot the logs<jupyter_code>log = pd.read_csv(os.path.join(model['path'], 'log.csv'))
fig, ax = plt.subplots()

ax.semilogy(log['epoch'], log['loss'], label='Train')
ax.semilogy(log['epoch'], log['val_loss'], label='Validation')

ax.legend()
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')<jupyter_output><empty_output><jupyter_text>*Finding*: Huge variance in validation loss is indicative of overfitting Plot the performance on the test set<jupyter_code>results = pd.read_csv(os.path.join(model['path'], 'test_predictions.csv'))
for m in ['r2_score', 'mean_squared_error', 'mean_absolute_error']:
    v = getattr(metrics, m)(results['y_true'], results['y_pred'])
    print(f'{m}: {v: .2f}')<jupyter_output>r2_score:  1.00
mean_squared_error:  1.99
mean_absolute_error:  1.03<jupyter_text>Plot the true vs predicted<jupyter_code>fig, ax = plt.subplots(figsize=(3.5, 3.5))

ax.set_title('Performance on hold-out set')
ax.scatter(results['y_true'], results['y_pred'], s=0.5, alpha=0.2)

ax.plot(ax.get_xlim(), ax.get_ylim(), 'k--')
ax.set_xlabel('$E$, True')
ax.set_ylabel('$E$, ML')

fig.set_size_inches(3.5, 3.5)<jupyter_output><empty_output><jupyter_text>Plot only the largest cluster size<jupyter_code>subset = results.query(f'n_waters == {results["n_waters"].max()}')
print(f'Scores for the {len(subset)} largest molecules with {results["n_waters"].max()} waters')
for m in ['r2_score', 'mean_squared_error', 'mean_absolute_error']:
    v = getattr(metrics, m)(subset['y_true'], subset['y_pred'])
    print(f'{m}: {v: .2f}')
fig, ax = plt.subplots()

errors = subset['y_pred'] - subset['y_true']
bins = np.linspace(-10, 10, 256)
ax.hist(errors, bins=bins, density=False)

ax.set_xlabel('Error (kcal/mol)')
ax.set_ylabel('Frequency')
fig.set_size_inches(3.5, 2)
fig, ax = plt.subplots(figsize=(3.5, 3.5))

ax.set_title('Clusters with 30 waters')
ax.scatter(subset['y_true'], subset['y_pred'], s=0.5, alpha=0.1)

ax.set_ylim(-340, -305)
ax.set_xlim(ax.get_ylim())
ax.set_ylim(ax.get_xlim())
ax.plot(ax.get_xlim(), ax.get_xlim(), 'k--')
ax.set_xlabel('$E$ (kcal/mol), True')
ax.set_ylabel('$E$ (kcal/mol), ML')

fig.tight_layout()<jupyter_output><empty_output><jupyter_text>Make a publication-ready figure<jupyter_code>fig, axs = plt.subplots(1, 3, figsize=(6.5, 2.5))

# Predicted vs actual plots
n_waters = results["n_waters"].max()
subset = results.query(f'n_waters == {n_waters}')
for d, ax, title in zip([results, subset], axs,
                        ['Full Dataset', '30-Water Clusters']):
    ax.set_title(title)
    ax.scatter(d['y_true'], d['y_pred'], s=0.7, alpha=0.2, edgecolor='none')

    max_ = max(ax.get_xlim()[1], ax.get_ylim()[1])
    min_ = min(ax.get_xlim()[0], ax.get_ylim()[0])
    ax.set_xlim([min_, max_])
    ax.set_ylim(ax.get_xlim())
    ax.plot(ax.get_xlim(), ax.get_xlim(), 'k--')
    ax.set_xlabel('$E$ (kcal/mol), True')
    ax.set_ylabel('$E$ (kcal/mol), ML')
    
    mae = metrics.mean_absolute_error(d['y_true'], d['y_pred'])
    r2 = metrics.r2_score(d['y_true'], d['y_pred'])
    ax.text(0.99, 0, f'MAE: {mae:.2f}\n$R^2$: {r2:.2f}',
            ha='right', va='bottom', transform=ax.transAxes,
            fontsize=10)
    
# Box and wisker plot
ax = axs[2]
error_stats = []
for s, subset in results.groupby('n_waters'):
    error = np.abs(subset['y_pred'] - subset['y_true']) / s
    error_stats.append({'size': s, 'mae': error.mean()})
error_stats = pd.DataFrame(error_stats)
ax.plot(error_stats['size'], error_stats['mae'], '--o', ms=3)
ax.set_xlabel('# Waters')
ax.set_ylabel('MAE (kcal/mol/water)')

# Add figure labels
for ax, l in zip(axs[:2], ['a', 'b']):
    ax.text(0.02, 0.9, f'({l})', transform=ax.transAxes)
axs[2].text(0.82, 0.9, '(c)', transform=axs[2].transAxes)
    
fig.tight_layout()
fig.savefig(os.path.join('figures', 'mpnn-performance.png'), dpi=320)<jupyter_output><empty_output><jupyter_text>Make the Box PlotTo match Jenna's<jupyter_code>results['abs_error_per_water'] = np.abs(results['y_true'] - results['y_pred']) / results['n_waters']
def make_box_plot(df, metric='abs_error_per_water'):       
    boxplot = df.query('n_waters >= 10 and n_waters <= 30').boxplot(metric, 'n_waters', grid=False, fontsize=20, figsize=(12,6), return_type='both')
    plt.ylim(-0.01,0.7)
    plt.ylabel('Absolute Error\n(kcal/mol/water)', fontsize=22, fontweight='bold', labelpad=15)
    plt.xlabel('Cluster Size', fontsize=22, fontweight='bold', labelpad=15)
    plt.xticks(range(1,23,2), ['10','12','14','16','18','20','22','24','26','28','30'])
    plt.xlim(0, 22)
    plt.suptitle('')
    plt.title('')
    plt.tight_layout()
    plt.savefig('figures/mpnn_boxplot-horz.png',dpi=600)
make_box_plot(results)<jupyter_output><empty_output><jupyter_text>Evaluate Hyperparameter SweepsWe did some manual hyperparameter tuning for the atomic model  Batch SizesEvaluate different batch sizes to get a tradeoff between accuracy and using the full GPU<jupyter_code>base_query = ('epochs==32 and shuffle_buffer_size==2097152 and activation=="sigmoid" '
              'and message_steps==4 and network_choice=="atomic" and dropout==0 and features==64')
model_info.query(base_query).sort_values('val_loss')[['batch_size', 'loss', 'val_loss', 'mean_squared_error', 'epoch_time']]<jupyter_output><empty_output><jupyter_text>*Finding*: We get decent accuracy with a batch size of 1024 and still use 90% of the GPU  Activation FunctionWe evaluated different activation functions for the message steps<jupyter_code>base_query = ('batch_size==1024 and epochs==32 and shuffle_buffer_size==2097152 '
              'and message_steps==4 and network_choice=="atomic" and dropout==0 and features==64')
model_info.query(base_query).sort_values('mean_squared_error')[['activation', 'loss', 'val_loss', 'mean_squared_error', 'epoch_time']]<jupyter_output><empty_output><jupyter_text>*Finding*: We should go with the softplus. Fastest and most accurate  Number of Message Passing LayersWe compared increasing the number of message passing layers<jupyter_code>base_query = ('hostname=="lambda3" and shuffle_buffer_size==2097152 and batch_size==1024 and activation=="softplus" and epochs==32 '
              'and network_choice=="atomic"')
model_info.query(base_query).sort_values('message_steps')[['network_choice', 'message_steps', 
                                                           'loss', 'val_loss', 'mean_squared_error',
                                                          'epoch_time']]
fig, ax = plt.subplots()

for label, subset in model_info.query(base_query).sort_values('message_steps').groupby('network_choice'):
    ax.plot(subset['message_steps'], subset['mean_absolute_error'], '-o', label=label)

ax.set_xscale('log', base=2)
ax.set_xlabel('Message Steps')
ax.set_ylabel('Mean Absolute Error')
ax.legend()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1715128184126464, 'fraction_numerical': 0.021459578999099026, 'mean_word_length': 5.529411764705882, 'pattern_counts': {'":': 0, '<': 67, '<?xml version=': 0, '>': 67, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '17', 'dir': 'jupyter-structured-clean-dedup', 'id': '11455', 'license': 'Apache-2.0', 'n_tokens_mistral': 4675, 'n_tokens_neox': 4449, 'n_words': 878, 'path': 'challenge-2/1_evaluate-models.ipynb', 'repo_name': 'exalearn/hydronet', 'size': '12209'}