Initial commit: FourCastNet source code only

utils/data_loader_multifiles.py

+#BSD 3-Clause License
+#
+#Copyright (c) 2022, FourCastNet authors
+#All rights reserved.
+#
+#Redistribution and use in source and binary forms, with or without
+#modification, are permitted provided that the following conditions are met:
+#
+#1. Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+#2. Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+#3. Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+#The code was authored by the following people:
+#
+#Jaideep Pathak - NVIDIA Corporation
+#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
+#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
+#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory 
+#Ashesh Chattopadhyay - Rice University 
+#Morteza Mardani - NVIDIA Corporation 
+#Thorsten Kurth - NVIDIA Corporation 
+#David Hall - NVIDIA Corporation 
+#Zongyi Li - California Institute of Technology, NVIDIA Corporation 
+#Kamyar Azizzadenesheli - Purdue University 
+#Pedram Hassanzadeh - Rice University 
+#Karthik Kashinath - NVIDIA Corporation 
+#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
+
+import logging
+import glob
+import torch
+import random
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch import Tensor
+import h5py
+import math
+#import cv2
+from utils.img_utils import reshape_fields, reshape_precip
+
+
+def get_data_loader(params, files_pattern, distributed, train):
+
+  dataset = GetDataset(params, files_pattern, train)
+  sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
+  
+  dataloader = DataLoader(dataset,
+                          batch_size=int(params.batch_size),
+                          num_workers=params.num_data_workers,
+                          shuffle=False, #(sampler is None),
+                          sampler=sampler if train else None,
+                          drop_last=True,
+                          pin_memory=torch.cuda.is_available())
+
+  if train:
+    return dataloader, dataset, sampler
+  else:
+    return dataloader, dataset
+
+class GetDataset(Dataset):
+  def __init__(self, params, location, train):
+    self.params = params
+    self.location = location
+    self.train = train
+    self.dt = params.dt
+    self.n_history = params.n_history
+    self.in_channels = np.array(params.in_channels)
+    self.out_channels = np.array(params.out_channels)
+    self.n_in_channels = len(self.in_channels)
+    self.n_out_channels = len(self.out_channels)
+    self.crop_size_x = params.crop_size_x
+    self.crop_size_y = params.crop_size_y
+    self.roll = params.roll
+    self._get_files_stats()
+    self.two_step_training = params.two_step_training
+    self.orography = params.orography
+    self.precip = True if "precip" in params else False
+    self.add_noise = params.add_noise if train else False
+
+    if self.precip:
+        path = params.precip+'/train' if train else params.precip+'/test'
+        self.precip_paths = glob.glob(path + "/*.h5")
+        self.precip_paths.sort()
+
+    try:
+        self.normalize = params.normalize
+    except:
+        self.normalize = True #by default turn on normalization if not specified in config
+
+    if self.orography:
+      self.orography_path = params.orography_path
+
+  def _get_files_stats(self):
+    self.files_paths = glob.glob(self.location + "/*.h5")
+    self.files_paths.sort()
+    self.n_years = len(self.files_paths)
+    with h5py.File(self.files_paths[0], 'r') as _f:
+        logging.info("Getting file stats from {}".format(self.files_paths[0]))
+        self.n_samples_per_year = _f['fields'].shape[0]
+        #original image shape (before padding)
+        self.img_shape_x = _f['fields'].shape[2] -1#just get rid of one of the pixels
+        self.img_shape_y = _f['fields'].shape[3]
+
+    self.n_samples_total = self.n_years * self.n_samples_per_year
+    self.files = [None for _ in range(self.n_years)]
+    self.precip_files = [None for _ in range(self.n_years)]
+    logging.info("Number of samples per year: {}".format(self.n_samples_per_year))
+    logging.info("Found data at path {}. Number of examples: {}. Image Shape: {} x {} x {}".format(self.location, self.n_samples_total, self.img_shape_x, self.img_shape_y, self.n_in_channels))
+    logging.info("Delta t: {} hours".format(6*self.dt))
+    logging.info("Including {} hours of past history in training at a frequency of {} hours".format(6*self.dt*self.n_history, 6*self.dt))
+
+
+  def _open_file(self, year_idx):
+    _file = h5py.File(self.files_paths[year_idx], 'r')
+    self.files[year_idx] = _file['fields']  
+    if self.orography:
+      _orog_file = h5py.File(self.orography_path, 'r')
+      self.orography_field = _orog_file['orog']
+    if self.precip:
+      self.precip_files[year_idx] = h5py.File(self.precip_paths[year_idx], 'r')['tp']
+    
+  
+  def __len__(self):
+    return self.n_samples_total
+
+
+  def __getitem__(self, global_idx):
+    year_idx = int(global_idx/self.n_samples_per_year) #which year we are on
+    local_idx = int(global_idx%self.n_samples_per_year) #which sample in that year we are on - determines indices for centering
+
+    y_roll = np.random.randint(0, 1440) if self.train else 0#roll image in y direction
+
+    #open image file
+    if self.files[year_idx] is None:
+        self._open_file(year_idx)
+
+    if not self.precip:
+      #if we are not at least self.dt*n_history timesteps into the prediction
+      if local_idx < self.dt*self.n_history:
+          local_idx += self.dt*self.n_history
+
+      #if we are on the last image in a year predict identity, else predict next timestep
+      step = 0 if local_idx >= self.n_samples_per_year-self.dt else self.dt
+    else:
+      inp_local_idx = local_idx
+      tar_local_idx = local_idx
+      #if we are on the last image in a year predict identity, else predict next timestep
+      step = 0 if tar_local_idx >= self.n_samples_per_year-self.dt else self.dt
+      # first year has 2 missing samples in precip (they are first two time points)
+      if year_idx == 0:
+        lim = 1458
+        local_idx = local_idx%lim 
+        inp_local_idx = local_idx + 2
+        tar_local_idx = local_idx
+        step = 0 if tar_local_idx >= lim-self.dt else self.dt
+
+    #if two_step_training flag is true then ensure that local_idx is not the last or last but one sample in a year
+    if self.two_step_training:
+        if local_idx >= self.n_samples_per_year - 2*self.dt:
+            #set local_idx to last possible sample in a year that allows taking two steps forward
+            local_idx = self.n_samples_per_year - 3*self.dt
+
+    if self.train and self.roll:
+      y_roll = random.randint(0, self.img_shape_y)
+    else:
+      y_roll = 0
+
+    if self.orography:
+        orog = self.orography_field[0:720] 
+    else:
+        orog = None
+
+    if self.train and (self.crop_size_x or self.crop_size_y):
+      rnd_x = random.randint(0, self.img_shape_x-self.crop_size_x)
+      rnd_y = random.randint(0, self.img_shape_y-self.crop_size_y)    
+    else: 
+      rnd_x = 0
+      rnd_y = 0
+      
+    if self.precip:
+      return reshape_fields(self.files[year_idx][inp_local_idx, self.in_channels], 'inp', self.crop_size_x, self.crop_size_y, rnd_x, rnd_y,self.params, y_roll, self.train), \
+                reshape_precip(self.precip_files[year_idx][tar_local_idx+step], 'tar', self.crop_size_x, self.crop_size_y, rnd_x, rnd_y, self.params, y_roll, self.train)
+    else:
+        if self.two_step_training:
+            return reshape_fields(self.files[year_idx][(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.in_channels], 'inp', self.crop_size_x, self.crop_size_y, rnd_x, rnd_y,self.params, y_roll, self.train, self.normalize, orog, self.add_noise), \
+                    reshape_fields(self.files[year_idx][local_idx + step:local_idx + step + 2, self.out_channels], 'tar', self.crop_size_x, self.crop_size_y, rnd_x, rnd_y, self.params, y_roll, self.train, self.normalize, orog)
+        else:
+            return reshape_fields(self.files[year_idx][(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.in_channels], 'inp', self.crop_size_x, self.crop_size_y, rnd_x, rnd_y,self.params, y_roll, self.train, self.normalize, orog, self.add_noise), \
+                    reshape_fields(self.files[year_idx][local_idx + step, self.out_channels], 'tar', self.crop_size_x, self.crop_size_y, rnd_x, rnd_y, self.params, y_roll, self.train, self.normalize, orog)
+
+
+
+
+
+    
+
+
+  
+    
+

utils/date_time_to_hours.py

+#BSD 3-Clause License
+#
+#Copyright (c) 2022, FourCastNet authors
+#All rights reserved.
+#
+#Redistribution and use in source and binary forms, with or without
+#modification, are permitted provided that the following conditions are met:
+#
+#1. Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+#2. Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+#3. Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+#The code was authored by the following people:
+#
+#Jaideep Pathak - NVIDIA Corporation
+#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
+#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
+#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory 
+#Ashesh Chattopadhyay - Rice University 
+#Morteza Mardani - NVIDIA Corporation 
+#Thorsten Kurth - NVIDIA Corporation 
+#David Hall - NVIDIA Corporation 
+#Zongyi Li - California Institute of Technology, NVIDIA Corporation 
+#Kamyar Azizzadenesheli - Purdue University 
+#Pedram Hassanzadeh - Rice University 
+#Karthik Kashinath - NVIDIA Corporation 
+#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
+
+import numpy as np
+from datetime import datetime
+
+
+#day_of_year = datetime.now().timetuple().tm_yday  # returns 1 for January 1st
+#time_tuple = datetime.now().timetuple()
+date_strings = ["2016-01-01 00:00:00", "2016-09-13 00:00:00", "2016-09-17 00:00:00", "2016-09-21 00:00:00", "2016-09-25 00:00:00", "2016-09-29 00:00:00", "2016-10-03 00:00:00", "2016-10-07 00:00:00"]
+
+ics = []
+
+for date_ in date_strings:
+    date_obj = datetime.strptime(date_, '%Y-%m-%d %H:%M:%S') #datetime.fromisoformat(date_) 
+    print(date_obj.timetuple())
+    day_of_year = date_obj.timetuple().tm_yday - 1
+    hour_of_day = date_obj.timetuple().tm_hour
+    hours_since_jan_01_epoch = 24*day_of_year + hour_of_day
+    ics.append(int(hours_since_jan_01_epoch/6))
+    print(day_of_year, hour_of_day)
+    print("hours = ", hours_since_jan_01_epoch )
+    print("steps = ", hours_since_jan_01_epoch/6) 
+
+
+print(ics)
+
+ics = []
+for date_ in date_strings:
+    date_obj = datetime.fromisoformat(date_) #datetime.strptime(date_, '%Y-%m-%d %H:%M:%S') #datetime.fromisoformat(date_) 
+    print(date_obj.timetuple())
+    day_of_year = date_obj.timetuple().tm_yday - 1
+    hour_of_day = date_obj.timetuple().tm_hour
+    hours_since_jan_01_epoch = 24*day_of_year + hour_of_day
+    ics.append(int(hours_since_jan_01_epoch/6))
+    print(day_of_year, hour_of_day)
+    print("hours = ", hours_since_jan_01_epoch )
+    print("steps = ", hours_since_jan_01_epoch/6) 
+
+
+print(ics)
+

utils/img_utils.py

+#BSD 3-Clause License
+#
+#Copyright (c) 2022, FourCastNet authors
+#All rights reserved.
+#
+#Redistribution and use in source and binary forms, with or without
+#modification, are permitted provided that the following conditions are met:
+#
+#1. Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+#2. Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+#3. Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+#The code was authored by the following people:
+#
+#Jaideep Pathak - NVIDIA Corporation
+#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
+#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
+#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory 
+#Ashesh Chattopadhyay - Rice University 
+#Morteza Mardani - NVIDIA Corporation 
+#Thorsten Kurth - NVIDIA Corporation 
+#David Hall - NVIDIA Corporation 
+#Zongyi Li - California Institute of Technology, NVIDIA Corporation 
+#Kamyar Azizzadenesheli - Purdue University 
+#Pedram Hassanzadeh - Rice University 
+#Karthik Kashinath - NVIDIA Corporation 
+#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
+
+import logging
+import glob
+from types import new_class
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import random
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch import Tensor
+import h5py
+import math
+import torchvision.transforms.functional as TF
+import matplotlib
+import matplotlib.pyplot as plt
+
+class PeriodicPad2d(nn.Module):
+    """ 
+        pad longitudinal (left-right) circular 
+        and pad latitude (top-bottom) with zeros
+    """
+    def __init__(self, pad_width):
+       super(PeriodicPad2d, self).__init__()
+       self.pad_width = pad_width
+
+    def forward(self, x):
+        # pad left and right circular
+        out = F.pad(x, (self.pad_width, self.pad_width, 0, 0), mode="circular") 
+        # pad top and bottom zeros
+        out = F.pad(out, (0, 0, self.pad_width, self.pad_width), mode="constant", value=0) 
+        return out
+
+def reshape_fields(img, inp_or_tar, crop_size_x, crop_size_y,rnd_x, rnd_y, params, y_roll, train, normalize=True, orog=None, add_noise=False):
+    #Takes in np array of size (n_history+1, c, h, w) and returns torch tensor of size ((n_channels*(n_history+1), crop_size_x, crop_size_y)
+
+    if len(np.shape(img)) ==3:
+      img = np.expand_dims(img, 0)
+
+    
+    img = img[:, :, 0:720] #remove last pixel
+    n_history = np.shape(img)[0] - 1
+    img_shape_x = np.shape(img)[-2]
+    img_shape_y = np.shape(img)[-1]
+    n_channels = np.shape(img)[1] #this will either be N_in_channels or N_out_channels
+    channels = params.in_channels if inp_or_tar =='inp' else params.out_channels
+    means = np.load(params.global_means_path)[:, channels]
+    stds = np.load(params.global_stds_path)[:, channels]
+    if crop_size_x == None:
+        crop_size_x = img_shape_x
+    if crop_size_y == None:
+        crop_size_y = img_shape_y
+
+    if normalize:
+        if params.normalization == 'minmax':
+          raise Exception("minmax not supported. Use zscore")
+        elif params.normalization == 'zscore':
+          img -=means
+          img /=stds
+
+    if params.add_grid:
+        if inp_or_tar == 'inp':
+            if params.gridtype == 'linear':
+                assert params.N_grid_channels == 2, "N_grid_channels must be set to 2 for gridtype linear"
+                x = np.meshgrid(np.linspace(-1, 1, img_shape_x))
+                y = np.meshgrid(np.linspace(-1, 1, img_shape_y))
+                grid_x, grid_y = np.meshgrid(y, x)
+                grid = np.stack((grid_x, grid_y), axis = 0)
+            elif params.gridtype == 'sinusoidal':
+                assert params.N_grid_channels == 4, "N_grid_channels must be set to 4 for gridtype sinusoidal"
+                x1 = np.meshgrid(np.sin(np.linspace(0, 2*np.pi, img_shape_x)))
+                x2 = np.meshgrid(np.cos(np.linspace(0, 2*np.pi, img_shape_x)))
+                y1 = np.meshgrid(np.sin(np.linspace(0, 2*np.pi, img_shape_y)))
+                y2 = np.meshgrid(np.cos(np.linspace(0, 2*np.pi, img_shape_y)))
+                grid_x1, grid_y1 = np.meshgrid(y1, x1)
+                grid_x2, grid_y2 = np.meshgrid(y2, x2)
+                grid = np.expand_dims(np.stack((grid_x1, grid_y1, grid_x2, grid_y2), axis = 0), axis = 0)
+            img = np.concatenate((img, grid), axis = 1 )
+
+    if params.orography and inp_or_tar == 'inp':
+        img = np.concatenate((img, np.expand_dims(orog, axis = (0,1) )), axis = 1)
+        n_channels += 1
+
+    if params.roll:
+        img = np.roll(img, y_roll, axis = -1)
+
+    if train and (crop_size_x or crop_size_y):
+        img = img[:,:,rnd_x:rnd_x+crop_size_x, rnd_y:rnd_y+crop_size_y]
+
+    if inp_or_tar == 'inp':
+        img = np.reshape(img, (n_channels*(n_history+1), crop_size_x, crop_size_y))
+    elif inp_or_tar == 'tar':
+        if params.two_step_training:
+            img = np.reshape(img, (n_channels*2, crop_size_x, crop_size_y))
+        else:
+            img = np.reshape(img, (n_channels, crop_size_x, crop_size_y))
+
+    if add_noise:
+        img = img + np.random.normal(0, scale=params.noise_std, size=img.shape)
+
+    return torch.as_tensor(img)
+          
+def reshape_precip(img, inp_or_tar, crop_size_x, crop_size_y,rnd_x, rnd_y, params, y_roll, train, normalize=True):
+
+    if len(np.shape(img)) ==2:
+      img = np.expand_dims(img, 0)
+
+    img = img[:,:720,:]
+    img_shape_x = img.shape[-2]
+    img_shape_y = img.shape[-1]
+    n_channels = 1
+    if crop_size_x == None:
+        crop_size_x = img_shape_x
+    if crop_size_y == None:
+        crop_size_y = img_shape_y
+
+    if normalize:
+        eps = params.precip_eps
+        img = np.log1p(img/eps)
+    if params.add_grid:
+        if inp_or_tar == 'inp':
+            if params.gridtype == 'linear':
+                assert params.N_grid_channels == 2, "N_grid_channels must be set to 2 for gridtype linear"
+                x = np.meshgrid(np.linspace(-1, 1, img_shape_x))
+                y = np.meshgrid(np.linspace(-1, 1, img_shape_y))
+                grid_x, grid_y = np.meshgrid(y, x)
+                grid = np.stack((grid_x, grid_y), axis = 0)
+            elif params.gridtype == 'sinusoidal':
+                assert params.N_grid_channels == 4, "N_grid_channels must be set to 4 for gridtype sinusoidal"
+                x1 = np.meshgrid(np.sin(np.linspace(0, 2*np.pi, img_shape_x)))
+                x2 = np.meshgrid(np.cos(np.linspace(0, 2*np.pi, img_shape_x)))
+                y1 = np.meshgrid(np.sin(np.linspace(0, 2*np.pi, img_shape_y)))
+                y2 = np.meshgrid(np.cos(np.linspace(0, 2*np.pi, img_shape_y)))
+                grid_x1, grid_y1 = np.meshgrid(y1, x1)
+                grid_x2, grid_y2 = np.meshgrid(y2, x2)
+                grid = np.expand_dims(np.stack((grid_x1, grid_y1, grid_x2, grid_y2), axis = 0), axis = 0)
+            img = np.concatenate((img, grid), axis = 1 )
+
+    if params.roll:
+        img = np.roll(img, y_roll, axis = -1)
+
+    if train and (crop_size_x or crop_size_y):
+        img = img[:,rnd_x:rnd_x+crop_size_x, rnd_y:rnd_y+crop_size_y]
+
+    img = np.reshape(img, (n_channels, crop_size_x, crop_size_y))
+    return torch.as_tensor(img)
+
+
+def vis_precip(fields):
+    pred, tar = fields
+    fig, ax = plt.subplots(1, 2, figsize=(24,12))
+    ax[0].imshow(pred, cmap="coolwarm")
+    ax[0].set_title("tp pred")
+    ax[1].imshow(tar, cmap="coolwarm")
+    ax[1].set_title("tp tar")
+    fig.tight_layout()
+    return fig
+
+
+    
+

utils/logging_utils.py

+import os
+import logging
+
+_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+
+def config_logger(log_level=logging.INFO):
+  logging.basicConfig(format=_format, level=log_level)
+
+def log_to_file(logger_name=None, log_level=logging.INFO, log_filename='tensorflow.log'):
+
+  if not os.path.exists(os.path.dirname(log_filename)):
+    os.makedirs(os.path.dirname(log_filename))
+
+  if logger_name is not None:
+    log = logging.getLogger(logger_name)
+  else:
+    log = logging.getLogger()
+
+  fh = logging.FileHandler(log_filename)
+  fh.setLevel(log_level)
+  fh.setFormatter(logging.Formatter(_format))
+  log.addHandler(fh)
+
+def log_versions():
+  import torch
+  import subprocess
+
+  logging.info('--------------- Versions ---------------')
+  logging.info('git branch: ' + str(subprocess.check_output(['git', 'branch']).strip()))
+  logging.info('git hash: ' + str(subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip()))
+  logging.info('Torch: ' + str(torch.__version__))
+  logging.info('----------------------------------------')

utils/weighted_acc_rmse.py

+#BSD 3-Clause License
+#
+#Copyright (c) 2022, FourCastNet authors
+#All rights reserved.
+#
+#Redistribution and use in source and binary forms, with or without
+#modification, are permitted provided that the following conditions are met:
+#
+#1. Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+#2. Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+#3. Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+#The code was authored by the following people:
+#
+#Jaideep Pathak - NVIDIA Corporation
+#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
+#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
+#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory 
+#Ashesh Chattopadhyay - Rice University 
+#Morteza Mardani - NVIDIA Corporation 
+#Thorsten Kurth - NVIDIA Corporation 
+#David Hall - NVIDIA Corporation 
+#Zongyi Li - California Institute of Technology, NVIDIA Corporation 
+#Kamyar Azizzadenesheli - Purdue University 
+#Pedram Hassanzadeh - Rice University 
+#Karthik Kashinath - NVIDIA Corporation 
+#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
+
+import os
+import time
+import numpy as np
+import argparse
+import h5py
+#from netCDF4 import Dataset as DS
+from collections import OrderedDict
+from utils import logging_utils
+logging_utils.config_logger()
+#from utils.YParams import YParams
+from utils.data_loader_multifiles import get_data_loader
+import wandb
+import torch
+import warnings
+
+def unlog_tp(x, eps=1E-5):
+#    return np.exp(x + np.log(eps)) - eps
+    return eps*(np.exp(x)-1)
+
+def unlog_tp_torch(x, eps=1E-5):
+#    return torch.exp(x + torch.log(eps)) - eps
+    return eps*(torch.exp(x)-1)
+
+def mean(x, axis = None):
+    #spatial mean
+    y = np.sum(x, axis) / np.size(x, axis)
+    return y
+
+def lat_np(j, num_lat):
+    return 90 - j * 180/(num_lat-1)
+
+def weighted_acc(pred,target, weighted  = True):
+    #takes in shape [1, num_lat, num_long]
+    if len(pred.shape) ==2:
+        pred = np.expand_dims(pred, 0)
+    if len(target.shape) ==2:
+        target = np.expand_dims(target, 0)
+    
+    num_lat = np.shape(pred)[1]
+    num_long = np.shape(target)[2]
+#    pred -= mean(pred)
+#    target -= mean(target)
+    s = np.sum(np.cos(np.pi/180* lat_np(np.arange(0, num_lat), num_lat)))
+    weight = np.expand_dims(latitude_weighting_factor(np.arange(0, num_lat), num_lat, s), -1) if weighted else 1
+    r = (weight*pred*target).sum() /np.sqrt((weight*pred*pred).sum() * (weight*target*target).sum())
+    return r
+
+def weighted_acc_masked(pred,target, weighted  = True, maskarray=1):
+    #takes in shape [1, num_lat, num_long]
+    if len(pred.shape) ==2:
+        pred = np.expand_dims(pred, 0)
+    if len(target.shape) ==2:
+        target = np.expand_dims(target, 0)
+
+    num_lat = np.shape(pred)[1]
+    num_long = np.shape(target)[2]
+    pred -= mean(pred)
+    target -= mean(target)
+    s = np.sum(np.cos(np.pi/180* lat(np.arange(0, num_lat), num_lat)))
+    weight = np.expand_dims(latitude_weighting_factor(np.arange(0, num_lat), num_lat, s), -1) if weighted else 1
+    r = (maskarray*weight*pred*target).sum() /np.sqrt((maskarray*weight*pred*pred).sum() * (maskarray*weight*target*target).sum())
+    return r
+
+def weighted_rmse(pred, target):
+    if len(pred.shape) ==2:
+        pred = np.expand_dims(pred, 0)
+    if len(target.shape) ==2:
+        target = np.expand_dims(target, 0)
+    #takes in arrays of size [1, h, w]  and returns latitude-weighted rmse
+    num_lat = np.shape(pred)[1]
+    num_long = np.shape(target)[2]
+    s = np.sum(np.cos(np.pi/180* lat_np(np.arange(0, num_lat), num_lat)))
+    weight = np.expand_dims(latitude_weighting_factor(np.arange(0, num_lat), num_lat, s), -1)
+    return np.sqrt(1/num_lat * 1/num_long  * np.sum(np.dot(weight.T, (pred[0] - target[0])**2)))
+
+def latitude_weighting_factor(j, num_lat, s):
+    return num_lat*np.cos(np.pi/180. * lat_np(j, num_lat))/s
+
+def top_quantiles_error(pred, target):
+    if len(pred.shape) ==2:
+        pred = np.expand_dims(pred, 0)
+    if len(target.shape) ==2:
+        target = np.expand_dims(target, 0)
+    qs = 100
+    qlim = 5
+    qcut = 0.1
+    qtile = 1. - np.logspace(-qlim, -qcut, num=qs)
+    P_tar = np.quantile(target, q=qtile, axis=(1,2))
+    P_pred = np.quantile(pred, q=qtile, axis=(1,2))
+    return np.mean(P_pred - P_tar, axis=0)
+
+
+# torch version for rmse comp
+@torch.jit.script
+def lat(j: torch.Tensor, num_lat: int) -> torch.Tensor:
+    return 90. - j * 180./float(num_lat-1)
+
+@torch.jit.script
+def latitude_weighting_factor_torch(j: torch.Tensor, num_lat: int, s: torch.Tensor) -> torch.Tensor:
+    return num_lat * torch.cos(3.1416/180. * lat(j, num_lat))/s
+
+@torch.jit.script
+def weighted_rmse_torch_channels(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    #takes in arrays of size [n, c, h, w]  and returns latitude-weighted rmse for each chann
+    num_lat = pred.shape[2]
+    #num_long = target.shape[2]
+    lat_t = torch.arange(start=0, end=num_lat, device=pred.device)
+
+    s = torch.sum(torch.cos(3.1416/180. * lat(lat_t, num_lat)))
+    weight = torch.reshape(latitude_weighting_factor_torch(lat_t, num_lat, s), (1, 1, -1, 1))
+    result = torch.sqrt(torch.mean(weight * (pred - target)**2., dim=(-1,-2)))
+    return result
+
+@torch.jit.script
+def weighted_rmse_torch(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    result = weighted_rmse_torch_channels(pred, target)
+    return torch.mean(result, dim=0)
+
+@torch.jit.script
+def weighted_acc_masked_torch_channels(pred: torch.Tensor, target: torch.Tensor, maskarray: torch.Tensor) -> torch.Tensor:
+    #takes in arrays of size [n, c, h, w]  and returns latitude-weighted acc
+    num_lat = pred.shape[2]
+    lat_t = torch.arange(start=0, end=num_lat, device=pred.device)
+    s = torch.sum(torch.cos(3.1416/180. * lat(lat_t, num_lat)))
+    weight = torch.reshape(latitude_weighting_factor_torch(lat_t, num_lat, s), (1, 1, -1, 1))
+    result = torch.sum(maskarray * weight * pred * target, dim=(-1,-2)) / torch.sqrt(torch.sum(maskarray * weight * pred * pred, dim=(-1,-2)) * torch.sum(maskarray * weight * target *  target, dim=(-1,-2)))
+    return result
+
+@torch.jit.script
+def weighted_acc_torch_channels(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    #takes in arrays of size [n, c, h, w]  and returns latitude-weighted acc
+    num_lat = pred.shape[2]
+    #num_long = target.shape[2]
+    lat_t = torch.arange(start=0, end=num_lat, device=pred.device)
+    s = torch.sum(torch.cos(3.1416/180. * lat(lat_t, num_lat)))
+    weight = torch.reshape(latitude_weighting_factor_torch(lat_t, num_lat, s), (1, 1, -1, 1))
+    result = torch.sum(weight * pred * target, dim=(-1,-2)) / torch.sqrt(torch.sum(weight * pred * pred, dim=(-1,-2)) * torch.sum(weight * target *
+    target, dim=(-1,-2)))
+    return result
+
+@torch.jit.script
+def weighted_acc_torch(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    result = weighted_acc_torch_channels(pred, target)
+    return torch.mean(result, dim=0)
+
+@torch.jit.script
+def unweighted_acc_torch_channels(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    result = torch.sum(pred * target, dim=(-1,-2)) / torch.sqrt(torch.sum(pred * pred, dim=(-1,-2)) * torch.sum(target *
+    target, dim=(-1,-2)))
+    return result
+
+@torch.jit.script
+def unweighted_acc_torch(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    result = unweighted_acc_torch_channels(pred, target)
+    return torch.mean(result, dim=0)
+
+@torch.jit.script
+def top_quantiles_error_torch(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    qs = 100
+    qlim = 3
+    qcut = 0.1
+    n, c, h, w = pred.size()
+    qtile = 1. - torch.logspace(-qlim, -qcut, steps=qs, device=pred.device)
+    P_tar = torch.quantile(target.view(n,c,h*w), q=qtile, dim=-1)
+    P_pred = torch.quantile(pred.view(n,c,h*w), q=qtile, dim=-1)
+    return torch.mean(P_pred - P_tar, dim=0)