Merge pull request #23 from GeoOcean/feature/nns

Feature/nns

Author: tausiaj

bluemath_tk/deeplearning/models/__init__.py

@@ -2,12 +2,12 @@
 from keras import layers
 from typing import List, Tuple
 
+
 def ResidualBlock(width: int) -> layers.Layer:
     def apply(x: layers.Layer) -> layers.Layer:
-        
         input_width = x.shape[3]
         residual = x if input_width == width else layers.Conv2D(width, kernel_size=1)(x)
-        
+
         x = layers.LayerNormalization(axis=-1, center=True, scale=True)(x)
         x = layers.Conv2D(
             width, kernel_size=3, padding="same", activation=keras.activations.swish
@@ -18,9 +18,11 @@ def apply(x: layers.Layer) -> layers.Layer:
 
     return apply
 
+
 def DownBlock(width: int, block_depth: int) -> layers.Layer:
-    def apply(x: Tuple[layers.Layer, List[layers.Layer]]) -> Tuple[layers.Layer, List[layers.Layer]]:
-        
+    def apply(
+        x: Tuple[layers.Layer, List[layers.Layer]],
+    ) -> Tuple[layers.Layer, List[layers.Layer]]:
         x, skips = x
         for _ in range(block_depth):
             x = ResidualBlock(width)(x)
@@ -30,26 +32,30 @@ def apply(x: Tuple[layers.Layer, List[layers.Layer]]) -> Tuple[layers.Layer, Lis
 
     return apply
 
+
 def UpBlock(width: int, block_depth: int) -> layers.Layer:
     def apply(x: Tuple[layers.Layer, List[layers.Layer]]) -> layers.Layer:
-        
         x, skips = x
         x = layers.UpSampling2D(size=2, interpolation="bilinear")(x)
         for _ in range(block_depth):
             x = layers.Concatenate()([x, skips.pop()])
             x = ResidualBlock(width)(x)
         return x
+
     return apply
 
-def get_model(image_height: int, 
-              image_width: int, 
-              input_frames: int, 
-              output_frames: int, 
-              down_widths: List[int] = [64, 128, 256], 
-              up_widths: List[int] = [256, 128, 64], 
-              block_depth: int = 2) -> keras.Model:
+
+def get_model(
+    image_height: int,
+    image_width: int,
+    input_frames: int,
+    output_frames: int,
+    down_widths: List[int] = [64, 128, 256],
+    up_widths: List[int] = [256, 128, 64],
+    block_depth: int = 2,
+) -> keras.Model:
     """Builds the U-Net like model with residual blocks and skip connections."""
-    
+
     inputs = keras.Input(shape=(image_height, image_width, input_frames))
     x = layers.Conv2D(down_widths[0], kernel_size=1)(inputs)
 
@@ -64,5 +70,5 @@ def get_model(image_height: int,
         x = UpBlock(width, block_depth)([x, skips])
 
     outputs = layers.Conv2D(output_frames, kernel_size=1, kernel_initializer="zeros")(x)
-    
+
     return keras.Model(inputs, outputs, name="residual_unet")

bluemath_tk/deeplearning/models/resnet_model.py

@@ -1,11 +1,13 @@
-name: bluemath
+name: bluemath-deep
 channels:
   - conda-forge
 dependencies:
   - python=3.12
   - numpy
   - pandas
   - xarray
+  - netcdf4
+  - scipy
   - scikit-learn
   - matplotlib
   - qt
@@ -14,6 +16,8 @@ dependencies:
   - cartopy
   - pytest
   - cdsapi
+  - keras
+  - tensorflow
   # PIP packages can be added below (avoid this whenever possible)
   - pip:
       - minisom

environment.yml

@@ -0,0 +1,59 @@
+import unittest
+from unittest.mock import Mock
+import keras
+
+
+class TestResNetTraining(unittest.TestCase):
+    def setUp(self):
+        # Mock the data generator
+        self.train_generator = Mock()
+        self.train_generator.num_images = 5000
+        self.train_generator.input_height = 64
+        self.train_generator.input_width = 64
+        self.train_generator.output_height = 64
+        self.train_generator.output_width = 64
+        self.train_generator.batch_size = 1
+
+        # Define optimizer
+        self.optimizer = keras.optimizers.AdamW
+
+        # Mock the model
+        self.model = Mock()
+        self.model.compile = Mock()
+        self.model.fit = Mock(return_value=Mock(history={"loss": [0.1, 0.05, 0.01]}))
+
+    def test_model_training(self):
+        # Compile the model
+        self.model.compile(
+            optimizer=self.optimizer(learning_rate=1e-4, weight_decay=1e-5),
+            loss=keras.losses.mean_squared_error,
+        )
+
+        # Assert that compile was called with correct parameters
+        self.model.compile.assert_called_with(
+            optimizer=self.optimizer(learning_rate=1e-4, weight_decay=1e-5),
+            loss=keras.losses.mean_squared_error,
+        )
+
+        # Start the train loop with the fit method
+        history = self.model.fit(
+            self.train_generator, initial_epoch=0, epochs=20, steps_per_epoch=500
+        )
+
+        # Assert that fit was called with correct parameters
+        self.model.fit.assert_called_with(
+            self.train_generator, initial_epoch=0, epochs=20, steps_per_epoch=500
+        )
+
+        # Assert that the training history is as expected
+        self.assertIn("loss", history.history)
+        self.assertEqual(len(history.history["loss"]), 3)
+        self.assertAlmostEqual(history.history["loss"][0], 0.1)
+        self.assertAlmostEqual(history.history["loss"][1], 0.05)
+        self.assertAlmostEqual(history.history["loss"][2], 0.01)
+
+        print("training complete")
+
+
+if __name__ == "__main__":
+    unittest.main()