added brown interator
This commit is contained in:
parent
dba5664cf3
commit
c01a323376
66
py/brown/brown.py
Normal file
66
py/brown/brown.py
Normal file
|
@ -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]
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user