diff --git a/py/brown/brown.py b/py/brown/brown.py new file mode 100644 index 0000000..320b128 --- /dev/null +++ b/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] + + +