New vectorized walk function

Author: ashwinvis

.gitignore

@@ -3,7 +3,9 @@
 /venv*
 .jupyter_cache
 jupyter_execute
+__pycache__
 
 # pixi environments
 .pixi
 *.egg-info
+_ipython-input*-profile

content/example/walk.py

@@ -1,20 +1,56 @@
+"""A 1-D random walk.
+
+See also:
+- https://lectures.scientific-python.org/intro/numpy/auto_examples/plot_randomwalk.html
+
+"""
 import numpy as np
 
+
 def step():
     import random
-    return 1. if random.random() > .5 else -1.
+    return 1.0 if random.random() > 0.5 else -1.0
+
+
+def walk(n: int, dx: float = 1.0):
+    """The for-loop version.
+
+    Parameters
+    ----------
+    n: int
+        Number of time steps
+
+    dx: float
+        Step size. Default step size is unity.
+
+    """
+    xs = np.zeros(n)
 
-def walk(n):
-    x = np.zeros(n)
-    dx = 1. / n
     for i in range(n - 1):
-        x_new = x[i] + dx * step()
-        if x_new > 5e-3:
-            x[i + 1] = 0.
-        else:
-            x[i + 1] = x_new
-    return x
+        x_new = xs[i] + dx * step()
+        xs[i + 1] = x_new
+
+    return xs
+
+
+def walk_vec(n: int, dx: float = 1.0):
+    """The vectorized version of :func:`walk` using numpy functions."""
+    import random
+    steps = np.array(random.sample([1, -1], k=n, counts=[10 * n, 10 * n]))
+
+    # steps = np.random.choice([1, -1], size=n)
+
+    dx_steps = dx * steps
+
+    # set initial condition to zero
+    dx_steps[0] = 0
+    # use cumulative sum to replicate time evolution of position x
+    xs = np.cumsum(dx_steps)
+
+    return xs
+
 
 if __name__ == "__main__":
-    n = 100000
-    x = walk(n)
+    n = 1_000_000
+    _ = walk(n)
+    _ = walk_vec(n)

content/example/walk_ensemble_plot.py

@@ -0,0 +1,28 @@
+import numpy as np
+import matplotlib.pyplot as plt
+# See walk.py for the implementation
+import walk
+
+
+n_stories = 1000
+t_max = 200
+positions = np.empty((n_stories, t_max))
+
+for i_story in range(n_stories):
+    positions[i_story, :] = walk.walk(t_max, 1)
+
+# Determine the time evolution of the root-mean-square distance.
+sq_distance = positions**2
+# Root mean square distance
+rms_distance = np.sqrt(np.mean(sq_distance, axis=0))
+
+t = np.arange(t_max)
+
+fig, ax = plt.subplots()
+ax.plot(t, rms_distance, "g", label="ensemble RMS distance")
+ax.plot(t, np.sqrt(t), "k--", label=r"theoretical $\sqrt{\langle (\delta x)^2 \rangle}$")
+ax.set(xlabel="time", ylabel="x")
+ax.legend()
+
+plt.tight_layout()
+plt.show()