Skip to content

Commit

Permalink
add notebook stuff
Browse files Browse the repository at this point in the history
  • Loading branch information
@svandenhaute
svandenhaute committed Nov 8, 2023
1 parent 5d88f0a commit 2cc1eaa
Show file tree
Hide file tree
Showing 3 changed files with 369 additions and 0 deletions.
128 changes: 128 additions & 0 deletions notebook/notebook_utils.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
import io
from pathlib import Path

import ipywidgets as widgets
import numpy as np
import py3Dmol
import pymbar
from ase import Atoms
from ase.io import write
from pymbar import timeseries


def visualize_trajectory(trajectory: list[Atoms]):
# Create a py3Dmol view object
block = py3Dmol.view(width=800, height=400)
block.setBackgroundColor("lightgray")
block.zoom(1.3)

def show_structure(frame):
atoms = trajectory[frame].copy()
cv = atoms.info["CV"]

# Convert the Atoms object to an XYZ formatted string
xyz_io = io.StringIO()
write(xyz_io, atoms, format="xyz")
xyz_str = xyz_io.getvalue()

block.removeAllModels()
block.addModel(xyz_str, format="xyz")
block.setStyle(
{},
{
"stick": {"colorscheme": "gray", "radius": 0.1},
"sphere": {"radius": 0.3},
},
)
block.show()
energy_label.value = f"CV value: {cv:.2f}"

# Create a slider for navigating the trajectory
# slider = widgets.IntSlider(min=0, max=len(trajectory) - 1, description="Frame")

# Create a label to display the energy
cv = trajectory[0].info["CV"]
energy_label = widgets.Label(value=f"CV value: {cv:.2f}")

# display(
# widgets.VBox(
# [
# slider,
# energy_label,
# widgets.interactive_output(show_structure, {"frame": slider}),
# ]
# )
# )


def generate_fes(
path_umbrellas: Path,
centers: np.ndarray,
kappa: float,
temperature: float,
nbins: int,
) -> tuple[np.ndarray, np.ndarray]:
kB = 1.381e-23 * 6.022e23 / 1000.0 # in kJ/mol/K

trajectories = []
i = 0
while True:
path = path_umbrellas / "{}.txt".format(i)
if path.exists():
trajectories.append(np.loadtxt(path))
i += 1
else:
break
assert i == len(centers)
N_k = [0] * len(centers)

# decorrelate
for i, trajectory in enumerate(trajectories):
indices = timeseries.subsample_correlated_data(trajectory)
N_k[i] = len(indices)
trajectories[i] = trajectory[indices]

# compute bin centers
bin_center_i = np.zeros([nbins])
cv_min = min([np.min(t) for t in trajectories])
cv_max = max([np.max(t) for t in trajectories])
bin_edges = np.linspace(cv_min, cv_max, nbins + 1)
for i in range(nbins):
bin_center_i[i] = 0.5 * (bin_edges[i] + bin_edges[i + 1])

# pad and put in array
N = max(N_k)
c_kn = np.zeros((len(centers), N))
for i, trajectory in enumerate(trajectories):
c_kn[i, : N_k[i]] = trajectory[:]

# bias energy from snapshot n in umbrella k evaluated at umbrella l
beta = 1 / (kB * temperature)
u_kln = np.zeros((len(centers), len(centers), c_kn.shape[1]))
for k in range(len(centers)):
for n in range(N_k[k]):
for l in range(len(centers)): # noqa: E741
u_kln[k, l, n] = beta * kappa / 2 * (c_kn[k, n] - centers[l]) ** 2

c_n = pymbar.utils.kn_to_n(c_kn, N_k=N_k)
fes = pymbar.FES(u_kln, N_k, verbose=True)
u_kn = np.zeros(c_kn.shape)

# compute free energy profile in unbiased potential (in units of kT).
histogram_parameters = {}
histogram_parameters["bin_edges"] = bin_edges
fes.generate_fes(
u_kn, c_n, fes_type="histogram", histogram_parameters=histogram_parameters
)
results = fes.get_fes(
bin_center_i, reference_point="from-lowest", uncertainty_method="analytical"
)
center_f_i = results["f_i"] * kB * temperature
center_df_i = results["df_i"] * kB * temperature

# Write out free energy profile
print("free energy profile [kJ/mol]")
print(f"{'bin':>8s} {'f':>8s} {'df':>8s}")
for i in range(nbins):
print(f"{bin_center_i[i]:8.1f} {center_f_i[i]:8.3f} {center_df_i[i]:8.3f}")
return bin_center_i, center_f_i
240 changes: 240 additions & 0 deletions notebook/zerokelvin.xyz
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,240 @@
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-1.487061 pbc="F F F"
O 1.04654576 -1.55962797 1.01544860
C 0.14921681 -2.30925077 0.71774831
O -0.85162780 -2.04467245 -0.08370986
H 0.08320651 -3.33257292 1.10876381
H -0.77498652 -1.13207313 -0.42019152
O -0.58268825 0.61409530 -1.03260613
C 0.31461602 1.36373408 -0.73487046
O 1.31543668 1.09916662 0.06662335
H 0.38061481 2.38706612 -1.12587079
H 1.23881398 0.18655913 0.40308220
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-1.336648 pbc="F F F"
O 1.06075706 -1.50467246 1.00187267
C 0.15812102 -2.25216241 0.69953770
O -0.83640174 -1.98373046 -0.09928500
H 0.09583813 -3.27462222 1.09272960
H -0.75871546 -1.05915424 -0.44011153
O -0.59694257 0.55916548 -1.01897060
C 0.30571059 1.30664402 -0.71665908
O 1.30023838 1.03821219 0.08215775
H 0.36800021 2.32909835 -1.10986054
H 1.22254238 0.11364574 0.42300653
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-1.183675 pbc="F F F"
O 1.07203526 -1.46542371 0.99288070
C 0.16459824 -2.21153554 0.68666828
O -0.82420775 -1.94005147 -0.10951986
H 0.10655214 -3.23307938 1.08239353
H -0.74613709 -1.00252088 -0.45561389
O -0.60818454 0.51990820 -1.00997904
C 0.29924675 1.26603012 -0.70377582
O 1.28804933 0.99453546 0.09241142
H 0.35725375 2.28756643 -1.09952904
H 1.20994191 0.05699478 0.43848123
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-1.029174 pbc="F F F"
O 1.08124020 -1.43746185 0.98726139
C 0.16827873 -2.18099814 0.67617082
O -0.81415866 -1.90485743 -0.11757911
H 0.11336653 -3.20184754 1.07374254
H -0.73533503 -0.95256187 -0.46947570
O -0.61740025 0.49193973 -1.00439235
C 0.29555327 1.23548115 -0.69329101
O 1.27797556 0.95935071 0.10048140
H 0.35047052 2.25632647 -1.09087233
H 1.19915711 0.00705277 0.45237185
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-0.873563 pbc="F F F"
O 1.09163096 -1.41473594 0.98451987
C 0.17346143 -2.15676842 0.66934480
O -0.80311002 -1.87745427 -0.12151467
H 0.12320269 -3.17665898 1.06966591
H -0.72555674 -0.90832079 -0.48151448
O -0.62772163 0.46917276 -1.00175377
C 0.29038434 1.21124451 -0.68648305
O 1.26685000 0.93199236 0.10452771
H 0.34067317 2.23111910 -1.08684145
H 1.18933379 -0.03716633 0.46446663
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-0.717027 pbc="F F F"
O 1.10503068 -1.39377750 0.98506094
C 0.18275253 -2.13687924 0.66726645
O -0.78889051 -1.85778174 -0.12028042
H 0.14026803 -3.15520838 1.07215818
H -0.71644164 -0.86926886 -0.49179405
O -0.64141512 0.44839610 -1.00186395
C 0.28107880 1.19136308 -0.68438546
O 1.25285285 0.91218288 0.10296708
H 0.32366707 2.20963032 -1.08942465
H 1.18024531 -0.07623266 0.47471338
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-0.559489 pbc="F F F"
O 1.12055414 -1.37612952 0.98770109
C 0.19382958 -2.12101651 0.66821281
O -0.77332964 -1.84314153 -0.11529707
H 0.16054646 -3.13731715 1.07879224
H -0.70849147 -0.83236767 -0.50183310
O -0.65668588 0.43059445 -1.00487178
C 0.26999608 1.17550815 -0.68532146
O 1.23714488 0.89763096 0.09820022
H 0.30325214 2.19183063 -1.09585597
H 1.17233170 -0.11316781 0.48469053
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-0.400803 pbc="F F F"
O 1.13522340 -1.36444017 0.99243225
C 0.20250837 -2.10942748 0.66985865
O -0.75912360 -1.83034649 -0.10951726
H 0.17736556 -3.12386505 1.08542476
H -0.70252742 -0.79386845 -0.51303507
O -0.67144221 0.41895704 -1.00948533
C 0.26134882 1.16389644 -0.68701844
O 1.22288509 0.88486960 0.09249483
H 0.28666449 2.17823256 -1.10282189
H 1.16624550 -0.15158402 0.49608502
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-0.241012 pbc="F F F"
O 1.15189452 -1.35463756 0.99974486
C 0.21355745 -2.10096092 0.67394092
O -0.74226062 -1.82174681 -0.10165206
H 0.19797621 -3.11314371 1.09532201
H -0.69773082 -0.75687358 -0.52562207
O -0.68812523 0.40916286 -1.01677347
C 0.25027670 1.15544454 -0.69106143
O 1.20612798 0.87620723 0.08448200
H 0.26589164 2.16760700 -1.11249010
H 1.16154017 -0.18863505 0.50852684
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=-0.080431 pbc="F F F"
O 1.16344998 -1.35499063 1.00776999
C 0.21637993 -2.09751519 0.67441660
O -0.73064666 -1.81089699 -0.09826717
H 0.20533579 -3.10861421 1.09848099
H -0.69525649 -0.71423524 -0.54213866
O -0.69963207 0.40948301 -1.02488662
C 0.24744339 1.15200130 -0.69153192
O 1.19446842 0.86538078 0.08115222
H 0.25851582 2.16310102 -1.11559127
H 1.15908987 -0.23128983 0.52501333
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=0.080431 pbc="F F F"
O 1.17516575 -1.35795168 1.01706225
C 0.21922737 -2.09684331 0.67607665
O -0.71883759 -1.80235536 -0.09380606
H 0.21271882 -3.10685899 1.10280883
H -0.69488332 -0.67189645 -0.55995183
O -0.71132991 0.41242842 -1.03418133
C 0.24460562 1.15132283 -0.69319720
O 1.18267427 0.85684110 0.07668485
H 0.25109069 2.16133710 -1.11992882
H 1.15871629 -0.27359965 0.54285017
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=0.241012 pbc="F F F"
O 1.18610457 -1.36398410 1.02714552
C 0.22121604 -2.09914941 0.67841176
O -0.70778750 -1.79611385 -0.08888495
H 0.21877263 -3.10820789 1.10749788
H -0.69657349 -0.62971975 -0.57915309
O -0.72227820 0.41847055 -1.04425715
C 0.24261453 1.15363267 -0.69552847
O 1.17161825 0.85059781 0.07176850
H 0.24505969 2.16268785 -1.12462049
H 1.16040148 -0.31578988 0.56203798
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=0.400806 pbc="F F F"
O 1.19634668 -1.37280836 1.03796968
C 0.22284858 -2.10462820 0.68188801
O -0.69746070 -1.79326004 -0.08303106
H 0.22424433 -3.11276332 1.11324832
H -0.70006217 -0.58830742 -0.59921164
O -0.73252328 0.42729949 -1.05507894
C 0.24098377 1.15911082 -0.69900727
O 1.16130080 0.84773738 0.06589774
H 0.23958333 2.16724899 -1.13036123
H 1.16388667 -0.35720533 0.58210389
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=0.559489 pbc="F F F"
O 1.20704562 -1.38319533 1.04977928
C 0.22601458 -2.11328276 0.68777175
O -0.68685223 -1.79554974 -0.07476869
H 0.23223049 -3.12026162 1.12194728
H -0.70508982 -0.54904530 -0.61945519
O -0.74321732 0.43768094 -1.06689388
C 0.23781481 1.16776768 -0.70488795
O 1.15068457 0.85003284 0.05764819
H 0.23159684 2.17474784 -1.13906048
H 1.16892046 -0.39647054 0.60233720
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=0.717024 pbc="F F F"
O 1.21674406 -1.39765668 1.06261611
C 0.22858100 -2.12607683 0.69499336
O -0.67741297 -1.80236176 -0.06538379
H 0.23911375 -3.13202070 1.13168997
H -0.71122661 -0.50996512 -0.64036564
O -0.75291501 0.45214240 -1.07973269
C 0.23524963 1.18056222 -0.71210830
O 1.14124382 0.85684581 0.04826651
H 0.22471769 2.18650580 -1.14880546
H 1.17505265 -0.43555114 0.62324743
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=0.873561 pbc="F F F"
O 1.22504234 -1.41762251 1.07680628
C 0.22974490 -2.14366702 0.70343505
O -0.66963547 -1.81358571 -0.05494933
H 0.24350320 -3.14886735 1.14200138
H -0.71815414 -0.46955791 -0.66236178
O -0.76110396 0.47204162 -1.09408415
C 0.23408653 1.19814995 -0.72055265
O 1.13328687 0.86817769 0.03809245
H 0.22042338 2.20329772 -1.15924331
H 1.18195433 -0.47594249 0.64527355
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=1.029177 pbc="F F F"
O 1.23509213 -1.44033839 1.09351937
C 0.23386609 -2.16580082 0.71572419
O -0.66008613 -1.83175418 -0.04083697
H 0.25203782 -3.16988373 1.15697234
H -0.72646501 -0.42926379 -0.68524195
O -0.77129238 0.49484341 -1.11059363
C 0.22996573 1.22028361 -0.73284259
O 1.12395679 0.88621042 0.02366065
H 0.21178336 2.22437548 -1.17406927
H 1.19028960 -0.51624801 0.66812535
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=1.183675 pbc="F F F"
O 1.24491427 -1.47007870 1.11293423
C 0.23777306 -2.19458753 0.73069497
O -0.65096848 -1.85633097 -0.02403017
H 0.25983715 -3.19774695 1.17421415
H -0.73577646 -0.38476987 -0.71062743
O -0.78109008 0.52456334 -1.13004383
C 0.22605875 1.24907002 -0.74781652
O 1.11480927 0.91083057 0.00690810
H 0.20398185 2.25222950 -1.19132573
H 1.19960868 -0.56075541 0.69350974
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=1.336649 pbc="F F F"
O 1.25583997 -1.50822689 1.13684701
C 0.24318120 -2.23250258 0.75058738
O -0.64074473 -1.89093694 -0.00234472
H 0.26971857 -3.23459863 1.19668304
H -0.74645907 -0.33368813 -0.73993147
O -0.79198564 0.56269984 -1.15400176
C 0.22064980 1.28698666 -0.76770179
O 1.10453546 0.94544002 -0.01471356
H 0.19413096 2.28907297 -1.21382065
H 1.21028148 -0.61182232 0.72281401
10
Properties=species:S:1:pos:R:3 reference_stdout=F reference_stderr=F reference_status=F CV=1.487062 pbc="F F F"
O 1.26955417 -1.56207770 1.16928613
C 0.25089024 -2.28577900 0.77862924
O -0.62779090 -1.94010017 0.02750839
H 0.28156192 -3.28693226 1.22717277
H -0.76045697 -0.26814575 -0.77753751
O -0.80572297 0.61656140 -1.18640469
C 0.21293911 1.34026388 -0.79574544
O 1.09161273 0.99459014 -0.04461365
H 0.18227442 2.34141265 -1.24429997
H 1.22428624 -0.67736919 0.76042223
1 change: 1 addition & 0 deletions tests/test_learning.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -177,3 +177,4 @@ def test_temperature_ramp(context):
assert not T == 500
T = apply_temperature_ramp(100, 500, 5, T)
assert T == 500
print(apply_temperature_ramp(50, 2000, 5))

0 comments on commit 2cc1eaa

Please sign in to comment.