full_test_fdf.py

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

Automatically generated by Colaboratory.

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

# from cloud
"""

# from google.colab import drive
# drive.mount('/content/drive')

## uncomment this to load in files from Cloud

# from google.colab import auth
# auth.authenticate_user()

# project_id = 'peaceful-impact-247117'
# bucket_name = 'ml_collisions-data1'

# !gcloud config set project {project_id}
# !gsutil -m cp -r gs://ml-collisions-data1/hdf5_data/ /content/

"""# 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
from torch.utils.data import Dataset, DataLoader
from scipy import stats
from google.colab import files
from torch.nn import functional as F
from torch.autograd import Variable

batch_size = 32
lr = 1e-4
momentum = 0.99
num_epochs = 50
percentage_train = 0.8
percentage_val = 0.1
lr_decay = 0.1
step_size = 2
loss_weights = [1,1e0,1e21,1e15]
#loss_weights = [1,0,0,0]
nphi = 1
plot_rate = 250 # how often the loss is plotted
plot = False # whether or not to plot contour of df
output_rate = 500 # how often the loss is printed 
val_rate = 2000 # every how many iterations the validation set is run 

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

"""# architectures

## unet
"""

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

"""## reseg"""

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

"""## hourglass"""

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('/content/hdf5_data/hdf_f.h5','r')
  hf_df = h5py.File('/content/hdf5_data/hdf_df.h5','r')
  
  i_f = hf_f['i_f'][iphi]
  e_f = hf_f['e_f'][iphi]  
  i_df = hf_df['i_df'][iphi]
  e_df = hf_df['e_df'][iphi] 

  ind1,ind2,ind3 = i_f.shape
  
  f = np.zeros([ind2,2,ind1,ind1])
  df = np.zeros([ind2,2,ind1,ind1])
  props = np.zeros([ind2,2,2])

  # get data into one f and df array that is [166676,2,32,32]
  
  # there is quite a bit of unnecessary confusion regarding the species numbers for ions and electrons
  # in XGC, electrons = 0, ions = 1, but for some reason I mixed them up here and started doing it the other
  # way around - as a result some of these indices are one way and others are another way
  # I'll try to comment out which index is which when needed
  
  # here ions = 0, electrons = 1
  
  for n in range(ind2):
    f[n,0,:,:-1] = i_f[:,n,:]
    f[n,1,:,:-1] = e_f[:,n,:]
    df[n,0,:,:-1] = i_df[:,n,:]
    df[n,1,:,:-1] = e_df[:,n,:]

    # need to change velocity grid from 32 x 31 to 32 x 32 to fit into NN
    f[n,0,:,-1] = i_f[:,n,-1]
    f[n,1,:,-1] = e_f[:,n,-1]
    df[n,0,:,-1] = i_df[:,n,-1]
    df[n,1,:,-1] = e_df[:,n,-1]
    
  del i_f,e_f,i_df,e_df
  
  # load in temperature and volume for calculation of conservation properties
  hf_cons = h5py.File('/content/hdf5_data/hdf_cons_fullvol.h5','r')
  
  #props[:,:,0] = hf_cons['f0_grid_vol_vonly'][...].transpose() 
  props[:,:,0] = hf_cons['f0_grid_vol'][...].transpose() 
  props[:,:,1] = hf_cons['f0_T_ev'][...].transpose()

  # load in mean and std dev for z normalization
  # for this data, ions = 0, electrons = 1
  hf_stats = h5py.File('/content/hdf5_data/hdf_stats.h5','r')
    
  std_f = hf_stats['std_f'][...]
  std_df = hf_stats['std_df'][...]
  std_fdf = hf_stats['std_fdf'][...]
  mean_f = hf_stats['mean_f'][...]
  mean_df = hf_stats['mean_df'][...]
  mean_fdf = hf_stats['mean_fdf'][...]
  
  # network learns better with targets = f + df
  df+=f
  
  for n in range(ind2):
    f[n] = (f[n]-mean_f)/std_f
#     df[n] = (df[n]-mean_df)/std_df
    df[n] = (df[n]-mean_fdf)/std_fdf
    
    
  # this part turns mean_f from [2,32,32] to [batch_size,2,32,32] for easier unnormalization in training
  mean_f = mean_f[np.newaxis]
  mean_df = mean_df[np.newaxis]
  mean_fdf = mean_fdf[np.newaxis]
  std_f = std_f[np.newaxis]
  std_df = std_df[np.newaxis]
  std_fdf = std_fdf[np.newaxis]
  
  for i in range(int(np.ceil(np.log(batch_size)/np.log(2)))):
    mean_f = np.concatenate((mean_f,mean_f),axis=0)
    mean_df = np.concatenate((mean_df,mean_df),axis=0)
    mean_fdf = np.concatenate((mean_fdf,mean_fdf),axis=0)
    std_f = np.concatenate((std_f,std_f),axis=0)
    std_df = np.concatenate((std_df,std_df),axis=0)  
    std_fdf = np.concatenate((std_fdf,std_fdf),axis=0)
  
  mean_f = torch.from_numpy(mean_f).to(device).float()
  mean_df = torch.from_numpy(mean_df).to(device).float()
  mean_fdf = torch.from_numpy(mean_fdf).to(device).float()
  std_f = torch.from_numpy(std_f).to(device).float()
  std_df = torch.from_numpy(std_df).to(device).float()
  std_fdf = torch.from_numpy(std_fdf).to(device).float()
       
  return f,df,props,ind2,std_f,std_fdf,mean_f,mean_fdf

"""# load cons_vars"""

def load_cons_vars():
  
  hf_cons = h5py.File('/content/hdf5_data/hdf_cons_fullvol.h5','r')
  
  class conservation_variables():
    
    def __init__(self, hf_cons):
      self.f0_grid_vol = hf_cons['f0_grid_vol'][...] # don't think you need this or f0_t_ev anymore
      self.f0_dsmu = hf_cons['f0_dsmu'][...]
      self.f0_dvp = hf_cons['f0_dvp'][...]
      self.f0_t_ev = hf_cons['f0_T_ev'][...] # since they're accessed above
      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
      
  cons = conservation_variables(hf_cons)
  
  return cons

"""# dataset class"""

class DistFuncDataset(Dataset):

    def __init__(self, f_array, df_array, props_array):
        self.data = torch.from_numpy(f_array).float()
        self.target = torch.from_numpy(df_array).float()
        self.props = torch.from_numpy(props_array).float()
        
    def __len__(self):
        return len(self.data)
      
    def __getitem__(self, index):
        x = self.data[index]
        y = self.target[index]
        y = y.view(-1,32,32)
        z = self.props[index]
            
        return x, y, z

"""# split data"""

def split_data(f,df,props,num_nodes):
    
    # get random indices for training/validation/testing sets
    inds = np.arange(num_nodes)
    #np.random.seed(0)
    np.random.shuffle(inds) 
    
    num_train=int(np.floor(percentage_train*num_nodes))
    num_val=int(np.floor(percentage_val*num_nodes))
    
    train_inds = inds[:num_train]
    val_inds = inds[num_train:num_train+num_val]
    test_inds = inds[num_train+num_val:]
 
    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
        
    props_train = props[train_inds]
    props_val = props[val_inds]
    props_test = props[test_inds]
    
    del props
    
    # create dataloaders for training and validation only
    trainset = DistFuncDataset(f_train, df_train, props_train)
    
    trainloader = DataLoader(trainset, batch_size=batch_size, 
                             shuffle=True, pin_memory=True, num_workers=4)
    
    del f_train, df_train, props_train
    
    valset = DistFuncDataset(f_val, df_val, props_val)
    
    valloader = DataLoader(valset, batch_size=batch_size, 
                           shuffle=True, pin_memory=True, num_workers=4)
    
    # return raw testing data to collect test data across different iphis 
    # only then is the data put into a dataloader
        
    return trainloader, valloader, f_test, df_test, props_test

"""# check props"""

# don't need this one anymore for this script
def check_properties(f_slice, vol):
    
    f_slice = f_slice.double()
       
    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
          
    f0_smu_max = 4
    f0_vp_max = 4
    
    vpar = np.linspace(-f0_vp_max,f0_vp_max,npar) # 2*f0_nvp + 1
    vperp = np.linspace(0,f0_smu_max,nperp) # f0_nmu + 1
    
    vpar = torch.tensor(vpar).float().to(device)
    vperp = torch.tensor(vperp).float().to(device)
      
    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,j -> ijk',ones_tensor,vperp)
    vpar_tensor = torch.einsum('ijk,k -> ijk',ones_tensor,vpar)
    
    mass_tensor = vol_tensor
    mom_tensor = vpar_tensor*vol_tensor
    energy_tensor = (vpar_tensor**2 + vperp_tensor**2)*vol_tensor
        
    mass_array, mom_array, energy_array = \
    mass_array.to(device), mom_array.to(device), energy_array.to(device)

    #mass = (torch.sum(f_slice*mass_array).float())/nbatch
    #momentum = (torch.sum(f_slice*mom_array).float())/nbatch
    #energy = (torch.sum(f_slice*energy_array).float())/nbatch
    
    mass = torch.sum(f_slice*mass_tensor, dim = (1,2))
    momentum = torch.sum(f_slice*mom_tensor, dim = (1,2))
    energy = torch.sum(f_slice*energy_tensor, dim = (1,2))
                
    return mass, momentum, energy

# don't need this one either
def check_properties4(fi,fe,dfi,dfe,voli,vole,cons,props):
  
  masse = cons.ptl_mass[0]
  massi = cons.ptl_mass[1]
  
#   vthi = np.sqrt(props[:,1,1]*cons.sml_ev2j/massi)
#   vthe = np.sqrt(props[:,1,0]*cons.sml_ev2j/massi)
  
  dni,dpi,dwi = check_properties(dfi,voli)
  dne,dpe,dwe = check_properties(dfe,vole)

#   f0_smu_max = 4
#   f0_vp_max = 4
  
  ni,momi,eni = check_properties(fi,voli)
  ne,mome,ene = check_properties(fe,vole)
  
  dni_n = torch.abs(dni/ni)
  dne_n = torch.abs(dne/ne)
  
  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 dni_n,dne_n,dp_p,dw_w

def check_properties3(f_slice, df_slice, props, cons):
  
  # just gets shape to help with vectorization
  if len(df_slice.shape) == 2:
    nperp, npar = df_slice.shape
    nbatch = 1
  elif len(df_slice.shape) == 3:  
    nbatch,nperp,npar = df_slice.shape
  
  # since this is from XGC, the ion index is 1
  mass = cons.ptl_mass[1]
  
  imu = np.linspace(0,cons.f0_nmu,cons.f0_nmu+1)
  ivp = np.linspace(-cons.f0_nvp,cons.f0_nvp,2*cons.f0_nvp+1)
  
  mu_vol = torch.ones(cons.f0_nmu+1).to(device)
  mu_vol[0], mu_vol[-1] = 0.5, 0.5
  
  mu = torch.tensor(imu*cons.f0_dsmu).float().to(device)
  vp = torch.tensor(ivp*cons.f0_dvp).float().to(device)
  
  # I thought I needed this ones tensor as a sort of template to multiply things onto
  # but it's probably not necessary
  ones_tensor = torch.ones(nbatch,nperp,npar).to(device)
  
  vol = torch.einsum('ijk,i,j -> ijk',ones_tensor,props[:,0],mu_vol).to(device)
  en_th = torch.einsum('ijk,i -> ijk',ones_tensor,props[:,1]).to(device)*cons.sml_ev2j
  vth = torch.sqrt(en_th/mass).to(device)
    
  mu_tensor = torch.einsum('ijk,j -> ijk',ones_tensor,mu).to(device)
  vp_tensor = torch.einsum('ijk,k -> ijk',ones_tensor,vp).to(device)
  
  en = 0.5*(vp_tensor**2 + mu_tensor**2)
  
  dens_tensor = vol.to(device)
  mom_tensor = vp_tensor*vth*mass*vol.to(device)
  energy_tensor = en*en_th*vol.to(device)
  
  # diag_port1 subroutine in XGC divides the properties by the density of f
  f_density = torch.sum(f_slice*dens_tensor, dim = (1,2)) 
  
  # tensors before the sum are [batch_size,32,31]
  # after summing over the velocity grid the output properties are vectors of length batch size
  density = torch.sum(df_slice*dens_tensor, dim = (1,2))/f_density
  momentum = torch.sum(df_slice*mom_tensor, dim = (1,2))/f_density
  energy = torch.sum(df_slice*energy_tensor, dim = (1,2))/f_density
  
  return density, momentum, energy

"""# train"""

# main function for training
# no loop over epochs here - only over data in the trainloader which is just for one phi
def train(trainloader,valloader,sp_flag,epoch,end,std_f,std_df,mean_f,mean_df,cons):
  
    # initialize vectors to which to append different losses
    props_before = []
    props_after = []
    train_loss_vector = []
    l2_loss_vector = []
    cons_loss_vector = []
    val_loss_vector = []
    
    # keep running losses - they will be reset every output 
    running_loss = 0.0
    running_l2_loss = 0.0
    running_cons_loss = 0.0
    timestart = timeit.default_timer()
    for i, (data, targets, props) in enumerate(trainloader):
        timeend = timeit.default_timer()
        #print(timeend-timestart)
     
        # data = f, targets = actual df, outputs = predicted df
        data, targets, props = data.to(device), targets.to(device), props.to(device)      
       
        # will have to see need for this in the future - ions are zero
        if sp_flag == 0:
            optimizer.zero_grad()
        else:
            optimizer_e.zero_grad()

        outputs = net(data)
        outputs = outputs.to(device)
                
        # get size of batch - needed for last iteration where len(data)!=batch_size
        # there's a better way of doing this that I've updated in the new script
        # unnormalize data for conservation calculation
        # in all of this ions = 0, electrons = 1
        nbatch = len(data)     
        if nbatch != batch_size:
          data_unnorm = data*std_f[:nbatch] + mean_f[:nbatch]
          targets_unnorm = targets*std_df[:nbatch,0] + mean_df[:nbatch,0]
          outputs_unnorm = outputs*std_df[:nbatch,0] + mean_df[:nbatch,0]
        
        else:
          data_unnorm = data*std_f + mean_f
          targets_unnorm = targets*std_df[:,0] + mean_df[:,0]
          outputs_unnorm = outputs*std_df[:,0] + mean_df[:,0]
          
        # subtract off f to make it just df in the conservation calculation  
        targets_nof = targets_unnorm - data_unnorm[:,0,:,:]
        outputs_nof = outputs_unnorm - data_unnorm[:,0,:,:]
          
        # returns 1D vectors for each property of length nbatch
        # ions = 0, electrons = 1
        # all arrays here have velocity grid dimensions of 32x32 - only pass in 32x31 for accurate
        # calculations
        mass_b,mom_b,energy_b = check_properties3(data_unnorm[:,0,:,:-1],\
                                                 targets_nof[:,0,:,:-1],props,cons) 
        mass_a,mom_a,energy_a = check_properties3(data_unnorm[:,0,:,:-1],\
                                                 outputs_nof[:,0,:,:-1],props,cons)
        
        # these will be used to produce the y=x graphs of the properties before and after
        props_before.append([(torch.sum(mass_b)/nbatch).item(),\
                             (torch.sum(mom_b)/nbatch).item(),\
                             (torch.sum(energy_b)/nbatch).item()])
        props_after.append([torch.sum((mass_a)/nbatch).item(),\
                             torch.sum((mom_a)/nbatch).item(),\
                             torch.sum((energy_a)/nbatch).item()])
        
        # get differences and take average across batch
        mass_loss = torch.sum(torch.abs(mass_a - mass_b)).float()/nbatch
        mom_loss = torch.sum(torch.abs(mom_a - mom_b)).float()/nbatch
        energy_loss = torch.sum(torch.abs(energy_a - energy_b)).float()/nbatch        
        
        # this is l2 loss normalized - maybe it would be good to have the L2 loss of the unnormalized 
        # data as well as a performance metric
        l2_loss = criterion(outputs,targets)  
                  
        # this is the overall loss that will be backpropped
        loss = l2_loss*loss_weights[0] \
             + mass_loss*loss_weights[1] \
             + mom_loss*loss_weights[2] \
             + energy_loss*loss_weights[3]    
        
        # this is just for outputting the loss
        cons_loss = mass_loss*loss_weights[1] \
                  + mom_loss*loss_weights[2] \
                  + energy_loss*loss_weights[3] \

        loss.backward()
        if sp_flag == 0:
            optimizer.step()
        else:
            optimizer_e.step()

        # keep running tally
        running_loss += loss.item()
        running_l2_loss += l2_loss.item()
        running_cons_loss += cons_loss.item()
        
        # currently set to print every 500
        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())
        
        # how often the loss is appended to the loss vectors - currently every 250
        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
            # contour plot if desired
            plot_df(targets_nof[0,0,:,:-1],outputs_nof[0,0,:,:-1],epoch) if plot == True else None
        
        # how often the validation will occur - currently every 2000
        if i % val_rate == val_rate-1:         
          val_loss = validate(valloader,props,cons,std_f,std_df,mean_f,mean_df)
          val_loss_vector.append(val_loss)
        
          # determines best model
          is_best = False
          if val_loss < np.min(val_loss_vector): ## check this
            is_best = True 
          
          # saves model to .tar file - for some reason I have this set to only save every 2 validation
          # runs - probably should save every one
          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()  
        
    # this end marker just keeps track of which iteration we're on in the case that more than one phi
    # is used
    end += i + 1
    
    # get the conservation array - would like it to be [n_iterations, 6], for 3 properties before
    # and 3 after - I don't think this is the correct shape, so there's some postprocessing I need to 
    # do before I can actually plot these
    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

"""# train one batch"""

# not needed - was used for one batch tests

def train_one(data,targets,props,std_f,std_df,mean_f,mean_df):

  train_loss_vector = []
  l2_loss_vector = []
  cons_loss_vector = [[],[],[]]

  running_loss = 0.0
  running_l2_loss = 0.0
  running_den = 0.0
  running_mom = 0.0
  running_en = 0.0
  
#   fi_mean = np.mean(data[0,0].cpu().numpy())
#   fe_mean = np.mean(data[0,1].cpu().numpy())
#   df_mean = np.mean(targets.cpu().numpy())
  
#   fi_std = np.std(data[0,0].cpu().numpy())
#   fe_std = np.std(data[0,1].cpu().numpy())
#   df_std = np.std(targets.cpu().numpy())
  
#   data[0,0] = (data[0,0]-fi_mean)/fi_std
#   data[0,1] = (data[0,1]-fe_mean)/fe_std
#   targets = (targets-df_mean)/df_std

  plt.contourf(data[0,0].cpu().numpy())
  plt.title('fi')
  plt.colorbar()
  plt.show()

  plt.contourf(data[0,1].cpu().numpy())
  plt.title('fe')
  plt.colorbar()
  plt.show()

  plt.contourf(targets[0,0].cpu().numpy())
  plt.title('df')
  plt.colorbar()
  plt.show()
  
  for epoch in range(2000):

    data, targets, props = data.to(device), targets.to(device), props.to(device)
     
    outputs = net(data)
    outputs = outputs.to(device)   

    data_unnorm = data*std_f + mean_f
    targets_unnorm = targets*std_df[:,0] + mean_df[:,0]
    outputs_unnorm = outputs*std_df[:,0] + mean_df[:,0]
    
    mass_b,mom_b,energy_b = check_properties(data_unnorm[:,0,:,:-1],\
                                         targets_unnorm[:,0,:,:-1],props,cons) 
    mass_a,mom_a,energy_a = check_properties(data_unnorm[:,0,:,:-1],\
                                         outputs_unnorm[:,0,:,:-1],props,cons)
#     print(mass_a,mass_b)  
      #plot_df(data[0,0,:,:-1],outputs[0,0,:,:-1],epoch)
      #plot_df(targets_unnorm[0,0,:,:-1],outputs_unnorm[0,0,:,:-1],epoch)

    l2_loss = criterion(outputs,targets) 
    
    nbatch = len(data)     
    mass_loss = torch.sum(torch.abs(mass_a - mass_b)).float()/nbatch
    mom_loss = torch.sum(torch.abs(mom_a - mom_b)).float()/nbatch
    energy_loss = torch.sum(torch.abs(energy_a - energy_b)).float()/nbatch     
#     print(mass_loss.item(),mom_loss.item(),energy_loss.item())

    if epoch < 1000:
      loss = l2_loss
    else:
      loss = l2_loss*5*loss_weights[0] \
        + mass_loss*loss_weights[1] \
        + mom_loss*loss_weights[2] \
        + energy_loss*loss_weights[3]   
    
#     loss = l2_loss*loss_weights[0] \
#          + mass_loss*loss_weights[1] \
#          + mom_loss*loss_weights[2] \
#          + energy_loss*loss_weights[3]   
  
    loss.backward()
    optimizer.step()

    running_loss += loss.item()
    running_l2_loss += l2_loss.item()
    running_den += mass_loss.item()*loss_weights[1]
    running_mom += mom_loss.item()*loss_weights[2]
    running_en += energy_loss.item()*loss_weights[3]

    if epoch % output_rate == output_rate-1:
      print('   [%d] loss: %.8f' % (epoch + 1, running_loss / output_rate)) 
      print('      L2 loss: %.8f' % (running_l2_loss / output_rate))
      print('      conservation loss: %.8f' % ((running_den+running_mom+running_en) / output_rate))
      
      #print('data',torch.mean(data).item())
      #print('targets',torch.mean(targets).item())
      #print('outputs',torch.mean(outputs).item())
      
      train_loss_vector.append(running_loss / output_rate)  
      l2_loss_vector.append(running_l2_loss / output_rate)  
      cons_loss_vector[0].append(running_den / output_rate)  
      cons_loss_vector[1].append(running_mom / output_rate)  
      cons_loss_vector[2].append(running_en / output_rate)  
      
      running_loss = 0.0
      running_l2_loss = 0.0
      running_den = 0.0
      running_mom = 0.0
      running_en = 0.0
      
#     if epoch > 1000:  
#       if running_den > 100:
#         plot_df(targets_unnorm[0,0,:,:-1],outputs_unnorm[0,0,:,:-1],epoch)
#       elif running_den < 30:
#         plot_df(targets_unnorm[0,0,:,:-1],outputs_unnorm[0,0,:,:-1],epoch)
      
    if epoch % 100 == 99:  
      plt.plot(l2_loss_vector)
      plt.plot(cons_loss_vector[0])
      plt.plot(cons_loss_vector[1])
      plt.plot(cons_loss_vector[2])

      plt.legend(['l2 loss','den','mom','energy'])
      #plt.yscale('log')
      plt.show()
      
      plot_df(targets_unnorm[0,0,:,:-1],outputs_unnorm[0,0,:,:-1],epoch)
    #scheduler.step(loss)
  
  plt.contourf(outputs[0,0].cpu().detach().numpy())
  plt.title('output df')
  plt.colorbar()
  plt.show()
#   data_unnorm = data*std_f + mean_f
#   targets_unnorm = targets*std_df[:,0] + mean_df[:,0]
#   outputs_unnorm = outputs*std_df[:,0] + mean_df[:,0]
#   plot_df(targets_unnorm[0,0,:,:-1],outputs_unnorm[0,0,:,:-1],epoch)

  return train_loss_vector,l2_loss_vector,cons_loss_vector

"""# validate"""

def validate(valloader,props,cons,std_f,std_df,mean_f,mean_df):
  
  print('      Running validation set')
  
  running_loss = 0.0
  
  # don't compute gradients - loss not backpropped
  with torch.no_grad():
    for i, (data, targets, props) in enumerate(valloader):
      
      # same general procedure as during training
      
      data, targets, props = data.to(device), targets.to(device), props.to(device)
      
      outputs = net(data)
      outputs = outputs.to(device)
      
      nbatch = len(data)
      if nbatch != batch_size:
        data_unnorm = data*std_f[:nbatch] + mean_f[:nbatch]
        targets_unnorm = targets*std_df[:nbatch,0] + mean_df[:nbatch,0]
        outputs_unnorm = outputs*std_df[:nbatch,0] + mean_df[:nbatch,0]

      else:
        data_unnorm = data*std_f + mean_f
        targets_unnorm = targets*std_df[:,0] + mean_df[:,0]
        outputs_unnorm = outputs*std_df[:,0] + mean_df[:,0]
        
      targets_nof = targets_unnorm - data_unnorm[:,0,:,:]
      outputs_nof = outputs_unnorm - data_unnorm[:,0,:,:]
        
      mass_b,mom_b,energy_b = check_properties3(data_unnorm[:,0,:,:-1],\
                                                targets_nof[:,0,:,:-1],props,cons)
      mass_a,mom_a,energy_a = check_properties3(data_unnorm[:,0,:,:-1],\
                                                outputs_nof[:,0,:,:-1],props,cons)
      
      mass_loss = torch.sum(torch.abs(mass_a - mass_b)).float()/nbatch
      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, targets)
      
      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()
  
  # get average loss over all iterations in validation set
  avg_loss = running_loss/(i+1)
  
  print('         Validation loss: %.3f' % (avg_loss))

  return avg_loss

"""# test"""

def test(f_test,df_test,props_test):
    
    # get raw test data across all phi and create test loader - probably not necessary at this point
    testset = DistFuncDataset(f_test, df_test, props_test)
    
    testloader = DataLoader(testset, batch_size=batch_size, 
                            shuffle=True, num_workers=4)
      
    props_before = []
    props_after = []
    
    l2_error=[]
    
    # these were used to count how many of the testing points had L2 error < 1
    lt1=0
    gt1=0
    with torch.no_grad():
        for (data, targets, props) in testloader:

            # same procedure as training/validation
            data, targets, props = data.to(device), targets.to(device), props.to(device)
          
            outputs = net(data)
            outputs = outputs.to(device)   
            
            if len(data) != batch_size:
              limit = len(data)
              data_unnorm = data*std_f[:limit] + mean_f[:limit]
              targets_unnorm = targets*std_df[:limit,0] + mean_df[:limit,0]
              outputs_unnorm = outputs*std_df[:limit,0] + mean_df[:limit,0]

            else:
              data_unnorm = data*std_f + mean_f
              targets_unnorm = targets*std_df[:,0] + mean_df[:,0]
              outputs_unnorm = outputs*std_df[:,0] + mean_df[:,0]
              
            targets_nof = targets_unnorm - data_unnorm[:,0,:,:]
            outputs_nof = outputs_unnorm - data_unnorm[:,0,:,:]
                       
            props_before.append([torch.sum(each_prop).item()\
                                 for each_prop in check_properties3(data_unnorm[:,0,:,:-1],\
                                                                    targets_nof[:,0,:,:-1],props,cons)])         
            props_after.append([torch.sum(each_prop).item()\
                                for each_prop in check_properties3(data_unnorm[:,0,:,:-1],\
                                                                   outputs_nof[:,0,:,:-1],props,cons)])
                                
            loss = criterion(outputs, targets)
            l2_error.append(loss.item()*100)
            # count how many have less than 1
            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])
    
    # percent error of conservation error
    cons_error[0] = np.abs((cons_array[:,3]-cons_array[:,0])/cons_array[:,0])
    cons_error[1] = np.abs((cons_array[:,4]-cons_array[:,1])/cons_array[:,1])
    cons_error[2] = np.abs((cons_array[:,5]-cons_array[:,2])/cons_array[:,2])
    
    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

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')

"""# plot"""

def plot_df(df_xgc,df_ml,epoch):
  
  # takes in torch tensors with 2D velocity grid of 2 dfs
  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('df{}'.format(epoch+1))
  
  plt.show()

"""# main one batch"""

## main script to run one batch 

# 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)
# #scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,'min')

# print('Loading data')
# f,df,props,num_nodes,std_f,std_df,mean_f,mean_df = load_data_hdf(0)
# cons = load_cons_vars()
# print('Splitting data')
# trainloader,valloader,f_test,df_test,props_test = split_data(f,df[:,0,:,:],props[:,1],num_nodes)
# trainiter = iter(trainloader)
# data, targets, props = trainiter.next()
# print('Starting training')
# a,b,c = train_one(data,targets,props,std_f,std_df,mean_f,mean_df)

"""# main"""

# main block that performs all function calls

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

# choose optimizer
#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)

# loop through epochs done externally of train function
for epoch in range(num_epochs):

  lr_epoch = [group['lr'] for group in optimizer.param_groups][0]

  print('Epoch: {} (lr = {})'.format(epoch+1,lr_epoch)) 
  
  epoch1 = timeit.default_timer() 
  # end variable used to keep track of number of iterations done if nphi > 1 within one batch
  end = 0
  
  # loop over iphi within one epoch
  for iphi in range(nphi):

    print('Beginning training iphi = {}'.format(iphi))
    print('   Loading data')
    load1 = timeit.default_timer()
    # load one iphi at a time - needed because RAM limit of Colab is 25GB
    f,df,props,num_nodes,std_f,std_df,mean_f,mean_df = load_data_hdf(iphi)
    cons = load_cons_vars()
    load2 = timeit.default_timer()
    print('      Loading time: %.3fs' % (load2-load1))

    print('   Creating training set')
    # currently only for ions - df has ions = 0, electrons = 1; props has electrons = 0, ions = 1
    # only get dataloader object for training and validation set - get raw test data for now
    trainloader,valloader,f_test,df_test,props_test = split_data(f,df[:,0,:,:],props[:,1],num_nodes)
    del f,df,props
    
    train1 = timeit.default_timer()
    # need conditional statements here to organize loss vectors and gather testing data
    # if first epoch take care of gathering testing data 
    if epoch == 0:
      # concatenate test data of iphi = 0 with test data from other phis
      # 
      if iphi == 0:
        f_all_test,df_all_test,props_all_test = f_test,df_test,props_test
        del f_test,df_test,props_test

        # get first loss vectors to in epoch = 0 and iphi = 0
        print('   Starting training')
        train_loss, l2_loss, cons_loss, val_loss, cons_array, end = \
                                   train(trainloader,valloader,0,epoch,end,std_f,std_df,mean_f,mean_df,cons)

      else:
        # concatenate raw test data
        f_all_test = np.vstack((f_all_test,f_test))
        df_all_test = np.vstack((df_all_test,df_test))
        props_all_test = np.vstack((props_all_test,props_test))
        del f_test,df_test,props_test

        # get loss vectors to append to existing ones
        print('   Starting training')
        train_loss_to_app, l2_loss_to_app, cons_loss_to_app, val_loss_to_app, cons_to_cat, end = \
                                   train(trainloader,valloader,0,epoch,end,std_f,std_df,mean_f,mean_df,cons)

        # added appending of cons/l2 loss since last version
        for loss1 in train_loss_to_app:
          train_loss.append(loss1)
        for loss2 in l2_loss_to_app:
          l2_loss.append(loss2)
        for loss3 in cons_loss_to_app:
          cons_loss.append(loss3)          
        for loss4 in val_loss_to_app:
          val_loss.append(loss4)    
        cons_array = np.concatenate((cons_array, cons_to_cat), axis=1)
    
    # if epoch != 0 no need to do anything with testing data
    else:
      del f_test,df_test,props_test
      print('   Starting training')
      # still need to append loss vectors to main ones
      train_loss_to_app, l2_loss_to_app, cons_loss_to_app, val_loss_to_app, cons_to_cat, end = \
                                 train(trainloader,valloader,0,epoch,end,std_f,std_df,mean_f,mean_df,cons)

      for loss1 in train_loss_to_app:
          train_loss.append(loss1)
      for loss2 in l2_loss_to_app:
          l2_loss.append(loss2)         
      for loss3 in cons_loss_to_app:
          cons_loss.append(loss3)          
      for loss4 in val_loss_to_app:
          val_loss.append(loss4)      
      cons_array = np.concatenate((cons_array, cons_to_cat), axis=0)
         
    train2 = timeit.default_timer()
    print('Finished tranining iphi = {}'.format(iphi))
    print('   Training time for iphi = %d: %.3fs' % (iphi,train2-train1))
  
  # plot at the end of each epoch
  train_iterations = np.linspace(1,len(train_loss),len(train_loss))
  val_iterations = np.linspace(2,len(train_loss),len(val_loss))
  
  plt.plot(train_iterations,train_loss,color='blue')
  plt.plot(val_iterations,val_loss,color='orange')
  plt.plot(train_iterations,l2_loss,color='red')
  plt.plot(train_iterations,cons_loss,color='green')
  plt.legend(['total','validation','l2','cons'])
  plt.yscale('log')
  plt.show()
  
  epoch2 = timeit.default_timer()
  #scheduler.step()
  print('Epoch time: {}s\n'.format(epoch2-epoch1))

# test and return arrays with information about l2/cons error 
print('Starting testing')
l2_error_i, cons_error_i = test(f_all_test,df_all_test,props_all_test)
print('Finished testing')

stop = timeit.default_timer()
print('Runtime: %.3fmins' % ((stop-start)/60))

## used for hyperparameter tuning - get train/validation loss for different learning rates
## I've been using a script I have saved locally to plot this data

# fid = open('/content/lr_new_tune.txt','a')
# fid.write('lr: '+str(lr)+'\n')
# #fid.write('runtime: '+str(stop-start)+'\n')
# fid.write('training\n')
# for tloss in train_loss:
#     fid.write(str(tloss)+'\n')   
# fid.write('validation\n')
# for vloss in val_loss:
#     fid.write(str(vloss)+'\n')
# fid.close()

## used to produce plots of specific nodes

# f,df,props,num_nodes,std_f,std_df,mean_f,mean_df = load_data_hdf(0)

# f_low = torch.from_numpy(np.expand_dims(f[51789],axis=0)).to(device).float()
# f_high = torch.from_numpy(np.expand_dims(f[141153],axis=0)).to(device).float()
                         
# df_low = torch.from_numpy(np.expand_dims(df[51789],axis=0)).to(device).float()
# df_high = torch.from_numpy(np.expand_dims(df[141153],axis=0)).to(device).float()

# out_low = net(f_low)
# out_high = net(f_high)

# out_low_unnorm = out_low*std_df[0] + mean_df[0] 
# df_low_unnorm = df_low*std_df[0] + mean_df[0]
# f_low_unnorm = f_low*std_f[0] + mean_f[0]

# out_high_unnorm = out_high*std_df[0,0] + mean_df[0,0] 
# df_high_unnorm = df_high*std_df[0,0] + mean_df[0,0]
# f_high_unnorm = f_high*std_f[0] + mean_f[0]

# out_low_unnorm = out_low_unnorm - f_low_unnorm[0,0,:,:]
# df_low_unnorm = df_low_unnorm - f_low_unnorm[0,0,:,:]

# out_high_unnorm = out_high_unnorm - f_high_unnorm[0,0,:,:]
# df_high_unnorm = df_high_unnorm - f_high_unnorm[0,0,:,:]

# plot_df(df_low_unnorm[0,0,:,:-1],out_low_unnorm[0,0,:,:-1],1)
# plot_df(df_high_unnorm[0,0,:,:-1],out_high_unnorm[0,0,:,:-1],2)

## prints conservation properties of specific nodes chosen above
 
# mass_b,mom_b,energy_b = check_properties3(f_high_unnorm[:,0,:,:-1],\
#                                           df_high_unnorm[:,0,:,:-1],props,cons) 
# mass_a,mom_a,energy_a = check_properties3(f_high_unnorm[:,0,:,:-1],\
#                                           out_high_unnorm[:,0,:,:-1],props,cons)

# print(mass_b.item(),mom_b.item(),energy_b.item())
# print(mass_a.item(),mom_a.item(),energy_a.item())

# generates training plots

vals = [5*i for i in range(6)]
pos = [80*i for i in range(6)]

plt.plot(train_iterations,train_loss,color='blue')
plt.plot(val_iterations,val_loss,color='orange')
plt.plot(train_iterations,l2_loss,color='red')
plt.plot(train_iterations,cons_loss,color='green')
plt.ylabel('loss')
plt.xlabel('epochs')
plt.xticks(pos,vals)
plt.title('ReSeg training, {} epochs, iphi=0'.format(num_epochs))
plt.legend(['Training','Validation','L2','Conservation'])
plt.yscale('log')
plt.savefig('few_epochs.png')
files.download('few_epochs.png')

# for epoch in range(num_epochs):
#   files.download('figs/dfs_{}.png'.format(epoch+1))

## few lines that change the shape of the cons_array returned during training - used for the y=x plots

#for i in range(10):
#  cons_array_np[4167*i:4167*(i+1)] = cons_array[:,i*6:(i+1)*6]
  
cons_array_np=cons_array.transpose()

# next few blocks produce the y=x plots for the conservation arrays - there may have to be some things
# done to the shape of the arrays to get the code to work

num_each_epoch = int(len(cons_array_np[0])/num_epochs)
for lim in range(num_epochs):
  plt.plot(cons_array_np[0,lim*num_each_epoch:(lim+1)*num_each_epoch],\
           cons_array_np[3,lim*num_each_epoch:(lim+1)*num_each_epoch],'o')

# min_dens = cons_array_np[0].min()
# max_dens = cons_array_np[0].max()

min_dens=-1.5
max_dens=1

plt.plot(np.linspace(min_dens,max_dens,100),np.linspace(min_dens,max_dens,100))
plt.ylabel('density change targets')
plt.xlabel('density change outputs')
plt.title('Density conservation, {} epochs, iphi=0'.format(num_epochs))
plt.legend(['epoch{}'.format(epoch+1) for epoch in range(num_epochs)])
plt.xlim([min_dens,max_dens])
plt.ylim([min_dens,max_dens])
plt.savefig('dens.png')
files.download('dens.png')

for lim in range(num_epochs):
  plt.plot(cons_array_np[1,lim*num_each_epoch:(lim+1)*num_each_epoch],\
           cons_array_np[4,lim*num_each_epoch:(lim+1)*num_each_epoch],'o')
  
# min_mom = cons_array_np[1].min()
# max_mom = cons_array_np[1].max()

min_mom = -0.3e-21
max_mom = 0.3e-21
  
plt.plot(np.linspace(min_mom,max_mom,100),np.linspace(min_mom,max_mom,100))
plt.ylabel('momentum change targets')
plt.xlabel('momentum change outputs')
plt.title('Momentum conservation, {} epochs, iphi=0'.format(num_epochs))
plt.legend(['epoch{}'.format(epoch+1) for epoch in range(num_epochs)])
plt.xlim([min_mom,max_mom])
plt.ylim([min_mom,max_mom])
plt.savefig('momentum.png')
files.download('momentum.png')

for lim in range(num_epochs):
  plt.plot(cons_array_np[2,lim*num_each_epoch:(lim+1)*num_each_epoch],\
           cons_array_np[5,lim*num_each_epoch:(lim+1)*num_each_epoch],'o')
  
# min_en = cons_array_np[2].min()
# max_en = cons_array_np[2].max()

min_en = -12.5e-16
max_en = 7.5e-16

plt.plot(np.linspace(min_en,max_en,100),np.linspace(min_en,max_en,100))
plt.ylabel('energy change targets')
plt.xlabel('energy change outputs')
plt.title('Energy conservation, {} epochs, iphi=0'.format(num_epochs))
plt.legend(['epoch{}'.format(epoch+1) for epoch in range(num_epochs)])
plt.xlim([min_en,max_en])
plt.ylim([min_en,max_en])
plt.savefig('energy.png')
files.download('energy.png')