使用Python绘制2D随机游走的问题

Question

I'm trying to create a 2D random walk with Python. The random walk is to occur within a square, and if the particle crosses any side of the square the particle is to appear on the other side -- in effect, the random walk is occurring on a torus.

Here's a copy of my code:

from random import randrange as rand
from math import cos, sin, radians
import matplotlib.pyplot as plt

N = 100   # Size of square as a multiple of the step size.
NSteps = 5000   # Number of steps in simulation.
xStart = 0   # x coordinate of starting location. Origin is at centre of square
yStart = 0   # y coordinate of starting location. Origin is at centre of square

s = 1   # Step number.
x = xStart   # x coordinate of point.
y = yStart   # y coordinate of point.
xList = []   # List of the x coordinates of all points visited.
yList = []   # List of the y coordinates of all points visited.

while s <= NSteps:
    angle = radians(rand(361))
    x += cos(angle)    
    if x > N/2:
        x -= N
    elif x < -N/2:
        x += N    
    xList += [x]
    y += sin(angle)
    if y > N/2:
        y -= N
    elif y < -N/2:
        y += N
    yList += [y]    
    s += 1

plt.figure(figsize=(13,8))
frame = plt.gca()
plt.plot(xList,yList,c="b")
plt.xlim(-N/2,N/2)
plt.ylim(-N/2,N/2)
frame.axes.get_xaxis().set_visible(False)
frame.axes.get_yaxis().set_visible(False)
plt.savefig("randomWalk.png", bbox_inches="tight")

This code produces a plot like this:

As you can see, whenever the particle crosses one of the sides I get these 'streaks' on the plot because plot() will connect two points no matter how far apart they are. Is there a way to prevent this from happening?

Answer 1

I also re-wrote your steping code a bit to make it easier (in my opinion) to read:

from random import randrange as rand
from numpy import cos, sin, radians
import numpy as np
import matplotlib.pyplot as plt

N = 100   # Size of square as a multiple of the step size.
NSteps = 5000   # Number of steps in simulation.
xStart = 0   # x coordinate of starting location. Origin is at centre of square
yStart = 0   # y coordinate of starting location. Origin is at centre of square

s = 1   # Step number.
x = xStart   # x coordinate of point.
y = yStart   # y coordinate of point.
xList = []   # List of the x coordinates of all points visited.
yList = []   # List of the y coordinates of all points visited.

def wrap(v, N):

    if v > N/2:
        return v - N, True
    elif v < -N/2:
        return v + N, True
    return v, False


for j in range(NSteps):
    angle = radians(rand(361))
    x, wrap_flag_x = wrap(x + cos(angle), N)
    y, wrap_flag_y = wrap(y + sin(angle), N)
    if wrap_flag_x or wrap_flag_y:
        xList.append(np.nan)
        yList.append(np.nan)
    xList.append(x)
    yList.append(y)    

fig, ax = plt.subplots()
ax.plot(xList,yList,c="b")
ax.set_xlim(-N/2,N/2)
ax.set_ylim(-N/2,N/2)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

They are putting np.nan (Not A Number, which is part of the float spec) into your list. When mpl is drawing lines it (with the default line style) connects all of the points. The points that are np.nan can't be drawn to the screen so the line from the last point to the np.nan point is not drawn and the line from the np.nan to the next point is not drawn, hence the break in your line.

As a side note, most of this simulation can be vectorized:

from numpy.random import randint
from numpy import cos, sin, radians, cumsum
import numpy as np
import matplotlib.pyplot as plt

N = 100   # Size of square as a multiple of the step size.
NSteps = 5000   # Number of steps in simulation.
x_start = 0   # x coordinate of starting location. Origin is at centre of square
y_start = 0   # y coordinate of starting location. Origin is at centre of square



# get all of the angles
angles = radians(randint(low=0, high=361, size=NSteps))

# get (unwrapped) positions
x = cumsum(cos(angles)) + x_start
y = cumsum(sin(angles)) + y_start

# find where the position crosses the boundary
x_wraps = np.where(np.diff((x + N/2) // N))[0]
y_wraps = np.where(np.diff((y + N/2) // N))[0]

# do the wrapping
x = x - N * ((x + N/2)//N)
y = y - N * ((y + N/2)//N)

I leave using the wrap locations to insert nans as an exercise for the reader ;)