diff --git a/examples/selector.py b/examples/selector.py index 2bf8903..4e48c7f 100644 --- a/examples/selector.py +++ b/examples/selector.py @@ -1,5 +1,6 @@ from functools import partial +import numpy as np import matplotlib.pyplot as plt from matplotlib import cm @@ -41,3 +42,36 @@ plt.xlabel(r"$t$ / ps") plt.ylabel(r"S_q(t)") plt.savefig(f"{data_dir}/selector.png", dpi=300, bbox_inches="tight") plt.show() + + +def multi_radial_selector(atoms, bins): + indices = [] + for i in range(len(bins) - 1): + index = md.coordinates.selector_radial_cylindrical( + atoms, r_min=bins[i], r_max=bins[i + 1] + ) + indices.append(index) + return indices + + +bins = np.arange(0.0, 1.6, 0.1) +r = (bins[:-1] + bins[1:]) / 2 + +time, results = md.correlation.shifted_correlation( + partial(md.correlation.isf, q=22.7), + oxygen_water, + selector=partial(multi_radial_selector, bins=bins), + segments=100, + skip=0.1, +) + +c = [cm.plasma(i) for i in np.linspace(0, 1, len(r))] +plt.figure() +for i, result in enumerate(results): + plt.plot(time, result, "-", c=c[i], label=round(r[i], 2)) +plt.legend(title=r"$r$ / nm", ncols=2) +plt.xscale("log") +plt.xlabel(r"$t$ / ps") +plt.ylabel(r"S_q(t)") +plt.savefig(f"{data_dir}/multi_selector.png", dpi=300, bbox_inches="tight") +plt.show()