Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement