<jupyter_start><jupyter_text>Fourier Feature Mapping on PytorchLightening<jupyter_code>%load_ext autoreload
%autoreload 2
import os

import numpy as np

import torch
import pytorch_lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from torch.utils.data import Dataset, DataLoader

import matplotlib.pyplot as plt
from tqdm.notebook import tqdm as tqdm

# verify cuda
torch.cuda.is_available()
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
print(device)
from turboflow.utils import phy_utils as phy
from turboflow.utils import dsp_utils as dsp<jupyter_output><empty_output><jupyter_text>Input data<jupyter_code>## SET UP DATA
path_to_data = '../data/2021-Turb2D_velocities.npy'
IMGs = np.load(path_to_data)
img_index = 42
X = IMGs[img_index,:,:,:2] / 255
U = IMGs[img_index,:,:,2]
V = IMGs[img_index,:,:,3]
t = np.arange(0,IMGs.shape[0],step=0.01)
Re = 3900

# normalize output
y = IMGs[img_index,:,:,2:4]
print('Y shape', y.shape)
print('Y min, max:', np.min(y), np.max(y))
y = y / np.max(np.abs(y))
print('after normalization, Y min, max:', np.min(y), np.max(y))

W = phy.compute_vorticity(U,V)
D = phy.compute_divergence(U,V)

plt.figure(figsize=(20,20))
plt.subplot(141)
plt.imshow(U)
plt.subplot(142)
plt.imshow(V)
plt.subplot(143)
plt.imshow(W)
plt.subplot(144)
plt.imshow(D)
plt.tight_layout()
plt.show()
Xtest = X
ytest = y

# downsampled resultion
scale = 1/4
X1ds = dsp.interpolate2D(X, X[:,:,0], scale=scale)[:,:,None]
X2ds = dsp.interpolate2D(X, X[:,:,1], scale=scale)[:,:,None]
y1ds = dsp.interpolate2D(X, y[:,:,0], scale=scale)[:,:,None]
y2ds = dsp.interpolate2D(X, y[:,:,1], scale=scale)[:,:,None]
Xtrain = np.concatenate([X1ds, X2ds], axis=-1)
ytrain = np.concatenate([y1ds, y2ds], axis=-1)

scale = 1/4
X1ds = dsp.interpolate2D(X, X[:,:,0], scale=scale)[:,:,None]
X2ds = dsp.interpolate2D(X, X[:,:,1], scale=scale)[:,:,None]
y1ds = dsp.interpolate2D(X, y[:,:,0], scale=scale)[:,:,None]
y2ds = dsp.interpolate2D(X, y[:,:,1], scale=scale)[:,:,None]

# fig, axs = plt.subplots(1,2)
# axs[0].imshow(ytest[:,:,0])
# axs[1].imshow(ytest[:,:,1])
# plt.show()

# fig, axs = plt.subplots(1,2)
# axs[0].imshow(ytrain[:,:,0])
# axs[1].imshow(ytrain[:,:,1])
# plt.show()

# upsampled resolution
yup = dsp.interpolate2D(Xtrain, ytrain[:,:,0], scale=4)[:,:,None]

K41 = np.linspace(0,256,1000)

# b1, b0, _r, _p_val, _stderr = stats.linregress(np.log10(mx), np.log10(my)) 

fig, axs = plt.subplots(1,1,figsize=(5,5))
axs.loglog(np.arange(len(phy.powerspec(ytrain[:,:,0]))), phy.powerspec(ytrain[:,:,0]), label='U train 64')
axs.loglog(np.arange(len(phy.powerspec(ytest[:,:,0]))), phy.powerspec(ytest[:,:,0]), label='U test 256')
axs.set_ylim(10**(-15), 10)
axs.set_xlim(1, 512)
axs.legend()
plt.tight_layout()
plt.show()
plt.figure(figsize=(10,5))
plt.subplot(141)
plt.title('[LR] X - res (%d,%d)' % Xtrain[:,:,0].shape)
plt.imshow(Xtrain[:,:,0])
plt.subplot(142)
plt.title('[LR] Y - res (%d,%d)' % Xtrain[:,:,1].shape)
plt.imshow(Xtrain[:,:,1])
plt.subplot(143)
plt.title('[LR] Ux  - res (%d,%d)' % ytrain[:,:,0].shape)
plt.imshow(ytrain[:,:,0])
plt.subplot(144)
plt.title('[LR] Uy - res (%d,%d)' % ytrain[:,:,1].shape)
plt.imshow(ytrain[:,:,1])
plt.tight_layout()
plt.savefig('../figures/turboflow_train.pdf', dpi=150)
plt.show()

plt.figure(figsize=(10,5))
plt.subplot(141)
plt.title('[HR] X - res (%d,%d)' % Xtest[:,:,0].shape)
plt.imshow(Xtest[:,:,0])
plt.subplot(142)
plt.title('[HR] Y - res (%d,%d)' % Xtest[:,:,1].shape)
plt.imshow(Xtest[:,:,1])
plt.subplot(143)
plt.title('[HR] Ux - res (%d,%d)' % ytest[:,:,0].shape)
plt.imshow(ytest[:,:,0])
plt.subplot(144)
plt.title('[HR] Uy - res (%d,%d)' % ytest[:,:,1].shape)
plt.imshow(ytest[:,:,1])
plt.tight_layout()
plt.savefig('../figures/turboflow_test.pdf', dpi=150)
plt.show()<jupyter_output><empty_output><jupyter_text>Dataset<jupyter_code>class MyDataset(torch.utils.data.Dataset):
  'Characterizes a dataset for PyTorch'
  def __init__(self, Xlow, ylow, Xmid, ymid):

        assert np.max(np.abs(y)) <= 1
        # assert np.max(np.abs(X)) <= 1

        'Initialization'
        try:
            assert Xlow.shape == ylow.shape
            assert Xmid.shape == ymid.shape
        except:
            print('Dimension error')
            print('Xlow', Xlow.shape)
            print('ylow', ylow.shape)
            print('Xmid', Xmid.shape)
            print('ymid', ymid.shape)

        # convert to torch and add an empty dimension for the batch
        self.inputs_low = torch.from_numpy(Xlow).float() #XY x 2 (x,y)
        self.target_low = torch.from_numpy(ylow).float()  #XxYx3 (RGB)
        
        self.inputs_mid = torch.from_numpy(Xmid).float() #XY x 2 (x,y)
        self.target_mid = torch.from_numpy(ymid).float()  #XxYx3 (RGB)
            
  def __len__(self):
        'Denotes the total number of samples'
        return 1
    
  def __getitem__(self, index):
        'Generates one sample of data'
        # Load data and get label
        Xlow = self.inputs_low
        ylow = self.target_low
        
        Xmid = self.inputs_mid
        ymid = self.target_mid

        return Xlow, ylow, Xmid, ymid


class DataModule(pl.LightningDataModule):
    def __init__(self, train_data, val_data, test_data, batch_size):
        
        super().__init__()
        self.batch_size = 1
        self.train_data = train_data
        self.test_data = test_data
        self.val_data = val_data
        
    def prepare_data(self):
        self.train_dataset = MyDataset(*self.train_data)
        self.val_dataset = MyDataset(*self.val_data)
        self.test_dataset = MyDataset(*self.test_data)
    
    def train_dataloader(self):
        return DataLoader(self.train_dataset, self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, self.batch_size)
low_scale = 1/4
mid_scale = 1/2
Xlow = np.concatenate([
            dsp.interpolate2D(X, X[:,:,0], scale=low_scale)[:,:,None],
            dsp.interpolate2D(X, X[:,:,1], scale=low_scale)[:,:,None]], 
       axis=-1).transpose(2,0,1)
Xmid = np.concatenate([
            dsp.interpolate2D(X, X[:,:,0], scale=mid_scale)[:,:,None],
            dsp.interpolate2D(X, X[:,:,1], scale=mid_scale)[:,:,None]],
        axis=-1).transpose(2,0,1)
ylow = np.concatenate([
            dsp.interpolate2D(X, y[:,:,0], scale=low_scale)[:,:,None],
            dsp.interpolate2D(X, y[:,:,1], scale=low_scale)[:,:,None]],
        axis=-1).transpose(2,0,1)
ymid = np.concatenate([
            dsp.interpolate2D(X, y[:,:,0], scale=mid_scale)[:,:,None],
            dsp.interpolate2D(X, y[:,:,1], scale=mid_scale)[:,:,None]],
        axis=-1).transpose(2,0,1)

dm = DataModule(
    train_data=[Xlow, ylow, Xmid, ymid],
    val_data=[Xlow, ylow, Xmid, ymid],
    test_data=[Xlow, ylow, Xmid, ymid],
    batch_size = 1,
)

## test
dm.prepare_data()
print('Train')
for batch in dm.train_dataloader():
    Xlow, ylow, Xmid, ymid = batch
    print(Xlow.shape, ylow.shape)
    print(Xmid.shape, ymid.shape)
    
print('Validation')
for batch in dm.val_dataloader():
    Xlow, ylow, Xmid, ymid = batch
    print(Xlow.shape, ylow.shape)
    print(Xmid.shape, ymid.shape)<jupyter_output>Train
torch.Size([1, 2, 64, 64]) torch.Size([1, 2, 64, 64])
torch.Size([1, 2, 128, 128]) torch.Size([1, 2, 128, 128])
Validation
torch.Size([1, 2, 64, 64]) torch.Size([1, 2, 64, 64])
torch.Size([1, 2, 128, 128]) torch.Size([1, 2, 128, 128])<jupyter_text>Model<jupyter_code>""" Full assembly of the parts to form the complete network """

import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)
    

class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)
    

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)
        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        print('bottleneck', x5.shape)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits
    
    
class fMLP(nn.Module):
    def __init__(self, layer_dimension, f_nfeatures, f_scale, device):
        super().__init__()
        
        layers = []
        num_layers = len(layer_dimension)
        
        blocks = []
        for l in range(num_layers-2):
            blocks.append(create_blockReLU(layer_dimension[l], layer_dimension[l+1]))
            
        blocks.append(nn.Linear(layer_dimension[-2], layer_dimension[-1]))
        blocks.append(nn.Tanh())
        
        self.rff = Fourier(f_nfeatures, f_scale)
        self.mlp = nn.Sequential(*blocks)
        
    
    def forward(self, x):
        x = self.rff(x)
        x = self.mlp(x)
        # in lightning, forward defines the prediction/inference actions
        return x
class UNetPl(pl.LightningModule):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.unet = UNet(in_channels, out_channels)
        self.fmlp = FMlp()
        
    def forward(self, x, q):
         = self.fmlp(x) 
        
    
    def training_step(self, batch, batch_idx):
        print('Traning step')
        # training_step defined the train loop.
        # It is independent of forward
        Xlow, ylow, Xmid, ymid = batch
        x = torch.cat([Xlow, ylow], axis=1) # concatenate over the channel dimension
        q = Xmid
        
        x_pred = self.forward(x, q)
        
        print(x_pred)
        1/0
    
        # reconstruction loss
        loss_rec = F.mse_loss(x_pred, x_true)

        # losses 
        loss = loss_rec
        
        # Logging to TensorBoard by default
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        print('Validation step')
        Xlow, ylow, Xmid, ymid = batch
        x = torch.cat([Xlow, ylow], axis=1) # concatenate over the channel dimension

        x_pred = self.forward(x)
        print(x_pred.shape)
        1/0
        loss = F.mse_loss(x_pred, x_true)
        self.log('valid_loss', loss, on_step=True)
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)
        return optimizer<jupyter_output><empty_output><jupyter_text>Check UNet implementation<jupyter_code># # creating a pseudo-type input tensor of size batch_size*in_channels*x_dim*y_dim
# batch_size = 1
# in_channels, x_dim = 2, 64
# out_channels = 2
# x = np.random.random((batch_size, in_channels, x_dim, x_dim))
# x = torch.from_numpy(x).float()
# y = np.random.random((batch_size, out_channels, x_dim, x_dim))
# y = torch.from_numpy(y).float()

# # running the UNet model for the given "depth" and desired number of "out_channels" on the input tensor
# model = UNet(n_channels=in_channels, n_classes=out_channels)
# pred = model(x)

# # checking the dimensions of the output and input
# ## output dimension is: batch_size*out_channels*x_dim*y_dim
# print("low resolution input image shape:", input.shape)
# print("high resolution output image shape:", model_output.shape)

# criterion = nn.MSELoss()
# loss = torch.sqrt(criterion(pred, y))
# loss.backward()<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>in_channels = 4
out_channels = 16

model = UNetPl(in_channels, out_channels)
early_stop_callback = EarlyStopping(monitor='valid_loss')

trainer = pl.Trainer(gpus=1, 
                     check_val_every_n_epoch=200, 
                     max_epochs=int(10e3), 
                     callbacks=[early_stop_callback])

trainer.fit(model, dm)
ngrid = 64
model.eval().to(device)
coords = np.linspace(0, 1, ngrid, endpoint=False)
coords = np.stack(np.meshgrid(coords, coords), -1) # X x Y x 2
tmp = coords[:,:,0].copy()
coords[:,:,0] = coords[:,:,1]
coords[:,:,1] = tmp
coords = coords.transpose(2,0,1)
pred = model(torch.from_numpy(coords).float().to(device)[None,...])
ypred = pred.cpu().detach().numpy().squeeze()

plt.imshow(ypred[0,:,:])
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15925413454569415, 'fraction_numerical': 0.032681030506687754, 'mean_word_length': 3.161776802851659, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 26, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '567954', 'license': 'Apache-2.0', 'n_tokens_mistral': 6112, 'n_tokens_neox': 5814, 'n_words': 1098, 'path': 'notebooks/[TurboNeuralFourier] MeshfreeFlow on PytorchLightening on Turbulent Data.ipynb', 'repo_name': 'Chutlhu/TurboSuperResultion', 'size': '15177'}