73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
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
|
|
|
|
delta_px, delta_py = self._interaction(self.x, self.y)
|
|
self.px = delta_px + self.px
|
|
self.py = delta_py + self.py
|
|
# XXX: We need the (-1)**i to make the problem
|
|
# symmetric.
|
|
# FIXME: is this necessary?
|
|
self.px[np.isnan(self.px)] = self.speed_of_light * (-1)**self._i
|
|
self.py[np.isnan(self.py)] = self.speed_of_light * (-1)**self._i
|
|
|
|
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]
|
|
|
|
|
|
|