added brown interator

py/brown/brown.py

@@ -0,0 +1,66 @@
+from brown.interaction import UFuncWrapper
+import numpy as np
+
+class BrownIterator(object):
+	def __init__(self
+			, m
+			, c
+			, x
+			, y
+			, px
+			, py
+			, borders_x
+			, borders_y
+			, speed_of_light=1e3 # The value that will replace NaN momenta
+			, border_dampening=1
+			, dt=0.1):
+		self._max_m = m
+		self._i = 0
+		self.c = c
+		self.x = x
+		self.y = y
+		self.px = px
+		self.py = py
+		self.borders_x = borders_x
+		self.borders_y = borders_y
+		self.dt = dt
+		self.speed_of_light = speed_of_light
+		self.border_dampening = border_dampening
+		self._interaction = UFuncWrapper(1, c)
+	def __iter__(self):
+		self._i = 0
+		return self
+	def __next__(self):
+		self._i += 1
+		if(self._i > self._max_m and self._max_m > 0):
+			raise StopIteration()
+		if(self._i == 1):
+			return self.x, self.y
+
+		self.px, self.py = self._interaction(self.x, self.y, self.px, self.py)
+		self.px[np.isnan(self.px)] = self.speed_of_light
+		self.py[np.isnan(self.py)] = self.speed_of_light
+		self._reflect_at_borders()
+
+		self.x += self.dt * self.px
+		self.y += self.dt * self.py
+
+		return self.x, self.y
+
+
+	def _reflect_at_borders(self):
+		if(len(self.borders_x) > 0):
+			self.px[self.x <= self.borders_x[0]] *= -self.border_dampening
+			self.x[self.x <= self.borders_x[0]] = self.borders_x[0]
+		if(len(self.borders_x) > 1):
+			self.px[self.x >= self.borders_x[1]] *= -self.border_dampening
+			self.x[self.x >= self.borders_x[1]] = self.borders_x[1]
+		if(len(self.borders_y) > 0):
+			self.py[self.y <= self.borders_y[0]] *= -self.border_dampening
+			self.y[self.y <= self.borders_y[0]] = self.borders_y[0]
+		if(len(self.borders_y) > 1):
+			self.py[self.y >= self.borders_y[1]] *= -self.border_dampening
+			self.y[self.y >= self.borders_y[1]] = self.borders_y[1]
+
+		
+