text_plain_output_1.png
from torch import nn, optim

from torch.nn.functional import softmax

from tqdm import tqdm

import numpy as np

import pandas as pd

import torch

nb_takes = 15
nb_reduced = int(300 / nb_takes)
nb_takes_phy = nb_takes
nb_reduced_phy = nb_reduced

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
a = 100
b = 100
dim_mesh = (a - 1) * (b - 1)
bias_bo = True
bias_bo_g = True
bias_bo_d = True
train_samples = 1200
end_samples = 1600
path = '../input/2stage-simon/f_set_nonl.npy'
ffine_all = 10000 * np.load(path)
ffine = ffine_all[:train_samples]
nb_samples = np.shape(ffine)[0]
nb_times = np.shape(ffine)[1]
fcoarse_reshape = np.reshape(ffine, (nb_samples, nb_times, a - 1, b - 1))
fcoarse_reshape_torch_before = torch.tensor(fcoarse_reshape, dtype=torch.float32, device=device)
ffine_test = ffine_all[train_samples:end_samples]
nb_samples_test = np.shape(ffine_test)[0]
nb_times_test = np.shape(ffine_test)[1]
fcoarse_reshape_test = np.reshape(ffine_test, (nb_samples_test, nb_times_test, a - 1, b - 1))
fcoarse_reshape_test_torch_before = torch.tensor(fcoarse_reshape_test, dtype=torch.float32, device=device)

path = '../input/cem100/cembasis_10.npy'
_R_ = np.load(path)
Rt = np.transpose(_R_)
RtFh = []
for jx in tqdm(range(nb_samples)):
    RtFh_t = []
    for ix in range(nb_times):
        RtFh_t.append(np.matmul(Rt, ffine[jx, ix, :]))
    RtFh.append(RtFh_t)
RtFh = np.array(RtFh)
RtFh0 = RtFh[:, :, 0::3]
Rtfh0_torch = torch.tensor(RtFh0, dtype=torch.float32, device=device)
nb_model_phy0 = Rtfh0_torch.size()[-1]
RtFh1 = RtFh[:, :, 1::3]
Rtfh1_torch = torch.tensor(RtFh1, dtype=torch.float32, device=device)
nb_model_phy1 = Rtfh1_torch.size()[-1]
RtFh2 = RtFh[:, :, 2::3]
Rtfh2_torch = torch.tensor(RtFh2, dtype=torch.float32, device=device)
nb_model_phy2 = Rtfh2_torch.size()[-1]
Rtfh_torch = torch.tensor(RtFh, dtype=torch.float32, device=device)
nb_model_phy = Rtfh_torch.size()[-1]
RtFh_test = []
for jx in tqdm(range(nb_samples_test)):
    RtFh_t = []
    for ix in range(nb_times_test):
        RtFh_t.append(np.matmul(Rt, ffine_test[jx, ix, :]))
    RtFh_test.append(RtFh_t)
RtFh_test = np.array(RtFh_test)
RtFh_test0 = RtFh_test[:, :, 0::3]
Rtfh0_torch_test = torch.tensor(RtFh_test0, dtype=torch.float32, device=device)
nb_model_phy0_test = Rtfh0_torch_test.size()[-1]
RtFh_test1 = RtFh_test[:, :, 1::3]
Rtfh1_torch_test = torch.tensor(RtFh_test1, dtype=torch.float32, device=device)
nb_model_phy1_test = Rtfh1_torch_test.size()[-1]
invRtR = LA.inv(np.matmul(Rt, _R_))
path = '../input/2stage-simon/sol_set_nonl.txt'
u_f_all = 10000 * pd.read_csv(path, sep=' ', header=None).values
u_f = u_f_all[:train_samples]
u_H = []
for ix in range(nb_samples):
    u_H.append(np.matmul(invRtR, np.matmul(Rt, u_f[ix])))
u_H = np.array(u_H)
nb_rdm = np.shape(u_H)[1]
u_f_test = u_f_all[train_samples:end_samples]
u_H_test = []
for ix in range(nb_samples_test):
    u_H_test.append(np.matmul(invRtR, np.matmul(Rt, u_f_test[ix])))
u_H_test = np.array(u_H_test)
nb_rdm_test = np.shape(u_H_test)[1]

class OneHeadAttention(nn.Module):
    def __init__(self, d_model, d_reduced, bo):
        super().__init__()
        self.v_linear = nn.Linear(d_model, d_reduced, bias=bias_bo).cuda()
        self.q_linear = nn.Linear(d_model, d_reduced, bias=bias_bo).cuda()
        self.k_linear = nn.Linear(d_model, d_reduced, bias=bias_bo).cuda()
    def forward(self, vv):
        v = self.v_linear(vv)
        k = self.k_linear(vv)
        q = self.q_linear(vv)
        qkt = torch.matmul(q, torch.transpose(k, 1, 2))
        sm_qkt = softmax(qkt, dim=-1)
        out = torch.matmul(sm_qkt, v)
        return out
model = OneHeadAttention(nb_model_phy, nb_reduced_phy, True)
def init_weights(m):
    if type(m) == nn.Linear:
        m.bias.data.uniform_(-1 / 100000, 1 / 100000)
model.apply(init_weights)
out = model(Rtfh_torch)
np.shape(out)

nb_heads = 6
class MultiHeads(nn.Module):
    def __init__(self, d_model, d_reduced, d_head, bo):
        super().__init__()
        self.head1 = OneHeadAttention(d_model, d_reduced, bo)
        self.head2 = OneHeadAttention(d_model, d_reduced, bo)
        self.head3 = OneHeadAttention(d_model, d_reduced, bo)
        self.head4 = OneHeadAttention(d_model, d_reduced, bo)
        self.head5 = OneHeadAttention(d_model, d_reduced, bo)
        self.head6 = OneHeadAttention(d_model, d_reduced, bo)
        self.linear = nn.Linear(d_reduced * d_head, d_reduced, bias=bias_bo).cuda()
    def forward(self, v):
        out1 = self.head1(v)
        out2 = self.head2(v)
        out3 = self.head3(v)
        out4 = self.head4(v)
        out5 = self.head5(v)
        out6 = self.head6(v)
        concat_out = torch.cat((out1, out2, out3, out4, out5, out6), dim=-1)
        out = self.linear(concat_out)
        return out
def init_weights(m):
    if type(m) == nn.Linear:
        m.bias.data.uniform_(-1 / 100000, 1 / 100000)
model = MultiHeads(nb_model_phy, nb_reduced_phy, nb_heads, True)
model.apply(init_weights)
out = model(Rtfh_torch)
print(out.size())