From 4a990c5ce8380566c06e486735c17ff20f33f3d4 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Thu, 2 Dec 2021 12:27:04 +0100
Subject: [PATCH] [study] add stellar parsing

---
 studyprocessor.py | 54 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 52 insertions(+), 2 deletions(-)

diff --git a/studyprocessor.py b/studyprocessor.py
index ce0e2d3..96ebc63 100644
--- a/studyprocessor.py
+++ b/studyprocessor.py
@@ -244,10 +244,9 @@ class StudyProcessor(Aggregator, HDF5Container):
         return slope
 
     def _extract_sinks_from_log(self, series, log_filename, run):
-        cmd_grep = "sed '/cpu.*/d' {} | grep 'Number of sink' -A 2".format(log_filename)
+        cmd_grep = f"grep 'Number of sink' {log_filename} -A 2"
         content = os.popen(cmd_grep).readlines()
         for i in range(0, len(content), 4):
-
             try:
                 nb_sink = np.int(content[i].split("=")[1])
                 mass_sink = np.float(content[i + 1].split("=")[1])
@@ -269,6 +268,42 @@ class StudyProcessor(Aggregator, HDF5Container):
                 )
         return series
 
+    def _extract_stellar_from_log(self, stellar_objects, log_filename, run):
+        cmd_grep = f"grep stellar {log_filename} -n"
+        content = os.popen(cmd_grep).readlines()
+        nb_stellar = list(map(lambda s: int(s.split()[1]), content))
+        line_numbers = list(map(lambda s: int(s.split(":")[0]), content))
+        current_line = 0
+        logfile = open(log_filename)
+        for i in range(0, len(line_numbers)):
+            try:
+                while current_line < line_numbers[i] + 3:
+                    logfile.readline()
+                    current_line += 1
+                for j in range(nb_stellar[i]):
+                    line_stellar = logfile.readline().split()
+                    current_line += 1
+                    mass = float(line_stellar[3])
+                    time = float(line_stellar[4])
+                    lifetime = float(line_stellar[5])
+                    id = int(line_stellar[6])
+
+                    stellar_objects["mass"][run].append(mass)
+                    stellar_objects["time"][run].append(time)
+                    stellar_objects["lifetime"][run].append(lifetime)
+                    stellar_objects["id"][run].append(id)
+
+            except (ValueError, IndexError):
+                self._log(
+                    "[stellar] Error encountered in parsing {} (grepped block {})".format(
+                        log_filename, i
+                    ),
+                    "WARNING",
+                )
+                raise
+        logfile.close()
+        return stellar_objects
+
     def _extract_sfr_from_log(self, series, log_filename, run):
         cmd_grep = "grep '\[SFR' {} ".format(log_filename)
         content = os.popen(cmd_grep).readlines()
@@ -669,6 +704,21 @@ class StudyProcessor(Aggregator, HDF5Container):
                     "mem_parts": U.none,
                 },
             ),
+            "stellar_from_log": Rule(
+                self,
+                partial(
+                    self._from_log,
+                    ["time", "mass", "lifetime", "id"],
+                    self._extract_stellar_from_log,
+                ),
+                group="/dataset",
+                unit={
+                    "time": "unit_time",
+                    "mass": "unit_mass",
+                    "lifetime": "unit_time",
+                    "id": U.none,
+                },
+            ),
             "turb_power": Rule(
                 self,
                 self._turb_power,