Skip to content

Commit

Permalink
text and unit format
Browse files Browse the repository at this point in the history
  • Loading branch information
@will-roscoe
will-roscoe committed Nov 25, 2023
1 parent 20d7d7e commit bdee3dc
Show file tree
Hide file tree
Showing 4 changed files with 134 additions and 69 deletions.
23 changes: 18 additions & 5 deletions main.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@




import nbody.core as nb

phys = nb.Engine(dt=1000)
phys.load_as('bodies')
for i,color in enumerate(['#F28322','#BFBEBD', '#D9B391', '#63BAA6', '#F27A5E', '#BFAE99', '#D9B779', '#95BBBF', '#789EBF']):
phys.bodies[i].color = color
phys.save_as('bodies', 'solarsystem_bodies')
phys.load_as('bodies','solarsystem_bodies')
#for i,color in enumerate(['#F28322','#BFBEBD', '#D9B391', '#63BAA6', '#F27A5E', '#BFAE99', '#D9B779', '#95BBBF', '#789EBF']):
#phys.bodies[i].color = color
for b in phys.bodies:
b.pos.units, b.vel.units, b.acc.units = 'metre', 'metre / second', 'metre / second^2'
phys.save_as('bodies','solarsystem_bodies')
sim = nb.mplVisual(phys, 'SS', phys.bodies[0],None, False,
show_grid= True,
show_shadows= False,
Expand All @@ -17,4 +22,12 @@
guistyle = 'dark',
do_picking = True,
show_info = True)
sim.start(fps=30, frameskip=1000, plotskip=200)
sim.start(fps=30, frameskip=1000, plotskip=200)
'''
s = 615634.456743
k = s * nb.ur.metre/nb.ur.second**2
print(k)
nb.ur.default_format = '~P'
k = (s * nb.ur.m/nb.ur.s**2).to_compact()
print(f'{k:.3f~P}')
'''
Binary file modified nbody/__pycache__/core.cpython-312.pyc
View file
Binary file not shown.
180 changes: 116 additions & 64 deletions nbody/core.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from genericpath import exists
import re
from datetime import datetime, timedelta
from multiprocessing import Pool, freeze_support
Expand All @@ -19,11 +20,11 @@
from cycler import cycler
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Line3D

from pint import UnitRegistry, Quantity

# JPL Querying Packages
from astroquery.jplhorizons import Horizons
import astropy.units as u
import astropy.units as ast


from . import errors as e
Expand Down Expand Up @@ -70,12 +71,12 @@ def __init__(self,mass,init_pos,init_vel=(0,0,0),

_id = self.identity

self.pos = HistoricVector(li=init_pos,identity=f'{_id}_pos',units_v='m')
self.vel = HistoricVector(li=init_vel,identity=f'{_id}_vel',units_v='ms^-1')
self.acc = HistoricVector(0,0,0,identity=f'{_id}_acc',units_v='ms^-2')
self.pos = HistoricVector(li=init_pos,identity=f'{_id}_pos',units_v='metre')
self.vel = HistoricVector(li=init_vel,identity=f'{_id}_vel',units_v='metre / second')
self.acc = HistoricVector(0,0,0,identity=f'{_id}_acc',units_v='metre / second^2')

self.mass = Variable(mass,identity=f'{_id}_mass',units='kg')
self.radius = Variable(radius,identity=f'{_id}_rad',units='m')
self.radius = Variable(radius,identity=f'{_id}_rad',units='metre')
self.bounce = bounce
self.color = color

Expand Down Expand Up @@ -199,8 +200,8 @@ def horizons_query(searchquery,observer='0',time='2023-11-03',num_type=float,ret
if m_unit == 'g':
mass /= 1000

x, y, z = [num_type(_tab[pos].quantity[0].to_value(u.m)) for pos in ('x', 'y', 'z')]
vx, vy, vz = [num_type(_tab[pos].quantity[0].to_value(u.m/u.s)) for pos in ('vx', 'vy', 'vz')]
x, y, z = [num_type(_tab[pos].quantity[0].to_value(ast.m)) for pos in ('x', 'y', 'z')]
vx, vy, vz = [num_type(_tab[pos].quantity[0].to_value(ast.m/ast.s)) for pos in ('vx', 'vy', 'vz')]

if return_type.lower() == 'print':
print(f'''Object: {name}\n***********\nMass: {mass} kg\nRadius: {rad} m\n
Expand All @@ -226,18 +227,37 @@ def horizons_batch(search_queries,observer='0',time='2023-11-03',num_type=float,















ur = UnitRegistry()
ur.default_format = '~P'
class Formatter:
def __init__(self, output_raw = False):
self.b = None
self.i = 0
{'identity':'','scalars':{'mass':0*ur.kg,'radius':0*ur.m,'energy':0*ur.joule,
'period':0*ur.s},'vectors':{'pos':[0*ur.m for n in range(3)],
'vel':[0*ur.m/ur.s for n in range(3)],'acc':[0*ur.m/ur.s**2 for n in range(3)]}}
def _quantities(self):
self.q = {'identity' : self.b.identity,
'scalars':{
'mass' : self.b.mass[self.i]*ur(self.b.mass.units),
'radius' : self.b.radius[self.i]*ur(self.b.radius.units),
'energy' : mplVisual.ke(self.b) * ur.joule,
'period' : mplVisual.period(self.b) * ur.s},
'vectors':{
'pos' : [self.b.pos[self.i][n]*ur(self.b.pos.units) for n in range(3)],
'vel' : [self.b.vel[self.i][n]*ur(self.b.vel.units) for n in range(3)],
'acc' : [self.b.acc[self.i][n]*ur(self.b.acc.units) for n in range(3)]}
}
return self.q
def _str(self):
return "{}\npos:{}\nvel:{}\
\nacc:{}\nmass:{}\nbody radius:{}\nKE:{}\
\nest period:{}"
def standardise(self):
if self.q['mass'] > 1.*ur.light_year:
self.q['mass'] = self.q['mass'].ito(ur.light_year)
if self.q['mass'] > 1.*ur.astronomical_unit:
self.q['mass'] = self.q['mass'].ito(ur.astronomical_unit)

class Engine:
def __init__(self,dt=1,checking_range=3):
Expand Down Expand Up @@ -385,19 +405,20 @@ def __init__(self, engine, name='NBody Simulation (Matplotib)',
autoscale=True,show_grid=True,show_shadows=False,
show_acceleration=False,show_velocity=False,vector_size=1,
labelling_type='legend',body_model='dots',guistyle='default',
do_picking = False, show_info=False):
do_picking = False, show_info=False, info_best_units = 'simple'):

(self.engine,self.focus_body,self.focus_range,self.autoscale,self.show_grid,
self.show_shadows,self.show_acceleration,self.show_velocity,self.vector_size,
self.labelling_type,self.body_model,self.guistyle,self.do_picking, self.show_info) = typecheck((
self.labelling_type,self.body_model,self.guistyle,self.do_picking, self.show_info, self.change_units) = typecheck((
(engine,Engine),(focus_body,(Body,NoneType)),(focus_range,(*NumType, NoneType)),
(autoscale,bool),(show_grid,bool),(show_shadows,bool),(show_acceleration,bool),
(show_velocity,bool),(vector_size,NumType),(labelling_type,str),(body_model,str),
(guistyle,str),(do_picking,bool), (show_info, bool)))
(guistyle,str),(do_picking,bool),(show_info, bool),(info_best_units, str)))

self.fig = plt.figure(name, figsize=(16,9))
self.ax = self.fig.add_subplot(projection='3d')
self.ax1 = self.fig.add_axes((0.05,0.25,0.05,0.5))

self.fig.subplots_adjust(left=0,bottom=0,right=1,top=1,wspace=0,hspace=0)
(self.zoom_slider,self._frameskip,self._plotskip,self.ax.computed_zorder, self.inf_b) = (
Slider(self.ax1,'Zoom',0.1,10,valinit=1,orientation='vertical'),1,1,False, self.focus_body)
Expand All @@ -414,7 +435,10 @@ def _axcolor(color):
for ax in (self.ax.xaxis,self.ax.yaxis,self.ax.zaxis):
ax.label.set_color(color)
mpl.rcParams['axes.labelcolor'] = color

self.do_fast_units = {'simple':False, 'fast':True, 'long':False}[self.change_units]
self.change_units = {'simple':True, 'fast':False, 'long':False}[self.change_units]


if self.guistyle == 'dark':
for artist in (self.fig,self.ax):
artist.set_facecolor('black')
Expand All @@ -423,7 +447,7 @@ def _axcolor(color):
mpl.rcParams['text.color'] = 'white'
self.tcolor = 'white'
self.zoom_slider.label.set_color(self.tcolor)

else:
self.tcolor = 'black'

Expand All @@ -439,31 +463,36 @@ def _axcolor(color):
def _draw_vectors(self,pos,other,c):
self.ax.quiver(*pos,*other,length=self.vector_size,color=c,zorder=8,clip_on=False)



def _animate(self,ind):
def sm_a(body):
def convert_to_str(self, quantity, simple=True):
try:
if simple == False:
if isinstance(quantity, Iterable):
return '('+''.join((f'{q:~P}' for q in quantity))+')'
else:
return f'{quantity:~P}'
elif simple == True:
if isinstance(quantity, Iterable):
return '('+''.join((f'{(q.to_compact()):.4f~P}' for q in quantity))+')'
else:
return f'{(quantity.to_compact()):.4f~P}'
except ValueError:
return 'NaN'
def sm_a(self,body,ind):
try:
a = max(list(Vector(body.pos[i]).magnitude() for i in range(0,ind,self._plotskip)))
return a
except ValueError:
return 'NaN'

def period(body):
a = sm_a(body)
if a != 'NaN':

return 2*math.pi*math.sqrt((a**3)/(G*self._mass_big))
else:
return a
def period(self,body):
a = self.sm_a(body)
if a != 'NaN':



def ke(body):
return Vector(body.vel[ind])*(body.vel[ind])*(1/2)*body.mass.c()



return 2*math.pi*math.sqrt((a**3)/(G*self._mass_big))
else:
return a
def ke(self,body,ind):
return Vector(body.vel[ind])*(body.vel[ind])*(1/2)*body.mass.c()
def _animate(self,ind):


co,pr = {'clip_on':False}, 10
Expand Down Expand Up @@ -503,7 +532,7 @@ def ke(body):
for b in self.engine.bodies:
cut = 0
try:
tau = (ind-(2*period(b)/self.engine.dt))
tau = (ind-(2*self.period(b)/self.engine.dt))
if tau > cut:
cut = math.ceil(tau)
except TypeError:
Expand Down Expand Up @@ -533,39 +562,62 @@ def ke(body):
len(_poshist[2]),color='black',zorder=1.5,**co)



if self.labelling_type == 'legend':
self.leg = self.ax.legend(draggable=True, facecolor='black', fancybox=True)
for text in self.leg.get_texts():
text.set_picker(pr)
text.set_color(self.tcolor)

if self.show_info:
self.ax.text2D(0.05, 0.2, f"{self.inf_b.identity}\npos:{self.inf_b.pos[ind]}m\nvel:{self.inf_b.vel[ind]}ms^-1\
\nacc:{self.inf_b.acc[ind]}ms^-2\nmass:{self.inf_b.mass.c()} kg\nbody radius:{self.inf_b.radius.c()}\nKE:{ke(self.inf_b)} J\
\nest period:{period(self.focus_body)} s", size='small', transform=self.ax.transAxes)


ident = self.inf_b.identity
if not self.do_fast_units:
dets = {'po' : [self.inf_b.pos[ind][n]*ur(self.inf_b.pos.units) for n in range(3)],
'v' : [self.inf_b.vel[ind][n]*ur(self.inf_b.vel.units) for n in range(3)],
'a': [self.inf_b.acc[ind][n]*ur(self.inf_b.acc.units) for n in range(3)],
'm' : self.inf_b.mass.c() * ur(self.inf_b.mass.units),
'r' : self.inf_b.radius.c() * ur(self.inf_b.radius.units),
'k' : ke(self.inf_b) * ur.joule,
'pe' : period(self.focus_body) * ur.s}
print(dets)
for key,quant in dets.items():
dets[key] = self.convert_to_str(quant, simple=self.change_units)
else:
dets = {'po' : f'({self.inf_b.pos[ind]}) m','v' : f'({self.inf_b.vel[ind]}) ms^-1',
'a' : f'({self.inf_b.acc[ind]}) ms^-2','m' : f'{self.inf_b.mass.c()} kg',
'r' : f'{self.inf_b.radius.c()} m','k' : f'{ke(self.inf_b)} J',
'pe' : f'{period(self.focus_body)} s'}
_str = f"{ident}\npos:{dets['po']}\nvel:{dets['v']}\
\nacc:{dets['a']}\nmass:{dets['m']}\nbody radius:{dets['r']}\nKE:{dets['k']}\
\nest period:{dets['pe']}"
self.ax.text2D(0.05, 0.2 ,_str , size='small', transform=self.ax.transAxes)
def _on_pick(self,event):
print(event.artist)
if isinstance(event.artist,(Line3D, Text)):
if isinstance(event.artist,Line3D):
identity = event.artist.get_label()
else:
identity = event.artist.get_text()
if isinstance(event.artist,Text):
identity = event.artist.get_text()
body = list(b for b in self.engine.bodies if b.identity == identity)[0]
if not body:
return
self.focus_body = body
self.inf_b = body



def start(self,fps=None,frameskip=1,plotskip=1,cache=False):
self._plotskip, inv = plotskip, (1/fps)/1000
elif isinstance(event.artist,Line3D):
identity = event.artist.get_label()
body = list(b for b in self.engine.bodies if b.identity == identity)[0]
if not body:
return
self.focus_body = body


def start(self,fps,frameskip=1,plotskip=1,cache=False, speed_control=False):
self._plotskip, self._fps = plotskip, fps
f = list(frameskip*x for x in range(int(len(self.engine)/frameskip)))
tqdm.write('Starting Visual Environment')
anim = animation.FuncAnimation(self.fig, func=self._animate, interval=inv, frames=f, cache_frame_data=cache)
tqdm.write('Starting Visual Environment')
anim = animation.FuncAnimation(self.fig, func=self._animate, interval=1000/self._fps, frames=f, cache_frame_data=cache)
if speed_control:
self.ax2 = self.fig.add_axes((0.1,0.25,0.05,0.5))
self.speed_slider = Slider(self.ax2,'FPS',0.1,120,valinit=self._fps,orientation='vertical')
self.speed_slider.label.set_color(self.tcolor)
def _sp_ud(val):
anim.event_source.interval = 1000/val
self.speed_slider.on_changed(_sp_ud)
if self.do_picking:
self.fig.canvas.mpl_connect('pick_event',self._on_pick)
plt.show()
Expand Down
Binary file modified solarsystem_bodies.npz
View file
Binary file not shown.

0 comments on commit bdee3dc

Please sign in to comment.