From 6385bfb8b872f9175d84c4c998099f813dc7fe70 Mon Sep 17 00:00:00 2001
From: Terrence Tricco <tstricco@mun.ca>
Date: Thu, 25 Apr 2024 14:36:55 -0230
Subject: [PATCH] GradSPH file reading

---
 sarracen/__init__.py             |  1 +
 sarracen/readers/read_gradsph.py | 72 ++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)
 create mode 100644 sarracen/readers/read_gradsph.py

diff --git a/sarracen/__init__.py b/sarracen/__init__.py
index 1c38c30..f7d6249 100644
--- a/sarracen/__init__.py
+++ b/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
 
diff --git a/sarracen/readers/read_gradsph.py b/sarracen/readers/read_gradsph.py
new file mode 100644
index 0000000..21f65f2
--- /dev/null
+++ b/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