load_model.py

# -*- coding: utf-8 -*-
"""full_both_sp.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1hQ629OzxIuMUyDMCIiqXJirxvdDKdNKc

"""# set vars"""

from __future__ import print_function, division
import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim
import timeit
import h5py
import sys
from torch.utils.data import Dataset, DataLoader
from scipy import stats
from torch.nn import functional as F
from torch.autograd import Variable
from matplotlib.tri import Triangulation

dir = sys.argv[1]

batch_size = 128
lr = 1e-5
momentum = 0.99
num_epochs = 100
percentage_train = 1
percentage_val = 0
lr_decay = 0.5
step_size = 5
# loss_weights = [1,1e0,1e21,1e15]
loss_weights = [1,0.5,0.5,0.5,0.5]
nphi = 4
plot_rate = 250
output_rate = 500
val_rate = 2000
datapath = '/scratch/gpfs/marcoam/ml_collisions/data/xgc1/ti272_JET_heat_load/'
run_num = '00094/'
lim = 150000
modelpath = '/home/marcoam/runs/'+str(dir)+'/checkpoint.pth.tar'

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

"""# nets"""

def conv3x3(in_planes, out_planes):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, padding=1, bias=True)

class UnetDownBlock(nn.Module):
   
    def __init__(self, inplanes, planes, predownsample_block):
        
        super(UnetDownBlock, self).__init__()
        
        self.predownsample_block = predownsample_block
        self.conv1 = conv3x3(inplanes, planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        
    def forward(self, x):
        
        x = self.predownsample_block(x)
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        
        return x
    
class UnetUpBlock(nn.Module):
   
    def __init__(self, inplanes, planes, postupsample_block=None):
        
        super(UnetUpBlock, self).__init__()
        
        self.conv1 = conv3x3(inplanes, planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        
        if postupsample_block is None: 
            
            self.postupsample_block = nn.ConvTranspose2d(in_channels=planes,
                                                         out_channels=planes//2,
                                                         kernel_size=2,
                                                         stride=2)
            
        else:
            
            self.postupsample_block = postupsample_block
        
    def forward(self, x):
        
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.postupsample_block(x)
        
        return x
    
    
class Unet(nn.Module):
    
    def __init__(self):
        
        super(Unet, self).__init__()
        
        self.predownsample_block = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.identity_block = nn.Sequential()
        
        self.block1 = UnetDownBlock(
                                    predownsample_block=self.identity_block,
                                    inplanes=2, planes=64
                                    )
        
        self.block2_down = UnetDownBlock(
                                         predownsample_block=self.predownsample_block,
                                         inplanes=64, planes=128
                                         )
        
        self.block3_down = UnetDownBlock(
                                         predownsample_block=self.predownsample_block,
                                         inplanes=128, planes=256
                                         )

        self.block4_down = UnetDownBlock(
                                         predownsample_block=self.predownsample_block,
                                         inplanes=256, planes=512
                                         )
        
        self.block5_down = UnetDownBlock(
                                         predownsample_block=self.predownsample_block,
                                         inplanes=512, planes=1024
                                         )
        
        self.block1_up = nn.ConvTranspose2d(in_channels=1024, out_channels=512,
                                                  kernel_size=2, stride=2)
        
        self.block2_up = UnetUpBlock(
                                     inplanes=1024, planes=512
                                     )
        
        self.block3_up = UnetUpBlock(
                                     inplanes=512, planes=256
                                     )
        
        self.block4_up = UnetUpBlock(
                                     inplanes=256, planes=128
                                     )
        
        self.block5 = UnetUpBlock(
                                  inplanes=128, planes=64,
                                  postupsample_block=self.identity_block
                                  )
        
        self.logit_conv = nn.Conv2d(
                                    in_channels=64, out_channels=1, kernel_size=1,
                                    )
        
        
    def forward(self, x):
        
        features_1s_down = self.block1(x)
        features_2s_down = self.block2_down(features_1s_down)
        features_4s_down = self.block3_down(features_2s_down)
        features_8s_down = self.block4_down(features_4s_down)
        
        features_16s = self.block5_down(features_8s_down)
        
        features_8s_up = self.block1_up(features_16s)
        features_8s_up = torch.cat([features_8s_down, features_8s_up],dim=1)
        
        features_4s_up = self.block2_up(features_8s_up)
        features_4s_up = torch.cat([features_4s_down, features_4s_up],dim=1)
        
        features_2s_up = self.block3_up(features_4s_up)
        features_2s_up = torch.cat([features_2s_down, features_2s_up],dim=1)
        
        features_1s_up = self.block4_up(features_2s_up)
        features_1s_up = torch.cat([features_1s_down, features_1s_up],dim=1)
        
        features_final = self.block5(features_1s_up)
        
        logits = self.logit_conv(features_final)
       
        return logits

class VGG(nn.Module):
    
    def __init__(self, features, num_classes=1000, init_weights=True):
        super(VGG, self).__init__()
        self.features = features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 2
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

cfgs = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs):
    model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
    return model

def vgg16(pretrained=False, progress=True, **kwargs):
    """VGG 16-layer model (configuration "D")

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs)

  
class VGG16(nn.Module):
    
    def __init__(self, n_layers, usegpu=True):
        super(VGG16,self).__init__()
        
        self.cnn = vgg16(pretrained=False)
        self.cnn = nn.Sequential(*list(self.cnn.children())[0])
        self.cnn = nn.Sequential(*list(self.cnn.children())[:n_layers])
        
    def __get_outputs(self,x):
        
        outputs = []
        for i, layer in enumerate(self.cnn.children()):
            x = layer(x)
            outputs.append(x)
            
        return outputs
    
    def forward(self,x):
        outputs = self.__get_outputs(x)
        
        return outputs[-1]
   
    
class SkipVGG16(nn.Module):
    
    def __init__(self, usegpu=True):
        super(SkipVGG16, self).__init__()
        
        self.outputs = [3,8]
        self.n_filters = [64,128]
        
        self.model = VGG16(n_layers=16, usegpu=usegpu)
        
    def forward(self,x):
        
        out = []
        for i, layer in enumerate(list(self.model.children())[0]):
            x = layer(x)
            if i in self.outputs:
                out.append(x)
        out.append(x)
        
        return out
    
    
class ReNet(nn.Module):
    
    def __init__(self, n_input, n_units, patch_size=(1, 1), usegpu=True):
        super(ReNet, self).__init__()
        
        self.usegpu=usegpu
        
        self.patch_size_height = int(patch_size[0])
        self.patch_size_width = int(patch_size[1])
        
        assert self.patch_size_height >= 1
        assert self.patch_size_width >= 1
        
        self.tiling = False if ((self.patch_size_height == 1) and (
            self.patch_size_width == 1)) else True
                
        rnn_hor_n_inputs = n_input * self.patch_size_height * \
            self.patch_size_width
            
        self.rnn_hor = nn.GRU(rnn_hor_n_inputs, n_units,
                              num_layers=1, batch_first=True,
                              bidirectional=True)
        
        self.rnn_ver = nn.GRU(n_units * 2, n_units,
                              num_layers=1, batch_first=True,
                              bidirectional=True)
        
    def __tile(self,x):

        if (x.size(2) % self.patch_size_height) == 0:
            n_height_padding = 0
        else:
            n_height_padding = self.patch_size_height - \
                x.size(2) % self.patch_size_height
        if (x.size(3) % self.patch_size_width) == 0:
            n_width_padding = 0
        else:
            n_width_padding = self.patch_size_width - \
                x.size(3) % self.patch_size_width

        n_top_padding = n_height_padding / 2
        n_bottom_padding = n_height_padding - n_top_padding

        n_left_padding = n_width_padding / 2
        n_right_padding = n_width_padding - n_left_padding

        x = F.pad(x, (n_left_padding, n_right_padding,
                      n_top_padding, n_bottom_padding))

        b, n_filters, n_height, n_width = x.size()

        assert n_height % self.patch_size_height == 0
        assert n_width % self.patch_size_width == 0

        new_height = n_height / self.patch_size_height
        new_width = n_width / self.patch_size_width

        x = x.view(b, n_filters, new_height, self.patch_size_height,
                   new_width, self.patch_size_width)
        x = x.permute(0, 2, 4, 1, 3, 5)
        x = x.contiguous()
        x = x.view(b, new_height, new_width, self.patch_size_height *
                   self.patch_size_width * n_filters)
        x = x.permute(0, 3, 1, 2)
        x = x.contiguous()

        return x
                
    def __swap_hw(self, x):

        # x : b, nf, h, w
        x = x.permute(0, 1, 3, 2)
        x = x.contiguous()
        #  x : b, nf, w, h

        return x
    
    def rnn_forward(self, x, hor_or_ver):

        # x : b, nf, h, w
        assert hor_or_ver in ['hor', 'ver']

        if hor_or_ver == 'ver':
            x = self.__swap_hw(x)

        x = x.permute(0, 2, 3, 1)
        x = x.contiguous()
        b, n_height, n_width, n_filters = x.size()
        # x : b, h, w, nf

        x = x.view(b * n_height, n_width, n_filters)
        # x : b * h, w, nf
        if hor_or_ver == 'hor':
            x, _ = self.rnn_hor(x)
        elif hor_or_ver == 'ver':
            x, _ = self.rnn_ver(x)
            
        x = x.contiguous()
        x = x.view(b, n_height, n_width, -1)
        # x : b, h, w, nf

        x = x.permute(0, 3, 1, 2)
        x = x.contiguous()
        # x : b, nf, h, w

        if hor_or_ver == 'ver':
            x = self.__swap_hw(x)

        return x
    
    def forward(self, x):

        # x : b, nf, h, w
        if self.tiling:
            x = self.__tile(x)

        x = self.rnn_forward(x, 'hor')
        x = self.rnn_forward(x, 'ver')

        return x
        

class ReSeg(nn.Module):
    
    def __init__(self, usegpu=True):
        super(ReSeg, self).__init__()
        
        self.cnn = SkipVGG16(usegpu=usegpu)
        
        self.renet1 = ReNet(256, 100, usegpu=usegpu)
        self.renet2 = ReNet(200, 100, usegpu=usegpu)
        
        self.upsampling1 = nn.ConvTranspose2d(200, 100,
                                              kernel_size=2,stride=2)
        self.relu1 = nn.ReLU()
        self.upsampling2 = nn.ConvTranspose2d(100 + self.cnn.n_filters[1], 100,
                                              kernel_size=2,stride=2)
        self.relu2 = nn.ReLU()
        
        self.final = nn.Conv2d(
                                in_channels = 100 + self.cnn.n_filters[0], 
                                out_channels = 1,
                                kernel_size=1,stride=1)
        
    def forward(self, x):
        
        first_skip, second_skip, x_enc = self.cnn(x)
        x_enc = self.renet1(x_enc)
        x_enc = self.renet2(x_enc)
        x_dec = self.relu1(self.upsampling1(x_enc))
        x_dec = torch.cat((x_dec, second_skip), dim=1)
        x_dec = self.relu2(self.upsampling2(x_dec))
        x_dec = torch.cat((x_dec, first_skip), dim=1)
        x_out = self.final(x_dec)
        
        return x_out

class VGG(nn.Module):
    
    def __init__(self, features, num_classes=1000, init_weights=True):
        super(VGG, self).__init__()
        self.features = features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 2
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

cfgs = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs):
    model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
    return model

def vgg16(pretrained=False, progress=True, **kwargs):
    """VGG 16-layer model (configuration "D")

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs)

  
class VGG16(nn.Module):
    
    def __init__(self, n_layers, usegpu=True):
        super(VGG16,self).__init__()
        
        self.cnn = vgg16(pretrained=False)
        self.cnn = nn.Sequential(*list(self.cnn.children())[0])
        self.cnn = nn.Sequential(*list(self.cnn.children())[:n_layers])
        
    def __get_outputs(self,x):
        
        outputs = []
        for i, layer in enumerate(self.cnn.children()):
            x = layer(x)
            outputs.append(x)
            
        return outputs
    
    def forward(self,x):
        outputs = self.__get_outputs(x)
        
        return outputs[-1]
  
    
class ReNet(nn.Module):
    
    def __init__(self, n_input, n_units, patch_size=(1, 1), usegpu=True):
        super(ReNet, self).__init__()
        
        self.usegpu=usegpu
        
        self.patch_size_height = int(patch_size[0])
        self.patch_size_width = int(patch_size[1])
        
        assert self.patch_size_height >= 1
        assert self.patch_size_width >= 1
        
        self.tiling = False if ((self.patch_size_height == 1) and (
            self.patch_size_width == 1)) else True
                
        rnn_hor_n_inputs = n_input * self.patch_size_height * \
            self.patch_size_width
            
        self.rnn_hor = nn.GRU(rnn_hor_n_inputs, n_units,
                              num_layers=1, batch_first=True,
                              bidirectional=True)
        
        self.rnn_ver = nn.GRU(n_units * 2, n_units,
                              num_layers=1, batch_first=True,
                              bidirectional=True)
        
    def __tile(self,x):

        if (x.size(2) % self.patch_size_height) == 0:
            n_height_padding = 0
        else:
            n_height_padding = self.patch_size_height - \
                x.size(2) % self.patch_size_height
        if (x.size(3) % self.patch_size_width) == 0:
            n_width_padding = 0
        else:
            n_width_padding = self.patch_size_width - \
                x.size(3) % self.patch_size_width

        n_top_padding = n_height_padding / 2
        n_bottom_padding = n_height_padding - n_top_padding

        n_left_padding = n_width_padding / 2
        n_right_padding = n_width_padding - n_left_padding

        x = F.pad(x, (n_left_padding, n_right_padding,
                      n_top_padding, n_bottom_padding))

        b, n_filters, n_height, n_width = x.size()

        assert n_height % self.patch_size_height == 0
        assert n_width % self.patch_size_width == 0

        new_height = n_height / self.patch_size_height
        new_width = n_width / self.patch_size_width

        x = x.view(b, n_filters, new_height, self.patch_size_height,
                   new_width, self.patch_size_width)
        x = x.permute(0, 2, 4, 1, 3, 5)
        x = x.contiguous()
        x = x.view(b, new_height, new_width, self.patch_size_height *
                   self.patch_size_width * n_filters)
        x = x.permute(0, 3, 1, 2)
        x = x.contiguous()

        return x
                
    def __swap_hw(self, x):

        # x : b, nf, h, w
        x = x.permute(0, 1, 3, 2)
        x = x.contiguous()
        #  x : b, nf, w, h

        return x
    
    def rnn_forward(self, x, hor_or_ver):

        # x : b, nf, h, w
        assert hor_or_ver in ['hor', 'ver']

        if hor_or_ver == 'ver':
            x = self.__swap_hw(x)

        x = x.permute(0, 2, 3, 1)
        x = x.contiguous()
        b, n_height, n_width, n_filters = x.size()
        # x : b, h, w, nf

        x = x.view(b * n_height, n_width, n_filters)
        # x : b * h, w, nf
        if hor_or_ver == 'hor':
            x, _ = self.rnn_hor(x)
        elif hor_or_ver == 'ver':
            x, _ = self.rnn_ver(x)
            
        x = x.contiguous()
        x = x.view(b, n_height, n_width, -1)
        # x : b, h, w, nf

        x = x.permute(0, 3, 1, 2)
        x = x.contiguous()
        # x : b, nf, h, w

        if hor_or_ver == 'ver':
            x = self.__swap_hw(x)

        return x
    
    def forward(self, x):

        # x : b, nf, h, w
        if self.tiling:
            x = self.__tile(x)

        x = self.rnn_forward(x, 'hor')
        x = self.rnn_forward(x, 'ver')

        return x


class ConvGRUCell(nn.Module):
    
    def __init__(self, input_size, hidden_size, kernel_size, usegpu=True):
        super(ConvGRUCell, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.usegpu = usegpu
        
        _n_inputs = self.input_size + self.hidden_size
        self.conv_gates = nn.Conv2d(_n_inputs,
                                    2 * self.hidden_size,
                                    self.kernel_size,
                                    padding=self.kernel_size // 2)

        self.conv_ct = nn.Conv2d(_n_inputs, self.hidden_size,
                                 self.kernel_size,
                                 padding=self.kernel_size // 2)
        
    def forward(self, x, hidden):
        
        batch_size, _, height, width = x.size()
        
        if hidden is None:
            size_h = [batch_size, self.hidden_size, height, width]
            hidden = Variable(torch.zeros(size_h))

            if self.usegpu:
                hidden = hidden.cuda()
                
        c1 = self.conv_gates(torch.cat((x, hidden), dim=1))
        rt, ut = c1.chunk(2, 1)

        reset_gate = torch.sigmoid(rt)
        update_gate = torch.sigmoid(ut)

        gated_hidden = torch.mul(reset_gate, hidden)

        ct = torch.tanh(self.conv_ct(torch.cat((x, gated_hidden), dim=1)))

        next_h = torch.mul(update_gate, hidden) + (1 - update_gate) * ct

        return next_h
    

class RecurrentHourglass(nn.Module):
    
    def __init__(self, input_n_filters, hidden_n_filters, kernel_size,
                 n_levels, embedding_size, usegpu=True):
        super(RecurrentHourglass, self).__init__()
            
        assert n_levels >= 1
    
        self.input_n_filters = input_n_filters
        self.hidden_n_filters = hidden_n_filters
        self.kernel_size = kernel_size
        self.n_levels = n_levels
        self.embedding_size = embedding_size
        self.usegpu = usegpu
        
        self.convgru_cell = ConvGRUCell(self.hidden_n_filters,
                                        self.hidden_n_filters,
                                        self.kernel_size,
                                        self.usegpu)
        
        self.__generate_pre_post_convs()
        
    def __generate_pre_post_convs(self):
        
        def __get_conv(input_n_filters, output_n_filters):
            return nn.Conv2d(input_n_filters, output_n_filters,
                             self.kernel_size,
                             padding=self.kernel_size // 2)
            
        self.pre_conv_layers = [__get_conv(self.input_n_filters,
                                           self.hidden_n_filters), ]
    
        for _ in range(self.n_levels - 1):
            self.pre_conv_layers.append(__get_conv(self.hidden_n_filters,
                                                   self.hidden_n_filters))
        self.pre_conv_layers = nn.ModuleList(self.pre_conv_layers)
    
        self.post_conv_layers = [__get_conv(self.hidden_n_filters,
                                            self.embedding_size), ]
        for _ in range(self.n_levels - 1):
            self.post_conv_layers.append(__get_conv(self.hidden_n_filters,
                                                    self.hidden_n_filters))
        self.post_conv_layers = nn.ModuleList(self.post_conv_layers)
    
    def forward_encoding(self, x):
        
        convgru_outputs = []
        hidden = None
        for i in range(self.n_levels):
            x = F.relu(self.pre_conv_layers[i](x))
            hidden = self.convgru_cell(x, hidden)
            convgru_outputs.append(hidden)
            
        return convgru_outputs
    
    def forward_decoding(self, convgru_outputs):
        
        _last_conv_layer = self.post_conv_layers[self.n_levels - 1]
        _last_output = convgru_outputs[self.n_levels - 1]
        
        post_feature_map = F.relu(_last_conv_layer(_last_output))
        for i in range(self.n_levels - 1)[::-1]:
            post_feature_map += convgru_outputs[i]
            post_feature_map = self.post_conv_layers[i](post_feature_map)
            post_feature_map = F.relu(post_feature_map)
            
        return post_feature_map
    
    def forward(self, x):
        
        x = self.forward_encoding(x)
        x = self.forward_decoding(x)
        
        return x
    

class StackedRecurrentHourglass(nn.Module):
    
    def __init__(self, usegpu=True):
        super(StackedRecurrentHourglass, self).__init__()
        
        self.usegpu = usegpu
        
        self.base_cnn = self.__generate_base_cnn()
        
        self.enc_stacked_hourglass = self.__generate_enc_stacked_hg(64,3)
        
        self.stacked_renet = self.__generate_stacked_renet(64,2)
        
        self.decoder = self.__generate_decoder(64)
        
    def __generate_base_cnn(self):
        
        base_cnn = VGG16(n_layers=4, usegpu=self.usegpu)
        
        return base_cnn
    
    def __generate_enc_stacked_hg(self, input_n_filters, n_levels):
        
        stacked_hourglass = nn.Sequential()
        stacked_hourglass.add_module('Hourglass_1',
                                     RecurrentHourglass(
                                         input_n_filters=input_n_filters,
                                         hidden_n_filters=64,
                                         kernel_size=3,
                                         n_levels=n_levels,
                                         embedding_size=64,
                                         usegpu=self.usegpu))
        stacked_hourglass.add_module('pool_1',
                                     nn.MaxPool2d(2, stride=2))
        stacked_hourglass.add_module('Hourglass_2',
                                     RecurrentHourglass(
                                         input_n_filters=64,
                                         hidden_n_filters=64,
                                         kernel_size=3,
                                         n_levels=n_levels,
                                         embedding_size=64,
                                         usegpu=self.usegpu))        
        stacked_hourglass.add_module('pool_2',
                                     nn.MaxPool2d(2, stride=2))    

        return stacked_hourglass

    def __generate_stacked_renet(self, input_n_filters, n_renets):

        assert n_renets >= 1
        
        renet = nn.Sequential()
        renet.add_module('ReNet_1', ReNet(input_n_filters, 32,
                                          patch_size=(1, 1),
                                          usegpu=self.usegpu))
        for i in range(1, n_renets):
            renet.add_module('ReNet_{}'.format(i + 1),
                             ReNet(32 * 2, 32, patch_size=(1, 1),
                                   usegpu=self.usegpu))
            
        return renet
    
    def __generate_decoder(self, input_n_filters):
        
        decoder = nn.Sequential()
        decoder.add_module('ConvTranspose_1',
                           nn.ConvTranspose2d(input_n_filters,
                                              64,
                                              kernel_size=(2, 2),
                                              stride=(2, 2)))
        decoder.add_module('ReLU_1', nn.ReLU())
        decoder.add_module('ConvTranspose_2',
                           nn.ConvTranspose2d(64,
                                              64,
                                              kernel_size=(2, 2),
                                              stride=(2, 2)))
        decoder.add_module('ReLU_2', nn.ReLU())
        decoder.add_module('Final',
                           nn.ConvTranspose2d(64,
                                              1,
                                              kernel_size=(1, 1),
                                              stride=(1, 1)))
        
        return decoder
            
    def forward(self, x):
        
        x = self.base_cnn(x)
        x = self.enc_stacked_hourglass(x)
        x = self.stacked_renet(x)
        x = self.decoder(x)
        
        return x

"""# choose"""

#net = Unet().to(device)
net = ReSeg().to(device)
#net = StackedRecurrentHourglass().to(device)

"""# load data"""

def load_data_hdf(iphi):
  
  hf_f = h5py.File(datapath+run_num+'hdf_f.h5','r')
  hf_df = h5py.File(datapath+run_num+'hdf_df.h5','r')
  
  e_f = hf_f['e_f'][iphi]  
  i_f = hf_f['i_f'][iphi]
  e_df = hf_df['e_df'][iphi] 
  i_df = hf_df['i_df'][iphi]
 
  ind1,ind2,ind3 = i_f.shape
  
  f = np.zeros([lim,2,ind1,ind1])
  df = np.zeros([lim,2,ind1,ind1])

  for n in range(lim):
    f[n,0,:,:-1] = e_f[:,n,:]
    f[n,1,:,:-1] = i_f[:,n,:]
    df[n,0,:,:-1] = e_df[:,n,:]
    df[n,1,:,:-1] = i_df[:,n,:]

    f[n,0,:,-1] = e_f[:,n,-1]
    f[n,1,:,-1] = i_f[:,n,-1]
    df[n,0,:,-1] = e_df[:,n,-1]
    df[n,1,:,-1] = i_df[:,n,-1]
    
  del i_f,e_f,i_df,e_df
  df+=f

  # instantiate variables for conservation properties and for normalization
  hf_cons = h5py.File(datapath+run_num+'hdf_cons_fullvol.h5','r')
  hf_vol = h5py.File(datapath+run_num+'hdf_vol.h5')
  cons = conservation_variables(hf_cons,hf_vol)
  
  hf_stats = h5py.File(datapath+run_num+'hdf_stats.h5','r')
  zvars = stats_variables(hf_stats)
  
  for n in range(lim):
    f[n] = (f[n]-zvars.mean_f)/zvars.std_f
#     df[n] = (df[n]-zvars.mean_df)/zvars.std_df
    df[n] = (df[n]-zvars.mean_fdf)/zvars.std_fdf
    
  zvars.mean_f = zvars.mean_f[np.newaxis]
  zvars.mean_df = zvars.mean_df[np.newaxis]
  zvars.mean_fdf = zvars.mean_fdf[np.newaxis]
  zvars.std_f = zvars.std_f[np.newaxis]
  zvars.std_df = zvars.std_df[np.newaxis]
  zvars.std_fdf = zvars.std_fdf[np.newaxis]
  
  for i in range(int(np.ceil(np.log(batch_size)/np.log(2)))):
    zvars.mean_f = np.concatenate((zvars.mean_f,zvars.mean_f),axis=0)
    zvars.mean_df = np.concatenate((zvars.mean_df,zvars.mean_df),axis=0)
    zvars.mean_fdf = np.concatenate((zvars.mean_fdf,zvars.mean_fdf),axis=0)
    zvars.std_f = np.concatenate((zvars.std_f,zvars.std_f),axis=0)
    zvars.std_df = np.concatenate((zvars.std_df,zvars.std_df),axis=0)  
    zvars.std_fdf = np.concatenate((zvars.std_fdf,zvars.std_fdf),axis=0)
  
  zvars.mean_f = torch.from_numpy(zvars.mean_f).to(device).float()
  zvars.mean_df = torch.from_numpy(zvars.mean_df).to(device).float()
  zvars.mean_fdf = torch.from_numpy(zvars.mean_fdf).to(device).float()
  zvars.std_f = torch.from_numpy(zvars.std_f).to(device).float()
  zvars.std_df = torch.from_numpy(zvars.std_df).to(device).float()
  zvars.std_fdf = torch.from_numpy(zvars.std_fdf).to(device).float()
       
  return f,df,lim,zvars,cons

class stats_variables():
  
  def __init__(self, hf_stats):
    self.std_f = hf_stats['std_f'][...]
    self.std_df = hf_stats['std_df'][...]
    self.std_fdf = hf_stats['std_fdf'][...]
    self.mean_f = hf_stats['mean_f'][...]
    self.mean_df = hf_stats['mean_df'][...]
    self.mean_fdf = hf_stats['mean_fdf'][...]

class conservation_variables():

  def __init__(self, hf_cons, hf_vol):
    self.f0_dsmu = hf_cons['f0_dsmu'][...]
    self.f0_dvp = hf_cons['f0_dvp'][...]
    self.f0_nvp = hf_cons['f0_nvp'][...]
    self.f0_nmu = hf_cons['f0_nmu'][...]
    self.ptl_mass = hf_cons['ptl_mass'][...]
    self.sml_ev2j = 1.6022e-19
    
    self.temp = hf_cons['f0_T_ev'][...]
    self.vol = np.zeros([self.temp.shape[0],self.f0_nmu+1,self.temp.shape[1]])
    self.vol[0] = hf_vol['vole'][0]
    self.vol[1] = hf_vol['voli'][0]

class DistFuncDataset(Dataset):

    def __init__(self, f_array, df_array, temp_array, vol_array):
        self.data = torch.from_numpy(f_array).float()
        self.target = torch.from_numpy(df_array).float()
        self.temp = torch.from_numpy(temp_array).float()
        self.vol = torch.from_numpy(vol_array).float()
        
    def __len__(self):
        return len(self.data)
      
    def __getitem__(self, index):
        a = self.data[index]
        b = self.target[index]
        b = b.view(-1,32,32)
        c = self.temp[index]
        d = self.vol[index]
            
        return a, b, c, d

def get_inds():
    
    filename = '/home/marcoam/runs/'+str(dir)+'/inds.txt'
    fid = open(filename,'r')

    train_inds = []
    val_inds = []
    test_inds = []

    train = False
    val = False
    test = False

    for line in fid.readlines():
      if line[:-1] == 'train':
        train = True
        val = False
        test = False

      if line[:-1] == 'val':
        train = False
        val = True
        test = False

      if line[:-1] == 'test':
        train = False
        val = False
        test = True

      if train == True and line[:-1] != 'train':
        train_inds.append(int(line[:-1]))
      elif val == True and line[:-1] != 'val':
        val_inds.append(int(line[:-1]))
      elif test == True and line[:-1] != 'test':
        test_inds.append(int(line[:-1]))

    return(train_inds,val_inds,test_inds)

"""# split data"""

def split_data(f,df,cons,train_inds,val_inds,test_inds):
    
    f_train = f[train_inds]
    f_val = f[val_inds]
    f_test = f[test_inds]
        
    del f  
      
    df_train = df[train_inds]
    df_val = df[val_inds]
    df_test = df[test_inds]
    
    del df
        
    # temperature and volume separate arrays here
    temp = np.einsum('ij -> ji', cons.temp)   
    temp_train = temp[train_inds]
    temp_val = temp[val_inds]
    temp_test = temp[test_inds]
    
    del temp
    
    vol = np.einsum('ijk -> kji', cons.vol)
    vol_train = vol[train_inds]
    vol_val = vol[val_inds]
    vol_test = vol[test_inds]
    
    
    trainset = DistFuncDataset(f_train, df_train, temp_train, vol_train)
    
    trainloader = DataLoader(trainset, batch_size=batch_size, 
                             shuffle=False, pin_memory=True, num_workers=4)
    
    del f_train, df_train, temp_train, vol_train
   
    valset = DistFuncDataset(f_val, df_val, temp_val, vol_val)
    
    valloader = DataLoader(valset, batch_size=batch_size, 
                           shuffle=True, pin_memory=True, num_workers=4)
 
    return trainloader, valloader, f_test, df_test, temp_test, vol_test

"""# check props"""

# same procedure as col_f_convergence_eval
def check_properties_each(f_slice, cons, temp, vol, sp):
    
    f_slice = f_slice.float()
       
    if len(f_slice.shape) == 2:
      nperp, npar = f_slice.shape
      nbatch = 1
    elif len(f_slice.shape) == 3:  
      nbatch,nperp,npar = f_slice.shape
      
    vth = torch.sqrt(temp*cons.sml_ev2j/cons.ptl_mass[sp]) 
    
    vpar = np.linspace(-cons.f0_nvp,cons.f0_nvp,2*cons.f0_nvp+1)*cons.f0_dvp
    vperp = np.linspace(0,cons.f0_nmu,cons.f0_nmu+1)*cons.f0_dsmu
    
    
#     print('vth',vth*torch.from_numpy(cons.f0_dsmu))
#     print('vpar',torch.from_numpy(vpar)*vth[0])
#     print('vperp',torch.from_numpy(vperp)*vth[0])
    
    vpar = torch.tensor(vpar).float().to(device)
    vperp = torch.tensor(vperp).float().to(device)
        
    mass = cons.ptl_mass[sp]
    conv_factor_notemp = 1/np.sqrt((2*np.pi*cons.sml_ev2j/mass)**3)
    temp_factor = 1/torch.sqrt(temp)
    
    smu_n = cons.f0_dsmu/3 # smu_n = f0_dsmu/f0_mu0_factor, f0_mu0_factor = 3
    f_slice_norm = torch.einsum('ijk,i -> ijk',f_slice,temp_factor)*conv_factor_notemp/smu_n
      
    ones_tensor = torch.ones(nbatch,nperp,npar).float().to(device)
      
    vol_tensor = torch.einsum('ijk,ij -> ijk',ones_tensor,vol)
    vperp_tensor = torch.einsum('ijk,i,j -> ijk',ones_tensor,vth,vperp)
    vpar_tensor = torch.einsum('ijk,i,k -> ijk',ones_tensor,vth,vpar)
    
    mass_tensor = vol_tensor
    mom_tensor = vpar_tensor*cons.ptl_mass[sp]*vol_tensor
    energy_tensor = (vpar_tensor**2 + vperp_tensor**2)*cons.ptl_mass[sp]*vol_tensor
        
    mass_tensor, mom_tensor, energy_tensor = \
    mass_tensor.to(device), mom_tensor.to(device), energy_tensor.to(device)                   
         
    mass = torch.sum(f_slice_norm*mass_tensor, dim = (1,2))
    momentum = torch.sum(f_slice_norm*mom_tensor, dim = (1,2))
    energy = torch.sum(f_slice_norm*energy_tensor, dim = (1,2))
                
    return mass, momentum, energy

# makes calls to individual df/f property calculation
# computes more useful quantities as in col_f_core_m after the calls to col_f_convergence_eval 
def check_properties_main(f,df,temp,vol,cons):
  
  masse = torch.from_numpy(np.array([cons.ptl_mass[0]])).to(device).float()
  massi = torch.from_numpy(np.array([cons.ptl_mass[1]])).to(device).float()
  
  #print('df')
  dne,dpe,dwe = check_properties_each(df[:,0],cons,temp[:,0],vol[:,:,0],0)
  dni,dpi,dwi = check_properties_each(df[:,1],cons,temp[:,1],vol[:,:,1],1)
  
  #print('f')  
  ne,mome,ene = check_properties_each(f[:,0],cons,temp[:,0],vol[:,:,0],0)
  ni,momi,eni = check_properties_each(f[:,1],cons,temp[:,1],vol[:,:,1],1) 
      
  dne_n = torch.abs(dne/ne)
  dni_n = torch.abs(dni/ni)
  
  dp_p = torch.abs(dpi + dpe)/torch.max(torch.abs(momi + mome),1e-3*torch.max(massi,masse)*ne)
  
  dw_w = torch.abs((dwi + dwe)/(eni + ene))
  
  return dne_n,dni_n,dp_p,dw_w

# dataiter = iter(trainloader)
# data, targets, temp, vol = dataiter.next()
# data, targets, temp, vol = data.to(device), targets.to(device), temp.to(device), vol.to(device)     

# outputs = net(data)
# outputs = outputs.to(device)

# nbatch = len(data)     

# data_unnorm = data*zvars.std_f[:nbatch] + zvars.mean_f[:nbatch]
# targets_unnorm = targets*zvars.std_fdf[:nbatch] + zvars.mean_fdf[:nbatch]
# outputs_unnorm = outputs[:nbatch,0]*zvars.std_fdf[:nbatch,1] + zvars.mean_fdf[:nbatch,1]

# targets_nof = targets_unnorm - data_unnorm
# outputs_nof = outputs_unnorm[:nbatch] - data_unnorm[:nbatch,1]

# outputs_nof_to_cat = outputs_nof[:nbatch].unsqueeze(1)
# targets_nof_to_cat = targets_nof[:nbatch,0].unsqueeze(1)

# outputs_nof = torch.cat((outputs_nof_to_cat,targets_nof_to_cat),1)
# check_properties_main(data_unnorm[:,:,:,:-1],outputs_nof[:,:,:,:-1],temp,vol,cons)

"""# train"""

def train(trainloader,valloader,sp_flag,epoch,end,zvars,cons):
  
    props_before = []
    props_after = []
    train_loss_vector = []
    l2_loss_vector = []
    cons_loss_vector = []
    val_loss_vector = []
    
    running_loss = 0.0
    running_l2_loss = 0.0
    running_cons_loss = 0.0
    timestart = timeit.default_timer()
    for i, (data, targets, temp, vol) in enumerate(trainloader):
        timeend = timeit.default_timer()
        #print(timeend-timestart)
     
        data, targets, temp, vol = data.to(device), targets.to(device), temp.to(device), vol.to(device)     
      
        if sp_flag == 0:
            optimizer.zero_grad()
        else:
            optimizer_e.zero_grad()

        outputs = net(data)
        outputs = outputs.to(device)
        
        nbatch = len(data)     
        
        data_unnorm = data*zvars.std_f[:nbatch] + zvars.mean_f[:nbatch]
        targets_unnorm = targets*zvars.std_fdf[:nbatch] + zvars.mean_fdf[:nbatch]
        outputs_unnorm = outputs[:,0]*zvars.std_fdf[:nbatch,1] + zvars.mean_fdf[:nbatch,1]
                   
        targets_nof = targets_unnorm - data_unnorm
        outputs_nof = outputs_unnorm[:nbatch] - data_unnorm[:nbatch,1]
        
        outputs_nof_to_cat = outputs_nof[:nbatch].unsqueeze(1)
        targets_nof_to_cat = targets_nof[:nbatch,0].unsqueeze(1)        
        
        # concatenate with actual dfe
        outputs_nof = torch.cat((outputs_nof_to_cat,targets_nof_to_cat),1)
                  
        # updated calls to check_properties with correct arguments  
        masse_b,massi_b,mom_b,energy_b = check_properties_main(data_unnorm[:,:,:,:-1],\
                                                               targets_nof[:,:,:,:-1],temp,vol,cons)

        masse_a,massi_a,mom_a,energy_a = check_properties_main(data_unnorm[:,:,:,:-1],\
                                                               outputs_nof[:,:,:,:-1],temp,vol,cons)

        props_before.append([(torch.sum(massi_b+masse_b)/nbatch).item(),\
                             (torch.sum(mom_b)/nbatch).item(),\
                             (torch.sum(energy_b)/nbatch).item()])
        props_after.append([torch.sum((massi_a+massi_b)/nbatch).item(),\
                             torch.sum((mom_a)/nbatch).item(),\
                             torch.sum((energy_a)/nbatch).item()])
       
        masse_loss = torch.sum(masse_a)/nbatch
        massi_loss = torch.sum(massi_a)/nbatch
        mom_loss = torch.sum(mom_a)/nbatch
        energy_loss = torch.sum(energy_a)/nbatch

        if i % 200 == 199:
            print('outputs',masse_loss.item(),massi_loss.item(),mom_loss.item(),energy_loss.item())
                
        #masse_loss = torch.sum(torch.abs(masse_a - masse_b)).float()/nbatch
        #massi_loss = torch.sum(torch.abs(massi_a - massi_b)).float()/nbatch
        #mass_loss = massi_loss + masse_loss
        #mom_loss = torch.sum(torch.abs(mom_a - mom_b)).float()/nbatch
        #energy_loss = torch.sum(torch.abs(energy_a - energy_b)).float()/nbatch    
                
        l2_loss = criterion(outputs[:,0],targets[:,1])  
                  
#        loss = l2_loss*loss_weights[0] \
#             + mass_loss*loss_weights[1] \
#             + mom_loss*loss_weights[2] \
#             + energy_loss*loss_weights[3]    
        
#        cons_loss = mass_loss*loss_weights[1] \
#                  + mom_loss*loss_weights[2] \
#                  + energy_loss*loss_weights[3] \

        loss = l2_loss*loss_weights[0]\
             + masse_loss*loss_weights[1]\
             + massi_loss*loss_weights[2]\
             + mom_loss*loss_weights[3]\
             + energy_loss*loss_weights[4]
  
        cons_loss = masse_loss*loss_weights[1]\
                  + massi_loss*loss_weights[2]\
                  + mom_loss*loss_weights[3]\
                  + energy_loss*loss_weights[4]
  
        loss.backward()
        if sp_flag == 0:
            optimizer.step()
        else:
            optimizer_e.step()

        running_loss += loss.item()
        running_l2_loss += l2_loss.item()
        running_cons_loss += cons_loss.item()
        
        if i % output_rate == output_rate-1:
            print('   [%d, %5d] loss: %.6f' %
                  (epoch + 1, end + i + 1, running_loss / output_rate))
            print('      L2 loss: %.6f' % (running_l2_loss / output_rate))
            print('      conservation loss: %.6f' % (running_cons_loss / output_rate))
            
#             print('outputs',mass_a[0].item(),mom_a[0].item(),energy_a[0].item())
#             print('targets',mass_b[0].item(),mom_b[0].item(),energy_b[0].item())
          
        if i % plot_rate == plot_rate-1:
            train_loss_vector.append(running_loss / output_rate)
            l2_loss_vector.append(running_l2_loss / output_rate)
            cons_loss_vector.append(running_cons_loss / output_rate)
            running_loss = 0.0
            running_l2_loss = 0.0
            running_cons_loss = 0.0
            #plot_df(targets_unnorm[0,0,:,:-1],outputs_unnorm[0,0,:,:-1],epoch)
        
        if i % val_rate == val_rate-1:         
          val_loss = validate(valloader,cons,zvars)
          val_loss_vector.append(val_loss)
        
          is_best = False
          if val_loss < np.min(val_loss_vector): ## check this
            is_best = True 

          if i % (2*val_rate) == (2*val_rate-1):
            save_checkpoint({
                             'epoch': epoch+1,
                             'state_dict': net.state_dict(),
                             'val_loss': val_loss,
                             'optimizer': optimizer.state_dict(),
                             }, is_best, lr)

        timestart = timeit.default_timer()  
    end += i + 1
    
    cons_array = np.concatenate((np.array(props_before),np.array(props_after)),axis=1)
    
    return train_loss_vector, l2_loss_vector, cons_loss_vector, val_loss_vector, cons_array, end

"""# validate"""

def validate(valloader,cons,zvars):
  
  print('      Running validation set')
  
  running_loss = 0.0
  
  with torch.no_grad():
    for i, (data, targets, temp, vol) in enumerate(valloader):
      
      data, targets, temp, vol = data.to(device), targets.to(device), temp.to(device), vol.to(device)
      
      outputs = net(data)
      outputs = outputs.to(device)
            
      nbatch = len(data)     

      data_unnorm = data*zvars.std_f[:nbatch] + zvars.mean_f[:nbatch]
      targets_unnorm = targets*zvars.std_fdf[:nbatch] + zvars.mean_fdf[:nbatch]
      outputs_unnorm = outputs[:nbatch,0]*zvars.std_fdf[:nbatch,1] + zvars.mean_fdf[:nbatch,1]

      targets_nof = targets_unnorm - data_unnorm
      outputs_nof = outputs_unnorm[:nbatch] - data_unnorm[:nbatch,1]
      
      outputs_nof_to_cat = outputs_nof[:nbatch].unsqueeze(1)
      targets_nof_to_cat = targets_nof[:nbatch,0].unsqueeze(1)        
        
      # concatenate with actual dfe
      outputs_nof = torch.cat((outputs_nof_to_cat,targets_nof_to_cat),1)

      masse_b,massi_b,mom_b,energy_b = check_properties_main(data_unnorm[:,:,:,:-1],\
                                                 targets_nof[:,:,:,:-1],temp,vol,cons) 
      masse_a,massi_a,mom_a,energy_a = check_properties_main(data_unnorm[:,:,:,:-1],\
                                                 outputs_nof[:,:,:,:-1],temp,vol,cons)  
      masse_loss = torch.sum(masse_a)
      massi_loss = torch.sum(massi_a)
      mom_loss = torch.sum(mom_a)
      energy_loss = torch.sum(energy_a)
                
      l2_loss = criterion(outputs[:,0],targets[:,1])  
                  
      loss = l2_loss*loss_weights[0]\
            + masse_loss*loss_weights[1]\
            + massi_loss*loss_weights[2]\
            + mom_loss*loss_weights[3]\
    	    + energy_loss*loss_weights[4]    
  
#      masse_loss = torch.sum(torch.abs(masse_a - masse_b)).float()/nbatch
#      massi_loss = torch.sum(torch.abs(massi_a - massi_b)).float()/nbatch
#      mass_loss = massi_loss + masse_loss
#      mom_loss = torch.sum(torch.abs(mom_a - mom_b)).float()/nbatch
#      energy_loss = torch.sum(torch.abs(energy_a - energy_b)).float()/nbatch   
     
#      l2_loss = criterion(outputs[:,0],targets[:,1])  
      
#      loss = l2_loss*loss_weights[0] \
#           + mass_loss*loss_weights[1] \
#           + mom_loss*loss_weights[2] \
#           + energy_loss*loss_weights[3] 
      
      running_loss += loss.item()
  
  #print(i+nbatch/batch_size)
  avg_loss = running_loss/(i+1)
  
  print('         Validation loss: %.3f' % (avg_loss))

  return avg_loss

"""# test"""

# Commented out IPython magic to ensure Python compatibility.
def test(f_test,df_test,temp_test,vol_test):
 
    testset = DistFuncDataset(f_test, df_test, temp_test, vol_test)
    
    testloader = DataLoader(testset, batch_size=batch_size, 
                            shuffle=True, num_workers=4)
      
    props_before = []
    props_after = []
    
    l2_error=[]
    lt1=0
    gt1=0
    with torch.no_grad():
        for (data, targets, temp, vol) in testloader:

            data, targets, temp, vol = data.to(device), targets.to(device), temp.to(device), vol.to(device)
          
            outputs = net(data)
            outputs = outputs.to(device)   
            
            nbatch = len(data)     

            data_unnorm = data*zvars.std_f[:nbatch] + zvars.mean_f[:nbatch]
            targets_unnorm = targets*zvars.std_fdf[:nbatch] + zvars.mean_fdf[:nbatch]
            outputs_unnorm = outputs[:nbatch,0]*zvars.std_fdf[:nbatch,1] + zvars.mean_fdf[:nbatch,1]

            targets_nof = targets_unnorm - data_unnorm
            outputs_nof = outputs_unnorm[:nbatch] - data_unnorm[:nbatch,1]

            outputs_nof_to_cat = outputs_nof[:nbatch].unsqueeze(1)
            targets_nof_to_cat = targets_nof[:nbatch,0].unsqueeze(1)        

            # concatenate with actual dfe
            outputs_nof = torch.cat((outputs_nof_to_cat,targets_nof_to_cat),1)
  
            props_before.append([torch.sum(each_prop).item()\
                                 for each_prop in check_properties_main(data_unnorm[:,:,:,:-1],\
                                                                   targets_nof[:,:,:,:-1],temp,vol,cons)])         
            props_after.append([torch.sum(each_prop).item()\
                                for each_prop in check_properties_main(data_unnorm[:,:,:,:-1],\
                                                                  outputs_nof[:,:,:,:-1],temp,vol,cons)])
                                
            loss = criterion(outputs, targets)
            l2_error.append(loss.item()*100)
            if loss.item()*100 < 1:
                lt1+=1
            else:
                gt1+=1
    
    cons_array = np.concatenate((np.array(props_before),np.array(props_after)),axis=1)

    num_error = len(cons_array)
    cons_error = np.zeros([3,num_error])
    
    cons_error[0] = np.abs(((cons_array[:,4]+cons_array[:,5])-(cons_array[:,0]+cons_array[:,1]))/(cons_array[:,0]+cons_array[:,1]))
    cons_error[1] = np.abs((cons_array[:,6]-cons_array[:,2])/cons_array[:,2])
    cons_error[2] = np.abs((cons_array[:,7]-cons_array[:,3])/cons_array[:,3])
    
    print('Finished testing')
    print('Percentage with MSE<1: %d %%' % (
            100 * lt1/(lt1+gt1)))
    print('Percent error in conservation properties:\nmass: \
#             %d %%\nmomentum: %d %%\nenergy: %d %%' % ( 
            100*cons_error[0].max(), 
            100*cons_error[1].max(), 
            100*cons_error[2].max()))
    
    return l2_error, cons_error, props_after

def save_checkpoint(state, is_best, lr, filename='checkpoint.pth.tar'): 
#   torch.save(state,'/content/checkpoints/'+str(lr)+'/'+filename)
  torch.save(state, filename)
  if is_best:
    shutil.copy(filename, 'model_best.pth.tar')

def test_load(trainloader,valloader,zvars,cons):
 
  fid1 = open('/home/marcoam/runs/'+str(dir)+'/vmasse.txt','a')
  fid2 = open('/home/marcoam/runs/'+str(dir)+'/vmassi.txt','a')
  fid3 = open('/home/marcoam/runs/'+str(dir)+'/vmom.txt','a')
  fid4 = open('/home/marcoam/runs/'+str(dir)+'/ven.txt','a')
  fid5 = open('/home/marcoam/runs/'+str(dir)+'/vcontour.txt','a')
  
  f_worst = {}
  s_worst = {}
  
  f_worst['error'] = 0
  s_worst['error'] = 0
 
  f_worst['ind'] = 0
  f_worst['fe'] = 0
  f_worst['fi'] = 0
  f_worst['dfxgc'] = 0
  f_worst['dfml'] = 0

#  for i in range(0):
  for i, (data, targets, temp, vol) in enumerate(valloader):	
#    print(i)    
    time1test = timeit.default_timer()

    data, targets, temp, vol = data.to(device), targets.to(device), temp.to(device), vol.to(device)
    
    checkpoint = torch.load(modelpath)
    net.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    epoch = checkpoint['epoch']
    net.eval()
   
    outputs = net(data)
    outputs = outputs.to(device)

    nbatch = len(data) 
        
    data_unnorm = data*zvars.std_f[:nbatch] + zvars.mean_f[:nbatch]
    targets_unnorm = targets*zvars.std_fdf[:nbatch] + zvars.mean_fdf[:nbatch]
    outputs_unnorm = outputs[:,0]*zvars.std_fdf[:nbatch,1] + zvars.mean_fdf[:nbatch,1]
                    
    targets_nof = targets_unnorm - data_unnorm
    outputs_nof = outputs_unnorm[:nbatch] - data_unnorm[:nbatch,1]
        
    outputs_nof_to_cat = outputs_nof[:nbatch].unsqueeze(1)
    targets_nof_to_cat = targets_nof[:nbatch,0].unsqueeze(1)        
        
    # concatenate with actual dfe
    outputs_nof = torch.cat((targets_nof_to_cat,outputs_nof_to_cat),1)
                  
    # updated calls to check_properties with correct arguments  
    masse_b,massi_b,mom_b,energy_b = check_properties_main(data_unnorm[:,:,:,:-1],\
                                                               targets_nof[:,:,:,:-1],temp,vol,cons)

    masse_a,massi_a,mom_a,energy_a = check_properties_main(data_unnorm[:,:,:,:-1],\
                                                               outputs_nof[:,:,:,:-1],temp,vol,cons)

#    l2_loss = criterion(outputs[:0],targets[:,1])
    masse_a,massi_a,mom_a,energy_a = masse_a.cpu(),massi_a.cpu(),mom_a.cpu(),energy_a.cpu()
    masse_b,massi_b,mom_b,energy_b = masse_b.cpu(),massi_b.cpu(),mom_b.cpu(),energy_b.cpu()

    # find worst and 2nd worst and write f and df
    for j in range(nbatch):
      if mom_a[j] > s_worst['error']:
        s_worst['error'] = mom_a[j]
        s_worst['ind'] = nbatch*i+j
        s_worst['fe'] = data_unnorm[j,0,:,:-1].cpu().detach().numpy()
        s_worst['fi'] = data_unnorm[j,1,:,:-1].cpu().detach().numpy()
        s_worst['dfxgc'] = targets_nof[j,1,:,:-1].cpu().detach().numpy()
        s_worst['dfml'] = outputs_nof[j,1,:,:-1].cpu().detach().numpy()

        if mom_a[j] > f_worst['error']:
          s_worst['error'] = f_worst['error']
          s_worst['ind'] = f_worst['ind']
          s_worst['fe'] = f_worst['fe']
          s_worst['fi'] = f_worst['fi']
          s_worst['dfxgc'] = f_worst['dfxgc']
          s_worst['dfml'] = f_worst['dfml']

          f_worst['error'] = mom_a[j]
          f_worst['ind'] = nbatch*i+j
          f_worst['fe'] = data_unnorm[j,0,:,:-1].cpu().detach().numpy()
          f_worst['fi'] = data_unnorm[j,1,:,:-1].cpu().detach().numpy()
          f_worst['dfxgc'] = targets_nof[j,1,:,:-1].cpu().detach().numpy()
          f_worst['dfml'] = outputs_nof[j,1,:,:-1].cpu().detach().numpy()

 
      fid1.write(str(masse_a[j].item())+' '+str(masse_b[j].item())+'\n')	
      fid2.write(str(massi_a[j].item())+' '+str(massi_b[j].item())+'\n')	
      fid3.write(str(mom_a[j].item())+' '+str(mom_b[j].item())+'\n')	
      fid4.write(str(energy_a[j].item())+' '+str(energy_b[j].item())+'\n')	
  fid1.close()
  fid2.close()
  fid3.close()
  fid4.close()

# write 1st worst and 2nd worst
  fid_f1 = open('/home/marcoam/runs/'+str(dir)+'/vfirst_worst_fe.txt','w')
  fid_f2 = open('/home/marcoam/runs/'+str(dir)+'/vfirst_worst_fi.txt','w')
  fid_f3 = open('/home/marcoam/runs/'+str(dir)+'/vfirst_worst_dfxgc.txt','w')
  fid_f4 = open('/home/marcoam/runs/'+str(dir)+'/vfirst_worst_dfml.txt','w')
  fid_f5 = open('/home/marcoam/runs/'+str(dir)+'/vfirst_worst_info.txt','w')
 
  for line in f_worst['fe']: #fe
    for entry in line:
      fid_f1.write(str(entry)+' ')
    fid_f1.write('\n')
  for line in f_worst['fi']: #fi
    for entry in line:
      fid_f2.write(str(entry)+' ')
    fid_f2.write('\n')
  for line in f_worst['dfxgc']: #dfxgc
    for entry in line:
      fid_f3.write(str(entry)+' ')
    fid_f3.write('\n')
  for line in f_worst['dfml']: #dfml
    for entry in line:
      fid_f4.write(str(entry)+' ')
    fid_f4.write('\n')
  fid_f5.write(str(f_worst['ind'])+'\n')
  fid_f5.write(str(f_worst['error'].item()))
  
  fid_s1 = open('/home/marcoam/runs/'+str(dir)+'/vsecond_worst_fe.txt','w')
  fid_s2 = open('/home/marcoam/runs/'+str(dir)+'/vsecond_worst_fi.txt','w')
  fid_s3 = open('/home/marcoam/runs/'+str(dir)+'/vsecond_worst_dfxgc.txt','w')
  fid_s4 = open('/home/marcoam/runs/'+str(dir)+'/vsecond_worst_dfml.txt','w')
  fid_s5 = open('/home/marcoam/runs/'+str(dir)+'/vsecond_worst_info.txt','w')
 
  for line in s_worst['fe']: #fe
    for entry in line:
      fid_s1.write(str(entry)+' ')
    fid_s1.write('\n')
  for line in s_worst['fi']: #fi
    for entry in line:
      fid_s2.write(str(entry)+' ')
    fid_s2.write('\n')
  for line in s_worst['dfxgc']: #dfxgc
    for entry in line:
      fid_s3.write(str(entry)+' ')
    fid_s3.write('\n')
  for line in s_worst['dfml']: #dfml
    for entry in line:
      fid_s4.write(str(entry)+' ')
    fid_s4.write('\n')
  fid_s5.write(str(s_worst['ind'])+'\n')
  fid_s5.write(str(s_worst['error'].item()))

  fm = h5py.File('/scratch/gpfs/marcoam/ml_collisions/data/xgc1/ti272_JET_heat_load/00094/xgc.mesh.h5','r')
  RZ = fm['RZ'][...]
  tri = fm['tri'][...]
  fid5.write('rz\n')
  for r in RZ:
    fid5.write(str(r[0])+'\n')
    fid5.write(str(r[1])+'\n')
  fid5.write('tri\n')
  for t in tri:
    fid5.write(str(t[0])+'\n')
    fid5.write(str(t[1])+'\n')
    fid5.write(str(t[2])+'\n')
  fid5.close()

  return None

"""# plot"""

def plot_df(df_xgc,df_ml,epoch):
  
  df_xgc = df_xgc.cpu().detach().numpy()
  df_ml = df_ml.cpu().detach().numpy()
  
  df_min = df_xgc.min().item()
  df_max = df_xgc.max().item()
  
  cbarticks = np.linspace(df_min,df_max,10)
  
  fig = plt.figure()
  fig.set_figheight(10)
  fig.set_figwidth(10)

  v_aspect=32/31
  ax1 = fig.add_subplot(1,2,1,aspect=v_aspect)
  ax2 = fig.add_subplot(1,2,2,aspect=v_aspect)
  ctr = ax1.contourf(df_xgc, vmin=df_min, vmax=df_max)
  ax2.contourf(df_ml)
  
  ax1.set_title('Actual df')
  ax2.set_title('Predicted df')
  
  fig.subplots_adjust(right=0.8)
  cbar_ax = fig.add_axes([0.85, 0.3, 0.05, 0.4])
  fig.colorbar(ctr, cax=cbar_ax, ticks=cbarticks)
  
#   fig.savefig('figs/dfs_{}'.format(epoch+1))
  
  plt.show()

"""# main"""

start = timeit.default_timer()
criterion = nn.MSELoss()

optimizer = optim.SGD(net.parameters(), lr=lr, momentum=momentum)
#optimizer = optim.Adam(net.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=step_size,gamma=lr_decay)

train_loss = []
val_loss = []

lr_epoch = [group['lr'] for group in optimizer.param_groups][0]
 
for iphi in range(nphi):

  print('Beginning training iphi = {}'.format(iphi))
  print('   Loading data')
  load1 = timeit.default_timer()
  f,df,num_nodes,zvars,cons = load_data_hdf(iphi)
  load2 = timeit.default_timer()
  print('      Loading time: %.3fs' % (load2-load1))

  print('   Creating training set')
  train_inds,val_inds,test_inds = get_inds()
  trainloader,valloader,f_test,df_test,temp_test,vol_test = split_data(f,df,cons,train_inds,val_inds,test_inds)
  del f,df

  print('   Running previously loaded model') 
  fid1 = open('/home/marcoam/runs/'+str(dir)+'/vmasse.txt','w')
  fid2 = open('/home/marcoam/runs/'+str(dir)+'/vmassi.txt','w')
  fid3 = open('/home/marcoam/runs/'+str(dir)+'/vmom.txt','w')
  fid4 = open('/home/marcoam/runs/'+str(dir)+'/ven.txt','w')
  fid5 = open('/home/marcoam/runs/'+str(dir)+'/vcontour.txt','w')
  test_load(trainloader,valloader,zvars,cons)