Merge branch 'master' into debug

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

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('<i4', '<f8', '<f8', '(6,)<f8')
                 plid[i] = record[0]
-                Mpl[i] = record[1]
+                GMpl[i] = record[1]
                 Rpl[i] = record[2]
                 pvec[:, i] = record[3]
         if ntp > 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)

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)