Add sheet11.py

Jaslo/sheet11.py

+import itertools
+from typing import Callable
+
+import numpy as np
+from numpy.typing import ArrayLike
+import matplotlib.pyplot as plt
+
+from tqdm import tqdm
+
+# Problem 11.1
+# a)
+
+def langevin(force: Callable,
+             r0: ArrayLike,
+             p0: ArrayLike,
+             m: float,
+             dt: float,
+             N: int,
+             T: float,
+             gamma: float,
+             L: float) -> (ArrayLike, ArrayLike, ArrayLike):
+    assert(r0.shape == p0.shape)
+
+    r = np.zeros((N, *r0.shape))
+    p = np.zeros((N, *p0.shape))
+    U = np.zeros(N)
+
+    r[0] = r0
+    p[0] = p0
+
+    f, U[0] = force(r[0], L)
+
+    for i in tqdm(range(1, N)):
+        # B
+        p[i] = p[i-1] + 1/2 * f * dt
+
+        # A
+        r[i] = r[i-1] + 1/2 * p[i] / m * dt
+        # enforce periodic boundaries
+        r[i] -= (r[i] > L / 2) * L
+        r[i] += (r[i] < -L / 2) * L
+
+        # O
+        p[i] = p[i] * np.exp(-gamma * dt) + np.sqrt(m * T * (1 - np.exp(-2 * gamma * dt))) * np.random.normal(size=p0.shape)
+
+        # A
+        r[i] = r[i] + 1/2 * p[i] / m * dt
+        # enforce periodic boundaries
+        r[i] -= (r[i] > L / 2) * L
+        r[i] += (r[i] < -L / 2) * L
+
+        # update force
+        f, U[i] = force(r[i], L)
+
+        # B
+        p[i] = p[i] + 1/2 * f * dt
+
+    return r, p, U
+
+N = 64
+d = 3
+sig = 1
+eps = 1
+L = 6 * sig
+m = 1
+tau = np.sqrt(m * sig**2 / eps)
+a = 1.5 * sig
+k_BT = 1.5 * eps
+dt = 0.002 * tau
+
+gamma = 1 # TODO: ??
+
+# create lattice
+n = int(np.round(N**(1/3)))
+x = np.linspace(-(n-1) / 2, (n-1) / 2, n)
+r0 = a * np.array(list(itertools.product(x, repeat=3)))
+
+p0 = np.random.normal(0, np.sqrt(m* k_BT), (N, d))
+
+t_eq = 20 * tau
+t_max = 100 * tau
+t = np.arange(0, t_max, dt)
+
+n_eq = int(t_eq / dt)
+
+def lj(r: ArrayLike, L: float) -> (ArrayLike, float):
+    N = len(r)
+    f = np.zeros(r.shape)
+    U = 0
+
+    sig6 = sig**6
+
+    for i in range(1, N):
+        dr = r[i:] - r[:N-i]
+        # periodic boundary conditions
+        dr -= (dr > L / 2) * L
+        dr += (dr < -L / 2) * L
+
+        r2 = 1 / np.sum(dr**2, axis=1) # scalar product per element
+        r6 = r2 * r2 * r2
+
+        U += 4 * eps * sig6 * np.sum(r6 * (sig6 * r6 - 1))
+        fval = np.transpose(-24 * eps * sig6 * r6 * r2 * (2 * sig6 * r6 - 1) * dr.T) # cursed scalar multiplication
+
+        f[:N-i] += fval
+        f[i:] -= fval
+
+    return f, U
+
+r, p, U = langevin(lj, r0, p0, m, dt, len(t), k_BT, gamma, L)
+
+# plt.plot(t, U / N)
+# plt.show()
+
+rp = r[n_eq:]
+
+msd = np.zeros(len(rp))
+for i in tqdm(range(1, len(msd))):
+    msd[i] = np.mean(np.sum((rp[i:] - rp[:-i])**2, axis=2))
+
+plt.loglog(t[n_eq:], msd)
+plt.show()
+
+# b)
+def verlet(force: Callable,
+           r0: ArrayLike,
+           p0: ArrayLike,
+           m: float,
+           dt: float,
+           N: int,
+           L: float) -> (ArrayLike, ArrayLike, ArrayLike):
+    assert(r0.shape == p0.shape)
+
+    r = np.zeros((N, *r0.shape))
+    p = np.zeros((N, *p0.shape))
+    U = np.zeros(N)
+
+    r[0] = r0
+    p[0] = p0
+
+    f, U[0] = force(r[0], L)
+
+    for i in tqdm(range(1, N)):
+        p[i] = p[i-1] + 1/2 * f * dt
+
+        r[i] = r[i-1] + p[i] / m * dt
+        # enforce periodic boundaries
+        r[i] -= (r[i] > L / 2) * L
+        r[i] += (r[i] < -L / 2) * L
+
+        # update force
+        f, U[i] = force(r[i], L)
+
+        # B
+        p[i] = p[i] + 1/2 * f * dt
+
+    return r, p, U
+
+rv, pv, Uv = verlet(lj, r[n_eq], p[n_eq], m, dt, len(t) - n_eq, L)
+
+Kv = np.sum(pv**2, axis=(1,2)) / (2 * m)
+
+print(f"k_B T = {2/3 * np.mean(Kv) / N}")
+
+msdv = np.zeros(len(rv))
+for i in tqdm(range(1, len(msdv))):
+    msdv[i] = np.mean(np.sum((rv[i:] - rv[:-i])**2, axis=2))
+
+plt.loglog(t[n_eq:], msd / t[n_eq:])
+plt.loglog(t[n_eq:], msdv / t[n_eq:])
+plt.show()