<jupyter_start><jupyter_text>load pre-trained<jupyter_code>device = torch.device('cuda:0')
model = torch.load('trained_model/dP_UFNO_4500_trained')
model.to(device)
# Prepare for plotting
dx = np.cumsum(3.5938*np.power(1.035012, range(200))) + 0.1
X, Y = np.meshgrid(dx, np.linspace(0,200,num = 96))

def pcolor(x):
    plt.jet()
    return plt.pcolor(X[:thickness, :], Y[:thickness, :], np.flipud(x), shading='auto')

def dnorm_dP(dP):
    dP = dP * 18.772821433027488
    dP = dP + 4.172939172019009 
    return dP

times = np.cumsum(np.power(1.421245, range(24)))
time_print = []
for t in range(times.shape[0]):
    if times[t]<365:
        title = str(int(times[t]))+' d'
        time_print = np.append(time_print,title)
    else:
        title = f'{round(int(times[t])/365, 1)} y' 
        time_print = np.append(time_print,title)

dnorm_inj = lambda a : (a * (3e6 - 3e5) + 3e5) / (1e6 /365*1000/1.862)
dnorm_temp = lambda a : a * (180 - 30) + 30
dnorm_P = lambda a : a * (300 - 100) + 100
dnorm_lam = lambda a : a * 0.4 + 0.3
dnorm_Swi = lambda a : a * 0.2 + 0.1
test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a[:1,...], 
                                                                         test_u[:1,...]), 
                                          batch_size=1, 
                                          shuffle=False)<jupyter_output><empty_output><jupyter_text>Visualize predictions<jupyter_code>x, y = next(iter(test_loader))
x, y = x.to(device), y.to(device)
pred = model(x).view(-1,96,200,24)

y_plot = y.cpu().detach().numpy()
x_plot = x.cpu().detach().numpy()
pred_plot = dnorm_dP(pred.cpu().detach().numpy())

# reservior have different thickness as marked in the permeability map
mask = x_plot[0,:,:,0,0] != 0 
thickness = sum(mask[:,0])

# extract input parameters
poro_map = x_plot[0,:,:,0,2][mask].reshape((thickness, -1))
kr_map = np.exp(x_plot[0,:,:,0,0][mask].reshape((thickness, -1))*15)
kz_map = np.exp(x_plot[0,:,:,0,1][mask].reshape((thickness, -1))*15)

inj_rate = dnorm_inj(x_plot[0,0,0,0,4])
temperature = dnorm_temp(x_plot[0,0,0,0,5])
pressure = dnorm_P(x_plot[0,0,0,0,6])
Swi = dnorm_Swi(x_plot[0,0,0,0,7])
lam = dnorm_lam(x_plot[0,0,0,0,8])

print(f'Params: injection rate: {inj_rate:.2f} MT/yr, temperature: {temperature:.1f} C, initial pressure: {pressure:.1f} bar, Swi: {Swi:.2f}, lan: {lam:.2f}')

t_lst = [14,20,23]
plt.figure(figsize=(15,6))    
for j, t in enumerate(t_lst):   
    plt.subplot(4,3,j+1)
    if j == 2:
        pcolor(poro_map)
        plt.title('$\phi$ (-)')
    elif j == 1:
        pcolor(kz_map)
        plt.title('$k_z$ (mD)')
    else:
        pcolor(kr_map)
        plt.title('$k_r$ (mD)')
    plt.colorbar(fraction=0.02)
    plt.xlim([0,3500])        

    plt.subplot(4,3,j+4)
    pcolor(y_plot[0,:,:,t][mask].reshape((thickness, -1)))
    plt.title('$dP$ (bar), '+f't={time_print[t]}')
    plt.colorbar(fraction=0.02)
    plt.xlim([0,3500])

    plt.subplot(4,3,j+7)
    pcolor(pred_plot[0,:,:,t][mask].reshape((thickness, -1)))
    plt.title('$\hat{dP}$ (bar), '+f't={time_print[t]}')
    plt.colorbar(fraction=0.02)
    plt.xlim([0,3500])

    plt.subplot(4,3,j+10)
    pcolor((pred_plot[0,:,:,t][mask].reshape((thickness, -1))-y_plot[0,:,:,t][mask].reshape((thickness, -1))))
    plt.colorbar(fraction=0.02)
    plt.title('|$dP-\hat{dP}$|, '+f't={time_print[t]}')
    plt.xlim([0,3500])
plt.tight_layout()
plt.show()<jupyter_output>Params: injection rate: 1.92 MT/yr, temperature: 121.2 C, initial pressure: 258.3 bar, Swi: 0.28, lan: 0.45
{'fraction_non_alphanumeric': 0.2060400790290714, 'fraction_numerical': 0.07451312447078746, 'mean_word_length': 3.4134495641344955, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '108142', 'license': 'MIT', 'n_tokens_mistral': 1679, 'n_tokens_neox': 1514, 'n_words': 233, 'path': 'eval_UFNO_pressure_buildup.ipynb', 'repo_name': 'gegewen/ufno', 'size': '3543'}