diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py index 17eb90d87..76c973533 100644 --- a/python/swiftest/swiftest/init_cond.py +++ b/python/swiftest/swiftest/init_cond.py @@ -215,4 +215,36 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds): plds = plxr.to_dataset(dim='vec') ds = xr.combine_by_coords([ds, plds]) + return ds + +def vec2xr(param, idvals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, Rhill=None, t=0.0): + + dims = ['time', 'id', 'vec'] + if GMpl is not None: + ispl = True + else: + ispl = False + + if ispl and Rpl is None: + print("Massive bodies need a radius value.") + return None + if ispl and Rhill is None and param['RHILL_PRESENT'] == 'YES': + print("Rhill is required.") + return None + + clab, plab, tlab = swiftest.io.make_swiftest_labels(param) + if ispl: + if param['RHILL_PRESENT'] == 'YES': + vec = np.vstack([v1, v2, v3, v4, v5, v6, GMpl, Rpl, Rhill]).T + else: + vec = np.vstack([v1, v2, v3, v4, v5, v6, GMpl, Rpl]).T + else: + vec = np.vstack([v1, v2, v3, v4, v5, v6]).T + bodyframe = np.expand_dims(vec, axis=0) + if ispl: + bodyxr = xr.DataArray(bodyframe, dims=dims, coords={'time': [t], 'id': tpid, 'vec': plab}) + else: + bodyxr = xr.DataArray(bodyframe, dims=dims, coords={'time': [t], 'id': tpid, 'vec': tlab}) + + ds = bodyxr.to_dataset(dim='vec') return ds \ No newline at end of file diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index e8601f133..ecc5051b0 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -351,13 +351,13 @@ def swifter_stream(f, param): tvec = np.empty((6, ntp)) tpid = np.empty(ntp, dtype='int') if npl > 0: - Mpl = np.empty(npl) + GMpl = np.empty(npl) Rpl = np.empty(npl) for i in range(npl): # Read single-line pl frame for record = f.read_record(' 0: @@ -384,7 +384,7 @@ def swifter_stream(f, param): plab = tlab.copy() plab.append('GMass') plab.append('Radius') - pvec = np.vstack([pvec, Mpl, Rpl]) + pvec = np.vstack([pvec, GMpl, Rpl]) yield t, npl, plid, pvec.T, plab, \ ntp, tpid, tvec.T, tlab @@ -499,7 +499,7 @@ def swiftest_stream(f, param): p4 = f.read_reals(np.float64) p5 = f.read_reals(np.float64) p6 = f.read_reals(np.float64) - Mpl = f.read_reals(np.float64) + GMpl = f.read_reals(np.float64) if param['RHILL_PRESENT'] == 'YES': Rhill = f.read_reals(np.float64) Rpl = f.read_reals(np.float64) @@ -525,7 +525,7 @@ def swiftest_stream(f, param): clab, plab, tlab = make_swiftest_labels(param) if npl > 0: - pvec = np.vstack([p1, p2, p3, p4, p5, p6, Mpl, Rpl]) + pvec = np.vstack([p1, p2, p3, p4, p5, p6, GMpl, Rpl]) else: pvec = np.empty((8, 0)) plid = np.empty(0) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index c602cbe8c..78f262df8 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -73,6 +73,35 @@ def add(self, plname, date=date.today().isoformat(), idval=None): return + def addp(self, idvals, t1, t2, t3, t4, t5, t6, Gmass=None, radius=None, Rhill=None, t=0.0): + """ + Adds a body (test particle or massive body) to the internal DataSet given a set up 6 vectors (orbital elements + or cartesian state vectors, depending on the value of self.param). Input all angles in degress + + Parameters + ---------- + idvals : Integer array of body index values. + t1 : xh for param['IN_FORM'] == "XV"; a for param['IN_FORM'] == "EL" + t2 : yh for param['IN_FORM'] == "XV"; e for param['IN_FORM'] == "EL" + t3 : zh for param['IN_FORM'] == "XV"; inc for param['IN_FORM'] == "EL" + t4 : vxh for param['IN_FORM'] == "XV"; capom for param['IN_FORM'] == "EL" + t5 : vyh for param['IN_FORM'] == "XV"; omega for param['IN_FORM'] == "EL" + t6 : vzh for param['IN_FORM'] == "XV"; capm for param['IN_FORM'] == "EL" + Gmass : Optional: Array of G*mass values if these are massive bodies + radius : Optional: Array radius values if these are massive bodies + Rhill : Optional: Array Rhill values if these are massive bodies + Returns + ------- + self.ds : xarray dataset + """ + dsnew = init_cond.vec2xr(self.param, idvals, t1, t2, t3, t4, t5, t6, Gmass, radius, Rhill) + if dsnew is not None: + self.ds = xr.combine_by_coords([self.ds, dsnew]) + return + + + + def read_param(self, param_file, codename="Swiftest"): if codename == "Swiftest": self.param = io.read_swiftest_param(param_file, self.param)