1 files changed, 83 insertions, 0 deletions
diff --git a/code/sunlab/common/mathlib/random_walks.py b/code/sunlab/common/mathlib/random_walks.py
new file mode 100644
index 0000000..3aa3bcb
--- /dev/null
+++ b/code/sunlab/common/mathlib/random_walks.py
@@ -0,0 +1,83 @@
+def get_levy_flight(T=50, D=2, t0=0.1, alpha=3, periodic=False):
+    from numpy import vstack
+    from mistree import get_levy_flight as get_flight
+
+    if D == 2:
+        x, y = get_flight(T, mode="2D", periodic=periodic, t_0=t0, alpha=alpha)
+        xy = vstack([x, y]).T
+    elif D == 3:
+        x, y, z = get_flight(T, mode="3D", periodic=periodic, t_0=t0, alpha=alpha)
+        xy = vstack([x, y, z]).T
+    else:
+        raise ValueError(f"Dimension {D} not supported!")
+    return xy
+
+
+def get_levy_flights(N=10, T=50, D=2, t0=0.1, alpha=3, periodic=False):
+    from numpy import moveaxis, array
+
+    trajectories = []
+    for _ in range(N):
+        xy = get_levy_flight(T=T, D=D, t0=t0, alpha=alpha, periodic=periodic)
+        trajectories.append(xy)
+    return moveaxis(array(trajectories), 0, 1)
+
+
+def get_jitter_levy_flights(
+    N=10, T=50, D=2, t0=0.1, alpha=3, periodic=False, noise=5e-2
+):
+    from numpy.random import randn
+
+    trajectories = get_levy_flights(
+        N=N, T=T, D=D, t0=t0, alpha=alpha, periodic=periodic
+    )
+    return trajectories + randn(*trajectories.shape) * noise
+
+
+def get_gaussian_random_walk(T=50, D=2, R=5, step_size=0.5, soft=None):
+    from numpy import array, sin, cos, exp, zeros, pi
+    from numpy.random import randn, uniform, rand
+    from numpy.linalg import norm
+
+    def is_in(x, R=1):
+        from numpy.linalg import norm
+
+        return norm(x) < R
+
+    X = zeros((T, D))
+    for t in range(1, T):
+        while True:
+            if D == 2:
+                angle = uniform(0, pi * 2)
+                step = randn(1) * step_size
+                X[t, :] = X[t - 1, :] + array([cos(angle), sin(angle)]) * step
+            else:
+                X[t, :] = X[t - 1, :] + randn(D) / D * step_size
+            if soft is None:
+                if is_in(X[t, :], R):
+                    break
+            elif rand() < exp(-(norm(X[t, :]) - R) * soft):
+                break
+    return X
+
+
+def get_gaussian_random_walks(N=10, T=50, D=2, R=5, step_size=0.5, soft=None):
+    from numpy import moveaxis, array
+
+    trajectories = []
+    for _ in range(N):
+        xy = get_gaussian_random_walk(T=T, D=D, R=R, step_size=step_size, soft=soft)
+        trajectories.append(xy)
+    return moveaxis(array(trajectories), 0, 1)
+
+
+def get_gaussian_sample(T=50, D=2):
+    from numpy.random import randn
+
+    return randn(T, D)
+
+
+def get_gaussian_samples(N=10, T=50, D=2, R=5, step_size=0.5):
+    from numpy.random import randn
+
+    return randn(T, N, D)