Merge pull request #24 from GeoOcean/feature/wrappers

Feature/wrappers

Author: tausiaj

bluemath_tk/deeplearning/generators/ncDataGenerator.py

@@ -0,0 +1,89 @@
+import numpy as np
+import xarray as xr
+import keras.utils
+
+
+class DataGenerator(keras.utils.Sequence):
+    def __init__(
+        self,
+        msl_path,
+        tp_path,
+        num_images,
+        sequential=False,
+        batch_size=1,
+    ):
+        # create memory-mapped files for high_res and low_res datasets
+
+        # inputs
+
+        self.msl = xr.open_dataarray(msl_path).values[:, :64, :64]
+
+        # outputs
+
+        self.tp = xr.open_dataarray(tp_path).values[:, :64, :64]
+
+        # set boolean for sequential or random dataset
+        self.sequential = sequential
+        # counter for keeping track of seuquential generator
+        self.counter = 0
+        self.num_images = num_images
+        # set sequence len
+        # flag for diffusion/unet
+        self.batch_size = batch_size
+        self.num_samples = self.msl.shape[0]
+
+    @property
+    def num_batches(self):
+        return int(np.floor(self.num_images / self.batch_size))
+
+    def __len__(self):
+        return self.num_batches
+
+    def min_max_normalize(self, arr, min, max):
+        normalized = (arr - min) / (max - min)
+        result = np.where(np.isnan(arr), np.nan, normalized)
+        return result
+
+    # must be called to restart the sequential
+    def counter_reset(self):
+        self.counter = 0
+
+    def __getitem__(self, idx):
+
+        # prepare the resulting array
+        inputs = np.zeros((self.batch_size, 64, 64, 1))
+        outputs = np.zeros((self.batch_size, 64, 64, 1))
+
+        # random path
+        if self.sequential == False:
+            # compose the batch one element at the time
+            for i in range(self.batch_size):
+                # get a random number in range
+                random = np.random.randint(0, self.num_samples - 1)
+
+                # inputs
+                inputs[i, :, :, 0] = self.min_max_normalize(
+                    self.msl[random], 95680, 104256
+                )
+
+                # outputs
+                outputs[i, :, :, 0] = self.min_max_normalize(
+                    self.tp[random], 0.0, 0.02197266
+                )
+
+        # sequential path
+        if self.sequential == True:
+            # compose the batch one element at the time
+            for i in range(self.batch_size):
+                # inputs
+                inputs[i, :, :, 0] = self.min_max_normalize(
+                    self.msl[self.counter], 95680, 104256
+                )
+
+                # outputs
+                outputs[i, :, :, 0] = self.min_max_normalize(
+                    self.tp[self.counter], 0.0, 0.02197266
+                )
+
+                self.counter = self.counter + 1
+        return inputs, outputs

bluemath_tk/deeplearning/resnet.py

@@ -1,6 +1,6 @@
 import keras
 from models import resnet_model
-from generators.mockDataGenerator import MockDataGenerator
+from generators.ncDataGenerator import DataGenerator
 
 # instantiate model class (load memory)
 model = resnet_model.get_model(
@@ -11,14 +11,17 @@
 print(model.summary())
 
 # instantiate generator class
-train_generator = MockDataGenerator(
-    num_images=5000,
-    input_height=64,
-    input_width=64,
-    output_height=64,
-    output_width=64,
+train_generator = DataGenerator(
+    msl_path="/home/tausiaj/DATA/Comparison-ERA5/msl_spain.nc",
+    tp_path="/home/tausiaj/DATA/Comparison-ERA5/tp_spain.nc",
+    num_images=8760,
+    sequential=False,
     batch_size=1,
 )
+
+a, b = train_generator.__getitem__(1)
+print(a.shape)
+print(b.shape)
 # define oprimizer
 optimizer = keras.optimizers.AdamW
 model.compile(

bluemath_tk/wrappers/_base_wrappers.py

@@ -1,13 +1,14 @@
 import os
 import itertools
 from typing import List
+import numpy as np
 from jinja2 import Environment, FileSystemLoader
 from ..core.models import BlueMathModel
 
 
 class BaseModelWrapper(BlueMathModel):
     """
-    Base class for model wrappers.
+    Base class for numerical models wrappers.
 
     Attributes
     ----------
@@ -21,6 +22,10 @@ class BaseModelWrapper(BlueMathModel):
         The directory where the output files will be saved.
     env : Environment
         The Jinja2 environment.
+    cases_dirs : List[str]
+        The list with cases directories.
+    cases_context : List[dict]
+        The list with cases context.
 
     Methods
     -------
@@ -32,8 +37,12 @@ class BaseModelWrapper(BlueMathModel):
         from the input dictionary.
     render_file_from_template(template_name, context, output_filename=None)
         Render a file from a template.
-    build_cases()
-        Build the cases.
+    write_array_in_file(array, filename)
+        Write an array in a file.
+    copy_files(src, dst)
+        Copy file(s) from source to destination.
+    build_cases(mode="all_combinations")
+        Create the cases folders and render the input files.
     run_cases()
         Run the cases.
     """
@@ -44,6 +53,7 @@ def __init__(
         templates_name: List[str],
         model_parameters: dict,
         output_dir: str,
+        default_parameters: dict = None,
     ):
         """
         Initialize the BaseModelWrapper.
@@ -58,14 +68,59 @@ def __init__(
             The parameters to be used in the templates.
         output_dir : str
             The directory where the output files will be saved.
+        default_parameters : dict, optional
+            The default parameters for the model. If None, the parameters will
+            not be checked.
+            Default is None.
         """
 
         super().__init__()
+        if default_parameters is not None:
+            self._check_parameters_type(
+                default_parameters=default_parameters, model_parameters=model_parameters
+            )
         self.templates_dir = templates_dir
         self.templates_name = templates_name
         self.model_parameters = model_parameters
         self.output_dir = output_dir
         self.env = Environment(loader=FileSystemLoader(self.templates_dir))
+        self.cases_dirs: List[str] = []
+        self.cases_context: List[dict] = []
+
+    def _check_parameters_type(self, default_parameters: dict, model_parameters: dict):
+        """
+        Check if the parameters have the correct type.
+
+        Parameters
+        ----------
+        default_parameters : dict
+            The default parameters for the model.
+        model_parameters : dict
+            The parameters to be used in the templates.
+
+        Raises
+        ------
+        ValueError
+            If a parameter has the wrong type.
+        """
+
+        for model_param, param_value in model_parameters.items():
+            if model_param not in default_parameters:
+                self.logger.warning(
+                    f"Parameter {model_param} is not in the default_parameters"
+                )
+            else:
+                if isinstance(param_value, list) and all(
+                    isinstance(item, default_parameters[model_param])
+                    for item in param_value
+                ):
+                    self.logger.info(
+                        f"Parameter {model_param} has the correct type: {type(default_parameters[model_param])}"
+                    )
+                else:
+                    raise ValueError(
+                        f"Parameter {model_param} has the wrong type: {type(default_parameters[model_param])}"
+                    )
 
     def create_cases_context_one_by_one(self):
         """
@@ -143,18 +198,89 @@ def render_file_from_template(
         with open(output_filename, "w") as f:
             f.write(rendered_content)
 
-    def write_array_in_file(self, array, filename):
+    def write_array_in_file(self, array: np.ndarray, filename: str):
         """
         Write an array in a file.
 
         Parameters
         ----------
-        array : np.array
-            The array to be written.
+        array : np.ndarray
+            The array to be written. Can be 1D or 2D.
         filename : str
             The name of the file.
         """
 
         with open(filename, "w") as f:
-            for item in array:
-                f.write(f"{item}\n")
+            if array.ndim == 1:
+                for item in array:
+                    f.write(f"{item}\n")
+            elif array.ndim == 2:
+                for row in array:
+                    f.write(" ".join(map(str, row)) + "\n")
+            else:
+                raise ValueError("Only 1D and 2D arrays are supported")
+
+    def copy_files(self, src: str, dst: str):
+        """
+        Copy file(s) from source to destination.
+
+        Parameters
+        ----------
+        src : str
+            The source file.
+        dst : str
+            The destination file.
+        """
+
+        if os.path.isdir(src):
+            os.makedirs(dst, exist_ok=True)
+            for file in os.listdir(src):
+                with open(file, "r") as f:
+                    content = f.read()
+                with open(os.path.join(dst, file), "w") as f:
+                    f.write(content)
+        else:
+            with open(src, "r") as f:
+                content = f.read()
+            with open(dst, "w") as f:
+                f.write(content)
+
+    def build_cases(self, mode: str = "all_combinations"):
+        """
+        Create the cases folders and render the input files.
+
+        Parameters
+        ----------
+        mode : str, optional
+            The mode to create the cases. Can be "all_combinations" or "one_by_one".
+            Default is "all_combinations".
+        """
+
+        if mode == "all_combinations":
+            self.cases_context = self.create_cases_context_all_combinations()
+        elif mode == "one_by_one":
+            self.cases_context = self.create_cases_context_one_by_one()
+        else:
+            raise ValueError(f"Invalid mode to create cases: {mode}")
+        for case_num, case_context in enumerate(self.cases_context):
+            case_dir = os.path.join(self.output_dir, f"{case_num:04}")
+            self.cases_dirs.append(case_dir)
+            os.makedirs(case_dir, exist_ok=True)
+            for template_name in self.templates_name:
+                self.render_file_from_template(
+                    template_name=template_name,
+                    context=case_context,
+                    output_filename=os.path.join(case_dir, template_name),
+                )
+        self.logger.info(
+            f"{len(self.cases_dirs)} cases created in {mode} mode and saved in {self.output_dir}"
+        )
+
+    def run_cases(self):
+        """
+        Run the cases.
+        """
+
+        if self.cases_dirs:
+            for case_dir in self.cases_dirs:
+                self.run_model(case_dir)