corrections

Author: paulaharder

README.md

@@ -4,6 +4,16 @@ This code belongs to a paper currently under review, a preprint can be found at:
 
 Abstract: *The availability of reliable, high-resolution climate and weather data is important to inform long-term decisions on climate adaptation and mitigation and to guide rapid responses to extreme events. Forecasting models are limited by computational costs and, therefore, often generate coarse-resolution predictions. Statistical downscaling can provide an efficient method of upsampling low-resolution data. In this field, deep learning has been applied successfully, often using image super-resolution methods from computer vision. However, despite achieving visually compelling results in some cases, such models frequently violate conservation laws when predicting physical variables. In order to conserve physical quantities, we develop methods that guarantee physical constraints are satisfied by a deep learning downscaling model while also improving their performance according to traditional metrics. We compare different constraining approaches and demonstrate their applicability across different neural architectures as well as a variety of climate and weather data sets. While our novel methodologies enable faster and more accurate climate predictions, we also show how they can improve super-resolution for satellite data and standard data sets.*
 
+## Setup
+
+Clone the repository and install the requirements
+```sh
+$ git clone https://github.com/RolnickLab/constrained-downscaling.git
+$ cd constrained-downscaling
+$ conda env create -f requirements.yml
+$ conda activate constrained-ds
+```
+
 ## Get the data
 
 One of our main data sets, ERA5 total columnt water, 4x upsampling, can be downloaded in a ML-ready form at: https://drive.google.com/file/d/1IENhP1-aTYyqOkRcnmCIvxXkvUW2Qbdx/view?usp=sharing
@@ -22,38 +32,29 @@ $ rm era5_sr_data.zip
 
 Other data sets are available upon request from the author or can be generated by using public sources for ERA5 (https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-single-levels?tab=form.) and NorESM (https://esgf-index1.ceda.ac.uk/search/cmip6-ceda/) data.
 
-## Setup
-
-Clone the repository and install the requirements
-```sh
-$ git clone https://github.com/RolnickLab/constrained-downscaling.git
-$ cd constrained-downscaling
-$ conda env create -f requirements.yml
-$ conda activate constrained_ds
-```
 
 ## Run training 
 
 To run our standard CNN withour constrained run
 
 ```sh
-$ python main.py --dataset era5_twc --model cnn --model_id twc_cnn_noconstraints --constraints none
+$ python main.py --dataset era5_sr_data --model cnn --model_id twc_cnn_noconstraints --constraints none
 ```
 
 to run with softmax constraining (hard constraining) run
 
 ```sh
-$ python main.py --dataset era5_twc --model cnn --model_id twc_cnn_softmaxconstraints --constraints softmax
+$ python main.py --dataset era5_sr_data --model cnn --model_id twc_cnn_softmaxconstraints --constraints softmax
 ```
 
 to run with soft constraining run, with a factor of alpha run
 
 ```sh
-$ python main.py --dataset era5_twc --model cnn --model_id twc_cnn_softconstraints --constraints soft --loss mass_constraints --alpha 0.99
+$ python main.py --dataset era5_sr_data --model cnn --model_id twc_cnn_softconstraints --constraints soft --loss mass_constraints --alpha 0.99
 ```
 
 For other setups: 
---model can be either cnn, gan, convgru, flowconvgru
+--model can be either cnn, gan, convgru, flowconvgru (last two require different data sets)
 --constraints can be none, softmax, gh, mult, add, soft
 other arguents are --epochs, --lr (learning rate), --number_residual_blocks, --weight_decay
 

main.py

@@ -2,6 +2,7 @@
 from utils import load_data
 import numpy as np
 import argparse
+import os
 import torch
 
 def add_arguments():
@@ -22,10 +23,15 @@ def add_arguments():
  parser.add_argument("--alpha", default=0.99, type=float)
  parser.add_argument("--test_val_train", default="val")
  parser.add_argument("--training_evalonly", default="training")
+ parser.add_argument("--dim_channels", default=1, type=int)
  return parser.parse_args()
 
 def main(args):
  #load data
+ if not os.path.exists('./models'):
+ os.makedirs('./models')
+ if not os.path.exists('./data/prediction'):
+ os.makedirs('./data/prediction')
  if args.training_evalonly == 'training':
  data = load_data(args)
  #run training

models.py

@@ -79,18 +79,15 @@ def forward(self, y, lr):
  return out 
 
 class SoftmaxConstraints(nn.Module):
- def __init__(self, upsampling_factor, cwindow_size, exp_factor=1):
+ def __init__(self, upsampling_factor, exp_factor=1):
  super(SoftmaxConstraints, self).__init__()
- self.pool = torch.nn.AvgPool2d(kernel_size=cwindow_size)
- self.lr_pool = torch.nn.AvgPool2d(kernel_size=int(cwindow_size/upsampling_factor))
  self.upsampling_factor = upsampling_factor
- self.cwindow_size = cwindow_size
  self.exp_factor = exp_factor
  def forward(self, y, lr):
  y = torch.exp(y*self.exp_factor)
  sum_y = self.pool(y)
  lr_sum = self.lr_pool(lr)
- out = y*torch.kron(lr_sum*1/sum_y, torch.ones((self.cwindow_size,self.cwindow_size)).to('cuda'))
+ out = y*torch.kron(lr_sum*1/sum_y, torch.ones((self.upsampling_factor,self.upsampling_factor)).to('cuda'))
  return out
 
 
@@ -110,7 +107,7 @@ def forward(self, y):
 
 
 class ResNet(nn.Module):
- def __init__(self, number_channels=64, number_residual_blocks=4, upsampling_factor=2, noise=False, constraints='none', dim=1, cwindow_size=4):
+ def __init__(self, number_channels=64, number_residual_blocks=4, upsampling_factor=2, noise=False, constraints='none', dim=1):
  super(ResNet, self).__init__()
  # First layer
  if noise:
@@ -136,7 +133,7 @@ def __init__(self, number_channels=64, number_residual_blocks=4, upsampling_fact
  #optional renomralization layer
  self.is_constraints = False
  if constraints == 'softmax':
- self.constraints = SoftmaxConstraints(upsampling_factor=upsampling_factor, cwindow_size=cwindow_size)
+ self.constraints = SoftmaxConstraints(upsampling_factor=upsampling_factor)
  self.is_constraints = True
  elif constraints == 'enforce_op':
  self.constraints = EnforcementOperator(upsampling_factor=upsampling_factor)

requirements.yml

@@ -1,4 +1,4 @@
-name: condtrained_ds
+name: constrained-ds
 channels:
  - pytorch
  - nvidia

training.py

@@ -1,4 +1,4 @@
-from utils import process_for_training, is_gan, is_noisegan, load_model, get_optimizer, get_criterion, process_for_eval, get_loss, load_data
+from utils import process_for_training, is_gan, load_model, get_optimizer, get_criterion, process_for_eval, get_loss, load_data
 import models
 import numpy as np
 from tqdm import tqdm
@@ -7,7 +7,9 @@
 import torchgeometry as tgm
 import csv
 import numpy as np
-from scoring import main_scoring
+from torch.utils.data import DataLoader, TensorDataset
+from torchmetrics.functional import multiscale_structural_similarity_index_measure, structural_similarity_index_measure
+from skimage import transform
 device = 'cuda'
 
 def run_training(args, data):
@@ -27,7 +29,6 @@ def run_training(args, data):
  running_loss = 0 
  running_discr_loss = 0
  running_adv_loss = 0
- running_mass_loss = 0
  for (inputs, targets) in data[0]: 
  inputs, targets = process_for_training(inputs, targets)
  if is_gan(args):
@@ -37,7 +38,6 @@ def run_training(args, data):
  else:
  loss = optimizer_step(model, optimizer, criterion, inputs, targets, data[0], args)
  running_loss += loss
- running_mass_loss += mass_loss
  loss = running_loss/len(data[0])
  if is_gan(args):
  dicsr_loss = running_discr_loss/len(data)
@@ -51,11 +51,10 @@ def run_training(args, data):
  val_loss = validate_model(model, criterion, data[1], best, epoch, args, discriminator_model, criterion_discr)
  else:
  val_loss = validate_model(model, criterion, data[1], best, epoch, args)
- val_losses.append(val_loss)
  print('Val loss: {:.5f}'.format(val_loss))
  checkpoint(model, val_loss, best, args, epoch)
  best = np.minimum(best, val_loss)
- data = load_data(args.test_val_train, args) 
+ data = load_data(args) 
  scores = evaluate_model( data, args)
 
 
@@ -110,7 +109,7 @@ def validate_model(model, criterion, data, best, epoch, args, discriminator_mode
  adversarial_loss = criterion_discr(fake_output.detach(), real_label)
  loss += args.adv_factor * adversarial_loss
  else:
-  outputs = model(inputs)
+ outputs = model(inputs)
  loss = get_loss(outputs, targets, inputs, args) 
  running_loss += loss.item()
  loss = running_loss/len(data)
@@ -120,12 +119,11 @@ def validate_model(model, criterion, data, best, epoch, args, discriminator_mode
 Tensor = torch.cuda.FloatTensor
 
 def checkpoint(model, val_loss, best, args, epoch):
- print(val_loss, best)
  if val_loss < best:
  checkpoint = {'model': model,'state_dict': model.state_dict()}
  torch.save(checkpoint, './models/'+args.model_id+'.pth')
 
-def evaluate_model(data, args, add_string=None):
+def evaluate_model(data, args):
  model = load_model(args)
  load_weights(model, args.model_id)
  model.eval()
@@ -149,38 +147,133 @@ def evaluate_model(data, args, add_string=None):
  else:
  torch.save(full_pred, './data/prediction/'+args.dataset+'_'+args.model_id+ '_' + args.test_val_train+'.pt')
  calculate_scores(args)
+ 
+def calculate_scores(args):
+ input_val = torch.load('./data/'+args.dataset+'/'+ args.test_val_train+'/input_'+ args.test_val_train+'.pt')
+ target_val = torch.load('./data/'+args.dataset+'/'+ args.test_val_train+'/target_'+ args.test_val_train+'.pt')
+ val_data = TensorDataset(input_val, target_val)
+ pred = np.zeros(target_val.shape)
+ max_val = target_val.max()
+ min_val = target_val.min()
+ mse = 0
+ mae = 0
+ ssim = 0
+ mean_bias = 0
+ mean_abs_bias = 0
+ mass_violation = 0
+ ms_ssim = 0
+ corr = 0
+ crps = 0
+ neg_mean = 0
+ neg_num = 0
+ 
+ l2_crit = nn.MSELoss()
+ l1_crit = nn.L1Loss()
+ 
+ if args.model == 'gan':
+ en_pred = torch.load('./data/prediction/'+args.dataset+'_'+args.model_id+ '_' + args.test_val_train+'_ensemble.pt')
+ pred = torch.mean(en_pred, dim=1)
+ en_pred = en_pred.detach().cpu().numpy()
+ else:
+ pred = torch.load('./data/prediction/'+args.dataset+'_'+args.model_id+ '_' + args.test_val_train+'.pt')
+
+ pred = pred.detach().cpu().numpy()
+ j = 0 
+ for i,(lr, hr) in enumerate(val_data):
+ im = lr.numpy()
+ mse += l2_crit(torch.Tensor(pred[i,j,...]), hr[j,...]).item()
+ mae += l1_crit(torch.Tensor(pred[i,j,...]), hr[j,...]).item()
+ mean_bias += torch.mean( hr[j,...]-torch.Tensor(pred[i,j,...]))
+ mean_abs_bias += torch.abs(torch.mean( hr[j,...]-torch.Tensor(pred[i,j,...])))
+ corr += pearsonr(torch.Tensor(pred[i,j,...]).flatten(), hr[j,...].flatten())
+ ms_ssim += multiscale_structural_similarity_index_measure(torch.Tensor(pred[i,j:j+1,...]), hr[j:j+1,...], data_range=max_val-min_val, kernel_size=11, betas=(0.2856, 0.3001, 0.2363))
+ ssim += structural_similarity_index_measure(torch.Tensor(pred[i,j:j+1,...]), hr[j:j+1,...] , data_range=max_val-min_val, kernel_size=11)
+ neg_num += np.sum(pred[i,j,...] < 0)
+ neg_mean += np.sum(pred[pred < 0])/(pred.shape[-1]*pred.shape[-1])
+ if args.model == 'gan':
+ crps_ens = crps_ensemble(hr[j,0,...].numpy(), en_pred[i,:,j,0,...])
+ crps += crps_ens
+
+ mass_violation += np.mean( np.abs(transform.downscale_local_mean(pred[i,j,...], (1,args.upsampling_factor,args.upsampling_factor)) -im[j,...]))
+
+ mse *= 1/input_val.shape[0]
+ mae *= 1/input_val.shape[0]
+ ssim *= 1/input_val.shape[0]
+ mean_bias *= 1/input_val.shape[0]
+ mean_abs_bias *= 1/input_val.shape[0]
+ corr *= 1/input_val.shape[0]
+ ms_ssim *= 1/input_val.shape[0]
+ crps *= 1/input_val.shape[0]
+ neg_mean *= 1/input_val.shape[0]
+ mass_violation *= 1/input_val.shape[0]
+ psnr = calculate_pnsr(mse, target_val.max() ) 
+ rmse = torch.sqrt(torch.Tensor([mse])).numpy()[0]
+ ssim = float(ssim.numpy())
+ ms_ssim =float( ms_ssim.numpy())
+ psnr = psnr.numpy()
+ corr = float(corr.numpy())
+ mean_bias = float(mean_bias.numpy())
+ mean_abs_bias = float(mean_abs_bias.numpy())
+ scores = {'MSE':mse, 'RMSE':rmse, 'PSNR': psnr[0], 'MAE':mae, 'SSIM':ssim, 'MS SSIM': ms_ssim, 'Pearson corr': corr, 'Mean bias': mean_bias, 'Mean abs bias': mean_abs_bias, 'Mass_violation': mass_violation, 'neg mean': neg_mean, 'neg num': neg_num,'CRPS': crps}
+ print(scores)
+ create_report(scores, args)
 
 
 def calculate_pnsr(mse, max_val):
  return 20 * torch.log10(max_val / torch.sqrt(torch.Tensor([mse])))
 
-def create_report(scores, args, add_string=None):
+def create_report(scores, args):
  args_dict = args_to_dict(args)
  #combine scorees and args dict
  args_scores_dict = args_dict | scores
  #save dict
- save_dict(args_scores_dict, args, add_string)
+ save_dict(args_scores_dict, args)
 
 def args_to_dict(args):
  return vars(args)
 
 
 def save_dict(dictionary, args):
-
- w = csv.writer(open('./data/score_log/'+args.model_id+'.csv', 'w'))
- 
+ w = csv.writer(open('./data/score_log/'+args.model_id+'.csv', 'w')) 
  # loop over dictionary keys and values
  for key, val in dictionary.items():
  # write every key and value to file
  w.writerow([key, val])
 
 def load_weights(model, model_id):
- PATH = '/home/harder/constraint_generative_ml/models/'+model_id+'.pth'
+ PATH = '/home/harder/constrained-downscaling/models/'+model_id+'.pth'
  checkpoint = torch.load(PATH) # ie, model_best.pth.tar
  model.load_state_dict(checkpoint['state_dict'])
  model.to('cuda')
  return model
 
+def pearsonr(x, y):
+ mean_x = torch.mean(x)
+ mean_y = torch.mean(y)
+ xm = x.sub(mean_x)
+ ym = y.sub(mean_y)
+ r_num = xm.dot(ym)
+ r_den = torch.norm(xm, 2) * torch.norm(ym, 2)
+ r_val = r_num / r_den
+ return r_val
+
+def crps_ensemble(observation, forecasts):
+ fc = forecasts.copy()
+ fc.sort(axis=0)
+ obs = observation
+ fc_below = fc<obs[None,...]
+ crps = np.zeros_like(obs)
+ for i in range(fc.shape[0]):
+ below = fc_below[i,...]
+ weight = ((i+1)**2 - i**2) / fc.shape[-1]**2
+ crps[below] += weight * (obs[below]-fc[i,...][below])
+
+ for i in range(fc.shape[0]-1,-1,-1):
+ above = ~fc_below[i,...]
+ k = fc.shape[0]-1-i
+ weight = ((k+1)**2 - k**2) / fc.shape[0]**2
+ crps[above] += weight * (fc[i,...][above]-obs[above])
+ return np.mean(crps)
 
 
 

utils.py

@@ -2,20 +2,18 @@
 import torch.optim as optim
 import torch.nn as nn
 import models
-import Learnable_basis
 from torch.utils.data import DataLoader, TensorDataset
 device = 'cuda'
 
 def load_data(args):
- input_train = torch.load('./data/train/input_train.pt')
- target_train = torch.load('./data/train/target_train.pt')
- 
+ input_train = torch.load('./data/'+args.dataset+'/train/input_train.pt')
+ target_train = torch.load('./data/'+args.dataset+'/train/target_train.pt')
  if args.test_val_train == 'test':
- input_val = torch.load('./data/test/input_test.pt')
- target_val = torch.load('./data/test/target_test.pt')
+ input_val = torch.load('./data/'+args.dataset+'/test/input_test.pt')
+ target_val = torch.load('./data/'+args.dataset+'/test/target_test.pt')
  elif args.test_val_train == 'val':
- input_val = torch.load('./data/val/input_val.pt')
- target_val = torch.load('./data/val/target_val.pt')
+ input_val = torch.load('./data/'+args.dataset+'/val/input_val.pt')
+ target_val = torch.load('./data/'+args.dataset+'/val/target_val.pt')
  elif args.test_val_train == 'train':
  input_val = input_train
  target_val = target_train
@@ -59,9 +57,10 @@ def load_model(args, discriminator=False):
  elif args.model == 'flowconvgru':
  model = models.TimeEndToEndModel( number_channels=args.number_channels, number_residual_blocks=args.number_residual_blocks, upsampling_factor=args.upsampling_factor, time_steps=3, constraints=args.constraints)
  elif args.model == 'gan':
- model = models.ResNet(number_channels=args.number_channels, number_residual_blocks=args.number_residual_blocks, upsampling_factor=args.upsampling_factor, noise=args.noise, constraints=args.constraints, dim=args.dim_channels)
+ model = models.ResNet(number_channels=args.number_channels, number_residual_blocks=args.number_residual_blocks, upsampling_factor=args.upsampling_factor, noise=(args.model=='gan'), constraints=args.constraints, dim=args.dim_channels)
  elif args.model == 'cnn':
- model = models.ResNet(number_channels=args.number_channels, number_residual_blocks=args.number_residual_blocks, upsampling_factor=args.upsampling_factor, noise=args.noise, constraints=args.constraints, dim=args.dim_channels, cwindow_size= args.constraints_window_size)
+ model = models.ResNet(number_channels=args.number_channels, number_residual_blocks=args.number_residual_blocks, upsampling_factor=args.upsampling_factor, noise=(args.model=='gan'), constraints=args.constraints, dim=args.dim_channels)
+ model = model.to(device)
  return model
 
 def get_optimizer(args, model):
@@ -91,7 +90,7 @@ def process_for_training(inputs, targets):
  return inputs, targets
 
 def process_for_eval(outputs, targets, mean, std, max_val, args): 
- if args.gan:
+ if args.model == 'gan':
  outputs[:,:,0,0,...] = outputs[:,0,0,...]*(max_val[0].to(device)-min_val[0].to(device))+min_val[0].to(device) 
  targets[:,0,0,...] = targets[:,0,0,...]*(max_val[0].to(device)-min_val[0].to(device))+min_val[0].to(device)
  else: