Address issue 44 and 63

src/gz21_ocean_momentum/analysis/utils.py

@@ -14,6 +14,7 @@
 import pandas as pd
 from analysis.analysis import TimeSeriesForPoint
 import xarray as xr
+from typing import Optional
 from scipy.ndimage import gaussian_filter
 from data.pangeo_catalog import get_patch, get_whole_data
 from cartopy.crs import PlateCarree
@@ -165,7 +166,7 @@ def onClick(event):
     fig.canvas.mpl_connect("button_press_event", onClick)
 
 
-def sample(data: np.ndarray, step_time: int = 1, nb_per_time: int = 5):
+def sample(data: np.ndarray, step_time: int = 1, nb_per_time: int = 5, random_state: Optional[int] = None):
     """Samples points from the data, where it is assumed that the data
     is 4-D, with the first dimension representing time , the second
     the channel, and the others representing spatial dimensions.
@@ -185,6 +186,9 @@ def sample(data: np.ndarray, step_time: int = 1, nb_per_time: int = 5):
     :nb_per_time: int,
         Number of points used (chosen randomly according to a uniform
         distribution over the spatial domain) for each image.
+
+    :random_state: int, optional,
+        Random state used for the random number generator.
 
 
     Returns
@@ -194,6 +198,7 @@ def sample(data: np.ndarray, step_time: int = 1, nb_per_time: int = 5):
     """
     if data.ndim != 4:
         raise ValueError("The data is expected to have 4 dimensions.")
+    np.random.seed(random_state)
     n_times, n_channels, n_x, n_y = data.shape
     time_indices = np.arange(0, n_times, step_time)
     x_indices = np.random.randint(0, n_x, (time_indices.shape[0], 2, nb_per_time))

src/gz21_ocean_momentum/models/fully_conv_net.py

@@ -76,6 +76,8 @@ def __init__(
         # store in_chans as attribute
         self._n_in_channels = in_chans
 
+        self._final_transformation = lambda x: x
+
 
     @staticmethod
     def _process_padding(padding: Optional[str] = None) -> Tuple[int, int]:

src/gz21_ocean_momentum/utils.py

@@ -6,8 +6,12 @@
 @author: arthur
 """
 
+import os
 import mlflow
 from mlflow.tracking import client
+import torch
+import random
+import numpy as np
 import pandas as pd
 import pickle
 import gz21_ocean_momentum.models as models
@@ -126,3 +130,14 @@ def pickle_artifact(run_id: str, path: str):
     file = client.download_artifacts(run_id, path)
     f = open(file, "rb")
     return pickle.load(f)
+
+
+def seed_all(seed: int = 0):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    # seed torch
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True

tests/models/test_fully_conv_net.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from gz21_ocean_momentum.models.fully_conv_net import *
+from gz21_ocean_momentum.utils import seed_all
 import torch
 import numpy as np
 
@@ -10,7 +11,6 @@ def test_construct_valid():
     Simple check migrated from `models.models1`.
     """
     net = FullyCNN()
-    net._final_transformation = lambda x: x
     input_ = torch.randint(0, 10, (17, 2, 35, 30)).to(dtype=torch.float)
     input_[0, 0, 0, 0] = np.nan
     output = net(input_)