def draw_energy_autocorr(dataset, axis=None):
if axis is None:
fig, new_axis = plt.subplots()
else:
new_axis = axis
tsteps = dataset["tsteps"]
en = dataset["energy"].values
en_k = np.fft.fft(en, axis=-1)
en_corr = np.fft.ifft((en_k * en_k.conj()), axis=-1).real
corr_ds = xr.DataArray(en_corr,
coords=dataset["energy"].coords,
attrs=dataset["energy"].attrs)
# TODO: StackOverflow - ? easier way to cast fft over dataset.
mean, stdev = gu.mean_std(corr_ds, dim=['run'])
# TODO: add fitting+plotting code to determine+display correlation time
if len(np.shape(mean)) != 1:
print("WARNING: draw_energy encountered unexpected dimensions"
" after averaging over runs.")
# new_axis.plot(np.log10(tsteps.values), np.log10(mean.values))
new_axis.errorbar(
tsteps.values,
mean.values, # yerr=stdev.values,
capsize=2.)
if axis is None:
return (fig, new_axis)
else:
return axis
def draw_bend_autocorr_rw(dataset, dims=2):
count = dataset["tsteps"].size
tmp = 100
display_counts = []
while tmp <= count:
display_counts.append(tmp)
tmp = int(tmp * np.sqrt(10000))
x, y = naive_curve(dataset, dims)
fig, axes = plt.subplots(ncols=len(display_counts),
sharey='row',
squeeze=False)
for ax, count in zip(axes[0], display_counts):
tmp = (dataset["bend_autocorr"].isel(run=0,
kickSize=0,
tsteps=slice(count)))
tmp_mean, tmp_std = gu.mean_std(tmp, dim='tsteps')
ax.errorbar(x,
tmp_mean.values,
yerr=tmp_std.values,
capsize=2.0,
label="RW",
color="red")
ax.plot(x, y, label="Naive", color="green")
ax.set_title("avg over {0} RWs".format(count))
ax.legend(loc="lower right")
fig.suptitle(
"#rods={0}, Correlation function averaged over random walks.".format(
dataset.attrs["rodCount"]))
plt.show(block=False)
return (fig, axes)
def draw_angle_profile(dataset, axis=None, total_only=False, show=None):
def tsteps_slice(ndraw):
nonlocal dataset
n = len(dataset["tsteps"])
if ndraw > n:
ndraw = n
print("WARNING: attempting to draw at more points than sampled.")
return dataset["tsteps"][::n // ndraw]
if axis is None:
fig, new_axis = plt.subplots()
else:
new_axis = axis
if show is None:
ndraw = 5
tsteps = tsteps_slice(ndraw)
elif isinstance(show, int):
ndraw = show
tsteps = tsteps_slice(ndraw)
elif show == "all":
tsteps = dataset["tsteps"]
else:
raise ValueError(
"show should be one of None, int (> 0) or the string 'all'.")
mean, stdev = gu.mean_std(dataset["angles"].sel(tsteps=tsteps) /
(2 * np.pi),
dim=['run'])
if len(np.shape(
mean)) != 3: # one for time, one for rods, one for 3 angles
print("WARNING: draw_angle_profile encountered unexpected dimensions"
" after averaging over runs.")
x = np.arange(dataset.attrs["rodCount"])
# if not total_only:
# for (i, a) in enumerate(ANGLES_STR):
# new_axis.errorbar(x, mean.isel(angle_str=i),
# yerr=stdev.isel(angle_str=i), label=a)
# TODO: replace phi + psi (as a proxy) with actual twist value
# ? How to compute total twist from phi, theta, psi values
for tstep in tsteps:
tmp_mean, tmp_std = mean.sel(tsteps=tstep), stdev.sel(tsteps=tstep)
new_axis.plot(
x,
tmp_mean.isel(angle_str=0) + tmp_mean.isel(angle_str=2),
# yerr=(tmp_std.isel(angle_str=0)**2 + tmp_std.isel(angle_str=2)**2)**0.5,
label=("t = " + str(tstep.values)))
new_axis.legend(loc="upper left")
new_axis.set_ylabel("Angle/2π radians")
new_axis.set_xlabel("Rod number")
if axis is None:
return (fig, new_axis)
else:
return axis
def draw_energy(dataset, axis=None, show=None):
flag = axis is None
if flag:
fig, axis = plt.subplots()
# else:
# new_axis = axis
tsteps = dataset["tsteps"]
mean, stdev = gu.mean_std(dataset["energy"].sel(tsteps=tsteps),
dim=['run'])
if len(np.shape(mean)) != 1:
print("WARNING: draw_energy encountered unexpected dimensions"
" after averaging over runs.")
axis.errorbar(tsteps.values, mean.values, yerr=stdev.values, capsize=2.)
if flag:
return (fig, axis)
return axis
def draw_acceptance(dataset, axis=None):
if axis is None:
fig, new_axis = plt.subplots()
else:
new_axis = axis
for (i, a) in enumerate(ANGLES_STR):
mean, stdev = gu.mean_std(dataset["acceptance"].isel(angle_str=i),
dim=["run"])
new_axis.errorbar(dataset["tsteps"].values,
mean.values,
yerr=stdev.values,
label=a)
new_axis.legend(loc="upper right")
new_axis.set_xlabel("Time (Monte Carlo steps)")
new_axis.set_ylabel("Acceptance probability")
if axis is None:
return (fig, new_axis)
else:
return axis
def hat_curve(dataset):
"""Draws a curve for z_bead vs turns."""
twists = dataset["twists"]
tsteps = dataset["tsteps"]
tslice = slice(0, tsteps.size, tsteps.size // twists.size)
z_mean, z_stdev = gu.mean_std(
dataset["extension"].isel(tsteps=tslice).sel(axis='z'), dim='run')
def f(x):
tmp = x.rename({"tsteps": "twists"})
tmp["twists"] = twists
return tmp
z_mean = f(z_mean).to_dataset().rename({"extension": "mean"})
z_stdev = f(z_stdev).to_dataset().rename({"extension": "stdev"})
z = xr.merge([z_mean, z_stdev])
fig, axis = plt.subplots()
axis.errorbar(z["twists"].values / (2 * np.pi),
z["mean"].values,
yerr=z["stdev"].values,
label='z')
plt.legend(loc='upper left')
plt.show(block=False)
def draw_bend_autocorr(dataset, energy=False, rw_dataset=None, dims=2):
fig, axes = plt.subplots(nrows=2 if energy else 1,
ncols=min(5, len(dataset["run"])),
sharey='row',
squeeze=False)
x, y = naive_curve(dataset, dims)
if rw_dataset is not None:
total_num_rw = rw_dataset["tsteps"].size
x2 = rw_dataset.attrs["rodLength"] * np.arange(
rw_dataset.attrs["rodCount"])
y2 = (rw_dataset["bend_autocorr"].isel(run=0,
kickSize=0).mean(dim='tsteps'))
test_num_rw = 25
def partial_rw_y(start):
return (rw_dataset["bend_autocorr"].isel(
run=0, kickSize=0,
tsteps=slice(start, start + test_num_rw)).mean(dim='tsteps'))
ys = [partial_rw_y(yi) for yi in [0, 50, 100, 250, 500, 750]]
start = max(500, int(dataset["tsteps"][0]))
step = int(dataset["tsteps"][1] - dataset["tsteps"][0])
stop = int(dataset["tsteps"][-1] + step)
for (i, ax) in enumerate(axes[0]):
for ks in dataset["kickSize"]:
tmp = (dataset["bend_autocorr"].sel(run=i,
kickSize=ks,
tsteps=np.arange(
start, stop, step)))
tmp_mean, tmp_std = gu.mean_std(tmp, dim='tsteps')
ax.errorbar(
x,
tmp_mean.values, # yerr=tmp_std.values,
capsize=2.0,
label="MC, ks={0}".format(ks.values),
color='blue')
# ax.set_yscale('log')
if rw_dataset is not None:
ax.plot(x2, y2, color='red', label="{0} RW".format(total_num_rw))
for yi in ys:
ax.plot(x2,
yi,
color='grey',
label="{0} RW".format(test_num_rw))
ax.plot(x, y, color='green', label="Naive")
else:
ax.plot(x, y, color='green', label="Naive")
ax.set_ylim(0, 1)
ax.set_title("Run# = {0}".format(i))
handles, labels = axes[0][-1].get_legend_handles_labels()
by_label = OrderedDict(zip(labels, handles))
axes[0][-1].legend(by_label.values(), by_label.keys(), loc="upper right")
axes[0][0].set_ylabel("(Tangent vector autocorrelation)")
axes[0][len(axes[0]) // 2].set_xlabel("Length (nm)")
if energy:
for (i, ax) in enumerate(axes[1]):
for ks in dataset["kickSize"]:
draw_energy_autocorr(
dataset.sel(kickSize=ks).isel(run=slice(i, i + 1)),
axis=ax)
axes[1][0].set_ylabel("Energy autocorrelation")
axes[1][len(axes[1]) // 2].set_xlabel("Time")
fig.suptitle(
"#rods={1}, Correlation function averaged over t={2} to t={0} in steps of {3}"
.format(int(dataset["tsteps"][-1]), dataset.attrs["rodCount"], start,
step))
plt.show(block=False)
return (fig, axes)
def draw_force_extension(dataset, acceptance=True):
kickSizes = dataset.coords["kickSize"].values
forces = dataset.coords["force"].values
runs = dataset.coords["run"].values
B = dataset.attrs["B"]
L = dataset.attrs["rodCount"]
T = dataset.attrs["temperature"]
Pinv = dataset.data_vars["Pinv"].values
pre_steps = dataset.attrs["pre_steps"]
extra_steps = dataset.attrs["extra_steps"]
fig, axes = plt.subplots(nrows=(2 if acceptance else 1),
ncols=kickSizes.size,
sharex="row",
sharey="row",
squeeze=False)
sns.set_style("darkgrid")
ms_curve_x, ms_curve_y = marko_siggia_curve(B, 740)
for (j, ks) in enumerate(kickSizes):
# 'axis' dimension is the last dimension
# there doesn't seem to be a simple way to broadcast np.linalg.norm
tmp = gu.norm(dataset["extension"].sel(kickSize=ks), dim='axis')
mean, stdev = gu.mean_std(tmp.groupby("force"))
print(mean.values)
print(forces)
axes[0, j].errorbar(mean.values,
forces,
xerr=stdev.values,
capsize=4.0,
linestyle='')
axes[0, j].plot(ms_curve_x, ms_curve_y)
axes[0, j].set_xlim(left=500, right=740)
axes[0, j].set_ylim(bottom=-0.5, top=10 + 0.5)
axes[0, j].set_title("kick size = {0}".format(ks))
axes[0, j].set_xlabel("")
axes[0, j].set_ylabel("")
axes[0, 0].set_ylabel("Force (pN)")
axes[0, kickSizes.size // 2].set_xlabel("Extension (nm)")
# Molecular dynamics simulation will not have acceptance values.
acceptance = acceptance and ("acceptance" in dataset)
if acceptance:
for (j, ks) in enumerate(kickSizes):
tmp = dataset["acceptance"].sel(kickSize=ks)
for (j_a, angle) in enumerate(ANGLES_STR):
mean, stdev = gu.mean_std(
tmp.isel(angle_str=j_a).groupby("force"))
axes[1, j].errorbar(forces,
mean.values,
yerr=stdev.values,
label=angle)
axes[1, j].set_ylim(bottom=0, top=1)
axes[1, 0].set_ylabel("Acceptance ratio")
axes[1, kickSizes.size // 2].set_xlabel("Force (pN)")
fig.suptitle(
("Contour length (straight) L_0 = 740 nm, #rods = {6}, T = {7:.0f} K\n"
"Effective persistence length Lp = {5:.1f} nm, "
"Intrinsic disorder persistence length P = {0:.1f} nm\n"
"Thermalization steps = {1}, extra steps = {2}, "
"#samples in extra steps = {3}, runs = {4}").format(
1 / Pinv if Pinv != 0 else np.inf, pre_steps, extra_steps,
dataset.coords["tsteps"].size, runs, B, L, T),
fontdict={"fontsize": 10})
plt.legend(loc="upper right")
plt.show(block=False)