Merge pull request #19 from GeoOcean/feature/nns

Feature/nns

Author: tausiaj

bluemath_tk/deeplearning/__init__.py

@@ -0,0 +1,206 @@
+name: deep
+channels:
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - anyio=4.2.0=py311h06a4308_0
+  - argon2-cffi=21.3.0=pyhd3eb1b0_0
+  - argon2-cffi-bindings=21.2.0=py311h5eee18b_0
+  - asttokens=2.0.5=pyhd3eb1b0_0
+  - async-lru=2.0.4=py311h06a4308_0
+  - attrs=23.1.0=py311h06a4308_0
+  - babel=2.11.0=py311h06a4308_0
+  - beautifulsoup4=4.12.2=py311h06a4308_0
+  - bleach=4.1.0=pyhd3eb1b0_0
+  - brotli-python=1.0.9=py311h6a678d5_7
+  - bzip2=1.0.8=h5eee18b_5
+  - ca-certificates=2024.3.11=h06a4308_0
+  - certifi=2024.2.2=py311h06a4308_0
+  - cffi=1.16.0=py311h5eee18b_0
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - comm=0.2.1=py311h06a4308_0
+  - debugpy=1.6.7=py311h6a678d5_0
+  - decorator=5.1.1=pyhd3eb1b0_0
+  - defusedxml=0.7.1=pyhd3eb1b0_0
+  - executing=0.8.3=pyhd3eb1b0_0
+  - idna=3.4=py311h06a4308_0
+  - ipykernel=6.28.0=py311h06a4308_0
+  - ipython=8.20.0=py311h06a4308_0
+  - jedi=0.18.1=py311h06a4308_1
+  - jinja2=3.1.3=py311h06a4308_0
+  - json5=0.9.6=pyhd3eb1b0_0
+  - jsonschema=4.19.2=py311h06a4308_0
+  - jsonschema-specifications=2023.7.1=py311h06a4308_0
+  - jupyter-lsp=2.2.0=py311h06a4308_0
+  - jupyter_client=8.6.0=py311h06a4308_0
+  - jupyter_core=5.5.0=py311h06a4308_0
+  - jupyter_events=0.8.0=py311h06a4308_0
+  - jupyter_server=2.10.0=py311h06a4308_0
+  - jupyter_server_terminals=0.4.4=py311h06a4308_1
+  - jupyterlab=4.0.11=py311h06a4308_0
+  - jupyterlab_pygments=0.1.2=py_0
+  - jupyterlab_server=2.25.1=py311h06a4308_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libsodium=1.0.18=h7b6447c_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libuuid=1.41.5=h5eee18b_0
+  - markupsafe=2.1.3=py311h5eee18b_0
+  - matplotlib-inline=0.1.6=py311h06a4308_0
+  - mistune=2.0.4=py311h06a4308_0
+  - nbclient=0.8.0=py311h06a4308_0
+  - nbconvert=7.10.0=py311h06a4308_0
+  - nbformat=5.9.2=py311h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.6.0=py311h06a4308_0
+  - notebook-shim=0.2.3=py311h06a4308_0
+  - openssl=3.0.13=h7f8727e_0
+  - overrides=7.4.0=py311h06a4308_0
+  - packaging=23.2=py311h06a4308_0
+  - pandocfilters=1.5.0=pyhd3eb1b0_0
+  - parso=0.8.3=pyhd3eb1b0_0
+  - pexpect=4.8.0=pyhd3eb1b0_3
+  - platformdirs=3.10.0=py311h06a4308_0
+  - prometheus_client=0.14.1=py311h06a4308_0
+  - prompt-toolkit=3.0.43=py311h06a4308_0
+  - prompt_toolkit=3.0.43=hd3eb1b0_0
+  - psutil=5.9.0=py311h5eee18b_0
+  - ptyprocess=0.7.0=pyhd3eb1b0_2
+  - pure_eval=0.2.2=pyhd3eb1b0_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pygments=2.15.1=py311h06a4308_1
+  - pysocks=1.7.1=py311h06a4308_0
+  - python=3.11.9=h955ad1f_0
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - python-fastjsonschema=2.16.2=py311h06a4308_0
+  - python-json-logger=2.0.7=py311h06a4308_0
+  - pytz=2024.1=py311h06a4308_0
+  - pyyaml=6.0.1=py311h5eee18b_0
+  - pyzmq=25.1.2=py311h6a678d5_0
+  - readline=8.2=h5eee18b_0
+  - referencing=0.30.2=py311h06a4308_0
+  - requests=2.31.0=py311h06a4308_1
+  - rfc3339-validator=0.1.4=py311h06a4308_0
+  - rfc3986-validator=0.1.1=py311h06a4308_0
+  - rpds-py=0.10.6=py311hb02cf49_0
+  - send2trash=1.8.2=py311h06a4308_0
+  - setuptools=68.2.2=py311h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - sniffio=1.3.0=py311h06a4308_0
+  - soupsieve=2.5=py311h06a4308_0
+  - sqlite=3.41.2=h5eee18b_0
+  - stack_data=0.2.0=pyhd3eb1b0_0
+  - terminado=0.17.1=py311h06a4308_0
+  - tinycss2=1.2.1=py311h06a4308_0
+  - tk=8.6.12=h1ccaba5_0
+  - tornado=6.3.3=py311h5eee18b_0
+  - traitlets=5.7.1=py311h06a4308_0
+  - typing-extensions=4.9.0=py311h06a4308_1
+  - typing_extensions=4.9.0=py311h06a4308_1
+  - urllib3=2.1.0=py311h06a4308_1
+  - wcwidth=0.2.5=pyhd3eb1b0_0
+  - webencodings=0.5.1=py311h06a4308_1
+  - websocket-client=0.58.0=py311h06a4308_4
+  - wheel=0.41.2=py311h06a4308_0
+  - xz=5.4.6=h5eee18b_0
+  - yaml=0.2.5=h7b6447c_0
+  - zeromq=4.3.5=h6a678d5_0
+  - zlib=1.2.13=h5eee18b_0
+  - pip:
+      - absl-py==2.1.0
+      - array-record==0.5.1
+      - astunparse==1.6.3
+      - cachetools==5.3.3
+      - cartopy==0.24.1
+      - cf-xarray==0.10.0
+      - cftime==1.6.4
+      - click==8.1.7
+      - contourpy==1.2.1
+      - cycler==0.12.1
+      - dm-tree==0.1.8
+      - etils==1.8.0
+      - filelock==3.13.4
+      - flatbuffers==24.3.25
+      - fonttools==4.51.0
+      - fsspec==2024.3.1
+      - gast==0.5.4
+      - gdown==5.1.0
+      - google-auth==2.29.0
+      - google-auth-oauthlib==1.2.0
+      - google-pasta==0.2.0
+      - googleapis-common-protos==1.63.0
+      - grpcio==1.62.2
+      - h5py==3.11.0
+      - importlib-resources==6.4.0
+      - joblib==1.4.2
+      - keras==3.3.3
+      - kiwisolver==1.4.5
+      - libclang==18.1.1
+      - llvmlite==0.43.0
+      - markdown==3.6
+      - markdown-it-py==3.0.0
+      - matplotlib==3.8.4
+      - mdurl==0.1.2
+      - ml-dtypes==0.2.0
+      - namex==0.0.8
+      - netcdf4==1.7.1.post2
+      - nltk==3.8.1
+      - numba==0.60.0
+      - numpy==1.26.4
+      - nvidia-cublas-cu12==12.3.4.1
+      - nvidia-cuda-cupti-cu12==12.3.101
+      - nvidia-cuda-nvcc-cu12==12.3.107
+      - nvidia-cuda-nvrtc-cu12==12.3.107
+      - nvidia-cuda-runtime-cu12==12.3.101
+      - nvidia-cudnn-cu12==8.9.7.29
+      - nvidia-cufft-cu12==11.0.12.1
+      - nvidia-curand-cu12==10.3.4.107
+      - nvidia-cusolver-cu12==11.5.4.101
+      - nvidia-cusparse-cu12==12.2.0.103
+      - nvidia-nccl-cu12==2.19.3
+      - nvidia-nvjitlink-cu12==12.3.101
+      - oauthlib==3.2.2
+      - opencv-python==4.10.0.84
+      - opt-einsum==3.3.0
+      - optree==0.11.0
+      - pandas==2.2.3
+      - pillow==10.3.0
+      - pip==24.0
+      - plotly==5.24.1
+      - promise==2.3
+      - protobuf==4.25.3
+      - pyasn1==0.6.0
+      - pyasn1-modules==0.4.0
+      - pyparsing==3.1.2
+      - pyproj==3.7.0
+      - pyshp==2.3.1
+      - regex==2024.5.10
+      - requests-oauthlib==2.0.0
+      - rich==13.7.1
+      - rsa==4.9
+      - scikit-learn==1.4.2
+      - scipy==1.13.0
+      - seaborn==0.13.2
+      - shapely==2.0.6
+      - sparse==0.15.4
+      - tenacity==9.0.0
+      - tensorboard==2.15.2
+      - tensorboard-data-server==0.7.2
+      - tensorflow==2.15.0.post1
+      - tensorflow-datasets==4.9.4
+      - tensorflow-estimator==2.15.0
+      - tensorflow-io-gcs-filesystem==0.36.0
+      - tensorflow-metadata==1.15.0
+      - termcolor==2.4.0
+      - threadpoolctl==3.5.0
+      - toml==0.10.2
+      - tqdm==4.66.2
+      - tzdata==2024.2
+      - werkzeug==3.0.2
+      - wrapt==1.14.1
+      - xarray==2024.9.0
+      - xesmf==0.8.8
+      - zipp==3.18.1

bluemath_tk/deeplearning/deep_environment.yml

@@ -0,0 +1,40 @@
+import numpy as np
+import keras.utils
+
+class MockDataGenerator(keras.utils.Sequence):
+    def __init__(self, 
+                 num_images: int, 
+                 input_frames: int = 1, 
+                 output_frames: int = 1, 
+                 batch_size: int = 8, 
+                 input_height: int = 256, 
+                 input_width: int = 256, 
+                 output_height: int = 256, 
+                 output_width: int = 256):
+        
+        self.input_height = input_height
+        self.input_width = input_width
+        self.output_height = output_height
+        self.output_width = output_width
+        self.input_frames = input_frames
+        self.output_frames = output_frames
+        self.batch_size = batch_size
+        self.num_images = num_images
+
+    @property
+    def num_batches(self):
+        return int(np.floor(self.num_images / self.batch_size))
+
+    def __len__(self) -> int:
+        """Returns the total number of batches."""
+        return self.num_batches
+
+    def __getitem__(self, idx: int) -> tuple[np.ndarray, np.ndarray]:
+        """Generates one batch of random data"""
+        # Generate random input and output data
+        inputs = np.random.rand(self.batch_size, self.input_height, self.input_width, self.input_frames)
+        outputs = np.random.rand(self.batch_size, self.output_height, self.output_width, self.output_frames)
+
+        return inputs, outputs
+
+

bluemath_tk/deeplearning/generators/__init__.py

@@ -0,0 +1,68 @@
+import keras
+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
+        )(x)
+        x = layers.Conv2D(width, kernel_size=3, padding="same")(x)
+        x = layers.Add()([x, residual])
+        return x
+
+    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]]:
+        
+        x, skips = x
+        for _ in range(block_depth):
+            x = ResidualBlock(width)(x)
+            skips.append(x)
+        x = layers.AveragePooling2D(pool_size=2)(x)
+        return x, skips
+
+    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:
+    """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)
+
+    skips = []
+    for width in down_widths[:-1]:
+        x, skips = DownBlock(width, block_depth)([x, skips])
+
+    for _ in range(block_depth):
+        x = ResidualBlock(down_widths[-1])(x)
+
+    for width in up_widths[1:]:
+        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/generators/mockDataGenerator.py

@@ -0,0 +1,39 @@
+import keras
+from models import resnet_model
+from generators.mockDataGenerator import MockDataGenerator
+
+# instantiate model class (load memory)
+model = resnet_model.get_model(
+              image_height = 64,
+              image_width = 64,
+              input_frames =  1,
+              output_frames =  1)
+
+# print summary of the model
+print(model.summary())
+
+# instantiate generator class
+train_generator = MockDataGenerator(num_images=5000,
+                                       input_height = 64,
+                                       input_width = 64,
+                                       output_height = 64,
+                                       output_width = 64,
+                                       batch_size=1)
+# define oprimizer
+optimizer=keras.optimizers.AdamW
+model.compile(
+    optimizer=optimizer(
+        learning_rate=1e-4, weight_decay=1e-5
+    ),
+    loss=keras.losses.mean_squared_error,
+)
+
+# start the train loop with the fit method
+history = model.fit(
+    train_generator,
+    initial_epoch = 0,
+    epochs=20,
+    steps_per_epoch=500)
+
+
+print("training complete")