diff --git a/.github/pyproject.toml b/.github/pyproject.toml index 3d182e4..c96ce0c 100644 --- a/.github/pyproject.toml +++ b/.github/pyproject.toml @@ -11,5 +11,5 @@ ignore = [ ] [tool.ruff.lint.per-file-ignores] -"__init__.py" = ["E402"] # ignore module imports -"**/tests/*" = ["E402"] +"**__init__.py" = ["F401", "F402"] # ignore module imports +"**/tests/*" = ["F401", "F402"] diff --git a/main.py b/main.py index faa6b65..2fa7db5 100644 --- a/main.py +++ b/main.py @@ -1,14 +1,12 @@ - -import cProfile -import pstats import nbody as nb -from line_profiler import profile + LINE_PROFILE=1 bodies = nb.horizons_batch(('10','199','299','399','499','599','699','799','899')) -for i,color in enumerate(['#F28322','#BFBEBD', '#D9B391', '#63BAA6', '#F27A5E', '#BFAE99', '#D9B779', '#95BBBF', '#789EBF']): +for i,color in enumerate(['#F28322','#BFBEBD', '#D9B391', '#63BAA6', '#F27A5E', + '#BFAE99', '#D9B779', '#95BBBF', '#789EBF']): bodies[i].color = color phys = nb.Engine(dt=1000) phys.attach_bodies(bodies) @@ -23,7 +21,8 @@ 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) + 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) ''' diff --git a/nbody/__init__.py b/nbody/__init__.py index 0182582..37417e6 100644 --- a/nbody/__init__.py +++ b/nbody/__init__.py @@ -1,5 +1,5 @@ # NBody __init__.py -''' +""" File Structure: ⸺ nbody/ ├── core/ @@ -43,10 +43,9 @@ .core.visual .tools.io .tools.horizons -''' +""" from .core.body import Body from .core.engine import Engine -from .core.visual import mplVisual +from .core.visual import MPLVisual from .tools.horizons import horizons_batch, horizons_query from .tools.io import obj_from, export_obj - diff --git a/nbody/build/config.py b/nbody/build/config.py deleted file mode 100644 index 49a04dc..0000000 --- a/nbody/build/config.py +++ /dev/null @@ -1,3 +0,0 @@ - -# Arithimetic Calculation -USE_FLOAT = False # whether to override number types to float, useful when gmpy2 is not available. diff --git a/nbody/build/gmpy-gmpy2-2.2.0a1.tar.gz b/nbody/build/gmpy-gmpy2-2.2.0a1.tar.gz deleted file mode 100644 index 648a554..0000000 Binary files a/nbody/build/gmpy-gmpy2-2.2.0a1.tar.gz and /dev/null differ diff --git a/nbody/build/gmpy2-2.0.8-cp27-cp27m-win32.whl b/nbody/build/gmpy2-2.0.8-cp27-cp27m-win32.whl deleted file mode 100644 index de47a35..0000000 Binary files a/nbody/build/gmpy2-2.0.8-cp27-cp27m-win32.whl and /dev/null differ diff --git a/nbody/build/gmpy2-2.0.8-cp27-cp27m-win_amd64.whl b/nbody/build/gmpy2-2.0.8-cp27-cp27m-win_amd64.whl deleted file mode 100644 index 20f8fde..0000000 Binary files a/nbody/build/gmpy2-2.0.8-cp27-cp27m-win_amd64.whl and /dev/null differ diff --git a/nbody/build/gmpy2-2.1.5-cp311-cp311-win_amd64.whl b/nbody/build/gmpy2-2.1.5-cp311-cp311-win_amd64.whl deleted file mode 100644 index 1b83b54..0000000 Binary files a/nbody/build/gmpy2-2.1.5-cp311-cp311-win_amd64.whl and /dev/null differ diff --git a/nbody/core/base.py b/nbody/core/base.py index 53b6e06..9d0db39 100644 --- a/nbody/core/base.py +++ b/nbody/core/base.py @@ -1,10 +1,10 @@ from __future__ import annotations # Python Builtins -from numpy import ndarray + from numbers import Number from collections.abc import Iterable from mpmath import mp, fp -import gmpy2 as g + # Local error definitions. from ..tools import errors as e @@ -16,13 +16,13 @@ mp.pretty, fp.pretty = True, True -def _O(obj): +def _O(obj): # noqa: N802 if isinstance(obj,(Vector, Variable, HistoricVariable, HistoricVector)): return obj.c() else: return obj -def _V(obj): +def _V(obj): # noqa: N802 if isinstance(obj, Iterable): if len(obj) == 3: return Vector(li=obj) @@ -60,13 +60,14 @@ def typecheck(argtype): class Variable: def __init__(self,init_var,identity='Variable',units=None): - (self.record,self.identity,self.units) = typecheck(((_ntype(init_var)(init_var),NumType),(identity,str),(units,(str,NoneType)))) + (self.record,self.identity,self.units) = typecheck(((_ntype(init_var)(init_var),NumType), + (identity,str),(units,(str,NoneType)))) self.units = (units if units else '') def c(self): return self.record - def next(self, val): + def next(self, val): # noqa: A003 self.record = _O(val) def __len__(self): @@ -184,7 +185,7 @@ def c(self): return self.record[-1] - def next(self,next_val): + def next(self,next_val): # noqa: A003 temp = _O(next_val) if isinstance(temp,NumType): self.record.append(self.type(temp)) @@ -339,7 +340,8 @@ def __init__(self,x=None,y=None,z=None,li=None,identity=None,units_v=None): li = ((x,y,z) if (isinstance(x,NumType) and isinstance(y,NumType) and isinstance(z,NumType)) else _O(li)) nt = _ntype(*li) if isinstance(li,(tuple,list)) and len(li) == 3: - (self.X,self.Y,self.Z) = list(HistoricVariable(vl,f'{self.identity}_{i}',self.units) for (vl,i) in ((nt(li[0]),'x'),(nt(li[1]),'y'),(nt(li[2]),'z'))) + (self.X,self.Y,self.Z) = list(HistoricVariable(vl,f'{self.identity}_{i}',self.units) for (vl,i) in + ((nt(li[0]),'x'),(nt(li[1]),'y'),(nt(li[2]),'z'))) else: e.raise_type_error('l,x,y,z',(Iterable,*NumType),(li,x,y,z)) self.type = nt @@ -359,7 +361,7 @@ def c(self, usage=None): except KeyError: e.raise_out_of_range('c()',usage) - def next(self,next_vals): + def next(self,next_vals): # noqa: A003 temp = _O(next_vals) if isinstance(temp,Iterable): diff --git a/nbody/core/body.py b/nbody/core/body.py index 4e39a7f..d8091eb 100644 --- a/nbody/core/body.py +++ b/nbody/core/body.py @@ -1,10 +1,12 @@ import math from ..tools import errors as e -from .base import NullVector, typecheck, _O, _V, Iterable, NumType, NoneType, HistoricVector, Variable, Vector, VectorType +from .base import (typecheck, _O, _V, + Iterable, NumType, NoneType, VectorType, + HistoricVector, Variable, Vector, NullVector) try: from scipy.constants import G -except: +except ModuleNotFoundError: G = 6.6743*10**(-11) class Body: @@ -51,17 +53,15 @@ def __getitem__(self, ind): return (self.pos[ind], self.vel[ind], self.acc[ind]) elif isinstance(ind, str): return {'pos':self.pos['all'], 'vel': self.vel['all'], 'acc':self.acc['all'], - 'current':list(d.c() for d in (self.pos, self.vel, self.acc)),'info':{'identity':self.identity, 'mass':self.mass, - 'radius':self.radius, 'color':self.color, 'bounce':self.bounce},'x':list(d.X for d in (self.pos, self.vel, self.acc)), - 'y':list(d.Y for d in (self.pos, self.vel, self.acc)),'z':list(d.Z for d in (self.pos, self.vel, self.acc)), - '_data_':{'pos':self.pos['all'], 'vel': self.vel['all'], 'acc':self.acc['all'],'identity':self.identity, 'mass':self.mass, - 'radius':self.radius, 'color':self.color, 'bounce':self.bounce}}[ind] + 'current':list(d.c() for d in (self.pos, self.vel, self.acc)),'info':{'identity':self.identity, 'mass':self.mass, + 'radius':self.radius, 'color':self.color, 'bounce':self.bounce},'x':list(d.X for d in + (self.pos, self.vel, self.acc)), + 'y':list(d.Y for d in (self.pos, self.vel, self.acc)),'z':list(d.Z for d in (self.pos, self.vel, self.acc)), + '_data_':{'pos':self.pos['all'], 'vel': self.vel['all'], 'acc':self.acc['all'],'identity':self.identity, + 'mass':self.mass, 'radius':self.radius, 'color':self.color, 'bounce':self.bounce}}[ind] def update(self,dt=1,vel_change=None,acc_change=None,vel_next=None): - if vel_next: - vel = _V(vel_next) - else: - vel = self.vel + vel = _V(vel_next) if vel_next else self.vel if acc_change is not None: acc_change = _V(acc_change) diff --git a/nbody/core/engine.py b/nbody/core/engine.py index a2c9720..915cc6e 100644 --- a/nbody/core/engine.py +++ b/nbody/core/engine.py @@ -29,7 +29,7 @@ def attach_bodies(self, new_bodies): def _loadeng(self,eng): # used for loading an engine from an npz file. - self = eng + self = eng #noqa: F841 def save_as(self,dump='engine',file_name='nbody_data'): @@ -40,11 +40,11 @@ def save_as(self,dump='engine',file_name='nbody_data'): def load_as(self,objects='engine',file_name='nbody_data'): - _loadobjs = {'bodies':("self.attach_bodies(objs['bodies'])",), - 'engine':("self._loadeng(objs['engine'])",)} + _loadobjs = {'bodies':("self.attach_bodies(_objs['bodies'])",), + 'engine':("self._loadeng(_objs['engine'])",)} with open(f'{file_name}.npz','rb') as file: - objs = np.load(file,allow_pickle=True) + _objs = np.load(file,allow_pickle=True) for func in _loadobjs[objects]: try: # usage limited to values in dict above. @@ -93,17 +93,19 @@ def _check_collision(self,body,co_restitution=0): # 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))) - #calc vel change: dv = n^ * [n^*[v2 -v1]*[1+e]*m"]/m1, n^ is normal unit vector, m" is effective mass as above. + # 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) 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. + # 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 = (True if body_dists[0]/body.radius <= 1.01 and body_dists[-1]/body.radius <= 1.01 else False) + on_plane = (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 @@ -136,7 +138,8 @@ def _gen_sim(self): 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'): + for _ in trange(intervals, + desc=f'«Engine» → Evaluating motion for each interval of {self.dt} seconds', unit='ints'): for body, col_vel, on_plane, acc_g in self._gen_sim(): @@ -147,7 +150,9 @@ def simulate(self,intervals): body.update(dt=self.dt,vel_next=(col_vel+fieldvel).c(),acc_change=acc_g) if intervals+1 == len(self.bodies[0].pos): - tqdm.write(f'«Engine» → Finished Evaluating {intervals} intervals, ~{len(self.bodies[0].pos)} total intervals.') + tqdm.write( + f'«Engine» → Finished Evaluating {intervals} intervals, ~{len(self.bodies[0].pos)} total intervals.', + ) def barycenter(self, index): mass_dist = NullVector() @@ -166,7 +171,7 @@ def __getitem__(self, ind): 'dofg':self.do_fieldgravity, 'planes':self.planes, 'fields':self.fields, - 'bodies':self.bodies + 'bodies':self.bodies, } \ No newline at end of file diff --git a/nbody/core/visual.py b/nbody/core/visual.py index afee88f..55a42b3 100644 --- a/nbody/core/visual.py +++ b/nbody/core/visual.py @@ -23,10 +23,10 @@ PICKRADIUS = 10 -def sphere(pos,radius,N=20): +def sphere(pos,radius,n=20): (c,r) = (pos,radius) # get a cubic mesh of points - u,v = np.mgrid[0:2*np.pi:N*1j, 0:np.pi:N*1j] + u,v = np.mgrid[0:2*np.pi:n*1j, 0:np.pi:n*1j] # compute x,y,z values from spherical mesh x = r*np.cos(u)*np.sin(v)+c[0] y = r*np.sin(u)*np.sin(v)+c[1] @@ -34,7 +34,7 @@ def sphere(pos,radius,N=20): return x,y,z -class mplVisual: +class MPLVisual: def __init__(self, engine, name='NBody Simulation (Matplotib)', @@ -56,12 +56,14 @@ def __init__(self, 'labelling_type':'legend', 'body_model':'dots', 'info_body':focus_body, - 'fmt_params':dict(), + 'fmt_params':dict(output_raw=False, + items=['identity','mass','radius','energy','period','pos','vel','acc', 'time'], + vector_pos=False, vector_vel=False, vector_acc=False), 'file':None, 'step_skip_frames':1, 'step_skip_points':1, 'max_period':2, - 'fps':30 + 'fps':30, } self.files = dict() self.engine = engine @@ -89,21 +91,27 @@ def __init__(self, # list of data intervals to create frames for flist = list(self.plt['frmstep']*x for x in range(int(len(self.engine)/self.plt['frmstep']))) # organise args for animation function. - self.anim_args = dict(interval=(5 if self.args['speed_control'] is True else 1000/self.args['fps']), frames=flist, cache_frame_data=self.args['anim_cache']) + self.anim_args = dict(interval=(5 if self.args['speed_control'] is True else 1000/self.args['fps']), + frames=flist, cache_frame_data=self.args['anim_cache']) # build data for trails self.trail_data = dict() - with tqdm(total = len(self.engine.bodies)*len(flist), desc='«mplVisual» → Building Trails', unit='items') as tbar: + with tqdm(total = len(self.engine.bodies)*len(flist), + desc='«mplVisual» → Building Trails', unit='items') as tbar: 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'], engine=self.engine)/self.engine.dt)) + tau = (f-(self.plt['frmstep']*b.get_('period', f, self.plt['ptstep'], + self.plt['major_body'],engine=self.engine)/self.engine.dt)) if tau > 0: lower = math.ceil(tau) - else: raise TypeError + 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)) + 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: @@ -113,11 +121,13 @@ def __init__(self, self.trail_data[b] = body_data # init a formatter to manage the info readout 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, engine=self.engine,plotskip=self.plt['ptstep'], c_mass=self.plt['major_body']) + self.fmt = Formatter(engine=self.engine,plotskip=self.plt['ptstep'], + c_mass=self.plt['major_body'], **self.args['fmt_params']) # if true precalculate all info readouts for each frame and body. if self.args['info_calc'] is True: self.info_data = dict() - with tqdm(total = len(self.engine.bodies)*len(flist), desc='«mplVisual» → Precomputing All Descriptions', unit='items') as pbar: + with tqdm(total = len(self.engine.bodies)*len(flist), + desc='«mplVisual» → Precomputing All Descriptions', unit='items') as pbar: for b in self.engine.bodies: body_data = dict() for f in flist: @@ -142,7 +152,8 @@ def __init__(self, # axes to contain speed control slider 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 = 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.args['color_dict']['text']) def _sp_ud(val): self.plt['interval'] = 1000/val @@ -175,7 +186,8 @@ def _draw_info(self, ind): else: # or read from precalculated. inf_string = self.info_data[self.args['info_body']][ind] - self.ax.text2D(s=inf_string, transform=self.ax.transAxes, x=0.05, y=0.2, size='small', horizontalalignment='left', + self.ax.text2D(s=inf_string, transform=self.ax.transAxes, + x=0.05, y=0.2, size='small', horizontalalignment='left', verticalalignment='bottom', color=self.args['color_dict']['text']) def _draw_legend(self): handles, labels = self.ax.get_legend_handles_labels() @@ -185,7 +197,11 @@ def _draw_legend(self): 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.args['color_dict']['bkgd'], fancybox=True, loc=1, fontsize='small') + self.leg = self.ax.legend(handle_list, + label_list, + draggable=True, + facecolor=self.args['color_dict']['bkgd'], + fancybox=True, loc=1, fontsize='small') for text in self.leg.get_texts(): text.set_picker(PICKRADIUS) text.set_color(self.args['color_dict']['text']) @@ -213,7 +229,8 @@ def _animate(self,ind): # set plot range to max distances from focus body (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))) + 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: limx,limy,limz=0,0,0 if self.args['focus_range'] is None: diff --git a/nbody/tests/test_base.py b/nbody/tests/test_base.py index 1c4b94a..f7f293e 100644 --- a/nbody/tests/test_base.py +++ b/nbody/tests/test_base.py @@ -20,9 +20,9 @@ def test_typecheck(self): with pytest.raises(TypeError): nb.typecheck(((123, str))) assert nb.typecheck(((2, int))) == 2 - def test_O(self): + def test_o(self): assert nb._O(nb.Variable(1)) == 1 and nb._O(nb.NullVector()) == (0,0,0) - def test_V(self): + def test_v(self): with pytest.raises((IndexError, TypeError)): nb._V([0]) assert nb._V(0) == 0 and nb._V((0,0,0)).c() == nb.Vector((0,0,0)).c() @@ -32,21 +32,21 @@ class TestVarVectorBasics: '''Testing __init__, length and type definitions ''' def test_variable(self): - var = nb.Variable(0) + _var = nb.Variable(0) def test_histvariable(self): - hvar = nb.HistoricVariable(0) + _hvar = nb.HistoricVariable(0) def test_vector(self): - vect = nb.Vector((0,0,0)) - vect2 = nb.Vector(x=0, y=0, z=0) + _vect = nb.Vector((0,0,0)) + _vect2 = nb.Vector(x=0, y=0, z=0) def test_histvector(self): - hvect = nb.HistoricVector(0,0,0) - hvect2 = nb.HistoricVector(li=(0,0,0)) + _hvect = nb.HistoricVector(0,0,0) + _hvect2 = nb.HistoricVector(li=(0,0,0)) def test_len(self): - varL = len(nb.Variable(0)) - hvarL = len(_hvar) - vectL = len(nb.Vector((0,0,0))) - hvectL = len(_hvect) - assert varL == vectL and hvarL == hvectL + var_l = len(nb.Variable(0)) + hvar_l = len(_hvar) + vect_l = len(nb.Vector((0,0,0))) + hvect_l = len(_hvect) + assert var_l == vect_l and hvar_l == hvect_l def test_types(self): assert (isinstance(nb.Variable(0), nb.VarType) and isinstance(nb.HistoricVariable(0), nb.VarType) and @@ -103,7 +103,9 @@ def test_var_to_hvar(self): and hvar.record[1] == 0. and nb.HistoricVariable(1.) + nb.Variable(-1.) == 0.) def test_full(self): - n = {'p':nb.HistoricVector(0.,0.,0.), 'v':nb.HistoricVector(2.,-0.5,0.), 'a':nb.HistoricVector(0.01,-0.03,0.5)} + n = {'p':nb.HistoricVector(0.,0.,0.), + 'v':nb.HistoricVector(2.,-0.5,0.), + 'a':nb.HistoricVector(0.01,-0.03,0.5)} time = nb.HistoricVariable(0.) for i in range(1,10): time.next(i*abs(n['v'].c(1))) diff --git a/nbody/tests/test_tools.py b/nbody/tests/test_tools.py index 364d9df..729dc6e 100644 --- a/nbody/tests/test_tools.py +++ b/nbody/tests/test_tools.py @@ -1,7 +1,11 @@ import pytest + +import shutil +import os + from .. import export_obj, obj_from, Engine, Body, mplVisual from ..tools.horizons import horizons_batch, horizons_query -import shutil, os + def cleanup(folder): if os.path.exists(folder): shutil.rmtree(folder) @@ -11,12 +15,12 @@ def cleanup(folder): 'body':'test_body.txt', 'bodyext':'test_bodyext.txt', 'eng':'test_engine.txt', - 'engrun':'test_enginerun.txt' + 'engrun':'test_enginerun.txt', } tloc = { 'body':['tmp_body','.txt'], 'bodyext':['tmp_bodyext','.txt'], - 'eng':['tmp_engine','.txt'] + 'eng':['tmp_engine','.txt'], } queries = ('10','199','299','399','499','599','699','799','899') diff --git a/nbody/tools/formatter.py b/nbody/tools/formatter.py index d8d3cdb..0431d30 100644 --- a/nbody/tools/formatter.py +++ b/nbody/tools/formatter.py @@ -1,6 +1,6 @@ from pint import Quantity, UnitRegistry -from ..core.base import Vector, Iterable, _V, _O +from ..core.base import Vector, Iterable, _O ur = UnitRegistry() Q_ = ur.Quantity ur.define('light_year = 9460730472580800 * meter = ly = lightyear') @@ -44,7 +44,8 @@ class Formatter: 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, engine=None, plotskip=0, c_mass=0): + 'period','pos','vel','acc','time'), + vector_pos=True, vector_vel=False, vector_acc = False, engine=None, plotskip=0, c_mass=0): # sort parameters into useful form self.par = {'raw':output_raw, 'ps': plotskip, 'cm': c_mass} self.items = items @@ -58,21 +59,23 @@ def __init__(self, output_raw = False, items=('identity','mass','radius','energy def _lookup(self): # dict containing pint quantities of target - return {'identity' : self.target[0].identity, - 'mass' : Q_(float(self.target[0].mass.c()), ur(self.target[0].mass.units)), - 'radius' : Q_(float(self.target[0].radius.c()), ur(self.target[0].radius.units)), - 'energy' :Q_(float(_O(self.target[0].get_('ke', self.target[1], self.par['ps'], self.par['cm']))), ur.joule), - 'period' : Q_(float(_O(self.target[0].get_('period', self.target[1], self.par['ps'], self.par['cm'], engine=self.engine))), ur.s), - 'pos' : ([Q_(float(self.target[0].pos[self.target[1]][n]),ur(self.target[0].pos.units)) - for n in range(3)] if self._m['pos'] - else Q_(float(_O(Vector(self.target[0].pos[self.target[1]]).magnitude())), ur(self.target[0].pos.units))), - 'vel' : ([Q_(float(self.target[0].vel[self.target[1]][n]), ur(self.target[0].vel.units)) - for n in range(3)] if self._m['vel'] - else Q_(_O(float(Vector(self.target[0].vel[self.target[1]]).magnitude())), ur(self.target[0].vel.units))), - 'acc' : ([Q_(float(self.target[0].acc[self.target[1]][n]), ur(self.target[0].acc.units)) - for n in range(3)] if self._m['acc'] - else Q_(float(_O(Vector(self.target[0].acc[self.target[1]]).magnitude())), ur(self.target[0].acc.units))), - 'time' : Q_(self.target[1]*self.engine.dt, ur.s), + return { + 'identity' : self.target[0].identity, + 'mass' : Q_(float(self.target[0].mass.c()), ur(self.target[0].mass.units)), + 'radius' : Q_(float(self.target[0].radius.c()), ur(self.target[0].radius.units)), + 'energy' :Q_(float(_O(self.target[0].get_('ke', self.target[1], self.par['ps'], self.par['cm']))), ur.joule), + 'period' : Q_(float(_O(self.target[0].get_('period', self.target[1], self.par['ps'], self.par['cm'], + engine=self.engine))), ur.s), + 'pos' : ([Q_(float(self.target[0].pos[self.target[1]][n]),ur(self.target[0].pos.units)) + for n in range(3)] if self._m['pos'] + else Q_(float(_O(Vector(self.target[0].pos[self.target[1]]).magnitude())), ur(self.target[0].pos.units))), + 'vel' : ([Q_(float(self.target[0].vel[self.target[1]][n]), ur(self.target[0].vel.units)) + for n in range(3)] if self._m['vel'] + else Q_(_O(float(Vector(self.target[0].vel[self.target[1]]).magnitude())), ur(self.target[0].vel.units))), + 'acc' : ([Q_(float(self.target[0].acc[self.target[1]][n]), ur(self.target[0].acc.units)) + for n in range(3)] if self._m['acc'] + else Q_(float(_O(Vector(self.target[0].acc[self.target[1]]).magnitude())), ur(self.target[0].acc.units))), + 'time' : Q_(self.target[1]*self.engine.dt, ur.s), } def _basetemplate(self): @@ -114,10 +117,7 @@ def _get_best_unit(self, quant, units): def convert(self, arg=None): conv = {} - if not arg: - args = self._quantities() - else: - args = arg + args = arg if arg else self._quantities() # get best units for each item we want for key,value in args.items(): if not isinstance(value, Iterable): @@ -130,10 +130,7 @@ def convert(self, arg=None): def _quantities_to_strings(self, arg=None): strings = {} - if arg == None: - args = self._quantities() - else: - args = arg + args = self._quantities() if arg == None else arg for key, value in args.items(): try: if isinstance(value, Iterable): diff --git a/nbody/tools/io.py b/nbody/tools/io.py index 4187a97..1cccdc0 100644 --- a/nbody/tools/io.py +++ b/nbody/tools/io.py @@ -14,7 +14,7 @@ def tup(iterable): return str(tuple(str(c) for c in tmp)).replace("'", "") -def obj_from(object, obj='engine'): +def obj_from(object, obj='engine'): # noqa: A002 if isinstance(object, str): if os.path.isfile(object): with open(object, 'r') as f: @@ -41,9 +41,8 @@ def obj_from(object, obj='engine'): key = line.strip('[\n').strip(']').strip('*').strip() body_strings[key] = i for s_i,s_line in enumerate(_input): - if s_i > i: - if ']' in s_line and isinstance(body_strings[key], NumType): - body_strings[key] = _input[i:s_i] + if s_i > i and ']' in s_line and isinstance(body_strings[key], NumType): + body_strings[key] = _input[i:s_i] # callables elif line.startswith('-*'): val = [p.strip() for p in line.strip('-*').strip('\n').split('=')] @@ -57,7 +56,7 @@ def obj_from(object, obj='engine'): for key,value in body_strings.items(): params = dict() for line in value: - p = list(l.strip() for l in line.strip().split('=')) + p = list(part.strip() for part in line.strip().split('=')) if p[0] in ('name', 'identity', 'id'): params['identity'] = p[1] elif p[0] in ('color', 'colour'): @@ -75,7 +74,9 @@ def obj_from(object, obj='engine'): elif '[' in p[0] or ']' in p[0] or '*' in p[0] or '#' in p[0]: pass else: - tqdm.write(f'«obj_from()» → [!] Couldn\'t parse "{line}" on line {i} so it has been skipped.') + tqdm.write( + f'«obj_from()» → [!] Couldn\'t parse "{line}" on line {i} so it has been skipped.', + ) print(params) bodies[key] = Body(**params) # init engine and add bodies @@ -99,13 +100,15 @@ def obj_from(object, obj='engine'): elif val[0] == 'sim' and len(val) == 2: engine.simulate(int(val[1])) else: - tqdm.write(f'«obj_from()» → [!] Couldn\'t parse function "{' '.join(val)}" so it has been skipped.') + tqdm.write( + f'«obj_from()» → [!] Couldn\'t parse function "{' '.join(val)}" so it has been skipped.', + ) return engine elif obj == 'body': # get body info and init params = {} for line in _input: - p = list(l.strip() for l in line.strip().split('=')) + p = list(part.strip() for part in line.strip().split('=')) if p[0] in ('name', 'identity', 'id'): params['identity'] = p[1] elif p[0] in ('color', 'colour'): @@ -138,7 +141,7 @@ def obj_from(object, obj='engine'): if key in ('name', 'identity', 'id'): result['identity'] = value elif key in ('color', 'colour'): - result['color'] = (value if value is not 'None' else None) + result['color'] = (value if value != 'None' else None) elif key in ('mass', 'radius'): result[key] = mpf(value) elif key == 'bounce': @@ -180,16 +183,12 @@ def obj_from(object, obj='engine'): else: e.raise_type_error('object', str, object) -def export_obj(object, loc, overwritefile=True, info_name=None, csvs=True, direxists=True): +def export_obj(object, loc, overwritefile=True, info_name=None, csvs=True, direxists=True): # noqa: A002 # make file directory if isinstance(object,(Engine, Body)): os.makedirs(loc, exist_ok=direxists) if isinstance(object, Engine): - if info_name == None: - fname = 'eng_info' - else: - # create eng specific info string - fname = info_name + fname = "eng_info" if info_name == None else info_name eng_info = [ f'dt = {object.dt}\n', f'checking_range = {object._rangechk}\n', @@ -197,7 +196,7 @@ def export_obj(object, loc, overwritefile=True, info_name=None, csvs=True, direx f'# bodies:({len(object.bodies)})\n', f'~ do_collisions = {object.do_collisions}\n', f'~ do_bodygravity = {object.do_bodygravity}\n', - f'~ do_fieldgravity = {object.do_fieldgravity}\n' + f'~ do_fieldgravity = {object.do_fieldgravity}\n', ] ids = [] # append info for each object @@ -234,15 +233,12 @@ def export_obj(object, loc, overwritefile=True, info_name=None, csvs=True, direx 'velZ':b.vel.Z.record, 'accX':b.acc.X.record, 'accY':b.acc.Y.record, - 'accZ':b.acc.Z.record + 'accZ':b.acc.Z.record, } d_df = pd.DataFrame.from_dict(info_dict, 'index').transpose() - d_df.to_csv(f'{loc}/{ids[i]}.csv', sep=',',index=False, mode=('w' if overwritefile is True else 'x')) + d_df.to_csv(f'{loc}/{ids[i]}.csv',sep=',',index=False, mode=('w' if overwritefile is True else 'x')) if isinstance(object, Body): - if info_name == None: - fname = 'body_info' - else: - fname = info_name + fname = "body_info" if info_name == None else info_name # make body info string obj_id = str(object.identity).replace(' ','_').split(':')[0].split('(Static)')[0] bod_info = [f'*{obj_id} [\n', @@ -267,7 +263,7 @@ def export_obj(object, loc, overwritefile=True, info_name=None, csvs=True, direx 'velZ':object.vel.Z.record, 'accX':object.acc.X.record, 'accY':object.acc.Y.record, - 'accZ':object.acc.Z.record + 'accZ':object.acc.Z.record, } d_df = pd.DataFrame.from_dict(info_dict, 'index').transpose() d_df.to_csv(f'{loc}/{fname}.csv', sep=',',index=False, mode=('w' if overwritefile is True else 'x'))