darcy_flow_main.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from data_load_darcy import *
import random
import matplotlib.pyplot as plt
from darcy_flow_uno2d import UNO_9, UNO_11
import operator
from functools import reduce
from functools import partial
from train_darcy import train_model
from timeit import default_timer
from utilities3 import *
from Adam import Adam
from torchsummary import summary
import gc
import math

plt.rcParams["figure.figsize"] = [6, 30]
plt.rcParams["image.interpolation"] = "nearest"
torch.manual_seed(10001)
random.seed(10001)

import sys
import logging


# hyper parameters
# weight decay = 0.001
# Learning rate = 0.001
# scheduler_gamma = 0.5

# Loading and Processing the data. file1 and file2 are generated data of Darcy Flow. Data can be generated by provided data generator.

train_a_1, train_u_1, test_a_1, test_u_1 = load_data_darcy(
    2, 800, 200, "Path to data file1"
)
train_a_2, train_u_2, test_a_2, test_u_2 = load_data_darcy(
    2, 800, 200, "Path to data file2"
)

sub = 2  # subsampling rate
S = 211  # Grid size/ resolution after sub-sampling
# single input and output
T_in = 1
T_f = 1

# number of train, test and validation samples
ntrain = 1500
nval = 250
ntest = 250
batch_size = 16
width = 32  #
inwidth = 3
epochs = 700

a = torch.cat([train_a_1, train_a_2, test_a_1, test_a_2], dim=0)
u = torch.cat([train_u_1, train_u_2, test_u_1, test_u_2], dim=0)

indexs = [i for i in range(a.shape[0])]
random.shuffle(indexs)

train_a, val_a, test_a = (
    a[indexs[:ntrain]],
    a[indexs[ntrain : ntrain + nval]],
    a[indexs[ntrain + nval :]],
)
train_u, val_u, test_u = (
    u[indexs[:ntrain]],
    u[indexs[ntrain : ntrain + nval]],
    u[indexs[ntrain + nval :]],
)
print(train_a.shape, val_a.shape, test_a.shape)
train_a = train_a.reshape(ntrain, S, S, T_in)
val_a = val_a.reshape(nval, S, S, T_in)
test_a = test_a.reshape(ntest, S, S, T_in)

gc.collect()


train_loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(train_a, train_u),
    batch_size=batch_size,
    shuffle=True,
)
test_loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(test_a, test_u), batch_size=batch_size, shuffle=True
)
val_loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(val_a, val_u), batch_size=batch_size, shuffle=True
)

model = UNO_9(inwidth, width, pad=12).cuda()
# model = UNO_11(inwidth,width,pad = 8).cuda()
summary(model, (S, S, 1))
gc.collect()

train_model(
    model,
    train_loader,
    val_loader,
    test_loader,
    ntrain,
    nval,
    ntest,
    S,
    "Darcy-D9-211.pt",
    T_f=T_f,
    batch_size=batch_size,
    epochs=epochs,
    learning_rate=0.001,
    scheduler_step=100,
    scheduler_gamma=0.5,
    weight_decay=0.001,
)