diff --git a/exam/ex08/main.py b/exam/ex08/main.py new file mode 100644 index 0000000..3f2d9a3 --- /dev/null +++ b/exam/ex08/main.py @@ -0,0 +1,49 @@ +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() + + +