50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
|
import numpy as np
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
|
||
|
|
class BrownIterator(object):
|
||
|
|
def __init__(self, N, m):
|
||
|
|
self._N = N
|
||
|
|
self._max_m = m
|
||
|
|
self._i = 0
|
||
|
|
self._xs = None
|
||
|
|
self._ys = None
|
||
|
|
def __iter__(self):
|
||
|
|
self._xs = np.zeros(self._N) + 1.5
|
||
|
|
self._ys = np.zeros(self._N)
|
||
|
|
self._i = 0
|
||
|
|
return self
|
||
|
|
def __next__(self):
|
||
|
|
self._i += 1
|
||
|
|
if(self._i > self._max_m):
|
||
|
|
raise StopIteration()
|
||
|
|
if(self._i == 1):
|
||
|
|
return self._xs, self._ys
|
||
|
|
|
||
|
|
theta = np.random.uniform(0, np.pi * 2, self._N)
|
||
|
|
|
||
|
|
self._xs = self._xs + np.cos(theta)
|
||
|
|
self._ys = self._ys + np.sin(theta)
|
||
|
|
|
||
|
|
# Reflect the particles
|
||
|
|
self._xs[self._xs > 3] += 1.5 * (3 - self._xs[self._xs > 3])
|
||
|
|
self._ys[self._ys > 3] += 1.5 * (3 - self._ys[self._ys > 3])
|
||
|
|
self._xs[self._xs < -3] += 1.5 * (-3 - self._xs[self._xs < -3])
|
||
|
|
self._ys[self._ys < -3] += 1.5 * (-3 - self._ys[self._ys < -3])
|
||
|
|
|
||
|
|
return self._xs, self._ys
|
||
|
|
|
||
|
|
if( __name__ == "__main__"):
|
||
|
|
data = np.array([i for i in BrownIterator(1000, 251)])
|
||
|
|
print(data)
|
||
|
|
|
||
|
|
p1, = plt.plot(data[5,0], data[5,1], "r.", label="t = 5")
|
||
|
|
p2, = plt.plot(data[25,0], data[25,1], "y.", label="t = 25")
|
||
|
|
p3, = plt.plot(data[50,0], data[50,1], "b.", label="t = 50")
|
||
|
|
p4, = plt.plot(data[250,0], data[250,1], "g.", label="t = 250")
|
||
|
|
|
||
|
|
plt.legend(handles=[p1, p2, p3, p4])
|
||
|
|
plt.show()
|
||
|
|
|
||
|
|
|
||
|
|
|