diff --git a/body.txt b/body.txt new file mode 100644 index 0000000..60b999d --- /dev/null +++ b/body.txt @@ -0,0 +1,4 @@ +name = jul +position = (-1.4,0,0) +mass = 1.23 +radius = 60000000 \ No newline at end of file diff --git a/main.py b/main.py index 2e60a70..2d1eb2b 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,7 @@ - - - import nbody.core as nb +from nbody.errors import raise_type_error import nbody.horizons as source phys = nb.Engine(dt=1000) ''' @@ -17,17 +15,13 @@ phys.do_fieldgravity = False phys.simulate(20000) phys.save_as('bodies','solarsystem_bodies') + ''' phys.load_as('bodies','solarsystem_bodies') -sim = nb.mplVisual(phys, 'SS', phys.bodies[0],None, False, - show_grid= True, - show_shadows= False, - show_acceleration = False, - show_velocity= False, - vector_size = 1, - labelling_type = 'legend', - body_model = 'dots', - guistyle = 'dark', - do_picking = True, - show_info = True) -sim.start(frameskip=1000, plotskip=200, speed_control=True, cache=True) + +sim = nb.mplVisual(engine=phys, + name='SS', + focus_body=phys.bodies[0], show_info=True, autoscale=False, frameskip=200, plotskip=200, max_p=1, max_pts=10, cache=False, do_picking=True) #do_picking=True, autoscale=False) + +sim.start() + diff --git a/nbody/__pycache__/core.cpython-312.pyc b/nbody/__pycache__/core.cpython-312.pyc index 383935b..14aaebb 100644 Binary files a/nbody/__pycache__/core.cpython-312.pyc and b/nbody/__pycache__/core.cpython-312.pyc differ diff --git a/nbody/__pycache__/styles.cpython-312.pyc b/nbody/__pycache__/styles.cpython-312.pyc new file mode 100644 index 0000000..8e34ca8 Binary files /dev/null and b/nbody/__pycache__/styles.cpython-312.pyc differ diff --git a/nbody/__pycache__/text.cpython-312.pyc b/nbody/__pycache__/text.cpython-312.pyc index 0bbf7e7..17fdff1 100644 Binary files a/nbody/__pycache__/text.cpython-312.pyc and b/nbody/__pycache__/text.cpython-312.pyc differ diff --git a/nbody/core.py b/nbody/core.py index ad7ed93..74106ad 100644 --- a/nbody/core.py +++ b/nbody/core.py @@ -3,21 +3,22 @@ import math from time import sleep +from matplotlib.lines import Line2D from tqdm import tqdm, trange import numpy as np - +import os # Plotting and animation Packages import matplotlib as mpl #mpl.use('QT5Agg') import matplotlib.pyplot as plt import matplotlib.animation as animation -from matplotlib.widgets import Slider +from matplotlib.widgets import Button, Slider from cycler import cycler from mpl_toolkits.mplot3d import Axes3D from mpl_toolkits.mplot3d.art3d import Line3D from matplotlib.text import Text - +from matplotlib.patches import Circle from . import errors as e @@ -30,16 +31,14 @@ except ModuleNotFoundError: G = 6.6743*10**(-11) - +PICKRADIUS = 10 def sphere(pos,radius,N=20): (c,r) = (pos,radius) u,v = np.mgrid[0:2*np.pi:N*1j, 0:np.pi:N*1j] - x = r*np.cos(u)*np.sin(v)+c[0] y = r*np.sin(u)*np.sin(v)+c[1] z = r*np.cos(v)+c[2] - return x,y,z @@ -47,7 +46,6 @@ class Body: def __init__(self,mass,init_pos,init_vel=(0,0,0), radius=0,bounce=0.999,color=None,identity=None): - if isinstance(identity,str): self.identity = identity elif identity is None: @@ -76,7 +74,7 @@ def __str__(self): mass="{self.mass.c()} {self.mass.units}",\n\ currentpos="{self.pos.c()} {self.pos.units}",\n\ currentvel="{self.vel.c()} {self.vel.units}",\n\ - currentvel="{self.acc.c()} {self.acc.units}")' + currentacc="{self.acc.c()} {self.acc.units}")' def __repr__(self): @@ -121,24 +119,23 @@ def update(self,dt=1,vel_change=None,acc_change=None,vel_next=None): def _reinitialise(self,init_pos=None,init_vel=None): - self.acc = HistoricVector(0,0,0,identity=f'{self.identity}_acc',units_v='ms^-2') + self.acc = HistoricVector(0,0,0,identity=f'{self.identity}_acc',units_v=self.acc.units) if init_pos != None: if isinstance(init_pos,(*Iterable,*VectorType)): - self.pos = HistoricVector(li=init_pos,identity=f'{self.identity}_pos',units_v='m') + self.pos = HistoricVector(li=init_pos,identity=f'{self.identity}_pos',units_v=self.pos.units) else: e.raise_type_error('init_pos',Iterable,init_pos) if init_vel != None: if isinstance(init_vel,(*Iterable,*VectorType)): - self.vel = HistoricVector(li=init_vel,identity=f'{self.identity}_vel',units_v='ms^-1') + self.vel = HistoricVector(li=init_vel,identity=f'{self.identity}_vel',units_v=self.vel.units) else: e.raise_type_error('init_vel',(*Iterable,*VectorType),init_vel) - def get_(self, item, ind, params): - (plotskip, c_mass) = params + def get_(self, item, ind, plotskip, c_mass): def sma(): try: a = max(list(Vector(self.pos[ind]).magnitude() for i in range(0,ind,plotskip))) @@ -159,8 +156,31 @@ def ke(): **dict.fromkeys(['ke', 'kinetic_energy'], ke)} return _get_lookup[item]() - - +def body_from(object): + if isinstance(object, str) and os.path.isfile(object): + params = dict() + with open(object, 'r') as f: + for line in f: + p = list(l.strip() for l in line.strip().split('=')) + if p[0] in ('name', 'identity', 'id'): + params['identity'] = p[1] + elif p[0] in ('color', 'color'): + params['color'] = p[1] + elif p[0] in ('mass', 'radius', 'bounce'): + params[p[0]] = float(p[1]) + elif 'pos' in p[0]: + tp = p[1].strip('()[]').split(',') + params['init_pos'] = list(float(t) for t in tp) + elif 'vel' in p[0]: + tp = p[1].strip('()[]').split(',') + params['init_vel'] = list(float(t) for t in tp) + else: + print(p) + return Body(**params) + if isinstance(object, dict): + return Body(**object) + else: + e.raise_type_error('object', (dict, str), object) class Engine: @@ -235,7 +255,7 @@ def create_plane(self, const_axis ='z', const_val = 0): e.raise_value_error('const_axis,const_val',(str,*NumType),(const_axis,const_val)) def _check_collision(self,body,co_restitution=0): - if self.do_collisions: + if self.do_collisions == True: returned_coll = False for bod in self.bodies: if body != bod: @@ -263,7 +283,7 @@ def _check_collision(self,body,co_restitution=0): def _find_gravity(self,body): res = NullVector() - if self.do_bodygravity: + if self.do_bodygravity == True: for bod2 in self.bodies: if bod2 != body: dist_12 = body.pos - bod2.pos @@ -300,200 +320,269 @@ def simulate(self,intervals): if intervals+1 == len(self.bodies[0].pos): tqdm.write(f'Finished Evaluating {intervals} intervals, ~{len(self.bodies[0].pos)} total intervals.') +# kwargs +''' +kwargs +info_calc +show_acc +show_vel +vector_size +speed_control +start_ind +as dict +linecolor backgroundcolor facecolor textcolor - - - - - -PICKRADIUS = 10 - -_ax_loc = dict(zoom=(0.05,0.25,0.05,0.5), #ax1 - fps=(0.1,0.25,0.05,0.5), #ax2 - subadj=(0,0,1,1,0,0), #fig.subplots - fsize=(16,9), b_as=(1,1,1)) #fig , fig.bbox - -_c = dict(w='white', b='black', cl=(0.,0.,0.,0.),) - - -_artists = dict(dot = dict(zorder=4, clip_on=False,picker=True, marker='o', pickradius=PICKRADIUS), - plane=dict(zorder=1, clip_on=False, color=('xkcd:azure', 0.5)), - trail=dict(zorder=7, clip_on=False, picker=True, pickradius=PICKRADIUS), - surf=dict(zorder=2, clip_on=False, pickradius=PICKRADIUS), - label=dict(zorder=10, clip_on=False), - shadw=dict(zorder=1.5, clip_on=False, color='black'), - vect=dict(zorder=8, clip_on=False)) - -_slider = dict(zoom=dict(label='Zoom', valmin=0.1, valmax=10, orientation='vertical'), - fps=dict(label='Interval',valmax=2,orientation='vertical')) - -_text = dict(info=dict(x=0.05, y=0.2, size='small'), - ax_labels=dict(xlabel='x',ylabel='y',zlabel='z'), - leg=dict(draggable=True, facecolor='black', fancybox=True)) - +''' + class mplVisual: - def __init__(self, engine, name='NBody Simulation (Matplotib)', - focus_body=None,focus_range=None, - 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, info_params = {'output_raw':True, - 'vector_pos':False, 'vector_vel':False, 'vector_acc':False}): - - (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, self.info_params) = 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),(info_params, dict))) # new options here - - self.fig = plt.figure(name, figsize=_ax_loc['fsize']) - self.ax = self.fig.add_subplot(projection='3d') - self.ax1 = self.fig.add_axes(_ax_loc['zoom']) - - self.fig.subplots_adjust(*_ax_loc['subadj']) - self.zoom_slider = Slider(self.ax1,valinit=1,**_slider['zoom']) - (self._frameskip,self._plotskip, self.ax.computed_zorder, self.inf_b) = ( - 1,1,False, self.focus_body) - - def _clearpanes(): - for ax in (self.ax.xaxis,self.ax.yaxis,self.ax.zaxis): - ax.set_pane_color(_c['cl']) - - def _axcolor(color): - for spine in self.ax.get_children(): - if isinstance(spine,mpl.spines.Spine): - spine.set_color(color) - self.ax.tick_params(color=color,labelcolor=color) - for ax in (self.ax.xaxis,self.ax.yaxis,self.ax.zaxis): - ax.label.set_color(color) - mpl.rcParams['axes.labelcolor'] = color - - - if self.guistyle == 'dark': - for artist in (self.fig,self.ax): - artist.set_facecolor(_c['b']) - _clearpanes() - _axcolor(_c['w']) - mpl.rcParams['text.color'] = _c['w'] - self.tcolor = _c['w'] - self.zoom_slider.label.set_color(self.tcolor) - + def __init__(self, + engine, + name='NBody Simulation (Matplotib)', + show_info=False, + show_grid=True, + focus_body=None, + do_picking = False, + fps=30, + step_skip_frames=1, + step_skip_points=1, + max_period=2, + max_pts=None, **kwargs): + self.vct, self.col = dict(), + self.args = { + 'color_dict': dict(line='black', face=(0,0,0,0), bkgd='white', text='black'), + 'speed_control': False, + 'vect_params': dict(vel=False, acc=False, size=1), + 'start_index':0, + 'info_calc':False, + 'focus_range':None, + 'anim_cache':False, + 'show_grid':True, + 'max_pts':None, + 'is_running':True, + 'labelling_type':'legend', + 'body_model':'dots', + 'info_body':focus_body, + 'fmt_params':dict()} + self.files = dict() + self.args.update(kwargs) + if self.args['show_grid'] == False: + self.col['line'] = (0,0,0,0) + self.col['face'] = (0,0,0,0) + + + + self.engine = engine + self.show_info=show_info + self.show_grid=show_grid + self.focus_body=focus_body + self.do_pick=do_picking + + self.plt = dict(ptstep=step_skip_points, + maxperiod=int(max_period), + maxpts=(int(max_pts) if max_pts is not None else None), + interval=1000/fps, + frmstep=step_skip_frames, + major_body=max((b.mass.c() for b in self.engine.bodies))) + flist = list(self.plt['frmstep']*x for x in range(int(len(self.engine)/self.plt['frmstep']))) + self.anim_args = dict(interval=(5 if self.args['speed_control'] is True else 1000/fps), frames=flist, cache_frame_data=self.args['anim_cache']) + + self.trail_data = dict() + for b in self.engine.bodies: + body_data = dict() + for f in flist: + try: + tau = (f-(self.plt['frmstep']*b.get_('period', f, self.plt['ptstep'], self.plt['major_body'])/self.engine.dt)) + if tau > 0: + lower = math.ceil(tau) + else: raise TypeError + except TypeError: + lower = self.args['start_index'] + body_data_f = list(list(float(m) for m in _b.record[lower:f:self.plt['ptstep']]) for _b in (b.pos.X,b.pos.Y,b.pos.Z)) + if self.plt['maxpts'] is not None: + while any(len(i) > self.plt['maxpts'] for i in body_data_f): + for i in body_data_f: + i.pop(0) + body_data[f] = body_data_f + self.trail_data[b] = body_data + + if show_info == True: + self.fmt = Formatter(output_raw=False,items=['identity','mass','radius','energy','period','pos','vel','acc', 'time'], vector_pos=False, vector_vel=False, vector_acc=False, dt=self.engine.dt, c_mass=self.plt['major_body']) + + if self.args['info_calc'] is True: + self.info_data = dict() + for b in self.engine.bodies: + body_data = dict() + for f in flist: + self.fmt.target = [b,f] + body_data[f] = str(self.fmt) + self.info_data[b] = body_data else: - self.tcolor = _c['b'] + self.info_data = None + + self.fig = plt.figure(name, figsize=(16,9)) - if not self.show_grid: + self.ax = self.fig.add_subplot(projection='3d', computed_zorder = False) + + self.zoom_ax = self.fig.add_axes((0.05,0.25,0.05,0.5)) + self.zoom_slider = Slider(self.zoom_ax,valinit=1,label='Zoom',valmin=0.1,valmax=10, orientation='vertical') + + self.plpa_ax = self.fig.add_axes((0.48,0.05,0.04, 0.05)) + self.playpause = Button(self.plpa_ax, label=' ▶ ▐▐ ', hovercolor='white', color=(0.5,0.5,0.5,0.5)) + self.playpause.label.set(fontstretch=0.6, fontsize='large', color='black') + self.playpause.on_clicked(self._toggle) + + if self.args['speed_control'] == True: + self.spd_ax = self.fig.add_axes((0.1,0.25,0.05,0.5)) + self.speed_slider = Slider(self.spd_ax,valinit=1000/self.plt['interval'], valmax=1000/5,label='Target\nFPS',valmin=0.1,orientation='vertical') + self.speed_slider.label.set_color(self.col['t']) + def _sp_ud(val): + self.plt['interval'] = 1000/val + self.speed_slider.on_changed(_sp_ud) + + self.fig.subplots_adjust(0,0,1,1,0,0) + self.ax.tick_params(color=self.col['line'],labelcolor=self.col['text']) + for artist in (self.fig,self.ax): + artist.set_facecolor(self.col['bkgd']) + for spine in self.ax.get_children(): + if isinstance(spine,mpl.spines.Spine): + spine.set_color(self.col['line']) + mpl.rcParams['text.color'] = self.col['text'] + for ax in (self.ax.xaxis,self.ax.yaxis, self.ax.zaxis): + ax.label.set_color((self.col['text'] if show_grid is True else (0,0,0,0))) + ax.set_pane_color(self.col['face']) + if show_grid == False: self.ax.grid(False) - _axcolor(_c['cl']) - _clearpanes() - self._cmass = 0 # changed big mass thing here - for bod in self.engine.bodies: - if bod.mass.c() > self._cmass: - self._cmass = bod.mass.c() - if self.show_info: - self.fm = Formatter(plotskip=1, c_mass=self._cmass, **self.info_params) - + self.ax.set_autoscale_on(False) + + + + + def _draw_info(self, ind): + if self.info_data is None: + self.fmt.target = [self.info_body, ind] + inf_string = str(self.fmt) + else: + inf_string = self.info_data[self.info_body][ind] + self.ax.text2D(s=inf_string, transform=self.ax.transAxes, x=0.05, y=0.2, size='small', horizontalalignment='left', + verticalalignment='bottom', color=self.col['t']) + def _draw_legend(self): + handles, labels = self.ax.get_legend_handles_labels() + handle_list, label_list = [], [] + for handle, label in zip(handles, labels): + if label not in label_list: + handle_list.append(handle) + label_list.append(label) + self.leg = self.ax.legend(handle_list, label_list, draggable=True, facecolor=self.col['b'], fancybox=True, loc=1, fontsize='small') + for text in self.leg.get_texts(): + text.set_picker(PICKRADIUS) + text.set_color(self.col['t']) + def _draw_vectors(self,pos,other,c): - self.ax.quiver(*pos,*other,length=self.vector_size,color=c,**_artists['vect']) + self.ax.quiver(*pos,*other,length=self.vct['size'],color=c,zorder=8, clip_on=False) + + def _animate(self,ind): - sleep(self._int) + if self.args['speed_control'] == True: + sleep((self.plt['interval']-5)/1000) self.ax.clear() + + self.ax.set(xlabel='$x$', + ylabel='$y$', + zlabel='$z$') - # figure building - - self.ax.set(**_text['ax_labels']) - if not self.autoscale: - self.ax.set_box_aspect(_ax_loc['b_as'], zoom=self.zoom_slider.val) - self.ax.set_autoscale_on(False) - - if self.focus_body is not None: - (limx,limy,limz) = (float(m) for m in self.focus_body.pos[ind]) - if self.focus_range is None: - self.focus_range = float(max((max(self.focus_body.pos - bod.pos) for bod in self.engine.bodies))) - else: - limx,limy,limz=0,0,0 - if self.focus_range is None: - self.focus_range = max(max(*bod.pos[ind]) for bod in self._engine.bodies) - - self.ax.set(xlim=((limx-self.focus_range),(limx+self.focus_range)),ylim=((limy-self.focus_range), - (limy+self.focus_range)),zlim=((limz-self.focus_range),(limz+self.focus_range))) + + self.ax.set_box_aspect((1,1,1),zoom=self.zoom_slider.val) + + + if self.focus_body is not None: + (limx,limy,limz) = (float(m) for m in self.focus_body.pos[ind]) + if self.args['focus_range'] is None: + self.args['focus_range'] = float(max((max(abs(x) for x in (self.focus_body.pos - bod.pos)) for bod in self.engine.bodies))) else: - self.ax.set_autoscale_on(True) - self.ax.set_box_aspect(None,zoom=self.zoom_slider.val) + limx,limy,limz=0,0,0 + if self.args['focus_range'] is None: + self.args['focus_range'] = max(max(*bod.pos[ind]) for bod in self.engine.bodies) + rng = self.args['focus_range'] + self.ax.set(xlim=((limx-rng),(limx+rng)), + ylim=((limy-rng),(limy+rng)), + zlim=((limz-rng),(limz+rng))) + + + # plot planes - + for plane in self.engine.planes: xl, yl, zl, = self.ax.get_xlim(), self.ax.get_ylim(), self.ax.get_zlim() pl_const = np.array([[plane[1], plane[1]], [plane[1], plane[1]]]) points = {'x':(pl_const,np.array([[yl[0],yl[1]],[yl[0],yl[1]]]), - np.array([[zl[0],zl[0]],[zl[1],zl[1]]])), - 'y':(np.array([[xl[0],xl[1]],[xl[0],xl[1]]]),pl_const, - np.array([[zl[0],zl[0]],[zl[1],zl[1]]])), - 'z':(np.array([[xl[0],xl[1]],[xl[0],xl[1]]]), - np.array([[yl[0],yl[0]],[yl[1],yl[1]]]),pl_const)} - self.ax.plot_surface(*points[plane[0]], **_artists['plane']) + np.array([[zl[0],zl[0]],[zl[1],zl[1]]])), + 'y':(np.array([[xl[0],xl[1]],[xl[0],xl[1]]]),pl_const, + np.array([[zl[0],zl[0]],[zl[1],zl[1]]])), + 'z':(np.array([[xl[0],xl[1]],[xl[0],xl[1]]]), + np.array([[yl[0],yl[0]],[yl[1],yl[1]]]),pl_const)} + self.ax.plot_surface(*points[plane[0]], zorder=1, clip_on=False, color=('xkcd:azure', 0.5)) # plot bodies for b in self.engine.bodies: - # cutting away any excessive looping orbits for performance - cut = 0 - try: - tau = (ind-(2*b.get_('period', ind, self._plotskip, self._cmass)/self.engine.dt)) - if tau > cut: - cut = math.ceil(tau) - except TypeError: - cut = 0 - # position and trail - _poshist = list(list(float(m) for m in _b.record[cut:ind:self._plotskip]) for _b in (b.pos.X,b.pos.Y,b.pos.Z)) + _poshist = self.trail_data[b][ind] _pos = [float(m) for m in b.pos[ind]] # draw vectors - if self.show_velocity: - self._draw_vectors(_pos,[float(m) for m in b.vel[ind]],'r') - if self.show_acceleration: - self._draw_vectors(_pos,[float(m) for m in b.acc[ind]],'g') + if self.vct['vel'] == True: + self._draw_vectors(_pos, + [float(m) for m in b.vel[ind]], + 'r') + if self.vct['acc'] == True: + self._draw_vectors(_pos, + [float(m) for m in b.acc[ind]], + 'g') # draw trail - self.ax.plot(*_poshist,label=f'{b.identity}',color=b.color, **_artists['trail']) + self.ax.plot(*_poshist, + label=f'{b.identity}', + color=b.color, + zorder=7, + clip_on=False, + picker=True, + pickradius=PICKRADIUS) # draw body - if self.body_model == 'dots': - self.ax.plot(*_pos,label=f'{b.identity}',color=b.color,**_artists['dot']) + if self.args['body_model'] == 'dots': + self.ax.plot(*_pos, + label=f'{b.identity}', + color=b.color, + zorder=4, + clip_on=False, + picker=True, + marker='o', + pickradius=PICKRADIUS) else: _sf = {'wireframe':self.ax.plot_wireframe, 'surface':self.ax.plot_surface} - _sf[self.body_model](*sphere(_pos,b.radius.c()),label=f'{b.identity}', color=b.color **_artists['surf']) + _sf[self.args['body_model']](*sphere(_pos,b.radius.c()), + label=f'{b.identity}', + color=b.color, + zorder=2, + clip_on=False, + pickradius=PICKRADIUS) # draw label - if self.labelling_type == 'label': - self.ax.text(*_pos,b.identity,color=b.color, **_artists['label']) - # draw shadows - if self.show_shadows: - self.ax.plot(*_poshist[0:2],[(_poshist[2][ind]-self.focus_range)]* - len(_poshist[2]), **_artists['shadw']) - # draw legend - if self.labelling_type == 'legend': - handles, labels = self.ax.get_legend_handles_labels() - handle_list, label_list = [], [] - for handle, label in zip(handles, labels): - if label not in label_list: - handle_list.append(handle) - label_list.append(label) - self.leg = self.ax.legend(handle_list, label_list, **_text['leg']) - - - - for text in self.leg.get_texts(): - text.set_picker(PICKRADIUS) - text.set_color(self.tcolor) + if self.gui['label'] == 'label': + for b in self.engine.bodies: + self.ax.text(*[float(m) for m in b.pos[ind]],b.identity, + color=b.color, + zorder=10, + clip_on=False) + else: + self._draw_legend() # draw info panel - if self.show_info: - self.fm.target = [self.inf_b, ind] # new all inside formatter object - self.ax.text2D(s=str(self.fm), transform=self.ax.transAxes, **_text['info']) - # object picking in interactive window + if self.show_info == True: + self._draw_info(ind) + + +# object picking in interactive window def _on_pick(self,event): print(event.artist) if isinstance(event.artist,Text): @@ -501,31 +590,47 @@ def _on_pick(self,event): body = list(b for b in self.engine.bodies if b.identity == identity)[0] if not body: return - self.inf_b = body - elif isinstance(event.artist,Line3D): + self.args['info_body'] = body + elif isinstance(event.artist,(Line3D, Line2D)): 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,interval=0.033,frameskip=1,plotskip=1,cache=False, speed_control=False): - self._plotskip, self.int = plotskip, interval - self.fm.par['ps'] = self._plotskip - f = list(frameskip*x for x in range(int(len(self.engine)/frameskip))) + self.focus_range = None + + def _toggle(self, event): + if self.args['is_running'] == True: + self.anim.event_source.stop() + self.args['is_running'] = False + else: + self.anim.event_source.start() + self.args['is_running'] = True + + + + def start(self, **viewparams): tqdm.write('Starting Visual Environment') - self._int = 0 - anim = animation.FuncAnimation(self.fig, func=self._animate, interval=self.int*1000, frames=f, cache_frame_data=cache, ) - if speed_control: - self.ax2 = self.fig.add_axes(_ax_loc['fps']) - self.speed_slider = Slider(self.ax2,valinit=self.int, valmin=self.int,**_slider['fps']) - self.speed_slider.label.set_color(self.tcolor) - def _sp_ud(val): - self._int = val-self.int - self.speed_slider.on_changed(_sp_ud) - if self.do_picking: + self.anim = animation.FuncAnimation(self.fig, func=self._animate, **self.anim_args) + if viewparams: + self.ax.view_init(**viewparams) + if self.do_pick: self.fig.canvas.mpl_connect('pick_event',self._on_pick) plt.show() + +from .text import Formatter + + + + + + + + + + + + + + -from .text import Formatter \ No newline at end of file diff --git a/nbody/dev/styles.py b/nbody/dev/styles.py new file mode 100644 index 0000000..3f1214f --- /dev/null +++ b/nbody/dev/styles.py @@ -0,0 +1,34 @@ +from mpl_toolkits.mplot3d.art3d import Text3D +from matplotlib.spines import Spine +from mpl_toolkits.mplot3d.axes3d import Axes3D +from mpl_toolkits.mplot3d.axis3d import XAxis, YAxis, ZAxis +from .core import mplVisual +axes_3Ds = (XAxis, YAxis, ZAxis) +class Style: + def __init__(self,gui:mplVisual, line_color, background_color, grid): + gui.fig.subplots_adjust(0,0,1,1,0,0) + if grid['show'] == False: + self.ax.grid(False) + ax_color = (0,0,0,0) + else: + ax_color = line_color + + for artist in (gui.fig, *gui.fig.get_children()): + artist.set_facecolor(background_color) + + for artist in gui.ax.get_children(): + if isinstance(artist, Spine): + artist.set_color(ax_color) + if isinstance(artist, axes_3Ds): + artist.set_pane_color((0.,0.,0.,0.)) + print(f'{artist}:{artist.get_children()}') + + + #mpl.rcParams['text.color'] = self.pcolor + + + def _axcolor(self, color): + self.ax.tick_params(color=color,labelcolor=color) + for ax in (self.ax.xaxis,self.ax.yaxis, self.ax.zaxis): + ax.label.set_color(color) + mpl.rcParams['axes.labelcolor'] = color diff --git a/nbody/dev/test.ipynb b/nbody/dev/test.ipynb new file mode 100644 index 0000000..1fd6601 --- /dev/null +++ b/nbody/dev/test.ipynb @@ -0,0 +1,60 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[['x0', 'x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7', 'x8', 'x9'], ['y0', 'y1', 'y2', 'y3', 'y4', 'y5', 'y6', 'y7', 'y8', 'y9'], ['z0', 'z1', 'z2', 'z3', 'z4', 'z5', 'z6', 'z7', 'z8', 'z9']]\n" + ] + } + ], + "source": [ + "list = [[f'{c}{n}' for n in range(10)] for c in ('x', 'y', 'z')]\n", + "print(list)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[['x0', 'x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7', 'x8', 'x9'], ['y0', 'y1', 'y2', 'y3', 'y4', 'y5', 'y6', 'y7', 'y8', 'y9'], ['z0', 'z1', 'z2', 'z3', 'z4', 'z5', 'z6', 'z7', 'z8', 'z9']]\n" + ] + } + ], + "source": [ + "print(list[n][1])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/nbody/text.py b/nbody/text.py index 42b69e6..6e36dc8 100644 --- a/nbody/text.py +++ b/nbody/text.py @@ -4,11 +4,11 @@ ur = UnitRegistry() -ur.define('light_year = 9460730472580800 * meter = ly') -ur.define('jupiter_mass = 1.898*10**27 * kg = jmass = M_J') -ur.define('solar_mass = 1.98892*10**30 * kg = smass = M_O') -ur.define('earth_mass = 5.9742*10**24 * kg = emass = M_E') -ur.define('astronomical_unit_per_year = astronomical unit / year = au_yr') +ur.define('light_year = 9460730472580800 * meter = ly = lightyear') +ur.define('jupiter_mass = 1.89818*10**27 * kg = $M_J$ = jmass') +ur.define('solar_mass = 1.9885*10**30 * kg = $M_O$ = smass') +ur.define('earth_mass = 5.97219*10**24 * kg = $M_E$ = emass') +ur.define('astronomical_unit_per_year = astronomical unit / year') ### conversion preferences: ## (in order of largest to smallest) @@ -39,77 +39,88 @@ 'period':time_conversions, 'pos':distance_conversions, 'vel':('check_mass', velocity_conversions), - 'acc':('check_mass', acceleration_conversions)} - + 'acc':('check_mass', acceleration_conversions), + 'time':time_conversions} ur.default_format = '~P' class Formatter: - def __init__(self, output_raw = False, vector_pos=True, vector_vel=False, vector_acc = False, plotskip=0, c_mass=0): + def __init__(self, output_raw = False, items=('identity','mass','radius','energy', + 'period','pos','vel','acc','time'), vector_pos=True, vector_vel=False, vector_acc = False, dt=None, plotskip=0, c_mass=0): + self.par = {'raw':output_raw, 'ps': plotskip, 'cm': c_mass} + self.items = items self.target = [None, 0] + self.dt = dt self._m = {'pos':vector_pos, 'vel':vector_vel, 'acc':vector_acc} self.q = {'identity':'','mass':0*ur.kg,'radius':0*ur.m,'energy':0*ur.joule, - 'period':0*ur.s,'pos':0*ur.m/ur.s,'vel':0*ur.m/ur.s,'acc':0*ur.m/ur.s**2} - - def _basetemplate(self): - return "{}\nmass:{}\nradius:{}\nKE:{}\nperiod:{}\npos:{}\nvel:{}\nacc:{}" - - def _quantities(self, body=None): - if body != None: self.target[0] = body - self.q = {'identity' : self.target[0].identity, + 'period':0*ur.s,'pos':0*ur.m,'vel':0*ur.m/ur.s,'acc':0*ur.m/ur.s**2} + def _lookup(self): + return {'identity' : self.target[0].identity, 'mass' : self.target[0].mass.c()*ur(self.target[0].mass.units), 'radius' : self.target[0].radius.c()*ur(self.target[0].radius.units), - 'energy' : self.target[0].get_('ke', self.target[1], (self.par['ps'], self.par['cm'])) * ur.joule, - 'period' : self.target[0].get_('period', self.target[1], (self.par['ps'], self.par['cm'])) * ur.s, + 'energy' : self.target[0].get_('ke', self.target[1], self.par['ps'], self.par['cm']) * ur.joule, + 'period' : self.target[0].get_('period', self.target[1], self.par['ps'], self.par['cm']) * ur.s, 'pos' : ([self.target[0].pos[self.target[1]][n]*ur(self.target[0].pos.units) - for n in range(3)] if self._m['pos'] - else Vector(self.target[0].pos[self.target[1]]).magnitude()*ur(self.target[0].pos.units)), + for n in range(3)] if self._m['pos'] + else Vector(self.target[0].pos[self.target[1]]).magnitude()*ur(self.target[0].pos.units)), 'vel' : ([self.target[0].vel[self.target[1]][n]*ur(self.target[0].vel.units) - for n in range(3)] if self._m['vel'] - else Vector(self.target[0].vel[self.target[1]]).magnitude()*ur(self.target[0].vel.units)), + for n in range(3)] if self._m['vel'] + else Vector(self.target[0].vel[self.target[1]]).magnitude()*ur(self.target[0].vel.units)), 'acc' : ([self.target[0].acc[self.target[1]][n]*ur(self.target[0].acc.units) - for n in range(3)] if self._m['acc'] - else Vector(self.target[0].acc[self.target[1]]).magnitude()*ur(self.target[0].acc.units)) - } + for n in range(3)] if self._m['acc'] + else Vector(self.target[0].acc[self.target[1]]).magnitude()*ur(self.target[0].acc.units)), + 'time' : self.target[1]*self.dt * ur.s, + } + + def _basetemplate(self): + return "{}\nmass:{}\nradius:{}\nKE:{}\nperiod:{}\npos:{}\nvel:{}\nacc:{}" + + def _quantities(self, body=None): + if body != None: self.target[0] = body + q = {} + for key in self.items: + self.q[key] = self._lookup()[key] return self.q + def _get_best_unit(self, quant, units): - if units is None: + if units is None or (isinstance(quant, Quantity) and quant.m=='NaN'): return quant if units[0] == 'check_mass': if self.q['mass'] > 0.001*ur.emass: return quant else: units = units[1] - for _unit in units: - if quant > 1.*_unit: - return quant.to(_unit) - else: - return quant + for _un in units: + if quant > 0.05*_un: + return quant.to(_un) + return quant + def convert(self, arg=None): + conv = {} if not arg: args = self._quantities() else: args = arg for key,value in args.items(): if not isinstance(value, Iterable): - value = self._get_best_unit(value, _units[key]) + conv[key] = self._get_best_unit(value, _units[key]) else: - value = [self._get_best_unit(v, _units[key]) for v in value] - - def quantities_to_strings(self, arg=None): + conv[key] = [self._get_best_unit(v, _units[key]) for v in value] + return conv + def _quantities_to_strings(self, arg=None): strings = [] - if not arg: + if arg == None: args = self._quantities() else: args = arg for key, value in args.items(): try: if isinstance(value, Iterable): - strings.append('('+''.join((f'{(q.to_compact()):.4f~P}' for q in value))+')') + strings.append('('+''.join((f'{q.to_compact():.4f~P}' for q in value))+')') elif isinstance(value, Quantity) and value.m != 'NaN': - strings.append(f'{(value.to_compact()):.4f~P}') + strings.append(f'{value.to_compact():.4f~P}') elif isinstance(value, Quantity) and value.m == 'NaN': strings.append('NaN') else: @@ -119,9 +130,12 @@ def quantities_to_strings(self, arg=None): return strings def __str__(self): - if not self.par['raw']: + if self.par['raw'] == False: return self._basetemplate().format(*self._quantities_to_strings(self.convert())) else: - return self._basetemplate().format(*self.quantities_to_strings()) + return self._basetemplate().format(*self._quantities_to_strings()) + from .core import Body + + diff --git a/solarsystem_bodies.npz b/solarsystem_bodies.npz index c6b5b42..b2a004f 100644 Binary files a/solarsystem_bodies.npz and b/solarsystem_bodies.npz differ