From 301d08dce7d454cc87608271f025cdc1ae4cd4c6 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Fri, 3 Mar 2023 11:24:58 +0100
Subject: [PATCH] Still improve loading

---
 snapshotprocessor.py | 50 ++++++++++++++++++++++++++++++++++----------
 studyprocessor.py    |  8 ++++---
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/snapshotprocessor.py b/snapshotprocessor.py
index e48dce4..f2ec817 100644
--- a/snapshotprocessor.py
+++ b/snapshotprocessor.py
@@ -448,20 +448,33 @@ class SnapshotProcessor(HDF5Container):
             "level": "levelp",
             "epoch": "epoch",
         }
-        hydro_file = open(f"{path}/output_{self.num:05}/hydro_file_descriptor.txt")
-        part_file = open(f"{path}/output_{self.num:05}/part_file_descriptor.txt")
+        hydro_filename = f"{path}/output_{self.num:05}/hydro_file_descriptor.txt"
+        if os.path.exists(hydro_filename):
+            hydro_file = open(hydro_filename)
+            hlines = hydro_file.readlines()
+            hydro_var = np.unique(
+                list(map(lambda s: s.split(",")[1][1:].split("_")[0], hlines[2:]))
+            )
+        else:
+            hydro_var = []
+
+        part_filename = f"{path}/output_{self.num:05}/part_file_descriptor.txt"
+        if os.path.exists(part_filename):
+            hydro_file = open(part_filename)
+            part_file = open(f"{path}/output_{self.num:05}/part_file_descriptor.txt")
+            has_grav = os.path.exists(
+                f"{path}/output_{self.num:05}/grav_{self.num:05}.out00000"
+            )
+            plines = part_file.readlines()
+            part_var = np.unique(
+                list(map(lambda s: s.split(",")[1][1:].split("_")[0], plines[2:]))
+            )
+        else:
+            part_var = []
+
         has_grav = os.path.exists(
             f"{path}/output_{self.num:05}/grav_{self.num:05}.out00000"
         )
-        # ugly parsing
-        hlines = hydro_file.readlines()
-        plines = part_file.readlines()
-        hydro_var = np.unique(
-            list(map(lambda s: s.split(",")[1][1:].split("_")[0], hlines[2:]))
-        )
-        part_var = np.unique(
-            list(map(lambda s: s.split(",")[1][1:].split("_")[0], plines[2:]))
-        )
 
         def is_available(available_vars, pymsesrc, var):
             if var in ["g", "phi"]:
@@ -500,6 +513,21 @@ class SnapshotProcessor(HDF5Container):
             )
         )
 
+        if (
+            "Br" in pymses.rcConfig.Ramses.amr_fields.field_name_list
+            and "B" not in hydro_var
+        ):
+            pymses.rcConfig.Ramses.amr_fields.remove_field("P")
+            pymses.rcConfig.Ramses.amr_fields.remove_field("Br")
+            pymses.rcConfig.Ramses.amr_fields.remove_field("Bl")
+            pymses.rcConfig.Ramses.amr_fields.add_scalar(
+                name="P", ivar=4, file_prefix="hydro"
+            )
+
+        if not has_grav and "g" in pymses.rcConfig.Ramses.amr_fields.field_name_list:
+            pymses.rcConfig.Ramses.amr_fields.remove_field("g")
+            pymses.rcConfig.Ramses.amr_fields.remove_field("phi")
+
         self._amr = self._ro.amr_source(self.params.pymses.variables)
         self._part = self._ro.particle_source(self.params.pymses.part_variables)
 
diff --git a/studyprocessor.py b/studyprocessor.py
index 30f8786..d690ca3 100644
--- a/studyprocessor.py
+++ b/studyprocessor.py
@@ -96,13 +96,15 @@ class StudyProcessor(Aggregator, HDF5Container):
             for run in self.runs:
                 nml_src = f"{self.path}/{run}/{self.params.input.nml_filename}"
                 nml_dest = f"{self.path_out}/{run}/{self.params.input.nml_filename}"
-                copy_file(nml_src, nml_dest, update=1)
+                if os.path.exists(nml_src):
+                    copy_file(nml_src, nml_dest, update=1)
 
                 logs = self.get_logs(run)
                 os.makedirs(f"{self.path_out}/{run}/logs", exist_ok=True)
                 for log in logs:
                     dest = f"{self.path_out}/{run}/logs/{os.path.basename(log)}"
-                    copy_file(log, dest, update=1)
+                    if os.path.exists(log):
+                        copy_file(log, dest, update=1)
 
         # Define rules
         self.def_rules()
@@ -751,7 +753,7 @@ class StudyProcessor(Aggregator, HDF5Container):
             "ssm": Rule(
                 self._surfacic_sink_mass,
                 group="/series/sinks_from_log",
-                unit=U.Msun / U.pc ** 2,
+                unit=U.Msun / U.pc**2,
                 description="Surfacic sink mass",
                 dependencies=["sinks_from_log"],
             ),