finalized v2

MATHTYPE

@@ -0,0 +1 @@
+mp

main.py

@@ -1,6 +1,6 @@
 
 import nbody as nb
-
+'''
 LINE_PROFILE=1
 
 bodies = nb.horizons_batch(('10','199','299','399','499','599','699','799','899'))
@@ -16,15 +16,23 @@
 phys.simulate(20000)
 #phys.save_as('bodies','solarsystem_bodies')
 
-
+'''
 #phys.load_as('bodies','solarsystem_bodies')
-
-sim = nb.mplVisual(engine=phys, 
-                name='SS',
-                focus_body=phys.bodies[0], show_info=True, autoscale=False,
-                step_skip_frames=100, step_skip_points=100, max_period=3, cache=False, do_picking=True)
-
-
+eng= nb.Engine(dt=0.05)
+bodies = (nb.Body(mass=5, init_pos=(-10.0,0,0), init_vel=(2.0,0,0), bounce=1, radius=2),
+          nb.Body(mass=10, init_pos=(10.0,0,0), init_vel=(-2.0,0,0), bounce=1, radius=2),
+            #Body(mass=1, init_pos=(15,1.,0), init_vel=(-0.5,0,0), bounce=1, radius=1)
+            )
+eng.do_bodygravity = False
+eng.do_collisions = True
+eng.do_fieldgravity = False
+eng.attach_bodies(bodies)
+eng.simulate(1000)
+sim = nb.MPLVisual(engine=eng, 
+                name='SS', show_info=True,
+                step_skip_frames=1, step_skip_points=1, max_pts=1, cache=False, body_model='surface')
+
+sim.start()
 '''
 import cProfile, pstats
 profiler = cProfile.Profile()

nbody/__init__.py

@@ -44,6 +44,7 @@
      .tools.io
      .tools.horizons
 """
+from .core.base import math_conf as CONFIG #noqa
 from .core.body import Body
 from .core.engine import Engine
 from .core.visual import MPLVisual

nbody/core/base.py

@@ -1,21 +1,25 @@
 from __future__ import annotations
 # Python Builtins
 
+
 from numbers import Number
-from collections.abc import Iterable
 from mpmath import mp, fp
 
 # Local error definitions.
 from ..tools import errors as e
 
+# CONFIGURE GLOBAL MATH FUNCTIONS
+from ..tools import _config as conf
+math_conf = conf.MathContext(use=mp)
 
-
+Iterable = (list, tuple)
 Any = object
 NoneType = type(None)
-NumType = (Number,type(mp.mpf(1)), type(fp.mpf(1)))
+NumType = (Number, int, float ,type(mp.mpf(1)),type(fp.mpf(1)))
 mp.pretty, fp.pretty = True, True
 
 
+
 def _O(obj):  # noqa: N802
     if isinstance(obj,(Vector, Variable, HistoricVariable, HistoricVector)):
         return obj.c()
@@ -34,13 +38,13 @@ def _V(obj):  # noqa: N802
 
 def _ntype(*objs):
     types = [type(_O(obj)) for obj in objs]
-    for ntype in types:
-        if ntype not in NumType:
-            return ntype
+    for i,obj in enumerate(objs):
+        if not isinstance(obj, NumType):
+            return types[i]
     if all(ntype == int for ntype in types):
         return int
     else:
-        return fp.mpf
+        return math_conf.type
 
 
 def typecheck(argtype):
@@ -94,49 +98,49 @@ def __repr__(self):
             return f'VarObj({self.c()} {self.units}, len={len(self)}, rec={self.record}, id={self.identity})'
 
     def __add__(self,other):
-        return Variable(fp.fadd(self.c(),_O(other)))
+        return Variable(math_conf.add(self.c(),_O(other)))
 
     __radd__ = __add__
 
     def __sub__(self,other):
-        return Variable(fp.fsub(self.c(),_O(other)))
+        return Variable(math_conf.chop(math_conf.sub(self.c(),_O(other))))
 
     def __mul__(self,other):
-        return Variable(fp.fmul(self.c(),_O(other)))
+        return Variable(math_conf.mul(self.c(),_O(other)))
 
     __rmul__ = __mul__
 
     def __truediv__(self,other):
-        return Variable(fp.fdiv(self.c(),_O(other)))
+        return Variable(math_conf.div(self.c(),_O(other)))
 
     def __iadd__(self,other):
-        self.next(fp.fadd(self.c(),_O(other)))
+        self.next(math_conf.add(self.c(),_O(other)))
         return self
 
     def __isub__(self,other):
-        self.next(fp.fsub(self.c(),_O(other)))
+        self.next(math_conf.chop(math_conf.sub(self.c(),_O(other))))
         return self
 
     def __imul__(self,other):
-        self.next(fp.fmul(self.c(),_O(other)))
+        self.next(math_conf.mul(self.c(),_O(other)))
         return self
 
     def __itruediv__(self,other):
-        self.next(fp.fdiv(self.c(),_O(other)))
+        self.next(math_conf.div(self.c(),_O(other)))
         return self
 
     def __rsub__(self,other):
-        return Variable(fp.chop(mp.fsub(_O(other), self.c())))
+        return Variable(math_conf.chop(mp.fsub(_O(other), self.c())))
 
 
     def __rtruediv__(self,other):
-        return Variable(fp.fdiv(_O(other), self.c()))
+        return Variable(math_conf.div(_O(other), self.c()))
 
     def __pow__(self, other):
-        return Variable(fp.power(self.c(),_O(other)))
+        return Variable(math_conf.pow(self.c(),_O(other)))
 
     def __rpow__(self, other):
-        return Variable(fp.power(_O(other), self.c()))
+        return Variable(math_conf.pow(_O(other), self.c()))
 
     def __eq__(self, other):
         temp = _O(other)
@@ -150,7 +154,7 @@ def __lt__(self, other):
 
 
     def __le__(self, other):
-        return self.__eq__(other) or self.__le__(other)
+        return (self.__eq__(other) or self.__lt__(other))
 
 
     def __ne__(self, other):
@@ -164,7 +168,7 @@ def __gt__(self, other):
 
 
     def __ge__(self, other):
-        return self.__eq__(other) or self.__gt__(other)
+        return (self.__eq__(other) or self.__gt__(other))
 
 class HistoricVariable(Variable):
     def __init__(self,init_var,identity='HistoricVariable',units=''):
@@ -250,7 +254,7 @@ def c(self,usage=None):
             e.raise_out_of_range('c()',usage)
 
     def magnitude(self):
-        return fp.norm(fp.matrix(self.c()))
+        return math_conf.norm(math_conf.matrix(self.c()))
 
     def unit(self):
         if float(self.magnitude()) == 0.:
@@ -281,7 +285,7 @@ def __iter__(self):
     def __add__(self,other):
         temp = _O(other)
         if len(temp) == 3:
-            return Vector([fp.fadd(self.c(i),temp[i]) for i in range(3)])
+            return Vector([math_conf.add(self.c(i),temp[i]) for i in range(3)])
         else:
             e.raise_component_error('other or temp',temp)
 
@@ -290,16 +294,16 @@ def __add__(self,other):
     def __sub__(self,other):
         temp = _O(other)
         if len(temp) == 3:
-            return Vector([fp.fsub(self.c(i),temp[i]) for i in range(3)])
+            return Vector([math_conf.sub(self.c(i),temp[i]) for i in range(3)])
         else:
             e.raise_component_error('other or temp',temp)
 
     def __mul__(self,other):
         temp = _O(other)
         if isinstance(temp,Iterable) and len(temp) == 3:
-            return Variable(mp.fsum([fp.fmul(val, temp[i]) for (i, val) in enumerate(self.c())]))
+            return Variable(math_conf.sum([math_conf.mul(val, temp[i]) for (i, val) in enumerate(self.c())]))
         elif isinstance(temp,NumType):
-            return Vector([fp.fmul(val,temp) for val in self.c()])
+            return Vector([math_conf.mul(val,temp) for val in self.c()])
         else:
             e.raise_component_error('other or temp',temp)
 
@@ -308,15 +312,15 @@ def __mul__(self,other):
     def __truediv__(self,other):
         temp = _O(other)
         if isinstance(temp,NumType):
-            return Vector([fp.fdiv(val,temp) for val in self.c()])
+            return Vector([math_conf.div(val,temp) for val in self.c()])
         else:
             e.raise_type_error('other',NumType,other)
 
 
     def __rsub__(self,other):
         temp = _O(other)
         if len(temp) == 3:
-            return Vector([fp.fsub(temp[i], self.c(i)) for i in range(3)])
+            return Vector([math_conf.sub(temp[i], self.c(i)) for i in range(3)])
         else:
             e.raise_component_error('other or temp',temp)
 

nbody/core/body.py

@@ -74,10 +74,7 @@ def update(self,dt=1,vel_change=None,acc_change=None,vel_next=None):
 
         if vel_change is not None:
             vel_change = _V(vel_change)
-            if isinstance(vel_change,Iterable) and len(vel_change.c()) == 3:
-                self.pos.next(self.pos + (vel + vel_change)*dt)
-            else:
-                raise RuntimeError
+            self.pos.next(self.pos + (vel + vel_change)*dt)
         else:
             self.pos.next(self.pos + vel*dt)
 
@@ -114,7 +111,7 @@ def per():
                 else:
                     return a
             def ke():
-                return (Vector(self.vel[ind])*(self.vel[ind])*(1/2)*self.mass).c()
+                return ((Vector(self.vel[ind]).magnitude()**2)*(1/2)*self.mass).c()
 
             _get_lookup = {**dict.fromkeys(['sma', 'semi_major_axis'],sma),
                        **dict.fromkeys(['per', 'period'],per),

nbody/core/engine.py

@@ -86,26 +86,31 @@ def _check_collision(self,body,co_restitution=0):
         if self.do_collisions == True:
             # has it found a collision?
             returned_coll = False
-            for bod in self.bodies:
-                if body != bod: 
+            for body2 in self.bodies:
+                if body != body2: 
                     # get distance between bodies
-                    body_dist = (bod.pos - body.pos).magnitude()             
+                    body_dist = (body2.pos - body.pos).magnitude()             
                     # check if distance is less than radii
-                    if body_dist < (bod.radius + body.radius).c():
-                        (returned_coll,n,meff) = (True,(bod.pos-body.pos)/body_dist,1/((1/bod.mass)+(1/body.mass)))
+                    if body_dist < (body2.radius + body.radius).c():
+                        returned_coll = True
+                        n = (body2.pos-body.pos).unit()
+                        meff = 1/((1/body2.mass)+(1/body.mass))
+                        rel_v = body2.vel - body.vel
+                        impulse =  n *(n*rel_v)*(1 + co_restitution) * meff
+                        dv = impulse/body.mass
                         # calc vel change: dv = n^ * [n^*[v2 -v1]*[1+e]*m"]/m1, n^ is normal unit vector,
                         # m" is effective mass as above.
-                        dv = n*(-1*((n*(body.vel - bod.vel))*(1+co_restitution)*meff)/body.mass)
+                        #dv = n * -1*(( (  n   *(body.vel-body2.vel) )*(   1+co_restitution)*meff     )/body.mass.c())
                         return (dv+body.vel), False
 
             unitcomps = {'x':Vector((1,0,0)),'y':Vector((0,1,0)),'z':Vector((0,0,1))}
             for pl in self.planes:
                 # check for plane collision. we estimate a range of times as the plane has no thickness,
                 # and time interval would have to be extremely miniscule compared to velocity to always catch.
                 body_dists = list(abs(((body.pos + body.vel*m*self.dt)[pl[0]] - pl[1])) for m in range(self._rangechk))
-                if any([(body_dists[i] < body.radius) for i in range(self._rangechk)]) or\
-                body_dists[0] < body.radius:
-                    on_plane = (body_dists[0]/body.radius <= 1.01 and body_dists[-1]/body.radius <= 1.01)
+                if any([(body_dists[i] < body.radius) for i in range(self._rangechk)]):
+                    on_plane = ((body_dists[0]/body.radius <= 1.01) if len(body_dists) == 1 
+                                else (body_dists[0]/body.radius <= 1.01 and body_dists[-1]/body.radius <= 1.01))
                     returned_coll = True
                     # similar to above, but normals are precalculated.
                     return (unitcomps[pl[0]]*(-2*(body.vel*unitcomps[pl[0]])) + body.vel)*co_restitution, on_plane
@@ -132,17 +137,19 @@ def _find_gravity(self,body):
 
         return res/body.mass
 
-    def _gen_sim(self):
-        for i,body in enumerate(self.bodies):
-            yield [body, *self._check_collision(body, body.bounce),self._find_gravity(body)]
+
+
 
     def simulate(self,intervals):
         if isinstance(intervals, int) and len(self.bodies) > 0:
             for _ in trange(intervals, 
                             desc=f'«Engine» → Evaluating motion for each interval of {self.dt} seconds', unit='ints'):
+                _temp = [[0,0,0] for _ in self.bodies]
+                for i,body in enumerate(self.bodies):
+                    _temp[i] = [*self._check_collision(body, body.bounce),self._find_gravity(body)]
 
-
-                for body, col_vel, on_plane, acc_g in self._gen_sim():
+                for i,body in enumerate(self.bodies):
+                    col_vel, on_plane, acc_g = _temp[i]
                     if not on_plane and self.do_fieldgravity:
                         fieldvel = list(sum(v.c(i) for v in self.fields) for i in range(3))
                     else: