GradSPH file reading

sarracen/__init__.py

@@ -2,6 +2,7 @@
 from .readers.read_marisa import read_marisa
 from .readers.read_csv import read_csv
 from .readers.read_phantom import read_phantom
+from .readers.read_gradsph import read_gradsph
 
 from .sarracen_dataframe import SarracenDataFrame
 

sarracen/readers/read_gradsph.py

@@ -0,0 +1,72 @@
+import pandas as pd
+
+from ..sarracen_dataframe import SarracenDataFrame
+
+
+def read_gradsph(filename: str, separate_types: str = 'sinks'):
+    """
+    Read data from a GradSPH dump file.
+
+    Global values stored in the dump file are stored within the data frame in
+    the dictionary ``params``.
+
+    Parameters
+    ----------
+    filename : str
+        Name of the file to be loaded.
+    separate_types : {None, 'sinks', 'all'}, default='sinks'
+        Whether to separate SPH particles and sink particles into separate
+        SarracenDataFrames. ``None`` returns all particle types in one
+        SarracenDataFrame. '`sinks`' and '`all`' separate sink particles into
+        a second SarracenDataFrame.
+
+    Returns
+    -------
+    SarracenDataFrame or list of SarracenDataFrame
+
+    Examples
+    --------
+    By default, SPH particles are grouped into one SarracenDataFrame and sink
+    particles into a second SarracenDataFrame.
+
+    >>> sdf, sdf_sinks = sarracen.read_gradsph('col3139')
+    """
+    with open(filename, 'r') as fp:
+
+        n, ninactive, nsink = fp.readline().split()
+        n, ninactive, nsink = int(n), int(ninactive), int(nsink)
+
+        t, gamma = fp.readline().split()
+        t, gamma = float(t), float(gamma)
+
+        params = {'n': n,
+                  'ninactive': ninactive,
+                  'nsink': nsink,
+                  't': t,
+                  'gamma': gamma}
+
+        sinks = []
+        for i in range(nsink):
+            sinks.append(fp.readline().split())
+
+        sink_header = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'mass']
+        df_sinks = pd.DataFrame(sinks, columns=sink_header, dtype=float)
+
+        parts = []
+        for i in range(n - ninactive):
+            parts.append(fp.readline().split())
+        part_header = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'mass', 'h', 'cs',
+                       'rho', 'temp']
+        df_parts = pd.DataFrame(parts, columns=part_header, dtype=float)
+
+        if separate_types == 'sinks' or separate_types == 'all':
+            df_list = [SarracenDataFrame(df_parts, params=params),
+                       SarracenDataFrame(df_sinks, params=params)]
+        else:
+            df_list = [SarracenDataFrame(pd.concat([df_parts, df_sinks],
+                                                   ignore_index=True),
+                                        params=params)]
+
+        df_list = df_list[0] if len(df_list) == 1 else df_list
+
+        return df_list