<jupyter_start><jupyter_text>Top-down learning a GNN potential for coarse-grained waterIn this notebook, we will learn a graph neural network (GNN) potential for coarse-grained water from experimental data.The coarse-grained particle will represent the oxygen atom of atomistic water.Analogous to the diamond example, we will directly apply the DiffTRe library without looking under its hood.We define all components necessary to initialize DiffTRe: experimental data, simulation box, GNN potential with prior, simulator, loss function and optimizer.<jupyter_code>from DiffTRe import io, difftre, custom_simulator, custom_quantity, custom_space, custom_energy
from jax_md import space, partition, simulate
import optax
import jax.numpy as jnp
from jax import device_get, random, jit, checkpoint, tree_util
import numpy as np
from scipy import interpolate as sci_interpolate
import matplotlib.pyplot as plt
from functools import partial
import time, pickle
import warnings
warnings.filterwarnings('ignore')  # disable warnings about float64 usage

energy_pickle_file_path = 'output/saved_models/CG_water_GNN.pkl'
optimization_pickle_file_path = 'output/saved_optimization_results/CG_water_optimization.pkl'
long_traj_pickle_file_path = 'output/saved_optimization_results/CG_water_long_traj.pkl'<jupyter_output><empty_output><jupyter_text>Experimental observablesLet's start by defining target observables. For top-down coarse-graining, we choose experimental data from the paper_Quantum Differences between Heavy and Light Water_ [(<NAME>, 2008)](https://doi.org/10.1103/PhysRevLett.101.065502).Specifically, we want to match the oxygen-oxygen radial distribution function (RDF), the oxygen-oxygen-oxygen angulardistribution function (ADF), and a pressure of 1 bar.<jupyter_code>pressure_conversion = 16.6054  # from kJ/mol nm^-3 to bar
pressure_target = 1. / pressure_conversion  # 1 bar in kJ / mol nm^3

# O-O RDF
rdf_bin_centers, rdf_bin_boundaries, sigma_RDF = custom_quantity.rdf_discretization(RDF_cut=1.0)  # cut RDF at 1nm
reference_rdf = np.loadtxt('data/experimental/O_O_RDF.csv')
rdf_spline = sci_interpolate.interp1d(reference_rdf[:, 0], reference_rdf[:, 1], kind='cubic')
reference_rdf = rdf_spline(rdf_bin_centers)
rdf_struct = custom_quantity.RDFParams(reference_rdf, rdf_bin_centers, rdf_bin_boundaries, sigma_RDF)

# O-O-O ADF
adf_bin_centers, sigma_ADF = custom_quantity.adf_discretization(nbins=200)
reference_adf = np.loadtxt('data/experimental/O_O_O_ADF.csv')
adf_spline = sci_interpolate.interp1d(reference_adf[:, 0], reference_adf[:, 1], kind='cubic')
reference_adf = adf_spline(adf_bin_centers)
# only triplets within 0.318 nm are considered, consistent with experimental data
adf_struct = custom_quantity.ADFParams(reference_adf, adf_bin_centers, sigma_ADF, r_outer=0.318, r_inner=0.)

target_dict = {'rdf': rdf_struct, 'adf': adf_struct, 'pressure': pressure_target}<jupyter_output><empty_output><jupyter_text>Simulation setupNext, we define the simulation parameters.Temperature and density are already determined by the experiment.<jupyter_code>system_temperature = 296.15  # Kelvin = 23 deg. Celsius
Boltzmann_constant = 0.0083145107  # in kJ / mol K
kbT = system_temperature * Boltzmann_constant

file = 'data/confs/Water_experimental.gro'  # 901 particles
R_init, v, box = io.load_configuration(file)  # initial configuration
N = R_init.shape[0]

mass = 18.0154  # in u: O + 2 * H
density = mass * N * 1.66054 / jnp.prod(box)
print('Model Density:', density, 'g/l. Experimental density: 997.87 g/l')

# simulation times in ps
time_step = 0.002
total_time = 70. 
t_equilib = 10.  # equilibration time: skip states before sampling observables
print_every = 0.1  # save state every 0.1 ps for computation of averages
timings_struct = custom_simulator.process_printouts(time_step, total_time, t_equilib, print_every)

# define displacement and shift functions from box: 
# use periodic_general as we take derivatives wrt. box for pressure
R_init, box_tensor = custom_space.scale_to_fractional_coordinates(R_init, box)
displacement, shift = space.periodic_general(box_tensor)<jupyter_output>Model Density: 998.28436 g/l. Experimental density: 997.87 g/l<jupyter_text>GNN PotentialNow we define the GNN as well as the prior potential. As our prior potential, we choose the repulsive part of theLennard-Jones (LJ) potential. As the prior length scale $\sigma$, we simply take the length scale of the LJ potential of the SPC water model $\sigma=0.3165$ nm and set the energy scale to $\epsilon = 1$ kJ / mol - reduced compared to the LJ potential as we neglect the attractive term.Both potentials are combined by adding the computed potential energy values. The `energy` function that sums both potentialsis initialized by the `energy_fn_template`. The template enables updating GNN weights during optimizationwithout compromising "jitability". You can define your own custom `energy_fn_template` following the in- and outputsignature shown below.<jupyter_code>key = random.PRNGKey(0)  # define random seed for initialization of model and simulation
model_init_key, simuation_init_key = random.split(key, 2)

r_cut = 0.5

# create neighbor list
box_nbrs = jnp.ones(3)
neighbor_fn = partition.neighbor_list(displacement, box_nbrs, r_cut, dr_threshold=0.05, capacity_multiplier=1.5,
                                      disable_cell_list=True)
nbrs_init = neighbor_fn(R_init)

prior_fn = custom_energy.generic_repulsion_neighborlist(displacement, sigma=0.3165, epsilon=1., exp=12,
                                                         initialize_neighbor_list=False)

init_fn, GNN_energy = custom_energy.DimeNetPP_neighborlist(displacement, R_init, nbrs_init, r_cut)
init_params = init_fn(model_init_key, R_init, neighbor=nbrs_init)

def energy_fn_template(energy_params):
    gnn_energy = partial(GNN_energy, energy_params)
    def energy(R, neighbor, **dynamic_kwargs):
        return gnn_energy(R, neighbor=neighbor, **dynamic_kwargs) + prior_fn(R, neighbor=neighbor, **dynamic_kwargs)
    return jit(energy)<jupyter_output><empty_output><jupyter_text>Loss functionAs some observables depend on the potential (e.g. the pressure in this case), we need the `energy_fn_template`to define the loss function. `DiffTRe_init` builds a mean-squared-error (MSE) loss from a `quantity_dict`.Building the `quantity_dict` requires to define for each observable $O_k(S_i, U_{\theta})$ a function to compute the respectiveobservable for each state $S_i$, the corresponding target value, and a scale $\gamma$ defining the relative importanceof each observable in the loss (see examples below).<jupyter_code>quantity_dict = {}
if 'rdf' in target_dict:
    rdf_struct = target_dict['rdf']
    rdf_fn = custom_quantity.initialize_radial_distribution_fun(box_tensor, displacement, rdf_struct)
    rdf_dict = {'compute_fn': checkpoint(rdf_fn), 'target': rdf_struct.reference_rdf, 'gamma': 1.}
    quantity_dict['rdf'] = rdf_dict
if 'adf' in target_dict:
    adf_struct = target_dict['adf']
    adf_fn = custom_quantity.initialize_angle_distribution_neighborlist(displacement, adf_struct,
                                                                         R_init=R_init, nbrs_init=nbrs_init)
    adf_target_dict = {'compute_fn': checkpoint(adf_fn), 'target': adf_struct.reference_adf, 'gamma': 1.}
    quantity_dict['adf'] = adf_target_dict
if 'pressure' in target_dict:
    pressure_fn = custom_quantity.init_pressure(energy_fn_template, box_tensor)
    pressure_target_dict = {'compute_fn': checkpoint(pressure_fn), 'target': target_dict['pressure'], 'gamma': 1.e-7}
    quantity_dict['pressure'] = pressure_target_dict<jupyter_output><empty_output><jupyter_text>SimulatorNow we define the simulator used to generate reference trajectories.Similar to the `energy_fn_template`, the `simulator_template` takes the current potential energy functionand returns a corresponding NVT simulator. This is achieved by simply "baking in" all constantsimulation parameters via `partial`. For this example, we select a Nose-Hoover thermostat with a chain length of 5.With a given simulator, we generate the initial simulation state from which the first trajectory is sampled.<jupyter_code>energy_fn_init = energy_fn_template(init_params)
simulator_template = partial(simulate.nvt_nose_hoover, shift_fn=shift, dt=time_step, kT=kbT,
                             chain_length=5, chain_steps=1)
init, _ = simulator_template(energy_fn_init)
state = init(simuation_init_key, R_init, mass=mass, neighbor=nbrs_init)
init_sim_state = (state, nbrs_init)<jupyter_output><empty_output><jupyter_text>OptimizerWe use an Adam optimizer with exponential learning rate decay.<jupyter_code>num_updates = 300
initial_lr = 0.003
lr_schedule = optax.exponential_decay(-initial_lr, 200, 0.01)
optimizer = optax.chain(
    optax.scale_by_adam(0.1, 0.4),
    optax.scale_by_schedule(lr_schedule)
)<jupyter_output><empty_output><jupyter_text>Initialize DiffTReWith all individual components defined, we can initialize DiffTRe. `update_fn` can be called iteratively to update potential parameters using the whole DiffTRepipeline. `trajectory_state` is the initial reference trajectory.<jupyter_code>simulation_funs = (simulator_template, energy_fn_template, neighbor_fn)
update_fn, trajectory_state = difftre.DiffTRe_init(simulation_funs, timings_struct, quantity_dict, kbT,
                                                   init_params, init_sim_state, optimizer)<jupyter_output>Time for a single trajectory generation: 8.614710903167724 mins<jupyter_text>Learning the GNN potentialThe optimization takes several hours. You can skip it and load the trained model in the next cell.<jupyter_code>loss_history, times_per_update, predicted_quantities = [], [], []

params = init_params
opt_state = optimizer.init(init_params)  # initialize optimizer state
for step in range(num_updates):
    start_time = time.time()
    params, opt_state, trajectory_state, loss_val, predictions = update_fn(step, params, opt_state, trajectory_state)
    loss_val.block_until_ready()
    step_time = time.time() - start_time

    times_per_update.append(step_time)
    loss_history.append(loss_val)
    predicted_quantities.append(predictions)
    print("Step {} in {:0.2f} sec".format(step, step_time), 'Loss = ', loss_val, '\n')

    if jnp.isnan(loss_val):  # stop learning when optimization diverged
        print('Loss is NaN. This was likely caused by divergence of the optimization or a bad model setup '
              'causing a NaN trajectory.')
        break

# save optimized energy_params for re-use
final_energy_params = device_get(params)
with open(energy_pickle_file_path, 'wb') as f:
    pickle.dump(final_energy_params, f)

with open(optimization_pickle_file_path, 'wb') as f:
    pickle.dump([loss_history, times_per_update, predicted_quantities], f)
# load results if necessary
final_energy_params = pickle.load(open(energy_pickle_file_path, "rb" ))
final_energy_params = tree_util.tree_map(jnp.array, final_energy_params)
loss_history, times_per_update, predicted_quantities = pickle.load(open(optimization_pickle_file_path, "rb" ))<jupyter_output><empty_output><jupyter_text>Visualize training processTo assess convergence, we can visualize the loss and compute-time per update. In the beginning, recomputation of the reference trajectory is necessary at each update, due to large parameter updates. Later, with decreased step sizes, referencetrajectories can be re-used frequently, resulting in a significant speed-up as only the reweighting step and the gradientneed to be computed.<jupyter_code>fig, ax1 = plt.subplots()
ax1.set_xlabel('Update Step')
ax1.set_ylabel('MSE Loss')
ax1.semilogy(loss_history, color='#3c5488ff', label='Loss')
ax2 = ax1.twinx()
ax2.plot(times_per_update, label='Time per update', color='#4dbbd5ff')
ax2.set_ylabel('$t$ in s')
fig.legend(loc="upper right", bbox_to_anchor=(0.38,0.2), bbox_transform=ax1.transAxes)
plt.savefig('Train_history_CG_water.png')
plt.show()<jupyter_output><empty_output><jupyter_text>Next, we visualize convergence of target observables.<jupyter_code>fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 4))

if 'rdf' in predicted_quantities[0]:
    rdf_series = [prediction_dict['rdf'] for prediction_dict in predicted_quantities]

    ax1.set_xlabel('$r$ in nm')
    ax1.set_ylabel('RDF')
    ax1.plot(rdf_bin_centers, rdf_series[0], label='initial guess', color='#00a087ff')
    ax1.plot(rdf_bin_centers, rdf_series[-1], label='predicted', color='#3c5488ff')
    ax1.plot(rdf_bin_centers, reference_rdf, label='target', linestyle='--', color='k')
    ax1.legend()


if 'adf' in predicted_quantities[0]:
    adf_series = [prediction_dict['adf'] for prediction_dict in predicted_quantities]

    ax2.set_xlabel(r'$\alpha$ in rad')
    ax2.set_ylabel('ADF')
    ax2.plot(adf_bin_centers, adf_series[0], label='initial guess', color='#00a087ff')
    ax2.plot(adf_bin_centers, adf_series[-1], label='predicted', color='#3c5488ff')
    ax2.plot(adf_bin_centers, reference_adf, label='target', linestyle='--', color='k')
    ax2.legend(loc=1)

if 'pressure' in predicted_quantities[0]:
    pressure_series = [prediction_dict['pressure'] for prediction_dict in predicted_quantities]

    ax3.set_xlabel('Update Step')
    ax3.set_ylabel('$p$ in $\mathrm{kJ / mol nm^3}$')
    ax3.plot(pressure_series, label='predicted', color='#3c5488ff')
    ax3.axhline(y=pressure_target, linestyle='--', color='k', label='target')
    ax3.legend()
    print('Pressure final epoch:', pressure_series[-1], 'kJ / mol nm^3')

plt.savefig('Predicted_CG_water_properties.png')
plt.show()<jupyter_output>Pressure final epoch: 2.4268165 kJ / mol nm^3<jupyter_text>Validate resultsGiven that we've trained the model on rather short trajectories, let's validate the model on a longer trajectory.This allows to detect if the model overfitted to initial conditions or drifts away from the targets.<jupyter_code>total_time_long = 11000.  # 11 ns
t_equilib_long = 1000.  # 1 ns
long_trajectory_struct = custom_simulator.process_printouts(time_step, total_time_long, t_equilib_long, print_every)

trajectory_generator = difftre.trajectory_generator_init(simulation_funs, long_trajectory_struct)
long_traj_state = trajectory_generator(final_energy_params, init_sim_state)
quantity_traj = difftre.compute_quantity_traj(long_traj_state, quantity_dict, neighbor_fn, final_energy_params)

long_traj_rdf = jnp.mean(quantity_traj['rdf'], axis=0)
long_traj_adf = jnp.mean(quantity_traj['adf'], axis=0)
long_traj_press = jnp.mean(quantity_traj['pressure'], axis=0)

with open(long_traj_pickle_file_path, 'wb') as f:
    pickle.dump([long_traj_rdf, long_traj_adf, long_traj_press], f)
# you can skip the long run above and just load results
long_traj_rdf, long_traj_adf, long_traj_press = pickle.load(open(long_traj_pickle_file_path, "rb" ))<jupyter_output><empty_output><jupyter_text>Comparing results from the long trajectory with our targets confirms that training was successful.<jupyter_code>print('Prediced pressure:', long_traj_press, 'kJ / mol nm^3; Target:', pressure_target, 'kJ / mol nm^3')

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), dpi=300)

ax1.set_xlabel('$r$ in nm')
ax1.set_ylabel('RDF')
ax1.plot(rdf_bin_centers, long_traj_rdf, label='predicted', color='#3c5488ff')
ax1.plot(rdf_bin_centers, reference_rdf, label='target', linestyle='--', color='k')
ax1.legend()

ax2.set_xlabel(r'$\alpha$ in rad')
ax2.set_ylabel('ADF')
ax2.plot(adf_bin_centers, long_traj_adf, label='predicted', color='#3c5488ff')
ax2.plot(adf_bin_centers, reference_adf, label='target', linestyle='--', color='k')
ax2.legend(loc=1)

plt.savefig('Predicted_water_properties_long_traj.png')
plt.show()<jupyter_output>Prediced pressure: 0.7751365 kJ / mol nm^3; Target: 0.060221373769978444 kJ / mol nm^3
{'fraction_non_alphanumeric': 0.09692940804051915, 'fraction_numerical': 0.02228553339664451, 'mean_word_length': 5.915936952714536, 'pattern_counts': {'":': 0, '<': 50, '<?xml version=': 0, '>': 50, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '600767', 'license': 'Apache-2.0', 'n_tokens_mistral': 5504, 'n_tokens_neox': 5146, 'n_words': 1322, 'path': 'CG_Water.ipynb', 'repo_name': 'moradza/difftre', 'size': '15806'}