Organize in submodules

studyprocessor.py

@@ -12,9 +12,9 @@ from scipy.stats import linregress
 from baseprocessor import Rule, HDF5Container
 from aggregator import Aggregator
 from snapshotprocessor import SnapshotProcessor
-from run_selector import RunSelector
-from params import default_params
-from units import U
+from runselector import RunSelector
+from utils.params import default_params
+from utils.units import U
 
 
 class StudyProcessor(Aggregator, HDF5Container):
@@ -452,6 +452,31 @@ class StudyProcessor(Aggregator, HDF5Container):
                 series["turb_energy"][run].append(np.nan)
         return series
 
+    def _extract_SN_Mom_from_log(self, series, log_filename, run):
+        cmd_grep = f"grep -e 'SN momentum' -e 'Fine step' {log_filename}"
+        content = os.popen(cmd_grep).readlines()
+        block_err = []
+        time = 0
+        for i in range(0, len(content)):
+            try:
+                if content[i][1:5] == "Fine":
+                    data = content[i].replace("=", " ").split()
+                    time = np.float(data[4])
+                elif content[i][1:3] == "SN"    :
+                    series["time"][run].append(time)
+                    series["SN_momentum"][run].append(np.float(content[i].split()[-1]))
+                else:
+                    raise ValueError("Wrong start of line")
+            except (AssertionError, ValueError, IndexError):
+                block_err.append(i)
+        
+        if len(block_err) > 0:
+            self.logger.warning(
+                    f"Error encountered in parsing {log_filename} (grepped blocks {block_err})"
+                )  
+        return series
+        
+
     def get_logs(self, run):
         glob_str = f"{self.path}/{run}/{self.params.input.log_prefix}*"
         logs = glob.glob(glob_str)
@@ -815,6 +840,15 @@ class StudyProcessor(Aggregator, HDF5Container):
                     "id": U.none,
                 },
             ),
+            "SN_momentum_from_log": Rule(
+                partial(self._from_log, ["time", "SN_momentum"], self._extract_SN_Mom_from_log),
+                group="/series",
+                unit={"time": "unit_time", "SN_momentum" : {"unit_mass" : 1, "unit_velocity" : 1}},
+                description={
+                    "time": "Time",
+                    "SN_momentum": "Injected momentum",
+                },
+            ),
             "turb_power": Rule(
                 self._turb_power,
                 group="/series/rms_from_log",