import numpy
import error
import math
#System specifications:
#System classes are responsible for defining the differential operator, and
#rendering the initial state to seed solvers
#System constructors convert data from human read/writeable form, to an
#internal representation
#System state consists of two parts: the time component, which is a float,
#and the w vector, which has some shape, subject only to the restriction that
#the differential operator returns an array with the same shape, when given a
#w vector and a time
def optional_w(f):
def f_opt(self, w=False):
if w == False:
w = self.initial[0]['w']
return f(self, w)
f_opt.__name__ = f.__name__
f_opt.__doc__ = f.__doc__
return f_opt
def optional_state(f):
def f_opt(self, state=False):
if state == False:
state = self.initial[0]
return f(self, state)
f_opt.__name__ = f.__name__
f_opt.__doc__ = f.__doc__
return f_opt
def statify(f):
def f_new(state=False):
if state == False:
return f()
return f(state['w'])
f_new.__name__ = f.__name__
f_new.__doc__ = f.__doc__
return f_new
class System(object):
"Dummy implementation of the System class."
w_type = NotImplemented
initial = NotImplemented
@optional_state
def differential(self, state):
raise NotImplementedError
def flat_differential(self, t=None, w=None):
if w is None or t is None:
return self.differential().view(dtype=numpy.float64)
else:
return self.differential(
{'t': t, 'w': numpy.rec.array(w,
dtype=self.initial[0]['w'].dtype)}).view(
dtype=numpy.float64)
class Polynomial(System):
def __init__(self, coefficients, t=0):
self.coefficients = coefficients
#self.initial =
def exact(self, t):
return sum([coefficients[i] * t**i for i in range(len(coefficients))])
@optional_state
def differential(self, state):
t = state['t']
return sum(i * [coefficients[i] * t**(i-1) for i in range(len(coefficients))])
class Exponential(System):
def __init__(self, initial=1, lambda_=-1, t=0):
self.lambda_ = lambda_
self.w_type = numpy.float
self.array_type = numpy.dtype([('t', numpy.float, 1),
('w', self.w_type, 1)])
self.initial = numpy.empty(1, self.array_type)
self.initial[0]['t'] = t
self.initial[0]['w'] = initial
def exact(self, t):
return self.initial[0]['w'] * math.exp(self.lambda_ * (t - self.initial[0]['t']))
@optional_state
def differential(self, state):
return state['w'] * self.lambda_
class Orbital(System):
def __init__(self, *bodies, **kwargs):
t = 0
if 't' in kwargs:
t = kwargs['t']
self.G = 1
if 'G' in kwargs:
self.G = kwargs['G']
self.size = len(bodies)
if self.size == 0: #Degenerate case in which nothing happens
return
if self.size == 1: #Degenerate case in which all methods have no error
return
self.masses = numpy.asarray([[body['mass']] for body in bodies])
self.mass = sum(self.masses)
self.potential_constants = -self.G*numpy.ma.masked_array(self.masses.transpose()*self.masses, numpy.eye(self.size))
self.D = len(bodies[0]['pos'])
self.zero = numpy.zeros(self.D)
if self.D != 2 and self.D != 3:
return
self.point_type = numpy.dtype([('pos', numpy.float, self.D),
('vel', numpy.float, self.D)])
self.w_type = numpy.dtype((self.point_type, self.size))
self.array_type = numpy.dtype([('t', numpy.float, 1),
('w', self.w_type, 1)])
self.initial = numpy.empty(1, self.array_type)
self.initial[0]['t'] = t
for i in range(self.size):
for j in range(self.D):
self.initial[0]['w'][i]['pos'][j] = bodies[i]['pos'][j]
self.initial[0]['w'][i]['vel'][j] = bodies[i]['vel'][j]
self.errors = {
"barycenter": error.Projection(statify(self.barycenter),
self.projected_barycenter,
self.momentum()/sum(self.masses),
t),
"energy": error.Invariant(statify(self.mechanical_energy)),
"momentum": error.InvariantCoordinate(statify(self.momentum)),
"angular momentum": (
self.D==3 and
error.InvariantCoordinate or
error.Invariant)(statify(self.angular_momentum))}
if (self.momentum() == self.zero).all():
print "WARNING: The system's momentum is zero. This is physically unlikely, and leads to degenerate behavior."
@optional_state
def pos(self, state):
return state['w']['pos']
@optional_state
def vel(self, state):
return state['w']['vel']
@optional_state
def tee(self, state):
return state['t']
@optional_w
def momenta(self, w):
return self.masses * w['vel']
def momentum(self, w=False):
return sum(self.momenta(w))
#@optional_w
def mechanical_energy(self, w=False):
return self.kinetic_energy(w) + self.potential_energy(w)
@optional_w
def kinetic_energy(self, w):
return 0.5 * sum(sum(self.masses * (w['vel'] * w['vel'])))
@optional_w
def barycenter(self, w):
return sum(self.masses * w['pos']) / self.mass
def projected_barycenter(self, t=None):
diff = 0
if t is not None:
diff = t - self.initial[0]['t']
return self.barycenter() + diff * self.momentum() / self.mass
@optional_w
def angular_momentum(self, w):
return sum(numpy.cross(w['pos']-numpy.repeat([self.barycenter(w)],
self.size, 0),
self.momenta(w)))
@optional_w
def displacements(self, w):
pos = numpy.repeat([w['pos']], self.size, 0)
#print pos.shape
return pos - pos.transpose(1, 0, 2)
#@optional_w
def distances_squared(self, w=False):
return (self.displacements(w)**2).sum(2)
#@optional_w
def distances(self, w=False):
return self.distances_squared(w)**0.5
#@optional_w
def potentials(self, w=False):
return self.potential_constants / self.distances(w)
#@optional_w
def potential_energy(self, w=False):
return 0.5*self.potentials(w).filled(0).sum()
#@optional_w
def forces(self, w=False):
return ((numpy.repeat(self.potentials(w) / self.distances_squared(w),
self.D).reshape((self.size, self.size, self.D)) *
self.displacements(w)).filled(0)).sum(0)
#@optional_w
def accelerations(self, w=False):
return self.forces(w)/self.masses
@optional_state
def differential(self, state):
w = state['w']
result = numpy.empty_like(w)
result['pos'] = w['vel']
result['vel'] = self.accelerations(w)
return result
Comments