Note

Go to the end to download the full example code.

TRAM on a 1D double well

This example shows how to use the transition-based reweighting analysis method (TRAM) to estimate the free energies and Markov model of a simple double-well potential, sampled using umbrella sampling.

For more information see the TRAM estimator and its respective TRAM tutorial.

plot tram
  0%|          | 0/1000 [00:00<?, ?it/s]
Initializing free energies using MBAR:   0%|          | 0/1000 [00:00<?, ?it/s]
Initializing free energies using MBAR - [inc: 3.4e-01]:   0%|          | 1/1000 [00:00<00:04, 247.82it/s]
Initializing free energies using MBAR - [inc: 9.0e-02]:   0%|          | 2/1000 [00:00<00:03, 259.81it/s]
Initializing free energies using MBAR - [inc: 3.3e-02]:   0%|          | 3/1000 [00:00<00:03, 269.21it/s]
Initializing free energies using MBAR - [inc: 1.7e-02]:   0%|          | 4/1000 [00:00<00:03, 275.89it/s]
Initializing free energies using MBAR - [inc: 1.5e-02]:   0%|          | 5/1000 [00:00<00:03, 280.07it/s]
Initializing free energies using MBAR - [inc: 1.0e-02]:   1%|          | 6/1000 [00:00<00:03, 275.34it/s]
Initializing free energies using MBAR - [inc: 6.5e-03]:   1%|          | 7/1000 [00:00<00:03, 271.34it/s]
Initializing free energies using MBAR - [inc: 4.2e-03]:   1%|          | 8/1000 [00:00<00:03, 269.40it/s]
Initializing free energies using MBAR - [inc: 2.7e-03]:   1%|          | 9/1000 [00:00<00:03, 270.66it/s]
Initializing free energies using MBAR - [inc: 1.8e-03]:   1%|          | 10/1000 [00:00<00:03, 272.97it/s]
Initializing free energies using MBAR - [inc: 1.1e-03]:   1%|          | 11/1000 [00:00<00:03, 274.98it/s]
Initializing free energies using MBAR - [inc: 7.4e-04]:   1%|          | 12/1000 [00:00<00:03, 276.85it/s]
Initializing free energies using MBAR - [inc: 4.8e-04]:   1%|▏         | 13/1000 [00:00<00:03, 278.72it/s]
Initializing free energies using MBAR - [inc: 3.1e-04]:   1%|▏         | 14/1000 [00:00<00:03, 280.18it/s]
Initializing free energies using MBAR - [inc: 2.0e-04]:   2%|▏         | 15/1000 [00:00<00:03, 281.40it/s]
Initializing free energies using MBAR - [inc: 1.3e-04]:   2%|▏         | 16/1000 [00:00<00:03, 280.89it/s]
Initializing free energies using MBAR - [inc: 8.3e-05]:   2%|▏         | 17/1000 [00:00<00:03, 281.98it/s]
Initializing free energies using MBAR - [inc: 5.4e-05]:   2%|▏         | 18/1000 [00:00<00:03, 283.01it/s]
Initializing free energies using MBAR - [inc: 3.5e-05]:   2%|▏         | 19/1000 [00:00<00:03, 283.99it/s]
Initializing free energies using MBAR - [inc: 2.2e-05]:   2%|▏         | 20/1000 [00:00<00:03, 284.81it/s]
Initializing free energies using MBAR - [inc: 1.5e-05]:   2%|▏         | 21/1000 [00:00<00:03, 285.59it/s]
Initializing free energies using MBAR - [inc: 9.4e-06]:   2%|▏         | 22/1000 [00:00<00:03, 286.19it/s]
Initializing free energies using MBAR - [inc: 6.1e-06]:   2%|▏         | 23/1000 [00:00<00:03, 286.83it/s]
Initializing free energies using MBAR - [inc: 3.9e-06]:   2%|▏         | 24/1000 [00:00<00:03, 287.41it/s]
Initializing free energies using MBAR - [inc: 2.5e-06]:   2%|▎         | 25/1000 [00:00<00:03, 287.94it/s]
Initializing free energies using MBAR - [inc: 1.6e-06]:   3%|▎         | 26/1000 [00:00<00:03, 288.50it/s]
Initializing free energies using MBAR - [inc: 1.1e-06]:   3%|▎         | 27/1000 [00:00<00:03, 288.99it/s]
Initializing free energies using MBAR - [inc: 6.9e-07]:   3%|▎         | 28/1000 [00:00<00:03, 289.22it/s]
Initializing free energies using MBAR - [inc: 6.9e-07]:   3%|▎         | 29/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 4.4e-07]:   3%|▎         | 29/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 2.9e-07]:   3%|▎         | 30/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 1.9e-07]:   3%|▎         | 31/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 1.2e-07]:   3%|▎         | 32/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 7.8e-08]:   3%|▎         | 33/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 5.0e-08]:   3%|▎         | 34/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 3.2e-08]:   4%|▎         | 35/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 2.1e-08]:   4%|▎         | 36/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 1.4e-08]:   4%|▎         | 37/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 8.8e-09]:   4%|▍         | 38/1000 [00:00<00:03, 289.37it/s]
Initializing free energies using MBAR - [inc: 8.8e-09]: 100%|██████████| 38/38 [00:00<00:00, 287.66it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]
Running TRAM estimate:   0%|          | 0/1000 [00:00<?, ?it/s]
Running TRAM estimate - [inc: 9.9e-01]:   0%|          | 1/1000 [00:00<00:03, 270.55it/s]
Running TRAM estimate - [inc: 3.3e-01]:   0%|          | 2/1000 [00:00<00:04, 237.58it/s]
Running TRAM estimate - [inc: 1.8e-01]:   0%|          | 3/1000 [00:00<00:04, 248.20it/s]
Running TRAM estimate - [inc: 1.1e-01]:   0%|          | 4/1000 [00:00<00:03, 255.82it/s]
Running TRAM estimate - [inc: 7.1e-02]:   0%|          | 5/1000 [00:00<00:03, 261.02it/s]
Running TRAM estimate - [inc: 4.6e-02]:   1%|          | 6/1000 [00:00<00:03, 265.48it/s]
Running TRAM estimate - [inc: 3.0e-02]:   1%|          | 7/1000 [00:00<00:03, 267.85it/s]
Running TRAM estimate - [inc: 2.0e-02]:   1%|          | 8/1000 [00:00<00:03, 270.64it/s]
Running TRAM estimate - [inc: 1.3e-02]:   1%|          | 9/1000 [00:00<00:03, 272.99it/s]
Running TRAM estimate - [inc: 8.4e-03]:   1%|          | 10/1000 [00:00<00:03, 274.82it/s]
Running TRAM estimate - [inc: 5.3e-03]:   1%|          | 11/1000 [00:00<00:03, 276.24it/s]
Running TRAM estimate - [inc: 3.3e-03]:   1%|          | 12/1000 [00:00<00:03, 276.43it/s]
Running TRAM estimate - [inc: 2.3e-03]:   1%|▏         | 13/1000 [00:00<00:03, 277.63it/s]
Running TRAM estimate - [inc: 1.7e-03]:   1%|▏         | 14/1000 [00:00<00:03, 278.72it/s]
Running TRAM estimate - [inc: 1.4e-03]:   2%|▏         | 15/1000 [00:00<00:03, 279.29it/s]
Running TRAM estimate - [inc: 1.2e-03]:   2%|▏         | 16/1000 [00:00<00:03, 279.97it/s]
Running TRAM estimate - [inc: 1.1e-03]:   2%|▏         | 17/1000 [00:00<00:03, 280.61it/s]
Running TRAM estimate - [inc: 1.0e-03]:   2%|▏         | 18/1000 [00:00<00:03, 281.23it/s]
Running TRAM estimate - [inc: 9.6e-04]:   2%|▏         | 19/1000 [00:00<00:03, 281.61it/s]
Running TRAM estimate - [inc: 9.6e-04]: 100%|██████████| 19/19 [00:00<00:00, 272.22it/s]


12 import numpy as np
13 import matplotlib.pyplot as plt
14
15 from deeptime.data import tmatrix_metropolis1d
16 from deeptime.markov.msm import MarkovStateModel, TRAM
17 from deeptime.clustering import ClusterModel
18
19 xs = np.linspace(-1.5, 1.5, num=100)
20 n_samples = 10000
21 bias_centers = [-1, -0.5, 0.0, 0.5, 1]
22
23
24 def harmonic(x0, x):
25     return 2 * (x - x0) ** 4
26
27
28 def plot_contour_with_colourbar(data, vmin=None, vmax=None):
29     if vmin is None:
30         vmin = np.min(data)
31     if vmax is None:
32         vmax = np.max(data)
33     fig, (ax1) = plt.subplots(1, figsize=(3.5, 3))
34     im = ax1.contourf(data, vmin=vmin, vmax=vmax, levels=50, cmap='jet')
35     plt.colorbar(im)
36     plt.show()
37
38
39 def get_bias_functions():
40     bias_functions = []
41     for i, bias_center in enumerate(bias_centers):
42         bias = lambda x, x0=bias_center: harmonic(x0, x)
43         bias_functions.append(bias)
44     return bias_functions
45
46
47 def sample_trajectories(bias_functions):
48     trajs = np.zeros((len(bias_centers), n_samples), dtype=np.int32)
49
50     for i, bias in enumerate(bias_functions):
51         biased_energies = (xs - 1) ** 4 * (xs + 1) ** 4 - 0.1 * xs + bias(xs)
52
53         biased_energies /= np.max(biased_energies)
54         transition_matrix = tmatrix_metropolis1d(biased_energies)
55
56         msm = MarkovStateModel(transition_matrix)
57         trajs[i] = msm.simulate(n_steps=n_samples)
58     return trajs
59
60
61 if __name__ == "__main__":
62     bias_functions = get_bias_functions()
63     trajectories = sample_trajectories(bias_functions)
64
65     # move from trajectory over 100 bins back to the space of the xs: (-1.5, 1.5)
66     trajectories = trajectories / 100 * 3 - 1.5
67
68     bias_matrices = np.zeros((len(bias_centers), n_samples, len(bias_centers)))
69     for i, traj in enumerate(trajectories):
70         for j, bias_function in enumerate(bias_functions):
71             bias_matrices[i, :, j] = bias_function(traj)
72
73     # discretize the trajectories into two Markov states (centered around the two wells)
74     clustering = ClusterModel(cluster_centers=np.asarray([-0.75, 0.75]), metric='euclidean')
75
76     dtrajs = clustering.transform(trajectories.flatten()).reshape((len(bias_matrices), n_samples))
77
78     from tqdm import tqdm
79     tram = TRAM(lagtime=1, maxiter=1000, maxerr=1e-3, progress=tqdm, init_strategy="MBAR")
80
81     # For every simulation frame seen in trajectory i and time step t, btrajs[i][t,k] is the
82     # bias energy of that frame evaluated in the k'th thermodynamic state (i.e. at the k'th
83     # Umbrella/Hamiltonian/temperature).
84     model = tram.fit_fetch((dtrajs, bias_matrices))
85
86     plot_contour_with_colourbar(model.biased_conf_energies)

Total running time of the script: (0 minutes 0.275 seconds)

Estimated memory usage: 580 MB