<jupyter_start><jupyter_text>Deep-TICA - Alanine, multithermalThis is a brief tutorial that describes how to train the Deep-TICA CVs from a biased simulation. We use a OPES-multithermal simulation of alanine dipeptide as example, using all distances between heavy atoms as input descriptors.This example is taken from: _Bonati, Piccini and Parrinello, Deep learning the slow modes for rare events sampling, PNAS (2021)_.  Import<jupyter_code>import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from mlcvs.utils.io import load_dataframe
from mlcvs.tica import DeepTICA_CV<jupyter_output><empty_output><jupyter_text>Load data We use the function `load_dataframe` from `mlcvs.utils` that takes a PLUMED COLVAR file and returns a pandas dataframe. First we load the COLVAR file containing the data about the OPES-multicanonical simulation and then the COLVAR_DRIVER which contains the input descriptors calculated along the trajectories (the second step can be avoided if they are already present in the first file).We also discard the initial part of the simulation to put ourselves in the quasi-static regime, which is very quickly reached by OPES (it can be monitored by looking at the DELTAFS file for the multicanonical simulation).---| Parameter | Type | Description || :- | :- | :- || folder | str |files location || from_time | float | discard initial part of the trajectories [time units] || descriptor_filter | str |filter descriptors based on names |<jupyter_code>#------------- PARAMETERS -------------
folder            = 'data/ala2_multithermal/'
from_time         = 15000
descriptor_filter = 'd'
#--------------------------------------

# Load data
opes     =  load_dataframe(folder+"COLVAR", start=from_time)
features =  load_dataframe(folder+"COLVAR_DRIVER", start=from_time)

colvar = pd.concat([opes, features[[i for i in features.columns if (i not in opes.columns)]]], axis=1)
# Select descriptors 
X = colvar.filter(regex=descriptor_filter).values
n_input = X.shape[1]

print(X.shape)<jupyter_output>(35001, 45)<jupyter_text>Compute weights for time rescaling Here we extract the time $t$, the energy $E$ (needed for the multicanonical reweight [1]) and the bias $V$ from the COLVAR file. We then calculate the weights as:\begin{cases}    w = e^{\beta\ V + (\beta_0-\beta)\ E} & \text{if multicanonical}\\    w = e^{\beta\ V}              & \text{otherwise}\end{cases}NB: note that if simulation temperature $\beta_0$ is equal to the reweighting one $\beta$ the multicanonical reweight coincides with the standard umbrella sampling-like case.Once we have computed the weights, we rescale the time at step $k$ by using the instantaneus acceleration:$$ dt'_k = w_k\ dt $$and then compute the cumulative rescaled time:$$ t'_k = \sum_{i=0} ^k dt'_i $$[1] Invernizzi, Piaggi, and Parrinello. "Unified approach to enhanced sampling." _Physical Review X_ 10.4 (2020): 041034.---| Parameter | Type | Description || :- | :- | :- || multicanonical | bool | flag to determine if using a standard reweight (false) or a multicanonical one (true) || temp | float | reweighting temperature || temp0 | float | simulation temperature (only needed if multicanonical == True) |<jupyter_code>#------------- PARAMETERS -------------
multicanonical    = True
temp              = 300.
temp0             = 300.
#--------------------------------------

# Calculate inverse temperature
kb=0.008314
beta=1./(kb*temp)
beta0=1./(kb*temp0)

# Extract cvs from df

t = colvar['time'].values # save time
ene = colvar['ene'].values.astype(np.float64) # store energy as long double
bias = colvar.filter(regex='.bias').values.sum(axis=1) # Load *.bias columns and sum them

# Compute log-weights for time reweighting
logweight = beta*bias

if multicanonical:
    ene -= np.mean(ene) #first shift energy by its mean value
    logweight += (beta0-beta)*ene<jupyter_output><empty_output><jupyter_text>Create dataset of time-lagged configurations In order to train the Deep-TICA CVs we will need to compute the time-lagged covariance matrices in the rescaled time $t'$. The standard way is to look for configurations which are distant a lag-time $\tau$ in the time series. However, in the rescaled time the time-series is _exponentially_ unevenly spaced. Hence, a naive search will lead to severe numerical issue. To address this, we use the algorithm proposed in [2]. In a nutshell, this method assume that the observable $O(t'_k)$ have the same value from scaled time $t'_k$ to $t'_{k+1}$. This leads to weighting each pair of configurations based both on the rescaled time around $t'_k$ and $t'_k+\tau$ (see supp. information of [2] for details). All of this is done under the hood by the function `mlcvs.utils.data.find_time_lagged_configurations`. To generate the the training and validation set, we use the function `create_time_lagged_dataset` which searches for the pairs of configurations and the corresponding weight. This is divided in training and validation data, and fed to a `FastTensorDataloader` which allows for an efficient training of the model.[2] Yang and Parrinello. "Refining collective coordinates and improving free energy representation in variational enhanced sampling." _Journal of chemical theory and computation_ 14.6 (2018): 2889-2894.---| Parameter | Type | Description || :- | :- | :- || lag_time | float | lag_time for the calculation of the covariance matrices [in rescaled time] || n_train | int | number of training configurations || n_valid | int | number of validation configurations |<jupyter_code>from mlcvs.utils.data import create_time_lagged_dataset, FastTensorDataLoader
from torch.utils.data import Subset

#------------- PARAMETERS -------------
lag_time = 0.1
n_train  = 25000 
n_valid  = 10000
#--------------------------------------

# create dataset
dataset = create_time_lagged_dataset(X,t=t,lag_time=lag_time,logweights=logweight)

# split train - valid 
train_data = Subset(dataset, np.arange(2*n_train))
valid_data = Subset(dataset, np.arange(2*n_train,len(dataset)))

# create dataloaders 
train_loader = FastTensorDataLoader(train_data, batch_size=len(train_data))
valid_loader = FastTensorDataLoader(valid_data, batch_size=len(valid_data))

print('Time-lagged pairs:\t',len(dataset))
print('Training data:\t\t',len(train_data))
print('Validation data:\t',len(valid_data))<jupyter_output>Time-lagged pairs:	 69998
Training data:		 50000
Validation data:	 19998<jupyter_text>Training Here we setup a few parameters and then train the Deep-TICA CVs. We first instantiate a object of the `DeepTICA_CV` class which defines the NN but also the loss function and a train loop. See class documentation for further details about parameters and methods.---| Parameter | Type | Description || :- | :- | :- || **Neural network** || nodes | list | NN architecture (last value equal to the number of hidden layers which are input of TICA)  || activ_type | string | Activation function (relu,tanh,elu,linear) || loss_type | string | Loss function operating on the TICA eigenvalues (sum2,sum,single)|| n_eig | int | Number of eigenvalues to optimize (or if loss_type=single which one to select) || **Optimization** || lrate | float | Learning rate || l2_reg | float | L2 regularization || num_epochs | int | Number of epochs || **Early Stopping** || earlystop | bool | Whether to use early stopping based on validation loss || es_patience | int | Number of epochs before stopping || es_consecutive | bool | Whether es_patience should count consecutive (True) or cumulative patience || **Log** || log_every | int | How often print the train/valid loss during training |<jupyter_code>#------------- PARAMETERS -------------
nodes             = [n_input,30,30,2]
activ_type        = 'tanh'
loss_type         = 'sum2'
n_eig             = nodes[-1]

lrate             = 1e-3
l2_reg            = 0.

num_epochs        = 1000
earlystop         = True
es_patience       = 10
es_consecutive    = False

log_every         = 10
#--------------------------------------

# DEVICE
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# MODEL
model = DeepTICA_CV(nodes)
model.to(device)

# OPTIMIZER
opt = torch.optim.Adam(model.parameters(), lr=lrate, weight_decay=l2_reg)

# REGULARIZATION
model.set_optimizer(opt)
model.set_earlystopping(patience=es_patience,min_delta=0.,consecutive=es_consecutive, save_best_model=True, log=False)

# TRAIN
model.fit(train_loader,valid_loader,
            standardize_inputs=True,
            standardize_outputs=True,
            loss_type=loss_type,
            n_eig=n_eig,
            nepochs=num_epochs,
            info=False, log_every=log_every)
## move the model back to cpu for convenience
model.to('cpu')<jupyter_output><empty_output><jupyter_text>In the next cells we analyze training and the resulting Deep-TICA CVs.  Plot learning curve<jupyter_code>fig, axs = plt.subplots(1,2,figsize=(12,5),dpi=100)

loss_train = [x.cpu() for x in model.loss_train]
loss_valid = [x.cpu() for x in model.loss_valid]

# Loss function
ax = axs[0]
ax.plot(loss_train,'-',label='Train')
ax.plot(loss_valid,'--',label='Valid')
ax.set_ylabel('Loss Function')

# Eigenvalues vs epoch
ax = axs[1]
with torch.no_grad():
    evals_train = np.asarray(torch.cat(model.evals_train).cpu())
for i in range(n_eig):
    ax.plot(evals_train[:,i],label='Eig. '+str(i+1))
ax.set_ylabel('Eigenvalues')

# Common setup
for ax in axs:
    if model.earlystopping_.early_stop:
        ax.axvline(model.earlystopping_.best_epoch,ls='dotted',color='grey',alpha=0.5,label='Early Stopping')
        ax.set_xlabel('#Epochs')
        ax.legend(ncol=2)

plt.tight_layout()
plt.show()<jupyter_output><empty_output><jupyter_text>Plot CVs isolines in $\phi-\psi$ space<jupyter_code># Hexbin plot in physical space    
fig,axs = plt.subplots(1,n_eig,figsize=(6*n_eig,5),dpi=100)

x = colvar['phi'].values
y = colvar['psi'].values

# compute cvs
with torch.no_grad():
    s = model(torch.Tensor(X)).numpy()
    
for i,ax in enumerate(axs):
    pp = ax.hexbin(x,y,C=s[:,i],gridsize=150,cmap='fessa')
    cbar = plt.colorbar(pp,ax=ax)
    ax.set_title('Deep-TICA '+str(i+1))
    ax.set_xlabel(r'$\phi$ [rad]')
    ax.set_ylabel(r'$\psi$ [rad]')
    cbar.set_label('Deep-TICA '+str(i+1))

plt.tight_layout()
plt.show()<jupyter_output><empty_output><jupyter_text>Plot FES<jupyter_code>from mlcvs.utils.fes import compute_fes

fig,axs = plt.subplots(1,n_eig,figsize=(6*n_eig,5),dpi=100)

for i in range(n_eig):
    fes,grid,bounds,error = compute_fes(s[:,i], weights=np.exp(logweight),
                                        blocks=3,
                                        bandwidth=0.02,scale_by='range',
                                        plot=True, plot_max_fes=100, ax = axs[i])
    axs[i].set_xlabel('Deep-TICA '+str(i+1))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12813164463344734, 'fraction_numerical': 0.01580844966256818, 'mean_word_length': 3.895022624434389, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '620423', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3509, 'n_tokens_neox': 3304, 'n_words': 1061, 'path': 'docs/notebooks/ala2_deeptica_multithermal.ipynb', 'repo_name': 'luigibonati/mlcvs', 'size': '10817'}