add mkdir commands where needed in datasets

Author: liellnima

emulator/src/data/climate_dataset.py

@@ -5,6 +5,7 @@
 import numpy as np
 import xarray as xr
 
+from pathlib import Path
 from typing import Dict, Optional, List, Tuple, Union
 
 from emulator.src.data.constants import (
@@ -64,6 +65,7 @@ def __init__(
 
         self.test_dir = output_save_dir
         self.output_save_dir = output_save_dir
+        Path(self.output_save_dir).mkdir(parents=True, exist_ok=True)
 
         self.channels_last = channels_last
         self.load_data_into_mem = load_data_into_mem
@@ -421,7 +423,8 @@ def __init__(  # inherits all the stuff from Base
     ):
         self.mode = mode
         self.output_save_dir = output_save_dir
-        self.root_dir = os.path.join(data_dir, "outputs/CMIP6")
+        self.root_dir = Path(data_dir) / "outputs" / "CMIP6"
+        self.root_dir.mkdir(parents=True, exist_ok=True)
 
         self.input_nc_files = []
         self.output_nc_files = []
@@ -442,7 +445,8 @@ def __init__(  # inherits all the stuff from Base
         )
 
         if isinstance(climate_model, str):
-            self.root_dir = os.path.join(self.root_dir, climate_model)
+            self.root_dir = self.root_dir / climate_model
+            self.root_dir.mkdir(parents=True, exist_ok=True)
         else:
             raise NotImplementedError("For loading multiple climate models, please make sure to use the Super Climate Dataset Class.")
 
@@ -461,6 +465,9 @@ def __init__(  # inherits all the stuff from Base
                 )  # Taking multiple ensemble members
                 if i == (num_ensembles - 1):
                     break  # if num_ensemble ==-1 we take all
+        
+        for ens_dir in self.ensemble_dir:
+            Path(ens_dir).mkdir(parents=True, exist_ok=True)
 
         # Check here if os.path.isfile($SCRATCH/data.npz) exists
         # if it does, use self._reload data(path)
@@ -584,8 +591,10 @@ def __init__(
         self.channels_last = channels_last
 
         self.mode = mode
-        self.root_dir = os.path.join(data_dir, "inputs/input4mips")
-        self.output_save_dir = output_save_dir
+        self.root_dir = Path(data_dir) / "inputs" / "input4mips"
+        self.root_dir.mkdir(parents=True, exist_ok=True)
+        self.output_save_dir = Path(output_save_dir)
+        self.output_save_dir.mkdir(parents=True, exist_ok=True)
         self.input_nc_files = []
         self.output_nc_files = []
 
@@ -710,7 +719,7 @@ def __init__(
                 self.norm_data = self.normalize_data(self.raw_data, stats)
 
             self.data_path = self.save_data_into_disk(
-                self.raw_data, fname, output_save_dir
+                self.raw_data, fname, str(output_save_dir)
             )
 
             self.copy_to_slurm(self.data_path)

emulator/src/data/super_climate_dataset.py

@@ -5,6 +5,7 @@
 import numpy as np
 import xarray as xr
 
+from pathlib import Path
 from abc import ABC, abstractmethod
 from typing import Dict, Optional, List, Tuple, Union
 
@@ -363,6 +364,7 @@ def __init__(
         super().__init__()
         self.index_manager = index_manager
         self.output_save_dir = output_save_dir
+        Path(self.output_save_dir).mkdir(parents=True, exist_ok=True)
         self.channels_last = channels_last
         self.load_data_into_mem = load_data_into_mem
 
@@ -545,6 +547,7 @@ def __init__(  # inherits all the stuff from Base
     ):
         self.mode = mode
         self.output_save_dir = output_save_dir
+        Path(output_save_dir).mkdir(parents=True, exist_ok=True)
 
         self.input_nc_files = []
         self.output_nc_files = []
@@ -604,6 +607,7 @@ def __init__(  # inherits all the stuff from Base
                         var_dir = os.path.join(
                             data_dir, exp, var, f"{CMIP6_NOM_RES}/{CMIP6_TEMP_RES}/{y}"
                         )
+                        Path(var_dir).mkdir(parents=True, exist_ok=True)
                         files = glob.glob(var_dir + f"/*.nc", recursive=True)
                         if len(files) == 0:
                             raise FileNotFoundError(f"No files for climate model {climate_model}, ensemble member {data_dir.split("/")[-1]}, var {var}, year {y}, scenario {exp}. Please check if climate model runs for this exact pairing actually exist.")
@@ -689,8 +693,10 @@ def __init__(  # inherits all the stuff from Base
         self.channels_last = channels_last
 
         self.mode = mode
-        self.root_dir = os.path.join(data_dir, "inputs/input4mips")
+        self.root_dir = Path(data_dir) / "inputs" / "input4mips"
+        self.root_dir.mkdir(parents=True, exist_ok=True)
         self.output_save_dir = output_save_dir
+        Path(self.output_save_dir).mkdir(parents=True, exist_ok=True)
         self.input_nc_files = []
         self.output_nc_files = []
 
@@ -749,6 +755,7 @@ def __init__(  # inherits all the stuff from Base
                             var,
                             f"{CMIP6_NOM_RES}/{CMIP6_TEMP_RES}/{y}",
                         )
+                        Path(var_dir).mkdir(parents=True, exist_ok=True)
 
                 output_nc_files = []
                 for exp in scenarios:
@@ -769,6 +776,7 @@ def __init__(  # inherits all the stuff from Base
                             var,
                             f"{CMIP6_NOM_RES}/{CMIP6_TEMP_RES}/{y}",
                         )
+                        Path(var_dir).mkdir(parents=True, exist_ok=True)
                         files = glob.glob(
                             var_dir + f"/**/*{filter_path_by}*.nc", recursive=True
                         )