diff --git a/examples/Basic_Simulation/initial_conditions.py b/examples/Basic_Simulation/initial_conditions.py index a577c6658..c69a5d26a 100644 --- a/examples/Basic_Simulation/initial_conditions.py +++ b/examples/Basic_Simulation/initial_conditions.py @@ -42,14 +42,10 @@ GM_pl = (np.array([6e23, 8e23, 1e24, 3e24, 5e24]) / sim.param['MU2KG']) * sim.GU R_pl = np.full(npl, (3 * (GM_pl / sim.GU) / (4 * np.pi * density_pl)) ** (1.0 / 3.0)) Rh_pl = a_pl * ((GM_pl) / (3 * sim.GU)) ** (1.0 / 3.0) -Ip1_pl = [0.4, 0.4, 0.4, 0.4, 0.4] -Ip2_pl = [0.4, 0.4, 0.4, 0.4, 0.4] -Ip3_pl = [0.4, 0.4, 0.4, 0.4, 0.4] -rotx_pl = [0.0, 0.0, 0.0, 0.0, 0.0] -roty_pl = [0.0, 0.0, 0.0, 0.0, 0.0] -rotz_pl = [0.0, 0.0, 0.0, 0.0, 0.0] +Ip_pl = np.full((npl,3),0.4,) +rot_pl = np.zeros((npl,3)) -sim.add_body(name=name_pl, v1=a_pl, v2=e_pl, v3=inc_pl, v4=capom_pl, v5=omega_pl, v6=capm_pl, Gmass=GM_pl, radius=R_pl, rhill=Rh_pl, Ip1=Ip1_pl, Ip2=Ip2_pl, Ip3=Ip3_pl, rotx=rotx_pl, roty=roty_pl, rotz=rotz_pl) +sim.add_body(name=name_pl, a=a_pl, e=e_pl, inc=inc_pl, capom=capom_pl, omega=omega_pl, capm=capm_pl, Gmass=GM_pl, radius=R_pl, rhill=Rh_pl, Ip=Ip_pl, rot=rot_pl) # Add 10 user-defined test particles ntp = 10 @@ -62,9 +58,9 @@ omega_tp = default_rng().uniform(0.0, 360.0, ntp) capm_tp = default_rng().uniform(0.0, 360.0, ntp) -sim.add_body(name=name_tp, v1=a_tp, v2=e_tp, v3=inc_tp, v4=capom_tp, v5=omega_tp, v6=capm_tp) +sim.add_body(name=name_tp, a=a_tp, e=e_tp, inc=inc_tp, capom=capom_tp, omega=omega_tp, capm=capm_tp) # Display the run configuration parameters sim.get_parameter() # Run the simulation -sim.run(tstart=0.0, tstop=1.0e2, dt=0.01, tstep_out=1.0e0, dump_cadence=0, ) +#sim.run(tstart=0.0, tstop=1.0e2, dt=0.01, tstep_out=1.0e0, dump_cadence=0, ) diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py index f93a89a8f..b24eff0d8 100644 --- a/python/swiftest/swiftest/init_cond.py +++ b/python/swiftest/swiftest/init_cond.py @@ -266,36 +266,35 @@ def vec2xr(param: Dict, **kwargs: Any): Ip: (n,3) array-like of flaot, optional Principal axes moments of inertia vectors if these are massive bodies with rotation enabled. This can be used instead of passing Ip1, Ip2, and Ip3 separately - t : array of floats + time : array of floats Time at start of simulation Returns ------- ds : xarray dataset """ + scalar_dims = ['id'] + vector_dims = ['id','space'] + space_coords = np.array(["x","y","z"]) + vector_vars = ["rh","vh","Ip","rot"] + scalar_vars = ["name","a","e","inc","capom","omega","capm","Gmass","radius","rhill","J2","J4"] + time_vars = ["rh","vh","Ip","rot","a","e","inc","capom","omega","capm","Gmass","radius","rhill","J2","J4"] - + # Check for valid keyword arguments kwargs = {k:kwargs[k] for k,v in kwargs.items() if v is not None} - if param['ROTATION']: - if "rot" not in kwargs: - warnings.warn("Rotation vectors must be given when rotation is enabled",stacklevel=2) + if "rot" not in kwargs and "Gmass" in kwargs: + warnings.warn("Rotation vectors must be given when rotation is enabled for massive bodies",stacklevel=2) return - if "Ip" not in kwargs: + if "Ip" not in kwargs and "rot" in kwargs: kwargs['Ip'] = np.full_like(rot, 0.4) - if "t" not in kwargs: - kwargs["t"] = np.array([0.0]) - - scalar_dims = ['id'] - vector_dims = ['id','space'] - time_dims =['time'] - space_coords = np.array(["x","y","z"]) + if "time" not in kwargs: + kwargs["time"] = np.array([0.0]) + valid_arguments = vector_vars + scalar_vars + ['time','id'] - vector_vars = ["rh","vh","Ip","rot"] - scalar_vars = ["name","a","e","inc","capom","omega","capm","Gmass","radius","rhill","J2","J4"] - time_vars = ["rh","vh","Ip","rot","a","e","inc","capom","omega","capm","Gmass","radius","rhill","J2","J4"] + kwargs = {k:v for k,v in kwargs.items() if k in valid_arguments} data_vars = {k:(scalar_dims,v) for k,v in kwargs.items() if k in scalar_vars} data_vars.update({k:(vector_dims,v) for k,v in kwargs.items() if k in vector_vars}) @@ -307,12 +306,6 @@ def vec2xr(param: Dict, **kwargs: Any): ) time_vars = [v for v in time_vars if v in ds] for v in time_vars: - ds[v] = ds[v].expand_dims({"time":1}).assign_coords({"time": kwargs['t']}) - - # Re-order dimension coordinates so that they are in the same order as the Fortran side - idx = ds.indexes - dim_order = ["time", "id", "space"] - idx = {index_name: idx[index_name] for index_name in dim_order} - ds = ds.reindex(idx) + ds[v] = ds[v].expand_dims({"time":1}).assign_coords({"time": kwargs['time']}) return ds \ No newline at end of file diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 446ee0e3d..8394c5e44 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -2085,9 +2085,9 @@ def add_solar_system_body(self, #Convert the list receieved from the solar_system_horizons output and turn it into arguments to vec2xr if len(body_list) == 1: - values = list(np.hsplit(np.array(body_list[0]),17)) + values = list(np.hsplit(np.array(body_list[0],dtype=np.dtype(object)),17)) else: - values = list(np.squeeze(np.hsplit(np.array(body_list),17))) + values = list(np.squeeze(np.hsplit(np.array(body_list,np.dtype(object)),17))) keys = ["id","name","a","e","inc","capom","omega","capm","rh","vh","Gmass","radius","rhill","Ip","rot","J2","J4"] kwargs = dict(zip(keys,values)) scalar_floats = ["a","e","inc","capom","omega","capm","Gmass","radius","rhill","J2","J4"] @@ -2107,7 +2107,7 @@ def add_solar_system_body(self, if np.all(np.isnan(kwargs[k])): kwargs[k] = None - kwargs['t'] = np.array([self.param['TSTART']]) + kwargs['time'] = np.array([self.param['TSTART']]) dsnew = init_cond.vec2xr(self.param,**kwargs) @@ -2370,7 +2370,7 @@ def input_to_array_3d(val,n=None): capom,nbodies = input_to_array(capom,"f",nbodies) omega,nbodies = input_to_array(omega,"f",nbodies) capm,nbodies = input_to_array(capm,"f",nbodies) - idvals,nbodies = input_to_array(id,"i",nbodies) + id,nbodies = input_to_array(id,"i",nbodies) mass,nbodies = input_to_array(mass,"f",nbodies) Gmass,nbodies = input_to_array(Gmass,"f",nbodies) rhill,nbodies = input_to_array(rhill,"f",nbodies) @@ -2388,50 +2388,18 @@ def input_to_array_3d(val,n=None): else: maxid = self.data.id.max().values[()] - if idvals is None: - idvals = np.arange(start=maxid+1,stop=maxid+1+nbodies,dtype=int) + if id is None: + id = np.arange(start=maxid+1,stop=maxid+1+nbodies,dtype=int) if name is None: - name=np.char.mod(f"Body%d",idvals) + name=np.char.mod(f"Body%d",id) if len(self.data) > 0: - dup_id = np.in1d(idvals, self.data.id) + dup_id = np.in1d(id, self.data.id) if any(dup_id): - raise ValueError(f"Duplicate ids detected: ", *idvals[dup_id]) + raise ValueError(f"Duplicate ids detected: ", *id[dup_id]) - t = self.param['TSTART'] - - if rh is not None: - if v1 is not None or v2 is not None or v3 is not None: - raise ValueError("Cannot use rh and v1,v2,v3 inputs simultaneously!") - else: - v1 = rh.T[0] - v2 = rh.T[1] - v3 = rh.T[2] - - if vh is not None: - if v4 is not None or v5 is not None or v6 is not None: - raise ValueError("Cannot use vh and v4,v5,v6 inputs simultaneously!") - else: - v4 = vh.T[0] - v5 = vh.T[1] - v6 = vh.T[2] - - if rot is not None: - if rotx is not None or roty is not None or rotz is not None: - raise ValueError("Cannot use rot and rotx,roty,rotz inputs simultaneously!") - else: - rotx = rot.T[0] - roty = rot.T[1] - rotz = rot.T[2] - - if Ip is not None: - if Ip1 is not None or Ip2 is not None or Ip3 is not None: - raise ValueError("Cannot use Ip and Ip1,Ip2,Ip3 inputs simultaneously!") - else: - Ip1 = Ip.T[0] - Ip2 = Ip.T[1] - Ip3 = Ip.T[2] + time = [self.param['TSTART']] if mass is not None: if Gmass is not None: @@ -2439,11 +2407,8 @@ def input_to_array_3d(val,n=None): else: Gmass = self.param['GU'] * mass - dsnew = init_cond.vec2xr(self.param, name, v1, v2, v3, v4, v5, v6, idvals, - GMpl=Gmass, Rpl=radius, rhill=rhill, - Ip1=Ip1, Ip2=Ip2, Ip3=Ip3, - rotx=rotx, roty=roty, rotz=rotz, - J2=J2, J4=J4,t=t) + dsnew = init_cond.vec2xr(self.param, name=name, a=a, e=e, inc=inc, capom=capom, omega=omega, capm=capm, id=id, + Gmass=Gmass, radius=radius, rhill=rhill, Ip=Ip, rh=rh, vh=vh,rot=rot, J2=J2, J4=J4, time=time) dsnew = self._combine_and_fix_dsnew(dsnew) self.save(verbose=False)