skip data augmentation if not done before

ersilia-os · Nov 19, 2022 · dc29c22 · dc29c22
1 parent a194115
commit dc29c22
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/zairachem/estimators/from_molmap/estimate.py b/zairachem/estimators/from_molmap/estimate.py
@@ -27,7 +27,10 @@ def __init__(self, path):
             self.get_output_dir(), ESTIMATORS_SUBFOLDER, ESTIMATORS_FAMILY_SUBFOLDER
         )
         if _USE_AUGMENTED:
-            self._data_filename = DATA_AUGMENTED_FILENAME
+            if os.path.exists(self.path, DATA_SUBFOLDER, DATA_AUGMENTED_FILENAME):
+                self._data_filename = DATA_AUGMENTED_FILENAME
+            else:
+                self._data_filename = DATA_FILENAME
         else:
             self._data_filename = DATA_FILENAME
 

diff --git a/zairachem/setup/training.py b/zairachem/setup/training.py
@@ -79,7 +79,9 @@ def __init__(
         self.time_budget = time_budget  # TODO
 
     def is_lazy(self):
-        with open(os.path.join(self.output_dir, DATA_SUBFOLDER, PRESETS_FILENAME), "w") as f:
+        with open(
+            os.path.join(self.output_dir, DATA_SUBFOLDER, PRESETS_FILENAME), "w"
+        ) as f:
             data = {"is_lazy": self._is_lazy}
             json.dump(data, f, indent=4)
         return self._is_lazy